[aot] Serialize built graph, deserialize and run. #5016

ailzhang · 2022-05-20T00:27:39Z

Stack from ghstack (oldest at bottom):

related: #4786

[Update]: based on an offline discussion with @k-ye, I've split the
original Graph class into GraphBuilder and CompiledGraph classes
in C++. Note that the implementation didn't follow exactly the builder
design pattern as our builder is slightly simpler as shown below.
The complexity in our problem is more in the need of serialization and
deserialization for the same graph representation intead of its
construction process. So IMHO it's good enough to separate the
GraphBuilder and Runner(CompiledGraph) as we discussed. Please feel
free to correct me if I'm wrong!

GraphBuilder
    |
 compile()
    |
    |
CompiledGraph <----  serialize/deserialize ----> file
    |
    |
   run()

This PR demonstrates a minimal example of serializing a built graph,
deserializing and running it.

related: #4786 This PR demonstrates a minimal example of serializing a built graph, deserializing and running it. [ghstack-poisoned]

k-ye · 2022-05-20T05:32:26Z

tests/cpp/aot/aot_save_load_test.cpp

+
+  // API based on proposal https://github.com/taichi-dev/taichi/issues/3642
+  // Initialize Vulkan program
+  taichi::uint64 *result_buffer{nullptr};


Oh, I feel like it would be great if these tests can made the same for all the archs supporting AOT! (cc @turbo0628 @jim19930609 ) A more realistic result probably is that we have different setup routines for each arch. But the graph part stays the same.

taichi/backends/vulkan/aot_module_loader_impl.h

k-ye · 2022-05-20T05:37:10Z

taichi/backends/vulkan/aot_module_builder_impl.h

