[LVI] Constant-propagate a zero extension of the switch condition value through case edges

[hjyamauchi]

Summary

(This is a second attempt as https://reviews.llvm.org/D34822 was reverted.)

LazyValueInfo currently computes the constant value of the switch condition through case edges, which allows the constant value to be propagated through the case edges.

But we have seen a case where a zero-extended value of the switch condition is used past case edges for which the constant propagation doesn't occur.

This patch adds a small logic to handle such a case in getEdgeValueLocal().

This is motivated by the Python 2.7 eval loop in PyEval_EvalFrameEx() where the lack of the constant propagation causes longer live ranges and more spill code than necessary.

With this patch, we see that the code size of PyEval_EvalFrameEx() decreases by ~5.4% and a performance test improves by ~4.6%.

[hjyamauchi]

The first diff is the same as the reverted https://reviews.llvm.org/D34822.

[hjyamauchi]

This second diff is the new one that fixes the bug (see test20).

[hjyamauchi]

The issue with D34822 was that it is not correct to propagate the constant range of the value that's derived from the switch condition to the *default* edge.

For example,

%12 = and i64 %11, 1
switch i64 %11, label %default [

i64 0, label %86
i64 3, label %87

]

...

%default:

%92 = icmp eq i64 %12, 0, !dbg !1921
br i1 %92, label %94, label %93, !dbg !1921

It's correct to propagate that %11 != 0 and %11 != 3 to the %default edge.

But it's not correct to propagate the corresponding value of %12, that is %12 != 0 (0 & 1) and %12 != 1 (3 & 1), on the %default edge (which isn't even logically possible.)

The fix is to restrict the default-case propagation to the "Val == Condition" case only.

[hjyamauchi]

This is ready for review.

Note the newly added @test20 demonstrates the bug in D34822.

[sanjoy]

Code wise this looks correct to me, but I'd prefer to phrase the solution a bit differently in the commit message and in the comments: IIUC the bug we had was that we were assuming Condition != T (in the default block) implied f(Condition) != f(T). However this isn't true for a non injective f (like f(i32 x) = and i32 x, 2).

For instance we could still do this transformation for

  x = cond + 2
  switch cond, label %left, [ i32 20, label %right]

left:
  ...
right:
 ...

In left we know that cond is not 20 and thus x is not 22, since f(x) = x + 2 is injective.

(I'm not suggesting we actually do the above btw, I think the case this patch is handling (f(x) = x) is enough for now.)

lib/Analysis/LazyValueInfo.cpp

158

Can you return None; here?

test/Transforms/CorrelatedValuePropagation/range.ll

587

Can we demonstrate the issue with just one switch?

[hjyamauchi]

Agreed on the phrasing and that there are cases where we could do the propagation/transformation.

I fixed the code comment there.

[hjyamauchi]

Addressed the inline comments.

[sanjoy]

lgtm

lib/Analysis/LazyValueInfo.cpp

1508–1513

I would be somewhat more explicit here:

We know Condition != EdgeVal in BBTo.  In some cases we can use this to infer Val == f(Condition) is != f(EdgeVal).  For now, we only do this when f is identity (i.e. Val == Condition), but we should be able to do this for any injective f.

The "There is no need to perform difference for those cases." comment is also somewhat misleading -- it is not just that there is no need, we cannot perform the difference in these cases. Can you please update that comment since you're touching nearby code anyway?

[hjyamauchi]

Addressed the comment.

[hjyamauchi]

lib/Analysis/LazyValueInfo.cpp

1508–1513

Done. I assumed by EdgeVal you meant CaseValue where EdgeVal = f(CaseValue).

[sanjoy]

lib/Analysis/LazyValueInfo.cpp

1508–1513

Yes, thanks!

lib/Analysis/LazyValueInfo.cpp

 //
 //===----------------------------------------------------------------------===//
 
 #include "llvm/Analysis/LazyValueInfo.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/Analysis/AssumptionCache.h"
 #include "llvm/Analysis/ConstantFolding.h"
+#include "llvm/Analysis/InstructionSimplify.h"
 #include "llvm/Analysis/TargetLibraryInfo.h"
 #include "llvm/Analysis/ValueTracking.h"
 #include "llvm/IR/AssemblyAnnotationWriter.h"
 #include "llvm/IR/CFG.h"
 #include "llvm/IR/ConstantRange.h"
 #include "llvm/IR/Constants.h"
 #include "llvm/IR/DataLayout.h"
 #include "llvm/IR/Dominators.h"
public:
   }
 
   const ConstantRange &getConstantRange() const {
     assert(isConstantRange() &&
            "Cannot get the constant-range of a non-constant-range!");
     return Range;
   }
 
