Handle memref.copy/subview ops & Add compiler pass pipeline

include/Conversion/ConversionPasses.h

@@ -23,6 +23,10 @@ std::unique_ptr<mlir::Pass> createLowerAffineToNeuraPass();
 // TaskFlow Conversion Passes.
 std::unique_ptr<mlir::Pass> createConvertAffineToTaskflowPass();
 std::unique_ptr<mlir::Pass> createConvertTaskflowToNeuraPass();
+
+// Memref SubView and Copy Conversion Passes.
+std::unique_ptr<mlir::Pass> createFoldSubViewPass();
+std::unique_ptr<mlir::Pass> createConvertCopyToAffineLoopsPass();
 #define GEN_PASS_REGISTRATION
 #include "Conversion/ConversionPasses.h.inc"
 

include/Conversion/ConversionPasses.td

@@ -93,4 +93,35 @@ def ConvertTaskflowToNeura : Pass<"convert-taskflow-to-neura", "ModuleOp">{
   ];
 }
 
+def FoldSubView : Pass<"fold-subview", "func::FuncOp"> {
+  let summary = "Folds memref.subview into affine load/store operations";
+  let description = [{
+    Eliminates memref.subview operations by folding them into their users
+    (affine.load and affine.store). Adjusts indices to account for subview
+    offsets and strides, enabling direct access to the source memref.
+  }];
+  let constructor = "mlir::createFoldSubViewPass()";
+  let dependentDialects = [
+    "mlir::affine::AffineDialect",
+    "mlir::memref::MemRefDialect",
+    "mlir::arith::ArithDialect",
+    "mlir::func::FuncDialect"
+  ];
+}
+
+def ConvertCopyToAffineLoops : Pass<"convert-copy-to-affine-loops", "func::FuncOp"> {
+  let summary = "Converts memref.copy to explicit affine loop nests";
+  let description = [{
+    Converts memref.copy operations into explicit affine.for loop nests
+    with affine.load and affine.store operations. This makes data movement
+    explicit for CGRA compilation with global addressing.
+  }];
+  let constructor = "mlir::createConvertCopyToAffineLoopsPass()";
+  let dependentDialects = [
+    "mlir::affine::AffineDialect",
+    "mlir::memref::MemRefDialect",
+    "mlir::func::FuncDialect"
+  ];
+}
+
 #endif // CONVERSION_PASSES_TD

include/TaskflowDialect/TaskflowPasses.h

