From ce5b88bf10ac5841ca0691ef3f2d87e62748dcd7 Mon Sep 17 00:00:00 2001
From: Antonio Frighetto <me@antoniofrighetto.com>
Date: Sat, 16 Sep 2023 12:11:10 +0200
Subject: [PATCH] [InstCombine] Handle constant arms in `select` of `srem` fold

Extend folding for `2^n` euclidean division remainder operations
on signed integers by handling the specific instance in which one
`select` arm has already been replaced by 1.

Reported-By: HypheX

Fixes: https://github.com/llvm/llvm-project/issues/66417.
---
 .../InstCombine/InstCombineSelect.cpp         | 31 +++++++++++++------
 .../Transforms/InstCombine/select-divrem.ll   | 22 +++++++++++++
 2 files changed, 44 insertions(+), 9 deletions(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
index e2965e7a5703..a62bfd2e89c4 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
@@ -2616,20 +2616,33 @@ static Instruction *foldSelectWithSRem(SelectInst &SI, InstCombinerImpl &IC,
   if (!TrueIfSigned)
     std::swap(TrueVal, FalseVal);
 
-  // We are matching a quite specific pattern here:
+  auto FoldToBitwiseAnd = [&](Value *Remainder) -> Instruction * {
+    Value *Add = Builder.CreateAdd(
+        Remainder, Constant::getAllOnesValue(RemRes->getType()));
+    return BinaryOperator::CreateAnd(Op, Add);
+  };
+
+  // Match the general case:
   // %rem = srem i32 %x, %n
   // %cnd = icmp slt i32 %rem, 0
   // %add = add i32 %rem, %n
   // %sel = select i1 %cnd, i32 %add, i32 %rem
-  if (!(match(TrueVal, m_Add(m_Value(RemRes), m_Value(Remainder))) &&
-        match(RemRes, m_SRem(m_Value(Op), m_Specific(Remainder))) &&
-        IC.isKnownToBeAPowerOfTwo(Remainder, /*OrZero*/ true) &&
-        FalseVal == RemRes))
-    return nullptr;
+  if (match(TrueVal, m_Add(m_Value(RemRes), m_Value(Remainder))) &&
+      match(RemRes, m_SRem(m_Value(Op), m_Specific(Remainder))) &&
+      IC.isKnownToBeAPowerOfTwo(Remainder, /*OrZero*/ true) &&
+      FalseVal == RemRes)
+    return FoldToBitwiseAnd(Remainder);
 
-  Value *Add = Builder.CreateAdd(Remainder,
-                                 Constant::getAllOnesValue(RemRes->getType()));
-  return BinaryOperator::CreateAnd(Op, Add);
+  // Match the case where the one arm has been replaced by constant 1:
+  // %rem = srem i32 %n, 2
+  // %cnd = icmp slt i32 %rem, 0
+  // %sel = select i1 %cnd, i32 1, i32 %rem
+  if (match(TrueVal, m_One()) &&
+      match(RemRes, m_SRem(m_Value(Op), m_SpecificInt(2))) &&
+      FalseVal == RemRes)
+    return FoldToBitwiseAnd(ConstantInt::get(RemRes->getType(), 2));
+
+  return nullptr;
 }
 
 static Value *foldSelectWithFrozenICmp(SelectInst &Sel, InstCombiner::BuilderTy &Builder) {
diff --git a/llvm/test/Transforms/InstCombine/select-divrem.ll b/llvm/test/Transforms/InstCombine/select-divrem.ll
index a5b56609d606..f007c53359ca 100644
--- a/llvm/test/Transforms/InstCombine/select-divrem.ll
+++ b/llvm/test/Transforms/InstCombine/select-divrem.ll
@@ -321,3 +321,25 @@ define i8 @rem_euclid_non_const_pow2(i8 %0, i8 %1) {
   %sel = select i1 %cond, i8 %add, i8 %rem
   ret i8 %sel
 }
+
+define i32 @rem_euclid_pow2_true_arm_folded(i32 %n) {
+; CHECK-LABEL: @rem_euclid_pow2_true_arm_folded(
+; CHECK-NEXT:    [[RES:%.*]] = and i32 [[N:%.*]], 1
+; CHECK-NEXT:    ret i32 [[RES]]
+;
+  %rem = srem i32 %n, 2
+  %neg = icmp slt i32 %rem, 0
+  %res = select i1 %neg, i32 1, i32 %rem
+  ret i32 %res
+}
+
+define i32 @rem_euclid_pow2_false_arm_folded(i32 %n) {
+; CHECK-LABEL: @rem_euclid_pow2_false_arm_folded(
+; CHECK-NEXT:    [[RES:%.*]] = and i32 [[N:%.*]], 1
+; CHECK-NEXT:    ret i32 [[RES]]
+;
+  %rem = srem i32 %n, 2
+  %nonneg = icmp sge i32 %rem, 0
+  %res = select i1 %nonneg, i32 %rem, i32 1
+  ret i32 %res
+}