Check the available VK extensions before using CoopVec APIs in GFX (#6849)

* Check the available VK extensions before using CoopVec APIs in GFX

* Remove a redundant request for cooperative vector extension for vk

9 files changed, 121 insertions, 21 deletions

diff --git a/tools/gfx/debug-layer/debug-device.cpp b/tools/gfx/debug-layer/debug-device.cpp
index c5798710a..2a3eb4c52 100644
--- a/tools/gfx/debug-layer/debug-device.cpp
+++ b/tools/gfx/debug-layer/debug-device.cpp
@@ -685,6 +685,14 @@ Result DebugDevice::getTextureRowAlignment(size_t* outAlignment)
     return baseObject->getTextureRowAlignment(outAlignment);
 }
 
+Result DebugDevice::getCooperativeVectorProperties(
+    CooperativeVectorProperties* properties,
+    uint32_t* propertyCount)
+{
+    SLANG_GFX_API_FUNC;
+    return baseObject->getCooperativeVectorProperties(properties, propertyCount);
+}
+
 Result DebugDevice::createShaderTable(const IShaderTable::Desc& desc, IShaderTable** outTable)
 {
     SLANG_GFX_API_FUNC;
diff --git a/tools/gfx/debug-layer/debug-device.h b/tools/gfx/debug-layer/debug-device.h
index a4debd2e7..d1e96e909 100644
--- a/tools/gfx/debug-layer/debug-device.h
+++ b/tools/gfx/debug-layer/debug-device.h
@@ -161,6 +161,11 @@ public:
         size_t* outSize,
         size_t* outAlignment) override;
     virtual SLANG_NO_THROW Result SLANG_MCALL getTextureRowAlignment(size_t* outAlignment) override;
+
+    virtual SLANG_NO_THROW Result SLANG_MCALL getCooperativeVectorProperties(
+        CooperativeVectorProperties* properties,
+        uint32_t* propertyCount) override;
+
     virtual SLANG_NO_THROW Result SLANG_MCALL
     createShaderTable(const IShaderTable::Desc& desc, IShaderTable** outTable) override;
 };
diff --git a/tools/gfx/renderer-shared.cpp b/tools/gfx/renderer-shared.cpp
index ad32a012a..cb25079fb 100644
--- a/tools/gfx/renderer-shared.cpp
+++ b/tools/gfx/renderer-shared.cpp
@@ -770,6 +770,14 @@ Result RendererBase::getTextureRowAlignment(Size* outAlignment)
     return SLANG_E_NOT_AVAILABLE;
 }
 
+Result RendererBase::getCooperativeVectorProperties(
+    CooperativeVectorProperties* properties,
+    uint32_t* propertyCount)
+{
+    *propertyCount = 0;
+    return SLANG_E_NOT_AVAILABLE;
+}
+
 Result RendererBase::getShaderObjectLayout(
     slang::ISession* session,
     slang::TypeReflection* type,
diff --git a/tools/gfx/renderer-shared.h b/tools/gfx/renderer-shared.h
index 18270c6cd..e1086aca9 100644
--- a/tools/gfx/renderer-shared.h
+++ b/tools/gfx/renderer-shared.h
@@ -1338,6 +1338,11 @@ public:
     // Provides a default implementation that returns SLANG_E_NOT_AVAILABLE.
     virtual SLANG_NO_THROW Result SLANG_MCALL getTextureRowAlignment(size_t* outAlignment) override;
 
+    // Provides a default implementation that returns SLANG_E_NOT_AVAILABLE.
+    virtual SLANG_NO_THROW Result SLANG_MCALL getCooperativeVectorProperties(
+        CooperativeVectorProperties* properties,
+        uint32_t* propertyCount) override;
+
     Result getEntryPointCodeFromShaderCache(
         slang::IComponentType* program,
         SlangInt entryPointIndex,
@@ -1392,6 +1397,7 @@ protected:
 
 protected:
     Slang::List<Slang::String> m_features;
+    std::vector<CooperativeVectorProperties> m_cooperativeVectorProperties;
 
 public:
     SlangContext slangContext;
diff --git a/tools/gfx/vulkan/vk-api.cpp b/tools/gfx/vulkan/vk-api.cpp
index 09b8e92f5..dbf32345f 100644
--- a/tools/gfx/vulkan/vk-api.cpp
+++ b/tools/gfx/vulkan/vk-api.cpp
@@ -86,17 +86,6 @@ Slang::Result VulkanApi::initPhysicalDevice(VkPhysicalDevice physicalDevice)
     vkGetPhysicalDeviceFeatures(m_physicalDevice, &m_deviceFeatures);
     vkGetPhysicalDeviceMemoryProperties(m_physicalDevice, &m_deviceMemoryProperties);
 
-    if (vkGetPhysicalDeviceCooperativeVectorPropertiesNV)
-    {
-        uint32_t nProps = 0;
-        vkGetPhysicalDeviceCooperativeVectorPropertiesNV(m_physicalDevice, &nProps, nullptr);
-        m_cooperativeVectorProperties.setCount(nProps);
-        vkGetPhysicalDeviceCooperativeVectorPropertiesNV(
-            m_physicalDevice,
-            &nProps,
-            m_cooperativeVectorProperties.getBuffer());
-    }
-
     return SLANG_OK;
 }
 
diff --git a/tools/gfx/vulkan/vk-device.cpp b/tools/gfx/vulkan/vk-device.cpp
index 981b5abac..b9ca76493 100644
--- a/tools/gfx/vulkan/vk-device.cpp
+++ b/tools/gfx/vulkan/vk-device.cpp
@@ -730,6 +730,12 @@ Result DeviceImpl::initVulkanInstanceAndDevice(
                 "ray-tracing-validation");
         }
 
