[RISCV] Implement computeKnownBitsForTargetNode for SHL_ADD (#159105)

Author: pfusik

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

@@ -21568,6 +21568,16 @@ void RISCVTargetLowering::computeKnownBitsForTargetNode(const SDValue Op,
     Known = Known.sext(BitWidth);
     break;
   }
+  case RISCVISD::SHL_ADD: {
+    KnownBits Known2;
+    Known = DAG.computeKnownBits(Op.getOperand(0), DemandedElts, Depth + 1);
+    unsigned ShAmt = Op.getConstantOperandVal(1);
+    Known <<= ShAmt;
+    Known.Zero.setLowBits(ShAmt); // the <<= operator left these bits unknown
+    Known2 = DAG.computeKnownBits(Op.getOperand(2), DemandedElts, Depth + 1);
+    Known = KnownBits::add(Known, Known2);
+    break;
+  }
   case RISCVISD::CTZW: {
     KnownBits Known2 = DAG.computeKnownBits(Op.getOperand(0), Depth + 1);
     unsigned PossibleTZ = Known2.trunc(32).countMaxTrailingZeros();

llvm/test/CodeGen/RISCV/rv32zba.ll

@@ -1265,3 +1265,38 @@ define i64 @select9i64(i1 zeroext %x) {
   %select = select i1 %x, i64 9, i64 0
   ret i64 %select
 }
+
+define ptr @shl_add_knownbits(ptr %p, i32 %i) {
+; RV32I-LABEL: shl_add_knownbits:
+; RV32I:       # %bb.0:
+; RV32I-NEXT:    slli a1, a1, 18
+; RV32I-NEXT:    srli a1, a1, 18
+; RV32I-NEXT:    slli a2, a1, 1
+; RV32I-NEXT:    slli a1, a1, 3
+; RV32I-NEXT:    sub a1, a1, a2
+; RV32I-NEXT:    srli a1, a1, 3
+; RV32I-NEXT:    add a0, a0, a1
+; RV32I-NEXT:    ret
+;
+; RV32ZBA-LABEL: shl_add_knownbits:
+; RV32ZBA:       # %bb.0:
+; RV32ZBA-NEXT:    slli a1, a1, 18
+; RV32ZBA-NEXT:    srli a1, a1, 18
+; RV32ZBA-NEXT:    sh1add a1, a1, a1
+; RV32ZBA-NEXT:    srli a1, a1, 2
+; RV32ZBA-NEXT:    add a0, a0, a1
+; RV32ZBA-NEXT:    ret
+;
+; RV32XANDESPERF-LABEL: shl_add_knownbits:
+; RV32XANDESPERF:       # %bb.0:
+; RV32XANDESPERF-NEXT:    nds.bfoz a1, a1, 13, 0
+; RV32XANDESPERF-NEXT:    nds.lea.h a1, a1, a1
+; RV32XANDESPERF-NEXT:    srli a1, a1, 2
+; RV32XANDESPERF-NEXT:    add a0, a0, a1
+; RV32XANDESPERF-NEXT:    ret
+  %and = and i32 %i, 16383
+  %mul = mul i32 %and, 6
+  %shr = lshr i32 %mul, 3
+  %r = getelementptr i8, ptr %p, i32 %shr
+  ret ptr %r
+}

llvm/test/CodeGen/RISCV/rv64zba.ll

@@ -4636,3 +4636,38 @@ define i32 @select9(i1 zeroext %x) {
   %select = select i1 %x, i32 9, i32 0
   ret i32 %select
 }
+
+define ptr @shl_add_knownbits(ptr %p, i64 %i) {
+; RV64I-LABEL: shl_add_knownbits:
+; RV64I:       # %bb.0:
+; RV64I-NEXT:    slli a1, a1, 50
+; RV64I-NEXT:    srli a1, a1, 50
+; RV64I-NEXT:    slli a2, a1, 1
+; RV64I-NEXT:    slli a1, a1, 3
+; RV64I-NEXT:    sub a1, a1, a2
+; RV64I-NEXT:    srli a1, a1, 3
+; RV64I-NEXT:    add a0, a0, a1
+; RV64I-NEXT:    ret
+;
+; RV64ZBA-LABEL: shl_add_knownbits:
+; RV64ZBA:       # %bb.0:
+; RV64ZBA-NEXT:    slli a1, a1, 50
+; RV64ZBA-NEXT:    srli a1, a1, 50
+; RV64ZBA-NEXT:    sh1add a1, a1, a1
+; RV64ZBA-NEXT:    srli a1, a1, 2
+; RV64ZBA-NEXT:    add a0, a0, a1
+; RV64ZBA-NEXT:    ret
+;
+; RV64XANDESPERF-LABEL: shl_add_knownbits:
+; RV64XANDESPERF:       # %bb.0:
+; RV64XANDESPERF-NEXT:    nds.bfoz a1, a1, 13, 0
+; RV64XANDESPERF-NEXT:    nds.lea.h a1, a1, a1
+; RV64XANDESPERF-NEXT:    srli a1, a1, 2
+; RV64XANDESPERF-NEXT:    add a0, a0, a1
+; RV64XANDESPERF-NEXT:    ret
+  %and = and i64 %i, 16383
+  %mul = mul i64 %and, 6
+  %shr = lshr i64 %mul, 3
+  %r = getelementptr i8, ptr %p, i64 %shr
+  ret ptr %r
+}