[RISCV] Fix crash a vector add has a 4x sext and zext operand.

[craig.topper]

Summary

We can narrow one of the extends and keep the other original by
using a vwaddu.wv or vwadd.wv.

We were previously forgetting to keep the original operand and
instead took the source of its extend. This resulted in a type
mismatch that later failed with an impossible physical register copy.

To fix this I've refactored some code to maintain information about
whether the source needs to be extended at all for longer so we could
use it in materialize.

[qcolombet]

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

8459

That part shouldn't be needed. The getSource() should do the right thing.
I believe the reason it doesn't is because the NodeExtensionHelper got out-of-sync with the actual operand, maybe when we commute the operand of VWADD_W?

Let me check something.

[craig.topper]

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

8459

I don't think it had anything to with commuting. I think we lost track of the fact that we were matching to VWADDU_W and didn't want to look through the VSEXT_VL that was on the LHS. In this function we were acting based only on the opcode of the operand and not what pattern we were matching.

[qcolombet]

Thanks for the quick fix.

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

8459

Ok well, got that wrong :).

You're right when we don't require any extension we just want to use the original value.

8738

That part of the change may not be needed to fix the issue, but I like your refactoring better anyway.
So no need to split it.

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

SDValue getSource() const {
     }
   }
 
   /// Check if this instance represents a splat.
   bool isSplat() const {
     return OrigOperand.getOpcode() == RISCVISD::VMV_V_X_VL;
   }
 
