Merge branch 'main' into fixup_eltwise_add_with_l2_missing_deallocs

Author: erwei-xilinx

mlir/include/air/Transform/AIRMiscPasses.h

@@ -35,7 +35,9 @@ std::unique_ptr<mlir::Pass> createAIRUnrollOuterPerfectlyNestedLoopsPass();
 std::unique_ptr<mlir::Pass> createAIRUnrollOuterPerfectlyNestedLoopsPass(
     AIRUnrollOuterPerfectlyNestedLoopsPassOptions options);
 std::unique_ptr<mlir::Pass> createAIRSplitL2MemrefForBufferConstraintPass();
-std::unique_ptr<Pass> createAIRForceL1MemrefInHerdPass();
+std::unique_ptr<Pass> createAIROverrideMemRefMemorySpacePass();
+std::unique_ptr<mlir::Pass> createAIROverrideMemRefMemorySpacePass(
+    AIROverrideMemRefMemorySpaceOptions options);
 
 } // namespace air
 } // namespace xilinx

mlir/include/air/Transform/PassDetail.h

@@ -72,7 +72,7 @@ namespace air {
 #define GEN_PASS_DEF_AIRSHRINKMEMREFSIZESBYACCESS
 #define GEN_PASS_DEF_AIRSPLITL2MEMREFFORBUFFERCONSTRAINTPASS
 #define GEN_PASS_DEF_DMATOCHANNEL
-#define GEN_PASS_DEF_AIRFORCEL1MEMREFINHERDPASS
+#define GEN_PASS_DEF_AIROVERRIDEMEMREFMEMORYSPACE
 #include "air/Transform/Passes.h.inc"
 
 } // namespace air

mlir/include/air/Transform/Passes.td

@@ -1503,12 +1503,19 @@ def DmaToChannel : Pass<"air-dma-to-channel", "ModuleOp"> {
   }];
 }
 
