[InstCombine] Distributive or+mul with const operand

[Allen]

Summary

We aleady support the transform: (X+C1)*CI -> X*CI+C1*CI
Here the case is a little special as the form of (X+C1)*CI is transformed into (X|C1)*CI,
so we should also support the transform: (X|C1)*CI -> X*CI+C1*CI
Fixes https://github.com/llvm/llvm-project/issues/57278

[bcl5980]

We may need a minimal test case.

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

232

Match 'or' then check haveNoCommonBitsSet should be better I think.

[spatel]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

229–232

The existing code is awkward. We should update it before adding to it. Use m_ImmConstant(C1) to do this transform but ignore a constant expression.

I updated the test that was intended to go with this code change here:
5260146a8a74084f3d38d8bb448ae3c5690b9084

[bcl5980]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

223–224

If we update current code, Op1 should be m_ImmConstant also

[Allen]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

229–232

hi @spatel

I tried with your new case, and it will fail on the check **if (!match(Mul, m_Mul(m_Value(), m_Value())))**, so the instcombine will not take active, is it within your expectations?

[Allen]

address the comment

[Allen]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

223–224

Done, thanks

232

Thanks very much, use haveNoCommonBitsSet make the code much clear.

[bcl5980]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

229–232

I think we can remove the condition:

if (!match(Mul, m_Mul(m_Value(), m_Value())))

But I don't know what happen if (INT_MIN * -1).
Based on current folding code, Constant (INTMIN / -1) is a poison value.

230–246

Maybe we should keep nuw flag in this change. Or add some tests and TODO for that.
https://alive2.llvm.org/ce/z/awsQrx

231

May be we need attach the context information for the function like:

haveNoCommonBitsSet(X, C1, DL, &AC, &I, &DT)

[bcl5980]

And please add at least 1 vector test also.

[Allen]

1、add a vector case
2、keep nuw flag
3、update with haveNoCommonBitsSet(X, C1, DL, &AC, &I, &DT)

[Allen]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

229–232

the above checking is deleted, and now i32 ptrtoint (i32* @g to i32) doesn't match the match(Op1, m_ImmConstant()) for that case, which I think it is within expectations ?

230–246

I make a try to propagation the flag, thanks

231

Apply your comment, thanks

[bcl5980]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

232

Only check nuw for mul works for pattern or.
But for the pattern add, we need check both instructions are nuw then can keep the flag.
And target pattern can be nuw for both instruction also.

[Allen]

update the checking logic for nuw flag

[Allen]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

232

Thank you for your patient guidance.

[RKSimon]

llvm/test/Transforms/InstCombine/mul.ll

690

Please can you add a case with undefs in the vector constants:

define <2 x i32> @PR57278_shl_vec_undef(<2 x i32> %v1) {

%shl = shl nuw <2 x i32> %v1, <i32 2, i32 undef>
%add = or <2 x i32> %shl, <i32 3, i32 undef>
%mul = mul nuw <2 x i32> %add, <i32 3, i32 undef>
ret <2 x i32> %mul

}

[Allen]

Add a new vector case with undef

[Allen]

llvm/test/Transforms/InstCombine/mul.ll

690

Done, Thanks

[bcl5980]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

236

Should be

if (I.hasNoUnsignedWrap() && OrNUW) {
  cast<BinaryOperator>(NewMul)->setHasNoUnsignedWrap();
  BO->setHasNoUnsignedWrap();
}

[spatel]

I think it would be better to make the 'nuw' change separately. That is missing from the current transform, so it could be done before this patch with 'add' alone, or make that a follow-up patch to this one.

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

229–232

Yes, that is the expectation - the transform should not be attempted with constant expressions.

229–232

It would be clearer to capture Op1 as a constant and create this constant directly:

