[SimplifyCFG] Propagating case value when turning switch range into icmp

[wwei]

Summary

When converting switch range into an icmp, the initial case value may be lost,
which may prevent instructions in successors from being optimized.
For example:

define i64 @demo(i64 %x) {
  entry:
      switch i64 %x, label %bb3 [
          i64 0, label %bb1
          i64 1, label %bb2
      ]
  bb1:
      ret i64 0
  bb2:
      ; this will necessarily be false because %x == 1
      %0 = icmp eq i64 %x, 100
      br i1 %0, label %bb4, label %bb5
  bb3:
      unreachable
  bb4:
      ret i64 200
  bb5:
      ret i64 10
}

Ideally, the result after SimplifyCFG should be as follows:

define i64 @demo(i64 %x) {
  %entry:
    %switch = icmp eq i64 %x, 0
    %. = select i1 %switch, i64 0, i64 10
    ret i64 %.
}

After applying this patch, the case value will be Propagated to the instructions in successors.
issue: SimplifyCFG Need to remove useless comparison after turning switch to icmp

[lebedev.ri]

It doesn't look like something SimplifyCFG should be dealing with.

[wwei]

In D119839#3322795, @lebedev.ri wrote:

It doesn't look like something SimplifyCFG should be dealing with.

@lebedev.ri The reason is that after switch range is converted to icmp, the original case value no longer exists in the IR, and there is no chance for other pass to eliminate this comparison instruction. Look at the IR after simplifycfg：

define i64 @demo(i64 %x) {
entry:
  %switch = icmp ult i64 %x, 1
  %0 = icmp eq i64 %x, 100     ; %x == 1?  no pass can get this result！
  %. = select i1 %0, i64 200, i64 10
  %common.ret.op = select i1 %switch, i64 0, i64 %.
  ret i64 %common.ret.op
}

So It's impossible to remove %0 = icmp eq i64 %x, 100. Maybe simplifycfg has the responsibility to do constant propagation if switch to icmp transformation is performed.
I can't find other solutions at the moment, do you have better suggestions or deas?

[lebedev.ri]

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

6577–6580

Perhaps TurnSwitchRangeIntoICmp() simply should have similar treatment to SwitchToLookupTable()?

[lebedev.ri]

Will post patch in a sec.

[lebedev.ri]

D119854

[lebedev.ri]

Abandon this?

[wwei]

In D119839#3328631, @lebedev.ri wrote:

Abandon this?

Sure.

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

if (!CasesA.empty() && CasesAreContiguous(CasesA)) {
     OtherDest = DestB;
   } else if (CasesAreContiguous(CasesB)) {
     ContiguousCases = &CasesB;
     ContiguousDest = DestB;
     OtherDest = DestA;
   } else
     return false;
 
+  // Try to propagate switch value to case successors.
+  SmallVector<const Use *, 8> Worklist;
+  Value *SwitchCond = SI->getCondition();
+  for (const Use &U : SwitchCond->uses())
+    Worklist.push_back(&U);
+  while (!Worklist.empty()) {
+    const Use *U = Worklist.pop_back_val();
+    auto UserInst = dyn_cast<Instruction>(U->getUser());
+    auto UserParent = UserInst->getParent();
+    if (UserParent == SI->getDefaultDest())
+      continue;
+    if (UserParent == DestA && DestA->getSinglePredecessor()) {
+      assert(CasesA.size() == 1);
+      Value *CaseAValue = cast<Value>(CasesA.back());
+      UserInst->replaceUsesOfWith(SwitchCond, CaseAValue);
+    }
+    if (UserParent == DestB && DestB->getSinglePredecessor()) {
+      assert(CasesB.size() == 1);
+      Value *CaseBValue = cast<Value>(CasesB.back());
+      UserInst->replaceUsesOfWith(SwitchCond, CaseBValue);
+    }
+  }
+
   // Start building the compare and branch.
 
   Constant *Offset = ConstantExpr::getNeg(ContiguousCases->back());
   Constant *NumCases =
       ConstantInt::get(Offset->getType(), ContiguousCases->size());
 
   Value *Sub = SI->getCondition();
   if (!Offset->isNullValue())
if (isValueEqualityComparison(SI)) {
 
     // If the block only contains the switch, see if we can fold the block
     // away into any preds.
     if (SI == &*BB->instructionsWithoutDebug(false).begin())
       if (FoldValueComparisonIntoPredecessors(SI, Builder))
         return requestResimplify();
   }
 
   // Try to transform the switch into an icmp and a branch.
   if (TurnSwitchRangeIntoICmp(SI, Builder))
     return requestResimplify();
 
   // Remove unreachable cases.
   if (eliminateDeadSwitchCases(SI, DTU, Options.AC, DL))
     return requestResimplify();
 
   if (switchToSelect(SI, Builder, DTU, DL, TTI))
     return requestResimplify();
 
   if (Options.ForwardSwitchCondToPhi && ForwardSwitchConditionToPHI(SI))

