vulkan: handle misalignment in conv_2d and conv_3d (#29365)

* vulkan: handle misalignment in conv_2d and conv_3d

* fix test-backend-ops print
This commit is contained in:
Ruben Ortlam
2026-09-24 14:08:26 +02:00
committed by GitHub
parent 3423f940e8
commit 6b790a9c29
4 changed files with 65 additions and 14 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]);
}
}
}
+19 -4
View File
@@ -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<int64_t, 4> ne_input = { 64, 64, 16, 1 },
std::array<int64_t, 4> 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<std::unique_ptr<test_case>> 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)