[InstCombine] ease use constraint in tryFactorization()

[spatel]

Summary

(x * y) + x --> x * (y + 1)
(x * y) - x --> x * (y - 1)

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

This is one of the IR transforms suggested in issue #57255.

This should be better in IR because it removes a use of a variable operand (we already fold the case with a constant multiply operand).
The backend should be able to re-distribute the multiply if that's better for the target.

[hiraditya]

llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp

1422

Can this cause a regression when Y is a power of two?
Some other optimization would rewrite x*y as y << n (when n == lg(y) at compile time)?

[bcl5980]

llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp

1422

https://alive2.llvm.org/ce/z/LRobTp
It looks the answer is yes.

[spatel]

llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp

1422

That is correct - this transform could obscure a power-of-2 transform.
But it could also expose a power-of-2 multiply that we miss today:
https://alive2.llvm.org/ce/z/ZUjR5S

In the typical instcombine sequence, we would process the multiply operand before the add user of that operand, so I haven't found a way to show a regression from this patch yet.

We need a more general solution if we are going to get the optimal IR for all patterns.

But the example shows a simpler pattern that I missed initially:
https://alive2.llvm.org/ce/z/eMhvQa

I think we should do that transform for the same reason stated here - it reduces the number of uses of a variable operand.

And it looks like we might not need this patch because -reassociate gives us that simpler pattern.

[spatel]

Patch updated:
Handle a simpler pattern - just looking for a mul with a common operand with an add.

[bcl5980]

llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp

1422

This looks like a special case of distributive law: x * y + x * 1
How about move this to InstCombinerImpl::SimplifyUsingDistributiveLaws or InstCombinerImpl::tryFactorization ?

[Allen]

llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp

1422

so, the decomposition the const of a MUL should be improved first?

1423

should we need consider the flag of NUW and NSW ?

[spatel]

llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp

1422

We do try to match this pattern in that code, but we miss it because the operand 'x' has multiple uses. I'll try to make it work if that isn't too ugly...but this explicit pattern-match might be the clearer way to achieve the transform.

1422

Ideally yes, we would have decomposition completed in the backend. But I also think this patch can uncover folds that we miss today (like the example where we created a power-of-2 multiply). So this patch should not be blocked on improving the backend. There could be both wins and losses.

1423

I only found one case so far where we can partially propagate nuw:
https://alive2.llvm.org/ce/z/aGSyrK
Everything else seems to be illegal. Propagating the flag would be another reason to code this match explicitly rather than trying to make it work through tryFactorization().

[bcl5980]

llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp

1423

If we don't move it to tryFactorization, do we need to add the similar pattern x*y - x --> (x-1) * y ?

1423

sorry, should be x*y -x --> x * (y-1)

[spatel]

Patch updated:
Adjust the use checks in tryFactorization() to handle the motivating patterns. This also changes some existing tests with multi-uses that we ignored before. Those look like good changes to me for the same reason cited before - it reduces uses of variables.

Also, I didn't remember that we have code that tries to propagate nsw/nuw through those folds. It is there, but incomplete as noted with the TODO items in the subtract tests.

[bcl5980]

LGTM. But please wait for someone else.

[Allen]

llvm/test/Transforms/InstCombine/sub.ll

2007 ↗

