[llvm] [Reassociate] Distribute multiply over add to enable factorization (PR #178201)

via llvm-commits llvm-commits at lists.llvm.org
Fri Apr 17 08:57:51 PDT 2026


https://github.com/hazarathayya updated https://github.com/llvm/llvm-project/pull/178201

>From 4960a859c01fd32f4f50dd24e87181644b18e840 Mon Sep 17 00:00:00 2001
From: hazarathayya <hazarathayyayallanki at gmail.com>
Date: Tue, 27 Jan 2026 04:10:58 -0800
Subject: [PATCH 1/6] [Reassociate] Distribute multiply over add to enable
 factorization

Fixes #167190

This patch adds distribution of (A+B)*C -> A*C + B*C when:
- The add has a single use (avoid code bloat)
- Both add operands are non-constant (avoid unprofitable cases)

This enables optimization of patterns like:
  (x*C1) - ((y+x)*C2) -> x*(C1-C2) - (y*C2)

The transformation works by:
1. Pre-processing multiplications to distribute when beneficial
2. Allowing subsequent factorization to extract common factors

Also modifies factorization heuristic to prefer non-constant factors
over constant factors when occurrence counts are equal.

Example:
  Before: (x*8697) - ((y+x)*6436)
  After:  x*2261 + y*(-6436)

Note: Distribution preserves nsw/nuw flags. Comprehensive flag
preservation across all reassociate passes can be addressed separately.
---
 llvm/lib/Transforms/Scalar/Reassociate.cpp    | 118 +++++++++++++++++-
 .../Reassociate/reassociate-distribute.ll     |  17 +++
 2 files changed, 130 insertions(+), 5 deletions(-)
 create mode 100644 llvm/test/Transforms/Reassociate/reassociate-distribute.ll

diff --git a/llvm/lib/Transforms/Scalar/Reassociate.cpp b/llvm/lib/Transforms/Scalar/Reassociate.cpp
index ea6c394740f22..82a501cbcac46 100644
--- a/llvm/lib/Transforms/Scalar/Reassociate.cpp
+++ b/llvm/lib/Transforms/Scalar/Reassociate.cpp
@@ -304,6 +304,84 @@ static BinaryOperator *LowerNegateToMultiply(Instruction *Neg) {
   return Res;
 }
 
+/// Try to distribute multiply over add when beneficial for reassociation
+/// Transforms: (A + B) * C  →  A*C + B*C
+/// This enables further reassociation opportunities
+static BinaryOperator *tryDistributeMul(BinaryOperator *I) {
+  if (I->getOpcode() != Instruction::Mul)
+    return nullptr;
+
+  Value *MulOp0 = I->getOperand(0);
+  Value *MulOp1 = I->getOperand(1);
+
+  // We need one constant and one add
+  ConstantInt *C = dyn_cast<ConstantInt>(MulOp1);
+  BinaryOperator *Add = dyn_cast<BinaryOperator>(MulOp0);
+
+  // Try swapped operands if needed
+  if (!C || !Add) {
+    C = dyn_cast<ConstantInt>(MulOp0);
+    Add = dyn_cast<BinaryOperator>(MulOp1);
+  }
+
+  if (!C || !Add || Add->getOpcode() != Instruction::Add)
+    return nullptr;
+
+  // Only distribute if the add has one use (avoid code bloat)
+  if (!Add->hasOneUse())
+    return nullptr;
+
+  // ===== NEW PROFITABILITY CHECK =====
+  // Only distribute if both add operands are non-constant
+  // This avoids distributing things like (x + 1) * 3 which don't help
+  Value *AddLHS = Add->getOperand(0);
+  Value *AddRHS = Add->getOperand(1);
+
+  bool AIsConstant = isa<Constant>(AddLHS);
+  bool BIsConstant = isa<Constant>(AddRHS);
+
+  // Don't distribute if either operand is a constant
+  // Distributing (x + 5) * 3 → x*3 + 15 doesn't create reassociation
+  // opportunities
+  if (AIsConstant || BIsConstant) {
+    return nullptr;
+  }
+  // ===== END PROFITABILITY CHECK =====
+
+  // DEBUG: Print what we found
+  dbgs() << "DISTRIBUTING: " << *I << "\n";
+  dbgs() << "  Add: " << *Add << "\n";
+  dbgs() << "  Constant: " << *C << "\n";
+
+  Value *A = Add->getOperand(0);
+  Value *B = Add->getOperand(1);
+
+  IRBuilder<> Builder(I);
+
+  // TODO: Reassociate pass needs comprehensive flag preservation.
+  // Distribution preserves flags, but later passes may drop them.
+  // Create with proper flags preserved
+  Value *AC = Builder.CreateMul(A, C);
+  Value *BC = Builder.CreateMul(B, C);
+
+  if (auto *MulInst = dyn_cast<BinaryOperator>(AC))
+    MulInst->copyIRFlags(I); // Copy nsw/nuw from original mul
+  if (auto *MulInst = dyn_cast<BinaryOperator>(BC))
+    MulInst->copyIRFlags(I);
+
+  Value *NewAdd = Builder.CreateAdd(AC, BC);
+  if (auto *AddInst = dyn_cast<BinaryOperator>(NewAdd))
+    AddInst->copyIRFlags(Add); // Copy nsw/nuw from original add
+
+  // After creating the new operations:
+  dbgs() << "Created AC: " << *AC << "\n";
+  dbgs() << "Created BC: " << *BC << "\n";
+  dbgs() << "Created NewAdd: " << *NewAdd << "\n";
+
+  return cast<BinaryOperator>(NewAdd);
+}
+
+
 using RepeatedValue = std::pair<Value *, uint64_t>;
 
 /// Given an associative binary expression, return the leaf
@@ -1603,10 +1681,13 @@ Value *ReassociatePass::OptimizeAdd(Instruction *I,
         continue;
 
       unsigned Occ = ++FactorOccurrences[Factor];
-      if (Occ > MaxOcc) {
+      if (Occ > MaxOcc || 
+          (Occ == MaxOcc && 
+           (isa<Instruction>(Factor) || isa<Argument>(Factor)) && 
+           isa<Constant>(MaxOccVal) && !isa<UndefValue>(MaxOccVal))) {
         MaxOcc = Occ;
         MaxOccVal = Factor;
-      }
+      }     
 
       // If Factor is a negative constant, add the negated value as a factor
       // because we can percolate the negate out.  Watch for minint, which
@@ -1617,7 +1698,10 @@ Value *ReassociatePass::OptimizeAdd(Instruction *I,
           if (!Duplicates.insert(Factor).second)
             continue;
           unsigned Occ = ++FactorOccurrences[Factor];
-          if (Occ > MaxOcc) {
+          if (Occ > MaxOcc || 
+              (Occ == MaxOcc && 
+               (isa<Instruction>(Factor) || isa<Argument>(Factor)) && 
+               isa<Constant>(MaxOccVal) && !isa<UndefValue>(MaxOccVal))) {
             MaxOcc = Occ;
             MaxOccVal = Factor;
           }
@@ -1630,10 +1714,13 @@ Value *ReassociatePass::OptimizeAdd(Instruction *I,
           if (!Duplicates.insert(Factor).second)
             continue;
           unsigned Occ = ++FactorOccurrences[Factor];
-          if (Occ > MaxOcc) {
+          if (Occ > MaxOcc || 
+              (Occ == MaxOcc && 
+               (isa<Instruction>(Factor) || isa<Argument>(Factor)) && 
+               isa<Constant>(MaxOccVal) && !isa<UndefValue>(MaxOccVal))) {
             MaxOcc = Occ;
             MaxOccVal = Factor;
-          }
+          }         
         }
       }
     }
@@ -2566,6 +2653,27 @@ PreservedAnalyses ReassociatePass::run(Function &F, FunctionAnalysisManager &) {
 
   MadeChange = false;
 
+  // Pre-process: Distribute multiplications to enable reassociation
+  for (BasicBlock *BI : RPOT) {
+    for (BasicBlock::iterator II = BI->begin(), IE = BI->end(); II != IE;) {
+      Instruction *Inst = &*II;
+      ++II; // Advance before modification
+
+      if (auto *Mul = dyn_cast<BinaryOperator>(Inst)) {
+        if (Mul->getOpcode() == Instruction::Mul) {
+          if (BinaryOperator *Dist = tryDistributeMul(Mul)) {
+            LLVM_DEBUG(dbgs()
+                       << "Distributed: " << *Mul << " -> " << *Dist << "\n");
+            Mul->replaceAllUsesWith(Dist);
+            Mul->eraseFromParent();
+            MadeChange = true;
+          }
+        }
+      }
+    }
+  }
+
+
   // Traverse the same blocks that were analysed by BuildRankMap.
   for (BasicBlock *BI : RPOT) {
     assert(RankMap.count(&*BI) && "BB should be ranked.");
diff --git a/llvm/test/Transforms/Reassociate/reassociate-distribute.ll b/llvm/test/Transforms/Reassociate/reassociate-distribute.ll
new file mode 100644
index 0000000000000..426061839f905
--- /dev/null
+++ b/llvm/test/Transforms/Reassociate/reassociate-distribute.ll
@@ -0,0 +1,17 @@
+; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
+; RUN: opt -passes=reassociate -S < %s | FileCheck %s
+
+; Verify that distribution doesn't happen when add has multiple uses
+; Instead, reassociate factors the common add operand
+define i16 @test_no_distribute_multiuse(i16 %x, i16 %y) {
+; CHECK-LABEL: @test_no_distribute_multiuse(
+; CHECK-NEXT:    [[ADD:%.*]] = add i16 [[Y:%.*]], [[X:%.*]]
+; CHECK-NEXT:    [[REASS_MUL:%.*]] = mul i16 [[ADD]], 30
+; CHECK-NEXT:    ret i16 [[REASS_MUL]]
+;
+  %add = add i16 %x, %y
+  %mul1 = mul i16 %add, 10
+  %mul2 = mul i16 %add, 20
+  %result = add i16 %mul1, %mul2
+  ret i16 %result
+}

>From 1e1b71961134a1a8964119cb4709bc31712ceae4 Mon Sep 17 00:00:00 2001
From: hazarathayya <hazarathayyayallanki at gmail.com>
Date: Thu, 19 Feb 2026 01:48:05 -0800
Subject: [PATCH 2/6] fixed clang-format corrections

---
 llvm/lib/Transforms/Scalar/Reassociate.cpp | 24 ++++++++++------------
 1 file changed, 11 insertions(+), 13 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/Reassociate.cpp b/llvm/lib/Transforms/Scalar/Reassociate.cpp
index 82a501cbcac46..cdad4bb2f092e 100644
--- a/llvm/lib/Transforms/Scalar/Reassociate.cpp
+++ b/llvm/lib/Transforms/Scalar/Reassociate.cpp
@@ -381,7 +381,6 @@ static BinaryOperator *tryDistributeMul(BinaryOperator *I) {
   return cast<BinaryOperator>(NewAdd);
 }
 
-
 using RepeatedValue = std::pair<Value *, uint64_t>;
 
 /// Given an associative binary expression, return the leaf
@@ -1681,13 +1680,13 @@ Value *ReassociatePass::OptimizeAdd(Instruction *I,
         continue;
 
       unsigned Occ = ++FactorOccurrences[Factor];
-      if (Occ > MaxOcc || 
-          (Occ == MaxOcc && 
-           (isa<Instruction>(Factor) || isa<Argument>(Factor)) && 
+      if (Occ > MaxOcc ||
+          (Occ == MaxOcc &&
+           (isa<Instruction>(Factor) || isa<Argument>(Factor)) &&
            isa<Constant>(MaxOccVal) && !isa<UndefValue>(MaxOccVal))) {
         MaxOcc = Occ;
         MaxOccVal = Factor;
-      }     
+      }
 
       // If Factor is a negative constant, add the negated value as a factor
       // because we can percolate the negate out.  Watch for minint, which
@@ -1698,9 +1697,9 @@ Value *ReassociatePass::OptimizeAdd(Instruction *I,
           if (!Duplicates.insert(Factor).second)
             continue;
           unsigned Occ = ++FactorOccurrences[Factor];
-          if (Occ > MaxOcc || 
-              (Occ == MaxOcc && 
-               (isa<Instruction>(Factor) || isa<Argument>(Factor)) && 
+          if (Occ > MaxOcc ||
+              (Occ == MaxOcc &&
+               (isa<Instruction>(Factor) || isa<Argument>(Factor)) &&
                isa<Constant>(MaxOccVal) && !isa<UndefValue>(MaxOccVal))) {
             MaxOcc = Occ;
             MaxOccVal = Factor;
@@ -1714,13 +1713,13 @@ Value *ReassociatePass::OptimizeAdd(Instruction *I,
           if (!Duplicates.insert(Factor).second)
             continue;
           unsigned Occ = ++FactorOccurrences[Factor];
-          if (Occ > MaxOcc || 
-              (Occ == MaxOcc && 
-               (isa<Instruction>(Factor) || isa<Argument>(Factor)) && 
+          if (Occ > MaxOcc ||
+              (Occ == MaxOcc &&
+               (isa<Instruction>(Factor) || isa<Argument>(Factor)) &&
                isa<Constant>(MaxOccVal) && !isa<UndefValue>(MaxOccVal))) {
             MaxOcc = Occ;
             MaxOccVal = Factor;
-          }         
+          }
         }
       }
     }
@@ -2673,7 +2672,6 @@ PreservedAnalyses ReassociatePass::run(Function &F, FunctionAnalysisManager &) {
     }
   }
 
-
   // Traverse the same blocks that were analysed by BuildRankMap.
   for (BasicBlock *BI : RPOT) {
     assert(RankMap.count(&*BI) && "BB should be ranked.");

>From a3b1302cb4a90ed354c822939ab893925ef670ce Mon Sep 17 00:00:00 2001
From: hazarathayya <hazarathayyayallanki at gmail.com>
Date: Thu, 12 Mar 2026 07:54:59 -0700
Subject: [PATCH 3/6] changed to optimizeinst, added more test cases

---
 llvm/lib/Transforms/Scalar/Reassociate.cpp    | 101 ++++++++---
 .../Reassociate/reassociate-distribute.ll     | 167 ++++++++++++++++--
 2 files changed, 225 insertions(+), 43 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/Reassociate.cpp b/llvm/lib/Transforms/Scalar/Reassociate.cpp
index cdad4bb2f092e..234275eac3fbf 100644
--- a/llvm/lib/Transforms/Scalar/Reassociate.cpp
+++ b/llvm/lib/Transforms/Scalar/Reassociate.cpp
@@ -331,7 +331,6 @@ static BinaryOperator *tryDistributeMul(BinaryOperator *I) {
   if (!Add->hasOneUse())
     return nullptr;
 
-  // ===== NEW PROFITABILITY CHECK =====
   // Only distribute if both add operands are non-constant
   // This avoids distributing things like (x + 1) * 3 which don't help
   Value *AddLHS = Add->getOperand(0);
@@ -346,12 +345,11 @@ static BinaryOperator *tryDistributeMul(BinaryOperator *I) {
   if (AIsConstant || BIsConstant) {
     return nullptr;
   }
-  // ===== END PROFITABILITY CHECK =====
 
   // DEBUG: Print what we found
-  dbgs() << "DISTRIBUTING: " << *I << "\n";
-  dbgs() << "  Add: " << *Add << "\n";
-  dbgs() << "  Constant: " << *C << "\n";
+  LLVM_DEBUG(dbgs() << "DISTRIBUTING: " << *I << "\n");
+  LLVM_DEBUG(dbgs() << "  Add: " << *Add << "\n");
+  LLVM_DEBUG(dbgs() << "  Constant: " << *C << "\n");
 
   Value *A = Add->getOperand(0);
   Value *B = Add->getOperand(1);
@@ -374,9 +372,9 @@ static BinaryOperator *tryDistributeMul(BinaryOperator *I) {
     AddInst->copyIRFlags(Add); // Copy nsw/nuw from original add
 
   // After creating the new operations:
-  dbgs() << "Created AC: " << *AC << "\n";
-  dbgs() << "Created BC: " << *BC << "\n";
-  dbgs() << "Created NewAdd: " << *NewAdd << "\n";
+  LLVM_DEBUG(dbgs() << "Created AC: " << *AC << "\n");
+  LLVM_DEBUG(dbgs() << "Created BC: " << *BC << "\n");
+  LLVM_DEBUG(dbgs() << "Created NewAdd: " << *NewAdd << "\n");
 
   return cast<BinaryOperator>(NewAdd);
 }
@@ -1040,6 +1038,61 @@ static BinaryOperator *convertOrWithNoCommonBitsToAdd(Instruction *Or) {
   return New;
 }
 
+static bool ShouldBreakUpDistribution(Instruction *Mul) {
+  int ConstSide = isa<Constant>(Mul->getOperand(0)) ? 0 : 1;
+  Value *MulConst   = Mul->getOperand(ConstSide);
+  Value *AddSubVal  = Mul->getOperand(1 - ConstSide);
+
+  if (!isa<Constant>(MulConst))
+    return false;
+
+  // Only handle integer mul — FMul requires FastMathFlags handling
+  if (isa<FPMathOperator>(Mul))
+    return false;
+
+  if (!isReassociableOp(AddSubVal, Instruction::Add, Instruction::FAdd) &&
+      !isReassociableOp(AddSubVal, Instruction::Sub, Instruction::FSub))
+    return false;
+
+  User *AddSubInst = cast<User>(AddSubVal);
+
+  for (Value *Operand : AddSubInst->operands()) {
+    for (User *U : Operand->users()) {
+      Instruction *I = dyn_cast<Instruction>(U);
+      if (!I || I == Mul)
+        continue;
+      if ((I->getOpcode() == Instruction::Mul ||
+          I->getOpcode() == Instruction::FMul) &&
+          (isa<Constant>(I->getOperand(0)) ||
+          isa<Constant>(I->getOperand(1)))){ 
+        return true; 
+      }
+   }
+  }
+  return false;
+}
+
+static BinaryOperator *BreakUpDistribute(Instruction *Mul,
+                                       ReassociatePass::OrderedSet &ToRedo) {
+
+  int ConstSide = isa<Constant>(Mul->getOperand(0)) ? 0 : 1;
+  Instruction *AddSub = cast<Instruction>(Mul->getOperand(1-ConstSide));
+
+  BinaryOperator *M1 = CreateMul(AddSub->getOperand(0), Mul->getOperand(ConstSide), "Mul1", Mul->getIterator(), Mul);
+  BinaryOperator *M2 = CreateMul(AddSub->getOperand(1), Mul->getOperand(ConstSide), "Mul2", Mul->getIterator(), Mul);
+  BinaryOperator *Result;
+  if (isReassociableOp(AddSub, Instruction::Add, Instruction::FAdd))
+      Result = CreateAdd(M1, M2, "DistAdd", Mul->getIterator(), Mul);
+  else {
+      Result = BinaryOperator::CreateSub(M1, M2, "DistSub", Mul->getIterator());
+  }
+
+  Mul->replaceAllUsesWith(Result);
+  Result->setDebugLoc(Mul->getDebugLoc());
+
+  return Result;
+}
+
 /// Return true if we should break up this subtract of X-Y into (X + -Y).
 static bool ShouldBreakUpSubtract(Instruction *Sub) {
   // If this is a negation, we can't split it up!
@@ -2228,6 +2281,7 @@ Instruction *ReassociatePass::canonicalizeNegFPConstants(Instruction *I) {
 /// instructions is not allowed.
 void ReassociatePass::OptimizeInst(Instruction *I) {
   // Only consider operations that we understand.
+  errs() << "----------------------- Optimizing " << *I << "-------------------------------------\n";
   if (!isa<UnaryOperator>(I) && !isa<BinaryOperator>(I))
     return;
 
@@ -2283,6 +2337,16 @@ void ReassociatePass::OptimizeInst(Instruction *I) {
     I = NI;
   }
 
+  if(I->getOpcode() == Instruction::Mul && 
+      ShouldBreakUpDistribution(I)){
+    // handle NSW/NUW flags
+    errs() << "we are getting this I  = "<< *I << "\n";
+    BinaryOperator *Result = BreakUpDistribute(I, RedoInsts);
+    RedoInsts.insert(I);
+    MadeChange = true;
+    I = Result;
+  }
+
   // If this is a subtract instruction which is not already in negate form,
   // see if we can convert it to X+-Y.
   if (I->getOpcode() == Instruction::Sub) {
@@ -2634,6 +2698,7 @@ PreservedAnalyses ReassociatePass::run(Function &F, FunctionAnalysisManager &) {
   // BuildRankMap to pre calculate ranks correctly. It also excludes dead basic
   // blocks (it has been seen that the analysis in this pass could hang when
   // analysing dead basic blocks).
+  errs() << "-----------------------Pass start -------------------------------------\n";
   ReversePostOrderTraversal<Function *> RPOT(&F);
 
   // Calculate the rank map for F.
@@ -2652,26 +2717,6 @@ PreservedAnalyses ReassociatePass::run(Function &F, FunctionAnalysisManager &) {
 
   MadeChange = false;
 
-  // Pre-process: Distribute multiplications to enable reassociation
-  for (BasicBlock *BI : RPOT) {
-    for (BasicBlock::iterator II = BI->begin(), IE = BI->end(); II != IE;) {
-      Instruction *Inst = &*II;
-      ++II; // Advance before modification
-
-      if (auto *Mul = dyn_cast<BinaryOperator>(Inst)) {
-        if (Mul->getOpcode() == Instruction::Mul) {
-          if (BinaryOperator *Dist = tryDistributeMul(Mul)) {
-            LLVM_DEBUG(dbgs()
-                       << "Distributed: " << *Mul << " -> " << *Dist << "\n");
-            Mul->replaceAllUsesWith(Dist);
-            Mul->eraseFromParent();
-            MadeChange = true;
-          }
-        }
-      }
-    }
-  }
-
   // Traverse the same blocks that were analysed by BuildRankMap.
   for (BasicBlock *BI : RPOT) {
     assert(RankMap.count(&*BI) && "BB should be ranked.");
diff --git a/llvm/test/Transforms/Reassociate/reassociate-distribute.ll b/llvm/test/Transforms/Reassociate/reassociate-distribute.ll
index 426061839f905..3889be51cc0c0 100644
--- a/llvm/test/Transforms/Reassociate/reassociate-distribute.ll
+++ b/llvm/test/Transforms/Reassociate/reassociate-distribute.ll
@@ -1,17 +1,154 @@
 ; NOTE: Assertions have been autogenerated by utils/update_test_checks.py
-; RUN: opt -passes=reassociate -S < %s | FileCheck %s
-
-; Verify that distribution doesn't happen when add has multiple uses
-; Instead, reassociate factors the common add operand
-define i16 @test_no_distribute_multiuse(i16 %x, i16 %y) {
-; CHECK-LABEL: @test_no_distribute_multiuse(
-; CHECK-NEXT:    [[ADD:%.*]] = add i16 [[Y:%.*]], [[X:%.*]]
-; CHECK-NEXT:    [[REASS_MUL:%.*]] = mul i16 [[ADD]], 30
-; CHECK-NEXT:    ret i16 [[REASS_MUL]]
-;
-  %add = add i16 %x, %y
-  %mul1 = mul i16 %add, 10
-  %mul2 = mul i16 %add, 20
-  %result = add i16 %mul1, %mul2
-  ret i16 %result
+; RUN: opt -passes=reassociate -S %s | FileCheck %s
+
+; ---- SHOULD transform ----
+
+; Basic case: a*8697 - (a+b)*6436 → a*2261 + b*(-6436)
+define i32 @test_basic(i32 %a, i32 %b) {
+; CHECK-LABEL: @test_basic(
+; CHECK-NEXT:    [[MUL1_NEG:%.*]] = mul i32 [[B:%.*]], -6436
+; CHECK-NEXT:    [[REASS_MUL:%.*]] = mul i32 [[A:%.*]], 2261
+; CHECK-NEXT:    [[SUB:%.*]] = add i32 [[REASS_MUL]], [[MUL1_NEG]]
+; CHECK-NEXT:    ret i32 [[SUB]]
+;
+  %mul  = mul nsw i32 %a, 8697
+  %add  = add nsw i32 %a, %b
+  %mul1 = mul nsw i32 %add, 6436
+  %sub  = sub nsw i32 %mul, %mul1
+  ret i32 %sub
+}
+
+; Constant on left side of mul: a*100 + 3*(a+b)
+define i32 @test_const_on_left(i32 %a, i32 %b) {
+; CHECK-LABEL: @test_const_on_left(
+; CHECK-NEXT:    [[MUL1:%.*]] = mul i32 [[B:%.*]], 7
+; CHECK-NEXT:    [[REASS_MUL:%.*]] = mul i32 [[A:%.*]], 107
+; CHECK-NEXT:    [[ADD2:%.*]] = add i32 [[MUL1]], [[REASS_MUL]]
+; CHECK-NEXT:    ret i32 [[ADD2]]
+;
+  %mul  = mul nsw i32 %a, 100
+  %add  = add nsw i32 %a, %b
+  %mul1 = mul nsw i32 7, %add
+  %add2 = add nsw i32 %mul, %mul1
+  ret i32 %add2
+}
+
+; Sub instead of add inside: a*500 - (a-b)*300
+define i32 @test_with_sub(i32 %a, i32 %b) {
+; CHECK-LABEL: @test_with_sub(
+; CHECK-NEXT:    [[MUL2_NEG_NEG:%.*]] = mul i32 [[B:%.*]], 300
+; CHECK-NEXT:    [[REASS_MUL:%.*]] = mul i32 [[A:%.*]], 200
+; CHECK-NEXT:    [[SUB:%.*]] = add i32 [[REASS_MUL]], [[MUL2_NEG_NEG]]
+; CHECK-NEXT:    ret i32 [[SUB]]
+;
+  %mul  = mul nsw i32 %a, 500
+  %sub1 = sub nsw i32 %a, %b
+  %mul1 = mul nsw i32 %sub1, 300
+  %sub  = sub nsw i32 %mul, %mul1
+  ret i32 %sub
+}
+
+; b operand matches: b*11 + (a+b)*4
+define i32 @test_b_matches(i32 %a, i32 %b) {
+; CHECK-LABEL: @test_b_matches(
+; CHECK-NEXT:    [[MUL2:%.*]] = mul i32 [[A:%.*]], 4
+; CHECK-NEXT:    [[REASS_MUL:%.*]] = mul i32 [[B:%.*]], 15
+; CHECK-NEXT:    [[ADD2:%.*]] = add i32 [[REASS_MUL]], [[MUL2]]
+; CHECK-NEXT:    ret i32 [[ADD2]]
+;
+  %mul  = mul nsw i32 %b, 11
+  %add  = add nsw i32 %a, %b
+  %mul1 = mul nsw i32 %add, 4
+  %add2 = add nsw i32 %mul, %mul1
+  ret i32 %add2
+}
+
+; i16 type: x*10 + (x+y)*20
+define i16 @test_i16(i16 %x, i16 %y) {
+; CHECK-LABEL: @test_i16(
+; CHECK-NEXT:    [[MUL1:%.*]] = mul i16 [[Y:%.*]], 20
+; CHECK-NEXT:    [[REASS_MUL:%.*]] = mul i16 [[X:%.*]], 30
+; CHECK-NEXT:    [[ADD2:%.*]] = add i16 [[MUL1]], [[REASS_MUL]]
+; CHECK-NEXT:    ret i16 [[ADD2]]
+;
+  %mul1 = mul nsw i16 %x, 10
+  %add  = add nsw i16 %x, %y
+  %mul2 = mul nsw i16 %add, 20
+  %add2 = add i16 %mul1, %mul2
+  ret i16 %add2
+}
+
+; NSW/NUW must be dropped on new instructions after distribution
+; (x+y) *nsw 20 does NOT imply (x *nsw 20) + (y *nsw 20)
+define i16 @test_nsw_dropped(i16 %x, i16 %y) {
+; CHECK-LABEL: @test_nsw_dropped(
+; CHECK-NEXT:    [[MUL1:%.*]] = mul i16 [[Y:%.*]], 25
+; CHECK-NEXT:    [[REASS_MUL:%.*]] = mul i16 [[X:%.*]], 40
+; CHECK-NEXT:    [[ADD2:%.*]] = add i16 [[MUL1]], [[REASS_MUL]]
+; CHECK-NEXT:    ret i16 [[ADD2]]
+;
+  %mul1 = mul nsw i16 %x, 15
+  %add  = add nsw i16 %x, %y
+  %mul2 = mul nsw i16 %add, 25
+  %add2 = add i16 %mul1, %mul2
+  ret i16 %add2
+}
+
+; ---- Should NOT transform ----
+
+; No sibling constant mul — nothing to combine with
+define i32 @test_no_match(i32 %a, i32 %b) {
+; CHECK-LABEL: @test_no_match(
+; CHECK-NEXT:    [[ADD:%.*]] = add nsw i32 [[B:%.*]], [[A:%.*]]
+; CHECK-NEXT:    [[MUL:%.*]] = mul nsw i32 [[ADD]], 999
+; CHECK-NEXT:    ret i32 [[MUL]]
+; CHECK-NOT:     Mul1
+; CHECK-NOT:     Mul2
+; CHECK-NOT:     DistAdd
+; CHECK-NOT:     DistSub;
+;
+  %add = add nsw i32 %a, %b
+  %mul = mul nsw i32 %add, 999
+  ret i32 %mul
+}
+
+; AddSub has multiple uses — oneUse required for distribution
+define i32 @test_multi_use_add(i32 %a, i32 %b) {
+; CHECK-LABEL: @test_multi_use_add(
+; CHECK-NEXT:    [[MUL1:%.*]] = mul i32 [[B:%.*]], -442
+; CHECK-NEXT:    [[REASS_MUL:%.*]] = mul i32 [[A:%.*]], 335
+; CHECK-NEXT:    [[RET:%.*]] = add i32 [[REASS_MUL]], [[MUL1]]
+; CHECK-NEXT:    ret i32 [[RET]]
+; CHECK-NOT:     Mul1
+; CHECK-NOT:     Mul2
+; CHECK-NOT:     DistAdd
+; CHECK-NOT:     DistSub
+;
+  %mul  = mul nsw i32 %a, 777
+  %add  = add nsw i32 %a, %b
+  %mul1 = mul nsw i32 %add, 444
+  %mul2 = mul nsw i32 %add, 2
+  %sub  = sub nsw i32 %mul, %mul1
+  %ret  = add nsw i32 %sub, %mul2
+  ret i32 %ret
+}
+
+; Sibling mul has no constant operand — should not trigger
+define i32 @test_no_const_on_other_mul(i32 %a, i32 %b, i32 %c) {
+; CHECK-LABEL: @test_no_const_on_other_mul(
+; CHECK-NEXT:    [[MUL:%.*]] = mul nsw i32 [[C:%.*]], [[A:%.*]]
+; CHECK-NEXT:    [[ADD:%.*]] = add nsw i32 [[B:%.*]], [[A]]
+; CHECK-NEXT:    [[MUL1:%.*]] = mul nsw i32 [[ADD]], 555
+; CHECK-NEXT:    [[SUB:%.*]] = sub nsw i32 [[MUL]], [[MUL1]]
+; CHECK-NEXT:    ret i32 [[SUB]]
+; CHECK-NOT:     Mul1
+; CHECK-NOT:     Mul2
+; CHECK-NOT:     DistAdd
+; CHECK-NOT:     DistSub
+;
+  %mul  = mul nsw i32 %a, %c
+  %add  = add nsw i32 %a, %b
+  %mul1 = mul nsw i32 %add, 555
+  %sub  = sub nsw i32 %mul, %mul1
+  ret i32 %sub
 }

>From 34cf837d9e6e41e13b6d23b6b1faac85e3032039 Mon Sep 17 00:00:00 2001
From: hazarathayya <hazarathayyayallanki at gmail.com>
Date: Thu, 12 Mar 2026 07:59:04 -0700
Subject: [PATCH 4/6] removed trydistriubted() fn, and test comments

---
 llvm/lib/Transforms/Scalar/Reassociate.cpp | 79 ----------------------
 1 file changed, 79 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/Reassociate.cpp b/llvm/lib/Transforms/Scalar/Reassociate.cpp
index 234275eac3fbf..3c283bae557f3 100644
--- a/llvm/lib/Transforms/Scalar/Reassociate.cpp
+++ b/llvm/lib/Transforms/Scalar/Reassociate.cpp
@@ -304,81 +304,6 @@ static BinaryOperator *LowerNegateToMultiply(Instruction *Neg) {
   return Res;
 }
 
-/// Try to distribute multiply over add when beneficial for reassociation
-/// Transforms: (A + B) * C  →  A*C + B*C
-/// This enables further reassociation opportunities
-static BinaryOperator *tryDistributeMul(BinaryOperator *I) {
-  if (I->getOpcode() != Instruction::Mul)
-    return nullptr;
-
-  Value *MulOp0 = I->getOperand(0);
-  Value *MulOp1 = I->getOperand(1);
-
-  // We need one constant and one add
-  ConstantInt *C = dyn_cast<ConstantInt>(MulOp1);
-  BinaryOperator *Add = dyn_cast<BinaryOperator>(MulOp0);
-
-  // Try swapped operands if needed
-  if (!C || !Add) {
-    C = dyn_cast<ConstantInt>(MulOp0);
-    Add = dyn_cast<BinaryOperator>(MulOp1);
-  }
-
-  if (!C || !Add || Add->getOpcode() != Instruction::Add)
-    return nullptr;
-
-  // Only distribute if the add has one use (avoid code bloat)
-  if (!Add->hasOneUse())
-    return nullptr;
-
-  // Only distribute if both add operands are non-constant
-  // This avoids distributing things like (x + 1) * 3 which don't help
-  Value *AddLHS = Add->getOperand(0);
-  Value *AddRHS = Add->getOperand(1);
-
-  bool AIsConstant = isa<Constant>(AddLHS);
-  bool BIsConstant = isa<Constant>(AddRHS);
-
-  // Don't distribute if either operand is a constant
-  // Distributing (x + 5) * 3 → x*3 + 15 doesn't create reassociation
-  // opportunities
-  if (AIsConstant || BIsConstant) {
-    return nullptr;
-  }
-
-  // DEBUG: Print what we found
-  LLVM_DEBUG(dbgs() << "DISTRIBUTING: " << *I << "\n");
-  LLVM_DEBUG(dbgs() << "  Add: " << *Add << "\n");
-  LLVM_DEBUG(dbgs() << "  Constant: " << *C << "\n");
-
-  Value *A = Add->getOperand(0);
-  Value *B = Add->getOperand(1);
-
-  IRBuilder<> Builder(I);
-
-  // TODO: Reassociate pass needs comprehensive flag preservation.
-  // Distribution preserves flags, but later passes may drop them.
-  // Create with proper flags preserved
-  Value *AC = Builder.CreateMul(A, C);
-  Value *BC = Builder.CreateMul(B, C);
-
-  if (auto *MulInst = dyn_cast<BinaryOperator>(AC))
-    MulInst->copyIRFlags(I); // Copy nsw/nuw from original mul
-  if (auto *MulInst = dyn_cast<BinaryOperator>(BC))
-    MulInst->copyIRFlags(I);
-
-  Value *NewAdd = Builder.CreateAdd(AC, BC);
-  if (auto *AddInst = dyn_cast<BinaryOperator>(NewAdd))
-    AddInst->copyIRFlags(Add); // Copy nsw/nuw from original add
-
-  // After creating the new operations:
-  LLVM_DEBUG(dbgs() << "Created AC: " << *AC << "\n");
-  LLVM_DEBUG(dbgs() << "Created BC: " << *BC << "\n");
-  LLVM_DEBUG(dbgs() << "Created NewAdd: " << *NewAdd << "\n");
-
-  return cast<BinaryOperator>(NewAdd);
-}
-
 using RepeatedValue = std::pair<Value *, uint64_t>;
 
 /// Given an associative binary expression, return the leaf
@@ -2281,7 +2206,6 @@ Instruction *ReassociatePass::canonicalizeNegFPConstants(Instruction *I) {
 /// instructions is not allowed.
 void ReassociatePass::OptimizeInst(Instruction *I) {
   // Only consider operations that we understand.
-  errs() << "----------------------- Optimizing " << *I << "-------------------------------------\n";
   if (!isa<UnaryOperator>(I) && !isa<BinaryOperator>(I))
     return;
 
@@ -2339,8 +2263,6 @@ void ReassociatePass::OptimizeInst(Instruction *I) {
 
   if(I->getOpcode() == Instruction::Mul && 
       ShouldBreakUpDistribution(I)){
-    // handle NSW/NUW flags
-    errs() << "we are getting this I  = "<< *I << "\n";
     BinaryOperator *Result = BreakUpDistribute(I, RedoInsts);
     RedoInsts.insert(I);
     MadeChange = true;
@@ -2698,7 +2620,6 @@ PreservedAnalyses ReassociatePass::run(Function &F, FunctionAnalysisManager &) {
   // BuildRankMap to pre calculate ranks correctly. It also excludes dead basic
   // blocks (it has been seen that the analysis in this pass could hang when
   // analysing dead basic blocks).
-  errs() << "-----------------------Pass start -------------------------------------\n";
   ReversePostOrderTraversal<Function *> RPOT(&F);
 
   // Calculate the rank map for F.

>From 60272820215dd56f3bbabe00aa77e527a7e2ce03 Mon Sep 17 00:00:00 2001
From: hazarathayya <hazarathayyayallanki at gmail.com>
Date: Tue, 17 Mar 2026 08:03:22 -0700
Subject: [PATCH 5/6] modified distribution calling to OptimizeInst() fn

---
 llvm/lib/Transforms/Scalar/Reassociate.cpp    | 38 ++++++++++++-------
 .../Reassociate/reassociate-distribute.ll     |  8 ++--
 2 files changed, 29 insertions(+), 17 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/Reassociate.cpp b/llvm/lib/Transforms/Scalar/Reassociate.cpp
index 3c283bae557f3..dbe141576eb31 100644
--- a/llvm/lib/Transforms/Scalar/Reassociate.cpp
+++ b/llvm/lib/Transforms/Scalar/Reassociate.cpp
@@ -963,6 +963,10 @@ static BinaryOperator *convertOrWithNoCommonBitsToAdd(Instruction *Or) {
   return New;
 }
 
+/// Return true if Mul is of the form (X+Y)*C or (X-Y)*C where C is a
+/// constant, and there exists a sibling instruction of the form X*C' or Y*C'
+/// in the same expression — indicating that distribution followed by
+/// factoring will reduce the instruction count.
 static bool ShouldBreakUpDistribution(Instruction *Mul) {
   int ConstSide = isa<Constant>(Mul->getOperand(0)) ? 0 : 1;
   Value *MulConst   = Mul->getOperand(ConstSide);
@@ -979,17 +983,20 @@ static bool ShouldBreakUpDistribution(Instruction *Mul) {
       !isReassociableOp(AddSubVal, Instruction::Sub, Instruction::FSub))
     return false;
 
-  User *AddSubInst = cast<User>(AddSubVal);
+  Instruction *AddSubInst = cast<Instruction>(AddSubVal);
 
   for (Value *Operand : AddSubInst->operands()) {
+    if (isa<Constant>(Operand))
+      continue;
     for (User *U : Operand->users()) {
       Instruction *I = dyn_cast<Instruction>(U);
-      if (!I || I == Mul)
+      if (!I || I == AddSubInst)
         continue;
-      if ((I->getOpcode() == Instruction::Mul ||
-          I->getOpcode() == Instruction::FMul) &&
-          (isa<Constant>(I->getOperand(0)) ||
-          isa<Constant>(I->getOperand(1)))){ 
+      if ((I->getOpcode() == Instruction::Mul) && !(I->use_empty()) &&
+          (isa<Constant>(I->getOperand(0)) || isa<Constant>(I->getOperand(1)))){ 
+        int SibConstSide = isa<Constant>(I->getOperand(0)) ? 0 : 1;
+        if (I->getOperand(SibConstSide) == Mul->getOperand(ConstSide))
+            continue;
         return true; 
       }
    }
@@ -997,20 +1004,23 @@ static bool ShouldBreakUpDistribution(Instruction *Mul) {
   return false;
 }
 
+/// Distribute Mul of the form (X+Y)*C into X*C + Y*C.
+/// For the sub case (X-Y)*C, the second term uses -C to avoid
+/// introducing a negation instruction.
 static BinaryOperator *BreakUpDistribute(Instruction *Mul,
                                        ReassociatePass::OrderedSet &ToRedo) {
 
   int ConstSide = isa<Constant>(Mul->getOperand(0)) ? 0 : 1;
   Instruction *AddSub = cast<Instruction>(Mul->getOperand(1-ConstSide));
+  Constant *C = cast<Constant>(Mul->getOperand(ConstSide));
+  Constant *C2 = (AddSub->getOpcode() == Instruction::Sub)
+                    ? cast<Constant>(ConstantExpr::getNeg(C))
+                    : C;
 
-  BinaryOperator *M1 = CreateMul(AddSub->getOperand(0), Mul->getOperand(ConstSide), "Mul1", Mul->getIterator(), Mul);
-  BinaryOperator *M2 = CreateMul(AddSub->getOperand(1), Mul->getOperand(ConstSide), "Mul2", Mul->getIterator(), Mul);
+  BinaryOperator *M1 = CreateMul(AddSub->getOperand(0), C, "Mul1", Mul->getIterator(), Mul);
+  BinaryOperator *M2 = CreateMul(AddSub->getOperand(1), C2, "Mul2", Mul->getIterator(), Mul);
   BinaryOperator *Result;
-  if (isReassociableOp(AddSub, Instruction::Add, Instruction::FAdd))
-      Result = CreateAdd(M1, M2, "DistAdd", Mul->getIterator(), Mul);
-  else {
-      Result = BinaryOperator::CreateSub(M1, M2, "DistSub", Mul->getIterator());
-  }
+  Result = CreateAdd(M1, M2, "DistAdd", Mul->getIterator(), Mul);
 
   Mul->replaceAllUsesWith(Result);
   Result->setDebugLoc(Mul->getDebugLoc());
@@ -2265,8 +2275,10 @@ void ReassociatePass::OptimizeInst(Instruction *I) {
       ShouldBreakUpDistribution(I)){
     BinaryOperator *Result = BreakUpDistribute(I, RedoInsts);
     RedoInsts.insert(I);
+    RedoInsts.insert(Result);
     MadeChange = true;
     I = Result;
+    return;
   }
 
   // If this is a subtract instruction which is not already in negate form,
diff --git a/llvm/test/Transforms/Reassociate/reassociate-distribute.ll b/llvm/test/Transforms/Reassociate/reassociate-distribute.ll
index 3889be51cc0c0..5426e84ebcd2d 100644
--- a/llvm/test/Transforms/Reassociate/reassociate-distribute.ll
+++ b/llvm/test/Transforms/Reassociate/reassociate-distribute.ll
@@ -23,7 +23,7 @@ define i32 @test_const_on_left(i32 %a, i32 %b) {
 ; CHECK-LABEL: @test_const_on_left(
 ; CHECK-NEXT:    [[MUL1:%.*]] = mul i32 [[B:%.*]], 7
 ; CHECK-NEXT:    [[REASS_MUL:%.*]] = mul i32 [[A:%.*]], 107
-; CHECK-NEXT:    [[ADD2:%.*]] = add i32 [[MUL1]], [[REASS_MUL]]
+; CHECK-NEXT:    [[ADD2:%.*]] = add i32 [[REASS_MUL]], [[MUL1]]
 ; CHECK-NEXT:    ret i32 [[ADD2]]
 ;
   %mul  = mul nsw i32 %a, 100
@@ -68,7 +68,7 @@ define i16 @test_i16(i16 %x, i16 %y) {
 ; CHECK-LABEL: @test_i16(
 ; CHECK-NEXT:    [[MUL1:%.*]] = mul i16 [[Y:%.*]], 20
 ; CHECK-NEXT:    [[REASS_MUL:%.*]] = mul i16 [[X:%.*]], 30
-; CHECK-NEXT:    [[ADD2:%.*]] = add i16 [[MUL1]], [[REASS_MUL]]
+; CHECK-NEXT:    [[ADD2:%.*]] = add i16 [[REASS_MUL]], [[MUL1]]
 ; CHECK-NEXT:    ret i16 [[ADD2]]
 ;
   %mul1 = mul nsw i16 %x, 10
@@ -84,7 +84,7 @@ define i16 @test_nsw_dropped(i16 %x, i16 %y) {
 ; CHECK-LABEL: @test_nsw_dropped(
 ; CHECK-NEXT:    [[MUL1:%.*]] = mul i16 [[Y:%.*]], 25
 ; CHECK-NEXT:    [[REASS_MUL:%.*]] = mul i16 [[X:%.*]], 40
-; CHECK-NEXT:    [[ADD2:%.*]] = add i16 [[MUL1]], [[REASS_MUL]]
+; CHECK-NEXT:    [[ADD2:%.*]] = add i16 [[REASS_MUL]], [[MUL1]]
 ; CHECK-NEXT:    ret i16 [[ADD2]]
 ;
   %mul1 = mul nsw i16 %x, 15
@@ -105,7 +105,7 @@ define i32 @test_no_match(i32 %a, i32 %b) {
 ; CHECK-NOT:     Mul1
 ; CHECK-NOT:     Mul2
 ; CHECK-NOT:     DistAdd
-; CHECK-NOT:     DistSub;
+; CHECK-NOT:     DistSub
 ;
   %add = add nsw i32 %a, %b
   %mul = mul nsw i32 %add, 999

>From b0e3476a1f8f80c6aba2d5c37d1d863f778d4072 Mon Sep 17 00:00:00 2001
From: hazarathayya <hazarathayyayallanki at gmail.com>
Date: Fri, 17 Apr 2026 08:57:07 -0700
Subject: [PATCH 6/6] updated shouldbreakupdistr to check only on parent's
 siblings from checking global user list

---
 llvm/lib/Transforms/Scalar/Reassociate.cpp | 47 ++++++++++++++--------
 1 file changed, 31 insertions(+), 16 deletions(-)

diff --git a/llvm/lib/Transforms/Scalar/Reassociate.cpp b/llvm/lib/Transforms/Scalar/Reassociate.cpp
index dbe141576eb31..27158709e2ccb 100644
--- a/llvm/lib/Transforms/Scalar/Reassociate.cpp
+++ b/llvm/lib/Transforms/Scalar/Reassociate.cpp
@@ -979,27 +979,42 @@ static bool ShouldBreakUpDistribution(Instruction *Mul) {
   if (isa<FPMathOperator>(Mul))
     return false;
 
-  if (!isReassociableOp(AddSubVal, Instruction::Add, Instruction::FAdd) &&
-      !isReassociableOp(AddSubVal, Instruction::Sub, Instruction::FSub))
+  BinaryOperator *AddSubInst =
+      isReassociableOp(AddSubVal, Instruction::Add, Instruction::FAdd);
+  if (!AddSubInst)
+    AddSubInst =
+        isReassociableOp(AddSubVal, Instruction::Sub, Instruction::FSub);
+  if (!AddSubInst)
     return false;
 
-  Instruction *AddSubInst = cast<Instruction>(AddSubVal);
+  // Collect non-constant operands of the add/sub (at most 2).
+  SmallPtrSet<Value *, 2> AddOps;
+  for (Value *Op : AddSubInst->operands())
+    if (!isa<Constant>(Op))
+      AddOps.insert(Op);
+  if (AddOps.empty())
+    return false;
 
-  for (Value *Operand : AddSubInst->operands()) {
-    if (isa<Constant>(Operand))
+  for (User *MulUser : Mul->users()) {
+    auto *Parent = dyn_cast<Instruction>(MulUser);
+    if (!Parent)
       continue;
-    for (User *U : Operand->users()) {
-      Instruction *I = dyn_cast<Instruction>(U);
-      if (!I || I == AddSubInst)
+    for (Value *Sibling : Parent->operands()) {
+      if (Sibling == Mul)
         continue;
-      if ((I->getOpcode() == Instruction::Mul) && !(I->use_empty()) &&
-          (isa<Constant>(I->getOperand(0)) || isa<Constant>(I->getOperand(1)))){ 
-        int SibConstSide = isa<Constant>(I->getOperand(0)) ? 0 : 1;
-        if (I->getOperand(SibConstSide) == Mul->getOperand(ConstSide))
-            continue;
-        return true; 
-      }
-   }
+      auto *SibMul = dyn_cast<BinaryOperator>(Sibling);
+      if (!SibMul || SibMul->getOpcode() != Instruction::Mul ||
+          SibMul->use_empty())
+        continue;
+      // Sibling must be NonConst * C'.
+      int SibCS    = isa<Constant>(SibMul->getOperand(0)) ? 0 : 1;
+      Value *SibC  = SibMul->getOperand(SibCS);
+      Value *SibNC = SibMul->getOperand(1 - SibCS);
+      if (!isa<Constant>(SibC))
+        continue;
+      if (AddOps.count(SibNC))
+        return true;
+    }
   }
   return false;
 }



More information about the llvm-commits mailing list