[Matrix] Refactor tiled loops in a struct. NFC

[thegameg]

Summary

The three loops have the same structure: index, header, latch.

[fhahn]

LGTM, thanks!

llvm/include/llvm/Transforms/Utils/MatrixUtils.h

28

nit: should this be ColumnLoop.Index?

[thegameg]

https://github.com/llvm/llvm-project/commit/2c6e8b4636700f22a76eeda01e4a3258692b80f3

llvm/include/llvm/Transforms/Utils/MatrixUtils.h

 class BasicBlock;
 class Value;
 class Loop;
 class LoopInfo;
 class IRBuilderBase;
 
 /// A helper struct to create IR loop nests for tiling in IR of the following
 /// form:
-///   for CurrentColumn = 0..NumColumns
-///     for CurrentRow = 0..NumRows
-///       for CurrentInner = 0..NumInner
+///   for Column.Index = 0..NumColumns
+///     for RowLoop.Index = 0..NumRows
+///       for KLoop.Index = 0..NumInner
 struct TileInfo {
   /// Number of rows of the matrix.
   unsigned NumRows;
 
   /// Number of columns of the matrix.
   unsigned NumColumns;
 
   /// Number of columns of the first matrix of a multiply /
   /// number of rows of the second matrix of a multiply.
   unsigned NumInner;
 
   /// Number of rows/columns in a tile.
   unsigned TileSize = -1;
 
-  /// Start row of the current tile to compute.
-  Value *CurrentRow;
-
-  /// Start column of the current tile to compute.
-  Value *CurrentCol;
-
-  /// Current tile offset during the tile computation.
-  Value *CurrentK;
+  /// Properties of a single loop used when generating the tiled loop nest.
+  struct MatrixLoop {
+    /// The index updated on every iteration.
+    Value *Index = nullptr;
+    /// The header and latch of the loop.
+    BasicBlock *Header = nullptr;
+    BasicBlock *Latch = nullptr;
+  };
 
-  /// Header of the outermost loop iterating from 0..NumColumns.
-  BasicBlock *ColumnLoopHeader = nullptr;
-
-  /// Header of the second loop iterating from 0..NumRows.
-  BasicBlock *RowLoopHeader = nullptr;
-  /// Latch of the second loop iterating from 0..NumRows.
-  BasicBlock *RowLoopLatch = nullptr;
-  /// Header of the innermost loop iterating from 0..NumInner.
-  BasicBlock *InnerLoopHeader = nullptr;
-  /// Latch of the innermost loop iterating from 0..NumInner.
-  BasicBlock *InnerLoopLatch = nullptr;
+  /// The loop iterating on the rows.
+  MatrixLoop RowLoop;
+  /// The loop iterating on the columns.
+  MatrixLoop ColumnLoop;
+  /// The loop iterating on k (inner dimension).
+  MatrixLoop KLoop;
 
   TileInfo(unsigned NumRows, unsigned NumColumns, unsigned NumInner,
            unsigned TileSize)
       : NumRows(NumRows), NumColumns(NumColumns), NumInner(NumInner),
         TileSize(TileSize) {}
 
   /// Creates an IR loop nests for tiling of the form below. Returns the block
   /// for the inner loop body and sets {Column,Row,Inner}LoopHeader/Latch
   /// fields.
   ///
-  /// for CurrentColumn = 0..NumColumns
-  ///   for CurrentRow = 0..NumRows
-  ///     for CurrentInner = 0..NumInner
+  /// for ColumnLoop.Index = 0..NumColumns
+  ///   for RowLoop.Index = 0..NumRows
+  ///     for InnerLoop.Index = 0..NumInner
   BasicBlock *CreateTiledLoops(BasicBlock *Start, BasicBlock *End,
                                IRBuilderBase &B, DomTreeUpdater &DTU,
                                LoopInfo &LI);
 
 private:
   /// Creates a new loop with header, body and latch blocks that iterates from
   /// [0, Bound). Updates \p Preheader to branch to the new header and uses \p
   /// Exit as exit block.  Adds the new loop blocks to \L and applies dominator

llvm/lib/Transforms/Scalar/LowerMatrixIntrinsics.cpp

BasicBlock *End =
         SplitBlock(InsertI->getParent(), InsertI, DT, LI, nullptr, "continue");
     IRBuilder<> Builder(MatMul);
     BasicBlock *InnerBody = TI.CreateTiledLoops(Start, End, Builder, DTU, *LI);
 
     Type *TileVecTy =
         FixedVectorType::get(MatMul->getType()->getScalarType(), TileSize);
     MatrixTy TileResult;
     // Insert in the inner loop header.
-    Builder.SetInsertPoint(TI.InnerLoopHeader->getTerminator());
+    Builder.SetInsertPoint(TI.KLoop.Header->getTerminator());
     // Create PHI nodes for the result columns to accumulate across iterations.
     SmallVector<PHINode *, 4> ColumnPhis;
     for (unsigned I = 0; I < TileSize; I++) {
       auto *Phi = Builder.CreatePHI(TileVecTy, 2, "result.vec." + Twine(I));
       Phi->addIncoming(ConstantAggregateZero::get(TileVecTy),
-                       TI.RowLoopHeader->getSingleSuccessor());
+                       TI.RowLoop.Header->getSingleSuccessor());
       TileResult.addVector(Phi);
       ColumnPhis.push_back(Phi);
     }
 
     // Insert in the inner loop body, which computes
     //   Res += Load(CurrentRow, K) * Load(K, CurrentColumn)
     Builder.SetInsertPoint(InnerBody->getTerminator());
     // Load tiles of the operands.
