[InstCombine] Canonicalize switch(X^C) expressions to switch(X)

[antoniofrighetto]

switch(X^C) expressions can be folded to switch(X). Minor opportunity to generalize simplifications in visitSwitchInst via an inverse function helper as well.

Proof: https://alive2.llvm.org/ce/z/TMRy_3.

[llvmbot]

@llvm/pr-subscribers-llvm-transforms

Author: Antonio Frighetto (antoniofrighetto)

Changes

switch(X^C) expressions can be folded to switch(X).

Proof: https://alive2.llvm.org/ce/z/TMRy_3.

---

Full diff: https://github.com/llvm/llvm-project/pull/143677.diff

3 Files Affected:

• (modified) llvm/lib/Transforms/InstCombine/InstructionCombining.cpp (+12)
• (modified) llvm/test/Transforms/InstCombine/narrow-switch.ll (+3-3)
• (added) llvm/test/Transforms/InstCombine/switch-xor.ll (+59)

diff --git a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
index e261807bbc035..3969030a0ad5a 100644
--- a/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstructionCombining.cpp
@@ -3948,6 +3948,18 @@ Instruction *InstCombinerImpl::visitSwitchInst(SwitchInst &SI) {
     }
   }
 
+  ConstantInt *XorRHS;
+  if (match(Cond, m_Xor(m_Value(Op0), m_ConstantInt(XorRHS)))) {
+    // Fold 'switch (X^C) case A' into 'switch (X) case A^C'.
+    for (auto &Case : SI.cases()) {
+      Constant *NewCase = ConstantExpr::getXor(Case.getCaseValue(), XorRHS);
+      assert(isa<ConstantInt>(NewCase) &&
+             "Result of expression should be constant");
+      Case.setValue(cast<ConstantInt>(NewCase));
+    }
+    return replaceOperand(SI, 0, Op0);
+  }
+
   // Fold switch(select cond, X, Y) into switch(X/Y) if possible
   if (auto *Select = dyn_cast<SelectInst>(Cond)) {
     if (Value *V =
diff --git a/llvm/test/Transforms/InstCombine/narrow-switch.ll b/llvm/test/Transforms/InstCombine/narrow-switch.ll
index 05a30b910e5ee..7d2d3ee94d49b 100644
--- a/llvm/test/Transforms/InstCombine/narrow-switch.ll
+++ b/llvm/test/Transforms/InstCombine/narrow-switch.ll
@@ -171,9 +171,9 @@ case124:
 define i32 @trunc32to16(i32 %a0) #0 {
 ; ALL-LABEL: @trunc32to16(
 ; ALL:         switch i16
-; ALL-NEXT:    i16 63, label %sw.bb
-; ALL-NEXT:    i16 1, label %sw.bb1
-; ALL-NEXT:    i16 100, label %sw.bb2
+; ALL-NEXT:    i16 15767, label %sw.bb
+; ALL-NEXT:    i16 15785, label %sw.bb1
+; ALL-NEXT:    i16 15820, label %sw.bb2
 ; ALL-NEXT:    ]
 ;
 entry:
diff --git a/llvm/test/Transforms/InstCombine/switch-xor.ll b/llvm/test/Transforms/InstCombine/switch-xor.ll
new file mode 100644
index 0000000000000..a7b65e406dfa2
--- /dev/null
+++ b/llvm/test/Transforms/InstCombine/switch-xor.ll
@@ -0,0 +1,59 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt < %s -passes=instcombine -S | FileCheck %s
+
+define i1 @test_switch_with_xor(i32 %x) {
+; CHECK-LABEL: define i1 @test_switch_with_xor(
+; CHECK-SAME: i32 [[X:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    switch i32 [[X]], label %[[SW_DEFAULT:.*]] [
+; CHECK-NEXT:      i32 3, label %[[SW_BB:.*]]
+; CHECK-NEXT:      i32 0, label %[[SW_BB]]
+; CHECK-NEXT:      i32 1, label %[[SW_BB]]
+; CHECK-NEXT:    ]
+; CHECK:       [[SW_BB]]:
+; CHECK-NEXT:    ret i1 true
+; CHECK:       [[SW_DEFAULT]]:
+; CHECK-NEXT:    ret i1 false
+;
+entry:
+  %xor = xor i32 %x, 2
+  switch i32 %xor, label %sw.default [
+  i32 1, label %sw.bb
+  i32 2, label %sw.bb
+  i32 3, label %sw.bb
+  ]
+
+sw.bb:
+  ret i1 true
+sw.default:
+  ret i1 false
+}
+
+define i1 @test_switch_with_xor_nonconstant_ops(i32 %x, i32 %y) {
+; CHECK-LABEL: define i1 @test_switch_with_xor_nonconstant_ops(
+; CHECK-SAME: i32 [[X:%.*]], i32 [[Y:%.*]]) {
+; CHECK-NEXT:  [[ENTRY:.*:]]
+; CHECK-NEXT:    [[XOR:%.*]] = xor i32 [[X]], [[Y]]
+; CHECK-NEXT:    switch i32 [[XOR]], label %[[SW_DEFAULT:.*]] [
+; CHECK-NEXT:      i32 1, label %[[SW_BB:.*]]
+; CHECK-NEXT:      i32 2, label %[[SW_BB]]
+; CHECK-NEXT:      i32 3, label %[[SW_BB]]
+; CHECK-NEXT:    ]
+; CHECK:       [[SW_BB]]:
+; CHECK-NEXT:    ret i1 true
+; CHECK:       [[SW_DEFAULT]]:
+; CHECK-NEXT:    ret i1 false
+;
+entry:
+  %xor = xor i32 %x, %y
+  switch i32 %xor, label %sw.default [
+  i32 1, label %sw.bb
+  i32 2, label %sw.bb
+  i32 3, label %sw.bb
+  ]
+
+sw.bb:
+  ret i1 true
+sw.default:
+  ret i1 false
+}

[nikic]

Can you please pre-commit a regeneration of this test file?

[nikic]

This is the third copy of this code... may be worthwhile to generalize via InverseFunction or something like that?

[antoniofrighetto]

Generalized this, thanks, hope it aligns with what you had in mind.

[nikic]

From llvm-opt-benchmark, a potential problem is that this can something turn small switch values into very large ones -- the common case is a xor by INT_MIN. I haven't checked how this affects codegen.

[antoniofrighetto]

From llvm-opt-benchmark, a potential problem is that this can something turn small switch values into very large ones -- the common case is a xor by INT_MIN. I haven't checked how this affects codegen.

When not invoking the middle-end (except for SimplifyCFG, was not expecting llc to invoke it), the codegen looks slightly worse both in the optimized and unoptimized case: https://llvm.godbolt.org/z/jceqnev4n. Not sure if we wish to proceed further with the canonicalization :(

[dtcxzyw]

From llvm-opt-benchmark, a potential problem is that this can something turn small switch values into very large ones -- the common case is a xor by INT_MIN. I haven't checked how this affects codegen.

When not invoking the middle-end (except for SimplifyCFG, was not expecting llc to invoke it), the codegen looks slightly worse both in the optimized and unoptimized case: https://llvm.godbolt.org/z/jceqnev4n. Not sure if we wish to proceed further with the canonicalization :(

I think it is ok to perform this canonicalization as it doesn't break the density of switch cases (i.e., reverting SimplifyCFG transforms). As for the regression, it can be fixed by canonicalizing sub X, INT_MIN -> xor X, INT_MIN: https://alive2.llvm.org/ce/z/UGblFI

[dtcxzyw]

Do we need to check whether the condition is only used by the switch? Absorbing constants into switch cases may not be profitable if the condition is also used in other places...

[antoniofrighetto]

Do we need to check whether the condition is only used by the switch? Absorbing constants into switch cases may not be profitable if the condition is also used in other places...

Sounds reasonable to me.

[antoniofrighetto]

@nikic Think we should proceed with this canonicalization? If so, should try xor X, INT_MIN -> sub X, INT_MIN first?

[antoniofrighetto]

From llvm-opt-benchmark, a potential problem is that this can something turn small switch values into very large ones -- the common case is a xor by INT_MIN. I haven't checked how this affects codegen.

When not invoking the middle-end (except for SimplifyCFG, was not expecting llc to invoke it), the codegen looks slightly worse both in the optimized and unoptimized case: https://llvm.godbolt.org/z/jceqnev4n. Not sure if we wish to proceed further with the canonicalization :(

I think it is ok to perform this canonicalization as it doesn't break the density of switch cases (i.e., reverting SimplifyCFG transforms). As for the regression, it can be fixed by canonicalizing sub X, INT_MIN -> xor X, INT_MIN: https://alive2.llvm.org/ce/z/UGblFI

It turns out it may be particularly hard to proceed with xor X, INT_MIN -> sub X, INT_MIN fold, as there are a lot of transforms that lean on the inverse (current) canonicalization. Conveniently (though maybe not that elegantly), we may prevent this by checking if the constant value is at the extremes (tests updated). May this look better?

[dtcxzyw]

LG