[SDAG] Add SimplifyDemandedBits support for ISD::SPLAT_VECTOR_PARTS

[luke]

Summary

Similar to https://reviews.llvm.org/D158839, this allows some shift and rotate
operations on RV32 to better select an immediate or scalar operand, due to the
upper bits of the splat being marked as undef.

[craig.topper]

llvm/test/CodeGen/RISCV/rvv/vnsrl-sdnode.ll

644 ↗

(On Diff #553495)

This looks like what we really have is a missing combine on trunc of splat_vector.

644 ↗

(On Diff #553495)

Not objecting to this patch just that it might not show the proper motivation.

[luke]

llvm/test/CodeGen/RISCV/rvv/vnsrl-sdnode.ll

644 ↗

(On Diff #553495)

Is the combine you're referring to something like:

      t4: i32,ch = CopyFromReg t0, Register:i32 %1
      t6: i32,ch = CopyFromReg t0, Register:i32 %2
    t8: i64 = build_pair t4, t6
  t11: nxv1i64 = splat_vector t8
t13: nxv1i32 = truncate t11

to

      t4: i32,ch = CopyFromReg t0, Register:i32 %1
t13: nxv1i32 = splat_vector t4

And not on splat_vector_parts after legalisation?

[craig.topper]

llvm/test/CodeGen/RISCV/rvv/vnsrl-sdnode.ll

644 ↗

(On Diff #553495)

Do we already have a combine for truncate of build_pair?

From a quick glance, I see that we do convert (truncate (build_vector X, Y, Z)) to (build_vector (trunc X), (trunc Y), (trunc Z)). here.

// Attempt to pre-truncate BUILD_VECTOR sources.                               
if (N0.getOpcode() == ISD::BUILD_VECTOR && !LegalOperations &&                 
    TLI.isTruncateFree(SrcVT.getScalarType(), VT.getScalarType()) &&           
    // Avoid creating illegal types if running after type legalizer.           
    (!LegalTypes || TLI.isTypeLegal(VT.getScalarType()))) {                    
  SDLoc DL(N);                                                                 
  EVT SVT = VT.getScalarType();                                                
  SmallVector<SDValue, 8> TruncOps;                                            
  for (const SDValue &Op : N0->op_values()) {                                  
    SDValue TruncOp = DAG.getNode(ISD::TRUNCATE, DL, SVT, Op);                 
    TruncOps.push_back(TruncOp);                                               
  }                                                                            
  return DAG.getBuildVector(VT, DL, TruncOps);                                 
}

Maybe we should do the same for splat_vector?

[reames]

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

1168

I think Scl is short for Scalar? If so, just write it. Scl isn't a common abbreviation.

llvm/test/CodeGen/RISCV/rvv/vnsrl-sdnode.ll

644 ↗

(On Diff #553495)

I agree with Craig here. We definitely should have a trunc(splat) to splat(trunc) transform. If we don't already, we should fix that.

To be clear, this patch can be landed. You need to add a bit more test coverage for cases which *aren't* truncates so that once the other patch lands we still have test coverage for this code, but that's about the only thing missing.

Or said differently, Craig and I are suggesting *additional* work, not *alternative* work.

[luke]

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

1168

Yeah I had never seen this abbreviation before, but it's what's used in the rest of this function so I just chose it for consistency. Happy to expand it if preferred, I'm not strongly opinionated.

llvm/test/CodeGen/RISCV/rvv/vnsrl-sdnode.ll

644 ↗

(On Diff #553495)

I'm in agreement too here, happy to submit a patch for that.

[luke]

Rebase on top of new test cases

[luke]

llvm/test/CodeGen/RISCV/rvv/splat-vector-split-i64-vl-sdnode.ll

117

This a0 is actually undef, it's just selected by coincidence because there's no combine after RISCVDAGToDAGISel::PreprocessISelDAG to remove the store of undef. Might it be worthwhile short-circuiting a store of undef directly in SelectionDAG::getStore?

llvm/lib/CodeGen/SelectionDAG/TargetLowering.cpp

case ISD::SPLAT_VECTOR: {
     if (SimplifyDemandedBits(Scl, DemandedSclBits, KnownScl, TLO, Depth + 1))
       return true;
 
     // Implicitly truncate the bits to match the official semantics of
     // SPLAT_VECTOR.
     Known = KnownScl.trunc(BitWidth);
     break;
   }
+  case ISD::SPLAT_VECTOR_PARTS: {
+    unsigned NumSclBits = Op.getOperand(0).getScalarValueSizeInBits();
+    assert(NumSclBits * Op.getNumOperands() == BitWidth &&
+           "Expected SPLAT_VECTOR_PARTS scalars to cover element width");
+    for (auto [I, Scl] : enumerate(Op->ops())) {
+      APInt DemandedSclBits =
+          DemandedBits.extractBits(NumSclBits, NumSclBits * I);
+      KnownBits KnownScl;
+      if (SimplifyDemandedBits(Scl, DemandedSclBits, KnownScl, TLO, Depth + 1))
+        return true;
+      Known.insertBits(KnownScl, NumSclBits * I);
+    }
+    break;
+  }
   case ISD::LOAD: {
     auto *LD = cast<LoadSDNode>(Op);
     if (getTargetConstantFromLoad(LD)) {
       Known = TLO.DAG.computeKnownBits(Op, DemandedElts, Depth);
       return false; // Don't fall through, will infinitely loop.
     }
     if (ISD::isZEXTLoad(Op.getNode()) && Op.getResNo() == 0) {
       // If this is a ZEXTLoad and we are looking at the loaded value.

llvm/test/CodeGen/RISCV/rvv/splat-vector-split-i64-vl-sdnode.ll

; CHECK-NEXT:    ret
     ret i32 %4
 }
 
 define <vscale x 2 x i64> @demanded_bits_lo(i64 %x) {
 ; CHECK-LABEL: demanded_bits_lo:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    addi sp, sp, -16
 ; CHECK-NEXT:    .cfi_def_cfa_offset 16
-; CHECK-NEXT:    sw a1, 4(sp)
-; CHECK-NEXT:    sw a0, 0(sp)
-; CHECK-NEXT:    mv a0, sp
-; CHECK-NEXT:    vsetvli a1, zero, e64, m2, ta, ma
-; CHECK-NEXT:    vlse64.v v8, (a0), zero
 ; CHECK-NEXT:    sw zero, 12(sp)
-; CHECK-NEXT:    li a0, -1
-; CHECK-NEXT:    sw a0, 8(sp)
-; CHECK-NEXT:    addi a0, sp, 8
-; CHECK-NEXT:    vlse64.v v10, (a0), zero
-; CHECK-NEXT:    vand.vv v8, v8, v10
+; CHECK-NEXT:    li a1, -1
+; CHECK-NEXT:    sw a1, 8(sp)
+; CHECK-NEXT:    addi a1, sp, 8
+; CHECK-NEXT:    vsetvli a2, zero, e64, m2, ta, ma
+; CHECK-NEXT:    vlse64.v v8, (a1), zero
+; CHECK-NEXT:    vand.vx v8, v8, a0
 ; CHECK-NEXT:    addi sp, sp, 16
 ; CHECK-NEXT:    ret
   %1 = insertelement <vscale x 2 x i64> poison, i64 %x, i32 0
   %2 = shufflevector <vscale x 2 x i64> %1, <vscale x 2 x i64> poison, <vscale x 2 x i32> zeroinitializer
   %3 = and <vscale x 2 x i64> %2, shufflevector(<vscale x 2 x i64> insertelement(<vscale x 2 x i64> poison, i64 u0xffffffff, i32 0), <vscale x 2 x i64> poison, <vscale x 2 x i32> zeroinitializer)
   ret <vscale x 2 x i64> %3
 }
 
 define <vscale x 2 x i64> @demanded_bits_hi(i64 %x) {
 ; CHECK-LABEL: demanded_bits_hi:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    addi sp, sp, -16
 ; CHECK-NEXT:    .cfi_def_cfa_offset 16
 ; CHECK-NEXT:    sw a1, 4(sp)
 ; CHECK-NEXT:    sw a0, 0(sp)
 ; CHECK-NEXT:    mv a0, sp
 ; CHECK-NEXT:    vsetvli a1, zero, e64, m2, ta, ma
 ; CHECK-NEXT:    vlse64.v v8, (a0), zero
 ; CHECK-NEXT:    li a0, -1
 ; CHECK-NEXT:    sw a0, 12(sp)
 ; CHECK-NEXT:    sw zero, 8(sp)
 ; CHECK-NEXT:    addi a0, sp, 8
 ; CHECK-NEXT:    vlse64.v v10, (a0), zero
 ; CHECK-NEXT:    vand.vv v8, v8, v10
 ; CHECK-NEXT:    addi sp, sp, 16
 ; CHECK-NEXT:    ret
   %1 = insertelement <vscale x 2 x i64> poison, i64 %x, i32 0
   %2 = shufflevector <vscale x 2 x i64> %1, <vscale x 2 x i64> poison, <vscale x 2 x i32> zeroinitializer
   %3 = and <vscale x 2 x i64> %2, shufflevector(<vscale x 2 x i64> insertelement(<vscale x 2 x i64> poison, i64 u0xffffffff00000000, i32 0), <vscale x 2 x i64> poison, <vscale x 2 x i32> zeroinitializer)
   ret <vscale x 2 x i64> %3
 }