-    MatrixTy A = loadMatrix(LPtr, {}, false, LShape, TI.CurrentRow, TI.CurrentK,
-                            {TileSize, TileSize}, EltType, Builder);
-    MatrixTy B = loadMatrix(RPtr, {}, false, RShape, TI.CurrentK, TI.CurrentCol,
-                            {TileSize, TileSize}, EltType, Builder);
+    MatrixTy A =
+        loadMatrix(LPtr, {}, false, LShape, TI.RowLoop.Index, TI.KLoop.Index,
+    MatrixTy B =
+        loadMatrix(RPtr, {}, false, RShape, TI.KLoop.Index, TI.ColumnLoop.Index,
     emitMatrixMultiply(TileResult, A, B, Builder, true, false,
                        getFastMathFlags(MatMul));
     // Store result after the inner loop is done.
-    Builder.SetInsertPoint(TI.RowLoopLatch->getTerminator());
+    Builder.SetInsertPoint(TI.RowLoop.Latch->getTerminator());
     storeMatrix(TileResult, Store->getPointerOperand(), Store->getAlign(),
                 Store->isVolatile(), {LShape.NumRows, RShape.NumColumns},
-                TI.CurrentRow, TI.CurrentCol, EltType, Builder);
+                TI.RowLoop.Index, TI.ColumnLoop.Index, EltType, Builder);
 
     for (unsigned I = 0; I < TileResult.getNumVectors(); I++)
-      ColumnPhis[I]->addIncoming(TileResult.getVector(I), TI.InnerLoopLatch);
+      ColumnPhis[I]->addIncoming(TileResult.getVector(I), TI.KLoop.Latch);
 
     // Force unrolling of a few iterations of the inner loop, to make sure there
     // is enough work per iteration.
     // FIXME: The unroller should make this decision directly instead, but
     // currently the cost-model is not up to the task.
     unsigned InnerLoopUnrollCount = std::min(10u, LShape.NumColumns / TileSize);
-    addStringMetadataToLoop(LI->getLoopFor(TI.InnerLoopHeader),
+    addStringMetadataToLoop(LI->getLoopFor(TI.KLoop.Header),
                             "llvm.loop.unroll.count", InnerLoopUnrollCount);
   }
 
   void emitSIMDTiling(CallInst *MatMul, LoadInst *LoadOp0, LoadInst *LoadOp1,
                       StoreInst *Store,
                       SmallPtrSetImpl<Instruction *> &FusedInsts) {
     assert(MatrixLayout == MatrixLayoutTy::ColumnMajor &&
            "Tiling only supported for column-major matrixes at the moment!");

llvm/lib/Transforms/Utils/MatrixUtils.cpp

 
 // Creates the following loop nest skeleton:
 //  for C = 0; C < NumColumns; C += TileSize
 //    for R = 0; R < NumRows; R += TileSize
 //      for K = 0; K < Inner ; K += TileSize
 BasicBlock *TileInfo::CreateTiledLoops(BasicBlock *Start, BasicBlock *End,
                                        IRBuilderBase &B, DomTreeUpdater &DTU,
                                        LoopInfo &LI) {
-  Loop *ColLoop = LI.AllocateLoop();
-  Loop *RowLoop = LI.AllocateLoop();
-  Loop *InnerLoop = LI.AllocateLoop();
-  RowLoop->addChildLoop(InnerLoop);
-  ColLoop->addChildLoop(RowLoop);
+  Loop *ColumnLoopInfo = LI.AllocateLoop();
+  Loop *RowLoopInfo = LI.AllocateLoop();
+  Loop *KLoopInfo = LI.AllocateLoop();
+  RowLoopInfo->addChildLoop(KLoopInfo);
+  ColumnLoopInfo->addChildLoop(RowLoopInfo);
   if (Loop *ParentL = LI.getLoopFor(Start))
-    ParentL->addChildLoop(ColLoop);
+    ParentL->addChildLoop(ColumnLoopInfo);
   else
-    LI.addTopLevelLoop(ColLoop);
+    LI.addTopLevelLoop(ColumnLoopInfo);
 
   BasicBlock *ColBody =
       CreateLoop(Start, End, B.getInt64(NumColumns), B.getInt64(TileSize),
-                 "cols", B, DTU, ColLoop, LI);
-  BasicBlock *ColLatch = ColBody->getSingleSuccessor();
+                 "cols", B, DTU, ColumnLoopInfo, LI);
+  ColumnLoop.Latch = ColBody->getSingleSuccessor();
   BasicBlock *RowBody =
-      CreateLoop(ColBody, ColLatch, B.getInt64(NumRows), B.getInt64(TileSize),
-                 "rows", B, DTU, RowLoop, LI);
-  RowLoopLatch = RowBody->getSingleSuccessor();
+      CreateLoop(ColBody, ColumnLoop.Latch, B.getInt64(NumRows),
+                 B.getInt64(TileSize), "rows", B, DTU, RowLoopInfo, LI);
+  RowLoop.Latch = RowBody->getSingleSuccessor();
 
   BasicBlock *InnerBody =
-      CreateLoop(RowBody, RowLoopLatch, B.getInt64(NumInner),
-                 B.getInt64(TileSize), "inner", B, DTU, InnerLoop, LI);
-  InnerLoopLatch = InnerBody->getSingleSuccessor();
-  ColumnLoopHeader = ColBody->getSinglePredecessor();
-  RowLoopHeader = RowBody->getSinglePredecessor();
-  InnerLoopHeader = InnerBody->getSinglePredecessor();
-  CurrentRow = &*RowLoopHeader->begin();
-  CurrentCol = &*ColumnLoopHeader->begin();
-  CurrentK = &*InnerLoopHeader->begin();
+      CreateLoop(RowBody, RowLoop.Latch, B.getInt64(NumInner),
+                 B.getInt64(TileSize), "inner", B, DTU, KLoopInfo, LI);
+  KLoop.Latch = InnerBody->getSingleSuccessor();
+  ColumnLoop.Header = ColBody->getSinglePredecessor();
+  RowLoop.Header = RowBody->getSinglePredecessor();
+  KLoop.Header = InnerBody->getSinglePredecessor();
+  RowLoop.Index = &*RowLoop.Header->begin();
+  ColumnLoop.Index = &*ColumnLoop.Header->begin();
+  KLoop.Index = &*KLoop.Header->begin();
 
   return InnerBody;
 }