@@ -21,6 +21,11 @@ class AotModuleBuilderImpl : public AotModuleBuilder {
  void dump(const std::string &output_dir,
            const std::string &filename) const override;

+  // FIXME: remove me once TaichiAotData is no longer backend specific.


I think it's fine to add ti_aot_data_ to the module builder base class. It's just that not every impl uses it, yet?

k-ye · 2022-05-20T05:51:58Z

taichi/backends/vulkan/aot_module_builder_impl.cpp

@@ -147,6 +149,10 @@ void AotModuleBuilderImpl::add_per_backend(const std::string &identifier,
  ti_aot_data_.spirv_codes.push_back(compiled.task_spirv_source_codes);
 }

+void AotModuleBuilderImpl::add_compiled_kernel(aot::Kernel *kernel) {
+  kernel->save_to_module(this);


I'm just thinking if we can put save_to_module's impl within each module builder impl. That way aot::Kernel doesn't need to know anything about the module.

Furthermore, aot::Kernel has a misleading namespace now. It's probably more suitable to call this CompiledKernel or something...

+1 on the misleading namespace, I'll leave that to a separate search and replace PR!

related: taichi-dev#4786 This PR demonstrates a minimal example of serializing a built graph, deserializing and running it. ghstack-source-id: dd72f446caf0fbb33491d165e64ade93d9d92978 Pull Request resolved: taichi-dev#5016

related: #4786 This PR demonstrates a minimal example of serializing a built graph, deserializing and running it. [ghstack-poisoned]

related: taichi-dev#4786 This PR demonstrates a minimal example of serializing a built graph, deserializing and running it. ghstack-source-id: e369b326734b3cffaf91262ed84806044e121c3c Pull Request resolved: taichi-dev#5016

@k-ye

related: taichi-dev#4786 [Update]: based on an offline discussion with @k-ye, I've split the original `Graph` class into `GraphBuilder` and `CompiledGraph` classes in C++. ``` GraphBuilder | compile() | | CompiledGraph <---- serialize/deserialize ----> file | | run() ``` This PR demonstrates a minimal example of serializing a built graph, deserializing and running it. ghstack-source-id: e369b326734b3cffaf91262ed84806044e121c3c Pull Request resolved: taichi-dev#5016

k-ye

LGTM!

k-ye · 2022-05-23T01:33:55Z

taichi/aot/module_builder.h

@@ -37,6 +38,10 @@ class AotModuleBuilder {
  virtual void dump(const std::string &output_dir,
                    const std::string &filename) const = 0;

+  void dump_graph(std::string output_dir) const;


OOC, why we need a separate dump_graph, instead of dumping everything in dump?

@k-ye Since dump_graph is backend independent so I wanted to keep it in the base class. Note it's called in the dump in every backend indeed we didn't call them separately.

doesnt need to be public?

aha nice catch! it was a legacy public api used in the first prototype. Removed in #5024! :D

related: #4786 [Update]: based on an offline discussion with k-ye, I've split the original `Graph` class into `GraphBuilder` and `CompiledGraph` classes in C++. Note that the implementation didn't follow exactly the builder design pattern as our builder is slightly simpler as shown below. The complexity in our problem is more in the need of serialization and deserialization for the same graph representation intead of its construction process. So IMHO it's good enough to separate the GraphBuilder and Runner(`CompiledGraph`) as we discussed. Please feel free to correct me if I'm wrong! ``` GraphBuilder | compile() | | CompiledGraph <---- serialize/deserialize ----> file | | run() ``` This PR demonstrates a minimal example of serializing a built graph, deserializing and running it. [ghstack-poisoned]

CLAassistant · 2022-05-23T02:11:17Z

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you all sign our Contributor License Agreement before we can accept your contribution.
1 out of 2 committers have signed the CLA.

✅ ailzhang
❌ Ailing Zhang

Ailing Zhang seems not to be a GitHub user. You need a GitHub account to be able to sign the CLA. If you have already a GitHub account, please add the email address used for this commit to your account.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}

related: #4786 [Update]: based on an offline discussion with k-ye, I've split the original `Graph` class into `GraphBuilder` and `CompiledGraph` classes in C++. Note that the implementation didn't follow exactly the builder design pattern as our builder is slightly simpler as shown below. The complexity in our problem is more in the need of serialization and deserialization for the same graph representation intead of its construction process. So IMHO it's good enough to separate the GraphBuilder and Runner(`CompiledGraph`) as we discussed. Please feel free to correct me if I'm wrong! ``` GraphBuilder | compile() | | CompiledGraph <---- serialize/deserialize ----> file | | run() ``` This PR demonstrates a minimal example of serializing a built graph, deserializing and running it. [ghstack-poisoned]

ailzhang · 2022-05-23T02:22:35Z

Agh got some CLA issue on my newly setup workstation, lemme see if I can hack it around...

[aot] Serialize built graph, deserialize and run.

014593e

related: #4786 This PR demonstrates a minimal example of serializing a built graph, deserializing and running it. [ghstack-poisoned]

This was referenced May 20, 2022

[test] Unify kernel setup for ndarray related tests #5014

Merged

[aot] Build and run graph without serialization #5015

Merged

[aot] Move ArgKind as first argument in Arg class #5017

Closed

Update on "[aot] Serialize built graph, deserialize and run."

0cc69e0

related: #4786 This PR demonstrates a minimal example of serializing a built graph, deserializing and running it. [ghstack-poisoned]

k-ye reviewed May 20, 2022

View reviewed changes

Update on "[aot] Serialize built graph, deserialize and run."

1c60025

related: #4786 This PR demonstrates a minimal example of serializing a built graph, deserializing and running it. [ghstack-poisoned]

ailzhang requested a review from k-ye May 20, 2022 11:23

k-ye approved these changes May 23, 2022

View reviewed changes

ailzhang mentioned this pull request May 23, 2022

[aot] Bind graph APIs to python and add mpm88 example #5023

Closed

ailzhang closed this May 23, 2022

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[aot] Serialize built graph, deserialize and run. #5016

[aot] Serialize built graph, deserialize and run. #5016

ailzhang commented May 20, 2022 •

edited

Loading

k-ye May 20, 2022

k-ye May 20, 2022

k-ye May 20, 2022

ailzhang May 20, 2022

k-ye left a comment

k-ye May 23, 2022

ailzhang May 23, 2022

k-ye May 23, 2022

ailzhang May 23, 2022

CLAassistant commented May 23, 2022

ailzhang commented May 23, 2022

[aot] Serialize built graph, deserialize and run. #5016

[aot] Serialize built graph, deserialize and run. #5016

Conversation

ailzhang commented May 20, 2022 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

k-ye left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

CLAassistant commented May 23, 2022

ailzhang commented May 23, 2022

ailzhang commented May 20, 2022 •

edited

Loading