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

.gitmodules

@@ -47,6 +47,3 @@
 [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,12 +92,7 @@ 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"
-  "taichi/backends/vulkan/snode_struct_compiler.cpp"
-  "taichi/backends/vulkan/snode_struct_compiler.h"
-)
+file(GLOB TAICHI_VULKAN_REQUIRED_SOURCE "taichi/backends/vulkan/runtime.h" "taichi/backends/vulkan/runtime.cpp")
 
 list(REMOVE_ITEM TAICHI_CORE_SOURCE ${TAICHI_BACKEND_SOURCE})
 
@@ -261,10 +256,6 @@ 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,9 +4,8 @@
 
 
 @ti.kernel
-def p() -> ti.f32:
+def p():
     print(42)
-    return 40 + 2
 
 
-print(p())
+p()

taichi/backends/device.h

@@ -12,27 +12,25 @@ 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,
-  // 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,
+  vk_has_spv_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::spirv_has_atomic_float_add) &&
+      if (device_->get_cap(DeviceCapability::vk_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::spirv_version) > 0x10300) {
+    if (device_->get_cap(DeviceCapability::vk_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,23 +1140,12 @@ 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)
@@ -1390,15 +1379,12 @@ 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,6 +187,8 @@ 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_;
@@ -215,6 +217,10 @@ 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;
@@ -375,10 +381,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::spirv_version, 0x10000);
+  ti_device_->set_cap(DeviceCapability::vk_spirv_version, 0x10000);
 
   if (physical_device_properties.apiVersion >= VK_API_VERSION_1_1) {
-    ti_device_->set_cap(DeviceCapability::spirv_version, 0x10300);
+    ti_device_->set_cap(DeviceCapability::vk_spirv_version, 0x10300);
   }
 
   // Detect extensions
@@ -421,7 +427,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::spirv_version, 0x10400);
+      ti_device_->set_cap(DeviceCapability::vk_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);
@@ -448,15 +454,15 @@ void EmbeddedVulkanDevice::create_logical_device() {
 
   if (device_supported_features.shaderInt16) {
     device_features.shaderInt16 = true;
-    ti_device_->set_cap(DeviceCapability::spirv_has_int16, true);
+    ti_device_->set_cap(DeviceCapability::vk_has_int16, true);
   }
   if (device_supported_features.shaderInt64) {
     device_features.shaderInt64 = true;
-    ti_device_->set_cap(DeviceCapability::spirv_has_int64, true);
+    ti_device_->set_cap(DeviceCapability::vk_has_int64, true);
   }
   if (device_supported_features.shaderFloat64) {
     device_features.shaderFloat64 = true;
-    ti_device_->set_cap(DeviceCapability::spirv_has_float64, true);
+    ti_device_->set_cap(DeviceCapability::vk_has_float64, true);
   }
   if (device_supported_features.wideLines) {
     device_features.wideLines = true;
@@ -493,7 +499,7 @@ void EmbeddedVulkanDevice::create_logical_device() {
 
       if (variable_ptr_feature.variablePointers &&
           variable_ptr_feature.variablePointersStorageBuffer) {
-        ti_device_->set_cap(DeviceCapability::spirv_has_variable_ptr, true);
+        ti_device_->set_cap(DeviceCapability::vk_has_spv_variable_ptr, true);
       }
       *pNextEnd = &variable_ptr_feature;
       pNextEnd = &variable_ptr_feature.pNext;
@@ -505,14 +511,13 @@ void EmbeddedVulkanDevice::create_logical_device() {
       vkGetPhysicalDeviceFeatures2KHR(physical_device_, &features2);
 
       if (shader_atomic_float_feature.shaderBufferFloat32AtomicAdd) {
-        ti_device_->set_cap(DeviceCapability::spirv_has_atomic_float_add, true);
+        ti_device_->set_cap(DeviceCapability::vk_has_atomic_float_add, true);
       } else if (shader_atomic_float_feature.shaderBufferFloat64AtomicAdd) {
-        ti_device_->set_cap(DeviceCapability::spirv_has_atomic_float64_add,
-                            true);
+        ti_device_->set_cap(DeviceCapability::vk_has_atomic_float64_add, true);
       } else if (shader_atomic_float_feature.shaderBufferFloat32Atomics) {
-        ti_device_->set_cap(DeviceCapability::spirv_has_atomic_float, true);
+        ti_device_->set_cap(DeviceCapability::vk_has_atomic_float, true);
       } else if (shader_atomic_float_feature.shaderBufferFloat64Atomics) {
-        ti_device_->set_cap(DeviceCapability::spirv_has_atomic_float64, true);
+        ti_device_->set_cap(DeviceCapability::vk_has_atomic_float64, true);
       }
       *pNextEnd = &shader_atomic_float_feature;
       pNextEnd = &shader_atomic_float_feature.pNext;
@@ -524,9 +529,9 @@ void EmbeddedVulkanDevice::create_logical_device() {
       vkGetPhysicalDeviceFeatures2KHR(physical_device_, &features2);
 
       if (shader_f16_i8_feature.shaderFloat16) {
-        ti_device_->set_cap(DeviceCapability::spirv_has_float16, true);
+        ti_device_->set_cap(DeviceCapability::vk_has_float16, true);
       } else if (shader_f16_i8_feature.shaderInt8) {
-        ti_device_->set_cap(DeviceCapability::spirv_has_int8, true);
+        ti_device_->set_cap(DeviceCapability::vk_has_int8, true);
       }
       *pNextEnd = &shader_f16_i8_feature;
       pNextEnd = &shader_f16_i8_feature.pNext;
@@ -549,6 +554,8 @@ 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,14 +61,36 @@ class EmbeddedVulkanDevice {
   explicit EmbeddedVulkanDevice(const Params &params);
   ~EmbeddedVulkanDevice();
 
-  const VulkanDevice *device() const {
-    return ti_device_.get();
+  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_;
+  }
+
+  Device *get_ti_device() const;
+
  private:
   void create_instance();
   void setup_debug_messenger();
@@ -81,12 +103,19 @@ 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,10 +40,6 @@ 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,8 +35,6 @@ class VulkanLoader {
   VkDevice vulkan_device_{VK_NULL_HANDLE};
 };
 
-bool is_vulkan_api_available();
-
 }  // namespace vulkan
 }  // namespace lang
 }  // namespace taichi