Skip to content

Commit

Permalink
[metal] Create helper methods for TLS codegen (#1982)
Browse files Browse the repository at this point in the history
  • Loading branch information
k-ye authored Oct 21, 2020
1 parent 523c68a commit 6e68872
Showing 1 changed file with 25 additions and 34 deletions.
59 changes: 25 additions & 34 deletions taichi/backends/metal/codegen_metal.cpp
Original file line number Diff line number Diff line change
@@ -931,17 +931,9 @@ class KernelCodegen : public IRVisitor {
// end_ = total_elems + begin_expr
emit("const int end_ = {} + {};", total_elems_name, begin_expr);

emit_runtime_and_memalloc_def();
if (used_tls) {
// Using TLS means we will access some SNodes within this kernel. The
// struct of an SNode needs Runtime and MemoryAllocator to construct.
emit_runtime_and_memalloc_def();
// Using |int32_t| because it aligns to 4bytes.
emit("// TLS prologue");
const std::string tls_bufi32_name = "tls_bufi32_";
emit("int32_t {}[{}];", tls_bufi32_name, (stmt->tls_size + 3) / 4);
emit("thread char* {} = reinterpret_cast<thread char*>({});",
kTlsBufferName, tls_bufi32_name);
stmt->tls_prologue->accept(this);
generate_tls_prologue(stmt);
}

emit("for (int ii = begin_; ii < end_; ii += {}) {{", kKernelGridSizeName);
@@ -964,12 +956,7 @@ class KernelCodegen : public IRVisitor {
emit("}}"); // closes for loop

if (used_tls) {
TI_ASSERT(stmt->tls_epilogue != nullptr);
inside_tls_epilogue_ = true;
emit("{{ // TLS epilogue");
stmt->tls_epilogue->accept(this);
inside_tls_epilogue_ = false;
emit("}}");
generate_tls_epilogue(stmt);
}

current_appender().pop_indent();
@@ -1010,17 +997,7 @@ class KernelCodegen : public IRVisitor {
emit_runtime_and_memalloc_def();

if (used_tls) {
// Using TLS means we will access some SNodes within this kernel. The
// struct of an SNode needs Runtime and MemoryAllocator to construct.
// Using |int32_t| because it aligns to 4bytes.
//
// TODO(k-ye): De-dupe TLS for range-for and struct-for.
emit("// TLS prologue");
const std::string tls_bufi32_name = "tls_bufi32_";
emit("int32_t {}[{}];", tls_bufi32_name, (stmt->tls_size + 3) / 4);
emit("thread char* {} = reinterpret_cast<thread char*>({});",
kTlsBufferName, tls_bufi32_name);
stmt->tls_prologue->accept(this);
generate_tls_prologue(stmt);
}

emit("ListManager parent_list;");
@@ -1079,19 +1056,33 @@ class KernelCodegen : public IRVisitor {
current_appender().pop_indent();

if (used_tls) {
// TODO(k-ye): De-dupe TLS for range-for and struct-for.
TI_ASSERT(stmt->tls_epilogue != nullptr);
inside_tls_epilogue_ = true;
emit("{{ // TLS epilogue");
stmt->tls_epilogue->accept(this);
inside_tls_epilogue_ = false;
emit("}}");
generate_tls_epilogue(stmt);
}
emit("}}\n"); // closes kernel

mtl_kernels_attribs()->push_back(ka);
}

void generate_tls_prologue(OffloadedStmt *stmt) {
TI_ASSERT(stmt->tls_prologue != nullptr);
emit("// TLS prologue");
const std::string tls_bufi32_name = "tls_bufi32_";
// Using |int32_t| because it aligns to 4bytes.
emit("int32_t {}[{}];", tls_bufi32_name, (stmt->tls_size + 3) / 4);
emit("thread char* {} = reinterpret_cast<thread char*>({});",
kTlsBufferName, tls_bufi32_name);
stmt->tls_prologue->accept(this);
}

void generate_tls_epilogue(OffloadedStmt *stmt) {
TI_ASSERT(stmt->tls_epilogue != nullptr);
inside_tls_epilogue_ = true;
emit("{{ // TLS epilogue");
stmt->tls_epilogue->accept(this);
inside_tls_epilogue_ = false;
emit("}}");
}

void add_runtime_list_op_kernel(OffloadedStmt *stmt,
const std::string &kernel_name) {
using Type = OffloadedStmt::TaskType;

0 comments on commit 6e68872

Please sign in to comment.