Skip to content

Commit 510b461

Browse files
nihuinihui
authored
discover VK_NV_cooperative_vector (#6117)
1 parent 9cdc02b commit 510b461

3 files changed

Lines changed: 176 additions & 0 deletions

File tree

‎src/gpu.cpp‎

Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -255,6 +255,9 @@ PFN_vkCreateAndroidSurfaceKHR vkCreateAndroidSurfaceKHR = 0;
255255
// VK_NV_cooperative_matrix
256256
PFN_vkGetPhysicalDeviceCooperativeMatrixPropertiesNV vkGetPhysicalDeviceCooperativeMatrixPropertiesNV = 0;
257257

258+
// VK_NV_cooperative_vector
259+
PFN_vkGetPhysicalDeviceCooperativeVectorPropertiesNV vkGetPhysicalDeviceCooperativeVectorPropertiesNV = 0;
260+
258261
class GpuInfoPrivate
259262
{
260263
public:
@@ -358,6 +361,7 @@ class GpuInfoPrivate
358361
int support_VK_ANDROID_external_memory_android_hardware_buffer;
359362
#endif // __ANDROID_API__ >= 26
360363
int support_VK_NV_cooperative_matrix;
364+
int support_VK_NV_cooperative_vector;
361365

362366
// extension features
363367
void* queryExtensionFeatures;
@@ -374,6 +378,7 @@ class GpuInfoPrivate
374378
VkPhysicalDeviceShaderSubgroupRotateFeaturesKHR queryShaderSubgroupRotateFeatures;
375379
VkPhysicalDeviceShaderAtomicFloatFeaturesEXT queryShaderAtomicFloatFeatures;
376380
VkPhysicalDeviceShaderAtomicFloat2FeaturesEXT queryShaderAtomicFloat2Features;
381+
VkPhysicalDeviceCooperativeVectorFeaturesNV queryCooperativeVectorFeaturesNV;
377382

378383
// extension properties
379384
void* queryExtensionProperties;
@@ -382,10 +387,12 @@ class GpuInfoPrivate
382387
VkPhysicalDeviceSubgroupProperties querySubgroupProperties;
383388
VkPhysicalDeviceDriverPropertiesKHR queryDriverProperties;
384389
VkPhysicalDeviceSubgroupSizeControlPropertiesEXT querySubgroupSizeControlProperties;
390+
VkPhysicalDeviceCooperativeVectorPropertiesNV queryCooperativeVectorPropertiesNV;
385391

386392
// extension sub properties
387393
std::vector<VkCooperativeMatrixPropertiesKHR> queryCooperativeMatrixSubProperties;
388394
std::vector<VkCooperativeMatrixPropertiesNV> queryCooperativeMatrixSubPropertiesNV;
395+
std::vector<VkCooperativeVectorPropertiesNV> queryCooperativeVectorSubPropertiesNV;
389396
};
390397

391398
void GpuInfoPrivate::query_features()
@@ -707,6 +714,7 @@ int GpuInfoPrivate::query_extensions()
707714
support_VK_ANDROID_external_memory_android_hardware_buffer = 0;
708715
#endif // __ANDROID_API__ >= 26
709716
support_VK_NV_cooperative_matrix = 0;
717+
support_VK_NV_cooperative_vector = 0;
710718
for (uint32_t j = 0; j < deviceExtensionPropertyCount; j++)
711719
{
712720
const VkExtensionProperties& exp = deviceExtensionProperties[j];
@@ -794,6 +802,8 @@ int GpuInfoPrivate::query_extensions()
794802
#endif // __ANDROID_API__ >= 26
795803
else if (strcmp(exp.extensionName, "VK_NV_cooperative_matrix") == 0)
796804
support_VK_NV_cooperative_matrix = exp.specVersion;
805+
else if (strcmp(exp.extensionName, "VK_NV_cooperative_vector") == 0)
806+
support_VK_NV_cooperative_vector = exp.specVersion;
797807
}
798808

799809
if (support_VK_KHR_buffer_device_address)
@@ -945,6 +955,16 @@ void GpuInfoPrivate::query_extension_features()
945955
queryExtensionFeatures = &queryShaderAtomicFloat2Features;
946956
}
947957