+        SIMPLE_EXTENSION_FEATURE(
+            extendedFeatures.cooperativeVectorFeatures,
+            cooperativeVector,
+            VK_NV_COOPERATIVE_VECTOR_EXTENSION_NAME,
+            "cooperative-vector");
+
 #undef SIMPLE_EXTENSION_FEATURE
 
         if (extendedFeatures.vulkan12Features.shaderBufferInt64Atomics)
@@ -752,16 +758,6 @@ Result DeviceImpl::initVulkanInstanceAndDevice(
             deviceCreateInfo.pNext = &extendedFeatures.vulkan12Features;
         }
 
-        if (extendedFeatures.cooperativeVectorFeatures.cooperativeVector)
-        {
-            deviceExtensions.add(VK_NV_COOPERATIVE_VECTOR_EXTENSION_NAME);
-
-            extendedFeatures.cooperativeVectorFeatures.pNext = (void*)deviceCreateInfo.pNext;
-            deviceCreateInfo.pNext = &extendedFeatures.cooperativeVectorFeatures;
-
-            m_features.add("cooperative-vector");
-        }
-
         VkPhysicalDeviceProperties2 extendedProps = {
             VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_PROPERTIES_2};
         VkPhysicalDeviceRayTracingPipelinePropertiesKHR rtProps = {
@@ -1548,6 +1544,47 @@ Result DeviceImpl::getTextureRowAlignment(Size* outAlignment)
     return SLANG_OK;
 }
 
+Result DeviceImpl::getCooperativeVectorProperties(
+    CooperativeVectorProperties* properties,
+    uint32_t* propertyCount)
+{
+    if (!m_api.m_extendedFeatures.cooperativeVectorFeatures.cooperativeVector ||
+        !m_api.vkGetPhysicalDeviceCooperativeVectorPropertiesNV)
+        return SLANG_E_NOT_AVAILABLE;
+
+    if (m_cooperativeVectorProperties.empty())
+    {
+        uint32_t vkPropertyCount = 0;
+        m_api.vkGetPhysicalDeviceCooperativeVectorPropertiesNV(
+            m_api.m_physicalDevice,
+            &vkPropertyCount,
+            nullptr);
+        std::vector<VkCooperativeVectorPropertiesNV> vkProperties(vkPropertyCount);
+        SLANG_VK_RETURN_ON_FAIL(m_api.vkGetPhysicalDeviceCooperativeVectorPropertiesNV(
+            m_api.m_physicalDevice,
+            &vkPropertyCount,
+            vkProperties.data()));
+        for (const auto& vkProps : vkProperties)
+        {
+            CooperativeVectorProperties props;
+            props.inputType =
+                VulkanUtil::translateCooperativeVectorComponentType(vkProps.inputType);
+            props.inputInterpretation =
+                VulkanUtil::translateCooperativeVectorComponentType(vkProps.inputInterpretation);
+            props.matrixInterpretation =
+                VulkanUtil::translateCooperativeVectorComponentType(vkProps.matrixInterpretation);
+            props.biasInterpretation =
+                VulkanUtil::translateCooperativeVectorComponentType(vkProps.biasInterpretation);
+            props.resultType =
+                VulkanUtil::translateCooperativeVectorComponentType(vkProps.resultType);
+            props.transpose = vkProps.transpose;
+            m_cooperativeVectorProperties.push_back(props);
+        }
+    }
+
+    return RendererBase::getCooperativeVectorProperties(properties, propertyCount);
+}
+
 Result DeviceImpl::createTextureResource(
     const ITextureResource::Desc& descIn,
     const ITextureResource::SubresourceData* initData,
diff --git a/tools/gfx/vulkan/vk-device.h b/tools/gfx/vulkan/vk-device.h
index 06e19ad7e..922117eb0 100644
--- a/tools/gfx/vulkan/vk-device.h
+++ b/tools/gfx/vulkan/vk-device.h
@@ -129,6 +129,10 @@ public:
 
     virtual SLANG_NO_THROW Result SLANG_MCALL getTextureRowAlignment(Size* outAlignment) override;
 
+    virtual SLANG_NO_THROW Result SLANG_MCALL getCooperativeVectorProperties(
+        CooperativeVectorProperties* properties,
+        uint32_t* propertyCount) override;
+
     virtual SLANG_NO_THROW Result SLANG_MCALL
     createFence(const IFence::Desc& desc, IFence** outFence) override;
 
diff --git a/tools/gfx/vulkan/vk-util.cpp b/tools/gfx/vulkan/vk-util.cpp
index 2490d48ce..4f45973f4 100644
--- a/tools/gfx/vulkan/vk-util.cpp
+++ b/tools/gfx/vulkan/vk-util.cpp
@@ -608,6 +608,46 @@ VkSamplerReductionMode VulkanUtil::translateReductionOp(TextureReductionOp op)
     }
 }
 
