mlx 1.0.0

data/mlx/mlx/backend/cuda/conv/gemm_grouped_conv.cu

@@ -0,0 +1,231 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/conv/conv.h"
+#include "mlx/backend/cuda/gemms/cublas_gemm.h"
+#include "mlx/backend/cuda/kernel_utils.cuh"
+#include "mlx/dtype_utils.h"
+
+#include <cooperative_groups.h>
+
+namespace mlx::core {
+
+namespace cu {
+
+namespace cg = cooperative_groups;
+
+template <typename T, int NDIM>
+__global__ void naive_grouped_unfold_transpose_nd(
+    const T* in,
+    T* out,
+    int filter_size,
+    int out_pixels,
+    const __grid_constant__ ConvParams<NDIM> params) {
+  auto block = cg::this_thread_block();
+  auto tid = block.group_index();
+  auto lid = block.thread_index();
+
+  int index_batch = tid.z / out_pixels; // [0, N)
+  int index_out_spatial = tid.z % out_pixels; // [0, H_out * W_out)
+  int index_wt_spatial =
+      tid.x * block.dim_threads().x + lid.x; // [0, H_wt * W_wt)
+
+  if (index_wt_spatial >= filter_size / params.C) {
+    return;
+  }
+
+  in += tid.y; // [0, C)
+  out += tid.z * filter_size + tid.y * (filter_size / params.C);
+
+  bool valid = index_batch < params.N;
+
+  // Get the coordinates in input.
+  int index_in[NDIM] = {};
+  int wt_stride = 1;
+#pragma unroll
+  for (int i = NDIM - 1; i >= 0; --i) {
+    int index_out = index_out_spatial % params.out_spatial_dims[i];
+    int index_wt = index_wt_spatial % params.wt_spatial_dims[i];
+    out += index_wt * wt_stride;
+
+    if (params.flip) {
+      index_wt = params.wt_spatial_dims[i] - index_wt - 1;
+    }
+
+    int index = index_out * params.strides[i] - params.padding[i] +
+        index_wt * params.kernel_dilation[i];
+    int index_max =
+        1 + params.input_dilation[i] * (params.in_spatial_dims[i] - 1);
+
+    valid &= (index >= 0) && (index < index_max) &&
+        (index % params.input_dilation[i] == 0);
+
+    index_in[i] = index / params.input_dilation[i];
+
+    index_out_spatial /= params.out_spatial_dims[i];
+    index_wt_spatial /= params.wt_spatial_dims[i];
+    wt_stride *= params.wt_spatial_dims[i];
+  }
+
+  if (valid) {
+    int in_offset = index_batch * params.in_strides[0];
+#pragma unroll
+    for (int i = 0; i < NDIM; ++i) {
+      in_offset += index_in[i] * params.in_strides[i + 1];
+    }
+    *out = in[in_offset];
+  } else {
+    *out = T{0};
+  }
+}
+
+} // namespace cu
+
+template <int NDIM>
+array grouped_unfold_transpose_inputs_nd(
+    cu::CommandEncoder& encoder,
+    const array& in,
+    int mat_M,
+    int mat_K,
+    int mat_N,
+    ConvParams<NDIM>& params) {
+  array unfolded({mat_M, mat_K * params.groups}, in.dtype(), nullptr, {});
+  unfolded.set_data(cu::malloc_async(unfolded.nbytes(), encoder));
+  encoder.add_temporary(unfolded);
+
+  int filter_size = params.C;
+#pragma unroll
+  for (int i = 0; i < NDIM; ++i) {
+    filter_size *= params.wt_spatial_dims[i];
+  }
+
+  int out_pixels = 1;
+#pragma unroll
+  for (int i = 0; i < NDIM; ++i) {
+    out_pixels *= params.out_spatial_dims[i];
+  }
+
+  int wt_spatial_size = (mat_K * params.groups) / params.C;
+  dim3 block_dims;
+  block_dims.x = std::min(std::max(wt_spatial_size, 32), 1024);
+  dim3 num_blocks;
+  num_blocks.x = cuda::ceil_div(wt_spatial_size, block_dims.x);
+  num_blocks.y = params.C;
+  num_blocks.z = mat_M;
+
+  encoder.set_input_array(in);
+  encoder.set_output_array(unfolded);
+  dispatch_float_types(in.dtype(), "unfold", [&](auto type_tag) {
+    using DataType = cuda_type_t<MLX_GET_TYPE(type_tag)>;
+    encoder.add_kernel_node(
+        cu::naive_grouped_unfold_transpose_nd<DataType, NDIM>,
+        num_blocks,
+        block_dims,
+        0,
+        gpu_ptr<DataType>(in),
+        gpu_ptr<DataType>(unfolded),
+        filter_size,
+        out_pixels,
+        params);
+  });
+
+  return unfolded;
+}
+
+template <int NDIM>
+void gemm_grouped_conv_nd(
+    cu::CommandEncoder& encoder,
+    const array& in,
+    const array& wt,
+    array& out,
+    ConvParams<NDIM>& params,
+    Stream s) {
+  // Get gemm shapes.
+  int C_per_group = params.C / params.groups;
+  int O_per_group = params.O / params.groups;
+  int mat_M = out.size() / params.O; // N * H_out * W_out
+  int mat_K = wt.size() / params.O; // C_per_group * H_wt * W_wt
+  int mat_N = O_per_group; // O_per_group
+
+  // Unfold input to (N * H_out * W_out, C * H_wt * W_wt) for gemm.
+  array in_unfolded = grouped_unfold_transpose_inputs_nd<NDIM>(
+      encoder, in, mat_M, mat_K, mat_N, params);
+
+  // Reshape weight to (O, C_per_group, H_wt * W_wt) for gemm.
+  int wt_spatial_size = (wt.size() / wt.shape(0)) / wt.shape(-1);
+  array wt_view(
+      {params.O, C_per_group, wt_spatial_size}, wt.dtype(), nullptr, {});
+  wt_view.copy_shared_buffer(
+      wt, {wt.strides(0), 1, C_per_group}, wt.flags(), wt.size());
+  array wt_reshaped = contiguous_copy_gpu(wt_view, s);
+
+  // Batch with size of groups.
+  Shape batch_shape{params.groups};
+  Strides a_batch_strides{mat_K};
+  Strides b_batch_strides{mat_N * mat_K};
+
+  // Run matmul.
+  CublasGemm gemm(
+      encoder.device(),
+      in.dtype(),
+      false, // a_transposed
+      mat_M, // a_rows
+      mat_K, // a_cols
+      mat_K * params.groups, // lda
+      true, // b_transposed
+      mat_K, // b_rows
+      mat_N, // b_cols
+      mat_K, // ldb
+      batch_shape.back(),
+      a_batch_strides.back(),
+      b_batch_strides.back());
+  gemm.set_out(
+      out.dtype(),
+      false, // out_transposed
+      mat_M, // out_rows
+      mat_N, // out_cols
+      mat_N * params.groups, // out_ld
+      params.groups, // batch_count
+      mat_N); // batch_stride
+  gemm.run(
+      encoder,
+      out,
+      in_unfolded,
+      wt_reshaped,
+      batch_shape,
+      a_batch_strides,
+      b_batch_strides);
+}
+
+void gemm_grouped_conv(
+    cu::CommandEncoder& encoder,
+    const array& in,
+    const array& wt,
+    array& out,
+    const std::vector<int>& strides,
+    const std::vector<int>& padding,
+    const std::vector<int>& kernel_dilation,
+    const std::vector<int>& input_dilation,
+    int groups,
+    bool flip,
+    Stream s) {
+  int conv_ndim = in.ndim() - 2;
+  if (conv_ndim < 1 || conv_ndim > 3) {
+    throw std::runtime_error(
+        fmt::format("[conv] Unsupported gemm_conv for {}D conv.", conv_ndim));
+  }
+  dispatch_1_2_3(conv_ndim, [&](auto ndim_constant) {
+    ConvParams<ndim_constant()> params(
+        in,
+        wt,
+        out,
+        strides,
+        padding,
+        kernel_dilation,
+        input_dilation,
+        groups,
+        flip);
+    gemm_grouped_conv_nd<ndim_constant()>(encoder, in, wt, out, params, s);
+  });
+}
+
+} // namespace mlx::core

