[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 created https://github.com/llvm/llvm-project/pull/86307

This solves #86161.

It is worth mentioning that, as @dtcxzyw pointed out, there is an inverse fold in
https://github.com/llvm/llvm-project/blob/90454a609894ab278a87be2b9f5c49714caba8df/llvm/lib/Transforms/Utils/SimplifyCFG.cpp#L6911-L6996
, so IR may be flickering in the optimization pipeline.
`SimplifyCFG` will have the upper hand, so there won't be any major problems. But in some rare cases, it might cause a regression (in @dtcxzyw's opt-benchmark).

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

---
 .../test/Transforms/InstCombine/switch-rol.ll | 34 +++++++++++++++++++
 1 file changed, 34 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..ee8c278bca2397
--- /dev/null
+++ b/llvm/test/Transforms/InstCombine/switch-rol.ll
@@ -0,0 +1,34 @@
+; 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
+}

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

---
 .../Transforms/InstCombine/InstructionCombining.cpp    | 10 ++++++++++
 llvm/test/Transforms/InstCombine/switch-rol.ll         |  5 ++---
 2 files changed, 12 insertions(+), 3 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 ee8c278bca2397..1cd55ff91c9492 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()