[llvm] [SimplifyCFG] Delete the unnecessary range check for small mask operation (PR #65835)

[via llvm-commits]

https://github.com/vfdff updated https://github.com/llvm/llvm-project/pull/65835

>From dd72ab92b9272d9d8c0074361fdf03efd300ab94 Mon Sep 17 00:00:00 2001
From: zhongyunde 00443407 <zhongyunde at huawei.com>
Date: Fri, 1 Sep 2023 23:05:45 -0400
Subject: [PATCH 1/2] [SimplifyCFG] Delete the unnecessary range check for
 small mask operation

When the small mask value little than 64, we can eliminate the checking
for upper limit of the range by enlarge the lookup table size to its next
pow2 value.

Fixes https://github.com/llvm/llvm-project/issues/65120
---
 llvm/lib/Transforms/Utils/SimplifyCFG.cpp     | 23 ++++++--
 .../Transforms/SimplifyCFG/switch_mask.ll     | 57 +++++++++++++++++++
 2 files changed, 76 insertions(+), 4 deletions(-)
 create mode 100644 llvm/test/Transforms/SimplifyCFG/switch_mask.ll

diff --git a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
index 7af366b47e9f56d..90e10a7a75067da 100644
--- a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
+++ b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
@@ -6552,10 +6552,25 @@ static bool SwitchToLookupTable(SwitchInst *SI, IRBuilder<> &Builder,
     // Note: We call removeProdecessor later since we need to be able to get the
     // PHI value for the default case in case we're using a bit mask.
   } else {
-    Value *Cmp = Builder.CreateICmpULT(
-        TableIndex, ConstantInt::get(MinCaseVal->getType(), TableSize));
-    RangeCheckBranch =
-        Builder.CreateCondBr(Cmp, LookupBB, SI->getDefaultDest());
+    int KeepDefaultBranch = true;
+    if (UseSwitchConditionAsTableIndex) {
+      ConstantRange CR = computeConstantRange(TableIndex, /* ForSigned */ false,
+                                              /* UseInstrInfo*/ true);
+      if (DL.fitsInLegalInteger(1 << CR.getActiveBits())) {
+        // The default branch is unreachable when we enlarge the lookup table.
+        TableSize = CR.getUpper().getZExtValue();
+        KeepDefaultBranch = false;
+      }
+    }
+
+    if (KeepDefaultBranch) {
+      Value *Cmp = Builder.CreateICmpULT(
+          TableIndex, ConstantInt::get(MinCaseVal->getType(), TableSize));
+      RangeCheckBranch =
+          Builder.CreateCondBr(Cmp, LookupBB, SI->getDefaultDest());
+    } else
+      RangeCheckBranch = Builder.CreateBr(LookupBB);
+
     if (DTU)
       Updates.push_back({DominatorTree::Insert, BB, LookupBB});
   }
diff --git a/llvm/test/Transforms/SimplifyCFG/switch_mask.ll b/llvm/test/Transforms/SimplifyCFG/switch_mask.ll
new file mode 100644
index 000000000000000..454814ec7070d42
--- /dev/null
+++ b/llvm/test/Transforms/SimplifyCFG/switch_mask.ll
@@ -0,0 +1,57 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -passes=simplifycfg --switch-to-lookup -S < %s | FileCheck %s
+
+target datalayout = "e-m:e-i64:64-f80:128-n8:16:32:64-S128"
+
+; https://alive2.llvm.org/ce/z/tuxLhJ
+define i1 @switch_lookup_with_small_mask(i64 %x) {
+; CHECK-LABEL: @switch_lookup_with_small_mask(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[AND:%.*]] = and i64 [[X:%.*]], 15
+; CHECK-NEXT:    [[SWITCH_CAST:%.*]] = trunc i64 [[AND]] to i16
+; CHECK-NEXT:    [[SWITCH_SHIFTAMT:%.*]] = mul nuw nsw i16 [[SWITCH_CAST]], 1
+; CHECK-NEXT:    [[SWITCH_DOWNSHIFT:%.*]] = lshr i16 1030, [[SWITCH_SHIFTAMT]]
+; CHECK-NEXT:    [[SWITCH_MASKED:%.*]] = trunc i16 [[SWITCH_DOWNSHIFT]] to i1
+; CHECK-NEXT:    ret i1 [[SWITCH_MASKED]]
+;
+entry:
+  %and = and i64 %x, 15
+  switch i64 %and, label %default [
+  i64 10, label %lor.end
+  i64 1, label %lor.end
+  i64 2, label %lor.end
+  ]
+
+default:                                          ; preds = %entry
+  br label %lor.end
+
+lor.end:                                          ; preds = %entry, %entry, %entry, %default
+  %0 = phi i1 [ true, %entry ], [ false, %default ], [ true, %entry ], [ true, %entry ]
+  ret i1 %0
+}
+
+define i1 @switch_lookup_with_small_urem(i64 %x) {
+; CHECK-LABEL: @switch_lookup_with_small_urem(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[AND:%.*]] = urem i64 [[X:%.*]], 42
+; CHECK-NEXT:    [[SWITCH_CAST:%.*]] = trunc i64 [[AND]] to i42
+; CHECK-NEXT:    [[SWITCH_SHIFTAMT:%.*]] = mul nuw nsw i42 [[SWITCH_CAST]], 1
+; CHECK-NEXT:    [[SWITCH_DOWNSHIFT:%.*]] = lshr i42 1030, [[SWITCH_SHIFTAMT]]
+; CHECK-NEXT:    [[SWITCH_MASKED:%.*]] = trunc i42 [[SWITCH_DOWNSHIFT]] to i1
+; CHECK-NEXT:    ret i1 [[SWITCH_MASKED]]
+;
+entry:
+  %and = urem i64 %x, 42
+  switch i64 %and, label %default [
+  i64 10, label %lor.end
+  i64 1, label %lor.end
+  i64 2, label %lor.end
+  ]
+
+default:                                          ; preds = %entry
+  br label %lor.end
+
+lor.end:                                          ; preds = %entry, %entry, %entry, %default
+  %0 = phi i1 [ true, %entry ], [ false, %default ], [ true, %entry ], [ true, %entry ]
+  ret i1 %0
+}

>From 9c120c10dc6b60a23c22280de015c26e1316eb5e Mon Sep 17 00:00:00 2001
From: Zhongyunde <zhongyunde at huawei.com>
Date: Mon, 18 Sep 2023 15:53:48 +0800
Subject: [PATCH 2/2] Fix UB when 1 << CR.getActiveBits() overflow

---
 llvm/lib/Transforms/Utils/SimplifyCFG.cpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
index 90e10a7a75067da..33b6e12e807f163 100644
--- a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
+++ b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
@@ -6556,7 +6556,7 @@ static bool SwitchToLookupTable(SwitchInst *SI, IRBuilder<> &Builder,
     if (UseSwitchConditionAsTableIndex) {
       ConstantRange CR = computeConstantRange(TableIndex, /* ForSigned */ false,
                                               /* UseInstrInfo*/ true);
-      if (DL.fitsInLegalInteger(1 << CR.getActiveBits())) {
+      if (DL.fitsInLegalInteger(CR.getUpper().getZExtValue())) {
         // The default branch is unreachable when we enlarge the lookup table.
         TableSize = CR.getUpper().getZExtValue();
         KeepDefaultBranch = false;