[LVI] Thread binop over select with constant arms

[dtcxzyw]

Motivating case from https://github.com/delta-io/delta-rs: https://alive2.llvm.org/ce/z/3mzr4C

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

@llvm/pr-subscribers-llvm-analysis

Author: Yingwei Zheng (dtcxzyw)

Changes

Motivating case from https://github.com/delta-io/delta-rs: https://alive2.llvm.org/ce/z/3mzr4C

---

Full diff: https://github.com/llvm/llvm-project/pull/110212.diff

2 Files Affected:

• (modified) llvm/lib/Analysis/LazyValueInfo.cpp (+58-2)
• (modified) llvm/test/Transforms/CorrelatedValuePropagation/cond-at-use.ll (+37)

diff --git a/llvm/lib/Analysis/LazyValueInfo.cpp b/llvm/lib/Analysis/LazyValueInfo.cpp
index c9849b86b664c8..3246bdded4f10c 100644
--- a/llvm/lib/Analysis/LazyValueInfo.cpp
+++ b/llvm/lib/Analysis/LazyValueInfo.cpp
@@ -924,18 +924,74 @@ LazyValueInfoImpl::solveBlockValueBinaryOpImpl(
     Instruction *I, BasicBlock *BB,
     std::function<ConstantRange(const ConstantRange &, const ConstantRange &)>
         OpFn) {
+  Value *LHS = I->getOperand(0);
+  Value *RHS = I->getOperand(1);
+
+  auto GetValueFromCondition =
+      [&](Value *V, Value *Cond,
+          bool CondIsTrue) -> std::optional<ConstantRange> {
+    std::optional<ValueLatticeElement> OptVal = getBlockValue(V, BB, I);
+    if (!OptVal)
+      return std::nullopt;
+    return OptVal->asConstantRange(V->getType());
+  };
+
+  auto ThreadBinOpOverSelect =
+      [&](Value *X, const ConstantRange &CRX, SelectInst *Y,
+          bool XIsLHS) -> std::optional<ValueLatticeElement> {
+    Value *Cond = Y->getCondition();
+    // Only handle selects with constant values.
+    Constant *TrueC = dyn_cast<Constant>(Y->getTrueValue());
+    if (!TrueC)
+      return std::nullopt;
+    Constant *FalseC = dyn_cast<Constant>(Y->getFalseValue());
+    if (!FalseC)
+      return std::nullopt;
+    if (!isGuaranteedNotToBeUndef(Cond, AC))
+      return std::nullopt;
+
+    ConstantRange TrueX =
+        CRX.intersectWith(getValueFromCondition(X, Cond, /*CondIsTrue=*/true,
+                                                /*UseBlockValue=*/false)
+                              ->asConstantRange(X->getType()));
+    ConstantRange FalseX =
+        CRX.intersectWith(getValueFromCondition(X, Cond, /*CondIsTrue=*/false,
+                                                /*UseBlockValue=*/false)
+                              ->asConstantRange(X->getType()));
+    ConstantRange TrueY = TrueC->toConstantRange();
+    ConstantRange FalseY = FalseC->toConstantRange();
+
+    if (XIsLHS)
+      return ValueLatticeElement::getRange(
+          OpFn(TrueX, TrueY).unionWith(OpFn(FalseX, FalseY)));
+    return ValueLatticeElement::getRange(
+        OpFn(TrueY, TrueX).unionWith(OpFn(FalseY, FalseX)));
+  };
+
   // Figure out the ranges of the operands.  If that fails, use a
   // conservative range, but apply the transfer rule anyways.  This
   // lets us pick up facts from expressions like "and i32 (call i32
   // @foo()), 32"
-  std::optional<ConstantRange> LHSRes = getRangeFor(I->getOperand(0), I, BB);
+  std::optional<ConstantRange> LHSRes = getRangeFor(LHS, I, BB);
   if (!LHSRes)
     return std::nullopt;
 
-  std::optional<ConstantRange> RHSRes = getRangeFor(I->getOperand(1), I, BB);
+  // Try to thread binop over rhs select
+  if (auto *SI = dyn_cast<SelectInst>(RHS)) {
+    if (auto Res = ThreadBinOpOverSelect(LHS, *LHSRes, SI, /*XIsLHS=*/true))
+      return *Res;
+  }
+
+  std::optional<ConstantRange> RHSRes = getRangeFor(RHS, I, BB);
   if (!RHSRes)
     return std::nullopt;
 
+  // Try to thread binop over lhs select
+  if (auto *SI = dyn_cast<SelectInst>(LHS)) {
+    if (auto Res = ThreadBinOpOverSelect(RHS, *RHSRes, SI, /*XIsLHS=*/false))
+      return *Res;
+  }
+
   const ConstantRange &LHSRange = *LHSRes;
   const ConstantRange &RHSRange = *RHSRes;
   return ValueLatticeElement::getRange(OpFn(LHSRange, RHSRange));
diff --git a/llvm/test/Transforms/CorrelatedValuePropagation/cond-at-use.ll b/llvm/test/Transforms/CorrelatedValuePropagation/cond-at-use.ll
index 3af4c70a5621c8..66787c76b41a68 100644
--- a/llvm/test/Transforms/CorrelatedValuePropagation/cond-at-use.ll
+++ b/llvm/test/Transforms/CorrelatedValuePropagation/cond-at-use.ll
@@ -630,3 +630,40 @@ define i64 @test_shl_nsw_at_use(i64 noundef %x) {
   %res = select i1 %cmp, i64 %shr, i64 0
   ret i64 %res
 }
+
+define i1 @test_icmp_mod(i64 noundef %x) {
+; CHECK-LABEL: @test_icmp_mod(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[REM:%.*]] = srem i64 [[X:%.*]], 86400
+; CHECK-NEXT:    [[CMP:%.*]] = icmp slt i64 [[REM]], 0
+; CHECK-NEXT:    [[COND:%.*]] = select i1 [[CMP]], i64 86400, i64 0
+; CHECK-NEXT:    [[ADD:%.*]] = add nsw i64 [[COND]], [[REM]]
+; CHECK-NEXT:    ret i1 false
+;
+entry:
+  %rem = srem i64 %x, 86400
+  %cmp = icmp slt i64 %rem, 0
+  %cond = select i1 %cmp, i64 86400, i64 0
+  %add = add nsw i64 %cond, %rem
+  %cmp1 = icmp ugt i64 %add, 86399
+  ret i1 %cmp1
+}
+
+define i1 @test_icmp_mod_undef(i64 %x) {
+; CHECK-LABEL: @test_icmp_mod_undef(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[REM:%.*]] = srem i64 [[X:%.*]], 86400
+; CHECK-NEXT:    [[CMP:%.*]] = icmp slt i64 [[REM]], 0
+; CHECK-NEXT:    [[COND:%.*]] = select i1 [[CMP]], i64 86400, i64 0
+; CHECK-NEXT:    [[ADD:%.*]] = add nsw i64 [[COND]], [[REM]]
+; CHECK-NEXT:    [[CMP1:%.*]] = icmp ugt i64 [[ADD]], 86399
+; CHECK-NEXT:    ret i1 [[CMP1]]
+;
+entry:
+  %rem = srem i64 %x, 86400
+  %cmp = icmp slt i64 %rem, 0
+  %cond = select i1 %cmp, i64 86400, i64 0
+  %add = add nsw i64 %cond, %rem
+  %cmp1 = icmp ugt i64 %add, 86399
+  ret i1 %cmp1
+}

[nikic]

LGTM