Add function (task)-level hardware target assignment pass for heterogeneous computing

[YanzhouTang]

Summary

This PR introduces a new pass AssignTaskTarget that operates at a higher abstraction level than existing partitioning mechanisms. It assigns hardware targets (CPU, CGRA, DOE) to compute functions before they are lowered to taskflow operations, enabling coarse-grained workload partitioning in heterogeneous computing systems.

Motivation

The existing PartitionTaskByTarget pass operates at the taskflow level (loop-to-CGRA mapping), which is fine-grained for certain use cases. We need a higher-level pass that can:

1. Partition at function granularity: Assign entire computational functions to different hardware units
2. Early hardware mapping: Make hardware decisions at the Linalg/func level before lowering to lower-level IRs
3. Enable heterogeneous orchestration: Support scenarios where different components run on different hardware accelerators (Now we support NeRF algorithm)

Changes

New Pass: AssignTaskTarget

• Location: lib/Conversion/AssignTaskTarget/
• Functionality:
  • Analyzes function names and assigns target.device attributes
  • Supports CPU, CGRA, and DOE (Data Orchestration Engine) hardware targets
  • Uses pattern matching for initial implementation (extensible to more sophisticated analysis)
• Usage: mlir-neura-opt --assign-task-target input.mlir

Example transformation:
// Before
func.func @hash_encoder_func(...) { ... }

// After
func.func @hash_encoder_func(...) attributes {target.device = "doe"} { ... }

[tancheng]

what is the principle to determine which accelerator should be assigned to which kernel?

[YanzhouTang]

what is the principle to determine which accelerator should be assigned to which kernel?

Currently, it's hardcoded for the NeRF use case with name-based pattern matching (e.g., sampler → CPU, encoder → DOE, mlp → CGRA).

We plan to add YAML configuration support in the next iteration to allow flexible user-specified mappings.

Fully automated partitioning based on workload analysis is a potential future direction, but we believe explicit configuration provides better predictability and user control for now.

[tancheng]

what is the principle to determine which accelerator should be assigned to which kernel?

Currently, it's hardcoded for the NeRF use case with name-based pattern matching (e.g., sampler → CPU, encoder → DOE, mlp → CGRA).

We plan to add YAML configuration support in the next iteration to allow flexible user-specified mappings.

Fully automated partitioning based on workload analysis is a potential future direction, but we believe explicit configuration provides better predictability and user control for now.

Sounds good on the naming pattern matching part.

However, plz discuss with @ShangkunLi about the granularity of assigning kernel to accelerator. It seems @ShangkunLi and @HobbitQia would perform mapping algorithm on kernel instead of func. So we either need to somehow transform func to neura kernel, or vice versa. WDYT, @ShangkunLi @HobbitQia @YanzhouTang?

[YanzhouTang]

what is the principle to determine which accelerator should be assigned to which kernel?

Currently, it's hardcoded for the NeRF use case with name-based pattern matching (e.g., sampler → CPU, encoder → DOE, mlp → CGRA).
We plan to add YAML configuration support in the next iteration to allow flexible user-specified mappings.
Fully automated partitioning based on workload analysis is a potential future direction, but we believe explicit configuration provides better predictability and user control for now.

Sounds good on the naming pattern matching part.

However, plz discuss with @ShangkunLi about the granularity of assigning kernel to accelerator. It seems @ShangkunLi and @HobbitQia would perform mapping algorithm on kernel instead of func. So we either need to somehow transform func to neura kernel, or vice versa. WDYT, @ShangkunLi @HobbitQia @YanzhouTang?

Thanks for raising this important point!

Before responding, I'd like to clarify a few technical details about the granularity and integration: @ShangkunLi @HobbitQia

1. What IR level does your mapping algorithm operate on? (taskflow, neura.kernel, affine.for, or others?)
2. Can you share an example MLIR input that your algorithm expects?

Let's align on these before finalizing the design. Happy to adjust the implementation based on your feedback!