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

[hjyamauchi]

Summary

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%.

[davidxl]

This patch needs a test case.

[hjyamauchi]

Added a test.

[hjyamauchi]

Wei, can you review this?

[hjyamauchi]

For more context:

Here's a test C program that demonstrates this issue (reduced from the Python 2.7 runtime eval loop code) :

void bar(int v);

int foo(unsigned char* p, int (*f)(int)) {
  int op;
  void* targets[256];
  int i;
  for (i = 0; i < 256; ++i) {
    targets[i] = &&unknown_op;
  }
  targets[93] = &&TARGET_93;
  targets[145] = &&TARGET_145;
  for (;;) {
    op = *p++;
 dispatch_op:
    switch (op) {
      TARGET_93:
      op = 93;
      case 93: {
        if (op != 93) {
          f(42);
        }
        goto *targets[*p++];
      }

      TARGET_145:
      op = 145;
      case 145: {
        op = *p++;
        goto dispatch_op;
      }

      unknown_op:
      default:
        bar(op);
        break;
    }
  }
  return 0;
}

There's a control merge between "TARGET_93" (that the computed goto may jumps to) and "case 93" (that the normal switch control flow may jump to).

The value of 'op' at "case 93:" would be constant 93 regardless of the control flow.

So, we could eliminate the call to 'f'.

GCC does this whereas LLVM does not.

[hjyamauchi]

I think this situation occurs due to the combination of the following two optimizations that shrink/narrow the width of operands.

They succeed narrowing the width of the switch condition operand "op" (%op.0.in in the test) and its address (%p.addr.0.pn in the test) from i32 to i8.

But since "op" still has a use as an i32 (passed to bar()), the zext remains.

Because of the zext, the current case edge value propagation doesn't perform the value propagation (op==93) on the case edge and the elimination of the call to f.

1. lib/Transforms/InstCombine/InstructionCombining.cpp

Instruction *InstCombiner::visitSwitchInst(SwitchInst &SI) {
  ...
  // Shrink the condition operand if the new type is smaller than the old type.
  // This may produce a non-standard type for the switch, but that's ok because
  // the backend should extend back to a legal type for the target.
  if (NewWidth > 0 && NewWidth < Known.getBitWidth()) {
    IntegerType *Ty = IntegerType::get(SI.getContext(), NewWidth);
    Builder.SetInsertPoint(&SI);
    Value *NewCond = Builder.CreateTrunc(Cond, Ty, "trunc");
    SI.setCondition(NewCond);

    for (auto Case : SI.cases()) {
      APInt TruncatedCase = Case.getCaseValue()->getValue().trunc(NewWidth);
      Case.setValue(ConstantInt::get(SI.getContext(), TruncatedCase));
    }
    return &SI;
  }

  return nullptr;
}

2. lib/Transforms/InstCombine/InstCombinePHI.cpp

Instruction *InstCombiner::FoldPHIArgZextsIntoPHI(PHINode &Phi) {
  ...
  // All incoming values are zexts or constants that are safe to truncate.
  // Create a new phi node of the narrow type, phi together all of the new
  // operands, and zext the result back to the original type.
  PHINode *NewPhi = PHINode::Create(NarrowType, NumIncomingValues,
                                    Phi.getName() + ".shrunk");
  for (unsigned i = 0; i != NumIncomingValues; ++i)
    NewPhi->addIncoming(NewIncoming[i], Phi.getIncomingBlock(i));

  InsertNewInstBefore(NewPhi, Phi);
  return CastInst::CreateZExtOrBitCast(NewPhi, Phi.getType());
}

[sanjoy]

lib/Analysis/LazyValueInfo.cpp

1397

s/extention/extension/

1409

This is oddly specific to zero extensions. Can you use some helpers in InstructionSimplify.h to make this more general over other kinds of operations? Something like: if Val has Condition as an operand and if Val is a CastInst or BinaryOperator (say), then invoke SimplifyBinOp / SimplifyFBinOp / SimplifyCastInst (depending on the opcode) with the operand that was Condition replaced with the appropriate constant.

[wmi]

I think we can find many similar problems.

long g;
void foo(int b) {
  int a = b;  
  g  = a;
  if (b == 3) {
    // assume we have `a = b` here, current LVI should be able to find out `a` equals to 3.
    __builtin_printf("%lx\n", a);
  }
}

