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

via llvm-commits llvm-commits at lists.llvm.org
Thu Feb 19 01:48:48 PST 2026


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

>From d6b333fb1170a11140c1d24f62c8c0aa258b86a3 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/2] [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 6d7ce36606083..eca634f3c5d98 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 11b6f5ce882883e749ef372ca2367f41ad4eff54 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/2] 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 eca634f3c5d98..6b4ff9d6d63b3 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.");



More information about the llvm-commits mailing list