[aot] C-API device capability query (taichi-dev#6549)

Issue: #

### Brief Summary

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

.github/workflows/scripts/aot-demo.sh

@@ -59,6 +59,11 @@ function prepare-unity-build-env {
     cmake --build .
     popd
     cp tu2-build/bin/libtaichi_unity.so Taichi-UnityExample/Assets/Plugins/Android
+
+    pushd Taichi-UnityExample
+    pip install /taichi-wheel/*.whl
+    python scripts/implicit_fem.cgraph.py --aot
+    popd
 }
 
 function build-unity-demo {

c_api/docs/taichi/taichi_core.h.md

@@ -244,9 +244,8 @@ A collection of Taichi kernels (a compute graph) to launch on the offload target
 
 `enumeration.error`
 
-Errors reported by the Taichi C-API. Enumerants greater than or equal to zero are success states.
+Errors reported by the Taichi C-API.
 
-- `enumeration.error.truncated`: The output data is truncated because the user-provided buffer is too small.
 - `enumeration.error.success`: The Taichi C-API invocation finished gracefully.
 - `enumeration.error.not_supported`: The invoked API, or the combination of parameters is not supported by the Taichi C-API.
 - `enumeration.error.corrupted_data`: Provided data is corrupted.
@@ -257,6 +256,7 @@ Errors reported by the Taichi C-API. Enumerants greater than or equal to zero ar
 - `enumeration.error.argument_not_found`: One or more kernel arguments are missing.
 - `enumeration.error.invalid_interop`: The intended interoperation is not possible on the current arch. For example, attempts to export a Vulkan object from a CUDA runtime are not allowed.
 - `enumeration.error.invalid_state`: The Taichi C-API enters an unrecoverable invalid state. Related Taichi objects are potentially corrupted. The users *should* release the contaminated resources for stability. Please feel free to file an issue if you encountered this error in a normal routine.
+- `enumeration.error.incompatible_module`: The AOT module is not compatible with the current runtime.
 
 `enumeration.arch`
 

c_api/include/taichi/cpp/taichi.hpp

@@ -3,6 +3,7 @@
 #include <cstring>
 #include <list>
 #include <vector>
+#include <map>
 #include <string>
 #include <utility>
 #include "taichi/taichi.h"
@@ -703,6 +704,19 @@ class Runtime {
     return *this;
   }
 
+  std::map<TiCapability, uint32_t> get_capabilities() const {
+    uint32_t n = 0;
+    ti_get_runtime_capabilities(runtime_, &n, nullptr);
+    std::vector<TiCapabilityLevelInfo> devcaps(n);
+    ti_get_runtime_capabilities(runtime_, &n, devcaps.data());
+
+    std::map<TiCapability, uint32_t> out{};
+    for (auto devcap : devcaps) {
+      out[devcap.capability] = devcap.level;
+    }
+    return out;
+  }
+
   Memory allocate_memory(const TiMemoryAllocateInfo &allocate_info) {
     TiMemory memory = ti_allocate_memory(runtime_, &allocate_info);
     return Memory(runtime_, memory, allocate_info.size, true);

c_api/include/taichi/taichi_core.h

@@ -317,11 +317,8 @@ typedef struct TiComputeGraph_t *TiComputeGraph;
 
 // Enumeration `TiError`
 //
-// Errors reported by the Taichi C-API. Enumerants greater than or equal to zero
-// are success states.
+// Errors reported by the Taichi C-API.
 typedef enum TiError {
-  // The output data is truncated because the user-provided buffer is too small.
-  TI_ERROR_TRUNCATED = 1,
   // The Taichi C-API invocation finished gracefully.
   TI_ERROR_SUCCESS = 0,
   // The invoked API, or the combination of parameters is not supported by the
@@ -351,6 +348,8 @@ typedef enum TiError {
   // contaminated resources for stability. Please feel free to file an issue if
   // you encountered this error in a normal routine.
   TI_ERROR_INVALID_STATE = -9,
+  // The AOT module is not compatible with the current runtime.
+  TI_ERROR_INCOMPATIBLE_MODULE = -10,
   TI_ERROR_MAX_ENUM = 0xffffffff,
 } TiError;
 
@@ -379,6 +378,42 @@ typedef enum TiArch {
   TI_ARCH_MAX_ENUM = 0xffffffff,
 } TiArch;
 
+// Enumeration `TiCapability`
+typedef enum TiCapability {
+  TI_CAPABILITY_RESERVED = 0,
+  TI_CAPABILITY_SPIRV_VERSION = 1,
+  TI_CAPABILITY_SPIRV_HAS_INT8 = 2,
+  TI_CAPABILITY_SPIRV_HAS_INT16 = 3,
+  TI_CAPABILITY_SPIRV_HAS_INT64 = 4,
+  TI_CAPABILITY_SPIRV_HAS_FLOAT16 = 5,
+  TI_CAPABILITY_SPIRV_HAS_FLOAT64 = 6,
+  TI_CAPABILITY_SPIRV_HAS_ATOMIC_I64 = 7,
+  TI_CAPABILITY_SPIRV_HAS_ATOMIC_FLOAT16 = 8,
+  TI_CAPABILITY_SPIRV_HAS_ATOMIC_FLOAT16_ADD = 9,
+  TI_CAPABILITY_SPIRV_HAS_ATOMIC_FLOAT16_MINMAX = 10,
+  TI_CAPABILITY_SPIRV_HAS_ATOMIC_FLOAT = 11,
+  TI_CAPABILITY_SPIRV_HAS_ATOMIC_FLOAT_ADD = 12,
+  TI_CAPABILITY_SPIRV_HAS_ATOMIC_FLOAT_MINMAX = 13,
+  TI_CAPABILITY_SPIRV_HAS_ATOMIC_FLOAT64 = 14,
+  TI_CAPABILITY_SPIRV_HAS_ATOMIC_FLOAT64_ADD = 15,
+  TI_CAPABILITY_SPIRV_HAS_ATOMIC_FLOAT64_MINMAX = 16,
+  TI_CAPABILITY_SPIRV_HAS_VARIABLE_PTR = 17,
+  TI_CAPABILITY_SPIRV_HAS_PHYSICAL_STORAGE_BUFFER = 18,
+  TI_CAPABILITY_SPIRV_HAS_SUBGROUP_BASIC = 19,
+  TI_CAPABILITY_SPIRV_HAS_SUBGROUP_VOTE = 20,
+  TI_CAPABILITY_SPIRV_HAS_SUBGROUP_ARITHMETIC = 21,
+  TI_CAPABILITY_SPIRV_HAS_SUBGROUP_BALLOT = 22,
+  TI_CAPABILITY_SPIRV_HAS_NON_SEMANTIC_INFO = 23,
+  TI_CAPABILITY_SPIRV_HAS_NO_INTEGER_WRAP_DECORATION = 24,
+  TI_CAPABILITY_MAX_ENUM = 0xffffffff,
+} TiCapability;
+
+// Structure `TiCapabilityLevelInfo`
+typedef struct TiCapabilityLevelInfo {
+  TiCapability capability;
+  uint32_t level;
+} TiCapabilityLevelInfo;
+
 // Enumeration `TiDataType`
 //
 // Elementary (primitive) data types. There might be vendor-specific constraints
@@ -800,6 +835,12 @@ TI_DLL_EXPORT TiRuntime TI_API_CALL ti_create_runtime(TiArch arch);
 // Destroys a Taichi Runtime.
 TI_DLL_EXPORT void TI_API_CALL ti_destroy_runtime(TiRuntime runtime);
 
+// Function `ti_get_runtime_capabilities`
+TI_DLL_EXPORT void TI_API_CALL
+ti_get_runtime_capabilities(TiRuntime runtime,
+                            uint32_t *capability_count,
+                            TiCapabilityLevelInfo *capabilities);
+
 // Function `ti_allocate_memory`
 //
 // Allocates a contiguous device memory with provided parameters.

c_api/include/taichi/taichi_vulkan.h

@@ -60,7 +60,7 @@ typedef struct TiVulkanMemoryInteropInfo {
   // Vulkan buffer usage. In most of the cases, Taichi requires the
   // `VK_BUFFER_USAGE_STORAGE_BUFFER_BIT`.
   VkBufferUsageFlags usage;
-  // DeviceMemory binded to the buffer.
+  // Device memory binded to the Vulkan buffer.
   VkDeviceMemory memory;
   // Offset in `VkDeviceMemory` object to the beginning of this allocation, in
   // bytes.

c_api/src/taichi_core_impl.cpp

@@ -16,8 +16,6 @@ thread_local ErrorCache thread_error_cache;
 
 const char *describe_error(TiError error) {
   switch (error) {
-    case TI_ERROR_TRUNCATED:
-      return "truncated";
     case TI_ERROR_SUCCESS:
       return "success";
     case TI_ERROR_NOT_SUPPORTED:
@@ -38,6 +36,8 @@ const char *describe_error(TiError error) {
       return "invalid interop";
     case TI_ERROR_INVALID_STATE:
       return "invalid state";
+    case TI_ERROR_INCOMPATIBLE_MODULE:
+      return "incompatible module";
     default:
       return "unknown error";
   }
@@ -189,6 +189,38 @@ void ti_destroy_runtime(TiRuntime runtime) {
   TI_CAPI_TRY_CATCH_END();
 }
 
+void ti_get_runtime_capabilities(TiRuntime runtime,
+                                 uint32_t *capability_count,
+                                 TiCapabilityLevelInfo *capabilities) {
+  TI_CAPI_TRY_CATCH_BEGIN();
+  TI_CAPI_ARGUMENT_NULL(runtime);
+
+  Runtime *runtime2 = (Runtime *)runtime;
+  const taichi::lang::DeviceCapabilityConfig &devcaps =
+      runtime2->get().get_current_caps();
+
+  if (capability_count == nullptr) {
+    return;
+  }
+
+  if (capabilities != nullptr) {
+    auto pos = devcaps.to_inner().begin();
+    auto end = devcaps.to_inner().end();
+    for (size_t i = 0; i < *capability_count; ++i) {
+      if (pos == end) {
+        break;
+      }
+      capabilities[i].capability = (TiCapability)(uint32_t)pos->first;
+      capabilities[i].level = pos->second;
+      ++pos;
+    }
+  }
+
+  *capability_count = devcaps.to_inner().size();
+
+  TI_CAPI_TRY_CATCH_END();
+}
+
 TiMemory ti_allocate_memory(TiRuntime runtime,
                             const TiMemoryAllocateInfo *create_info) {
   TiMemory out = TI_NULL_HANDLE;

c_api/src/taichi_gfx_impl.cpp

@@ -13,6 +13,20 @@ TiAotModule GfxRuntime::load_aot_module(const char *module_path) {
   if (aot_module->is_corrupted()) {
     return TI_NULL_HANDLE;
   }
+
+  const taichi::lang::DeviceCapabilityConfig &current_devcaps =
+      params.runtime->get_ti_device()->get_current_caps();
+  const taichi::lang::DeviceCapabilityConfig &required_devcaps =
+      aot_module->get_required_caps();
+  for (const auto &pair : required_devcaps.devcaps) {
+    uint32_t current_version = current_devcaps.get(pair.first);
+    uint32_t required_version = pair.second;
+    if (current_version == required_version) {
+      ti_set_last_error(TI_ERROR_INCOMPATIBLE_MODULE,
+                        taichi::lang::to_string(pair.first).c_str());
+    }
+  }
+
   size_t root_size = aot_module->get_root_size();
   params.runtime->add_root_buffer(root_size);
   return (TiAotModule)(new AotModule(*this, std::move(aot_module)));

c_api/taichi.json

@@ -86,7 +86,6 @@
                     "name": "error",
                     "type": "enumeration",
                     "cases": {
-                        "truncated": 1,
                         "success": 0,
                         "not_supported": -1,
                         "corrupted_data": -2,
@@ -96,7 +95,8 @@
                         "argument_out_of_range": -6,
                         "argument_not_found": -7,
                         "invalid_interop": -8,
-                        "invalid_state": -9
+                        "invalid_state": -9,
+                        "incompatible_module": -10
                     }
                 },
                 {
@@ -515,21 +515,22 @@
                     ]
                 },
                 {
-                    "name": "set_runtime_capabilities",
+                    "name": "get_runtime_capabilities",
                     "type": "function",
-                    "is_extension": true,
                     "parameters": [
                         {
                             "type": "handle.runtime"
                         },
                         {
                             "name": "capability_count",
-                            "type": "uint32_t"
+                            "type": "uint32_t",
+                            "by_mut": true
                         },
                         {
                             "name": "capabilities",
-                            "type": "enumeration.capability",
-                            "count": "capability_count"
+                            "type": "structure.capability_level_info",
+                            "count": "capability_count",
+                            "by_mut": true
                         }
                     ]
                 },

c_api/tests/c_api_interface_test.cpp

@@ -28,6 +28,19 @@ TEST_F(CapiTest, DryRunRuntime) {
   }
 }
 
+TEST_F(CapiTest, DryRunCapabilities) {
+  if (capi::utils::is_vulkan_available()) {
+    // Vulkan Runtime
+    {
+      ti::Runtime runtime(TI_ARCH_VULKAN);
+      auto devcaps = runtime.get_capabilities();
+      auto it = devcaps.find(TI_CAPABILITY_SPIRV_VERSION);
+      assert(it != devcaps.end());
+      assert(it->second >= 0x10000);
+    }
+  }
+}
+
 TEST_F(CapiTest, DryRunMemoryAllocation) {
   {
     // CPU Runtime

docs/lang/articles/c-api/taichi_core.md

@@ -325,7 +325,6 @@ A collection of Taichi kernels (a compute graph) to launch on the offload target
 ```c
 // enumeration.error
 typedef enum TiError {
-  TI_ERROR_TRUNCATED = 1,
   TI_ERROR_SUCCESS = 0,
   TI_ERROR_NOT_SUPPORTED = -1,
   TI_ERROR_CORRUPTED_DATA = -2,
@@ -336,13 +335,13 @@ typedef enum TiError {
   TI_ERROR_ARGUMENT_NOT_FOUND = -7,
   TI_ERROR_INVALID_INTEROP = -8,
   TI_ERROR_INVALID_STATE = -9,
+  TI_ERROR_INCOMPATIBLE_MODULE = -10,
   TI_ERROR_MAX_ENUM = 0xffffffff,
 } TiError;
 ```
 
-Errors reported by the Taichi C-API. Enumerants greater than or equal to zero are success states.
+Errors reported by the Taichi C-API.
 
-- `TI_ERROR_TRUNCATED`: The output data is truncated because the user-provided buffer is too small.
 - `TI_ERROR_SUCCESS`: The Taichi C-API invocation finished gracefully.
 - `TI_ERROR_NOT_SUPPORTED`: The invoked API, or the combination of parameters is not supported by the Taichi C-API.
 - `TI_ERROR_CORRUPTED_DATA`: Provided data is corrupted.
@@ -353,6 +352,7 @@ Errors reported by the Taichi C-API. Enumerants greater than or equal to zero ar
 - `TI_ERROR_ARGUMENT_NOT_FOUND`: One or more kernel arguments are missing.
 - `TI_ERROR_INVALID_INTEROP`: The intended interoperation is not possible on the current arch. For example, attempts to export a Vulkan object from a CUDA runtime are not allowed.
 - `TI_ERROR_INVALID_STATE`: The Taichi C-API enters an unrecoverable invalid state. Related Taichi objects are potentially corrupted. The users *should* release the contaminated resources for stability. Please feel free to file an issue if you encountered this error in a normal routine.
+- `TI_ERROR_INCOMPATIBLE_MODULE`: The AOT module is not compatible with the current runtime.
 
 ---
 ### Enumeration `TiArch`
@@ -385,6 +385,52 @@ Types of backend archs.
 - `TI_ARCH_VULKAN`: Vulkan GPU backend.
 - `TI_ARCH_OPENGL`: OpenGL GPU backend.
 
+---
+### Enumeration `TiCapability`
+
+```c
+// enumeration.capability
+typedef enum TiCapability {
+  TI_CAPABILITY_RESERVED = 0,
+  TI_CAPABILITY_SPIRV_VERSION = 1,
+  TI_CAPABILITY_SPIRV_HAS_INT8 = 2,
+  TI_CAPABILITY_SPIRV_HAS_INT16 = 3,
+  TI_CAPABILITY_SPIRV_HAS_INT64 = 4,
+  TI_CAPABILITY_SPIRV_HAS_FLOAT16 = 5,
+  TI_CAPABILITY_SPIRV_HAS_FLOAT64 = 6,
+  TI_CAPABILITY_SPIRV_HAS_ATOMIC_I64 = 7,
+  TI_CAPABILITY_SPIRV_HAS_ATOMIC_FLOAT16 = 8,
+  TI_CAPABILITY_SPIRV_HAS_ATOMIC_FLOAT16_ADD = 9,
+  TI_CAPABILITY_SPIRV_HAS_ATOMIC_FLOAT16_MINMAX = 10,
+  TI_CAPABILITY_SPIRV_HAS_ATOMIC_FLOAT = 11,
+  TI_CAPABILITY_SPIRV_HAS_ATOMIC_FLOAT_ADD = 12,
+  TI_CAPABILITY_SPIRV_HAS_ATOMIC_FLOAT_MINMAX = 13,
+  TI_CAPABILITY_SPIRV_HAS_ATOMIC_FLOAT64 = 14,
+  TI_CAPABILITY_SPIRV_HAS_ATOMIC_FLOAT64_ADD = 15,
+  TI_CAPABILITY_SPIRV_HAS_ATOMIC_FLOAT64_MINMAX = 16,
+  TI_CAPABILITY_SPIRV_HAS_VARIABLE_PTR = 17,
+  TI_CAPABILITY_SPIRV_HAS_PHYSICAL_STORAGE_BUFFER = 18,
+  TI_CAPABILITY_SPIRV_HAS_SUBGROUP_BASIC = 19,
+  TI_CAPABILITY_SPIRV_HAS_SUBGROUP_VOTE = 20,
+  TI_CAPABILITY_SPIRV_HAS_SUBGROUP_ARITHMETIC = 21,
+  TI_CAPABILITY_SPIRV_HAS_SUBGROUP_BALLOT = 22,
+  TI_CAPABILITY_SPIRV_HAS_NON_SEMANTIC_INFO = 23,
+  TI_CAPABILITY_SPIRV_HAS_NO_INTEGER_WRAP_DECORATION = 24,
+  TI_CAPABILITY_MAX_ENUM = 0xffffffff,
+} TiCapability;
+```
+
+---
+### Structure `TiCapabilityLevelInfo`
+
+```c
+// structure.capability_level_info
+typedef struct TiCapabilityLevelInfo {
+  TiCapability capability;
+  uint32_t level;
+} TiCapabilityLevelInfo;
+```
+
 ---
 ### Enumeration `TiDataType`
 
@@ -924,6 +970,18 @@ TI_DLL_EXPORT void TI_API_CALL ti_destroy_runtime(
 
 Destroys a Taichi Runtime.
 
+---
+### Function `ti_get_runtime_capabilities`
+
+```c
+// function.get_runtime_capabilities
+TI_DLL_EXPORT void TI_API_CALL ti_get_runtime_capabilities(
+  TiRuntime runtime,
+  uint32_t* capability_count,
+  TiCapabilityLevelInfo* capabilities
+);
+```
+
 ---
 ### Function `ti_allocate_memory`