958+
// query nv cooperative vector
959+
memset(&queryCooperativeVectorFeaturesNV, 0, sizeof(queryCooperativeVectorFeaturesNV));
960+
queryCooperativeVectorFeaturesNV.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_COOPERATIVE_VECTOR_FEATURES_NV;
961+
queryCooperativeVectorFeaturesNV.pNext = 0;
962+
if (support_VK_NV_cooperative_vector)
963+
{
964+
queryCooperativeVectorFeaturesNV.pNext = queryExtensionFeatures;
965+
queryExtensionFeatures = &queryCooperativeVectorFeaturesNV;
966+
}
967+
948968
if (support_VK_KHR_get_physical_device_properties2)
949969
{
950970
VkPhysicalDeviceFeatures2KHR queryFeatures;
@@ -1050,6 +1070,16 @@ void GpuInfoPrivate::query_extension_properties()
10501070
queryExtensionProperties = &querySubgroupSizeControlProperties;
10511071
}
10521072

1073+
// query nv cooperative vector
1074+
memset(&queryCooperativeVectorPropertiesNV, 0, sizeof(queryCooperativeVectorPropertiesNV));
1075+
queryCooperativeVectorPropertiesNV.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_COOPERATIVE_VECTOR_PROPERTIES_NV;
1076+
queryCooperativeVectorPropertiesNV.pNext = 0;
1077+
if (support_VK_NV_cooperative_vector)
1078+
{
1079+
queryCooperativeVectorPropertiesNV.pNext = queryExtensionProperties;
1080+
queryExtensionProperties = &queryCooperativeVectorPropertiesNV;
1081+
}
1082+
10531083
if (support_VK_KHR_get_physical_device_properties2)
10541084
{
10551085
VkPhysicalDeviceProperties2KHR queryProperties;
@@ -1197,6 +1227,37 @@ void GpuInfoPrivate::query_extension_properties()
11971227
}
11981228
}
11991229

1230+
// query supported cooperative vector types and operations
1231+
queryCooperativeVectorSubPropertiesNV.clear();
1232+
if (support_VK_NV_cooperative_vector && queryCooperativeVectorFeaturesNV.cooperativeVector)
1233+
{
1234+
uint32_t propertyCount = 0;
1235+
VkResult ret = vkGetPhysicalDeviceCooperativeVectorPropertiesNV(physicalDevice, &propertyCount, 0);
1236+
if (ret != VK_SUCCESS)
1237+
{
1238+
NCNN_LOGE("vkGetPhysicalDeviceCooperativeVectorPropertiesNV failed %d", ret);
1239+
}
1240+
1241+
queryCooperativeVectorSubPropertiesNV.resize(propertyCount);
1242+
for (uint32_t j = 0; j < propertyCount; j++)
1243+
{
1244+
memset(&queryCooperativeVectorSubPropertiesNV[j], 0, sizeof(queryCooperativeVectorSubPropertiesNV[j]));
1245+
queryCooperativeVectorSubPropertiesNV[j].sType = VK_STRUCTURE_TYPE_COOPERATIVE_VECTOR_PROPERTIES_NV;
1246+
queryCooperativeVectorSubPropertiesNV[j].pNext = 0;
1247+
}
1248+
ret = vkGetPhysicalDeviceCooperativeVectorPropertiesNV(physicalDevice, &propertyCount, queryCooperativeVectorSubPropertiesNV.data());
1249+
if (ret != VK_SUCCESS)
1250+
{
1251+
NCNN_LOGE("vkGetPhysicalDeviceCooperativeVectorPropertiesNV failed %d", ret);
1252+
}
1253+
1254+
for (uint32_t j = 0; j < propertyCount; j++)
1255+
{
1256+
const VkCooperativeVectorPropertiesNV& cvp = queryCooperativeVectorSubPropertiesNV[j];
1257+
// NCNN_LOGE("cvp %d %d %d %d %d %d", cvp.inputType, cvp.inputInterpretation, cvp.matrixInterpretation, cvp.biasInterpretation, cvp.resultType, cvp.transpose);
1258+
}
1259+
}
1260+
12001261
if (queryDriverProperties.driverID == VK_DRIVER_ID_MESA_TURNIP)
12011262
{
12021263
// turnip crash when compiling large shader with full subgroup
@@ -1766,6 +1827,11 @@ int GpuInfo::support_VK_NV_cooperative_matrix() const
17661827
return d->support_VK_NV_cooperative_matrix;
17671828
}
17681829

