[demo] Add code for AOT tutorial

.gitignore

@@ -12,3 +12,4 @@ __pycache__
 imgui.ini
 2_mpm88/assets/mpm88/*
 3_implicit_fem/assets/implicit_fem/*
+0_tutorial/assets/tutorial/*

0_tutorial/CMakeLists.txt

@@ -0,0 +1,6 @@
+set(TAICHI_TUTORIAL_DEMO_NAME "0_tutorial_demo")
+message("-- Building ${TAICHI_TUTORIAL_DEMO_NAME}")
+add_executable(${TAICHI_TUTORIAL_DEMO_NAME} ${CMAKE_CURRENT_SOURCE_DIR}/app.cpp)
+target_include_directories(${TAICHI_TUTORIAL_DEMO_NAME} PUBLIC ${TAICHI_C_API_INSTALL_DIR}/include)
+target_link_libraries(${TAICHI_TUTORIAL_DEMO_NAME} ${taichi_c_api})
+set_target_properties(${TAICHI_TUTORIAL_DEMO_NAME} PROPERTIES RUNTIME_OUTPUT_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR})

0_tutorial/app.cpp

@@ -0,0 +1,51 @@
+#include <cmath>
+#include <iostream>
+#include <numeric>
+#include <taichi/cpp/taichi.hpp>
+
+struct App0_tutorial {
+  static const uint32_t NPARTICLE = 8192 * 2;
+
+  ti::Runtime runtime_;
+  ti::AotModule module_;
+  ti::ComputeGraph g_demo_;
+  ti::NdArray<float> x_;
+
+  App0_tutorial() {
+    runtime_ = ti::Runtime(TI_ARCH_VULKAN);
+    module_ = runtime_.load_aot_module("0_tutorial/assets/tutorial");
+    g_demo_ = module_.get_compute_graph("demo_graph");
+    x_ = runtime_.allocate_ndarray<float>({NPARTICLE}, {}, true);
+    std::cout << "Initialized!" << std::endl;
+  }
+
+  bool run() {
+    float base = 0.2;
+
+    g_demo_["x"] = x_;
+    g_demo_["base"] = base;
+
+    g_demo_.launch();
+    runtime_.wait();
+
+    std::vector<float> dst(NPARTICLE);
+    x_.read(dst);
+
+    float sum = std::accumulate(dst.begin(), dst.end(), 0.);
+    float expected = NPARTICLE * 50 * base;
+    if (std::abs(sum - expected) < 0.1) {
+      std::cout << "Passed" << std::endl;
+    } else {
+      std::string err_str = "Expected: " + std::to_string(expected) + ", Got: " + std::to_string(sum);
+      throw std::runtime_error(err_str);
+    }
+
+    return true;
+  }
+};
+
+int main(int argc, const char** argv) {
+  App0_tutorial app;
+  app.run();
+  return 0;
+}

0_tutorial/assets/tutorial.py

@@ -0,0 +1,72 @@
+import argparse
+import numpy as np
+import os
+
+import taichi as ti
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--arch", type=str)
+parser.add_argument("--aot", action='store_true', default=False)
+args = parser.parse_args()
+
+curr_dir = os.path.dirname(os.path.realpath(__file__))
+
+def compile_demo(arch, aot):
+    ti.init(arch=arch)
+
+    if ti.lang.impl.current_cfg().arch != arch:
+        return
+
+    n_particles = 8192
+
+    @ti.kernel
+    def init(x: ti.types.ndarray(field_dim=1)):
+        for i in x:
+            x[i] = 0
+
+    @ti.kernel
+    def add_base(x: ti.types.ndarray(field_dim=1), base: ti.f32):
+        for i in range(x.shape[0]):
+            x[i] += base
+
+
+    N_ITER = 50
+
+    sym_x = ti.graph.Arg(ti.graph.ArgKind.NDARRAY,
+                            'x',
+                            ti.f32,
+                            field_dim=1,
+                            element_shape=())
+
+    sym_base = ti.graph.Arg(ti.graph.ArgKind.SCALAR, 'base', ti.f32)
+
+    g_demo_builder = ti.graph.GraphBuilder()
+    g_demo_builder.dispatch(init, sym_x)
+
+    substep = g_demo_builder.create_sequential()
+    substep.dispatch(add_base, sym_x, sym_base)
+
+    for i in range(N_ITER):
+        g_demo_builder.append(substep)
+
+    g_demo = g_demo_builder.compile()
+
+    x = ti.ndarray(ti.f32, shape=(n_particles))
+
+    if aot:
+        mod = ti.aot.Module(arch)
+        mod.add_graph('demo_graph', g_demo)
+
+        save_dir = os.path.join(curr_dir, "tutorial")
+        os.makedirs(save_dir, exist_ok=True)
+        mod.save(save_dir, '')
+        print(f'Save compiled artifact to {save_dir}')
+    else:
+        print('Running the graph in Python')
+        g_demo.run({'x': x, 'base': 0.1})
+        assert np.allclose(x.to_numpy().sum(), n_particles * N_ITER * 0.1)
+        print('Passed verification!')
+
+if __name__ == "__main__":
+    assert args.arch == "vulkan"
+    compile_demo(arch=ti.vulkan, aot=args.aot)

CMakeLists.txt

@@ -223,3 +223,8 @@ endif()
 foreach(PYTHON_SCRIPT ${TAICHI_AOT_PYTHON_LIST})
     execute_process(COMMAND ${PYTHON_EXECUTABLE} ${PYTHON_SCRIPT} --arch vulkan)
 endforeach()
+
+# 0_tutorial is intended to be minimal code for demonstration so it's independent of framework.
+add_subdirectory("0_tutorial")
+
+execute_process(COMMAND ${PYTHON_EXECUTABLE} "${CMAKE_CURRENT_SOURCE_DIR}/0_tutorial/assets/tutorial.py" --arch vulkan --aot)