From e764f4e45a4ebb767138d64faa0b53bee474109f Mon Sep 17 00:00:00 2001
From: Vissidarte-Herman <93570324+Vissidarte-Herman@users.noreply.github.com>
Date: Mon, 21 Mar 2022 11:34:38 +0800
Subject: [PATCH 1/8] Update differences_between_taichi_and_python_programs.md
 (#4583)

Fixed a broken link.
---
 .../basic/differences_between_taichi_and_python_programs.md     | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/docs/lang/articles/basic/differences_between_taichi_and_python_programs.md b/docs/lang/articles/basic/differences_between_taichi_and_python_programs.md
index 783d4105a40d0..6efe9aae760ba 100644
--- a/docs/lang/articles/basic/differences_between_taichi_and_python_programs.md
+++ b/docs/lang/articles/basic/differences_between_taichi_and_python_programs.md
@@ -155,7 +155,7 @@ Taichi partially supports list comprehension and dictionary comprehension,
 but does not support set comprehension.
 
 For list comprehensions and dictionary comprehensions, the `if`s and `for`s in them are evaluated at compile time.
-The iterators and conditions are implicitly in [static scope](/lang/articles/advanced/meta#static-scope).
+The iterators and conditions are implicitly in [static scope](/lang/articles/advanced/meta.md#static-scope).
 
 ### Operator `is`
 

From 39402d9479f9e25d30348272b7a29cf4cb56a6a7 Mon Sep 17 00:00:00 2001
From: Yi Xu <xy_xuyi@foxmail.com>
Date: Mon, 21 Mar 2022 15:30:30 +0800
Subject: [PATCH 2/8] [test] Fix ill-formed test_binary_func_ret (#4587)

* [test] Fix ill-formed test_binary_func_ret

* Auto Format

Co-authored-by: Taichi Gardener <taichigardener@gmail.com>
---
 tests/python/test_return.py | 14 ++++++--------
 1 file changed, 6 insertions(+), 8 deletions(-)

diff --git a/tests/python/test_return.py b/tests/python/test_return.py
index 30bcc273dfcde..e4a8b6d88d6f7 100644
--- a/tests/python/test_return.py
+++ b/tests/python/test_return.py
@@ -29,8 +29,13 @@ def func2() -> ti.i32:
     assert func2() == 3
 
 
+@pytest.mark.parametrize("dt1,dt2,dt3,castor",
+                         [(ti.i32, ti.f32, ti.f32, float),
+                          (ti.f32, ti.i32, ti.f32, float),
+                          (ti.i32, ti.f32, ti.i32, int),
+                          (ti.f32, ti.i32, ti.i32, int)])
 @test_utils.test()
-def _test_binary_func_ret(dt1, dt2, dt3, castor):
+def test_binary_func_ret(dt1, dt2, dt3, castor):
     @ti.kernel
     def func(a: dt1, b: dt2) -> dt3:
         return a * b
@@ -49,13 +54,6 @@ def func(a: dt1, b: dt2) -> dt3:
         assert func(x, y) == test_utils.approx(castor(x * y))
 
 
-def test_binary_func_ret():
-    _test_binary_func_ret(ti.i32, ti.f32, ti.f32, float)
-    _test_binary_func_ret(ti.f32, ti.i32, ti.f32, float)
-    _test_binary_func_ret(ti.i32, ti.f32, ti.i32, int)
-    _test_binary_func_ret(ti.f32, ti.i32, ti.i32, int)
-
-
 @test_utils.test()
 def test_return_in_static_if():
     @ti.kernel

From b6e7e76af3010c30527c94140246efbcd775c58a Mon Sep 17 00:00:00 2001
From: Ye Kuang <k-ye@users.noreply.github.com>
Date: Mon, 21 Mar 2022 16:34:19 +0800
Subject: [PATCH 3/8] [gui] Make GGUI VBO configurable (#4575)

* wip

* wip

* [gui] Make VBO configurable

* fix

* fix

* fix

* revert build

* add missing filie

* newline

* tweak

* Auto Format

* fix vertex copy

* Auto Format

* constexpr

* fix

* Auto Format

Co-authored-by: Taichi Gardener <taichigardener@gmail.com>
---
 .../backends/vulkan/aot_module_loader_impl.h  |  5 +-
 taichi/ui/backends/vulkan/renderable.cpp      | 39 +++++++++++----
 taichi/ui/backends/vulkan/renderable.h        | 50 +++++++++++--------
 .../backends/vulkan/renderables/circles.cpp   |  9 ++--
 .../ui/backends/vulkan/renderables/circles.h  |  6 ++-
 .../ui/backends/vulkan/renderables/lines.cpp  |  2 +-
 taichi/ui/backends/vulkan/renderables/lines.h |  2 +-
 .../ui/backends/vulkan/renderables/mesh.cpp   |  2 +-
 taichi/ui/backends/vulkan/renderables/mesh.h  |  2 +-
 .../backends/vulkan/renderables/particles.cpp |  2 +-
 .../backends/vulkan/renderables/particles.h   |  2 +-
 .../backends/vulkan/renderables/set_image.cpp | 27 +++++++---
 .../backends/vulkan/renderables/set_image.h   |  2 +-
 .../backends/vulkan/renderables/triangles.cpp |  2 +-
 .../backends/vulkan/renderables/triangles.h   |  2 +-
 taichi/ui/backends/vulkan/renderer.cpp        | 27 ++++++----
 taichi/ui/backends/vulkan/renderer.h          |  5 +-
 taichi/ui/backends/vulkan/vertex.cpp          | 27 ++++++++++
 taichi/ui/backends/vulkan/vertex.h            | 44 +++++++++++++++-
 taichi/ui/common/canvas_base.h                |  5 +-
 taichi/ui/common/renderable_info.h            |  5 +-
 taichi/ui/common/scene_base.h                 |  7 +--
 22 files changed, 198 insertions(+), 76 deletions(-)
 create mode 100644 taichi/ui/backends/vulkan/vertex.cpp

diff --git a/taichi/backends/vulkan/aot_module_loader_impl.h b/taichi/backends/vulkan/aot_module_loader_impl.h
index 5329e889420c2..7d32d991f2e89 100644
--- a/taichi/backends/vulkan/aot_module_loader_impl.h
+++ b/taichi/backends/vulkan/aot_module_loader_impl.h
@@ -1,5 +1,6 @@
 #pragma once
 
+#include <any>
 #include <string>
 #include <vector>
 
@@ -15,12 +16,12 @@ namespace vulkan {
 
 class VkRuntime;
 
-struct AotModuleParams {
+struct TI_DLL_EXPORT AotModuleParams {
   std::string module_path;
   VkRuntime *runtime{nullptr};
 };
 
-std::unique_ptr<aot::Module> make_aot_module(std::any mod_params);
+TI_DLL_EXPORT std::unique_ptr<aot::Module> make_aot_module(std::any mod_params);
 
 }  // namespace vulkan
 }  // namespace lang
diff --git a/taichi/ui/backends/vulkan/renderable.cpp b/taichi/ui/backends/vulkan/renderable.cpp
index 4ba2da75ffed1..1c327bac3dc60 100644
--- a/taichi/ui/backends/vulkan/renderable.cpp
+++ b/taichi/ui/backends/vulkan/renderable.cpp
@@ -1,4 +1,5 @@
 #include "taichi/ui/backends/vulkan/renderable.h"
+
 #include "taichi/program/program.h"
 #include "taichi/ui/utils/utils.h"
 
@@ -39,6 +40,7 @@ void Renderable::init_buffers() {
 }
 
 void Renderable::update_data(const RenderableInfo &info) {
+  TI_ASSERT(info.vbo_attrs == config_.vbo_attrs);
   // We might not have a current program if GGUI is used in external apps to
   // load AOT modules
   Program *prog = app_context_->prog();
@@ -78,7 +80,7 @@ void Renderable::update_data(const RenderableInfo &info) {
     vbo_dev_ptr = get_device_ptr(prog, info.vbo.snode);
   }
 
-  uint64_t vbo_size = sizeof(Vertex) * num_vertices;
+  const uint64_t vbo_size = config_.vbo_size() * num_vertices;
 
   Device::MemcpyCapability memcpy_cap = Device::check_memcpy_capability(
       vertex_buffer_.get_ptr(), vbo_dev_ptr, vbo_size);
@@ -140,20 +142,37 @@ void Renderable::create_graphics_pipeline() {
   raster_params.depth_test = true;
   raster_params.depth_write = true;
 
-  std::vector<VertexInputBinding> vertex_inputs = {{0, sizeof(Vertex), false}};
+  std::vector<VertexInputBinding> vertex_inputs = {
+      {/*binding=*/0, config_.vbo_size(), /*instance=*/false}};
   // TODO: consider using uint8 for colors and normals
-  std::vector<VertexInputAttribute> vertex_attribs = {
-      {0, 0, BufferFormat::rgb32f, offsetof(Vertex, pos)},
-      {1, 0, BufferFormat::rgb32f, offsetof(Vertex, normal)},
-      {2, 0, BufferFormat::rg32f, offsetof(Vertex, texCoord)},
-      {3, 0, BufferFormat::rgba32f, offsetof(Vertex, color)}};
+  std::vector<VertexInputAttribute> vertex_attribs;
+  if (VboHelpers::has_attr(config_.vbo_attrs, VertexAttributes::kPos)) {
+    vertex_attribs.push_back({/*location=*/0, /*binding=*/0,
+                              /*format=*/BufferFormat::rgb32f,
+                              /*offset=*/offsetof(Vertex, pos)});
+  }
+  if (VboHelpers::has_attr(config_.vbo_attrs, VertexAttributes::kNormal)) {
+    vertex_attribs.push_back({/*location=*/1, /*binding=*/0,
+                              /*format=*/BufferFormat::rgb32f,
+                              /*offset=*/offsetof(Vertex, normal)});
+  }
+  if (VboHelpers::has_attr(config_.vbo_attrs, VertexAttributes::kUv)) {
+    vertex_attribs.push_back({/*location=*/2, /*binding=*/0,
+                              /*format=*/BufferFormat::rg32f,
+                              /*offset=*/offsetof(Vertex, tex_coord)});
+  }
+  if (VboHelpers::has_attr(config_.vbo_attrs, VertexAttributes::kColor)) {
+    vertex_attribs.push_back({/*location=*/3, /*binding=*/0,
+                              /*format=*/BufferFormat::rgba32f,
+                              /*offset=*/offsetof(Vertex, color)});
+  }
 
   pipeline_ = app_context_->device().create_raster_pipeline(
       source, raster_params, vertex_inputs, vertex_attribs);
 }
 
 void Renderable::create_vertex_buffer() {
-  size_t buffer_size = sizeof(Vertex) * config_.max_vertices_count;
+  const size_t buffer_size = config_.vbo_size() * config_.max_vertices_count;
 
   Device::AllocParams vb_params{buffer_size, false, false,
                                 app_context_->requires_export_sharing(),
@@ -181,7 +200,7 @@ void Renderable::create_index_buffer() {
 }
 
 void Renderable::create_uniform_buffers() {
-  size_t buffer_size = config_.ubo_size;
+  const size_t buffer_size = config_.ubo_size;
   if (buffer_size == 0) {
     return;
   }
@@ -192,7 +211,7 @@ void Renderable::create_uniform_buffers() {
 }
 
 void Renderable::create_storage_buffers() {
-  size_t buffer_size = config_.ssbo_size;
+  const size_t buffer_size = config_.ssbo_size;
   if (buffer_size == 0) {
     return;
   }
diff --git a/taichi/ui/backends/vulkan/renderable.h b/taichi/ui/backends/vulkan/renderable.h
index 16e5fae6045ce..6fc9ab87ffa6b 100644
--- a/taichi/ui/backends/vulkan/renderable.h
+++ b/taichi/ui/backends/vulkan/renderable.h
@@ -1,39 +1,45 @@
 #pragma once
 
-#include <iostream>
-#include <fstream>
-#include <stdexcept>
 #include <algorithm>
+#include <array>
 #include <chrono>
-#include <vector>
-#include <cstring>
-#include <cstdlib>
 #include <cstdint>
-#include <array>
+#include <cstdlib>
+#include <cstring>
+#include <fstream>
+#include <iostream>
 #include <optional>
 #include <set>
-#include "taichi/ui/utils/utils.h"
-#include "taichi/ui/backends/vulkan/vertex.h"
+#include <stdexcept>
+#include <vector>
+
+#include "taichi/backends/vulkan/vulkan_device.h"
 #include "taichi/ui/backends/vulkan/app_context.h"
 #include "taichi/ui/backends/vulkan/swap_chain.h"
-
+#include "taichi/ui/backends/vulkan/vertex.h"
 #include "taichi/ui/common/renderable_info.h"
-#include "taichi/backends/vulkan/vulkan_device.h"
-
-TI_UI_NAMESPACE_BEGIN
+#include "taichi/ui/utils/utils.h"
 
+namespace taichi {
+namespace ui {
 namespace vulkan {
 
 struct RenderableConfig {
-  int max_vertices_count;
-  int max_indices_count;
-  int vertices_count;
-  int indices_count;
-  size_t ubo_size;
-  size_t ssbo_size;
+  int max_vertices_count{0};
+  int max_indices_count{0};
+  int vertices_count{0};
+  int indices_count{0};
+  size_t ubo_size{0};
+  size_t ssbo_size{0};
   std::string vertex_shader_path;
   std::string fragment_shader_path;
-  taichi::lang::TopologyType topology_type;
+  taichi::lang::TopologyType topology_type{
+      taichi::lang::TopologyType::Triangles};
+  VertexAttributes vbo_attrs{VboHelpers::all()};
+
+  size_t vbo_size() const {
+    return VboHelpers::size(vbo_attrs);
+  }
 };
 
 class Renderable {
@@ -95,5 +101,5 @@ class Renderable {
 };
 
 }  // namespace vulkan
-
-TI_UI_NAMESPACE_END
+}  // namespace ui
+}  // namespace taichi
diff --git a/taichi/ui/backends/vulkan/renderables/circles.cpp b/taichi/ui/backends/vulkan/renderables/circles.cpp
index 9b2a36e695620..973dd663003f5 100644
--- a/taichi/ui/backends/vulkan/renderables/circles.cpp
+++ b/taichi/ui/backends/vulkan/renderables/circles.cpp
@@ -16,7 +16,9 @@ void Circles::update_data(const CirclesInfo &info) {
              info.radius);
 }
 
-void Circles::init_circles(AppContext *app_context, int vertices_count) {
+void Circles::init_circles(AppContext *app_context,
+                           int vertices_count,
+                           VertexAttributes vbo_attrs) {
   RenderableConfig config = {
       vertices_count,
       1,
@@ -27,14 +29,15 @@ void Circles::init_circles(AppContext *app_context, int vertices_count) {
       app_context->config.package_path + "/shaders/Circles_vk_vert.spv",
       app_context->config.package_path + "/shaders/Circles_vk_frag.spv",
       TopologyType::Points,
+      vbo_attrs,
   };
 
   Renderable::init(config, app_context);
   Renderable::init_render_resources();
 }
 
-Circles::Circles(AppContext *app_context) {
-  init_circles(app_context, 1);
+Circles::Circles(AppContext *app_context, VertexAttributes vbo_attrs) {
+  init_circles(app_context, /*vertices_count=*/1, vbo_attrs);
 }
 
 void Circles::update_ubo(glm::vec3 color,
diff --git a/taichi/ui/backends/vulkan/renderables/circles.h b/taichi/ui/backends/vulkan/renderables/circles.h
index 4507920d8bf9a..52fc4a0625d74 100644
--- a/taichi/ui/backends/vulkan/renderables/circles.h
+++ b/taichi/ui/backends/vulkan/renderables/circles.h
@@ -27,7 +27,7 @@ namespace vulkan {
 
 class Circles final : public Renderable {
  public:
-  Circles(AppContext *app_context);
+  Circles(AppContext *app_context, VertexAttributes vbo_attrs);
   void update_data(const CirclesInfo &info);
 
  private:
@@ -37,7 +37,9 @@ class Circles final : public Renderable {
     float radius;
   };
 
-  void init_circles(AppContext *app_context, int vertices_count);
+  void init_circles(AppContext *app_context,
+                    int vertices_count,
+                    VertexAttributes vbo_attrs);
 
   void update_ubo(glm::vec3 color, bool use_per_vertex_color, float radius);
 
diff --git a/taichi/ui/backends/vulkan/renderables/lines.cpp b/taichi/ui/backends/vulkan/renderables/lines.cpp
index 773bb0e12fb4d..71d402842e413 100644
--- a/taichi/ui/backends/vulkan/renderables/lines.cpp
+++ b/taichi/ui/backends/vulkan/renderables/lines.cpp
@@ -38,7 +38,7 @@ void Lines::init_lines(AppContext *app_context,
   Renderable::init_render_resources();
 }
 
-Lines::Lines(AppContext *app_context) {
+Lines::Lines(AppContext *app_context, VertexAttributes vbo_attrs) {
   init_lines(app_context, 4, 6);
 }
 
diff --git a/taichi/ui/backends/vulkan/renderables/lines.h b/taichi/ui/backends/vulkan/renderables/lines.h
index d747fd2104e1b..b26bbeaeaca02 100644
--- a/taichi/ui/backends/vulkan/renderables/lines.h
+++ b/taichi/ui/backends/vulkan/renderables/lines.h
@@ -27,7 +27,7 @@ namespace vulkan {
 
 class Lines final : public Renderable {
  public:
-  Lines(AppContext *app_context);
+  Lines(AppContext *app_context, VertexAttributes vbo_attrs);
 
   void update_data(const LinesInfo &info);
 
diff --git a/taichi/ui/backends/vulkan/renderables/mesh.cpp b/taichi/ui/backends/vulkan/renderables/mesh.cpp
index b632ec0e35cf8..a0cf8cea661c2 100644
--- a/taichi/ui/backends/vulkan/renderables/mesh.cpp
+++ b/taichi/ui/backends/vulkan/renderables/mesh.cpp
@@ -9,7 +9,7 @@ namespace vulkan {
 
 using namespace taichi::lang;
 
-Mesh::Mesh(AppContext *app_context) {
+Mesh::Mesh(AppContext *app_context, VertexAttributes vbo_attrs) {
   init_mesh(app_context, 3, 3);
 }
 
diff --git a/taichi/ui/backends/vulkan/renderables/mesh.h b/taichi/ui/backends/vulkan/renderables/mesh.h
index 5ceba5dd76a99..a88380cd761e9 100644
--- a/taichi/ui/backends/vulkan/renderables/mesh.h
+++ b/taichi/ui/backends/vulkan/renderables/mesh.h
@@ -27,7 +27,7 @@ namespace vulkan {
 
 class Mesh final : public Renderable {
  public:
-  Mesh(AppContext *app_context);
+  Mesh(AppContext *app_context, VertexAttributes vbo_attrs);
 
   void update_data(const MeshInfo &info, const Scene &scene);
 
diff --git a/taichi/ui/backends/vulkan/renderables/particles.cpp b/taichi/ui/backends/vulkan/renderables/particles.cpp
index 699bb9410cf0a..2fa0ba1269023 100644
--- a/taichi/ui/backends/vulkan/renderables/particles.cpp
+++ b/taichi/ui/backends/vulkan/renderables/particles.cpp
@@ -9,7 +9,7 @@ namespace vulkan {
 using namespace taichi::lang;
 using namespace taichi::lang::vulkan;
 
-Particles::Particles(AppContext *app_context) {
+Particles::Particles(AppContext *app_context, VertexAttributes vbo_attrs) {
   init_particles(app_context, 1);
 }
 
diff --git a/taichi/ui/backends/vulkan/renderables/particles.h b/taichi/ui/backends/vulkan/renderables/particles.h
index d75f4b1bdb88a..b71e89d4a670c 100644
--- a/taichi/ui/backends/vulkan/renderables/particles.h
+++ b/taichi/ui/backends/vulkan/renderables/particles.h
@@ -27,7 +27,7 @@ namespace vulkan {
 
 class Particles final : public Renderable {
  public:
-  Particles(AppContext *app_context);
+  Particles(AppContext *app_context, VertexAttributes vbo_attrs);
 
   void update_data(const ParticlesInfo &info, const Scene &scene);
 
diff --git a/taichi/ui/backends/vulkan/renderables/set_image.cpp b/taichi/ui/backends/vulkan/renderables/set_image.cpp
index 316297849396e..8c3831b6ec345 100644
--- a/taichi/ui/backends/vulkan/renderables/set_image.cpp
+++ b/taichi/ui/backends/vulkan/renderables/set_image.cpp
@@ -79,7 +79,7 @@ void SetImage::update_data(const SetImageInfo &info) {
   stream->submit_synced(cmd_list.get());
 }
 
-SetImage::SetImage(AppContext *app_context) {
+SetImage::SetImage(AppContext *app_context, VertexAttributes vbo_attrs) {
   init_set_image(app_context, 1, 1);
 }
 
@@ -154,19 +154,30 @@ void SetImage::update_vertex_buffer_() {
       {{1.f, 1.f, 0.f}, {0.f, 0.f, 1.f}, {1.f, 0.f}, {1.f, 1.f, 1.f}},
       {{1.f, -1.f, 0.f}, {0.f, 0.f, 1.f}, {1.f, 1.f}, {1.f, 1.f, 1.f}},
   };
-
+  // Our actual VBO might only use the first several attributes in `Vertex`,
+  // therefore this slicing & copying for each Vertex.
   {
-    Vertex *mapped_vbo =
-        (Vertex *)app_context_->device().map(staging_vertex_buffer_);
-
-    memcpy(mapped_vbo, vertices.data(),
-           (size_t)config_.vertices_count * sizeof(Vertex));
+    char *mapped_vbo =
+        static_cast<char *>(app_context_->device().map(staging_vertex_buffer_));
+    for (int i = 0; i < vertices.size(); ++i) {
+      const char *src = reinterpret_cast<const char *>(&vertices[i]);
+      for (auto a : VboHelpers::kOrderedAttrs) {
+        const auto a_sz = VboHelpers::size(a);
+        if (VboHelpers::has_attr(config_.vbo_attrs, a)) {
+          memcpy(mapped_vbo, src, a_sz);
+          mapped_vbo += a_sz;
+        }
+        // Pointer to the full Vertex attributes needs to be advanced
+        // unconditionally.
+        src += a_sz;
+      }
+    }
     app_context_->device().unmap(staging_vertex_buffer_);
   }
 
   app_context_->device().memcpy_internal(
       vertex_buffer_.get_ptr(0), staging_vertex_buffer_.get_ptr(0),
-      config_.vertices_count * sizeof(Vertex));
+      config_.vertices_count * config_.vbo_size());
 }
 
 void SetImage::update_index_buffer_() {
diff --git a/taichi/ui/backends/vulkan/renderables/set_image.h b/taichi/ui/backends/vulkan/renderables/set_image.h
index df2a46a91ca48..22db823c5f1aa 100644
--- a/taichi/ui/backends/vulkan/renderables/set_image.h
+++ b/taichi/ui/backends/vulkan/renderables/set_image.h
@@ -37,7 +37,7 @@ class SetImage final : public Renderable {
     float y_factor{1.0};
   };
 
-  SetImage(AppContext *app_context);
+  SetImage(AppContext *app_context, VertexAttributes vbo_attrs);
 
   void update_data(const SetImageInfo &info);
 
diff --git a/taichi/ui/backends/vulkan/renderables/triangles.cpp b/taichi/ui/backends/vulkan/renderables/triangles.cpp
index af802c0d0cadb..4db136f9b38f9 100644
--- a/taichi/ui/backends/vulkan/renderables/triangles.cpp
+++ b/taichi/ui/backends/vulkan/renderables/triangles.cpp
@@ -34,7 +34,7 @@ void Triangles::init_triangles(AppContext *app_context,
   Renderable::init_render_resources();
 }
 
-Triangles::Triangles(AppContext *app_context) {
+Triangles::Triangles(AppContext *app_context, VertexAttributes vbo_attrs) {
   init_triangles(app_context, 3, 3);
 }
 
diff --git a/taichi/ui/backends/vulkan/renderables/triangles.h b/taichi/ui/backends/vulkan/renderables/triangles.h
index f0600c72fd43a..aa7c7f5097dbf 100644
--- a/taichi/ui/backends/vulkan/renderables/triangles.h
+++ b/taichi/ui/backends/vulkan/renderables/triangles.h
@@ -27,7 +27,7 @@ namespace vulkan {
 
 class Triangles final : public Renderable {
  public:
-  Triangles(AppContext *app_context);
+  Triangles(AppContext *app_context, VertexAttributes vbo_attrs);
 
   void update_data(const TrianglesInfo &info);
 
diff --git a/taichi/ui/backends/vulkan/renderer.cpp b/taichi/ui/backends/vulkan/renderer.cpp
index d5c89785fbcee..1d937a4868c2d 100644
--- a/taichi/ui/backends/vulkan/renderer.cpp
+++ b/taichi/ui/backends/vulkan/renderer.cpp
@@ -1,4 +1,5 @@
 #include "renderer.h"
+
 #include "taichi/ui/utils/utils.h"
 
 using taichi::lang::Program;
@@ -18,18 +19,19 @@ void Renderer::init(Program *prog,
 }
 
 template <typename T>
-std::unique_ptr<Renderable> get_new_renderable(AppContext *app_context) {
-  return std::unique_ptr<Renderable>{new T(app_context)};
+std::unique_ptr<Renderable> get_new_renderable(AppContext *app_context,
+                                               VertexAttributes vbo_attrs) {
+  return std::unique_ptr<Renderable>{new T(app_context, vbo_attrs)};
 }
 
 template <typename T>
-T *Renderer::get_renderable_of_type() {
+T *Renderer::get_renderable_of_type(VertexAttributes vbo_attrs) {
   if (next_renderable_ >= renderables_.size()) {
-    renderables_.push_back(get_new_renderable<T>(&app_context_));
+    renderables_.push_back(get_new_renderable<T>(&app_context_, vbo_attrs));
   } else if (dynamic_cast<T *>(renderables_[next_renderable_].get()) ==
              nullptr) {
     renderables_.insert(renderables_.begin() + next_renderable_,
-                        get_new_renderable<T>(&app_context_));
+                        get_new_renderable<T>(&app_context_, vbo_attrs));
   }
 
   if (T *t = dynamic_cast<T *>(renderables_[next_renderable_].get())) {
@@ -43,37 +45,40 @@ void Renderer::set_background_color(const glm::vec3 &color) {
 }
 
 void Renderer::set_image(const SetImageInfo &info) {
-  SetImage *s = get_renderable_of_type<SetImage>();
+  SetImage *s = get_renderable_of_type<SetImage>(VboHelpers::all());
   s->update_data(info);
   next_renderable_ += 1;
 }
 
 void Renderer::triangles(const TrianglesInfo &info) {
-  Triangles *triangles = get_renderable_of_type<Triangles>();
+  Triangles *triangles =
+      get_renderable_of_type<Triangles>(info.renderable_info.vbo_attrs);
   triangles->update_data(info);
   next_renderable_ += 1;
 }
 
 void Renderer::lines(const LinesInfo &info) {
-  Lines *lines = get_renderable_of_type<Lines>();
+  Lines *lines = get_renderable_of_type<Lines>(info.renderable_info.vbo_attrs);
   lines->update_data(info);
   next_renderable_ += 1;
 }
 
 void Renderer::circles(const CirclesInfo &info) {
-  Circles *circles = get_renderable_of_type<Circles>();
+  Circles *circles =
+      get_renderable_of_type<Circles>(info.renderable_info.vbo_attrs);
   circles->update_data(info);
   next_renderable_ += 1;
 }
 
 void Renderer::mesh(const MeshInfo &info, Scene *scene) {
-  Mesh *mesh = get_renderable_of_type<Mesh>();
+  Mesh *mesh = get_renderable_of_type<Mesh>(info.renderable_info.vbo_attrs);
   mesh->update_data(info, *scene);
   next_renderable_ += 1;
 }
 
 void Renderer::particles(const ParticlesInfo &info, Scene *scene) {
-  Particles *particles = get_renderable_of_type<Particles>();
+  Particles *particles =
+      get_renderable_of_type<Particles>(info.renderable_info.vbo_attrs);
   particles->update_data(info, *scene);
   next_renderable_ += 1;
 }
diff --git a/taichi/ui/backends/vulkan/renderer.h b/taichi/ui/backends/vulkan/renderer.h
index 362489dd8bf18..53656592d1d86 100644
--- a/taichi/ui/backends/vulkan/renderer.h
+++ b/taichi/ui/backends/vulkan/renderer.h
@@ -12,6 +12,8 @@
 #include <array>
 #include <optional>
 #include <set>
+#include <memory>
+
 #include "taichi/ui/utils/utils.h"
 #include "taichi/ui/backends/vulkan/vertex.h"
 #include "taichi/ui/backends/vulkan/scene.h"
@@ -21,7 +23,6 @@
 #include "taichi/ui/common/canvas_base.h"
 
 #include "gui.h"
-#include <memory>
 
 #include "renderables/set_image.h"
 #include "renderables/triangles.h"
@@ -80,7 +81,7 @@ class TI_DLL_EXPORT Renderer {
   AppContext app_context_;
 
   template <typename T>
-  T *get_renderable_of_type();
+  T *get_renderable_of_type(VertexAttributes vbo_attrs);
 };
 
 }  // namespace vulkan
diff --git a/taichi/ui/backends/vulkan/vertex.cpp b/taichi/ui/backends/vulkan/vertex.cpp
new file mode 100644
index 0000000000000..95083f9013515
--- /dev/null
+++ b/taichi/ui/backends/vulkan/vertex.cpp
@@ -0,0 +1,27 @@
+#include "taichi/ui/backends/vulkan/vertex.h"
+
+#include <stddef.h>
+
+namespace taichi {
+namespace ui {
+
+// static
+size_t VboHelpers::size(VertexAttributes va) {
+  size_t res = 0;
+  if (VboHelpers::has_attr(va, VertexAttributes::kPos)) {
+    res += sizeof(Vertex::pos);
+  }
+  if (VboHelpers::has_attr(va, VertexAttributes::kNormal)) {
+    res += sizeof(Vertex::normal);
+  }
+  if (VboHelpers::has_attr(va, VertexAttributes::kUv)) {
+    res += sizeof(Vertex::tex_coord);
+  }
+  if (VboHelpers::has_attr(va, VertexAttributes::kColor)) {
+    res += sizeof(Vertex::color);
+  }
+  return res;
+}
+
+}  // namespace ui
+}  // namespace taichi
diff --git a/taichi/ui/backends/vulkan/vertex.h b/taichi/ui/backends/vulkan/vertex.h
index 7e53a9c9d5a28..663adb21f2935 100644
--- a/taichi/ui/backends/vulkan/vertex.h
+++ b/taichi/ui/backends/vulkan/vertex.h
@@ -1,5 +1,9 @@
 #pragma once
 
+#include <stddef.h>
+
+#include <type_traits>
+
 namespace taichi {
 namespace ui {
 
@@ -21,9 +25,47 @@ struct Vertex {
   };
   vec3 pos;
   vec3 normal;
-  vec2 texCoord;
+  vec2 tex_coord;
   vec4 color;
 };
 
+enum class VertexAttributes : char {
+  kPos = 0b0001,
+  kNormal = 0b0010,
+  kUv = 0b0100,
+  kColor = 0b1000,
+};
+
+constexpr inline VertexAttributes operator|(VertexAttributes src,
+                                            VertexAttributes a) {
+  using UT = std::underlying_type_t<VertexAttributes>;
+  return static_cast<VertexAttributes>(UT(src) | UT(a));
+}
+
+class VboHelpers {
+ public:
+  constexpr static VertexAttributes kOrderedAttrs[] = {
+      VertexAttributes::kPos,
+      VertexAttributes::kNormal,
+      VertexAttributes::kUv,
+      VertexAttributes::kColor,
+  };
+
+  constexpr static VertexAttributes empty() {
+    return static_cast<VertexAttributes>(0);
+  }
+  constexpr static VertexAttributes all() {
+    return VertexAttributes::kPos | VertexAttributes::kNormal |
+           VertexAttributes::kUv | VertexAttributes::kColor;
+  }
+
+  static size_t size(VertexAttributes va);
+
+  static bool has_attr(VertexAttributes src, VertexAttributes attr) {
+    using UT = std::underlying_type_t<VertexAttributes>;
+    return UT(src) & UT(attr);
+  }
+};
+
 }  // namespace ui
 }  // namespace taichi
diff --git a/taichi/ui/common/canvas_base.h b/taichi/ui/common/canvas_base.h
index 44198b9e83d31..2d224cd5b0acc 100644
--- a/taichi/ui/common/canvas_base.h
+++ b/taichi/ui/common/canvas_base.h
@@ -1,4 +1,5 @@
 #pragma once
+
 #include "taichi/ui/common/field_info.h"
 #include "taichi/ui/common/scene_base.h"
 #include "taichi/ui/common/renderable_info.h"
@@ -18,13 +19,13 @@ struct TrianglesInfo {
 struct CirclesInfo {
   RenderableInfo renderable_info;
   glm::vec3 color;
-  float radius;
+  float radius{0};
 };
 
 struct LinesInfo {
   RenderableInfo renderable_info;
   glm::vec3 color;
-  float width;
+  float width{0};
 };
 
 class CanvasBase {
diff --git a/taichi/ui/common/renderable_info.h b/taichi/ui/common/renderable_info.h
index aff474cc5e0b8..585ff40ec0afc 100644
--- a/taichi/ui/common/renderable_info.h
+++ b/taichi/ui/common/renderable_info.h
@@ -1,4 +1,6 @@
 #pragma once
+
+#include "taichi/ui/backends/vulkan/vertex.h"
 #include "taichi/ui/common/field_info.h"
 #include "taichi/ui/utils/utils.h"
 
@@ -7,7 +9,8 @@ TI_UI_NAMESPACE_BEGIN
 struct RenderableInfo {
   FieldInfo vbo;
   FieldInfo indices;
-  bool has_per_vertex_color;
+  bool has_per_vertex_color{false};
+  VertexAttributes vbo_attrs{VboHelpers::all()};
 };
 
 TI_UI_NAMESPACE_END
diff --git a/taichi/ui/common/scene_base.h b/taichi/ui/common/scene_base.h
index e7a0c46f79e63..c88efb6a0d12c 100644
--- a/taichi/ui/common/scene_base.h
+++ b/taichi/ui/common/scene_base.h
@@ -1,6 +1,7 @@
 #pragma once
 
 #include <vector>
+
 #include "taichi/ui/common/field_info.h"
 #include "taichi/ui/common/renderable_info.h"
 #include "taichi/ui/common/camera.h"
@@ -17,14 +18,14 @@ struct MeshInfo {
   RenderableInfo renderable_info;
   glm::vec3 color;
   bool two_sided{false};
-  int object_id;
+  int object_id{0};
 };
 
 struct ParticlesInfo {
   RenderableInfo renderable_info;
   glm::vec3 color;
-  float radius;
-  int object_id;
+  float radius{0};
+  int object_id{0};
 };
 
 class SceneBase {

From 816b52b1a33658b84233a221dd12a5e93bfaf21d Mon Sep 17 00:00:00 2001
From: Jiasheng Zhang <leonz5288@hotmail.com>
Date: Mon, 21 Mar 2022 16:59:19 +0800
Subject: [PATCH 4/8] [misc] Write version info right after creation of uuid
 (#4589)

* [misc] Write version info right after creation of uuid

* Auto Format

Co-authored-by: Taichi Gardener <taichigardener@gmail.com>
---
 python/taichi/_version_check.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/python/taichi/_version_check.py b/python/taichi/_version_check.py
index 432b1f27462ac..9bbec699c3335 100644
--- a/python/taichi/_version_check.py
+++ b/python/taichi/_version_check.py
@@ -91,6 +91,8 @@ def try_check_version():
                                    cur_date)
         else:
             cur_uuid = str(uuid.uuid4())
+            write_version_info({'status': 0}, cur_uuid, version_info_path,
+                               cur_date)
             response = check_version(cur_uuid)
             write_version_info(response, cur_uuid, version_info_path, cur_date)
     # Wildcard exception to catch potential file writing errors.

From 8f21da92e8e4df052ef1ba20cb96ee480674fb89 Mon Sep 17 00:00:00 2001
From: Jiasheng Zhang <leonz5288@hotmail.com>
Date: Mon, 21 Mar 2022 19:57:44 +0800
Subject: [PATCH 5/8] [ci] Release script add test for tests/python/examples
 (#4590)

---
 ci/scripts/release_test.sh | 33 +++++++++++++++++++++++++++++++++
 1 file changed, 33 insertions(+)

diff --git a/ci/scripts/release_test.sh b/ci/scripts/release_test.sh
index d64140b4f53bb..5a2e294e79e00 100644
--- a/ci/scripts/release_test.sh
+++ b/ci/scripts/release_test.sh
@@ -236,6 +236,36 @@ function taichi::test::voxel_editor {
     cd "${WORKDIR}"
 }
 
+function taichi::test::generate_videos {
+    local WORKDIR=${1}
+    local PATTERN="test_*.py"
+    local ORG="taichi-dev"
+    local REPO="taichi"
+
+    # divider
+    taichi::utils::line
+    taichi::utils::logger::info "Generating examples videos"
+    
+    # clone the repo
+    taichi::utils::git_clone "${ORG}" "${REPO}"
+    # mkdir "${REPO}/misc/output_videos"
+
+    # run tests
+    cd "${REPO}/tests/python/examples"
+    for directory in $(find ./ -mindepth 1 -maxdepth 1 -name "*" ! -name "__*" -type d); do
+        cd "${directory}"
+        for match in $(find ./ -maxdepth 1 -name "${PATTERN}" -type f); do
+            pytest -v "${match}"
+            taichi::utils::line
+            # taichi::utils::pause
+        done
+        cd ..
+    done
+
+    # go back to workdir
+    cd "${WORKDIR}"
+}
+
 function taichi::test::main {
     # set debugging flag
     DEBUG="false"
@@ -267,6 +297,9 @@ function taichi::test::main {
 
     # voxel editor tests 
     taichi::test::voxel_editor "${WORKDIR}"
+
+    # generating example videos
+    taichi::test::generate_videos "${WORKDIR}"
 }
 
 taichi::test::main

From f2f79d62d6252484b4328722779c2cd6c109143d Mon Sep 17 00:00:00 2001
From: Ailing <ailzhang@users.noreply.github.com>
Date: Mon, 21 Mar 2022 20:51:18 +0800
Subject: [PATCH 6/8] [aot] Support return in vulkan aot (#4593)

Vulkan aot return didn't work since it was trying to serialize
`Datatype` which was a raw pointer. This PR fixes the issue by changing
it to a `PrimitiveTypeID` which is a `int`. Note it's possible that we
limit general `Kernel` args and rets to be primitive types only but leaving that
to a followup PR to avoid making this PR too complicated.
---
 .../vulkan/aot_module_builder_impl.cpp        | 32 ++-----------------
 .../backends/vulkan/aot_module_builder_impl.h |  2 --
 taichi/backends/vulkan/runtime.cpp            | 22 ++++++-------
 taichi/codegen/spirv/kernel_utils.cpp         | 25 +++++++++------
 taichi/codegen/spirv/kernel_utils.h           | 10 ++++--
 taichi/codegen/spirv/spirv_codegen.cpp        | 14 ++++----
 taichi/program/context.h                      |  5 +++
 tests/cpp/aot/aot_save_load_test.cpp          | 29 ++++++++++++++++-
 8 files changed, 79 insertions(+), 60 deletions(-)

diff --git a/taichi/backends/vulkan/aot_module_builder_impl.cpp b/taichi/backends/vulkan/aot_module_builder_impl.cpp
index c57d785d7d82b..60cc5f7aa266f 100644
--- a/taichi/backends/vulkan/aot_module_builder_impl.cpp
+++ b/taichi/backends/vulkan/aot_module_builder_impl.cpp
@@ -46,12 +46,12 @@ class AotDataConverter {
     for (const auto &arg : in.ctx_attribs.args()) {
       if (!arg.is_array) {
         aot::ScalarArg scalar_arg{};
-        scalar_arg.dtype_name = arg.dt.to_string();
+        scalar_arg.dtype_name = PrimitiveType::get(arg.dtype).to_string();
         scalar_arg.offset_in_args_buf = arg.offset_in_mem;
         res.scalar_args[arg.index] = scalar_arg;
       } else {
         aot::ArrayArg arr_arg{};
-        arr_arg.dtype_name = arg.dt.to_string();
+        arr_arg.dtype_name = PrimitiveType::get(arg.dtype).to_string();
         arr_arg.field_dim = arg.field_dim;
         arr_arg.element_shape = arg.element_shape;
         arr_arg.shape_offset_in_args_buf = arg.index * sizeof(int32_t);
@@ -105,32 +105,6 @@ AotModuleBuilderImpl::AotModuleBuilderImpl(
   }
 }
 
-uint32_t AotModuleBuilderImpl::to_vk_dtype_enum(DataType dt) {
-  if (dt == PrimitiveType::u64) {
-    return 0;
-  } else if (dt == PrimitiveType::i64) {
-    return 1;
-  } else if (dt == PrimitiveType::u32) {
-    return 2;
-  } else if (dt == PrimitiveType::i32) {
-    return 3;
-  } else if (dt == PrimitiveType::u16) {
-    return 4;
-  } else if (dt == PrimitiveType::i16) {
-    return 5;
-  } else if (dt == PrimitiveType::u8) {
-    return 6;
-  } else if (dt == PrimitiveType::i8) {
-    return 7;
-  } else if (dt == PrimitiveType::f64) {
-    return 8;
-  } else if (dt == PrimitiveType::f32) {
-    return 9;
-  } else {
-    TI_NOT_IMPLEMENTED
-  }
-}
-
 std::string AotModuleBuilderImpl::write_spv_file(
     const std::string &output_dir,
     const TaskAttributes &k,
@@ -194,7 +168,7 @@ void AotModuleBuilderImpl::add_field_per_backend(const std::string &identifier,
   aot::CompiledFieldData field_data;
   field_data.field_name = identifier;
   field_data.is_scalar = is_scalar;
-  field_data.dtype = to_vk_dtype_enum(dt);
+  field_data.dtype = static_cast<int>(dt->cast<PrimitiveType>()->type);
   field_data.dtype_name = dt.to_string();
   field_data.shape = shape;
   field_data.mem_offset_in_parent = dense_desc.mem_offset_in_parent_cell;
diff --git a/taichi/backends/vulkan/aot_module_builder_impl.h b/taichi/backends/vulkan/aot_module_builder_impl.h
index 408e24b1f3dc8..bbd6b40e4df48 100644
--- a/taichi/backends/vulkan/aot_module_builder_impl.h
+++ b/taichi/backends/vulkan/aot_module_builder_impl.h
@@ -40,8 +40,6 @@ class AotModuleBuilderImpl : public AotModuleBuilder {
                              const TaskAttributes &k,
                              const std::vector<uint32_t> &source_code) const;
 
-  uint32_t to_vk_dtype_enum(DataType dt);
-
   const std::vector<CompiledSNodeStructs> &compiled_structs_;
   TaichiAotData ti_aot_data_;
   std::unique_ptr<Device> aot_target_device_;
diff --git a/taichi/backends/vulkan/runtime.cpp b/taichi/backends/vulkan/runtime.cpp
index 410c40f552151..e500e8921bc87 100644
--- a/taichi/backends/vulkan/runtime.cpp
+++ b/taichi/backends/vulkan/runtime.cpp
@@ -67,16 +67,15 @@ class HostDeviceContextBlitter {
     char *const device_base =
         reinterpret_cast<char *>(device_->map(*device_args_buffer_));
 
-#define TO_DEVICE(short_type, type)                    \
-  if (dt->is_primitive(PrimitiveTypeID::short_type)) { \
-    auto d = host_ctx_->get_arg<type>(i);              \
-    reinterpret_cast<type *>(device_ptr)[0] = d;       \
-    break;                                             \
+#define TO_DEVICE(short_type, type)               \
+  if (arg.dtype == PrimitiveTypeID::short_type) { \
+    auto d = host_ctx_->get_arg<type>(i);         \
+    reinterpret_cast<type *>(device_ptr)[0] = d;  \
+    break;                                        \
   }
 
     for (int i = 0; i < ctx_attribs_->args().size(); ++i) {
       const auto &arg = ctx_attribs_->args()[i];
-      const auto dt = arg.dt;
       char *device_ptr = device_base + arg.offset_in_mem;
       do {
         if (arg.is_array) {
@@ -118,13 +117,14 @@ class HostDeviceContextBlitter {
           TO_DEVICE(f64, float64)
         }
         if (device_->get_cap(DeviceCapability::spirv_has_float16)) {
-          if (dt->is_primitive(PrimitiveTypeID::f16)) {
+          if (arg.dtype == PrimitiveTypeID::f16) {
             auto d = fp16_ieee_from_fp32_value(host_ctx_->get_arg<float>(i));
             reinterpret_cast<uint16 *>(device_ptr)[0] = d;
             break;
           }
         }
-        TI_ERROR("Vulkan does not support arg type={}", data_type_name(arg.dt));
+        TI_ERROR("Vulkan does not support arg type={}",
+                 PrimitiveType::get(arg.dtype).to_string());
       } while (0);
     }
 
@@ -196,8 +196,8 @@ class HostDeviceContextBlitter {
       // *arg* on the host context.
       const auto &ret = ctx_attribs_->rets()[i];
       char *device_ptr = device_base + ret.offset_in_mem;
-      const auto dt = ret.dt;
-      const auto num = ret.stride / data_type_size(ret.dt);
+      const auto dt = PrimitiveType::get(ret.dtype);
+      const auto num = ret.stride / data_type_size(dt);
       for (int j = 0; j < num; ++j) {
         if (device_->get_cap(DeviceCapability::spirv_has_int8)) {
           TO_HOST(i8, int8, j)
@@ -227,7 +227,7 @@ class HostDeviceContextBlitter {
           }
         }
         TI_ERROR("Vulkan does not support return value type={}",
-                 data_type_name(ret.dt));
+                 data_type_name(PrimitiveType::get(ret.dtype)));
       }
     }
 #undef TO_HOST
diff --git a/taichi/codegen/spirv/kernel_utils.cpp b/taichi/codegen/spirv/kernel_utils.cpp
index b29e03b460372..759061a6823cf 100644
--- a/taichi/codegen/spirv/kernel_utils.cpp
+++ b/taichi/codegen/spirv/kernel_utils.cpp
@@ -53,10 +53,13 @@ KernelContextAttributes::KernelContextAttributes(const Kernel &kernel)
       rets_bytes_(0),
       extra_args_bytes_(RuntimeContext::extra_args_size) {
   arg_attribs_vec_.reserve(kernel.args.size());
+  // TODO: We should be able to limit Kernel args and rets to be primitive types
+  // as well but let's leave that as a followup up PR.
   for (const auto &ka : kernel.args) {
     ArgAttributes aa;
-    aa.dt = ka.dt;
-    const size_t dt_bytes = data_type_size(aa.dt);
+    TI_ASSERT(ka.dt->is<PrimitiveType>());
+    aa.dtype = ka.dt->cast<PrimitiveType>()->type;
+    const size_t dt_bytes = data_type_size(ka.dt);
     aa.is_array = ka.is_array;
     if (aa.is_array) {
       aa.field_dim = ka.total_dim - ka.element_shape.size();
@@ -70,13 +73,16 @@ KernelContextAttributes::KernelContextAttributes(const Kernel &kernel)
     RetAttributes ra;
     size_t dt_bytes{0};
     if (auto tensor_type = kr.dt->cast<TensorType>()) {
-      ra.dt = tensor_type->get_element_type();
-      dt_bytes = data_type_size(ra.dt);
+      auto tensor_dtype = tensor_type->get_element_type();
+      TI_ASSERT(tensor_dtype->is<PrimitiveType>());
+      ra.dtype = tensor_dtype->cast<PrimitiveType>()->type;
+      dt_bytes = data_type_size(tensor_dtype);
       ra.is_array = true;
       ra.stride = tensor_type->get_num_elements() * dt_bytes;
     } else {
-      ra.dt = kr.dt;
-      dt_bytes = data_type_size(ra.dt);
+      TI_ASSERT(kr.dt->is<PrimitiveType>());
+      ra.dtype = kr.dt->cast<PrimitiveType>()->type;
+      dt_bytes = data_type_size(kr.dt);
       ra.is_array = false;
       ra.stride = dt_bytes;
     }
@@ -88,9 +94,10 @@ KernelContextAttributes::KernelContextAttributes(const Kernel &kernel)
     size_t bytes = offset;
     for (int i = 0; i < vec->size(); ++i) {
       auto &attribs = (*vec)[i];
-      const size_t dt_bytes = (attribs.is_array && !is_ret)
-                                  ? sizeof(uint64_t)
-                                  : data_type_size(attribs.dt);
+      const size_t dt_bytes =
+          (attribs.is_array && !is_ret)
+              ? sizeof(uint64_t)
+              : data_type_size(PrimitiveType::get(attribs.dtype));
       // Align bytes to the nearest multiple of dt_bytes
       bytes = (bytes + dt_bytes - 1) / dt_bytes * dt_bytes;
       attribs.offset_in_mem = bytes;
diff --git a/taichi/codegen/spirv/kernel_utils.h b/taichi/codegen/spirv/kernel_utils.h
index 21d528dba12c3..95a8aa71196e3 100644
--- a/taichi/codegen/spirv/kernel_utils.h
+++ b/taichi/codegen/spirv/kernel_utils.h
@@ -140,12 +140,18 @@ class KernelContextAttributes {
     size_t offset_in_mem{0};
     // Index of the input arg or the return value in the host `Context`
     int index{-1};
-    DataType dt;
+    PrimitiveTypeID dtype{PrimitiveTypeID::unknown};
     bool is_array{false};
     std::vector<int> element_shape;
     std::size_t field_dim{0};
 
-    TI_IO_DEF(stride, offset_in_mem, index, is_array, element_shape, field_dim);
+    TI_IO_DEF(stride,
+              offset_in_mem,
+              index,
+              dtype,
+              is_array,
+              element_shape,
+              field_dim);
   };
 
  public:
diff --git a/taichi/codegen/spirv/spirv_codegen.cpp b/taichi/codegen/spirv/spirv_codegen.cpp
index 821ccb12b91f6..4c5adb7faed85 100644
--- a/taichi/codegen/spirv/spirv_codegen.cpp
+++ b/taichi/codegen/spirv/spirv_codegen.cpp
@@ -502,7 +502,7 @@ class TaskCodegen : public IRVisitor {
       //    ir_->int_immediate_number(ir_->i32_type(), offset_in_mem);
       // ir_->register_value(stmt->raw_name(), val);
     } else {
-      const auto dt = arg_attribs.dt;
+      const auto dt = PrimitiveType::get(arg_attribs.dtype);
       const auto val_type = ir_->get_primitive_type(dt);
       spirv::Value buffer_val = ir_->make_value(
           spv::OpAccessChain,
@@ -1806,9 +1806,9 @@ class TaskCodegen : public IRVisitor {
                                         "arg_ptr" + std::to_string(arg.index),
                                         arg.offset_in_mem);
       } else {
-        struct_components_.emplace_back(ir_->get_primitive_type(arg.dt),
-                                        "arg" + std::to_string(arg.index),
-                                        arg.offset_in_mem);
+        struct_components_.emplace_back(
+            ir_->get_primitive_type(PrimitiveType::get(arg.dtype)),
+            "arg" + std::to_string(arg.index), arg.offset_in_mem);
       }
     }
     // A compromise for use in constants buffer
@@ -1833,7 +1833,8 @@ class TaskCodegen : public IRVisitor {
     // Now we only have one ret
     TI_ASSERT(ctx_attribs_->rets().size() == 1);
     for (auto &ret : ctx_attribs_->rets()) {
-      if (auto tensor_type = ret.dt->cast<TensorType>()) {
+      if (auto tensor_type =
+              PrimitiveType::get(ret.dtype)->cast<TensorType>()) {
         struct_components_.emplace_back(
             ir_->get_array_type(
                 ir_->get_primitive_type(tensor_type->get_element_type()),
@@ -1841,7 +1842,8 @@ class TaskCodegen : public IRVisitor {
             "ret" + std::to_string(ret.index), ret.offset_in_mem);
       } else {
         struct_components_.emplace_back(
-            ir_->get_array_type(ir_->get_primitive_type(ret.dt), 1),
+            ir_->get_array_type(
+                ir_->get_primitive_type(PrimitiveType::get(ret.dtype)), 1),
             "ret" + std::to_string(ret.index), ret.offset_in_mem);
       }
     }
diff --git a/taichi/program/context.h b/taichi/program/context.h
index 0d6d7eef413a6..4d7562054f53a 100644
--- a/taichi/program/context.h
+++ b/taichi/program/context.h
@@ -50,6 +50,11 @@ struct RuntimeContext {
   void set_device_allocation(int i, bool is_device_allocation_) {
     is_device_allocation[i] = is_device_allocation_;
   }
+
+  template <typename T>
+  T get_ret(int i) {
+    return taichi_union_cast_with_different_sizes<T>(result_buffer[i]);
+  }
 #endif
 };
 
diff --git a/tests/cpp/aot/aot_save_load_test.cpp b/tests/cpp/aot/aot_save_load_test.cpp
index 6b66213121e00..798af91e54926 100644
--- a/tests/cpp/aot/aot_save_load_test.cpp
+++ b/tests/cpp/aot/aot_save_load_test.cpp
@@ -1,6 +1,7 @@
 #include "gtest/gtest.h"
 #include "taichi/ir/ir_builder.h"
 #include "taichi/ir/statements.h"
+#include "taichi/inc/constants.h"
 #include "taichi/program/program.h"
 #ifdef TI_WITH_VULKAN
 #include "taichi/backends/vulkan/aot_module_loader_impl.h"
@@ -29,7 +30,24 @@ using namespace lang;
 
   auto aot_builder = program.make_aot_module_builder(Arch::vulkan);
 
-  std::unique_ptr<Kernel> kernel_init, kernel_ret;
+  std::unique_ptr<Kernel> kernel_init, kernel_ret, kernel_simple_ret;
+
+  {
+    /*
+    @ti.kernel
+    def ret() -> ti.f32:
+      sum = 0.2
+      return sum
+    */
+    IRBuilder builder;
+    auto *sum = builder.create_local_var(PrimitiveType::f32);
+    builder.create_local_store(sum, builder.get_float32(0.2));
+    builder.create_return(builder.create_local_load(sum));
+
+    kernel_simple_ret =
+        std::make_unique<Kernel>(program, builder.extract_ir(), "simple_ret");
+    kernel_simple_ret->insert_ret(PrimitiveType::f32);
+  }
 
   {
     /*
@@ -79,6 +97,7 @@ using namespace lang;
     kernel_ret->insert_ret(PrimitiveType::i32);
   }
 
+  aot_builder->add("simple_ret", kernel_simple_ret.get());
   aot_builder->add_field("place", place, true, place->dt, {n}, 1, 1);
   aot_builder->add("init", kernel_init.get());
   aot_builder->add("ret", kernel_ret.get());
@@ -103,6 +122,7 @@ TEST(AotSaveLoad, Vulkan) {
       std::make_unique<taichi::lang::MemoryPool>(Arch::vulkan, nullptr);
   result_buffer = (taichi::uint64 *)memory_pool->allocate(
       sizeof(taichi::uint64) * taichi_result_buffer_entries, 8);
+  host_ctx.result_buffer = result_buffer;
 
   // Create Taichi Device for computation
   lang::vulkan::VulkanDeviceCreator::Params evd_params;
@@ -132,6 +152,13 @@ TEST(AotSaveLoad, Vulkan) {
   EXPECT_EQ(root_size, 64);
   vulkan_runtime->add_root_buffer(root_size);
 
+  auto simple_ret_kernel = vk_module->get_kernel("simple_ret");
+  EXPECT_TRUE(simple_ret_kernel);
+
+  simple_ret_kernel->launch(&host_ctx);
+  vulkan_runtime->synchronize();
+  EXPECT_FLOAT_EQ(host_ctx.get_ret<float>(0), 0.2);
+
   auto init_kernel = vk_module->get_kernel("init");
   EXPECT_TRUE(init_kernel);
 

From 673f74c14700d7b37e00aed3684a52e79af92957 Mon Sep 17 00:00:00 2001
From: PGZXB <420254146@qq.com>
Date: Tue, 22 Mar 2022 10:57:05 +0800
Subject: [PATCH 7/8] [bug] Fix bug that caching kernels with same AST will
 fail (#4582)

* Fix bug that caching same-ast kernels will fail

* Add {}

Co-authored-by: Ye Kuang <k-ye@users.noreply.github.com>

* Auto Format

* Fix bugs

Co-authored-by: Ye Kuang <k-ye@users.noreply.github.com>
Co-authored-by: Taichi Gardener <taichigardener@gmail.com>
---
 taichi/llvm/llvm_program.cpp | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/taichi/llvm/llvm_program.cpp b/taichi/llvm/llvm_program.cpp
index 027a7ebba5290..c79cf1914a83d 100644
--- a/taichi/llvm/llvm_program.cpp
+++ b/taichi/llvm/llvm_program.cpp
@@ -714,7 +714,9 @@ void LlvmProgramImpl::cache_kernel(
     const std::string &kernel_key,
     llvm::Module *module,
     std::vector<std::string> &&offloaded_task_name_list) {
-  TI_ASSERT(cache_data_.kernels.find(kernel_key) == cache_data_.kernels.end());
+  if (cache_data_.kernels.find(kernel_key) != cache_data_.kernels.end()) {
+    return;
+  }
   auto &kernel_cache = cache_data_.kernels[kernel_key];
   kernel_cache.kernel_key = kernel_key;
   kernel_cache.owned_module = llvm::CloneModule(*module);

From 3ab838d9e5057918ff58830f1ccd9b360ac3e395 Mon Sep 17 00:00:00 2001
From: Bo Qiao <boqiao@taichi.graphics>
Date: Tue, 22 Mar 2022 14:41:04 +0800
Subject: [PATCH 8/8] [lang] Fix vector matrix ndarray to numpy layout (#4597)

* [lang] Fix matrix ndarray to numpy

* Format

* Use static

* Use parametrize
---
 python/taichi/_kernels.py      | 22 ++++++++++++++++++----
 python/taichi/lang/_ndarray.py | 10 ++++++----
 python/taichi/lang/matrix.py   |  8 ++++----
 tests/python/test_ndarray.py   | 28 ++++++++++++++++++++++++++++
 4 files changed, 56 insertions(+), 12 deletions(-)

diff --git a/python/taichi/_kernels.py b/python/taichi/_kernels.py
index afbb8dc68f021..d202aa85b86a0 100644
--- a/python/taichi/_kernels.py
+++ b/python/taichi/_kernels.py
@@ -44,14 +44,21 @@ def ndarray_to_ext_arr(ndarray: any_arr(), arr: ext_arr()):
 
 @kernel
 def ndarray_matrix_to_ext_arr(ndarray: any_arr(), arr: ext_arr(),
+                              layout_is_aos: template(),
                               as_vector: template()):
     for I in grouped(ndarray):
         for p in static(range(ndarray[I].n)):
             for q in static(range(ndarray[I].m)):
                 if static(as_vector):
-                    arr[I, p] = ndarray[I][p]
+                    if static(layout_is_aos):
+                        arr[I, p] = ndarray[I][p]
+                    else:
+                        arr[p, I] = ndarray[I][p]
                 else:
-                    arr[I, p, q] = ndarray[I][p, q]
+                    if static(layout_is_aos):
+                        arr[I, p, q] = ndarray[I][p, q]
+                    else:
+                        arr[p, q, I] = ndarray[I][p, q]
 
 
 @kernel
@@ -124,14 +131,21 @@ def ext_arr_to_ndarray(arr: ext_arr(), ndarray: any_arr()):
 
 @kernel
 def ext_arr_to_ndarray_matrix(arr: ext_arr(), ndarray: any_arr(),
+                              layout_is_aos: template(),
                               as_vector: template()):
     for I in grouped(ndarray):
         for p in static(range(ndarray[I].n)):
             for q in static(range(ndarray[I].m)):
                 if static(as_vector):
-                    ndarray[I][p] = arr[I, p]
+                    if static(layout_is_aos):
+                        ndarray[I][p] = arr[I, p]
+                    else:
+                        ndarray[I][p] = arr[p, I]
                 else:
-                    ndarray[I][p, q] = arr[I, p, q]
+                    if static(layout_is_aos):
+                        ndarray[I][p, q] = arr[I, p, q]
+                    else:
+                        ndarray[I][p, q] = arr[p, q, I]
 
 
 @kernel
diff --git a/python/taichi/lang/_ndarray.py b/python/taichi/lang/_ndarray.py
index a1bd856dabf39..9c04d03a68e87 100644
--- a/python/taichi/lang/_ndarray.py
+++ b/python/taichi/lang/_ndarray.py
@@ -90,7 +90,7 @@ def _ndarray_to_numpy(self):
         impl.get_runtime().sync()
         return arr
 
-    def _ndarray_matrix_to_numpy(self, as_vector):
+    def _ndarray_matrix_to_numpy(self, layout, as_vector):
         """Converts matrix ndarray to a numpy array.
 
         Returns:
@@ -99,7 +99,8 @@ def _ndarray_matrix_to_numpy(self, as_vector):
         arr = np.zeros(shape=self.arr.shape, dtype=to_numpy_type(self.dtype))
         from taichi._kernels import \
             ndarray_matrix_to_ext_arr  # pylint: disable=C0415
-        ndarray_matrix_to_ext_arr(self, arr, as_vector)
+        layout_is_aos = 1 if layout == Layout.AOS else 0
+        ndarray_matrix_to_ext_arr(self, arr, layout_is_aos, as_vector)
         impl.get_runtime().sync()
         return arr
 
@@ -122,7 +123,7 @@ def _ndarray_from_numpy(self, arr):
         ext_arr_to_ndarray(arr, self)
         impl.get_runtime().sync()
 
-    def _ndarray_matrix_from_numpy(self, arr, as_vector):
+    def _ndarray_matrix_from_numpy(self, arr, layout, as_vector):
         """Loads all values from a numpy array.
 
         Args:
@@ -139,7 +140,8 @@ def _ndarray_matrix_from_numpy(self, arr, as_vector):
 
         from taichi._kernels import \
             ext_arr_to_ndarray_matrix  # pylint: disable=C0415
-        ext_arr_to_ndarray_matrix(arr, self, as_vector)
+        layout_is_aos = 1 if layout == Layout.AOS else 0
+        ext_arr_to_ndarray_matrix(arr, self, layout_is_aos, as_vector)
         impl.get_runtime().sync()
 
     @python_scope
diff --git a/python/taichi/lang/matrix.py b/python/taichi/lang/matrix.py
index 7f442afa0001e..40549b64d17db 100644
--- a/python/taichi/lang/matrix.py
+++ b/python/taichi/lang/matrix.py
@@ -1416,11 +1416,11 @@ def __getitem__(self, key):
 
     @python_scope
     def to_numpy(self):
-        return self._ndarray_matrix_to_numpy(as_vector=0)
+        return self._ndarray_matrix_to_numpy(self.layout, as_vector=0)
 
     @python_scope
     def from_numpy(self, arr):
-        self._ndarray_matrix_from_numpy(arr, as_vector=0)
+        self._ndarray_matrix_from_numpy(arr, self.layout, as_vector=0)
 
     def __deepcopy__(self, memo=None):
         ret_arr = MatrixNdarray(self.n, self.m, self.dtype, self.shape,
@@ -1474,11 +1474,11 @@ def __getitem__(self, key):
 
     @python_scope
     def to_numpy(self):
-        return self._ndarray_matrix_to_numpy(as_vector=1)
+        return self._ndarray_matrix_to_numpy(self.layout, as_vector=1)
 
     @python_scope
     def from_numpy(self, arr):
-        self._ndarray_matrix_from_numpy(arr, as_vector=1)
+        self._ndarray_matrix_from_numpy(arr, self.layout, as_vector=1)
 
     def __deepcopy__(self, memo=None):
         ret_arr = VectorNdarray(self.n, self.dtype, self.shape, self.layout)
diff --git a/tests/python/test_ndarray.py b/tests/python/test_ndarray.py
index 4cb6ebd70b3d2..50419e17b56c3 100644
--- a/tests/python/test_ndarray.py
+++ b/tests/python/test_ndarray.py
@@ -335,6 +335,34 @@ def test_ndarray_numpy_io():
     _test_ndarray_numpy_io()
 
 
+def _test_ndarray_matrix_numpy_io(layout):
+    n = 5
+    m = 2
+
+    x = ti.Vector.ndarray(n, ti.i32, (m, ), layout)
+    if layout == ti.Layout.AOS:
+        x_np = 1 + np.arange(n * m).reshape(m, n).astype(np.int32)
+    else:
+        x_np = 1 + np.arange(n * m).reshape(n, m).astype(np.int32)
+    x.from_numpy(x_np)
+    assert (x_np.flatten() == x.to_numpy().flatten()).all()
+
+    k = 2
+    x = ti.Matrix.ndarray(m, k, ti.i32, n, layout)
+    if layout == ti.Layout.AOS:
+        x_np = 1 + np.arange(m * k * n).reshape(n, m, k).astype(np.int32)
+    else:
+        x_np = 1 + np.arange(m * k * n).reshape(m, k, n).astype(np.int32)
+    x.from_numpy(x_np)
+    assert (x_np.flatten() == x.to_numpy().flatten()).all()
+
+
+@pytest.mark.parametrize('layout', layouts)
+@test_utils.test(arch=supported_archs_taichi_ndarray)
+def test_ndarray_matrix_numpy_io(layout):
+    _test_ndarray_matrix_numpy_io(layout)
+
+
 def _test_matrix_ndarray_python_scope(layout):
     a = ti.Matrix.ndarray(2, 2, ti.i32, 5, layout=layout)
     for i in range(5):