CONV_SHAPE_128x128,
CONV_SHAPE_64x32,
CONV_SHAPE_32x256,
+ CONV_SHAPE_64x128,
CONV_SHAPE_COUNT,
};
{ 128, 128, 16 }, // CONV_SHAPE_128x128
{ 64, 32, 32 }, // CONV_SHAPE_64x32
{ 32, 256, 16 }, // CONV_SHAPE_32x256
+ { 64, 128, 16 }, // CONV_SHAPE_64x128
};
enum dmmv_wg_sizes {
};
struct vk_conv2d_pipeline_state {
- vk_conv2d_pipeline_state(uint32_t s0, uint32_t s1, uint32_t p0, uint32_t p1, uint32_t d0, uint32_t d1, uint32_t KW, uint32_t KH)
- : s0(s0), s1(s1), p0(p0), p1(p1), d0(d0), d1(d1), KW(KW), KH(KH) {}
+ vk_conv2d_pipeline_state(uint32_t s0, uint32_t s1, uint32_t p0, uint32_t p1, uint32_t d0, uint32_t d1, uint32_t KW, uint32_t KH, uint32_t aligned)
+ : s0(s0), s1(s1), p0(p0), p1(p1), d0(d0), d1(d1), KW(KW), KH(KH), aligned(aligned) {}
uint32_t s0, s1, p0, p1, d0, d1, KW, KH;
+ // when set, shader can skip K/CRS/NPQ bounds checks and address clamps
+ uint32_t aligned;
bool operator<(const vk_conv2d_pipeline_state &b) const {
- return std::tie(s0, s1, p0, p1, d0, d1, KW, KH) <
- std::tie(b.s0, b.s1, b.p0, b.p1, b.d0, b.d1, b.KW, b.KH);
+ return std::tie(s0, s1, p0, p1, d0, d1, KW, KH, aligned) <
+ std::tie(b.s0, b.s1, b.p0, b.p1, b.d0, b.d1, b.KW, b.KH, b.aligned);
}
};
// conv2d, conv_transpose_2d
for (uint32_t s = 0; s < CONV_SHAPE_COUNT; ++s) {
- uint32_t conv2d_WG_SIZE = 256;
+ // smaller WG for the small-tile fallback gives more concurrent WGs per SM
+ uint32_t conv2d_WG_SIZE = (s == CONV_SHAPE_64x32) ? 128 : 256;
uint32_t use_collectives = 0; // Enables subgroup ops for preventing the re-calculation of indices.
uint32_t conv2d_TS_K = (s == CONV_SHAPE_64x32) ? 4 : 8;
uint32_t conv2d_SHMEM_PAD = 4;
conv2d_BS.CRS); // CRS block size should be capped at subgroup size for correctness when shuffle is used.
}
- uint32_t conv2d_shmem_req =
- (conv2d_BS.K * (conv2d_BS.CRS + conv2d_SHMEM_PAD) + conv2d_BS.CRS * (conv2d_BS.NPQ + conv2d_SHMEM_PAD)) * sizeof(float);
- if (device->properties.limits.maxComputeSharedMemorySize < conv2d_shmem_req) {
+ // cm1 is used only when cm2 is unavailable; capped at 64x128 (due to shared memory size).
+ // Requires 16x16x16 f16-acc since that's the fragment shape hard-coded in the shader.
+ // Subgroup size must be 32 or 64 (to keep WG_SIZE sane) and we need
+ // subgroup_size_control to force the driver to actually use it.
+ bool conv2d_use_cm1 = false;
+#if defined(VK_KHR_cooperative_matrix) && defined(GGML_VULKAN_COOPMAT_GLSLC_SUPPORT)
+ conv2d_use_cm1 = !device->coopmat2 &&
+ device->coopmat_support && device->coopmat_support_16x16x16_f16acc &&
+ device->subgroup_size_control &&
+ (device->subgroup_size == 32 || device->subgroup_size == 64) &&
+ s != CONV_SHAPE_128x128;
+#endif
+
+ const uint32_t conv2d_cm1_shmem_pad = 8;
+
+ auto shmem_req = [&](uint32_t pad, bool csh_store, bool fp16_shmem) {
+ const uint32_t elem_size = fp16_shmem ? (uint32_t)sizeof(uint16_t) : (uint32_t)sizeof(float);
+ const uint32_t csh_elems = csh_store ? conv2d_BS.K * conv2d_BS.NPQ : 0u;
+ return (conv2d_BS.K * (conv2d_BS.CRS + pad) + conv2d_BS.CRS * (conv2d_BS.NPQ + pad) + csh_elems) * elem_size;
+ };
+
+ // coopmat1 needs to store the output through shared memory, so check up front
+ // whether it'll fit and disable it before applying coopmat1 parameters.
+ if (conv2d_use_cm1 && device->properties.limits.maxComputeSharedMemorySize < shmem_req(conv2d_cm1_shmem_pad, true, true)) {
+ conv2d_use_cm1 = false;
+ }
+
+ uint32_t conv2d_WM = 16, conv2d_WN = 16; // cm1 subgroup tile, ignored otherwise
+ if (conv2d_use_cm1) {
+ conv2d_SHMEM_PAD = conv2d_cm1_shmem_pad;
+ // 16x16x16 fragments; pick WM/WN to keep WG_SIZE at 256
+ // (i.e. 8 subgroups for sg=32, 4 subgroups for sg=64).
+ const bool sg64 = (device->subgroup_size == 64);
+ switch (s) {
+ case CONV_SHAPE_64x32: conv2d_WM = sg64 ? 32 : 16; conv2d_WN = 16; break;
+ case CONV_SHAPE_64x128: conv2d_WM = 32; conv2d_WN = sg64 ? 64 : 32; break;
+ case CONV_SHAPE_32x256: conv2d_WM = sg64 ? 16 : 32; conv2d_WN = sg64 ? 128 : 32; break;
+ default: break;
+ }
+ const uint32_t warps_M = conv2d_BS.K / conv2d_WM;
+ const uint32_t warps_N = conv2d_BS.NPQ / conv2d_WN;
+ conv2d_WG_SIZE = warps_M * warps_N * device->subgroup_size;
+ }
+
+ // stage cm2 accumulator through shmem for coalesced global stores;
+ // skipped on 128x128 where the extra Csh footprint hurts occupancy.
+ // cm1 always uses the staged path.
+ uint32_t conv2d_csh_store = (device->coopmat2 && s != CONV_SHAPE_128x128) ? 1u : 0u;
+ if (conv2d_use_cm1) {
+ conv2d_csh_store = 1;
+ }
+
+ // shmem is fp16 on cm2/cm1 (matches Csh), fp32 on scalar
+ const bool conv2d_use_fp16_shmem = device->coopmat2 || conv2d_use_cm1;
+
+ // shrink CRS if the non-cm1 config still doesn't fit
+ if (device->properties.limits.maxComputeSharedMemorySize < shmem_req(conv2d_SHMEM_PAD, conv2d_csh_store, conv2d_use_fp16_shmem)) {
+ GGML_ASSERT(!conv2d_use_cm1);
conv2d_BS.CRS = 8;
if (use_collectives) {
conv2d_BS.CRS = std::min(device->subgroup_size, conv2d_BS.CRS);
}
+ conv2d_csh_store = 0;
}
std::array<uint32_t, 3> wg_denoms = { conv2d_BS.K, 1, 1 };
std::vector<uint32_t> spec_constants = { conv2d_WG_SIZE, conv2d_BS.K, conv2d_BS.CRS, conv2d_BS.NPQ, conv2d_TS_K, use_collectives, conv2d_SHMEM_PAD };
+ // cm1 needs a fixed subgroup width to match the WG_SIZE we computed
+ const uint32_t conv2d_required_subgroup_size = conv2d_use_cm1 ? device->subgroup_size : 0;
+
#define CREATE_CONV(name, type_suffix, spv_suffix) \
for (auto &c : device->pipeline_##name##type_suffix[s]) { \
const vk_conv2d_pipeline_state &state = c.first; \
spec_constants_cpy.push_back(state.d1); \
spec_constants_cpy.push_back(state.KW); \
spec_constants_cpy.push_back(state.KH); \
+ spec_constants_cpy.push_back(state.aligned); \
+ spec_constants_cpy.push_back(conv2d_csh_store); \
+ spec_constants_cpy.push_back(conv2d_WM); \
+ spec_constants_cpy.push_back(conv2d_WN); \
ggml_vk_create_pipeline( \
device, c.second, #name #type_suffix, \
name##type_suffix##spv_suffix##_len, name##type_suffix##spv_suffix##_data, "main", 3, \
- sizeof(vk_op_conv2d_push_constants), wg_denoms, spec_constants_cpy, 1, true, use_collectives); \
+ sizeof(vk_op_conv2d_push_constants), wg_denoms, spec_constants_cpy, 1, true, use_collectives || conv2d_required_subgroup_size, conv2d_required_subgroup_size); \
}
#define CREATE_CONVS(spv_suffix) \
CREATE_CONV(conv2d, _f32, spv_suffix) \
if (device->coopmat2) {
CREATE_CONVS(_cm2)
} else
+#endif
+#if defined(VK_KHR_cooperative_matrix) && defined(GGML_VULKAN_COOPMAT_GLSLC_SUPPORT)
+ if (conv2d_use_cm1) {
+ CREATE_CONVS(_cm1)
+ } else
#endif
if (conv2d_UNROLL) {
CREATE_CONVS(_unroll)
// so small convolutions will still choose a smaller tile.
const uint32_t shader_core_count = ctx->device->shader_core_count > 0 ? ctx->device->shader_core_count : 32;
- if (K > 64 && n_tiles(CONV_SHAPE_128x128) >= shader_core_count * 2) {
+ // 128x128 isn't used with cm1 due to shared memory size; fall through to a smaller tile.
+ bool allow_128x128 = true;
+#if defined(VK_KHR_cooperative_matrix) && defined(GGML_VULKAN_COOPMAT_GLSLC_SUPPORT)
+ if (!ctx->device->coopmat2 && ctx->device->coopmat_support && ctx->device->coopmat_support_16x16x16_f16acc) {
+ allow_128x128 = false;
+ }
+#endif
+
+ if (allow_128x128 && K > 64 && n_tiles(CONV_SHAPE_128x128) >= shader_core_count * 2) {
return CONV_SHAPE_128x128;
} else if (K <= 32 && n_tiles(CONV_SHAPE_32x256) >= shader_core_count * 2) {
return CONV_SHAPE_32x256;
+ } else if (K <= 64 && n_tiles(CONV_SHAPE_64x128) >= shader_core_count * 2) {
+ return CONV_SHAPE_64x128;
+ } else if (!allow_128x128 && K > 64 && n_tiles(CONV_SHAPE_64x128) >= shader_core_count * 2) {
+ // cm1 fallback for large K when 128x128 isn't available
+ return CONV_SHAPE_64x128;
} else {
return CONV_SHAPE_64x32;
}
uint32_t p1 = !transpose ? (uint32_t)ggml_get_op_params_i32(dst, 3) : 0;
uint32_t d0 = !transpose ? (uint32_t)ggml_get_op_params_i32(dst, 4) : 1;
uint32_t d1 = !transpose ? (uint32_t)ggml_get_op_params_i32(dst, 5) : 1;
- vk_conv2d_pipeline_state conv2d_pipeline_state(s0, s1, p0, p1, d0, d1, KW, KH);
+
+ // tile-aligned shapes let the shader skip bounds checks
+ const uint32_t Cin = (uint32_t)src1->ne[2];
+ const uint32_t CRS = Cin * KW * KH;
+ const uint32_t BS_K = vk_conv_block_sizes[shape].K;
+ const uint32_t BS_CRS = vk_conv_block_sizes[shape].CRS;
+ const uint32_t BS_NPQ = vk_conv_block_sizes[shape].NPQ;
+ const uint32_t aligned = ((K % BS_K == 0) &&
+ (CRS % BS_CRS == 0) &&
+ (NPQ % BS_NPQ == 0)) ? 1u : 0u;
+
+ vk_conv2d_pipeline_state conv2d_pipeline_state(s0, s1, p0, p1, d0, d1, KW, KH, aligned);
std::map<vk_conv2d_pipeline_state, vk_pipeline> *pipelines = nullptr;
if (op == GGML_OP_CONV_2D) {
#extension GL_KHR_memory_scope_semantics : enable
#endif
+#ifdef COOPMAT
+#extension GL_KHR_cooperative_matrix : enable
+#extension GL_KHR_shader_subgroup_basic : enable
+#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
+#extension GL_KHR_memory_scope_semantics : enable
+#endif
+
#ifdef USE_COLLECTIVES
# extension GL_KHR_shader_subgroup_shuffle : enable
#endif
// Kernel spatial sizes
layout(constant_id = 13) const uint KW = 1;
layout(constant_id = 14) const uint KH = 1;
+// when set, skip bounds checks and address clamps (K/CRS/NPQ are tile-aligned)
+layout(constant_id = 15) const uint aligned = 0;
+// stage cm2 result through shmem (Csh) for coalesced stores. cm1 always does this.
+layout(constant_id = 16) const uint csh_store = 0;
+
+#ifdef COOPMAT
+// cm1 subgroup tile: each subgroup computes a WM x WN region as a grid of
+// TM x TN x TK fragments. Requires WM%TM == WN%TN == BS_K%WM == BS_NPQ%WN ==
+// BS_CRS%TK == 0, and WG_SIZE == (BS_K/WM) * (BS_NPQ/WN) * subgroup_size.
+layout(constant_id = 17) const uint WM = 32;
+layout(constant_id = 18) const uint WN = 32;
+const uint TM = 16;
+const uint TN = 16;
+const uint TK = 16;
+const uint cms_per_row = WM / TM;
+const uint cms_per_col = WN / TN;
+const uint warps_M = BS_K / WM;
+const uint warps_N = BS_NPQ / WN;
+#endif
+
+// without padding, H_idx/W_idx are in bounds by construction (non-TRANSPOSE only)
+#ifdef TRANSPOSE
+const bool hw_in_bounds = false;
+#else
+const bool hw_in_bounds = (p0 == 0) && (p1 == 0);
+#endif
+
+// TRANSPOSE stride alignment is trivially satisfied for stride 1
+#ifdef TRANSPOSE
+const bool stride_in_bounds = (s0 == 1) && (s1 == 1);
+#else
+const bool stride_in_bounds = true;
+#endif
uint32_t tid = gl_LocalInvocationID.x;
const uint32_t WG_SIZE = gl_WorkGroupSize.x;
// Number of blocktiles per input
uint32_t NB_CRS = splitWork(CRS, BS_CRS);
-#ifdef COOPMAT2
+#if defined(COOPMAT2) || defined(COOPMAT)
#define SHMEM_TYPE float16_t
#else
#define SHMEM_TYPE float
shared SHMEM_TYPE Ash[Ash_len]; // K x CRS
shared SHMEM_TYPE Bsh[Bsh_len]; // CRS x NPQ
+#if defined(COOPMAT2) || defined(COOPMAT)
+// stage matC through shmem so global stores are row-major (NPQ-contiguous)
+const uint32_t Csh_stride = BS_NPQ;
+#ifdef COOPMAT
+const uint32_t Csh_len = BS_K * Csh_stride;
+#else
+const uint32_t Csh_len = csh_store != 0 ? BS_K * Csh_stride : 1;
+#endif
+shared SHMEM_TYPE Csh[Csh_len]; // K x NPQ
+#endif
+
// Threadtile sizes
const uint32_t TS_NPQ = BS_K * BS_NPQ / WG_SIZE / TS_K;
uint32_t OH_idx = fastdiv(NPQ_idx - N_idx * p.OH * p.OW, p.OWmp, p.OWL); // divide by p.OW;
uint32_t OW_idx = NPQ_idx - N_idx * p.OH * p.OW - OH_idx * p.OW;
uint32_t dst_idx = OW_idx + OH_idx * p.nb1 + K_idx * p.nb2 + N_idx * p.nb3;
- if (K_idx < K && NPQ_idx < NPQ) {
+ if (aligned != 0 || (K_idx < K && NPQ_idx < NPQ)) {
dst_data[dst_idx] = D_TYPE(elem);
}
return elem;
#ifdef COOPMAT2
coopmat<ACC_TYPE, gl_ScopeWorkgroup, BS_K, BS_NPQ, gl_MatrixUseAccumulator> matC;
matC = coopmat<ACC_TYPE, gl_ScopeWorkgroup, BS_K, BS_NPQ, gl_MatrixUseAccumulator>(0.0);
+#elif defined(COOPMAT)
+ coopmat<float16_t, gl_ScopeSubgroup, TM, TN, gl_MatrixUseAccumulator> sums[cms_per_row * cms_per_col];
+ [[unroll]] for (uint i = 0; i < cms_per_row * cms_per_col; i++) {
+ sums[i] = coopmat<float16_t, gl_ScopeSubgroup, TM, TN, gl_MatrixUseAccumulator>(0.0);
+ }
+ const uint warp_r = gl_SubgroupID / warps_N;
+ const uint warp_c = gl_SubgroupID % warps_N;
#else
float regC[TS_K][TS_NPQ];
for (uint32_t T_ly = 0; T_ly < TS_K; T_ly++) {
uint32_t B_lx = Ac;
uint32_t K_idx = B_idx_K * BS_K + B_ly; /* Global K_idx (row index of A)*/
#ifdef TRANSPOSE
- uint32_t knl_idx = min(KW_idx_a + KH_idx_a * p.nb01 + K_idx * p.nb02 + Cin_idx_a * p.nb03, K * CRS - 1);
+ uint32_t knl_idx = KW_idx_a + KH_idx_a * p.nb01 + K_idx * p.nb02 + Cin_idx_a * p.nb03;
#else
- uint32_t knl_idx = min(KW_idx_a + KH_idx_a * p.nb01 + Cin_idx_a * p.nb02 + K_idx * p.nb03, K * CRS - 1);
+ uint32_t knl_idx = KW_idx_a + KH_idx_a * p.nb01 + Cin_idx_a * p.nb02 + K_idx * p.nb03;
#endif
+ if (aligned == 0) {
+ knl_idx = min(knl_idx, K * CRS - 1);
+ }
float val = knl_data[knl_idx];
- if (K_idx >= K || CRS_idx_a >= CRS) {
+ if (aligned == 0 && (K_idx >= K || CRS_idx_a >= CRS)) {
val = 0.0;
}
Ash[B_ly * Ash_stride + B_lx] = SHMEM_TYPE(val);
uint32_t H_idx = OH_idx * s1 + KH_idx_b * d1 - p1;
uint32_t W_idx = OW_idx * s0 + KW_idx_b * d0 - p0;
#endif
- uint32_t src_idx =
- min(max(W_idx + H_idx * p.nb11 + Cin_idx_b * p.nb12 + N_idx * p.nb13, 0), p.Cin * p.N * p.W * p.H - 1);
+ uint32_t src_idx = W_idx + H_idx * p.nb11 + Cin_idx_b * p.nb12 + N_idx * p.nb13;
+ // skip clamp when address can't go OOB
+ if (aligned == 0 || !hw_in_bounds || !stride_in_bounds) {
+ src_idx = min(max(src_idx, 0), p.Cin * p.N * p.W * p.H - 1);
+ }
float val = src_data[src_idx];
- if (CRS_idx_b >= CRS || NPQ_idx >= NPQ
- || H_idx >= p.H || W_idx >= p.W // Lower bound checks aren't necessary. (idx >= 0x80000000 for such case)
+ bool oob = false;
+ if (aligned == 0 && (CRS_idx_b >= CRS || NPQ_idx >= NPQ)) {
+ oob = true;
+ }
+ // also catches lower-bound underflow (idx wraps to 0x80000000+)
+ if (!hw_in_bounds && (H_idx >= p.H || W_idx >= p.W)) {
+ oob = true;
+ }
#ifdef TRANSPOSE
- || (H_idx_x_s1 - H_idx * s1 != 0) || (W_idx_x_s0 - W_idx * s0 != 0)
+ if (!stride_in_bounds &&
+ ((H_idx_x_s1 - H_idx * s1 != 0) || (W_idx_x_s0 - W_idx * s0 != 0))) {
+ oob = true;
+ }
#endif
- ) {
+ if (oob) {
val = 0.0;
}
Bsh[B_ly * Bsh_stride + B_lx] = SHMEM_TYPE(val);
coopMatLoad(matA, Ash, 0, Ash_stride, gl_CooperativeMatrixLayoutRowMajor);
coopMatLoad(matB, Bsh, 0, Bsh_stride, gl_CooperativeMatrixLayoutRowMajor);
matC = coopMatMulAdd(matA, matB, matC);
+#elif defined(COOPMAT)
+ // each subgroup multiplies its grid of fragments per TK-sized CRS chunk
+ [[unroll]] for (uint k_step = 0; k_step < BS_CRS / TK; k_step++) {
+ coopmat<float16_t, gl_ScopeSubgroup, TM, TK, gl_MatrixUseA> cache_a[cms_per_row];
+ [[unroll]] for (uint cm_row = 0; cm_row < cms_per_row; cm_row++) {
+ const uint a_off = (warp_r * WM + cm_row * TM) * Ash_stride + k_step * TK;
+ coopMatLoad(cache_a[cm_row], Ash, a_off, Ash_stride, gl_CooperativeMatrixLayoutRowMajor);
+ }
+ [[unroll]] for (uint cm_col = 0; cm_col < cms_per_col; cm_col++) {
+ coopmat<float16_t, gl_ScopeSubgroup, TK, TN, gl_MatrixUseB> cache_b;
+ const uint b_off = k_step * TK * Bsh_stride + warp_c * WN + cm_col * TN;
+ coopMatLoad(cache_b, Bsh, b_off, Bsh_stride, gl_CooperativeMatrixLayoutRowMajor);
+ [[unroll]] for (uint cm_row = 0; cm_row < cms_per_row; cm_row++) {
+ sums[cm_col * cms_per_row + cm_row] = coopMatMulAdd(cache_a[cm_row], cache_b, sums[cm_col * cms_per_row + cm_row]);
+ }
+ }
+ }
#else
if (T_y * TS_K < K) {
UNROLL for (uint32_t CRS_lidx = 0; CRS_lidx < BS_CRS; CRS_lidx++) {
barrier();
}
/* Save C* */
+#if defined(COOPMAT2) || defined(COOPMAT)
+ // stage matC into Csh, then write to dst with coalesced NPQ-contiguous stores
+#ifdef COOPMAT
+ const bool use_staged_store = true;
+#else
+ const bool use_staged_store = (csh_store != 0);
+#endif
+ if (use_staged_store) {
+#ifdef COOPMAT
+ // cm1: each subgroup stores its fragment grid into its Csh slot
+ [[unroll]] for (uint cm_row = 0; cm_row < cms_per_row; cm_row++) {
+ [[unroll]] for (uint cm_col = 0; cm_col < cms_per_col; cm_col++) {
+ const uint csh_off = (warp_r * WM + cm_row * TM) * Csh_stride + warp_c * WN + cm_col * TN;
+ coopMatStore(sums[cm_col * cms_per_row + cm_row], Csh, csh_off, Csh_stride, gl_CooperativeMatrixLayoutRowMajor);
+ }
+ }
+#else
+ coopMatStore(matC, Csh, 0, Csh_stride, gl_CooperativeMatrixLayoutRowMajor);
+#endif
+ barrier();
+
+ // cooperative shmem->global: WG threads spread across BS_NPQ (the
+ // contiguous direction of dst), each iter covers store_rows_per_iter K-rows
+ const uint32_t store_rows_per_iter = WG_SIZE / BS_NPQ;
+ const uint32_t store_iters = BS_K / store_rows_per_iter;
+ const uint32_t k_thread_offset = tid / BS_NPQ;
+ const uint32_t npq_thread = tid % BS_NPQ;
+ [[unroll]] for (uint32_t i = 0; i < store_iters; i++) {
+ uint32_t k_local = i * store_rows_per_iter + k_thread_offset;
+ uint32_t K_idx = B_idx_K * BS_K + k_local;
+ uint32_t NPQ_idx = B_idx_NPQ * BS_NPQ + npq_thread;
+ uint32_t N_idx = fastdiv(NPQ_idx, p.OWOHmp, p.OWOHL);
+ uint32_t OH_idx = fastdiv(NPQ_idx - N_idx * p.OH * p.OW, p.OWmp, p.OWL);
+ uint32_t OW_idx = NPQ_idx - N_idx * p.OH * p.OW - OH_idx * p.OW;
+ uint32_t dst_idx = OW_idx + OH_idx * p.nb1 + K_idx * p.nb2 + N_idx * p.nb3;
+ if (aligned != 0 || (K_idx < K && NPQ_idx < NPQ)) {
+ dst_data[dst_idx] = D_TYPE(Csh[k_local * Csh_stride + npq_thread]);
+ }
+ }
+ }
#ifdef COOPMAT2
- coopMatPerElementNV(matC, matC, perElemOpStore);
+ else {
+ coopMatPerElementNV(matC, matC, perElemOpStore);
+ }
+#endif
#else
if (T_y * TS_K < K) {
for (uint32_t T_ly = 0; T_ly < TS_K; T_ly++) {
uint32_t OH_idx = fastdiv(NPQ_idx - N_idx * p.OH * p.OW, p.OWmp, p.OWL); // divide by p.OW;
uint32_t OW_idx = NPQ_idx - N_idx * p.OH * p.OW - OH_idx * p.OW;
uint32_t dst_idx = OW_idx + OH_idx * p.nb1 + K_idx * p.nb2 + N_idx * p.nb3;
- if (K_idx < K && NPQ_idx < NPQ) {
+ if (aligned != 0 || (K_idx < K && NPQ_idx < NPQ)) {
dst_data[dst_idx] = regC[T_ly][T_lx];
}
}
string_to_spv(name + (unroll ? "_unroll" : ""), "conv2d_mm.comp", defines);
#if defined(GGML_VULKAN_COOPMAT2_GLSLC_SUPPORT)
if (unroll) {
- defines["COOPMAT2"] = "1";
- string_to_spv(name, "conv2d_mm.comp", defines, true, false, true);
+ auto cm2_defines = defines;
+ cm2_defines["COOPMAT2"] = "1";
+ string_to_spv(name, "conv2d_mm.comp", cm2_defines, true, false, true);
+ }
+#endif
+#if defined(GGML_VULKAN_COOPMAT_GLSLC_SUPPORT)
+ if (unroll) {
+ auto cm1_defines = defines;
+ cm1_defines["COOPMAT"] = "1";
+ string_to_spv(name, "conv2d_mm.comp", cm1_defines, true, true, false);
}
#endif
}