[DAGCombiner][RISCV] Add target hook to decide hoisting LogicOp with extension

llvm/include/llvm/CodeGen/TargetLowering.h

@@ -4462,6 +4462,16 @@ class TargetLowering : public TargetLoweringBase {
     return false;
   }
 
+  /// Return true if it is profitable to hoist a LogicOp where both operands
+  /// have the same extension op. This transformation may not be desirable if
+  /// it disrupts a particularly auspicious target-specific tree (e.g.
+  /// (or disjoint (zext A), (zext B)) -> vwaddu.wv on RISC-V). By default it
+  /// returns true.
+  virtual bool isDesirableToHoistLogicOpWithExt(const SDNode *LogicOp,
+                                                unsigned ExtOp) const {
+    return true;
+  }
+
   /// Return true if the target supports swifterror attribute. It optimizes
   /// loads and stores to reading and writing a specific register.
   virtual bool supportSwiftError() const {

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

@@ -5981,6 +5981,9 @@ SDValue DAGCombiner::hoistLogicOpWithSameOpcodeHands(SDNode *N) {
          HandOpcode == ISD::ANY_EXTEND_VECTOR_INREG) &&
         LegalTypes && !TLI.isTypeDesirableForOp(LogicOpcode, XVT))
       return SDValue();
+    // If it is not desirable to hoist LogicOp with extension.
+    if (!TLI.isDesirableToHoistLogicOpWithExt(N, HandOpcode))
+      return SDValue();
     // logic_op (hand_op X), (hand_op Y) --> hand_op (logic_op X, Y)
     SDValue Logic = DAG.getNode(LogicOpcode, DL, XVT, X, Y);
     if (HandOpcode == ISD::SIGN_EXTEND_INREG)

llvm/lib/Target/RISCV/RISCVISelLowering.cpp

@@ -19893,6 +19893,14 @@ bool RISCVTargetLowering::isDesirableToCommuteWithShift(
   return true;
 }
 
+bool RISCVTargetLowering::isDesirableToHoistLogicOpWithExt(
+    const SDNode *LogicOp, unsigned ExtOp) const {
+  if ((ExtOp == ISD::ZERO_EXTEND || ExtOp == ISD::SIGN_EXTEND) &&
+      NodeExtensionHelper::isSupportedRoot(LogicOp, Subtarget))
+    return false;
+  return true;
+}
+
 bool RISCVTargetLowering::targetShrinkDemandedConstant(
     SDValue Op, const APInt &DemandedBits, const APInt &DemandedElts,
     TargetLoweringOpt &TLO) const {

llvm/lib/Target/RISCV/RISCVISelLowering.h

@@ -747,6 +747,9 @@ class RISCVTargetLowering : public TargetLowering {
   bool isDesirableToCommuteWithShift(const SDNode *N,
                                      CombineLevel Level) const override;
 
+  bool isDesirableToHoistLogicOpWithExt(const SDNode *LogicOp,
+                                        unsigned ExtOp) const override;
+
   /// If a physical register, this returns the register that receives the
   /// exception address on entry to an EH pad.
   Register

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

@@ -1417,31 +1417,25 @@ define <vscale x 2 x i32> @vwaddu_vv_disjoint_or_add(<vscale x 2 x i8> %x.i8, <v
   ret <vscale x 2 x i32> %add
 }
 
-; TODO: We could select vwaddu.vv, but when both arms of the or are the same
-; DAGCombiner::hoistLogicOpWithSameOpcodeHands moves the zext above the or.
 define <vscale x 2 x i32> @vwaddu_vv_disjoint_or(<vscale x 2 x i16> %x.i16, <vscale x 2 x i16> %y.i16) {
 ; CHECK-LABEL: vwaddu_vv_disjoint_or:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    vsetvli a0, zero, e16, mf2, ta, ma
-; CHECK-NEXT:    vor.vv v9, v8, v9
-; CHECK-NEXT:    vsetvli zero, zero, e32, m1, ta, ma
-; CHECK-NEXT:    vzext.vf2 v8, v9
+; CHECK-NEXT:    vwaddu.vv v10, v8, v9
+; CHECK-NEXT:    vmv1r.v v8, v10
 ; CHECK-NEXT:    ret
   %x.i32 = zext <vscale x 2 x i16> %x.i16 to <vscale x 2 x i32>
   %y.i32 = zext <vscale x 2 x i16> %y.i16 to <vscale x 2 x i32>
   %or = or disjoint <vscale x 2 x i32> %x.i32, %y.i32
   ret <vscale x 2 x i32> %or
 }
 
-; TODO: We could select vwadd.vv, but when both arms of the or are the same
-; DAGCombiner::hoistLogicOpWithSameOpcodeHands moves the zext above the or.
 define <vscale x 2 x i32> @vwadd_vv_disjoint_or(<vscale x 2 x i16> %x.i16, <vscale x 2 x i16> %y.i16) {
 ; CHECK-LABEL: vwadd_vv_disjoint_or:
 ; CHECK:       # %bb.0:
 ; CHECK-NEXT:    vsetvli a0, zero, e16, mf2, ta, ma
-; CHECK-NEXT:    vor.vv v9, v8, v9
-; CHECK-NEXT:    vsetvli zero, zero, e32, m1, ta, ma
-; CHECK-NEXT:    vsext.vf2 v8, v9
+; CHECK-NEXT:    vwadd.vv v10, v8, v9
+; CHECK-NEXT:    vmv1r.v v8, v10
 ; CHECK-NEXT:    ret
   %x.i32 = sext <vscale x 2 x i16> %x.i16 to <vscale x 2 x i32>
   %y.i32 = sext <vscale x 2 x i16> %y.i16 to <vscale x 2 x i32>