vulkan: handle misalignment in conv_2d and conv_3d (llama/29365)

* vulkan: handle misalignment in conv_2d and conv_3d

* fix test-backend-ops print
This commit is contained in:
Ruben Ortlam
2026-10-06 10:34:31 +03:00
committed by Georgi Gerganov
parent bf8128ec03
commit f938f31926
3 changed files with 46 additions and 10 deletions
@@ -742,6 +742,10 @@ struct vk_op_conv2d_push_constants {
// init_fastdiv_values constants for dividing by OW, OW*OH
uint32_t OWmp; uint32_t OWL;
uint32_t OWOHmp; uint32_t OWOHL;
uint32_t knl_offset;
uint32_t src_offset;
uint32_t dst_offset;
};
template <> inline void init_pushconst_fastdiv(vk_op_conv2d_push_constants &p) {
@@ -777,6 +781,10 @@ struct vk_op_conv3d_push_constants {
uint32_t OWmp; uint32_t OWL;
uint32_t OWOHmp; uint32_t OWOHL;
uint32_t OWOHODmp; uint32_t OWOHODL;
uint32_t knl_offset;
uint32_t src_offset;
uint32_t dst_offset;
};
template <> inline void init_pushconst_fastdiv(vk_op_conv3d_push_constants &p) {
@@ -1032,6 +1040,24 @@ template <> inline void init_pushconst_tensor_offsets(ggml_backend_vk_context *
GGML_UNUSED(src3);
}
template <> inline void init_pushconst_tensor_offsets(ggml_backend_vk_context * ctx, vk_op_conv2d_push_constants &p, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, const ggml_tensor * src3, ggml_tensor * dst) {
p.knl_offset = get_misalign_bytes(ctx, src0) / ggml_type_size(src0->type);
p.src_offset = get_misalign_bytes(ctx, src1) / ggml_type_size(src1->type);
p.dst_offset = get_misalign_bytes(ctx, dst) / ggml_type_size(dst->type);
GGML_UNUSED(src2);
GGML_UNUSED(src3);
}
template <> inline void init_pushconst_tensor_offsets(ggml_backend_vk_context * ctx, vk_op_conv3d_push_constants &p, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, const ggml_tensor * src3, ggml_tensor * dst) {
p.knl_offset = get_misalign_bytes(ctx, src0) / ggml_type_size(src0->type);
p.src_offset = get_misalign_bytes(ctx, src1) / ggml_type_size(src1->type);
p.dst_offset = get_misalign_bytes(ctx, dst) / ggml_type_size(dst->type);
GGML_UNUSED(src2);
GGML_UNUSED(src3);
}
template <> inline void init_pushconst_tensor_offsets(ggml_backend_vk_context * ctx, vk_op_im2col_3d_push_constants &p, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, const ggml_tensor * src3, ggml_tensor * dst) {
const uint32_t a_offset = get_misalign_bytes(ctx, src1) / ggml_type_size(src1->type);
const uint32_t d_offset = get_misalign_bytes(ctx, dst) / ggml_type_size(dst->type);
@@ -62,6 +62,11 @@ layout(push_constant) uniform parameter {
// fastdiv helper values
uint32_t OWmp; uint32_t OWL;
uint32_t OWOHmp; uint32_t OWOHL;
// element offsets for misaligned buffer bindings
uint32_t knl_offset;
uint32_t src_offset;
uint32_t dst_offset;
}
p;
@@ -206,7 +211,7 @@ ACC_TYPE perElemOpStore(const in uint32_t r, const in uint32_t c, const in ACC_T
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(elem);
dst_data[dst_idx + p.dst_offset] = D_TYPE(elem);
}
return elem;
}
@@ -286,7 +291,7 @@ void main() {
if (aligned == 0) {
knl_idx = min(knl_idx, K * CRS - 1);
}
float val = knl_data[knl_idx];
float val = knl_data[knl_idx + p.knl_offset];
if (aligned == 0 && (K_idx >= K || CRS_idx_a >= CRS)) {
val = 0.0;
}
@@ -341,7 +346,7 @@ void main() {
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];
float val = src_data[src_idx + p.src_offset];
bool oob = false;
if (aligned == 0 && (CRS_idx_b >= CRS || NPQ_idx >= NPQ)) {
oob = true;
@@ -444,7 +449,7 @@ void main() {
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]);
dst_data[dst_idx + p.dst_offset] = D_TYPE(Csh[k_local * Csh_stride + npq_thread]);
}
}
}
@@ -464,7 +469,7 @@ void main() {
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] = regC[T_ly][T_lx];
dst_data[dst_idx + p.dst_offset] = regC[T_ly][T_lx];
}
}
}
@@ -61,6 +61,11 @@ layout(push_constant) uniform parameter {
uint32_t OWmp; uint32_t OWL;
uint32_t OWOHmp; uint32_t OWOHL;
uint32_t OWOHODmp; uint32_t OWOHODL;
// element offsets for misaligned buffer bindings
uint32_t knl_offset;
uint32_t src_offset;
uint32_t dst_offset;
}
p;
@@ -214,7 +219,7 @@ ACC_TYPE perElemOpStore(const in uint32_t r, const in uint32_t c, const in ACC_T
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);
dst_data[dst_idx + p.dst_offset] = D_TYPE(elem);
}
return elem;
}
@@ -261,7 +266,7 @@ void main() {
if (aligned == 0) {
knl_idx = min(knl_idx, K * CRS - 1);
}
float val = knl_data[knl_idx];
float val = knl_data[knl_idx + p.knl_offset];
if (aligned == 0 && (K_idx >= K || CRS_idx_a >= CRS)) {
val = 0.0;
}
@@ -294,7 +299,7 @@ void main() {
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];
float val = src_data[src_idx + p.src_offset];
bool oob = false;
if (aligned == 0 && (CRS_idx_b >= CRS || NPQ_idx >= NPQ)) {
oob = true;
@@ -393,7 +398,7 @@ void main() {
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]);
dst_data[dst_idx + p.dst_offset] = D_TYPE(Csh[k_local * Csh_stride + npq_thread]);
}
}
}
@@ -415,7 +420,7 @@ void main() {
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]);
dst_data[dst_idx + p.dst_offset] = D_TYPE(regC[T_ly][T_lx]);
}
}
}