Purge AIRDealiasMemref pass as it is obsolete (Xilinx#920)

erwei-xilinx · web-flow · commit 2d3e948ff8e0 · 2025-02-27T19:49:23.000-08:00
diff --git a/mlir/include/air/Transform/AIRDependencyScheduleOpt.h b/mlir/include/air/Transform/AIRDependencyScheduleOpt.h
@@ -49,8 +49,6 @@ std::unique_ptr<mlir::Pass> createAIRUnrollChannelByFactorPattern();
 
 std::unique_ptr<mlir::Pass> createAIREnforceLoopCarriedMemrefDeallocPattern();
 
-std::unique_ptr<mlir::Pass> createAIRDeAliasMemref();
-
 std::unique_ptr<mlir::Pass> createAIRFuseChannels();
 std::unique_ptr<OperationPass<ModuleOp>>
 createAIRFuseChannels(const AIRFuseChannelsOptions &);
diff --git a/mlir/include/air/Transform/Passes.td b/mlir/include/air/Transform/Passes.td
@@ -1010,16 +1010,6 @@ def AIREnforceLoopCarriedMemrefDeallocPattern: Pass<"air-enforce-loop-carried-me
   }];
 }
 
-def AIRDeAliasMemref: Pass<"air-dealias-memref", "ModuleOp"> {
-  let summary = "De-alias a memref into multiple memrefs";
-  let constructor = "xilinx::air::createAIRDeAliasMemref()";
-  let description = [{
-    This pass detects memrefs which can de-alias into multiple memrefs over time, and
-    generates extra copies of this memref. This process can improve the stability of the
-    memref-to-buffer mappings.
-  }];
-}
-
 def AIRFuseChannels: Pass<"air-fuse-channels", "ModuleOp"> {
   let summary = "Fuse multiple air.channel ops into one";
   let constructor = "xilinx::air::createAIRFuseChannels()";
diff --git a/mlir/lib/Transform/AIRDependencyScheduleOpt.cpp b/mlir/lib/Transform/AIRDependencyScheduleOpt.cpp
@@ -3257,178 +3257,6 @@ class AIREnforceLoopCarriedMemrefDeallocPattern
 private:
 };
 
-// A pass which de-alias a memref with multiple channel accesses over time, into
-// multiple memrefs. Note that this implementation is temporary and not generic.
-// TODO: Rewrite as a graph partitioning problem.
-class AIRDeAliasMemref
-    : public xilinx::air::impl::AIRDeAliasMemrefBase<AIRDeAliasMemref> {
-
-public:
-  AIRDeAliasMemref() = default;
-  AIRDeAliasMemref(const AIRDeAliasMemref &pass) {}
-
-  void getDependentDialects(::mlir::DialectRegistry &registry) const override {
-    registry.insert<scf::SCFDialect, air::airDialect>();
-  }
-
-  void runOnFunction(func::FuncOp f) {
-
-    std::vector<memref::AllocOp> allocs;
-    f.walk([&](memref::AllocOp alloc) { allocs.push_back(alloc); });
-
-    // Count air.channel references
-    for (auto alloc : allocs) {
-      Value memref = nullptr;
-      if (auto exec = alloc->getParentOfType<air::ExecuteOp>()) {
-        memref = exec->getResult(1);
-      } else
-        memref = alloc.getMemref();
-      std::vector<air::ChannelInterface> chan_puts_gets;
-      for (auto user : memref.getUsers()) {
-        if (auto putget = dyn_cast<air::ChannelInterface>(user))
-          if (putget.getMemref() == memref)
-            chan_puts_gets.push_back(putget);
-      }
-
-      // Partition the subgraph
-      std::vector<int> partition_cuts;
-      if (!chan_puts_gets.empty()) {
-        for (unsigned i = 0; i < chan_puts_gets.size() - 1; i++) {
-          if (isa<air::ChannelGetOp>(chan_puts_gets[i].getOperation()) &&
-              isa<air::ChannelPutOp>(chan_puts_gets[i + 1].getOperation())) {
-            partition_cuts.push_back(i + 1);
-          }
-        }
-      }
-
-      // Allocate new memref per cut
-      std::vector<Operation *> new_memallocs;
-      for (unsigned i = 0; i < partition_cuts.size(); i++) {
-        OpBuilder builder(alloc);
-        Operation *new_op = nullptr;
-        if (auto exec = alloc->getParentOfType<air::ExecuteOp>()) {
-          builder.setInsertionPoint(exec);
-          new_op = builder.clone(*exec.getOperation());
-        } else
-          new_op = builder.clone(*alloc.getOperation());
-        new_memallocs.push_back(new_op);
-
-        // Create deallocs for the new memref
-        Value new_memref = isa<air::ExecuteOp>(new_op) ? new_op->getResult(1)
-                                                       : new_op->getResult(0);
-        for (auto user : memref.getUsers()) {
-          if (isa<memref::DeallocOp>(user)) {
-            if (isa<air::ExecuteOp>(new_op)) {
-              builder.setInsertionPoint(
-                  user->getParentOfType<air::ExecuteOp>());
-              // Async. dealloc
-              auto async_exec = builder.create<xilinx::air::ExecuteOp>(
-                  user->getLoc(), air::AsyncTokenType::get(alloc->getContext()),
-                  SmallVector<Value>{});
-              Block *async_exec_bb =
-                  builder.createBlock(&async_exec.getRegion());
-              builder.setInsertionPointToStart(async_exec_bb);
-              builder.create<memref::DeallocOp>(user->getLoc(), new_memref);
-              builder.create<air::ExecuteTerminatorOp>(user->getLoc());
-            } else {
-              builder.setInsertionPoint(user);
-              // Sync. dealloc
-              builder.create<memref::DeallocOp>(user->getLoc(), new_memref);
-            }
-          }
-        }
-      }
-
-      // Update references
-      partition_cuts.insert(partition_cuts.end(), chan_puts_gets.size());
-      for (unsigned i = 0; i < partition_cuts.size() - 1; i++) {
-        for (int j = partition_cuts[i]; j < partition_cuts[i + 1]; j++) {
-          if (auto old_put = dyn_cast<air::ChannelPutOp>(
-                  chan_puts_gets[j].getOperation())) {
-            Value new_memref = isa<air::ExecuteOp>(new_memallocs[i])
-                                   ? new_memallocs[i]->getResult(1)
-                                   : new_memallocs[i]->getResult(0);
-            OpBuilder builder(old_put);
-            replaceChannelPutOp(builder, old_put, new_memref);
-          } else if (auto old_get = dyn_cast<air::ChannelGetOp>(
-                         chan_puts_gets[j].getOperation())) {
-            Value new_memref = isa<air::ExecuteOp>(new_memallocs[i])
-                                   ? new_memallocs[i]->getResult(1)
-                                   : new_memallocs[i]->getResult(0);
-            OpBuilder builder(old_get);
-            replaceChannelGetOp(builder, old_get, new_memref);
-          }
-        }
-      }
-    }
-  }
-
-  void runOnOperation() override {
-    auto module = getOperation();
-
-    SmallVector<func::FuncOp, 4> funcOps;
-    module.walk([&](func::FuncOp op) { funcOps.push_back(op); });
-    for (auto f : funcOps) {
-      runOnFunction(f);
-    }
-  }
-
-private:
-  Operation *replaceChannelPutOp(OpBuilder builder, air::ChannelPutOp old,
-                                 Value new_memref) {
-    builder.setInsertionPoint(old);
-    SmallVector<Type, 1> tys;
-    if (old.getAsyncToken()) {
-      tys.push_back(air::AsyncTokenType::get(old->getContext()));
-    }
-    SmallVector<Value, 4> deps = old.getAsyncDependencies();
-    auto new_op = builder.create<air::ChannelPutOp>(
-        old->getLoc(), tys, deps, old.getChanName(), old.getIndices(),
-        new_memref, old.getSrcOffsets(), old.getSrcSizes(),
-        old.getSrcStrides());
-    if (old.getAsyncToken()) {
-      old.getAsyncToken().replaceAllUsesWith(new_op.getAsyncToken());
-      // Add dependence to the new memref
-      new_op.addAsyncDependency(
-          dyn_cast<air::ExecuteOp>(new_memref.getDefiningOp()).getAsyncToken());
-    }
-    if (old.getId() != -1) {
-      new_op->setAttr("id", mlir::IntegerAttr::get(
-                                mlir::IntegerType::get(old->getContext(), 32),
-                                old.getId()));
-    }
-    old->erase();
-    return new_op.getOperation();
-  }
-  Operation *replaceChannelGetOp(OpBuilder builder, air::ChannelGetOp old,
-                                 Value new_memref) {
-    builder.setInsertionPoint(old);
-    SmallVector<Type, 1> tys;
-    if (old.getAsyncToken()) {
-      tys.push_back(air::AsyncTokenType::get(old->getContext()));
-    }
-    SmallVector<Value, 4> deps = old.getAsyncDependencies();
-    auto new_op = builder.create<air::ChannelGetOp>(
-        old->getLoc(), tys, deps, old.getChanName(), old.getIndices(),
-        new_memref, old.getDstOffsets(), old.getDstSizes(),
-        old.getDstStrides());
-    new_op->setAttrs(old->getDiscardableAttrDictionary());
-    if (old.getAsyncToken()) {
-      old.getAsyncToken().replaceAllUsesWith(new_op.getAsyncToken());
-      // Add dependence to the new memref
-      new_op.addAsyncDependency(
-          dyn_cast<air::ExecuteOp>(new_memref.getDefiningOp()).getAsyncToken());
-    }
-    if (old.getId() != -1) {
-      new_op->setAttr("id", mlir::IntegerAttr::get(
-                                mlir::IntegerType::get(old->getContext(), 32),
-                                old.getId()));
-    }
-    old->erase();
-    return new_op.getOperation();
-  }
-};
-
 // A pass which transform multiple channel ops into one, where the data movement
 // is time-multiplexed.
 class AIRFuseChannels
