[refactor] Refactor llvm-offloaded-task-name mangling (#4418)

taichi-dev · Mar 2, 2022 · 1a8ec6e · 1a8ec6e
1 parent f05bc50
commit 1a8ec6e
Show file tree

Hide file tree

Showing 4 changed files with 12 additions and 10 deletions.
diff --git a/taichi/codegen/codegen_llvm.cpp b/taichi/codegen/codegen_llvm.cpp
@@ -124,9 +124,6 @@ CodeGenStmtGuard make_while_after_loop_guard(CodeGenLLVM *cg) {
 }  // namespace
 
 // CodeGenLLVM
-
-uint64 CodeGenLLVM::task_counter = 0;
-
 void CodeGenLLVM::visit(Block *stmt_list) {
   for (auto &stmt : stmt_list->statements) {
     stmt->accept(this);
@@ -1620,9 +1617,9 @@ std::string CodeGenLLVM::init_offloaded_task_function(OffloadedStmt *stmt,
       llvm::FunctionType::get(llvm::Type::getVoidTy(*llvm_context),
                               {llvm::PointerType::get(context_ty, 0)}, false);
 
-  auto task_kernel_name = fmt::format("{}_{}_{}{}", kernel_name, task_counter,
-                                      stmt->task_name(), suffix);
-  task_counter += 1;
+  auto task_kernel_name =
+      fmt::format("{}_{}_{}{}", kernel_name, kernel->get_next_task_id(),
+                  stmt->task_name(), suffix);
   func = llvm::Function::Create(task_function_type,
                                 llvm::Function::ExternalLinkage,
                                 task_kernel_name, module.get());

diff --git a/taichi/codegen/codegen_llvm.h b/taichi/codegen/codegen_llvm.h
@@ -49,8 +49,6 @@ class FunctionCreationGuard {
 
 class CodeGenLLVM : public IRVisitor, public LLVMModuleBuilder {
  public:
-  static uint64 task_counter;
-
   Kernel *kernel;
   IRNode *ir;
   Program *prog;

diff --git a/taichi/program/kernel.h b/taichi/program/kernel.h
@@ -113,6 +113,10 @@ class TI_DLL_EXPORT Kernel : public Callable {
 
   void account_for_offloaded(OffloadedStmt *stmt);
 
+  uint64 get_next_task_id() {
+    return task_counter_++;
+  }
+
   [[nodiscard]] std::string get_name() const override;
   /**
    * Whether the given |arch| is supported in the lower() method.
@@ -136,6 +140,7 @@ class TI_DLL_EXPORT Kernel : public Callable {
   // lower inital AST all the way down to a bunch of
   // OffloadedStmt for async execution
   bool lowered_{false};
+  std::atomic<uint64> task_counter_{0};
 };
 
 TLANG_NAMESPACE_END
diff --git a/taichi/program/program.cpp b/taichi/program/program.cpp
@@ -398,7 +398,8 @@ Kernel &Program::get_snode_writer(SNode *snode) {
 }
 
 Kernel &Program::get_ndarray_reader(Ndarray *ndarray) {
-  auto kernel_name = fmt::format("ndarray_reader");
+  static uint64 ndarray_reader_counter = 0;
+  auto kernel_name = fmt::format("ndarray_reader_{}", ndarray_reader_counter++);
   NdarrayRwKeys keys{ndarray->num_active_indices, ndarray->dtype};
   auto &ker = kernel([keys, this] {
     ExprGroup indices;
@@ -422,7 +423,8 @@ Kernel &Program::get_ndarray_reader(Ndarray *ndarray) {
 }
 
 Kernel &Program::get_ndarray_writer(Ndarray *ndarray) {
-  auto kernel_name = fmt::format("ndarray_writer");
+  static uint64 ndarray_writer_counter = 0;
+  auto kernel_name = fmt::format("ndarray_writer_{}", ndarray_writer_counter++);
   NdarrayRwKeys keys{ndarray->num_active_indices, ndarray->dtype};
   auto &ker = kernel([keys, this] {
     ExprGroup indices;