[DAG] vector div/rem with any zero element in divisor is undef

[spatel]

Summary

This is the backend counterpart to:
https://reviews.llvm.org/rL297390
https://reviews.llvm.org/rL297409
and follow-up to:
https://reviews.llvm.org/rL297384

It surprised me that we need to duplicate the check in FoldConstantArithmetic and FoldConstantVectorArithmetic, but one or the other doesn't catch all of the test cases. There is an existing code comment about merging those someday. If we fix more cases like this to properly recognize undef (signed overflow in sdiv, oversized shifts, etc), we should fix the code structure or at least make a helper function to avoid repeating code.

[efriedma]

I agree it's weird that there are two functions for constant folding... but you can leave it for now, I guess.

Maybe check for null or undef, rather than just null?

lib/CodeGen/SelectionDAG/SelectionDAG.cpp

3858

I think Ops[1] must be a BUILD_VECTOR here, but it isn't obvious from the structure of the code. Maybe explicitly cast<> it to BuildVectorSDNode?

[spatel]

Patch updated:

1. The code duplication was already bugging me, and we know it's only going to grow over time, so I'm proposing a single place to house the undef checking. Fixing the duplication between FoldConstantArithmetic and FoldConstantVectorArithmetic is another step.
2. Fixed to check for zero or undef elements in the build vector.

[efriedma]

LGTM with a minor tweak.

include/llvm/CodeGen/SelectionDAG.h

744

This is a little unclear... maybe "return true if the result of this operation is always undefined"?

[spatel]

In D30826#698045, @efriedma wrote:

I agree it's weird that there are two functions for constant folding... but you can leave it for now, I guess.

Maybe check for null or undef, rather than just null?

include/llvm/CodeGen/SelectionDAG.h

744

Yep - that's better. Thanks!

include/llvm/CodeGen/SelectionDAG.h

SDValue getCALLSEQ_END(SDValue Chain, SDValue Op1, SDValue Op2,
     Ops.push_back(Chain);
     Ops.push_back(Op1);
     Ops.push_back(Op2);
     if (InGlue.getNode())
       Ops.push_back(InGlue);
     return getNode(ISD::CALLSEQ_END, DL, NodeTys, Ops);
   }
 
+  /// For the given opcode and operand values, return true if any operand value
+  /// makes this operation undefined.
+  bool isUndef(unsigned Opcode, ArrayRef<SDValue> Ops);
+
   /// Return an UNDEF node. UNDEF does not have a useful SDLoc.
   SDValue getUNDEF(EVT VT) {
     return getNode(ISD::UNDEF, SDLoc(), VT);
   }
 
   /// Return a GLOBAL_OFFSET_TABLE node. This does not have a useful SDLoc.
   SDValue getGLOBAL_OFFSET_TABLE(EVT VT) {
     return getNode(ISD::GLOBAL_OFFSET_TABLE, SDLoc(), VT);

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

 }
 
 static SDValue simplifyDivRem(SDNode *N, SelectionDAG &DAG) {
   SDValue N0 = N->getOperand(0);
   SDValue N1 = N->getOperand(1);
   EVT VT = N->getValueType(0);
   SDLoc DL(N);
 
-  // X / undef -> undef
-  // X % undef -> undef
-  if (N1.isUndef())
-    return N1;
-
-  // X / 0 --> undef
-  // X % 0 --> undef
-  // We don't need to preserve faults!
-  if (isNullConstantOrNullSplatConstant(N1))
+  if (DAG.isUndef(N->getOpcode(), {N0, N1}))
     return DAG.getUNDEF(VT);
 
   // undef / X -> 0
   // undef % X -> 0
   if (N0.isUndef())
     return DAG.getConstant(0, DL, VT);
 
   return SDValue();

lib/CodeGen/SelectionDAG/SelectionDAG.cpp

 }
 
 SDValue SelectionDAG::FoldConstantArithmetic(unsigned Opcode, const SDLoc &DL,
                                              EVT VT, const ConstantSDNode *Cst1,
                                              const ConstantSDNode *Cst2) {
   if (Cst1->isOpaque() || Cst2->isOpaque())
     return SDValue();
 
-  // Division/remainder with a zero divisor is undefined behavior.
-  if ((Opcode == ISD::SDIV || Opcode == ISD::UDIV ||
-       Opcode == ISD::SREM || Opcode == ISD::UREM) &&
-      Cst2->isNullValue())
-    return getUNDEF(VT);
-
   std::pair<APInt, bool> Folded = FoldValue(Opcode, Cst1->getAPIntValue(),
                                             Cst2->getAPIntValue());
   if (!Folded.second)
     return SDValue();
   return getConstant(Folded.first, DL, VT);
 }
 
 SDValue SelectionDAG::FoldSymbolOffset(unsigned Opcode, EVT VT,
SDValue SelectionDAG::FoldSymbolOffset(unsigned Opcode, EVT VT,
   case ISD::ADD: break;
   case ISD::SUB: Offset = -uint64_t(Offset); break;
   default: return SDValue();
   }
   return getGlobalAddress(GA->getGlobal(), SDLoc(Cst2), VT,
                           GA->getOffset() + uint64_t(Offset));
 }
 
