[RISCV] Fix assertion failure from performBUILD_VECTORCombine when the binop is a shift.

[topperc]

The RHS of a shift can have a different type than the LHS. If there are undefs in the vector, we need the undef added to the RHS to match the type of any shift amounts that are also added to the vector.

For now just don't add shifts if their RHS and LHS don't match.

[llvmbot]

@llvm/pr-subscribers-backend-risc-v

Author: Craig Topper (topperc)

Changes

The RHS of a shift can have a different type than the LHS. If there are undefs in the vector, we need the undef added to the RHS to match the type of any shift amounts that are also added to the vector.

For now just don't add shifts if their RHS and LHS don't match.

---

Full diff: https://github.com/llvm/llvm-project/pull/69349.diff

2 Files Affected:

• (modified) llvm/lib/Target/RISCV/RISCVISelLowering.cpp (+5-1)
• (modified) llvm/test/CodeGen/RISCV/rvv/fixed-vectors-buildvec-of-binop.ll (+12)

diff --git a/llvm/lib/Target/RISCV/RISCVISelLowering.cpp b/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
index e8f001e491cdcaa..fe90a6a5a28d248 100644
--- a/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
+++ b/llvm/lib/Target/RISCV/RISCVISelLowering.cpp
@@ -13717,7 +13717,7 @@ static SDValue performSELECTCombine(SDNode *N, SelectionDAG &DAG,
   return tryFoldSelectIntoOp(N, DAG, FalseVal, TrueVal, /*Swapped*/true);
 }
 
-/// IF we have a build_vector where each lane is binop X, C, where C
+/// If we have a build_vector where each lane is binop X, C, where C
 /// is a constant (but not necessarily the same constant on all lanes),
 /// form binop (build_vector x1, x2, ...), (build_vector c1, c2, c3, ..).
 /// We assume that materializing a constant build vector will be no more
@@ -13759,6 +13759,10 @@ static SDValue performBUILD_VECTORCombine(SDNode *N, SelectionDAG &DAG,
     if (!isa<ConstantSDNode>(Op.getOperand(1)) &&
         !isa<ConstantFPSDNode>(Op.getOperand(1)))
       return SDValue();
+    // FIXME: Return failure if the RHS type doesn't match the LHS. Shifts may
+    // have different LHS and RHS types.
+    if (Op.getOperand(0).getValueType() != Op.getOperand(1).getValueType())
+      return SDValue();
     RHSOps.push_back(Op.getOperand(1));
   }
 
diff --git a/llvm/test/CodeGen/RISCV/rvv/fixed-vectors-buildvec-of-binop.ll b/llvm/test/CodeGen/RISCV/rvv/fixed-vectors-buildvec-of-binop.ll
index 2bbc04172bd1421..717dfb1bfd00537 100644
--- a/llvm/test/CodeGen/RISCV/rvv/fixed-vectors-buildvec-of-binop.ll
+++ b/llvm/test/CodeGen/RISCV/rvv/fixed-vectors-buildvec-of-binop.ll
@@ -442,3 +442,15 @@ define <4 x i32> @add_general_splat(i32 %a, i32 %b, i32 %c, i32 %d, i32 %e) {
   %v3 = insertelement <4 x i32> %v2, i32 %e3, i32 3
   ret <4 x i32> %v3
 }
+
+; This test previously failed with an assertion failure because constant shift
+; amounts are type legalized early.
+define void @buggy(i32 %0) #0 {
+entry:
+  %mul.us.us.i.3 = shl i32 %0, 1
+  %1 = insertelement <4 x i32> zeroinitializer, i32 %mul.us.us.i.3, i64 0
+  %2 = or <4 x i32> %1, <i32 1, i32 1, i32 1, i32 1>
+  %3 = shufflevector <4 x i32> %2, <4 x i32> zeroinitializer, <4 x i32> zeroinitializer
+  store <4 x i32> %3, ptr null, align 16
+  ret void
+}

[lukel97]

LGTM