(On Diff #454992)

sorry for the naive question, does the nsw of mul implies that of add? As you just say the mul could retain nsw.

[bcl5980]

llvm/test/Transforms/InstCombine/sub.ll

2007 ↗

(On Diff #454992)

https://alive2.llvm.org/ce/z/WjpNoD
add pattern can't retain nsw. overflow comes from -1 * INT_MIN

[Allen]

llvm/test/Transforms/InstCombine/sub.ll

2007 ↗

(On Diff #454992)

Thanks @bcl5980, as your case showed, the mul pattern also can't retain nsw ?

[bcl5980]

llvm/test/Transforms/InstCombine/sub.ll

2007 ↗

(On Diff #454992)

Yeah, these two patterns are the only two cases we can keep mul flag. All add/sub flags will be lost.
https://alive2.llvm.org/ce/z/HDyrwb

[spatel]

llvm/lib/Transforms/InstCombine/InstructionCombining.cpp

692 ↗

(On Diff #454992)

For reference, this is the block where we try propagating overflow flags. It can be adjusted independently of this patch.

[arsenm]

llvm/test/Transforms/InstCombine/add.ll

1759

Do we have a backend combine to undo this already?

[arsenm]

llvm/test/Transforms/InstCombine/add.ll

1759

Specifically, this is bad for us because it breaks integer mad matching

[bcl5980]

llvm/test/Transforms/InstCombine/add.ll

1759

There is a similar pattern in DAGCombiner but the mul second op is also constant:
https://github.com/llvm/llvm-project/blob/main/llvm/lib/CodeGen/SelectionDAG/DAGCombiner.cpp#L4112

// fold (mul (add x, c1), c2) -> (add (mul x, c2), c1*c2)
if (DAG.isConstantIntBuildVectorOrConstantInt(N1) &&
    N0.getOpcode() == ISD::ADD &&
    DAG.isConstantIntBuildVectorOrConstantInt(N0.getOperand(1)) &&
    isMulAddWithConstProfitable(N, N0, N1))
  return DAG.getNode(
      ISD::ADD, DL, VT,
      DAG.getNode(ISD::MUL, SDLoc(N0), VT, N0.getOperand(0), N1),
      DAG.getNode(ISD::MUL, SDLoc(N1), VT, N0.getOperand(1), N1));

And it looks AArch64 backend already has the pattern:
https://github.com/llvm/llvm-project/blob/main/llvm/lib/Target/AArch64/AArch64ISelLowering.cpp#L15066

// Canonicalize X*(Y+1) -> X*Y+X and (X+1)*Y -> X*Y+Y,
// and in MachineCombiner pass, add+mul will be combined into madd.
// Similarly, X*(1-Y) -> X - X*Y and (1-Y)*X -> X - Y*X.

AMDGPU can add similar code, maybe you need to consider more for that (divergence? data type?).
https://alive2.llvm.org/ce/z/gJRkR5

llvm/lib/Transforms/InstCombine/InstCombineAddSub.cpp

Instruction *InstCombinerImpl::visitAdd(BinaryOperator &I) {
   if (match(&I, m_c_BinOp(m_Or(m_Value(A), m_Value(B)),
                           m_c_And(m_Deferred(A), m_Deferred(B))))) {
     // Replacing operands in-place to preserve nuw/nsw flags.
     replaceOperand(I, 0, A);
     replaceOperand(I, 1, B);
     return &I;
   }
 
+  // TODO: Match the related pattern with sub instead of inner add.
+  Value *X, *Y, *Z;
+  if (match(&I, m_c_Add(m_OneUse(m_Mul(m_Value(X), m_Value(Y))),
+                        m_OneUse(m_c_Add(m_Deferred(X), m_Value(Z)))))) {
+    // (X * Y) + (X + Z) --> ((Y + 1) * X) + Z
+    Value *IncY = Builder.CreateAdd(Y, ConstantInt::get(I.getType(), 1));
+    Value *NewMul = Builder.CreateMul(IncY, X);
+    return BinaryOperator::CreateAdd(NewMul, Z);
+  }
+  if (match(&I, m_c_Add(m_OneUse(m_Mul(m_Value(X), m_Value(Y))),
+                        m_OneUse(m_c_Add(m_Deferred(Y), m_Value(Z)))))) {
+    // (X * Y) + (Y + Z) --> ((X + 1) * Y) + Z
+    Value *IncX = Builder.CreateAdd(X, ConstantInt::get(I.getType(), 1));
+    Value *NewMul = Builder.CreateMul(IncX, Y);
+    return BinaryOperator::CreateAdd(NewMul, Z);
+  }
+
   // TODO(jingyue): Consider willNotOverflowSignedAdd and
   // willNotOverflowUnsignedAdd to reduce the number of invocations of
   // computeKnownBits.
   bool Changed = false;
   if (!I.hasNoSignedWrap() && willNotOverflowSignedAdd(LHS, RHS, I)) {
     Changed = true;
     I.setHasNoSignedWrap(true);
   }

llvm/test/Transforms/InstCombine/add.ll

 ; CHECK-NEXT:    ret i32 [[G]]
 ;
   %E = add i32 %B, %A
   %F = add i32 %C, %E
   %G = add i32 %D, %F
   ret i32 %G
 }
 
 define i8 @mul_add_common_factor_commute1(i8 %x, i8 %y, i8 %z) {
 ; CHECK-LABEL: @mul_add_common_factor_commute1(
-; CHECK-NEXT:    [[M:%.*]] = mul i8 [[X:%.*]], [[Y:%.*]]
-; CHECK-NEXT:    [[A:%.*]] = add i8 [[X]], [[Z:%.*]]
-; CHECK-NEXT:    [[R:%.*]] = add i8 [[M]], [[A]]
+; CHECK-NEXT:    [[TMP1:%.*]] = add i8 [[Y:%.*]], 1
+; CHECK-NEXT:    [[TMP2:%.*]] = mul i8 [[TMP1]], [[X:%.*]]
+; CHECK-NEXT:    [[R:%.*]] = add i8 [[TMP2]], [[Z:%.*]]
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %m = mul i8 %x, %y
   %a = add i8 %x, %z
   %r = add i8 %m, %a
   ret i8 %r
 }
 
 define i8 @mul_add_common_factor_commute2(i8 %x, i8 %y, i8 %z) {
 ; CHECK-LABEL: @mul_add_common_factor_commute2(
-; CHECK-NEXT:    [[M:%.*]] = mul i8 [[Y:%.*]], [[X:%.*]]
-; CHECK-NEXT:    [[A:%.*]] = add i8 [[X]], [[Z:%.*]]
-; CHECK-NEXT:    [[R:%.*]] = add i8 [[M]], [[A]]
+; CHECK-NEXT:    [[TMP1:%.*]] = add i8 [[Y:%.*]], 1
+; CHECK-NEXT:    [[TMP2:%.*]] = mul i8 [[TMP1]], [[X:%.*]]
+; CHECK-NEXT:    [[R:%.*]] = add i8 [[TMP2]], [[Z:%.*]]
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %m = mul i8 %y, %x
   %a = add i8 %x, %z
   %r = add i8 %m, %a
   ret i8 %r
 }
 
 define i8 @mul_add_common_factor_commute3(i8 %x, i8 %y, i8 %z) {
 ; CHECK-LABEL: @mul_add_common_factor_commute3(
-; CHECK-NEXT:    [[M:%.*]] = mul i8 [[X:%.*]], 42
-; CHECK-NEXT:    [[A:%.*]] = add i8 [[Z:%.*]], [[X]]
-; CHECK-NEXT:    [[R:%.*]] = add i8 [[M]], [[A]]
+; CHECK-NEXT:    [[TMP1:%.*]] = mul i8 [[X:%.*]], 43
+; CHECK-NEXT:    [[R:%.*]] = add i8 [[TMP1]], [[Z:%.*]]
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %m = mul i8 %x, 42
   %a = add i8 %z, %x
   %r = add i8 %m, %a
   ret i8 %r
 }
 
 define <2 x i8> @mul_add_common_factor_commute4(<2 x i8> %x, <2 x i8> %y, <2 x i8> %z) {
 ; CHECK-LABEL: @mul_add_common_factor_commute4(
-; CHECK-NEXT:    [[M:%.*]] = mul <2 x i8> [[Y:%.*]], [[X:%.*]]
-; CHECK-NEXT:    [[A:%.*]] = add <2 x i8> [[Z:%.*]], [[X]]
-; CHECK-NEXT:    [[R:%.*]] = add <2 x i8> [[M]], [[A]]
+; CHECK-NEXT:    [[TMP1:%.*]] = add <2 x i8> [[Y:%.*]], <i8 1, i8 1>
+; CHECK-NEXT:    [[TMP2:%.*]] = mul <2 x i8> [[TMP1]], [[X:%.*]]
+; CHECK-NEXT:    [[R:%.*]] = add <2 x i8> [[TMP2]], [[Z:%.*]]
 ; CHECK-NEXT:    ret <2 x i8> [[R]]
 ;
   %m = mul <2 x i8> %y, %x
   %a = add <2 x i8> %z, %x
   %r = add <2 x i8> %m, %a
   ret <2 x i8> %r
 }
 
 define i8 @mul_add_common_factor_commute5(i8 %x, i8 %y, i8 %z) {
 ; CHECK-LABEL: @mul_add_common_factor_commute5(
-; CHECK-NEXT:    [[M:%.*]] = mul i8 [[X:%.*]], [[Y:%.*]]
-; CHECK-NEXT:    [[A:%.*]] = add i8 [[X]], 43
-; CHECK-NEXT:    [[R:%.*]] = add i8 [[A]], [[M]]
+; CHECK-NEXT:    [[TMP1:%.*]] = add i8 [[Y:%.*]], 1
+; CHECK-NEXT:    [[TMP2:%.*]] = mul i8 [[TMP1]], [[X:%.*]]
+; CHECK-NEXT:    [[R:%.*]] = add i8 [[TMP2]], 43
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %m = mul i8 %x, %y
   %a = add i8 %x, 43
   %r = add i8 %a, %m
   ret i8 %r
 }
 
 define i8 @mul_add_common_factor_commute6(i8 %x, i8 %y, i8 %z) {
 ; CHECK-LABEL: @mul_add_common_factor_commute6(
-; CHECK-NEXT:    [[M:%.*]] = mul i8 [[Y:%.*]], [[X:%.*]]
-; CHECK-NEXT:    [[A:%.*]] = add i8 [[X]], [[Z:%.*]]
-; CHECK-NEXT:    [[R:%.*]] = add i8 [[A]], [[M]]
+; CHECK-NEXT:    [[TMP1:%.*]] = add i8 [[Y:%.*]], 1
+; CHECK-NEXT:    [[TMP2:%.*]] = mul i8 [[TMP1]], [[X:%.*]]
+; CHECK-NEXT:    [[R:%.*]] = add i8 [[TMP2]], [[Z:%.*]]
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %m = mul i8 %y, %x
   %a = add i8 %x, %z
   %r = add i8 %a, %m
   ret i8 %r
 }
 
 define i8 @mul_add_common_factor_commute7(i8 %x, i8 %y, i8 %z) {
 ; CHECK-LABEL: @mul_add_common_factor_commute7(
-; CHECK-NEXT:    [[M:%.*]] = mul i8 [[X:%.*]], [[Y:%.*]]
-; CHECK-NEXT:    [[A:%.*]] = add i8 [[Z:%.*]], [[X]]
-; CHECK-NEXT:    [[R:%.*]] = add i8 [[A]], [[M]]
+; CHECK-NEXT:    [[TMP1:%.*]] = add i8 [[Y:%.*]], 1
+; CHECK-NEXT:    [[TMP2:%.*]] = mul i8 [[TMP1]], [[X:%.*]]
+; CHECK-NEXT:    [[R:%.*]] = add i8 [[TMP2]], [[Z:%.*]]
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %m = mul i8 %x, %y
   %a = add i8 %z, %x
   %r = add i8 %a, %m
   ret i8 %r
 }
 
 define i8 @mul_add_common_factor_commute8(i8 %x, i8 %y, i8 %z) {
 ; CHECK-LABEL: @mul_add_common_factor_commute8(
-; CHECK-NEXT:    [[M:%.*]] = mul i8 [[Y:%.*]], [[X:%.*]]
-; CHECK-NEXT:    [[A:%.*]] = add i8 [[Z:%.*]], [[X]]
-; CHECK-NEXT:    [[R:%.*]] = add i8 [[A]], [[M]]
+; CHECK-NEXT:    [[TMP1:%.*]] = add i8 [[Y:%.*]], 1
+; CHECK-NEXT:    [[TMP2:%.*]] = mul i8 [[TMP1]], [[X:%.*]]
+; CHECK-NEXT:    [[R:%.*]] = add i8 [[TMP2]], [[Z:%.*]]
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %m = mul i8 %y, %x
   %a = add i8 %z, %x
   %r = add i8 %a, %m
   ret i8 %r
 }
 
+; negative test - extra use
+
 define i8 @mul_add_common_factor_use1(i8 %x, i8 %y, i8 %z) {
 ; CHECK-LABEL: @mul_add_common_factor_use1(
 ; CHECK-NEXT:    [[M:%.*]] = mul i8 [[X:%.*]], [[Y:%.*]]
 ; CHECK-NEXT:    call void @use(i8 [[M]])
 ; CHECK-NEXT:    [[A:%.*]] = add i8 [[X]], [[Z:%.*]]
 ; CHECK-NEXT:    [[R:%.*]] = add i8 [[M]], [[A]]
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %m = mul i8 %x, %y
   call void @use(i8 %m)
   %a = add i8 %x, %z
   %r = add i8 %m, %a
   ret i8 %r
 }
 
+; negative test - extra use
+
 define i8 @mul_add_common_factor_use2(i8 %x, i8 %y, i8 %z) {
 ; CHECK-LABEL: @mul_add_common_factor_use2(
 ; CHECK-NEXT:    [[M:%.*]] = mul i8 [[X:%.*]], [[Y:%.*]]
 ; CHECK-NEXT:    [[A:%.*]] = add i8 [[X]], [[Z:%.*]]
 ; CHECK-NEXT:    call void @use(i8 [[A]])
 ; CHECK-NEXT:    [[R:%.*]] = add i8 [[M]], [[A]]
 ; CHECK-NEXT:    ret i8 [[R]]
 ;
   %m = mul i8 %x, %y
   %a = add i8 %x, %z
   call void @use(i8 %a)
   %r = add i8 %m, %a
   ret i8 %r
 }

llvm/test/Transforms/PGOProfile/chr.ll

 ; With no phi at the exit, but the exit needs a phi inserted after CHR.
 ; Roughly,
 ; t0 = *i
 ; if ((t0 & 1) != 0) // Likely true
 ;   foo()
 ; t1 = *j
 ; if ((t1 & 2) != 0) // Likely true
 ;   foo()
-; return (t1 * 42) - (t1 - 99)
+; return (t1 * 42) + (t1 - 99)
 ; ->
 ; t0 = *i
 ; if ((t0 & 3) == 3) { // Likely true
 ;   foo()
 ;   t1 = *j
 ;   foo()
 ; } else {
 ;   if ((t0 & 1) != 0)
 ;     foo()
 ;   if ((t0 & 2) != 0) {
 ;     t1 = *j
 ;     foo()
 ;   }
 ; }
 ; // A new phi for t1 is inserted here.
-; return (t1 * 42) - (t1 - 99)
+; return (t1 * 43) - 99
 define i32 @test_chr_10(ptr %i, ptr %j) !prof !14 {
 ; CHECK-LABEL: @test_chr_10(
 ; CHECK-NEXT:  entry:
 ; CHECK-NEXT:    [[TMP0:%.*]] = load i32, ptr [[I:%.*]], align 4
 ; CHECK-NEXT:    [[TMP1:%.*]] = and i32 [[TMP0]], 3
 ; CHECK-NEXT:    [[TMP2:%.*]] = icmp eq i32 [[TMP1]], 3
 ; CHECK-NEXT:    br i1 [[TMP2]], label [[BB0:%.*]], label [[ENTRY_SPLIT_NONCHR:%.*]], !prof [[PROF15]]
 ; CHECK:       bb0:

 ; CHECK-NEXT:    [[TMP6:%.*]] = and i32 [[TMP0]], 2
 ; CHECK-NEXT:    [[TMP7:%.*]] = icmp eq i32 [[TMP6]], 0
 ; CHECK-NEXT:    br i1 [[TMP7]], label [[BB3]], label [[BB2_NONCHR:%.*]], !prof [[PROF16]]
 ; CHECK:       bb2.nonchr:
 ; CHECK-NEXT:    call void @foo()
 ; CHECK-NEXT:    br label [[BB3]]
 ; CHECK:       bb3:
 ; CHECK-NEXT:    [[TMP8:%.*]] = phi i32 [ [[TMP3]], [[BB0]] ], [ [[TMP5]], [[BB2_NONCHR]] ], [ [[TMP5]], [[BB1_NONCHR]] ]
-; CHECK-NEXT:    [[TMP9:%.*]] = mul i32 [[TMP8]], 42
-; CHECK-NEXT:    [[TMP10:%.*]] = add i32 [[TMP8]], -99
-; CHECK-NEXT:    [[TMP11:%.*]] = add i32 [[TMP9]], [[TMP10]]
-; CHECK-NEXT:    ret i32 [[TMP11]]
+; CHECK-NEXT:    [[TMP9:%.*]] = mul i32 [[TMP8]], 43
+; CHECK-NEXT:    [[TMP10:%.*]] = add i32 [[TMP9]], -99
+; CHECK-NEXT:    ret i32 [[TMP10]]
 ;
 entry:
   %0 = load i32, ptr %i
   %1 = and i32 %0, 1
   %2 = icmp eq i32 %1, 0
   br i1 %2, label %bb1, label %bb0, !prof !15
 
 bb0: