diff --git a/ggml/src/ggml-vulkan/ggml-vulkan-push-constants.h b/ggml/src/ggml-vulkan/ggml-vulkan-push-constants.h index 8446e313c..f066d1064 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 c64004cdc..5ed15a258 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 d5ce4290b..6919861b3 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]); } } }