[PATCH] D148420: [InstCombine] Enhance select icmp and folding

[Peixin Qiao via Phabricator via llvm-commits]

peixin created this revision.
peixin added a project: LLVM.
Herald added a subscriber: hiraditya.
Herald added a project: All.
peixin requested review of this revision.
Herald added a subscriber: llvm-commits.

This folds (a << k) ? 2^k * a : 0 to 2^k * a.

Fix #62155.


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D148420

Files:
  llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
  llvm/test/Transforms/InstCombine/select-icmp-and.ll


Index: llvm/test/Transforms/InstCombine/select-icmp-and.ll
===================================================================
--- llvm/test/Transforms/InstCombine/select-icmp-and.ll
+++ llvm/test/Transforms/InstCombine/select-icmp-and.ll
@@ -618,3 +618,16 @@
   ret i8 %t3
 }
 
+; (x << k) ? 2^k * x : 0 --> 2^k * x
+
+define i32 @test_select_icmp_and_shl(i32 %x) {
+; CHECK-LABEL: @test_select_icmp_and_shl(
+; CHECK-NEXT:    [[T:%.*]] = shl i32 [[X:%.*]], 2
+; CHECK-NEXT:    ret i32 [[T]]
+;
+  %shl.mask = and i32 %x, 1073741823
+  %tobool.not = icmp eq i32 %shl.mask, 0
+  %mul = shl i32 %x, 2
+  %cond = select i1 %tobool.not, i32 0, i32 %mul
+  ret i32 %cond
+}
Index: llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
===================================================================
--- llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
+++ llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
@@ -227,6 +227,42 @@
   return V;
 }
 
+/// This folds:
+///  select (icmp eq (and X, C), 0), 0, (shl X, K)
+///    iff C is one number like "0..01..1" and the number of leading zeros is
+///    equal to K.
+/// To the following:
+///  shl X, K
+static Value *foldSelectICmpAndZeroShl(SelectInst &Sel, ICmpInst *Cmp,
+                                       InstCombiner::BuilderTy &Builder) {
+  // TODO: vector select or vector compare.
+  if (Sel.getType()->isVectorTy() || Cmp->getType()->isVectorTy())
+    return nullptr;
+
+  const APInt *K;
+  Value *X;
+  if (!match(Sel.getTrueValue(), m_Zero()) ||
+      !match(Sel.getFalseValue(), m_Shl(m_Value(X), m_APInt(K))) ||
+      !ICmpInst::isEquality(Cmp->getPredicate()) ||
+      !match(Cmp->getOperand(1), m_Zero()))
+    return nullptr;
+
+  Value *V;
+  const APInt *C;
+  if (!match(Cmp->getOperand(0), m_And(m_Value(V), m_APInt(C))))
+    return nullptr;
+
+  if (X != V)
+    return nullptr;
+
+  uint32_t BitWidth = V->getType()->getScalarSizeInBits();
+  if (C->countLeadingZeros() + C->countTrailingOnes() != BitWidth ||
+      (int64_t)C->countLeadingZeros() != K->getSExtValue())
+    return nullptr;
+
+  return Sel.getFalseValue();
+}
+
 /// We want to turn code that looks like this:
 ///   %C = or %A, %B
 ///   %D = select %cond, %C, %A
@@ -1722,6 +1758,9 @@
   if (Value *V = foldSelectICmpAnd(SI, ICI, Builder))
     return replaceInstUsesWith(SI, V);
 
+  if (Value *V = foldSelectICmpAndZeroShl(SI, ICI, Builder))
+    return replaceInstUsesWith(SI, V);
+
   // NOTE: if we wanted to, this is where to detect integer MIN/MAX
   Value *TrueVal = SI.getTrueValue();
   Value *FalseVal = SI.getFalseValue();


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D148420.513865.patch
Type: text/x-patch
Size: 2597 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20230415/ccbcf652/attachment.bin>