Merge pull request coredac#298 from coredac/fix/assertion-and-performance

guosran · web-flow · commit b90a129933bb · 2026-03-22T23:09:54.000-07:00
Fix Assertion Crash and Optimize Compile Time
diff --git a/lib/NeuraDialect/Architecture/Architecture.cpp b/lib/NeuraDialect/Architecture/Architecture.cpp
@@ -314,7 +314,15 @@ void Architecture::applyTileOverrides(
   for (const auto &override : tile_overrides) {
     Tile *tile = nullptr;
     if (override.tile_x >= 0 && override.tile_y >= 0) {
-      tile = getTile(override.tile_x, override.tile_y);
+      // cloneWithNewDimensions() copies tile_overrides from the original
+      // (larger) architecture and replays them on a smaller cloned grid.
+      // An override may reference coordinates that exist in the original but
+      // fall outside the cloned grid's bounds.  Skips such overrides rather
+      // than asserting, so the clone only applies overrides that are relevant
+      // to its own coordinate space.
+      auto it = coord_to_tile_.find({override.tile_x, override.tile_y});
+      if (it != coord_to_tile_.end())
+        tile = it->second;
     }
 
     if (tile) {
@@ -518,8 +526,19 @@ void Architecture::applyLinkOverrides(
     }
     // Handles link removal.
     else {
-      Tile *src_tile = getTile(override.src_tile_x, override.src_tile_y);
-      Tile *dst_tile = getTile(override.dst_tile_x, override.dst_tile_y);
+      // cloneWithNewDimensions() replays link_overrides from the original
+      // architecture on a smaller cloned grid.  Either endpoint of a link
+      // override may reference a tile that was not included in the clone.
+      // Skip the override in that case to avoid a crash; the link simply
+      // does not exist in the cloned architecture and needs no override.
+      auto src_it = coord_to_tile_.find(
+          {override.src_tile_x, override.src_tile_y});
+      auto dst_it = coord_to_tile_.find(
+          {override.dst_tile_x, override.dst_tile_y});
+      if (src_it == coord_to_tile_.end() || dst_it == coord_to_tile_.end())
+        continue; // One or both tiles do not exist in this architecture.
+      Tile *src_tile = src_it->second;
+      Tile *dst_tile = dst_it->second;
 
       if (src_tile && dst_tile) {
         bool link_already_exists = linkExists(src_tile, dst_tile);
@@ -653,11 +672,12 @@ void Architecture::removeTile(int tile_id) {
 
 Link *Architecture::getLink(int src_tile_x, int src_tile_y, int dst_tile_x,
                             int dst_tile_y) {
-  Tile *src_tile = getTile(src_tile_x, src_tile_y);
-  Tile *dst_tile = getTile(dst_tile_x, dst_tile_y);
-  if (!src_tile || !dst_tile) {
+  auto src_it = coord_to_tile_.find({src_tile_x, src_tile_y});
+  auto dst_it = coord_to_tile_.find({dst_tile_x, dst_tile_y});
+  if (src_it == coord_to_tile_.end() || dst_it == coord_to_tile_.end())
     return nullptr; // One of the tiles does not exist.
-  }
+  Tile *src_tile = src_it->second;
+  Tile *dst_tile = dst_it->second;
 
   for (const auto &[id, link] : link_storage_) {
     if (link && link->getSrcTile() == src_tile &&
@@ -718,12 +738,11 @@ void Architecture::removeLink(Tile *src_tile, Tile *dst_tile) {
 
 void Architecture::removeLink(int src_tile_x, int src_tile_y, int dst_tile_x,
                               int dst_tile_y) {
-  Tile *src_tile = getTile(src_tile_x, src_tile_y);
-  Tile *dst_tile = getTile(dst_tile_x, dst_tile_y);
-  if (!src_tile || !dst_tile) {
+  auto src_it = coord_to_tile_.find({src_tile_x, src_tile_y});
+  auto dst_it = coord_to_tile_.find({dst_tile_x, dst_tile_y});
+  if (src_it == coord_to_tile_.end() || dst_it == coord_to_tile_.end())
     return; // One of the tiles does not exist.
-  }
-  removeLink(src_tile, dst_tile);
+  removeLink(src_it->second, dst_it->second);
 }
 
 bool Architecture::canSupportCounter() const {
diff --git a/lib/TaskflowDialect/Transforms/Optimizations/ResourceAwareTaskOptimizationPass.cpp b/lib/TaskflowDialect/Transforms/Optimizations/ResourceAwareTaskOptimizationPass.cpp
@@ -347,30 +347,12 @@ class TaskDependencyGraph {
     return total;
   }
 
-  // Public wrapper for profileTask: used by UtilizationFuser to re-profile
-  // fused tasks with the real downstream Neura pipeline.
-  // When skip_mapper=true, only ResMII/RecMII analytical estimates are used
-  // (no MapToAcceleratorPass). This is safe for speculative balance checks
-  // where the mapper may backtrack indefinitely on larger tile arrays.
-  void profileTaskPublic(TaskGraphNode *node, TaskflowTaskOp task,
-                         bool skip_mapper = false) {
-    profileTask(node, task, skip_mapper);
-  }
-
-private:
-  llvm::DenseSet<std::pair<TaskGraphNode *, TaskGraphNode *>> edge_set;
-
-  void addEdge(TaskGraphNode *from, TaskGraphNode *to) {
-    auto key = std::make_pair(from, to);
-    if (edge_set.insert(key).second) {
-      from->successors.push_back(to);
-      to->predecessors.push_back(from);
-    }
-  }
-
   // Profiles a single TaskflowTaskOp: clones the task, wraps the kernel in a
   // standalone func, and runs InsertDataMov + MapToAcceleratorPass to obtain
   // ii.  skip_mapper: use only ResMII/RecMII analytical estimates.
+  // When skip_mapper=true, only ResMII/RecMII analytical estimates are used
+  // (no MapToAcceleratorPass). This is safe for speculative balance checks
+  // where the mapper may backtrack indefinitely on larger tile arrays.
   void profileTask(TaskGraphNode *node, TaskflowTaskOp task,
                    bool skip_mapper = false) {
     MLIRContext *ctx = task.getContext();
@@ -456,6 +438,17 @@ class TaskDependencyGraph {
     tmp_mod.erase();
   }
 
+private:
+  llvm::DenseSet<std::pair<TaskGraphNode *, TaskGraphNode *>> edge_set;
+
+  void addEdge(TaskGraphNode *from, TaskGraphNode *to) {
+    auto key = std::make_pair(from, to);
+    if (edge_set.insert(key).second) {
+      from->successors.push_back(to);
+      to->predecessors.push_back(from);
+    }
+  }
+
   // Wraps a neura.kernel into a standalone func in dst_module, runs
   // InsertDataMov + mapper, and returns compiled_ii / cp_depth.
   // x_tiles/y_tiles: multi-CGRA tile grid dimensions.
@@ -1725,7 +1718,7 @@ struct ResourceAwareTaskOptimizationPass
       // mapper is skipped entirely (only ResMII/RecMII estimates are used).
       auto profile_fn = [&graph, use_analytical](TaskGraphNode *node,
                                                  TaskflowTaskOp task) {
-        graph.profileTaskPublic(node, task, /*skip_mapper=*/use_analytical);
+        graph.profileTask(node, task, /*skip_mapper=*/use_analytical);
       };
       bool fuse_changed = fuser.fuse(func, graph, profile_fn);
 
@@ -1739,11 +1732,13 @@ struct ResourceAwareTaskOptimizationPass
       }
 
       // Phase 2: Latency-Aware Pipeline Balance.
-      // Balance probes use analytical-only profiling by default.
-      bool balance_skip = use_analytical || balanceSkipMapper.getValue();
-      auto balance_profile_fn = [&graph, balance_skip](TaskGraphNode *node,
+      // Balance probes always use analytical-only profiling (skip_mapper=true)
+      // to avoid exponential backtracking blowup during speculative probing.
+      // The balance-skip-mapper flag now only controls whether a final
+      // verification mapper run is performed after convergence (see below).
+      auto balance_profile_fn = [&graph](TaskGraphNode *node,
                                                        TaskflowTaskOp task) {
-        graph.profileTaskPublic(node, task, /*skip_mapper=*/balance_skip);
+        graph.profileTask(node, task, /*skip_mapper=*/true);
       };
       PipelineBalancer balancer;
       bool balance_changed = balancer.balance(graph, balance_profile_fn);
@@ -1777,19 +1772,22 @@ struct ResourceAwareTaskOptimizationPass
                    << graph.getTotalAllocatedCGRAs() << "\n";
 
       if (!balance_changed && !fuse_changed) {
-        // Converged — writes ALL attributes (cgra_count, ii, steps) to IR
-        // for every task. Non-fused tasks only got cgra_count written during
-        // intermediate iterations; ii, steps, and trip_count live only in the
-        // graph node and must be persisted here.
-        //
-        // Note: no re-profiling is done here.  When balance-skip-mapper=true
-        // (the default), the balance phase uses analytical estimates; those
-        // are the values written to the final IR.  When
-        // balance-skip-mapper=false, the balance phase already ran the real
-        // mapper for each speculative probe, so the graph already contains
-        // accurate compiled_ii / steps values.  Either way, the converged
-        // graph state is authoritative and written directly to IR.
+        // Converged — optionally re-profile with the real mapper for accuracy.
+        // When balance-skip-mapper=false, runs the mapper once per task that
+        // had its cgra_count changed, to get authoritative compiled_ii/steps.
+        bool balance_skip_mapper = use_analytical || balanceSkipMapper.getValue();
+        if (!balance_skip_mapper) {
+          llvm::errs() << "[ResourceAware] Running final mapper verification "
+                       << "for converged allocation...\n";
+          for (auto &node : graph.nodes) {
+            // Re-profile with the real mapper to get accurate II/steps.
+            node->shape = pickBestShape(node->cgra_count);
+            graph.profileTask(node.get(), node->op,
+                                    /*skip_mapper=*/false);
+          }
+        }
 
+        // Writes ALL attributes (cgra_count, ii, steps) to IR for every task.
         for (auto &node : graph.nodes) {
           OpBuilder b(node->op);
           node->shape = pickBestShape(node->cgra_count);
diff --git a/test/multi-cgra/taskflow/irregular-loop/irregular-loop.mlir b/test/multi-cgra/taskflow/irregular-loop/irregular-loop.mlir
@@ -382,19 +382,7 @@ module attributes {} {
 // PLACEMENT:      taskflow.task @Task_2
 // PLACEMENT-SAME: task_mapping_info = {cgra_positions = [{col = 0 : i32, row = 1 : i32}], read_sram_locations = [{col = 1 : i32, row = 1 : i32}], write_sram_locations = [{col = 0 : i32, row = 1 : i32}]}
 
-// CGRA Tile Occupation after RESOPT (4x4 grid, col x row):
-// +---+---+---+---+
-// | 0 | 1 | . | . |   Task_0_Task_1_utilfused (1x1, cgra_count=1)
-// +---+---+---+---+   Task_2 (1x1, cgra_count=1)
-// | . | . | . | . |
-// +---+---+---+---+
-// | . | . | . | . |
-// +---+---+---+---+
-// | . | . | . | . |
-// +---+---+---+---+
-// 0=Task_0_Task_1_utilfused, 1=Task_2; 2/16 CGRAs used
-
 // RESOPT:      taskflow.task @Task_0_Task_1_utilfused
-// RESOPT-SAME: {cgra_count = 1 : i32, compiled_ii = 3 : i32, steps = 5 : i32, tile_shape = "1x1", trip_count = 32 : i32}
+// RESOPT-SAME: {cgra_count = 2 : i32, compiled_ii = 3 : i32, steps = 5 : i32, tile_shape = "1x2", trip_count = 32 : i32}
 // RESOPT:      taskflow.task @Task_2
-// RESOPT-SAME: {cgra_count = 1 : i32, compiled_ii = 2 : i32, steps = 7 : i32, tile_shape = "1x1", trip_count = 32 : i32}
+// RESOPT-SAME: {cgra_count = 2 : i32, compiled_ii = 2 : i32, steps = 7 : i32, tile_shape = "1x2", trip_count = 32 : i32}
diff --git a/test/multi-cgra/taskflow/multi-nested/multi-nested.mlir b/test/multi-cgra/taskflow/multi-nested/multi-nested.mlir
@@ -534,21 +534,10 @@ module attributes {} {
 // PLACEMENT-SAME: task_mapping_info = {cgra_positions = [{col = 1 : i32, row = 1 : i32}], read_sram_locations = [{col = 1 : i32, row = 1 : i32}, {col = 2 : i32, row = 1 : i32}], write_sram_locations = [{col = 1 : i32, row = 1 : i32}]}
 
 // RESOPT:      taskflow.task @Task_1
-// RESOPT:      cgra_count = 1 : i32, compiled_ii = 2 : i32, steps = 4 : i32, tile_shape = "1x1", trip_count = 160 : i32
+// RESOPT:      cgra_count = 2 : i32, compiled_ii = 2 : i32, steps = 4 : i32, tile_shape = "1x2", trip_count = 160 : i32
 // RESOPT:      taskflow.task @Task_0_Task_2_fused_Task_3_utilfused
-// RESOPT:      cgra_count = 1 : i32, compiled_ii = 2 : i32, steps = 5 : i32, tile_shape = "1x1", trip_count = 192 : i32
+// RESOPT:      cgra_count = 2 : i32, compiled_ii = 2 : i32, steps = 5 : i32, tile_shape = "1x2", trip_count = 192 : i32
 // RESOPT:      taskflow.task @Task_4
-// RESOPT:      cgra_count = 1 : i32, compiled_ii = 2 : i32, steps = 4 : i32, tile_shape = "1x1", trip_count = 36 : i32
+// RESOPT:      cgra_count = 2 : i32, compiled_ii = 2 : i32, steps = 4 : i32, tile_shape = "1x2", trip_count = 36 : i32
 // RESOPT:      return
 
-// CGRA Tile Occupation after RESOPT (4x4 grid, col x row):
-// +---+---+---+---+
-// | 0 | 1 | 2 | . |   Task_1 (1x1, cgra_count=1)
-// +---+---+---+---+   Task_0_Task_2_fused_Task_3_utilfused (1x1, cgra_count=1)
-// | . | . | . | . |   Task_4 (1x1, cgra_count=1)
-// +---+---+---+---+
-// | . | . | . | . |
-// +---+---+---+---+
-// | . | . | . | . |
-// +---+---+---+---+
-// 0=Task_1, 1=Task_0_Task_2_fused_Task_3_utilfused, 2=Task_4; 3/16 CGRAs used
diff --git a/test/multi-cgra/taskflow/parallel-nested/parallel-nested.mlir b/test/multi-cgra/taskflow/parallel-nested/parallel-nested.mlir
@@ -161,17 +161,6 @@ module {
 // PLACEMENT-SAME: task_mapping_info = {cgra_positions = [{col = 1 : i32, row = 0 : i32}], read_sram_locations = [{col = 1 : i32, row = 0 : i32}, {col = 1 : i32, row = 0 : i32}], write_sram_locations = [{col = 1 : i32, row = 0 : i32}]}
 
 // RESOPT:      taskflow.task @Task_0_Task_1_utilfused
-// RESOPT:      cgra_count = 1 : i32, compiled_ii = 2 : i32, steps = 4 : i32, tile_shape = "1x1", trip_count = 64 : i32
+// RESOPT:      cgra_count = 2 : i32, compiled_ii = 2 : i32, steps = 4 : i32, tile_shape = "1x2", trip_count = 64 : i32
 // RESOPT:      return
 
-// CGRA Tile Occupation after RESOPT (4x4 grid, col x row):
-// +---+---+---+---+
-// | 0 | . | . | . |   row=0: Task_0_Task_1_utilfused (1x1, cgra_count=1)
-// +---+---+---+---+
-// | . | . | . | . |
-// +---+---+---+---+
-// | . | . | . | . |
-// +---+---+---+---+
-// | . | . | . | . |
-// +---+---+---+---+
-// 0=Task_0_Task_1_utilfused; 1/16 CGRAs used
diff --git a/test/multi-cgra/taskflow/resnet/simple_resnet_tosa.mlir b/test/multi-cgra/taskflow/resnet/simple_resnet_tosa.mlir
@@ -35,11 +35,11 @@
 // RUN: --canonicalize-live-in \
 // RUN: --leverage-predicated-value \
 // RUN: --transform-ctrl-to-data-flow \
-// RUN: --fold-constant 
-// RU: '--resource-aware-task-optimization=balance-skip-mapper=false' \
-// RU: --architecture-spec=%S/../../../arch_spec/architecture_with_counter.yaml \
-// RU: -o %t.resopt.mlir
-// RU: FileCheck %s --input-file=%t.resopt.mlir --check-prefixes=RESOPT
+// RUN: --fold-constant \
+// RUN: '--resource-aware-task-optimization=balance-skip-mapper=false' \
+// RUN: --architecture-spec=%S/../../../arch_spec/architecture_with_counter.yaml \
+// RUN: -o %t.resopt.mlir
+// RUN: FileCheck %s --input-file=%t.resopt.mlir --check-prefixes=RESOPT
 
 
 module attributes {torch.debug_module_name = "SimpleResNetBlock"} {
@@ -706,30 +706,15 @@ module attributes {torch.debug_module_name = "SimpleResNetBlock"} {
 
 
 // RESOPT:      taskflow.task @Task_1_Task_0_Task_2_utilfused_utilfused
-// RESOPT-SAME: {cgra_count = 1 : i32, compiled_ii = 4 : i32, steps = 3 : i32, tile_shape = "1x1", trip_count = 6400 : i32}
+// RESOPT-SAME: {cgra_count = 2 : i32, compiled_ii = 5 : i32, steps = 3 : i32, tile_shape = "1x2", trip_count = 6400 : i32}
 // RESOPT:      taskflow.task @Task_3
-// RESOPT-SAME: {cgra_count = 1 : i32, compiled_ii = 2 : i32, steps = 6 : i32, tile_shape = "1x1", trip_count = 2359296 : i32}
+// RESOPT-SAME: {cgra_count = 2 : i32, compiled_ii = 3 : i32, steps = 6 : i32, tile_shape = "1x2", trip_count = 2359296 : i32}
 // RESOPT:      taskflow.task @Task_4_Task_5_fused_Task_7_utilfused
-// RESOPT-SAME: {cgra_count = 1 : i32, compiled_ii = 2 : i32, steps = 7 : i32, tile_shape = "1x1", trip_count = 6400 : i32}
+// RESOPT-SAME: {cgra_count = 2 : i32, compiled_ii = 4 : i32, steps = 7 : i32, tile_shape = "1x2", trip_count = 6400 : i32}
 // RESOPT:      taskflow.task @Task_6_Task_8_utilfused
-// RESOPT-SAME: {cgra_count = 1 : i32, compiled_ii = 2 : i32, steps = 3 : i32, tile_shape = "1x1", trip_count = 4096 : i32}
+// RESOPT-SAME: {cgra_count = 2 : i32, compiled_ii = 3 : i32, steps = 3 : i32, tile_shape = "1x2", trip_count = 4096 : i32}
 // RESOPT:      taskflow.task @Task_9
-// RESOPT-SAME: {cgra_count = 1 : i32, compiled_ii = 2 : i32, steps = 6 : i32, tile_shape = "1x1", trip_count = 2359296 : i32}
+// RESOPT-SAME: {cgra_count = 2 : i32, compiled_ii = 3 : i32, steps = 6 : i32, tile_shape = "1x2", trip_count = 2359296 : i32}
 // RESOPT:      taskflow.task @Task_10_Task_11_Task_12_fused_fused
-// RESOPT-SAME: {cgra_count = 1 : i32, compiled_ii = 2 : i32, steps = 8 : i32, tile_shape = "1x1", trip_count = 4096 : i32}
+// RESOPT-SAME: {cgra_count = 2 : i32, compiled_ii = 7 : i32, steps = 8 : i32, tile_shape = "1x2", trip_count = 4096 : i32}
 // RESOPT:      return
-
-
-// CGRA Tile Occupation after RESOPT (4x4 grid, col x row):
-// +---+---+---+---+
-// | 0 | 1 | 2 | 3 |   row=0: 0=Task_1_..._utilfused, 1=Task_3, 2=Task_4_..._utilfused, 3=Task_6_Task_8_utilfused
-// +---+---+---+---+
-// | 4 | 5 | . | . |   row=1: 4=Task_9, 5=Task_10_..._fused_fused
-// +---+---+---+---+
-// | . | . | . | . |
-// +---+---+---+---+
-// | . | . | . | . |
-// +---+---+---+---+
-// 0=Task_1_Task_0_Task_2_utilfused_utilfused, 1=Task_3, 2=Task_4_Task_5_fused_Task_7_utilfused
-// 3=Task_6_Task_8_utilfused, 4=Task_9, 5=Task_10_Task_11_Task_12_fused_fused
-// 6/16 CGRAs used
