From 53eede597e8423410ec48b7c335c76f3b6c4c714 Mon Sep 17 00:00:00 2001
From: Sanjay Patel <spatel@rotateright.com>
Date: Mon, 12 Sep 2022 15:00:08 -0400
Subject: [PATCH] [InstCombine] look through 'not' of ctlz/cttz op with
 0-is-undef

https://alive2.llvm.org/ce/z/MNsC1S

This pattern was flagged at:
https://discourse.llvm.org/t/instcombines-select-optimizations-dont-trigger-reliably/64927
---
 .../Transforms/InstCombine/InstCombineSelect.cpp   | 15 ++++--
 .../Transforms/InstCombine/select-cmp-cttz-ctlz.ll | 63 ++++++++++++++++++++--
 2 files changed, 69 insertions(+), 9 deletions(-)

diff --git a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
index 98cf450..fc1a435 100644
--- a/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
+++ b/llvm/lib/Transforms/InstCombine/InstCombineSelect.cpp
@@ -952,8 +952,8 @@ static Value *foldSelectCttzCtlz(ICmpInst *ICI, Value *TrueVal, Value *FalseVal,
   Value *CmpLHS = ICI->getOperand(0);
   Value *CmpRHS = ICI->getOperand(1);
 
-  // Check if the condition value compares a value for equality against zero.
-  if (!ICI->isEquality() || !match(CmpRHS, m_Zero()))
+  // Check if the select condition compares a value for equality.
+  if (!ICI->isEquality())
     return nullptr;
 
   Value *SelectArg = FalseVal;
@@ -969,8 +969,15 @@ static Value *foldSelectCttzCtlz(ICmpInst *ICI, Value *TrueVal, Value *FalseVal,
 
   // Check that 'Count' is a call to intrinsic cttz/ctlz. Also check that the
   // input to the cttz/ctlz is used as LHS for the compare instruction.
-  if (!match(Count, m_Intrinsic<Intrinsic::cttz>(m_Specific(CmpLHS))) &&
-      !match(Count, m_Intrinsic<Intrinsic::ctlz>(m_Specific(CmpLHS))))
+  Value *X;
+  if (!match(Count, m_Intrinsic<Intrinsic::cttz>(m_Value(X))) &&
+      !match(Count, m_Intrinsic<Intrinsic::ctlz>(m_Value(X))))
+    return nullptr;
+
+  // (X == 0) ? BitWidth : ctz(X)
+  // (X == -1) ? BitWidth : ctz(~X)
+  if ((X != CmpLHS || !match(CmpRHS, m_Zero())) &&
+      (!match(X, m_Not(m_Specific(CmpLHS))) || !match(CmpRHS, m_AllOnes())))
     return nullptr;
 
   IntrinsicInst *II = cast<IntrinsicInst>(Count);
diff --git a/llvm/test/Transforms/InstCombine/select-cmp-cttz-ctlz.ll b/llvm/test/Transforms/InstCombine/select-cmp-cttz-ctlz.ll
index 8fe2a00..4dff0cc 100644
--- a/llvm/test/Transforms/InstCombine/select-cmp-cttz-ctlz.ll
+++ b/llvm/test/Transforms/InstCombine/select-cmp-cttz-ctlz.ll
@@ -287,13 +287,45 @@ define i32 @test5d(i64 %x) {
 define i32 @not_op_ctlz(i64 %x) {
 ; CHECK-LABEL: @not_op_ctlz(
 ; CHECK-NEXT:    [[N:%.*]] = xor i64 [[X:%.*]], -1
+; CHECK-NEXT:    [[CT:%.*]] = tail call i64 @llvm.ctlz.i64(i64 [[N]], i1 false), !range [[RNG2]]
+; CHECK-NEXT:    [[CAST:%.*]] = trunc i64 [[CT]] to i32
+; CHECK-NEXT:    ret i32 [[CAST]]
+;
+  %n = xor i64 %x, -1
+  %ct = tail call i64 @llvm.ctlz.i64(i64 %n, i1 true)
+  %cast = trunc i64 %ct to i32
+  %tobool = icmp eq i64 %x, -1
+  %r = select i1 %tobool, i32 64, i32 %cast
+  ret i32 %r
+}
+
+define i32 @not_op_cttz(i64 %x) {
+; CHECK-LABEL: @not_op_cttz(
+; CHECK-NEXT:    [[N:%.*]] = xor i64 [[X:%.*]], -1
+; CHECK-NEXT:    [[CT:%.*]] = tail call i64 @llvm.cttz.i64(i64 [[N]], i1 false), !range [[RNG2]]
+; CHECK-NEXT:    [[CAST:%.*]] = trunc i64 [[CT]] to i32
+; CHECK-NEXT:    ret i32 [[CAST]]
+;
+  %n = xor i64 %x, -1
+  %ct = tail call i64 @llvm.cttz.i64(i64 %n, i1 true)
+  %cast = trunc i64 %ct to i32
+  %tobool = icmp eq i64 %x, -1
+  %r = select i1 %tobool, i32 64, i32 %cast
+  ret i32 %r
+}
+
+; negative test
+
+define i32 @not_op_ctlz_wrong_xor_op1(i64 %x) {
+; CHECK-LABEL: @not_op_ctlz_wrong_xor_op1(
+; CHECK-NEXT:    [[N:%.*]] = xor i64 [[X:%.*]], -2
 ; CHECK-NEXT:    [[CT:%.*]] = tail call i64 @llvm.ctlz.i64(i64 [[N]], i1 true), !range [[RNG2]]
 ; CHECK-NEXT:    [[CAST:%.*]] = trunc i64 [[CT]] to i32
 ; CHECK-NEXT:    [[TOBOOL:%.*]] = icmp eq i64 [[X]], -1
 ; CHECK-NEXT:    [[R:%.*]] = select i1 [[TOBOOL]], i32 64, i32 [[CAST]]
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
-  %n = xor i64 %x, -1
+  %n = xor i64 %x, -2
   %ct = tail call i64 @llvm.ctlz.i64(i64 %n, i1 true)
   %cast = trunc i64 %ct to i32
   %tobool = icmp eq i64 %x, -1
@@ -301,19 +333,40 @@ define i32 @not_op_ctlz(i64 %x) {
   ret i32 %r
 }
 
-define i32 @not_op_cttz(i64 %x) {
-; CHECK-LABEL: @not_op_cttz(
+; negative test
+
+define i32 @not_op_ctlz_wrong_xor_op0(i64 %x, i64 %y) {
+; CHECK-LABEL: @not_op_ctlz_wrong_xor_op0(
+; CHECK-NEXT:    [[N:%.*]] = xor i64 [[Y:%.*]], -1
+; CHECK-NEXT:    [[CT:%.*]] = tail call i64 @llvm.ctlz.i64(i64 [[N]], i1 true), !range [[RNG2]]
+; CHECK-NEXT:    [[CAST:%.*]] = trunc i64 [[CT]] to i32
+; CHECK-NEXT:    [[TOBOOL:%.*]] = icmp eq i64 [[X:%.*]], -1
+; CHECK-NEXT:    [[R:%.*]] = select i1 [[TOBOOL]], i32 64, i32 [[CAST]]
+; CHECK-NEXT:    ret i32 [[R]]
+;
+  %n = xor i64 %y, -1
+  %ct = tail call i64 @llvm.ctlz.i64(i64 %n, i1 true)
+  %cast = trunc i64 %ct to i32
+  %tobool = icmp eq i64 %x, -1
+  %r = select i1 %tobool, i32 64, i32 %cast
+  ret i32 %r
+}
+
+; negative test
+
+define i32 @not_op_cttz_wrong_cmp(i64 %x) {
+; CHECK-LABEL: @not_op_cttz_wrong_cmp(
 ; CHECK-NEXT:    [[N:%.*]] = xor i64 [[X:%.*]], -1
 ; CHECK-NEXT:    [[CT:%.*]] = tail call i64 @llvm.cttz.i64(i64 [[N]], i1 true), !range [[RNG2]]
 ; CHECK-NEXT:    [[CAST:%.*]] = trunc i64 [[CT]] to i32
-; CHECK-NEXT:    [[TOBOOL:%.*]] = icmp eq i64 [[X]], -1
+; CHECK-NEXT:    [[TOBOOL:%.*]] = icmp eq i64 [[X]], 0
 ; CHECK-NEXT:    [[R:%.*]] = select i1 [[TOBOOL]], i32 64, i32 [[CAST]]
 ; CHECK-NEXT:    ret i32 [[R]]
 ;
   %n = xor i64 %x, -1
   %ct = tail call i64 @llvm.cttz.i64(i64 %n, i1 true)
   %cast = trunc i64 %ct to i32
-  %tobool = icmp eq i64 %x, -1
+  %tobool = icmp eq i64 %x, 0
   %r = select i1 %tobool, i32 64, i32 %cast
   ret i32 %r
 }
-- 
2.7.4