Use pipelining for outer loop of attention

Author: dhernandez0

mlir/include/mlir/Dialect/Rock/Passes.h

@@ -50,6 +50,7 @@ namespace rock {
 #define GEN_PASS_DECL_ROCKFINDFIRSTGEMMINDEXPASS
 #define GEN_PASS_DECL_ROCKREMOVEOUTPUTALLOCPASS
 #define GEN_PASS_DECL_ROCKBLOCKWISELOADTILETOTHREADWISEPASS
+#define GEN_PASS_DECL_ROCKPREPAREPIPELINEPASS
 #define GEN_PASS_DECL_ROCKANNOTATELIVENESSPASS
 
 #define GEN_PASS_REGISTRATION

mlir/include/mlir/Dialect/Rock/Passes.td

@@ -210,4 +210,10 @@ def RockAnnotateLivenessPass
   let dependentDialects = ["rock::RockDialect", "memref::MemRefDialect"];
 }
 
+def RockPreparePipelinePass
+    : Pass<"rock-prepare-pipeline", "::mlir::func::FuncOp"> {
+  let summary = "This pass prepares ops for pipelining";
+  let dependentDialects = ["rock::RockDialect", "func::FuncDialect"];
+}
+
 #endif // MLIR_DIALECT_ROCK_PASSES

mlir/lib/Dialect/Rock/IR/RockDialect.cpp

@@ -92,10 +92,16 @@ template <>
 struct rank<0> {};
 
 template <typename OpType>
+<<<<<<< HEAD
 static auto
 getGemmEffects(rank<1>, OpType &op,
                SmallVectorImpl<MemoryEffects::EffectInstance> &effects)
     -> decltype(void(op.getScaleA()), void(op.getScaleB())) {
+=======
+static void
+getGemmEffects(OpType &op,
+               SmallVectorImpl<MemoryEffects::EffectInstance> &effects) {
+>>>>>>> 64ae93840ed4 (fixes)
   auto *read = MemoryEffects::Read::get();
   auto *write = MemoryEffects::Write::get();
 

mlir/lib/Dialect/Rock/Pipelines/Pipelines.cpp

@@ -164,6 +164,7 @@ void rock::buildKernelPipeline(OpPassManager &pm,
   funcPm.addPass(rock::createRockShuffleGemmForReductions());
   funcPm.addPass(rock::createRockGridwiseGemmToBlockwisePass());
   funcPm.addPass(rock::createRockBlockwiseLoadTileToThreadwisePass());
+  funcPm.addPass(rock::createRockPreparePipelinePass());
 
   // We want to delay blockwise lowering in the fusion cases
   // until after linalg align pass because with reduction fusion

mlir/lib/Dialect/Rock/Transforms/BlockwiseLoadTileToThreadwise.cpp

@@ -206,11 +206,13 @@ class LoweringBlockwiseLoadTileOp final
     // we want to insert all allocs and transforms before the loop
     Operation *parentOp = op->getParentOp();
     assert(parentOp && "BlockwiseLoadTileOp must have a parent op");
-    if (isa<LoopLikeOpInterface>(parentOp))
-      b.setInsertionPoint(parentOp);
-    else
-      b.setInsertionPoint(op);
 
+    // let's add the allocs to the begginging of the loop
+    if (auto parentLoop = dyn_cast<LoopLikeOpInterface>(parentOp)) {
+      // workaround for parentLoop.getBody()
+      assert(parentLoop->getRegions().size() == 1);
+      b.setInsertionPointToStart(&parentLoop->getRegion(0).front());
+    }
     Value loadBuffer, storeBuffer;
     if (loadType == GemmLoadTileType::BypassLDS) {
       auto privateMemoryAddressSpace = b.getAttr<gpu::AddressSpaceAttr>(
@@ -233,10 +235,6 @@ class LoweringBlockwiseLoadTileOp final
     SmallVector<int64_t, 3> bidGridLengths = {G, mBlocks, nBlocks};
     SmallVector<StringRef, 3> bidGridOrder = {"g_block", "m_block", "n_block"};
 
-    // Create the stages for the blockwise load tile op
-    if (isa<LoopLikeOpInterface>(parentOp))
-      b.setInsertionPoint(op);
-
     auto [stageGlobalRead, stageGlobalReadNew] =
         createOrGetStage(b, loc, "GlobalRead", parentOp);
     {

mlir/lib/Dialect/Rock/Transforms/CMakeLists.txt

@@ -32,6 +32,7 @@ add_rocmlir_dialect_library(MLIRRockTransforms
   FindFirstGemmIndex.cpp
   RemoveOutputAlloc.cpp
   BlockwiseLoadTileToThreadwise.cpp
+  PreparePipeline.cpp
   AnnotateLiveness.cpp
 
   ADDITIONAL_HEADER_DIRS