AIROverrideMemrefMemorySpace: operating on ModuleOp (#1099)

Author: erwei-xilinx

mlir/include/air/Transform/Passes.td

@@ -1533,7 +1533,7 @@ def DmaToChannel : Pass<"air-dma-to-channel", "ModuleOp"> {
   }];
 }
 
-def AIROverrideMemRefMemorySpace : Pass<"air-override-memref-memory-space", "func::FuncOp"> {
+def AIROverrideMemRefMemorySpace : Pass<"air-override-memref-memory-space", "ModuleOp"> {
   let summary = "Force all memrefs allocated within code region to have the specified memory space.";
   let constructor = "xilinx::air::createAIROverrideMemRefMemorySpacePass()";
   let description = [{
@@ -1544,7 +1544,7 @@ def AIROverrideMemRefMemorySpace : Pass<"air-override-memref-memory-space", "fun
            "Memory space to override to.">,
     Option<"clScope", "scope", "std::string",
             /*default=*/"\"launch\"",
-            "AIR hierarchy scope to perform the transform under. Must be one of [herd, segment, launch].">
+            "AIR hierarchy scope to perform the transform under. Must be one of [herd, segment, launch, func].">
   ];
 }
 

mlir/lib/Transform/AIRMiscPasses.cpp

@@ -1234,8 +1234,9 @@ FailureOr<Value> tileChannelOpByFactor(
   Value zeroIdx = rewriter.create<arith::ConstantIndexOp>(loc, 0);
   // Create and apply affine map onto the split channel ops.
   SmallVector<Value> tokens;
-  int memorySpace = dyn_cast<MemRefType>(originalChanOp.getMemref().getType())
-                        .getMemorySpaceAsInt();
+  int memorySpace =
+      dyn_cast<BaseMemRefType>(originalChanOp.getMemref().getType())
+          .getMemorySpaceAsInt();
   for (int i = 0; i < factor; i++) {
     // Get affine map and split size from splitInfo.
     auto &[splitInfoDimOnOffsets, splitInfoAffineMap, splitInfoSplitOffset,
@@ -2354,9 +2355,11 @@ struct OverrideMemorySpacePattern : public OpRewritePattern<memref::AllocOp> {
       parent = alloc->getParentOfType<air::SegmentOp>();
     else if (clScope == "launch")
       parent = alloc->getParentOfType<air::LaunchOp>();
+    else if (clScope == "func")
+      parent = alloc->getParentOfType<func::FuncOp>();
     else
       return alloc->emitOpError(
-          "Invalid clScope value: expected one of herd/segment/launch");
+          "Invalid clScope value: expected one of herd/segment/launch/func");
 
     if (!parent)
       return failure();
@@ -2440,12 +2443,12 @@ class AIROverrideMemRefMemorySpacePass
 };
 
 void AIROverrideMemRefMemorySpacePass::runOnOperation() {
-  func::FuncOp funcOp = getOperation();
+  auto moduleOp = getOperation();
   MLIRContext *context = &getContext();
 
   RewritePatternSet patterns(context);
   patterns.add<OverrideMemorySpacePattern>(context, clScope, clMemorySpace);
-  (void)applyPatternsGreedily(funcOp, std::move(patterns));
+  (void)applyPatternsGreedily(moduleOp, std::move(patterns));
   RewritePatternSet fixResTypePatterns(context);
   if (clScope == "herd") {
     fixResTypePatterns.add<correctViewLikeOpIOMemorySpacesInScope<air::HerdOp>>(
@@ -2456,8 +2459,11 @@ void AIROverrideMemRefMemorySpacePass::runOnOperation() {
   } else if (clScope == "launch") {
     fixResTypePatterns
         .add<correctViewLikeOpIOMemorySpacesInScope<air::LaunchOp>>(context);
+  } else if (clScope == "func") {
+    fixResTypePatterns
+        .add<correctViewLikeOpIOMemorySpacesInScope<func::FuncOp>>(context);
   }
-  (void)applyPatternsGreedily(funcOp, std::move(fixResTypePatterns));
+  (void)applyPatternsGreedily(moduleOp, std::move(fixResTypePatterns));
 }
 
 } // namespace xilinx

test/xrt/31_triton_blk_ptr_eltwise_mul/run.py

@@ -64,7 +64,7 @@
                 "buffer-results-to-out-params",
                 "air-par-to-herd{depth=-1}",
                 "air-insert-launch-around-herd{insert-segment=false}",
-                "func.func(air-override-memref-memory-space{scope=herd memory-space=2})",
+                "air-override-memref-memory-space{scope=herd memory-space=2}",
                 "air-copy-to-dma",
                 "canonicalize",
                 "cse",

test/xrt/32_triton_matmul/run.py

@@ -62,7 +62,7 @@
                 "one-shot-bufferize",
                 f"func.func(air-wrap-func-with-parallel{{loop-bounds={launch_size[0]},{launch_size[1]},{launch_size[2]}}})",
                 "air-par-to-launch{depth=0 has-air-segment=true}",
-                "func.func(air-override-memref-memory-space{scope=launch memory-space=1})",
+                "air-override-memref-memory-space{scope=launch memory-space=1}",
             ]
         )
         + ")"

test/xrt/33_triton_matmul_ver2/run.py

@@ -67,7 +67,7 @@
                 "cse",
                 "air-par-to-herd{depth=-1}",
                 "air-insert-launch-around-herd{insert-segment=false}",
-                "func.func(air-override-memref-memory-space{scope=herd memory-space=2})",
+                "air-override-memref-memory-space{scope=herd memory-space=2}",
                 "air-copy-to-dma",
                 "canonicalize",
                 "cse",