[InstCombine] Handle mul in maintainNoSignedWrap (#123299)

Alive2: https://alive2.llvm.org/ce/z/Kgamks
Closes #123175.

For `@foo1`, the nsw flag is propagated because we first convert it into
`mul nsw nuw (shl nsw nuw X, 1), 3`.

Author: dtcxzyw

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

@@ -281,28 +281,33 @@ bool InstCombinerImpl::shouldChangeType(Type *From, Type *To) const {
 // Return true, if No Signed Wrap should be maintained for I.
 // The No Signed Wrap flag can be kept if the operation "B (I.getOpcode) C",
 // where both B and C should be ConstantInts, results in a constant that does
-// not overflow. This function only handles the Add and Sub opcodes. For
+// not overflow. This function only handles the Add/Sub/Mul opcodes. For
 // all other opcodes, the function conservatively returns false.
 static bool maintainNoSignedWrap(BinaryOperator &I, Value *B, Value *C) {
   auto *OBO = dyn_cast<OverflowingBinaryOperator>(&I);
   if (!OBO || !OBO->hasNoSignedWrap())
     return false;
 
-  // We reason about Add and Sub Only.
-  Instruction::BinaryOps Opcode = I.getOpcode();
-  if (Opcode != Instruction::Add && Opcode != Instruction::Sub)
-    return false;
-
   const APInt *BVal, *CVal;
   if (!match(B, m_APInt(BVal)) || !match(C, m_APInt(CVal)))
     return false;
 
+  // We reason about Add/Sub/Mul Only.
   bool Overflow = false;
-  if (Opcode == Instruction::Add)
+  switch (I.getOpcode()) {
+  case Instruction::Add:
     (void)BVal->sadd_ov(*CVal, Overflow);
-  else
+    break;
+  case Instruction::Sub:
     (void)BVal->ssub_ov(*CVal, Overflow);
-
+    break;
+  case Instruction::Mul:
+    (void)BVal->smul_ov(*CVal, Overflow);
+    break;
+  default:
+    // Conservatively return false for other opcodes.
+    return false;
+  }
   return !Overflow;
 }
 

llvm/test/Transforms/InstCombine/nsw.ll

@@ -415,3 +415,63 @@ define i8 @neg_nsw_mul_missing_nsw_on_mul(i8 %a1, i8 %a2, i8 %b) {
   %neg = sub nsw i8 0, %shl
   ret i8 %neg
 }
+
+; This could propagate nsw.
+
+define i16 @mul_nsw_reassoc_prop(i16 %x) {
+; CHECK-LABEL: @mul_nsw_reassoc_prop(
+; CHECK-NEXT:    [[B:%.*]] = mul nsw i16 [[X:%.*]], 6
+; CHECK-NEXT:    ret i16 [[B]]
+;
+  %a = mul nsw i16 %x, 3
+  %b = mul nsw i16 %a, 2
+  ret i16 %b
+}
+
+; This could propagate nsw.
+
+define i16 @mul_nsw_reassoc_prop_neg(i16 %x) {
+; CHECK-LABEL: @mul_nsw_reassoc_prop_neg(
+; CHECK-NEXT:    [[B:%.*]] = mul nsw i16 [[X:%.*]], -2201
+; CHECK-NEXT:    ret i16 [[B]]
+;
+  %a = mul nsw i16 %x, -71
+  %b = mul nsw i16 %a, 31
+  ret i16 %b
+}
+
+; Must not propagate nsw.
+
+define i16 @mul_nsw_reassoc_prop_no_nsw1(i16 %x) {
+; CHECK-LABEL: @mul_nsw_reassoc_prop_no_nsw1(
+; CHECK-NEXT:    [[B:%.*]] = mul i16 [[X:%.*]], 6
+; CHECK-NEXT:    ret i16 [[B]]
+;
+  %a = mul i16 %x, 3
+  %b = mul nsw i16 %a, 2
+  ret i16 %b
+}
+
+; Must not propagate nsw.
+
+define i16 @mul_nsw_reassoc_prop_no_nsw2(i16 %x) {
+; CHECK-LABEL: @mul_nsw_reassoc_prop_no_nsw2(
+; CHECK-NEXT:    [[B:%.*]] = mul i16 [[X:%.*]], 6
+; CHECK-NEXT:    ret i16 [[B]]
+;
+  %a = mul nsw i16 %x, 3
+  %b = mul i16 %a, 2
+  ret i16 %b
+}
+
+; Must not propagate nsw.
+
+define i16 @mul_nsw_reassoc_prop_overflow(i16 %x) {
+; CHECK-LABEL: @mul_nsw_reassoc_prop_overflow(
+; CHECK-NEXT:    [[B:%.*]] = mul i16 [[X:%.*]], -31777
+; CHECK-NEXT:    ret i16 [[B]]
+;
+  %a = mul nsw i16 %x, 1023
+  %b = mul nsw i16 %a, 33
+  ret i16 %b
+}