diff --git a/mlx/backend/metal/conv.cpp b/mlx/backend/metal/conv.cpp index 49f87d9a5b..17c7cb1373 100644 --- a/mlx/backend/metal/conv.cpp +++ b/mlx/backend/metal/conv.cpp @@ -30,6 +30,23 @@ ensure_row_contiguous(const array& x, metal::Device& d, const Stream& s) { return result; } +inline int max_unfold_rows(metal::Device& d, size_t row_bytes, int total_rows) { + size_t max_buffer = d.mtl_device()->maxBufferLength(); + size_t max_rows = row_bytes == 0 ? total_rows : max_buffer / row_bytes; + // Force a smaller tile (tests, or to bound the unfold buffer further). + if (int forced = env::get_var("MLX_CONV_UNFOLD_TILE_ROWS", 0); forced > 0) { + max_rows = std::min(max_rows, static_cast(forced)); + } + if (max_rows == 0) { + std::ostringstream msg; + msg << "[conv] A single unfolding row needs " << row_bytes + << " bytes, which exceeds the maximum Metal buffer size of " + << max_buffer << " bytes."; + throw std::runtime_error(msg.str()); + } + return static_cast(std::min(max_rows, static_cast(total_rows))); +} + template void explicit_gemm_conv_ND_gpu( const Stream& s, @@ -42,36 +59,12 @@ void explicit_gemm_conv_ND_gpu( int implicit_M = safe_cast(out.size() / conv_params.O, "conv"); int implicit_K = safe_cast(wt.size() / conv_params.O, "conv"); int implicit_N = conv_params.O; - // Prepare unfolding array - Shape unfolded_shape{implicit_M, implicit_K}; - array in_unfolded(unfolded_shape, in.dtype(), nullptr, {}); - - in_unfolded.set_data(allocator::malloc(in_unfolded.nbytes())); - // Prepare unfolding kernel std::string kname; kname.reserve(32); - concatenate(kname, "naive_unfold_nd_", type_to_name(in_unfolded), "_", N); + concatenate(kname, "naive_unfold_nd_", type_to_name(in), "_", N); auto& compute_encoder = metal::get_command_encoder(s); auto kernel = d.get_kernel(kname); - compute_encoder.set_compute_pipeline_state(kernel); - - compute_encoder.set_input_array(in, 0); - compute_encoder.set_output_array(in_unfolded, 1); - - compute_encoder.set_bytes(conv_params, 2); - - // Launch unfolding kernel - size_t tgp_x = std::min(conv_params.C, 64); - tgp_x = 32 * ((tgp_x + 32 - 1) / 32); - size_t tgp_y = 256 / tgp_x; - - MTL::Size grid_dims = MTL::Size( - conv_params.C, unfolded_shape[1] / conv_params.C, unfolded_shape[0]); - MTL::Size group_dims = MTL::Size( - std::min(tgp_x, grid_dims.width), std::min(tgp_y, grid_dims.height), 1); - - compute_encoder.dispatch_threads(grid_dims, group_dims); // Reshape weight Shape wt_reshape{implicit_K, implicit_N}; @@ -82,23 +75,73 @@ void explicit_gemm_conv_ND_gpu( wt_flags.col_contiguous = true; wt_reshaped.copy_shared_buffer(wt, wt_restride, wt_flags, wt.data_size()); - // Perform gemm - std::vector copies = {in_unfolded}; - return steel_matmul( - s, - d, - /*a = */ in_unfolded, - /*b = */ wt_reshaped, - /*c = */ out, - /*M = */ implicit_M, - /*N = */ implicit_N, - /*K = */ implicit_K, - /*batch_size_out = */ 1, - /*a_cols = */ implicit_K, - /*b_cols = */ implicit_K, - /*a_transposed = */ false, - /*b_transposed = */ true, - /*copies = */ copies); + // 2D view of the output; each tile writes a row window of it. + Strides out_2d_strides{ + out.strides()[out.ndim() - 2], out.strides()[out.ndim() - 1]}; + array out_2d({implicit_M, implicit_N}, out.dtype(), nullptr, {}); + out_2d.copy_shared_buffer(out, out_2d_strides, out.flags(), out.data_size()); + + // The full unfold buffer can exceed maxBufferLength, so unfold and gemm in + // row tiles reusing one buffer, keeping peak memory at a single tile. + size_t row_bytes = static_cast(implicit_K) * in.itemsize(); + int max_rows = max_unfold_rows(d, row_bytes, implicit_M); + + array in_unfolded({max_rows, implicit_K}, in.dtype(), nullptr, {}); + in_unfolded.set_data(allocator::malloc(in_unfolded.nbytes())); + + for (int row_offset = 0; row_offset < implicit_M; row_offset += max_rows) { + int tile_rows = std::min(max_rows, implicit_M - row_offset); + + // Tile view of the reused unfold buffer. + array in_tile({tile_rows, implicit_K}, in.dtype(), nullptr, {}); + in_tile.copy_shared_buffer( + in_unfolded, + in_unfolded.strides(), + in_unfolded.flags(), + in_tile.size()); + + compute_encoder.set_compute_pipeline_state(kernel); + compute_encoder.set_input_array(in, 0); + compute_encoder.set_output_array(in_tile, 1); + compute_encoder.set_bytes(conv_params, 2); + compute_encoder.set_bytes(row_offset, 3); + + size_t tgp_x = std::min(conv_params.C, 64); + tgp_x = 32 * ((tgp_x + 32 - 1) / 32); + size_t tgp_y = 256 / tgp_x; + + MTL::Size grid_dims = + MTL::Size(conv_params.C, implicit_K / conv_params.C, tile_rows); + MTL::Size group_dims = MTL::Size( + std::min(tgp_x, grid_dims.width), std::min(tgp_y, grid_dims.height), 1); + compute_encoder.dispatch_threads(grid_dims, group_dims); + + // Gemm the tile into its output rows. + array out_tile({tile_rows, implicit_N}, out.dtype(), nullptr, {}); + out_tile.copy_shared_buffer( + out_2d, + out_2d_strides, + out_2d.flags(), + out_tile.size(), + static_cast(row_offset) * out_2d_strides[0]); + + std::vector copies = {in_tile}; + steel_matmul( + s, + d, + /*a = */ in_tile, + /*b = */ wt_reshaped, + /*c = */ out_tile, + /*M = */ tile_rows, + /*N = */ implicit_N, + /*K = */ implicit_K, + /*batch_size_out = */ 1, + /*a_cols = */ implicit_K, + /*b_cols = */ implicit_K, + /*a_transposed = */ false, + /*b_transposed = */ true, + /*copies = */ copies); + } } template @@ -122,36 +165,12 @@ void explicit_gemm_conv_group_ND_gpu( kernel_size *= conv_params.wS[i]; } - // Prepare unfolding array - Shape unfolded_shape{implicit_M, implicit_K * groups}; - array in_unfolded(unfolded_shape, in.dtype(), nullptr, {}); - in_unfolded.set_data(allocator::malloc(in_unfolded.nbytes())); - // Prepare unfolding kernel std::string kname; kname.reserve(32); - concatenate( - kname, "naive_unfold_transpose_nd_", type_to_name(in_unfolded), "_", N); + concatenate(kname, "naive_unfold_transpose_nd_", type_to_name(in), "_", N); auto& compute_encoder = metal::get_command_encoder(s); auto kernel = d.get_kernel(kname); - compute_encoder.set_compute_pipeline_state(kernel); - - compute_encoder.set_input_array(in, 0); - compute_encoder.set_output_array(in_unfolded, 1); - - compute_encoder.set_bytes(conv_params, 2); - - // Launch unfolding kernel - size_t tgp_x = std::min(conv_params.C, 64); - tgp_x = 32 * ((tgp_x + 32 - 1) / 32); - size_t tgp_y = 256 / tgp_x; - - MTL::Size grid_dims = MTL::Size( - conv_params.C, unfolded_shape[1] / conv_params.C, unfolded_shape[0]); - MTL::Size group_dims = MTL::Size( - std::min(tgp_x, grid_dims.width), std::min(tgp_y, grid_dims.height), 1); - - compute_encoder.dispatch_threads(grid_dims, group_dims); // Transpose kernel weights so that we can slice them by contiguous chunks // of channel groups. @@ -163,29 +182,79 @@ void explicit_gemm_conv_group_ND_gpu( // Materialize array wt_transpose = contiguous_copy_gpu(wt_view, s); - // Perform gemm - std::vector copies = {in_unfolded, wt_transpose}; - return steel_matmul_regular( - /* const Stream& s = */ s, - /* Device& d = */ d, - /* const array& a = */ in_unfolded, - /* const array& b = */ wt_transpose, - /* array& c = */ out, - /* int M = */ implicit_M, - /* int N = */ implicit_N, - /* int K = */ implicit_K, - /* int batch_size_out = */ groups, - /* int lda = */ implicit_K * groups, - /* int ldb = */ implicit_K, - /* int ldd = */ implicit_N * groups, - /* bool transpose_a = */ false, - /* bool transpose_b = */ true, - /* std::vector& copies = */ copies, - /* Shape batch_shape = */ {1}, - /* Strides batch_strides = */ {0}, - /* int64_t A_batch_strides = */ int64_t(implicit_K), - /* int64_t B_batch_strides = */ int64_t(implicit_N) * implicit_K, - /* int64_t matrix_stride_out = */ int64_t(implicit_N)); + // 2D view of the output; each tile writes a row window of it. + Strides out_2d_strides{ + out.strides()[out.ndim() - 2], out.strides()[out.ndim() - 1]}; + array out_2d({implicit_M, conv_params.O}, out.dtype(), nullptr, {}); + out_2d.copy_shared_buffer(out, out_2d_strides, out.flags(), out.data_size()); + + // The full unfold buffer can exceed maxBufferLength, so unfold and gemm in + // row tiles reusing one buffer, keeping peak memory at a single tile. + size_t row_bytes = static_cast(implicit_K) * groups * in.itemsize(); + int max_rows = max_unfold_rows(d, row_bytes, implicit_M); + + array in_unfolded({max_rows, implicit_K * groups}, in.dtype(), nullptr, {}); + in_unfolded.set_data(allocator::malloc(in_unfolded.nbytes())); + + for (int row_offset = 0; row_offset < implicit_M; row_offset += max_rows) { + int tile_rows = std::min(max_rows, implicit_M - row_offset); + + // Tile view of the reused unfold buffer. + array in_tile({tile_rows, implicit_K * groups}, in.dtype(), nullptr, {}); + in_tile.copy_shared_buffer( + in_unfolded, + in_unfolded.strides(), + in_unfolded.flags(), + in_tile.size()); + + compute_encoder.set_compute_pipeline_state(kernel); + compute_encoder.set_input_array(in, 0); + compute_encoder.set_output_array(in_tile, 1); + compute_encoder.set_bytes(conv_params, 2); + compute_encoder.set_bytes(row_offset, 3); + + size_t tgp_x = std::min(conv_params.C, 64); + tgp_x = 32 * ((tgp_x + 32 - 1) / 32); + size_t tgp_y = 256 / tgp_x; + + MTL::Size grid_dims = MTL::Size( + conv_params.C, (implicit_K * groups) / conv_params.C, tile_rows); + MTL::Size group_dims = MTL::Size( + std::min(tgp_x, grid_dims.width), std::min(tgp_y, grid_dims.height), 1); + compute_encoder.dispatch_threads(grid_dims, group_dims); + + // Gemm the tile into its output rows. + array out_tile({tile_rows, conv_params.O}, out.dtype(), nullptr, {}); + out_tile.copy_shared_buffer( + out_2d, + out_2d_strides, + out_2d.flags(), + out_tile.size(), + static_cast(row_offset) * out_2d_strides[0]); + + std::vector copies = {in_tile, wt_transpose}; + steel_matmul_regular( + /* const Stream& s = */ s, + /* Device& d = */ d, + /* const array& a = */ in_tile, + /* const array& b = */ wt_transpose, + /* array& c = */ out_tile, + /* int M = */ tile_rows, + /* int N = */ implicit_N, + /* int K = */ implicit_K, + /* int batch_size_out = */ groups, + /* int lda = */ implicit_K * groups, + /* int ldb = */ implicit_K, + /* int ldd = */ implicit_N * groups, + /* bool transpose_a = */ false, + /* bool transpose_b = */ true, + /* std::vector& copies = */ copies, + /* Shape batch_shape = */ {1}, + /* Strides batch_strides = */ {0}, + /* int64_t A_batch_strides = */ int64_t(implicit_K), + /* int64_t B_batch_strides = */ int64_t(implicit_N) * implicit_K, + /* int64_t matrix_stride_out = */ int64_t(implicit_N)); + } } void implicit_gemm_conv_2D_gpu( diff --git a/mlx/backend/metal/kernels/conv.metal b/mlx/backend/metal/kernels/conv.metal index 220afc307d..6f3ce0cd43 100644 --- a/mlx/backend/metal/kernels/conv.metal +++ b/mlx/backend/metal/kernels/conv.metal @@ -20,6 +20,7 @@ template const device T* in [[buffer(0)]], device T* out [[buffer(1)]], const constant MLXConvParams* params [[buffer(2)]], + const constant int& row_offset [[buffer(3)]], uint3 gid [[thread_position_in_grid]]) { int filter_size = params->C; for (short i = 0; i < N; i++) @@ -39,8 +40,9 @@ template // gid.y: wS (Filter location to unfold input) // gid.x: C (channel) - int n = (gid.z) / out_pixels; - int oS = (gid.z) % out_pixels; + int global_row = row_offset + int(gid.z); + int n = global_row / out_pixels; + int oS = global_row % out_pixels; int wS = gid.y; bool valid = n < params->N; @@ -83,6 +85,7 @@ template const device T* in [[buffer(0)]], device T* out [[buffer(1)]], const constant MLXConvParams* params [[buffer(2)]], + const constant int& row_offset [[buffer(3)]], uint3 gid [[thread_position_in_grid]]) { int filter_size = params->C; for (short i = 0; i < N; i++) @@ -103,8 +106,9 @@ template // gid.y: wS (Filter location to unfold input) // gid.x: C (channel) - int n = (gid.z) / out_pixels; - int oS = (gid.z) % out_pixels; + int global_row = row_offset + int(gid.z); + int n = global_row / out_pixels; + int oS = global_row % out_pixels; int wS = gid.y; bool valid = n < params->N; @@ -150,6 +154,7 @@ template const device itype* in [[buffer(0)]], \ device itype* out [[buffer(1)]], \ const constant MLXConvParams* params [[buffer(2)]], \ + const constant int& row_offset [[buffer(3)]], \ uint3 gid [[thread_position_in_grid]]); \ template \ [[host_name("naive_unfold_transpose_nd_" #name "_" #n)]] [[kernel]] void \ @@ -157,6 +162,7 @@ template const device itype* in [[buffer(0)]], \ device itype* out [[buffer(1)]], \ const constant MLXConvParams* params [[buffer(2)]], \ + const constant int& row_offset [[buffer(3)]], \ uint3 gid [[thread_position_in_grid]]); #define instantiate_naive_unfold_nd_dims(name, itype) \ diff --git a/python/tests/test_conv_transpose.py b/python/tests/test_conv_transpose.py index 7289955ed4..e6def081a7 100644 --- a/python/tests/test_conv_transpose.py +++ b/python/tests/test_conv_transpose.py @@ -1,6 +1,7 @@ # Copyright © 2023-2024 Apple Inc. import math +import os import unittest from itertools import permutations @@ -805,6 +806,43 @@ def run_conv_transpose_3d_output_padding( dtype=dtype, ) + @unittest.skipIf(not mx.metal.is_available(), "requires Metal") + def test_conv_transpose_unfold_tiling(self): + # The explicit-GEMM conv path unfolds into one buffer that can exceed + # maxBufferLength for large outputs; it unfolds and runs the gemm in row + # tiles instead (issue #3082). + key = "MLX_CONV_UNFOLD_TILE_ROWS" + prev = os.environ.get(key) + cases = ( + (mx.conv_transpose1d, (2, 9, 4), (5, 3, 4), {"stride": 2}), + (mx.conv_transpose2d, (2, 5, 5, 4), (5, 3, 3, 4), {"stride": 2}), + (mx.conv_transpose3d, (1, 4, 4, 4, 2), (3, 2, 2, 2, 2), {"stride": 2}), + (mx.conv_transpose1d, (2, 9, 4), (6, 3, 2), {"stride": 2, "groups": 2}), + ) + try: + # tile_rows=1 forces uniform tiles; 7 does not divide any output, so + # it exercises a partial final tile too. + for conv, in_shape, wt_shape, kwargs in cases: + for tile_rows in (1, 7): + with self.subTest( + conv=conv.__name__, tile_rows=tile_rows, **kwargs + ): + np.random.seed(0) + x = mx.array(np.random.normal(size=in_shape).astype(np.float32)) + w = mx.array(np.random.normal(size=wt_shape).astype(np.float32)) + os.environ.pop(key, None) + untiled = conv(x, w, **kwargs) + mx.eval(untiled) + os.environ[key] = str(tile_rows) + tiled = conv(x, w, **kwargs) + mx.eval(tiled) + self.assertTrue(np.allclose(untiled, tiled, atol=1e-4)) + finally: + if prev is None: + os.environ.pop(key, None) + else: + os.environ[key] = prev + if __name__ == "__main__": mlx_tests.MLXTestRunner()