diff --git a/tests/cpp/backends/llvm/field_aot_cpu.py b/tests/cpp/backends/llvm/field_aot_cpu.py index 4556d9ca10ae1..2c227d1bbed7a 100644 --- a/tests/cpp/backends/llvm/field_aot_cpu.py +++ b/tests/cpp/backends/llvm/field_aot_cpu.py @@ -1,108 +1,6 @@ -import os +from utils import compile_field_aot import taichi as ti - -def compile_aot(): - # Make sure "debug" mode is on - # in both python & C++ tests - ti.init(arch=ti.x64, debug=True) - - x = ti.field(ti.i32) - y = ti.field(ti.i32) - - common = ti.root.dense(ti.i, 4) - common.dense(ti.i, 8).place(x) - - p = common.pointer(ti.i, 2) - p.dense(ti.i, 8).place(y) - - @ti.kernel - def init_fields(base: int): - # Dense SNode - for i in range(4 * 8): - x[i] = base + i - - # Pointer SNode - y[32] = 4 - y[33] = 5 - y[9] = 10 - - @ti.kernel - def check_init_x(base: int): - # Check numerical accuracy for Dense SNodes - for i in range(4 * 8): - assert (x[i] == base + i) - - @ti.kernel - def check_init_y(): - # Check sparsity for Pointer SNodes - for i in range(8): - if i == 1 or i == 4: - assert (ti.is_active(p, [i])) - else: - assert (not ti.is_active(p, [i])) - - # Check numerical accuracy for Pointer SNodes - for i in range(8, 8 + 8): - if i == 9: - assert (y[i] == 10) - else: - assert (y[i] == 0) - - for i in range(32, 32 + 8): - if i == 32: - assert (y[i] == 4) - elif i == 33: - assert (y[i] == 5) - else: - assert (y[i] == 0) - - @ti.kernel - def deactivate_pointer_fields(): - ti.deactivate(p, [1]) - ti.deactivate(p, [4]) - - @ti.kernel - def activate_pointer_fields(): - ti.activate(p, [7]) - ti.activate(p, [3]) - - @ti.kernel - def check_deactivate_pointer_fields(): - assert (not ti.is_active(p, [1])) - assert (not ti.is_active(p, [4])) - - @ti.kernel - def check_activate_pointer_fields(): - assert (ti.is_active(p, [7])) - assert (ti.is_active(p, [3])) - - for i in range(7 * 8, 7 * 8 + 8): - assert (y[i] == 0) - - for i in range(3 * 8, 3 * 8 + 8): - assert (y[i] == 0) - - assert "TAICHI_AOT_FOLDER_PATH" in os.environ.keys() - dir_name = str(os.environ["TAICHI_AOT_FOLDER_PATH"]) - - m = ti.aot.Module(ti.x64) - - m.add_kernel(init_fields, template_args={}) - m.add_kernel(check_init_x, template_args={}) - m.add_kernel(check_init_y, template_args={}) - - m.add_kernel(deactivate_pointer_fields, template_args={}) - m.add_kernel(activate_pointer_fields, template_args={}) - - m.add_kernel(check_deactivate_pointer_fields, template_args={}) - m.add_kernel(check_activate_pointer_fields, template_args={}) - - m.add_field("x", x) - m.add_field("y", y) - - m.save(dir_name, 'x64-aot') - - -compile_aot() +if __name__ == "__main__": + compile_field_aot(arch=ti.x64) diff --git a/tests/cpp/backends/llvm/field_aot_cuda.py b/tests/cpp/backends/llvm/field_aot_cuda.py new file mode 100644 index 0000000000000..281c79c4b3831 --- /dev/null +++ b/tests/cpp/backends/llvm/field_aot_cuda.py @@ -0,0 +1,6 @@ +from utils import compile_field_aot + +import taichi as ti + +if __name__ == "__main__": + compile_field_aot(arch=ti.cuda) diff --git a/tests/cpp/backends/llvm/field_aot_test.cpp b/tests/cpp/backends/llvm/field_aot_test.cpp index 1af9c5aea4a7c..70bca3cc821c3 100644 --- a/tests/cpp/backends/llvm/field_aot_test.cpp +++ b/tests/cpp/backends/llvm/field_aot_test.cpp @@ -16,29 +16,9 @@ namespace taichi { namespace lang { -TEST(LlvmAOTTest, Field) { - CompileConfig cfg; - cfg.arch = Arch::x64; - cfg.kernel_profiler = false; - constexpr KernelProfilerBase *kNoProfiler = nullptr; - LlvmProgramImpl prog{cfg, kNoProfiler}; - auto *compute_device = prog.get_compute_device(); - - // Must have handled all the arch fallback logic by this point. - auto memory_pool = std::make_unique(cfg.arch, compute_device); - prog.initialize_host(); - uint64 *result_buffer{nullptr}; - prog.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.program = &prog; - std::unique_ptr mod = cpu::make_aot_module(aot_params); - +void run_field_tests(aot::Module *mod, + LlvmProgramImpl *prog, + uint64 *result_buffer) { aot::Kernel *k_init_fields = mod->get_kernel("init_fields"); aot::Kernel *k_check_init_x = mod->get_kernel("check_init_x"); aot::Kernel *k_check_init_y = mod->get_kernel("check_init_y"); @@ -57,15 +37,15 @@ TEST(LlvmAOTTest, Field) { aot::Field *field_x = mod->get_field("0" /*snode_tree_id*/); aot::Field *field_y = mod->get_field("0" /*snode_tree_id*/); - finalize_aot_field(mod.get(), field_x, result_buffer); - finalize_aot_field(mod.get(), field_y, result_buffer); + finalize_aot_field(mod, field_x, result_buffer); + finalize_aot_field(mod, field_y, result_buffer); int base_value = 10; /* -------- Test Case 1 ------ */ // Kernel: init_fields(int) { RuntimeContext ctx; - ctx.runtime = prog.get_llvm_runtime(); + ctx.runtime = prog->get_llvm_runtime(); ctx.set_arg(0, base_value); k_init_fields->launch(&ctx); } @@ -73,14 +53,14 @@ TEST(LlvmAOTTest, Field) { // Kernel: check_init_x(int) { RuntimeContext ctx; - ctx.runtime = prog.get_llvm_runtime(); + ctx.runtime = prog->get_llvm_runtime(); ctx.set_arg(0, base_value); k_check_init_x->launch(&ctx); } // Kernel: check_init_y() { RuntimeContext ctx; - ctx.runtime = prog.get_llvm_runtime(); + ctx.runtime = prog->get_llvm_runtime(); k_check_init_y->launch(&ctx); } @@ -88,13 +68,13 @@ TEST(LlvmAOTTest, Field) { // Kernel: deactivate_pointer_fields() { RuntimeContext ctx; - ctx.runtime = prog.get_llvm_runtime(); + ctx.runtime = prog->get_llvm_runtime(); k_deactivate_pointer_fields->launch(&ctx); } // Kernel: check_deactivate_pointer_fields() { RuntimeContext ctx; - ctx.runtime = prog.get_llvm_runtime(); + ctx.runtime = prog->get_llvm_runtime(); k_check_deactivate_pointer_fields->launch(&ctx); } @@ -102,18 +82,70 @@ TEST(LlvmAOTTest, Field) { // Kernel: activate_pointer_fields() { RuntimeContext ctx; - ctx.runtime = prog.get_llvm_runtime(); + ctx.runtime = prog->get_llvm_runtime(); k_activate_pointer_fields->launch(&ctx); } // Kernel: check_activate_pointer_fields() { RuntimeContext ctx; - ctx.runtime = prog.get_llvm_runtime(); + ctx.runtime = prog->get_llvm_runtime(); k_check_activate_pointer_fields->launch(&ctx); } // Check assertion error from ti.kernel - prog.check_runtime_error(result_buffer); + prog->check_runtime_error(result_buffer); +} + +TEST(LlvmAotTest, CpuField) { + CompileConfig cfg; + cfg.arch = Arch::x64; + cfg.kernel_profiler = false; + constexpr KernelProfilerBase *kNoProfiler = nullptr; + LlvmProgramImpl prog{cfg, kNoProfiler}; + auto *compute_device = prog.get_compute_device(); + + // Must have handled all the arch fallback logic by this point. + auto memory_pool = std::make_unique(cfg.arch, compute_device); + prog.initialize_host(); + uint64 *result_buffer{nullptr}; + prog.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.program = &prog; + std::unique_ptr mod = cpu::make_aot_module(aot_params); + + run_field_tests(mod.get(), &prog, result_buffer); +} + +TEST(LlvmAotTest, CudaField) { + if (is_cuda_api_available()) { + CompileConfig cfg; + cfg.arch = Arch::cuda; + cfg.kernel_profiler = false; + constexpr KernelProfilerBase *kNoProfiler = nullptr; + LlvmProgramImpl prog{cfg, kNoProfiler}; + + // Must have handled all the arch fallback logic by this point. + prog.initialize_host(); + uint64 *result_buffer{nullptr}; + prog.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.program = &prog; + auto mod = cuda::make_aot_module(aot_params); + + run_field_tests(mod.get(), &prog, result_buffer); + } } } // namespace lang diff --git a/tests/cpp/backends/llvm/kernel_aot_cpu.py b/tests/cpp/backends/llvm/kernel_aot_cpu.py index 11c3a582b08cb..1be461fc65d7f 100644 --- a/tests/cpp/backends/llvm/kernel_aot_cpu.py +++ b/tests/cpp/backends/llvm/kernel_aot_cpu.py @@ -1,24 +1,6 @@ -import os +from utils import compile_kernel_aot import taichi as ti - -def compile_aot(): - ti.init(arch=ti.x64) - - @ti.kernel - def run(base: int, arr: ti.types.ndarray()): - for i in arr: - arr[i] = base + i - - arr = ti.ndarray(int, shape=16) - - assert "TAICHI_AOT_FOLDER_PATH" in os.environ.keys() - dir_name = str(os.environ["TAICHI_AOT_FOLDER_PATH"]) - - m = ti.aot.Module(ti.x64) - m.add_kernel(run, template_args={'arr': arr}) - m.save(dir_name, 'x64-aot') - - -compile_aot() +if __name__ == "__main__": + compile_kernel_aot(arch=ti.x64) diff --git a/tests/cpp/backends/llvm/kernel_aot_cuda.py b/tests/cpp/backends/llvm/kernel_aot_cuda.py index 74396f46c869d..efaa5f2d03c4d 100644 --- a/tests/cpp/backends/llvm/kernel_aot_cuda.py +++ b/tests/cpp/backends/llvm/kernel_aot_cuda.py @@ -1,24 +1,6 @@ -import os +from utils import compile_kernel_aot import taichi as ti - -def compile_aot(): - ti.init(arch=ti.cuda) - - @ti.kernel - def run(base: int, arr: ti.types.ndarray()): - for i in arr: - arr[i] = base + i - - arr = ti.ndarray(int, shape=16) - - assert "TAICHI_AOT_FOLDER_PATH" in os.environ.keys() - dir_name = str(os.environ["TAICHI_AOT_FOLDER_PATH"]) - - m = ti.aot.Module(ti.cuda) - m.add_kernel(run, template_args={'arr': arr}) - m.save(dir_name, 'cuda-aot') - - -compile_aot() +if __name__ == "__main__": + compile_kernel_aot(arch=ti.cuda) diff --git a/tests/cpp/backends/llvm/kernel_aot_test.cpp b/tests/cpp/backends/llvm/kernel_aot_test.cpp index dca0aa2ea1631..e1e5fc43ec776 100644 --- a/tests/cpp/backends/llvm/kernel_aot_test.cpp +++ b/tests/cpp/backends/llvm/kernel_aot_test.cpp @@ -15,7 +15,7 @@ namespace taichi { namespace lang { -TEST(LlvmProgramTest, FullPipeline) { +TEST(LlvmAotTest, CpuKernel) { CompileConfig cfg; cfg.arch = Arch::x64; cfg.kernel_profiler = false; @@ -55,7 +55,7 @@ TEST(LlvmProgramTest, FullPipeline) { } } -TEST(LlvmProgramTest, FullPipelineCUDA) { +TEST(LlvmAotTest, CudaKernel) { if (is_cuda_api_available()) { CompileConfig cfg; cfg.arch = Arch::cuda; diff --git a/tests/cpp/backends/llvm/utils.py b/tests/cpp/backends/llvm/utils.py new file mode 100644 index 0000000000000..1db7ba9c54d8b --- /dev/null +++ b/tests/cpp/backends/llvm/utils.py @@ -0,0 +1,123 @@ +import os + +import taichi as ti + + +def compile_field_aot(arch): + # Make sure "debug" mode is on + # in both python & C++ tests + ti.init(arch=arch, debug=True) + + x = ti.field(ti.i32) + y = ti.field(ti.i32) + + common = ti.root.dense(ti.i, 4) + common.dense(ti.i, 8).place(x) + + p = common.pointer(ti.i, 2) + p.dense(ti.i, 8).place(y) + + @ti.kernel + def init_fields(base: int): + # Dense SNode + for i in range(4 * 8): + x[i] = base + i + + # Pointer SNode + y[32] = 4 + y[33] = 5 + y[9] = 10 + + @ti.kernel + def check_init_x(base: int): + # Check numerical accuracy for Dense SNodes + for i in range(4 * 8): + assert (x[i] == base + i) + + @ti.kernel + def check_init_y(): + # Check sparsity for Pointer SNodes + for i in range(8): + if i == 1 or i == 4: + assert (ti.is_active(p, [i])) + else: + assert (not ti.is_active(p, [i])) + + # Check numerical accuracy for Pointer SNodes + for i in range(8, 8 + 8): + if i == 9: + assert (y[i] == 10) + else: + assert (y[i] == 0) + + for i in range(32, 32 + 8): + if i == 32: + assert (y[i] == 4) + elif i == 33: + assert (y[i] == 5) + else: + assert (y[i] == 0) + + @ti.kernel + def deactivate_pointer_fields(): + ti.deactivate(p, [1]) + ti.deactivate(p, [4]) + + @ti.kernel + def activate_pointer_fields(): + ti.activate(p, [7]) + ti.activate(p, [3]) + + @ti.kernel + def check_deactivate_pointer_fields(): + assert (not ti.is_active(p, [1])) + assert (not ti.is_active(p, [4])) + + @ti.kernel + def check_activate_pointer_fields(): + assert (ti.is_active(p, [7])) + assert (ti.is_active(p, [3])) + + for i in range(7 * 8, 7 * 8 + 8): + assert (y[i] == 0) + + for i in range(3 * 8, 3 * 8 + 8): + assert (y[i] == 0) + + assert "TAICHI_AOT_FOLDER_PATH" in os.environ.keys() + dir_name = str(os.environ["TAICHI_AOT_FOLDER_PATH"]) + + m = ti.aot.Module(arch) + + m.add_kernel(init_fields, template_args={}) + m.add_kernel(check_init_x, template_args={}) + m.add_kernel(check_init_y, template_args={}) + + m.add_kernel(deactivate_pointer_fields, template_args={}) + m.add_kernel(activate_pointer_fields, template_args={}) + + m.add_kernel(check_deactivate_pointer_fields, template_args={}) + m.add_kernel(check_activate_pointer_fields, template_args={}) + + m.add_field("x", x) + m.add_field("y", y) + + m.save(dir_name, 'whatever') + + +def compile_kernel_aot(arch): + ti.init(arch=arch) + + @ti.kernel + def run(base: int, arr: ti.types.ndarray()): + for i in arr: + arr[i] = base + i + + arr = ti.ndarray(int, shape=16) + + assert "TAICHI_AOT_FOLDER_PATH" in os.environ.keys() + dir_name = str(os.environ["TAICHI_AOT_FOLDER_PATH"]) + + m = ti.aot.Module(arch) + m.add_kernel(run, template_args={'arr': arr}) + m.save(dir_name, 'whatever') diff --git a/tests/test_utils.py b/tests/test_utils.py index 74cd5c7f0fba0..84b12fa30c4ca 100644 --- a/tests/test_utils.py +++ b/tests/test_utils.py @@ -13,12 +13,14 @@ import taichi as ti __aot_test_cases = { - "LlvmProgramTest.FullPipeline": + "LlvmAotTest.CpuKernel": os.path.join('cpp', 'backends', 'llvm', 'kernel_aot_cpu.py'), - "LlvmProgramTest.FullPipelineCUDA": + "LlvmAotTest.CudaKernel": os.path.join('cpp', 'backends', 'llvm', 'kernel_aot_cuda.py'), - "LlvmAOTTest.Field": + "LlvmAotTest.CpuField": os.path.join('cpp', 'backends', 'llvm', 'field_aot_cpu.py'), + "LlvmAotTest.CudaField": + os.path.join('cpp', 'backends', 'llvm', 'field_aot_cuda.py'), "AotLoadGraph.Mpm88": os.path.join('cpp', 'aot', 'mpm88_graph_aot.py'), }