Resource aware task optimization

include/NeuraDialect/Architecture/Architecture.h

@@ -494,6 +494,23 @@ class Architecture {
   // Checks if the architecture supports counter operations.
   bool canSupportCounter() const;
 
+  /// Clones the architecture but with new per-cgra dimensions.
+  /// The provided tile_overrides will be appended to the existing ones.
+  ///
+  /// Example — create an 8×4 tile array (2×1 CGRA rectangle) with all tiles
+  /// present:
+  ///   auto arch_2x1 = getArchitecture().cloneWithNewDimensions(8, 4);
+  ///
+  /// Example — create a 12×8 bounding box for a T-shape (4 CGRAs) where only
+  /// specific tiles are valid:
+  ///   std::vector<TileOverride> overrides;
+  ///   // First mark all tiles as non-existent, then mark valid ones existent.
+  ///   // (see MapToAcceleratorPass for the full valid_tiles parsing logic)
+  ///   auto arch_T = getArchitecture().cloneWithNewDimensions(8, 12, overrides);
+  std::unique_ptr<Architecture> cloneWithNewDimensions(
+      int new_per_cgra_rows, int new_per_cgra_columns,
+      const std::vector<TileOverride> &additional_overrides = {}) const;
+
 private:
   // Helper methods for constructor initialization.
   void initializeTiles(int rows, int columns);
@@ -532,6 +549,13 @@ class Architecture {
   int per_cgra_rows_;
   int per_cgra_columns_;
   int max_ctrl_mem_items_;
+
+  BaseTopology multi_cgra_base_topology_;
+  BaseTopology per_cgra_base_topology_;
+  TileDefaults tile_defaults_;
+  std::vector<TileOverride> tile_overrides_;
+  LinkDefaults link_defaults_;
+  std::vector<LinkOverride> link_overrides_;
 };
 
 // Function for getting the architecture object.

include/NeuraDialect/NeuraPasses.h

@@ -23,7 +23,10 @@ std::unique_ptr<mlir::Pass> createInsertCtrlMovPass();
 std::unique_ptr<mlir::Pass> createAssignAcceleratorPass();
 std::unique_ptr<mlir::Pass> createTransformCtrlToDataFlowPass();
 std::unique_ptr<mlir::Pass> createLeveragePredicatedValuePass();
-std::unique_ptr<mlir::Pass> createMapToAcceleratorPass();
+/// Creates the MapToAccelerator pass. Tile dimensions default to 0 (use
+/// architecture singleton) when not specified via options.
+std::unique_ptr<mlir::Pass> createMapToAcceleratorPass(
+    const MapToAcceleratorOptions &options = MapToAcceleratorOptions{});
 std::unique_ptr<mlir::Pass> createGenerateCodePass();
 std::unique_ptr<mlir::Pass> createCanonicalizeReturnPass();
 std::unique_ptr<mlir::Pass> createCanonicalizeLiveInPass();

include/NeuraDialect/NeuraPasses.td

@@ -54,7 +54,41 @@ def MapToAccelerator : Pass<"map-to-accelerator", "ModuleOp"> {
   let summary = "Map Neura operations onto a given accelerator";
   let description = [{
     This pass performs mapping from Neura operations to accelerator.
+
+    x-tiles and y-tiles specify the **tile** dimensions of the target array
+    (not the CGRA count).  Each CGRA contains a per_cgra_rows × per_cgra_cols
+    tile grid (currently 4×4).  So for a single CGRA, x-tiles=4 y-tiles=4;
+    for a 1×2 rectangular pair, x-tiles=8 y-tiles=4; etc.
+
+    When x-tiles=0 and y-tiles=0 (the default), the global Architecture
+    singleton determines the tile grid — this is equivalent to a single CGRA.
+
+    Examples:
+      Single CGRA (default):
+        --map-to-accelerator
+      1×3 rectangular (3 CGRAs in a row):
+        --map-to-accelerator x-tiles=12 y-tiles=4
+      T-shape (4 CGRAs: top row of 3 + centre below):
+        --map-to-accelerator x-tiles=12 y-tiles=8 \
+            valid-tiles="0_0,1_0,2_0,3_0,4_0,5_0,6_0,7_0,8_0,9_0,10_0,11_0,\
+                         4_1,5_1,6_1,7_1,4_4,5_4,6_4,7_4,4_5,5_5,6_5,7_5"
   }];
+  let options = [
+    Option<"x_tiles", "x-tiles", "int", /*default=*/"0",
+           "Total number of tiles in the X dimension of the target array "
+           "(not the number of CGRAs).  Each CGRA contributes per_cgra_cols "
+           "tiles.  0 means use the global Architecture singleton (1 CGRA).">,
+    Option<"y_tiles", "y-tiles", "int", /*default=*/"0",
+           "Total number of tiles in the Y dimension of the target array "
+           "(not the number of CGRAs).  Each CGRA contributes per_cgra_rows "
+           "tiles.  0 means use the global Architecture singleton (1 CGRA).">,
+    Option<"valid_tiles", "valid-tiles", "std::string", /*default=*/"\"\"",
+           "Comma-separated list of tile coordinates (x_y) that are actually "
+           "present in the array, used for non-rectangular CGRA shapes such as "
+           "L-blocks or T-blocks. Empty string means all tiles in the "
+           "x-tiles x y-tiles rectangle are valid. "
+           "Example: 0_0,1_0,0_1 selects three tiles forming an L-shape.">
+  ];
   let constructor = "neura::createMapToAcceleratorPass()";
 }
 

include/TaskflowDialect/TaskflowPasses.h

@@ -29,6 +29,7 @@ std::unique_ptr<mlir::Pass> createMapTaskOnCgraPass();
 std::unique_ptr<mlir::Pass> createAffineLoopTreeSerializationPass();
 std::unique_ptr<mlir::Pass> createAffineLoopPerfectionPass();
 std::unique_ptr<mlir::Pass> createMemoryAccessStreamingFusionPass();
+std::unique_ptr<mlir::Pass> createResourceAwareTaskOptimizationPass();
 
 #define GEN_PASS_REGISTRATION
 #include "TaskflowDialect/TaskflowPasses.h.inc"

include/TaskflowDialect/TaskflowPasses.td

@@ -86,4 +86,32 @@ def MemoryAccessStreamingFusion : Pass<"memory-access-streaming-fusion", "func::
   }];
   let constructor = "taskflow::createMemoryAccessStreamingFusionPass()";
 }