1830+
int GpuInfo::support_VK_NV_cooperative_vector() const
1831+
{
1832+
return d->support_VK_NV_cooperative_vector;
1833+
}
1834+
17691835
const void* GpuInfo::queryExtensionFeatures() const
17701836
{
17711837
return d->queryExtensionFeatures;
@@ -1841,6 +1907,11 @@ const VkPhysicalDeviceShaderAtomicFloat2FeaturesEXT& GpuInfo::queryShaderAtomicF
18411907
return d->queryShaderAtomicFloat2Features;
18421908
}
18431909

1910+
const VkPhysicalDeviceCooperativeVectorFeaturesNV& GpuInfo::queryCooperativeVectorFeaturesNV() const
1911+
{
1912+
return d->queryCooperativeVectorFeaturesNV;
1913+
}
1914+
18441915
const void* GpuInfo::queryExtensionProperties() const
18451916
{
18461917
return d->queryExtensionProperties;
@@ -1866,6 +1937,11 @@ const VkPhysicalDeviceSubgroupSizeControlPropertiesEXT& GpuInfo::querySubgroupSi
18661937
return d->querySubgroupSizeControlProperties;
18671938
}
18681939

1940+
const VkPhysicalDeviceCooperativeVectorPropertiesNV& GpuInfo::queryCooperativeVectorPropertiesNV() const
1941+
{
1942+
return d->queryCooperativeVectorPropertiesNV;
1943+
}
1944+
18691945
const std::vector<VkCooperativeMatrixPropertiesKHR>& GpuInfo::queryCooperativeMatrixSubProperties() const
18701946
{
18711947
return d->queryCooperativeMatrixSubProperties;
@@ -1876,6 +1952,11 @@ const std::vector<VkCooperativeMatrixPropertiesNV>& GpuInfo::queryCooperativeMat
18761952
return d->queryCooperativeMatrixSubPropertiesNV;
18771953
}
18781954

1955+
const std::vector<VkCooperativeVectorPropertiesNV>& GpuInfo::queryCooperativeVectorSubPropertiesNV() const
1956+
{
1957+
return d->queryCooperativeVectorSubPropertiesNV;
1958+
}
1959+
18791960
static int init_instance_core()
18801961
{
18811962
vkAllocateCommandBuffers = (PFN_vkAllocateCommandBuffers)vkGetInstanceProcAddr(g_instance, "vkAllocateCommandBuffers");
@@ -2026,6 +2107,11 @@ static int init_instance_extension()
20262107
vkGetPhysicalDeviceCooperativeMatrixPropertiesNV = (PFN_vkGetPhysicalDeviceCooperativeMatrixPropertiesNV)vkGetInstanceProcAddr(g_instance, "vkGetPhysicalDeviceCooperativeMatrixPropertiesNV");
20272108
}
20282109

2110+
// VK_NV_cooperative_vector
2111+
{
2112+
vkGetPhysicalDeviceCooperativeVectorPropertiesNV = (PFN_vkGetPhysicalDeviceCooperativeVectorPropertiesNV)vkGetInstanceProcAddr(g_instance, "vkGetPhysicalDeviceCooperativeVectorPropertiesNV");
2113+
}
2114+
20292115
return 0;
20302116
}
20312117

@@ -4068,6 +4154,12 @@ int VulkanDevice::init_device_extension()
40684154
}
40694155
#endif // __ANDROID_API__ >= 26
40704156

4157+
if (info.support_VK_NV_cooperative_vector())
4158+
{
4159+
vkCmdConvertCooperativeVectorMatrixNV = (PFN_vkCmdConvertCooperativeVectorMatrixNV)vkGetDeviceProcAddr(d->device, "vkCmdConvertCooperativeVectorMatrixNV");
4160+
vkConvertCooperativeVectorMatrixNV = (PFN_vkConvertCooperativeVectorMatrixNV)vkGetDeviceProcAddr(d->device, "vkConvertCooperativeVectorMatrixNV");
4161+
}
4162+
40714163
return 0;
40724164
}
40734165

‎src/gpu.h‎

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -182,6 +182,9 @@ extern PFN_vkCreateAndroidSurfaceKHR vkCreateAndroidSurfaceKHR;
182182
// VK_NV_cooperative_matrix
183183
extern PFN_vkGetPhysicalDeviceCooperativeMatrixPropertiesNV vkGetPhysicalDeviceCooperativeMatrixPropertiesNV;
184184