-def AIRForceL1MemrefInHerdPass : Pass<"air-force-l1-memref-in-herd", "func::FuncOp"> {
-  let summary = "Force all memrefs allocated within air.herd to have memory space L1.";
-  let constructor = "xilinx::air::createAIRForceL1MemrefInHerdPass()";
+def AIROverrideMemRefMemorySpace : Pass<"air-override-memref-memory-space", "func::FuncOp"> {
+  let summary = "Force all memrefs allocated within code region to have the specified memory space.";
+  let constructor = "xilinx::air::createAIROverrideMemRefMemorySpacePass()";
   let description = [{
-    Experimental pass. Force all memrefs allocated within air.herd to have memory space L1.
+    Experimental pass. Force all memrefs allocated within a specified code region to have the specified memory space.
   }];
+  let options = [
+    Option<"clMemorySpace", "memory-space", "unsigned", /*default=*/"0",
+           "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].">
+  ];
 }
 
 #endif // AIR_CONVERSION_PASSES

mlir/lib/Transform/AIRMiscPasses.cpp

@@ -1956,81 +1956,128 @@ void AIRSplitL2MemrefForBufferConstraintPass::runOnOperation() {
   air::renumberMemcpyIfOps(&func.getBody());
 }
 
-// Experimental: This pattern forces all memrefs allocated within the air.herd
-// to be L1.
-struct ForceL1MemrefInHerdPattern : public OpRewritePattern<memref::AllocOp> {
-  using OpRewritePattern<memref::AllocOp>::OpRewritePattern;
+// Experimental pattern to override the memory space of `memref.alloc`
+// operations when they appear inside a specified parent scope (e.g. herd,
+// segment).
+struct OverrideMemorySpacePattern : public OpRewritePattern<memref::AllocOp> {
+  OverrideMemorySpacePattern(MLIRContext *ctx, StringRef scope, int memSpace)
+      : OpRewritePattern<memref::AllocOp>(ctx), clScope(scope),
+        clMemorySpace(memSpace) {}
 
   LogicalResult matchAndRewrite(memref::AllocOp alloc,
                                 PatternRewriter &rewriter) const override {
-    auto parentHerdOp = alloc->getParentOfType<air::HerdOp>();
-    if (!parentHerdOp)
+    Operation *parent = nullptr;
+
+    if (clScope == "herd")
+      parent = alloc->getParentOfType<air::HerdOp>();
+    else if (clScope == "segment")
+      parent = alloc->getParentOfType<air::SegmentOp>();
+    else if (clScope == "launch")
+      parent = alloc->getParentOfType<air::LaunchOp>();
+    else
+      return alloc->emitOpError(
+          "Invalid clScope value: expected one of herd/segment/launch");
+
+    if (!parent)
       return failure();
 
-    auto memref = dyn_cast<MemRefType>(alloc.getMemref().getType());
-    if (!memref)
+    auto memrefTy = dyn_cast<MemRefType>(alloc.getMemref().getType());
+    if (!memrefTy)
       return failure();
-    if (memref.getMemorySpaceAsInt() == (int)air::MemorySpace::L1)
+    if ((int)memrefTy.getMemorySpaceAsInt() == clMemorySpace)
       return failure();
 
     auto newMemrefType =
-        MemRefType::get(memref.getShape(), memref.getElementType(),
-                        memref.getLayout().getAffineMap(),
-                        rewriter.getI32IntegerAttr((int)air::MemorySpace::L1));
+        MemRefType::get(memrefTy.getShape(), memrefTy.getElementType(),
+                        memrefTy.getLayout().getAffineMap(),
+                        rewriter.getI32IntegerAttr(clMemorySpace));
 
     rewriter.replaceOpWithNewOp<memref::AllocOp>(alloc, newMemrefType);
 
     return success();
   }
+
+private:
+  StringRef clScope; // Parent operation type to match
+  int clMemorySpace; // Target memory space value to assign
 };
-struct correctMemrefSubviewIOMemorySpaces
-    : public OpRewritePattern<memref::SubViewOp> {
-  using OpRewritePattern<memref::SubViewOp>::OpRewritePattern;
 
-  LogicalResult matchAndRewrite(memref::SubViewOp subview,
-                                PatternRewriter &rewriter) const override {
-    auto srcTy = dyn_cast<MemRefType>(subview.getViewSource().getType());
-    auto destTy = dyn_cast<MemRefType>(subview.getResult().getType());
-
-    auto subviewOutputType =
-        llvm::cast<MemRefType>(memref::SubViewOp::inferResultType(
-            srcTy, subview.getMixedOffsets(), subview.getMixedSizes(),
-            subview.getMixedStrides()));
-    if (destTy == subviewOutputType)
-      return failure();
+// Pattern to correct memory spaces of view-like operations within a given
+// scope, following the application of OverrideMemorySpacePattern.
+template <typename OpTy>
+struct correctViewLikeOpIOMemorySpacesInScope : public OpRewritePattern<OpTy> {
+  using OpRewritePattern<OpTy>::OpRewritePattern;
 
-    rewriter.replaceOpWithNewOp<memref::SubViewOp>(
-        subview, subviewOutputType, subview.getViewSource(),
-        subview.getMixedOffsets(), subview.getMixedSizes(),
-        subview.getMixedStrides());
+  LogicalResult matchAndRewrite(OpTy IFAOp,
+                                PatternRewriter &rewriter) const override {
 
+    if (!IFAOp->template hasTrait<OpTrait::IsIsolatedFromAbove>())
+      return failure();
+    llvm::DenseMap<ViewLikeOpInterface, SmallVector<OpResult>> viewLikeOpsToRes;
+    IFAOp->template walk([&](ViewLikeOpInterface viewLike) {
+      auto srcTy = dyn_cast<MemRefType>(viewLike.getViewSource().getType());
+      if (!srcTy)
+        return;
+      for (auto res : viewLike->getResults()) {
+        auto destTy = dyn_cast<MemRefType>(res.getType());
+        if (!destTy)
+          return;
+        if (srcTy.getMemorySpaceAsInt() == destTy.getMemorySpaceAsInt())
+          continue;
+        viewLikeOpsToRes[viewLike].push_back(res);
+      }
+    });
+    for (auto [viewLike, results] : viewLikeOpsToRes) {
+      for (OpResult res : results) {
+        auto srcTy = dyn_cast<MemRefType>(viewLike.getViewSource().getType());
+        auto destTy = dyn_cast<MemRefType>(res.getType());
+        MemRefType::Builder builder(destTy);
+        builder.setMemorySpace(srcTy.getMemorySpace());
+        rewriter.modifyOpInPlace(viewLike, [&]() { res.setType(builder); });
+      }
+    }
     return success();
   }
 };
 
-// An experimental pass forcing all memrefs allocated within an air.herd to have
-// memory space L1.
-class AIRForceL1MemrefInHerdPass
-    : public air::impl::AIRForceL1MemrefInHerdPassBase<
-          AIRForceL1MemrefInHerdPass> {
+// An experimental pass forcing all memrefs allocated within a specified air
+// code region to have the specified memory space.
+class AIROverrideMemRefMemorySpacePass
+    : public air::impl::AIROverrideMemRefMemorySpaceBase<
+          AIROverrideMemRefMemorySpacePass> {
 
 public:
-  AIRForceL1MemrefInHerdPass() = default;
-  AIRForceL1MemrefInHerdPass(const AIRForceL1MemrefInHerdPass &pass){};
+  AIROverrideMemRefMemorySpacePass() = default;
+  AIROverrideMemRefMemorySpacePass(
+      const AIROverrideMemRefMemorySpacePass &pass){};
+  AIROverrideMemRefMemorySpacePass(
+      const ::xilinx::air::AIROverrideMemRefMemorySpaceOptions &options)
+      : AIROverrideMemRefMemorySpaceBase(options) {}
 
   void runOnOperation() override;
 
 private:
 };
 
-void AIRForceL1MemrefInHerdPass::runOnOperation() {
+void AIROverrideMemRefMemorySpacePass::runOnOperation() {
   func::FuncOp funcOp = getOperation();
   MLIRContext *context = &getContext();
 
   RewritePatternSet patterns(context);
-  patterns.add<ForceL1MemrefInHerdPattern, correctMemrefSubviewIOMemorySpaces>(
-      context);
+  patterns.add<OverrideMemorySpacePattern>(context, clScope, clMemorySpace);
   (void)applyPatternsGreedily(funcOp, std::move(patterns));
+  RewritePatternSet fixResTypePatterns(context);
+  if (clScope == "herd") {
+    fixResTypePatterns.add<correctViewLikeOpIOMemorySpacesInScope<air::HerdOp>>(
+        context);
+  } else if (clScope == "segment") {
+    fixResTypePatterns
+        .add<correctViewLikeOpIOMemorySpacesInScope<air::SegmentOp>>(context);
+  } else if (clScope == "launch") {
+    fixResTypePatterns
+        .add<correctViewLikeOpIOMemorySpacesInScope<air::LaunchOp>>(context);
+  }
+  (void)applyPatternsGreedily(funcOp, std::move(fixResTypePatterns));
 }
 
 } // anonymous namespace
@@ -2092,8 +2139,12 @@ std::unique_ptr<Pass> createAIRSplitL2MemrefForBufferConstraintPass() {
   return std::make_unique<AIRSplitL2MemrefForBufferConstraintPass>();
 }
 
-std::unique_ptr<Pass> createAIRForceL1MemrefInHerdPass() {
-  return std::make_unique<AIRForceL1MemrefInHerdPass>();
+std::unique_ptr<Pass> createAIROverrideMemRefMemorySpacePass() {
+  return std::make_unique<AIROverrideMemRefMemorySpacePass>();
+}
+std::unique_ptr<Pass> createAIROverrideMemRefMemorySpacePass(
+    AIROverrideMemRefMemorySpaceOptions options) {
+  return std::make_unique<AIROverrideMemRefMemorySpacePass>(options);
 }
 
 } // namespace air

mlir/test/Transform/AIRMiscPasses/air_force_l1_memref_in_herd.mlir

@@ -0,0 +1,77 @@
+//===- air_override_memref_memory_space.mlir -------------------*- MLIR -*-===//
+//
+// Copyright (C) 2025, Advanced Micro Devices, Inc. All rights reserved.
+// SPDX-License-Identifier: MIT
+//
+//===----------------------------------------------------------------------===//
+
+// RUN: air-opt %s -air-override-memref-memory-space="scope=herd memory-space=2" | FileCheck %s
+// RUN: air-opt %s -air-override-memref-memory-space="scope=launch memory-space=2" | FileCheck %s --check-prefix=LAUNCH
+// RUN: air-opt %s -air-override-memref-memory-space="scope=launch memory-space=1" | FileCheck %s --check-prefix=MS1
+
+module {
+
+  // CHECK-LABEL: func.func @func0
+  // CHECK: memref.alloc() : memref<32x64xf32, 2 : i32>
+  // LAUNCH-LABEL: func.func @func0
+  // LAUNCH: memref.alloc() : memref<32x64xf32, 2 : i32>
+  // MS1-LABEL: func.func @func0
+  // MS1: memref.alloc() : memref<32x64xf32, 1 : i32>
+
+  func.func @func0(%arg0: memref<*xf32> {tt.divisibility = 16 : i32}, %arg1: memref<*xf32> {tt.divisibility = 16 : i32}, %arg2: memref<*xf32> {tt.divisibility = 16 : i32}, %arg3: i32, %arg4: i32, %arg5: i32, %arg6: i32, %arg7: i32, %arg8: i32) {
+    %c2 = arith.constant 2 : index
+    %c1 = arith.constant 1 : index
+    %c1_0 = arith.constant 1 : index
+    air.launch (%arg9, %arg10) in (%arg11=%c1, %arg12=%c1_0) args(%arg13=%arg0, %arg14=%arg1, %arg15=%arg2) : memref<*xf32>, memref<*xf32>, memref<*xf32> {
+      air.segment @bare_matmul_0  args(%arg16=%arg9, %arg17=%arg10, %arg18=%arg11, %arg19=%arg12, %arg20=%arg13, %arg21=%arg14, %arg22=%arg15) : index, index, index, index, memref<*xf32>, memref<*xf32>, memref<*xf32> {
+        %c2_1 = arith.constant 2 : index
+        %c2_2 = arith.constant 2 : index
+        air.herd @herd_0  tile (%arg23, %arg24) in (%arg25=%c2_1, %arg26=%c2_2) args(%arg27=%arg16, %arg28=%arg17, %arg29=%arg18, %arg30=%arg19, %arg31=%arg20, %arg32=%arg21, %arg33=%arg22) : index, index, index, index, memref<*xf32>, memref<*xf32>, memref<*xf32> {
+          %c32 = arith.constant 32 : index
+          %c2048 = arith.constant 2048 : index
+          %c64 = arith.constant 64 : index
+          %1 = arith.muli %arg23, %c2048 : index
+          %reinterpret_cast = memref.reinterpret_cast %arg31 to offset: [%1], sizes: [32, 64], strides: [%c64, 1] : memref<*xf32> to memref<32x64xf32, strided<[?, 1], offset: ?>>
+          %alloc = memref.alloc() : memref<32x64xf32>
+          memref.copy %reinterpret_cast, %alloc : memref<32x64xf32, strided<[?, 1], offset: ?>> to memref<32x64xf32>
+        }
+      }
+    }
+    return
+  }
+
+  // LAUNCH-LABEL: func.func @func1
+  // LAUNCH: memref.alloc() : memref<8x4x4x8xf32, 2 : i32>
+  // LAUNCH: memref.collapse_shape {{.*}} : memref<8x4x4x8xf32, 2 : i32> into memref<32x32xf32, 2 : i32>
+  // LAUNCH: memref.alloc() : memref<4x8x8x4xf32, 2 : i32>
+  // LAUNCH: memref.collapse_shape {{.*}} : memref<4x8x8x4xf32, 2 : i32> into memref<32x32xf32, 2 : i32>
+  // LAUNCH: memref.alloc() : memref<32x32xf32, 2 : i32>
+  // LAUNCH: memref.expand_shape {{.*}} : memref<32x32xf32, 2 : i32> into memref<8x4x8x4xf32, 2 : i32>
+  // LAUNCH: memref.alloc() : memref<8x8x4x4xf32, 2 : i32>
+  // MS1-LABEL: func.func @func1
+  // MS1: memref.alloc() : memref<8x4x4x8xf32, 1 : i32>
+  // MS1: memref.collapse_shape {{.*}} : memref<8x4x4x8xf32, 1 : i32> into memref<32x32xf32, 1 : i32>
+  // MS1: memref.alloc() : memref<4x8x8x4xf32, 1 : i32>
+  // MS1: memref.collapse_shape {{.*}} : memref<4x8x8x4xf32, 1 : i32> into memref<32x32xf32, 1 : i32>
+  // MS1: memref.alloc() : memref<32x32xf32, 1 : i32>
+  // MS1: memref.expand_shape {{.*}} : memref<32x32xf32, 1 : i32> into memref<8x4x8x4xf32, 1 : i32>
+  // MS1: memref.alloc() : memref<8x8x4x4xf32, 1 : i32>
+
+  func.func @func1(%arg0: memref<*xf32> {tt.divisibility = 16 : i32}, %arg1: memref<*xf32> {tt.divisibility = 16 : i32}, %arg2: memref<*xf32> {tt.divisibility = 16 : i32}, %arg3: i32, %arg4: i32, %arg5: i32, %arg6: i32, %arg7: i32, %arg8: i32) {
+    %c1 = arith.constant 1 : index
+    %c2 = arith.constant 2 : index
+    air.launch (%arg9, %arg10, %arg11) in (%arg12=%c2, %arg13=%c2, %arg14=%c1) args(%arg15=%arg0, %arg16=%arg1, %arg17=%arg2) : memref<*xf32>, memref<*xf32>, memref<*xf32> {
+      %cst = arith.constant 0.000000e+00 : f32
+      %alloc_2 = memref.alloc() {alignment = 64 : i64} : memref<8x4x4x8xf32>
+      %collapse_shape = memref.collapse_shape %alloc_2 [[0, 1], [2, 3]] : memref<8x4x4x8xf32> into memref<32x32xf32>
+      %alloc_3 = memref.alloc() {alignment = 64 : i64} : memref<4x8x8x4xf32>
+      %collapse_shape_4 = memref.collapse_shape %alloc_3 [[0, 1], [2, 3]] : memref<4x8x8x4xf32> into memref<32x32xf32>
+      %alloc_5 = memref.alloc() {alignment = 64 : i64} : memref<32x32xf32>
+      linalg.matmul ins(%collapse_shape, %collapse_shape_4 : memref<32x32xf32>, memref<32x32xf32>) outs(%alloc_5 : memref<32x32xf32>)
+      %expand_shape = memref.expand_shape %alloc_5 [[0, 1], [2, 3]] output_shape [8, 4, 8, 4] : memref<32x32xf32> into memref<8x4x8x4xf32>
+      %alloc_6 = memref.alloc() {alignment = 64 : i64} : memref<8x8x4x4xf32>
+      linalg.transpose ins(%expand_shape : memref<8x4x8x4xf32>) outs(%alloc_6 : memref<8x8x4x4xf32>) permutation = [0, 2, 1, 3] 
+    }
+    return
+  }
+}

mlir/test/Transform/AIRMiscPasses/air_override_memref_memory_space.mlir

@@ -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-force-l1-memref-in-herd)",
+                "func.func(air-override-memref-memory-space{scope=herd memory-space=2})",
                 "air-copy-to-dma",
                 "canonicalize",
                 "cse",