Constant *MulC;
if (match(Op1, m_ImmConstant(MulC))) {
  ...
  Constant *NewC = ConstantExpr::getMul(C1, MulC);

Please adjust the code comment above here to match whatever variable names you choose. So it may look like this:

// (X + C1) * MulC --> (X * MulC) + (C1 * MulC)

[Allen]

In D132658#3751724, @spatel wrote:

I think it would be better to make the 'nuw' change separately. That is missing from the current transform, so it could be done before this patch with 'add' alone, or make that a follow-up patch to this one.

Thanks for your suggestion, addressed with D132777

[Allen]

rebase after D132777

[spatel]

Please pre-commit the baseline tests. If something is not transforming, you can add a comment in this patch to make it clear that it is a negative test or TODO item.

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

234

Is there a reason to not use the more direct ConstExpr::getMul() API and make it explicit that NewC is a Constant *?

llvm/test/Transforms/InstCombine/mul.ll

746

Replace "undef" with "poison" in this test - we are transitioning away from undef in IR.

[Allen]

In D132658#3754136, @spatel wrote:

Please pre-commit the baseline tests. If something is not transforming, you can add a comment in this patch to make it clear that it is a negative test or TODO item.

Addressed in the link D132820, and I'll rebase this commit after D132820 merged, thanks!

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

234

Thanks your suggestion, done

llvm/test/Transforms/InstCombine/mul.ll

746

Thanks, apply with your comment

[bcl5980]

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

234

@spatel, based on nikic's RFC, maybe we need avoid use ConstExpr.
https://discourse.llvm.org/t/rfc-remove-most-constant-expressions/63179

239–240

If NewMulBO is constant we also can set BO to nuw. So the logic should be this

if (I.hasNoUnsignedWrap() && Op0NUW && NewMulBO) {
  if (auto *NewMulBO = dyn_cast<BinaryOperator>(NewMul))
    NewMulBO->setHasNoUnsignedWrap();
  BO->setHasNoUnsignedWrap();
}

239–240

Sorry should be this one:

if (I.hasNoUnsignedWrap() && Op0NUW) {
  if (auto *NewMulBO = dyn_cast<BinaryOperator>(NewMul))
    NewMulBO->setHasNoUnsignedWrap();
  BO->setHasNoUnsignedWrap();
}

[Allen]

a) rebase to the top as the baseline test case merged
b) update the logic of propagate nuw
c) Still use the Buillder.CreateMul as RFC https://discourse.llvm.org/t/rfc-remove-most-constant-expressions/63179

[bcl5980]

LGTM, but please wait until someone else to approve it.

[spatel]

LGTM

llvm/lib/Transforms/InstCombine/InstCombineMulDivRem.cpp

