Add support for 8x8x16 cooperative matrices (Intel Arc)

Author: sergcpp

internal/RendererVK.cpp

@@ -213,9 +213,9 @@ bool Ray::Vk::Context::Init(ILog *log, const VulkanDevice &vk_device, const Vulk
         return false;
     }
 
-    CheckVkPhysicalDeviceFeatures(api_, physical_device_, device_properties_, mem_properties_, graphics_family_index_,
-                                  raytracing_supported_, ray_query_supported_, fp16_supported_, int64_supported_,
-                                  int64_atomics_supported_, coop_matrix_size_, pageable_memory_supported_);
+    CheckVkPhysicalDeviceFeatures(api_, physical_device_, device_properties_, mem_properties_, coop_mat_properties_,
+                                  graphics_family_index_, raytracing_supported_, ray_query_supported_, fp16_supported_,
+                                  int64_supported_, int64_atomics_supported_, pageable_memory_supported_);
 
     // mask out unsupported stages
     if (!raytracing_supported_) {
@@ -227,7 +227,7 @@ bool Ray::Vk::Context::Init(ILog *log, const VulkanDevice &vk_device, const Vulk
 
     if (!external_ && !InitVkDevice(api_, device_, physical_device_, graphics_family_index_, raytracing_supported_,
                                     ray_query_supported_, fp16_supported_, int64_supported_, int64_atomics_supported_,
-                                    coop_matrix_size_[0] != -1, pageable_memory_supported_, log)) {
+                                    coop_mat_properties_.MSize != 0, pageable_memory_supported_, log)) {
         return false;
     }
 
@@ -559,14 +559,12 @@ bool Ray::Vk::Context::ChooseVkPhysicalDevice(const Api &api, VkPhysicalDevice &
     return true;
 }
 
-void Ray::Vk::Context::CheckVkPhysicalDeviceFeatures(const Api &api, VkPhysicalDevice &physical_device,
-                                                     VkPhysicalDeviceProperties &out_device_properties,
-                                                     VkPhysicalDeviceMemoryProperties &out_mem_properties,
-                                                     uint32_t &out_graphics_family_index,
-                                                     bool &out_raytracing_supported, bool &out_ray_query_supported,
-                                                     bool &out_shader_fp16_supported, bool &out_shader_int64_supported,
-                                                     bool &out_int64_atomics_supported, int out_coop_matrix_size[3],
-                                                     bool &out_pageable_memory_supported) {
+void Ray::Vk::Context::CheckVkPhysicalDeviceFeatures(
+    const Api &api, VkPhysicalDevice &physical_device, VkPhysicalDeviceProperties &out_device_properties,
+    VkPhysicalDeviceMemoryProperties &out_mem_properties, VkCooperativeMatrixPropertiesKHR &out_coop_mat_properties,
+    uint32_t &out_graphics_family_index, bool &out_raytracing_supported, bool &out_ray_query_supported,
+    bool &out_shader_fp16_supported, bool &out_shader_int64_supported, bool &out_int64_atomics_supported,
+    bool &out_pageable_memory_supported) {
     api.vkGetPhysicalDeviceProperties(physical_device, &out_device_properties);
     api.vkGetPhysicalDeviceMemoryProperties(physical_device, &out_mem_properties);
 
@@ -594,7 +592,7 @@ void Ray::Vk::Context::CheckVkPhysicalDeviceFeatures(const Api &api, VkPhysicalD
          shader_buf_int64_atomics_supported = false, memory_priority_supported = false,
          pageable_memory_supported = false;
 
-    int coop_matrix_size[3] = {-1, -1, -1};
+    VkCooperativeMatrixPropertiesKHR coop_mat_properties = {};
 
     { // check for features support
         uint32_t extension_count;
@@ -671,18 +669,28 @@ void Ray::Vk::Context::CheckVkPhysicalDeviceFeatures(const Api &api, VkPhysicalD
                                                                   coop_matrix_props.data());
 
             bool found = false;
+            // We try to find 16x16x16 size and F16 accumulator first (NV and AMD)
             for (const VkCooperativeMatrixPropertiesKHR &p : coop_matrix_props) {
                 if (p.AType == VK_COMPONENT_TYPE_FLOAT16_KHR && p.BType == VK_COMPONENT_TYPE_FLOAT16_KHR &&
                     p.CType == VK_COMPONENT_TYPE_FLOAT16_KHR && p.ResultType == VK_COMPONENT_TYPE_FLOAT16_KHR &&
                     p.MSize == 16 && p.NSize == 16 && p.KSize == 16 && p.scope == VK_SCOPE_SUBGROUP_KHR) {
-                    coop_matrix_size[0] = 16;
-                    coop_matrix_size[1] = 16;
-                    coop_matrix_size[2] = 16;
+                    coop_mat_properties = p;
                     found = true;
                     break;
                 }
             }
-            coop_matrix_supported &= found;
+            if (!found) {
+                // Try to find 8x8x16 size and F32 accumulator (Intel)
+                for (const VkCooperativeMatrixPropertiesKHR &p : coop_matrix_props) {
+                    if (p.AType == VK_COMPONENT_TYPE_FLOAT16_KHR && p.BType == VK_COMPONENT_TYPE_FLOAT16_KHR &&
+                        p.CType == VK_COMPONENT_TYPE_FLOAT32_KHR && p.ResultType == VK_COMPONENT_TYPE_FLOAT32_KHR &&
+                        p.MSize == 8 && p.NSize == 8 && p.KSize == 16 && p.scope == VK_SCOPE_SUBGROUP_KHR) {
+                        coop_mat_properties = p;
+                        found = true;
+                        break;
+                    }
+                }
+            }
         }
     }
 
@@ -691,7 +699,7 @@ void Ray::Vk::Context::CheckVkPhysicalDeviceFeatures(const Api &api, VkPhysicalD
     out_shader_fp16_supported = (shader_fp16_supported && storage_fp16_supported);
     out_shader_int64_supported = shader_int64_supported;
     out_int64_atomics_supported = shader_buf_int64_atomics_supported;
-    memcpy(out_coop_matrix_size, coop_matrix_size, 3 * sizeof(int));
+    out_coop_mat_properties = coop_mat_properties;
     out_pageable_memory_supported = (memory_priority_supported && pageable_memory_supported);
 }
 

internal/Vk/ContextVK.cpp

@@ -27,6 +27,7 @@ class Context {
     VkPhysicalDeviceLimits phys_device_limits_ = {};
     VkPhysicalDeviceProperties device_properties_ = {};
     VkPhysicalDeviceMemoryProperties mem_properties_ = {};
+    VkCooperativeMatrixPropertiesKHR coop_mat_properties_ = {};
     uint32_t graphics_family_index_ = 0;
 
     VkDevice device_ = {};
@@ -46,8 +47,6 @@ class Context {
 
     bool subgroup_supported_ = false;
 
-    int coop_matrix_size_[3] = {-1, -1, -1};
-
     bool pageable_memory_supported_ = false;
 
     uint32_t supported_stages_mask_ = 0xffffffff;
@@ -94,14 +93,14 @@ class Context {
     bool int64_supported() const { return int64_supported_; }
     bool int64_atomics_supported() const { return int64_atomics_supported_; }
     bool subgroup_supported() const { return subgroup_supported_; }
-    const int *coop_matrix_size() const { return coop_matrix_size_; }
 
     uint32_t supported_stages_mask() const { return supported_stages_mask_; };
     bool image_blit_supported() const { return true; }
 
     const VkPhysicalDeviceLimits &phys_device_limits() const { return phys_device_limits_; }
     const VkPhysicalDeviceProperties &device_properties() const { return device_properties_; }
     const VkPhysicalDeviceMemoryProperties &mem_properties() const { return mem_properties_; }
+    const VkCooperativeMatrixPropertiesKHR &coop_mat_properties() const { return coop_mat_properties_; }
 
     const VkPhysicalDeviceRayTracingPipelinePropertiesKHR &rt_props() const { return rt_props_; }
 
@@ -153,10 +152,11 @@ class Context {
     static void CheckVkPhysicalDeviceFeatures(const Api &api, VkPhysicalDevice &physical_device,
                                               VkPhysicalDeviceProperties &device_properties,
                                               VkPhysicalDeviceMemoryProperties &mem_properties,
+                                              VkCooperativeMatrixPropertiesKHR &coop_mat_properties,
                                               uint32_t &graphics_family_index, bool &out_raytracing_supported,
                                               bool &out_ray_query_supported, bool &out_shader_fp16_supported,
                                               bool &out_shader_int64_supported, bool &out_int64_atomics_supported,
-                                              int out_coop_matrix_size[3], bool &out_pageable_memory_supported);
+                                              bool &out_pageable_memory_supported);
     static bool InitVkDevice(const Api &api, VkDevice &device, VkPhysicalDevice physical_device,
                              uint32_t graphics_family_index, bool enable_raytracing, bool enable_ray_query,
                              bool enable_fp16, bool enable_int64, bool enable_int64_atomics, bool enable_coop_matrix,

internal/Vk/ContextVK.h

@@ -17,6 +17,13 @@
     #define f16vec4 vec4
 #endif
 
+#ifndef C_TYPE
+    #define C_TYPE float16_t
+    #define C_CONVERT 0
+#else
+    #define C_CONVERT 1
+#endif
+
 layout(push_constant) uniform UniformParams {
     Params g_params;
 };
@@ -232,6 +239,11 @@ shared float16_t g_mat_staging1[16 * 16];
 shared float16_t g_mat_staging2[16 * 16];
 shared float16_t g_mat_staging3[16 * 16];
 
+#ifdef COOP_M
+shared C_TYPE g_mat_staging_C0[COOP_M * 16];
+shared C_TYPE g_mat_staging_C1[COOP_M * 16];
+#endif
+
 void main() {
     ivec3 tile_id = ivec3(gl_WorkGroupID), li = ivec3(gl_LocalInvocationID);
 
@@ -248,22 +260,22 @@ void main() {
         return;
     }
 
-    coopmat<float16_t, gl_ScopeSubgroup, COOP_M, COOP_N, gl_MatrixUseAccumulator> C0[C_ROWS][C_COLS], C1[C_ROWS][C_COLS];
+    coopmat<C_TYPE, gl_ScopeSubgroup, COOP_M, COOP_N, gl_MatrixUseAccumulator> C0[C_ROWS][C_COLS], C1[C_ROWS][C_COLS];
     for (int i = 0; i < C_COLS; ++i) {
         const int ii = int(gl_LocalInvocationIndex);
         for (int jj = 0; jj < COOP_M && ii < COOP_N; ++jj) {
-            g_mat_staging0[jj * 16 + ii] = float16_t(0.0);
+            g_mat_staging_C0[jj * 16 + ii] = C_TYPE(0.0);
             if (ii < OUT_CHANNELS) {
-                g_mat_staging0[jj * 16 + ii] = g_biases[c + i * COOP_N + ii];
+                g_mat_staging_C0[jj * 16 + ii] = C_TYPE(g_biases[c + i * COOP_N + ii]);
             }
             // zero out shared memory to avoid NANs later
-            g_mat_staging1[jj * 16 + ii] = g_mat_staging2[jj * 16 + ii] = g_mat_staging3[jj * 16 + ii] = float16_t(0.0);
+            g_mat_staging0[jj * 16 + ii] = g_mat_staging1[jj * 16 + ii] = g_mat_staging2[jj * 16 + ii] = g_mat_staging3[jj * 16 + ii] = float16_t(0.0);
         }
         groupMemoryBarrier(); barrier();
 
         for (int j = 0; j < C_ROWS; ++j) {
-            coopMatLoad(C0[j][i], g_mat_staging0, 0u, 16u, gl_CooperativeMatrixLayoutRowMajor);
-            coopMatLoad(C1[j][i], g_mat_staging0, 0u, 16u, gl_CooperativeMatrixLayoutRowMajor);
+            coopMatLoad(C0[j][i], g_mat_staging_C0, 0u, 16u, gl_CooperativeMatrixLayoutRowMajor);
+            coopMatLoad(C1[j][i], g_mat_staging_C0, 0u, 16u, gl_CooperativeMatrixLayoutRowMajor);
         }
     }
 
@@ -501,34 +513,34 @@ void main() {
     for (int j = 0; j < C_ROWS; ++j) {
         for (int i = 0; i < cols_count; ++i) {
             for (int k = 0; k < C0[j][i].length(); ++k) {
-                C0[j][i][k] = max(max(C0[j][i][k], C1[j][i][k]), float16_t(0.0));
+                C0[j][i][k] = max(max(C0[j][i][k], C1[j][i][k]), C_TYPE(0.0));
             }
 
-            coopMatStore(C0[j][i], g_mat_staging0, 0u, 16u, gl_CooperativeMatrixLayoutRowMajor);
+            coopMatStore(C0[j][i], g_mat_staging_C0, 0u, 16u, gl_CooperativeMatrixLayoutRowMajor);
             groupMemoryBarrier(); barrier();
 
             for (int jj = 0; jj < COOP_M; jj += 2) {
                 for (int ii = 0; ii < COOP_N; ++ii) {
-                    const float16_t out_val = max(g_mat_staging0[(jj + 0) * 16 + ii], g_mat_staging0[(jj + 1) * 16 + ii]);
-                    g_out_buf[OUT_CHANNELS * ((y / 2 + 1) * g_params.output_stride + x / 2 + j * COOP_M / 2 + jj / 2 + 1) + c + i * COOP_N + ii] = max(out_val, float16_t(0.0));
+                    const C_TYPE out_val = max(max(g_mat_staging_C0[(jj + 0) * 16 + ii], g_mat_staging_C0[(jj + 1) * 16 + ii]), C_TYPE(0.0));
+                    g_out_buf[OUT_CHANNELS * ((y / 2 + 1) * g_params.output_stride + x / 2 + j * COOP_M / 2 + jj / 2 + 1) + c + i * COOP_N + ii] = float16_t(out_val);
                 }
             }
         }
     }
 #elif OUT_IMG
     for (int j = 0; j < rows_count && c == 0; ++j) {
         for (int k = 0; k < C0[j][0].length(); ++k) {
-            C0[j][0][k] = max(C0[j][0][k], float16_t(0.0));
-            C1[j][0][k] = max(C1[j][0][k], float16_t(0.0));
+            C0[j][0][k] = max(C0[j][0][k], C_TYPE(0.0));
+            C1[j][0][k] = max(C1[j][0][k], C_TYPE(0.0));
         }
 
-        coopMatStore(C0[j][0], g_mat_staging0, 0u, 16u, gl_CooperativeMatrixLayoutRowMajor);
-        coopMatStore(C1[j][0], g_mat_staging1, 0u, 16u, gl_CooperativeMatrixLayoutRowMajor);
+        coopMatStore(C0[j][0], g_mat_staging_C0, 0u, 16u, gl_CooperativeMatrixLayoutRowMajor);
+        coopMatStore(C1[j][0], g_mat_staging_C1, 0u, 16u, gl_CooperativeMatrixLayoutRowMajor);
         groupMemoryBarrier(); barrier();
 
         for (int jj = 0; jj < COOP_M; ++jj) {
-            vec4 val0 = vec4(g_mat_staging0[jj * 16 + 0], g_mat_staging0[jj * 16 + 1], g_mat_staging0[jj * 16 + 2], 1.0),
-                 val1 = vec4(g_mat_staging1[jj * 16 + 0], g_mat_staging1[jj * 16 + 1], g_mat_staging1[jj * 16 + 2], 1.0);
+            vec4 val0 = vec4(g_mat_staging_C0[jj * 16 + 0], g_mat_staging_C0[jj * 16 + 1], g_mat_staging_C0[jj * 16 + 2], 1.0),
+                 val1 = vec4(g_mat_staging_C1[jj * 16 + 0], g_mat_staging_C1[jj * 16 + 1], g_mat_staging_C1[jj * 16 + 2], 1.0);
             val0.xyz = transfer_output(val0.xyz);
             val1.xyz = transfer_output(val1.xyz);
             imageStore(g_out_img, ivec2(x + j * COOP_M + jj, y), val0);
@@ -554,14 +566,29 @@ void main() {
     for (int j = 0; j < rows_count; ++j) {
         for (int i = 0; i < cols_count; ++i) {
             for (int k = 0; k < C0[j][i].length(); ++k) {
-                C0[j][i][k] = max(C0[j][i][k], float16_t(0.0));
-                C1[j][i][k] = max(C1[j][i][k], float16_t(0.0));
+                C0[j][i][k] = max(C0[j][i][k], C_TYPE(0.0));
+                C1[j][i][k] = max(C1[j][i][k], C_TYPE(0.0));
             }
 
+#if C_CONVERT
+            coopMatStore(C0[j][i], g_mat_staging_C0, 0u, 16u, gl_CooperativeMatrixLayoutRowMajor);
+            coopMatStore(C1[j][i], g_mat_staging_C1, 0u, 16u, gl_CooperativeMatrixLayoutRowMajor);
+            groupMemoryBarrier(); barrier();
+
+            for (int jj = 0; jj < COOP_M; ++jj) {
+                for (int ii = 0; ii < COOP_N; ++ii) {
+                    g_out_buf[OUT_CHANNELS * ((y + 0 + 1) * g_params.output_stride + x + j * COOP_M + jj + 1) + c + i * COOP_N + ii] = float16_t(g_mat_staging_C0[jj * 16 + ii]);
+                    if (y + 1 < int(g_params.out_dims[1])) {
+                        g_out_buf[OUT_CHANNELS * ((y + 1 + 1) * g_params.output_stride + x + j * COOP_M + jj + 1) + c + i * COOP_N + ii] = float16_t(g_mat_staging_C1[jj * 16 + ii]);
+                    }
+                }
+            }
+#else // C_CONVERT
             coopMatStore(C0[j][i], g_out_buf, OUT_CHANNELS * ((y + 0 + 1) * g_params.output_stride + x + j * COOP_M + 1) + c + i * COOP_N, OUT_CHANNELS, gl_CooperativeMatrixLayoutRowMajor);
             if (y + 1 < int(g_params.out_dims[1])) {
                 coopMatStore(C1[j][i], g_out_buf, OUT_CHANNELS * ((y + 1 + 1) * g_params.output_stride + x + j * COOP_M + 1) + c + i * COOP_N, OUT_CHANNELS, gl_CooperativeMatrixLayoutRowMajor);
             }
+#endif // C_CONVERT
         }
     }
 #endif // OUT_IMG