@@ -12,6 +12,10 @@
 #include <memory>
 namespace mlir {
 namespace taskflow {
+
+void registerTaskflowConversionPassPipeline();
+void registerTosaToAffineConversionPassPipeline();
+
 // Passes defined in TaskflowPasses.td
 #define GEN_PASS_DECL
 #include "TaskflowDialect/TaskflowPasses.h.inc"

lib/Conversion/AffineToTaskflow/CMakeLists.txt

@@ -2,6 +2,8 @@ include_directories(${CMAKE_CURRENT_BINARY_DIR})
 
 add_mlir_conversion_library(MLIRAffineToTaskflowPass
   AffineToTaskflowPass.cpp
+  FoldSubViewPass.cpp
+  ConvertCopyToAffineLoopsPass.cpp
 
   DEPENDS
   MLIRConversionIncGen

lib/Conversion/AffineToTaskflow/ConvertCopyToAffineLoopsPass.cpp

@@ -0,0 +1,99 @@
+//===- ConvertCopyToAffineLoopsPass.cpp - Converts memref.copy to loops ---===//
+//
+// This pass converts memref.copy operations into explicit affine loop nests.
+//
+//===----------------------------------------------------------------------===//
+
+#include "Conversion/ConversionPasses.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+
+using namespace mlir;
+
+namespace {
+// Converts memref.copy to nested affine loops with affine.load/store.
+struct CopyOpLoweringPattern : public OpRewritePattern<memref::CopyOp> {
+  using OpRewritePattern<memref::CopyOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(memref::CopyOp copy,
+                                PatternRewriter &rewriter) const override {
+    // Checks if the target has any users besides this copy.
+    // If the target (e.g., a subview) is only used by this copy and nothing
+    // else, this copy is dead code and should be removed without conversion.
+    Value target = copy.getTarget();
+    bool has_other_users = false;
+    for (auto *user : target.getUsers()) {
+      if (user != copy.getOperation()) {
+        has_other_users = true;
+        break;
+      }
+    }
+
+    if (!has_other_users) {
+      // Target has no users besides this copy, so just erase the copy.
+      rewriter.eraseOp(copy);
+      return success();
+    }
+
+    // Target has other users, convert copy to affine loops.
+    rewriter.setInsertionPoint(copy);
+    auto loc = copy.getLoc();
+    MemRefType memref_type = dyn_cast<MemRefType>(copy.getSource().getType());
+
+    // Creates explicit memory copy using an affine loop nest.
+    SmallVector<Value> ivs;
+    for (auto dim_size : memref_type.getShape()) {
+      auto loop = rewriter.create<affine::AffineForOp>(loc, 0, dim_size);
+      rewriter.setInsertionPointToStart(loop.getBody());
+      ivs.push_back(loop.getInductionVar());
+    }
+
+    // Creates affine load from source and store to target.
+    Value value =
+        rewriter.create<affine::AffineLoadOp>(loc, copy.getSource(), ivs);
+    rewriter.create<affine::AffineStoreOp>(loc, value, copy.getTarget(), ivs);
+
+    rewriter.eraseOp(copy);
+    return success();
+  }
+};
+} // namespace
+
+//===----------------------------------------------------------------------===//
+// Pass implementation
+//===----------------------------------------------------------------------===//
+
+namespace {
+struct ConvertCopyToAffineLoopsPass
+    : public PassWrapper<ConvertCopyToAffineLoopsPass,
+                         OperationPass<func::FuncOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(ConvertCopyToAffineLoopsPass)
+
+  StringRef getArgument() const final { return "convert-copy-to-affine-loops"; }
+
+  StringRef getDescription() const final {
+    return "Convert memref.copy to explicit affine loop nests";
+  }
+
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<affine::AffineDialect, memref::MemRefDialect,
+                    func::FuncDialect>();
+  }
+
+  void runOnOperation() override {
+    RewritePatternSet patterns(&getContext());
+    patterns.add<CopyOpLoweringPattern>(&getContext());
+    if (failed(applyPatternsGreedily(getOperation(), std::move(patterns)))) {
+      signalPassFailure();
+    }
+  }
+};
+} // namespace
+
+std::unique_ptr<Pass> mlir::createConvertCopyToAffineLoopsPass() {
+  return std::make_unique<ConvertCopyToAffineLoopsPass>();
+}

lib/Conversion/AffineToTaskflow/FoldSubViewPass.cpp

@@ -0,0 +1,191 @@
+//===- FoldSubViewPass.cpp - Fold memref.subview into load/store ---------===//
+//
+// This pass folds memref.subview operations into their affine.load and
+// affine.store users by adjusting the access indices. Designed for CGRA
+// systems with global addressing.
+//
+//===----------------------------------------------------------------------===//
+
+#include "Conversion/ConversionPasses.h"
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "mlir/Pass/Pass.h"
+#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
+#include "llvm/ADT/SmallBitVector.h"
+
+using namespace mlir;
+
+//===----------------------------------------------------------------------===//
+// Utility functions.
+//===----------------------------------------------------------------------===//
+
+// Resolves the source indices for a load/store operation that accesses a
+// subview. Computes the adjusted indices to access the source memref directly.
+//
+// For example:
+//   %subview = memref.subview %source[%offset0, %offset1][...][%stride0,
+//   %stride1] %val = affine.load %subview[%i, %j]
+// becomes:
+//   %val = affine.load %source[%i * %stride0 + %offset0, %j * %stride1 +
+//   %offset1]
+static LogicalResult
+resolveSourceIndices(Location loc, PatternRewriter &rewriter,
+                     memref::SubViewOp sub_view_op, ValueRange indices,
+                     SmallVectorImpl<Value> &source_indices) {
+  SmallVector<OpFoldResult> mixed_offsets = sub_view_op.getMixedOffsets();
+  SmallVector<OpFoldResult> mixed_sizes = sub_view_op.getMixedSizes();
+  SmallVector<OpFoldResult> mixed_strides = sub_view_op.getMixedStrides();
+
+  SmallVector<Value> use_indices;
+  // Handles rank-reducing subviews: for every unit-dim size, adds a zero index.
+  unsigned result_dim = 0;
+  llvm::SmallBitVector unused_dims = sub_view_op.getDroppedDims();
+  for (auto dim :
+       llvm::seq<unsigned>(0, sub_view_op.getSourceType().getRank())) {
+    if (unused_dims.test(dim)) {
+      use_indices.push_back(rewriter.create<arith::ConstantIndexOp>(loc, 0));
+    } else {
+      use_indices.push_back(indices[result_dim++]);
+    }
+  }
+
+  if (use_indices.size() != mixed_offsets.size()) {
+    return failure();
+  }
+
+  source_indices.resize(use_indices.size());
+  for (auto index : llvm::seq<size_t>(0, mixed_offsets.size())) {
+    SmallVector<Value> dynamic_operands;
+    AffineExpr expr = rewriter.getAffineDimExpr(0);
+    unsigned num_symbols = 0;
+    dynamic_operands.push_back(use_indices[index]);
+
+    // Multiplies by stride: index * stride.
+    if (auto attr = mixed_strides[index].dyn_cast<Attribute>()) {
+      int64_t stride_val = dyn_cast<IntegerAttr>(attr).getInt();
+      if (stride_val != 1) {
+        expr = expr * stride_val;
+      }
+    } else {
+      dynamic_operands.push_back(dyn_cast<Value>(mixed_strides[index]));
+      expr = expr * rewriter.getAffineSymbolExpr(num_symbols++);
+    }
+
+    // Adds offset: index * stride + offset.
+    if (auto attr = mixed_offsets[index].dyn_cast<Attribute>()) {
+      int64_t offset_val = dyn_cast<IntegerAttr>(attr).getInt();
+      if (offset_val != 0) {
+        expr = expr + offset_val;
+      }
+    } else {
+      dynamic_operands.push_back(dyn_cast<Value>(mixed_offsets[index]));
+      expr = expr + rewriter.getAffineSymbolExpr(num_symbols++);
+    }
+
+    source_indices[index] = rewriter.create<affine::AffineApplyOp>(
+        loc, AffineMap::get(1, num_symbols, expr), dynamic_operands);
+  }
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
+// Patterns
+//===----------------------------------------------------------------------===//
+
+namespace {
+/// Folds affine.load from a subview into a load from the source memref.
+class LoadOpOfSubViewFolder final
+    : public OpRewritePattern<affine::AffineLoadOp> {
+public:
+  using OpRewritePattern<affine::AffineLoadOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(affine::AffineLoadOp loadOp,
+                                PatternRewriter &rewriter) const override {
+    auto subViewOp = loadOp.getMemRef().getDefiningOp<memref::SubViewOp>();
+    if (!subViewOp)
+      return failure();
+
+    SmallVector<Value, 4> sourceIndices;
+    if (failed(resolveSourceIndices(loadOp.getLoc(), rewriter, subViewOp,
+                                    loadOp.getIndices(), sourceIndices)))
+      return failure();
+
+    rewriter.replaceOpWithNewOp<affine::AffineLoadOp>(
+        loadOp, subViewOp.getSource(), sourceIndices);
+    return success();
+  }
+};
+
+/// Folds affine.store to a subview into a store to the source memref.
+class StoreOpOfSubViewFolder final
+    : public OpRewritePattern<affine::AffineStoreOp> {
+public:
+  using OpRewritePattern<affine::AffineStoreOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(affine::AffineStoreOp storeOp,
+                                PatternRewriter &rewriter) const override {
+    auto subViewOp = storeOp.getMemRef().getDefiningOp<memref::SubViewOp>();
+    if (!subViewOp)
+      return failure();
+
+    SmallVector<Value, 4> sourceIndices;
+    if (failed(resolveSourceIndices(storeOp.getLoc(), rewriter, subViewOp,
+                                    storeOp.getIndices(), sourceIndices)))
+      return failure();
+
+    rewriter.replaceOpWithNewOp<affine::AffineStoreOp>(
+        storeOp, storeOp.getValue(), subViewOp.getSource(), sourceIndices);
+    return success();
+  }
+};
+} // namespace
+
+//===----------------------------------------------------------------------===//
+// Pass implementation
+//===----------------------------------------------------------------------===//
+
+namespace {
+struct FoldSubViewPass
+    : public PassWrapper<FoldSubViewPass, OperationPass<func::FuncOp>> {
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(FoldSubViewPass)
+
+  StringRef getArgument() const final { return "fold-subview"; }
+
+  StringRef getDescription() const final {
+    return "Fold memref.subview into affine load/store operations";
+  }
+
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<affine::AffineDialect, memref::MemRefDialect,
+                    arith::ArithDialect, func::FuncDialect>();
+  }
+
+  void runOnOperation() override {
+    // Step 1: Folds subviews into their load/store users.
+    RewritePatternSet patterns(&getContext());
+    patterns.add<LoadOpOfSubViewFolder, StoreOpOfSubViewFolder>(&getContext());
+    if (failed(applyPatternsGreedily(getOperation(), std::move(patterns)))) {
+      signalPassFailure();
+    }
+
+    // Step 2: Cleans up dead subview operations that have no remaining users.
+    SmallVector<memref::SubViewOp> dead_sub_views;
+    getOperation().walk([&](memref::SubViewOp sub_view_op) {
+      if (sub_view_op->use_empty()) {
+        dead_sub_views.push_back(sub_view_op);
+      }
+    });
+
+    for (auto sub_view_op : dead_sub_views) {
+      sub_view_op.erase();
+    }
+  }
+};
+} // namespace
+
+std::unique_ptr<Pass> mlir::createFoldSubViewPass() {
+  return std::make_unique<FoldSubViewPass>();
+}