From 29b3fa736803cd482b8d70554cd5b85176359d73 Mon Sep 17 00:00:00 2001
From: PENGUINLIONG <admin@penguinliong.moe>
Date: Wed, 14 Sep 2022 14:06:07 +0800
Subject: [PATCH 1/3] Dump required device capability in AOT module meta

---
 taichi/aot/module_data.h                       | 4 +++-
 taichi/rhi/device.h                            | 5 +++++
 taichi/runtime/gfx/aot_module_builder_impl.cpp | 2 ++
 taichi/runtime/gfx/aot_utils.h                 | 4 +++-
 4 files changed, 13 insertions(+), 2 deletions(-)

diff --git a/taichi/aot/module_data.h b/taichi/aot/module_data.h
index f0c6ea740ae06..824de73fe6283 100644
--- 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
index 3e8bce1d35118..f0204ffc6e76c 100644
--- 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
index afd6c10e707dc..364a1657dfb91 100644
--- 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
index e8c1f5b0ea150..2bfdc367f57f1 100644
--- 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

From 4995cee69daa8bfa89b97f64fcc950c287d3f149 Mon Sep 17 00:00:00 2001
From: "pre-commit-ci[bot]"
 <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Date: Wed, 14 Sep 2022 06:11:04 +0000
Subject: [PATCH 2/3] [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci
---
 taichi/rhi/device.h | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

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

From bb1553cd99cbadd824fbd63a0d98acc09f8fd9b8 Mon Sep 17 00:00:00 2001
From: PENGUINLIONG <admin@penguinliong.moe>
Date: Wed, 14 Sep 2022 16:05:41 +0800
Subject: [PATCH 3/3] Redirect AOT modules to upgrade branch

---
 .github/workflows/scripts/aot-demo.sh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/scripts/aot-demo.sh b/.github/workflows/scripts/aot-demo.sh
index aeff7dac6dfd3..9ee8c0b127dd2 100755
--- 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