185+
// VK_NV_cooperative_vector
186+
extern PFN_vkGetPhysicalDeviceCooperativeVectorPropertiesNV vkGetPhysicalDeviceCooperativeVectorPropertiesNV;
187+
185188
// get info
186189
NCNN_EXPORT int get_gpu_count();
187190
NCNN_EXPORT int get_default_gpu_index();
@@ -347,6 +350,7 @@ class NCNN_EXPORT GpuInfo
347350
int support_VK_ANDROID_external_memory_android_hardware_buffer() const;
348351
#endif // __ANDROID_API__ >= 26
349352
int support_VK_NV_cooperative_matrix() const;
353+
int support_VK_NV_cooperative_vector() const;
350354

351355
// extension features
352356
const void* queryExtensionFeatures() const;
@@ -363,6 +367,7 @@ class NCNN_EXPORT GpuInfo
363367
const VkPhysicalDeviceShaderSubgroupRotateFeaturesKHR& queryShaderSubgroupRotateFeatures() const;
364368
const VkPhysicalDeviceShaderAtomicFloatFeaturesEXT& queryShaderAtomicFloatFeatures() const;
365369
const VkPhysicalDeviceShaderAtomicFloat2FeaturesEXT& queryShaderAtomicFloat2Features() const;
370+
const VkPhysicalDeviceCooperativeVectorFeaturesNV& queryCooperativeVectorFeaturesNV() const;
366371

367372
// extension properties
368373
const void* queryExtensionProperties() const;
@@ -371,10 +376,12 @@ class NCNN_EXPORT GpuInfo
371376
const VkPhysicalDeviceSubgroupProperties& querySubgroupProperties() const;
372377
const VkPhysicalDeviceDriverPropertiesKHR& queryDriverProperties() const;
373378
const VkPhysicalDeviceSubgroupSizeControlPropertiesEXT& querySubgroupSizeControlProperties() const;
379+
const VkPhysicalDeviceCooperativeVectorPropertiesNV& queryCooperativeVectorPropertiesNV() const;
374380

375381
// extension sub properties
376382
const std::vector<VkCooperativeMatrixPropertiesKHR>& queryCooperativeMatrixSubProperties() const;
377383
const std::vector<VkCooperativeMatrixPropertiesNV>& queryCooperativeMatrixSubPropertiesNV() const;
384+
const std::vector<VkCooperativeVectorPropertiesNV>& queryCooperativeVectorSubPropertiesNV() const;
378385

379386
private:
380387
GpuInfo(const GpuInfo&);
@@ -499,6 +506,10 @@ class NCNN_EXPORT VulkanDevice
499506
PFN_vkGetMemoryAndroidHardwareBufferANDROID vkGetMemoryAndroidHardwareBufferANDROID;
500507
#endif // __ANDROID_API__ >= 26
501508

509+
// VK_NV_cooperative_vector
510+
PFN_vkCmdConvertCooperativeVectorMatrixNV vkCmdConvertCooperativeVectorMatrixNV;
511+
PFN_vkConvertCooperativeVectorMatrixNV vkConvertCooperativeVectorMatrixNV;
512+
502513
protected:
503514
// device extension
504515
int init_device_extension();

‎src/vulkan_header_fix.h‎

Lines changed: 73 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1502,4 +1502,77 @@ typedef struct VkPhysicalDeviceShaderFloatControls2Features
15021502
typedef VkPhysicalDeviceShaderFloatControls2Features VkPhysicalDeviceShaderFloatControls2FeaturesKHR;
15031503
#endif // VK_KHR_shader_float_controls2
15041504

