Resource aware task optimization

[guosran]

Overview

This PR introduces ResourceAwareTaskOptimizationPass, a two-phase MLIR pass that optimizes CGRA resource allocation for the Neura taskflow dialect on a 4×4 CGRA grid (16 CGRAs total).

Phase 1: Utilization Fusion

Merges independent tasks (no SSA or memory dependency edges in either direction) into a single fused task, sequentially concatenating their loop bodies. This frees up CGRA budget that Phase 2 can reallocate to critical-path bottlenecks.

Phase 2: Latency-Aware Pipeline Balance

Uses the pipelined latency model:

latency(task) = II × (⌈trip_count / cgra_count⌉ − 1) + steps

Iteratively finds the critical-path bottleneck (minimum slack node with highest individual latency) and allocates one additional CGRA to it, repeating until the 16-CGRA budget is exhausted or no improvement is possible.

The outer loop (max 10 iterations) alternates fusion and balance until convergence (no change in either phase).

---

Speculative Profiling for compiled_ii and steps

To obtain accurate II and steps without waiting for full compilation:

1. Phase 1 (Taskflow → Neura): Clone the parent func::FuncOp, strip all tasks except the target, run ConstructHyperblockFromTask → ClassifyCounters → ConvertTaskflowToNeura on the clone to produce neura.kernel ops.
2. Phase 2 (Neura pipeline): Clone each kernel body into a standalone func::FuncOp tagged accelerator="neura", then run the full Neura lowering pipeline (LowerAffinePass → ConvertSCFToCFPass → AssignAccelerator → LowerMemRefToNeura → LowerArithToNeura → ... → InsertDataMovPass).

compiled_ii Extraction — Trade-offs

Source	When used	Accuracy
MapToAcceleratorPass → mapping_info.compiled_ii	All ops are DataMov-wrapped AND total ops ≤ 150	Highest (real modulo scheduler result)
max(ResMII, RecMII)	Mapper skipped (size guard or DataMov guard fails)	Lower bound, conservative
Default ii=1, steps=1	Phase 1 or 2 pipeline fails entirely	Pessimistic fallback

Guard conditions for the mapper:

• DataMov completeness: All non-reserve operand producers must be neura.data_mov. If InsertDataMovPass didn't fully wrap all operands (happens for kernels with complex control flow), the mapper asserts.
• Op count limit (kMapperOpLimit = 150): Prevents exponential backtracking in the modulo scheduler during speculative profiling of large kernels.

Split-Profile for Fused Tasks

After fusion, the fused task body contains N sequential loop nests. ConvertTaskflowToNeuraPass asserts hyperblock_count == 1, so we cannot profile the fused task directly. Instead we:

1. Create a temporary single-loop wrapper task for each top-level loop nest.
2. Profile each independently.
3. Assign max(ii) and sum(steps) to the fused task.

---

Test Coverage

Test	Tasks	Fusions	Result
irregular-loop	3 (incl. reduction)	0 (value output guard)	1+7+8=16 CGRAs
parallel-nested	2→1 (fused)	1	cgra_count=10, total=16
multi-nested	4→3 (one fusion)	1	Task_1: 6, fused: 9, Task_4: 1
resnet	13→6 (7 fusions)	7	Task_3: 6, Task_9: 6, others: 1

---

Known Limitations

1. Perfectly-nested assumption for trip_count: For non-perfectly-nested loops inside a task body, computeTripCount multiplies inner-loop counts of each top-level loop structure. This is accurate for the current workloads (convolutions, matmuls).
2. kMapperOpLimit = 150: Large kernels skip MapToAcceleratorPass and fall back to ResMII/RecMII bounds. This is a deliberate performance vs. accuracy trade-off for speculative profiling.
3. Fusion limited to write-output tasks: Tasks with value outputs (reductions) are excluded from utilization fusion. Full support would require tracking value-output flow across the fused task boundary.

[Copilot]

Pull request overview

This PR adds a new MLIR optimization pass that fuses independent Taskflow tasks and balances CGRA allocation using a pipelined latency model, plus updates several multi-CGRA tests to exercise the new behavior.

Changes:

• Introduces ResourceAwareTaskOptimizationPass implementing utilization fusion + latency-aware CGRA rebalancing with speculative profiling.
• Wires the new pass into build/registration (CMake + Passes.td/h).
• Extends Taskflow MLIR tests with --resource-aware-task-optimization RUN lines and RESOPT FileCheck assertions.

Reviewed changes

Copilot reviewed 9 out of 10 changed files in this pull request and generated 7 comments.

Show a summary per file

File	Description
lib/TaskflowDialect/Transforms/Optimizations/ResourceAwareTaskOptimizationPass.cpp	Implements the new two-phase optimization pass and speculative profiling pipeline
lib/TaskflowDialect/Transforms/Optimizations/CMakeLists.txt	Builds/links the new pass into the optimization library
include/TaskflowDialect/TaskflowPasses.td	Registers the new pass and its summary/description
include/TaskflowDialect/TaskflowPasses.h	Exposes the factory method for the new pass
test/multi-cgra/taskflow/resnet/simple_resnet_tosa.mlir	Adds RUN + FileCheck coverage for RESOPT expectations
test/multi-cgra/taskflow/parallel-nested/parallel-nested.mlir	Adds RUN + RESOPT checks (but currently duplicated)
test/multi-cgra/taskflow/multi-nested/multi-nested.mlir	Adds RUN + RESOPT checks
test/multi-cgra/taskflow/irregular-loop/irregular-loop.mlir	Adds RUN + RESOPT checks
test/benchmark/Zeonica_Testbench	Updates submodule pointer
debug.log	Adds a debug artifact containing a crash backtrace/logs

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[tancheng]

Shouldn't we fix #260 first to align the task/func/kernel?

[ShangkunLi]

Shouldn't we fix #260 first to align the task/func/kernel?

I think they are orthogonal. This pr is trying to do some optimizations on the task dependency graph, regardless of how we construct this graph.

For now, we build the task dependency graph based on the affine loops within one func. We can further extend it so that we can create a task dependency graph from multiple funcs.