-  /// Get or create a value that can feed \p Root with the given \p ExtOpc.
-  /// If \p ExtOpc is None, this returns the source of this operand.
+  /// Get or create a value that can feed \p Root with the given extension \p
+  /// SExt. If \p SExt is None, this returns the source of this operand.
   /// \see ::getSource().
   SDValue getOrCreateExtendedOp(const SDNode *Root, SelectionDAG &DAG,
-                                Optional<unsigned> ExtOpc) const {
+                                Optional<bool> SExt) const {
+    if (!SExt.has_value())
+      return OrigOperand;
+
+    MVT NarrowVT = getNarrowType(Root);
+
     SDValue Source = getSource();
-    if (!ExtOpc)
+    if (Source.getValueType() == NarrowVT)
       return Source;
 
-    MVT NarrowVT = getNarrowType(Root);
+    unsigned ExtOpc = *SExt ? RISCVISD::VSEXT_VL : RISCVISD::VZEXT_VL;
+
     // If we need an extension, we should be changing the type.
-    assert(Source.getValueType() != NarrowVT && "Needless extension");
     SDLoc DL(Root);
     auto [Mask, VL] = getMaskAndVL(Root);
     switch (OrigOperand.getOpcode()) {
     case RISCVISD::VSEXT_VL:
     case RISCVISD::VZEXT_VL:
-      return DAG.getNode(*ExtOpc, DL, NarrowVT, Source, Mask, VL);
+      return DAG.getNode(ExtOpc, DL, NarrowVT, Source, Mask, VL);
     case RISCVISD::VMV_V_X_VL:
       return DAG.getNode(RISCVISD::VMV_V_X_VL, DL, NarrowVT,
                          DAG.getUNDEF(NarrowVT), Source.getOperand(1), VL);
     default:
       // Other opcodes can only come from the original LHS of VW(ADD|SUB)_W_VL
       // and that operand should already have the right NarrowVT so no
       // extension should be required at this point.
       llvm_unreachable("Unsupported opcode");
struct NodeExtensionHelper {
   static SmallVector<CombineToTry> getSupportedFoldings(const SDNode *Root);
 };
 
 /// Helper structure that holds all the necessary information to materialize a
 /// combine that does some extension folding.
 struct CombineResult {
   /// Opcode to be generated when materializing the combine.
   unsigned TargetOpcode;
-  /// Extension opcode to be applied to the source of LHS when materializing
-  /// TargetOpcode.
-  /// \see NodeExtensionHelper::getSource().
-  Optional<unsigned> LHSExtOpc;
-  /// Extension opcode to be applied to the source of RHS when materializing
-  /// TargetOpcode.
-  Optional<unsigned> RHSExtOpc;
+  // No value means no extension is needed. If extension is needed, the value
+  // indicates if it needs to be sign extended.
+  Optional<bool> SExtLHS;
+  Optional<bool> SExtRHS;
   /// Root of the combine.
   SDNode *Root;
   /// LHS of the TargetOpcode.
   NodeExtensionHelper LHS;
   /// RHS of the TargetOpcode.
   NodeExtensionHelper RHS;
 
   CombineResult(unsigned TargetOpcode, SDNode *Root,
                 const NodeExtensionHelper &LHS, Optional<bool> SExtLHS,
                 const NodeExtensionHelper &RHS, Optional<bool> SExtRHS)
-      : TargetOpcode(TargetOpcode), Root(Root), LHS(LHS), RHS(RHS) {
-    MVT NarrowVT = NodeExtensionHelper::getNarrowType(Root);
-    if (SExtLHS && LHS.getSource().getValueType() != NarrowVT)
-      LHSExtOpc = *SExtLHS ? RISCVISD::VSEXT_VL : RISCVISD::VZEXT_VL;
-    if (SExtRHS && RHS.getSource().getValueType() != NarrowVT)
-      RHSExtOpc = *SExtRHS ? RISCVISD::VSEXT_VL : RISCVISD::VZEXT_VL;
-  }
+      : TargetOpcode(TargetOpcode), SExtLHS(SExtLHS), SExtRHS(SExtRHS),
+        Root(Root), LHS(LHS), RHS(RHS) {}
 
   /// Return a value that uses TargetOpcode and that can be used to replace
   /// Root.
   /// The actual replacement is *not* done in that method.
   SDValue materialize(SelectionDAG &DAG) const {
     SDValue Mask, VL, Merge;
     std::tie(Mask, VL) = NodeExtensionHelper::getMaskAndVL(Root);
     Merge = Root->getOperand(2);
     return DAG.getNode(TargetOpcode, SDLoc(Root), Root->getValueType(0),
-                       LHS.getOrCreateExtendedOp(Root, DAG, LHSExtOpc),
-                       RHS.getOrCreateExtendedOp(Root, DAG, RHSExtOpc), Merge,
+                       LHS.getOrCreateExtendedOp(Root, DAG, SExtLHS),
+                       RHS.getOrCreateExtendedOp(Root, DAG, SExtRHS), Merge,
                        Mask, VL);
   }
 };
 
 /// Check if \p Root follows a pattern Root(ext(LHS), ext(RHS))
 /// where `ext` is the same for both LHS and RHS (i.e., both are sext or both
 /// are zext) and LHS and RHS can be folded into Root.
 /// AllowSExt and AllozZExt define which form `ext` can take in this pattern.

llvm/test/CodeGen/RISCV/rvv/fixed-vectors-vwaddu.ll

; RV64-NEXT:    ret
   %a = load <2 x i32>, <2 x i32>* %x
   %b = load i64, i64* %y
   %d = insertelement <2 x i64> poison, i64 %b, i64 0
   %e = shufflevector <2 x i64> %d, <2 x i64> poison, <2 x i32> zeroinitializer
   %f = zext <2 x i32> %a to <2 x i64>
   %g = add <2 x i64> %e, %f
   ret <2 x i64> %g
 }
+
+define <4 x i64> @crash(<4 x i16> %x, <4 x i16> %y) {
+; CHECK-LABEL: crash:
+; CHECK:       # %bb.0:
+; CHECK-NEXT:    vsetivli zero, 4, e64, m2, ta, ma
+; CHECK-NEXT:    vsext.vf4 v10, v8
+; CHECK-NEXT:    vsetvli zero, zero, e32, m1, ta, ma
+; CHECK-NEXT:    vzext.vf2 v8, v9
+; CHECK-NEXT:    vwaddu.wv v10, v10, v8
+; CHECK-NEXT:    vmv2r.v v8, v10
+; CHECK-NEXT:    ret
+  %a = sext <4 x i16> %x to <4 x i64>
+  %b = zext <4 x i16> %y to <4 x i64>
+  %c = add <4 x i64> %a, %b
+  ret <4 x i64> %c
+}