[DAGCombiner] Make it so that `udiv` can be folded with `(select c, NonZero, 1)`

[goldstein.w.n]

Summary

This is done by allowing speculation of udiv if we can prove the
denominator is non-zero.

https://alive2.llvm.org/ce/z/VNCt_q

[goldstein.w.n]

Rebase

[pengfei]

llvm/include/llvm/CodeGen/SelectionDAG.h

2330

Why don't refactor it directly?

[goldstein.w.n]

llvm/include/llvm/CodeGen/SelectionDAG.h

2330

I figured they have seperate purposes. The Opcode version is checking whether an opcode is unconditionally safe to speculate, node is asking about this specific case. But I'll check the other usages and see if the refactored version makes sense.

[goldstein.w.n]

just refactor isSafeToSpeculativelyExecute

[goldstein.w.n]

Did general refactror.

[craig.topper]

In D149198#4297802, @goldstein.w.n wrote:

Did general refactror.

Does that increase the review burden? Now we have to audit the other transforms to make sure the don't change operand 1 in a way that would be unsafe?

Plus we should have additional tests?

[pengfei]

In D149198#4297846, @craig.topper wrote:

In D149198#4297802, @goldstein.w.n wrote:

Did general refactror.

Does that increase the review burden? Now we have to audit the other transforms to make sure the don't change operand 1 in a way that would be unsafe?

Plus we should have additional tests?

Maybe we can use another NFC patch to refactor the interface only?

[goldstein.w.n]

In D149198#4298034, @pengfei wrote:

In D149198#4297846, @craig.topper wrote:

In D149198#4297802, @goldstein.w.n wrote:

Did general refactror.

Does that increase the review burden? Now we have to audit the other transforms to make sure the don't change operand 1 in a way that would be unsafe?

Plus we should have additional tests?

Maybe we can use another NFC patch to refactor the interface only?

Sure. I'll leave as previous version in this patch and make the refactor a follow up.
Edit: Looking at some of the other locations more closely. I think it makes
sense to leave as a seperate API. isSafeToSpeculativelyExecute is used to query
about the opcode without knowing the operands so serves a different purpose.

[craig.topper]

In D149198#4298034, @pengfei wrote:

In D149198#4297846, @craig.topper wrote:

In D149198#4297802, @goldstein.w.n wrote:

Did general refactror.

Does that increase the review burden? Now we have to audit the other transforms to make sure the don't change operand 1 in a way that would be unsafe?

Plus we should have additional tests?

Maybe we can use another NFC patch to refactor the interface only?

It wouldn't be NFC. It would make some UDIV instructions supported by additional transforms.

[goldstein.w.n]

Go back to original logic

[goldstein.w.n]

Rebase

[RKSimon]

llvm/include/llvm/CodeGen/SelectionDAG.h

2344

Add doxygen description

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

2389–2390

update comment

[goldstein.w.n]

Fix comments

[goldstein.w.n]

n/a

[goldstein.w.n]

Rebase

[pengfei]

LGTM.

llvm/include/llvm/CodeGen/SelectionDAG.h

2344

safe

[goldstein.w.n]

Rebase

[goldstein.w.n]

Rebase

[xbolva00]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

2391

but UREM is not handled above..

[goldstein.w.n]

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

2391

Well the comment still holds true, but indeed could include urem in this patch. Will update.

[goldstein.w.n]

Rebase

[goldstein.w.n]

Sorry for long delay on this, going to push this shortly.

[goldstein.w.n]

Rebase

llvm/include/llvm/CodeGen/SelectionDAG.h

public:
 
   /// Get the (commutative) neutral element for the given opcode, if it exists.
   SDValue getNeutralElement(unsigned Opcode, const SDLoc &DL, EVT VT,
                             SDNodeFlags Flags);
 
   /// Some opcodes may create immediate undefined behavior when used with some
   /// values (integer division-by-zero for example). Therefore, these operations
   /// are not generally safe to move around or change.
   bool isSafeToSpeculativelyExecute(unsigned Opcode) const {
     switch (Opcode) {
     case ISD::SDIV:
     case ISD::SREM:
     case ISD::SDIVREM:
     case ISD::UDIV:
     case ISD::UREM:
     case ISD::UDIVREM:
       return false;
     default:
       return true;
     }
   }
 