1505+
#ifndef VK_NV_cooperative_vector
1506+
#define VK_NV_cooperative_vector 1
1507+
#define VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_COOPERATIVE_VECTOR_FEATURES_NV (VkStructureType)1000491000
1508+
#define VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_COOPERATIVE_VECTOR_PROPERTIES_NV (VkStructureType)1000491001
1509+
#define VK_STRUCTURE_TYPE_COOPERATIVE_VECTOR_PROPERTIES_NV (VkStructureType)1000491002
1510+
#define VK_STRUCTURE_TYPE_CONVERT_COOPERATIVE_VECTOR_MATRIX_INFO_NV (VkStructureType)1000491004
1511+
typedef enum VkCooperativeVectorMatrixLayoutNV
1512+
{
1513+
VK_COOPERATIVE_VECTOR_MATRIX_LAYOUT_ROW_MAJOR_NV = 0,
1514+
VK_COOPERATIVE_VECTOR_MATRIX_LAYOUT_COLUMN_MAJOR_NV = 1,
1515+
VK_COOPERATIVE_VECTOR_MATRIX_LAYOUT_INFERENCING_OPTIMAL_NV = 2,
1516+
VK_COOPERATIVE_VECTOR_MATRIX_LAYOUT_TRAINING_OPTIMAL_NV = 3,
1517+
VK_COOPERATIVE_VECTOR_MATRIX_LAYOUT_MAX_ENUM_NV = 0x7FFFFFFF
1518+
} VkCooperativeVectorMatrixLayoutNV;
1519+
typedef union VkDeviceOrHostAddressConstKHR
1520+
{
1521+
VkDeviceAddress deviceAddress;
1522+
const void* hostAddress;
1523+
} VkDeviceOrHostAddressConstKHR;
1524+
typedef union VkDeviceOrHostAddressKHR
1525+
{
1526+
VkDeviceAddress deviceAddress;
1527+
void* hostAddress;
1528+
} VkDeviceOrHostAddressKHR;
1529+
typedef struct VkPhysicalDeviceCooperativeVectorPropertiesNV
1530+
{
1531+
VkStructureType sType;
1532+
void* pNext;
1533+
VkShaderStageFlags cooperativeVectorSupportedStages;
1534+
VkBool32 cooperativeVectorTrainingFloat16Accumulation;
1535+
VkBool32 cooperativeVectorTrainingFloat32Accumulation;
1536+
uint32_t maxCooperativeVectorComponents;
1537+
} VkPhysicalDeviceCooperativeVectorPropertiesNV;
1538+
typedef struct VkPhysicalDeviceCooperativeVectorFeaturesNV
1539+
{
1540+
VkStructureType sType;
1541+
void* pNext;
1542+
VkBool32 cooperativeVector;
1543+
VkBool32 cooperativeVectorTraining;
1544+
} VkPhysicalDeviceCooperativeVectorFeaturesNV;
1545+
typedef struct VkCooperativeVectorPropertiesNV
1546+
{
1547+
VkStructureType sType;
1548+
void* pNext;
1549+
VkComponentTypeKHR inputType;
1550+
VkComponentTypeKHR inputInterpretation;
1551+
VkComponentTypeKHR matrixInterpretation;
1552+
VkComponentTypeKHR biasInterpretation;
1553+
VkComponentTypeKHR resultType;
1554+
VkBool32 transpose;
1555+
} VkCooperativeVectorPropertiesNV;
1556+
typedef struct VkConvertCooperativeVectorMatrixInfoNV
1557+
{
1558+
VkStructureType sType;
1559+
const void* pNext;
1560+
size_t srcSize;
1561+
VkDeviceOrHostAddressConstKHR srcData;
1562+
size_t* pDstSize;
1563+
VkDeviceOrHostAddressKHR dstData;
1564+
VkComponentTypeKHR srcComponentType;
1565+
VkComponentTypeKHR dstComponentType;
1566+
uint32_t numRows;
1567+
uint32_t numColumns;
1568+
VkCooperativeVectorMatrixLayoutNV srcLayout;
1569+
size_t srcStride;
1570+
VkCooperativeVectorMatrixLayoutNV dstLayout;
1571+
size_t dstStride;
1572+
} VkConvertCooperativeVectorMatrixInfoNV;
1573+
typedef VkResult(VKAPI_PTR* PFN_vkGetPhysicalDeviceCooperativeVectorPropertiesNV)(VkPhysicalDevice physicalDevice, uint32_t* pPropertyCount, VkCooperativeVectorPropertiesNV* pProperties);
1574+
typedef VkResult(VKAPI_PTR* PFN_vkConvertCooperativeVectorMatrixNV)(VkDevice device, const VkConvertCooperativeVectorMatrixInfoNV* pInfo);
1575+
typedef void(VKAPI_PTR* PFN_vkCmdConvertCooperativeVectorMatrixNV)(VkCommandBuffer commandBuffer, uint32_t infoCount, const VkConvertCooperativeVectorMatrixInfoNV* pInfos);
1576+
#endif // VK_NV_cooperative_vector
1577+
15051578
#endif // NCNN_VULKAN_HEADER_FIX_H

0 commit comments

Comments
 (0)