Added more test cases

taichi-dev · Jul 18, 2022 · d2622f6 · d2622f6
1 parent d3dfe01
commit d2622f6
Show file tree

Hide file tree

Showing 5 changed files with 85 additions and 0 deletions.
diff --git a/c_api/src/c_api_test_utils.cpp b/c_api/src/c_api_test_utils.cpp
@@ -1,4 +1,5 @@
 #include "c_api_test_utils.h"
+#include "taichi/platform/cuda/detect_cuda.h"
 
 #ifdef TI_WITH_VULKAN
 #include "taichi/rhi/vulkan/vulkan_loader.h"
@@ -15,5 +16,9 @@ bool is_vulkan_available() {
 #endif
 }
 
+bool is_cuda_available() {
+  return taichi::is_cuda_api_available();
+}
+
 }  // namespace utils
 }  // namespace capi
diff --git a/c_api/src/c_api_test_utils.h b/c_api/src/c_api_test_utils.h
@@ -5,6 +5,7 @@ namespace capi {
 namespace utils {
 
 TI_DLL_EXPORT bool TI_API_CALL is_vulkan_available();
+TI_DLL_EXPORT bool TI_API_CALL is_cuda_available();
 
 }  // namespace utils
 }  // namespace capi
diff --git a/c_api/tests/c_api_interface_test.cpp b/c_api/tests/c_api_interface_test.cpp
@@ -16,6 +16,13 @@ TEST(CapiDryRun, Runtime) {
     TiRuntime runtime = ti_create_runtime(arch);
     ti_destroy_runtime(runtime);
   }
+
+  if (capi::utils::is_cuda_available()) {
+    // Vulkan Runtime
+    TiArch arch = TiArch::TI_ARCH_CUDA;
+    TiRuntime runtime = ti_create_runtime(arch);
+    ti_destroy_runtime(runtime);
+  }
 }
 
 TEST(CapiDryRun, MemoryAllocation) {
@@ -38,6 +45,28 @@ TEST(CapiDryRun, MemoryAllocation) {
 
     ti_destroy_runtime(runtime);
   }
+
+  if (capi::utils::is_vulkan_available()) {
+    // Vulkan Runtime
+    TiArch arch = TiArch::TI_ARCH_VULKAN;
+    TiRuntime runtime = ti_create_runtime(arch);
+
+    TiMemory memory = ti_allocate_memory(runtime, &alloc_info);
+    ti_free_memory(runtime, memory);
+
+    ti_destroy_runtime(runtime);
+  }
+
+  if (capi::utils::is_cuda_available()) {
+    // Cuda Runtime
+    TiArch arch = TiArch::TI_ARCH_CUDA;
+    TiRuntime runtime = ti_create_runtime(arch);
+
+    TiMemory memory = ti_allocate_memory(runtime, &alloc_info);
+    ti_free_memory(runtime, memory);
+
+    ti_destroy_runtime(runtime);
+  }
 }
 
 TEST(CapiDryRun, CpuAotModule) {
@@ -57,3 +86,45 @@ TEST(CapiDryRun, CpuAotModule) {
     ti_destroy_runtime(runtime);
   }
 }
+
+TEST(CapiDryRun, VulkanAotModule) {
+  if (capi::utils::is_vulkan_available()) {
+    const auto folder_dir = getenv("TAICHI_AOT_FOLDER_PATH");
+
+    std::stringstream aot_mod_ss;
+    aot_mod_ss << folder_dir;
+
+    {
+      // Vulkan Runtime
+      TiArch arch = TiArch::TI_ARCH_VULKAN;
+      TiRuntime runtime = ti_create_runtime(arch);
+
+      TiAotModule aot_mod =
+          ti_load_aot_module(runtime, aot_mod_ss.str().c_str());
+      ti_destroy_aot_module(aot_mod);
+
+      ti_destroy_runtime(runtime);
+    }
+  }
+}
+
+TEST(CapiDryRun, CudaAotModule) {
+  if (capi::utils::is_cuda_available()) {
+    const auto folder_dir = getenv("TAICHI_AOT_FOLDER_PATH");
+
+    std::stringstream aot_mod_ss;
+    aot_mod_ss << folder_dir;
+
+    {
+      // Vulkan Runtime
+      TiArch arch = TiArch::TI_ARCH_CUDA;
+      TiRuntime runtime = ti_create_runtime(arch);
+
+      TiAotModule aot_mod =
+          ti_load_aot_module(runtime, aot_mod_ss.str().c_str());
+      ti_destroy_aot_module(aot_mod);
+
+      ti_destroy_runtime(runtime);
+    }
+  }
+}
diff --git a/tests/cpp/aot/llvm/kernel_aot_test.py b/tests/cpp/aot/llvm/kernel_aot_test.py
@@ -13,5 +13,7 @@
         compile_kernel_aot(arch=ti.cpu)
     elif args.arch == "cuda":
         compile_kernel_aot(arch=ti.cuda)
+    elif args.arch == "vulkan":
+        compile_kernel_aot(arch=ti.vulkan)
     else:
         assert False
diff --git a/tests/test_utils.py b/tests/test_utils.py
@@ -36,6 +36,12 @@
 __capi_aot_test_cases = {
     "CapiDryRun.CpuAotModule":
     [os.path.join('cpp', 'aot', 'llvm', 'kernel_aot_test.py'), "--arch=cpu"],
+    "CapiDryRun.VulkanAotModule": [
+        os.path.join('cpp', 'aot', 'llvm', 'kernel_aot_test.py'),
+        "--arch=vulkan"
+    ],
+    "CapiDryRun.CudaAotModule":
+    [os.path.join('cpp', 'aot', 'llvm', 'kernel_aot_test.py'), "--arch=cuda"],
 }