Fix copying scalars by adding fill_gpu (#1402)

* fix copying scalars by adding fill_gpu

* Another copy scalar changed to fill

---------

Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>

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),

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

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

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);
     }
   }
 

mlx/backend/metal/matmul.cpp

@@ -526,7 +526,7 @@ void Matmul::eval_gpu(const std::vector<array>& 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<array>& 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<array>& 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;

mlx/backend/metal/slicing.cpp

@@ -85,7 +85,7 @@ void pad_gpu(
     std::vector<int> 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;

tests/ops_tests.cpp

@@ -2463,6 +2463,28 @@ TEST_CASE("test pad") {
   CHECK_EQ(pad(x, 1).shape(), std::vector<int>{3, 4, 5});
   CHECK_EQ(pad(x, {0, 1}).shape(), std::vector<int>{2, 3, 4});
   CHECK_EQ(pad(x, {{1, 1}, {1, 2}, {3, 1}}).shape(), std::vector<int>{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<bool>());
 }
 
 TEST_CASE("test power") {