Skip to content

fix metal conv bounds checks and integer overflows #4672

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

Merged
zcbenz merged 2 commits into from
Oct 12, 2026
Merged

fix metal conv bounds checks and integer overflows #4672

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
372330f
fix metal conv bounds checks and integer overflows
Prudctual Oct 10, 2026
d8f41d9
nit
zcbenz Oct 12, 2026
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: 3 additions & 3 deletions mlx/backend/metal/conv.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -356,7 +356,7 @@ void implicit_gemm_conv_2D_gpu(
}

int tn = (implicit_N + bn - 1) / bn;
int tm = (implicit_M + bm - 1) / bm;
int tm = (static_cast<int64_t>(implicit_M) + bm - 1) / bm;
int swizzle_log = 0;

// Fix small channel specialization
Expand Down Expand Up @@ -554,7 +554,7 @@ void implicit_gemm_conv_2D_general_gpu(
int bk = 16;

int tn = (implicit_N + bn - 1) / bn;
int tm = (adj_implicit_m + bm - 1) / bm;
int tm = (static_cast<int64_t>(adj_implicit_m) + bm - 1) / bm;
int swizzle_log = 0;

// Get channel iteration info
Expand Down Expand Up @@ -678,7 +678,7 @@ void implicit_gemm_conv_3D_gpu(
}

int tn = (implicit_N + bn - 1) / bn;
int tm = (implicit_M + bm - 1) / bm;
int tm = (static_cast<int64_t>(implicit_M) + bm - 1) / bm;
int swizzle_log = 0;

bool small_filter =
Expand Down
4 changes: 2 additions & 2 deletions mlx/backend/metal/kernels/conv.metal
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -418,8 +418,8 @@ winograd_conv_2d_weight_transform(
// wt_out is stored transposed (A x A x C x O)
short ohw_0 = sm * 8 + sn;
short ohw_1 = sm * 8 + sn + 1;
device T* wt_out_0 = wt_out + ohw_0 * C * O + ko;
device T* wt_out_1 = wt_out + ohw_1 * C * O + ko;
device T* wt_out_0 = wt_out + size_t(ohw_0) * C * O + ko;
device T* wt_out_1 = wt_out + size_t(ohw_1) * C * O + ko;

// Prepare shared memory
threadgroup T Ws[BO][R][R][BC];
Expand Down
2 changes: 1 addition & 1 deletion mlx/backend/metal/kernels/steel/conv/kernels/steel_conv_general.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -186,7 +186,7 @@ implicit_gemm_conv_2d_general(
if (offset_n >= gemm_params->N)
return;

short diff = gemm_params->N - offset_n;
int diff = gemm_params->N - offset_n;

STEEL_PRAGMA_UNROLL
for (int i = 0; i < mma_t::TM; i++) {
Expand Down
28 changes: 16 additions & 12 deletions mlx/backend/metal/kernels/steel/conv/loaders/loader_channel_l.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -47,8 +47,8 @@ struct Conv2DInputBlockLoaderLargeFilter {
const constant MLXConvParams<2>* params;
const constant ImplicitGemmConv2DParams* gemm_params;

short weight_h;
short weight_w;
int weight_h;
int weight_w;

const device T* src[n_rows];

Expand Down Expand Up @@ -197,8 +197,8 @@ struct Conv2DInputBlockLoaderSmallFilter {
const constant MLXConvParams<2>* params;
const constant ImplicitGemmConv2DParams* gemm_params;

short weight_h;
short weight_w;
int weight_h;
int weight_w;

const device T* src[n_rows];

Expand Down Expand Up @@ -384,6 +384,7 @@ struct Conv2DWeightBlockLoader {
int weight_step;

const int read_n;
const int N;
const bool do_read;

/* Constructor */
Expand All @@ -405,6 +406,7 @@ struct Conv2DWeightBlockLoader {
weight_hw(0),
weight_step(params->C / params->groups),
read_n(offsets.y + bi),
N(gemm_params_->N),
do_read(read_n + n_rows * TROWS <= gemm_params_->N) {}

/* Load from device memory into threadgroup memory - without bound checking */
Expand All @@ -419,7 +421,7 @@ struct Conv2DWeightBlockLoader {
}
} else {
for (short i = 0; i < BN; i += TROWS) {
if ((read_n + i) < params->O) {
if ((read_n + i) < N) {
STEEL_PRAGMA_UNROLL
for (short j = 0; j < vec_size; j++) {
dst[i * dst_ld + j] = src[i * src_ld + j];
Expand Down Expand Up @@ -481,9 +483,9 @@ struct Conv3DInputBlockLoaderLargeFilter {
const constant MLXConvParams<3>* params;
const constant ImplicitGemmConv3DParams* gemm_params;

short weight_d;
short weight_h;
short weight_w;
int weight_d;
int weight_h;
int weight_w;

short kdil_d;
short kdil_h;
Expand Down Expand Up @@ -665,9 +667,9 @@ struct Conv3DInputBlockLoaderSmallFilter {
const constant MLXConvParams<3>* params;
const constant ImplicitGemmConv3DParams* gemm_params;

short weight_d;
short weight_h;
short weight_w;
int weight_d;
int weight_h;
int weight_w;

const device T* src[n_rows];

Expand Down Expand Up @@ -887,6 +889,7 @@ struct Conv3DWeightBlockLoader {
int weight_step;

const int read_n;
const int N;
const bool do_read;

/* Constructor */
Expand All @@ -908,6 +911,7 @@ struct Conv3DWeightBlockLoader {
weight_dhw(0),
weight_step(params->C / params->groups),
read_n(offsets.y + bi),
N(gemm_params_->N),
do_read(read_n + n_rows * TROWS <= gemm_params_->N) {}

/* Load from device memory into threadgroup memory - without bound checking */
Expand All @@ -922,7 +926,7 @@ struct Conv3DWeightBlockLoader {
}
} else {
for (short i = 0; i < BN; i += TROWS) {
if ((read_n + i) < params->O) {
if ((read_n + i) < N) {
STEEL_PRAGMA_UNROLL
for (short j = 0; j < vec_size; j++) {
dst[i * dst_ld + j] = src[i * src_ld + j];
Expand Down
6 changes: 4 additions & 2 deletions mlx/backend/metal/kernels/steel/conv/loaders/loader_channel_n.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -234,6 +234,7 @@ struct Conv2DWeightBlockLoaderSmallChannels {
int weight_hw;

const int read_n;
const int N;
const bool do_read;

/* Constructor */
Expand All @@ -254,14 +255,15 @@ struct Conv2DWeightBlockLoaderSmallChannels {
params(params_),
weight_hw(thread_idx % TCOLS),
read_n(offsets.y + bi),
N(gemm_params_->N),
do_read(read_n + BN <= gemm_params_->N) {}

/* Load from device memory into threadgroup memory - without bound checking */
METAL_FUNC void load_unsafe() const thread {
if (bi >= BROWS || bj >= BCOLS)
return;

if (read_n >= params->O || weight_hw >= params->wS[1] * params->wS[0]) {
if (read_n >= N || weight_hw >= params->wS[1] * params->wS[0]) {
STEEL_PRAGMA_UNROLL
for (short i = 0; i < BROWS; i += TROWS) {
STEEL_PRAGMA_UNROLL
Expand Down Expand Up @@ -290,7 +292,7 @@ struct Conv2DWeightBlockLoaderSmallChannels {
}
} else {
for (short i = 0; i < BROWS; i += TROWS) {
if (((read_n + i) < params->O)) {
if (((read_n + i) < N)) {
STEEL_PRAGMA_UNROLL
for (short j = 0; j < n_channels; j++) {
dst[i * dst_ld + j] = curr_src[i * src_ld + j];
Expand Down
8 changes: 4 additions & 4 deletions mlx/backend/metal/kernels/steel/conv/loaders/loader_general.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,8 +48,8 @@ struct Conv2DInputBlockLoaderGeneral {
const short base_wh;
const short base_ww;

short weight_h;
short weight_w;
int weight_h;
int weight_w;

const device T* src[n_rows];

Expand Down Expand Up @@ -248,8 +248,8 @@ struct Conv2DWeightBlockLoaderGeneral {
const short base_wh;
const short base_ww;

short weight_h;
short weight_w;
int weight_h;
int weight_w;

const int start_row;

Expand Down
17 changes: 17 additions & 0 deletions python/tests/test_conv.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1617,6 +1617,23 @@ def test_grouped_conv_padding_stays_in_output(self):
mx.eval(wide_y)
self.assertEqual(tuple(wide_y.shape), (1, 1, 500000, 2))

@unittest.skipIf(not mx.is_available(mx.gpu), "requires GPU")
def test_grouped_depthwise_conv_padding_bounds(self):
x = mx.random.normal((1, 16, 64))
w = mx.random.normal((64, 3, 1))
y = mx.conv1d(x, w, padding=1, groups=64)
mx.eval(y)
self.assertEqual(tuple(y.shape), (1, 16, 64))

@unittest.skipIf(not mx.is_available(mx.gpu), "requires GPU")
def test_conv2d_large_output_channels(self):
x = mx.random.normal((1, 33, 33, 5))
w = mx.random.normal((32800, 1, 1, 5))
yg = mx.conv2d(x, w, stride=(2, 2), stream=mx.gpu)
yc = mx.conv2d(x, w, stride=(2, 2), stream=mx.cpu)
mx.eval(yg, yc)
self.assertTrue(mx.allclose(yg, yc, atol=1e-3))


if __name__ == "__main__":
mlx_tests.MLXTestRunner()
Loading