From 62e8b8d3ac2ce2f5593c5f622736de9bb250ce5d Mon Sep 17 00:00:00 2001 From: JasonHonKL Date: Sat, 5 Sep 2026 13:57:14 +0800 Subject: [PATCH 1/2] [metal] Refactor ensure innermost contiguous function --- mlx/backend/metal/normalization.cpp | 52 ++--------------------------- mlx/backend/metal/softmax.cpp | 21 +----------- mlx/backend/metal/utils.h | 31 +++++++++++++++++ 3 files changed, 34 insertions(+), 70 deletions(-) diff --git a/mlx/backend/metal/normalization.cpp b/mlx/backend/metal/normalization.cpp index f3370f087d..4fc2f10b58 100644 --- a/mlx/backend/metal/normalization.cpp +++ b/mlx/backend/metal/normalization.cpp @@ -22,31 +22,7 @@ void RMSNorm::eval_gpu( auto& out = outputs[0]; // Make sure that the last dimension is contiguous - auto set_output = [&s, &out](const array& x) { - bool no_copy = x.flags().contiguous && x.strides()[x.ndim() - 1] == 1; - if (no_copy && x.ndim() > 1) { - auto s = x.strides()[x.ndim() - 2]; - no_copy &= (s == 0 || s == x.shape().back() || x.shape(-2) == 1); - } - if (no_copy) { - if (x.is_donatable()) { - out.copy_shared_buffer(x); - } else { - out.set_data( - allocator::malloc(x.data_size() * x.itemsize()), - x.data_size(), - x.strides(), - x.flags()); - } - return x; - } else { - array x_copy = contiguous_copy_gpu(x, s); - out.copy_shared_buffer(x_copy); - return x_copy; - } - }; - - const array x = set_output(inputs[0]); + const array x = ensure_innermost_contiguous(out, inputs[0], s, true); const array& w = inputs[1]; auto axis_size = static_cast(x.shape().back()); @@ -224,31 +200,7 @@ void LayerNorm::eval_gpu( auto& out = outputs[0]; // Make sure that the last dimension is contiguous - auto set_output = [&s, &out](const array& x) { - bool no_copy = x.flags().contiguous && x.strides()[x.ndim() - 1] == 1; - if (no_copy && x.ndim() > 1) { - auto s = x.strides()[x.ndim() - 2]; - no_copy &= (s == 0 || s == x.shape().back() || x.shape(-2) == 1); - } - if (no_copy) { - if (x.is_donatable()) { - out.copy_shared_buffer(x); - } else { - out.set_data( - allocator::malloc(x.data_size() * x.itemsize()), - x.data_size(), - x.strides(), - x.flags()); - } - return x; - } else { - array x_copy = contiguous_copy_gpu(x, s); - out.copy_shared_buffer(x_copy); - return x_copy; - } - }; - - const array x = set_output(inputs[0]); + const array x = ensure_innermost_contiguous(out, inputs[0], s, true); const array& w = inputs[1]; const array& b = inputs[2]; diff --git a/mlx/backend/metal/softmax.cpp b/mlx/backend/metal/softmax.cpp index f455cf2d1f..70642a9397 100644 --- a/mlx/backend/metal/softmax.cpp +++ b/mlx/backend/metal/softmax.cpp @@ -22,26 +22,7 @@ void Softmax::eval_gpu(const std::vector& inputs, array& out) { auto& d = metal::device(s.device); // Make sure that the last dimension is contiguous - auto set_output = [&s, &out](const array& x) { - if (x.flags().contiguous && x.strides()[x.ndim() - 1] == 1) { - if (x.is_donatable()) { - out.copy_shared_buffer(x); - } else { - out.set_data( - allocator::malloc(x.data_size() * x.itemsize()), - x.data_size(), - x.strides(), - x.flags()); - } - return x; - } else { - array x_copy = contiguous_copy_gpu(x, s); - out.copy_shared_buffer(x_copy); - return x_copy; - } - }; - - const array in = set_output(inputs[0]); + const array in = ensure_innermost_contiguous(out, inputs[0], s); int axis_size = in.shape().back(); int n_rows = in.data_size() / axis_size; diff --git a/mlx/backend/metal/utils.h b/mlx/backend/metal/utils.h index c4cef8cb6c..09ec3815bd 100644 --- a/mlx/backend/metal/utils.h +++ b/mlx/backend/metal/utils.h @@ -5,6 +5,7 @@ #include #include "mlx/array.h" +#include "mlx/backend/gpu/copy.h" #include "mlx/backend/metal/device.h" #include "mlx/primitives.h" @@ -81,6 +82,36 @@ inline size_t ceildiv(size_t n, size_t m) { return (n + m - 1) / m; } +// Ensure the last dimension of x is contiguous, donating or copying the +// buffer into out. With relax_prev_stride, a second-to-last stride of 0, +// D, or a size-1 axis is also accepted. +inline array ensure_innermost_contiguous( + array& out, + const array& x, + const Stream& s, + bool relax_prev_stride = false) { + bool no_copy = x.flags().contiguous && x.strides()[x.ndim() - 1] == 1; + if (relax_prev_stride && no_copy && x.ndim() > 1) { + auto stride = x.strides()[x.ndim() - 2]; + no_copy &= (stride == 0 || stride == x.shape().back() || x.shape(-2) == 1); + } + if (no_copy) { + if (x.is_donatable()) { + out.copy_shared_buffer(x); + } else { + out.set_data( + allocator::malloc(x.data_size() * x.itemsize()), + x.data_size(), + x.strides(), + x.flags()); + } + return x; + } + array x_copy = contiguous_copy_gpu(x, s); + out.copy_shared_buffer(x_copy); + return x_copy; +} + inline void check_kernel_threadgroup_size( const MTL::ComputePipelineState* kernel, MTL::Size group_dims, From 70c0a1bfd12d92c342ea1fdd41130bca86477dd1 Mon Sep 17 00:00:00 2001 From: JasonHonKL Date: Sat, 5 Sep 2026 14:07:21 +0800 Subject: [PATCH 2/2] remove unnecessary comment --- mlx/backend/metal/utils.h | 3 --- 1 file changed, 3 deletions(-) diff --git a/mlx/backend/metal/utils.h b/mlx/backend/metal/utils.h index 09ec3815bd..348ca67afe 100644 --- a/mlx/backend/metal/utils.h +++ b/mlx/backend/metal/utils.h @@ -82,9 +82,6 @@ inline size_t ceildiv(size_t n, size_t m) { return (n + m - 1) / m; } -// Ensure the last dimension of x is contiguous, donating or copying the -// buffer into out. With relax_prev_stride, a second-to-last stride of 0, -// D, or a size-1 axis is also accepted. inline array ensure_innermost_contiguous( array& out, const array& x,