[Matrix] Hoist finalizeLowering into caller. NFC #143038

jroelofs · 2025-06-05T21:45:27Z

No description provided.

llvmbot · 2025-06-05T21:46:07Z

@llvm/pr-subscribers-llvm-transforms

Author: Jon Roelofs (jroelofs)

Changes

Full diff: https://github.com/llvm/llvm-project/pull/143038.diff

1 Files Affected:

(modified) llvm/lib/Transforms/Scalar/LowerMatrixIntrinsics.cpp (+55-65)

diff --git a/llvm/lib/Transforms/Scalar/LowerMatrixIntrinsics.cpp b/llvm/lib/Transforms/Scalar/LowerMatrixIntrinsics.cpp
index 20279bf69dd59..c2a89f3c97aa1 100644
--- a/llvm/lib/Transforms/Scalar/LowerMatrixIntrinsics.cpp
+++ b/llvm/lib/Transforms/Scalar/LowerMatrixIntrinsics.cpp
@@ -1054,24 +1054,26 @@ class LowerMatrixIntrinsics {
       if (FusedInsts.count(Inst))
         continue;
 
-      IRBuilder<> Builder(Inst);
-
       const ShapeInfo &SI = ShapeMap.at(Inst);
 
       Value *Op1;
       Value *Op2;
+      MatrixTy Result;
       if (auto *BinOp = dyn_cast<BinaryOperator>(Inst))
-        VisitBinaryOperator(BinOp, SI);
+        Result = VisitBinaryOperator(BinOp, SI);
       else if (auto *UnOp = dyn_cast<UnaryOperator>(Inst))
-        VisitUnaryOperator(UnOp, SI);
-      else if (CallInst *CInst = dyn_cast<CallInst>(Inst))
-        VisitCallInst(CInst);
+        Result = VisitUnaryOperator(UnOp, SI);
+      else if (auto *Intr = dyn_cast<IntrinsicInst>(Inst))
+        Result = VisitIntrinsicInst(Intr, SI);
       else if (match(Inst, m_Load(m_Value(Op1))))
-        VisitLoad(cast<LoadInst>(Inst), SI, Op1, Builder);
+        Result = VisitLoad(cast<LoadInst>(Inst), SI, Op1);
       else if (match(Inst, m_Store(m_Value(Op1), m_Value(Op2))))
-        VisitStore(cast<StoreInst>(Inst), SI, Op1, Op2, Builder);
+        Result = VisitStore(cast<StoreInst>(Inst), SI, Op1, Op2);
       else
         continue;
+
+      IRBuilder<> Builder(Inst);
+      finalizeLowering(Inst, Result, Builder);
       Changed = true;
     }
 
@@ -1111,27 +1113,24 @@ class LowerMatrixIntrinsics {
   }
 
   /// Replace intrinsic calls.
