diff --git a/mlx/backend/metal/conv.cpp b/mlx/backend/metal/conv.cpp index 5fda6c3e8..2974d010b 100644 --- a/mlx/backend/metal/conv.cpp +++ b/mlx/backend/metal/conv.cpp @@ -552,7 +552,7 @@ void winograd_conv_2D_gpu( // Fill with zeros array zero_arr = array(0, in.dtype()); - copy_gpu(zero_arr, in_padded, CopyType::Scalar, s); + fill_gpu(zero_arr, in_padded, s); copies_w.push_back(zero_arr); // Pick input slice from padded @@ -571,7 +571,6 @@ void winograd_conv_2D_gpu( copies_w.push_back(in_padded_slice); copies_w.push_back(in_padded); - copies_w.push_back(zero_arr); MLXConvParams<2> conv_params_updated{ /* const int N = */ in_padded.shape(0), diff --git a/mlx/backend/metal/copy.cpp b/mlx/backend/metal/copy.cpp index 2ba38c10d..0fc99220d 100644 --- a/mlx/backend/metal/copy.cpp +++ b/mlx/backend/metal/copy.cpp @@ -174,4 +174,31 @@ void copy_gpu_inplace( in, out, in.shape(), istride, ostrides, ioffset, 0, ctype, s); } +void fill_gpu(const array& val, array& out, const Stream& s) { + if (out.size() == 0) { + return; + } + out.set_data(allocator::malloc_or_wait(out.nbytes())); + bool use_2d = out.data_size() > UINT32_MAX; + auto& d = metal::device(s.device); + std::string kernel_name = std::string(use_2d ? "s2" : "s") + "_copy" + + type_to_name(val) + type_to_name(out); + auto kernel = get_copy_kernel(d, kernel_name, val, out); + auto& compute_encoder = d.get_command_encoder(s.index); + compute_encoder->setComputePipelineState(kernel); + + compute_encoder.set_input_array(val, 0); + compute_encoder.set_output_array(out, 1); + + size_t nthreads = out.data_size(); + MTL::Size grid_dims = use_2d ? get_2d_grid_dims(out.shape(), out.strides()) + : MTL::Size(nthreads, 1, 1); + NS::UInteger thread_group_size = kernel->maxTotalThreadsPerThreadgroup(); + if (thread_group_size > nthreads) { + thread_group_size = nthreads; + } + MTL::Size group_dims = MTL::Size(thread_group_size, 1, 1); + compute_encoder.dispatchThreads(grid_dims, group_dims); +} + } // namespace mlx::core diff --git a/mlx/backend/metal/copy.h b/mlx/backend/metal/copy.h index c810868f5..3042c714e 100644 --- a/mlx/backend/metal/copy.h +++ b/mlx/backend/metal/copy.h @@ -37,4 +37,7 @@ void copy_gpu_inplace( CopyType ctype, const Stream& s); +// Fill the output with the scalar val +void fill_gpu(const array& val, array& out, const Stream& s); + } // namespace mlx::core diff --git a/mlx/backend/metal/custom_kernel.cpp b/mlx/backend/metal/custom_kernel.cpp index 92cb55dcd..b63f35c1c 100644 --- a/mlx/backend/metal/custom_kernel.cpp +++ b/mlx/backend/metal/custom_kernel.cpp @@ -17,9 +17,8 @@ void CustomKernel::eval_gpu( for (auto& out : outputs) { out.set_data(allocator::malloc_or_wait(out.nbytes())); if (init_value_) { - array init = array(init_value_.value(), out.dtype()); - copy_gpu(init, out, CopyType::Scalar, s); - copies.push_back(init); + copies.emplace_back(init_value_.value(), out.dtype()); + fill_gpu(copies.back(), out, s); } } diff --git a/mlx/backend/metal/matmul.cpp b/mlx/backend/metal/matmul.cpp index 5edb66197..2cfeb2b26 100644 --- a/mlx/backend/metal/matmul.cpp +++ b/mlx/backend/metal/matmul.cpp @@ -526,7 +526,7 @@ void Matmul::eval_gpu(const std::vector& inputs, array& out) { // Return 0s if either input is empty if (a_pre.size() == 0 || b_pre.size() == 0) { array zero = array(0, a_pre.dtype()); - copy_gpu(zero, out, CopyType::Scalar, s); + fill_gpu(zero, out, s); auto command_buffer = d.get_command_buffer(s.index); command_buffer->addCompletedHandler([zero](MTL::CommandBuffer*) {}); return; @@ -1156,7 +1156,7 @@ void BlockMaskedMM::eval_gpu(const std::vector& inputs, array& out) { // Return 0s if either input is empty if (a_pre.size() == 0 || b_pre.size() == 0) { array zero = array(0, a_pre.dtype()); - copy_gpu(zero, out, CopyType::Scalar, s); + fill_gpu(zero, out, s); auto command_buffer = d.get_command_buffer(s.index); command_buffer->addCompletedHandler([zero](MTL::CommandBuffer*) {}); return; @@ -1565,7 +1565,7 @@ void GatherMM::eval_gpu(const std::vector& inputs, array& out) { // Return 0s if either input is empty if (a_pre.size() == 0 || b_pre.size() == 0) { array zero = array(0, a_pre.dtype()); - copy_gpu(zero, out, CopyType::Scalar, s); + fill_gpu(zero, out, s); auto command_buffer = d.get_command_buffer(s.index); command_buffer->addCompletedHandler([zero](MTL::CommandBuffer*) {}); return; diff --git a/mlx/backend/metal/slicing.cpp b/mlx/backend/metal/slicing.cpp index 05f051438..a95337fd5 100644 --- a/mlx/backend/metal/slicing.cpp +++ b/mlx/backend/metal/slicing.cpp @@ -85,7 +85,7 @@ void pad_gpu( std::vector low_pad_size, const Stream& s) { // Fill output with val - copy_gpu(val, out, CopyType::Scalar, s); + fill_gpu(val, out, s); // Find offset for start of input values size_t data_offset = 0; diff --git a/tests/ops_tests.cpp b/tests/ops_tests.cpp index 0333c04e4..fb6507267 100644 --- a/tests/ops_tests.cpp +++ b/tests/ops_tests.cpp @@ -2463,6 +2463,28 @@ TEST_CASE("test pad") { CHECK_EQ(pad(x, 1).shape(), std::vector{3, 4, 5}); CHECK_EQ(pad(x, {0, 1}).shape(), std::vector{2, 3, 4}); CHECK_EQ(pad(x, {{1, 1}, {1, 2}, {3, 1}}).shape(), std::vector{3, 5, 7}); + + x = array({1.0f, 2.0f, 3.0f, 4.0f}, {2, 2}); + auto padded_x = pad(x, 1); + auto expected = array( + {0.0f, + 0.0f, + 0.0f, + 0.0f, + 0.0f, + 1.0f, + 2.0f, + 0.0f, + 0.0f, + 3.0f, + 4.0f, + 0.0f, + 0.0f, + 0.0f, + 0.0f, + 0.0f}, + {4, 4}); + CHECK(array_equal(padded_x, expected).item()); } TEST_CASE("test power") {