zhangyiss/mlx
1
0
Fork 0
mirror of https://github.com/ml-explore/mlx.git synced 2025-06-24 17:31:16 +08:00
Code Issues Actions Packages Projects Releases Wiki Activity

Fix copying scalars by adding fill_gpu (#1402)

Browse Source 
* fix copying scalars by adding fill_gpu

* Another copy scalar changed to fill

---------

Co-authored-by: Angelos Katharopoulos <a_katharopoulos@apple.com>
This commit is contained in:
Awni Hannun 2024-09-09 15:54:08 -07:00 committed by GitHub
parent 3ae6aabe9f
commit e7e59c6f05
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
7 changed files with 59 additions and 9 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

3
mlx/backend/metal/conv.cpp
View File

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

27
mlx/backend/metal/copy.cpp
View File

@ -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

3
mlx/backend/metal/copy.h
View File

@ -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

5
mlx/backend/metal/custom_kernel.cpp
View File

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

6
mlx/backend/metal/matmul.cpp
View File

@ -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;

2
mlx/backend/metal/slicing.cpp
View File

@ -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;

22
tests/ops_tests.cpp
View File

@ -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") {