data/mlx/mlx/backend/cuda/conv.cpp

@@ -0,0 +1,403 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/cuda/conv/conv.h"
+#include "mlx/backend/cuda/cudnn_utils.h"
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/lru_cache.h"
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/primitives.h"
+
+#include <nvtx3/nvtx3.hpp>
+
+#include <cassert>
+
+namespace mlx::core {
+
+namespace {
+
+enum ConvBackendType {
+  CONV_FALLBACK,
+  CONV_FORWARD,
+  CONV_BACKWARD_INPUT,
+  CONV_BACKWARD_WEIGHT,
+};
+
+struct ConvCacheKey {
+  int device_id;
+  fe::DataType_t cudnn_dtype;
+  std::array<int, MAX_NDIM> input_shape;
+  std::array<int, MAX_NDIM> weight_shape;
+  std::array<int, MAX_NDIM> stride;
+  std::array<int, MAX_NDIM> padding_lo;
+  std::array<int, MAX_NDIM> padding_hi;
+  std::array<int, MAX_NDIM> dilation;
+  int groups;
+  bool flip;
+  uint8_t input_alignment;
+  uint8_t weight_alignment;
+  uint8_t output_alignment;
+};
+
+auto& conv_cache() {
+  static LRUBytesKeyCache<
+      ConvCacheKey,
+      std::pair<ConvBackendType, std::optional<DnnGraph>>>
+      cache("MLX_CUDA_CONV_CACHE_SIZE", /* default_capacity */ 128);
+  return cache;
+}
+
+auto get_conv_settings(
+    ConvBackendType backend_type,
+    array& x,
+    array& w,
+    array& y,
+    const std::vector<int>& kernel_strides,
+    const std::vector<int>& padding_lo_,
+    const std::vector<int>& padding_hi_,
+    const std::vector<int>& kernel_dilation,
+    const std::vector<int>& input_dilation) {
+  auto padding_lo = convert_vector<int64_t>(padding_lo_);
+  auto padding_hi = convert_vector<int64_t>(padding_hi_);
+
+  if (backend_type == CONV_BACKWARD_INPUT) {
+    for (int i = 0; i < padding_lo.size(); ++i) {
+      int wt_size = 1 + kernel_dilation[i] * (w.shape(1 + i) - 1);
+      padding_lo[i] = wt_size - padding_lo[i] - 1;
+      int in_size = 1 + kernel_strides[i] * (y.shape(1 + i) - 1);
+      int out_size = 1 + input_dilation[i] * (x.shape(1 + i) - 1);
+      padding_hi[i] = out_size - in_size + padding_hi[i];
+    }
+    return std::make_tuple(
+        convert_vector<int64_t>(input_dilation),
+        std::move(padding_lo),
+        std::move(padding_hi),
+        convert_vector<int64_t>(kernel_dilation));
+
+  } else if (backend_type == CONV_BACKWARD_WEIGHT) {
+    padding_hi = padding_lo;
+    return std::make_tuple(
+        convert_vector<int64_t>(kernel_dilation),
+        std::move(padding_lo),
+        std::move(padding_hi),
+        convert_vector<int64_t>(kernel_strides));
+
+  } else {
+    return std::make_tuple(
+        convert_vector<int64_t>(kernel_strides),
+        std::move(padding_lo),
+        std::move(padding_hi),
+        convert_vector<int64_t>(kernel_dilation));
+  }
+}
+
+std::optional<DnnGraph> build_conv_graph(
+    cu::CommandEncoder& encoder,
+    ConvBackendType backend_type,
+    Dtype dtype,
+    array& x,
+    array& w,
+    array& y,
+    const std::vector<int64_t>& stride,
+    const std::vector<int64_t>& padding_lo,
+    const std::vector<int64_t>& padding_hi,
+    const std::vector<int64_t>& dilation) {
+  auto compute_dtype =
+      (dtype == float16 || dtype == bfloat16) ? float32 : dtype;
+  DnnGraph graph(encoder.device().get_cudnn_handle(), dtype, compute_dtype);
+  auto x_ = graph.tensor_nchw("X", 'x', x);
+  auto w_ = graph.tensor_nchw("W", 'w', w);
+
+  auto set_options = [&](auto& options) {
+    options.set_compute_data_type(dtype_to_cudnn_type(compute_dtype))
+        .set_convolution_mode(fe::ConvolutionMode_t::CROSS_CORRELATION)
+        .set_stride(stride)
+        .set_pre_padding(padding_lo)
+        .set_post_padding(padding_hi)
+        .set_dilation(dilation);
+  };
+
+  std::shared_ptr<fe::graph::Tensor_attributes> y_;
+  if (backend_type == CONV_FORWARD) {
+    auto options = fe::graph::Conv_fprop_attributes();
+    set_options(options);
+    y_ = graph.conv_fprop(x_, w_, options);
+  } else if (backend_type == CONV_BACKWARD_INPUT) {
+    auto options = fe::graph::Conv_dgrad_attributes();
+    set_options(options);
+    y_ = graph.conv_dgrad(x_, w_, options);
+  } else if (backend_type == CONV_BACKWARD_WEIGHT) {
+    auto options = fe::graph::Conv_wgrad_attributes();
+    set_options(options);
+    y_ = graph.conv_wgrad(w_, x_, options);
+  }
+  graph.tensor_nchw(y_, 'y', y)->set_output(true);
+
+  if (graph.prepare().is_bad()) {
+    return std::nullopt;
+  }
+  graph.deselect_numeric_notes({fe::NumericalNote_t::DOWN_CONVERT_INPUTS});
+  if (dtype == float32 && !env::enable_tf32()) {
+    graph.deselect_numeric_notes({fe::NumericalNote_t::TENSOR_CORE});
+  }
+  CHECK_CUDNN_FE_ERROR(graph.build());
+  return graph;
+}
+
+// Transpose from (C_out, H, W, C_in / groups) to (C_in, H, W, C_out / groups).
+array group_transpose(
+    const array& x,
+    int groups,
+    int group_dim,
+    int axis1,
+    int axis2,
+    Stream s) {
+  if (groups == 1) {
+    return swapaxes_in_eval(x, axis1, axis2);
+  }
+  int ndim = x.ndim();
+  if (group_dim < 0) {
+    group_dim += ndim;
+  }
+  if (axis1 < 0) {
+    axis1 += ndim;
+  }
+  if (axis2 < 0) {
+    axis2 += ndim;
+  }
+  if (group_dim <= axis1) {
+    axis1 += 1;
+  }
+  if (group_dim <= axis2) {
+    axis2 += 1;
+  }
+  auto shape = x.shape();
+  shape.insert(shape.begin() + group_dim, groups);
+  shape[group_dim + 1] = shape[group_dim + 1] / groups;
+  array x_trans = reshape_in_eval(x, std::move(shape), s);
+  x_trans = swapaxes_in_eval(x_trans, axis1, axis2);
+  x_trans = flatten_in_eval(x_trans, group_dim, group_dim + 1, s);
+  return x_trans;
+}
+
+// Do necessary transposes and copies to prepare the inputs and outputs for
+// building the cuDNN conv op. It is safe to be called multiple times in one
+// eval_gpu, with cost of possible redundant copies.
+std::tuple<array, array, array> prepare_args(
+    cu::CommandEncoder& encoder,
+    ConvBackendType backend_type,
+    array in,
+    array wt,
+    array out,
+    int groups,
+    Stream s) {
+  // Transpose the args depending on the backend type.
+  // TODO: Handle groups.
+  if (backend_type == CONV_BACKWARD_INPUT) {
+    wt = group_transpose(wt, groups, 0, 0, -1, s);
+  } else if (backend_type == CONV_BACKWARD_WEIGHT) {
+    in = group_transpose(in, groups, -1, 0, -1, s);
+    wt = swapaxes_in_eval(wt, 0, -1);
+    // Create a contiguous array that shares the data with |out|, but with dim
+    // C_in and C_out swapped.
+    Shape shape(out.shape());
+    std::swap(shape.front(), shape.back());
+    Strides strides(shape.size(), 1);
+    for (int i = shape.size() - 2; i >= 0; --i) {
+      strides[i] = shape[i + 1] * strides[i + 1];
+    }
+    array intermediate(std::move(shape), out.dtype(), nullptr, {});
+    intermediate.copy_shared_buffer(
+        out, std::move(strides), {true, true, false}, out.data_size());
+    out = intermediate;
+  }
+
+  // cuDNN requires contiguous input.
+  if (!in.flags().row_contiguous) {
+    in = contiguous_copy_gpu(in, s);
+    encoder.add_temporary(in);
+  }
+  if (!wt.flags().row_contiguous) {
+    wt = contiguous_copy_gpu(wt, s);
+    encoder.add_temporary(wt);
+  }
+
+  return {std::move(in), std::move(wt), std::move(out)};
+}
+
+// Register inputs and outputs before actually running conv op. Can only be
+// called once per eval_gpu.
+void register_args(
+    cu::CommandEncoder& encoder,
+    ConvBackendType backend_type,
+    array& in,
+    array& wt,
+    array& intermediate_out,
+    array& final_out) {
+  encoder.set_input_array(in);
+  encoder.set_input_array(wt);
+  encoder.set_output_array(final_out);
+
+  if (backend_type == CONV_BACKWARD_WEIGHT) {
+    // Turn |out| into a strided array, which will have C_in and C_out swapped
+    // in vjp and the final |grad_weight| will then be contiguous.
+    Strides strides = intermediate_out.strides();
+    std::swap(strides.front(), strides.back());
+    final_out.copy_shared_buffer(
+        intermediate_out,
+        std::move(strides),
+        {false, false, false},
+        intermediate_out.data_size());
+  }
+}
+
+} // namespace
+
+void Convolution::eval_gpu(const std::vector<array>& inputs, array& out_) {
+  nvtx3::scoped_range r("Convolution::eval_gpu");
+  if (out_.size() == 0) {
+    return;
+  }
+  auto& s = stream();
+  auto& encoder = cu::get_command_encoder(s);
+
+  assert(inputs.size() == 2);
+  array in = inputs[0];
+  array wt = inputs[1];
+  array out = out_;
+  out.set_data(cu::malloc_async(out.nbytes(), encoder));
+  Dtype dtype = out.dtype();
+
+  // Search cache.
+  BytesKey<ConvCacheKey> cache_key;
+  cache_key.pod = {
+      encoder.device().cuda_device(),
+      dtype_to_cudnn_type(dtype),
+      vector_key(in.shape()),
+      vector_key(wt.shape()),
+      vector_key(kernel_strides_),
+      vector_key(padding_lo_),
+      vector_key(padding_hi_),
+      vector_key(kernel_dilation_),
+      groups_,
+      flip_,
+      get_alignment(in),
+      get_alignment(wt),
+      get_alignment(out)};
+  if (auto it = conv_cache().find(cache_key); it != conv_cache().end()) {
+    auto& [backend_type, graph] = it->second;
+    if (graph) {
+      // Run cached graph.
+      std::tie(in, wt, out) =
+          prepare_args(encoder, backend_type, in, wt, out, groups_, s);
+      register_args(encoder, backend_type, in, wt, out, out_);
+      CHECK_CUDNN_FE_ERROR(graph->encode_capturing(
+          encoder,
+          {
+              {'x', gpu_ptr<void>(in)},
+              {'w', gpu_ptr<void>(wt)},
+              {'y', gpu_ptr<void>(out)},
+          }));
+    } else {
+      // Run fallback kernel.
+      gemm_conv(
+          encoder,
+          in,
+          wt,
+          out,
+          kernel_strides_,
+          padding_lo_,
+          kernel_dilation_,
+          input_dilation_,
+          groups_,
+          flip_,
+          s);
+    }
+    return;
+  }
+
+  // There is no reliable way to deduce the proper cuDNN backend for the
+  // convolution, so we make a best guess and then try.
+  SmallVector<ConvBackendType, 2> try_backends;
+  if (flip_) {
+    // When weight is flipped, we assume it is backward input convolution.
+    try_backends.push_back(CONV_BACKWARD_INPUT);
+  } else {
+    // Otherwise it could be backward weight convolution or forward convolution,
+    // mathematically there is no difference so we have to use heuristics.
+    // Empirically backward convolutions have large kernel dimensions, and
+    // usually have |in| and |wt| transposed.
+    if (!in.flags().row_contiguous && !wt.flags().row_contiguous &&
+        wt.shape(2) > out.shape(2)) {
+      try_backends = {CONV_BACKWARD_WEIGHT, CONV_FORWARD};
+    } else {
+      try_backends = {CONV_FORWARD, CONV_BACKWARD_WEIGHT};
+    }
+  }
+
+  // Try to build op graph.
+  ConvBackendType backend_type;
+  std::optional<DnnGraph> graph;
+  for (auto try_backend : try_backends) {
+    auto [x, w, y] =
+        prepare_args(encoder, try_backend, in, wt, out, groups_, s);
+    auto [stride, padding_lo, padding_hi, dilation] = get_conv_settings(
+        try_backend,
+        x,
+        w,
+        y,
+        kernel_strides_,
+        padding_lo_,
+        padding_hi_,
+        kernel_dilation_,
+        input_dilation_);
+    graph = build_conv_graph(
+        encoder,
+        try_backend,
+        dtype,
+        x,
+        w,
+        y,
+        stride,
+        padding_lo,
+        padding_hi,
+        dilation);
+    if (graph) {
+      backend_type = try_backend;
+      in = std::move(x);
+      wt = std::move(w);
+      out = std::move(y);
+      break;
+    }
+  }
+
+  if (graph) {
+    register_args(encoder, backend_type, in, wt, out, out_);
+    CHECK_CUDNN_FE_ERROR(graph->encode_capturing(
+        encoder,
+        {
+            {'x', gpu_ptr<void>(in)},
+            {'w', gpu_ptr<void>(wt)},
+            {'y', gpu_ptr<void>(out)},
+        }));
+    conv_cache().emplace(
+        cache_key, std::make_pair(backend_type, std::move(*graph)));
+    return;
+  }
+
+  // Use fallback kernel for settings not supported by cuDNN.
+  gemm_conv(
+      encoder,
+      in,
+      wt,
+      out,
+      kernel_strides_,
+      padding_lo_,
+      kernel_dilation_,
+      input_dilation_,
+      groups_,
+      flip_,
+      s);
+  conv_cache().emplace(cache_key, std::make_pair(CONV_FALLBACK, std::nullopt));
+}
+
+} // namespace mlx::core

