Teach SimplifyCFG to fold switches into lookup tables in more cases.

[resistor]

Summary

In particular, it couldn't handle cases where lookup table constant
expressions involved bitcasts. This does not seem to come up
frequently in C++, but comes up reasonably often in Rust via
#[derive(Debug)].

Originally reported by pcwalton.

[pcwalton]

Looks fine to me. Note that to trigger on Debug we also need the frontend changes at https://github.com/rust-lang/rust/pull/88832.

[lebedev.ri]

Won't this also skip non-zero-offset GEP's?
Are you sure you don't want stripPointerCasts()/stripPointerCastsAndAliases()/stripPointerCastsSameRepresentation() instead?

[nikic]

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

5060

Note that GEPs are already handled below, so we should either drop that handling, or handle bitcasts in the same way. The current GEP handling uses "no notional overindexing" as the condition, which I agree is an unnecessary requirement here -- your choice of "inbounds with constant offsets" is the right one. Whether there is notional overindexing really shouldn't matter in this context, the relocation only cares about offsets.

However, something your implementation changes is that shouldBuildLookupTablesForConstant() no longer gets called on the original constant, only on the stripped one. It looks like only ARM uses this hook and wouldn't be affected by the change (because the stripped bitcasts/GEPs do not affect whether the constant requires relocation), but I think it would be in the spirit of the hook to continue calling it on the original constant.

[resistor]

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

5060

Good point. I will update to account for this.

[resistor]

Update based on comments to preserve original behavior of shouldBuildLookupTablesForConstant().

[resistor]

In D109565#2995449, @lebedev.ri wrote:

Won't this also skip non-zero-offset GEP's?

Yes, that's intentional. See the comment and the other thread of discussion.

[nikic]

llvm/test/Transforms/SimplifyCFG/switch-to-lookup-bitcast.ll

3 ↗

