Skip to content

metal: support non-zero left padding for PAD #25827

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
flyingtimes wants to merge 1 commit into from
Closed

metal: support non-zero left padding for PAD #25827

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 4 additions & 2 deletions ggml/src/ggml-metal/ggml-metal-device.m
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1214,8 +1214,10 @@ bool ggml_metal_device_supports_op(ggml_metal_device_t dev, const struct ggml_te
return false;
}

return (ggml_get_op_params_i32(op, 0) == 0) && (ggml_get_op_params_i32(op, 2) == 0) &&
(ggml_get_op_params_i32(op, 4) == 0) && (ggml_get_op_params_i32(op, 6) == 0);
// kernel_pad_impl now supports non-zero left padding on all dims,
// so we no longer require lp0/lp1/lp2/lp3 == 0. Element type is
// constrained separately by the pipeline name lookup.
return true;
case GGML_OP_PAD_REFLECT_1D:
case GGML_OP_TIMESTEP_EMBEDDING:
case GGML_OP_LEAKY_RELU:
Expand Down
4 changes: 4 additions & 0 deletions ggml/src/ggml-metal/ggml-metal-impl.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1053,6 +1053,10 @@ typedef struct {
uint64_t nb1;
uint64_t nb2;
uint64_t nb3;
int32_t lp0; // left padding per dim (new: support non-zero left pad)
int32_t lp1;
int32_t lp2;
int32_t lp3;
} ggml_metal_kargs_pad;

typedef struct {
Expand Down
6 changes: 5 additions & 1 deletion ggml/src/ggml-metal/ggml-metal-ops.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4284,7 +4284,11 @@ int ggml_metal_op_pad(ggml_metal_op_t ctx, int idx) {
/*.nb0 =*/ nb0,
/*.nb1 =*/ nb1,
/*.nb2 =*/ nb2,
/*.nb3 =*/ nb3
/*.nb3 =*/ nb3,
/*.lp0 =*/ ggml_get_op_params_i32(op, 0),
/*.lp1 =*/ ggml_get_op_params_i32(op, 2),
/*.lp2 =*/ ggml_get_op_params_i32(op, 4),
/*.lp3 =*/ ggml_get_op_params_i32(op, 6)
};

auto pipeline = ggml_metal_library_get_pipeline_pad(lib, op);
Expand Down
18 changes: 13 additions & 5 deletions ggml/src/ggml-metal/ggml-metal.metal
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5842,9 +5842,15 @@ kernel void kernel_pad_impl(
const int32_t k0 = tgpig.x/args.ne1;
const int32_t i1 = tgpig.x - k0*args.ne1;

const int32_t i03 = i3;
const int32_t i02 = i2;
const int32_t i01 = i1;
// Source coords = dst coords minus left padding. If a source coord falls
// outside [0, ne0x) the whole row/column is in the padding region -> fill 0.
const int32_t i03 = i3 - args.lp3;
const int32_t i02 = i2 - args.lp2;
const int32_t i01 = i1 - args.lp1;

const bool row_in_src = (i01 >= 0 && i01 < args.ne01)
&& (i02 >= 0 && i02 < args.ne02)
&& (i03 >= 0 && i03 < args.ne03);

device const T * src0_ptr = (device const T *) (src0 + i03*args.nb03 + i02*args.nb02 + i01*args.nb01);
device T * dst_ptr = (device T *) (dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1);
Expand All @@ -5855,8 +5861,10 @@ kernel void kernel_pad_impl(
break;
}

if (i0 < args.ne00 && i1 < args.ne01 && i2 < args.ne02 && i3 < args.ne03) {
dst_ptr[i0] = src0_ptr[i0];
// dim 0 source coord, with left padding offset
const int32_t i00 = i0 - args.lp0;
if (row_in_src && i00 >= 0 && i00 < args.ne00) {
dst_ptr[i0] = src0_ptr[i00];
} else {
dst_ptr[i0] = 0.0f;
}
Expand Down
6 changes: 6 additions & 0 deletions tests/test-backend-ops.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -9390,6 +9390,12 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
test_cases.emplace_back(new test_pad());
test_cases.emplace_back(new test_pad(GGML_TYPE_F32, {33, 17, 2, 1}, 4, 3, true)); // circular
test_cases.emplace_back(new test_pad_ext());
// left-padding only (lp > 0, rp == 0). Regression coverage for the Metal
// pad kernel, which previously only handled right padding (lp == 0).
test_cases.emplace_back(new test_pad_ext(GGML_TYPE_F32, {101, 1024, 1, 1}, 2, 0, 0, 0, 0, 0, 0, 0)); // mirrors CosyVoice's flow PAD
test_cases.emplace_back(new test_pad_ext(GGML_TYPE_F32, {512, 512, 1, 1}, 4, 0, 0, 0, 0, 0, 0, 0)); // dim0 left only
test_cases.emplace_back(new test_pad_ext(GGML_TYPE_F32, {512, 512, 1, 1}, 0, 0, 3, 0, 0, 0, 0, 0)); // dim1 left only
test_cases.emplace_back(new test_pad_ext(GGML_TYPE_F32, {64, 64, 3, 1}, 2, 3, 1, 0, 0, 0, 0, 0)); // mixed left+right
test_cases.emplace_back(new test_pad(GGML_TYPE_F32, {1024, 1, 1, 1}, 1, 0, false));
test_cases.emplace_back(new test_pad(GGML_TYPE_F32, {1024, 2, 1, 1}, 1, 0, false));
test_cases.emplace_back(new test_pad(GGML_TYPE_F32, {1024, 16, 1, 1}, 0, 1, false));
Expand Down