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

[via llvm-commits]

================
@@ -862,4 +862,130 @@ define i32 @ashr_mul_times_5_div_4_exact_2(i32 %x) {
   ret i32 %ashr
 }
 
+define i32 @ashr_shift_mul(i32 %x) {
+; CHECK-LABEL: @ashr_shift_mul(
+; CHECK-NEXT:    [[TMP1:%.*]] = freeze i32 [[X:%.*]]
+; CHECK-NEXT:    [[TMP2:%.*]] = ashr exact i32 [[TMP1]], 3
+; CHECK-NEXT:    [[RES:%.*]] = add i32 [[TMP1]], [[TMP2]]
+; 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 %x) {
+; CHECK-LABEL: @ashr_shift_mul_nuw(
+; CHECK-NEXT:    [[TMP1:%.*]] = lshr exact i32 [[X:%.*]], 3
+; CHECK-NEXT:    [[RES:%.*]] = add nuw i32 [[TMP1]], [[X]]
+; 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 %x) {
+; CHECK-LABEL: @ashr_shift_mul_nsw(
+; CHECK-NEXT:    [[TMP1:%.*]] = freeze i32 [[X:%.*]]
+; CHECK-NEXT:    [[TMP2:%.*]] = ashr exact i32 [[TMP1]], 3
+; CHECK-NEXT:    [[RES:%.*]] = add nsw i32 [[TMP1]], [[TMP2]]
+; CHECK-NEXT:    ret i32 [[RES]]
+;
+  %a = ashr exact i32 %x, 3
+  %res = mul nsw i32 %a, 9
+  ret i32 %res
+}
----------------
goldsteinn wrote:

I don't think you have solved the problem that you can't propagate `nsw` to the add if the ShiftAmt may be `BitWidth - 1` or `MulAp` may be INT_MIN. This case might get simplified beforehand, but it still shouldn't break w/ your code.

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