[llvm] [aot] Add unit tests for Dynamic SNodes with LLVM AOT

tests/cpp/aot/llvm/dynamic_aot_test.cpp

@@ -0,0 +1,115 @@
+#include "gtest/gtest.h"
+
+#include "taichi/program/kernel_profiler.h"
+#include "taichi/runtime/llvm/llvm_runtime_executor.h"
+#include "taichi/system/memory_pool.h"
+#include "taichi/runtime/cpu/aot_module_loader_impl.h"
+#include "taichi/runtime/cuda/aot_module_loader_impl.h"
+#include "taichi/runtime/llvm/llvm_aot_module_loader.h"
+#include "taichi/rhi/cuda/cuda_driver.h"
+#include "taichi/platform/cuda/detect_cuda.h"
+
+#define TI_RUNTIME_HOST
+#include "taichi/program/context.h"
+#undef TI_RUNTIME_HOST
+
+namespace taichi {
+namespace lang {
+
+static void run_dynamic_tests(aot::Module *mod,
+                              LlvmRuntimeExecutor *exec,
+                              uint64 *result_buffer) {
+  aot::Kernel *k_activate = mod->get_kernel("activate");
+  aot::Kernel *k_check_value_0 = mod->get_kernel("check_value_0");
+  aot::Kernel *k_deactivate = mod->get_kernel("deactivate");
+  aot::Kernel *k_check_value_1 = mod->get_kernel("check_value_1");
+
+  // Initialize SNodeTree
+  aot::Field *snode_tree_0 = mod->get_snode_tree("0" /*snode_tree_id*/);
+  allocate_aot_snode_tree_type(mod, snode_tree_0, result_buffer);
+
+  /* -------- Test Case 1 ------ */
+  // Kernel: activate()
+  {
+    RuntimeContext ctx;
+    ctx.runtime = exec->get_llvm_runtime();
+    k_activate->launch(&ctx);
+  }
+
+  // Kernel: check_value_0()
+  {
+    RuntimeContext ctx;
+    ctx.runtime = exec->get_llvm_runtime();
+    k_check_value_0->launch(&ctx);
+  }
+
+  /* -------- Test Case 2 ------ */
+  // Kernel: deactivate()
+  {
+    RuntimeContext ctx;
+    ctx.runtime = exec->get_llvm_runtime();
+    k_deactivate->launch(&ctx);
+  }
+  // Kernel: check_value_1()
+  {
+    RuntimeContext ctx;
+    ctx.runtime = exec->get_llvm_runtime();
+    k_check_value_1->launch(&ctx);
+  }
+
+  // Check assertion error from ti.kernel
+  exec->check_runtime_error(result_buffer);
+}
+
+TEST(LlvmAotTest, CpuDynamic) {
+  CompileConfig cfg;
+  cfg.arch = Arch::x64;
+  cfg.kernel_profiler = false;
+  constexpr KernelProfilerBase *kNoProfiler = nullptr;
+  LlvmRuntimeExecutor exec{cfg, kNoProfiler};
+  auto *compute_device = exec.get_compute_device();
+
+  // Must have handled all the arch fallback logic by this point.
+  auto memory_pool = std::make_unique<MemoryPool>(cfg.arch, compute_device);
+  uint64 *result_buffer{nullptr};
+  exec.materialize_runtime(memory_pool.get(), kNoProfiler, &result_buffer);
+
+  cpu::AotModuleParams aot_params;
+  const auto folder_dir = getenv("TAICHI_AOT_FOLDER_PATH");
+
+  std::stringstream aot_mod_ss;
+  aot_mod_ss << folder_dir;
+  aot_params.module_path = aot_mod_ss.str();
+  aot_params.executor_ = &exec;
+  std::unique_ptr<aot::Module> mod = cpu::make_aot_module(aot_params);
+
+  run_dynamic_tests(mod.get(), &exec, result_buffer);
+}
+
+TEST(LlvmAotTest, CudaDynamic) {
+  if (is_cuda_api_available()) {
+    CompileConfig cfg;
+    cfg.arch = Arch::cuda;
+    cfg.kernel_profiler = false;
+    constexpr KernelProfilerBase *kNoProfiler = nullptr;
+    LlvmRuntimeExecutor exec{cfg, kNoProfiler};
+
+    // Must have handled all the arch fallback logic by this point.
+    uint64 *result_buffer{nullptr};
+    exec.materialize_runtime(nullptr, kNoProfiler, &result_buffer);
+
+    cuda::AotModuleParams aot_params;
+    const auto folder_dir = getenv("TAICHI_AOT_FOLDER_PATH");
+
+    std::stringstream aot_mod_ss;
+    aot_mod_ss << folder_dir;
+    aot_params.module_path = aot_mod_ss.str();
+    aot_params.executor_ = &exec;
+    auto mod = cuda::make_aot_module(aot_params);
+
+    run_dynamic_tests(mod.get(), &exec, result_buffer);
+  }
+}
+
+}  // namespace lang
+}  // namespace taichi

