[llvm] [InstCombine] Fold `switch(rol(x, C1)) case C2:` to `switch(x) case rol(C2, -C1):` (PR #86307)

[via llvm-commits]

https://github.com/YanWQ-monad updated https://github.com/llvm/llvm-project/pull/86307

>From 87a72040a928d602a06040f6090295639c10125a Mon Sep 17 00:00:00 2001
From: YanWQ-monad <YanWQmonad at gmail.com>
Date: Sat, 23 Mar 2024 00:47:03 +0800
Subject: [PATCH 1/2] Pre-commit: add test

---
 .../test/Transforms/InstCombine/switch-rol.ll | 66 +++++++++++++++++++
 1 file changed, 66 insertions(+)
 create mode 100644 llvm/test/Transforms/InstCombine/switch-rol.ll

diff --git a/llvm/test/Transforms/InstCombine/switch-rol.ll b/llvm/test/Transforms/InstCombine/switch-rol.ll
new file mode 100644
index 00000000000000..cd0d066e3e0ce3
--- /dev/null
+++ b/llvm/test/Transforms/InstCombine/switch-rol.ll
@@ -0,0 +1,66 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 4
+; RUN: opt < %s -passes=instcombine -S | FileCheck %s
+
+declare void @dummy()
+
+define i32 @switch_rol(i32 %a) #0 {
+; CHECK-LABEL: define i32 @switch_rol(
+; CHECK-SAME: i32 [[A:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[ROL:%.*]] = call i32 @llvm.fshl.i32(i32 [[A]], i32 [[A]], i32 30)
+; CHECK-NEXT:    switch i32 [[ROL]], label [[DEFAULT:%.*]] [
+; CHECK-NEXT:      i32 0, label [[TRAP_EXIT:%.*]]
+; CHECK-NEXT:      i32 5, label [[TRAP_EXIT]]
+; CHECK-NEXT:    ]
+; CHECK:       default:
+; CHECK-NEXT:    call void @dummy()
+; CHECK-NEXT:    br label [[TRAP_EXIT]]
+; CHECK:       trap.exit:
+; CHECK-NEXT:    ret i32 0
+;
+entry:
+  %rol = call i32 @llvm.fshl.i32(i32 %a, i32 %a, i32 30)
+  switch i32 %rol, label %default [
+  i32 0, label %trap.exit
+  i32 5, label %trap.exit
+  ]
+
+default:
+  call void @dummy()
+  br label %trap.exit
+
+trap.exit:
+  ret i32 0
+}
+
+define i32 @switch_rol_2(i32 %a) #0 {
+; CHECK-LABEL: define i32 @switch_rol_2(
+; CHECK-SAME: i32 [[A:%.*]]) {
+; CHECK-NEXT:  entry:
+; CHECK-NEXT:    [[TMP0:%.*]] = add i32 [[A]], -7
+; CHECK-NEXT:    [[ROL:%.*]] = call i32 @llvm.fshl.i32(i32 [[TMP0]], i32 [[TMP0]], i32 30)
+; CHECK-NEXT:    switch i32 [[ROL]], label [[DEFAULT:%.*]] [
+; CHECK-NEXT:      i32 0, label [[TRAP_EXIT:%.*]]
+; CHECK-NEXT:      i32 5, label [[TRAP_EXIT]]
+; CHECK-NEXT:    ]
+; CHECK:       default:
+; CHECK-NEXT:    call void @dummy()
+; CHECK-NEXT:    br label [[TRAP_EXIT]]
+; CHECK:       trap.exit:
+; CHECK-NEXT:    ret i32 0
+;
+entry:
+  %1 = sub i32 %a, 7
+  %rol = call i32 @llvm.fshl.i32(i32 %1, i32 %1, i32 30)
+  switch i32 %rol, label %default [
+  i32 0, label %trap.exit
+  i32 5, label %trap.exit
+  ]
+
+default:
+  call void @dummy()
+  br label %trap.exit
+
+trap.exit:
+  ret i32 0
+}

>From e52ac48b2563ad5d1c41b5f4f6ad5861345cbf6d Mon Sep 17 00:00:00 2001
From: YanWQ-monad <YanWQmonad at gmail.com>
Date: Sat, 23 Mar 2024 00:47:45 +0800
Subject: [PATCH 2/2] Fold 'switch(rol(x, C1))'

---
 .../Transforms/InstCombine/InstructionCombining.cpp | 10 ++++++++++
 llvm/test/Transforms/InstCombine/switch-rol.ll      | 13 +++++--------
 2 files changed, 15 insertions(+), 8 deletions(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
index 7c40fb4fc86082..b6611cfbbfc1f4 100644
--- a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
@@ -3645,6 +3645,16 @@ Instruction *InstCombinerImpl::visitSwitchInst(SwitchInst &SI) {
     }
   }
 
+  // Fold 'switch(rol(x, C1)) case C2:' to 'switch(x) case rol(C2, -C1):'
+  if (match(Cond,
+            m_FShl(m_Value(Op0), m_Deferred(Op0), m_ConstantInt(ShiftAmt)))) {
+    for (auto &Case : SI.cases()) {
+      const APInt NewCase = Case.getCaseValue()->getValue().rotr(ShiftAmt);
+      Case.setValue(ConstantInt::get(SI.getContext(), NewCase));
+    }
+    return replaceOperand(SI, 0, Op0);
+  }
+
   KnownBits Known = computeKnownBits(Cond, 0, &SI);
   unsigned LeadingKnownZeros = Known.countMinLeadingZeros();
   unsigned LeadingKnownOnes = Known.countMinLeadingOnes();
diff --git a/llvm/test/Transforms/InstCombine/switch-rol.ll b/llvm/test/Transforms/InstCombine/switch-rol.ll
index cd0d066e3e0ce3..7557ad5dfe276e 100644
--- a/llvm/test/Transforms/InstCombine/switch-rol.ll
+++ b/llvm/test/Transforms/InstCombine/switch-rol.ll
@@ -7,10 +7,9 @@ define i32 @switch_rol(i32 %a) #0 {
 ; CHECK-LABEL: define i32 @switch_rol(
 ; CHECK-SAME: i32 [[A:%.*]]) {
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[ROL:%.*]] = call i32 @llvm.fshl.i32(i32 [[A]], i32 [[A]], i32 30)
-; CHECK-NEXT:    switch i32 [[ROL]], label [[DEFAULT:%.*]] [
+; CHECK-NEXT:    switch i32 [[A]], label [[DEFAULT:%.*]] [
 ; CHECK-NEXT:      i32 0, label [[TRAP_EXIT:%.*]]
-; CHECK-NEXT:      i32 5, label [[TRAP_EXIT]]
+; CHECK-NEXT:      i32 20, label [[TRAP_EXIT]]
 ; CHECK-NEXT:    ]
 ; CHECK:       default:
 ; CHECK-NEXT:    call void @dummy()
@@ -37,11 +36,9 @@ define i32 @switch_rol_2(i32 %a) #0 {
 ; CHECK-LABEL: define i32 @switch_rol_2(
 ; CHECK-SAME: i32 [[A:%.*]]) {
 ; CHECK-NEXT:  entry:
-; CHECK-NEXT:    [[TMP0:%.*]] = add i32 [[A]], -7
-; CHECK-NEXT:    [[ROL:%.*]] = call i32 @llvm.fshl.i32(i32 [[TMP0]], i32 [[TMP0]], i32 30)
-; CHECK-NEXT:    switch i32 [[ROL]], label [[DEFAULT:%.*]] [
-; CHECK-NEXT:      i32 0, label [[TRAP_EXIT:%.*]]
-; CHECK-NEXT:      i32 5, label [[TRAP_EXIT]]
+; CHECK-NEXT:    switch i32 [[A]], label [[DEFAULT:%.*]] [
+; CHECK-NEXT:      i32 7, label [[TRAP_EXIT:%.*]]
+; CHECK-NEXT:      i32 27, label [[TRAP_EXIT]]
 ; CHECK-NEXT:    ]
 ; CHECK:       default:
 ; CHECK-NEXT:    call void @dummy()