Make hic_executor_index_ an attribute

so cloning automatically works without a cloning constructor.
NVIDIA · Jan 15, 2025 · ecc60b6 · ecc60b6
1 parent 3cce103
commit ecc60b6
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 8 deletions.
diff --git a/csrc/host_ir/container.h b/csrc/host_ir/container.h
@@ -46,8 +46,8 @@ class HostIrContainer final : public Fusion {
     return kernel_executors_.push_back(std::move(ke));
   }
 
-  KernelExecutor* getKernelExecutor(int64_t index) {
-    return kernel_executors_[index].get();
+  KernelExecutor* getKernelExecutor(int64_t index) const {
+    return kernel_executors_.at(index).get();
   }
 
   Stream* getDefaultStream();

diff --git a/csrc/host_ir/host_ir.cpp b/csrc/host_ir/host_ir.cpp
@@ -124,8 +124,9 @@ LaunchKernel::LaunchKernel(
     int64_t hic_executor_index,
     const std::vector<Val*>& inputs,
     const std::vector<Val*>& outputs)
-    : Expr(passkey, inputs, outputs, {}),
-      hic_executor_index_(hic_executor_index) {}
+    : Expr(passkey, inputs, outputs, {}) {
+  addDataAttribute(hic_executor_index);
+}
 
 NVFUSER_DEFINE_CLONE_AND_CREATE(LaunchKernel)
 

diff --git a/csrc/host_ir/host_ir.h b/csrc/host_ir/host_ir.h
@@ -140,11 +140,8 @@ class LaunchKernel : public Expr {
   }
 
   int64_t getIndex() const {
-    return hic_executor_index_;
+    return attribute<int64_t>(0);
   }
-
- private:
-  const int64_t hic_executor_index_;
 };
 
 class Stream : public Val {