[OpenGL] Add mem_offset_in_parent to serialized file in AOT.

taichi/backends/opengl/aot_data.h

@@ -27,6 +27,7 @@ struct CompiledFieldData {
   std::string field_name;
   uint32_t dtype;
   std::string dtype_name;
+  size_t mem_offset_in_parent{0};
   std::vector<int> shape;
   bool is_scalar{false};
   int row_num{0};
@@ -35,6 +36,7 @@ struct CompiledFieldData {
   TI_IO_DEF(field_name,
             dtype,
             dtype_name,
+            mem_offset_in_parent,
             shape,
             is_scalar,
             row_num,

taichi/backends/opengl/aot_module_builder_impl.cpp

@@ -65,6 +65,16 @@ void AotModuleBuilderImpl::add_per_backend(const std::string &identifier,
   aot_data_.kernels.push_back({compiled, identifier});
 }
 
+size_t AotModuleBuilderImpl::get_snode_base_address(const SNode *snode) {
+  if (snode->type == SNodeType::root)
+    return 0;
+  int chid = find_children_id(snode);
+  const auto &parent_meta =
+      compiled_structs_.snode_map.at(snode->parent->node_type_name);
+  auto choff = parent_meta.children_offsets[chid];
+  return choff + get_snode_base_address(snode->parent);
+}
+
 void AotModuleBuilderImpl::add_field_per_backend(const std::string &identifier,
                                                  const SNode *rep_snode,
                                                  bool is_scalar,
@@ -98,7 +108,15 @@ void AotModuleBuilderImpl::add_field_per_backend(const std::string &identifier,
     TI_NOT_IMPLEMENTED
   }
 
-  aot_data_.fields.push_back({identifier, gl_dtype_enum, dt.to_string(), shape,
+  // Note that currently we only support adding dense fields in AOT for all
+  // backends. In opengl backend we only error out when a non dense field is
+  // added to the aot module, but in metal backend we error out earlier when
+  // constructing aot module. Ideally we will unify this behavior but it doesn't
+  // matter too much for now.
+  TI_ERROR_IF(!all_fields_are_dense_in_container(rep_snode->parent),
+              "AOT: only supports dense field");
+  aot_data_.fields.push_back({identifier, gl_dtype_enum, dt.to_string(),
+                              get_snode_base_address(rep_snode), shape,
                               is_scalar, row_num, column_num});
 }
 

taichi/backends/opengl/aot_module_builder_impl.h

@@ -33,8 +33,9 @@ class AotModuleBuilderImpl : public AotModuleBuilder {
                             Kernel *kernel) override;
 
  private:
-  StructCompiledResult &compiled_structs_;
+  size_t get_snode_base_address(const SNode *snode);
 
+  StructCompiledResult &compiled_structs_;
   AotData aot_data_;
   bool allow_nv_shader_extension_ = false;
 };

taichi/backends/opengl/opengl_kernel_launcher.h

@@ -2,6 +2,7 @@
 
 #include "taichi/lang_util.h"
 #include "taichi/backends/device.h"
+#include "taichi/ir/snode.h"
 
 #include <vector>
 
@@ -30,6 +31,7 @@ struct OpenGlRuntime {
 using SNodeId = std::string;
 
 struct SNodeInfo {
+  const SNode *snode;
   size_t stride;
   size_t length;
   std::vector<size_t> children_offsets;

taichi/backends/opengl/struct_opengl.cpp

@@ -37,7 +37,8 @@ void OpenglStructCompiler::generate_types(const SNode &snode) {
   const auto &node_name = snode.node_type_name;
   const auto child_name = node_name + "_ch";
   auto &snode_info = snode_map_[node_name];
-  auto &snode_child_info = snode_map_[child_name];
+  snode_info.snode = &snode;
+  SNodeInfo snode_child_info;
   if (!is_place) {
     size_t stride_num = 0;
     snode_info.children_offsets.resize(snode.ch.size());

taichi/program/aot_module_builder.cpp

@@ -31,5 +31,31 @@ void AotModuleBuilder::add_kernel_template(const std::string &identifier,
   add_per_backend_tmpl(identifier, key, kernel);
 }
 
+bool AotModuleBuilder::all_fields_are_dense_in_container(
+    const SNode *container) {
+  for (const auto &ch : container->ch) {
+    if (ch->type != SNodeType::place) {
+      return false;
+    }
+  }
+  const auto *parent = container->parent;
+  if (!parent) {
+    return false;
+  }
+  if (parent->type != SNodeType::root) {
+    return false;
+  }
+  return true;
+}
+
+int AotModuleBuilder::find_children_id(const SNode *snode) {
+  auto parent = snode->parent;
+  for (int i = 0; i < parent->ch.size(); i++) {
+    if (parent->ch[i].get() == snode)
+      return i;
+  }
+  TI_ERROR("Child not found in parent!");
+}
+
 }  // namespace lang
 }  // namespace taichi

taichi/program/aot_module_builder.h

@@ -48,6 +48,10 @@ class AotModuleBuilder {
   virtual void add_per_backend_tmpl(const std::string &identifier,
                                     const std::string &key,
                                     Kernel *kernel) = 0;
+
+  bool all_fields_are_dense_in_container(const SNode *container);
+
+  int find_children_id(const SNode *snode);
 };
 
 }  // namespace lang

tests/python/test_aot.py

@@ -2,6 +2,8 @@
 import os
 import tempfile
 
+import pytest
+
 import taichi as ti
 
 
@@ -55,3 +57,18 @@ def foo(n: ti.template()):
         with open(os.path.join(tmpdir,
                                f'{filename}_metadata.json')) as json_file:
             json.load(json_file)
+
+
+@ti.test(arch=ti.opengl)
+def test_non_dense_snode():
+    n = 8
+    x = ti.field(dtype=ti.f32)
+    y = ti.field(dtype=ti.f32)
+    blk = ti.root.dense(ti.i, n)
+    blk.place(x)
+    blk.dense(ti.i, n).place(y)
+
+    with pytest.raises(RuntimeError, match='AOT: only supports dense field'):
+        m = ti.aot.Module(ti.opengl)
+        m.add_field('x', x)
+        m.add_field('y', y)