@@ -6175,10 +6003,6 @@ std::unique_ptr<Pass> createAIREnforceLoopCarriedMemrefDeallocPattern() {
   return std::make_unique<AIREnforceLoopCarriedMemrefDeallocPattern>();
 }
 
-std::unique_ptr<Pass> createAIRDeAliasMemref() {
-  return std::make_unique<AIRDeAliasMemref>();
-}
-
 std::unique_ptr<Pass> createAIRFuseChannels() {
   return std::make_unique<AIRFuseChannels>();
 }
diff --git a/mlir/test/Transform/AIRDependencyScheduleOpt/dealias_memref.mlir b/mlir/test/Transform/AIRDependencyScheduleOpt/dealias_memref.mlir
diff --git a/test/xrt/23_ctrlpkt_config/aie.py b/test/xrt/23_ctrlpkt_config/aie.py
@@ -142,12 +142,26 @@ def forward(lhs, rhs):
             "air-dependency-canonicalize",
             "canonicalize",
             "cse",
+            "air-isolate-async-dma-loop-nests",
+            "canonicalize",
+            "cse",
+            "air-fuse-channels",
+            "canonicalize",
+            "cse",
+            ### Scaling to 4 AIE columns
+            "func.func(air-split-l2-memref)",
+            "canonicalize",
+            "cse",
+            "air-isolate-async-dma-loop-nests",
+            ###
+            "canonicalize",
+            "cse",
+            "func.func(air-fuse-alloc-dealloc)",
+            "func.func(air-shrink-memref-sizes-by-access)",
             "air-label-scf-for-to-ping-pong",
             "air-ping-pong-transform",
-            "air-dealias-memref",
             "canonicalize",
             "cse",
-            "air-isolate-async-dma-loop-nests",
             "func.func(air-opt-memtile-dma-bds{device=npu1_4col})",
             "canonicalize",
             "cse",
diff --git a/test/xrt/24_ctrlpkt_config_2gemms_4x4/aie.py b/test/xrt/24_ctrlpkt_config_2gemms_4x4/aie.py
@@ -185,10 +185,14 @@
                 "cse",
                 ### Scaling to 4 AIE columns
                 "func.func(air-split-l2-memref)",
+                "canonicalize",
+                "cse",
                 "air-isolate-async-dma-loop-nests",
                 ###
                 "canonicalize",
                 "cse",
+                "func.func(air-fuse-alloc-dealloc)",
+                "func.func(air-shrink-memref-sizes-by-access)",
                 "func.func(air-loop-fusion)",
                 "air-label-scf-for-to-ping-pong",
                 "air-ping-pong-transform",
diff --git a/test/xrt/24_ctrlpkt_config_2gemms_4x4/aie2.py b/test/xrt/24_ctrlpkt_config_2gemms_4x4/aie2.py
@@ -190,6 +190,8 @@
                 ###
                 "canonicalize",
                 "cse",
+                "func.func(air-fuse-alloc-dealloc)",
+                "func.func(air-shrink-memref-sizes-by-access)",
                 "func.func(air-loop-fusion)",
                 "air-label-scf-for-to-ping-pong",
                 "air-ping-pong-transform",