Instruction *InstCombinerImpl::visitMul(BinaryOperator &I) {
   }
 
   if (Instruction *FoldedMul = foldBinOpIntoSelectOrPhi(I))
     return FoldedMul;
 
   if (Value *FoldedMul = foldMulSelectToNegate(I, Builder))
     return replaceInstUsesWith(I, FoldedMul);
 
   // Simplify mul instructions with a constant RHS.
   Constant *MulC;
   if (match(Op1, m_ImmConstant(MulC))) {
     // Canonicalize (X+C1)*MulC -> X*MulC+C1*MulC.
+    // Canonicalize (X|C1)*MulC -> X*MulC+C1*MulC.
     Value *X;
     Constant *C1;
-    if (match(Op0, m_OneUse(m_Add(m_Value(X), m_ImmConstant(C1))))) {
+    if ((match(Op0, m_OneUse(m_Add(m_Value(X), m_ImmConstant(C1))))) ||
+        (match(Op0, m_OneUse(m_Or(m_Value(X), m_ImmConstant(C1)))) &&
+         haveNoCommonBitsSet(X, C1, DL, &AC, &I, &DT))) {
       // C1*MulC simplifies to a tidier constant.
       Value *NewC = Builder.CreateMul(C1, MulC);
       auto *BOp0 = cast<BinaryOperator>(Op0);
+      bool Op0NUW =
+          (BOp0->getOpcode() == Instruction::Or || BOp0->hasNoUnsignedWrap());
       Value *NewMul = Builder.CreateMul(X, MulC);
       auto *BO = BinaryOperator::CreateAdd(NewMul, NewC);
-      auto *NewMulBO = dyn_cast<BinaryOperator>(NewMul);
-      if (I.hasNoUnsignedWrap() && BOp0->hasNoUnsignedWrap() && NewMulBO) {
+      if (I.hasNoUnsignedWrap() && Op0NUW) {
+        // If NewMulBO is constant we also can set BO to nuw.
+        if (auto *NewMulBO = dyn_cast<BinaryOperator>(NewMul))
           NewMulBO->setHasNoUnsignedWrap();
         BO->setHasNoUnsignedWrap();
       }
       return BO;
     }
   }
 
   // abs(X) * abs(X) -> X * X
   // nabs(X) * nabs(X) -> X * X
   if (Op0 == Op1) {
     Value *X, *Y;
     SelectPatternFlavor SPF = matchSelectPattern(Op0, X, Y).Flavor;

llvm/test/Transforms/InstCombine/mul.ll

 ; CHECK-NEXT:    [[MUL:%.*]] = add i32 [[TMP1]], 9
 ; CHECK-NEXT:    ret i32 [[MUL]]
 ;
   %add = add i32 %a, 3
   %mul = mul nuw i32 %add, 3
   ret i32 %mul
 }
 
-; https://alive2.llvm.org/ce/z/jJ8rZd
+; Don't propagate nsw flag in this change, https://alive2.llvm.org/ce/z/jJ8rZd
 define i32 @PR57278_shl(i32 %a) {
 ; CHECK-LABEL: @PR57278_shl(
-; CHECK-NEXT:    [[SHL:%.*]] = shl nsw i32 [[A:%.*]], 2
-; CHECK-NEXT:    [[ADD:%.*]] = or i32 [[SHL]], 3
-; CHECK-NEXT:    [[MUL:%.*]] = mul nsw i32 [[ADD]], 3
+; CHECK-NEXT:    [[TMP1:%.*]] = mul i32 [[A:%.*]], 12
+; CHECK-NEXT:    [[MUL:%.*]] = add i32 [[TMP1]], 9
 ; CHECK-NEXT:    ret i32 [[MUL]]
 ;
   %shl = shl nsw i32 %a, 2
   %add = or i32 %shl, 3
   %mul = mul nsw i32 %add, 3
   ret i32 %mul
 }
 
 ; Negative test: Have common bits set
 define i32 @PR57278_shl_1(i32 %a) {
 ; CHECK-LABEL: @PR57278_shl_1(
 ; CHECK-NEXT:    [[SHL:%.*]] = shl nsw i32 [[A:%.*]], 2
 ; CHECK-NEXT:    [[ADD:%.*]] = or i32 [[SHL]], 4
 ; CHECK-NEXT:    [[MUL:%.*]] = mul nsw i32 [[ADD]], 3
 ; CHECK-NEXT:    ret i32 [[MUL]]
 ;
   %shl = shl nsw i32 %a, 2
   %add = or i32 %shl, 4
   %mul = mul nsw i32 %add, 3
   ret i32 %mul
 }
 
 ; Keep nuw flag in this change, https://alive2.llvm.org/ce/z/awsQrx
 define i32 @PR57278_mul(i32 %a) {
 ; CHECK-LABEL: @PR57278_mul(
-; CHECK-NEXT:    [[MUL0:%.*]] = mul nuw i32 [[A:%.*]], 12
-; CHECK-NEXT:    [[ADD:%.*]] = or i32 [[MUL0]], 3
-; CHECK-NEXT:    [[MUL:%.*]] = mul nuw i32 [[ADD]], 3
+; CHECK-NEXT:    [[TMP1:%.*]] = mul nuw i32 [[A:%.*]], 36
+; CHECK-NEXT:    [[MUL:%.*]] = add nuw i32 [[TMP1]], 9
 ; CHECK-NEXT:    ret i32 [[MUL]]
 ;
   %mul0 = mul nuw i32 %a, 12
   %add = or i32 %mul0, 3
   %mul = mul nuw i32 %add, 3
   ret i32 %mul
 }
 
-; Negative test: Have common bits set
+; Negative test: Have common bits set, https://alive2.llvm.org/ce/z/bHZRh5
 define i32 @PR57278_mul_1(i32 %a) {
 ; CHECK-LABEL: @PR57278_mul_1(
 ; CHECK-NEXT:    [[MUL0:%.*]] = mul nuw i32 [[A:%.*]], 12
 ; CHECK-NEXT:    [[ADD:%.*]] = or i32 [[MUL0]], 4
 ; CHECK-NEXT:    [[MUL:%.*]] = mul nuw i32 [[ADD]], 3
 ; CHECK-NEXT:    ret i32 [[MUL]]
 ;
   %mul0 = mul nuw i32 %a, 12
   %add = or i32 %mul0, 4
   %mul = mul nuw i32 %add, 3
   ret i32 %mul
 }
 
-; https://alive2.llvm.org/ce/z/AXKBjK
+; Test the haveNoCommonBitsSet with assume, https://alive2.llvm.org/ce/z/AXKBjK
 define i32 @PR57278_mul_assume(i32 %a) {
 ; CHECK-LABEL: @PR57278_mul_assume(
 ; CHECK-NEXT:    [[COMBITS:%.*]] = and i32 [[A:%.*]], 3
 ; CHECK-NEXT:    [[NOCOMBITS:%.*]] = icmp eq i32 [[COMBITS]], 0
 ; CHECK-NEXT:    call void @llvm.assume(i1 [[NOCOMBITS]])
-; CHECK-NEXT:    [[ADD:%.*]] = or i32 [[A]], 3
-; CHECK-NEXT:    [[MUL:%.*]] = mul i32 [[ADD]], 5
+; CHECK-NEXT:    [[TMP1:%.*]] = mul i32 [[A]], 5
+; CHECK-NEXT:    [[MUL:%.*]] = add i32 [[TMP1]], 15
 ; CHECK-NEXT:    ret i32 [[MUL]]
 ;
   %combits = and i32 %a , 3
   %nocombits = icmp eq i32 %combits, 0
   call void @llvm.assume(i1 %nocombits)
 
   %add = or i32 %a, 3
   %mul = mul i32 %add, 5
   ret i32 %mul
 }
 
 declare void @llvm.assume(i1)
 
 ; https://alive2.llvm.org/ce/z/XYpv9q
 define <2 x i32> @PR57278_shl_vec(<2 x i32> %v1) {
 ; CHECK-LABEL: @PR57278_shl_vec(
-; CHECK-NEXT:    [[SHL:%.*]] = shl nuw <2 x i32> [[V1:%.*]], <i32 2, i32 3>
-; CHECK-NEXT:    [[ADD:%.*]] = or <2 x i32> [[SHL]], <i32 3, i32 3>
-; CHECK-NEXT:    [[MUL:%.*]] = mul nuw <2 x i32> [[ADD]], <i32 3, i32 3>
+; CHECK-NEXT:    [[TMP1:%.*]] = mul nuw <2 x i32> [[V1:%.*]], <i32 12, i32 24>
+; CHECK-NEXT:    [[MUL:%.*]] = add nuw <2 x i32> [[TMP1]], <i32 9, i32 9>
 ; CHECK-NEXT:    ret <2 x i32> [[MUL]]
 ;
   %shl = shl nuw <2 x i32> %v1, <i32 2, i32 3>
   %add = or <2 x i32> %shl, <i32 3, i32 3>
   %mul = mul nuw <2 x i32> %add, <i32 3, i32 3>
   ret <2 x i32> %mul
 }
 
+; TODO: vector with poison should also be supported, https://alive2.llvm.org/ce/z/XYpv9q
 define <2 x i32> @PR57278_shl_vec_poison(<2 x i32> %v1) {
 ; CHECK-LABEL: @PR57278_shl_vec_poison(
 ; CHECK-NEXT:    [[SHL:%.*]] = shl nuw <2 x i32> [[V1:%.*]], <i32 2, i32 poison>
 ; CHECK-NEXT:    [[ADD:%.*]] = or <2 x i32> [[SHL]], <i32 3, i32 poison>
 ; CHECK-NEXT:    [[MUL:%.*]] = mul nuw <2 x i32> [[ADD]], <i32 3, i32 poison>
 ; CHECK-NEXT:    ret <2 x i32> [[MUL]]
 ;
   %shl = shl nuw <2 x i32> %v1, <i32 2, i32 poison>