From 6b790a9c291b5d7af3312bbf9f0c558aa023b13e Mon Sep 17 00:00:00 2001 From: Ruben Ortlam Date: Thu, 24 Sep 2026 14:08:26 +0200 Subject: [PATCH] vulkan: handle misalignment in conv_2d and conv_3d (#29365) * vulkan: handle misalignment in conv_2d and conv_3d * fix test-backend-ops print --- .../ggml-vulkan/ggml-vulkan-push-constants.h | 26 +++++++++++++++++++ .../ggml-vulkan/vulkan-shaders/conv2d_mm.comp | 15 +++++++---- .../ggml-vulkan/vulkan-shaders/conv3d_mm.comp | 15 +++++++---- tests/test-backend-ops.cpp | 23 +++++++++++++--- 4 files changed, 65 insertions(+), 14 deletions(-) diff --git a/ggml/src/ggml-vulkan/ggml-vulkan-push-constants.h b/ggml/src/ggml-vulkan/ggml-vulkan-push-constants.h index 8446e313c2..f066d10644 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan-push-constants.h +++ b/ggml/src/ggml-vulkan/ggml-vulkan-push-constants.h @@ -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); diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/conv2d_mm.comp b/ggml/src/ggml-vulkan/vulkan-shaders/conv2d_mm.comp index c64004cdc4..5ed15a258d 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/conv2d_mm.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/conv2d_mm.comp @@ -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]; } } } diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/conv3d_mm.comp b/ggml/src/ggml-vulkan/vulkan-shaders/conv3d_mm.comp index d5ce4290b9..6919861b3b 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/conv3d_mm.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/conv3d_mm.comp @@ -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]); } } } diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index 4c67805d45..e198d09221 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -6258,9 +6258,10 @@ struct test_conv_2d : public test_case { const int dilation1; // Whether the inputs are contiguous in the channel dim or the width dim const bool cwhn; + const int kernel_offset; std::string vars() override { - return VARS_TO_STR10(ne_input, ne_kernel, type_kernel, stride0, stride1, padding0, padding1, dilation0, dilation1, cwhn); + return VARS_TO_STR11(ne_input, ne_kernel, type_kernel, stride0, stride1, padding0, padding1, dilation0, dilation1, cwhn, kernel_offset); } double max_nmse_err() override { @@ -6296,7 +6297,8 @@ struct test_conv_2d : public test_case { test_conv_2d(std::array ne_input = { 64, 64, 16, 1 }, std::array ne_kernel = { 3, 3, 1, 16 }, ggml_type type_kernel = GGML_TYPE_F32, int stride0 = 1, - int stride1 = 1, int padding0 = 0, int padding1 = 0, int dilation0 = 1, int dilation1 = 1, bool cwhn = false) : + int stride1 = 1, int padding0 = 0, int padding1 = 0, int dilation0 = 1, int dilation1 = 1, bool cwhn = false, + int kernel_offset = 0) : ne_input(ne_input), ne_kernel(ne_kernel), type_kernel(type_kernel), @@ -6306,13 +6308,25 @@ struct test_conv_2d : public test_case { padding1(padding1), dilation0(dilation0), dilation1(dilation1), - cwhn(cwhn) {} + cwhn(cwhn), + kernel_offset(kernel_offset) {} ggml_tensor * build_graph(ggml_context * ctx) override { ggml_tensor * input = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne_input.data()); ggml_set_name(input, "input"); - ggml_tensor * kernel = ggml_new_tensor(ctx, type_kernel, 4, ne_kernel.data()); + ggml_tensor * kernel; + if (kernel_offset == 0) { + kernel = ggml_new_tensor(ctx, type_kernel, 4, ne_kernel.data()); + } else { + const int64_t nelem = ne_kernel[0] * ne_kernel[1] * ne_kernel[2] * ne_kernel[3]; + ggml_tensor * storage = ggml_new_tensor_1d(ctx, type_kernel, nelem + kernel_offset); + const size_t element_size = ggml_type_size(type_kernel); + kernel = ggml_view_4d(ctx, storage, ne_kernel[0], ne_kernel[1], ne_kernel[2], ne_kernel[3], + ne_kernel[0] * element_size, ne_kernel[0] * ne_kernel[1] * element_size, + ne_kernel[0] * ne_kernel[1] * ne_kernel[2] * element_size, + kernel_offset * element_size); + } ggml_set_name(kernel, "kernel"); if (cwhn) { @@ -9389,6 +9403,7 @@ static std::vector> make_test_cases_eval() { test_cases.emplace_back(new test_conv_2d({ 19, 17, 8, 2 }, { 3, 3, 8, 65 }, GGML_TYPE_F16, 1, 1, 1, 1, 1, 1)); test_cases.emplace_back(new test_conv_2d({ 19, 17, 16, 3 }, { 3, 3, 16, 33 }, GGML_TYPE_F16, 2, 3, 4, 2, 2, 1)); test_cases.emplace_back(new test_conv_2d({ 13, 11, 16, 3 }, { 1, 1, 16, 33 }, GGML_TYPE_F16, 1, 1, 0, 0, 1, 1)); + test_cases.emplace_back(new test_conv_2d({ 19, 17, 8, 2 }, { 3, 3, 8, 17 }, GGML_TYPE_F16, 1, 1, 1, 1, 1, 1, false, 1)); // sycl backend will limit task global_range < MAX_INT // test cases for 2D im2col with large input W and H (occurs in stable-diffusion)