diff --git a/mlx/backend/cpu/cholesky.cpp b/mlx/backend/cpu/cholesky.cpp index 3c5bbbc93d..2194b8b580 100644 --- a/mlx/backend/cpu/cholesky.cpp +++ b/mlx/backend/cpu/cholesky.cpp @@ -33,7 +33,7 @@ void cholesky_impl(const array& a, array& factor, bool upper, Stream stream) { N = a.shape(-1), size = a.size()]() mutable { char uplo = (upper) ? 'L' : 'U'; - size_t num_matrices = size / (N * N); + size_t num_matrices = (N > 0) ? size / (size_t(N) * N) : 0; for (int i = 0; i < num_matrices; i++) { // Compute Cholesky factorization. int info; diff --git a/mlx/backend/metal/CMakeLists.txt b/mlx/backend/metal/CMakeLists.txt index 5eb8f543c7..de61c62340 100644 --- a/mlx/backend/metal/CMakeLists.txt +++ b/mlx/backend/metal/CMakeLists.txt @@ -44,6 +44,7 @@ if(MLX_METAL_JIT) make_jit_source(binary_two) make_jit_source(fft kernels/fft/radix.h kernels/fft/readwrite.h) make_jit_source(logsumexp) + make_jit_source(cholesky) make_jit_source(ternary) make_jit_source(softmax) make_jit_source(scan) @@ -106,6 +107,7 @@ target_sources( mlx PRIVATE ${CMAKE_CURRENT_SOURCE_DIR}/allocator.cpp ${CMAKE_CURRENT_SOURCE_DIR}/binary.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/cholesky.cpp ${CMAKE_CURRENT_SOURCE_DIR}/device_info.cpp ${CMAKE_CURRENT_SOURCE_DIR}/compiled.cpp ${CMAKE_CURRENT_SOURCE_DIR}/conv.cpp diff --git a/mlx/backend/metal/cholesky.cpp b/mlx/backend/metal/cholesky.cpp new file mode 100644 index 0000000000..31198d157f --- /dev/null +++ b/mlx/backend/metal/cholesky.cpp @@ -0,0 +1,194 @@ +// Copyright © 2025 Apple Inc. + +#include +#include + +#include "mlx/backend/gpu/copy.h" +#include "mlx/backend/metal/device.h" +#include "mlx/backend/metal/kernels.h" +#include "mlx/backend/metal/matmul.h" +#include "mlx/backend/metal/utils.h" +#include "mlx/primitives.h" + +namespace mlx::core { + +void Cholesky::eval_gpu(const std::vector& inputs, array& out) { + assert(inputs.size() == 1); + + if (inputs[0].dtype() != float32) { + throw std::runtime_error( + "[Cholesky::eval_gpu] Metal Cholesky only supports float32."); + } + + auto& s = stream(); + auto& d = metal::device(s.device); + auto& compute_encoder = metal::get_command_encoder(s); + + // The factorization runs in place, so copy the input into the output. + const array& in = inputs[0]; + copy_gpu( + in, + out, + in.flags().row_contiguous ? CopyType::Vector : CopyType::General, + s); + + int N = out.shape(-1); + size_t num_matrices = (N == 0) ? 0 : out.size() / (size_t(N) * N); + if (num_matrices == 0 || N == 0) { + return; + } + + // Blocked right-looking factorization. Each strip of nb_outer columns is + // factored 32 columns at a time (cholesky_potf2 + cholesky_trsm + a + // strip-bounded cholesky_syrk), and the trailing update beyond the strip, + // where nearly all the FLOPs are for a large matrix, runs on the steel + // GEMM. A final fixup pass zeroes the strict upper triangle, or transposes + // in place for the upper factor. The phases are separate dispatches + // ordered by buffer barriers. + constexpr int NB = 32; + constexpr int nb_outer = 512; + int upper = upper_ ? 1 : 0; + std::string tname = type_to_name(out); + auto potf2_kernel = get_cholesky_kernel(d, "cholesky_potf2_" + tname, out); + auto trsm_kernel = get_cholesky_kernel(d, "cholesky_trsm_" + tname, out); + auto syrk32_kernel = get_cholesky_kernel(d, "cholesky_syrk32_" + tname, out); + auto syrk64_kernel = get_cholesky_kernel(d, "cholesky_syrk64_" + tname, out); + auto fixup_kernel = get_cholesky_kernel(d, "cholesky_fixup_" + tname, out); + + // Scratch for the inverted diagonal blocks. + array linv({static_cast(num_matrices), NB, NB}, float32, nullptr, {}); + linv.set_data(allocator::malloc(linv.nbytes())); + compute_encoder.add_temporary(linv); + + // Factor the 32 columns at p: POTF2 (+ TRTRI) then the TRSM below. + auto factor32 = [&](int p) { + compute_encoder.set_compute_pipeline_state(potf2_kernel); + compute_encoder.set_output_array(out, 0); + compute_encoder.set_bytes(N, 1); + compute_encoder.set_bytes(p, 2); + compute_encoder.set_output_array(linv, 3); + compute_encoder.dispatch_threadgroups( + MTL::Size(num_matrices, 1, 1), MTL::Size(256, 1, 1)); + compute_encoder.barrier(); + int m = N - p - NB; + if (m > 0) { + compute_encoder.set_compute_pipeline_state(trsm_kernel); + compute_encoder.set_output_array(out, 0); + compute_encoder.set_bytes(N, 1); + compute_encoder.set_bytes(p, 2); + compute_encoder.set_input_array(linv, 3); + int nrb = (m + 127) / 128; + compute_encoder.dispatch_threadgroups( + MTL::Size(nrb, num_matrices, 1), MTL::Size(256, 1, 1)); + compute_encoder.barrier(); + } + }; + // Rank-kd update bounded to the next ncols columns of the strip. + auto strip_syrk = + [&](MTL::ComputePipelineState* kernel, int p, int kd, int ncols) { + int ht = N - p - kd; + int cb = std::min(ht, ncols); + if (ht <= 0 || cb <= 0) { + return; + } + int ntj = (cb + NB - 1) / NB; + int nti = (ht + NB - 1) / NB; + compute_encoder.set_compute_pipeline_state(kernel); + compute_encoder.set_output_array(out, 0); + compute_encoder.set_bytes(N, 1); + compute_encoder.set_bytes(p, 2); + compute_encoder.set_bytes(cb, 3); + compute_encoder.dispatch_threadgroups( + MTL::Size(size_t(nti) * ntj, num_matrices, 1), + MTL::Size(128, 1, 1)); + compute_encoder.barrier(); + }; + + auto strided_flags = out.flags(); + strided_flags.row_contiguous = false; + strided_flags.col_contiguous = false; + strided_flags.contiguous = false; + + for (int po = 0; po < N; po += nb_outer) { + int W = std::min(nb_outer, N - po); + int strip_end = po + W; + + // Factor the strip with paired 32-column panels: prime the second + // panel's columns with a rank-32 update, factor it, then apply a single + // rank-64 update to the rest of the strip. + int pi = po; + while (pi < strip_end) { + factor32(pi); + int rem = strip_end - pi - NB; + if (rem <= 0) { + pi += NB; + } else if (rem <= NB) { + strip_syrk(syrk32_kernel, pi, NB, rem); + factor32(pi + NB); + pi += 2 * NB; + } else { + strip_syrk(syrk32_kernel, pi, NB, NB); + factor32(pi + NB); + strip_syrk(syrk64_kernel, pi, 2 * NB, rem - NB); + pi += 2 * NB; + } + } + + // Trailing update beyond the strip on the steel GEMM: C -= L21 L21ᵀ, in + // place on strided submatrix views. + int m = N - strip_end; + if (m > 0) { + Strides sub_strides{int64_t(N) * N, int64_t(N), 1}; + array l21({static_cast(num_matrices), m, W}, float32, nullptr, {}); + l21.copy_shared_buffer( + out, + sub_strides, + strided_flags, + out.data_size(), + int64_t(strip_end) * N + po); + array csub({static_cast(num_matrices), m, m}, float32, nullptr, {}); + csub.copy_shared_buffer( + out, + sub_strides, + strided_flags, + out.data_size(), + int64_t(strip_end) * N + strip_end); + std::vector copies; + steel_matmul_regular_axpby( + /* s = */ s, + /* d = */ d, + /* a = */ l21, + /* b = */ l21, + /* c = */ csub, + /* out = */ csub, + /* M = */ m, + /* N = */ m, + /* K = */ W, + /* batch_size_out = */ static_cast(num_matrices), + /* lda = */ N, + /* ldb = */ N, + /* ldd = */ N, + /* transpose_a = */ false, + /* transpose_b = */ true, + /* copies = */ copies, + /* batch_shape = */ {1}, + /* batch_strides = */ {0}, + /* A_batch_stride = */ int64_t(N) * N, + /* B_batch_stride = */ int64_t(N) * N, + /* matrix_stride_out = */ int64_t(N) * N, + /* C_batch_stride = */ int64_t(N) * N, + /* alpha = */ -1.0f, + /* beta = */ 1.0f); + compute_encoder.barrier(); + } + } + + compute_encoder.set_compute_pipeline_state(fixup_kernel); + compute_encoder.set_output_array(out, 0); + compute_encoder.set_bytes(N, 1); + compute_encoder.set_bytes(upper, 2); + compute_encoder.dispatch_threads( + MTL::Size(N, num_matrices, 1), MTL::Size(std::min(N, 256), 1, 1)); +} + +} // namespace mlx::core diff --git a/mlx/backend/metal/jit/includes.h b/mlx/backend/metal/jit/includes.h index ac9fb81e26..640693c0ce 100644 --- a/mlx/backend/metal/jit/includes.h +++ b/mlx/backend/metal/jit/includes.h @@ -23,6 +23,7 @@ const char* gather_axis(); const char* gather_front(); const char* hadamard(); const char* logsumexp(); +const char* cholesky(); const char* quantized_utils(); const char* quantized(); const char* fp_quantized(); diff --git a/mlx/backend/metal/jit_kernels.cpp b/mlx/backend/metal/jit_kernels.cpp index 9ed9ce43e0..6d52821d51 100644 --- a/mlx/backend/metal/jit_kernels.cpp +++ b/mlx/backend/metal/jit_kernels.cpp @@ -326,6 +326,34 @@ MTL::ComputePipelineState* get_logsumexp_kernel( return d.get_kernel(kernel_name, lib); } +MTL::ComputePipelineState* get_cholesky_kernel( + metal::Device& d, + const std::string& kernel_name, + const array& out) { + // kernel_name is "cholesky_{phase}_{type}", e.g. "cholesky_potf2_float32". + // Build all the phases into one library per type. + std::string tname = kernel_name.substr(kernel_name.rfind('_') + 1); + std::string lib_name = "cholesky_" + tname; + auto lib = d.get_library(lib_name, [&] { + auto t_str = get_type_string(out.dtype()); + std::string kernel_source; + kernel_source = metal::utils(); + kernel_source += metal::cholesky(); + kernel_source += get_template_definition( + "cholesky_potf2_" + tname, "cholesky_potf2", t_str); + kernel_source += get_template_definition( + "cholesky_trsm_" + tname, "cholesky_trsm", t_str); + kernel_source += get_template_definition( + "cholesky_syrk32_" + tname, "cholesky_syrk", t_str, 32); + kernel_source += get_template_definition( + "cholesky_syrk64_" + tname, "cholesky_syrk", t_str, 64); + kernel_source += get_template_definition( + "cholesky_fixup_" + tname, "cholesky_fixup", t_str); + return kernel_source; + }); + return d.get_kernel(kernel_name, lib); +} + MTL::ComputePipelineState* get_scan_kernel( metal::Device& d, const std::string& kernel_name, diff --git a/mlx/backend/metal/kernels.h b/mlx/backend/metal/kernels.h index cb289fc1e7..0fd5799645 100644 --- a/mlx/backend/metal/kernels.h +++ b/mlx/backend/metal/kernels.h @@ -64,6 +64,11 @@ MTL::ComputePipelineState* get_logsumexp_kernel( const std::string& kernel_name, const array& out); +MTL::ComputePipelineState* get_cholesky_kernel( + metal::Device& d, + const std::string& kernel_name, + const array& out); + MTL::ComputePipelineState* get_scan_kernel( metal::Device& d, const std::string& kernel_name, diff --git a/mlx/backend/metal/kernels/CMakeLists.txt b/mlx/backend/metal/kernels/CMakeLists.txt index 2bd99d2375..541d8d3f17 100644 --- a/mlx/backend/metal/kernels/CMakeLists.txt +++ b/mlx/backend/metal/kernels/CMakeLists.txt @@ -47,6 +47,7 @@ function(build_kernel KERNEL) endfunction(build_kernel) build_kernel(arg_reduce) +build_kernel(cholesky cholesky.h) build_kernel(conv steel/conv/params.h) build_kernel(layer_norm) build_kernel(random) diff --git a/mlx/backend/metal/kernels/cholesky.h b/mlx/backend/metal/kernels/cholesky.h new file mode 100644 index 0000000000..5db15623a6 --- /dev/null +++ b/mlx/backend/metal/kernels/cholesky.h @@ -0,0 +1,333 @@ +// Copyright © 2025 Apple Inc. + +#include +#include + +// Cholesky factorization kernels. +// +// The host (backend/metal/cholesky.cpp) factors the matrix in place in `out` +// with a blocked right-looking algorithm: strips of 512 columns are factored +// 32 columns at a time (cholesky_potf2 + cholesky_trsm + a strip-bounded +// cholesky_syrk), the large trailing update runs on the steel GEMM, and a +// final cholesky_fixup pass zeroes the strict upper triangle (or transposes +// in place for the upper factor). Only the lower triangle is read or written +// until the fixup pass, so the same kernels serve both triangles. +// +// The three phases are separate dispatches ordered by buffer barriers: every +// location is written by at most one dispatch and read by later ones, so +// there are no cross-threadgroup hazards inside any dispatch. + +// Factor the (up to) 32x32 diagonal block at (p, p). One threadgroup per +// matrix: simdgroup 0 keeps one row per lane in registers and broadcasts the +// pivot column with simd_shuffle, writes the factored block back, and stores +// the inverse of the block (TRTRI, one column per lane) to `linv` for the +// TRSM dispatch that follows. +template +[[kernel]] void cholesky_potf2( + device T* out [[buffer(0)]], + const constant int& N [[buffer(1)]], + const constant int& p [[buffer(2)]], + device T* linv [[buffer(3)]], + uint gid [[threadgroup_position_in_grid]], + uint lid [[thread_position_in_threadgroup]], + uint tgsize [[threads_per_threadgroup]], + uint sgid [[simdgroup_index_in_threadgroup]], + uint slane [[thread_index_in_simdgroup]]) { + constexpr int NB = 32; + threadgroup T Ldiag[NB * NB]; + + device T* A = out + size_t(gid) * N * N; + int W = metal::min(NB, N - p); + + for (int idx = int(lid); idx < W * W; idx += int(tgsize)) { + int r = idx / W; + int c = idx % W; + Ldiag[r * NB + c] = (c <= r) ? A[size_t(p + r) * N + (p + c)] : T(0); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Right-looking factorization in registers: after step j every lane i + // holds the final L[i][j] in r[j] (i >= j). Lanes run in lockstep within + // the simdgroup, so no barriers are needed. + if (sgid == 0) { + int lane = int(slane); + T r[NB] = {T(0)}; + if (lane < W) { + for (int k = 0; k <= lane; ++k) { + r[k] = Ldiag[lane * NB + k]; + } + } + for (int j = 0; j < W; ++j) { + T piv = metal::sqrt(simd_shuffle(r[j], ushort(j))); + if (lane == j) { + r[j] = piv; + } else if (lane > j) { + r[j] /= piv; + } + for (int k = j + 1; k < W; ++k) { + T lkj = simd_shuffle(r[j], ushort(k)); + if (lane >= k) { + r[k] -= r[j] * lkj; + } + } + } + if (lane < W) { + for (int k = 0; k <= lane; ++k) { + Ldiag[lane * NB + k] = r[k]; + } + } + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + for (int idx = int(lid); idx < W * W; idx += int(tgsize)) { + int r = idx / W; + int c = idx % W; + if (c <= r) { + A[size_t(p + r) * N + (p + c)] = Ldiag[r * NB + c]; + } + } + + // TRTRI for the TRSM below the block (only needed when such rows exist). + if (N - p > NB && sgid == 0) { + int j = int(slane); + T x[NB] = {T(0)}; + MLX_MTL_PRAGMA_UNROLL + for (int i = 0; i < NB; ++i) { + T s = (i == j) ? T(1) : T(0); + MLX_MTL_PRAGMA_UNROLL + for (int k = 0; k < i; ++k) { + s -= Ldiag[i * NB + k] * x[k]; + } + x[i] = (i >= j) ? s / Ldiag[i * NB + i] : T(0); + } + device T* lv = linv + size_t(gid) * NB * NB; + for (int i = 0; i < NB; ++i) { + lv[i * NB + j] = x[i]; + } + } +} + +// TRSM: X_new = X * Linv^T for the panel rows below the diagonal block, +// 128 rows per threadgroup (grid is row blocks x batch). Each simdgroup owns +// one 16-row slice as a 2x4 arrangement of 8x8 simdgroup-matrix fragments; +// rows past the last full slice fall back to one scalar product per lane. +template +[[kernel]] void cholesky_trsm( + device T* out [[buffer(0)]], + const constant int& N [[buffer(1)]], + const constant int& p [[buffer(2)]], + const device T* linv [[buffer(3)]], + uint2 tg_pos [[threadgroup_position_in_grid]], + uint2 lid2 [[thread_position_in_threadgroup]], + uint2 tgsize2 [[threads_per_threadgroup]], + uint sgid [[simdgroup_index_in_threadgroup]], + uint slane [[thread_index_in_simdgroup]]) { + constexpr int NB = 32; + threadgroup T Linv[NB * NB]; + + uint lid = lid2.x; + uint tgsize = tgsize2.x; + device T* A = out + size_t(tg_pos.y) * N * N; + const device T* lv = linv + size_t(tg_pos.y) * NB * NB; + + for (int idx = int(lid); idx < NB * NB; idx += int(tgsize)) { + Linv[idx] = lv[idx]; + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + int r0 = p + NB + int(tg_pos.x) * 128 + int(sgid) * 16; + if (r0 + 16 <= N) { + metal::simdgroup_matrix acc[2][4]; + MLX_MTL_PRAGMA_UNROLL + for (int ra = 0; ra < 2; ++ra) { + MLX_MTL_PRAGMA_UNROLL + for (int c = 0; c < 4; ++c) { + acc[ra][c] = metal::simdgroup_matrix(T(0)); + } + } + MLX_MTL_PRAGMA_UNROLL + for (int kb = 0; kb < NB; kb += 8) { + metal::simdgroup_matrix a[2]; + MLX_MTL_PRAGMA_UNROLL + for (int ra = 0; ra < 2; ++ra) { + simdgroup_load(a[ra], A + size_t(r0 + ra * 8) * N + (p + kb), ulong(N)); + } + MLX_MTL_PRAGMA_UNROLL + for (int c = 0; c < 4; ++c) { + metal::simdgroup_matrix b; + simdgroup_load( + b, Linv + c * 8 * NB + kb, ulong(NB), ulong2(0, 0), true); + MLX_MTL_PRAGMA_UNROLL + for (int ra = 0; ra < 2; ++ra) { + simdgroup_multiply_accumulate(acc[ra][c], a[ra], b, acc[ra][c]); + } + } + } + MLX_MTL_PRAGMA_UNROLL + for (int ra = 0; ra < 2; ++ra) { + MLX_MTL_PRAGMA_UNROLL + for (int c = 0; c < 4; ++c) { + simdgroup_store( + acc[ra][c], A + size_t(r0 + ra * 8) * N + (p + c * 8), ulong(N)); + } + } + } else if (r0 < N) { + // Scalar tail: each lane handles one row. + int r = r0 + int(slane); + if (r < N) { + T x[NB]; + MLX_MTL_PRAGMA_UNROLL + for (int k = 0; k < NB; ++k) { + x[k] = A[size_t(r) * N + (p + k)]; + } + MLX_MTL_PRAGMA_UNROLL + for (int j = 0; j < NB; ++j) { + T s = T(0); + MLX_MTL_PRAGMA_UNROLL + for (int k = 0; k < NB; ++k) { + s += x[k] * Linv[j * NB + k]; + } + A[size_t(r) * N + (p + j)] = s; + } + } + } +} + +// Trailing (SYRK) update: C -= L21 * L21^T over 32x32 tiles, using the KD +// panel columns p..p+KD-1 (the trailing matrix starts at q = p + KD). Only +// the first ncols trailing columns are updated: the host bounds the update +// to the current strip and leaves the rest of the trailing matrix to the +// steel GEMM. The grid is (row tiles x column tiles) x batch; each +// threadgroup of 4 simdgroups computes one tile as a 2x2 arrangement of +// 16x16 simdgroup-matrix products, and tiles crossing the row or column +// bound take a scalar path. +template +[[kernel]] void cholesky_syrk( + device T* out [[buffer(0)]], + const constant int& N [[buffer(1)]], + const constant int& p [[buffer(2)]], + const constant int& ncols [[buffer(3)]], + uint2 tg_pos [[threadgroup_position_in_grid]], + uint2 lid2 [[thread_position_in_threadgroup]], + uint2 tgsize2 [[threads_per_threadgroup]], + uint sgid [[simdgroup_index_in_threadgroup]]) { + constexpr int NB = 32; + threadgroup T As[NB * KD]; + threadgroup T Bs[NB * KD]; + + uint lid = lid2.x; + uint tgsize = tgsize2.x; + device T* A = out + size_t(tg_pos.y) * N * N; + int q = p + KD; + int Ht = N - q; + + // Tiles above the diagonal have nothing to update; the whole threadgroup + // returns before the barrier below. + int ntj = (ncols + NB - 1) / NB; + int t = int(tg_pos.x); + int ti = t / ntj; + int tj = t % ntj; + if (ti < tj) { + return; + } + + // Stage the two panel blocks this tile needs (rows q + ti*32 / q + tj*32 of + // panel columns p..p+KD-1), zero-padded past N. + for (int idx = int(lid); idx < NB * KD; idx += int(tgsize)) { + int r = idx / KD; + int c = idx % KD; + int ra = q + ti * NB + r; + int rb = q + tj * NB + r; + As[idx] = (ra < N) ? A[size_t(ra) * N + (p + c)] : T(0); + Bs[idx] = (rb < N) ? A[size_t(rb) * N + (p + c)] : T(0); + } + threadgroup_barrier(mem_flags::mem_threadgroup); + + if ((ti + 1) * NB <= Ht && (tj + 1) * NB <= ncols) { + // Full tile: 4 simdgroups, each a 16x16 block of the 32x32 tile. + int sr = int(sgid) / 2; + int sc = int(sgid) % 2; + metal::simdgroup_matrix acc[2][2]; + MLX_MTL_PRAGMA_UNROLL + for (int fr = 0; fr < 2; ++fr) { + MLX_MTL_PRAGMA_UNROLL + for (int fc = 0; fc < 2; ++fc) { + acc[fr][fc] = metal::simdgroup_matrix(T(0)); + } + } + MLX_MTL_PRAGMA_UNROLL + for (int kb = 0; kb < KD; kb += 8) { + metal::simdgroup_matrix a[2]; + metal::simdgroup_matrix b[2]; + MLX_MTL_PRAGMA_UNROLL + for (int f = 0; f < 2; ++f) { + simdgroup_load(a[f], As + (sr * 16 + f * 8) * KD + kb, ulong(KD)); + simdgroup_load( + b[f], + Bs + (sc * 16 + f * 8) * KD + kb, + ulong(KD), + ulong2(0, 0), + true); + } + MLX_MTL_PRAGMA_UNROLL + for (int fr = 0; fr < 2; ++fr) { + MLX_MTL_PRAGMA_UNROLL + for (int fc = 0; fc < 2; ++fc) { + simdgroup_multiply_accumulate(acc[fr][fc], a[fr], b[fc], acc[fr][fc]); + } + } + } + device T* C = + A + size_t(q + ti * NB + sr * 16) * N + (q + tj * NB + sc * 16); + MLX_MTL_PRAGMA_UNROLL + for (int fr = 0; fr < 2; ++fr) { + MLX_MTL_PRAGMA_UNROLL + for (int fc = 0; fc < 2; ++fc) { + device T* Cf = C + size_t(fr) * 8 * N + fc * 8; + metal::simdgroup_matrix c; + simdgroup_load(c, Cf, ulong(N)); + c.thread_elements()[0] -= acc[fr][fc].thread_elements()[0]; + c.thread_elements()[1] -= acc[fr][fc].thread_elements()[1]; + simdgroup_store(c, Cf, ulong(N)); + } + } + } else { + // Boundary tile: scalar with bounds checks. + for (int e = int(lid); e < NB * NB; e += int(tgsize)) { + int er = e / NB; + int ec = e % NB; + if (ti * NB + er < Ht && tj * NB + ec < ncols) { + T s = T(0); + for (int k = 0; k < KD; ++k) { + s += As[er * KD + k] * Bs[ec * KD + k]; + } + A[size_t(q + ti * NB + er) * N + (q + tj * NB + ec)] -= s; + } + } + } +} + +// Epilogue: the factor lives in the lower triangle. Zero the strict upper +// triangle, or for upper = true move it in place (R = L^T) and zero the +// strict lower. One thread per (matrix, row): each thread streams contiguous +// elements of its own row, and for upper only ever writes the strict upper +// triangle while reading its own lower row, so threads never collide. +template +[[kernel]] void cholesky_fixup( + device T* out [[buffer(0)]], + const constant int& N [[buffer(1)]], + const constant int& upper [[buffer(2)]], + uint2 pos [[thread_position_in_grid]]) { + int r = int(pos.x); + device T* A = out + size_t(pos.y) * N * N; + if (upper == 0) { + for (int c = r + 1; c < N; ++c) { + A[size_t(r) * N + c] = T(0); + } + } else { + for (int c = 0; c < r; ++c) { + A[size_t(c) * N + r] = A[size_t(r) * N + c]; + A[size_t(r) * N + c] = T(0); + } + } +} diff --git a/mlx/backend/metal/kernels/cholesky.metal b/mlx/backend/metal/kernels/cholesky.metal new file mode 100644 index 0000000000..8e46be56d7 --- /dev/null +++ b/mlx/backend/metal/kernels/cholesky.metal @@ -0,0 +1,21 @@ +// Copyright © 2025 Apple Inc. + +#include +#include +#include +#include + +using namespace metal; + +// clang-format off +#include "mlx/backend/metal/kernels/utils.h" +#include "mlx/backend/metal/kernels/cholesky.h" + +#define instantiate_cholesky(name, itype) \ + instantiate_kernel("cholesky_potf2_" #name, cholesky_potf2, itype) \ + instantiate_kernel("cholesky_trsm_" #name, cholesky_trsm, itype) \ + instantiate_kernel("cholesky_syrk32_" #name, cholesky_syrk, itype, 32) \ + instantiate_kernel("cholesky_syrk64_" #name, cholesky_syrk, itype, 64) \ + instantiate_kernel("cholesky_fixup_" #name, cholesky_fixup, itype) + +instantiate_cholesky(float32, float) // clang-format on diff --git a/mlx/backend/metal/nojit_kernels.cpp b/mlx/backend/metal/nojit_kernels.cpp index 64cfa39ed5..2386a4db72 100644 --- a/mlx/backend/metal/nojit_kernels.cpp +++ b/mlx/backend/metal/nojit_kernels.cpp @@ -79,6 +79,13 @@ MTL::ComputePipelineState* get_logsumexp_kernel( return d.get_kernel(kernel_name); } +MTL::ComputePipelineState* get_cholesky_kernel( + metal::Device& d, + const std::string& kernel_name, + const array&) { + return d.get_kernel(kernel_name); +} + MTL::ComputePipelineState* get_scan_kernel( metal::Device& d, const std::string& kernel_name, diff --git a/mlx/backend/metal/primitives.cpp b/mlx/backend/metal/primitives.cpp index d5bbf797e4..a06cd92bed 100644 --- a/mlx/backend/metal/primitives.cpp +++ b/mlx/backend/metal/primitives.cpp @@ -216,11 +216,6 @@ void Inverse::eval_gpu(const std::vector& inputs, array& output) { throw std::runtime_error("[Inverse::eval_gpu] Metal inversion NYI."); } -void Cholesky::eval_gpu(const std::vector& inputs, array& out) { - throw std::runtime_error( - "[Cholesky::eval_gpu] Metal Cholesky decomposition NYI."); -} - void Eig::eval_gpu( const std::vector& inputs, std::vector& outputs) { diff --git a/mlx/linalg.cpp b/mlx/linalg.cpp index bdf05fb80b..f3679b3465 100644 --- a/mlx/linalg.cpp +++ b/mlx/linalg.cpp @@ -335,7 +335,6 @@ array cholesky( const array& a, bool upper /* = false */, StreamOrDevice s /* = {} */) { - check_cpu_stream(s, "[linalg::cholesky]"); check_float(a.dtype(), "[linalg::cholesky]"); if (a.ndim() < 2) { std::ostringstream msg; @@ -350,6 +349,7 @@ array cholesky( "[linalg::cholesky] Cholesky decomposition is only defined for square " "matrices."); } + return array( a.shape(), a.dtype(), diff --git a/python/tests/test_linalg.py b/python/tests/test_linalg.py index afdf75d7c8..4526891569 100644 --- a/python/tests/test_linalg.py +++ b/python/tests/test_linalg.py @@ -276,6 +276,58 @@ def test_cholesky(self): for M, L in zip(AB, Ls): self.assertTrue(mx.allclose(L @ L.T, M, rtol=1e-5, atol=1e-7)) + def test_cholesky_gpu(self): + if not mx.is_available(mx.gpu): + return + mx.random.seed(7) + + # Sizes cover the 32-column panel and 512-column strip boundaries of + # the blocked factorization (multiples, non-multiples, and sizes + # smaller than one panel), with single and batched inputs. + for batch, n in [ + (8, 8), + (600, 32), + (600, 48), + (300, 96), + (150, 200), + (120, 256), + (1, 1024), + (2, 1300), + (1, 2048), + ]: + X = mx.random.normal((batch, n, n)) + A = X @ X.swapaxes(-1, -2) + n * mx.eye(n) + for upper in (False, True): + L = mx.linalg.cholesky(A, upper=upper, stream=mx.gpu) + rec = L.swapaxes(-1, -2) @ L if upper else L @ L.swapaxes(-1, -2) + self.assertTrue(mx.allclose(rec, A, rtol=1e-4, atol=1e-4)) + # The other triangle is exactly zero + if upper: + self.assertTrue(mx.all(L == mx.triu(L))) + else: + self.assertTrue(mx.all(L == mx.tril(L))) + + # Non-contiguous input + X = mx.random.normal((2, 128, 128)) + A = X @ X.swapaxes(-1, -2) + 128 * mx.eye(128) + At = mx.swapaxes(A, -1, -2) + L = mx.linalg.cholesky(At, stream=mx.gpu) + Lc = mx.linalg.cholesky(At, stream=mx.cpu) + self.assertTrue(mx.allclose(L, Lc, rtol=1e-4, atol=1e-4)) + + # float64 is not supported on the GPU; the op does not fall back to the + # CPU, so a GPU stream raises. + Ad = mx.eye(8, dtype=mx.float64, stream=mx.cpu) + with self.assertRaises(ValueError): + mx.eval(mx.linalg.cholesky(Ad, stream=mx.gpu)) + + # Zero-size inputs + for shape in [(0, 4, 4), (3, 0, 0)]: + for stream in (mx.cpu, mx.gpu): + L = mx.linalg.cholesky(mx.zeros(shape), stream=stream) + mx.eval(L) + self.assertEqual(L.shape, shape) + def test_pseudo_inverse(self): A = mx.array([[1, 2, 3], [6, -5, 4], [-9, 8, 7]], dtype=mx.float32) A_plus = mx.linalg.pinv(A, stream=mx.cpu)