+  Optional<APInt> asConstantInteger() const {
+    if (isConstant() && isa<ConstantInt>(Val)) {
+      return Val->getUniqueInteger();
+    } else if (isConstantRange() && Range.isSingleElement()) {
+      return *Range.getSingleElement();
+    }
+    return None;
+  }
+
 private:
   void markOverdefined() {
     if (isOverdefined())
       return;
     Tag = overdefined;
   }
 
   void markConstant(Constant *V) {

 }
 
 LVILatticeVal getValueFromCondition(Value *Val, Value *Cond, bool isTrueDest) {
   assert(Cond && "precondition");
   DenseMap<Value*, LVILatticeVal> Visited;
   return getValueFromCondition(Val, Cond, isTrueDest, Visited);
 }
 
+// Return true if Usr has Op as an operand, otherwise false.
+static bool usesOperand(User *Usr, Value *Op) {
+  return find(Usr->operands(), Op) != Usr->op_end();
+}
+
+// Check if Val can be simplified to an integer constant when the value of one
+// of its operands Op is an integer constant OpConstVal. If so, return it as an
+// lattice value range with a single element or otherwise return an overdefined
+// lattice value.
+static LVILatticeVal constantFoldUser(Value *Val, Value *Op,
+                                      const APInt &OpConstVal,
+                                      const DataLayout &DL) {
+  Constant* OpConst = Constant::getIntegerValue(Op->getType(), OpConstVal);
+  // Check if Val can be simplified to a constant.
+  if (auto *CI = dyn_cast<CastInst>(Val)) {
+    assert(CI->getOperand(0) == Op && "Operand 0 isn't Op");
+    if (auto *C = dyn_cast_or_null<ConstantInt>(
+            SimplifyCastInst(CI->getOpcode(), OpConst,
+                             CI->getDestTy(), DL))) {
+      return LVILatticeVal::getRange(ConstantRange(C->getUniqueInteger()));
+    }
+  } else if (auto *BO = dyn_cast<BinaryOperator>(Val)) {
+    bool Op0Match = BO->getOperand(0) == Op;
+    bool Op1Match = BO->getOperand(1) == Op;
+    assert((Op0Match || Op1Match) &&
+           "Operand 0 nor Operand 1 isn't a match");
+    Value *LHS = Op0Match ? OpConst : BO->getOperand(0);
+    Value *RHS = Op1Match ? OpConst : BO->getOperand(1);
+    if (auto *C = dyn_cast_or_null<ConstantInt>(
+            SimplifyBinOp(BO->getOpcode(), LHS, RHS, DL))) {
+      return LVILatticeVal::getRange(ConstantRange(C->getUniqueInteger()));
+    }
+  }
+  return LVILatticeVal::getOverdefined();
+}
+
 /// \brief Compute the value of Val on the edge BBFrom -> BBTo. Returns false if
 /// Val is not constrained on the edge.  Result is unspecified if return value
 /// is false.
 static bool getEdgeValueLocal(Value *Val, BasicBlock *BBFrom,
                               BasicBlock *BBTo, LVILatticeVal &Result) {
   // TODO: Handle more complex conditionals. If (v == 0 || v2 < 1) is false, we
   // know that v != 0.
   if (BranchInst *BI = dyn_cast<BranchInst>(BBFrom->getTerminator())) {
     // If this is a conditional branch and only one successor goes to BBTo, then
     // we may be able to infer something from the condition.
     if (BI->isConditional() &&
         BI->getSuccessor(0) != BI->getSuccessor(1)) {
       bool isTrueDest = BI->getSuccessor(0) == BBTo;
       assert(BI->getSuccessor(!isTrueDest) == BBTo &&
              "BBTo isn't a successor of BBFrom");
+      Value *Condition = BI->getCondition();
 
       // If V is the condition of the branch itself, then we know exactly what
       // it is.
-      if (BI->getCondition() == Val) {
+      if (Condition == Val) {
         Result = LVILatticeVal::get(ConstantInt::get(
                               Type::getInt1Ty(Val->getContext()), isTrueDest));
         return true;
       }
 
       // If the condition of the branch is an equality comparison, we may be
       // able to infer the value.
-      Result = getValueFromCondition(Val, BI->getCondition(), isTrueDest);
+      Result = getValueFromCondition(Val, Condition, isTrueDest);
+      if (!Result.isOverdefined())
+        return true;
+
+      if (User *Usr = dyn_cast<User>(Val)) {
+        assert(Result.isOverdefined() && "Result isn't overdefined");
+        if (isa<IntegerType>(Val->getType())) {
+          const DataLayout &DL = BBTo->getModule()->getDataLayout();
+          if (usesOperand(Usr, Condition)) {
+            // If Val has Condition as an operand and Val can be folded into a
+            // constant with either Condition == true or Condition == false,
+            // propagate the constant.
+            // eg.
+            //   ; %Val is true on the edge to %then.
+            //   %Val = and i1 %Condition, true.
+            //   br %Condition, label %then, label %else
+            APInt ConditionVal(1, isTrueDest ? 1 : 0);
+            Result = constantFoldUser(Val, Condition, ConditionVal, DL);
+          } else {
+            // If one of Val's operand has an inferred value, we may be able to
+            // infer the value of Val.
+            // eg.
+            //    ; %Val is 94 on the edge to %then.
+            //    %Val = add i8 %Op, 1
+            //    %Condition = icmp eq i8 %Op, 93
+            //    br i1 %Condition, label %then, label %else
+            for (unsigned i = 0; i < Usr->getNumOperands(); ++i) {
+              Value *Op = Usr->getOperand(i);
+              LVILatticeVal OpLatticeVal =
+                  getValueFromCondition(Op, Condition, isTrueDest);
+              if (Optional<APInt> OpConst = OpLatticeVal.asConstantInteger()) {
+                Result = constantFoldUser(Val, Op, OpConst.getValue(), DL);
+                break;
+              }
+            }
+          }
+        }
+      }
       if (!Result.isOverdefined())
         return true;
     }
   }
 
   // If the edge was formed by a switch on the value, then we may know exactly
   // what it is.
   if (SwitchInst *SI = dyn_cast<SwitchInst>(BBFrom->getTerminator())) {
-    if (SI->getCondition() != Val)
+    Value *Condition = SI->getCondition();
+    if (!isa<IntegerType>(Val->getType()))
+      return false;
+    bool ValUsesCondition = false;
+    if (Condition != Val) {
+      // Check if Val has Condition as an operand.
+      if (User *Usr = dyn_cast<User>(Val))
+        ValUsesCondition = usesOperand(Usr, Condition);
+      if (!ValUsesCondition)
         return false;
+    }
+    assert((Condition == Val || ValUsesCondition) &&
+           "Condition != Val nor Val doesn't use Condition");
 
     bool DefaultCase = SI->getDefaultDest() == BBTo;
     unsigned BitWidth = Val->getType()->getIntegerBitWidth();
     ConstantRange EdgesVals(BitWidth, DefaultCase/*isFullSet*/);
 
     for (auto Case : SI->cases()) {
-      ConstantRange EdgeVal(Case.getCaseValue()->getValue());
+      APInt CaseValue = Case.getCaseValue()->getValue();
+      ConstantRange EdgeVal(CaseValue);
+      if (ValUsesCondition) {
+        const DataLayout &DL = BBTo->getModule()->getDataLayout();
+        LVILatticeVal EdgeLatticeVal =
+            constantFoldUser(Val, Condition, CaseValue, DL);
+        if (EdgeLatticeVal.isOverdefined())
+          return false;
+        EdgeVal = EdgeLatticeVal.getConstantRange();
+      }
       if (DefaultCase) {
         // It is possible that the default destination is the destination of
-        // some cases. There is no need to perform difference for those cases.
-        if (Case.getCaseSuccessor() != BBTo)
+        // some cases. We cannot perform difference for those cases.
+        // We know Condition != CaseValue in BBTo.  In some cases we can use
+        // this to infer Val == f(Condition) is != f(CaseValue).  For now, we
+        // only do this when f is identity (i.e. Val == Condition), but we
+        // should be able to do this for any injective f.
+        if (Case.getCaseSuccessor() != BBTo && Condition == Val)
           EdgesVals = EdgesVals.difference(EdgeVal);
       } else if (Case.getCaseSuccessor() == BBTo)
         EdgesVals = EdgesVals.unionWith(EdgeVal);
     }
     Result = LVILatticeVal::getRange(std::move(EdgesVals));
     return true;
   }
   return false;

test/Transforms/CorrelatedValuePropagation/range.ll

 
 then:
   %result = icmp eq i32 %a, 255
   ret i1 %result
 
 else:
   ret i1 false
 }
+
+define i32 @test16(i8 %a) {
+entry:
+  %b = zext i8 %a to i32
+  br label %dispatch
+
+dispatch:
+  %cmp = icmp eq i8 %a, 93
+  br i1 %cmp, label %target93, label %dispatch
+
+; CHECK-LABEL: @test16(
+; CHECK: target93:
+; CHECK-NEXT: ret i32 93
+target93:
+  ret i32 %b
+}
+
+define i32 @test16_i1(i1 %a) {
+entry:
+  %b = zext i1 %a to i32
+  br label %dispatch
+
+dispatch:
+  br i1 %a, label %true, label %dispatch
+
+; CHECK-LABEL: @test16_i1(
+; CHECK: true:
+; CHECK-NEXT: ret i32 1
+true:
+  ret i32 %b
+}
+
+define i8 @test17(i8 %a) {
+entry:
+  %c = add i8 %a, 3
+  br label %dispatch
+
+dispatch:
+  %cmp = icmp eq i8 %a, 93
+  br i1 %cmp, label %target93, label %dispatch
+
+; CHECK-LABEL: @test17(
+; CHECK: target93:
+; CHECK-NEXT: ret i8 96
+target93:
+  ret i8 %c
+}
+
+define i8 @test17_2(i8 %a) {
+entry:
+  %c = add i8 %a, %a
+  br label %dispatch
+
+dispatch:
+  %cmp = icmp eq i8 %a, 93
+  br i1 %cmp, label %target93, label %dispatch
+
+; CHECK-LABEL: @test17_2(
+; CHECK: target93:
+; CHECK-NEXT: ret i8 -70
+target93:
+  ret i8 %c
+}
+
+define i1 @test17_i1(i1 %a) {
+entry:
+  %c = and i1 %a, true
+  br label %dispatch
+
+dispatch:
+  br i1 %a, label %true, label %dispatch
+
+; CHECK-LABEL: @test17_i1(
+; CHECK: true:
+; CHECK-NEXT: ret i1 true
+true:
+  ret i1 %c
+}
+
+define i32 @test18(i8 %a) {
+entry:
+  %b = zext i8 %a to i32
+  br label %dispatch
+
+dispatch:
+  switch i8 %a, label %dispatch [
+    i8 93, label %target93
+    i8 -111, label %dispatch
+  ]
+
+; CHECK-LABEL: @test18(
+; CHECK: target93:
+; CHECK-NEXT: ret i32 93
+target93:
+  ret i32 %b
+}
+
+define i8 @test19(i8 %a) {
+entry:
+  %c = add i8 %a, 3
+  br label %dispatch
+
+dispatch:
+  switch i8 %a, label %dispatch [
+    i8 93, label %target93
+    i8 -111, label %dispatch
+  ]
+
+; CHECK-LABEL: @test19(
+; CHECK: target93:
+; CHECK-NEXT: ret i8 96
+target93:
+  ret i8 %c
+}
+
+define i1 @test20(i64 %a) {
+entry:
+  %b = and i64 %a, 7
+  br label %dispatch
+
+dispatch:
+  switch i64 %a, label %default [
+    i64 0, label %exit2
+    i64 -2147483647, label %exit2
+  ]
+
+default:
+  %c = icmp eq i64 %b, 0
+  br label %exit
+
+exit:
+; Negative test. Shouldn't be incorrectly optimized to "ret i1 false".
+; CHECK-LABEL: @test20(
+; CHECK: exit:
+; CHECK-NOT: ret i1 false
+; CHECK: exit2:
+  ret i1 %c
+
+exit2:
+  ret i1 false
+}