Skip to content

Commit

Permalink
[Metal] Support multiple SNode roots in codegen/runtime (taichi-dev#3090
Browse files Browse the repository at this point in the history
)

* [Metal] Support multiple SNode roots in codegen/runtime

* Auto Format

* fix

* Auto Format

* fix

* Update taichi/backends/metal/kernel_utils.cpp

Co-authored-by: Ailing  <[email protected]>

Co-authored-by: Taichi Gardener <[email protected]>
Co-authored-by: Ailing  <[email protected]>
  • Loading branch information
3 people authored Oct 9, 2021
1 parent 64f6220 commit bac269e
Show file tree
Hide file tree
Showing 11 changed files with 310 additions and 131 deletions.
8 changes: 4 additions & 4 deletions taichi/backends/metal/aot_module_builder_impl.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -11,10 +11,10 @@ namespace metal {

AotModuleBuilderImpl::AotModuleBuilderImpl(
const CompiledRuntimeModule *compiled_runtime_module,
const CompiledStructs *compiled_structs,
const std::vector<CompiledStructs> &compiled_snode_trees,
const BufferMetaData &buffer_meta_data)
: compiled_runtime_module_(compiled_runtime_module),
compiled_structs_(compiled_structs),
compiled_snode_trees_(compiled_snode_trees),
buffer_meta_data_(buffer_meta_data) {
ti_aot_data_.metadata = buffer_meta_data;
}
Expand Down Expand Up @@ -53,7 +53,7 @@ void AotModuleBuilderImpl::dump(const std::string &output_dir,

void AotModuleBuilderImpl::add_per_backend(const std::string &identifier,
Kernel *kernel) {
auto compiled = run_codegen(compiled_runtime_module_, compiled_structs_,
auto compiled = run_codegen(compiled_runtime_module_, compiled_snode_trees_,
kernel, &strtab_, /*offloaded=*/nullptr);
compiled.kernel_name = identifier;
ti_aot_data_.kernels.push_back(std::move(compiled));
Expand All @@ -79,7 +79,7 @@ void AotModuleBuilderImpl::add_per_backend_field(const std::string &identifier,
void AotModuleBuilderImpl::add_per_backend_tmpl(const std::string &identifier,
const std::string &key,
Kernel *kernel) {
auto compiled = run_codegen(compiled_runtime_module_, compiled_structs_,
auto compiled = run_codegen(compiled_runtime_module_, compiled_snode_trees_,
kernel, &strtab_, /*offloaded=*/nullptr);
for (auto &k : ti_aot_data_.tmpl_kernels) {
if (k.kernel_bundle_name == identifier) {
Expand Down
4 changes: 2 additions & 2 deletions taichi/backends/metal/aot_module_builder_impl.h
Original file line number Diff line number Diff line change
Expand Up @@ -16,7 +16,7 @@ class AotModuleBuilderImpl : public AotModuleBuilder {
public:
explicit AotModuleBuilderImpl(
const CompiledRuntimeModule *compiled_runtime_module,
const CompiledStructs *compiled_structs,
const std::vector<CompiledStructs> &compiled_snode_trees,
const BufferMetaData &buffer_meta_data);

void dump(const std::string &output_dir,
Expand All @@ -36,7 +36,7 @@ class AotModuleBuilderImpl : public AotModuleBuilder {

private:
const CompiledRuntimeModule *compiled_runtime_module_;
const CompiledStructs *compiled_structs_;
const std::vector<CompiledStructs> &compiled_snode_trees_;
BufferMetaData buffer_meta_data_;
PrintStringTable strtab_;
TaichiAotData ti_aot_data_;
Expand Down
Loading

0 comments on commit bac269e

Please sign in to comment.