[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

3802

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 ↗

(On Diff #91560)

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 ↗

(On Diff #91560)

Yep - that's better. Thanks!

lib/CodeGen/SelectionDAG/DAGCombiner.cpp

if (N1.isUndef())
     return N1;
 
   // X / 0 --> undef
   // X % 0 --> undef
   // We don't need to preserve faults!
   if (isNullConstantOrNullSplatConstant(N1))
     return DAG.getUNDEF(VT);
 
+  // If any element of a constant divisor vector is zero, the whole op is undef.
+  if (ISD::isBuildVectorOfConstantSDNodes(N1.getNode()) &&
+      any_of(N1->op_values(), isNullConstant))
+    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,
   // fold them individually.
   BuildVectorSDNode *BV1 = dyn_cast<BuildVectorSDNode>(Cst1);
   BuildVectorSDNode *BV2 = dyn_cast<BuildVectorSDNode>(Cst2);
   if (!BV1 || !BV2)
     return SDValue();
 
   assert(BV1->getNumOperands() == BV2->getNumOperands() && "Out of sync!");
 
+  // If any element of a constant divisor vector is zero, the whole op is undef.
+  if ((Opcode == ISD::SDIV || Opcode == ISD::UDIV ||
+       Opcode == ISD::SREM || Opcode == ISD::UREM) &&
+      any_of(BV2->op_values(), isNullConstant))
+    return getUNDEF(VT);
+
   EVT SVT = VT.getScalarType();
   SmallVector<SDValue, 4> Outputs;
   for (unsigned I = 0, E = BV1->getNumOperands(); I != E; ++I) {
     SDValue V1 = BV1->getOperand(I);
     SDValue V2 = BV2->getOperand(I);
 
     // Avoid BUILD_VECTOR nodes that perform implicit truncation.
     // FIXME: This is valid and could be handled by truncation.
SDValue SelectionDAG::FoldConstantVectorArithmetic(unsigned Opcode,
 
   // All operands must be vector types with the same number of elements as
   // the result type and must be either UNDEF or a build vector of constant
   // or UNDEF scalars.
   if (!all_of(Ops, IsConstantBuildVectorOrUndef) ||
       !all_of(Ops, IsScalarOrSameVectorSize))
     return SDValue();
 
+  // If any element of a constant divisor vector is zero, the whole op is undef.
+  if ((Opcode == ISD::SDIV || Opcode == ISD::UDIV ||
+       Opcode == ISD::SREM || Opcode == ISD::UREM) &&
+      any_of(Ops[1]->op_values(), isNullConstant))
+    return getUNDEF(VT);
+
   // 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()) {

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
 }