+
+def ResourceAwareTaskOptimization : Pass<"resource-aware-task-optimization", "func::FuncOp"> {
+  let summary = "Balances pipeline latency and fuses independent tasks for CGRA utilization.";
+  let description = [{
+    Two-phase optimization:
+    1. Utilization Fusion: merges independent (no-edge) tasks, selecting pairs
+       that minimize |trip_count_a - trip_count_b| for balanced utilization.
+    2. Pipeline Balance: allocates extra CGRAs to critical-path bottleneck tasks.
+       More CGRAs combine tile arrays into larger arrays for mapping, potentially
+       lowering compiled_ii.  Latency model: II * (trip_count - 1) + steps.
+    Targets a 4x4 CGRA grid (16 CGRAs total, one CGRA per cell).
+    Currently a single task may be allocated at most 4 CGRAs.
+    Supported CGRA array shapes for a task (all fit within the 4×4 grid):
+      - rect : a perfect rectangle, e.g. 1×1, 1×2, 2×1, 1×3, 3×1, 2×2, 1×4, 4×1.
+      - L    : an L-shaped block of 3 or 4 CGRAs, e.g.
+                 3 CGRAs: (0,0)(1,0)(0,1) — two in a row + one below-left.
+                 4 CGRAs: (0,0)(0,1)(0,2)(1,2) — three in a column + one offset.
+      - T    : a T-shaped block of 4 CGRAs, e.g.
+                 (0,0)(1,0)(2,0)(1,1) — three in a row + one below centre.
+    Non-rectangular shapes are represented by their bounding box plus an
+    explicit tile list that enumerates only the occupied CGRA positions.
+    Compiled_ii must come from the downstream Neura pipeline (asserts on failure).
+  }];
+  let constructor = "taskflow::createResourceAwareTaskOptimizationPass()";
+  let dependentDialects = [
+    "mlir::affine::AffineDialect",
+    "mlir::func::FuncDialect"];
+}
 #endif // TASKFLOW_PASSES_TD

lib/NeuraDialect/Architecture/Architecture.cpp