data/mlx/mlx/backend/cuda/copy/copy.cuh

@@ -0,0 +1,55 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/device/cast_op.cuh"
+#include "mlx/backend/cuda/kernel_utils.cuh"
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/dtype_utils.h"
+
+namespace mlx::core {
+
+void copy_contiguous(
+    cu::CommandEncoder& encoder,
+    CopyType ctype,
+    const array& in,
+    array& out,
+    int64_t offset_in,
+    int64_t offset_out);
+
+void copy_general(
+    cu::CommandEncoder& encoder,
+    CopyType ctype,
+    const array& in,
+    array& out,
+    int64_t offset_in,
+    int64_t offset_out,
+    const Shape& shape,
+    const Strides& strides_in,
+    const Strides& strides_out);
+
+void copy_general_dynamic(
+    cu::CommandEncoder& encoder,
+    CopyType ctype,
+    const array& in,
+    array& out,
+    int64_t offset_in,
+    int64_t offset_out,
+    const Shape& shape,
+    const Strides& strides_in,
+    const Strides& strides_out,
+    const array& dynamic_offset_in,
+    const array& dynamic_offset_out);
+
+void copy_general_input(
+    cu::CommandEncoder& encoder,
+    CopyType ctype,
+    const array& in,
+    array& out,
+    int64_t offset_in,
+    int64_t offset_out,
+    const Shape& shape,
+    const Strides& strides_in);
+
+} // namespace mlx::core