[InstCombine] Infer exact for lshr by cttz (#136696)

Infer the 'exact' flag on an 'lshr' or 'ashr' instruction when the shift
amount is computed via a 'cttz' intrinsic on the same operand.
 
Proof: https://alive2.llvm.org/ce/z/CQR2PG

Fixes #131444.

Author: houngkoungting

llvm/lib/Transforms/InstCombine/InstCombineShifts.cpp

@@ -994,6 +994,12 @@ static bool setShiftFlags(BinaryOperator &I, const SimplifyQuery &Q) {
       I.setIsExact();
       return true;
     }
+    // Infer 'exact' flag if shift amount is cttz(x) on the same operand.
+    if (match(I.getOperand(1), m_Intrinsic<Intrinsic::cttz>(
+                                   m_Specific(I.getOperand(0)), m_Value()))) {
+      I.setIsExact();
+      return true;
+    }
   }
 
   // Compute what we know about shift count.

llvm/test/Transforms/InstCombine/cttz-shift-exact.ll

@@ -0,0 +1,41 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py UTC_ARGS: --version 5
+; RUN: opt -passes=instcombine -S < %s | FileCheck %s
+
+declare i32 @llvm.cttz.i32(i32, i1)
+
+define i32 @test_cttz_lshr(i32 %x) {
+; CHECK-LABEL: define i32 @test_cttz_lshr(
+; CHECK-SAME: i32 [[X:%.*]]) {
+; CHECK-NEXT:    [[CTTZ:%.*]] = call range(i32 0, 33) i32 @llvm.cttz.i32(i32 [[X]], i1 true)
+; CHECK-NEXT:    [[SH:%.*]] = lshr exact i32 [[X]], [[CTTZ]]
+; CHECK-NEXT:    ret i32 [[SH]]
+;
+  %cttz = call i32 @llvm.cttz.i32(i32 %x, i1 false)
+  %sh = lshr i32 %x, %cttz
+  ret i32 %sh
+}
+
+define i32 @test_cttz_ashr(i32 %x) {
+; CHECK-LABEL: define i32 @test_cttz_ashr(
+; CHECK-SAME: i32 [[X:%.*]]) {
+; CHECK-NEXT:    [[CTTZ:%.*]] = call range(i32 0, 33) i32 @llvm.cttz.i32(i32 [[X]], i1 true)
+; CHECK-NEXT:    [[SH:%.*]] = ashr exact i32 [[X]], [[CTTZ]]
+; CHECK-NEXT:    ret i32 [[SH]]
+;
+  %cttz = call i32 @llvm.cttz.i32(i32 %x, i1 true)
+  %sh = ashr i32 %x, %cttz
+  ret i32 %sh
+}
+
+define i32 @test_cttz_diff_operand(i32 %x, i32 %y) {
+; CHECK-LABEL: define i32 @test_cttz_diff_operand(
+; CHECK-SAME: i32 [[X:%.*]], i32 [[Y:%.*]]) {
+; CHECK-NEXT:    [[CTTZ:%.*]] = call range(i32 0, 33) i32 @llvm.cttz.i32(i32 [[Y]], i1 true)
+; CHECK-NEXT:    [[SH:%.*]] = lshr i32 [[X]], [[CTTZ]]
+; CHECK-NEXT:    ret i32 [[SH]]
+;
+  %cttz = call i32 @llvm.cttz.i32(i32 %y, i1 true)
+  %sh = lshr i32 %x, %cttz
+  ret i32 %sh
+}
+

llvm/test/Transforms/InstCombine/select.ll

@@ -2736,7 +2736,7 @@ define i32 @pr47322_more_poisonous_replacement(i32 %arg) {
 ; CHECK-LABEL: @pr47322_more_poisonous_replacement(
 ; CHECK-NEXT:    [[CMP:%.*]] = icmp eq i32 [[ARG:%.*]], 0
 ; CHECK-NEXT:    [[TRAILING:%.*]] = call range(i32 0, 33) i32 @llvm.cttz.i32(i32 [[ARG]], i1 true)
-; CHECK-NEXT:    [[SHIFTED:%.*]] = lshr i32 [[ARG]], [[TRAILING]]
+; CHECK-NEXT:    [[SHIFTED:%.*]] = lshr exact i32 [[ARG]], [[TRAILING]]
 ; CHECK-NEXT:    [[R1_SROA_0_1:%.*]] = select i1 [[CMP]], i32 0, i32 [[SHIFTED]]
 ; CHECK-NEXT:    ret i32 [[R1_SROA_0_1]]
 ;