llvm/test/Transforms/SimplifyCFG/switch-to-icmp-with-value-propagate.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt < %s -simplifycfg -S | FileCheck %s
+
+; propagate value to bb2.
+define i64 @test1(i64 %x) {
+; CHECK-LABEL: @test1(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[SWITCH:%.*]] = icmp ult i64 [[X:%.*]], 1
+; CHECK-NEXT:    [[TMP0:%.*]] = icmp eq i64 1, 100
+; CHECK-NEXT:    [[DOT:%.*]] = select i1 [[TMP0]], i64 200, i64 10
+; CHECK-NEXT:    [[COMMON_RET_OP:%.*]] = select i1 [[SWITCH]], i64 0, i64 [[DOT]]
+; CHECK-NEXT:    ret i64 [[COMMON_RET_OP]]
+;
+entry:
+  switch i64 %x, label %bb3 [
+  i64 0, label %bb1
+  i64 1, label %bb2
+  ]
+bb1:
+  ret i64 0
+bb2:
+  %0 = icmp eq i64 %x, 100
+  br i1 %0, label %bb4, label %bb5
+bb3:
+  unreachable
+bb4:
+  ret i64 200
+bb5:
+  ret i64 10
+}
+
+; propagate value both to bb1 and bb2.
+define i64 @test2(i64 %x) {
+; CHECK-LABEL: @test2(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[X_OFF:%.*]] = add i64 [[X:%.*]], -1
+; CHECK-NEXT:    [[SWITCH:%.*]] = icmp ult i64 [[X_OFF]], 1
+; CHECK-NEXT:    br i1 [[SWITCH]], label [[BB1:%.*]], label [[BB2:%.*]]
+; CHECK:       bb1:
+; CHECK-NEXT:    [[TMP0:%.*]] = icmp eq i64 1, 100
+; CHECK-NEXT:    br i1 [[TMP0]], label [[BB4:%.*]], label [[COMMON_RET:%.*]]
+; CHECK:       common.ret:
+; CHECK-NEXT:    [[COMMON_RET_OP:%.*]] = phi i64 [ 200, [[BB4]] ], [ 0, [[BB1]] ], [ 10, [[BB2]] ]
+; CHECK-NEXT:    ret i64 [[COMMON_RET_OP]]
+; CHECK:       bb2:
+; CHECK-NEXT:    [[TMP1:%.*]] = icmp eq i64 2, 101
+; CHECK-NEXT:    br i1 [[TMP1]], label [[BB4]], label [[COMMON_RET]]
+; CHECK:       bb4:
+; CHECK-NEXT:    br label [[COMMON_RET]]
+;
+entry:
+  switch i64 %x, label %bb3 [
+  i64 1, label %bb1
+  i64 2, label %bb2
+  ]
+bb1:
+  %0 = icmp eq i64 %x, 100
+  br i1 %0, label %bb4, label %return
+return:
+  ret i64 0
+bb2:
+  %1 = icmp eq i64 %x, 101
+  br i1 %1, label %bb4, label %bb5
+bb3:
+  unreachable
+bb4:
+  ret i64 200
+bb5:
+  ret i64 10
+}
+
+define i64 @test3(i64 %x) {
+; CHECK-LABEL: @test3(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[COND:%.*]] = icmp eq i64 [[X:%.*]], 1
+; CHECK-NEXT:    [[SPEC_SELECT:%.*]] = select i1 [[COND]], i64 10, i64 0
+; CHECK-NEXT:    ret i64 [[SPEC_SELECT]]
+;
+entry:
+  switch i64 %x, label %bb1 [
+  i64 1, label %bb2
+  ]
+bb1:
+  ret i64 0
+bb2:
+  %0 = icmp eq i64 %x, 100
+  br i1 %0, label %bb4, label %bb5
+bb4:
+  ret i64 200
+bb5:
+  ret i64 10
+}
+
+; bb2 has two predecessors with case value 1 and 2.
+define i64 @test_fail1(i64 %x) {
+; CHECK-LABEL: @test_fail1(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[SWITCH:%.*]] = icmp ult i64 [[X:%.*]], 1
+; CHECK-NEXT:    [[TMP0:%.*]] = icmp eq i64 [[X]], 100
+; CHECK-NEXT:    [[DOT:%.*]] = select i1 [[TMP0]], i64 200, i64 10
+; CHECK-NEXT:    [[COMMON_RET_OP:%.*]] = select i1 [[SWITCH]], i64 0, i64 [[DOT]]
+; CHECK-NEXT:    ret i64 [[COMMON_RET_OP]]
+;
+entry:
+  switch i64 %x, label %bb3 [
+  i64 0, label %bb1
+  i64 1, label %bb2
+  i64 2, label %bb2
+  ]
+bb1:
+  ret i64 0
+bb2:
+  %0 = icmp eq i64 %x, 100
+  br i1 %0, label %bb4, label %bb5
+bb3:
+  unreachable
+bb4:
+  ret i64 200
+bb5:
+  ret i64 10
+}
+
+; return block has two predecessors.
+define i64 @test_fail2(i64 %x) {
+; CHECK-LABEL: @test_fail2(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[SWITCH:%.*]] = icmp ult i64 [[X:%.*]], 1
+; CHECK-NEXT:    [[SPEC_SELECT:%.*]] = select i1 [[SWITCH]], i64 2, i64 [[X]]
+; CHECK-NEXT:    ret i64 [[SPEC_SELECT]]
+;
+entry:
+  switch i64 %x, label %bb2 [
+  i64 0, label %bb1
+  i64 1, label %return
+  ]
+bb1:
+  br label %return
+return:
+  %retval.0 = phi i64 [ %x, %entry ], [ 2, %bb1 ]
+  ret i64 %retval.0
+bb2:
+  unreachable
+}