[aot] Dump required device capability in AOT module meta (#6056)

taichi-dev · Sep 14, 2022 · 5046716 · 5046716
1 parent cf4e710
commit 5046716
Show file tree

Hide file tree

Showing 5 changed files with 14 additions and 3 deletions.
diff --git a/.github/workflows/scripts/aot-demo.sh b/.github/workflows/scripts/aot-demo.sh
@@ -48,7 +48,7 @@ function prepare-unity-build-env {
     cd taichi
 
     # Dependencies
-    git clone --reference-if-able /var/lib/git-cache https://github.com/taichi-dev/Taichi-UnityExample
+    git clone --reference-if-able /var/lib/git-cache -b upgrade-modules1 https://github.com/taichi-dev/Taichi-UnityExample
 
     python misc/generate_unity_language_binding.py
     cp c_api/unity/*.cs Taichi-UnityExample/Assets/Taichi/Generated

diff --git a/taichi/aot/module_data.h b/taichi/aot/module_data.h
@@ -3,6 +3,7 @@
 #include <string>
 #include <vector>
 
+#include "taichi/rhi/device.h"
 #include "taichi/common/core.h"
 #include "taichi/common/serialization.h"
 
@@ -120,6 +121,7 @@ struct ModuleData {
   std::unordered_map<std::string, CompiledTaichiKernel> kernels;
   std::unordered_map<std::string, CompiledTaichiKernel> kernel_tmpls;
   std::vector<aot::CompiledFieldData> fields;
+  std::map<DeviceCapability, uint32_t> required_caps;
 
   size_t root_buffer_size;
 
@@ -129,7 +131,7 @@ struct ModuleData {
     ts.write_to_file(path);
   }
 
-  TI_IO_DEF(kernels, kernel_tmpls, fields, root_buffer_size);
+  TI_IO_DEF(kernels, kernel_tmpls, fields, required_caps, root_buffer_size);
 };
 
 }  // namespace aot

diff --git a/taichi/rhi/device.h b/taichi/rhi/device.h
@@ -418,6 +418,11 @@ class Device {
       dest.set_cap(k, v);
     }
   }
+  void clone_caps(std::map<DeviceCapability, uint32_t> &dest) const {
+    for (const auto &[k, v] : caps_) {
+      dest[k] = v;
+    }
+  }
 
   void print_all_cap() const;
 

diff --git a/taichi/runtime/gfx/aot_module_builder_impl.cpp b/taichi/runtime/gfx/aot_module_builder_impl.cpp
@@ -29,6 +29,7 @@ class AotDataConverter {
       res.kernels[ker.name] = val;
     }
     res.fields = in.fields;
+    res.required_caps = in.required_caps;
     res.root_buffer_size = in.root_buffer_size;
     return res;
   }
@@ -110,6 +111,7 @@ AotModuleBuilderImpl::AotModuleBuilderImpl(
   aot_target_device_ =
       target_device ? std::move(target_device)
                     : std::make_unique<aot::TargetDevice>(device_api_backend_);
+  aot_target_device_->clone_caps(ti_aot_data_.required_caps);
   if (!compiled_structs.empty()) {
     ti_aot_data_.root_buffer_size = compiled_structs[0].root_size;
   }

diff --git a/taichi/runtime/gfx/aot_utils.h b/taichi/runtime/gfx/aot_utils.h
@@ -1,6 +1,7 @@
 #pragma once
 
 #include <vector>
+#include <map>
 
 #include "taichi/codegen/spirv/kernel_utils.h"
 #include "taichi/aot/module_loader.h"
@@ -17,9 +18,10 @@ struct TaichiAotData {
   std::vector<std::vector<std::vector<uint32_t>>> spirv_codes;
   std::vector<spirv::TaichiKernelAttributes> kernels;
   std::vector<aot::CompiledFieldData> fields;
+  std::map<DeviceCapability, uint32_t> required_caps;
   size_t root_buffer_size{0};
 
-  TI_IO_DEF(kernels, fields, root_buffer_size);
+  TI_IO_DEF(kernels, fields, required_caps, root_buffer_size);
 };
 
 }  // namespace gfx