diff --git a/mlx/backend/metal/conv.cpp b/mlx/backend/metal/conv.cpp index 4fe62a716c..3a7c394d48 100644 --- a/mlx/backend/metal/conv.cpp +++ b/mlx/backend/metal/conv.cpp @@ -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(implicit_M) + bm - 1) / bm; int swizzle_log = 0; // Fix small channel specialization @@ -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(adj_implicit_m) + bm - 1) / bm; int swizzle_log = 0; // Get channel iteration info @@ -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(implicit_M) + bm - 1) / bm; int swizzle_log = 0; bool small_filter = diff --git a/mlx/backend/metal/kernels/conv.metal b/mlx/backend/metal/kernels/conv.metal index 6f3ce0cd43..a16ae8a275 100644 --- a/mlx/backend/metal/kernels/conv.metal +++ b/mlx/backend/metal/kernels/conv.metal @@ -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]; diff --git a/mlx/backend/metal/kernels/steel/conv/kernels/steel_conv_general.h b/mlx/backend/metal/kernels/steel/conv/kernels/steel_conv_general.h index 38250f9b81..bb896711c4 100644 --- a/mlx/backend/metal/kernels/steel/conv/kernels/steel_conv_general.h +++ b/mlx/backend/metal/kernels/steel/conv/kernels/steel_conv_general.h @@ -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++) { diff --git a/mlx/backend/metal/kernels/steel/conv/loaders/loader_channel_l.h b/mlx/backend/metal/kernels/steel/conv/loaders/loader_channel_l.h index 15dc13cccf..62881c755a 100644 --- a/mlx/backend/metal/kernels/steel/conv/loaders/loader_channel_l.h +++ b/mlx/backend/metal/kernels/steel/conv/loaders/loader_channel_l.h @@ -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]; @@ -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]; @@ -384,6 +384,7 @@ struct Conv2DWeightBlockLoader { int weight_step; const int read_n; + const int N; const bool do_read; /* Constructor */ @@ -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 */ @@ -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]; @@ -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; @@ -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]; @@ -887,6 +889,7 @@ struct Conv3DWeightBlockLoader { int weight_step; const int read_n; + const int N; const bool do_read; /* Constructor */ @@ -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 */ @@ -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]; diff --git a/mlx/backend/metal/kernels/steel/conv/loaders/loader_channel_n.h b/mlx/backend/metal/kernels/steel/conv/loaders/loader_channel_n.h index 1562ea0b3a..452f344af0 100644 --- a/mlx/backend/metal/kernels/steel/conv/loaders/loader_channel_n.h +++ b/mlx/backend/metal/kernels/steel/conv/loaders/loader_channel_n.h @@ -234,6 +234,7 @@ struct Conv2DWeightBlockLoaderSmallChannels { int weight_hw; const int read_n; + const int N; const bool do_read; /* Constructor */ @@ -254,6 +255,7 @@ 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 */ @@ -261,7 +263,7 @@ struct Conv2DWeightBlockLoaderSmallChannels { 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 @@ -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]; diff --git a/mlx/backend/metal/kernels/steel/conv/loaders/loader_general.h b/mlx/backend/metal/kernels/steel/conv/loaders/loader_general.h index cf7794235d..6bc90e8694 100644 --- a/mlx/backend/metal/kernels/steel/conv/loaders/loader_general.h +++ b/mlx/backend/metal/kernels/steel/conv/loaders/loader_general.h @@ -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]; @@ -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; diff --git a/python/tests/test_conv.py b/python/tests/test_conv.py index 445958ad2a..684074889b 100644 --- a/python/tests/test_conv.py +++ b/python/tests/test_conv.py @@ -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()