[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 4763fbec1089a1b4fee5f67bc4ac63817efd536a Mon Sep 17 00:00:00 2001
From: Zhongyunde <zhongyunde at huawei.com>
Date: Sat, 9 Sep 2023 13:18:34 +0800
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     | 14 ++++++++
 .../Transforms/SimplifyCFG/switch_mask.ll     | 33 +++++++++++++++++++
 2 files changed, 47 insertions(+)
 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..ee04b881bb0292d 100644
--- a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
+++ b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
@@ -6552,6 +6552,20 @@ 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 {
+    const APInt *MaskC;
+    if (UseSwitchConditionAsTableIndex &&
+        match(TableIndex, m_And(m_Value(), m_APInt(MaskC)))) {
+      APInt SCVPlus1 = *MaskC + 1;
+      const uint64_t *RawData = MaskC->getRawData();
+      // Expand the size of the table when the MaskC < 64.
+      if (SCVPlus1.isPowerOf2() && SCVPlus1.logBase2() <= 6 && RawData) {
+        assert(TableSize <= RawData[0] && "unexpect mask value");
+        // The range check will be deleted later when we enlarge the lookup
+        // table because the check always return true.
+        TableSize = RawData[0] + 1;
+      }
+    }
+
     Value *Cmp = Builder.CreateICmpULT(
         TableIndex, ConstantInt::get(MinCaseVal->getType(), TableSize));
     RangeCheckBranch =
diff --git a/llvm/test/Transforms/SimplifyCFG/switch_mask.ll b/llvm/test/Transforms/SimplifyCFG/switch_mask.ll
new file mode 100644
index 000000000000000..393448359b07cf5
--- /dev/null
+++ b/llvm/test/Transforms/SimplifyCFG/switch_mask.ll
@@ -0,0 +1,33 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --scrub-attributes
+; 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/2bNvhg
+define noundef i1 @switch_lookup_with_small_mask(i64 noundef %x) {
+; CHECK-LABEL: @switch_lookup_with_small_mask(
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[AND:%.*]] = and i64 [[X:%.*]], 15
+; CHECK-NEXT:    [[TMP0:%.*]] = icmp ult i64 [[AND]], 16
+; 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:    [[TMP1:%.*]] = select i1 [[TMP0]], i1 [[SWITCH_MASKED]], i1 false
+; CHECK-NEXT:    ret i1 [[TMP1]]
+;
+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
+}

>From ccf797bb49f7999206bfdf8b960545ba9f5e39e6 Mon Sep 17 00:00:00 2001
From: Zhongyunde <zhongyunde at huawei.com>
Date: Sat, 9 Sep 2023 21:16:04 +0800
Subject: [PATCH 2/2] Fix comment, use SCVPlus1.ule(64)

---
 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 ee04b881bb0292d..571f30ffc08a522 100644
--- a/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
+++ b/llvm/lib/Transforms/Utils/SimplifyCFG.cpp
@@ -6558,7 +6558,7 @@ static bool SwitchToLookupTable(SwitchInst *SI, IRBuilder<> &Builder,
       APInt SCVPlus1 = *MaskC + 1;
       const uint64_t *RawData = MaskC->getRawData();
       // Expand the size of the table when the MaskC < 64.
-      if (SCVPlus1.isPowerOf2() && SCVPlus1.logBase2() <= 6 && RawData) {
+      if (SCVPlus1.isPowerOf2() && SCVPlus1.ule(64) && RawData) {
         assert(TableSize <= RawData[0] && "unexpect mask value");
         // The range check will be deleted later when we enlarge the lookup
         // table because the check always return true.