From b6aa03e5b8be1e033a2ff4e4c0ab336b3df24057 Mon Sep 17 00:00:00 2001 From: Nripesh Niketan <86844847+NripeshN@users.noreply.github.com> Date: Mon, 26 Jan 2026 14:57:04 +0000 Subject: [PATCH] Update pre-commit hooks and versions for clang-format, black, and isort (#3059) --- .pre-commit-config.yaml | 6 +- mlx/array.cpp | 25 +-- mlx/array.h | 7 +- mlx/backend/cpu/compiled.cpp | 14 +- mlx/backend/cpu/jit_compiler.cpp | 31 ++-- mlx/backend/cuda/binary/binary.cuh | 11 +- mlx/backend/cuda/binary_two.cu | 11 +- mlx/backend/cuda/compiled.cpp | 46 +++-- mlx/backend/cuda/cublas_utils.h | 5 +- mlx/backend/cuda/cudnn_utils.h | 20 +- mlx/backend/cuda/cutlass_utils.cuh | 9 +- mlx/backend/cuda/device.cpp | 7 +- mlx/backend/cuda/event.cu | 7 +- mlx/backend/cuda/gemms/cublas_gemm.cpp | 5 +- mlx/backend/cuda/indexing.cpp | 68 +++---- mlx/backend/cuda/jit_module.cpp | 5 +- mlx/backend/cuda/lru_cache.h | 11 +- mlx/backend/cuda/scan.cu | 11 +- mlx/backend/cuda/sort.cu | 2 +- mlx/backend/cuda/unary/unary.cuh | 11 +- mlx/backend/metal/kernels/conv.metal | 68 +++---- mlx/backend/metal/kernels/gemv.metal | 8 +- mlx/backend/metal/kernels/gemv_masked.h | 4 +- mlx/backend/metal/kernels/steel/attn/mma.h | 4 +- .../kernels/steel/conv/kernels/steel_conv.h | 2 +- .../steel/conv/kernels/steel_conv_general.h | 2 +- .../steel/gemm/kernels/steel_gemm_gather.h | 4 +- .../gemm/kernels/steel_gemm_gather_nax.h | 2 +- .../steel/gemm/kernels/steel_gemm_masked.h | 4 +- .../steel/gemm/kernels/steel_gemm_segmented.h | 2 +- .../steel/gemm/kernels/steel_gemm_splitk.h | 2 +- mlx/backend/metal/kernels/steel/gemm/mma.h | 4 +- mlx/backend/metal/matmul.cpp | 171 +++++++++--------- mlx/distributed/jaccl/jaccl.cpp | 11 +- mlx/distributed/ring/ring.cpp | 81 +++++---- python/src/ops.cpp | 20 +- tests/linalg_tests.cpp | 10 +- 37 files changed, 376 insertions(+), 335 deletions(-) diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 9457405c..0345848e 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -6,17 +6,17 @@ repos: # - id: end-of-file-fixer # - id: trailing-whitespace - repo: https://github.com/pre-commit/mirrors-clang-format - rev: v19.1.7 + rev: v21.1.8 hooks: - id: clang-format # Using this mirror lets us use mypyc-compiled black, which is about 2x faster - repo: https://github.com/psf/black-pre-commit-mirror - rev: 25.1.0 + rev: 26.1.0 hooks: - id: black - repo: https://github.com/pycqa/isort - rev: 6.0.0 + rev: 7.0.0 hooks: - id: isort args: diff --git a/mlx/array.cpp b/mlx/array.cpp index 500b6cc9..1769f0f8 100644 --- a/mlx/array.cpp +++ b/mlx/array.cpp @@ -21,11 +21,12 @@ array::array( Dtype dtype, std::shared_ptr primitive, std::vector inputs) - : array_desc_(std::make_shared( - std::move(shape), - dtype, - std::move(primitive), - std::move(inputs))) { + : array_desc_( + std::make_shared( + std::move(shape), + dtype, + std::move(primitive), + std::move(inputs))) { if (has_primitive() && this->primitive().stream().device == Device::gpu) { for (auto& in : this->inputs()) { if (in.dtype() == float64) { @@ -69,16 +70,18 @@ array array::unsafe_weak_copy(const array& other) { } array::array(std::initializer_list data) - : array_desc_(std::make_shared( - Shape{static_cast(data.size())}, - float32)) { + : array_desc_( + std::make_shared( + Shape{static_cast(data.size())}, + float32)) { init(data.begin()); } array::array(std::initializer_list data, Dtype dtype) - : array_desc_(std::make_shared( - Shape{static_cast(data.size())}, - dtype)) { + : array_desc_( + std::make_shared( + Shape{static_cast(data.size())}, + dtype)) { init(data.begin()); } diff --git a/mlx/array.h b/mlx/array.h index 9b0f6575..6ea61c1a 100644 --- a/mlx/array.h +++ b/mlx/array.h @@ -542,9 +542,10 @@ template array::array( std::initializer_list data, Dtype dtype /* = TypeToDtype() */) - : array_desc_(std::make_shared( - Shape{static_cast(data.size())}, - dtype)) { + : array_desc_( + std::make_shared( + Shape{static_cast(data.size())}, + dtype)) { init(data.begin()); } diff --git a/mlx/backend/cpu/compiled.cpp b/mlx/backend/cpu/compiled.cpp index 53d6c47e..0806368e 100644 --- a/mlx/backend/cpu/compiled.cpp +++ b/mlx/backend/cpu/compiled.cpp @@ -119,13 +119,15 @@ void* compile( source_file.close(); try { - JitCompiler::exec(JitCompiler::build_command( - output_dir, source_file_name, shared_lib_name)); + JitCompiler::exec( + JitCompiler::build_command( + output_dir, source_file_name, shared_lib_name)); } catch (const std::exception& error) { - throw std::runtime_error(fmt::format( - "[Compile::eval_cpu] Failed to compile function {0}: {1}", - kernel_name, - error.what())); + throw std::runtime_error( + fmt::format( + "[Compile::eval_cpu] Failed to compile function {0}: {1}", + kernel_name, + error.what())); } } diff --git a/mlx/backend/cpu/jit_compiler.cpp b/mlx/backend/cpu/jit_compiler.cpp index fee31f4c..267e8fd8 100644 --- a/mlx/backend/cpu/jit_compiler.cpp +++ b/mlx/backend/cpu/jit_compiler.cpp @@ -36,10 +36,11 @@ struct VisualStudioInfo { // Get path of Visual Studio. // Use -latest to get only the most recent installation when multiple // versions are installed, avoiding path concatenation issues. - std::string vs_path = JitCompiler::exec(fmt::format( - "\"{0}\\Microsoft Visual Studio\\Installer\\vswhere.exe\"" - " -latest -property installationPath", - std::getenv("ProgramFiles(x86)"))); + std::string vs_path = JitCompiler::exec( + fmt::format( + "\"{0}\\Microsoft Visual Studio\\Installer\\vswhere.exe\"" + " -latest -property installationPath", + std::getenv("ProgramFiles(x86)"))); if (vs_path.empty()) { throw std::runtime_error("Can not find Visual Studio."); } @@ -52,10 +53,11 @@ struct VisualStudioInfo { .base(), vs_path.end()); // Read the envs from vcvarsall. - std::string envs = JitCompiler::exec(fmt::format( - "\"{0}\\VC\\Auxiliary\\Build\\vcvarsall.bat\" {1} >NUL && set", - vs_path, - arch)); + std::string envs = JitCompiler::exec( + fmt::format( + "\"{0}\\VC\\Auxiliary\\Build\\vcvarsall.bat\" {1} >NUL && set", + vs_path, + arch)); for (const std::string& line : str_split(envs, '\n')) { // Each line is in the format "ENV_NAME=values". auto pos = line.find_first_of('='); @@ -150,12 +152,13 @@ std::string JitCompiler::exec(const std::string& cmd) { int code = WEXITSTATUS(status); #endif if (code != 0) { - throw std::runtime_error(fmt::format( - "Failed to execute command with return code {0}: \"{1}\", " - "the output is: {2}", - code, - cmd, - ret)); + throw std::runtime_error( + fmt::format( + "Failed to execute command with return code {0}: \"{1}\", " + "the output is: {2}", + code, + cmd, + ret)); } return ret; } diff --git a/mlx/backend/cuda/binary/binary.cuh b/mlx/backend/cuda/binary/binary.cuh index 6c1ac248..d43027c1 100644 --- a/mlx/backend/cuda/binary/binary.cuh +++ b/mlx/backend/cuda/binary/binary.cuh @@ -346,11 +346,12 @@ void binary_op_gpu_inplace( }); } } else { - throw std::runtime_error(fmt::format( - "Can not do binary op {} on inputs of {} with result of {}.", - op, - dtype_to_string(a.dtype()), - dtype_to_string(out.dtype()))); + throw std::runtime_error( + fmt::format( + "Can not do binary op {} on inputs of {} with result of {}.", + op, + dtype_to_string(a.dtype()), + dtype_to_string(out.dtype()))); } }); }); diff --git a/mlx/backend/cuda/binary_two.cu b/mlx/backend/cuda/binary_two.cu index 1f242848..4acd3048 100644 --- a/mlx/backend/cuda/binary_two.cu +++ b/mlx/backend/cuda/binary_two.cu @@ -376,11 +376,12 @@ void binary_two_op_gpu_inplace( }); } } else { - throw std::runtime_error(fmt::format( - "Can not do binary op {} on inputs of {} with result of {}.", - op, - dtype_to_string(a.dtype()), - dtype_to_string(out_a.dtype()))); + throw std::runtime_error( + fmt::format( + "Can not do binary op {} on inputs of {} with result of {}.", + op, + dtype_to_string(a.dtype()), + dtype_to_string(out_a.dtype()))); } }); }); diff --git a/mlx/backend/cuda/compiled.cpp b/mlx/backend/cuda/compiled.cpp index adae38e5..29cf2e4f 100644 --- a/mlx/backend/cuda/compiled.cpp +++ b/mlx/backend/cuda/compiled.cpp @@ -36,14 +36,16 @@ struct FusedKernelBuilder { params.push_back( fmt::format("const {}* {}", dtype_to_cuda_type(x.dtype()), xname)); if (!is_scalar(x) && !contiguous) { - params.push_back(fmt::format( - "const __grid_constant__ cuda::std::array {}_strides", - xname)); + params.push_back( + fmt::format( + "const __grid_constant__ cuda::std::array {}_strides", + xname)); } } for (const auto& x : outputs) { - params.push_back(fmt::format( - "{}* {}", dtype_to_cuda_type(x.dtype()), namer.get_name(x))); + params.push_back( + fmt::format( + "{}* {}", dtype_to_cuda_type(x.dtype()), namer.get_name(x))); } if (!contiguous) { params.push_back( @@ -250,20 +252,30 @@ void Compiled::eval_gpu( builder.os += "\n} // namespace mlx::core::cu\n"; // Build kernel names. std::vector kernel_names; - kernel_names.push_back(fmt::format( - "mlx::core::cu::{}_contiguous", - lib_name(), - work_per_thread)); - kernel_names.push_back(fmt::format( - "mlx::core::cu::{}_contiguous", - lib_name(), - work_per_thread)); + kernel_names.push_back( + fmt::format( + "mlx::core::cu::{}_contiguous", + lib_name(), + work_per_thread)); + kernel_names.push_back( + fmt::format( + "mlx::core::cu::{}_contiguous", + lib_name(), + work_per_thread)); for (int wpt : {1, work_per_thread}) { for (int i = 1; i <= MAX_NDIM; ++i) { - kernel_names.push_back(fmt::format( - "mlx::core::cu::{}_strided<{}, uint32_t, {}>", lib_name(), i, wpt)); - kernel_names.push_back(fmt::format( - "mlx::core::cu::{}_strided<{}, int64_t, {}>", lib_name(), i, wpt)); + kernel_names.push_back( + fmt::format( + "mlx::core::cu::{}_strided<{}, uint32_t, {}>", + lib_name(), + i, + wpt)); + kernel_names.push_back( + fmt::format( + "mlx::core::cu::{}_strided<{}, int64_t, {}>", + lib_name(), + i, + wpt)); } } diff --git a/mlx/backend/cuda/cublas_utils.h b/mlx/backend/cuda/cublas_utils.h index 371ceeb7..56702b12 100644 --- a/mlx/backend/cuda/cublas_utils.h +++ b/mlx/backend/cuda/cublas_utils.h @@ -34,8 +34,9 @@ inline cudaDataType_t dtype_to_cublas_type(Dtype dtype, std::string_view tag) { case complex64: return CUDA_C_32F; default: - throw std::runtime_error(fmt::format( - "Unsupported dtype in {}: {}.", tag, dtype_to_string(dtype))); + throw std::runtime_error( + fmt::format( + "Unsupported dtype in {}: {}.", tag, dtype_to_string(dtype))); } } diff --git a/mlx/backend/cuda/cudnn_utils.h b/mlx/backend/cuda/cudnn_utils.h index 2975f675..ef930213 100644 --- a/mlx/backend/cuda/cudnn_utils.h +++ b/mlx/backend/cuda/cudnn_utils.h @@ -62,8 +62,9 @@ inline fe::DataType_t dtype_to_cudnn_type(Dtype dtype) { case float64: return fe::DataType_t::DOUBLE; default: - throw std::runtime_error(fmt::format( - "Unsupported dtype in cuDNN: {}.", dtype_to_string(dtype))); + throw std::runtime_error( + fmt::format( + "Unsupported dtype in cuDNN: {}.", dtype_to_string(dtype))); } } @@ -124,13 +125,14 @@ class DnnGraph : public fe::graph::Graph { // Create a cuDNN tensor for scalar. auto scalar(const char* name, int64_t uid, Dtype dtype) { - return Graph::tensor(fe::graph::Tensor_attributes() - .set_name(name) - .set_uid(uid) - .set_dim({1, 1, 1, 1}) - .set_stride({1, 1, 1, 1}) - .set_is_pass_by_value(true) - .set_data_type(dtype_to_cudnn_type(dtype))); + return Graph::tensor( + fe::graph::Tensor_attributes() + .set_name(name) + .set_uid(uid) + .set_dim({1, 1, 1, 1}) + .set_stride({1, 1, 1, 1}) + .set_is_pass_by_value(true) + .set_data_type(dtype_to_cudnn_type(dtype))); } // Call this before setting notes. diff --git a/mlx/backend/cuda/cutlass_utils.cuh b/mlx/backend/cuda/cutlass_utils.cuh index 45e1c166..c9bbf0d4 100644 --- a/mlx/backend/cuda/cutlass_utils.cuh +++ b/mlx/backend/cuda/cutlass_utils.cuh @@ -13,10 +13,11 @@ namespace mlx::core { // Throw exception if the cutlass API does not succeed. inline void check_cutlass_error(const char* name, cutlass::Status status) { if (status != cutlass::Status::kSuccess) { - throw std::runtime_error(fmt::format( - "{} failed with code: {}.", - name, - cutlass::cutlassGetStatusString(status))); + throw std::runtime_error( + fmt::format( + "{} failed with code: {}.", + name, + cutlass::cutlassGetStatusString(status))); } } diff --git a/mlx/backend/cuda/device.cpp b/mlx/backend/cuda/device.cpp index d101b252..a987ac7c 100644 --- a/mlx/backend/cuda/device.cpp +++ b/mlx/backend/cuda/device.cpp @@ -42,9 +42,10 @@ Device::Device(int device) : device_(device) { CHECK_CUDA_ERROR(cudaDeviceGetAttribute( &attr, cudaDevAttrConcurrentManagedAccess, device_)); if (attr != 1) { - throw std::runtime_error(fmt::format( - "Device {} does not support synchronization in managed memory.", - device_)); + throw std::runtime_error( + fmt::format( + "Device {} does not support synchronization in managed memory.", + device_)); } // The cublasLt handle is used by matmul. diff --git a/mlx/backend/cuda/event.cu b/mlx/backend/cuda/event.cu index 10446938..6880a904 100644 --- a/mlx/backend/cuda/event.cu +++ b/mlx/backend/cuda/event.cu @@ -119,9 +119,10 @@ void CudaEvent::init_pool() { class CopyableCudaEvent { public: explicit CopyableCudaEvent(Device& d) - : event_(std::make_shared( - d, - cudaEventDisableTiming | cudaEventBlockingSync)) {} + : event_( + std::make_shared( + d, + cudaEventDisableTiming | cudaEventBlockingSync)) {} void wait() { event_->wait(); diff --git a/mlx/backend/cuda/gemms/cublas_gemm.cpp b/mlx/backend/cuda/gemms/cublas_gemm.cpp index ee33c78a..78340457 100644 --- a/mlx/backend/cuda/gemms/cublas_gemm.cpp +++ b/mlx/backend/cuda/gemms/cublas_gemm.cpp @@ -27,8 +27,9 @@ cublasComputeType_t dtype_to_compute_type(Dtype dtype) { return mlx::core::env::enable_tf32() ? CUBLAS_COMPUTE_32F_FAST_TF32 : CUBLAS_COMPUTE_32F; default: - throw std::runtime_error(fmt::format( - "Unsupported dtype in CublasGemm: {}.", dtype_to_string(dtype))); + throw std::runtime_error( + fmt::format( + "Unsupported dtype in CublasGemm: {}.", dtype_to_string(dtype))); } } diff --git a/mlx/backend/cuda/indexing.cpp b/mlx/backend/cuda/indexing.cpp index 709cb754..f855cd85 100644 --- a/mlx/backend/cuda/indexing.cpp +++ b/mlx/backend/cuda/indexing.cpp @@ -86,13 +86,14 @@ void Gather::eval_gpu(const std::vector& inputs, array& out) { std::vector kernel_names; for (int ndim = 0; ndim <= MAX_NDIM; ++ndim) { for (int large = 0; large <= 1; ++large) { - kernel_names.push_back(fmt::format( - "mlx::core::cu::gather<{}, {}, {}, {}, {}>", - dtype_to_cuda_type(out.dtype()), - dtype_to_cuda_type(idx_dtype), - nidx, - ndim, - large ? "int64_t" : "int32_t")); + kernel_names.push_back( + fmt::format( + "mlx::core::cu::gather<{}, {}, {}, {}, {}>", + dtype_to_cuda_type(out.dtype()), + dtype_to_cuda_type(idx_dtype), + nidx, + ndim, + large ? "int64_t" : "int32_t")); } } return std::make_tuple(false, jit_source_gather, std::move(kernel_names)); @@ -179,14 +180,15 @@ void Scatter::eval_gpu(const std::vector& inputs, array& out) { std::vector kernel_names; for (int ndim = 0; ndim <= MAX_NDIM; ++ndim) { for (int large = 0; large <= 1; ++large) { - kernel_names.push_back(fmt::format( - "mlx::core::cu::scatter<{}, {}, mlx::core::cu::Scatter{}, {}, {}, {}>", - dtype_to_cuda_type(out.dtype()), - dtype_to_cuda_type(idx_dtype), - op, - nidx, - ndim, - large ? "int64_t" : "int32_t")); + kernel_names.push_back( + fmt::format( + "mlx::core::cu::scatter<{}, {}, mlx::core::cu::Scatter{}, {}, {}, {}>", + dtype_to_cuda_type(out.dtype()), + dtype_to_cuda_type(idx_dtype), + op, + nidx, + ndim, + large ? "int64_t" : "int32_t")); } } return std::make_tuple(false, jit_source_scatter, std::move(kernel_names)); @@ -258,14 +260,15 @@ void GatherAxis::eval_gpu(const std::vector& inputs, array& out) { for (int ndim = 0; ndim <= MAX_NDIM; ++ndim) { for (int contiguous = 0; contiguous < 4; ++contiguous) { for (int large = 0; large <= 1; ++large) { - kernel_names.push_back(fmt::format( - "mlx::core::cu::gather_axis<{}, {}, {}, {}, {}, {}>", - dtype_to_cuda_type(out.dtype()), - dtype_to_cuda_type(idx.dtype()), - ndim, - contiguous & 1 ? true : false, - contiguous & 2 ? true : false, - large ? "int64_t" : "int32_t")); + kernel_names.push_back( + fmt::format( + "mlx::core::cu::gather_axis<{}, {}, {}, {}, {}, {}>", + dtype_to_cuda_type(out.dtype()), + dtype_to_cuda_type(idx.dtype()), + ndim, + contiguous & 1 ? true : false, + contiguous & 2 ? true : false, + large ? "int64_t" : "int32_t")); } } } @@ -360,15 +363,16 @@ void ScatterAxis::eval_gpu(const std::vector& inputs, array& out) { for (int ndim = 0; ndim <= MAX_NDIM; ++ndim) { for (int contiguous = 0; contiguous < 4; ++contiguous) { for (int large = 0; large <= 1; ++large) { - kernel_names.push_back(fmt::format( - "mlx::core::cu::scatter_axis<{}, {}, mlx::core::cu::Scatter{}, {}, {}, {}, {}>", - dtype_to_cuda_type(out.dtype()), - dtype_to_cuda_type(idx.dtype()), - op, - ndim, - contiguous & 1 ? true : false, - contiguous & 2 ? true : false, - large ? "int64_t" : "int32_t")); + kernel_names.push_back( + fmt::format( + "mlx::core::cu::scatter_axis<{}, {}, mlx::core::cu::Scatter{}, {}, {}, {}, {}>", + dtype_to_cuda_type(out.dtype()), + dtype_to_cuda_type(idx.dtype()), + op, + ndim, + contiguous & 1 ? true : false, + contiguous & 2 ? true : false, + large ? "int64_t" : "int32_t")); } } } diff --git a/mlx/backend/cuda/jit_module.cpp b/mlx/backend/cuda/jit_module.cpp index 5db55bdd..39d5e878 100644 --- a/mlx/backend/cuda/jit_module.cpp +++ b/mlx/backend/cuda/jit_module.cpp @@ -330,8 +330,9 @@ void load_module( CUresult jit_result = cuModuleLoadDataEx( &module_, ptx.data(), std::size(options), options, values); if (jit_result != CUDA_SUCCESS) { - throw std::runtime_error(fmt::format( - "Failed to load compiled {} kernel: {}.", module_name, jit_log)); + throw std::runtime_error( + fmt::format( + "Failed to load compiled {} kernel: {}.", module_name, jit_log)); } // Load kernels. diff --git a/mlx/backend/cuda/lru_cache.h b/mlx/backend/cuda/lru_cache.h index 94a96a9d..f23f3ff0 100644 --- a/mlx/backend/cuda/lru_cache.h +++ b/mlx/backend/cuda/lru_cache.h @@ -89,11 +89,12 @@ class LRUCache { } if (env_name_ && ++cache_misses_ > 2 * capacity_) { - throw std::runtime_error(fmt::format( - "Cache thrashing is happening, please set the environment variable " - "{} to a larger value than {} to fix degraded performance.", - env_name_, - capacity_)); + throw std::runtime_error( + fmt::format( + "Cache thrashing is happening, please set the environment variable " + "{} to a larger value than {} to fix degraded performance.", + env_name_, + capacity_)); } vlist_.emplace_front(key, std::forward(value)); diff --git a/mlx/backend/cuda/scan.cu b/mlx/backend/cuda/scan.cu index 6a1e3ef1..b4d6cf36 100644 --- a/mlx/backend/cuda/scan.cu +++ b/mlx/backend/cuda/scan.cu @@ -454,11 +454,12 @@ void Scan::eval_gpu(const std::vector& inputs, array& out) { }); }); } else { - throw std::runtime_error(fmt::format( - "Can not do scan op {} on inputs of {} with result of {}.", - op_to_string(), - dtype_to_string(in.dtype()), - dtype_to_string(out.dtype()))); + throw std::runtime_error( + fmt::format( + "Can not do scan op {} on inputs of {} with result of {}.", + op_to_string(), + dtype_to_string(in.dtype()), + dtype_to_string(out.dtype()))); } }); }); diff --git a/mlx/backend/cuda/sort.cu b/mlx/backend/cuda/sort.cu index acc8deb7..2c07be48 100644 --- a/mlx/backend/cuda/sort.cu +++ b/mlx/backend/cuda/sort.cu @@ -1073,4 +1073,4 @@ void Partition::eval_gpu(const std::vector& inputs, array& out) { gpu_sort(stream(), inputs[0], out, axis_, false); } -} // namespace mlx::core +} // namespace mlx::core \ No newline at end of file diff --git a/mlx/backend/cuda/unary/unary.cuh b/mlx/backend/cuda/unary/unary.cuh index dda7d784..cade3f81 100644 --- a/mlx/backend/cuda/unary/unary.cuh +++ b/mlx/backend/cuda/unary/unary.cuh @@ -191,11 +191,12 @@ void unary_op_gpu_inplace( } }); } else { - throw std::runtime_error(fmt::format( - "Can not do unary op {} on input of {} with output of {}.", - op, - dtype_to_string(in.dtype()), - dtype_to_string(out.dtype()))); + throw std::runtime_error( + fmt::format( + "Can not do unary op {} on input of {} with output of {}.", + op, + dtype_to_string(in.dtype()), + dtype_to_string(out.dtype()))); } }); }); diff --git a/mlx/backend/metal/kernels/conv.metal b/mlx/backend/metal/kernels/conv.metal index fdec515e..50a14ad3 100644 --- a/mlx/backend/metal/kernels/conv.metal +++ b/mlx/backend/metal/kernels/conv.metal @@ -445,20 +445,20 @@ winograd_conv_2d_weight_transform( } } -#define instantiate_winograd_conv_2d_weight_transform_base(name, itype, bc) \ - template [[host_name("winograd_conv_2d_weight_transform_" #name \ - "_bc" #bc)]] [[kernel]] void \ - winograd_conv_2d_weight_transform( \ - const device itype* wt_in [[buffer(0)]], \ - device itype* wt_out [[buffer(1)]], \ - const constant int& C [[buffer(2)]], \ - const constant int& O [[buffer(3)]], \ - uint tid [[threadgroup_position_in_grid]], \ - uint simd_group_id [[simdgroup_index_in_threadgroup]], \ +#define instantiate_winograd_conv_2d_weight_transform_base(name, itype, bc) \ + template [[host_name( \ + "winograd_conv_2d_weight_transform_" #name "_bc" #bc)]] [[kernel]] void \ + winograd_conv_2d_weight_transform( \ + const device itype* wt_in [[buffer(0)]], \ + device itype* wt_out [[buffer(1)]], \ + const constant int& C [[buffer(2)]], \ + const constant int& O [[buffer(3)]], \ + uint tid [[threadgroup_position_in_grid]], \ + uint simd_group_id [[simdgroup_index_in_threadgroup]], \ uint simd_lane_id [[thread_index_in_simdgroup]]); template -[[kernel, max_total_threads_per_threadgroup(WM* WN * 32)]] void +[[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] void winograd_conv_2d_input_transform( const device T* inp_in [[buffer(0)]], device T* inp_out [[buffer(1)]], @@ -555,21 +555,21 @@ winograd_conv_2d_input_transform( } } -#define instantiate_winograd_conv_2d_input_transform(name, itype, bc) \ - template [[host_name("winograd_conv_2d_input_transform_" #name \ - "_bc" #bc)]] [[kernel]] void \ - winograd_conv_2d_input_transform( \ - const device itype* inp_in [[buffer(0)]], \ - device itype* inp_out [[buffer(1)]], \ - const constant MLXConvParams<2>& params [[buffer(2)]], \ - uint3 tid [[threadgroup_position_in_grid]], \ - uint3 lid [[thread_position_in_threadgroup]], \ - uint3 tgp_per_grid [[threadgroups_per_grid]], \ - uint simd_group_id [[simdgroup_index_in_threadgroup]], \ +#define instantiate_winograd_conv_2d_input_transform(name, itype, bc) \ + template [[host_name( \ + "winograd_conv_2d_input_transform_" #name "_bc" #bc)]] [[kernel]] void \ + winograd_conv_2d_input_transform( \ + const device itype* inp_in [[buffer(0)]], \ + device itype* inp_out [[buffer(1)]], \ + const constant MLXConvParams<2>& params [[buffer(2)]], \ + uint3 tid [[threadgroup_position_in_grid]], \ + uint3 lid [[thread_position_in_threadgroup]], \ + uint3 tgp_per_grid [[threadgroups_per_grid]], \ + uint simd_group_id [[simdgroup_index_in_threadgroup]], \ uint simd_lane_id [[thread_index_in_simdgroup]]); template -[[kernel, max_total_threads_per_threadgroup(WM* WN * 32)]] void +[[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] void winograd_conv_2d_output_transform( const device T* out_in [[buffer(0)]], device T* out_out [[buffer(1)]], @@ -676,17 +676,17 @@ winograd_conv_2d_output_transform( } } -#define instantiate_winograd_conv_2d_output_transform(name, itype, bo) \ - template [[host_name("winograd_conv_2d_output_transform_" #name \ - "_bo" #bo)]] [[kernel]] void \ - winograd_conv_2d_output_transform( \ - const device itype* out_in [[buffer(0)]], \ - device itype* out_out [[buffer(1)]], \ - const constant MLXConvParams<2>& params [[buffer(2)]], \ - uint3 tid [[threadgroup_position_in_grid]], \ - uint3 lid [[thread_position_in_threadgroup]], \ - uint3 tgp_per_grid [[threadgroups_per_grid]], \ - uint simd_group_id [[simdgroup_index_in_threadgroup]], \ +#define instantiate_winograd_conv_2d_output_transform(name, itype, bo) \ + template [[host_name( \ + "winograd_conv_2d_output_transform_" #name "_bo" #bo)]] [[kernel]] void \ + winograd_conv_2d_output_transform( \ + const device itype* out_in [[buffer(0)]], \ + device itype* out_out [[buffer(1)]], \ + const constant MLXConvParams<2>& params [[buffer(2)]], \ + uint3 tid [[threadgroup_position_in_grid]], \ + uint3 lid [[thread_position_in_threadgroup]], \ + uint3 tgp_per_grid [[threadgroups_per_grid]], \ + uint simd_group_id [[simdgroup_index_in_threadgroup]], \ uint simd_lane_id [[thread_index_in_simdgroup]]); // clang-format off diff --git a/mlx/backend/metal/kernels/gemv.metal b/mlx/backend/metal/kernels/gemv.metal index 0914d802..277b68aa 100644 --- a/mlx/backend/metal/kernels/gemv.metal +++ b/mlx/backend/metal/kernels/gemv.metal @@ -445,7 +445,7 @@ template < const int TN, /* Thread cols (in elements) */ const bool kDoNCBatch, /* Batch ndim > 1 */ const bool kDoAxpby> /* Do out = alpha * out + beta * bias */ -[[kernel, max_total_threads_per_threadgroup(BM* BN * 32)]] void gemv( +[[kernel, max_total_threads_per_threadgroup(BM * BN * 32)]] void gemv( const device T* mat [[buffer(0)]], const device T* in_vec [[buffer(1)]], const device T* bias [[buffer(2)]], @@ -553,7 +553,7 @@ template < const int SN, /* Simdgroup cols (in threads) */ const int TM, /* Thread rows (in elements) */ const int TN> /* Thread cols (in elements) */ -[[kernel, max_total_threads_per_threadgroup(BM* BN * 32)]] void gemv_gather( +[[kernel, max_total_threads_per_threadgroup(BM * BN * 32)]] void gemv_gather( const device T* mat [[buffer(0)]], const device T* in_vec [[buffer(1)]], const device T* bias [[buffer(2)]], @@ -666,7 +666,7 @@ template < const int TN, /* Thread cols (in elements) */ const bool kDoNCBatch, /* Batch ndim > 1 */ const bool kDoAxpby> /* Do out = alpha * out + beta * bias */ -[[kernel, max_total_threads_per_threadgroup(BM* BN * 32)]] void gemv_t( +[[kernel, max_total_threads_per_threadgroup(BM * BN * 32)]] void gemv_t( const device T* mat [[buffer(0)]], const device T* in_vec [[buffer(1)]], const device T* bias [[buffer(2)]], @@ -764,7 +764,7 @@ template < const int SN, /* Simdgroup cols (in threads) */ const int TM, /* Thread rows (in elements) */ const int TN> /* Thread cols (in elements) */ -[[kernel, max_total_threads_per_threadgroup(BM* BN * 32)]] void gemv_t_gather( +[[kernel, max_total_threads_per_threadgroup(BM * BN * 32)]] void gemv_t_gather( const device T* mat [[buffer(0)]], const device T* in_vec [[buffer(1)]], const device T* bias [[buffer(2)]], diff --git a/mlx/backend/metal/kernels/gemv_masked.h b/mlx/backend/metal/kernels/gemv_masked.h index 96b0c282..407d14bb 100644 --- a/mlx/backend/metal/kernels/gemv_masked.h +++ b/mlx/backend/metal/kernels/gemv_masked.h @@ -641,7 +641,7 @@ template < const int TM, /* Thread rows (in elements) */ const int TN, /* Thread cols (in elements) */ const bool kDoNCBatch> /* Batch ndim > 1 */ -[[kernel, max_total_threads_per_threadgroup(BM* BN * 32)]] void gemv_masked( +[[kernel, max_total_threads_per_threadgroup(BM * BN * 32)]] void gemv_masked( const device T* mat [[buffer(0)]], const device T* in_vec [[buffer(1)]], device T* out_vec [[buffer(3)]], @@ -741,7 +741,7 @@ template < const int TM, /* Thread rows (in elements) */ const int TN, /* Thread cols (in elements) */ const bool kDoNCBatch> /* Batch ndim > 1 */ -[[kernel, max_total_threads_per_threadgroup(BM* BN * 32)]] void gemv_t_masked( +[[kernel, max_total_threads_per_threadgroup(BM * BN * 32)]] void gemv_t_masked( const device T* mat [[buffer(0)]], const device T* in_vec [[buffer(1)]], device T* out_vec [[buffer(3)]], diff --git a/mlx/backend/metal/kernels/steel/attn/mma.h b/mlx/backend/metal/kernels/steel/attn/mma.h index b11a111b..737e930d 100644 --- a/mlx/backend/metal/kernels/steel/attn/mma.h +++ b/mlx/backend/metal/kernels/steel/attn/mma.h @@ -68,8 +68,8 @@ struct BaseMMAFrag { template using dtype_frag_t = typename metal::vec; - METAL_FUNC static constexpr short2 get_coord(ushort simd_lane_id - [[thread_index_in_simdgroup]]) { + METAL_FUNC static constexpr short2 get_coord( + ushort simd_lane_id [[thread_index_in_simdgroup]]) { const short qid = simd_lane_id / 4; const short fm = (qid & 4) + ((simd_lane_id / 2) % 4); const short fn = (qid & 2) * 2 + (simd_lane_id % 2) * 2; diff --git a/mlx/backend/metal/kernels/steel/conv/kernels/steel_conv.h b/mlx/backend/metal/kernels/steel/conv/kernels/steel_conv.h index 6f822c1d..850ec15b 100644 --- a/mlx/backend/metal/kernels/steel/conv/kernels/steel_conv.h +++ b/mlx/backend/metal/kernels/steel/conv/kernels/steel_conv.h @@ -13,7 +13,7 @@ template < int WN, int N_CHANNELS = 0, bool SMALL_FILTER = false> -[[kernel, max_total_threads_per_threadgroup(WM* WN * 32)]] void +[[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] void implicit_gemm_conv_2d( const device T* A [[buffer(0)]], const device T* B [[buffer(1)]], diff --git a/mlx/backend/metal/kernels/steel/conv/kernels/steel_conv_general.h b/mlx/backend/metal/kernels/steel/conv/kernels/steel_conv_general.h index 9afebd30..1241f773 100644 --- a/mlx/backend/metal/kernels/steel/conv/kernels/steel_conv_general.h +++ b/mlx/backend/metal/kernels/steel/conv/kernels/steel_conv_general.h @@ -13,7 +13,7 @@ template < int WN, typename AccumType = float, typename Epilogue = TransformNone> -[[kernel, max_total_threads_per_threadgroup(WM* WN * 32)]] void +[[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] void implicit_gemm_conv_2d_general( const device T* A [[buffer(0)]], const device T* B [[buffer(1)]], diff --git a/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_gather.h b/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_gather.h index 4493375c..4c055e69 100644 --- a/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_gather.h +++ b/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_gather.h @@ -17,7 +17,7 @@ template < bool transpose_a, bool transpose_b, typename AccumType = float> -[[kernel, max_total_threads_per_threadgroup(WM* WN * 32)]] void gather_mm_rhs( +[[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] void gather_mm_rhs( const device T* A [[buffer(0)]], const device T* B [[buffer(1)]], const device uint32_t* rhs_indices [[buffer(2)]], @@ -248,7 +248,7 @@ template < bool transpose_a, bool transpose_b, typename AccumType = float> -[[kernel, max_total_threads_per_threadgroup(WM* WN * 32)]] void gather_mm( +[[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] void gather_mm( const device T* A [[buffer(0)]], const device T* B [[buffer(1)]], const device uint32_t* lhs_indices [[buffer(2)]], diff --git a/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_gather_nax.h b/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_gather_nax.h index 29285833..0f83ed18 100644 --- a/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_gather_nax.h +++ b/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_gather_nax.h @@ -16,7 +16,7 @@ template < bool transpose_a, bool transpose_b, typename AccumType = float> -[[kernel, max_total_threads_per_threadgroup(WM* WN * 32)]] void +[[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] void gather_mm_rhs_nax( const device T* A [[buffer(0)]], const device T* B [[buffer(1)]], diff --git a/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_masked.h b/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_masked.h index c8ffe2b8..cc3ddd93 100644 --- a/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_masked.h +++ b/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_masked.h @@ -49,7 +49,7 @@ template < bool transpose_b, bool MN_aligned, bool K_aligned> -[[kernel, max_total_threads_per_threadgroup(WM* WN * 32)]] void +[[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] void block_masked_gemm( const device T* A [[buffer(0)]], const device T* B [[buffer(1)]], @@ -435,7 +435,7 @@ template < bool MN_aligned, bool K_aligned, bool has_operand_mask = false> -[[kernel, max_total_threads_per_threadgroup(WM* WN * 32)]] void +[[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] void block_masked_gemm( const device T* A [[buffer(0)]], const device T* B [[buffer(1)]], diff --git a/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_segmented.h b/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_segmented.h index b915eb34..5a43e223 100644 --- a/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_segmented.h +++ b/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_segmented.h @@ -16,7 +16,7 @@ template < bool transpose_a, bool transpose_b, typename AccumType = float> -[[kernel, max_total_threads_per_threadgroup(WM* WN * 32)]] void segmented_mm( +[[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] void segmented_mm( const device T* A [[buffer(0)]], const device T* B [[buffer(1)]], const device uint32_t* segments [[buffer(2)]], diff --git a/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_splitk.h b/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_splitk.h index 1ff97ea4..a372e939 100644 --- a/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_splitk.h +++ b/mlx/backend/metal/kernels/steel/gemm/kernels/steel_gemm_splitk.h @@ -18,7 +18,7 @@ template < bool transpose_b, bool MN_aligned, bool K_aligned> -[[kernel, max_total_threads_per_threadgroup(WM* WN * 32)]] void gemm_splitk( +[[kernel, max_total_threads_per_threadgroup(WM * WN * 32)]] void gemm_splitk( const device T* A [[buffer(0)]], const device T* B [[buffer(1)]], device U* C [[buffer(2)]], diff --git a/mlx/backend/metal/kernels/steel/gemm/mma.h b/mlx/backend/metal/kernels/steel/gemm/mma.h index 74151a95..11319940 100644 --- a/mlx/backend/metal/kernels/steel/gemm/mma.h +++ b/mlx/backend/metal/kernels/steel/gemm/mma.h @@ -46,8 +46,8 @@ struct BaseMMAFrag { typedef metal::simdgroup_matrix mat_type; typedef metal::vec frag_type; - METAL_FUNC static constexpr short2 get_coord(ushort simd_lane_id - [[thread_index_in_simdgroup]]) { + METAL_FUNC static constexpr short2 get_coord( + ushort simd_lane_id [[thread_index_in_simdgroup]]) { const short qid = simd_lane_id / 4; const short fm = (qid & 4) + ((simd_lane_id / 2) % 4); const short fn = (qid & 2) * 2 + (simd_lane_id % 2) * 2; diff --git a/mlx/backend/metal/matmul.cpp b/mlx/backend/metal/matmul.cpp index bdc3ed34..5b80a4d8 100644 --- a/mlx/backend/metal/matmul.cpp +++ b/mlx/backend/metal/matmul.cpp @@ -270,21 +270,20 @@ void steel_matmul_regular_axpby_nax( int swizzle_log = tm <= 3 ? 0 : 1; // Prepare steel matmul params - GEMMParams params{ - /* const int M = */ M, - /* const int N = */ N, - /* const int K = */ K, - /* const int lda = */ lda, - /* const int ldb = */ ldb, - /* const int ldd = */ ldd, - /* const int tiles_n = */ tn, - /* const int tiles_m = */ tm, - /* const int64_t batch_stride_a = */ A_batch_stride, - /* const int64_t batch_stride_b = */ B_batch_stride, - /* const int64_t batch_stride_d = */ matrix_stride_out, - /* const int swizzle_log = */ swizzle_log, - /* const int gemm_k_iterations_aligned = */ (K / bk), - /* const int batch_ndim = */ int(batch_shape.size())}; + GEMMParams params{/* const int M = */ M, + /* const int N = */ N, + /* const int K = */ K, + /* const int lda = */ lda, + /* const int ldb = */ ldb, + /* const int ldd = */ ldd, + /* const int tiles_n = */ tn, + /* const int tiles_m = */ tm, + /* const int64_t batch_stride_a = */ A_batch_stride, + /* const int64_t batch_stride_b = */ B_batch_stride, + /* const int64_t batch_stride_d = */ matrix_stride_out, + /* const int swizzle_log = */ swizzle_log, + /* const int gemm_k_iterations_aligned = */ (K / bk), + /* const int batch_ndim = */ int(batch_shape.size())}; // Prepare launch grid params int tile = 1 << swizzle_log; @@ -310,12 +309,11 @@ void steel_matmul_regular_axpby_nax( int ldc = c.strides()[c.ndim() - 2]; int fdc = c.strides()[c.ndim() - 1]; - GEMMAddMMParams params{ - /* const int ldc = */ ldc, - /* const int fdc = */ fdc, - /* const int64_t batch_stride_c = */ C_batch_stride, - /* const float alpha = */ alpha, - /* const float beta = */ beta}; + GEMMAddMMParams params{/* const int ldc = */ ldc, + /* const int fdc = */ fdc, + /* const int64_t batch_stride_c = */ C_batch_stride, + /* const float alpha = */ alpha, + /* const float beta = */ beta}; compute_encoder.set_input_array(c, 2); compute_encoder.set_bytes(params, 5); @@ -457,21 +455,20 @@ void steel_matmul_regular_axpby( int swizzle_log = 0; // tm >= 6 ? 3 : (tm <= 3 ? 0 : 2); // Prepare steel matmul params - GEMMParams params{ - /* const int M = */ M, - /* const int N = */ N, - /* const int K = */ K, - /* const int lda = */ lda, - /* const int ldb = */ ldb, - /* const int ldd = */ ldd, - /* const int tiles_n = */ tn, - /* const int tiles_m = */ tm, - /* const int64_t batch_stride_a = */ A_batch_stride, - /* const int64_t batch_stride_b = */ B_batch_stride, - /* const int64_t batch_stride_d = */ matrix_stride_out, - /* const int swizzle_log = */ swizzle_log, - /* const int gemm_k_iterations_aligned = */ (K / bk), - /* const int batch_ndim = */ int(batch_shape.size())}; + GEMMParams params{/* const int M = */ M, + /* const int N = */ N, + /* const int K = */ K, + /* const int lda = */ lda, + /* const int ldb = */ ldb, + /* const int ldd = */ ldd, + /* const int tiles_n = */ tn, + /* const int tiles_m = */ tm, + /* const int64_t batch_stride_a = */ A_batch_stride, + /* const int64_t batch_stride_b = */ B_batch_stride, + /* const int64_t batch_stride_d = */ matrix_stride_out, + /* const int swizzle_log = */ swizzle_log, + /* const int gemm_k_iterations_aligned = */ (K / bk), + /* const int batch_ndim = */ int(batch_shape.size())}; // Prepare launch grid params int tile = 1 << swizzle_log; @@ -497,12 +494,11 @@ void steel_matmul_regular_axpby( int ldc = c.strides()[c.ndim() - 2]; int fdc = c.strides()[c.ndim() - 1]; - GEMMAddMMParams params{ - /* const int ldc = */ ldc, - /* const int fdc = */ fdc, - /* const int64_t batch_stride_c = */ C_batch_stride, - /* const float alpha = */ alpha, - /* const float beta = */ beta}; + GEMMAddMMParams params{/* const int ldc = */ ldc, + /* const int fdc = */ fdc, + /* const int64_t batch_stride_c = */ C_batch_stride, + /* const float alpha = */ alpha, + /* const float beta = */ beta}; compute_encoder.set_input_array(c, 2); compute_encoder.set_bytes(params, 5); @@ -1554,21 +1550,20 @@ void BlockMaskedMM::eval_gpu(const std::vector& inputs, array& out) { int swizzle_log = 0; // tm >= 6 ? 3 : (tm <= 3 ? 0 : 2); // Prepare steel matmul params - GEMMParams params{ - /* const int M = */ M, - /* const int N = */ N, - /* const int K = */ K, - /* const int lda = */ lda, - /* const int ldb = */ ldb, - /* const int ldd = */ N, - /* const int tiles_n = */ tn, - /* const int tiles_m = */ tm, - /* const int64_t batch_stride_a = */ A_batch_str, - /* const int64_t batch_stride_b = */ B_batch_str, - /* const int64_t batch_stride_d = */ matrix_stride_out, - /* const int swizzle_log = */ swizzle_log, - /* const int gemm_k_iterations_aligned = */ (K / bk), - /* const int batch_ndim = */ int(batch_shape.size())}; + GEMMParams params{/* const int M = */ M, + /* const int N = */ N, + /* const int K = */ K, + /* const int lda = */ lda, + /* const int ldb = */ ldb, + /* const int ldd = */ N, + /* const int tiles_n = */ tn, + /* const int tiles_m = */ tm, + /* const int64_t batch_stride_a = */ A_batch_str, + /* const int64_t batch_stride_b = */ B_batch_str, + /* const int64_t batch_stride_d = */ matrix_stride_out, + /* const int swizzle_log = */ swizzle_log, + /* const int gemm_k_iterations_aligned = */ (K / bk), + /* const int batch_ndim = */ int(batch_shape.size())}; // Prepare launch grid params int tile = 1 << swizzle_log; @@ -2113,23 +2108,22 @@ void gather_mm( compute_encoder.set_compute_pipeline_state(kernel); // Prepare the matmul params - steel::GEMMParams params{ - /* const int M = */ M, - /* const int N = */ N, - /* const int K = */ K, - /* const int lda = */ static_cast(lda), - /* const int ldb = */ static_cast(ldb), - /* const int ldd = */ N, - /* const int tiles_n = */ (N + bn - 1) / bn, - /* const int tiles_m = */ (M + bm - 1) / bm, - /* const int64_t batch_stride_a = */ - (batch_ndim > 0) ? lhs_indices.strides()[0] : 0, - /* const int64_t batch_stride_b = */ - (batch_ndim > 0) ? rhs_indices.strides()[0] : 0, - /* const int64_t batch_stride_d = */ M * N, - /* const int swizzle_log = */ 0, - /* const int gemm_k_iterations_aligned = */ (K / bk), - /* const int batch_ndim = */ batch_ndim}; + steel::GEMMParams params{/* const int M = */ M, + /* const int N = */ N, + /* const int K = */ K, + /* const int lda = */ static_cast(lda), + /* const int ldb = */ static_cast(ldb), + /* const int ldd = */ N, + /* const int tiles_n = */ (N + bn - 1) / bn, + /* const int tiles_m = */ (M + bm - 1) / bm, + /* const int64_t batch_stride_a = */ + (batch_ndim > 0) ? lhs_indices.strides()[0] : 0, + /* const int64_t batch_stride_b = */ + (batch_ndim > 0) ? rhs_indices.strides()[0] : 0, + /* const int64_t batch_stride_d = */ M * N, + /* const int swizzle_log = */ 0, + /* const int gemm_k_iterations_aligned = */ (K / bk), + /* const int batch_ndim = */ batch_ndim}; // Prepare the grid MTL::Size group_dims = MTL::Size(32, wn, wm); @@ -2317,21 +2311,20 @@ void segmented_mm( compute_encoder.set_compute_pipeline_state(kernel); // Prepare the matmul params - steel::GEMMParams params{ - /* const int M = */ M, - /* const int N = */ N, - /* const int K = */ K, - /* const int lda = */ static_cast(lda), - /* const int ldb = */ static_cast(ldb), - /* const int ldd = */ N, - /* const int tiles_n = */ (N + bn - 1) / bn, - /* const int tiles_m = */ (M + bm - 1) / bm, - /* const int64_t batch_stride_a = */ 0, - /* const int64_t batch_stride_b = */ 0, - /* const int64_t batch_stride_d = */ M * N, - /* const int swizzle_log = */ 0, - /* const int gemm_k_iterations_aligned = */ 0, - /* const int batch_ndim = */ 0}; + steel::GEMMParams params{/* const int M = */ M, + /* const int N = */ N, + /* const int K = */ K, + /* const int lda = */ static_cast(lda), + /* const int ldb = */ static_cast(ldb), + /* const int ldd = */ N, + /* const int tiles_n = */ (N + bn - 1) / bn, + /* const int tiles_m = */ (M + bm - 1) / bm, + /* const int64_t batch_stride_a = */ 0, + /* const int64_t batch_stride_b = */ 0, + /* const int64_t batch_stride_d = */ M * N, + /* const int swizzle_log = */ 0, + /* const int gemm_k_iterations_aligned = */ 0, + /* const int batch_ndim = */ 0}; // Prepare the grid MTL::Size group_dims = MTL::Size(32, wn, wm); diff --git a/mlx/distributed/jaccl/jaccl.cpp b/mlx/distributed/jaccl/jaccl.cpp index fd5b0d1d..a9d28f4d 100644 --- a/mlx/distributed/jaccl/jaccl.cpp +++ b/mlx/distributed/jaccl/jaccl.cpp @@ -468,11 +468,12 @@ class SideChannel { } } } else { - sockets_.push_back(detail::TCPSocket::connect( - IBV_TAG, address, 4, 1000, [](int attempt, int wait) { - std::cerr << IBV_TAG << " Connection attempt " << attempt - << " waiting " << wait << " ms" << std::endl; - })); + sockets_.push_back( + detail::TCPSocket::connect( + IBV_TAG, address, 4, 1000, [](int attempt, int wait) { + std::cerr << IBV_TAG << " Connection attempt " << attempt + << " waiting " << wait << " ms" << std::endl; + })); sockets_[0].send(IBV_TAG, reinterpret_cast(&rank_), sizeof(int)); } } diff --git a/mlx/distributed/ring/ring.cpp b/mlx/distributed/ring/ring.cpp index 9679d9ff..ea400428 100644 --- a/mlx/distributed/ring/ring.cpp +++ b/mlx/distributed/ring/ring.cpp @@ -353,23 +353,24 @@ std::vector make_connections( int success; for (auto& address : addresses) { - sockets.push_back(detail::TCPSocket::connect( - RING_TAG, - address, - CONN_ATTEMPTS, - CONN_WAIT, - [verbose](int attempt, int wait) { - log_info( - verbose, - "Attempt", - attempt, - "waiting", - wait, - "ms (error:", - errno, - ")"); - }) - .detach()); + sockets.push_back( + detail::TCPSocket::connect( + RING_TAG, + address, + CONN_ATTEMPTS, + CONN_WAIT, + [verbose](int attempt, int wait) { + log_info( + verbose, + "Attempt", + attempt, + "waiting", + wait, + "ms (error:", + errno, + ")"); + }) + .detach()); } return sockets; @@ -510,17 +511,18 @@ class RingGroup : public GroupImpl { std::vector> all_gathers; for (int i = 0; i < n_gathers; i++) { auto offset = i * bytes_per_gather; - all_gathers.emplace_back(pool_.enqueue(std::bind( - &RingGroup::all_gather_impl, - this, - input_ptr + offset, - output_ptr + offset, - nbytes, - offset + bytes_per_gather > nbytes ? nbytes - offset - : bytes_per_gather, - sockets_right_[i / 2], - sockets_left_[i / 2], - (i % 2) ? -1 : 1))); + all_gathers.emplace_back(pool_.enqueue( + std::bind( + &RingGroup::all_gather_impl, + this, + input_ptr + offset, + output_ptr + offset, + nbytes, + offset + bytes_per_gather > nbytes ? nbytes - offset + : bytes_per_gather, + sockets_right_[i / 2], + sockets_left_[i / 2], + (i % 2) ? -1 : 1))); } for (auto& f : all_gathers) { f.wait(); @@ -634,17 +636,18 @@ class RingGroup : public GroupImpl { std::vector> all_sums; for (int i = 0; i < n_reduces; i++) { - all_sums.emplace_back(pool_.enqueue(std::bind( - &RingGroup::all_reduce_impl, - this, - reinterpret_cast( - buffers_.data() + i * ALL_SUM_SIZE * ALL_SUM_BUFFERS), - reinterpret_cast(out_ptr) + i * step, - std::min(size, (i + 1) * step) - i * step, - sockets_right_[i / 2], - sockets_left_[i / 2], - (i % 2) ? -1 : 1, - reduce_op))); + all_sums.emplace_back(pool_.enqueue( + std::bind( + &RingGroup::all_reduce_impl, + this, + reinterpret_cast( + buffers_.data() + i * ALL_SUM_SIZE * ALL_SUM_BUFFERS), + reinterpret_cast(out_ptr) + i * step, + std::min(size, (i + 1) * step) - i * step, + sockets_right_[i / 2], + sockets_left_[i / 2], + (i % 2) ? -1 : 1, + reduce_op))); } for (auto& f : all_sums) { f.wait(); diff --git a/python/src/ops.cpp b/python/src/ops.cpp index 71604685..de51ad50 100644 --- a/python/src/ops.cpp +++ b/python/src/ops.cpp @@ -53,8 +53,9 @@ void init_ops(nb::module_& m) { "shape"_a, nb::kw_only(), "stream"_a = nb::none(), - nb::sig("def reshape(a: array, /, shape: Sequence[int], *, stream: " - "Union[None, Stream, Device] = None) -> array"), + nb::sig( + "def reshape(a: array, /, shape: Sequence[int], *, stream: " + "Union[None, Stream, Device] = None) -> array"), R"pbdoc( Reshape an array while preserving the size. @@ -80,8 +81,9 @@ void init_ops(nb::module_& m) { "end_axis"_a = -1, nb::kw_only(), "stream"_a = nb::none(), - nb::sig("def flatten(a: array, /, start_axis: int = 0, end_axis: int = " - "-1, *, stream: Union[None, Stream, Device] = None) -> array"), + nb::sig( + "def flatten(a: array, /, start_axis: int = 0, end_axis: int = " + "-1, *, stream: Union[None, Stream, Device] = None) -> array"), R"pbdoc( Flatten an array. @@ -182,8 +184,9 @@ void init_ops(nb::module_& m) { "axis"_a, nb::kw_only(), "stream"_a = nb::none(), - nb::sig("def expand_dims(a: array, /, axis: Union[int, Sequence[int]], " - "*, stream: Union[None, Stream, Device] = None) -> array"), + nb::sig( + "def expand_dims(a: array, /, axis: Union[int, Sequence[int]], " + "*, stream: Union[None, Stream, Device] = None) -> array"), R"pbdoc( Add a size one dimension at the given axis. @@ -1675,8 +1678,9 @@ void init_ops(nb::module_& m) { }, nb::arg(), "dtype"_a = nb::none(), - nb::sig("def asarray(a: Union[scalar, array, Sequence], dtype: " - "Optional[Dtype] = None) -> array"), + nb::sig( + "def asarray(a: Union[scalar, array, Sequence], dtype: " + "Optional[Dtype] = None) -> array"), R"pbdoc( Convert the input to an array. diff --git a/tests/linalg_tests.cpp b/tests/linalg_tests.cpp index 041aedaa..683e7584 100644 --- a/tests/linalg_tests.cpp +++ b/tests/linalg_tests.cpp @@ -394,12 +394,14 @@ TEST_CASE("test matrix cholesky") { linalg::cholesky(array({0.0, 1.0}), /* upper = */ false, Device::cpu)); // Unsupported types throw - CHECK_THROWS(linalg::cholesky( - array({0, 1}, {1, 2}), /* upper = */ false, Device::cpu)); + CHECK_THROWS( + linalg::cholesky( + array({0, 1}, {1, 2}), /* upper = */ false, Device::cpu)); // Non-square throws. - CHECK_THROWS(linalg::cholesky( - array({1, 2, 3, 4, 5, 6}, {2, 3}), /* upper = */ false, Device::cpu)); + CHECK_THROWS( + linalg::cholesky( + array({1, 2, 3, 4, 5, 6}, {2, 3}), /* upper = */ false, Device::cpu)); const auto prng_key = random::key(220398); const auto sqrtA = random::normal({5, 5}, prng_key);