Revert "[misc] Revert #3175 #3164 (#3185)" (#3212)

This reverts commit 30b45f8.

.gitmodules

@@ -47,3 +47,6 @@
 [submodule "external/SPIRV-Tools"]
 	path = external/SPIRV-Tools
 	url = https://github.com/KhronosGroup/SPIRV-Tools
+[submodule "external/SPIRV-Cross"]
+	path = external/SPIRV-Cross
+	url = https://github.com/KhronosGroup/SPIRV-Cross

cmake/TaichiCore.cmake

@@ -92,7 +92,12 @@ file(GLOB TAICHI_OPENGL_REQUIRED_SOURCE
   "taichi/backends/opengl/codegen_opengl.*"
   "taichi/backends/opengl/struct_opengl.*"
 )
-file(GLOB TAICHI_VULKAN_REQUIRED_SOURCE "taichi/backends/vulkan/runtime.h" "taichi/backends/vulkan/runtime.cpp")
+file(GLOB TAICHI_VULKAN_REQUIRED_SOURCE
+  "taichi/backends/vulkan/runtime.h"
+  "taichi/backends/vulkan/runtime.cpp"
+  "taichi/backends/vulkan/snode_struct_compiler.cpp"
+  "taichi/backends/vulkan/snode_struct_compiler.h"
+)
 
 list(REMOVE_ITEM TAICHI_CORE_SOURCE ${TAICHI_BACKEND_SOURCE})
 
@@ -256,6 +261,10 @@ else()
     message(STATUS "TI_WITH_CUDA_TOOLKIT = OFF")
 endif()
 
+add_subdirectory(external/SPIRV-Cross)
+target_include_directories(${CORE_LIBRARY_NAME} PRIVATE external/SPIRV-Cross)
+target_link_libraries(${CORE_LIBRARY_NAME} spirv-cross-glsl spirv-cross-core)
+
 if (TI_WITH_VULKAN)
     # Vulkan libs
     # https://cmake.org/cmake/help/latest/module/FindVulkan.html

examples/minimal.py

@@ -4,8 +4,9 @@
 
 
 @ti.kernel
-def p():
+def p() -> ti.f32:
     print(42)
+    return 40 + 2
 
 
-p()
+print(p())

taichi/backends/device.h

@@ -12,25 +12,27 @@ namespace lang {
 // Or the backend runtime itself
 // Capabilities are per-device
 enum class DeviceCapability : uint32_t {
+  // Vulkan Caps
   vk_api_version,
-  vk_spirv_version,
   vk_has_physical_features2,
-  vk_has_int8,
-  vk_has_int16,
-  vk_has_int64,
-  vk_has_float16,
-  vk_has_float64,
   vk_has_external_memory,
-  vk_has_atomic_i64,
-  vk_has_atomic_float,  // load, store, exchange
-  vk_has_atomic_float_add,
-  vk_has_atomic_float_minmax,
-  vk_has_atomic_float64,  // load, store, exchange
-  vk_has_atomic_float64_add,
-  vk_has_atomic_float64_minmax,
   vk_has_surface,
   vk_has_presentation,
-  vk_has_spv_variable_ptr,
+  // SPIR-V Caps
+  spirv_version,
+  spirv_has_int8,
+  spirv_has_int16,
+  spirv_has_int64,
+  spirv_has_float16,
+  spirv_has_float64,
+  spirv_has_atomic_i64,
+  spirv_has_atomic_float,  // load, store, exchange
+  spirv_has_atomic_float_add,
+  spirv_has_atomic_float_minmax,
+  spirv_has_atomic_float64,  // load, store, exchange
+  spirv_has_atomic_float64_add,
+  spirv_has_atomic_float64_minmax,
+  spirv_has_variable_ptr,
 };
 
 class Device;

taichi/backends/vulkan/codegen_vulkan.cpp

@@ -825,7 +825,7 @@ class TaskCodegen : public IRVisitor {
     spirv::Value data = ir_->query_value(stmt->val->raw_name());
     spirv::Value val;
     if (dt->is_primitive(PrimitiveTypeID::f32)) {
-      if (device_->get_cap(DeviceCapability::vk_has_atomic_float_add) &&
+      if (device_->get_cap(DeviceCapability::spirv_has_atomic_float_add) &&
           stmt->op_type == AtomicOpType::add && is_compiled_struct) {
         val = ir_->make_value(
             spv::OpAtomicFAddEXT, ir_->get_primitive_type(dt), addr_ptr,
@@ -1024,7 +1024,7 @@ class TaskCodegen : public IRVisitor {
         task_attribs_.advisory_num_threads_per_group, 1, 1};
     ir_->set_work_group_size(group_size);
     std::vector<spirv::Value> buffers;
-    if (device_->get_cap(DeviceCapability::vk_spirv_version) > 0x10300) {
+    if (device_->get_cap(DeviceCapability::spirv_version) > 0x10300) {
       for (const auto &bb : task_attribs_.buffer_binds) {
         const auto it = buffer_value_map_.find(bb.buffer);
         if (it != buffer_value_map_.end()) {
@@ -1140,12 +1140,23 @@ class TaskCodegen : public IRVisitor {
     // For now, |total_invocs_name| is equal to |total_elems|. Once we support
     // dynamic range, they will be different.
     // https://www.khronos.org/opengl/wiki/Compute_Shader#Inputs
+
+    // HLSL & WGSL cross compilers do not support this builtin
+    /*
     spirv::Value total_invocs = ir_->cast(
         ir_->i32_type(),
         ir_->mul(ir_->get_num_work_groups(0),
                  ir_->uint_immediate_number(
                      ir_->u32_type(),
                      task_attribs_.advisory_num_threads_per_group, true)));
+                     */
+    const int group_x = (task_attribs_.advisory_total_num_threads +
+                         task_attribs_.advisory_num_threads_per_group - 1) /
+                        task_attribs_.advisory_num_threads_per_group;
+    spirv::Value total_invocs = ir_->uint_immediate_number(
+        ir_->i32_type(), group_x * task_attribs_.advisory_num_threads_per_group,
+        false);
+
     ir_->debug(spv::OpName, total_invocs, total_invocs_name);
 
     // Must get init label after making value(to make sure they are correct)
@@ -1379,12 +1390,15 @@ class KernelCodegen {
 
       // Enable to dump SPIR-V assembly of kernels
 #if 0
-       std::string spirv_asm;
-       spirv_tools_->Disassemble(optimized_spv, &spirv_asm);
-      TI_WARN("SPIR-V Assembly dump for {} :\n{}\n\n",params_.ti_kernel_name, spirv_asm);
-
-      std::ofstream fout((params_.ti_kernel_name).c_str(), std::ios::binary | std::ios::out);
-      fout.write(reinterpret_cast<const char*>(task_res.spirv_code.data()), task_res.spirv_code.size() * sizeof(uint32_t));
+      std::string spirv_asm;
+      spirv_tools_->Disassemble(optimized_spv, &spirv_asm);
+      TI_WARN("SPIR-V Assembly dump for {} :\n{}\n\n", params_.ti_kernel_name,
+              spirv_asm);
+
+      std::ofstream fout((params_.ti_kernel_name).c_str(),
+                         std::ios::binary | std::ios::out);
+      fout.write(reinterpret_cast<const char *>(task_res.spirv_code.data()),
+                 task_res.spirv_code.size() * sizeof(uint32_t));
       fout.close();
 #endif
 

taichi/backends/vulkan/embedded_device.cpp

@@ -187,8 +187,6 @@ EmbeddedVulkanDevice::EmbeddedVulkanDevice(
   pick_physical_device();
   create_logical_device();
 
-  // TODO: Change the ownership hierarchy, the taichi Device class should be at
-  // the top level
   {
     VulkanDevice::Params params;
     params.instance = instance_;
@@ -217,10 +215,6 @@ EmbeddedVulkanDevice::~EmbeddedVulkanDevice() {
   vkDestroyInstance(instance_, kNoVkAllocCallbacks);
 }
 
-Device *EmbeddedVulkanDevice::get_ti_device() const {
-  return ti_device_.get();
-}
-
 void EmbeddedVulkanDevice::create_instance() {
   VkApplicationInfo app_info{};
   app_info.sType = VK_STRUCTURE_TYPE_APPLICATION_INFO;
@@ -381,10 +375,10 @@ void EmbeddedVulkanDevice::create_logical_device() {
   vkGetPhysicalDeviceProperties(physical_device_, &physical_device_properties);
   ti_device_->set_cap(DeviceCapability::vk_api_version,
                       physical_device_properties.apiVersion);
-  ti_device_->set_cap(DeviceCapability::vk_spirv_version, 0x10000);
+  ti_device_->set_cap(DeviceCapability::spirv_version, 0x10000);
 
   if (physical_device_properties.apiVersion >= VK_API_VERSION_1_1) {
-    ti_device_->set_cap(DeviceCapability::vk_spirv_version, 0x10300);
+    ti_device_->set_cap(DeviceCapability::spirv_version, 0x10300);
   }
 
   // Detect extensions
@@ -427,7 +421,7 @@ void EmbeddedVulkanDevice::create_logical_device() {
     } else if (name == VK_KHR_SYNCHRONIZATION_2_EXTENSION_NAME) {
       enabled_extensions.push_back(ext.extensionName);
     } else if (name == VK_KHR_SPIRV_1_4_EXTENSION_NAME) {
-      ti_device_->set_cap(DeviceCapability::vk_spirv_version, 0x10400);
+      ti_device_->set_cap(DeviceCapability::spirv_version, 0x10400);
       enabled_extensions.push_back(ext.extensionName);
     } else if (name == VK_KHR_EXTERNAL_MEMORY_CAPABILITIES_EXTENSION_NAME) {
       ti_device_->set_cap(DeviceCapability::vk_has_external_memory, true);
@@ -454,15 +448,15 @@ void EmbeddedVulkanDevice::create_logical_device() {
 
   if (device_supported_features.shaderInt16) {
     device_features.shaderInt16 = true;
-    ti_device_->set_cap(DeviceCapability::vk_has_int16, true);
+    ti_device_->set_cap(DeviceCapability::spirv_has_int16, true);
   }
   if (device_supported_features.shaderInt64) {
     device_features.shaderInt64 = true;
-    ti_device_->set_cap(DeviceCapability::vk_has_int64, true);
+    ti_device_->set_cap(DeviceCapability::spirv_has_int64, true);
   }
   if (device_supported_features.shaderFloat64) {
     device_features.shaderFloat64 = true;
-    ti_device_->set_cap(DeviceCapability::vk_has_float64, true);
+    ti_device_->set_cap(DeviceCapability::spirv_has_float64, true);
   }
   if (device_supported_features.wideLines) {
     device_features.wideLines = true;
@@ -499,7 +493,7 @@ void EmbeddedVulkanDevice::create_logical_device() {
 
       if (variable_ptr_feature.variablePointers &&
           variable_ptr_feature.variablePointersStorageBuffer) {
-        ti_device_->set_cap(DeviceCapability::vk_has_spv_variable_ptr, true);
+        ti_device_->set_cap(DeviceCapability::spirv_has_variable_ptr, true);
       }
       *pNextEnd = &variable_ptr_feature;
       pNextEnd = &variable_ptr_feature.pNext;
@@ -511,13 +505,14 @@ void EmbeddedVulkanDevice::create_logical_device() {
       vkGetPhysicalDeviceFeatures2KHR(physical_device_, &features2);
 
       if (shader_atomic_float_feature.shaderBufferFloat32AtomicAdd) {
-        ti_device_->set_cap(DeviceCapability::vk_has_atomic_float_add, true);
+        ti_device_->set_cap(DeviceCapability::spirv_has_atomic_float_add, true);
       } else if (shader_atomic_float_feature.shaderBufferFloat64AtomicAdd) {
-        ti_device_->set_cap(DeviceCapability::vk_has_atomic_float64_add, true);
+        ti_device_->set_cap(DeviceCapability::spirv_has_atomic_float64_add,
+                            true);
       } else if (shader_atomic_float_feature.shaderBufferFloat32Atomics) {
-        ti_device_->set_cap(DeviceCapability::vk_has_atomic_float, true);
+        ti_device_->set_cap(DeviceCapability::spirv_has_atomic_float, true);
       } else if (shader_atomic_float_feature.shaderBufferFloat64Atomics) {
-        ti_device_->set_cap(DeviceCapability::vk_has_atomic_float64, true);
+        ti_device_->set_cap(DeviceCapability::spirv_has_atomic_float64, true);
       }
       *pNextEnd = &shader_atomic_float_feature;
       pNextEnd = &shader_atomic_float_feature.pNext;
@@ -529,9 +524,9 @@ void EmbeddedVulkanDevice::create_logical_device() {
       vkGetPhysicalDeviceFeatures2KHR(physical_device_, &features2);
 
       if (shader_f16_i8_feature.shaderFloat16) {
-        ti_device_->set_cap(DeviceCapability::vk_has_float16, true);
+        ti_device_->set_cap(DeviceCapability::spirv_has_float16, true);
       } else if (shader_f16_i8_feature.shaderInt8) {
-        ti_device_->set_cap(DeviceCapability::vk_has_int8, true);
+        ti_device_->set_cap(DeviceCapability::spirv_has_int8, true);
       }
       *pNextEnd = &shader_f16_i8_feature;
       pNextEnd = &shader_f16_i8_feature.pNext;
@@ -554,8 +549,6 @@ void EmbeddedVulkanDevice::create_logical_device() {
   if (params_.is_for_ui) {
     vkGetDeviceQueue(device_, queue_family_indices_.graphics_family.value(), 0,
                      &graphics_queue_);
-    vkGetDeviceQueue(device_, queue_family_indices_.graphics_family.value(), 0,
-                     &present_queue_);
   }
 
   vkGetDeviceQueue(device_, queue_family_indices_.compute_family.value(), 0,

taichi/backends/vulkan/embedded_device.h

@@ -61,36 +61,14 @@ class EmbeddedVulkanDevice {
   explicit EmbeddedVulkanDevice(const Params &params);
   ~EmbeddedVulkanDevice();
 
-  VkInstance instance() {
-    return instance_;
-  }
-
-  VulkanDevice *device() {
-    return ti_device_.get();
-  }
-
   const VulkanDevice *device() const {
     return ti_device_.get();
   }
 
-  VkPhysicalDevice physical_device() const {
-    return physical_device_;
-  }
-
-  VkSurfaceKHR surface() const {
-    return surface_;
-  }
-
-  VkInstance instance() const {
-    return instance_;
-  }
-
-  const VulkanQueueFamilyIndices &queue_family_indices() const {
-    return queue_family_indices_;
+  VulkanDevice *device() {
+    return ti_device_.get();
   }
 
-  Device *get_ti_device() const;
-
  private:
   void create_instance();
   void setup_debug_messenger();
@@ -103,19 +81,12 @@ class EmbeddedVulkanDevice {
   VkPhysicalDevice physical_device_{VK_NULL_HANDLE};
   VulkanQueueFamilyIndices queue_family_indices_;
   VkDevice device_{VK_NULL_HANDLE};
-  // TODO: It's probably not right to put these per-queue things here. However,
-  // in Taichi we only use a single queue on a single device (i.e. a single CUDA
-  // stream), so it doesn't make a difference.
+
   VkQueue compute_queue_{VK_NULL_HANDLE};
   VkQueue graphics_queue_{VK_NULL_HANDLE};
-  VkQueue present_queue_{VK_NULL_HANDLE};
 
   VkSurfaceKHR surface_{VK_NULL_HANDLE};
 
-  // TODO: Shall we have dedicated command pools for COMPUTE and TRANSFER
-  // commands, respectively?
-  VkCommandPool command_pool_{VK_NULL_HANDLE};
-
   std::unique_ptr<VulkanDevice> ti_device_{nullptr};
 
   Params params_;

taichi/backends/vulkan/loader.cpp

@@ -40,6 +40,10 @@ PFN_vkVoidFunction VulkanLoader::load_function(const char *name) {
   return result;
 }
 
+bool is_vulkan_api_available() {
+  return VulkanLoader::instance().init();
+}
+
 }  // namespace vulkan
 }  // namespace lang
 }  // namespace taichi

taichi/backends/vulkan/loader.h

@@ -35,6 +35,8 @@ class VulkanLoader {
   VkDevice vulkan_device_{VK_NULL_HANDLE};
 };
 
+bool is_vulkan_api_available();
+
 }  // namespace vulkan
 }  // namespace lang
 }  // namespace taichi