(On Diff #372023)

This target triple either needs to be dropped (if the test works without it), or the test should be moved in to the X86 subdirectory.

16 ↗

(On Diff #372023)

It would be good to reduce the number of switch cases in this test -- we probably don't need 12 to produce a lookup table.

llvm/test/Transforms/SimplifyCFG/switch-to-lookup-gep.ll

80 ↗

(On Diff #372023)

This only tests a GEP with a zero offset, which is a degenerate case. It would be better to test a non-zero offset.

We could also have a negative test that no lookup table is created if a GEP index is not a constant integer (but rather something like ptrtoint i8* @g to i64.

[resistor]

Update for comments.

llvm/test/Transforms/SimplifyCFG/switch-to-lookup-bitcast.ll

3 ↗

(On Diff #372023)

The test does not work without it.

llvm/test/Transforms/SimplifyCFG/switch-to-lookup-gep.ll

80 ↗

(On Diff #372023)

Updated test to use non-zero offsets. The proposed negative test doesn't fail, and it's not clear to me that it should fail.

[resistor]

Update for comments.

[nikic]

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

5072

Shouldn't this condition be StrippedC == C ||? What you're doing now is effectively allowing all constant expressions for which nothing can be stripped.

[nikic]

llvm/test/Transforms/SimplifyCFG/switch-to-lookup-gep.ll

80 ↗

(On Diff #372023)

This would effectively correspond to a "global1 + global2" expression, which is not a generally (or at all?) supported relocation type.

[resistor]

Update for comments.

[resistor]

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

5072

Done

[resistor]

llvm/test/Transforms/SimplifyCFG/switch-to-lookup-gep.ll

80 ↗

(On Diff #372023)

The proposed test is dicey. I'm not sure that "global1 + global2" can ever be correctly marked as inbounds, in which case it's a semantically invalid input.

[nikic]

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

5072

This still needs a negative test for unsupported constant expressions.

[resistor]

Add a negative test for non-inbounds GEPs.

[nikic]

llvm/test/Transforms/SimplifyCFG/switch-to-lookup-gep.ll

80 ↗

(On Diff #372023)

Unless LLVM can prove that it's not inbounds, it doesn't really matter. Say global2 is a weak symbol and you know that it will actually always be NULL and thus the addition is inbounds. We just use ptrtoint of globals as an easy way to produce an unfoldable expression.

[resistor]

llvm/test/Transforms/SimplifyCFG/switch-to-lookup-gep.ll

80 ↗

(On Diff #372023)

In that case we have to weaken the stripping - inbounds isn't a strong enough condition to correctness.

[nikic]

llvm/test/Transforms/SimplifyCFG/switch-to-lookup-gep.ll

80 ↗

(On Diff #372023)

The condition is already "inbounds with constant offsets". While not obvious from the name, "constant offsets" here means ConstantInt offsets and already excludes constant expression GEP operands.

[resistor]

Add negative test for ptrtoint.

[resistor]

llvm/test/Transforms/SimplifyCFG/switch-to-lookup-gep.ll

80 ↗

(On Diff #372023)

Added a test for this scenario.

[nikic]

LGTM, thanks!

llvm/lib/Transforms/Utils/SimplifyCFG.cpp

static bool ForwardSwitchConditionToPHI(SwitchInst *SI) {
   return Changed;
 }
 
 /// Return true if the backend will be able to handle
 /// initializing an array of constants like C.
 static bool ValidLookupTableConstant(Constant *C, const TargetTransformInfo &TTI) {
   if (C->isThreadDependent())
     return false;
   if (C->isDLLImportDependent())
     return false;
 
   if (!isa<ConstantFP>(C) && !isa<ConstantInt>(C) &&
       !isa<ConstantPointerNull>(C) && !isa<GlobalValue>(C) &&
       !isa<UndefValue>(C) && !isa<ConstantExpr>(C))
     return false;
 
   if (ConstantExpr *CE = dyn_cast<ConstantExpr>(C)) {
-    if (!CE->isGEPWithNoNotionalOverIndexing())
-      return false;
-    if (!ValidLookupTableConstant(CE->getOperand(0), TTI))
+    // Pointer casts and in-bounds GEPs will not prohibit the backend from
+    // materializing the array of constants.
+    Constant *StrippedC = cast<Constant>(CE->stripInBoundsConstantOffsets());
+    if (StrippedC != C && !ValidLookupTableConstant(StrippedC, TTI))
       return false;
   }
 
   if (!TTI.shouldBuildLookupTablesForConstant(C))
     return false;
 
   return true;
 }

llvm/test/Transforms/SimplifyCFG/X86/switch-to-lookup-bitcast.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --scrub-attributes
+; RUN: opt -simplifycfg --switch-to-lookup -S < %s | FileCheck %s
+target triple = "x86_64-unknown-linux-gnu"
+
+@alloc0 = private unnamed_addr constant <{ [1 x i8] }> <{ [1 x i8] c"A" }>, align 1
+@alloc1 = private unnamed_addr constant <{ [1 x i8] }> <{ [1 x i8] c"B" }>, align 1
+@alloc2 = private unnamed_addr constant <{ [1 x i8] }> <{ [1 x i8] c"C" }>, align 1
+
+define { [0 x i8]*, i64 } @switch_to_lookup_bitcast(i8 %0) unnamed_addr {
+; CHECK-LABEL: @switch_to_lookup_bitcast(
+; CHECK-NEXT:  start:
+; CHECK-NEXT:    [[SWITCH_GEP:%.*]] = getelementptr inbounds [3 x [0 x i8]*], [3 x [0 x i8]*]* @switch.table.switch_to_lookup_bitcast, i32 0, i8 [[TMP0:%.*]]
+; CHECK-NEXT:    [[SWITCH_LOAD:%.*]] = load [0 x i8]*, [0 x i8]** [[SWITCH_GEP]], align 8
+; CHECK-NEXT:    [[TMP1:%.*]] = insertvalue { [0 x i8]*, i64 } undef, [0 x i8]* [[SWITCH_LOAD]], 0
+; CHECK-NEXT:    [[TMP2:%.*]] = insertvalue { [0 x i8]*, i64 } [[TMP1]], i64 1, 1
+; CHECK-NEXT:    ret { [0 x i8]*, i64 } [[TMP2]]
+;
+start:
+  switch i8 %0, label %default [
+  i8 0, label %bb0
+  i8 1, label %bb1
+  i8 2, label %bb2
+  ]
+
+bb0:
+  br label %end
+
+bb1:
+  br label %end
+
+bb2:
+  br label %end
+
+default:
+  unreachable
+
+end:
+  %.sroa.0.0 = phi [0 x i8]* [ bitcast (<{ [1 x i8] }>* @alloc0 to [0 x i8]*), %bb0 ], [ bitcast (<{ [1 x i8] }>* @alloc1 to [0 x i8]*), %bb1 ], [ bitcast (<{ [1 x i8] }>* @alloc2 to [0 x i8]*), %bb2 ]
+  %1 = insertvalue { [0 x i8]*, i64 } undef, [0 x i8]* %.sroa.0.0, 0
+  %2 = insertvalue { [0 x i8]*, i64 } %1, i64 1, 1
+  ret { [0 x i8]*, i64 } %2
+}

llvm/test/Transforms/SimplifyCFG/X86/switch-to-lookup-gep.ll

+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --scrub-attributes
+; RUN: opt -simplifycfg --switch-to-lookup -S < %s | FileCheck %s
+target triple = "x86_64-unknown-linux-gnu"
+
+@alloc0 = private unnamed_addr constant <{ [2 x i8] }> <{ [2 x i8] c"A1" }>, align 1
+@alloc1 = private unnamed_addr constant <{ [2 x i8] }> <{ [2 x i8] c"B2" }>, align 1
+@alloc2 = private unnamed_addr constant <{ [2 x i8] }> <{ [2 x i8] c"C3" }>, align 1
+
+define { i8*, i64 } @switch_to_lookup_gep(i8 %0) unnamed_addr {
+; CHECK-LABEL: @switch_to_lookup_gep(
+; CHECK-NEXT:  start:
+; CHECK-NEXT:    [[SWITCH_GEP:%.*]] = getelementptr inbounds [3 x i8*], [3 x i8*]* @switch.table.switch_to_lookup_gep, i32 0, i8 [[TMP0:%.*]]
+; CHECK-NEXT:    [[SWITCH_LOAD:%.*]] = load i8*, i8** [[SWITCH_GEP]], align 8
+; CHECK-NEXT:    [[TMP1:%.*]] = insertvalue { i8*, i64 } undef, i8* [[SWITCH_LOAD]], 0
+; CHECK-NEXT:    [[TMP2:%.*]] = insertvalue { i8*, i64 } [[TMP1]], i64 1, 1
+; CHECK-NEXT:    ret { i8*, i64 } [[TMP2]]
+;
+start:
+  switch i8 %0, label %default [
+  i8 0, label %bb0
+  i8 1, label %bb1
+  i8 2, label %bb2
+  ]
+
+bb0:
+  br label %end
+
+bb1:
+  br label %end
+
+bb2:
+  br label %end
+
+default:
+  unreachable
+
+end:
+  %.sroa.0.0 = phi i8* [ getelementptr inbounds (<{ [2 x i8] }>, <{ [2 x i8] }>* @alloc0, i32 0, i32 0, i32 1), %bb0 ], [ getelementptr inbounds (<{ [2 x i8] }>, <{ [2 x i8] }>* @alloc1, i32 0, i32 0, i32 1), %bb1 ], [ getelementptr inbounds (<{ [2 x i8] }>, <{ [2 x i8] }>* @alloc2, i32 0, i32 0, i32 1), %bb2 ]
+  %1 = insertvalue { i8*, i64 } undef, i8* %.sroa.0.0, 0
+  %2 = insertvalue { i8*, i64 } %1, i64 1, 1
+  ret { i8*, i64 } %2
+}