[llvm] [aot] Added taichi_sparse unit test to C-API for CUDA backend

c_api/tests/c_api_aot_test.cpp

@@ -2,7 +2,7 @@
 #include "c_api_test_utils.h"
 #include "taichi/taichi_core.h"
 
-void kernel_aot_test(TiArch arch) {
+static void kernel_aot_test(TiArch arch) {
   uint32_t kArrLen = 32;
   int arg0_val = 0;
 
@@ -57,7 +57,7 @@ void kernel_aot_test(TiArch arch) {
   ti_destroy_runtime(runtime);
 }
 
-void field_aot_test(TiArch arch) {
+static void field_aot_test(TiArch arch) {
   int base_val = 10;
 
   const auto folder_dir = getenv("TAICHI_AOT_FOLDER_PATH");

c_api/tests/taichi_sparse_test.cpp

@@ -0,0 +1,54 @@
+#include "gtest/gtest.h"
+#include "c_api_test_utils.h"
+#include "taichi/taichi_core.h"
+
+static void taichi_sparse_test(TiArch arch) {
+  const auto folder_dir = getenv("TAICHI_AOT_FOLDER_PATH");
+
+  std::stringstream aot_mod_ss;
+  aot_mod_ss << folder_dir;
+
+  TiRuntime runtime = ti_create_runtime(arch);
+
+  // Load Aot and Kernel
+  TiAotModule aot_mod = ti_load_aot_module(runtime, aot_mod_ss.str().c_str());
+
+  TiKernel k_fill_img = ti_get_aot_module_kernel(aot_mod, "fill_img");
+  TiKernel k_block1_deactivate_all =
+      ti_get_aot_module_kernel(aot_mod, "block1_deactivate_all");
+  TiKernel k_activate = ti_get_aot_module_kernel(aot_mod, "activate");
+  TiKernel k_paint = ti_get_aot_module_kernel(aot_mod, "paint");
+  TiKernel k_check_img_value =
+      ti_get_aot_module_kernel(aot_mod, "check_img_value");
+
+  constexpr uint32_t arg_count = 1;
+  TiArgument args[arg_count];
+
+  ti_launch_kernel(runtime, k_fill_img, 0, &args[0]);
+  for (int i = 0; i < 100; i++) {
+    float val = 0.05f * i;
+    TiArgument base_arg = {.type = TiArgumentType::TI_ARGUMENT_TYPE_F32,
+                           .value = {.f32 = val}};
+    args[0] = std::move(base_arg);
+
+    ti_launch_kernel(runtime, k_block1_deactivate_all, 0, &args[0]);
+    ti_launch_kernel(runtime, k_activate, arg_count, &args[0]);
+    ti_launch_kernel(runtime, k_paint, 0, &args[0]);
+  }
+
+  // Accuracy Check
+  ti_launch_kernel(runtime, k_check_img_value, 0, &args[0]);
+
+  // Check Results
+  capi::utils::check_runtime_error(runtime);
+
+  ti_destroy_aot_module(aot_mod);
+  ti_destroy_runtime(runtime);
+}
+
+TEST(CapiTaichiSparseTest, Cuda) {
+  if (capi::utils::is_cuda_available()) {
+    TiArch arch = TiArch::TI_ARCH_CUDA;
+    taichi_sparse_test(arch);
+  }
+}

tests/cpp/aot/llvm/taichi_sparse_test.py

@@ -0,0 +1,93 @@
+import argparse
+import os
+
+from taichi.examples.patterns import taichi_logo
+
+import taichi as ti
+
+ti.init(arch=ti.cuda, debug=True)
+
+n = 512
+x = ti.field(dtype=ti.i32)
+res = n + n // 4 + n // 16 + n // 64
+img = ti.field(dtype=ti.f32, shape=(res, res))
+
+block1 = ti.root.pointer(ti.ij, n // 64)
+block2 = block1.pointer(ti.ij, 4)
+block3 = block2.pointer(ti.ij, 4)
+block3.dense(ti.ij, 4).place(x)
+
+
+@ti.kernel
+def check_img_value():
+    s: ti.f32 = 0
+    for i in ti.static(range(20)):
+        s += img[i, i]
+    assert s == 15.25
+
+
+@ti.kernel
+def fill_img():
+    img.fill(0.05)
+
+
+@ti.kernel
+def block1_deactivate_all():
+    for I in ti.grouped(block3):
+        ti.deactivate(block3, I)
+
+    for I in ti.grouped(block2):
+        ti.deactivate(block2, I)
+
+    for I in ti.grouped(block1):
+        ti.deactivate(block1, I)
+
+
+@ti.kernel
+def activate(t: ti.f32):
+    for i, j in ti.ndrange(n, n):
+        p = ti.Vector([i, j]) / n
+        p = ti.Matrix.rotation2d(ti.sin(t)) @ (p - 0.5) + 0.5
+
+        if taichi_logo(p) == 0:
+            x[i, j] = 1
+
+
+@ti.func
+def scatter(i):
+    return i + i // 4 + i // 16 + i // 64 + 2
+
+
+@ti.kernel
+def paint():
+    for i, j in ti.ndrange(n, n):
+        t = x[i, j]
+        block1_index = ti.rescale_index(x, block1, [i, j])
+        block2_index = ti.rescale_index(x, block2, [i, j])
+        block3_index = ti.rescale_index(x, block3, [i, j])
+        t += ti.is_active(block1, block1_index)
+        t += ti.is_active(block2, block2_index)
+        t += ti.is_active(block3, block3_index)
+        img[scatter(i), scatter(j)] = 1 - t / 4
+
+
+def save_kernels(arch):
+    m = ti.aot.Module(arch)
+
+    m.add_kernel(fill_img, template_args={})
+    m.add_kernel(block1_deactivate_all, template_args={})
+    m.add_kernel(activate, template_args={})
+    m.add_kernel(paint, template_args={})
+    m.add_kernel(check_img_value, template_args={})
+
+    m.add_field("x", x)
+    m.add_field("img", img)
+
+    assert "TAICHI_AOT_FOLDER_PATH" in os.environ.keys()
+    dir_name = str(os.environ["TAICHI_AOT_FOLDER_PATH"])
+
+    m.save(dir_name, 'whatever')
+
+
+if __name__ == '__main__':
+    save_kernels(arch=ti.cuda)

tests/test_utils.py

@@ -42,6 +42,8 @@
 }
 
 __capi_aot_test_cases = {
+    "CapiTaichiSparseTest.Cuda":
+    [os.path.join('cpp', 'aot', 'llvm', 'taichi_sparse_test.py'), ""],
     "CapiAotTest.CpuField":
     [os.path.join('cpp', 'aot', 'llvm', 'field_aot_test.py'), "--arch=cpu"],
     "CapiAotTest.CudaField":