+CooperativeVectorComponentType VulkanUtil::translateCooperativeVectorComponentType(
+    VkComponentTypeKHR type)
+{
+    switch (type)
+    {
+    case VK_COMPONENT_TYPE_FLOAT16_KHR:
+        return CooperativeVectorComponentType::Float16;
+    case VK_COMPONENT_TYPE_FLOAT32_KHR:
+        return CooperativeVectorComponentType::Float32;
+    case VK_COMPONENT_TYPE_FLOAT64_KHR:
+        return CooperativeVectorComponentType::Float64;
+    case VK_COMPONENT_TYPE_SINT8_KHR:
+        return CooperativeVectorComponentType::SInt8;
+    case VK_COMPONENT_TYPE_SINT16_KHR:
+        return CooperativeVectorComponentType::SInt16;
+    case VK_COMPONENT_TYPE_SINT32_KHR:
+        return CooperativeVectorComponentType::SInt32;
+    case VK_COMPONENT_TYPE_SINT64_KHR:
+        return CooperativeVectorComponentType::SInt64;
+    case VK_COMPONENT_TYPE_UINT8_KHR:
+        return CooperativeVectorComponentType::UInt8;
+    case VK_COMPONENT_TYPE_UINT16_KHR:
+        return CooperativeVectorComponentType::UInt16;
+    case VK_COMPONENT_TYPE_UINT32_KHR:
+        return CooperativeVectorComponentType::UInt32;
+    case VK_COMPONENT_TYPE_UINT64_KHR:
+        return CooperativeVectorComponentType::UInt64;
+    case VK_COMPONENT_TYPE_SINT8_PACKED_NV:
+        return CooperativeVectorComponentType::SInt8Packed;
+    case VK_COMPONENT_TYPE_UINT8_PACKED_NV:
+        return CooperativeVectorComponentType::UInt8Packed;
+    case VK_COMPONENT_TYPE_FLOAT_E4M3_NV:
+        return CooperativeVectorComponentType::FloatE4M3;
+    case VK_COMPONENT_TYPE_FLOAT_E5M2_NV:
+        return CooperativeVectorComponentType::FloatE5M2;
+    default:
+        return CooperativeVectorComponentType(0);
+    }
+}
+
 /* static */ Slang::Result VulkanUtil::handleFail(VkResult res)
 {
     if (res != VK_SUCCESS)
diff --git a/tools/gfx/vulkan/vk-util.h b/tools/gfx/vulkan/vk-util.h
index 4bad710a0..458979a2f 100644
--- a/tools/gfx/vulkan/vk-util.h
+++ b/tools/gfx/vulkan/vk-util.h
@@ -130,6 +130,9 @@ struct VulkanUtil
     static VkStencilOpState translateStencilState(DepthStencilOpDesc desc);
 
     static VkSamplerReductionMode translateReductionOp(TextureReductionOp op);
+
+    static CooperativeVectorComponentType translateCooperativeVectorComponentType(
+        VkComponentTypeKHR type);
 };
 
 struct AccelerationStructureBuildGeometryInfoBuilder