[llvm] [InstCombine] Fold mul (shr exact (X, N)), 2^N + 1 -> add (X , shr exact (X, N)) (PR #112407)

[Yingwei Zheng via llvm-commits]

================
@@ -1077,4 +1077,170 @@ entry:
   ret i32 %shr
 }
 
+define i32 @ashr_shift_mul(i32 noundef %x) {
+; CHECK-LABEL: @ashr_shift_mul(
+; CHECK-NEXT:    [[A:%.*]] = ashr exact i32 [[X:%.*]], 3
+; CHECK-NEXT:    [[RES:%.*]] = add i32 [[X]], [[A]]
+; CHECK-NEXT:    ret i32 [[RES]]
+;
+  %a = ashr exact i32 %x, 3
+  %res = mul i32 %a, 9
+  ret i32 %res
+}
+
+define i32 @ashr_shift_mul_nuw(i32 noundef %x) {
+; CHECK-LABEL: @ashr_shift_mul_nuw(
+; CHECK-NEXT:    [[TMP1:%.*]] = lshr exact i32 [[X:%.*]], 3
+; CHECK-NEXT:    [[RES:%.*]] = add nuw i32 [[X]], [[TMP1]]
+; CHECK-NEXT:    ret i32 [[RES]]
+;
+  %a = ashr exact i32 %x, 3
+  %res = mul nuw i32 %a, 9
+  ret i32 %res
+}
+
+define i32 @ashr_shift_mul_nsw(i32 noundef %x) {
+; CHECK-LABEL: @ashr_shift_mul_nsw(
+; CHECK-NEXT:    [[A:%.*]] = ashr exact i32 [[X:%.*]], 3
+; CHECK-NEXT:    [[RES:%.*]] = add nsw i32 [[X]], [[A]]
+; CHECK-NEXT:    ret i32 [[RES]]
+;
+  %a = ashr exact i32 %x, 3
+  %res = mul nsw i32 %a, 9
+  ret i32 %res
+}
+
+define i32 @lshr_shift_mul_nuw(i32 noundef %x) {
+; CHECK-LABEL: @lshr_shift_mul_nuw(
+; CHECK-NEXT:    [[A:%.*]] = lshr exact i32 [[X:%.*]], 3
+; CHECK-NEXT:    [[RES:%.*]] = add nuw i32 [[X]], [[A]]
+; CHECK-NEXT:    ret i32 [[RES]]
+;
+  %a = lshr exact i32 %x, 3
+  %res = mul nuw i32 %a, 9
+  ret i32 %res
+}
+
+define i32 @lshr_shift_mul(i32 noundef %x) {
+; CHECK-LABEL: @lshr_shift_mul(
+; CHECK-NEXT:    [[A:%.*]] = lshr exact i32 [[X:%.*]], 3
+; CHECK-NEXT:    [[RES:%.*]] = add i32 [[X]], [[A]]
+; CHECK-NEXT:    ret i32 [[RES]]
+;
+  %a = lshr exact i32 %x, 3
+  %res = mul i32 %a, 9
+  ret i32 %res
+}
+
+define i32 @lshr_shift_mul_nsw(i32 noundef %x) {
+; CHECK-LABEL: @lshr_shift_mul_nsw(
+; CHECK-NEXT:    [[A:%.*]] = lshr exact i32 [[X:%.*]], 3
+; CHECK-NEXT:    [[RES:%.*]] = add nsw i32 [[X]], [[A]]
+; CHECK-NEXT:    ret i32 [[RES]]
+;
+  %a = lshr exact i32 %x, 3
+  %res = mul nsw i32 %a, 9
+  ret i32 %res
+}
+
+; Negative test
+
+define i32 @lshr_no_exact(i32 %x) {
+; CHECK-LABEL: @lshr_no_exact(
+; CHECK-NEXT:    [[A:%.*]] = lshr i32 [[X:%.*]], 3
+; CHECK-NEXT:    [[RES:%.*]] = mul nuw nsw i32 [[A]], 9
+; CHECK-NEXT:    ret i32 [[RES]]
+;
+  %a = lshr i32 %x, 3
+  %res = mul nsw i32 %a, 9
+  ret i32 %res
+}
+
+; Negative test
+
+define i32 @ashr_no_exact(i32 %x) {
+; CHECK-LABEL: @ashr_no_exact(
+; CHECK-NEXT:    [[A:%.*]] = ashr i32 [[X:%.*]], 3
+; CHECK-NEXT:    [[RES:%.*]] = mul nsw i32 [[A]], 9
+; CHECK-NEXT:    ret i32 [[RES]]
+;
+  %a = ashr i32 %x, 3
+  %res = mul nsw i32 %a, 9
+  ret i32 %res
+}
+
+define i32 @lshr_no_undef(i32 %x) {
----------------
dtcxzyw wrote:

Please remove these `*_no_undef` tests.


https://github.com/llvm/llvm-project/pull/112407