[aot] C-API to get available archs (taichi-dev#6766)

This PR provides a new API to query all available APIs on the current platform. It is basically a public version of the test utils APIs (`is_vulkan_available`, etc.) @jim19930609 introduced before. Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
quadpixels · May 13, 2023 · 8684f88 · 8684f88
1 parent 8c018d8
commit 8684f88
Show file tree

Hide file tree

Showing 23 changed files with 216 additions and 91 deletions.
diff --git a/c_api/docs/taichi/taichi_core.h.md b/c_api/docs/taichi/taichi_core.h.md
@@ -440,16 +440,25 @@ A named argument value to feed compute graphs.
 - `structure.named_argument.name`: Name of the argument.
 - `structure.named_argument.argument`: Argument body.
 
+`function.get_available_archs`
+
+Gets a list of available archs on the current platform. An arch is only available if:
+
+1. The Runtime library is compiled with its support;
+2. The current platform is installed with a capable hardware or an emulation software.
+
+An available arch has at least one device available, i.e., device index 0 is always available. If an arch is not available on the current platform, a call to `function.create_runtime` with that arch is guaranteed failing.
+
 `function.get_last_error`
 
-Get the last error raised by Taichi C-API invocations. Returns the semantical error code.
+Gets the last error raised by Taichi C-API invocations. Returns the semantical error code.
 
 - `function.get_last_error.message_size`: Size of textual error message in `function.get_last_error.message`
 - `function.get_last_error.message`: Text buffer for the textual error message. Ignored when `message_size` is 0.
 
 `function.set_last_error`
 
-Set the provided error as the last error raised by Taichi C-API invocations. It can be useful in extended validation procedures in Taichi C-API wrappers and helper libraries.
+Sets the provided error as the last error raised by Taichi C-API invocations. It can be useful in extended validation procedures in Taichi C-API wrappers and helper libraries.
 
 - `function.set_last_error.error`: Semantical error code.
 - `function.set_last_error.message`: A null-terminated string of the textual error message or `nullptr` for empty error message.

diff --git a/c_api/include/taichi/cpp/taichi.hpp b/c_api/include/taichi/cpp/taichi.hpp
@@ -1,5 +1,6 @@
 // C++ wrapper of Taichi C-API
 #pragma once
+#include <cstddef>
 #include <cstring>
 #include <list>
 #include <vector>
@@ -10,6 +11,23 @@
 
 namespace ti {
 
+inline std::vector<TiArch> get_available_archs() {
+  uint32_t narch = 0;
+  ti_get_available_archs(&narch, nullptr);
+  std::vector<TiArch> archs(narch);
+  ti_get_available_archs(&narch, archs.data());
+  return archs;
+}
+inline bool is_arch_available(TiArch arch) {
+  std::vector<TiArch> archs = get_available_archs();
+  for (size_t i = 0; i < archs.size(); ++i) {
+    if (archs.at(i) == arch) {
+      return true;
+    }
+  }
+  return false;
+}
+
 // Token type for half-precision floats.
 struct half {
   uint16_t _;

diff --git a/c_api/include/taichi/taichi_core.h b/c_api/include/taichi/taichi_core.h
@@ -802,10 +802,26 @@ typedef struct TiNamedArgument {
   TiArgument argument;
 } TiNamedArgument;
 
+// Function `ti_get_available_archs`
+//
+// Gets a list of available archs on the current platform. An arch is only
+// available if:
+//
+// 1. The Runtime library is compiled with its support;
+// 2. The current platform is installed with a capable hardware or an emulation
+// software.
+//
+// An available arch has at least one device available, i.e., device index 0 is
+// always available. If an arch is not available on the current platform, a call
+// to [`ti_create_runtime`](#function-ti_create_runtime) with that arch is
+// guaranteed failing.
+TI_DLL_EXPORT void TI_API_CALL ti_get_available_archs(uint32_t *arch_count,
+                                                      TiArch *archs);
+
 // Function `ti_get_last_error`
 //
-// Get the last error raised by Taichi C-API invocations. Returns the semantical
-// error code.
+// Gets the last error raised by Taichi C-API invocations. Returns the
+// semantical error code.
 TI_DLL_EXPORT TiError TI_API_CALL ti_get_last_error(
     // Size of textual error message in `function.get_last_error.message`
     uint64_t message_size,
@@ -815,7 +831,7 @@ TI_DLL_EXPORT TiError TI_API_CALL ti_get_last_error(
 
 // Function `ti_set_last_error`
 //
-// Set the provided error as the last error raised by Taichi C-API invocations.
+// Sets the provided error as the last error raised by Taichi C-API invocations.
 // It can be useful in extended validation procedures in Taichi C-API wrappers
 // and helper libraries.
 TI_DLL_EXPORT void TI_API_CALL ti_set_last_error(

diff --git a/c_api/include/taichi/taichi_opengl.h b/c_api/include/taichi/taichi_opengl.h
@@ -1,9 +1,5 @@
 #pragma once
 
-#ifndef TI_WITH_OPENGL
-#define TI_WITH_OPENGL 1
-#endif  // TI_WITH_OPENGL
-
 #include <taichi/taichi.h>
 
 #ifdef __cplusplus

diff --git a/c_api/include/taichi/taichi_vulkan.h b/c_api/include/taichi/taichi_vulkan.h
@@ -6,10 +6,6 @@
 //
 #pragma once
 
-#ifndef TI_WITH_VULKAN
-#define TI_WITH_VULKAN 1
-#endif  // TI_WITH_VULKAN
-
 #include <taichi/taichi.h>
 
 #ifdef __cplusplus

diff --git a/c_api/src/c_api_test_utils.cpp b/c_api/src/c_api_test_utils.cpp
@@ -1,6 +1,5 @@
 #include "c_api_test_utils.h"
 #include "taichi_llvm_impl.h"
-#include "taichi/platform/cuda/detect_cuda.h"
 
 #ifdef TI_WITH_CUDA
 #include "taichi/rhi/cuda/cuda_driver.h"
@@ -37,26 +36,6 @@ bool check_cuda_value(void *ptr, double value) {
   return check_cuda_value_impl(ptr, value);
 }
 
-bool is_vulkan_available() {
-#ifdef TI_WITH_VULKAN
-  return taichi::lang::vulkan::is_vulkan_api_available();
-#else
-  return false;
-#endif
-}
-
-bool is_opengl_available() {
-#ifdef TI_WITH_OPENGL
-  return taichi::lang::opengl::is_opengl_api_available();
-#else
-  return false;
-#endif
-}
-
-bool is_cuda_available() {
-  return taichi::is_cuda_api_available();
-}
-
 void check_runtime_error(TiRuntime runtime) {
 #ifdef TI_WITH_LLVM
   auto *llvm_runtime = dynamic_cast<capi::LlvmRuntime *>((Runtime *)runtime);

diff --git a/c_api/src/c_api_test_utils.h b/c_api/src/c_api_test_utils.h
@@ -4,9 +4,6 @@
 namespace capi {
 namespace utils {
 
-TI_DLL_EXPORT bool TI_API_CALL is_vulkan_available();
-TI_DLL_EXPORT bool TI_API_CALL is_opengl_available();
-TI_DLL_EXPORT bool TI_API_CALL is_cuda_available();
 TI_DLL_EXPORT void TI_API_CALL check_runtime_error(TiRuntime runtime);
 
 TI_DLL_EXPORT bool TI_API_CALL check_cuda_value(void *ptr, float value);

diff --git a/c_api/src/taichi_core_impl.cpp b/c_api/src/taichi_core_impl.cpp
@@ -6,6 +6,48 @@
 #include "taichi/program/texture.h"
 #include "taichi/common/virtual_dir.h"
 
+bool is_vulkan_available() {
+#ifdef TI_WITH_VULKAN
+  return taichi::lang::vulkan::is_vulkan_api_available();
+#else
+  return false;
+#endif
+}
+
+bool is_opengl_available() {
+#ifdef TI_WITH_OPENGL
+  return taichi::lang::opengl::is_opengl_api_available();
+#else
+  return false;
+#endif
+}
+
+bool is_cuda_available() {
+#ifdef TI_WITH_CUDA
+  return taichi::is_cuda_api_available();
+#else
+  return false;
+#endif
+}
+
+bool is_x64_available() {
+#if defined(TI_WITH_LLVM) &&                                           \
+    (defined(__x86_64__) || defined(__x86_64) || defined(__amd64__) || \
+     defined(__amd64) || defined(_M_X64))
+  return true;
+#else
+  return false;
+#endif
+}
+
+bool is_arm64_available() {
+#if defined(TI_WITH_LLVM) && (defined(__arm64__) || defined(__aarch64__))
+  return true;
+#else
+  return false;
+#endif
+}
+
 struct ErrorCache {
   TiError error{TI_ERROR_SUCCESS};
   std::string message{};
@@ -107,6 +149,39 @@ Runtime &Event::runtime() {
 
 // -----------------------------------------------------------------------------
 
+void ti_get_available_archs(uint32_t *arch_count, TiArch *archs) {
+  if (arch_count == nullptr) {
+    return;
+  }
+
+  thread_local std::vector<TiArch> AVAILABLE_ARCHS{};
+  if (AVAILABLE_ARCHS.empty()) {
+    if (is_vulkan_available()) {
+      AVAILABLE_ARCHS.emplace_back(TI_ARCH_VULKAN);
+    }
+    if (is_opengl_available()) {
+      AVAILABLE_ARCHS.emplace_back(TI_ARCH_OPENGL);
+    }
+    if (is_cuda_available()) {
+      AVAILABLE_ARCHS.emplace_back(TI_ARCH_CUDA);
+    }
+    if (is_x64_available()) {
+      AVAILABLE_ARCHS.emplace_back(TI_ARCH_X64);
+    }
+    if (is_arm64_available()) {
+      AVAILABLE_ARCHS.emplace_back(TI_ARCH_ARM64);
+    }
+  }
+
+  size_t n = std::min((size_t)*arch_count, AVAILABLE_ARCHS.size());
+  *arch_count = (uint32_t)n;
+  if (archs != nullptr) {
+    for (size_t i = 0; i < n; ++i) {
+      archs[i] = AVAILABLE_ARCHS.at(i);
+    }
+  }
+}
+
 TiError ti_get_last_error(uint64_t message_size, char *message) {
   TiError out = TI_ERROR_INVALID_STATE;
   TI_CAPI_TRY_CATCH_BEGIN();

diff --git a/c_api/src/taichi_llvm_impl.cpp b/c_api/src/taichi_llvm_impl.cpp
@@ -136,8 +136,6 @@ void LlvmRuntime::wait() {
 
 }  // namespace capi
 
-#endif  // TI_WITH_LLVM
-
 // function.export_cpu_runtime
 void ti_export_cpu_memory(TiRuntime runtime,
                           TiMemory memory,
@@ -197,3 +195,5 @@ void ti_export_cuda_memory(TiRuntime runtime,
   TI_NOT_IMPLEMENTED;
 #endif
 }
+
+#endif  // TI_WITH_LLVM
diff --git a/c_api/src/taichi_llvm_impl.h b/c_api/src/taichi_llvm_impl.h
@@ -3,6 +3,10 @@
 
 #include "taichi_core_impl.h"
 
+#ifdef TI_WITH_CUDA
+#include "taichi/platform/cuda/detect_cuda.h"
+#endif
+
 namespace taichi::lang {
 class LlvmRuntimeExecutor;
 class MemoryPool;

diff --git a/c_api/src/taichi_opengl_impl.cpp b/c_api/src/taichi_opengl_impl.cpp
@@ -1,5 +1,5 @@
-#include "taichi_opengl_impl.h"
 #ifdef TI_WITH_OPENGL
+#include "taichi_opengl_impl.h"
 
 OpenglRuntime::OpenglRuntime()
     : GfxRuntime(taichi::Arch::opengl),

diff --git a/c_api/src/taichi_opengl_impl.h b/c_api/src/taichi_opengl_impl.h
@@ -2,6 +2,7 @@
 #ifdef TI_WITH_OPENGL
 
 #include "taichi_gfx_impl.h"
+#include "taichi/rhi/opengl/opengl_api.h"
 #include "taichi/rhi/opengl/opengl_device.h"
 
 class OpenglRuntime : public GfxRuntime {

diff --git a/c_api/src/taichi_vulkan_impl.h b/c_api/src/taichi_vulkan_impl.h
@@ -3,6 +3,7 @@
 
 #include "taichi_core_impl.h"
 #include "taichi_gfx_impl.h"
+#include "taichi/rhi/vulkan/vulkan_loader.h"
 #include "taichi/rhi/vulkan/vulkan_device.h"
 #include "taichi/rhi/vulkan/vulkan_device_creator.h"
 

diff --git a/c_api/taichi.json b/c_api/taichi.json
@@ -459,6 +459,23 @@
                         }
                     ]
                 },
+                {
+                    "name": "get_available_archs",
+                    "type": "function",
+                    "parameters": [
+                        {
+                            "name": "arch_count",
+                            "type": "uint32_t",
+                            "by_mut": true
+                        },
+                        {
+                            "name": "archs",
+                            "type": "enumeration.arch",
+                            "count": "arch_count",
+                            "by_mut": true
+                        }
+                    ]
+                },
                 {
                     "name": "get_last_error",
                     "type": "function",
@@ -973,12 +990,6 @@
         },
         {
             "name": "taichi/taichi_cuda.h",
-            "default_definitions": [
-                {
-                    "name": "TI_WITH_CUDA",
-                    "value": "1"
-                }
-            ],
             "required_modules": [
                 "taichi/taichi_core.h"
             ],
@@ -1018,12 +1029,6 @@
         },
         {
             "name": "taichi/taichi_vulkan.h",
-            "default_definitions": [
-                {
-                    "name": "TI_WITH_VULKAN",
-                    "value": "1"
-                }
-            ],
             "doc": "taichi/taichi_vulkan.h.md",
             "required_modules": [
                 "taichi/taichi_core.h"
@@ -1326,12 +1331,6 @@
         },
         {
             "name": "taichi/taichi_opengl.h",
-            "default_definitions": [
-                {
-                    "name": "TI_WITH_OPENGL",
-                    "value": "1"
-                }
-            ],
             "required_modules": [
                 "taichi/taichi_core.h"
             ],

diff --git a/c_api/tests/c_api_aot_test.cpp b/c_api/tests/c_api_aot_test.cpp
@@ -85,7 +85,7 @@ TEST_F(CapiTest, AotTestCpuField) {
 }
 
 TEST_F(CapiTest, AotTestCudaField) {
-  if (capi::utils::is_cuda_available()) {
+  if (ti::is_arch_available(TI_ARCH_CUDA)) {
     TiArch arch = TiArch::TI_ARCH_CUDA;
     field_aot_test(arch);
   }
@@ -97,21 +97,21 @@ TEST_F(CapiTest, AotTestCpuKernel) {
 }
 
 TEST_F(CapiTest, AotTestCudaKernel) {
-  if (capi::utils::is_cuda_available()) {
+  if (ti::is_arch_available(TI_ARCH_CUDA)) {
     TiArch arch = TiArch::TI_ARCH_CUDA;
     kernel_aot_test(arch);
   }
 }
 
 TEST_F(CapiTest, AotTestVulkanKernel) {
-  if (capi::utils::is_vulkan_available()) {
+  if (ti::is_arch_available(TI_ARCH_VULKAN)) {
     TiArch arch = TiArch::TI_ARCH_VULKAN;
     kernel_aot_test(arch);
   }
 }
 
 TEST_F(CapiTest, AotTestOpenglKernel) {
-  if (capi::utils::is_opengl_available()) {
+  if (ti::is_arch_available(TI_ARCH_OPENGL)) {
     TiArch arch = TiArch::TI_ARCH_OPENGL;
     kernel_aot_test(arch);
   }