Is there a way to make LVI more context sensitive, so that when we try to get value info for a and call LazyValueInfoImpl::solveBlockValueCast for a = b to find out the value info for b, we know the context is in the branch instead of the entry block? I just throw out the question and I don't know it is possible or not.

[• dberlin]

NewGVN, and other things will already get such cases.
I have a bug open for LVI to use predicateinfo, which would make it sensitive to this, as well as make it much faster, but that is a big amount of work.

[sanjoy]

lib/Analysis/LazyValueInfo.cpp

151

Can we return an Optional<APInt> here?

1360

You can use User::operands() and llvm::find here.

1366

I think we can return Optional<APInt> here.

Also: how about dropping Edge from the name, since we're not doing anything specific to bblock edges here? How about constantFoldUser?

1371

s/Condition/Op/

1384

In LLVM * aligns with the value. clang-format should fix this.

1394

I think this can be an assert -- switch only allows integer conditions.

1396

Does this logic subsume the logic in line 1431? And can we write the switch case case (no pun intended! :) ) in the same way? If so, I'd suggest pulling this logic out into a function and calling it once here and once from the switch-case block.

1404

Usually LLVM code does not use { for single statement ifs.

[hjyamauchi]

Addressed the comments.

lib/Analysis/LazyValueInfo.cpp

1394

I assume 'Val' can be any value live through the switch for which we're trying to find its value on this edge, and can be of any type?

1396

No, this logic does not subsume the logic in line 1431 (I verified that by commenting it out and running the tests.) This is because the 1431 logic handles the cases where 'Condition' is an operand of 'Val' whereas this logic handles the cases where 'Condition' is *not* an operand of 'Val', but one ('Op') of the operands of 'Val' can be inferred through the value of 'Condition'. Note this logic calls getValueFromCondition for 'Op'.

I moved the 1431 logic here and combined them.

I also factored out some common logic out to constantFoldUser() and simplified its call sites.

[sanjoy]

lgtm with minor comments.

lib/Analysis/LazyValueInfo.cpp

151

I didn't notice this the first time around, but I'd calls this function asConstantInteger (or something else that does not start with get). The getXXX in this class are simple accessors (which this function is not) so it is useful to differentiate this.

153

I think we can have direct returns here, instead of assigning to Result and then return Result;.

1374

This is minor, but LLVM tends to use auto * for assignments whose types are obvious from RHS.

1393–1399

The spacing is still off here, clang-format should fix this.

1394

Yes, you're right.

1396

Ok, SGTM.

1415

I suspect you can do

if (Optional<APInt> OpConst = OpLatticeVal.getConstantInteger())

[hjyamauchi]

Addressed the comments.

lib/Analysis/LazyValueInfo.cpp

public:
   }
 
   const ConstantRange &getConstantRange() const {
     assert(isConstantRange() &&
            "Cannot get the constant-range of a non-constant-range!");
     return Range;
   }
 
 private:
   void markOverdefined() {
     if (isOverdefined())
       return;
     Tag = overdefined;
   }
 
   void markConstant(Constant *V) {
     assert(V && "Marking constant with NULL");
     if (ConstantInt *CI = dyn_cast<ConstantInt>(V)) {
       markConstantRange(ConstantRange(CI->getValue()));

 LVILatticeVal getValueFromCondition(Value *Val, Value *Cond, bool isTrueDest) {
   assert(Cond && "precondition");
   DenseMap<Value*, LVILatticeVal> Visited;
   return getValueFromCondition(Val, Cond, isTrueDest, Visited);
 }
 
 /// \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");
 
       // If V is the condition of the branch itself, then we know exactly what
       // it is.
       if (BI->getCondition() == 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);
       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();
+    bool ZeroExtendedCondition = false;
+    if (isa<ZExtInst>(Val) &&
+        cast<ZExtInst>(Val)->getOperand(0) == Condition) {
+      // Val is an zero extention of the switch condition value.
+      ZeroExtendedCondition = true;
+    } else if (Condition != Val) {
       return false;
-
+    }
     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();
+      if (ZeroExtendedCondition) {
+        CaseValue = CaseValue.zext(BitWidth);
+      }
+      ConstantRange EdgeVal(CaseValue);
       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)
           EdgesVals = EdgesVals.difference(EdgeVal);
       } else if (Case.getCaseSuccessor() == BBTo)
         EdgesVals = EdgesVals.unionWith(EdgeVal);
     }
     Result = LVILatticeVal::getRange(std::move(EdgesVals));
     return true;
   }
   return false;