tests/cpp/aot/llvm/dynamic_aot_test.py

@@ -0,0 +1,17 @@
+import argparse
+
+from utils import compile_dynamic_aot
+
+import taichi as ti
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--arch", type=str)
+args = parser.parse_args()
+
+if __name__ == "__main__":
+    if args.arch == "cpu":
+        compile_dynamic_aot(arch=ti.cpu)
+    elif args.arch == "cuda":
+        compile_dynamic_aot(arch=ti.cuda)
+    else:
+        assert False

tests/cpp/aot/llvm/utils.py

@@ -259,3 +259,73 @@ def check_value_1():
     m.add_field("x", x)
 
     m.save(dir_name, 'whatever')
+
+
+def compile_dynamic_aot(arch):
+    ti.init(arch=arch)
+
+    x = ti.field(ti.i32)
+    block = ti.root.dense(ti.i, 5)
+    pixel = block.dynamic(ti.j, 5)
+    pixel.place(x)
+
+    @ti.kernel
+    def activate():
+        ti.append(x.parent(), 0, 1)
+        ti.append(x.parent(), 0, 3)
+        ti.append(x.parent(), 0, 5)
+        ti.append(x.parent(), 1, 7)
+        ti.append(x.parent(), 1, 7)
+        ti.append(x.parent(), 2, 9)
+        ti.append(x.parent(), 3, 12)
+
+    @ti.kernel
+    def deactivate():
+        ti.deactivate(x.parent(), 0)
+        x[1, 0] += 2
+
+    @ti.kernel
+    def check_value_0():
+        assert ti.length(x.parent(), 0) == 3
+        assert ti.length(x.parent(), 1) == 2
+        assert ti.length(x.parent(), 2) == 1
+        assert ti.length(x.parent(), 3) == 1
+        assert ti.length(x.parent(), 4) == 0
+
+        assert x[0, 0] == 1
+        assert x[0, 1] == 3
+        assert x[0, 2] == 5
+        assert x[1, 0] == 7
+        assert x[1, 1] == 7
+        assert x[2, 0] == 9
+        assert x[3, 0] == 12
+
+    @ti.kernel
+    def check_value_1():
+        assert ti.length(x.parent(), 0) == 0
+        assert ti.length(x.parent(), 1) == 2
+        assert ti.length(x.parent(), 2) == 1
+        assert ti.length(x.parent(), 3) == 1
+        assert ti.length(x.parent(), 4) == 0
+
+        assert x[0, 0] == 0
+        assert x[0, 1] == 0
+        assert x[0, 2] == 0
+        assert x[1, 0] == 9
+        assert x[1, 1] == 7
+        assert x[2, 0] == 9
+        assert x[3, 0] == 12
+
+    m = ti.aot.Module(arch)
+
+    m.add_kernel(activate, template_args={})
+    m.add_kernel(check_value_0, template_args={})
+    m.add_kernel(deactivate, template_args={})
+    m.add_kernel(check_value_1, template_args={})
+
+    m.add_field("x", x)
+
+    assert "TAICHI_AOT_FOLDER_PATH" in os.environ.keys()
+    tmpdir = str(os.environ["TAICHI_AOT_FOLDER_PATH"])
+
+    m.save(tmpdir, 'whatever')

tests/test_utils.py

@@ -21,6 +21,10 @@
     [os.path.join('cpp', 'aot', 'llvm', 'field_aot_test.py'), "--arch=cpu"],
     "LlvmAotTest.CudaField":
     [os.path.join('cpp', 'aot', 'llvm', 'field_aot_test.py'), "--arch=cuda"],
+    "LlvmAotTest.CpuDynamic":
+    [os.path.join('cpp', 'aot', 'llvm', 'dynamic_aot_test.py'), "--arch=cpu"],
+    "LlvmAotTest.CudaDynamic":
+    [os.path.join('cpp', 'aot', 'llvm', 'dynamic_aot_test.py'), "--arch=cuda"],
     "LlvmAotTest.CpuBitmasked": [
         os.path.join('cpp', 'aot', 'llvm', 'bitmasked_aot_test.py'),
         "--arch=cpu"