mlx 1.0.0

data/mlx/mlx/backend/metal/sort.cpp

@@ -0,0 +1,368 @@
+// Copyright © 2023-2024 Apple Inc.
+
+#include <algorithm>
+
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/backend/metal/device.h"
+#include "mlx/backend/metal/kernels.h"
+#include "mlx/backend/metal/utils.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+
+namespace {
+
+void single_block_sort(
+    const Stream& s,
+    metal::Device& d,
+    const array& in,
+    array& out,
+    int axis,
+    int bn,
+    int tn,
+    bool argsort) {
+  // Prepare shapes
+  int n_rows = in.size() / in.shape(axis);
+
+  auto in_nc_str = in.strides();
+  in_nc_str.erase(in_nc_str.begin() + axis);
+
+  auto out_nc_str = out.strides();
+  out_nc_str.erase(out_nc_str.begin() + axis);
+
+  auto nc_shape = in.shape();
+  nc_shape.erase(nc_shape.begin() + axis);
+
+  int nc_dim = nc_shape.size();
+
+  int size_sorted_axis = in.shape(axis);
+  int in_stride_sorted_axis = in.strides()[axis];
+  int out_stride_sorted_axis = out.strides()[axis];
+
+  // We can only use the contiguous kernel if the sorted axis
+  // has the largest or smallest stride.
+  // We also need the input to be contiguous
+  bool contiguous = in.flags().contiguous;
+  auto check_strides = [](array x, int sort_stride) {
+    int min_stride = *std::min_element(x.strides().begin(), x.strides().end());
+    int max_stride = *std::max_element(x.strides().begin(), x.strides().end());
+    return sort_stride == min_stride || sort_stride == max_stride;
+  };
+  contiguous &= check_strides(in, in_stride_sorted_axis);
+  contiguous &= check_strides(out, out_stride_sorted_axis);
+
+  // Prepare kernel name
+  std::ostringstream kname;
+  kname << (contiguous ? "c" : "nc");
+  if (argsort) {
+    kname << "arg";
+  }
+
+  kname << "_block_sort_" << type_to_name(in) << "_" << type_to_name(out)
+        << "_bn" << bn << "_tn" << tn;
+  auto kernel = get_sort_kernel(d, kname.str(), in, out, bn, tn);
+
+  // Prepare command encoder
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  compute_encoder.set_compute_pipeline_state(kernel);
+
+  // Set inputs
+  compute_encoder.set_input_array(in, 0);
+  compute_encoder.set_output_array(out, 1);
+  compute_encoder.set_bytes(size_sorted_axis, 2);
+  compute_encoder.set_bytes(in_stride_sorted_axis, 3);
+  compute_encoder.set_bytes(out_stride_sorted_axis, 4);
+
+  if (contiguous) {
+    int in_stride_segment_axis = INT32_MAX;
+    int out_stride_segment_axis = INT32_MAX;
+    for (int i = 0; i < in_nc_str.size(); i++) {
+      if (nc_shape[i] == 1) {
+        continue;
+      }
+      if (in_nc_str[i] > INT32_MAX || out_nc_str[i] > INT32_MAX) {
+        throw std::runtime_error("[Sort::eval_gpu] Stride too large.");
+      }
+      in_stride_segment_axis =
+          std::min(in_stride_segment_axis, static_cast<int>(in_nc_str[i]));
+      out_stride_segment_axis =
+          std::min(out_stride_segment_axis, static_cast<int>(out_nc_str[i]));
+    }
+    compute_encoder.set_bytes(in_stride_segment_axis, 5);
+    compute_encoder.set_bytes(out_stride_segment_axis, 6);
+  } else {
+    compute_encoder.set_bytes(nc_dim, 5);
+    if (nc_shape.empty()) {
+      int shape = 0;
+      int64_t stride = 0;
+      compute_encoder.set_bytes(shape, 6);
+      compute_encoder.set_bytes(stride, 7);
+      compute_encoder.set_bytes(stride, 8);
+    } else {
+      compute_encoder.set_vector_bytes(nc_shape, 6);
+      compute_encoder.set_vector_bytes(in_nc_str, 7);
+      compute_encoder.set_vector_bytes(out_nc_str, 8);
+    }
+  }
+
+  MTL::Size group_dims = MTL::Size(bn, 1, 1);
+  MTL::Size grid_dims = MTL::Size(1, n_rows, 1);
+
+  compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
+}
+
+void multi_block_sort(
+    const Stream& s,
+    metal::Device& d,
+    const array& in,
+    array& out,
+    int axis,
+    int bn,
+    int tn,
+    int n_blocks,
+    bool argsort) {
+  // Prepare shapes
+  int n_rows = in.size() / in.shape(axis);
+
+  auto nc_str = in.strides();
+  nc_str.erase(nc_str.begin() + axis);
+
+  auto nc_shape = in.shape();
+  nc_shape.erase(nc_shape.begin() + axis);
+
+  int nc_dim = nc_shape.size();
+
+  if (nc_dim == 0) {
+    nc_shape = {0};
+    nc_str = {1};
+  }
+
+  int size_sorted_axis = in.shape(axis);
+  int stride_sorted_axis = in.strides()[axis];
+
+  // Make temporary copies
+  array dev_vals_0({n_rows, size_sorted_axis}, in.dtype(), nullptr, {});
+  array dev_vals_1({n_rows, size_sorted_axis}, in.dtype(), nullptr, {});
+
+  array dev_idxs_0({n_rows, size_sorted_axis}, uint32, nullptr, {});
+  array dev_idxs_1({n_rows, size_sorted_axis}, uint32, nullptr, {});
+
+  array block_partitions({n_rows, n_blocks + 1}, uint32, nullptr, {});
+
+  // Do allocations
+  dev_vals_0.set_data(allocator::malloc(dev_vals_0.nbytes()));
+  dev_vals_1.set_data(allocator::malloc(dev_vals_1.nbytes()));
+  dev_idxs_0.set_data(allocator::malloc(dev_idxs_0.nbytes()));
+  dev_idxs_1.set_data(allocator::malloc(dev_idxs_1.nbytes()));
+  block_partitions.set_data(allocator::malloc(block_partitions.nbytes()));
+
+  std::vector<array> copies = {
+      dev_vals_0, dev_vals_1, dev_idxs_0, dev_idxs_1, block_partitions};
+
+  // Prepare command encoder
+  auto& compute_encoder = d.get_command_encoder(s.index);
+
+  // Do blockwise sort
+  {
+    std::ostringstream kname;
+    kname << "sort_mbsort_" << type_to_name(dev_vals_0) << "_"
+          << type_to_name(dev_idxs_0) << "_bn" << bn << "_tn" << tn;
+    auto kernel =
+        get_mb_sort_kernel(d, kname.str(), dev_vals_0, dev_idxs_0, bn, tn);
+    compute_encoder.set_compute_pipeline_state(kernel);
+
+    compute_encoder.set_input_array(in, 0);
+    compute_encoder.set_output_array(dev_vals_0, 1);
+    compute_encoder.set_output_array(dev_idxs_0, 2);
+    compute_encoder.set_bytes(size_sorted_axis, 3);
+    compute_encoder.set_bytes(stride_sorted_axis, 4);
+    compute_encoder.set_bytes(nc_dim, 5);
+    compute_encoder.set_vector_bytes(nc_shape, 6);
+    compute_encoder.set_vector_bytes(nc_str, 7);
+
+    MTL::Size group_dims = MTL::Size(bn, 1, 1);
+    MTL::Size grid_dims = MTL::Size(n_blocks, n_rows, 1);
+
+    compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
+  }
+
+  // Do merges
+  bool ping = false;
+  array dev_vals_in = dev_vals_0;
+  array dev_idxs_in = dev_idxs_0;
+  array dev_vals_out = dev_vals_1;
+  array dev_idxs_out = dev_idxs_1;
+
+  int n_thr_per_group = (n_blocks + 1) < 1024 ? (n_blocks + 1) : 1024;
+
+  for (int merge_tiles = 2; (merge_tiles / 2) < n_blocks; merge_tiles *= 2) {
+    dev_vals_in = ping ? dev_vals_1 : dev_vals_0;
+    dev_idxs_in = ping ? dev_idxs_1 : dev_idxs_0;
+    dev_vals_out = ping ? dev_vals_0 : dev_vals_1;
+    dev_idxs_out = ping ? dev_idxs_0 : dev_idxs_1;
+    ping = !ping;
+
+    // Do partition
+    {
+      std::ostringstream kname;
+      kname << "partition_mbsort_" << type_to_name(dev_vals_in) << "_"
+            << type_to_name(dev_idxs_in) << "_bn" << bn << "_tn" << tn;
+
+      auto kernel =
+          get_mb_sort_kernel(d, kname.str(), dev_vals_0, dev_idxs_0, bn, tn);
+      compute_encoder.set_compute_pipeline_state(kernel);
+
+      compute_encoder.set_output_array(block_partitions, 0);
+      compute_encoder.set_input_array(dev_vals_in, 1);
+      compute_encoder.set_input_array(dev_idxs_in, 2);
+      compute_encoder.set_bytes(size_sorted_axis, 3);
+      compute_encoder.set_bytes(merge_tiles, 4);
+      compute_encoder.set_bytes(n_blocks, 5);
+
+      MTL::Size group_dims = MTL::Size(n_thr_per_group, 1, 1);
+      MTL::Size grid_dims = MTL::Size(1, n_rows, 1);
+
+      compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
+    }
+
+    // Do merge
+    {
+      std::ostringstream kname;
+      kname << "merge_mbsort_" << type_to_name(dev_vals_in) << "_"
+            << type_to_name(dev_idxs_in) << "_bn" << bn << "_tn" << tn;
+
+      auto kernel =
+          get_mb_sort_kernel(d, kname.str(), dev_vals_0, dev_idxs_0, bn, tn);
+      compute_encoder.set_compute_pipeline_state(kernel);
+
+      compute_encoder.set_input_array(block_partitions, 0);
+      compute_encoder.set_input_array(dev_vals_in, 1);
+      compute_encoder.set_input_array(dev_idxs_in, 2);
+      compute_encoder.set_output_array(dev_vals_out, 3);
+      compute_encoder.set_output_array(dev_idxs_out, 4);
+      compute_encoder.set_bytes(size_sorted_axis, 5);
+      compute_encoder.set_bytes(merge_tiles, 6);
+      compute_encoder.set_bytes(n_blocks, 7);
+
+      MTL::Size group_dims = MTL::Size(bn, 1, 1);
+      MTL::Size grid_dims = MTL::Size(n_blocks, n_rows, 1);
+
+      compute_encoder.dispatch_threadgroups(grid_dims, group_dims);
+    }
+  }
+
+  // Copy outputs with appropriate strides
+  auto strides = out.strides();
+  for (int ax = axis + 1; ax < strides.size(); ax++) {
+    strides[ax] *= out.shape(axis);
+  }
+  strides[axis] = 1;
+  copy_gpu_inplace(
+      (argsort) ? dev_idxs_out : dev_vals_out,
+      out,
+      out.shape(),
+      strides,
+      out.strides(),
+      0,
+      0,
+      (axis == in.ndim() - 1) ? CopyType::Vector : CopyType::General,
+      s);
+
+  d.add_temporaries(std::move(copies), s.index);
+}
+
+void gpu_merge_sort(
+    const Stream& s,
+    metal::Device& d,
+    const array& in,
+    array& out,
+    int axis_,
+    bool argsort) {
+  // Get size info
+  int axis = axis_ < 0 ? axis_ + in.ndim() : axis_;
+  int size_sorted_axis = in.shape(axis);
+
+  // Get kernel size
+  int tn = 4;
+  int potential_bn = (size_sorted_axis + tn - 1) / tn;
+
+  int bn;
+  if (potential_bn > 256) {
+    bn = 512;
+  } else if (potential_bn > 128) {
+    bn = 256;
+  } else if (potential_bn > 64) {
+    bn = 128;
+  } else if (potential_bn > 32) {
+    bn = 64;
+  } else {
+    bn = 32;
+  }
+
+  if (bn == 512 && size_of(in.dtype()) > 4) {
+    bn = 256;
+  }
+
+  int n_per_block = bn * tn;
+  int n_blocks = (size_sorted_axis + n_per_block - 1) / n_per_block;
+
+  if (n_blocks > 1) {
+    return multi_block_sort(s, d, in, out, axis, bn, tn, n_blocks, argsort);
+  } else {
+    return single_block_sort(s, d, in, out, axis, bn, tn, argsort);
+  }
+}
+
+} // namespace
+
+void ArgSort::eval_gpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+
+  out.set_data(allocator::malloc(out.nbytes()));
+
+  auto& s = stream();
+  auto& d = metal::device(s.device);
+  auto& in = inputs[0];
+
+  gpu_merge_sort(s, d, in, out, axis_, true);
+}
+
+void Sort::eval_gpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+
+  out.set_data(allocator::malloc(out.nbytes()));
+
+  auto& s = stream();
+  auto& d = metal::device(s.device);
+  auto& in = inputs[0];
+
+  gpu_merge_sort(s, d, in, out, axis_, false);
+}
+
+void ArgPartition::eval_gpu(const std::vector<array>& inputs, array& out) {
+  // We direct arg partition to sort for now
+  assert(inputs.size() == 1);
+
+  out.set_data(allocator::malloc(out.nbytes()));
+
+  auto& s = stream();
+  auto& d = metal::device(s.device);
+  auto& in = inputs[0];
+
+  gpu_merge_sort(s, d, in, out, axis_, true);
+}
+
+void Partition::eval_gpu(const std::vector<array>& inputs, array& out) {
+  // We direct partition to sort for now
+  assert(inputs.size() == 1);
+
+  out.set_data(allocator::malloc(out.nbytes()));
+
+  auto& s = stream();
+  auto& d = metal::device(s.device);
+  auto& in = inputs[0];
+
+  gpu_merge_sort(s, d, in, out, axis_, false);
+}
+
+} // namespace mlx::core