+  /// Check if the provided node is save to speculatively executed given its
+  /// current arguments. So, while `udiv` the opcode is not safe to
+  /// speculatively execute, a given `udiv` node may be if the denominator is
+  /// known nonzero.
+  bool isSafeToSpeculativelyExecuteNode(const SDNode *N) const {
+    switch (N->getOpcode()) {
+    case ISD::UDIV:
+      return isKnownNeverZero(N->getOperand(1));
+    default:
+      return isSafeToSpeculativelyExecute(N->getOpcode());
+    }
+  }
+
 private:
   void InsertNode(SDNode *N);
   bool RemoveNodeFromCSEMaps(SDNode *N);
   void AddModifiedNodeToCSEMaps(SDNode *N);
   SDNode *FindModifiedNodeSlot(SDNode *N, SDValue Op, void *&InsertPos);
   SDNode *FindModifiedNodeSlot(SDNode *N, SDValue Op1, SDValue Op2,
                                void *&InsertPos);
   SDNode *FindModifiedNodeSlot(SDNode *N, ArrayRef<SDValue> Ops,

llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp

static SDValue foldSelectWithIdentityConstant(SDNode *N, SelectionDAG &DAG,
   SDValue N0 = N->getOperand(0);
   SDValue N1 = N->getOperand(1);
   if (ShouldCommuteOperands)
     std::swap(N0, N1);
 
   // TODO: Should this apply to scalar select too?
   if (N1.getOpcode() != ISD::VSELECT || !N1.hasOneUse())
     return SDValue();
 
-  // We can't hoist div/rem because of immediate UB (not speculatable).
-  unsigned Opcode = N->getOpcode();
-  if (!DAG.isSafeToSpeculativelyExecute(Opcode))
+  // We can't hoist all instructions because of immediate UB (not speculatable).
+  // For example div/rem by zero.
+  if (!DAG.isSafeToSpeculativelyExecuteNode(N))
     return SDValue();
 
+  unsigned Opcode = N->getOpcode();
   EVT VT = N->getValueType(0);
   SDValue Cond = N1.getOperand(0);
   SDValue TVal = N1.getOperand(1);
   SDValue FVal = N1.getOperand(2);
 
   // This transform increases uses of N0, so freeze it to be safe.
   // binop N0, (vselect Cond, IDC, FVal) --> vselect Cond, N0, (binop N0, FVal)
   unsigned OpNo = ShouldCommuteOperands ? 0 : 1;

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

; CHECK-NEXT:    ret
   %head2 = insertelement <vscale x 8 x i32> poison, i32 %b, i32 0
   %splat = shufflevector <vscale x 8 x i32> %head2, <vscale x 8 x i32> poison, <vscale x 8 x i32> zeroinitializer
   %vs = select <vscale x 8 x i1> %mask, <vscale x 8 x i32> %splat, <vscale x 8 x i32> %one
   %vc = udiv <vscale x 8 x i32> %va, %vs
   ret <vscale x 8 x i32> %vc
 }
 
 define <vscale x 8 x i32> @vdivu_vi_mask_nxv8i32(<vscale x 8 x i32> %va, <vscale x 8 x i1> %mask) {
-; CHECK-LABEL: vdivu_vi_mask_nxv8i32:
-; CHECK:       # %bb.0:
-; CHECK-NEXT:    vsetvli a0, zero, e32, m4, ta, ma
-; CHECK-NEXT:    vmv.v.i v12, 1
-; CHECK-NEXT:    vmerge.vim v12, v12, 7, v0
-; CHECK-NEXT:    vdivu.vv v8, v8, v12
-; CHECK-NEXT:    ret
+; RV32-LABEL: vdivu_vi_mask_nxv8i32:
+; RV32:       # %bb.0:
+; RV32-NEXT:    lui a0, 149797
+; RV32-NEXT:    addi a0, a0, -1755
+; RV32-NEXT:    vsetvli a1, zero, e32, m4, ta, mu
+; RV32-NEXT:    vmulhu.vx v12, v8, a0
+; RV32-NEXT:    vsub.vv v16, v8, v12
+; RV32-NEXT:    vsrl.vi v16, v16, 1
+; RV32-NEXT:    vadd.vv v12, v16, v12
+; RV32-NEXT:    vsrl.vi v8, v12, 2, v0.t
+; RV32-NEXT:    ret
+;
+; RV64-LABEL: vdivu_vi_mask_nxv8i32:
+; RV64:       # %bb.0:
+; RV64-NEXT:    lui a0, 149797
+; RV64-NEXT:    addiw a0, a0, -1755
+; RV64-NEXT:    vsetvli a1, zero, e32, m4, ta, mu
+; RV64-NEXT:    vmulhu.vx v12, v8, a0
+; RV64-NEXT:    vsub.vv v16, v8, v12
+; RV64-NEXT:    vsrl.vi v16, v16, 1
+; RV64-NEXT:    vadd.vv v12, v16, v12
+; RV64-NEXT:    vsrl.vi v8, v12, 2, v0.t
+; RV64-NEXT:    ret
   %head1 = insertelement <vscale x 8 x i32> poison, i32 1, i32 0
   %one = shufflevector <vscale x 8 x i32> %head1, <vscale x 8 x i32> poison, <vscale x 8 x i32> zeroinitializer
   %head2 = insertelement <vscale x 8 x i32> poison, i32 7, i32 0
   %splat = shufflevector <vscale x 8 x i32> %head2, <vscale x 8 x i32> poison, <vscale x 8 x i32> zeroinitializer
   %vs = select <vscale x 8 x i1> %mask, <vscale x 8 x i32> %splat, <vscale x 8 x i32> %one
   %vc = udiv <vscale x 8 x i32> %va, %vs
   ret <vscale x 8 x i32> %vc
 }

llvm/test/CodeGen/X86/divrem-by-select.ll

 ; CHECK-X64-V3-NEXT:    vmovq %rax, %xmm1
 ; CHECK-X64-V3-NEXT:    vpunpcklqdq {{.*#+}} xmm0 = xmm1[0],xmm0[0]
 ; CHECK-X64-V3-NEXT:    retq
 ;
 ; CHECK-X64-V4-LABEL: udiv_identity_const:
 ; CHECK-X64-V4:       # %bb.0:
 ; CHECK-X64-V4-NEXT:    vpsllq $63, %xmm0, %xmm0
 ; CHECK-X64-V4-NEXT:    vpmovq2m %xmm0, %k1
-; CHECK-X64-V4-NEXT:    vpbroadcastq {{.*#+}} xmm0 = [1,1]
-; CHECK-X64-V4-NEXT:    vpbroadcastq {{\.?LCPI[0-9]+_[0-9]+}}(%rip), %xmm0 {%k1}
-; CHECK-X64-V4-NEXT:    vpextrq $1, %xmm0, %rcx
-; CHECK-X64-V4-NEXT:    vpextrq $1, %xmm1, %rax
-; CHECK-X64-V4-NEXT:    xorl %edx, %edx
-; CHECK-X64-V4-NEXT:    divq %rcx
-; CHECK-X64-V4-NEXT:    movq %rax, %rcx
-; CHECK-X64-V4-NEXT:    vmovq %xmm0, %rsi
-; CHECK-X64-V4-NEXT:    vmovq %xmm1, %rax
-; CHECK-X64-V4-NEXT:    xorl %edx, %edx
-; CHECK-X64-V4-NEXT:    divq %rsi
-; CHECK-X64-V4-NEXT:    vmovq %rcx, %xmm0
-; CHECK-X64-V4-NEXT:    vmovq %rax, %xmm1
-; CHECK-X64-V4-NEXT:    vpunpcklqdq {{.*#+}} xmm0 = xmm1[0],xmm0[0]
+; CHECK-X64-V4-NEXT:    vpextrq $1, %xmm1, %rdx
+; CHECK-X64-V4-NEXT:    movabsq $3353953467947191203, %rax # imm = 0x2E8BA2E8BA2E8BA3
+; CHECK-X64-V4-NEXT:    mulxq %rax, %rcx, %rcx
+; CHECK-X64-V4-NEXT:    vmovq %rcx, %xmm0
+; CHECK-X64-V4-NEXT:    vmovq %xmm1, %rdx
+; CHECK-X64-V4-NEXT:    mulxq %rax, %rax, %rax
+; CHECK-X64-V4-NEXT:    vmovq %rax, %xmm2
+; CHECK-X64-V4-NEXT:    vpunpcklqdq {{.*#+}} xmm0 = xmm2[0],xmm0[0]
+; CHECK-X64-V4-NEXT:    vpsrlq $1, %xmm0, %xmm1 {%k1}
+; CHECK-X64-V4-NEXT:    vmovdqa %xmm1, %xmm0
 ; CHECK-X64-V4-NEXT:    retq
   %d = select <2 x i1> %c, <2 x i64> <i64 11, i64 11>, <2 x i64> <i64 1, i64 1>
   %r = udiv <2 x i64> %x, %d
   ret <2 x i64> %r
 }
 
 
 define <2 x i64> @udiv_identity_const_todo_getter_nonzero(<2 x i1> %c, <2 x i64> %x) {