[llvm] [RISCV] Fix assertion failure from performBUILD_VECTORCombine when the binop is a shift. (PR #69349)

[Craig Topper via llvm-commits]

https://github.com/topperc created https://github.com/llvm/llvm-project/pull/69349

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.

>From c1b1469dd9cabfae7a3e89df70729d28eb38fe2d Mon Sep 17 00:00:00 2001
From: Craig Topper <craig.topper at sifive.com>
Date: Tue, 17 Oct 2023 08:57:14 -0700
Subject: [PATCH] [RISCV] Fix assertion failure from performBUILD_VECTORCombine
 when the binop is a shift.

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.
---
 llvm/lib/Target/RISCV/RISCVISelLowering.cpp          |  6 +++++-
 .../RISCV/rvv/fixed-vectors-buildvec-of-binop.ll     | 12 ++++++++++++
 2 files changed, 17 insertions(+), 1 deletion(-)

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
+}