[LVI] Extend select handling to catch min/max/clamp idioms

[reames]

Summary

Most of this is fairly straight forward. I was a bit surprised to find we didn't have smin/umin on ConstantRange. Careful review of that would be appreciated.

Note that I'm only handling two constant ranges at this point. It would be reasonable to consider treating overdefined as a full range if the instruction is typed as an integer, but that should be a separate change.

[ab]

lib/IR/ConstantRange.cpp

718 ↗

(On Diff #47762)

smax -> smin? Same below

[reames]

lib/IR/ConstantRange.cpp

718 ↗

(On Diff #47762)

Yep, stale comment. Will fix, thanks.

[sanjoy]

lib/Analysis/LazyValueInfo.cpp

972 ↗

(On Diff #47762)

Nit: should be AddConstants.

979 ↗

(On Diff #47762)

This would be easier to follow if you summarize the transform you're doing / pattern you're matching, like

select (A == CIBase) _ (A + CIAdded)

982 ↗

(On Diff #47762)

How can you infer this without looking at the SI->getTrueValue()? Can't the select be select (A == CIBase) (CIBase + CIAdded) (A + CIAdded)? Or is there a check somewhere that I missed?

lib/IR/ConstantRange.cpp

717 ↗

(On Diff #47762)

These should get tests in unittests/IR/ConstantRange.cpp. I'd say split these out, add tests and land them -- the smin and umin implementations lgtm in their current form (with the edit @ab suggested).

[reames]

update patches, rebased after separating and landing constant range parts per Sanjoy's comment.

lib/Analysis/LazyValueInfo.cpp

972 ↗

(On Diff #47762)

Per style guide:
Function names should be verb phrases (as they represent actions), and command-like function should be imperative. The name should be camel case, and start with a lower case letter (e.g. openFile() or isFoo()).

While a lambda is technically a variable, it's really just a locally scoped function in this usage.

979 ↗

(On Diff #47762)

See updated patch. Is that better?

982 ↗

(On Diff #47762)

If A != CIBase, then the RHS can't contribute CIBase + CIAdded. The LHS can, but we're only constraining the input ranges here.

[sanjoy]

lgtm

lib/Analysis/LazyValueInfo.cpp

922 ↗

(On Diff #49242)

Is it possible for LHS to be the false value and RHS to be the true value?

981 ↗

(On Diff #49242)

Makes sense.

[reames]

lib/Analysis/LazyValueInfo.cpp

922 ↗

(On Diff #49242)

This is defensive programming in case matchSelectPattern get's smarter.

llvm/trunk/lib/Analysis/LazyValueInfo.cpp

bool LazyValueInfoCache::solveBlockValueSelect(LVILatticeVal &BBLV,
   LVILatticeVal FalseVal = getBlockValue(SI->getFalseValue(), BB);
   // If we hit overdefined, don't ask more queries.  We want to avoid poisoning
   // extra slots in the table if we can.
   if (FalseVal.isOverdefined()) {
     BBLV.markOverdefined();
     return true;
   }
 
+  if (TrueVal.isConstantRange() && FalseVal.isConstantRange()) {
+    ConstantRange TrueCR = TrueVal.getConstantRange();
+    ConstantRange FalseCR = FalseVal.getConstantRange();
+    Value *LHS = nullptr;
+    Value *RHS = nullptr;
+    SelectPatternResult SPR = matchSelectPattern(SI, LHS, RHS);
+    // Is this a min specifically of our two inputs?  (Avoid the risk of
+    // ValueTracking getting smarter looking back past our immediate inputs.)
+    if (SelectPatternResult::isMinOrMax(SPR.Flavor) &&
+        LHS == SI->getTrueValue() && RHS == SI->getFalseValue()) {
+      switch (SPR.Flavor) {
+      default:
+        llvm_unreachable("unexpected minmax type!");
+      case SPF_SMIN:                   /// Signed minimum
+        BBLV.markConstantRange(TrueCR.smin(FalseCR));
+        return true;
+      case SPF_UMIN:                   /// Unsigned minimum
+        BBLV.markConstantRange(TrueCR.umin(FalseCR));
+        return true;
+      case SPF_SMAX:                   /// Signed maximum
+        BBLV.markConstantRange(TrueCR.smax(FalseCR));
+        return true;
+      case SPF_UMAX:                   /// Unsigned maximum        
+        BBLV.markConstantRange(TrueCR.umax(FalseCR));
+        return true;
+      };
+    }
+    
+    // TODO: ABS, NABS from the SelectPatternResult
+  }
+
   // Can we constrain the facts about the true and false values by using the
   // condition itself?  This shows up with idioms like e.g. select(a > 5, a, 5).
   // TODO: We could potentially refine an overdefined true value above.
   if (auto *ICI = dyn_cast<ICmpInst>(SI->getCondition())) {
     LVILatticeVal TrueValTaken, FalseValTaken;
     if (!getValueFromFromCondition(SI->getTrueValue(), ICI,
                                    TrueValTaken, true))
       TrueValTaken.markOverdefined();
     if (!getValueFromFromCondition(SI->getFalseValue(), ICI,
                                    FalseValTaken, false))
       FalseValTaken.markOverdefined();
 
     TrueVal = intersect(TrueVal, TrueValTaken);
     FalseVal = intersect(FalseVal, FalseValTaken);
-  }
 
-  // TODO: handle idioms like min & max where we can use a more precise merge
-  // when our inputs are constant ranges.
+
+    // Handle clamp idioms such as:
+    //   %24 = constantrange<0, 17>
+    //   %39 = icmp eq i32 %24, 0
+    //   %40 = add i32 %24, -1
+    //   %siv.next = select i1 %39, i32 16, i32 %40
+    //   %siv.next = constantrange<0, 17> not <-1, 17>
+    // In general, this can handle any clamp idiom which tests the edge
+    // condition via an equality or inequality.  
+    ICmpInst::Predicate Pred = ICI->getPredicate();
+    Value *A = ICI->getOperand(0);
+    if (ConstantInt *CIBase = dyn_cast<ConstantInt>(ICI->getOperand(1))) {
+      auto addConstants = [](ConstantInt *A, ConstantInt *B) {
+        assert(A->getType() == B->getType());
+        return ConstantInt::get(A->getType(), A->getValue() + B->getValue());
+      };
+      // See if either input is A + C2, subject to the constraint from the
+      // condition that A != C when that input is used.  We can assume that
+      // that input doesn't include C + C2.
+      ConstantInt *CIAdded;
+      switch (Pred) {
+      case ICmpInst::ICMP_EQ:
+        if (match(SI->getFalseValue(), m_Add(m_Specific(A),
+                                             m_ConstantInt(CIAdded)))) {
+          auto ResNot = addConstants(CIBase, CIAdded);
+          FalseVal = intersect(FalseVal,
+                               LVILatticeVal::getNot(ResNot));
+        }
+        break;
+      case ICmpInst::ICMP_NE:
+        if (match(SI->getTrueValue(), m_Add(m_Specific(A),
+                                            m_ConstantInt(CIAdded)))) {
+          auto ResNot = addConstants(CIBase, CIAdded);
+          TrueVal = intersect(TrueVal,
+                              LVILatticeVal::getNot(ResNot));
+        }
+        break;
+      };
+    }
+  }
   
   LVILatticeVal Result;  // Start Undefined.
   Result.mergeIn(TrueVal, DL);
   Result.mergeIn(FalseVal, DL);
   BBLV = Result;
   return true;
 }
 

llvm/trunk/test/Transforms/CorrelatedValuePropagation/basic.ll

 ; CHECK: br label %next
 
 out:
   ret void
 
 next:
   ret void
 }
+
+define i1 @umin(i32 %a, i32 %b) {
+; CHECK-LABEL: @umin(
+entry:
+  %cmp = icmp ult i32 %a, 5
+  br i1 %cmp, label %a_guard, label %out
+
+a_guard:
+  %cmp2 = icmp ult i32 %b, 20
+  br i1 %cmp2, label %b_guard, label %out
+
+b_guard:
+  %sel_cmp = icmp ult i32 %a, %b
+  %min = select i1 %sel_cmp, i32 %a, i32 %b
+  %res = icmp eq i32 %min, 7
+  br label %next
+next:
+; CHECK: next:
+; CHECK: ret i1 false
+  ret i1 %res
+out:
+  ret i1 false
+}
+
+define i1 @smin(i32 %a, i32 %b) {
+; CHECK-LABEL: @smin(
+entry:
+  %cmp = icmp ult i32 %a, 5
+  br i1 %cmp, label %a_guard, label %out
+
+a_guard:
+  %cmp2 = icmp ult i32 %b, 20
+  br i1 %cmp2, label %b_guard, label %out
+
+b_guard:
+  %sel_cmp = icmp sle i32 %a, %b
+  %min = select i1 %sel_cmp, i32 %a, i32 %b
+  %res = icmp eq i32 %min, 7
+  br label %next
+next:
+; CHECK: next:
+; CHECK: ret i1 false
+  ret i1 %res
+out:
+  ret i1 false
+}
+
+define i1 @smax(i32 %a, i32 %b) {
+; CHECK-LABEL: @smax(
+entry:
+  %cmp = icmp sgt i32 %a, 5
+  br i1 %cmp, label %a_guard, label %out
+
+a_guard:
+  %cmp2 = icmp sgt i32 %b, 20
+  br i1 %cmp2, label %b_guard, label %out
+
+b_guard:
+  %sel_cmp = icmp sge i32 %a, %b
+  %max = select i1 %sel_cmp, i32 %a, i32 %b
+  %res = icmp eq i32 %max, 7
+  br label %next
+next:
+; CHECK: next:
+; CHECK: ret i1 false
+  ret i1 %res
+out:
+  ret i1 false
+}
+
+define i1 @umax(i32 %a, i32 %b) {
+; CHECK-LABEL: @umax(
+entry:
+  %cmp = icmp sgt i32 %a, 5
+  br i1 %cmp, label %a_guard, label %out
+
+a_guard:
+  %cmp2 = icmp sgt i32 %b, 20
+  br i1 %cmp2, label %b_guard, label %out
+
+b_guard:
+  %sel_cmp = icmp uge i32 %a, %b
+  %max = select i1 %sel_cmp, i32 %a, i32 %b
+  %res = icmp eq i32 %max, 7
+  br label %next
+next:
+; CHECK: next:
+; CHECK: ret i1 false
+  ret i1 %res
+out:
+  ret i1 false
+}
+
+define i1 @clamp_low1(i32 %a) {
+; CHECK-LABEL: @clamp_low1(
+entry:
+  %cmp = icmp sge i32 %a, 5
+  br i1 %cmp, label %a_guard, label %out
+
+a_guard:
+  %sel_cmp = icmp eq i32 %a, 5
+  %add = add i32 %a, -1
+  %sel = select i1 %sel_cmp, i32 5, i32 %a
+  %res = icmp eq i32 %sel, 4
+  br label %next
+next:
+; CHECK: next:
+; CHECK: ret i1 false
+  ret i1 %res
+out:
+  ret i1 false
+}
+
+define i1 @clamp_low2(i32 %a) {
+; CHECK-LABEL: @clamp_low2(
+entry:
+  %cmp = icmp sge i32 %a, 5
+  br i1 %cmp, label %a_guard, label %out
+
+a_guard:
+  %sel_cmp = icmp ne i32 %a, 5
+  %add = add i32 %a, -1
+  %sel = select i1 %sel_cmp, i32 %a, i32 5
+  %res = icmp eq i32 %sel, 4
+  br label %next
+next:
+; CHECK: next:
+; CHECK: ret i1 false
+  ret i1 %res
+out:
+  ret i1 false
+}
+
+define i1 @clamp_high1(i32 %a) {
+; CHECK-LABEL: @clamp_high1(
+entry:
+  %cmp = icmp sle i32 %a, 5
+  br i1 %cmp, label %a_guard, label %out
+
+a_guard:
+  %sel_cmp = icmp eq i32 %a, 5
+  %add = add i32 %a, 1
+  %sel = select i1 %sel_cmp, i32 5, i32 %a
+  %res = icmp eq i32 %sel, 6
+  br label %next
+next:
+; CHECK: next:
+; CHECK: ret i1 false
+  ret i1 %res
+out:
+  ret i1 false
+}
+
+define i1 @clamp_high2(i32 %a) {
+; CHECK-LABEL: @clamp_high2(
+entry:
+  %cmp = icmp sle i32 %a, 5
+  br i1 %cmp, label %a_guard, label %out
+
+a_guard:
+  %sel_cmp = icmp ne i32 %a, 5
+  %add = add i32 %a, 1
+  %sel = select i1 %sel_cmp, i32 %a, i32 5
+  %res = icmp eq i32 %sel, 6
+  br label %next
+next:
+; CHECK: next:
+; CHECK: ret i1 false
+  ret i1 %res
+out:
+  ret i1 false
+}
+
+; Just showing arbitrary constants work, not really a clamp
+define i1 @clamp_high3(i32 %a) {
+; CHECK-LABEL: @clamp_high3(
+entry:
+  %cmp = icmp sle i32 %a, 5
+  br i1 %cmp, label %a_guard, label %out
+
+a_guard:
+  %sel_cmp = icmp ne i32 %a, 5
+  %add = add i32 %a, 100
+  %sel = select i1 %sel_cmp, i32 %a, i32 5
+  %res = icmp eq i32 %sel, 105
+  br label %next
+next:
+; CHECK: next:
+; CHECK: ret i1 false
+  ret i1 %res
+out:
+  ret i1 false
+}