[llvm] [RISCV][ISel] Combine scalable vector add/sub/mul with zero/sign extension (PR #72340)

[via llvm-commits]

================
@@ -12694,19 +12705,45 @@ struct NodeExtensionHelper {
 
   using CombineToTry = std::function<std::optional<CombineResult>(
       SDNode * /*Root*/, const NodeExtensionHelper & /*LHS*/,
-      const NodeExtensionHelper & /*RHS*/)>;
+      const NodeExtensionHelper & /*RHS*/, SelectionDAG &,
+      const RISCVSubtarget &)>;
 
   /// Check if this node needs to be fully folded or extended for all users.
   bool needToPromoteOtherUsers() const { return EnforceOneUse; }
 
   /// Helper method to set the various fields of this struct based on the
   /// type of \p Root.
-  void fillUpExtensionSupport(SDNode *Root, SelectionDAG &DAG) {
+  void fillUpExtensionSupport(SDNode *Root, SelectionDAG &DAG,
+                              const RISCVSubtarget &Subtarget) {
     SupportsZExt = false;
     SupportsSExt = false;
     EnforceOneUse = true;
     CheckMask = true;
     switch (OrigOperand.getOpcode()) {
+    case ISD::ZERO_EXTEND: {
+      SupportsZExt = true;
+      SDLoc DL(Root);
+      MVT VT = Root->getSimpleValueType(0);
+      if (VT.isFixedLengthVector()) {
+        MVT ContainerVT = getContainerForFixedLengthVector(DAG, VT, Subtarget);
+        std::tie(Mask, VL) =
+            getDefaultVLOps(VT, ContainerVT, DL, DAG, Subtarget);
+      } else if (VT.isVector())
+        std::tie(Mask, VL) = getDefaultScalableVLOps(VT, DL, DAG, Subtarget);
+      break;
+    }
+    case ISD::SIGN_EXTEND: {
+      SupportsSExt = true;
+      SDLoc DL(Root);
+      MVT VT = Root->getSimpleValueType(0);
+      if (VT.isFixedLengthVector()) {
+        MVT ContainerVT = getContainerForFixedLengthVector(DAG, VT, Subtarget);
+        std::tie(Mask, VL) =
+            getDefaultVLOps(VT, ContainerVT, DL, DAG, Subtarget);
+      } else if (VT.isVector())
+        std::tie(Mask, VL) = getDefaultScalableVLOps(VT, DL, DAG, Subtarget);
+      break;
+    }
----------------
qcolombet wrote:

As far as I can tell, only `Supports(Z|S)Ext` changes between these two cases.
Could you merge them and set `Supports(Z|S)Ext` accordingly.
```
SupportsSExt = Opc == ISD::SIGN_EXTEND;
SupportsZExt = Opc == ISD::ZERO_EXTEND;
// then same code.
```

https://github.com/llvm/llvm-project/pull/72340