[llvm] [SimplifyCFG] Add optimization for switches of powers of two (PR #70977)

[via llvm-commits]

================
@@ -0,0 +1,139 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 3
+; RUN: opt %s -passes='simplifycfg<switch-to-lookup>' -simplifycfg-require-and-preserve-domtree=1 -S | FileCheck %s
+
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+
+; Check that the range of switch of powers of two is reduced and switch itself is lowered to jump-table
+define i32 @switch_of_powers(i32 %x) {
+; CHECK-LABEL: define i32 @switch_of_powers(
+; CHECK-SAME: i32 [[X:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = call i32 @llvm.cttz.i32(i32 [[X]], i1 true)
+; CHECK-NEXT:    [[SWITCH_MASKINDEX:%.*]] = trunc i32 [[TMP0]] to i8
+; CHECK-NEXT:    [[SWITCH_SHIFTED:%.*]] = lshr i8 121, [[SWITCH_MASKINDEX]]
+; CHECK-NEXT:    [[SWITCH_LOBIT:%.*]] = trunc i8 [[SWITCH_SHIFTED]] to i1
+; CHECK-NEXT:    call void @llvm.assume(i1 [[SWITCH_LOBIT]])
+; CHECK-NEXT:    [[SWITCH_GEP:%.*]] = getelementptr inbounds [7 x i32], ptr @switch.table.switch_of_powers, i32 0, i32 [[TMP0]]
+; CHECK-NEXT:    [[SWITCH_LOAD:%.*]] = load i32, ptr [[SWITCH_GEP]], align 4
+; CHECK-NEXT:    ret i32 [[SWITCH_LOAD]]
+;
+
+entry:
+  switch i32 %x, label %default_case [
+  i32 1,  label %bb1
+  i32 8,  label %bb2
+  i32 16, label %bb3
+  i32 32, label %bb4
+  i32 64, label %bb5
+  ]
+
+
+default_case: unreachable
+bb1: br label %return
+bb2: br label %return
+bb3: br label %return
+bb4: br label %return
+bb5: br label %return
+
+return:
+  %p = phi i32 [ 3, %bb1 ], [ 2, %bb2 ], [ 1, %bb3 ], [ 0, %bb4 ], [ 42, %bb5 ]
+  ret i32 %p
+}
+
+; Check that switch's of powers of two range is not reduced if default case is reachable
+define i32 @switch_of_powers_reachable_default(i32 %x) {
+; CHECK-LABEL: define i32 @switch_of_powers_reachable_default(
+; CHECK-SAME: i32 [[X:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    switch i32 [[X]], label [[RETURN:%.*]] [
+; CHECK-NEXT:    i32 1, label [[BB1:%.*]]
+; CHECK-NEXT:    i32 8, label [[BB2:%.*]]
+; CHECK-NEXT:    i32 16, label [[BB3:%.*]]
+; CHECK-NEXT:    i32 32, label [[BB4:%.*]]
+; CHECK-NEXT:    i32 64, label [[BB5:%.*]]
+; CHECK-NEXT:    ]
+; CHECK:       bb1:
+; CHECK-NEXT:    br label [[RETURN]]
+; CHECK:       bb2:
+; CHECK-NEXT:    br label [[RETURN]]
+; CHECK:       bb3:
+; CHECK-NEXT:    br label [[RETURN]]
+; CHECK:       bb4:
+; CHECK-NEXT:    br label [[RETURN]]
+; CHECK:       bb5:
+; CHECK-NEXT:    br label [[RETURN]]
+; CHECK:       return:
+; CHECK-NEXT:    [[P:%.*]] = phi i32 [ 3, [[BB1]] ], [ 2, [[BB2]] ], [ 1, [[BB3]] ], [ 0, [[BB4]] ], [ 42, [[BB5]] ], [ -1, [[ENTRY:%.*]] ]
+; CHECK-NEXT:    ret i32 [[P]]
+;
+
+entry:
+  switch i32 %x, label %default_case [
+  i32 1,  label %bb1
+  i32 8,  label %bb2
+  i32 16, label %bb3
+  i32 32, label %bb4
+  i32 64, label %bb5
+  ]
+
+
+default_case: br label %return
+bb1: br label %return
+bb2: br label %return
+bb3: br label %return
+bb4: br label %return
+bb5: br label %return
+
+return:
+  %p = phi i32 [ 3, %bb1 ], [ 2, %bb2 ], [ 1, %bb3 ], [ 0, %bb4 ], [ 42, %bb5 ], [-1, %default_case]
+  ret i32 %p
+}
+
+; Check that switch with zero case is not considered as switch of powers of two
+define i32 @switch_of_non_powers(i32 %x) {
+; CHECK-LABEL: define i32 @switch_of_non_powers(
+; CHECK-SAME: i32 [[X:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    switch i32 [[X]], label [[DEFAULT_CASE:%.*]] [
+; CHECK-NEXT:    i32 0, label [[RETURN:%.*]]
+; CHECK-NEXT:    i32 1, label [[BB2:%.*]]
+; CHECK-NEXT:    i32 16, label [[BB3:%.*]]
+; CHECK-NEXT:    i32 32, label [[BB4:%.*]]
+; CHECK-NEXT:    i32 64, label [[BB5:%.*]]
+; CHECK-NEXT:    ]
+; CHECK:       default_case:
+; CHECK-NEXT:    unreachable
+; CHECK:       bb2:
+; CHECK-NEXT:    br label [[RETURN]]
+; CHECK:       bb3:
+; CHECK-NEXT:    br label [[RETURN]]
+; CHECK:       bb4:
+; CHECK-NEXT:    br label [[RETURN]]
+; CHECK:       bb5:
+; CHECK-NEXT:    br label [[RETURN]]
+; CHECK:       return:
+; CHECK-NEXT:    [[P:%.*]] = phi i32 [ 2, [[BB2]] ], [ 1, [[BB3]] ], [ 0, [[BB4]] ], [ 42, [[BB5]] ], [ 3, [[ENTRY:%.*]] ]
+; CHECK-NEXT:    ret i32 [[P]]
+;
+
+entry:
+  switch i32 %x, label %default_case [
+  i32 0,  label %bb1
+  i32 1,  label %bb2
+  i32 16, label %bb3
+  i32 32, label %bb4
+  i32 64, label %bb5
+  ]
+
+
+default_case: unreachable
+bb1: br label %return
+bb2: br label %return
+bb3: br label %return
+bb4: br label %return
+bb5: br label %return
+
+return:
+  %p = phi i32 [ 3, %bb1 ], [ 2, %bb2 ], [ 1, %bb3 ], [ 0, %bb4 ], [ 42, %bb5 ]
+  ret i32 %p
+}
----------------
DKay7 wrote:

1. I found out that pre-checking if switch is already dense is useless because in that case it would be transformed to a jump table, so I removed this part.
2. Added test case covering the case when we don't apply transformation because it won't make the switch dense.
3. I can't come up with the case when switch will be transformed (i.e. become dense) but jump table will not be generated.

Thanks

https://github.com/llvm/llvm-project/pull/70977