Annotate liveness pass (+ reuseLDS pass changes) (#2015)

* Annotate liveness making some assumptions (write + read pattern for LDS). Also update ReuseLDS pass accordingly.

* Addressing PR comments

* clang format

Author: dhernandez0

mlir/include/mlir/Dialect/Rock/IR/RockOps.td

@@ -569,17 +569,25 @@ def Rock_GpuAllocOp:
   let hasVerifier = 1;
 }
 
-// Annotate lifetime of memory allocation on GPU memory hierachy.
-def Rock_GpuDeallocOp:
-    Rock_Op<"dealloc", [MemoryEffects<[MemFree<DefaultResource>]>]>,
-    Arguments<(ins AnyMemRef:$memref)> {
-  let summary = "Annotate lifetime of memory allocation on GPU";
+// Annotate the start of lifetime of a memory allocation on GPU.
+def Rock_LiveInOp : Rock_Op<"live_in">, Arguments<(ins AnyMemRef:$memref)> {
+  let summary =
+      "Annotate the start of lifetime of a LDS memory allocation on GPU";
   let description = [{
-    The `rock.dealloc` op annotates lifetime of memory allocation memory on GPU.
-    - Address space 0 : global.
-    - Address space 3 : LDS.
-    - Address space 5 : private (VGPR).
-    All other values would be considered as allocation on global.
+    The `rock.live_in` op annotates the start of lifetime of a LDS memory allocation on GPU.
+  }];
+  let assemblyFormat = [{
+    $memref attr-dict `:` type($memref)
+  }];
+  let hasVerifier = 1;
+}
+
+// Annotate the end of lifetime of a memory allocation on GPU.
+def Rock_LiveOutOp : Rock_Op<"live_out">, Arguments<(ins AnyMemRef:$memref)> {
+  let summary =
+      "Annotate the end of lifetime of a LDS memory allocation on GPU";
+  let description = [{
+    The `rock.live_out` op annotates the end of lifetime of a LDS memory allocation on GPU.
   }];
   let assemblyFormat = [{
     $memref attr-dict `:` type($memref)

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_ROCKANNOTATELIVENESSPASS
 
 #define GEN_PASS_REGISTRATION
 #include "mlir/Dialect/Rock/Passes.h.inc"

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

@@ -159,7 +159,8 @@ def RockOutputSwizzlePass : Pass<"rock-output-swizzle", "::mlir::func::FuncOp">
 }
 
 def RockReuseLDSPass : Pass<"rock-reuse-lds", "::mlir::func::FuncOp"> {
-  let summary = "This pass re-uses LDS memory by using the lifetime annotations (rock.dealloc)";
+  let summary = "This pass re-uses LDS memory by using the lifetime "
+                "annotations (rock.live_in, rock.live_out)";
   let dependentDialects = ["rock::RockDialect", "memref::MemRefDialect"];
 }
 
@@ -202,4 +203,11 @@ def RockBlockwiseLoadTileToThreadwisePass
                            "arith::ArithDialect", "memref::MemRefDialect"];
 }
 
+def RockAnnotateLivenessPass
+    : Pass<"rock-annotate-liveness", "::mlir::func::FuncOp"> {
+  let summary = "This pass annotates LDS memory with liveness ops "
+                "(rock.live_in, rock.live_out)";
+  let dependentDialects = ["rock::RockDialect", "memref::MemRefDialect"];
+}
+
 #endif // MLIR_DIALECT_ROCK_PASSES

mlir/include/mlir/Dialect/Rock/utility/loweringUtils.h

@@ -258,6 +258,9 @@ FailureOr<VectorDimInfo> getVectorDim(Location loc, Value matrix, Type elemType,
                                       int64_t dPerBlock, int64_t kpack,
                                       bool directToLDS);
 
+// Get the LDS size of the memref
+std::optional<int64_t> getWorkgroupMemorySize(MemRefType type);
+
 } // end namespace rock
 } // end namespace mlir
 #endif

mlir/lib/Conversion/RockToGPU/RockToGPU.cpp

@@ -102,18 +102,6 @@ struct MIGPUAllocRewritePattern : public OpRewritePattern<rock::GpuAllocOp> {
   }
 };
 