@@ -561,12 +561,15 @@ Architecture::Architecture(int multi_cgra_rows, int multi_cgra_columns,
                            const std::vector<LinkOverride> &link_overrides) {
   this->multi_cgra_rows_ = multi_cgra_rows;
   this->multi_cgra_columns_ = multi_cgra_columns;
-  // TODO: Support multi-CGRA topology in the future:
-  // https://github.com/coredac/dataflow/issues/163.
-  // this->multi_cgra_base_topology_ = multi_cgra_base_topology;
+  this->multi_cgra_base_topology_ = multi_cgra_base_topology;
   this->per_cgra_rows_ = per_cgra_rows;
   this->per_cgra_columns_ = per_cgra_columns;
+  this->per_cgra_base_topology_ = per_cgra_base_topology;
   this->max_ctrl_mem_items_ = max_ctrl_mem_items;
+  this->tile_defaults_ = tile_defaults;
+  this->tile_overrides_ = tile_overrides;
+  this->link_defaults_ = link_defaults;
+  this->link_overrides_ = link_overrides;
 
   // Initializes architecture components using helper methods.
   initializeTiles(per_cgra_rows, per_cgra_columns);
@@ -576,6 +579,20 @@ Architecture::Architecture(int multi_cgra_rows, int multi_cgra_columns,
   applyLinkOverrides(link_overrides);
 }
 
+std::unique_ptr<Architecture> Architecture::cloneWithNewDimensions(
+    int new_per_cgra_rows, int new_per_cgra_columns,
+    const std::vector<TileOverride> &additional_overrides) const {
+
+  std::vector<TileOverride> merged_overrides = tile_overrides_;
+  merged_overrides.insert(merged_overrides.end(), additional_overrides.begin(), additional_overrides.end());
+
+  return std::make_unique<Architecture>(
+      multi_cgra_rows_, multi_cgra_columns_, multi_cgra_base_topology_,
+      new_per_cgra_rows, new_per_cgra_columns, max_ctrl_mem_items_,
+      per_cgra_base_topology_, tile_defaults_, merged_overrides,
+      link_defaults_, link_overrides_);
+}
+
 Tile *Architecture::getTile(int id) {
   auto it = id_to_tile_.find(id);
   assert(it != id_to_tile_.end() && "Tile with given ID not found");

lib/NeuraDialect/Transforms/InsertDataMovPass.cpp

@@ -23,8 +23,14 @@ struct InsertDataMovForNeuraOps : public RewritePattern {
 
   LogicalResult matchAndRewrite(Operation *op,
                                 PatternRewriter &rewriter) const override {
+    // Only processes operations from the neura dialect. Operations from
+    // other dialects (arith, math, etc.) should have been lowered to neura
+    // ops by earlier passes (LowerArithToNeura, etc.) before this pass runs.
     if (op->getDialect()->getNamespace() != accel::kNeuraTarget ||
-        isa<neura::DataMovOp>(op)) {
+        isa<neura::DataMovOp>(op) ||
+        isa<neura::ReserveOp>(op) ||
+        isa<neura::KernelOp>(op) ||
+        isa<neura::FusedOp>(op)) {
       return failure();
     }
 

lib/NeuraDialect/Transforms/MapToAcceleratorPass.cpp

@@ -46,6 +46,11 @@ struct MapToAcceleratorPass
   }
 
   MapToAcceleratorPass() = default;
+  MapToAcceleratorPass(const MapToAcceleratorOptions &options) : MapToAcceleratorPass() {
+    this->x_tiles = options.x_tiles;
+    this->y_tiles = options.y_tiles;
+    this->valid_tiles = options.valid_tiles;
+  }
   MapToAcceleratorPass(const MapToAcceleratorPass &pass)
       : PassWrapper<MapToAcceleratorPass, OperationPass<ModuleOp>>(pass) {}
   Option<std::string> mappingStrategy{
@@ -72,6 +77,18 @@ struct MapToAcceleratorPass
       llvm::cl::desc(
           "Dump the resource allocation table after mapping (default: true)"),
       llvm::cl::init(true)};
+  Option<int> x_tiles{
+      *this, "x-tiles",
+      llvm::cl::desc("Override number of tiles in X dimension (0 = default)."),
+      llvm::cl::init(0)};
+  Option<int> y_tiles{
+      *this, "y-tiles",
+      llvm::cl::desc("Override number of tiles in Y dimension (0 = default)."),
+      llvm::cl::init(0)};
+  Option<std::string> valid_tiles{
+      *this, "valid-tiles",
+      llvm::cl::desc("Comma separated list of valid tile coords x_y,x_y to support non-rectangular shapes."),
+      llvm::cl::init("")};
 
   // Configures mapping strategy and mode based on command-line options.
   bool configureMappingStrategy(StringRef mapping_strategy_opt,
@@ -239,15 +256,15 @@ struct MapToAcceleratorPass
     }
 
     // Filters out operations inside fused_op regions.
-    // Only map the fused_op itself, not the operations within its region
+    // Only maps the fused_op itself, not the operations within its region.
     std::vector<Operation *> filtered_ops;
     int skipped_count = 0;
     for (Operation *op : topologically_sorted_ops) {
       Operation *parent_op = op->getParentOp();
-      // Check if parent is a fused_op by checking operation name
+      // Checks if the parent is a fused_op by inspecting the operation name.
       if (parent_op &&
           parent_op->getName().getStringRef().contains(attr::val::kOpFused)) {
-        // Skip operations inside fused_op region
+        // Skips operations inside a fused_op region.
         llvm::outs() << "[MapToAcceleratorPass] Skipping op inside fused_op: "
                      << *op << "\n";
         skipped_count++;
@@ -290,9 +307,9 @@ struct MapToAcceleratorPass
       MappingState mapping_state(architecture, ii, is_spatial_only);
       if (mapping_strategy->map(sorted_ops_with_alap_levels, critical_ops,
                                 architecture, mapping_state)) {
-        // success
+        // Success.
         if (dumpMappingTable) {
-          // logs to stderr
+          // Logs to stderr.
           mapping_state.dumpOpToLocs();
         }
         mapping_state.encodeMappingState();
@@ -355,7 +372,55 @@ struct MapToAcceleratorPass
       return;
     }
 
-    const Architecture &architecture = mlir::neura::getArchitecture();
+    const Architecture &global_arch = mlir::neura::getArchitecture();
+    std::unique_ptr<Architecture> custom_arch;
+    const Architecture *target_arch = &global_arch;
+
+    if (x_tiles.getValue() > 0 && y_tiles.getValue() > 0) {
+      std::vector<TileOverride> additional_overrides;
+      if (!valid_tiles.getValue().empty()) {
+        llvm::SmallVector<llvm::StringRef, 4> coords;
+        llvm::StringRef(valid_tiles.getValue()).split(coords, ',');
+
+        // Default: mark all tiles as non-existent first if valid_tiles provided.
+        for (int y = 0; y < y_tiles.getValue(); ++y) {
+          for (int x = 0; x < x_tiles.getValue(); ++x) {
+            TileOverride to;
+            to.tile_x = x;
+            to.tile_y = y;
+            to.existence = false;
+            additional_overrides.push_back(to);
+          }
+        }
+
+        // Then mark the valid ones as existent.
+        for (llvm::StringRef coord : coords) {
+          auto pair = coord.split('_');
+          int x, y;
+          if (!pair.first.getAsInteger(10, x) && !pair.second.getAsInteger(10, y)) {
+            TileOverride to;
+            to.tile_x = x;
+            to.tile_y = y;
+            to.existence = true;
+            additional_overrides.push_back(to);
+          }
+        }
+      }
+
+      // Builds a custom architecture with the requested tile dimensions.
+      // For non-rectangular shapes, tiles marked existence=false are removed
+      // by applyTileOverrides (which calls removeTile) BEFORE createLinks
+      // runs.  createMeshLinks/createLinkIfValid then check coord_to_tile_
+      // and skip absent tiles, so no links are ever created to/from
+      // non-existent tiles — boundary connectivity is handled correctly.
+      custom_arch = global_arch.cloneWithNewDimensions(
+        y_tiles.getValue(), x_tiles.getValue(), additional_overrides);
+      target_arch = custom_arch.get();
+      llvm::errs() << "[MapToAcceleratorPass] Overriding architecture dimensions to "
+                   << y_tiles.getValue() << "x" << x_tiles.getValue() << " tiles.\n";
+    }
+
+    const Architecture &architecture = *target_arch;
 
     // Maps kernels.
     module.walk([&](neura::KernelOp kernel_op) {
@@ -398,8 +463,9 @@ struct MapToAcceleratorPass
 
 namespace mlir::neura {
 
-std::unique_ptr<Pass> createMapToAcceleratorPass() {
-  return std::make_unique<MapToAcceleratorPass>();
+std::unique_ptr<Pass> createMapToAcceleratorPass(
+    const MapToAcceleratorOptions &options) {
+  return std::make_unique<MapToAcceleratorPass>(options);
 }
 
 } // namespace mlir::neura

lib/TaskflowDialect/Transforms/Optimizations/CMakeLists.txt

@@ -4,17 +4,23 @@ add_mlir_conversion_library(MLIRTaskflowOptimization
   AffineLoopTreeSerializationPass.cpp
   AffineLoopPerfectionPass.cpp
   MemoryAccessStreamingFusion.cpp
+  ResourceAwareTaskOptimizationPass.cpp
 
   DEPENDS
   MLIRTaskflowTransformsIncGen
 
   LINK_LIBS PUBLIC
   MLIRTaskflow
+  MLIRTaskflowTransforms
+  MLIRAffineDialect
   MLIRArithDialect
   MLIRFuncDialect
   MLIRLinalgDialect
   MLIRIR
   MLIRPass
   MLIRTransforms
+  MLIRNeura
+  MLIRNeuraTransforms
+  MLIRConversion
   MLIRSupport
 )