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

include/Conversion/ConversionPasses.h

@@ -21,6 +21,7 @@ std::unique_ptr<mlir::Pass> createLowerBuiltinToNeuraPass();
 std::unique_ptr<mlir::Pass> createLowerAffineToNeuraPass();
 
 // TaskFlow Conversion Passes.
+std::unique_ptr<mlir::Pass> createAssignTaskTargetPass();
 std::unique_ptr<mlir::Pass> createConvertAffineToTaskflowPass();
 std::unique_ptr<mlir::Pass> createConvertTaskflowToNeuraPass();
 #define GEN_PASS_REGISTRATION

include/Conversion/ConversionPasses.td

@@ -48,6 +48,29 @@ def LowerAffineToNeura : Pass<"lower-affine-to-neura", "func::FuncOp">{
 // TaskFlow Conversion Passes.
 //=========================================================//
 
+def AssignTaskTarget : Pass<"assign-task-target", "ModuleOp">{
+  let summary = "Assign hardware targets to compute tasks (functions)";
+  let description = [{
+    This pass assigns hardware target attributes (target.device) to functions
+    based on their names. It enables heterogeneous workload partitioning across
+    different hardware units such as CPU, CGRA, and DOE.
+
+    The pass applies simple pattern matching rules:
+    - Functions containing "ray_sampler" or "sampler" -> CPU
+    - Functions containing "hash_encoder" or "encoder" -> DOE
+    - Functions containing "nerf_mlp" or "mlp" -> CGRA
+    - Top-level orchestrator functions (e.g., "nerf_forward") -> CPU
+    - Default -> CPU
+
+    Example output:
+      func.func @ray_sampler_func(...) attributes {target.device = "cpu"} { ... }
+      func.func @hash_encoder_func(...) attributes {target.device = "doe"} { ... }
+      func.func @nerf_mlp_func(...) attributes {target.device = "cgra"} { ... }
+  }];
+  let constructor = "mlir::createAssignTaskTargetPass()";
+  let dependentDialects = ["mlir::func::FuncDialect"];
+}
+
 def ConvertAffineToTaskflow : Pass<"convert-affine-to-taskflow", "ModuleOp">{
   let summary = "Convert top-level affine.for operations to Taskflow dialect";
   let description = [{

lib/Conversion/AssignTaskTarget/AssignTaskTargetPass.cpp

@@ -0,0 +1,133 @@
+//===- AssignTaskTargetPass.cpp - Assign hardware targets to tasks --------===//
+//
+// This pass assigns hardware target attributes to compute tasks (functions)
+// based on task names. It helps partition the workload across different
+// hardware units (CPU, CGRA, DOE, etc.) in heterogeneous computing systems.
+//
+//===----------------------------------------------------------------------===//
+
+#include "Conversion/ConversionPasses.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"
+#include "mlir/IR/Builders.h"
+#include "mlir/IR/BuiltinOps.h"
+#include "mlir/Pass/Pass.h"
+#include "llvm/Support/raw_ostream.h"
+
+using namespace mlir;
+
+namespace {
+
+//===----------------------------------------------------------------------===//
+// Helper Functions
+//===----------------------------------------------------------------------===//
+
+/// Determines the hardware target for a given function based on its name.
+/// This function implements a simple pattern-matching strategy:
+///   - ray_sampler* -> CPU
+///   - hash_encoder* -> DOE
+///   - nerf_mlp* -> CGRA
+///   - nerf_forward (top-level) -> CPU
+///   - default -> CPU
+static StringRef matchHardwareTarget(StringRef funcName) {
+  // Top-level function: runs on CPU as coordinator
+  if (funcName == "nerf_forward") {
+    return "cpu";
+  }
+
+  // Pattern matching for compute tasks
+  if (funcName.contains("ray_sampler") || funcName.contains("sampler")) {
+    return "cpu";
+  }
+
+  if (funcName.contains("hash_encoder") || funcName.contains("encoder")) {
+    return "doe";
+  }
+
+  if (funcName.contains("nerf_mlp") || funcName.contains("mlp")) {
+    return "cgra";
+  }
+
+  // Default target
+  return "cpu";
+}
+
+//===----------------------------------------------------------------------===//
+// AssignTaskTarget Pass
+//===----------------------------------------------------------------------===//
+
+struct AssignTaskTargetPass
+    : public PassWrapper<AssignTaskTargetPass, OperationPass<ModuleOp>> {
+
+  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(AssignTaskTargetPass)
+
+  StringRef getArgument() const final { return "assign-task-target"; }
+
+  StringRef getDescription() const final {
+    return "Assign hardware targets to compute tasks (functions) based on "
+           "task names";
+  }
+
+  void runOnOperation() override {
+    ModuleOp module = getOperation();
+    OpBuilder builder(&getContext());
+
+    // Statistics
+    unsigned totalFuncs = 0;
+    unsigned assignedFuncs = 0;
+    llvm::DenseMap<StringRef, unsigned> targetStats;
+
+    llvm::errs() << "\n";
+    llvm::errs() << "========================================\n";
+    llvm::errs() << "AssignTaskTarget Pass\n";
+    llvm::errs() << "========================================\n\n";
+
+    // Walk through all functions in the module
+    module.walk([&](func::FuncOp funcOp) {
+      totalFuncs++;
+      StringRef funcName = funcOp.getName();
+
+      // Determine hardware target based on function name
+      StringRef target = matchHardwareTarget(funcName);
+
+      // Set the target.device attribute
+      funcOp->setAttr("target.device", builder.getStringAttr(target));
+
+      assignedFuncs++;
+      targetStats[target]++;
+
+      llvm::errs() << "  [ASSIGN] " << funcName << " -> " << target << "\n";
+    });
+
+    // Print summary
+    llvm::errs() << "\n";
+    llvm::errs() << "========================================\n";
+    llvm::errs() << "Summary\n";
+    llvm::errs() << "========================================\n";
+    llvm::errs() << "Total functions:    " << totalFuncs << "\n";
+    llvm::errs() << "Assigned functions: " << assignedFuncs << "\n";
+
+    if (!targetStats.empty()) {
+      llvm::errs() << "\nTarget distribution:\n";
+      for (auto &entry : targetStats) {
+        llvm::errs() << "  " << entry.first << ": " << entry.second
+                     << " function(s)\n";
+      }
+    }
+
+    llvm::errs() << "========================================\n\n";
+  }
+};
+
+}  // namespace
+
+//===----------------------------------------------------------------------===//
+// Pass Registration
+//===----------------------------------------------------------------------===//
+
+namespace mlir {
+
+std::unique_ptr<Pass> createAssignTaskTargetPass() {
+  return std::make_unique<AssignTaskTargetPass>();
+}
+
+}  // namespace mlir

lib/Conversion/AssignTaskTarget/CMakeLists.txt

@@ -0,0 +1,14 @@
+include_directories(${CMAKE_CURRENT_BINARY_DIR})
+
+add_mlir_conversion_library(MLIRAssignTaskTargetPass
+  AssignTaskTargetPass.cpp
+
+  DEPENDS
+  MLIRConversionIncGen
+
+  LINK_LIBS PUBLIC
+  MLIRFuncDialect
+  MLIRIR
+  MLIRPass
+  MLIRSupport
+)

lib/Conversion/CMakeLists.txt

@@ -5,6 +5,7 @@ add_subdirectory(AffineToNeura)
 add_subdirectory(LlvmToNeura)
 add_subdirectory(MemRefToNeura)
 add_subdirectory(BuiltinToNeura)
+add_subdirectory(AssignTaskTarget)
 add_subdirectory(AffineToTaskflow)
 add_subdirectory(TaskflowToNeura)
 
@@ -23,6 +24,7 @@ target_link_libraries(MLIRConversion INTERFACE
   MLIRNeuraLlvmToNeuraPass
   MLIRNeuraMemRefToNeuraPass
   MLIRNeuraBuiltinToNeuraPass
+  MLIRAssignTaskTargetPass
   MLIRAffineToTaskflowPass
   MLIRTaskflowToNeuraPass
   ${dialect_libs}

lib/TaskflowDialect/Transforms/CMakeLists.txt

@@ -3,6 +3,7 @@ get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
 add_mlir_library(MLIRTaskflowTransforms
     ConstructHyperblockFromTaskPass.cpp
     CanonicalizeTaskPass.cpp
+    PartitionTaskByTarget.cpp
     ClassifyCountersPass.cpp
 
     DEPENDS