-struct MIGPUDeallocRewritePattern
-    : public OpRewritePattern<rock::GpuDeallocOp> {
-  using OpRewritePattern<rock::GpuDeallocOp>::OpRewritePattern;
-
-  LogicalResult matchAndRewrite(rock::GpuDeallocOp op,
-                                PatternRewriter &b) const override {
-
-    b.eraseOp(op);
-    return mlir::success();
-  }
-};
-
 template <typename Tmi, typename Tgpu>
 struct MIOpRewritePattern : public OpRewritePattern<Tmi> {
   using OpRewritePattern<Tmi>::OpRewritePattern;
@@ -345,7 +333,7 @@ void LowerRockOpsToGPUPass::runOnOperation() {
     RewritePatternSet patterns(ctx);
 
     // rock-lowering
-    patterns.add<MIGPUAllocRewritePattern, MIGPUDeallocRewritePattern,
+    patterns.add<MIGPUAllocRewritePattern,
                  MIOpRewritePattern<rock::WorkgroupBarrierOp, gpu::BarrierOp>,
                  MIOpRewritePattern<rock::LDSBarrierOp, amdgpu::LDSBarrierOp>,
                  WorkgroupIdRewritePattern,

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

@@ -1138,20 +1138,43 @@ LogicalResult GpuAllocOp::verify() {
 }
 
 //===-----------------------------------------------------===//
-// GpuDeallocOp
+// LiveInOp
 //===-----------------------------------------------------===//
 
-LogicalResult GpuDeallocOp::verify() {
+LogicalResult LiveInOp::verify() {
   // Make sure the input memref defining operation is a GpuAllocOp
-  if (auto gpuAlloc = dyn_cast<GpuAllocOp>(getMemref().getDefiningOp())) {
-    // Make sure the size is bigger than 0
-    if (getByteSize(getMemref().getType()) > 0) {
-      return success();
-    }
-    return emitError("The size of rock.dealloc should be greather than zero.");
-  }
-  return emitError("The operand of rock.dealloc must be the result of a "
-                   "rock.alloc operation.");
+  if (!isa<GpuAllocOp>(getMemref().getDefiningOp()))
+    return emitError("The operand of rock.live_in must be the result of a "
+                     "rock.alloc operation.");
+
+  auto memSpace = dyn_cast_or_null<gpu::AddressSpaceAttr>(
+      getMemref().getType().getMemorySpace());
+  if (!memSpace ||
+      (memSpace &&
+       memSpace.getValue() != gpu::GPUDialect::getWorkgroupAddressSpace()))
+    return emitError("The operand of rock.live_in must a LDS memref");
+
+  return success();
+}
+
+//===-----------------------------------------------------===//
+// LiveOutOp
+//===-----------------------------------------------------===//
+
+LogicalResult LiveOutOp::verify() {
+  // Make sure the input memref defining operation is a GpuAllocOp
+  if (!isa<GpuAllocOp>(getMemref().getDefiningOp()))
+    return emitError("The operand of rock.live_out must be the result of a "
+                     "rock.alloc operation.");
+
+  auto memSpace = dyn_cast_or_null<gpu::AddressSpaceAttr>(
+      getMemref().getType().getMemorySpace());
+  if (!memSpace ||
+      (memSpace &&
+       memSpace.getValue() != gpu::GPUDialect::getWorkgroupAddressSpace()))
+    return emitError("The operand of rock.live_out must a LDS memref");
+
+  return success();
 }
 
 //===-----------------------------------------------------===//

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

@@ -183,8 +183,14 @@ void rock::buildKernelPipeline(OpPassManager &pm,
     funcPm.addPass(createConvertLinalgToAffineLoopsPass());
     funcPm.addPass(rock::createRockVectorizeFusionsPass());
   }
+  // We run reuse LDS before the output swizzle pass because it uses a heuristic
+  // to determine whether to swizzle or not, and that heuristic needs the actual
+  // LDS usage. After running output swizzle, we'll create a new LDS buffer and
+  // we need to run reuse LDS again to be able to reuse LDS memory.
+  funcPm.addPass(rock::createRockAnnotateLivenessPass());
   funcPm.addPass(rock::createRockReuseLDSPass());
   funcPm.addPass(rock::createRockOutputSwizzlePass());
+  funcPm.addPass(rock::createRockAnnotateLivenessPass());
   funcPm.addPass(rock::createRockReuseLDSPass());
 
   if (!options.enableApplicability) {

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

@@ -1519,9 +1519,6 @@ static LogicalResult insertBlockwiseReduction(
       /*extraViews=*/nullptr,
       getBlockSize(reduceOp->getParentOfType<func::FuncOp>()).value());
 
-  ViewLikeOpInterface viewOp =
-      ldsWorkspace.getDefiningOp<ViewLikeOpInterface>();
-  GpuDeallocOp::create(rewriter, loc, viewOp.getViewSource());
   // Create partial reduction views
   ArrayAttr paddedReducedTrStack;
   {