[CUDA] implement Hadamard transform (#3179)

mlx/backend/cuda/CMakeLists.txt

@@ -30,6 +30,7 @@ target_sources(
           ${CMAKE_CURRENT_SOURCE_DIR}/gemms/gemv.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/gemms/cublas_gemm.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/gemms/grouped_gemm_unaligned.cu
+          ${CMAKE_CURRENT_SOURCE_DIR}/hadamard.cu
           ${CMAKE_CURRENT_SOURCE_DIR}/jit_module.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/indexing.cpp
           ${CMAKE_CURRENT_SOURCE_DIR}/kernel_utils.cu

mlx/backend/cuda/device/hadamard.cuh

@@ -0,0 +1,184 @@
+// Copyright © 2025 Apple Inc.
+
+#pragma once
+
+#include "mlx/backend/cuda/device/utils.cuh"
+
+namespace mlx::core::cu {
+
+__device__ __forceinline__ void hadamard_radix_m(float* x);
+
+template <int N>
+struct Pow2Log2 {
+  static_assert(
+      (N > 0) && ((N & (N - 1)) == 0),
+      "N must be a positive power of two.");
+  static constexpr int value = 1 + Pow2Log2<N / 2>::value;
+};
+
+template <>
+struct Pow2Log2<1> {
+  static constexpr int value = 0;
+};
+
+template <int R>
+__device__ __forceinline__ void hadamard_radix_pow2(float* x) {
+  constexpr int kLogR = Pow2Log2<R>::value;
+  int h = 1;
+#pragma unroll
+  for (int s = 0; s < kLogR; ++s) {
+#pragma unroll
+    for (int i = 0; i < R / 2; ++i) {
+      int k = i & (h - 1);
+      int j = ((i - k) << 1) + k;
+      float a = x[j];
+      float b = x[j + h];
+      x[j] = a + b;
+      x[j + h] = a - b;
+    }
+    h <<= 1;
+  }
+}
+
+template <typename T, int N, int max_radix, int read_width, int stride = 1>
+__global__ void
+hadamard_n(const T* in, T* out, float scale, long long num_transforms) {
+  constexpr int kNumThreads = N / max_radix;
+  constexpr int kLogN = Pow2Log2<N>::value;
+  constexpr int kLogR = Pow2Log2<max_radix>::value;
+  constexpr int kNumSteps = kLogN / kLogR;
+  constexpr int kLogFinal = kLogN % kLogR;
+  constexpr int kFinalRadix = 1 << kLogFinal;
+
+  if (threadIdx.x >= kNumThreads) {
+    return;
+  }
+
+  __shared__ T buf[N];
+  int i = threadIdx.x;
+
+  for (long long transform = blockIdx.x; transform < num_transforms;
+       transform += gridDim.x) {
+    long long base = (transform / stride) * static_cast<long long>(N) * stride +
+        (transform % stride);
+
+    if constexpr (stride == 1) {
+#pragma unroll
+      for (int j = 0; j < max_radix / read_width; ++j) {
+        int index = j * read_width * kNumThreads + i * read_width;
+#pragma unroll
+        for (int r = 0; r < read_width; ++r) {
+          buf[index + r] = in[base + index + r];
+        }
+      }
+    } else {
+#pragma unroll
+      for (int j = 0; j < max_radix; ++j) {
+        buf[j * kNumThreads + i] = in[base + (j * kNumThreads + i) * stride];
+      }
+    }
+    __syncthreads();
+
+    float x[max_radix];
+    int h = 1;
+
+#pragma unroll
+    for (int s = 0; s < kNumSteps; ++s) {
+      int k = i & (h - 1);
+      int j = ((i - k) << kLogR) + k;
+
+#pragma unroll
+      for (int r = 0; r < max_radix; ++r) {
+        x[r] = static_cast<float>(buf[j + h * r]);
+      }
+
+      hadamard_radix_pow2<max_radix>(x);
+
+#pragma unroll
+      for (int r = 0; r < max_radix; ++r) {
+        buf[j + h * r] = static_cast<T>(x[r]);
+      }
+
+      h <<= kLogR;
+      __syncthreads();
+    }
+
+    if constexpr (kFinalRadix > 1) {
+#pragma unroll
+      for (int t = 0; t < max_radix / kFinalRadix; ++t) {
+        int index = i + t * kNumThreads;
+        int k = index & (h - 1);
+        int j = ((index - k) << kLogFinal) + k;
+#pragma unroll
+        for (int r = 0; r < kFinalRadix; ++r) {
+          x[r] = static_cast<float>(buf[j + h * r]);
+        }
+
+        hadamard_radix_pow2<kFinalRadix>(x);
+
+#pragma unroll
+        for (int r = 0; r < kFinalRadix; ++r) {
+          buf[j + h * r] = static_cast<T>(x[r]);
+        }
+      }
+      __syncthreads();
+    }
+
+    if constexpr (stride == 1) {
+#pragma unroll
+      for (int j = 0; j < max_radix / read_width; ++j) {
+        int index = j * read_width * kNumThreads + i * read_width;
+#pragma unroll
+        for (int r = 0; r < read_width; ++r) {
+          float val = static_cast<float>(buf[index + r]);
+          out[base + index + r] = static_cast<T>(val * scale);
+        }
+      }
+    } else {
+#pragma unroll
+      for (int j = 0; j < max_radix; ++j) {
+        out[base + (j * kNumThreads + i) * stride] = buf[j * kNumThreads + i];
+      }
+    }
+
+    __syncthreads();
+  }
+}
+
+template <typename T, int N, int M, int read_width>
+__global__ void
+hadamard_m(const T* in, T* out, float scale, long long num_tasks) {
+  constexpr int kTasksPerBatch = N / read_width;
+
+  for (long long task = blockIdx.x * blockDim.x + threadIdx.x; task < num_tasks;
+       task += blockDim.x * gridDim.x) {
+    long long i = task % kTasksPerBatch;
+    long long batch = task / kTasksPerBatch;
+    long long base = batch * static_cast<long long>(M) * N;
+
+    float x[read_width][M];
+#pragma unroll
+    for (int c = 0; c < M; ++c) {
+#pragma unroll
+      for (int r = 0; r < read_width; ++r) {
+        x[r][c] = static_cast<float>(in[base + c * N + i * read_width + r]);
+      }
+    }
+
+#pragma unroll
+    for (int r = 0; r < read_width; ++r) {
+      hadamard_radix_m(x[r]);
+    }
+
+#pragma unroll
+    for (int c = 0; c < M; ++c) {
+#pragma unroll
+      for (int r = 0; r < read_width; ++r) {
+        out[base + c * N + i * read_width + r] =
+            static_cast<T>(x[r][c] * scale);
+      }
+    }
+  }
+}
+
+} // namespace mlx::core::cu

mlx/backend/cuda/hadamard.cu

@@ -0,0 +1,245 @@
+// Copyright © 2025 Apple Inc.
+
+#include "mlx/backend/common/hadamard.h"
+#include "mlx/backend/cuda/device.h"
+#include "mlx/backend/cuda/jit_module.h"
+#include "mlx/backend/gpu/copy.h"
+#include "mlx/dtype_utils.h"
+#include "mlx/primitives.h"
+
+#include <fmt/format.h>
+#include <nvtx3/nvtx3.hpp>
+
+#include <algorithm>
+#include <cassert>
+#include <sstream>
+#include <stdexcept>
+#include <string_view>
+
+namespace mlx::core {
+
+namespace {
+
+constexpr int MAX_HADAMARD_THREADS_PER_BLOCK = 256;
+
+std::string gen_hadamard_codelet(int m) {
+  std::ostringstream source;
+  source << "namespace mlx::core::cu {\n";
+  source << "__device__ __forceinline__ void hadamard_radix_m(float* x) {\n";
+  if (m == 1) {
+    source << "}\n";
+    source << "} // namespace mlx::core::cu\n";
+    return source.str();
+  }
+
+  auto h_matrices = hadamard_matrices();
+  auto it = h_matrices.find(m);
+  if (it == h_matrices.end()) {
+    throw std::runtime_error("[hadamard] Invalid radix m.");
+  }
+  auto& matrix = it->second;
+
+  source << "  float tmp[" << m << "];\n";
+  auto start = 1;
+  auto end = matrix.find('\n', start);
+  int row_idx = 0;
+  while (end != std::string_view::npos) {
+    auto row = matrix.substr(start, end - start);
+    source << "  tmp[" << row_idx << "] =";
+    for (int i = 0; i < row.length(); ++i) {
+      source << " " << row[i] << " x[" << i << "]";
+    }
+    source << ";\n";
+    start = end + 1;
+    end = matrix.find('\n', start);
+    row_idx++;
+  }
+  source << "  #pragma unroll\n";
+  source << "  for (int i = 0; i < " << m << "; ++i) { x[i] = tmp[i]; }\n";
+  source << "}\n";
+  source << "} // namespace mlx::core::cu\n";
+  return source.str();
+}
+
+std::string hadamard_n_kernel_name(
+    const Dtype& dtype,
+    int n,
+    int max_radix,
+    int read_width,
+    int stride) {
+  return fmt::format(
+      "mlx::core::cu::hadamard_n<{}, {}, {}, {}, {}>",
+      dtype_to_cuda_type(dtype),
+      n,
+      max_radix,
+      read_width,
+      stride);
+}
+
+std::string
+hadamard_m_kernel_name(const Dtype& dtype, int n, int m, int read_width) {
+  return fmt::format(
+      "mlx::core::cu::hadamard_m<{}, {}, {}, {}>",
+      dtype_to_cuda_type(dtype),
+      n,
+      m,
+      read_width);
+}
+
+void hadamard_mn_contiguous(
+    const array& x,
+    array& y,
+    int m,
+    int n1,
+    int n2,
+    float scale,
+    const Stream& s) {
+  const int n = n1 * n2;
+  const int read_width_n1 = (n1 == 2) ? 2 : 4;
+  const int read_width_n2 = (n2 == 2) ? 2 : 4;
+  const int read_width_m = (n == 2 || m == 28) ? 2 : 4;
+  const int max_radix_1 = std::min(n1, 16);
+  const int max_radix_2 = std::min(n2, 16);
+  const float scale_n1 = 1.0f;
+  const float scale_n2 = (m == 1) ? scale : 1.0f;
+  const float scale_m = scale;
+
+  const std::string n1_kernel_name =
+      hadamard_n_kernel_name(x.dtype(), n1, max_radix_1, read_width_n1, n2);
+  const std::string n2_kernel_name =
+      hadamard_n_kernel_name(x.dtype(), n2, max_radix_2, read_width_n2, 1);
+  const std::string m_kernel_name =
+      hadamard_m_kernel_name(x.dtype(), n, m, read_width_m);
+
+  const std::string module_name = fmt::format(
+      "hadamard_{}_{}_{}_{}_{}_{}_{}_{}",
+      dtype_to_string(x.dtype()),
+      n,
+      m,
+      n1,
+      n2,
+      read_width_n1,
+      read_width_n2,
+      read_width_m);
+
+  cu::JitModule& mod = cu::get_jit_module(s.device, module_name, [&]() {
+    std::vector<std::string> kernel_names = {n2_kernel_name};
+    if (n1 > 1) {
+      kernel_names.push_back(n1_kernel_name);
+    }
+    if (m > 1) {
+      kernel_names.push_back(m_kernel_name);
+    }
+
+    std::string source = R"(
+        #include "mlx/backend/cuda/device/utils.cuh"
+    )";
+    source += gen_hadamard_codelet(m);
+    source += R"(
+        #include "mlx/backend/cuda/device/hadamard.cuh"
+    )";
+
+    return std::make_tuple(false, std::move(source), std::move(kernel_names));
+  });
+
+  auto& encoder = cu::get_command_encoder(s);
+
+  if (n1 > 1) {
+    const int64_t num_transforms = x.size() / n1;
+    const uint32_t num_blocks =
+        static_cast<uint32_t>(std::min<int64_t>(num_transforms, 65535));
+
+    encoder.set_input_array(x);
+    encoder.set_output_array(y);
+
+    cu::KernelArgs args;
+    args.append(x);
+    args.append(y);
+    args.append(scale_n1);
+    args.append(num_transforms);
+
+    auto kernel = mod.get_kernel(n1_kernel_name);
+    encoder.add_kernel_node_raw(
+        kernel, num_blocks, n1 / max_radix_1, {}, 0, args.args());
+  }
+
+  {
+    const auto& in = (n1 > 1) ? y : x;
+    const int64_t num_transforms = x.size() / n2;
+    const uint32_t num_blocks =
+        static_cast<uint32_t>(std::min<int64_t>(num_transforms, 65535));
+
+    encoder.set_input_array(in);
+    encoder.set_output_array(y);
+
+    cu::KernelArgs args;
+    args.append(in);
+    args.append(y);
+    args.append(scale_n2);
+    args.append(num_transforms);
+
+    auto kernel = mod.get_kernel(n2_kernel_name);
+    encoder.add_kernel_node_raw(
+        kernel, num_blocks, n2 / max_radix_2, {}, 0, args.args());
+  }
+
+  if (m > 1) {
+    const int64_t num_tasks = x.size() / (m * read_width_m);
+    const uint32_t block_dim = static_cast<uint32_t>(
+        std::min<int64_t>(num_tasks, MAX_HADAMARD_THREADS_PER_BLOCK));
+    const uint32_t num_blocks = static_cast<uint32_t>(
+        std::min<int64_t>((num_tasks + block_dim - 1) / block_dim, 65535));
+
+    encoder.set_input_array(y);
+    encoder.set_output_array(y);
+
+    cu::KernelArgs args;
+    args.append(y);
+    args.append(y);
+    args.append(scale_m);
+    args.append(num_tasks);
+
+    auto kernel = mod.get_kernel(m_kernel_name);
+    encoder.add_kernel_node_raw(
+        kernel, num_blocks, block_dim, {}, 0, args.args());
+  }
+}
+
+} // namespace
+
+void Hadamard::eval_gpu(const std::vector<array>& inputs, array& out) {
+  nvtx3::scoped_range r("Hadamard::eval_gpu");
+  assert(inputs.size() == 1);
+
+  auto& in = inputs[0];
+  if (in.dtype() != float16 && in.dtype() != bfloat16 &&
+      in.dtype() != float32) {
+    throw std::invalid_argument("[hadamard] Unsupported type.");
+  }
+
+  // n = m * 2^k where m in (1, 12, 20, 28)
+  auto [n, m] = decompose_hadamard(in.shape().back());
+  int n1 = 1;
+  int n2 = n;
+  if (n > 8192) {
+    for (n2 = 2; n2 * n2 < n; n2 *= 2) {
+    }
+    n1 = n / n2;
+  }
+
+  auto& s = stream();
+  auto& encoder = cu::get_command_encoder(s);
+  if (in.flags().row_contiguous) {
+    if (in.is_donatable()) {
+      out.copy_shared_buffer(in);
+    } else {
+      out.set_data(cu::malloc_async(out.nbytes(), encoder));
+    }
+    hadamard_mn_contiguous(in, out, m, n1, n2, scale_, s);
+  } else {
+    copy_gpu(in, out, CopyType::General, s);
+    hadamard_mn_contiguous(out, out, m, n1, n2, scale_, s);
+  }
+}
+
+} // namespace mlx::core

