}
};
+struct vk_conv3d_pipeline_state {
+ vk_conv3d_pipeline_state(uint32_t s0, uint32_t s1, uint32_t s2, uint32_t p0, uint32_t p1, uint32_t p2,
+ uint32_t d0, uint32_t d1, uint32_t d2, uint32_t KW, uint32_t KH, uint32_t KD, uint32_t aligned)
+ : s0(s0), s1(s1), s2(s2), p0(p0), p1(p1), p2(p2), d0(d0), d1(d1), d2(d2), KW(KW), KH(KH), KD(KD), aligned(aligned) {}
+
+ uint32_t s0, s1, s2, p0, p1, p2, d0, d1, d2, KW, KH, KD;
+ uint32_t aligned;
+
+ bool operator<(const vk_conv3d_pipeline_state &b) const {
+ return std::tie(s0, s1, s2, p0, p1, p2, d0, d1, d2, KW, KH, KD, aligned) <
+ std::tie(b.s0, b.s1, b.s2, b.p0, b.p1, b.p2, b.d0, b.d1, b.d2, b.KW, b.KH, b.KD, b.aligned);
+ }
+};
+
struct vk_solve_tri_pipeline_state {
vk_solve_tri_pipeline_state(uint32_t N, uint32_t K)
: N(N), K(K) {}
std::map<vk_conv2d_pipeline_state, vk_pipeline> pipeline_conv2d_f16_f32[CONV_SHAPE_COUNT];
std::map<vk_conv2d_pipeline_state, vk_pipeline> pipeline_conv_transpose_2d_f32[CONV_SHAPE_COUNT];
std::map<vk_conv2d_pipeline_state, vk_pipeline> pipeline_conv_transpose_2d_f16_f32[CONV_SHAPE_COUNT];
+ std::map<vk_conv3d_pipeline_state, vk_pipeline> pipeline_conv3d_f32[CONV_SHAPE_COUNT];
+ std::map<vk_conv3d_pipeline_state, vk_pipeline> pipeline_conv3d_f16_f32[CONV_SHAPE_COUNT];
vk_pipeline pipeline_conv2d_dw_whcn_f32, pipeline_conv2d_dw_whcn_f16_f32;
vk_pipeline pipeline_conv2d_dw_cwhn_f32, pipeline_conv2d_dw_cwhn_f16_f32;
init_fastdiv_values(p.OW*p.OH, p.OWOHmp, p.OWOHL);
}
+struct vk_op_conv3d_push_constants {
+ uint32_t OC;
+ uint32_t IC;
+ uint32_t N;
+
+ uint32_t IW;
+ uint32_t IH;
+ uint32_t ID;
+ uint32_t OW;
+ uint32_t OH;
+ uint32_t OD;
+
+ uint32_t nb01;
+ uint32_t nb02;
+ uint32_t nb03;
+
+ uint32_t nb11;
+ uint32_t nb12;
+ uint32_t nb13;
+
+ uint32_t nb1;
+ uint32_t nb2;
+ uint32_t nb3;
+
+ uint32_t OWmp; uint32_t OWL;
+ uint32_t OWOHmp; uint32_t OWOHL;
+ uint32_t OWOHODmp; uint32_t OWOHODL;
+};
+
+template <> void init_pushconst_fastdiv(vk_op_conv3d_push_constants &p) {
+ init_fastdiv_values(p.OW, p.OWmp, p.OWL);
+ init_fastdiv_values(p.OW*p.OH, p.OWOHmp, p.OWOHL);
+ init_fastdiv_values(p.OW*p.OH*p.OD, p.OWOHODmp, p.OWOHODL);
+}
+
struct vk_op_conv2d_dw_push_constants {
uint32_t ne;
uint32_t batches;
ggml_vk_create_pipeline(device, device->pipeline_opt_step_sgd_f32, "opt_step_sgd_f32", opt_step_sgd_f32_len, opt_step_sgd_f32_data, "main", 3, sizeof(vk_op_push_constants), {512, 1, 1}, {}, 1);
- // conv2d, conv_transpose_2d
+ // conv2d, conv_transpose_2d, conv3d
for (uint32_t s = 0; s < CONV_SHAPE_COUNT; ++s) {
// smaller WG for the small-tile fallback gives more concurrent WGs per SM
uint32_t conv2d_WG_SIZE = (s == CONV_SHAPE_64x32) ? 128 : 256;
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.
+ // 2D, transpose-2D, and 3D conv use the same KxCRS @ CRSxNPQ shmem
+ // layout. cm1 needs Csh for output, so check before applying cm1 params.
if (conv2d_use_cm1 && device->properties.limits.maxComputeSharedMemorySize < shmem_req(conv2d_cm1_shmem_pad, true, true)) {
conv2d_use_cm1 = false;
}
}
#undef CREATE_CONV
#undef CREATE_CONVS
+
+ std::vector<uint32_t> conv3d_spec_constants = { conv2d_WG_SIZE, conv2d_BS.K, conv2d_BS.CRS, conv2d_BS.NPQ, conv2d_TS_K, conv2d_SHMEM_PAD };
+#define CREATE_CONV3D(type_suffix, spv_suffix) \
+ for (auto &c : device->pipeline_conv3d##type_suffix[s]) { \
+ const vk_conv3d_pipeline_state &state = c.first; \
+ std::vector<uint32_t> spec_constants_cpy = conv3d_spec_constants; \
+ spec_constants_cpy.push_back(state.s0); \
+ spec_constants_cpy.push_back(state.s1); \
+ spec_constants_cpy.push_back(state.s2); \
+ spec_constants_cpy.push_back(state.p0); \
+ spec_constants_cpy.push_back(state.p1); \
+ spec_constants_cpy.push_back(state.p2); \
+ spec_constants_cpy.push_back(state.d0); \
+ spec_constants_cpy.push_back(state.d1); \
+ spec_constants_cpy.push_back(state.d2); \
+ spec_constants_cpy.push_back(state.KW); \
+ spec_constants_cpy.push_back(state.KH); \
+ spec_constants_cpy.push_back(state.KD); \
+ 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, "conv3d" #type_suffix, \
+ conv3d##type_suffix##spv_suffix##_len, conv3d##type_suffix##spv_suffix##_data, "main", 3, \
+ sizeof(vk_op_conv3d_push_constants), wg_denoms, spec_constants_cpy, 1, true, conv2d_required_subgroup_size != 0, conv2d_required_subgroup_size); \
+ }
+#if defined(GGML_VULKAN_COOPMAT2_GLSLC_SUPPORT)
+ if (device->coopmat2) {
+ CREATE_CONV3D(_f32, _cm2)
+ CREATE_CONV3D(_f16_f32, _cm2)
+ } else
+#endif
+#if defined(VK_KHR_cooperative_matrix) && defined(GGML_VULKAN_COOPMAT_GLSLC_SUPPORT)
+ if (conv2d_use_cm1) {
+ CREATE_CONV3D(_f32, _cm1)
+ CREATE_CONV3D(_f16_f32, _cm1)
+ } else
+#endif
+ if (conv2d_UNROLL) {
+ CREATE_CONV3D(_f32, _unroll)
+ CREATE_CONV3D(_f16_f32, _unroll)
+ } else {
+ CREATE_CONV3D(_f32, )
+ CREATE_CONV3D(_f16_f32, )
+ }
+#undef CREATE_CONV3D
}
ggml_vk_create_pipeline(device, device->pipeline_conv2d_dw_whcn_f32, "conv2d_dw_whcn_f32", conv2d_dw_whcn_f32_len, conv2d_dw_whcn_f32_data, "main", 3, sizeof(vk_op_conv2d_dw_push_constants), {512, 1, 1}, {}, 1);
}
}
return nullptr;
+ case GGML_OP_CONV_3D:
+ if (src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) {
+ const uint32_t OC = (uint32_t)ggml_get_op_params_i32(dst, 11);
+ const uint32_t IC = (uint32_t)ggml_get_op_params_i32(dst, 9);
+ const uint32_t N = (uint32_t)ggml_get_op_params_i32(dst, 10);
+ const uint32_t NPQ = N * dst->ne[2] * dst->ne[1] * dst->ne[0];
+ const vk_conv_shapes shape = ggml_vk_conv_select_shape(ctx, OC, NPQ);
+
+ const uint32_t KW = (uint32_t)src0->ne[0];
+ const uint32_t KH = (uint32_t)src0->ne[1];
+ const uint32_t KD = (uint32_t)src0->ne[2];
+ const uint32_t s0 = (uint32_t)ggml_get_op_params_i32(dst, 0);
+ const uint32_t s1 = (uint32_t)ggml_get_op_params_i32(dst, 1);
+ const uint32_t s2 = (uint32_t)ggml_get_op_params_i32(dst, 2);
+ const uint32_t p0 = (uint32_t)ggml_get_op_params_i32(dst, 3);
+ const uint32_t p1 = (uint32_t)ggml_get_op_params_i32(dst, 4);
+ const uint32_t p2 = (uint32_t)ggml_get_op_params_i32(dst, 5);
+ const uint32_t d0 = (uint32_t)ggml_get_op_params_i32(dst, 6);
+ const uint32_t d1 = (uint32_t)ggml_get_op_params_i32(dst, 7);
+ const uint32_t d2 = (uint32_t)ggml_get_op_params_i32(dst, 8);
+
+ const uint32_t CRS = IC * KW * KH * KD;
+ 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 = ((OC % BS_K == 0) &&
+ (CRS % BS_CRS == 0) &&
+ (NPQ % BS_NPQ == 0)) ? 1u : 0u;
+
+ vk_conv3d_pipeline_state conv3d_pipeline_state(s0, s1, s2, p0, p1, p2, d0, d1, d2, KW, KH, KD, aligned);
+
+ std::map<vk_conv3d_pipeline_state, vk_pipeline> *pipelines = nullptr;
+ if (src0->type == GGML_TYPE_F32) {
+ pipelines = &ctx->device->pipeline_conv3d_f32[shape];
+ } else if (src0->type == GGML_TYPE_F16) {
+ pipelines = &ctx->device->pipeline_conv3d_f16_f32[shape];
+ } else {
+ return nullptr;
+ }
+
+ vk_pipeline pipeline = nullptr;
+
+ {
+ std::lock_guard<std::mutex> guard(ctx->device->compile_mutex);
+ auto it = pipelines->find(conv3d_pipeline_state);
+ if (it != pipelines->end()) {
+ pipeline = it->second;
+ } else {
+ (*pipelines)[conv3d_pipeline_state] = pipeline = std::make_shared<vk_pipeline_struct>();
+ }
+ }
+
+ return pipeline;
+ }
+ return nullptr;
case GGML_OP_ADD1:
if (src0->type == GGML_TYPE_F16 && src1->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F16) {
return ctx->device->pipeline_add1_f16_f16;
GGML_ABORT("invalid push constant type for CONV_2D");
}
break;
+ case GGML_OP_CONV_3D:
+ if constexpr (std::is_same_v<PC, vk_op_conv3d_push_constants>) {
+ const uint32_t NPQ = pc.N * pc.OD * pc.OH * pc.OW;
+ const vk_conv_shapes shape = ggml_vk_conv_select_shape(ctx, pc.OC, NPQ);
+ const uint32_t NPQ_blocks = CEIL_DIV(NPQ, vk_conv_block_sizes[shape].NPQ);
+
+ elements = { pc.OC, NPQ_blocks, 1 };
+ if (elements[1] > 512) {
+ elements[2] = CEIL_DIV(elements[1], 512);
+ elements[1] = 512;
+ }
+ } else {
+ GGML_ABORT("invalid push constant type for CONV_3D");
+ }
+ break;
case GGML_OP_ADD:
case GGML_OP_SUB:
case GGML_OP_DIV:
ggml_vk_op_f32(ctx, subctx, src0, src1, nullptr, nullptr, dst, dst->op, std::move(p));
}
+static void ggml_vk_conv_3d(ggml_backend_vk_context * ctx, vk_context & subctx, const ggml_tensor * src0,
+ const ggml_tensor * src1, ggml_tensor * dst) {
+ GGML_ASSERT(src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16);
+ GGML_ASSERT(src1->type == GGML_TYPE_F32);
+ GGML_ASSERT(dst->type == GGML_TYPE_F32);
+
+ GGML_TENSOR_BINARY_OP_LOCALS
+ GGML_ASSERT(nb00 == sizeof(float) || nb00 == sizeof(ggml_fp16_t));
+ GGML_ASSERT(nb10 == sizeof(float));
+ GGML_ASSERT(nb0 == sizeof(float));
+
+ vk_op_conv3d_push_constants p{};
+ p.IC = static_cast<uint32_t>(ggml_get_op_params_i32(dst, 9));
+ p.N = static_cast<uint32_t>(ggml_get_op_params_i32(dst, 10));
+ p.OC = static_cast<uint32_t>(ggml_get_op_params_i32(dst, 11));
+ GGML_ASSERT(src0->ne[3] == (int64_t)p.IC * p.OC);
+ GGML_ASSERT(src1->ne[3] == (int64_t)p.IC * p.N);
+ GGML_ASSERT(dst->ne[3] == (int64_t)p.OC * p.N);
+
+ p.IW = static_cast<uint32_t>(ne10);
+ p.IH = static_cast<uint32_t>(ne11);
+ p.ID = static_cast<uint32_t>(ne12);
+ p.OW = static_cast<uint32_t>(ne0);
+ p.OH = static_cast<uint32_t>(ne1);
+ p.OD = static_cast<uint32_t>(ne2);
+
+ // the shader clamps src addresses to p.IC * p.N * p.IW * p.IH * p.ID - 1 in uint32, so the
+ // total input element count must fit in a uint32.
+ GGML_ASSERT((uint64_t)p.IC * p.N * p.IW * p.IH * p.ID <= 0xFFFFFFFFull);
+
+ p.nb01 = static_cast<uint32_t>(nb01 / nb00);
+ p.nb02 = static_cast<uint32_t>(nb02 / nb00);
+ p.nb03 = static_cast<uint32_t>(nb03 / nb00);
+
+ p.nb11 = static_cast<uint32_t>(nb11 / nb10);
+ p.nb12 = static_cast<uint32_t>(nb12 / nb10);
+ p.nb13 = static_cast<uint32_t>(nb13 / nb10);
+
+ p.nb1 = static_cast<uint32_t>(nb1 / nb0);
+ p.nb2 = static_cast<uint32_t>(nb2 / nb0);
+ p.nb3 = static_cast<uint32_t>(nb3 / nb0);
+
+ ggml_vk_op_f32(ctx, subctx, src0, src1, nullptr, nullptr, dst, GGML_OP_CONV_3D, std::move(p));
+}
+
static void ggml_vk_conv_2d_dw(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
vk_op_conv2d_dw_push_constants p{};
p.ne = ggml_nelements(dst);
case GGML_OP_CONV_TRANSPOSE_2D:
ggml_vk_conv_2d(ctx, compute_ctx, src0, src1, node);
+ break;
+ case GGML_OP_CONV_3D:
+ ggml_vk_conv_3d(ctx, compute_ctx, src0, src1, node);
+
break;
case GGML_OP_CONV_2D_DW:
ggml_vk_conv_2d_dw(ctx, compute_ctx, src0, src1, node);
ggml_is_contiguous(op->src[1]) &&
ggml_is_contiguous(op));
}
+ case GGML_OP_CONV_3D:
+ return (op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_F16) &&
+ op->src[1]->type == GGML_TYPE_F32 &&
+ op->type == GGML_TYPE_F32 &&
+ ggml_is_contiguous(op->src[0]) &&
+ ggml_is_contiguous(op->src[1]) &&
+ ggml_is_contiguous(op);
default:
return false;
}
const int32_t d0 = tensor->op_params[4];
const int32_t d1 = tensor->op_params[5];
tensor_clone = ggml_conv_2d(ggml_ctx, src_clone[0], src_clone[1], s0, s1, p0, p1, d0, d1);
+ } else if (tensor->op == GGML_OP_CONV_3D) {
+ const int32_t s0 = tensor->op_params[0];
+ const int32_t s1 = tensor->op_params[1];
+ const int32_t s2 = tensor->op_params[2];
+ const int32_t p0 = tensor->op_params[3];
+ const int32_t p1 = tensor->op_params[4];
+ const int32_t p2 = tensor->op_params[5];
+ const int32_t d0 = tensor->op_params[6];
+ const int32_t d1 = tensor->op_params[7];
+ const int32_t d2 = tensor->op_params[8];
+ const int32_t IC = tensor->op_params[9];
+ const int32_t N = tensor->op_params[10];
+ const int32_t OC = tensor->op_params[11];
+ tensor_clone = ggml_conv_3d_direct(ggml_ctx, src_clone[0], src_clone[1], s0, s1, s2, p0, p1, p2, d0, d1, d2, IC, N, OC);
} else if (tensor->op == GGML_OP_CONV_2D_DW) {
const int32_t s0 = tensor->op_params[0];
const int32_t s1 = tensor->op_params[1];
--- /dev/null
+#version 450
+
+#extension GL_EXT_control_flow_attributes : enable
+#ifdef COOPMAT2
+#extension GL_NV_cooperative_matrix2 : enable
+#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require
+#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
+
+#include "types.glsl"
+
+// shape notation: [dim(N), ..., dim(0)] -- stride(dim(j)) >= stride(dim(i)) if i > j
+layout(binding = 0) readonly buffer A {
+ A_TYPE knl_data[];
+}; // src0 - kernel: [KW, KH, KD, IC*OC]
+
+layout(binding = 1) readonly buffer B {
+ B_TYPE src_data[];
+}; // src1 - input: [IW, IH, ID, IC*N] -- channel_first format
+
+layout(binding = 2) writeonly buffer D {
+ D_TYPE dst_data[];
+}; // dst - result: [OW, OH, OD, OC*N]
+
+layout(push_constant) uniform parameter {
+ // I/O channels, batch size
+ uint32_t OC;
+ uint32_t IC;
+ uint32_t N;
+
+ // Tensor spatial sizes: input, output
+ uint32_t IW;
+ uint32_t IH;
+ uint32_t ID;
+ uint32_t OW;
+ uint32_t OH;
+ uint32_t OD;
+
+ // Strides in elements
+ uint32_t nb01;
+ uint32_t nb02;
+ uint32_t nb03;
+
+ uint32_t nb11;
+ uint32_t nb12;
+ uint32_t nb13;
+
+ uint32_t nb1;
+ uint32_t nb2;
+ uint32_t nb3;
+
+ // fastdiv helper values
+ uint32_t OWmp; uint32_t OWL;
+ uint32_t OWOHmp; uint32_t OWOHL;
+ uint32_t OWOHODmp; uint32_t OWOHODL;
+}
+
+p;
+
+layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in;
+// Blocktile sizes
+layout(constant_id = 1) const uint BS_K = 128;
+layout(constant_id = 2) const uint BS_CRS = 16;
+layout(constant_id = 3) const uint BS_NPQ = 128;
+// Thread-tile sizes
+layout(constant_id = 4) const uint TS_K = 8;
+layout(constant_id = 5) const uint SHMEM_PAD = 4;
+// Stride, padding, dilation
+layout(constant_id = 6) const uint s0 = 1;
+layout(constant_id = 7) const uint s1 = 1;
+layout(constant_id = 8) const uint s2 = 1;
+layout(constant_id = 9) const uint p0 = 0;
+layout(constant_id = 10) const uint p1 = 0;
+layout(constant_id = 11) const uint p2 = 0;
+layout(constant_id = 12) const uint d0 = 1;
+layout(constant_id = 13) const uint d1 = 1;
+layout(constant_id = 14) const uint d2 = 1;
+// Kernel spatial sizes
+layout(constant_id = 15) const uint KW = 1;
+layout(constant_id = 16) const uint KH = 1;
+layout(constant_id = 17) const uint KD = 1;
+// when set, skip bounds checks and address clamps (K/CRS/NPQ are tile-aligned)
+layout(constant_id = 18) const uint aligned = 0;
+// stage cm2 result through shmem (Csh) for coalesced stores. cm1 always does this.
+layout(constant_id = 19) 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 = 20) const uint WM = 32;
+layout(constant_id = 21) 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, ID_idx/IH_idx/IW_idx are in bounds by construction
+const bool dhw_in_bounds = (p0 == 0) && (p1 == 0) && (p2 == 0);
+
+uint32_t tid = gl_LocalInvocationID.x;
+const uint32_t WG_SIZE = gl_WorkGroupSize.x;
+
+uint splitWork(uint work_size, uint block_size) {
+ return (block_size + work_size - 1) / block_size;
+}
+
+uint32_t K = p.OC;
+uint32_t CRS = p.IC * KD * KH * KW;
+uint32_t NPQ = p.N * p.OD * p.OH * p.OW;
+
+// Number of blocktiles per input
+uint32_t NB_CRS = splitWork(CRS, BS_CRS);
+
+#if defined(COOPMAT2) || defined(COOPMAT)
+#define SHMEM_TYPE float16_t
+#else
+#define SHMEM_TYPE float
+#endif
+
+const uint32_t Ash_stride = BS_CRS + SHMEM_PAD;
+const uint32_t Bsh_stride = BS_NPQ + SHMEM_PAD;
+
+const uint32_t Ash_len = BS_K * Ash_stride;
+const uint32_t Bsh_len = BS_CRS * Bsh_stride;
+
+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;
+
+// Number of threadtiles per blocktile
+const uint32_t NT_NPQ = BS_NPQ / TS_NPQ;
+
+/*
+Compute
+KxCRS @ CRSxNPQ = K x NPQ
+K=OC
+C=IC
+D,R,S=KD,KH,KW
+Z,P,Q=OD,OH,OW
+*/
+
+uint32_t B_idx_K = gl_WorkGroupID.x;
+uint32_t B_idx_NPQ = gl_WorkGroupID.y + gl_WorkGroupID.z * 512;
+
+uint32_t T_y = tid / NT_NPQ;
+uint32_t T_x = tid % NT_NPQ;
+
+uint32_t Ar = tid / BS_CRS;
+uint32_t Ac = tid % BS_CRS;
+const uint32_t ArpWg = WG_SIZE / BS_CRS;
+
+uint32_t Br = tid / BS_NPQ;
+uint32_t Bc = tid % BS_NPQ;
+const uint32_t BrpWg = WG_SIZE / BS_NPQ;
+
+// see init_fastdiv_values in ggml-vulkan.cpp
+uint fastdiv(uint n, uint mp, uint L) {
+ uint msbs, lsbs;
+ // msbs = mulhi(n, mp)
+ umulExtended(n, mp, msbs, lsbs);
+ return (msbs + n) >> L;
+}
+
+void split_crs(uint32_t crs_idx, out uint32_t ic, out uint32_t kd, out uint32_t kh, out uint32_t kw) {
+ const uint32_t KHKW = KH * KW;
+ const uint32_t KDKHKW = KD * KHKW;
+ ic = crs_idx / KDKHKW;
+ uint32_t rem = crs_idx - ic * KDKHKW;
+ kd = rem / KHKW;
+ rem = rem - kd * KHKW;
+ kh = rem / KW;
+ kw = rem - kh * KW;
+}
+
+void split_npq(uint32_t npq_idx, out uint32_t n, out uint32_t od, out uint32_t oh, out uint32_t ow) {
+ const uint32_t OWOH = p.OW * p.OH;
+ n = fastdiv(npq_idx, p.OWOHODmp, p.OWOHODL);
+ uint32_t rem = npq_idx - n * p.OD * OWOH;
+ od = fastdiv(rem, p.OWOHmp, p.OWOHL);
+ rem = rem - od * OWOH;
+ oh = fastdiv(rem, p.OWmp, p.OWL);
+ ow = rem - oh * p.OW;
+}
+
+#ifdef COOPMAT2
+#define ACC_TYPE float16_t
+
+ACC_TYPE perElemOpStore(const in uint32_t r, const in uint32_t c, const in ACC_TYPE elem)
+{
+ uint32_t K_idx = B_idx_K * BS_K + r;
+ uint32_t NPQ_idx = B_idx_NPQ * BS_NPQ + c;
+ uint32_t N_idx;
+ uint32_t OD_idx;
+ uint32_t OH_idx;
+ uint32_t OW_idx;
+ split_npq(NPQ_idx, N_idx, OD_idx, OH_idx, OW_idx);
+ uint32_t dst_idx = OW_idx + OH_idx * p.nb1 + OD_idx * p.nb2 + (N_idx * p.OC + K_idx) * p.nb3;
+ if (aligned != 0 || (K_idx < K && NPQ_idx < NPQ)) {
+ dst_data[dst_idx] = D_TYPE(elem);
+ }
+ return elem;
+}
+#endif
+
+void main() {
+ if (B_idx_NPQ * BS_NPQ >= NPQ) {
+ return;
+ }
+
+#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++) {
+ for (uint32_t T_lx = 0; T_lx < TS_NPQ; T_lx++) {
+ regC[T_ly][T_lx] = 0.0;
+ }
+ }
+#endif
+ /* Advance block in CRS dim */
+ [[dont_unroll]] for (uint32_t B_idx_CRS = 0; B_idx_CRS < NB_CRS; B_idx_CRS++) {
+ uint32_t CRS_idx_a = B_idx_CRS * BS_CRS + Ac;
+ uint32_t IC_idx_a;
+ uint32_t KD_idx_a;
+ uint32_t KH_idx_a;
+ uint32_t KW_idx_a;
+ split_crs(CRS_idx_a, IC_idx_a, KD_idx_a, KH_idx_a, KW_idx_a);
+
+ /* Load kernel to A_block: (BS_K x BS_CRS)*/
+ UNROLL for (uint32_t r_offset = 0; r_offset < BS_K; r_offset += ArpWg) {
+ uint32_t B_ly = r_offset + Ar;
+ uint32_t B_lx = Ac;
+ uint32_t K_idx = B_idx_K * BS_K + B_ly; /* Global K_idx (row index of A)*/
+ uint32_t knl_idx = KW_idx_a + KH_idx_a * p.nb01 + KD_idx_a * p.nb02 + (K_idx * p.IC + IC_idx_a) * p.nb03;
+ if (aligned == 0) {
+ knl_idx = min(knl_idx, K * CRS - 1);
+ }
+ float val = knl_data[knl_idx];
+ if (aligned == 0 && (K_idx >= K || CRS_idx_a >= CRS)) {
+ val = 0.0;
+ }
+ Ash[B_ly * Ash_stride + B_lx] = SHMEM_TYPE(val);
+ }
+ /* Load input to B_block: (BS_CRS x BS_NPQ) */
+ UNROLL for (uint32_t r_offset = 0; r_offset < BS_CRS; r_offset += BrpWg) {
+ uint32_t B_ly = r_offset + Br; /* Row index of B block */
+ uint32_t B_lx = Bc;
+ uint32_t NPQ_idx = B_idx_NPQ * BS_NPQ + B_lx; /* Global NPQ index (column index of B) */
+ uint32_t N_idx;
+ uint32_t OD_idx;
+ uint32_t OH_idx;
+ uint32_t OW_idx;
+ split_npq(NPQ_idx, N_idx, OD_idx, OH_idx, OW_idx);
+
+ uint32_t CRS_idx_b = B_idx_CRS * BS_CRS + B_ly;
+ uint32_t IC_idx_b;
+ uint32_t KD_idx_b;
+ uint32_t KH_idx_b;
+ uint32_t KW_idx_b;
+ split_crs(CRS_idx_b, IC_idx_b, KD_idx_b, KH_idx_b, KW_idx_b);
+
+ uint32_t ID_idx = OD_idx * s2 + KD_idx_b * d2 - p2;
+ uint32_t IH_idx = OH_idx * s1 + KH_idx_b * d1 - p1;
+ uint32_t IW_idx = OW_idx * s0 + KW_idx_b * d0 - p0;
+
+ uint32_t src_idx = IW_idx + IH_idx * p.nb11 + ID_idx * p.nb12 + (N_idx * p.IC + IC_idx_b) * p.nb13;
+ // skip clamp when address can't go OOB
+ if (aligned == 0 || !dhw_in_bounds) {
+ src_idx = min(src_idx, p.IC * p.N * p.IW * p.IH * p.ID - 1);
+ }
+ float val = src_data[src_idx];
+ 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 (!dhw_in_bounds && (ID_idx >= p.ID || IH_idx >= p.IH || IW_idx >= p.IW)) {
+ oob = true;
+ }
+ if (oob) {
+ val = 0.0;
+ }
+ Bsh[B_ly * Bsh_stride + B_lx] = SHMEM_TYPE(val);
+ }
+ barrier();
+#ifdef COOPMAT2
+ coopmat<float16_t, gl_ScopeWorkgroup, BS_K, BS_CRS, gl_MatrixUseA> matA;
+ coopmat<float16_t, gl_ScopeWorkgroup, BS_CRS, BS_NPQ, gl_MatrixUseB> matB;
+
+ 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++) {
+ float regA[TS_K];
+ float regB[TS_NPQ];
+ for (uint32_t T_ly = 0; T_ly < TS_K; T_ly++) {
+ regA[T_ly] = Ash[(T_y * TS_K + T_ly) * Ash_stride + CRS_lidx];
+ }
+ for (uint32_t T_lx = 0; T_lx < TS_NPQ; T_lx++) {
+ regB[T_lx] = Bsh[CRS_lidx * Bsh_stride + T_x * TS_NPQ + T_lx];
+ }
+ for (uint32_t T_ly = 0; T_ly < TS_K; T_ly++) {
+ for (uint32_t T_lx = 0; T_lx < TS_NPQ; T_lx++) {
+ regC[T_ly][T_lx] = fma(regA[T_ly], regB[T_lx], regC[T_ly][T_lx]);
+ }
+ }
+ }
+ }
+#endif
+ 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;
+ uint32_t OD_idx;
+ uint32_t OH_idx;
+ uint32_t OW_idx;
+ split_npq(NPQ_idx, N_idx, OD_idx, OH_idx, OW_idx);
+ uint32_t dst_idx = OW_idx + OH_idx * p.nb1 + OD_idx * p.nb2 + (N_idx * p.OC + K_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
+ 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++) {
+ for (uint32_t T_lx = 0; T_lx < TS_NPQ; T_lx++) {
+ uint32_t K_idx = B_idx_K * BS_K + T_y * TS_K + T_ly;
+ uint32_t NPQ_idx = B_idx_NPQ * BS_NPQ + T_x * TS_NPQ + T_lx;
+ uint32_t N_idx;
+ uint32_t OD_idx;
+ uint32_t OH_idx;
+ uint32_t OW_idx;
+ split_npq(NPQ_idx, N_idx, OD_idx, OH_idx, OW_idx);
+ uint32_t dst_idx = OW_idx + OH_idx * p.nb1 + OD_idx * p.nb2 + (N_idx * p.OC + K_idx) * p.nb3;
+ if (aligned != 0 || (K_idx < K && NPQ_idx < NPQ)) {
+ dst_data[dst_idx] = D_TYPE(regC[T_ly][T_lx]);
+ }
+ }
+ }
+ }
+#endif
+}
}
}
+ for (auto unroll : {false, true}) {
+ for (auto a_f16 : {false, true}) {
+ std::map<std::string, std::string> defines = {
+ {"A_TYPE", a_f16 ? "float16_t" : "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"},
+ {"UNROLL", unroll ? "[[unroll]]" : ""},
+ };
+ std::string name = std::string("conv3d") + (a_f16 ? "_f16" : "") + "_f32";
+ string_to_spv(name + (unroll ? "_unroll" : ""), "conv3d_mm.comp", defines);
+#if defined(GGML_VULKAN_COOPMAT2_GLSLC_SUPPORT)
+ if (unroll) {
+ auto cm2_defines = defines;
+ cm2_defines["COOPMAT2"] = "1";
+ string_to_spv(name, "conv3d_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, "conv3d_mm.comp", cm1_defines, true, true, false);
+ }
+#endif
+ }
+ }
+
string_to_spv("conv2d_dw_whcn_f32", "conv2d_dw.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"WHCN", "1"}}));
string_to_spv("conv2d_dw_cwhn_f32", "conv2d_dw.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"CWHN", "1"}}));
string_to_spv("conv2d_dw_whcn_f16_f32", "conv2d_dw.comp", merge_maps(base_dict, {{"A_TYPE", "float16_t"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}, {"WHCN", "1"}}));
}
}
+ struct conv3d_perf_case {
+ int N, IC, ID, IH, IW, OC, KD, KH, KW, s0, s1, s2, p0, p1, p2, d0, d1, d2;
+ };
+
+ const std::vector<conv3d_perf_case> conv3d_cases = {
+ {1, 320, 8, 38, 26, 1280, 3, 3, 3, 1, 1, 1, 1, 1, 1, 1, 1, 1},
+ {1, 1280, 8, 38, 26, 1280, 3, 3, 3, 1, 1, 1, 1, 1, 1, 1, 1, 1},
+ {1, 320, 8, 76, 52, 1280, 3, 3, 3, 1, 1, 1, 1, 1, 1, 1, 1, 1},
+ {1, 1280, 8, 76, 52, 1280, 3, 3, 3, 1, 1, 1, 1, 1, 1, 1, 1, 1},
+ {1, 320, 8, 152, 104, 1280, 3, 3, 3, 1, 1, 1, 1, 1, 1, 1, 1, 1},
+#if 0
+ // too slow on some devices
+ {1, 1280, 8, 152, 104, 1280, 3, 3, 3, 1, 1, 1, 1, 1, 1, 1, 1, 1},
+ {1, 320, 4, 304, 208, 1280, 3, 3, 3, 1, 1, 1, 1, 1, 1, 1, 1, 1},
+ {1, 640, 4, 304, 208, 1280, 3, 3, 3, 1, 1, 1, 1, 1, 1, 1, 1, 1},
+#endif
+ };
+
+ for (ggml_type kernel_type : {GGML_TYPE_F32, GGML_TYPE_F16}) {
+ for (const conv3d_perf_case & c : conv3d_cases) {
+ test_cases.emplace_back(new test_conv_3d(
+ c.N, c.IC, c.ID, c.IH, c.IW,
+ c.OC, c.KD, c.KH, c.KW,
+ c.s0, c.s1, c.s2, c.p0, c.p1, c.p2, c.d0, c.d1, c.d2,
+ kernel_type));
+ }
+ }
+
test_cases.emplace_back(new test_bin_bcast(ggml_add, GGML_TYPE_F32, {4096, 1, 1, 1}, {1, 1, 1, 1}));
test_cases.emplace_back(new test_bin_bcast(ggml_add, GGML_TYPE_F32, {4096, 1, 1, 1}, {1, 512, 1, 1}));
overview / pkg / ggml / sources / llama.cpp / commitdiff