data/mlx/mlx/backend/metal/ternary.cpp

@@ -0,0 +1,160 @@
+// Copyright © 2024 Apple Inc.
+
+#include "mlx/backend/common/ternary.h"
+#include "mlx/backend/metal/device.h"
+#include "mlx/backend/metal/kernels.h"
+#include "mlx/backend/metal/utils.h"
+#include "mlx/primitives.h"
+
+namespace mlx::core {
+
+void ternary_op_gpu_inplace(
+    const std::vector<array>& inputs,
+    array& out,
+    const char* op,
+    const Stream& s) {
+  assert(inputs.size() == 3);
+  auto& a = inputs[0];
+  auto& b = inputs[1];
+  auto& c = inputs[2];
+  TernaryOpType topt = get_ternary_op_type(a, b, c);
+
+  if (out.size() == 0) {
+    return;
+  }
+
+  // Try to collapse contiguous dims
+  auto maybe_collapse = [topt, &a, &b, &c, &out]() {
+    if (topt == TernaryOpType::General) {
+      auto [shape, strides] = collapse_contiguous_dims(a, b, c, out);
+      return std::make_tuple(
+          shape, strides[0], strides[1], strides[2], strides[3]);
+    } else {
+      Strides e;
+      return std::make_tuple(Shape{}, e, e, e, e);
+    }
+  };
+  auto [shape, strides_a, strides_b, strides_c, strides_out] = maybe_collapse();
+
+  bool large;
+  auto ndim = shape.size();
+  int work_per_thread;
+  if (topt == TernaryOpType::General) {
+    large = a.data_size() > INT32_MAX || b.data_size() > INT32_MAX ||
+        c.data_size() > INT32_MAX || out.size() > INT32_MAX;
+    work_per_thread = large ? 4 : 2;
+  } else {
+    large = out.data_size() > INT32_MAX;
+    work_per_thread = get_work_per_thread(b.dtype(), out.data_size());
+  }
+  std::string kernel_name;
+  if (topt == TernaryOpType::General) {
+    kernel_name = "g";
+    if (shape.size() <= 3) {
+      kernel_name += std::to_string(shape.size());
+    } else if (work_per_thread > 1) {
+      concatenate(kernel_name, "n", std::to_string(work_per_thread));
+    }
+    if (large) {
+      kernel_name += "large";
+    }
+  } else {
+    if (topt == TernaryOpType::VectorScalarVector) {
+      kernel_name = "sv";
+    } else if (topt == TernaryOpType::VectorVectorScalar) {
+      kernel_name = "vs";
+    } else {
+      kernel_name = "v";
+    }
+    if (large) {
+      kernel_name += "2";
+    } else if (work_per_thread > 1) {
+      kernel_name += "n";
+    }
+  }
+  concatenate(kernel_name, "_", op, type_to_name(b));
+
+  auto& d = metal::device(s.device);
+
+  auto kernel = get_ternary_kernel(d, kernel_name, out.dtype(), op);
+
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  compute_encoder.set_compute_pipeline_state(kernel);
+  compute_encoder.set_input_array(a, 0);
+  compute_encoder.set_input_array(b, 1);
+  compute_encoder.set_input_array(c, 2);
+  compute_encoder.set_output_array(out, 3);
+
+  auto thread_group_size = kernel->maxTotalThreadsPerThreadgroup();
+  if (topt == TernaryOpType::General) {
+    // Launch up to 3D grid of threads
+    size_t dim0 = ndim > 0 ? shape[ndim - 1] : 1;
+    size_t dim1 = ndim > 1 ? shape[ndim - 2] : 1;
+    size_t rest = out.size() / (dim0 * dim1);
+
+    if (ndim > 3) {
+      compute_encoder.set_vector_bytes(shape, 4);
+      compute_encoder.set_vector_bytes(strides_a, 5);
+      compute_encoder.set_vector_bytes(strides_b, 6);
+      compute_encoder.set_vector_bytes(strides_c, 7);
+
+      compute_encoder.set_bytes(ndim, 8);
+      dim0 = (dim0 + work_per_thread - 1) / work_per_thread;
+    } else {
+      // The shape is implicit in the grid for <= 3D
+      compute_encoder.set_vector_bytes(strides_a, 4);
+      compute_encoder.set_vector_bytes(strides_b, 5);
+      compute_encoder.set_vector_bytes(strides_c, 6);
+    }
+
+    if (thread_group_size != 1024) {
+      throw std::runtime_error("[Metal::ternary] Must use 1024 sized block");
+    }
+    MTL::Size group_dims = get_block_dims(dim0, dim1, rest);
+    MTL::Size grid_dims = MTL::Size(dim0, dim1, rest);
+    compute_encoder.dispatch_threads(grid_dims, group_dims);
+  } else {
+    // Launch a 1D or 2D grid of threads
+    size_t nthreads = ceildiv(out.data_size(), work_per_thread);
+    if (thread_group_size > nthreads) {
+      thread_group_size = nthreads;
+    }
+    MTL::Size group_dims = MTL::Size(thread_group_size, 1, 1);
+    MTL::Size grid_dims;
+    if (large) {
+      compute_encoder.set_bytes<int64_t>(out.data_size(), 4);
+      grid_dims = get_2d_grid_dims(out.shape(), out.strides(), work_per_thread);
+    } else {
+      compute_encoder.set_bytes<int>(out.data_size(), 4);
+      grid_dims = MTL::Size(nthreads, 1, 1);
+    }
+    compute_encoder.dispatch_threads(grid_dims, group_dims);
+  }
+}
+
+void ternary_op_gpu(
+    const std::vector<array>& inputs,
+    array& out,
+    const char* op,
+    const Stream& s) {
+  auto& a = inputs[0];
+  auto& b = inputs[1];
+  auto& c = inputs[2];
+  TernaryOpType topt = get_ternary_op_type(a, b, c);
+  set_ternary_op_output_data(a, b, c, out, topt);
+  ternary_op_gpu_inplace(inputs, out, op, s);
+}
+
+void ternary_op_gpu(
+    const std::vector<array>& inputs,
+    array& out,
+    const char* op) {
+  auto& s = out.primitive().stream();
+  ternary_op_gpu(inputs, out, op, s);
+}
+
+void Select::eval_gpu(const std::vector<array>& inputs, array& out) {
+  ternary_op_gpu(inputs, out, name());
+}
+
+} // namespace mlx::core

