[CVP] Don't introduce zero after truncating when narrowUDivOrSDiv

[StephenFan]

Summary

The constant range computed at use point does not mean the input
constant range or in other words, the constant range computed at define point,
is the same range as computed at use point. For intructions that is
always speculative, it is safe. However, for instruction UDiv or SDiv,
and we want to narrow the data type of these, zero may be introduced
after truncating the divisor and causing an UB.

So this patch checks that if the constant range at define point does not
contain zero, then it also must not contain zero after truncating.

[StephenFan]

Delete blank line

[goldstein.w.n]

llvm/lib/Transforms/Scalar/CorrelatedValuePropagation.cpp

813

narrowSDivOrSRem i think

llvm/lib/Transforms/Scalar/CorrelatedValuePropagation.cpp

static bool expandUDivOrURem(BinaryOperator *Instr, const ConstantRange &XCR,
   Instr->replaceAllUsesWith(ExpandedOp);
   Instr->eraseFromParent();
   ++NumUDivURemsNarrowedExpanded;
   return true;
 }
 
 /// Try to shrink a udiv/urem's width down to the smallest power of two that's
 /// sufficient to contain its operands.
-static bool narrowUDivOrURem(BinaryOperator *Instr, const ConstantRange &XCR,
-                             const ConstantRange &YCR) {
+static bool narrowUDivOrURem(BinaryOperator *Instr,
+                             const ConstantRange &XCRAtUse,
+                             const ConstantRange &YCRAtUse,
+                             const ConstantRange &YCRAtDef) {
   assert(Instr->getOpcode() == Instruction::UDiv ||
          Instr->getOpcode() == Instruction::URem);
   assert(!Instr->getType()->isVectorTy());
 
   // Find the smallest power of two bitwidth that's sufficient to hold Instr's
   // operands.
 
   // What is the smallest bit width that can accommodate the entire value ranges
   // of both of the operands?
-  unsigned MaxActiveBits = std::max(XCR.getActiveBits(), YCR.getActiveBits());
+  unsigned MaxActiveBits =
+      std::max(XCRAtUse.getActiveBits(), YCRAtUse.getActiveBits());
   // Don't shrink below 8 bits wide.
   unsigned NewWidth = std::max<unsigned>(PowerOf2Ceil(MaxActiveBits), 8);
 
   // NewWidth might be greater than OrigWidth if OrigWidth is not a power of
   // two.
   if (NewWidth >= Instr->getType()->getIntegerBitWidth())
     return false;
 
+  // For divisor, if the constant range computed at define point does not
+  // contain zero, then we should ensure the constant range after truncating
+  // also does not contain zero. Otherwise we may introduce an UB.
+  if (!YCRAtDef.contains(APInt::getZero(YCRAtDef.getBitWidth())) &&
+      YCRAtDef.truncate(NewWidth).contains(APInt::getZero(NewWidth)))
+    return false;
+
   ++NumUDivURemsNarrowed;
   IRBuilder<> B{Instr};
   auto *TruncTy = Type::getIntNTy(Instr->getContext(), NewWidth);
   auto *LHS = B.CreateTruncOrBitCast(Instr->getOperand(0), TruncTy,
                                      Instr->getName() + ".lhs.trunc");
   auto *RHS = B.CreateTruncOrBitCast(Instr->getOperand(1), TruncTy,
                                      Instr->getName() + ".rhs.trunc");
   auto *BO = B.CreateBinOp(Instr->getOpcode(), LHS, RHS, Instr->getName());
   auto *Zext = B.CreateZExt(BO, Instr->getType(), Instr->getName() + ".zext");
   if (auto *BinOp = dyn_cast<BinaryOperator>(BO))
     if (BinOp->getOpcode() == Instruction::UDiv)
       BinOp->setIsExact(Instr->isExact());
 
   Instr->replaceAllUsesWith(Zext);
   Instr->eraseFromParent();
   return true;
 }
 
 static bool processUDivOrURem(BinaryOperator *Instr, LazyValueInfo *LVI) {
   assert(Instr->getOpcode() == Instruction::UDiv ||
          Instr->getOpcode() == Instruction::URem);
   if (Instr->getType()->isVectorTy())
     return false;
 
-  ConstantRange XCR = LVI->getConstantRangeAtUse(Instr->getOperandUse(0));
-  ConstantRange YCR = LVI->getConstantRangeAtUse(Instr->getOperandUse(1));
-  if (expandUDivOrURem(Instr, XCR, YCR))
+  ConstantRange YCRAtDef = LVI->getConstantRange(Instr->getOperand(1), Instr);
+  ConstantRange XCRAtUse = LVI->getConstantRangeAtUse(Instr->getOperandUse(0));
+  ConstantRange YCRAtUse = LVI->getConstantRangeAtUse(Instr->getOperandUse(1));
+  if (expandUDivOrURem(Instr, XCRAtUse, YCRAtUse))
     return true;
 
-  return narrowUDivOrURem(Instr, XCR, YCR);
+  return narrowUDivOrURem(Instr, XCRAtUse, YCRAtUse, YCRAtDef);
 }
 
 static bool processSRem(BinaryOperator *SDI, const ConstantRange &LCR,
                         const ConstantRange &RCR, LazyValueInfo *LVI) {
   assert(SDI->getOpcode() == Instruction::SRem);
   assert(!SDI->getType()->isVectorTy());
 
   if (LCR.abs().icmp(CmpInst::ICMP_ULT, RCR.abs())) {

llvm/test/Transforms/CorrelatedValuePropagation/udiv.ll

 bb:
 ; CHECK: udiv i8
   %div = sdiv i32 %n, 100
   br label %exit
 
 exit:
   ret void
 }
+

llvm/test/Transforms/CorrelatedValuePropagation/urem.ll

else:
   call void @use(i1 %c3.2)
   %c4.2 = icmp uge i8 %x, 1
   call void @use(i1 %c4.2)
   %c5.2 = icmp ugt i8 %x, 1
   call void @use(i1 %c5.2)
   ret void
 }
 
+define i32 @udiv_do_not_truncate(i32 noundef %call, i32 %v) {
+; CHECK-LABEL: @udiv_do_not_truncate(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[CALL_NON_NULL:%.*]] = or i32 [[CALL:%.*]], 1
+; CHECK-NEXT:    [[DIV:%.*]] = urem i32 8192, [[CALL_NON_NULL]]
+; CHECK-NEXT:    [[CMP:%.*]] = icmp ugt i32 [[CALL_NON_NULL]], 8192
+; CHECK-NEXT:    [[SELECT:%.*]] = select i1 [[CMP]], i32 1, i32 [[DIV]]
+; CHECK-NEXT:    ret i32 [[SELECT]]
+;
+entry:
+  %call_non_null = or i32 %call, 1
+  %div = urem i32 8192, %call_non_null
+  %cmp = icmp ugt i32 %call_non_null, 8192
+  %select = select i1 %cmp, i32 1, i32 %div
+  ret i32 %select
+}
+
 declare void @use(i1)