+bool SelectionDAG::isUndef(unsigned Opcode, ArrayRef<SDValue> Ops) {
+  switch (Opcode) {
+  case ISD::SDIV:
+  case ISD::UDIV:
+  case ISD::SREM:
+  case ISD::UREM: {
+    // If a divisor is zero/undef or any element of a divisor vector is
+    // zero/undef, the whole op is undef.
+    assert(Ops.size() == 2 && "Div/rem should have 2 operands");
+    SDValue Divisor = Ops[1];
+    if (Divisor.isUndef() || isNullConstant(Divisor))
+      return true;
+
+    return ISD::isBuildVectorOfConstantSDNodes(Divisor.getNode()) &&
+           any_of(Divisor->op_values(),
+                  [](SDValue V) { return V.isUndef() || isNullConstant(V); });
+    // TODO: Handle signed overflow.
+  }
+  // TODO: Handle oversized shifts.
+  default:
+    return false;
+  }
+}
+
 SDValue SelectionDAG::FoldConstantArithmetic(unsigned Opcode, const SDLoc &DL,
                                              EVT VT, SDNode *Cst1,
                                              SDNode *Cst2) {
   // If the opcode is a target-specific ISD node, there's nothing we can
   // do here and the operand rules may not line up with the below, so
   // bail early.
   if (Opcode >= ISD::BUILTIN_OP_END)
     return SDValue();
 
+  if (isUndef(Opcode, {SDValue(Cst1, 0), SDValue(Cst2, 0)}))
+    return getUNDEF(VT);
+
   // Handle the case of two scalars.
   if (const ConstantSDNode *Scalar1 = dyn_cast<ConstantSDNode>(Cst1)) {
     if (const ConstantSDNode *Scalar2 = dyn_cast<ConstantSDNode>(Cst2)) {
       SDValue Folded = FoldConstantArithmetic(Opcode, DL, VT, Scalar1, Scalar2);
       assert((!Folded || !VT.isVector()) &&
              "Can't fold vectors ops with scalar operands");
       return Folded;
     }
SDValue SelectionDAG::FoldConstantVectorArithmetic(unsigned Opcode,
                                                    ArrayRef<SDValue> Ops,
                                                    const SDNodeFlags *Flags) {
   // If the opcode is a target-specific ISD node, there's nothing we can
   // do here and the operand rules may not line up with the below, so
   // bail early.
   if (Opcode >= ISD::BUILTIN_OP_END)
     return SDValue();
 
+  if (isUndef(Opcode, Ops))
+    return getUNDEF(VT);
+
   // We can only fold vectors - maybe merge with FoldConstantArithmetic someday?
   if (!VT.isVector())
     return SDValue();
 
   unsigned NumElts = VT.getVectorNumElements();
 
   auto IsScalarOrSameVectorSize = [&](const SDValue &Op) {
     return !Op.getValueType().isVector() ||
SDValue SelectionDAG::FoldConstantVectorArithmetic(unsigned Opcode,
   // or UNDEF scalars.
   if (!all_of(Ops, IsConstantBuildVectorOrUndef) ||
       !all_of(Ops, IsScalarOrSameVectorSize))
     return SDValue();
 
   // If we are comparing vectors, then the result needs to be a i1 boolean
   // that is then sign-extended back to the legal result type.
   EVT SVT = (Opcode == ISD::SETCC ? MVT::i1 : VT.getScalarType());
 
   // Find legal integer scalar type for constant promotion and
   // ensure that its scalar size is at least as large as source.
   EVT LegalSVT = VT.getScalarType();
   if (NewNodesMustHaveLegalTypes && LegalSVT.isInteger()) {
     LegalSVT = TLI->getTypeToTransformTo(*getContext(), LegalSVT);
     if (LegalSVT.bitsLT(VT.getScalarType()))
       return SDValue();
   }

test/CodeGen/X86/div-rem-simplify.ll

; CHECK-NEXT:    retq
   ret <4 x i32> %div
 }
 
 ; If any element of a constant divisor vector is zero, the whole op is undef.
 
 define <4 x i32> @sdiv0elt_vec(<4 x i32> %x) {
 ; CHECK-LABEL: sdiv0elt_vec:
 ; CHECK:       # BB#0:
-; CHECK-NEXT:    movaps {{.*#+}} xmm0 = <u,12,u,4294967292>
 ; CHECK-NEXT:    retq
   %zero = and <4 x i32> %x, <i32 0, i32 0, i32 0, i32 0>
   %some_ones = or <4 x i32> %zero, <i32 0, i32 -1, i32 0, i32 3>
   %div = sdiv <4 x i32> <i32 -11, i32 -12, i32 -13, i32 -14>, %some_ones
   ret <4 x i32> %div
 }
 
 define <4 x i32> @udiv0elt_vec(<4 x i32> %x) {
 ; CHECK-LABEL: udiv0elt_vec:
 ; CHECK:       # BB#0:
-; CHECK-NEXT:    movaps {{.*#+}} xmm0 = <u,4,3,u>
 ; CHECK-NEXT:    retq
   %div = udiv <4 x i32> <i32 11, i32 12, i32 13, i32 14>, <i32 0, i32 3, i32 4, i32 0>
   ret <4 x i32> %div
 }
 
 define <4 x i32> @urem0elt_vec(<4 x i32> %x) {
 ; CHECK-LABEL: urem0elt_vec:
 ; CHECK:       # BB#0:
-; CHECK-NEXT:    movaps {{.*#+}} xmm0 = <u,u,u,2>
 ; CHECK-NEXT:    retq
   %zero = and <4 x i32> %x, <i32 0, i32 0, i32 0, i32 0>
   %some_ones = or <4 x i32> %zero, <i32 0, i32 0, i32 0, i32 3>
   %rem = urem <4 x i32> <i32 11, i32 12, i32 13, i32 14>, %some_ones
   ret <4 x i32> %rem
 }
 
 define <4 x i32> @srem0elt_vec(<4 x i32> %x) {
 ; CHECK-LABEL: srem0elt_vec:
 ; CHECK:       # BB#0:
-; CHECK-NEXT:    movl $-2, %eax
-; CHECK-NEXT:    movd %eax, %xmm0
 ; CHECK-NEXT:    retq
   %rem = srem <4 x i32> <i32 -11, i32 -12, i32 -13, i32 -14>, <i32 -3, i32 -3, i32 0, i32 2>
   ret <4 x i32> %rem
 }
 

test/CodeGen/X86/vec_sdiv_to_shift.ll

 ; AVX2-NEXT:    vpaddw %ymm1, %ymm0, %ymm0
 ; AVX2-NEXT:    vpsraw $2, %ymm0, %ymm0
 ; AVX2-NEXT:    retq
 entry:
   %a0 = sdiv <16 x i16> %var, <i16 4, i16 4, i16 4, i16 4, i16 4, i16 4, i16 4, i16 4, i16 4, i16 4, i16 4, i16 4, i16 4, i16 4, i16 4, i16 4>
   ret <16 x i16> %a0
 }
 
-; TODO: The div-by-0 lanes are folded away, so we use scalar ops. Would it be better to keep this in the vector unit?
+; Div-by-0 in any lane is UB.
 
 define <4 x i32> @sdiv_non_splat(<4 x i32> %x) {
 ; SSE-LABEL: sdiv_non_splat:
 ; SSE:       # BB#0:
-; SSE-NEXT:    movd %xmm0, %eax
-; SSE-NEXT:    movl %eax, %ecx
-; SSE-NEXT:    shrl $31, %ecx
-; SSE-NEXT:    addl %eax, %ecx
-; SSE-NEXT:    sarl %ecx
-; SSE-NEXT:    movd %ecx, %xmm0
 ; SSE-NEXT:    retq
 ;
 ; AVX-LABEL: sdiv_non_splat:
 ; AVX:       # BB#0:
-; AVX-NEXT:    vmovd %xmm0, %eax
-; AVX-NEXT:    movl %eax, %ecx
-; AVX-NEXT:    shrl $31, %ecx
-; AVX-NEXT:    addl %eax, %ecx
-; AVX-NEXT:    sarl %ecx
-; AVX-NEXT:    vmovd %ecx, %xmm0
 ; AVX-NEXT:    retq
   %y = sdiv <4 x i32> %x, <i32 2, i32 0, i32 0, i32 0>
   ret <4 x i32> %y
 }