-  void VisitCallInst(CallInst *Inst) {
+  MatrixTy VisitIntrinsicInst(IntrinsicInst *Inst, const ShapeInfo &SI) {
     assert(Inst->getCalledFunction() &&
            Inst->getCalledFunction()->isIntrinsic());
 
     switch (Inst->getCalledFunction()->getIntrinsicID()) {
     case Intrinsic::matrix_multiply:
-      LowerMultiply(Inst);
-      break;
+      return LowerMultiply(Inst);
     case Intrinsic::matrix_transpose:
-      LowerTranspose(Inst);
-      break;
+      return LowerTranspose(Inst);
     case Intrinsic::matrix_column_major_load:
-      LowerColumnMajorLoad(Inst);
-      break;
+      return LowerColumnMajorLoad(Inst);
     case Intrinsic::matrix_column_major_store:
-      LowerColumnMajorStore(Inst);
-      break;
+      return LowerColumnMajorStore(Inst);
     default:
-      llvm_unreachable(
-          "only intrinsics supporting shape info should be seen here");
+      break;
     }
+    llvm_unreachable(
+        "only intrinsics supporting shape info should be seen here");
   }
 
   /// Compute the alignment for a column/row \p Idx with \p Stride between them.
@@ -1197,26 +1196,24 @@ class LowerMatrixIntrinsics {
   }
 
   /// Lower a load instruction with shape information.
-  void LowerLoad(Instruction *Inst, Value *Ptr, MaybeAlign Align, Value *Stride,
-                 bool IsVolatile, ShapeInfo Shape) {
+  MatrixTy LowerLoad(Instruction *Inst, Value *Ptr, MaybeAlign Align,
+                     Value *Stride, bool IsVolatile, ShapeInfo Shape) {
     IRBuilder<> Builder(Inst);
-    finalizeLowering(Inst,
-                     loadMatrix(Inst->getType(), Ptr, Align, Stride, IsVolatile,
-                                Shape, Builder),
-                     Builder);
+    return loadMatrix(Inst->getType(), Ptr, Align, Stride, IsVolatile, Shape,
+                      Builder);
   }
 
   /// Lowers llvm.matrix.column.major.load.
   ///
   /// The intrinsic loads a matrix from memory using a stride between columns.
-  void LowerColumnMajorLoad(CallInst *Inst) {
+  MatrixTy LowerColumnMajorLoad(CallInst *Inst) {
     assert(MatrixLayout == MatrixLayoutTy::ColumnMajor &&
            "Intrinsic only supports column-major layout!");
     Value *Ptr = Inst->getArgOperand(0);
     Value *Stride = Inst->getArgOperand(1);
-    LowerLoad(Inst, Ptr, Inst->getParamAlign(0), Stride,
-              cast<ConstantInt>(Inst->getArgOperand(2))->isOne(),
-              {Inst->getArgOperand(3), Inst->getArgOperand(4)});
+    return LowerLoad(Inst, Ptr, Inst->getParamAlign(0), Stride,
+                     cast<ConstantInt>(Inst->getArgOperand(2))->isOne(),
+                     {Inst->getArgOperand(3), Inst->getArgOperand(4)});
   }
 
   /// Stores a sub-matrix \p StoreVal into the \p R x \p C matrix starting at \p
@@ -1259,28 +1256,27 @@ class LowerMatrixIntrinsics {
   }
 
   /// Lower a store instruction with shape information.
-  void LowerStore(Instruction *Inst, Value *Matrix, Value *Ptr, MaybeAlign A,
-                  Value *Stride, bool IsVolatile, ShapeInfo Shape) {
+  MatrixTy LowerStore(Instruction *Inst, Value *Matrix, Value *Ptr,
+                      MaybeAlign A, Value *Stride, bool IsVolatile,
+                      ShapeInfo Shape) {
     IRBuilder<> Builder(Inst);
     auto StoreVal = getMatrix(Matrix, Shape, Builder);
-    finalizeLowering(Inst,
-                     storeMatrix(Matrix->getType(), StoreVal, Ptr, A, Stride,
-                                 IsVolatile, Builder),
-                     Builder);
+    return storeMatrix(Matrix->getType(), StoreVal, Ptr, A, Stride, IsVolatile,
+                       Builder);
   }
 
   /// Lowers llvm.matrix.column.major.store.
   ///
   /// The intrinsic store a matrix back memory using a stride between columns.
-  void LowerColumnMajorStore(CallInst *Inst) {
+  MatrixTy LowerColumnMajorStore(CallInst *Inst) {
     assert(MatrixLayout == MatrixLayoutTy::ColumnMajor &&
            "Intrinsic only supports column-major layout!");
     Value *Matrix = Inst->getArgOperand(0);
     Value *Ptr = Inst->getArgOperand(1);
     Value *Stride = Inst->getArgOperand(2);
-    LowerStore(Inst, Matrix, Ptr, Inst->getParamAlign(1), Stride,
-               cast<ConstantInt>(Inst->getArgOperand(3))->isOne(),
-               {Inst->getArgOperand(4), Inst->getArgOperand(5)});
+    return LowerStore(Inst, Matrix, Ptr, Inst->getParamAlign(1), Stride,
+                      cast<ConstantInt>(Inst->getArgOperand(3))->isOne(),
+                      {Inst->getArgOperand(4), Inst->getArgOperand(5)});
   }
 
   // Set elements I..I+NumElts-1 to Block
@@ -2045,7 +2041,7 @@ class LowerMatrixIntrinsics {
   }
 
   /// Lowers llvm.matrix.multiply.
-  void LowerMultiply(CallInst *MatMul) {
+  MatrixTy LowerMultiply(CallInst *MatMul) {
     IRBuilder<> Builder(MatMul);
     auto *EltType = cast<FixedVectorType>(MatMul->getType())->getElementType();
     ShapeInfo LShape(MatMul->getArgOperand(2), MatMul->getArgOperand(3));
@@ -2067,11 +2063,11 @@ class LowerMatrixIntrinsics {
 
     emitMatrixMultiply(Result, Lhs, Rhs, Builder, false, false,
                        getFastMathFlags(MatMul));
-    finalizeLowering(MatMul, Result, Builder);
+    return Result;
   }
 
   /// Lowers llvm.matrix.transpose.
-  void LowerTranspose(CallInst *Inst) {
+  MatrixTy LowerTranspose(CallInst *Inst) {
     MatrixTy Result;
     IRBuilder<> Builder(Inst);
     Value *InputVal = Inst->getArgOperand(0);
@@ -2101,28 +2097,26 @@ class LowerMatrixIntrinsics {
     // TODO: Improve estimate of operations needed for transposes. Currently we
     // just count the insertelement/extractelement instructions, but do not
     // account for later simplifications/combines.
-    finalizeLowering(
-        Inst,
-        Result.addNumComputeOps(2 * ArgShape.NumRows * ArgShape.NumColumns)
-            .addNumExposedTransposes(1),
-        Builder);
+    return Result.addNumComputeOps(2 * ArgShape.NumRows * ArgShape.NumColumns)
+        .addNumExposedTransposes(1);
   }
 
   /// Lower load instructions.
-  void VisitLoad(LoadInst *Inst, const ShapeInfo &SI, Value *Ptr,
-                 IRBuilder<> &Builder) {
-    LowerLoad(Inst, Ptr, Inst->getAlign(), Builder.getInt64(SI.getStride()),
-              Inst->isVolatile(), SI);
+  MatrixTy VisitLoad(LoadInst *Inst, const ShapeInfo &SI, Value *Ptr) {
+    IRBuilder<> Builder(Inst);
+    return LowerLoad(Inst, Ptr, Inst->getAlign(),
+                     Builder.getInt64(SI.getStride()), Inst->isVolatile(), SI);
   }
 
-  void VisitStore(StoreInst *Inst, const ShapeInfo &SI, Value *StoredVal,
-                  Value *Ptr, IRBuilder<> &Builder) {
-    LowerStore(Inst, StoredVal, Ptr, Inst->getAlign(),
-               Builder.getInt64(SI.getStride()), Inst->isVolatile(), SI);
+  MatrixTy VisitStore(StoreInst *Inst, const ShapeInfo &SI, Value *StoredVal,
+                      Value *Ptr) {
+    IRBuilder<> Builder(Inst);
+    return LowerStore(Inst, StoredVal, Ptr, Inst->getAlign(),
+                      Builder.getInt64(SI.getStride()), Inst->isVolatile(), SI);
   }
 
   /// Lower binary operators.
-  void VisitBinaryOperator(BinaryOperator *Inst, const ShapeInfo &SI) {
+  MatrixTy VisitBinaryOperator(BinaryOperator *Inst, const ShapeInfo &SI) {
     Value *Lhs = Inst->getOperand(0);
     Value *Rhs = Inst->getOperand(1);
 
@@ -2141,14 +2135,12 @@ class LowerMatrixIntrinsics {
       Result.addVector(Builder.CreateBinOp(Inst->getOpcode(), A.getVector(I),
                                            B.getVector(I)));
 
-    finalizeLowering(Inst,
-                     Result.addNumComputeOps(getNumOps(Result.getVectorTy()) *
-                                             Result.getNumVectors()),
-                     Builder);
+    return Result.addNumComputeOps(getNumOps(Result.getVectorTy()) *
+                                   Result.getNumVectors());
   }
 
   /// Lower unary operators.
-  void VisitUnaryOperator(UnaryOperator *Inst, const ShapeInfo &SI) {
+  MatrixTy VisitUnaryOperator(UnaryOperator *Inst, const ShapeInfo &SI) {
     Value *Op = Inst->getOperand(0);
 
     IRBuilder<> Builder(Inst);
@@ -2171,10 +2163,8 @@ class LowerMatrixIntrinsics {
     for (unsigned I = 0; I < SI.getNumVectors(); ++I)
       Result.addVector(BuildVectorOp(M.getVector(I)));
 
-    finalizeLowering(Inst,
-                     Result.addNumComputeOps(getNumOps(Result.getVectorTy()) *
-                                             Result.getNumVectors()),
-                     Builder);
+    return Result.addNumComputeOps(getNumOps(Result.getVectorTy()) *
+                                   Result.getNumVectors());
   }
 
   /// Helper to linearize a matrix expression tree into a string. Currently

fhahn

LGTM, thanks!

…-finalize

github-actions · 2025-06-10T17:07:21Z

✅ With the latest revision this PR passed the C/C++ code formatter.

…-finalize

lei137 · 2025-06-13T22:28:50Z

Hi @jroelofs, This patch broke us downstream. Can you please provide some information as to what the benefits are provided by this change?

jroelofs · 2025-06-13T22:34:00Z

Hi @jroelofs, This patch broke us downstream. Can you please provide some information as to what the benefits are provided by this change?

The idea for it came from review feedback from @fhahn in #141869 (comment): we call finalizeLowering in all of the Visit* methods, so it should be NFC to hoist that into the caller. Can you share a small reproducer?

hoist finalizeLowering into caller

e96975e

jroelofs requested review from fhahn and anemet June 5, 2025 21:45

llvmbot added the llvm:transforms label Jun 5, 2025

fhahn approved these changes Jun 10, 2025

View reviewed changes

Merge remote-tracking branch 'origin/main' into jroelofs/lower-matrix…

0ed8c07

…-finalize

jroelofs added 2 commits June 10, 2025 10:26

clang-format

c0c68f4

Merge remote-tracking branch 'origin/main' into jroelofs/lower-matrix…

c3baab9

…-finalize

jroelofs merged commit 8345d62 into llvm:main Jun 10, 2025
5 of 7 checks passed

jroelofs deleted the jroelofs/lower-matrix-finalize branch June 10, 2025 19:36

rorth pushed a commit to rorth/llvm-project that referenced this pull request Jun 11, 2025

[Matrix] Hoist finalizeLowering into caller. NFC (llvm#143038)

2197b92

tomtor pushed a commit to tomtor/llvm-project that referenced this pull request Jun 14, 2025

[Matrix] Hoist finalizeLowering into caller. NFC (llvm#143038)

d1339a8

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[Matrix] Hoist finalizeLowering into caller. NFC #143038

[Matrix] Hoist finalizeLowering into caller. NFC #143038

Uh oh!

jroelofs commented Jun 5, 2025

Uh oh!

llvmbot commented Jun 5, 2025

Uh oh!

fhahn left a comment

Uh oh!

github-actions bot commented Jun 10, 2025 •

edited

Loading

Uh oh!

Uh oh!

lei137 commented Jun 13, 2025

Uh oh!

jroelofs commented Jun 13, 2025

Uh oh!

Uh oh!

[Matrix] Hoist finalizeLowering into caller. NFC #143038

[Matrix] Hoist finalizeLowering into caller. NFC #143038

Uh oh!

Conversation

jroelofs commented Jun 5, 2025

Uh oh!

llvmbot commented Jun 5, 2025

Uh oh!

fhahn left a comment

Choose a reason for hiding this comment

Uh oh!

github-actions bot commented Jun 10, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

lei137 commented Jun 13, 2025

Uh oh!

jroelofs commented Jun 13, 2025

Uh oh!

Uh oh!

github-actions bot commented Jun 10, 2025 •

edited

Loading