test/Analysis/LazyValueAnalysis/lvi-zeroextended-switch-condition.ll

+; RUN: opt < %s -jump-threading -print-lazy-value-info -disable-output 2>&1 | FileCheck %s
+
+declare void @bar(i32)
+
+define i32 @foo_switch(i8* nocapture readonly %p, i32 (i32)* nocapture %f) #0 {
+entry:
+  %targets = alloca [256 x i8*], align 16
+  br label %for.cond
+
+for.cond:                                         ; preds = %for.body, %entry
+  %i.0 = phi i32 [ 0, %entry ], [ %inc, %for.body ]
+  %cmp = icmp slt i32 %i.0, 256
+  br i1 %cmp, label %for.body, label %for.end
+
+for.body:                                         ; preds = %for.cond
+  %idxprom = sext i32 %i.0 to i64
+  %arrayidx = getelementptr inbounds [256 x i8*], [256 x i8*]* %targets, i64 0, i64 %idxprom
+  store i8* blockaddress(@foo_switch, %sw.default), i8** %arrayidx, align 8
+  %inc = add nsw i32 %i.0, 1
+  br label %for.cond
+
+for.end:                                          ; preds = %for.cond
+  %arrayidx1 = getelementptr inbounds [256 x i8*], [256 x i8*]* %targets, i64 0, i64 93
+  store i8* blockaddress(@foo_switch, %if.end), i8** %arrayidx1, align 8
+  %arrayidx2 = getelementptr inbounds [256 x i8*], [256 x i8*]* %targets, i64 0, i64 145
+  store i8* blockaddress(@foo_switch, %dispatch_op), i8** %arrayidx2, align 8
+  br label %for.cond3
+
+for.cond3:                                        ; preds = %sw.default, %for.end
+  %p.addr.0 = phi i8* [ %p, %for.end ], [ %p.addr.4, %sw.default ]
+  br label %dispatch_op
+
+dispatch_op:                                      ; preds = %dispatch_op, %if.end, %for.cond3
+  %p.addr.0.pn = phi i8* [ %p.addr.0, %for.cond3 ], [ %p.addr.1, %dispatch_op ], [ %incdec.ptr6, %if.end ]
+  %op.0.in = load i8, i8* %p.addr.0.pn, align 1
+  %p.addr.1 = getelementptr inbounds i8, i8* %p.addr.0.pn, i64 1
+  %op.0 = zext i8 %op.0.in to i32
+  switch i8 %op.0.in, label %sw.default [
+    i8 93, label %if.end
+    i8 -111, label %dispatch_op
+  ]
+
+; Check that lazy value info propagates the zero-extended switch condition value through the case
+; edge and figures that %op.1 is 93.
+;
+; CHECK-LABEL: if.end:
+; CHECK:       ; LatticeVal for: '  %op.1 = phi i32 [ %op.0, %dispatch_op ], [ 93, %if.end ]' in BB: '%if.end' is: constantrange<93, 94>
+if.end:                                           ; preds = %dispatch_op, %if.end
+  %p.addr.2 = phi i8* [ %p.addr.1, %dispatch_op ], [ %incdec.ptr6, %if.end ]
+  %op.1 = phi i32 [ %op.0, %dispatch_op ], [ 93, %if.end ]
+  %incdec.ptr6 = getelementptr inbounds i8, i8* %p.addr.2, i64 1
+  %0 = load i8, i8* %p.addr.2, align 1
+  %idxprom7 = zext i8 %0 to i64
+  %arrayidx8 = getelementptr inbounds [256 x i8*], [256 x i8*]* %targets, i64 0, i64 %idxprom7
+  %1 = load i8*, i8** %arrayidx8, align 8
+  indirectbr i8* %1, [label %sw.default, label %if.end, label %dispatch_op]
+
+sw.default:                                       ; preds = %if.end, %dispatch_op
+  %p.addr.4 = phi i8* [ %p.addr.1, %dispatch_op ], [ %incdec.ptr6, %if.end ]
+  %op.2 = phi i32 [ %op.0, %dispatch_op ], [ %op.1, %if.end ]
+  call void @bar(i32 %op.2) #2
+  br label %for.cond3
+}