mlx/backend/cuda/jit_module.cpp

@@ -231,6 +231,7 @@ constexpr const char* g_include_names[] = {
     INCLUDE_PREFIX "config.h",
     INCLUDE_PREFIX "complex.cuh",
     INCLUDE_PREFIX "fp16_math.cuh",
+    INCLUDE_PREFIX "hadamard.cuh",
     INCLUDE_PREFIX "indexing.cuh",
     INCLUDE_PREFIX "scatter_ops.cuh",
     INCLUDE_PREFIX "unary_ops.cuh",
@@ -247,6 +248,7 @@ constexpr const char* g_headers[] = {
     jit_source_config,
     jit_source_complex,
     jit_source_fp16_math,
+    jit_source_hadamard,
     jit_source_indexing,
     jit_source_scatter_ops,
     jit_source_unary_ops,

mlx/backend/cuda/primitives.cpp

@@ -27,7 +27,6 @@ namespace mlx::core {
 NO_GPU(BlockMaskedMM)
 NO_GPU(FFT)
 NO_GPU(GatherQMM)
-NO_GPU(Hadamard)
 NO_GPU_MULTI(LUF)
 NO_GPU_MULTI(QRF)
 NO_GPU(SegmentedMM)

python/tests/cuda_skip.py

@@ -8,9 +8,6 @@ cuda_skip = {
     "TestBlas.test_gather_mm_sorted_vjp",
     # Segmented matmul NYI
     "TestBlas.test_segmented_mm",
-    # Hadamard NYI
-    "TestOps.test_hadamard",
-    "TestOps.test_hadamard_grad_vmap",
     # FFTs NYI
     "TestFFT.test_fft",
     "TestFFT.test_fft_big_powers_of_two",