data/mlx/mlx/backend/metal/ternary.h

@@ -0,0 +1,21 @@
+// Copyright © 2024 Apple Inc.
+
+#pragma once
+
+#include "mlx/array.h"
+
+namespace mlx::core {
+
+void ternary_op_gpu(
+    const std::vector<array>& inputs,
+    array& out,
+    const char* op,
+    const Stream& s);
+
+void ternary_op_gpu_inplace(
+    const std::vector<array>& inputs,
+    array& out,
+    const char* op,
+    const Stream& s);
+
+} // namespace mlx::core

data/mlx/mlx/backend/metal/unary.cpp

@@ -0,0 +1,161 @@
+// Copyright © 2024 Apple Inc.
+
+#include "mlx/backend/common/unary.h"
+#include "mlx/backend/metal/device.h"
+#include "mlx/backend/metal/kernels.h"
+#include "mlx/backend/metal/utils.h"
+#include "mlx/primitives.h"
+
+#define UNARY_GPU(func)                                               \
+  void func::eval_gpu(const std::vector<array>& inputs, array& out) { \
+    unary_op_gpu(inputs, out, name());                                \
+  }
+
+namespace mlx::core {
+
+void unary_op_gpu_inplace(
+    const std::vector<array>& inputs,
+    array& out,
+    const char* op,
+    const Stream& s) {
+  auto& in = inputs[0];
+  bool contig = in.flags().contiguous;
+  if (in.size() == 0) {
+    return;
+  }
+
+  auto& d = metal::device(s.device);
+
+  auto maybe_collapse = [contig, &in]() {
+    if (!contig) {
+      return collapse_contiguous_dims(in);
+    } else {
+      return std::make_pair(Shape{}, Strides{});
+    }
+  };
+  auto [shape, strides] = maybe_collapse();
+  int ndim = shape.size();
+  bool large;
+  if (!contig) {
+    large = in.data_size() > INT32_MAX || out.size() > INT32_MAX;
+  } else {
+    large = in.data_size() > UINT32_MAX;
+  }
+  int work_per_thread;
+  std::string kernel_name;
+  if (contig) {
+    work_per_thread = get_work_per_thread(in.dtype(), in.data_size());
+    kernel_name = (large ? "v2" : (work_per_thread > 1 ? "vn" : "v"));
+  } else {
+    work_per_thread = large ? 4 : 1;
+    kernel_name = "gn" + std::to_string(work_per_thread);
+    if (large) {
+      kernel_name += "large";
+    }
+  }
+  concatenate(kernel_name, "_", op, type_to_name(in), type_to_name(out));
+  auto kernel = get_unary_kernel(d, kernel_name, in.dtype(), out.dtype(), op);
+
+  auto thread_group_size = kernel->maxTotalThreadsPerThreadgroup();
+  auto& compute_encoder = d.get_command_encoder(s.index);
+  compute_encoder.set_compute_pipeline_state(kernel);
+  compute_encoder.set_input_array(in, 0);
+  compute_encoder.set_output_array(out, 1);
+  if (!contig) {
+    // Launch up to 3D grid of threads
+    size_t dim0 = ndim > 0 ? shape[ndim - 1] : 1;
+    size_t dim1 = ndim > 1 ? shape[ndim - 2] : 1;
+    size_t rest = out.size() / (dim0 * dim1);
+    compute_encoder.set_vector_bytes(shape, 2);
+    compute_encoder.set_vector_bytes(strides, 3);
+    compute_encoder.set_bytes(ndim, 4);
+    if (thread_group_size != 1024) {
+      throw std::runtime_error("[Metal::unary] Must use 1024 sized block");
+    }
+    dim0 = (dim0 + work_per_thread - 1) / work_per_thread;
+    auto group_dims = get_block_dims(dim0, dim1, rest);
+    MTL::Size grid_dims = MTL::Size(dim0, dim1, rest);
+    compute_encoder.dispatch_threads(grid_dims, group_dims);
+  } else {
+    size_t nthreads = ceildiv(in.data_size(), work_per_thread);
+    if (thread_group_size > nthreads) {
+      thread_group_size = nthreads;
+    }
+
+    MTL::Size group_dims = MTL::Size(thread_group_size, 1, 1);
+    MTL::Size grid_dims;
+    if (large) {
+      compute_encoder.set_bytes<int64_t>(in.data_size(), 2);
+      grid_dims = get_2d_grid_dims(out.shape(), out.strides(), work_per_thread);
+    } else {
+      compute_encoder.set_bytes<int>(in.data_size(), 2);
+      grid_dims = MTL::Size(nthreads, 1, 1);
+    }
+    compute_encoder.dispatch_threads(grid_dims, group_dims);
+  }
+}
+
+void unary_op_gpu(
+    const std::vector<array>& inputs,
+    array& out,
+    const char* op,
+    const Stream& s) {
+  set_unary_output_data(inputs[0], out);
+  unary_op_gpu_inplace(inputs, out, op, s);
+}
+
+void unary_op_gpu(
+    const std::vector<array>& inputs,
+    array& out,
+    const char* op) {
+  auto& s = out.primitive().stream();
+  unary_op_gpu(inputs, out, op, s);
+}
+
+UNARY_GPU(Abs)
+UNARY_GPU(ArcCos)
+UNARY_GPU(ArcCosh)
+UNARY_GPU(ArcSin)
+UNARY_GPU(ArcSinh)
+UNARY_GPU(ArcTan)
+UNARY_GPU(ArcTanh)
+UNARY_GPU(BitwiseInvert)
+UNARY_GPU(Conjugate)
+UNARY_GPU(Cos)
+UNARY_GPU(Cosh)
+UNARY_GPU(Erf)
+UNARY_GPU(ErfInv)
+UNARY_GPU(Exp)
+UNARY_GPU(Expm1)
+UNARY_GPU(Imag)
+UNARY_GPU(Log1p)
+UNARY_GPU(LogicalNot)
+UNARY_GPU(Floor)
+UNARY_GPU(Ceil)
+UNARY_GPU(Negative)
+UNARY_GPU(Real)
+UNARY_GPU(Sigmoid)
+UNARY_GPU(Sign)
+UNARY_GPU(Sin)
+UNARY_GPU(Sinh)
+UNARY_GPU(Square)
+UNARY_GPU(Sqrt)
+UNARY_GPU(Tan)
+UNARY_GPU(Tanh)
+
+void Log::eval_gpu(const std::vector<array>& inputs, array& out) {
+  unary_op_gpu(inputs, out, name());
+}
+
+void Round::eval_gpu(const std::vector<array>& inputs, array& out) {
+  assert(inputs.size() == 1);
+  const auto& in = inputs[0];
+  if (issubdtype(in.dtype(), inexact)) {
+    unary_op_gpu(inputs, out, name());
+  } else {
+    // No-op integer types
+    out.copy_shared_buffer(in);
+  }
+}
+
+} // namespace mlx::core