diff --git a/metal/src/kernels/conv.rs b/metal/src/kernels/conv.rs deleted file mode 100644 index 5e18c53f4a..0000000000 --- a/metal/src/kernels/conv.rs +++ /dev/null @@ -1,114 +0,0 @@ -use crate::encoder::EncoderExt; -use crate::{LibraryName, MetalStream}; -use metal::MTLSize; -use tract_core::internal::*; -use tract_core::ops::cnn::Conv; -use tract_gpu::tensor::DeviceTensor; - -pub fn kernel_name(hw_rank: usize, dt: DatumType) -> TractResult { - let dt_name = if dt == DatumType::F16 { "f16" } else { "f32" }; - Ok(format!("conv{hw_rank}d_{dt_name}_generic")) -} - -pub fn metal_conv_dispatch( - stream: &MetalStream, - op: &Conv, - input: &DeviceTensor, - weights: &DeviceTensor, - bias: Option<&DeviceTensor>, - output: &DeviceTensor, -) -> TractResult<()> { - stream.retain_tensor(input); - stream.retain_tensor(weights); - if let Some(b) = bias { - stream.retain_tensor(b); - } - stream.retain_tensor(output); - - let input_shape = op.pool_spec.data_format.shape(input.shape())?; - let hw_rank = input_shape.hw_rank(); - let func_name = kernel_name(hw_rank, input.datum_type())?; - let pipeline = stream.load_pipeline(LibraryName::ConvOps, &func_name)?; - - let co_per_group = op.pool_spec.output_channels / op.group; - let ci_per_group = op.pool_spec.input_channels / op.group; - - // in_shape: [N, C, spatial...] - let in_n = *input_shape.n().unwrap_or(&1); - let in_c = *input_shape.c(); - let mut in_shape_buf: TVec = tvec![in_n as i32, in_c as i32]; - in_shape_buf.extend(input_shape.hw_dims().iter().map(|&d| d as i32)); - - let mut in_strides_buf: TVec = - tvec![*input_shape.n_stride().unwrap_or(&0) as i32, *input_shape.c_stride() as i32]; - in_strides_buf.extend(input_shape.hw_strides().iter().map(|&s| s as i32)); - - // ker_params: [groups, co_per_group, ci_per_group, ker_spatial...] - let mut ker_params: TVec = - tvec![op.group as i32, co_per_group as i32, ci_per_group as i32]; - ker_params.extend(weights.shape()[2..].iter().map(|&d| d as i32)); - - // ker_strides: [g_stride, o_stride, i_stride, spatial...] - let group_stride = weights.strides()[0] as usize * co_per_group; - let mut ker_strides: TVec = tvec![group_stride as i32]; - ker_strides.extend(weights.strides().iter().map(|&s| s as i32)); - - // padding - let padding = op.pool_spec.computed_padding(input_shape.hw_dims()); - let pad_buf: TVec = padding.iter().map(|p| p.pad_before as i32).collect(); - - let strides = op.pool_spec.strides(); - let strides_buf: TVec = strides.iter().map(|&s| s as i32).collect(); - - let dilations = op.pool_spec.dilations(); - let dilations_buf: TVec = dilations.iter().map(|&d| d as i32).collect(); - - let output_shape = op.pool_spec.data_format.shape(output.shape())?; - let out_n = *output_shape.n().unwrap_or(&1); - let out_c = *output_shape.c(); - let mut out_shape_buf: TVec = tvec![out_n as i32, out_c as i32]; - out_shape_buf.extend(output_shape.hw_dims().iter().map(|&d| d as i32)); - - let mut out_strides_buf: TVec = - tvec![*output_shape.n_stride().unwrap_or(&0) as i32, *output_shape.c_stride() as i32]; - out_strides_buf.extend(output_shape.hw_strides().iter().map(|&s| s as i32)); - - // bias_stride: -1 means no bias, 0 means scalar broadcast, 1 means per-channel - let bias_stride: i32 = if let Some(b) = bias { if b.rank() == 0 { 0 } else { 1 } } else { -1 }; - - let spatial_out: usize = output_shape.hw_dims().iter().product(); - let threads_per_group = 32usize; - - let command_buffer = stream.command_buffer(); - command_buffer.encode(|encoder| { - encoder.set_compute_pipeline_state(&pipeline); - encoder.set_metal_tensor(0, input, metal::MTLResourceUsage::Read); - encoder.set_slice(1, &in_shape_buf); - encoder.set_slice(2, &in_strides_buf); - encoder.set_metal_tensor(3, weights, metal::MTLResourceUsage::Read); - encoder.set_slice(4, &ker_params); - encoder.set_slice(5, &ker_strides); - if let Some(b) = bias { - encoder.set_metal_tensor(6, b, metal::MTLResourceUsage::Read); - } else { - // Empty buffer — kernel checks bias_stride < 0 - encoder.set_bytes(6, 0, std::ptr::null()); - } - encoder.set_slice(7, &[bias_stride]); - encoder.set_slice(8, &pad_buf); - encoder.set_slice(9, &strides_buf); - encoder.set_slice(10, &dilations_buf); - encoder.set_metal_tensor(11, output, metal::MTLResourceUsage::Write); - encoder.set_slice(12, &out_shape_buf); - encoder.set_slice(13, &out_strides_buf); - - let grid_size = MTLSize { - width: spatial_out.div_ceil(threads_per_group) as _, - height: out_c as _, - depth: out_n as _, - }; - let group_size = MTLSize { width: threads_per_group as _, height: 1, depth: 1 }; - encoder.dispatch_thread_groups(grid_size, group_size); - }); - Ok(()) -} diff --git a/metal/src/kernels/conv.metal b/metal/src/kernels/conv/direct_conv.metal similarity index 100% rename from metal/src/kernels/conv.metal rename to metal/src/kernels/conv/direct_conv.metal diff --git a/metal/src/kernels/conv/mlx_conv.metal b/metal/src/kernels/conv/mlx_conv.metal new file mode 100644 index 0000000000..cf03fc6204 --- /dev/null +++ b/metal/src/kernels/conv/mlx_conv.metal @@ -0,0 +1,4595 @@ +// Implicit-GEMM 2D convolution, ported from Apple MLX +// (https://github.com/ml-explore/mlx) @ fb5133e1049cc1482a330cd3a7135fda62a74b0d, +// MIT License, Copyright (c) 2023-2025 Apple Inc. +// +// Mechanical flattening of the MLX include closure for +// steel/conv/kernels/steel_conv_general.metal (each section verbatim, +// `#pragma once` and intra-repo includes stripped), minus the bf16 +// instantiation, since tract has no bf16 datum type. Resync against upstream +// by re-flattening; do not hand-edit the sections. +// +// clang-format off + +// ===== mlx/backend/metal/kernels/steel/conv/kernels/steel_conv_general.metal ===== +// Copyright © 2024 Apple Inc. + +#include + +// clang-format off + +// ===== mlx/backend/metal/kernels/utils.h ===== +// Copyright © 2023-2024 Apple Inc. + + +#include + + +// ===== mlx/backend/metal/kernels/bf16.h ===== +// Copyright © 2023 Apple Inc. + + +#include + +using namespace metal; + +typedef bfloat bfloat16_t; +inline uint16_t bfloat16_to_uint16(const bfloat16_t x) { + return as_type(x); +} + +inline bfloat16_t uint16_to_bfloat16(const uint16_t x) { + return as_type(x); +} + + +// ===== mlx/backend/metal/kernels/bf16_math.h ===== +// Copyright © 2023 Apple Inc. + + +/////////////////////////////////////////////////////////////////////////////// +// Metal math for bfloat16 +/////////////////////////////////////////////////////////////////////////////// + +/* + +Following the Metal Shading Language Specification (Metal 3.1) + +"bfloat is an extended itypeing point type that only allows implicit conversion + to a type of greater itypeing point rank. While bfloat can be implicitly + converted to itype, it cannot be implicitly converted to half, and neither + itype nor half can be implicitly converted to bfloat." + +Further, as far as I can tell, the stdlib math/simd functions are not defined +for bfloat and calling with an argument of type bfloat will result in that +argument getting implicitly converted to itype which then returns an output +that is (likely) a itype which cannot be implicitly converted into a bfloat + +This leads to situations where +bfloat a = 5.0bf; +bfloat b = metal::abs(a); // this will throw an error since abs return itype +bfloat c = static_cast(metal::abs(a)); // this is fine + +For the moment, I will be adding overloaded instantiations of the math +functions to accordingly automatically handle the casting + +*/ + +#define instantiate_metal_math_funcs(itype, otype, ctype, mfast) \ + \ + METAL_FUNC otype abs(itype x) { \ + return static_cast(__metal_fabs(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype acos(itype x) { \ + return static_cast(__metal_acos(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype acosh(itype x) { \ + return static_cast(__metal_acosh(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype asin(itype x) { \ + return static_cast(__metal_asin(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype asinh(itype x) { \ + return static_cast(__metal_asinh(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype atan(itype y_over_x) { \ + return static_cast( \ + __metal_atan(static_cast(y_over_x), mfast)); \ + } \ + METAL_FUNC otype atan2(itype y, itype x) { \ + return static_cast( \ + __metal_atan2(static_cast(y), static_cast(x), mfast)); \ + } \ + METAL_FUNC otype atanh(itype x) { \ + return static_cast(__metal_atanh(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype ceil(itype x) { \ + return static_cast(__metal_ceil(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype cos(itype x) { \ + return static_cast(__metal_cos(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype cosh(itype x) { \ + return static_cast(__metal_cosh(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype cospi(itype x) { \ + return static_cast(__metal_cospi(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype divide(itype x, itype y) { \ + return static_cast( \ + __metal_divide(static_cast(x), static_cast(y), mfast)); \ + } \ + METAL_FUNC otype exp(itype x) { \ + return static_cast(__metal_exp(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype exp10(itype x) { \ + return static_cast(__metal_exp10(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype exp2(itype x) { \ + return static_cast(__metal_exp2(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype fabs(itype x) { \ + return static_cast(__metal_fabs(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype fdim(itype x, itype y) { \ + ctype t = static_cast(x - y); \ + return static_cast(select(t, ctype(0), t < ctype(0) || x == y)); \ + } \ + METAL_FUNC otype floor(itype x) { \ + return static_cast(__metal_floor(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype fma(itype x, itype y, itype z) { \ + return static_cast(__metal_fma( \ + static_cast(x), static_cast(y), static_cast(z))); \ + } \ + METAL_FUNC otype fmax(itype x, itype y) { \ + return static_cast( \ + __metal_fmax(static_cast(x), static_cast(y), mfast)); \ + } \ + METAL_FUNC otype fmax3(itype x, itype y, itype z) { \ + return static_cast(__metal_fmax3( \ + static_cast(x), \ + static_cast(y), \ + static_cast(z), \ + mfast)); \ + } \ + METAL_FUNC otype fmedian3(itype x, itype y, itype z) { \ + return static_cast(__metal_fmedian3( \ + static_cast(x), \ + static_cast(y), \ + static_cast(z), \ + mfast)); \ + } \ + METAL_FUNC otype fmin(itype x, itype y) { \ + return static_cast( \ + __metal_fmin(static_cast(x), static_cast(y), mfast)); \ + } \ + METAL_FUNC otype fmin3(itype x, itype y, itype z) { \ + return static_cast(__metal_fmin3( \ + static_cast(x), \ + static_cast(y), \ + static_cast(z), \ + mfast)); \ + } \ + METAL_FUNC otype fmod(itype x, itype y) { \ + return static_cast( \ + __metal_fmod(static_cast(x), static_cast(y), mfast)); \ + } \ + METAL_FUNC otype fract(itype x) { \ + return static_cast(__metal_fract(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype frexp(itype x, thread int& exp) { \ + return static_cast(__metal_frexp(static_cast(x), &exp)); \ + } \ + METAL_FUNC otype ldexp(itype x, int k) { \ + return static_cast(__metal_ldexp(static_cast(x), k, mfast)); \ + } \ + METAL_FUNC otype log(itype x) { \ + return static_cast(__metal_log(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype log10(itype x) { \ + return static_cast(__metal_log10(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype log2(itype x) { \ + return static_cast(__metal_log2(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype max(itype x, itype y) { \ + return static_cast( \ + __metal_fmax(static_cast(x), static_cast(y), mfast)); \ + } \ + METAL_FUNC otype max3(itype x, itype y, itype z) { \ + return static_cast(__metal_fmax3( \ + static_cast(x), \ + static_cast(y), \ + static_cast(z), \ + mfast)); \ + } \ + METAL_FUNC otype median3(itype x, itype y, itype z) { \ + return static_cast(__metal_fmedian3( \ + static_cast(x), \ + static_cast(y), \ + static_cast(z), \ + mfast)); \ + } \ + METAL_FUNC otype min(itype x, itype y) { \ + return static_cast( \ + __metal_fmin(static_cast(x), static_cast(y), mfast)); \ + } \ + METAL_FUNC otype min3(itype x, itype y, itype z) { \ + return static_cast(__metal_fmin3( \ + static_cast(x), \ + static_cast(y), \ + static_cast(z), \ + mfast)); \ + } \ + METAL_FUNC otype nextafter(itype x, itype y) { \ + return static_cast( \ + __metal_nextafter(static_cast(x), static_cast(y))); \ + } \ + METAL_FUNC otype pow(itype x, itype y) { \ + return static_cast( \ + __metal_pow(static_cast(x), static_cast(y), mfast)); \ + } \ + METAL_FUNC otype powr(itype x, itype y) { \ + return static_cast( \ + __metal_powr(static_cast(x), static_cast(y), mfast)); \ + } \ + METAL_FUNC otype rint(itype x) { \ + return static_cast(__metal_rint(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype round(itype x) { \ + return static_cast(__metal_round(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype rsqrt(itype x) { \ + return static_cast(__metal_rsqrt(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype sin(itype x) { \ + return static_cast(__metal_sin(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype sinh(itype x) { \ + return static_cast(__metal_sinh(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype sinpi(itype x) { \ + return static_cast(__metal_sinpi(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype sqrt(itype x) { \ + return static_cast(__metal_sqrt(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype tan(itype x) { \ + return static_cast(__metal_tan(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype tanh(itype x) { \ + return static_cast(__metal_tanh(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype tanpi(itype x) { \ + return static_cast(__metal_tanpi(static_cast(x), mfast)); \ + } \ + METAL_FUNC otype trunc(itype x) { \ + return static_cast(__metal_trunc(static_cast(x), mfast)); \ + } + +namespace metal { + +instantiate_metal_math_funcs( + bfloat16_t, + bfloat16_t, + float, + __METAL_MAYBE_FAST_MATH__); + +namespace fast { + +instantiate_metal_math_funcs( + bfloat16_t, + bfloat16_t, + float, + __METAL_FAST_MATH__); + +} // namespace fast + +namespace precise { + +instantiate_metal_math_funcs( + bfloat16_t, + bfloat16_t, + float, + __METAL_PRECISE_MATH__); + +} // namespace precise + +} // namespace metal + +/////////////////////////////////////////////////////////////////////////////// +// Metal simd for bfloat16 +/////////////////////////////////////////////////////////////////////////////// + +#define instantiate_metal_simd_comm_funcs( \ + itype, otype, ctype, itype_to_ctype, ctype_to_otype) \ + \ + METAL_FUNC otype simd_broadcast(itype data, ushort broadcast_lane_id) { \ + return ctype_to_otype( \ + __metal_simd_broadcast(itype_to_ctype(data), broadcast_lane_id)); \ + } \ + \ + METAL_FUNC otype simd_shuffle(itype data, ushort simd_lane_id) { \ + return ctype_to_otype( \ + __metal_simd_shuffle(itype_to_ctype(data), simd_lane_id)); \ + } \ + \ + METAL_FUNC otype simd_shuffle_and_fill_down( \ + itype data, itype filling_data, ushort delta, ushort modulo) { \ + return ctype_to_otype(__metal_simd_shuffle_and_fill_down( \ + itype_to_ctype(data), itype_to_ctype(filling_data), delta, modulo)); \ + } \ + \ + METAL_FUNC otype simd_shuffle_and_fill_down( \ + itype data, itype filling_data, ushort delta) { \ + return ctype_to_otype(__metal_simd_shuffle_and_fill_down( \ + itype_to_ctype(data), \ + itype_to_ctype(filling_data), \ + delta, \ + __metal_get_simdgroup_size(ushort()))); \ + } \ + \ + METAL_FUNC otype simd_shuffle_and_fill_up( \ + itype data, itype filling_data, ushort delta, ushort modulo) { \ + return ctype_to_otype(__metal_simd_shuffle_and_fill_up( \ + itype_to_ctype(data), itype_to_ctype(filling_data), delta, modulo)); \ + } \ + \ + METAL_FUNC otype simd_shuffle_and_fill_up( \ + itype data, itype filling_data, ushort delta) { \ + return ctype_to_otype(__metal_simd_shuffle_and_fill_up( \ + itype_to_ctype(data), \ + itype_to_ctype(filling_data), \ + delta, \ + __metal_get_simdgroup_size(ushort()))); \ + } \ + \ + METAL_FUNC otype simd_shuffle_down(itype data, ushort delta) { \ + return ctype_to_otype( \ + __metal_simd_shuffle_down(itype_to_ctype(data), delta)); \ + } \ + \ + METAL_FUNC otype simd_shuffle_rotate_down(itype data, ushort delta) { \ + return ctype_to_otype( \ + __metal_simd_shuffle_rotate_down(itype_to_ctype(data), delta)); \ + } \ + \ + METAL_FUNC otype simd_shuffle_rotate_up(itype data, ushort delta) { \ + return ctype_to_otype( \ + __metal_simd_shuffle_rotate_up(itype_to_ctype(data), delta)); \ + } \ + \ + METAL_FUNC otype simd_shuffle_up(itype data, ushort delta) { \ + return ctype_to_otype( \ + __metal_simd_shuffle_up(itype_to_ctype(data), delta)); \ + } \ + \ + METAL_FUNC otype simd_shuffle_xor(itype data, ushort mask) { \ + return ctype_to_otype( \ + __metal_simd_shuffle_xor(itype_to_ctype(data), mask)); \ + } + +#define instantiate_metal_simd_reduction_funcs(itype, otype, ctype) \ + \ + METAL_FUNC otype simd_max(itype data) { \ + return static_cast(__metal_simd_max(static_cast(data))); \ + } \ + \ + METAL_FUNC otype simd_min(itype data) { \ + return static_cast(__metal_simd_min(static_cast(data))); \ + } \ + \ + METAL_FUNC otype simd_prefix_exclusive_product(itype data) { \ + return static_cast( \ + __metal_simd_prefix_exclusive_product(static_cast(data))); \ + } \ + \ + METAL_FUNC otype simd_prefix_exclusive_sum(itype data) { \ + return static_cast( \ + __metal_simd_prefix_exclusive_sum(static_cast(data))); \ + } \ + \ + METAL_FUNC otype simd_prefix_inclusive_product(itype data) { \ + return static_cast( \ + __metal_simd_prefix_inclusive_product(static_cast(data))); \ + } \ + \ + METAL_FUNC otype simd_prefix_inclusive_sum(itype data) { \ + return static_cast( \ + __metal_simd_prefix_inclusive_sum(static_cast(data))); \ + } \ + \ + METAL_FUNC otype simd_product(itype data) { \ + return static_cast(__metal_simd_product(static_cast(data))); \ + } \ + \ + METAL_FUNC otype simd_sum(itype data) { \ + return static_cast(__metal_simd_sum(static_cast(data))); \ + } \ + \ + METAL_FUNC otype simd_xor(itype data) { \ + return static_cast(__metal_simd_xor(static_cast(data))); \ + } + +namespace metal { + +instantiate_metal_simd_comm_funcs( + bfloat16_t, + bfloat16_t, + uint16_t, + bfloat16_to_uint16, + uint16_to_bfloat16); +instantiate_metal_simd_reduction_funcs(bfloat16_t, bfloat16_t, float); + +} // namespace metal + + +// ===== mlx/backend/metal/kernels/complex.h ===== +// Copyright © 2023 Apple Inc. + + +#include + +using namespace metal; + +struct complex64_t; + +template +static constexpr constant bool can_convert_to_complex64 = + !is_same_v && is_convertible_v; + +template +static constexpr constant bool can_convert_from_complex64 = + !is_same_v && + (is_convertible_v || is_convertible_v); + +struct complex64_t { + float real; + float imag; + + // Constructors + constexpr complex64_t(float real, float imag) : real(real), imag(imag) {}; + constexpr complex64_t() : real(0), imag(0) {}; + constexpr complex64_t() threadgroup : real(0), imag(0) {}; + + // Conversions to complex64_t + template < + typename T, + typename = typename enable_if>::type> + constexpr complex64_t(T x) thread : real(x), imag(0) {} + + template < + typename T, + typename = typename enable_if>::type> + constexpr complex64_t(T x) threadgroup : real(x), imag(0) {} + + template < + typename T, + typename = typename enable_if>::type> + constexpr complex64_t(T x) device : real(x), imag(0) {} + + template < + typename T, + typename = typename enable_if>::type> + constexpr complex64_t(T x) constant : real(x), imag(0) {} + + // Conversions from complex64_t + template < + typename T, + typename = typename enable_if>::type> + constexpr operator T() const thread { + return static_cast(real); + } + + template < + typename T, + typename = typename enable_if>::type> + constexpr operator T() const threadgroup { + return static_cast(real); + } + + template < + typename T, + typename = typename enable_if>::type> + constexpr operator T() const device { + return static_cast(real); + } + + template < + typename T, + typename = typename enable_if>::type> + constexpr operator T() const constant { + return static_cast(real); + } +}; + +constexpr complex64_t operator-(complex64_t x) { + return {-x.real, -x.imag}; +} + +constexpr bool operator>=(complex64_t a, complex64_t b) { + return (a.real > b.real) || (a.real == b.real && a.imag >= b.imag); +} + +constexpr bool operator>(complex64_t a, complex64_t b) { + return (a.real > b.real) || (a.real == b.real && a.imag > b.imag); +} + +constexpr bool operator<=(complex64_t a, complex64_t b) { + return operator>=(b, a); +} + +constexpr bool operator<(complex64_t a, complex64_t b) { + return operator>(b, a); +} + +constexpr bool operator==(complex64_t a, complex64_t b) { + return a.real == b.real && a.imag == b.imag; +} + +constexpr complex64_t operator+(complex64_t a, complex64_t b) { + return {a.real + b.real, a.imag + b.imag}; +} + +constexpr thread complex64_t& operator+=(thread complex64_t& a, complex64_t b) { + a.real += b.real; + a.imag += b.imag; + return a; +} + +constexpr threadgroup complex64_t& operator+=( + threadgroup complex64_t& a, + complex64_t b) { + a.real += b.real; + a.imag += b.imag; + return a; +} + +constexpr device complex64_t& operator+=(device complex64_t& a, complex64_t b) { + a.real += b.real; + a.imag += b.imag; + return a; +} + +constexpr complex64_t operator+(float a, complex64_t b) { + return {a + b.real, b.imag}; +} +constexpr complex64_t operator+(complex64_t a, float b) { + return {a.real + b, a.imag}; +} + +constexpr complex64_t operator-(complex64_t a, complex64_t b) { + return {a.real - b.real, a.imag - b.imag}; +} +constexpr complex64_t operator-(float a, complex64_t b) { + return {a - b.real, -b.imag}; +} +constexpr complex64_t operator-(complex64_t a, float b) { + return {a.real - b, a.imag}; +} + +constexpr complex64_t operator*(complex64_t a, complex64_t b) { + return {a.real * b.real - a.imag * b.imag, a.real * b.imag + a.imag * b.real}; +} + +constexpr complex64_t operator/(complex64_t a, complex64_t b) { + auto denom = b.real * b.real + b.imag * b.imag; + auto x = a.real * b.real + a.imag * b.imag; + auto y = a.imag * b.real - a.real * b.imag; + return {x / denom, y / denom}; +} + +constexpr complex64_t operator/(float a, complex64_t b) { + auto denom = b.real * b.real + b.imag * b.imag; + auto x = a * b.real; + auto y = -a * b.imag; + return {x / denom, y / denom}; +} + +constexpr complex64_t operator%(complex64_t a, complex64_t b) { + auto real = a.real - (b.real * static_cast(a.real / b.real)); + auto imag = a.imag - (b.imag * static_cast(a.imag / b.imag)); + if (real != 0 && (real < 0 != b.real < 0)) { + real += b.real; + } + if (imag != 0 && (imag < 0 != b.imag < 0)) { + imag += b.imag; + } + return {real, imag}; +} + + +// ===== mlx/backend/metal/kernels/defines.h ===== +// Copyright © 2023 Apple Inc. + + +#if defined __METAL__ || defined MLX_METAL_JIT +#define MTL_CONST constant +#else +#define MTL_CONST +#endif + +static MTL_CONST constexpr int MAX_REDUCE_SPECIALIZED_DIMS = 4; +static MTL_CONST constexpr int REDUCE_N_READS = 4; +static MTL_CONST constexpr int REDUCE_N_WRITES = 4; +static MTL_CONST constexpr int SOFTMAX_N_READS = 4; +static MTL_CONST constexpr int RMS_N_READS = 4; +static MTL_CONST constexpr int RMS_LOOPED_LIMIT = 4096; + +// Instantiate a templated kernel. +// Extra args are used as template parameters: +// e.g. instantiate_kernel(binary_int, binary, a, b) -> +// [[host_name(binary_int)]] [kernel] binary +#define instantiate_kernel(name, func, ...) \ + template [[host_name( \ + name)]] [[kernel]] decltype(func<__VA_ARGS__>) func<__VA_ARGS__>; + + +// ===== mlx/backend/metal/kernels/logging.h ===== +// Copyright © 2025 Apple Inc. + + +#if defined(__METAL_VERSION__) && (__METAL_VERSION__ >= 320) +#include + +namespace mlx { +using os_log = metal::os_log; +} // namespace mlx + +#else + +namespace mlx { +struct os_log { + constexpr os_log(constant char*, constant char*) constant {} + + template + void log_debug(constant char*, Args...) const {} + + template + void log_debug(constant char*, Args...) const constant {} +}; +} // namespace mlx + +#endif + +typedef half float16_t; + +// Work per thread values for different types. The values here are expected to +// match get_work_per_thread in mlx/backend/metal/utils.h +template +struct WorkPerThread { + static_assert(sizeof(U) <= 8, "Type too large"); + static constexpr int constant n = 8 / sizeof(U); +}; + +/////////////////////////////////////////////////////////////////////////////// +// Type limits utils +/////////////////////////////////////////////////////////////////////////////// + +template +struct Limits { + static const constant U max = metal::numeric_limits::max(); + static const constant U min = metal::numeric_limits::min(); + static const constant U finite_max = metal::numeric_limits::max(); + static const constant U finite_min = metal::numeric_limits::min(); +}; + +#define instantiate_default_limit(type) \ + template <> \ + struct Limits { \ + static constexpr constant type max = metal::numeric_limits::max(); \ + static constexpr constant type min = metal::numeric_limits::min(); \ + static constexpr constant type finite_max = \ + metal::numeric_limits::max(); \ + static constexpr constant type finite_min = \ + metal::numeric_limits::min(); \ + }; + +instantiate_default_limit(uint8_t); +instantiate_default_limit(uint16_t); +instantiate_default_limit(uint32_t); +instantiate_default_limit(uint64_t); +instantiate_default_limit(int8_t); +instantiate_default_limit(int16_t); +instantiate_default_limit(int32_t); +instantiate_default_limit(int64_t); + +#define instantiate_float_limit(type) \ + template <> \ + struct Limits { \ + static constexpr constant type max = \ + metal::numeric_limits::infinity(); \ + static constexpr constant type min = \ + -metal::numeric_limits::infinity(); \ + static constexpr constant type finite_max = \ + metal::numeric_limits::max(); \ + static constexpr constant type finite_min = \ + -metal::numeric_limits::max(); \ + }; + +instantiate_float_limit(half); +instantiate_float_limit(float); +instantiate_float_limit(bfloat16_t); + +template <> +struct Limits { + static constexpr constant bool max = true; + static constexpr constant bool min = false; +}; + +template <> +struct Limits { + static constexpr constant complex64_t max = complex64_t( + metal::numeric_limits::infinity(), + metal::numeric_limits::infinity()); + static constexpr constant complex64_t min = complex64_t( + -metal::numeric_limits::infinity(), + -metal::numeric_limits::infinity()); +}; + +/////////////////////////////////////////////////////////////////////////////// +// Indexing utils +/////////////////////////////////////////////////////////////////////////////// + +#define MLX_MTL_PRAGMA_UNROLL _Pragma("clang loop unroll(full)") + +/////////////////////////////////////////////////////////////////////////////// +// Single Array with generic dims + +template +METAL_FUNC IdxT elem_to_loc( + IdxT elem, + constant const int* shape, + constant const int64_t* strides, + int ndim) { + IdxT loc = 0; + for (int i = ndim - 1; i >= 0 && elem > 0; --i) { + loc += (elem % shape[i]) * IdxT(strides[i]); + elem /= shape[i]; + } + return loc; +} + +// Non templated version to handle arbitrary dims +template +METAL_FUNC IdxT elem_to_loc( + uint3 elem, + constant const int* shape, + constant const int64_t* strides, + int ndim) { + IdxT loc = + elem.x * IdxT(strides[ndim - 1]) + elem.y * IdxT(strides[ndim - 2]); + for (int d = ndim - 3; d >= 0; --d) { + loc += (elem.z % shape[d]) * IdxT(strides[d]); + elem.z /= shape[d]; + } + return loc; +} + +/////////////////////////////////////////////////////////////////////////////// +// Single Array with fixed N dims + +template +METAL_FUNC IdxT elem_to_loc_1(uint elem, constant const int64_t& stride) { + return elem * IdxT(stride); +} + +template +METAL_FUNC IdxT elem_to_loc_2(uint2 elem, constant const int64_t strides[2]) { + return elem.x * IdxT(strides[1]) + elem.y * IdxT(strides[0]); +} + +template +METAL_FUNC IdxT elem_to_loc_3(uint3 elem, constant const int64_t strides[3]) { + return elem.x * IdxT(strides[2]) + elem.y * IdxT(strides[1]) + + elem.z * IdxT(strides[0]); +} + +/////////////////////////////////////////////////////////////////////////////// +// Multiple Arrays with generic dims + +template +METAL_FUNC vec elem_to_loc_2_nd( + uint3 elem, + constant const int* shape, + constant const int64_t* a_strides, + constant const int64_t* b_strides, + int ndim) { + vec loc = { + IdxT( + elem.x * IdxT(a_strides[ndim - 1]) + + IdxT(elem.y) * IdxT(a_strides[ndim - 2])), + IdxT( + elem.x * IdxT(b_strides[ndim - 1]) + + elem.y * IdxT(b_strides[ndim - 2]))}; + for (int d = ndim - 3; d >= 0; --d) { + uint l = elem.z % shape[d]; + loc.x += l * IdxT(a_strides[d]); + loc.y += l * IdxT(b_strides[d]); + elem.z /= shape[d]; + } + return loc; +} + +template +METAL_FUNC vec elem_to_loc_3_nd( + uint3 elem, + constant const int* shape, + constant const int64_t* a_strides, + constant const int64_t* b_strides, + constant const int64_t* c_strides, + int ndim) { + vec loc = { + IdxT(elem.x * IdxT(a_strides[ndim - 1])) + + IdxT(elem.y * IdxT(a_strides[ndim - 2])), + IdxT(elem.x * IdxT(b_strides[ndim - 1])) + + IdxT(elem.y * IdxT(b_strides[ndim - 2])), + IdxT(elem.x * IdxT(c_strides[ndim - 1])) + + IdxT(elem.y * IdxT(c_strides[ndim - 2]))}; + for (int d = ndim - 3; d >= 0; --d) { + uint l = elem.z % shape[d]; + loc.x += l * IdxT(a_strides[d]); + loc.y += l * IdxT(b_strides[d]); + loc.z += l * IdxT(c_strides[d]); + elem.z /= shape[d]; + } + return loc; +} + +/////////////////////////////////////////////////////////////////////////////// +// Elem to loc in a loop utils +/////////////////////////////////////////////////////////////////////////////// + +template +struct LoopedElemToLoc { + int dim; + LoopedElemToLoc inner_looper; + OffsetT offset{0}; + int index{0}; + + LoopedElemToLoc(int dim) : dim(dim), inner_looper(dim - 1) {} + + void next(const constant int* shape, const constant int64_t* strides) { + if (dim == 0) { + return; + } + index++; + offset += OffsetT(strides[dim - 1]); + if (index >= shape[dim - 1]) { + index = 0; + inner_looper.next(shape, strides); + offset = inner_looper.offset; + } + } + + void next(int n, const constant int* shape, const constant int64_t* strides) { + if (dim == 0) { + return; + } + index += n; + offset += n * OffsetT(strides[dim - 1]); + + if (index >= shape[dim - 1]) { + int extra = index - shape[dim - 1]; + if (extra >= shape[dim - 1]) { + inner_looper.next(1 + extra / shape[dim - 1], shape, strides); + extra = extra % shape[dim - 1]; + } else { + inner_looper.next(shape, strides); + } + index = 0; + offset = inner_looper.offset; + if (extra > 0) { + next(extra, shape, strides); + } + } + } + + OffsetT location() { + return offset; + } +}; + +template +struct LoopedElemToLoc<1, OffsetT, true> { + int dim; + OffsetT offset{0}; + uint index{0}; + + LoopedElemToLoc(int dim) : dim(dim) {} + + void next(const constant int* shape, const constant int64_t* strides) { + index++; + if (dim > 1) { + offset = elem_to_loc(index, shape, strides, dim); + } else { + offset += OffsetT(strides[0]); + } + } + + void next(int n, const constant int* shape, const constant int64_t* strides) { + index += n; + if (dim > 1) { + offset = elem_to_loc(index, shape, strides, dim); + } else { + offset = index * OffsetT(strides[0]); + } + } + + OffsetT location() { + return offset; + } +}; + +template +struct LoopedElemToLoc<1, OffsetT, false> { + OffsetT offset{0}; + + LoopedElemToLoc(int) {} + + void next(const constant int*, const constant int64_t* strides) { + offset += OffsetT(strides[0]); + } + + void next(int n, const constant int*, const constant int64_t* strides) { + offset += n * OffsetT(strides[0]); + } + + OffsetT location() { + return offset; + } +}; + +/////////////////////////////////////////////////////////////////////////////// +// Calculation utils +/////////////////////////////////////////////////////////////////////////////// + +/** Compute ceil((float)N/(float)M) */ +template +inline T ceildiv(T N, U M) { + return (N + M - 1) / M; +} + +// https://docs.oracle.com/cd/E19957-01/806-3568/ncg_goldberg.html#1202 +inline float log1p(float x) { + float xp1 = 1.0f + x; + if (xp1 == Limits::max) { + return Limits::max; + } + if (xp1 == 1.0f) { + return x; + } + + return x * (metal::log(xp1) / (xp1 - 1.0f)); +} + +inline bfloat16_t log1p(bfloat16_t x) { + float xp1 = 1.0f + static_cast(x); + if (xp1 == Limits::max) { + return Limits::max; + } + if (xp1 == 1.0f) { + return x; + } + + return bfloat16_t(x * (metal::log(xp1) / (xp1 - 1.0f))); +} + +inline complex64_t log1p(complex64_t in) { + float x = in.real; + float y = in.imag; + float zabs = metal::precise::sqrt(x * x + y * y); + float theta = metal::atan2(y, x + 1); + if (zabs < 0.5f) { + float r = x * (2 + x) + y * y; + if (r == 0) { // handle underflow + return {x, theta}; + } + return {0.5f * log1p(r), theta}; + } else { + auto z0 = metal::sqrt((x + 1) * (x + 1) + y * y); + return {metal::log(z0), theta}; + } +} + +/////////////////////////////////////////////////////////////////////////////// +// SIMD shuffle ops +/////////////////////////////////////////////////////////////////////////////// + +inline uint64_t simd_shuffle_down(uint64_t data, uint16_t delta) { + return as_type( + metal::simd_shuffle_down(as_type(data), delta)); +} + +inline int64_t simd_shuffle_down(int64_t data, uint16_t delta) { + return as_type( + metal::simd_shuffle_down(as_type(data), delta)); +} + +inline bool simd_shuffle_down(bool data, uint16_t delta) { + return simd_shuffle_down(static_cast(data), delta); +} + +inline complex64_t simd_shuffle_down(complex64_t data, uint16_t delta) { + return complex64_t( + simd_shuffle_down(data.real, delta), simd_shuffle_down(data.imag, delta)); +} + +inline uint64_t simd_shuffle_up(uint64_t data, uint16_t delta) { + return as_type(metal::simd_shuffle_up(as_type(data), delta)); +} + +inline int64_t simd_shuffle_up(int64_t data, uint16_t delta) { + return as_type(metal::simd_shuffle_up(as_type(data), delta)); +} + +inline bool simd_shuffle_up(bool data, uint16_t delta) { + return simd_shuffle_up(static_cast(data), delta); +} + +inline complex64_t simd_shuffle_up(complex64_t data, uint16_t delta) { + return complex64_t( + simd_shuffle_up(data.real, delta), simd_shuffle_up(data.imag, delta)); +} + +inline uint64_t +simd_shuffle_and_fill_up(uint64_t data, uint64_t filling, uint16_t delta) { + return as_type(metal::simd_shuffle_and_fill_up( + as_type(data), as_type(filling), delta)); +} + +inline int64_t +simd_shuffle_and_fill_up(int64_t data, int64_t filling, uint16_t delta) { + return as_type(metal::simd_shuffle_and_fill_up( + as_type(data), as_type(filling), delta)); +} + +inline bool simd_shuffle_and_fill_up(bool data, bool filling, uint16_t delta) { + return simd_shuffle_and_fill_up( + static_cast(data), static_cast(filling), delta); +} + +inline complex64_t simd_shuffle_and_fill_up( + complex64_t data, + complex64_t filling, + uint16_t delta) { + return complex64_t( + simd_shuffle_and_fill_up(data.real, filling.real, delta), + simd_shuffle_and_fill_up(data.imag, filling.imag, delta)); +} + +inline uint64_t simd_shuffle(uint64_t data, uint16_t lane) { + return as_type(metal::simd_shuffle(as_type(data), lane)); +} + +inline int64_t simd_shuffle(int64_t data, uint16_t lane) { + return as_type(metal::simd_shuffle(as_type(data), lane)); +} + +inline bool simd_shuffle(bool data, uint16_t lane) { + return simd_shuffle(static_cast(data), lane); +} + +inline complex64_t simd_shuffle(complex64_t data, uint16_t lane) { + return complex64_t( + simd_shuffle(data.real, lane), simd_shuffle(data.imag, lane)); +} + +// std::conditional is not included with Metal +template +struct ConditionalType { + using type = U; +}; + +template +struct ConditionalType { + using type = T; +}; + + +// ===== mlx/backend/metal/kernels/steel/gemm/mma.h ===== +// Copyright © 2024 Apple Inc. + + +#include +#include +#include + + +// ===== mlx/backend/metal/kernels/steel/defines.h ===== +// Copyright © 2024 Apple Inc. + + +#define STEEL_CONST static constant constexpr const +#define STEEL_PRAGMA_UNROLL _Pragma("clang loop unroll(full)") +#define STEEL_PRAGMA_NO_UNROLL _Pragma("clang loop unroll(disable)") + + +// ===== mlx/backend/metal/kernels/steel/gemm/transforms.h ===== +// Copyright © 2024 Apple Inc. + + + +// ===== mlx/backend/metal/kernels/steel/utils.h ===== +// Copyright © 2024 Apple Inc. + + +#include + +METAL_FUNC ulong2 elem_to_loc_broadcast( + uint elem, + constant const int* shape, + constant const int64_t* a_strides, + constant const int64_t* b_strides, + int ndim) { + ulong loc_a{0}; + ulong loc_b{0}; + for (int i = ndim - 1; i >= 0 && elem > 0; --i) { + int pos_in_dim = (elem % shape[i]); + elem /= shape[i]; + loc_a += pos_in_dim * a_strides[i]; + loc_b += pos_in_dim * b_strides[i]; + } + return ulong2(loc_a, loc_b); +} + +METAL_FUNC ulong3 elem_to_loc_broadcast( + uint elem, + constant const int* shape, + constant const int64_t* a_strides, + constant const int64_t* b_strides, + constant const int64_t* c_strides, + int ndim) { + ulong loc_a{0}; + ulong loc_b{0}; + ulong loc_c{0}; + for (int i = ndim - 1; i >= 0 && elem > 0; --i) { + int pos_in_dim = (elem % shape[i]); + elem /= shape[i]; + loc_a += pos_in_dim * a_strides[i]; + loc_b += pos_in_dim * b_strides[i]; + loc_c += pos_in_dim * c_strides[i]; + } + return ulong3(loc_a, loc_b, loc_c); +} + + +/////////////////////////////////////////////////////////////////////////////// +// Transforms and Epilogues +/////////////////////////////////////////////////////////////////////////////// + +namespace mlx { +namespace steel { + +template +struct TransformNone { + static METAL_FUNC OutT apply(InT x) { + return static_cast(x); + } + + static METAL_FUNC OutT apply(InT x, OutT) { + return static_cast(x); + } +}; + +template +struct TransformAdd { + TransformAdd(const float, const float) {} + + static METAL_FUNC OutT apply(InT x) { + return static_cast(x); + } + + static METAL_FUNC OutT apply(InT x, OutT c) { + return static_cast(x) + c; + } +}; + +template +struct TransformAxpby { + const float alpha; + const float beta; + + TransformAxpby(const float alpha_, const float beta_) + : alpha(alpha_), beta(beta_) {} + + static METAL_FUNC OutT apply(InT x) { + return static_cast(x); + } + + METAL_FUNC OutT apply(InT x, OutT c) const { + return static_cast( + x * static_cast(alpha) + (static_cast(beta) * c)); + } +}; + +template +struct AccumHelper { + typedef float accum_type; +}; + +struct BlockSwizzle { + static METAL_FUNC int2 + swizzle(uint3 tid [[threadgroup_position_in_grid]], const int swizzle_log) { + const int tid_x = (tid.x) >> swizzle_log; + const int tid_y = + ((tid.y) << swizzle_log) + ((tid.x) & ((1 << swizzle_log) - 1)); + return int2(tid_x, tid_y); + } +}; + +} // namespace steel +} // namespace mlx + +// ===== mlx/backend/metal/kernels/steel/utils/integral_constant.h ===== +// Copyright © 2024 Apple Inc. + + +#include + +// ===== mlx/backend/metal/kernels/steel/utils/type_traits.h ===== +// Copyright © 2024 Apple Inc. + + +#include + +#pragma METAL internals : enable + +namespace metal { + +template +struct is_empty : metal::bool_constant<__is_empty(T)> {}; + +#ifdef __cpp_variable_templates +template +constexpr constant bool is_empty_v = is_empty::value; +#endif + +template +struct make_void { + typedef void type; +}; + +template +using void_t = typename make_void::type; + +template +struct is_static : metal::bool_constant>::value> {}; + +template +struct pointer_element {}; + +template +struct pointer_element { + using type = remove_cv_t; +}; +template +struct pointer_element { + using type = remove_cv_t; +}; +template +struct pointer_element { + using type = remove_cv_t; +}; +template +struct pointer_element { + using type = remove_cv_t; +}; + +template +using pointer_element_t = typename pointer_element>::type; + +} // namespace metal + +#pragma METAL internals : disable + +#pragma METAL internals : enable + +namespace mlx { +namespace steel { + +/////////////////////////////////////////////////////////////////////////////// +// Integral constant with casting +/////////////////////////////////////////////////////////////////////////////// + +template +struct integral_constant { + static constexpr constant T value = v; + using value_type = T; + using type = integral_constant; + + METAL_FUNC constexpr operator value_type() const noexcept { + return value; + } +}; + +template +using bool_constant = integral_constant; +using true_type = bool_constant; +using false_type = bool_constant; + +template +struct is_integral : bool_constant::value> {}; + +template +struct is_integral> + : bool_constant::value> {}; + +template +constexpr constant bool is_integral_v = is_integral::value; + +template +using Int = integral_constant; + +/////////////////////////////////////////////////////////////////////////////// +// Binary Operators on Integral constants +/////////////////////////////////////////////////////////////////////////////// + +#define integral_const_binop(__op__, __operator__) \ + template \ + METAL_FUNC constexpr auto __operator__( \ + integral_constant, integral_constant) { \ + constexpr auto res = tv __op__ uv; \ + return integral_constant{}; \ + } + +integral_const_binop(+, operator+); +integral_const_binop(-, operator-); +integral_const_binop(*, operator*); +integral_const_binop(/, operator/); + +integral_const_binop(==, operator==); +integral_const_binop(!=, operator!=); +integral_const_binop(<, operator<); +integral_const_binop(>, operator>); +integral_const_binop(<=, operator<=); +integral_const_binop(>=, operator>=); + +integral_const_binop(&&, operator&&); +integral_const_binop(||, operator||); + +template >> +METAL_FUNC constexpr auto operator||(true_type, T) { + return true_type{}; +} +template >> +METAL_FUNC constexpr auto operator||(T, true_type) { + return true_type{}; +} + +template >> +METAL_FUNC constexpr auto operator&&(false_type, T) { + return false_type{}; +} + +template >> +METAL_FUNC constexpr auto operator&&(T, false_type) { + return false_type{}; +} + +// Dispatch utilities +template +void dispatch_bool(bool v, F f) { + if (v) { + f(true_type{}); + } else { + f(false_type{}); + } +} + +template +constexpr void const_for_loop(F f) { + if constexpr (start < stop) { + constexpr auto idx = Int{}; + f(idx); + const_for_loop(f); + } +} + +#undef integral_const_binop + +/////////////////////////////////////////////////////////////////////////////// +// Reduction operators +/////////////////////////////////////////////////////////////////////////////// + +template +METAL_FUNC constexpr T sum(T x) { + return x; +} + +template +METAL_FUNC constexpr auto sum(T x, Us... us) { + return x + sum(us...); +} + +} // namespace steel +} // namespace mlx + +#pragma METAL internals : disable + + +using namespace metal; + +/////////////////////////////////////////////////////////////////////////////// +// MMA helper +/////////////////////////////////////////////////////////////////////////////// + +namespace mlx { +namespace steel { + +template +struct BaseMMAFrag { + static_assert( + kFragRows_ == 8, + "Only 8 x 8 fragment matrices are currently supported"); + static_assert( + kFragCols_ == 8, + "Only 8 x 8 fragment matrices are currently supported"); +}; + +template +struct BaseMMAFrag { + STEEL_CONST int kFragRows = 8; + STEEL_CONST int kFragCols = 8; + + STEEL_CONST int kElemsPerFrag = (kFragRows * kFragCols) / 32; + + STEEL_CONST int kElemRows = 1; + STEEL_CONST int kElemCols = 2; + + static_assert( + kElemRows * kElemCols == kElemsPerFrag, + "MMAFrag shape is not consistent with MMAFrag size"); + + 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]]) { + 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; + return short2{fn, fm}; + } + + template + METAL_FUNC static constexpr void + load(thread frag_type& dst, SrcPtrType src, StrX str_x, StrY str_y) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemRows; i++) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kElemCols; j++) { + dst[i * kElemCols + j] = static_cast(src[i * str_x + j * str_y]); + } + } + } + + template < + typename SrcPtrType, + typename StrX, + typename StrY, + typename LimX, + typename LimY, + typename OffX, + typename OffY> + METAL_FUNC static constexpr void load_safe( + thread frag_type& dst, + SrcPtrType src, + StrX str_x, + StrY str_y, + LimX lim_x, + LimY lim_y, + OffX off_x = Int<0>{}, + OffY off_y = Int<0>{}) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemRows; i++) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kElemCols; j++) { + if ((off_x + i) < lim_x && (off_y + j) < lim_y) { + dst[i * kElemCols + j] = + static_cast(src[(off_x + i) * str_x + (off_y + j) * str_y]); + } else { + dst[i * kElemCols + j] = T(0); + } + } + } + } + + template + METAL_FUNC static constexpr void + store(const thread frag_type& src, DstPtrType dst, StrX str_x, StrY str_y) { + using U = pointer_element_t; + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemRows; i++) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kElemCols; j++) { + dst[i * str_x + j * str_y] = static_cast(src[i * kElemCols + j]); + } + } + } + + template < + typename DstPtrType, + typename StrX, + typename StrY, + typename LimX, + typename LimY, + typename OffX, + typename OffY> + METAL_FUNC static constexpr void store_safe( + const thread frag_type& src, + DstPtrType dst, + StrX str_x, + StrY str_y, + LimX lim_x, + LimY lim_y, + OffX off_x = Int<0>{}, + OffY off_y = Int<0>{}) { + using U = pointer_element_t; + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemRows; i++) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kElemCols; j++) { + if ((off_x + i) < lim_x && (off_y + j) < lim_y) { + dst[(off_x + i) * str_x + (off_y + j) * str_y] = + static_cast(src[i * kElemCols + j]); + } + } + } + } + + template < + typename DstPtrType, + typename StrX, + typename StrY, + typename StartX, + typename StopX, + typename StartY, + typename StopY, + typename OffX, + typename OffY> + METAL_FUNC static constexpr void store_slice( + const thread frag_type& src, + DstPtrType dst, + StrX str_x, + StrY str_y, + StartX start_x, + StopX stop_x, + StartY start_y, + StopY stop_y, + OffX off_x = Int<0>{}, + OffY off_y = Int<0>{}) { + using U = pointer_element_t; + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kElemRows; i++) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kElemCols; j++) { + if ((off_x + i) < stop_x && (off_x + i) >= start_x && + (off_y + j) < stop_y && (off_y + j) >= start_y) { + dst[(off_x + i) * str_x + (off_y + j) * str_y] = + static_cast(src[i * kElemCols + j]); + } + } + } + } + + METAL_FUNC static constexpr void mma( + thread frag_type& D, + thread frag_type& A, + thread frag_type& B, + thread frag_type& C) { + mat_type D_mat; + mat_type A_mat; + mat_type B_mat; + mat_type C_mat; + + reinterpret_cast(A_mat.thread_elements()) = A; + reinterpret_cast(B_mat.thread_elements()) = B; + reinterpret_cast(C_mat.thread_elements()) = C; + + mma(D_mat, A_mat, B_mat, C_mat); + + D = reinterpret_cast(D_mat.thread_elements()); + } + + METAL_FUNC static constexpr void mma( + thread mat_type& D, + thread mat_type& A, + thread mat_type& B, + thread mat_type& C) { + simdgroup_multiply_accumulate(D, A, B, C); + } +}; + +template < + typename T, + int kTileRows_, + int kTileCols_, + class MMAFrag_ = BaseMMAFrag> +struct MMATile { + using MMAFrag_t = MMAFrag_; + using elem_type = T; + STEEL_CONST int kFragRows = MMAFrag_t::kFragRows; + STEEL_CONST int kFragCols = MMAFrag_t::kFragCols; + STEEL_CONST int kElemsPerFrag = MMAFrag_t::kElemsPerFrag; + + STEEL_CONST int kTileRows = kTileRows_; + STEEL_CONST int kTileCols = kTileCols_; + + STEEL_CONST int kRows = kTileRows * kFragRows; + STEEL_CONST int kCols = kTileCols * kFragCols; + + STEEL_CONST int kNumFrags = kTileRows * kTileCols; + STEEL_CONST int kElemsPerTile = kNumFrags * kElemsPerFrag; + + typedef typename MMAFrag_t::mat_type mat_type; + typedef typename MMAFrag_t::frag_type frag_type; + + frag_type val_frags[kNumFrags] = {frag_type(0)}; + + METAL_FUNC MMATile() thread {} + + METAL_FUNC constexpr void clear() { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kNumFrags; ++i) { + val_frags[i] = frag_type(0); + } + } + + METAL_FUNC constexpr thread frag_type& frag_at(const short i, const short j) { + return val_frags[i * kTileCols + j]; + } + + METAL_FUNC constexpr const thread frag_type& frag_at( + const short i, + const short j) const { + return val_frags[i * kTileCols + j]; + } + + METAL_FUNC mat_type mat_at(const short i, const short j) { + mat_type val_mat; + STEEL_PRAGMA_UNROLL + for (short ii = 0; ii < kElemsPerFrag; ++ii) { + val_mat.thread_elements()[ii] = frag_at(i, j)[ii]; + } + return val_mat; + } + + METAL_FUNC thread elem_type* elems() { + return reinterpret_cast(val_frags); + } + + METAL_FUNC const thread elem_type* elems() const { + return reinterpret_cast(val_frags); + } + + template + METAL_FUNC void load(const threadgroup U* src) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kTileRows; ++i) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kTileCols; ++j) { + MMAFrag_t::load( + frag_at(i, j), + &( + src[(i * kFragRows) * w_x * str_x + + (j * kFragCols) * w_y * str_y]), + Int{}, + Int{}); + } + } + } + + template + METAL_FUNC void store(threadgroup U* dst) const { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kTileRows; ++i) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kTileCols; ++j) { + MMAFrag_t::store( + frag_at(i, j), + &( + dst[(i * kFragRows) * w_x * str_x + + (j * kFragCols) * w_y * str_y]), + Int{}, + Int{}); + } + } + } + + template + METAL_FUNC void load(const device U* src, const int ld) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kTileRows; ++i) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kTileCols; ++j) { + MMAFrag_t::load( + frag_at(i, j), + &(src[(i * kFragRows) * w_x * ld + (j * kFragCols) * w_y]), + ld, + Int<1>{}); + } + } + } + + template + METAL_FUNC void store(device U* dst, const int ld) const { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < kTileRows; ++i) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < kTileCols; ++j) { + MMAFrag_t::store( + frag_at(i, j), + &(dst[(i * kFragRows) * w_x * ld + (j * kFragCols) * w_y]), + ld, + Int<1>{}); + } + } + } + + template + METAL_FUNC void + load_safe(const device U* src, const int ld, const short2 src_tile_dims) { + STEEL_PRAGMA_UNROLL + for (int i = 0; i < kTileRows; ++i) { + STEEL_PRAGMA_UNROLL + for (int j = 0; j < kTileCols; ++j) { + MMAFrag_t::load_safe( + frag_at(i, j), + src, + ld, + Int<1>{}, + src_tile_dims.y, + src_tile_dims.x, + (i * kFragRows) * w_x, + (j * kFragCols) * w_y); + } + } + } + + template + METAL_FUNC void + store_safe(device U* dst, const int ld, const short2 dst_tile_dims) const { + STEEL_PRAGMA_UNROLL + for (int i = 0; i < kTileRows; ++i) { + STEEL_PRAGMA_UNROLL + for (int j = 0; j < kTileCols; ++j) { + MMAFrag_t::store_safe( + frag_at(i, j), + dst, + ld, + Int<1>{}, + dst_tile_dims.y, + dst_tile_dims.x, + (i * kFragRows) * w_x, + (j * kFragCols) * w_y); + } + } + } + + template + METAL_FUNC void store_slice( + device U* dst, + const int ld, + const short2 start, + const short2 stop) const { + STEEL_PRAGMA_UNROLL + for (int i = 0; i < kTileRows; ++i) { + STEEL_PRAGMA_UNROLL + for (int j = 0; j < kTileCols; ++j) { + MMAFrag_t::store_slice( + frag_at(i, j), + dst, + ld, + Int<1>{}, + start.y, + stop.y, + start.x, + stop.x, + (i * kFragRows) * w_x, + (j * kFragCols) * w_y); + } + } + } +}; + +template +METAL_FUNC void tile_matmad( + thread MMATile& D, + thread MMATile& A, + thread MMATile& B, + thread MMATile& C) { + STEEL_PRAGMA_UNROLL + for (short m = 0; m < M; ++m) { + STEEL_PRAGMA_UNROLL + for (short n = 0; n < N; ++n) { + short n_serp = (m % 2) ? (N - 1 - n) : n; + STEEL_PRAGMA_UNROLL + for (short k = 0; k < K; ++k) { + MMATile::MMAFrag_t::mma( + D.frag_at(m, n_serp), + A.frag_at(m, k), + B.frag_at(k, n_serp), + C.frag_at(m, n_serp)); + } + } + } +} + +template +struct TransformNone { + static METAL_FUNC complex64_t apply(complex64_t x) { + return x; + } + static METAL_FUNC complex64_t apply(complex64_t x, complex64_t) { + return x; + } +}; + +template < + typename T, + typename U, + int BM, + int BN, + int BK, + int WM, + int WN, + bool transpose_a, + bool transpose_b, + short lda_tgp, + short ldb_tgp, + typename AccumType = float, + typename Epilogue = TransformNone> +struct BlockMMA { + // MMAFrag size + STEEL_CONST short kFragSize = 8; + using MMAFrag_acc_t = BaseMMAFrag; + + // Warp tile simdgroup matrix strides along M + STEEL_CONST short TM_stride = kFragSize * WM; + // Warp tile simdgroup matrix strides along M + STEEL_CONST short TN_stride = kFragSize * WN; + + // Warp tile size along M + STEEL_CONST short TM = BM / (kFragSize * WM); + // Warp tile size along N + STEEL_CONST short TN = BN / (kFragSize * WN); + + // Threadgroup A strides + STEEL_CONST short A_str_m = transpose_a ? 1 : lda_tgp; // M + STEEL_CONST short A_str_k = transpose_a ? lda_tgp : 1; // K + + // Threadgroup B strides + STEEL_CONST short B_str_k = transpose_b ? 1 : ldb_tgp; // K + STEEL_CONST short B_str_n = transpose_b ? ldb_tgp : 1; // N + + // Threadgroup strides along K + STEEL_CONST short tile_stride_a = kFragSize * A_str_k; + STEEL_CONST short tile_stride_b = kFragSize * B_str_k; + + // Simdgroup matrices + MMATile Atile; + MMATile Btile; + MMATile Ctile; + + // Offsets within threadgroup + short sm; + short sn; + + short As_offset; + short Bs_offset; + + /* Constructor */ + METAL_FUNC BlockMMA( + ushort simd_group_id [[simdgroup_index_in_threadgroup]], + ushort simd_lane_id [[thread_index_in_simdgroup]]) { + // Determine thread position in simdgroup matrix + short tm = kFragSize * (simd_group_id / WN); + short tn = kFragSize * (simd_group_id % WN); + + short2 simd_coord = MMAFrag_acc_t::get_coord(simd_lane_id); + sm = simd_coord.y; + sn = simd_coord.x; + + // Determine thread and simdgroup offset + As_offset = (tm + sm) * A_str_m + (sn)*A_str_k; // M, K + Bs_offset = (sm)*B_str_k + (tn + sn) * B_str_n; // K, N + + sm += tm; + sn += tn; + } + + /* (BM, BK) X (BK, BN) multiply accumulate function */ + METAL_FUNC void mma(const threadgroup T* As, const threadgroup T* Bs) { + // Adjust for simdgroup and thread location + As += As_offset; + Bs += Bs_offset; + + // Iterate over BK in blocks of kFragSize + STEEL_PRAGMA_UNROLL + for (short kk = 0; kk < BK; kk += kFragSize) { + simdgroup_barrier(mem_flags::mem_none); + + Atile.template load(As); + + simdgroup_barrier(mem_flags::mem_none); + + Btile.template load(Bs); + + simdgroup_barrier(mem_flags::mem_none); + + tile_matmad(Ctile, Atile, Btile, Ctile); + + // Progress to next simdgroup tile + As += tile_stride_a; + Bs += tile_stride_b; + } + } + + /* Store results from simdgroup_matrix results into device memory */ + METAL_FUNC void store_result(device U* D, const int ldd) { + // Apply epilogue + STEEL_PRAGMA_UNROLL + for (short i = 0; i < decltype(Ctile)::kElemsPerTile; i++) { + Ctile.elems()[i] = Epilogue::apply(Ctile.elems()[i]); + } + + // Adjust for simdgroup and thread location + D += sm * ldd + sn; + + Ctile.template store(D, ldd); + } + + METAL_FUNC void + store_result_slice(device U* D, const int ldd, short2 start, short2 stop) { + // Apply epilogue + STEEL_PRAGMA_UNROLL + for (short i = 0; i < decltype(Ctile)::kElemsPerTile; i++) { + Ctile.elems()[i] = Epilogue::apply(Ctile.elems()[i]); + } + + D += sm * ldd + sn; + start -= short2(sn, sm); + stop -= short2(sn, sm); + + // TODO: Check the start as well + if (stop.y <= 0 || stop.x <= 0) { + return; + } + + Ctile.template store_slice(D, ldd, start, stop); + } + + METAL_FUNC void + store_result_safe(device U* D, const int ldd, short2 dst_tile_dims) { + // Apply epilogue + STEEL_PRAGMA_UNROLL + for (short i = 0; i < decltype(Ctile)::kElemsPerTile; i++) { + Ctile.elems()[i] = Epilogue::apply(Ctile.elems()[i]); + } + + // Adjust for simdgroup and thread location + D += sm * ldd + sn; + dst_tile_dims -= short2(sn, sm); + + if (dst_tile_dims.x <= 0 || dst_tile_dims.y <= 0) + return; + + Ctile.template store_safe(D, ldd, dst_tile_dims); + } + + /* Apply epilogue */ + template + METAL_FUNC void apply_epilogue(thread const UnaryEpilogue& epilogue_op) { + // Loop over all simdgroup tiles + STEEL_PRAGMA_UNROLL + for (short i = 0; i < decltype(Ctile)::kElemsPerTile; i++) { + Ctile.elems()[i] = epilogue_op.apply(Ctile.elems()[i]); + } + } + + /* Apply epilogue */ + template + METAL_FUNC void apply_epilogue( + const device U* C, + const int ldc, + const int fdc, + thread const BinaryEpilogue& epilogue_op) { + // Adjust for simdgroup and thread location + C += (sm)*ldc + (sn)*fdc; + + // Loop over all simdgroup tiles + STEEL_PRAGMA_UNROLL + for (short i = 0; i < TM; i++) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < TN; j++) { + // Get accumulated result and associated offset in C + thread auto& accum = Ctile.frag_at(i, j); + int offset_c = (i * TM_stride) * ldc + (j * TN_stride) * fdc; + + // Apply epilogue + STEEL_PRAGMA_UNROLL + for (short k = 0; k < decltype(Ctile)::kElemsPerFrag; k++) { + accum[k] = epilogue_op.apply(accum[k], C[offset_c + k * fdc]); + } + } + } + } + + /* Apply epilogue */ + template + METAL_FUNC void apply_epilogue_safe( + const device U* C, + const int ldc, + const int fdc, + short2 dst_tile_dims, + thread const BinaryEpilogue& epilogue_op) { + // Adjust for simdgroup and thread location + C += (sm)*ldc + (sn)*fdc; + dst_tile_dims -= short2(sn, sm); + + if (dst_tile_dims.x <= 0 || dst_tile_dims.y <= 0) + return; + + // Loop over all simdgroup tiles + STEEL_PRAGMA_UNROLL + for (short i = 0; i < TM; i++) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < TN; j++) { + // Get accumulated result and associated offset in C + thread auto& accum = Ctile.frag_at(i, j); + int offset_c = (i * TM_stride) * ldc + (j * TN_stride) * fdc; + + constexpr short kelems = decltype(Ctile)::kElemsPerFrag; + + // Read C + U c_elems[kelems] = {0}; + + STEEL_PRAGMA_UNROLL + for (short k = 0; k < kelems; k++) { + if ((j * TN_stride + k) < dst_tile_dims.x) { + c_elems[k] = C[offset_c + k * fdc]; + } + } + + // Apply epilogue + STEEL_PRAGMA_UNROLL + for (short k = 0; k < kelems; k++) { + accum[k] = epilogue_op.apply(accum[k], c_elems[k]); + } + } + } + } + + /* Store results from simdgroup_matrix results into device memory */ + METAL_FUNC void store_result( + device U* D, + const int ldd, + const device U* C, + const int ldc, + const int fdc, + thread const Epilogue& epilogue_op) const { + // Adjust for simdgroup and thread location + C += (sm)*ldc + (sn)*fdc; + D += (sm)*ldd + sn; + + constexpr short kelems = decltype(Ctile)::kElemsPerFrag; + + // Loop over all simdgroup tiles + STEEL_PRAGMA_UNROLL + for (short i = 0; i < TM; i++) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < TN; j++) { + // Get accumulated result and associated offset in C + thread const auto& accum = Ctile.frag_at(i, j); + int offset_c = (i * TM_stride) * ldc + (j * TN_stride) * fdc; + int offset_d = (i * TM_stride) * ldd + (j * TN_stride); + + // Apply epilogue + STEEL_PRAGMA_UNROLL + for (short k = 0; k < kelems; k++) { + D[offset_d + k] = epilogue_op.apply(accum[k], C[offset_c + k * fdc]); + } + } + } + } + + METAL_FUNC void store_result_safe( + device U* D, + const int ldd, + const device U* C, + const int ldc, + const int fdc, + short2 dst_tile_dims, + thread const Epilogue& epilogue_op) const { + // Adjust for simdgroup and thread location + C += (sm)*ldc + (sn)*fdc; + D += (sm)*ldd + sn; + dst_tile_dims -= short2(sn, sm); + + if (dst_tile_dims.x <= 0 || dst_tile_dims.y <= 0) + return; + + constexpr short kelems = decltype(Ctile)::kElemsPerFrag; + + STEEL_PRAGMA_UNROLL + for (int i = 0; i < TM; i++) { + if (i * TM_stride < dst_tile_dims.y) { + STEEL_PRAGMA_UNROLL + for (int j = 0; j < TN; j++) { + // Get accumulated result and associated offset in C + thread const auto& accum = Ctile.frag_at(i, j); + int offset_c = (i * TM_stride) * ldc + (j * TN_stride) * fdc; + int offset_d = (i * TM_stride) * ldd + (j * TN_stride); + + // Apply epilogue + STEEL_PRAGMA_UNROLL + for (short k = 0; k < kelems; k++) { + if ((j * TN_stride + k) < dst_tile_dims.x) { + D[offset_d + k] = + epilogue_op.apply(accum[k], C[offset_c + k * fdc]); + } + } + } + } + } + } +}; + +template < + typename U, + int BM, + int BN, + int BK, + int WM, + int WN, + bool transpose_a, + bool transpose_b, + short lda_tgp, + short ldb_tgp, + typename AccumType, + typename Epilogue> +struct BlockMMA< + complex64_t, + U, + BM, + BN, + BK, + WM, + WN, + transpose_a, + transpose_b, + lda_tgp, + ldb_tgp, + AccumType, + Epilogue> { + static_assert( + metal::is_same_v, + "BlockMMA expects float accumulators"); + static_assert( + metal::is_same_v, + "For complex BlockMMA, U must be complex64_t; use a different epilogue for projections"); + // MMAFrag size + STEEL_CONST short kFragSize = 8; + using MMAFrag_acc_t = BaseMMAFrag; + + // Warp tile simdgroup matrix strides along M + STEEL_CONST short TM_stride = kFragSize * WM; + // Warp tile simdgroup matrix strides along M + STEEL_CONST short TN_stride = kFragSize * WN; + + // Warp tile size along M + STEEL_CONST short TM = BM / (kFragSize * WM); + // Warp tile size along N + STEEL_CONST short TN = BN / (kFragSize * WN); + + // Threadgroup A strides + STEEL_CONST short A_str_m = transpose_a ? 1 : lda_tgp; // M + STEEL_CONST short A_str_k = transpose_a ? lda_tgp : 1; // K + + // Threadgroup B strides + STEEL_CONST short B_str_k = transpose_b ? 1 : ldb_tgp; // K + STEEL_CONST short B_str_n = transpose_b ? ldb_tgp : 1; // N + + // Threadgroup strides along K + STEEL_CONST short tile_stride_a = kFragSize * A_str_k; + STEEL_CONST short tile_stride_b = kFragSize * B_str_k; + + // When indexing complex as float[2] + STEEL_CONST short A_str_m_f = A_str_m * 2; + STEEL_CONST short A_str_k_f = A_str_k * 2; + STEEL_CONST short B_str_k_f = B_str_k * 2; + STEEL_CONST short B_str_n_f = B_str_n * 2; + STEEL_CONST short tile_stride_a_f = tile_stride_a * 2; + STEEL_CONST short tile_stride_b_f = tile_stride_b * 2; + + // Accumulators (real/imag) + MMATile Ctile_r; + MMATile Ctile_i; + + // Offsets within threadgroup + short sm, sn; + short As_offset, Bs_offset; + + /* Constructor */ + METAL_FUNC BlockMMA( + ushort simd_group_id [[simdgroup_index_in_threadgroup]], + ushort simd_lane_id [[thread_index_in_simdgroup]]) { + // Determine thread position in simdgroup matrix + short tm = kFragSize * (simd_group_id / WN); + short tn = kFragSize * (simd_group_id % WN); + + short2 simd_coord = MMAFrag_acc_t::get_coord(simd_lane_id); + sm = simd_coord.y; + sn = simd_coord.x; + + // Determine thread and simdgroup offset + As_offset = (tm + sm) * A_str_m + (sn)*A_str_k; // (M,K) + Bs_offset = (sm)*B_str_k + (tn + sn) * B_str_n; // (K,N) + + sm += tm; + sn += tn; + } + + /* Karatsuba MMA: 3 real MMAs per K-chunk */ + METAL_FUNC void mma( + const threadgroup complex64_t* As, + const threadgroup complex64_t* Bs) { + // Adjust for simdgroup and thread location + As += As_offset; + Bs += Bs_offset; + threadgroup const float* As_f = + reinterpret_cast(As); + threadgroup const float* Bs_f = + reinterpret_cast(Bs); + + // Iterate over BK in blocks of kFragSize + STEEL_PRAGMA_UNROLL + for (short kk = 0; kk < BK; kk += kFragSize) { + simdgroup_barrier(mem_flags::mem_none); + + MMATile Ar, Ai; + Ar.template load(As_f + 0); + Ai.template load(As_f + 1); + + simdgroup_barrier(mem_flags::mem_none); + + MMATile Br, Bi; + Br.template load(Bs_f + 0); + Bi.template load(Bs_f + 1); + + simdgroup_barrier(mem_flags::mem_none); + + // P = Ar*Br ; Q = Ai*Bi ; R = (Ar+Ai)*(Br+Bi) + MMATile P, Q, R; + + tile_matmad(P, Ar, Br, P); + tile_matmad(Q, Ai, Bi, Q); + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < decltype(Ar)::kElemsPerTile; ++i) + Ar.elems()[i] += Ai.elems()[i]; + STEEL_PRAGMA_UNROLL + for (short i = 0; i < decltype(Br)::kElemsPerTile; ++i) + Br.elems()[i] += Bi.elems()[i]; + + tile_matmad(R, Ar, Br, R); + + // C_r += P - Q ; C_i -= Q + STEEL_PRAGMA_UNROLL + for (short i = 0; i < decltype(Ctile_r)::kElemsPerTile; ++i) { + const auto p = P.elems()[i]; + const auto q = Q.elems()[i]; + const auto r = R.elems()[i]; + Ctile_r.elems()[i] += (p - q); + Ctile_i.elems()[i] += (r - p - q); + } + + // Progress to next simdgroup tile + As_f += tile_stride_a_f; + Bs_f += tile_stride_b_f; + } + } + + /* Store results from simdgroup_matrix results into device memory */ + METAL_FUNC void store_result(device U* D, const int ldd) { + // Adjust for simdgroup and thread location + D += sm * ldd + sn; + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < TM; i++) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < TN; j++) { + thread const auto& r = Ctile_r.frag_at(i, j); + thread const auto& im = Ctile_i.frag_at(i, j); + int off = (i * TM_stride) * ldd + (j * TN_stride); + STEEL_PRAGMA_UNROLL + for (short k = 0; k < decltype(Ctile_r)::kElemsPerFrag; k++) { + D[off + k] = Epilogue::apply(complex64_t(r[k], im[k])); + } + } + } + } + + METAL_FUNC void + store_result_slice(device U* D, const int ldd, short2 start, short2 stop) { + D += sm * ldd + sn; + start -= short2(sn, sm); + stop -= short2(sn, sm); + + if (stop.y <= 0 || stop.x <= 0) + return; + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < TM; ++i) { + const int row = i * TM_stride; + if (row >= start.y && row < stop.y) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < TN; ++j) { + const int off = row * ldd + (j * TN_stride); + thread const auto& r = Ctile_r.frag_at(i, j); + thread const auto& im = Ctile_i.frag_at(i, j); + + STEEL_PRAGMA_UNROLL + for (short k = 0; k < decltype(Ctile_r)::kElemsPerFrag; ++k) { + const int col = j * TN_stride + k; + if (col >= start.x && col < stop.x) { + D[off + k] = Epilogue::apply(complex64_t(r[k], im[k])); + } + } + } + } + } + } + + METAL_FUNC void + store_result_safe(device U* D, const int ldd, short2 dst_tile_dims) { + D += sm * ldd + sn; + dst_tile_dims -= short2(sn, sm); + if (dst_tile_dims.x <= 0 || dst_tile_dims.y <= 0) + return; + STEEL_PRAGMA_UNROLL + for (short i = 0; i < TM; i++) { + if (i * TM_stride < dst_tile_dims.y) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < TN; j++) { + int off = (i * TM_stride) * ldd + (j * TN_stride); + thread const auto& r = Ctile_r.frag_at(i, j); + thread const auto& im = Ctile_i.frag_at(i, j); + STEEL_PRAGMA_UNROLL + for (short k = 0; k < decltype(Ctile_r)::kElemsPerFrag; k++) { + if ((j * TN_stride + k) < dst_tile_dims.x) { + D[off + k] = Epilogue::apply(complex64_t(r[k], im[k])); + } + } + } + } + } + } + + /* Apply epilogue */ + template + METAL_FUNC void apply_epilogue(thread const UnaryEpilogue& epilogue_op) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < decltype(Ctile_r)::kElemsPerTile; i++) { + complex64_t out = epilogue_op.apply( + complex64_t(Ctile_r.elems()[i], Ctile_i.elems()[i])); + Ctile_r.elems()[i] = out.real; + Ctile_i.elems()[i] = out.imag; + } + } + + /* Apply epilogue */ + template + METAL_FUNC void apply_epilogue( + const device U* C, + const int ldc, + const int fdc, + thread const BinaryEpilogue& epilogue_op) { + // Adjust for simdgroup and thread location + C += (sm)*ldc + (sn)*fdc; + + // Loop over all simdgroup tiles + STEEL_PRAGMA_UNROLL + for (short i = 0; i < TM; i++) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < TN; j++) { + // Get accumulated result and associated offset in Cr, Ci + thread auto& r = Ctile_r.frag_at(i, j); + thread auto& im = Ctile_i.frag_at(i, j); + int offset_c = (i * TM_stride) * ldc + (j * TN_stride) * fdc; + + STEEL_PRAGMA_UNROLL + for (short k = 0; k < decltype(Ctile_r)::kElemsPerFrag; k++) { + complex64_t out = epilogue_op.apply( + complex64_t(r[k], im[k]), C[offset_c + k * fdc]); + r[k] = out.real; + im[k] = out.imag; + } + } + } + } + + /* Apply epilogue */ + template + METAL_FUNC void apply_epilogue_safe( + const device U* C, + const int ldc, + const int fdc, + short2 dst_tile_dims, + thread const BinaryEpilogue& epilogue_op) { + // Adjust for simdgroup and thread location + C += (sm)*ldc + (sn)*fdc; + dst_tile_dims -= short2(sn, sm); + + if (dst_tile_dims.x <= 0 || dst_tile_dims.y <= 0) + return; + + // Loop over all simdgroup tiles + STEEL_PRAGMA_UNROLL + for (short i = 0; i < TM; i++) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < TN; j++) { + // Get accumulated result and associated offset in Cr, Ci + thread auto& r = Ctile_r.frag_at(i, j); + thread auto& im = Ctile_i.frag_at(i, j); + int offset_c = (i * TM_stride) * ldc + (j * TN_stride) * fdc; + + constexpr short kelems = decltype(Ctile_r)::kElemsPerFrag; + complex64_t tmp[kelems]; + + STEEL_PRAGMA_UNROLL + for (short k = 0; k < kelems; k++) { + if ((j * TN_stride + k) < dst_tile_dims.x && + (i * TM_stride) < dst_tile_dims.y) { + tmp[k] = C[offset_c + k * fdc]; + } else { + tmp[k] = complex64_t(0.0f, 0.0f); + } + } + + // Apply epilogue + STEEL_PRAGMA_UNROLL + for (short k = 0; k < kelems; k++) { + complex64_t out = epilogue_op.apply(complex64_t(r[k], im[k]), tmp[k]); + r[k] = out.real; + im[k] = out.imag; + } + } + } + } + + /* Store results from simdgroup_matrix results into device memory */ + METAL_FUNC void store_result( + device U* D, + const int ldd, + const device U* C, + const int ldc, + const int fdc, + thread const Epilogue& epilogue_op) const { + // Adjust for simdgroup and thread location + C += (sm)*ldc + (sn)*fdc; + D += (sm)*ldd + sn; + + constexpr short kelems = decltype(Ctile_r)::kElemsPerFrag; + + // Loop over all simdgroup tiles + STEEL_PRAGMA_UNROLL + for (short i = 0; i < TM; i++) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < TN; j++) { + // Get accumulated result and associated offset in Cr, Ci + thread const auto& r = Ctile_r.frag_at(i, j); + thread const auto& im = Ctile_i.frag_at(i, j); + int off_c = (i * TM_stride) * ldc + (j * TN_stride) * fdc; + int off_d = (i * TM_stride) * ldd + (j * TN_stride); + + // Apply epilogue + STEEL_PRAGMA_UNROLL + for (short k = 0; k < kelems; k++) { + D[off_d + k] = + epilogue_op.apply(complex64_t(r[k], im[k]), C[off_c + k * fdc]); + } + } + } + } + + METAL_FUNC void store_result_safe( + device U* D, + const int ldd, + const device U* C, + const int ldc, + const int fdc, + short2 dst_tile_dims, + thread const Epilogue& epilogue_op) const { + // Adjust for simdgroup and thread location + C += (sm)*ldc + (sn)*fdc; + D += (sm)*ldd + sn; + dst_tile_dims -= short2(sn, sm); + + if (dst_tile_dims.x <= 0 || dst_tile_dims.y <= 0) + return; + + constexpr short kelems = decltype(Ctile_r)::kElemsPerFrag; + + STEEL_PRAGMA_UNROLL + for (int i = 0; i < TM; i++) { + if (i * TM_stride < dst_tile_dims.y) { + STEEL_PRAGMA_UNROLL + for (int j = 0; j < TN; j++) { + // Get accumulated result and associated offset in Cr, Ci + thread const auto& r = Ctile_r.frag_at(i, j); + thread const auto& im = Ctile_i.frag_at(i, j); + int off_c = (i * TM_stride) * ldc + (j * TN_stride) * fdc; + int off_d = (i * TM_stride) * ldd + (j * TN_stride); + + // Apply epilogue + STEEL_PRAGMA_UNROLL + for (short k = 0; k < kelems; k++) { + if ((j * TN_stride + k) < dst_tile_dims.x) { + D[off_d + k] = epilogue_op.apply( + complex64_t(r[k], im[k]), C[off_c + k * fdc]); + } + } + } + } + } + } +}; + +} // namespace steel +} // namespace mlx + + +// ===== mlx/backend/metal/kernels/steel/conv/conv.h ===== +// Copyright © 2024 Apple Inc. + + +// (already included: mlx/backend/metal/kernels/steel/defines.h) +// (already included: mlx/backend/metal/kernels/steel/utils.h) + + +// ===== mlx/backend/metal/kernels/steel/conv/loader.h ===== +// Copyright © 2024 Apple Inc. + + + +// ===== mlx/backend/metal/kernels/steel/conv/loaders/loader_channel_l.h ===== +// Copyright © 2024 Apple Inc. + + +// (already included: mlx/backend/metal/kernels/steel/utils.h) + + +// ===== mlx/backend/metal/kernels/steel/conv/params.h ===== +// Copyright © 2024 Apple Inc. + + +template +struct MLXConvParams { + int N; // Batch size + int C; // In channels + int O; // Out channels + int iS[NDIM]; // Input spatial dim + int wS[NDIM]; // Weight spatial dim + int oS[NDIM]; // Output spatial dim + int str[NDIM]; // Kernel strides + int pad[NDIM]; // Input padding + int kdil[NDIM]; // Kernel dilation + int idil[NDIM]; // Input dilation + int64_t in_strides[NDIM + 2]; // In strides + int64_t wt_strides[NDIM + 2]; // Wt strides + int64_t out_strides[NDIM + 2]; // Out strides + int groups; // Input channel groups + bool flip; + + static MLXConvParams + with_padded_channels(MLXConvParams other, int pad_out, int pad_in) { + MLXConvParams params = other; + + // Update strides + for (int i = 0; i < NDIM + 1; i++) { + params.in_strides[i] = + (params.in_strides[i] / params.C) * (params.C + pad_in); + params.wt_strides[i] = + (params.wt_strides[i] / params.C) * (params.C + pad_in); + params.out_strides[i] = + (params.out_strides[i] / params.O) * (params.O + pad_out); + } + params.in_strides[NDIM + 1] = 1; + params.wt_strides[NDIM + 1] = 1; + params.out_strides[NDIM + 1] = 1; + + // Update channels + params.C += pad_in; + params.O += pad_out; + + return params; + }; +}; + +namespace mlx { +namespace steel { + +struct ImplicitGemmConv2DParams { + const int M; + const int N; + const int K; + + const int gemm_k_iterations; + + const int inp_jump_w; + const int inp_jump_h; + const int inp_jump_c; + + const int tiles_n; + const int tiles_m; + const int swizzle_log; +}; + +struct ImplicitGemmConv3DParams { + const int M; + const int N; + const int K; + + const int gemm_k_iterations; + + const int inp_jump_w; + const int inp_jump_h; + const int inp_jump_d; + const int inp_jump_c; + + const int tiles_n; + const int tiles_m; + const int swizzle_log; +}; + +struct Conv2DGeneralJumpParams { + const int f_wgt_jump_h; + const int f_wgt_jump_w; + + const int f_out_jump_h; + const int f_out_jump_w; + + const int adj_out_h; + const int adj_out_w; + const int adj_out_hw; + const int adj_implicit_m; +}; + +struct Conv2DGeneralBaseInfo { + int weight_base; + int weight_size; +}; + +} // namespace steel +} // namespace mlx + + +/////////////////////////////////////////////////////////////////////////////// +// Loading helper +/////////////////////////////////////////////////////////////////////////////// + +namespace mlx { +namespace steel { + +template < + typename T, + short BM, + short BN, + short BK, + short tgp_size, + short tgp_padding = 0> +struct Conv2DInputBlockLoaderLargeFilter { + // Destination dimensions + STEEL_CONST short BROWS = BM; + STEEL_CONST short BCOLS = BK; + + // Read dimensions + STEEL_CONST short dst_ld = BCOLS + tgp_padding; + STEEL_CONST short vec_size = tgp_size / (BROWS * BCOLS) >= 8 ? 8 : 4; + + // Thread read shape + STEEL_CONST short TCOLS = BCOLS / vec_size; + STEEL_CONST short TROWS = tgp_size / TCOLS; + + // Rows / strided reads within the block + STEEL_CONST short n_rows = BROWS / TROWS; + + // Thread location indices + const short thread_idx; + const short bi; + const short bj; + + // threadgroup and device memory + threadgroup T* dst; + + const constant MLXConvParams<2>* params; + const constant ImplicitGemmConv2DParams* gemm_params; + + short weight_h; + short weight_w; + + const device T* src[n_rows]; + + int read_n[n_rows]; + int read_ih[n_rows]; + int read_iw[n_rows]; + + /* Constructor */ + METAL_FUNC Conv2DInputBlockLoaderLargeFilter( + const device T* src_, + threadgroup T* dst_, + const int2 offsets, + const constant MLXConvParams<2>* params_, + const constant ImplicitGemmConv2DParams* gemm_params_, + uint simd_group_id [[simdgroup_index_in_threadgroup]], + uint simd_lane_id [[thread_index_in_simdgroup]]) + : thread_idx(simd_group_id * 32 + simd_lane_id), + bi(thread_idx / TCOLS), + bj(vec_size * (thread_idx % TCOLS)), + dst(dst_ + bi * dst_ld + bj), + params(params_), + gemm_params(gemm_params_), + weight_h(0), + weight_w(0) { + int out_n_pixels = params->oS[0] * params->oS[1]; + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < n_rows; ++i) { + int offset_nhw = offsets.y + bi + i * TROWS; + int n = offset_nhw / out_n_pixels; + int hw = offset_nhw % out_n_pixels; + int oh = hw / params->oS[1]; + int ow = hw % params->oS[1]; + + int ih = oh * params->str[0] - params->pad[0]; + int iw = ow * params->str[1] - params->pad[1]; + + read_n[i] = n; + read_ih[i] = ih; + read_iw[i] = iw; + + // Adjust for flip + if (params->flip) { + ih += (params->wS[0] - 1) * params->kdil[0]; + iw += (params->wS[1] - 1) * params->kdil[1]; + } + + // Read from input if in bounds + src[i] = src_ + n * params->in_strides[0] + ih * params->in_strides[1] + + iw * params->in_strides[2] + bj; + } + } + + /* Load from device memory into threadgroup memory - without bound checking */ + METAL_FUNC void load_unsafe() const { + STEEL_PRAGMA_UNROLL + for (short i = 0, is = 0; i < n_rows; ++i, is += TROWS) { + // Find bounds + int n = read_n[i]; + int ih = read_ih[i] + weight_h * params->kdil[0]; + int iw = read_iw[i] + weight_w * params->kdil[1]; + + // Read from input if in bounds + if ((n < params->N) && (ih >= 0 && ih < params->iS[0]) && + (iw >= 0 && iw < params->iS[1])) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < vec_size; ++j) { + dst[is * dst_ld + j] = src[i][j]; + } + } + + // Zero pad otherwise + else { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < vec_size; ++j) { + dst[is * dst_ld + j] = T(0); + } + } + } + } + + /* Iteration helper */ + METAL_FUNC void next() { + if (++weight_w < params->wS[1]) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < n_rows; i++) { + src[i] += gemm_params->inp_jump_w; + } + + return; + } + + weight_w = 0; + + if (++weight_h < params->wS[0]) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < n_rows; i++) { + src[i] += gemm_params->inp_jump_h; + } + + return; + } + + weight_h = 0; + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < n_rows; i++) { + src[i] += gemm_params->inp_jump_c; + } + } +}; + +template < + typename T, + short BM, + short BN, + short BK, + short tgp_size, + short tgp_padding = 0> +struct Conv2DInputBlockLoaderSmallFilter { + // Destination dimensions + STEEL_CONST short BROWS = BM; + STEEL_CONST short BCOLS = BK; + + // Read dimensions + STEEL_CONST short dst_ld = BCOLS + tgp_padding; + STEEL_CONST short vec_size = tgp_size / (BROWS * BCOLS) >= 8 ? 8 : 4; + + // Thread read shape + STEEL_CONST short TCOLS = BCOLS / vec_size; + STEEL_CONST short TROWS = tgp_size / TCOLS; + + // Rows / strided reads within the block + STEEL_CONST short n_rows = BROWS / TROWS; + + using mask_t = short; + + // Thread location indices + const short thread_idx; + const short bi; + const short bj; + + // threadgroup and device memory + threadgroup T* dst; + + const constant MLXConvParams<2>* params; + const constant ImplicitGemmConv2DParams* gemm_params; + + short weight_h; + short weight_w; + + const device T* src[n_rows]; + + mask_t mask_h[n_rows]; + mask_t mask_w[n_rows]; + + /* Constructor */ + METAL_FUNC Conv2DInputBlockLoaderSmallFilter( + const device T* src_, + threadgroup T* dst_, + const int2 offsets, + const constant MLXConvParams<2>* params_, + const constant ImplicitGemmConv2DParams* gemm_params_, + uint simd_group_id [[simdgroup_index_in_threadgroup]], + uint simd_lane_id [[thread_index_in_simdgroup]]) + : thread_idx(simd_group_id * 32 + simd_lane_id), + bi(thread_idx / TCOLS), + bj(vec_size * (thread_idx % TCOLS)), + dst(dst_ + bi * dst_ld + bj), + params(params_), + gemm_params(gemm_params_), + weight_h(0), + weight_w(0) { + int out_n_pixels = params->oS[0] * params->oS[1]; + + int read_n[n_rows]; + int read_ih[n_rows]; + int read_iw[n_rows]; + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < n_rows; ++i) { + int offset_nhw = offsets.y + bi + i * TROWS; + int n = offset_nhw / out_n_pixels; + int hw = offset_nhw % out_n_pixels; + int oh = hw / params->oS[1]; + int ow = hw % params->oS[1]; + + int ih = oh * params->str[0] - params->pad[0]; + int iw = ow * params->str[1] - params->pad[1]; + + read_n[i] = n; + read_ih[i] = ih; + read_iw[i] = iw; + + // Adjust for flip + if (params->flip) { + ih += (params->wS[0] - 1) * params->kdil[0]; + iw += (params->wS[1] - 1) * params->kdil[1]; + } + + // Read from input if in bounds + src[i] = src_ + n * params->in_strides[0] + ih * params->in_strides[1] + + iw * params->in_strides[2] + bj; + } + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < n_rows; ++i) { + mask_h[i] = 0; + mask_w[i] = 0; + } + + for (short kh = 0; kh < params->wS[0]; kh++) { + short flip_h = params->flip ? params->wS[0] - kh - 1 : kh; + STEEL_PRAGMA_UNROLL + for (short i = 0; i < n_rows; ++i) { + int n = read_n[i]; + int ih = read_ih[i] + flip_h * params->kdil[0]; + + bool in_bounds = n < params->N && ih >= 0 && ih < params->iS[0]; + + mask_h[i] |= (in_bounds << kh); + } + } + + for (short kw = 0; kw < params->wS[1]; kw++) { + short flip_w = params->flip ? params->wS[1] - kw - 1 : kw; + STEEL_PRAGMA_UNROLL + for (short i = 0; i < n_rows; ++i) { + int iw = read_iw[i] + flip_w * params->kdil[1]; + + bool in_bounds = iw >= 0 && iw < params->iS[1]; + + mask_w[i] |= (in_bounds << kw); + } + } + } + + /* Load from device memory into threadgroup memory - without bound checking */ + METAL_FUNC void load_unsafe() const { + mask_t h_mask = mask_t(1) << weight_h; + mask_t w_mask = mask_t(1) << weight_w; + + STEEL_PRAGMA_UNROLL + for (short i = 0, is = 0; i < n_rows; ++i, is += TROWS) { + // Read from input if in bounds + if ((mask_h[i] & h_mask) && (mask_w[i] & w_mask)) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < vec_size; ++j) { + dst[is * dst_ld + j] = src[i][j]; + } + } + + // Zero pad otherwise + else { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < vec_size; ++j) { + dst[is * dst_ld + j] = T(0); + } + } + } + } + + /* Iteration helper */ + METAL_FUNC void next() { + if (++weight_w < params->wS[1]) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < n_rows; i++) { + src[i] += gemm_params->inp_jump_w; + } + + return; + } + + weight_w = 0; + + if (++weight_h < params->wS[0]) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < n_rows; i++) { + src[i] += gemm_params->inp_jump_h; + } + + return; + } + + weight_h = 0; + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < n_rows; i++) { + src[i] += gemm_params->inp_jump_c; + } + } +}; + +template < + typename T, + short BM, + short BN, + short BK, + short tgp_size, + short tgp_padding = 0> +struct Conv2DWeightBlockLoader { + // Destination dimensions + STEEL_CONST short BROWS = BN; + STEEL_CONST short BCOLS = BK; + + // Read dimensions + STEEL_CONST short dst_ld = BCOLS + tgp_padding; + STEEL_CONST short vec_size = + (BN == 8) ? 1 : (tgp_size / (BROWS * BCOLS) >= 8 ? 8 : 4); + + // Thread read shape + STEEL_CONST short TCOLS = BCOLS / vec_size; + STEEL_CONST short TROWS = tgp_size / TCOLS; + + // Rows / strided reads within the block + STEEL_CONST short n_rows = BROWS / TROWS; + + // Leading dimension for src + const int src_ld; + + // Thread location indices + const short thread_idx; + const short bi; + const short bj; + + // threadgroup and device memory + threadgroup T* dst; + const device T* src; + + const constant MLXConvParams<2>* params; + + int weight_hw; + int weight_step; + + const int read_n; + const bool do_read; + + /* Constructor */ + METAL_FUNC Conv2DWeightBlockLoader( + const device T* src_, + threadgroup T* dst_, + const int2 offsets, + const constant MLXConvParams<2>* params_, + const constant ImplicitGemmConv2DParams* gemm_params_, + uint simd_group_id [[simdgroup_index_in_threadgroup]], + uint simd_lane_id [[thread_index_in_simdgroup]]) + : src_ld(params_->wt_strides[0]), + thread_idx(simd_group_id * 32 + simd_lane_id), + bi(thread_idx / TCOLS), + bj(vec_size * (thread_idx % TCOLS)), + dst(dst_ + bi * dst_ld + bj), + src(src_ + bi * src_ld + bj), + params(params_), + weight_hw(0), + weight_step(params->C / params->groups), + read_n(offsets.y + bi), + do_read(read_n + n_rows * TROWS <= gemm_params_->N) {} + + /* Load from device memory into threadgroup memory - without bound checking */ + METAL_FUNC void load_unsafe() const { + if (BN != 8 || do_read) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BN; i += TROWS) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < vec_size; j++) { + dst[i * dst_ld + j] = src[i * src_ld + j]; + } + } + } else { + for (short i = 0; i < BN; i += TROWS) { + if ((read_n + i) < params->O) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < vec_size; j++) { + dst[i * dst_ld + j] = src[i * src_ld + j]; + } + } else { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < vec_size; j++) { + dst[i * dst_ld + j] = T(0); + } + } + } + } + } + + /* Iteration helper */ + METAL_FUNC void next() { + if (++weight_hw < (params->wS[1] * params->wS[0])) { + src += weight_step; + return; + } + + weight_hw = 0; + + src += BK - (params->wS[1] * params->wS[0] - 1) * weight_step; + } +}; + +template < + typename T, + short BM, + short BN, + short BK, + short tgp_size, + short tgp_padding = 0> +struct Conv3DInputBlockLoaderLargeFilter { + // Destination dimensions + STEEL_CONST short BROWS = BM; + STEEL_CONST short BCOLS = BK; + + // Read dimensions + STEEL_CONST short dst_ld = BCOLS + tgp_padding; + STEEL_CONST short vec_size = tgp_size / (BROWS * BCOLS) >= 8 ? 8 : 4; + + // Thread read shape + STEEL_CONST short TCOLS = BCOLS / vec_size; + STEEL_CONST short TROWS = tgp_size / TCOLS; + + // Rows / strided reads within the block + STEEL_CONST short n_rows = BROWS / TROWS; + + // Thread location indices + const short thread_idx; + const short bi; + const short bj; + + // threadgroup and device memory + threadgroup T* dst; + + const constant MLXConvParams<3>* params; + const constant ImplicitGemmConv3DParams* gemm_params; + + short weight_d; + short weight_h; + short weight_w; + + short kdil_d; + short kdil_h; + short kdil_w; + + const device T* src[n_rows]; + + int read_n[n_rows]; + int read_id[n_rows]; + int read_ih[n_rows]; + int read_iw[n_rows]; + + /* Constructor */ + METAL_FUNC Conv3DInputBlockLoaderLargeFilter( + const device T* src_, + threadgroup T* dst_, + const int2 offsets, + const constant MLXConvParams<3>* params_, + const constant ImplicitGemmConv3DParams* gemm_params_, + uint simd_group_id [[simdgroup_index_in_threadgroup]], + uint simd_lane_id [[thread_index_in_simdgroup]]) + : thread_idx(simd_group_id * 32 + simd_lane_id), + bi(thread_idx / TCOLS), + bj(vec_size * (thread_idx % TCOLS)), + dst(dst_ + bi * dst_ld + bj), + params(params_), + gemm_params(gemm_params_), + weight_d(0), + weight_h(0), + weight_w(0), + kdil_d(params_->flip ? -params_->kdil[0] : params_->kdil[0]), + kdil_h(params_->flip ? -params_->kdil[1] : params_->kdil[1]), + kdil_w(params_->flip ? -params_->kdil[2] : params_->kdil[2]) { + int out_n_pixels = params->oS[0] * params->oS[1] * params->oS[2]; + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < n_rows; ++i) { + int offset_ndhw = offsets.y + bi + i * TROWS; + int n = offset_ndhw / out_n_pixels; + int dhw = offset_ndhw % out_n_pixels; + int od = dhw / (params->oS[1] * params->oS[2]); + int hw = dhw % (params->oS[1] * params->oS[2]); + int oh = hw / params->oS[2]; + int ow = hw % params->oS[2]; + + int id = od * params->str[0] - params->pad[0]; + int ih = oh * params->str[1] - params->pad[1]; + int iw = ow * params->str[2] - params->pad[2]; + + read_n[i] = n; + + if (params->flip) { + read_id[i] = id + (params->wS[0] - 1) * params->kdil[0]; + read_ih[i] = ih + (params->wS[1] - 1) * params->kdil[1]; + read_iw[i] = iw + (params->wS[2] - 1) * params->kdil[2]; + } else { + read_id[i] = id; + read_ih[i] = ih; + read_iw[i] = iw; + } + + // Adjust for flip + if (params->flip) { + id += (params->wS[0] - 1) * params->kdil[0]; + ih += (params->wS[1] - 1) * params->kdil[1]; + iw += (params->wS[2] - 1) * params->kdil[2]; + } + + // Read from input if in bounds + src[i] = src_ + n * params->in_strides[0] + id * params->in_strides[1] + + ih * params->in_strides[2] + iw * params->in_strides[3] + bj; + } + } + + /* Load from device memory into threadgroup memory - without bound checking */ + METAL_FUNC void load_unsafe() const { + STEEL_PRAGMA_UNROLL + for (short i = 0, is = 0; i < n_rows; ++i, is += TROWS) { + // Find bounds + int n = read_n[i]; + int id = read_id[i] + weight_d * kdil_d; + int ih = read_ih[i] + weight_h * kdil_h; + int iw = read_iw[i] + weight_w * kdil_w; + + // Read from input if in bounds + if ((n < params->N) && (id >= 0 && id < params->iS[0]) && + (ih >= 0 && ih < params->iS[1]) && (iw >= 0 && iw < params->iS[2])) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < vec_size; ++j) { + dst[is * dst_ld + j] = src[i][j]; + } + } + + // Zero pad otherwise + else { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < vec_size; ++j) { + dst[is * dst_ld + j] = T(0); + } + } + } + } + + /* Iteration helper */ + METAL_FUNC void next() { + if (++weight_w < params->wS[2]) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < n_rows; i++) { + src[i] += gemm_params->inp_jump_w; + } + + return; + } + + weight_w = 0; + + if (++weight_h < params->wS[1]) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < n_rows; i++) { + src[i] += gemm_params->inp_jump_h; + } + + return; + } + + weight_h = 0; + + if (++weight_d < params->wS[0]) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < n_rows; i++) { + src[i] += gemm_params->inp_jump_d; + } + + return; + } + + weight_d = 0; + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < n_rows; i++) { + src[i] += gemm_params->inp_jump_c; + } + } +}; + +template < + typename T, + short BM, + short BN, + short BK, + short tgp_size, + short tgp_padding = 0> +struct Conv3DInputBlockLoaderSmallFilter { + // Destination dimensions + STEEL_CONST short BROWS = BM; + STEEL_CONST short BCOLS = BK; + + // Read dimensions + STEEL_CONST short dst_ld = BCOLS + tgp_padding; + STEEL_CONST short vec_size = tgp_size / (BROWS * BCOLS) >= 8 ? 8 : 4; + + // Thread read shape + STEEL_CONST short TCOLS = BCOLS / vec_size; + STEEL_CONST short TROWS = tgp_size / TCOLS; + + // Rows / strided reads within the block + STEEL_CONST short n_rows = BROWS / TROWS; + + using mask_t = short; + + // Thread location indices + const short thread_idx; + const short bi; + const short bj; + + // threadgroup and device memory + threadgroup T* dst; + + const constant MLXConvParams<3>* params; + const constant ImplicitGemmConv3DParams* gemm_params; + + short weight_d; + short weight_h; + short weight_w; + + const device T* src[n_rows]; + + mask_t mask_d[n_rows]; + mask_t mask_h[n_rows]; + mask_t mask_w[n_rows]; + + /* Constructor */ + METAL_FUNC Conv3DInputBlockLoaderSmallFilter( + const device T* src_, + threadgroup T* dst_, + const int2 offsets, + const constant MLXConvParams<3>* params_, + const constant ImplicitGemmConv3DParams* gemm_params_, + uint simd_group_id [[simdgroup_index_in_threadgroup]], + uint simd_lane_id [[thread_index_in_simdgroup]]) + : thread_idx(simd_group_id * 32 + simd_lane_id), + bi(thread_idx / TCOLS), + bj(vec_size * (thread_idx % TCOLS)), + dst(dst_ + bi * dst_ld + bj), + params(params_), + gemm_params(gemm_params_), + weight_d(0), + weight_h(0), + weight_w(0) { + int out_n_pixels = params->oS[0] * params->oS[1] * params->oS[2]; + + int read_n[n_rows]; + int read_id[n_rows]; + int read_ih[n_rows]; + int read_iw[n_rows]; + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < n_rows; ++i) { + int offset_ndhw = offsets.y + bi + i * TROWS; + int n = offset_ndhw / out_n_pixels; + int dhw = offset_ndhw % out_n_pixels; + int od = dhw / (params->oS[1] * params->oS[2]); + int hw = dhw % (params->oS[1] * params->oS[2]); + int oh = hw / params->oS[2]; + int ow = hw % params->oS[2]; + + int id = od * params->str[0] - params->pad[0]; + int ih = oh * params->str[1] - params->pad[1]; + int iw = ow * params->str[2] - params->pad[2]; + + read_n[i] = n; + read_id[i] = id; + read_ih[i] = ih; + read_iw[i] = iw; + + // Adjust for flip + if (params->flip) { + id += (params->wS[0] - 1) * params->kdil[0]; + ih += (params->wS[1] - 1) * params->kdil[1]; + iw += (params->wS[2] - 1) * params->kdil[2]; + } + + // Read from input if in bounds + src[i] = src_ + n * params->in_strides[0] + id * params->in_strides[1] + + ih * params->in_strides[2] + iw * params->in_strides[3] + bj; + } + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < n_rows; ++i) { + mask_d[i] = 0; + mask_h[i] = 0; + mask_w[i] = 0; + } + + for (short kd = 0; kd < params->wS[0]; kd++) { + short flip_d = params->flip ? params->wS[0] - kd - 1 : kd; + STEEL_PRAGMA_UNROLL + for (short i = 0; i < n_rows; ++i) { + int n = read_n[i]; + int id = read_id[i] + flip_d * params->kdil[0]; + + bool in_bounds = n < params->N && id >= 0 && id < params->iS[0]; + + mask_d[i] |= (in_bounds << kd); + } + } + + for (short kh = 0; kh < params->wS[1]; kh++) { + short flip_h = params->flip ? params->wS[1] - kh - 1 : kh; + STEEL_PRAGMA_UNROLL + for (short i = 0; i < n_rows; ++i) { + int ih = read_ih[i] + flip_h * params->kdil[1]; + + bool in_bounds = ih >= 0 && ih < params->iS[1]; + + mask_h[i] |= (in_bounds << kh); + } + } + + for (short kw = 0; kw < params->wS[2]; kw++) { + short flip_w = params->flip ? params->wS[2] - kw - 1 : kw; + STEEL_PRAGMA_UNROLL + for (short i = 0; i < n_rows; ++i) { + int iw = read_iw[i] + flip_w * params->kdil[2]; + + bool in_bounds = iw >= 0 && iw < params->iS[2]; + + mask_w[i] |= (in_bounds << kw); + } + } + } + + /* Load from device memory into threadgroup memory - without bound checking */ + METAL_FUNC void load_unsafe() const { + mask_t d_mask = mask_t(1) << weight_d; + mask_t h_mask = mask_t(1) << weight_h; + mask_t w_mask = mask_t(1) << weight_w; + + STEEL_PRAGMA_UNROLL + for (short i = 0, is = 0; i < n_rows; ++i, is += TROWS) { + // Read from input if in bounds + if ((mask_d[i] & d_mask) && (mask_h[i] & h_mask) && + (mask_w[i] & w_mask)) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < vec_size; ++j) { + dst[is * dst_ld + j] = src[i][j]; + } + } + + // Zero pad otherwise + else { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < vec_size; ++j) { + dst[is * dst_ld + j] = T(0); + } + } + } + } + + /* Iteration helper */ + METAL_FUNC void next() { + if (++weight_w < params->wS[2]) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < n_rows; i++) { + src[i] += gemm_params->inp_jump_w; + } + + return; + } + + weight_w = 0; + + if (++weight_h < params->wS[1]) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < n_rows; i++) { + src[i] += gemm_params->inp_jump_h; + } + + return; + } + + weight_h = 0; + + if (++weight_d < params->wS[0]) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < n_rows; i++) { + src[i] += gemm_params->inp_jump_d; + } + + return; + } + + weight_d = 0; + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < n_rows; i++) { + src[i] += gemm_params->inp_jump_c; + } + } +}; + +template < + typename T, + short BM, + short BN, + short BK, + short tgp_size, + short tgp_padding = 0> +struct Conv3DWeightBlockLoader { + // Destination dimensions + STEEL_CONST short BROWS = BN; + STEEL_CONST short BCOLS = BK; + + // Read dimensions + STEEL_CONST short dst_ld = BCOLS + tgp_padding; + STEEL_CONST short vec_size = + (BN == 8) ? 1 : (tgp_size / (BROWS * BCOLS) >= 8 ? 8 : 4); + + // Thread read shape + STEEL_CONST short TCOLS = BCOLS / vec_size; + STEEL_CONST short TROWS = tgp_size / TCOLS; + + // Rows / strided reads within the block + STEEL_CONST short n_rows = BROWS / TROWS; + + // Leading dimension for src + const int src_ld; + + // Thread location indices + const short thread_idx; + const short bi; + const short bj; + + // threadgroup and device memory + threadgroup T* dst; + const device T* src; + + const constant MLXConvParams<3>* params; + + int weight_dhw; + int weight_step; + + const int read_n; + const bool do_read; + + /* Constructor */ + METAL_FUNC Conv3DWeightBlockLoader( + const device T* src_, + threadgroup T* dst_, + const int2 offsets, + const constant MLXConvParams<3>* params_, + const constant ImplicitGemmConv3DParams* gemm_params_, + uint simd_group_id [[simdgroup_index_in_threadgroup]], + uint simd_lane_id [[thread_index_in_simdgroup]]) + : src_ld(params_->wt_strides[0]), + thread_idx(simd_group_id * 32 + simd_lane_id), + bi(thread_idx / TCOLS), + bj(vec_size * (thread_idx % TCOLS)), + dst(dst_ + bi * dst_ld + bj), + src(src_ + bi * src_ld + bj), + params(params_), + weight_dhw(0), + weight_step(params->C / params->groups), + read_n(offsets.y + bi), + do_read(read_n + n_rows * TROWS <= gemm_params_->N) {} + + /* Load from device memory into threadgroup memory - without bound checking */ + METAL_FUNC void load_unsafe() const { + if (BN != 8 || do_read) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BN; i += TROWS) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < vec_size; j++) { + dst[i * dst_ld + j] = src[i * src_ld + j]; + } + } + } else { + for (short i = 0; i < BN; i += TROWS) { + if ((read_n + i) < params->O) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < vec_size; j++) { + dst[i * dst_ld + j] = src[i * src_ld + j]; + } + } else { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < vec_size; j++) { + dst[i * dst_ld + j] = T(0); + } + } + } + } + } + + /* Iteration helper */ + METAL_FUNC void next() { + if (++weight_dhw < (params->wS[0] * params->wS[1] * params->wS[2])) { + src += weight_step; + return; + } + + weight_dhw = 0; + + src += + BK - (params->wS[0] * params->wS[1] * params->wS[2] - 1) * weight_step; + } +}; + +} // namespace steel +} // namespace mlx + + +// ===== mlx/backend/metal/kernels/steel/conv/loaders/loader_channel_n.h ===== +// Copyright © 2024 Apple Inc. + + +// (already included: mlx/backend/metal/kernels/steel/utils.h) + +// (already included: mlx/backend/metal/kernels/steel/conv/params.h) + +/////////////////////////////////////////////////////////////////////////////// +// Loading helper +/////////////////////////////////////////////////////////////////////////////// + +namespace mlx { +namespace steel { + +template +struct ChannelHelper { + STEEL_CONST short n_channels = n_channels_; + STEEL_CONST short vec_size = n_channels_ <= 4 ? 4 : 8; + STEEL_CONST short excess = vec_size - n_channels_; +}; + +template <> +struct ChannelHelper<1> { + STEEL_CONST short n_channels = 1; + STEEL_CONST short vec_size = 1; + STEEL_CONST short excess = 0; +}; + +template <> +struct ChannelHelper<2> { + STEEL_CONST short n_channels = 2; + STEEL_CONST short vec_size = 2; + STEEL_CONST short excess = 0; +}; + +template <> +struct ChannelHelper<3> { + STEEL_CONST short n_channels = 3; + STEEL_CONST short vec_size = 4; + STEEL_CONST short excess = 1; +}; + +template <> +struct ChannelHelper<4> { + STEEL_CONST short n_channels = 4; + STEEL_CONST short vec_size = 4; + STEEL_CONST short excess = 0; +}; + +template < + typename T, + short BM, + short BN, + short BK, + short tgp_size, + short n_channels, + short tgp_padding = 0> +struct Conv2DInputBlockLoaderSmallChannels { + // Destination dimensions + STEEL_CONST short BROWS = BM; + STEEL_CONST short BCOLS = BK; + + // Read dimensions + STEEL_CONST short dst_ld = BCOLS + tgp_padding; + STEEL_CONST short vec_size = ChannelHelper::vec_size; + + // Thread read shape + STEEL_CONST short TCOLS = BCOLS / vec_size; + STEEL_CONST short TROWS = tgp_size / TCOLS; + + // Rows / strided reads within the block + STEEL_CONST short n_rows = BROWS / TROWS; + + // Thread location indices + const short thread_idx; + const short bi; + const short bj; + + // threadgroup and device memory + threadgroup T* dst; + + const constant MLXConvParams<2>* params; + const constant ImplicitGemmConv2DParams* gemm_params; + + int weight_hw; + + const device T* src[n_rows]; + + int read_n[n_rows]; + int read_ih[n_rows]; + int read_iw[n_rows]; + + /* Constructor */ + METAL_FUNC Conv2DInputBlockLoaderSmallChannels( + const device T* src_, + threadgroup T* dst_, + const int2 offsets, + const constant MLXConvParams<2>* params_, + const constant ImplicitGemmConv2DParams* gemm_params_, + uint simd_group_id [[simdgroup_index_in_threadgroup]], + uint simd_lane_id [[thread_index_in_simdgroup]]) + : thread_idx(simd_group_id * 32 + simd_lane_id), + bi(thread_idx / TCOLS), + bj(vec_size * (thread_idx % TCOLS)), + dst(dst_ + bi * dst_ld + bj), + params(params_), + gemm_params(gemm_params_), + weight_hw(thread_idx % TCOLS) { + int out_n_pixels = params->oS[0] * params->oS[1]; + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < n_rows; ++i) { + int offset_nhw = offsets.y + bi + i * TROWS; + int n = offset_nhw / out_n_pixels; + int hw = offset_nhw % out_n_pixels; + int oh = hw / params->oS[1]; + int ow = hw % params->oS[1]; + + int ih = oh * params->str[0] - params->pad[0]; + int iw = ow * params->str[1] - params->pad[1]; + + // Read from input if in bounds + src[i] = src_ + n * params->in_strides[0] + ih * params->in_strides[1] + + iw * params->in_strides[2]; + + read_n[i] = n; + read_ih[i] = ih; + read_iw[i] = iw; + } + } + + /* Load from device memory into threadgroup memory - without bound checking */ + METAL_FUNC void load_unsafe() const { + if (weight_hw >= params->wS[1] * params->wS[0]) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BROWS; i += TROWS) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < vec_size; j++) { + dst[i * dst_ld + j] = T(0); + } + } + return; + } + + int wh = (weight_hw / params->wS[1]); + int ww = (weight_hw % params->wS[1]); + + int flip_h = params->flip ? params->wS[0] - wh - 1 : wh; + int flip_w = params->flip ? params->wS[1] - ww - 1 : ww; + + int weight_h = flip_h * params->kdil[0]; + int weight_w = flip_w * params->kdil[1]; + + STEEL_PRAGMA_UNROLL + for (short i = 0, is = 0; i < n_rows; ++i, is += TROWS) { + // Find bounds + int n = read_n[i]; + int ih = read_ih[i] + weight_h; + int iw = read_iw[i] + weight_w; + + // Read from input if in bounds + if ((n < params->N) && (ih >= 0 && ih < params->iS[0]) && + (iw >= 0 && iw < params->iS[1])) { + const device T* curr_src = src[i] + weight_h * params->in_strides[1] + + weight_w * params->in_strides[2]; + + STEEL_PRAGMA_UNROLL + for (short j = 0; j < n_channels; ++j) { + dst[is * dst_ld + j] = curr_src[j]; + } + + STEEL_PRAGMA_UNROLL + for (short j = n_channels; j < vec_size; ++j) { + dst[is * dst_ld + j] = T(0); + } + } + + // Zero pad otherwise + else { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < vec_size; ++j) { + dst[is * dst_ld + j] = T(0); + } + } + } + } + + /* Iteration helper */ + METAL_FUNC void next() { + weight_hw += TCOLS; + } +}; + +template < + typename T, + short BM, + short BN, + short BK, + short tgp_size, + short n_channels, + short tgp_padding = 0> +struct Conv2DWeightBlockLoaderSmallChannels { + // Destination dimensions + STEEL_CONST short BROWS = BN; + STEEL_CONST short BCOLS = BK; + + // Read dimensions + STEEL_CONST short dst_ld = BCOLS + tgp_padding; + STEEL_CONST short vec_size = ChannelHelper::vec_size; + + // Thread read shape + STEEL_CONST short TCOLS = BCOLS / vec_size; + STEEL_CONST short TROWS = tgp_size / TCOLS; + + // Rows / strided reads within the block + STEEL_CONST short n_rows = BROWS / TROWS; + + // Leading dimension for src + const int src_ld; + + // Thread location indices + const short thread_idx; + const short bi; + const short bj; + + // threadgroup and device memory + threadgroup T* dst; + const device T* src; + + const constant MLXConvParams<2>* params; + + int weight_hw; + + const int read_n; + const bool do_read; + + /* Constructor */ + METAL_FUNC Conv2DWeightBlockLoaderSmallChannels( + const device T* src_, + threadgroup T* dst_, + const int2 offsets, + const constant MLXConvParams<2>* params_, + const constant ImplicitGemmConv2DParams* gemm_params_, + uint simd_group_id [[simdgroup_index_in_threadgroup]], + uint simd_lane_id [[thread_index_in_simdgroup]]) + : src_ld(params_->wt_strides[0]), + thread_idx(simd_group_id * 32 + simd_lane_id), + bi(thread_idx / TCOLS), + bj(vec_size * (thread_idx % TCOLS)), + dst(dst_ + bi * dst_ld + bj), + src(src_ + bi * src_ld), + params(params_), + weight_hw(thread_idx % TCOLS), + read_n(offsets.y + bi), + do_read(read_n + BN <= gemm_params_->N) {} + + /* Load from device memory into threadgroup memory - without bound checking */ + METAL_FUNC void load_unsafe() const { + if (bi >= BROWS || bj >= BCOLS) + return; + + if (read_n >= params->O || weight_hw >= params->wS[1] * params->wS[0]) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BROWS; i += TROWS) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < vec_size; j++) { + dst[i * dst_ld + j] = T(0); + } + } + + return; + } + + const device T* curr_src = src + weight_hw * (params->C / params->groups); + + if (BN != 8 || do_read) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BROWS; i += TROWS) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < n_channels; j++) { + dst[i * dst_ld + j] = curr_src[i * src_ld + j]; + } + + STEEL_PRAGMA_UNROLL + for (short j = n_channels; j < vec_size; j++) { + dst[i * dst_ld + j] = T(0); + } + } + } else { + for (short i = 0; i < BROWS; i += TROWS) { + if (((read_n + i) < params->O)) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < n_channels; j++) { + dst[i * dst_ld + j] = curr_src[i * src_ld + j]; + } + + STEEL_PRAGMA_UNROLL + for (short j = n_channels; j < vec_size; j++) { + dst[i * dst_ld + j] = T(0); + } + } else { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < vec_size; j++) { + dst[i * dst_ld + j] = T(0); + } + } + } + } + } + + /* Iteration helper */ + METAL_FUNC void next() { + weight_hw += TCOLS; + } +}; + +} // namespace steel +} // namespace mlx + +// (already included: mlx/backend/metal/kernels/steel/conv/params.h) +// (already included: mlx/backend/metal/kernels/steel/gemm/mma.h) + +using namespace metal; +using namespace mlx::steel; + +// (already included: mlx/backend/metal/kernels/steel/conv/params.h) +// (already included: mlx/backend/metal/kernels/steel/utils.h) + +// ===== mlx/backend/metal/kernels/steel/conv/kernels/steel_conv_general.h ===== +// Copyright © 2024 Apple Inc. + + +// ===== mlx/backend/metal/kernels/steel/conv/loaders/loader_general.h ===== +// Copyright © 2024 Apple Inc. + + +// (already included: mlx/backend/metal/kernels/steel/defines.h) + +/////////////////////////////////////////////////////////////////////////////// +// Loading helper +/////////////////////////////////////////////////////////////////////////////// + +namespace mlx { +namespace steel { + +template < + typename T, + short BM, + short BN, + short BK, + short tgp_size, + short tgp_padding = 0> +struct Conv2DInputBlockLoaderGeneral { + // Destination dimensions + STEEL_CONST short BROWS = BM; + STEEL_CONST short BCOLS = BK; + + // Read dimensions + STEEL_CONST short dst_ld = BCOLS + tgp_padding; + STEEL_CONST short vec_size = tgp_size / (BROWS * BCOLS) >= 8 ? 8 : 4; + + // Thread read shape + STEEL_CONST short TCOLS = BCOLS / vec_size; + STEEL_CONST short TROWS = tgp_size / TCOLS; + + // Rows / strided reads within the block + STEEL_CONST short n_rows = BROWS / TROWS; + + // Thread location indices + const short thread_idx; + const short bi; + const short bj; + + // threadgroup and device memory + threadgroup T* dst; + + const constant MLXConvParams<2>* params; + const constant Conv2DGeneralJumpParams* jump_params; + + const short base_wh; + const short base_ww; + + short weight_h; + short weight_w; + + const device T* src[n_rows]; + + int read_n[n_rows]; + int read_ih[n_rows]; + int read_iw[n_rows]; + + /* Constructor */ + METAL_FUNC Conv2DInputBlockLoaderGeneral( + const device T* src_, + threadgroup T* dst_, + const int4 offsets, + const constant MLXConvParams<2>* params_, + const constant Conv2DGeneralJumpParams* jump_params_, + const short base_wh_, + const short base_ww_, + uint simd_group_id [[simdgroup_index_in_threadgroup]], + uint simd_lane_id [[thread_index_in_simdgroup]]) + : thread_idx(simd_group_id * 32 + simd_lane_id), + bi(thread_idx / TCOLS), + bj(vec_size * (thread_idx % TCOLS)), + dst(dst_ + bi * dst_ld + bj), + params(params_), + jump_params(jump_params_), + base_wh(base_wh_), + base_ww(base_ww_), + weight_h(base_wh_), + weight_w(base_ww_) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < n_rows; ++i) { + int offset_nhw = offsets.y + bi + i * TROWS; + int n = offset_nhw / jump_params->adj_out_hw; + int hw = offset_nhw % jump_params->adj_out_hw; + int oh = + (hw / jump_params->adj_out_w) * jump_params->f_out_jump_h + offsets.z; + int ow = + (hw % jump_params->adj_out_w) * jump_params->f_out_jump_w + offsets.w; + + int ih = oh * params->str[0] - params->pad[0]; + int iw = ow * params->str[1] - params->pad[1]; + + read_n[i] = n; + read_ih[i] = ih; + read_iw[i] = iw; + + // Read from input if in bounds + src[i] = src_ + n * params->in_strides[0] + bj; + } + } + + /* Load from device memory into threadgroup memory - without bound checking */ + METAL_FUNC void load_unsafe() const { + STEEL_PRAGMA_UNROLL + for (short i = 0, is = 0; i < n_rows; ++i, is += TROWS) { + // Find bounds + int n = read_n[i]; + + int h_flip = params->flip ? params->wS[0] - weight_h - 1 : weight_h; + int w_flip = params->flip ? params->wS[1] - weight_w - 1 : weight_w; + + int ih_dil = read_ih[i] + h_flip * params->kdil[0]; + int iw_dil = read_iw[i] + w_flip * params->kdil[1]; + + int ih = ih_dil / params->idil[0]; + int iw = iw_dil / params->idil[1]; + + size_t offset = ih * params->in_strides[1] + iw * params->in_strides[2]; + + // Read from input if in bounds + if ((n < params->N) && (ih_dil >= 0 && ih < params->iS[0]) && + (iw_dil >= 0 && iw < params->iS[1])) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < vec_size; ++j) { + dst[is * dst_ld + j] = (src[i])[offset + j]; + } + } + + // Zero pad otherwise + else { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < vec_size; ++j) { + dst[is * dst_ld + j] = T(0); + } + } + } + } + + METAL_FUNC void load_safe(const short remaining_k) const { + STEEL_PRAGMA_UNROLL + for (short i = 0, is = 0; i < n_rows; ++i, is += TROWS) { + // Find bounds + int n = read_n[i]; + + int h_flip = params->flip ? params->wS[0] - weight_h - 1 : weight_h; + int w_flip = params->flip ? params->wS[1] - weight_w - 1 : weight_w; + + int ih_dil = read_ih[i] + h_flip * params->kdil[0]; + int iw_dil = read_iw[i] + w_flip * params->kdil[1]; + + int ih = ih_dil / params->idil[0]; + int iw = iw_dil / params->idil[1]; + + size_t offset = ih * params->in_strides[1] + iw * params->in_strides[2]; + + // Read from input if in bounds + if ((n < params->N) && (ih_dil >= 0 && ih < params->iS[0]) && + (iw_dil >= 0 && iw < params->iS[1])) { + if (bj + vec_size <= remaining_k) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < vec_size; ++j) { + dst[is * dst_ld + j] = (src[i])[offset + j]; + } + } else { + for (short j = 0; j < vec_size; ++j) { + if (bj + j < remaining_k) { + dst[is * dst_ld + j] = (src[i])[offset + j]; + } else { + dst[is * dst_ld + j] = T(0); + } + } + } + } + + // Zero pad otherwise + else { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < vec_size; ++j) { + dst[is * dst_ld + j] = T(0); + } + } + } + } + + /* Iteration helper */ + METAL_FUNC void next() { + weight_w += jump_params->f_wgt_jump_w; + if (weight_w < params->wS[1]) { + return; + } + + weight_w = base_ww; + + weight_h += jump_params->f_wgt_jump_h; + if (weight_h < params->wS[0]) { + return; + } + + weight_h = base_wh; + + STEEL_PRAGMA_UNROLL + for (short i = 0; i < n_rows; i++) { + src[i] += BK; + } + } +}; + +template < + typename T, + short BM, + short BN, + short BK, + short tgp_size, + short tgp_padding = 0> +struct Conv2DWeightBlockLoaderGeneral { + // Destination dimensions + STEEL_CONST short BROWS = BN; + STEEL_CONST short BCOLS = BK; + + // Read dimensions + STEEL_CONST short dst_ld = BCOLS + tgp_padding; + STEEL_CONST short vec_size = + (BN == 8) ? 1 : (tgp_size / (BROWS * BCOLS) >= 8 ? 8 : 4); + + // Thread read shape + STEEL_CONST short TCOLS = BCOLS / vec_size; + STEEL_CONST short TROWS = tgp_size / TCOLS; + + // Rows / strided reads within the block + STEEL_CONST short n_rows = BROWS / TROWS; + + // Leading dimension for src + const int src_ld; + + // Thread location indices + const short thread_idx; + const short bi; + const short bj; + + // threadgroup and device memory + threadgroup T* dst; + const device T* src; + + const constant MLXConvParams<2>* params; + const constant Conv2DGeneralJumpParams* jump_params; + + const short base_wh; + const short base_ww; + + short weight_h; + short weight_w; + + const int start_row; + + /* Constructor */ + METAL_FUNC Conv2DWeightBlockLoaderGeneral( + const device T* src_, + threadgroup T* dst_, + const int2 offsets, + const constant MLXConvParams<2>* params_, + const constant Conv2DGeneralJumpParams* jump_params_, + const short base_wh_, + const short base_ww_, + uint simd_group_id [[simdgroup_index_in_threadgroup]], + uint simd_lane_id [[thread_index_in_simdgroup]]) + : src_ld(params_->wt_strides[0]), + thread_idx(simd_group_id * 32 + simd_lane_id), + bi(thread_idx / TCOLS), + bj(vec_size * (thread_idx % TCOLS)), + dst(dst_ + bi * dst_ld + bj), + src(src_ + bi * src_ld + bj), + params(params_), + jump_params(jump_params_), + base_wh(base_wh_), + base_ww(base_ww_), + weight_h(base_wh_), + weight_w(base_ww_), + start_row(offsets.y + bi) {} + + /* Load from device memory into threadgroup memory - without bound checking */ + METAL_FUNC void load_unsafe() const { + const device T* curr_src = src + weight_h * params->wt_strides[1] + + weight_w * params->wt_strides[2]; + + if ((start_row + BN <= params->O)) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BN; i += TROWS) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < vec_size; j++) { + dst[i * dst_ld + j] = curr_src[i * src_ld + j]; + } + } + } else { + for (short i = 0; i < BN; i += TROWS) { + if ((start_row + i) < params->O) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < vec_size; j++) { + dst[i * dst_ld + j] = curr_src[i * src_ld + j]; + } + } else { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < vec_size; j++) { + dst[i * dst_ld + j] = T(0); + } + } + } + } + } + + METAL_FUNC void load_safe(const short remaining_k) const { + const device T* curr_src = src + weight_h * params->wt_strides[1] + + weight_w * params->wt_strides[2]; + + if ((start_row + BN <= params->O)) { + STEEL_PRAGMA_UNROLL + for (short i = 0; i < BN; i += TROWS) { + if (bj + vec_size <= remaining_k) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < vec_size; j++) { + dst[i * dst_ld + j] = curr_src[i * src_ld + j]; + } + } else { + for (short j = 0; j < vec_size; j++) { + if (bj + j < remaining_k) { + dst[i * dst_ld + j] = curr_src[i * src_ld + j]; + } else { + dst[i * dst_ld + j] = T(0); + } + } + } + } + } else { + for (short i = 0; i < BN; i += TROWS) { + if ((start_row + i) < params->O) { + if (bj + vec_size <= remaining_k) { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < vec_size; j++) { + dst[i * dst_ld + j] = curr_src[i * src_ld + j]; + } + } else { + for (short j = 0; j < vec_size; j++) { + if (bj + j < remaining_k) { + dst[i * dst_ld + j] = curr_src[i * src_ld + j]; + } else { + dst[i * dst_ld + j] = T(0); + } + } + } + } else { + STEEL_PRAGMA_UNROLL + for (short j = 0; j < vec_size; j++) { + dst[i * dst_ld + j] = T(0); + } + } + } + } + } + + /* Iteration helper */ + METAL_FUNC void next() { + weight_w += jump_params->f_wgt_jump_w; + if (weight_w < params->wS[1]) { + return; + } + + weight_w = base_ww; + + weight_h += jump_params->f_wgt_jump_h; + if (weight_h < params->wS[0]) { + return; + } + + weight_h = base_wh; + + src += BK; + } +}; + +} // namespace steel +} // namespace mlx + + +constant bool align_C [[function_constant(200)]]; + +template < + typename T, + int BM, + int BN, + int BK, + int WM, + int WN, + typename AccumType = float, + typename Epilogue = TransformNone> +[[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)]], + device T* C [[buffer(2)]], + const constant MLXConvParams<2>* params [[buffer(3)]], + const constant ImplicitGemmConv2DParams* gemm_params [[buffer(4)]], + const constant Conv2DGeneralJumpParams* jump_params [[buffer(5)]], + const constant Conv2DGeneralBaseInfo* base_h [[buffer(6)]], + const constant Conv2DGeneralBaseInfo* base_w [[buffer(7)]], + uint3 tid [[threadgroup_position_in_grid]], + uint3 lid [[thread_position_in_threadgroup]], + uint simd_gid [[simdgroup_index_in_threadgroup]], + uint simd_lid [[thread_index_in_simdgroup]]) { + (void)lid; + + constexpr bool transpose_a = false; + constexpr bool transpose_b = true; + constexpr short tgp_padding_a = 16 / sizeof(T); + constexpr short tgp_padding_b = 16 / sizeof(T); + + constexpr short shape_a_cols = (transpose_a ? BM : BK) + tgp_padding_a; + constexpr short shape_b_cols = (transpose_b ? BK : BN) + tgp_padding_b; + constexpr short shape_a_rows = (transpose_a ? BK : BM); + constexpr short shape_b_rows = (transpose_b ? BN : BK); + constexpr short tgp_mem_size_a = shape_a_cols * shape_a_rows; + constexpr short tgp_mem_size_b = shape_b_cols * shape_b_rows; + + constexpr short tgp_size = WM * WN * 32; + + // Input loader + using loader_a_t = + Conv2DInputBlockLoaderGeneral; + + // Weight loader + using loader_b_t = + Conv2DWeightBlockLoaderGeneral; + + using mma_t = BlockMMA< + T, + T, + BM, + BN, + BK, + WM, + WN, + transpose_a, + transpose_b, + shape_a_cols, + shape_b_cols>; + + threadgroup T As[tgp_mem_size_a]; + threadgroup T Bs[tgp_mem_size_b]; + + const int tid_y = ((tid.y) << gemm_params->swizzle_log) + + ((tid.x) & ((1 << gemm_params->swizzle_log) - 1)); + const int tid_x = (tid.x) >> gemm_params->swizzle_log; + + if (gemm_params->tiles_n <= tid_x || gemm_params->tiles_m <= tid_y) { + return; + } + + const int tid_z = tid.z; + + const int base_oh = tid_z / jump_params->f_out_jump_w; + const int base_ow = tid_z % jump_params->f_out_jump_w; + + const int base_wh = base_h[base_oh].weight_base; + const int base_ww = base_w[base_ow].weight_base; + + const int base_wh_size = base_h[base_oh].weight_size; + const int base_ww_size = base_w[base_ow].weight_size; + + const int c_row = tid_y * BM; + const int c_col = tid_x * BN; + const int K = gemm_params->K; + + B += c_col * K; + + const int4 offsets_a(0, c_row, base_oh, base_ow); + const int2 offsets_b(0, c_col); + + // Prepare threadgroup loading operations + loader_a_t loader_a( + A, + As, + offsets_a, + params, + jump_params, + base_wh, + base_ww, + simd_gid, + simd_lid); + loader_b_t loader_b( + B, + Bs, + offsets_b, + params, + jump_params, + base_wh, + base_ww, + simd_gid, + simd_lid); + + // Prepare threadgroup mma operation + mma_t mma_op(simd_gid, simd_lid); + + if (align_C) { + int gemm_k_iterations = + base_wh_size * base_ww_size * gemm_params->gemm_k_iterations; + + for (int k = 0; k < gemm_k_iterations; k++) { + threadgroup_barrier(mem_flags::mem_threadgroup); + // Load elements into threadgroup + loader_a.load_unsafe(); + loader_b.load_unsafe(); + + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Multiply and accumulate threadgroup elements + mma_op.mma(As, Bs); + + // Prepare for next iteration + loader_a.next(); + loader_b.next(); + } + } + + else { + for (int k = 1; k < gemm_params->gemm_k_iterations; k++) { + for (int j = 0; j < base_wh_size * base_ww_size; j++) { + threadgroup_barrier(mem_flags::mem_threadgroup); + // Load elements into threadgroup + loader_a.load_unsafe(); + loader_b.load_unsafe(); + + threadgroup_barrier(mem_flags::mem_threadgroup); + + // Multiply and accumulate threadgroup elements + mma_op.mma(As, Bs); + + // Prepare for next iteration + loader_a.next(); + loader_b.next(); + } + } + const short remaining_k = params->C % BK; + for (int j = 0; j < base_wh_size * base_ww_size; j++) { + // Load elements into threadgroup + threadgroup_barrier(mem_flags::mem_threadgroup); + loader_a.load_safe(remaining_k); + loader_b.load_safe(remaining_k); + threadgroup_barrier(mem_flags::mem_threadgroup); + // Multiply and accumulate threadgroup elements + mma_op.mma(As, Bs); + // Prepare for next iteration + loader_a.next(); + loader_b.next(); + } + } + + threadgroup_barrier(mem_flags::mem_none); + + // Store results to device memory + { + // Adjust for simdgroup and thread location + int offset_m = c_row + mma_op.sm; + int offset_n = c_col + mma_op.sn; + C += offset_n; + + if (offset_n >= gemm_params->N) + return; + + short diff = gemm_params->N - offset_n; + + STEEL_PRAGMA_UNROLL + for (int i = 0; i < mma_t::TM; i++) { + int cm = offset_m + i * mma_t::TM_stride; + + int n = cm / jump_params->adj_out_hw; + int hw = cm % jump_params->adj_out_hw; + int oh = + (hw / jump_params->adj_out_w) * jump_params->f_out_jump_h + base_oh; + int ow = + (hw % jump_params->adj_out_w) * jump_params->f_out_jump_w + base_ow; + + if (n < params->N && oh < params->oS[0] && ow < params->oS[1]) { + size_t offset_cm = static_cast(n) * params->out_strides[0] + + oh * params->out_strides[1] + ow * params->out_strides[2]; + + STEEL_PRAGMA_UNROLL + for (int j = 0; j < mma_t::TN; j++) { + // Get accumulated result and associated offset in C + thread const auto& accum = mma_op.Ctile.frag_at(i, j); + size_t offset = offset_cm + (j * mma_t::TN_stride); + + constexpr short kelems = decltype(mma_op.Ctile)::kElemsPerFrag; + + // Apply epilogue and output C + STEEL_PRAGMA_UNROLL + for (short k = 0; k < kelems; k++) { + if ((j * mma_t::TN_stride + k) < diff) { + C[offset + k] = Epilogue::apply(accum[k]); + } + } + } + } + } + } +} + + +using namespace metal; +using namespace mlx::steel; + +#define instantiate_implicit_conv_2d(name, itype, bm, bn, bk, wm, wn) \ + template \ + [[host_name("implicit_gemm_conv_2d_general_" #name "_bm" #bm "_bn" #bn \ + "_bk" #bk "_wm" #wm "_wn" #wn)]] [[kernel]] void \ + implicit_gemm_conv_2d_general( \ + const device itype* A [[buffer(0)]], \ + const device itype* B [[buffer(1)]], \ + device itype* C [[buffer(2)]], \ + const constant MLXConvParams<2>* params [[buffer(3)]], \ + const constant ImplicitGemmConv2DParams* gemm_params [[buffer(4)]], \ + const constant Conv2DGeneralJumpParams* jump_params [[buffer(5)]], \ + const constant Conv2DGeneralBaseInfo* base_h [[buffer(6)]], \ + const constant Conv2DGeneralBaseInfo* base_w [[buffer(7)]], \ + uint3 tid [[threadgroup_position_in_grid]], \ + uint3 lid [[thread_position_in_threadgroup]], \ + uint simd_gid [[simdgroup_index_in_threadgroup]], \ + uint simd_lid [[thread_index_in_simdgroup]]); + +#define instantiate_implicit_2d_filter(name, itype, bm, bn, bk, wm, wn) \ + instantiate_implicit_conv_2d(name, itype, bm, bn, bk, wm, wn) + +#define instantiate_implicit_2d_blocks(name, itype) \ + instantiate_implicit_2d_filter(name, itype, 32, 8, 16, 4, 1) \ + instantiate_implicit_2d_filter(name, itype, 64, 8, 16, 4, 1) \ + instantiate_implicit_2d_filter(name, itype, 32, 32, 16, 2, 2) \ + instantiate_implicit_2d_filter(name, itype, 32, 64, 16, 2, 2) \ + instantiate_implicit_2d_filter(name, itype, 64, 32, 16, 2, 2) \ + instantiate_implicit_2d_filter(name, itype, 64, 64, 16, 2, 2) + +instantiate_implicit_2d_blocks(float32, float); +instantiate_implicit_2d_blocks(float16, half); +// clang-format on + diff --git a/metal/src/kernels/conv/mlx_conv.rs b/metal/src/kernels/conv/mlx_conv.rs new file mode 100644 index 0000000000..792ceab799 --- /dev/null +++ b/metal/src/kernels/conv/mlx_conv.rs @@ -0,0 +1,275 @@ +//! Implicit-GEMM 2D convolution via the ported MLX kernel (see mlx_conv.metal). +//! +//! Covers NHWC f16/f32 2D convolutions with a single group and an `OHWI` +//! kernel — the layout the kernel indexes directly. Everything else is left to +//! the direct kernel by `mlx_conv_supported`. + +use crate::encoder::EncoderExt; +use crate::{ConstantValues, LibraryName, MetalStream, Value}; +use anyhow::ensure; +use metal::MTLSize; +use tract_core::internal::*; +use tract_core::ops::cnn::{Conv, KernelFormat}; +use tract_gpu::tensor::DeviceTensor; + +/// Mirror of MLX `MLXConvParams<2>` (steel/conv/params.h) — keep field order in sync. +#[repr(C)] +#[derive(Debug, Clone, Copy)] +struct MlxConvParams2D { + n: i32, + c: i32, + o: i32, + i_s: [i32; 2], + w_s: [i32; 2], + o_s: [i32; 2], + str: [i32; 2], + pad: [i32; 2], + kdil: [i32; 2], + idil: [i32; 2], + in_strides: [i64; 4], + wt_strides: [i64; 4], + out_strides: [i64; 4], + groups: i32, + flip: bool, +} + +/// Mirror of MLX `ImplicitGemmConv2DParams`. +#[repr(C)] +#[derive(Debug, Clone, Copy)] +struct ImplicitGemmConv2DParams { + m: i32, + n: i32, + k: i32, + gemm_k_iterations: i32, + inp_jump_w: i32, + inp_jump_h: i32, + inp_jump_c: i32, + tiles_n: i32, + tiles_m: i32, + swizzle_log: i32, +} + +/// Mirror of MLX `Conv2DGeneralJumpParams`. +#[repr(C)] +#[derive(Debug, Clone, Copy)] +struct Conv2DGeneralJumpParams { + f_wgt_jump_h: i32, + f_wgt_jump_w: i32, + f_out_jump_h: i32, + f_out_jump_w: i32, + adj_out_h: i32, + adj_out_w: i32, + adj_out_hw: i32, + adj_implicit_m: i32, +} + +/// Mirror of MLX `Conv2DGeneralBaseInfo`. +#[repr(C)] +#[derive(Debug, Clone, Copy)] +struct Conv2DGeneralBaseInfo { + weight_base: i32, + weight_size: i32, +} + +fn gcd(a: i32, b: i32) -> i32 { + if b == 0 { a } else { gcd(b, a % b) } +} + +fn lcm(a: i32, b: i32) -> i32 { + a / gcd(a, b) * b +} + +/// Whether the ported kernel can take this convolution: NHWC f16/f32, one +/// group, rank-2 spatial, `OHWI` weights, no output padding tricks. +pub fn mlx_conv_eligible(op: &Conv, in_facts: &[&TypedFact]) -> bool { + if op.group != 1 || op.q_params.is_some() { + return false; + } + if !matches!(in_facts[0].datum_type, DatumType::F16 | DatumType::F32) { + return false; + } + if in_facts[0].datum_type != in_facts[1].datum_type { + return false; + } + let Ok(shape) = op.pool_spec.data_format.shape(in_facts[0].shape.to_tvec()) else { + return false; + }; + if shape.hw_rank() != 2 || !op.pool_spec.data_format.c_is_last() { + return false; + } + in_facts.iter().all(|f| f.shape.as_concrete().is_some()) +} + +/// `out[N, oH, oW, O] = conv(in[N, iH, iW, C], wt[O, kH, kW, C])`, mirroring +/// mlx `implicit_gemm_conv_2D_general_gpu`. +pub fn dispatch_mlx_conv_2d( + stream: &MetalStream, + op: &Conv, + input: &DeviceTensor, + weights: &DeviceTensor, + output: &DeviceTensor, +) -> TractResult<()> { + let dt = input.datum_type(); + let tname = match dt { + DatumType::F32 => "float32", + DatumType::F16 => "float16", + _ => bail!("MLX conv: F32/F16 only, got {dt:?}"), + }; + let in_shape = op.pool_spec.data_format.shape(input.shape())?; + let out_shape = op.pool_spec.data_format.shape(output.shape())?; + ensure!(in_shape.hw_rank() == 2, "MLX conv is 2D only"); + + let n = *in_shape.n().unwrap_or(&1) as i32; + let c = *in_shape.c() as i32; + let o = *out_shape.c() as i32; + let i_s = [in_shape.hw_dims()[0] as i32, in_shape.hw_dims()[1] as i32]; + let w_s = [weights.shape()[1] as i32, weights.shape()[2] as i32]; + let o_s = [out_shape.hw_dims()[0] as i32, out_shape.hw_dims()[1] as i32]; + let strides = op.pool_spec.strides(); + let dilations = op.pool_spec.dilations(); + let str = [strides[0] as i32, strides[1] as i32]; + let kdil = [dilations[0] as i32, dilations[1] as i32]; + let idil = [1i32, 1]; + let padding = op.pool_spec.computed_padding(in_shape.hw_dims()); + let pad = [padding[0].pad_before as i32, padding[1].pad_before as i32]; + + let to4 = |s: &[isize]| -> [i64; 4] { [s[0] as i64, s[1] as i64, s[2] as i64, s[3] as i64] }; + let ceil_div = |a: i32, b: i32| (a + b - 1) / b; + let conv_params = MlxConvParams2D { + n, + c, + o, + i_s, + w_s, + o_s, + str, + pad, + kdil, + idil, + in_strides: to4(input.strides()), + wt_strides: to4(weights.strides()), + out_strides: to4(output.strides()), + groups: 1, + flip: false, + }; + + let implicit_m = n * o_s[0] * o_s[1]; + let implicit_n = o; + let (wm, wn) = (2i32, 2i32); + + let f_wgt_jump_h = lcm(idil[0], kdil[0]) / kdil[0]; + let f_wgt_jump_w = lcm(idil[1], kdil[1]) / kdil[1]; + let f_out_jump_h = lcm(idil[0], str[0]) / str[0]; + let f_out_jump_w = lcm(idil[1], str[1]) / str[1]; + let adj_out_h = ceil_div(o_s[0], f_out_jump_h); + let adj_out_w = ceil_div(o_s[1], f_out_jump_w); + let adj_out_hw = adj_out_h * adj_out_w; + let adj_implicit_m = n * adj_out_hw; + let jump_params = Conv2DGeneralJumpParams { + f_wgt_jump_h, + f_wgt_jump_w, + f_out_jump_h, + f_out_jump_w, + adj_out_h, + adj_out_w, + adj_out_hw, + adj_implicit_m, + }; + + let base_of = |jumps: i32, stride: i32, pad: i32, ws: i32, kdil: i32, wgt_jump: i32| { + (0..jumps) + .map(|i| { + let mut loop_pos = i * stride - pad; + let mut base = 0; + while base < ws && loop_pos % idil[0] != 0 { + base += 1; + loop_pos += kdil; + } + Conv2DGeneralBaseInfo { + weight_base: base, + weight_size: (ws - base + wgt_jump - 1) / wgt_jump, + } + }) + .collect::>() + }; + let base_h = base_of(f_out_jump_h, str[0], pad[0], w_s[0], kdil[0], f_wgt_jump_h); + let base_w = base_of(f_out_jump_w, str[1], pad[1], w_s[1], kdil[1], f_wgt_jump_w); + + let bm = if adj_implicit_m >= 8192 && c >= 64 { 64 } else { 32 }; + let bn = if bm == 64 && implicit_n >= 64 { 64 } else { 32 }; + let bk = 16i32; + let tn = ceil_div(implicit_n, bn); + let tm = ceil_div(adj_implicit_m, bm); + let swizzle_log = 0i32; + let align_c = c % bk == 0; + + let ijw = conv_params.in_strides[2] as i32 * kdil[1]; + let ijh = conv_params.in_strides[1] as i32 * kdil[0]; + let gemm_params = ImplicitGemmConv2DParams { + m: implicit_m, + n: implicit_n, + k: w_s[0] * w_s[1] * c, + gemm_k_iterations: ceil_div(c, bk), + inp_jump_w: ijw, + inp_jump_h: ijh - (w_s[1] - 1) * ijw, + inp_jump_c: bk - (w_s[0] - 1) * ijh - (w_s[1] - 1) * ijw, + tiles_n: tn, + tiles_m: tm, + swizzle_log, + }; + + let name = format!("implicit_gemm_conv_2d_general_{tname}_bm{bm}_bn{bn}_bk{bk}_wm{wm}_wn{wn}"); + let constants = Some(ConstantValues::new(vec![(200, Value::Bool(align_c))])); + let pipeline = stream.load_pipeline_with_constants(LibraryName::MlxConv, &name, constants)?; + + stream.retain_tensor(input); + stream.retain_tensor(weights); + stream.retain_tensor(output); + + let command_buffer = stream.command_buffer(); + command_buffer.encode(|encoder| { + encoder.set_compute_pipeline_state(&pipeline); + encoder.set_metal_tensor(0, input, metal::MTLResourceUsage::Read); + encoder.set_metal_tensor(1, weights, metal::MTLResourceUsage::Read); + encoder.set_metal_tensor(2, output, metal::MTLResourceUsage::Write); + encoder.set_slice(3, std::slice::from_ref(&conv_params)); + encoder.set_slice(4, std::slice::from_ref(&gemm_params)); + encoder.set_slice(5, std::slice::from_ref(&jump_params)); + encoder.set_slice(6, &base_h); + encoder.set_slice(7, &base_w); + let tile = 1 << swizzle_log; + let grid = MTLSize { + width: (tn * tile) as _, + height: ceil_div(tm, tile) as _, + depth: (f_out_jump_h * f_out_jump_w) as _, + }; + let group = MTLSize { width: 32, height: wn as _, depth: wm as _ }; + encoder.dispatch_thread_groups(grid, group); + }); + Ok(()) +} + +/// Runtime counterpart of `mlx_conv_supported`, checked against the device +/// tensors actually being dispatched. +pub fn mlx_conv_dispatchable(op: &Conv, input: &DeviceTensor, weights: &DeviceTensor) -> bool { + if op.group != 1 || op.q_params.is_some() || op.kernel_fmt != KernelFormat::OHWI { + return false; + } + if !matches!(input.datum_type(), DatumType::F16 | DatumType::F32) + || input.datum_type() != weights.datum_type() + { + return false; + } + if !op.pool_spec.data_format.c_is_last() || input.rank() != 4 || weights.rank() != 4 { + return false; + } + let natural = |t: &DeviceTensor| { + let mut s = 1isize; + t.shape().iter().rev().zip(t.strides().iter().rev()).all(|(&d, &st)| { + let ok = st == s; + s *= d as isize; + ok + }) + }; + natural(input) && natural(weights) +} diff --git a/metal/src/kernels/conv/mod.rs b/metal/src/kernels/conv/mod.rs new file mode 100644 index 0000000000..231d66f162 --- /dev/null +++ b/metal/src/kernels/conv/mod.rs @@ -0,0 +1,600 @@ +pub mod mlx_conv; + +use crate::encoder::EncoderExt; +use crate::{LibraryName, MetalStream}; +use metal::MTLSize; +use tract_core::internal::*; +use tract_core::ops::cnn::Conv; +use tract_gpu::tensor::DeviceTensor; + +pub fn kernel_name(hw_rank: usize, dt: DatumType) -> TractResult { + let dt_name = if dt == DatumType::F16 { "f16" } else { "f32" }; + Ok(format!("conv{hw_rank}d_{dt_name}_generic")) +} + +pub fn metal_conv_direct( + stream: &MetalStream, + op: &Conv, + input: &DeviceTensor, + weights: &DeviceTensor, + bias: Option<&DeviceTensor>, + output: &DeviceTensor, +) -> TractResult<()> { + metal_conv_dispatch_inner(stream, op, input, weights, bias, output) +} + +pub fn metal_conv_dispatch( + stream: &MetalStream, + op: &Conv, + input: &DeviceTensor, + weights: &DeviceTensor, + bias: Option<&DeviceTensor>, + output: &DeviceTensor, +) -> TractResult<()> { + // The MLX kernel takes bias separately (it is added by a following op), so + // it only handles the bias-free dispatch here. + if bias.is_none() && mlx_conv::mlx_conv_dispatchable(op, input, weights) { + return mlx_conv::dispatch_mlx_conv_2d(stream, op, input, weights, output); + } + metal_conv_dispatch_inner(stream, op, input, weights, bias, output) +} + +fn metal_conv_dispatch_inner( + stream: &MetalStream, + op: &Conv, + input: &DeviceTensor, + weights: &DeviceTensor, + bias: Option<&DeviceTensor>, + output: &DeviceTensor, +) -> TractResult<()> { + stream.retain_tensor(input); + stream.retain_tensor(weights); + if let Some(b) = bias { + stream.retain_tensor(b); + } + stream.retain_tensor(output); + + let input_shape = op.pool_spec.data_format.shape(input.shape())?; + let hw_rank = input_shape.hw_rank(); + let func_name = kernel_name(hw_rank, input.datum_type())?; + let pipeline = stream.load_pipeline(LibraryName::ConvOps, &func_name)?; + + let co_per_group = op.pool_spec.output_channels / op.group; + let ci_per_group = op.pool_spec.input_channels / op.group; + + // in_shape: [N, C, spatial...] + let in_n = *input_shape.n().unwrap_or(&1); + let in_c = *input_shape.c(); + let mut in_shape_buf: TVec = tvec![in_n as i32, in_c as i32]; + in_shape_buf.extend(input_shape.hw_dims().iter().map(|&d| d as i32)); + + let mut in_strides_buf: TVec = + tvec![*input_shape.n_stride().unwrap_or(&0) as i32, *input_shape.c_stride() as i32]; + in_strides_buf.extend(input_shape.hw_strides().iter().map(|&s| s as i32)); + + // ker_params: [groups, co_per_group, ci_per_group, ker_spatial...] + let mut ker_params: TVec = + tvec![op.group as i32, co_per_group as i32, ci_per_group as i32]; + ker_params.extend(weights.shape()[2..].iter().map(|&d| d as i32)); + + // ker_strides: [g_stride, o_stride, i_stride, spatial...] + let group_stride = weights.strides()[0] as usize * co_per_group; + let mut ker_strides: TVec = tvec![group_stride as i32]; + ker_strides.extend(weights.strides().iter().map(|&s| s as i32)); + + // padding + let padding = op.pool_spec.computed_padding(input_shape.hw_dims()); + let pad_buf: TVec = padding.iter().map(|p| p.pad_before as i32).collect(); + + let strides = op.pool_spec.strides(); + let strides_buf: TVec = strides.iter().map(|&s| s as i32).collect(); + + let dilations = op.pool_spec.dilations(); + let dilations_buf: TVec = dilations.iter().map(|&d| d as i32).collect(); + + let output_shape = op.pool_spec.data_format.shape(output.shape())?; + let out_n = *output_shape.n().unwrap_or(&1); + let out_c = *output_shape.c(); + let mut out_shape_buf: TVec = tvec![out_n as i32, out_c as i32]; + out_shape_buf.extend(output_shape.hw_dims().iter().map(|&d| d as i32)); + + let mut out_strides_buf: TVec = + tvec![*output_shape.n_stride().unwrap_or(&0) as i32, *output_shape.c_stride() as i32]; + out_strides_buf.extend(output_shape.hw_strides().iter().map(|&s| s as i32)); + + // bias_stride: -1 means no bias, 0 means scalar broadcast, 1 means per-channel + let bias_stride: i32 = if let Some(b) = bias { if b.rank() == 0 { 0 } else { 1 } } else { -1 }; + + let spatial_out: usize = output_shape.hw_dims().iter().product(); + let threads_per_group = 32usize; + + let command_buffer = stream.command_buffer(); + command_buffer.encode(|encoder| { + encoder.set_compute_pipeline_state(&pipeline); + encoder.set_metal_tensor(0, input, metal::MTLResourceUsage::Read); + encoder.set_slice(1, &in_shape_buf); + encoder.set_slice(2, &in_strides_buf); + encoder.set_metal_tensor(3, weights, metal::MTLResourceUsage::Read); + encoder.set_slice(4, &ker_params); + encoder.set_slice(5, &ker_strides); + if let Some(b) = bias { + encoder.set_metal_tensor(6, b, metal::MTLResourceUsage::Read); + } else { + // Empty buffer — kernel checks bias_stride < 0 + encoder.set_bytes(6, 0, std::ptr::null()); + } + encoder.set_slice(7, &[bias_stride]); + encoder.set_slice(8, &pad_buf); + encoder.set_slice(9, &strides_buf); + encoder.set_slice(10, &dilations_buf); + encoder.set_metal_tensor(11, output, metal::MTLResourceUsage::Write); + encoder.set_slice(12, &out_shape_buf); + encoder.set_slice(13, &out_strides_buf); + + let grid_size = MTLSize { + width: spatial_out.div_ceil(threads_per_group) as _, + height: out_c as _, + depth: out_n as _, + }; + let group_size = MTLSize { width: threads_per_group as _, height: 1, depth: 1 }; + encoder.dispatch_thread_groups(grid_size, group_size); + }); + Ok(()) +} + +#[cfg(test)] +mod mlx_conv_tests { + use crate::LibraryName; + use crate::utils::with_borrowed_metal_stream; + use tract_core::internal::*; + + use super::mlx_conv::dispatch_mlx_conv_2d; + use tract_core::ops::cnn::{Conv, KernelFormat, PaddingSpec, PoolSpec}; + use tract_core::ops::nn::DataFormat; + use tract_gpu::tensor::{DeviceTensor, IntoDevice}; + + #[test] + fn mlx_conv_library_compiles() -> TractResult<()> { + with_borrowed_metal_stream(|stream| { + stream.load_library(LibraryName::MlxConv)?; + Ok(()) + }) + } + + // Each kernel gets the weight layout it expects (direct: OIHW, MLX: OHWI) + // for the same convolution, so the two are comparable. + fn ohwi_to_oihw(w: &Tensor) -> TractResult { + w.clone().move_axis(3, 1) // O,H,W,I -> O,I,H,W + } + + fn ramp(dt: DatumType, shape: &[usize], seed: usize) -> TractResult { + let len: usize = shape.iter().product(); + let v: Vec = + (0..len).map(|i| (((i * 13 + seed * 7) % 23) as f32 - 11.0) / 32.0).collect(); + Ok(Tensor::from_shape(shape, &v)?.cast_to_dt(dt)?.into_owned()) + } + + #[allow(clippy::too_many_arguments)] + fn check_conv( + dt: DatumType, + n: usize, + ih: usize, + iw: usize, + c: usize, + o: usize, + kh: usize, + kw: usize, + stride: usize, + dil: usize, + padding: PaddingSpec, + ) -> TractResult<()> { + let pool_spec = PoolSpec::new( + DataFormat::NHWC, + tvec![kh, kw], + padding.clone(), + Some(tvec![dil, dil]), + Some(tvec![stride, stride]), + c, + o, + ); + let op = Conv { pool_spec, kernel_fmt: KernelFormat::OHWI, group: 1, q_params: None }; + let input = ramp(dt, &[n, ih, iw, c], 1)?; + let weights = ramp(dt, &[o, kh, kw, c], 2)?; + let bias = Tensor::zero_dt(dt, &[o])?; + let expected = op + .eval_out_of_plan(tvec![ + input.clone().into_tvalue(), + weights.clone().into_tvalue(), + bias.into_tvalue() + ])? + .unwrap() + .remove(0) + .into_tensor(); + let got = with_borrowed_metal_stream(|stream| { + let i = input.clone().into_device()?; + let w = weights.clone().into_device()?; + let out = unsafe { DeviceTensor::uninitialized_dt(dt, expected.shape())? }; + dispatch_mlx_conv_2d(stream, &op, &i, &w, &out)?; + stream.wait_until_completed()?; + Ok(out.to_host()?.into_tensor()) + })?; + expected.close_enough(&got, Approximation::Approximate).with_context(|| { + format!("dt={dt:?} n={n} {ih}x{iw}x{c} -> {o} k={kh}x{kw} s={stride} d={dil}") + }) + } + + #[test] + fn mlx_conv_1x1() -> TractResult<()> { + check_conv(DatumType::F32, 1, 8, 8, 16, 32, 1, 1, 1, 1, PaddingSpec::Valid) + } + + #[test] + fn mlx_conv_3x3_valid() -> TractResult<()> { + check_conv(DatumType::F32, 1, 10, 10, 16, 16, 3, 3, 1, 1, PaddingSpec::Valid) + } + + #[test] + fn mlx_conv_3x3_same() -> TractResult<()> { + check_conv(DatumType::F32, 2, 9, 7, 32, 16, 3, 3, 1, 1, PaddingSpec::SameUpper) + } + + #[test] + fn mlx_conv_strided() -> TractResult<()> { + check_conv(DatumType::F32, 1, 16, 16, 16, 32, 3, 3, 2, 1, PaddingSpec::SameUpper) + } + + #[test] + fn mlx_conv_dilated() -> TractResult<()> { + check_conv(DatumType::F32, 1, 16, 16, 16, 16, 3, 3, 1, 2, PaddingSpec::Valid) + } + + #[test] + fn mlx_conv_f16() -> TractResult<()> { + check_conv(DatumType::F16, 1, 12, 12, 32, 64, 3, 3, 1, 1, PaddingSpec::SameUpper) + } + + #[test] + fn mlx_conv_unaligned_channels() -> TractResult<()> { + check_conv(DatumType::F32, 1, 8, 8, 5, 7, 3, 3, 1, 1, PaddingSpec::Valid) + } + + // MLX implicit-GEMM conv against the direct kernel, on shapes a vision + // model actually runs. + // cargo test -p tract-metal bench_conv -- --ignored --nocapture + #[test] + #[ignore] + fn bench_conv() -> TractResult<()> { + use std::time::Instant; + println!("\n shape (N,H,W,C -> O, kHxkW, s) direct ms mlx ms gain"); + for &(n, ih, iw, c, o, k, stride) in &[ + (1usize, 112usize, 112usize, 32usize, 64usize, 3usize, 1usize), + (1, 56, 56, 64, 128, 3, 1), + (1, 28, 28, 128, 256, 3, 1), + (1, 56, 56, 64, 64, 1, 1), + (1, 224, 224, 3, 32, 3, 2), + (8, 28, 28, 128, 128, 3, 1), + ] { + for dt in [DatumType::F32, DatumType::F16] { + let pool_spec = PoolSpec::new( + DataFormat::NHWC, + tvec![k, k], + PaddingSpec::SameUpper, + Some(tvec![1, 1]), + Some(tvec![stride, stride]), + c, + o, + ); + let op = Conv { + pool_spec: pool_spec.clone(), + kernel_fmt: KernelFormat::OHWI, + group: 1, + q_params: None, + }; + let op_dir = + Conv { pool_spec, kernel_fmt: KernelFormat::OIHW, group: 1, q_params: None }; + let input = ramp(dt, &[n, ih, iw, c], 1)?; + let weights = ramp(dt, &[o, k, k, c], 2)?; + let weights_oihw = ohwi_to_oihw(&weights)?; + let oh = (ih + stride - 1) / stride; + let ow = (iw + stride - 1) / stride; + let (direct, mlx) = with_borrowed_metal_stream(|stream| { + let i = input.clone().into_device()?; + let w = weights.clone().into_device()?; + let w_dir = weights_oihw.clone().into_device()?; + let out = unsafe { DeviceTensor::uninitialized_dt(dt, &[n, oh, ow, o])? }; + let time = |f: &dyn Fn() -> TractResult<()>| -> TractResult { + for _ in 0..3 { + f()?; + } + stream.wait_until_completed()?; + let mut best = f64::MAX; + for _ in 0..5 { + let t = Instant::now(); + for _ in 0..10 { + f()?; + } + stream.wait_until_completed()?; + best = best.min(t.elapsed().as_secs_f64() / 10.0); + } + Ok(best) + }; + let d = time(&|| { + super::metal_conv_direct(stream, &op_dir, &i, &w_dir, None, &out) + })?; + let m = + time(&|| super::mlx_conv::dispatch_mlx_conv_2d(stream, &op, &i, &w, &out))?; + Ok((d, m)) + })?; + println!( + " {dt:?} {n}x{ih}x{iw}x{c} -> {o}, {k}x{k}, s{stride} {:9.4} {:9.4} {:6.2}x", + direct * 1e3, + mlx * 1e3, + direct / mlx + ); + } + } + Ok(()) + } + + // Sanity: both kernels must agree with the CPU op, so the bench compares + // two correct implementations. + #[test] + fn direct_and_mlx_agree_with_cpu() -> TractResult<()> { + let dt = DatumType::F32; + let (n, ih, iw, c, o, k) = (1usize, 16usize, 16usize, 32usize, 32usize, 3usize); + let spec = |fmt: KernelFormat| { + let pool_spec = PoolSpec::new( + DataFormat::NHWC, + tvec![k, k], + PaddingSpec::SameUpper, + Some(tvec![1, 1]), + Some(tvec![1, 1]), + c, + o, + ); + Conv { pool_spec, kernel_fmt: fmt, group: 1, q_params: None } + }; + let op_mlx = spec(KernelFormat::OHWI); + let op_dir = spec(KernelFormat::OIHW); + let input = ramp(dt, &[n, ih, iw, c], 1)?; + let w_ohwi = ramp(dt, &[o, k, k, c], 2)?; + let w_oihw = ohwi_to_oihw(&w_ohwi)?; + let bias = Tensor::zero_dt(dt, &[o])?; + let expected = op_mlx + .eval_out_of_plan(tvec![ + input.clone().into_tvalue(), + w_ohwi.clone().into_tvalue(), + bias.clone().into_tvalue() + ])? + .unwrap() + .remove(0) + .into_tensor(); + // the transposed weights must describe the same convolution + let via_oihw = op_dir + .eval_out_of_plan(tvec![ + input.clone().into_tvalue(), + w_oihw.clone().into_tvalue(), + bias.into_tvalue() + ])? + .unwrap() + .remove(0) + .into_tensor(); + expected.close_enough(&via_oihw, Approximation::Approximate)?; + + with_borrowed_metal_stream(|stream| { + let i = input.clone().into_device()?; + let wd = w_oihw.clone().into_device()?; + let wm = w_ohwi.clone().into_device()?; + let od = unsafe { DeviceTensor::uninitialized_dt(dt, expected.shape())? }; + super::metal_conv_direct(stream, &op_dir, &i, &wd, None, &od)?; + stream.wait_until_completed()?; + expected + .close_enough(&od.to_host()?.into_tensor(), Approximation::Approximate) + .context("direct kernel")?; + let om = unsafe { DeviceTensor::uninitialized_dt(dt, expected.shape())? }; + super::mlx_conv::dispatch_mlx_conv_2d(stream, &op_mlx, &i, &wm, &om)?; + stream.wait_until_completed()?; + expected + .close_enough(&om.to_host()?.into_tensor(), Approximation::Approximate) + .context("mlx kernel")?; + Ok(()) + }) + } + + // End to end through the metal transform: an NHWC conv must reach the MLX + // kernel and match the CPU result. + #[test] + fn conv_routes_through_metal_transform() -> TractResult<()> { + use crate::MetalTransform; + use tract_core::transform::ModelTransform; + let dt = DatumType::F32; + let (n, ih, iw, c, o, k) = (1usize, 12usize, 12usize, 16usize, 32usize, 3usize); + let pool_spec = PoolSpec::new( + DataFormat::NHWC, + tvec![k, k], + PaddingSpec::SameUpper, + Some(tvec![1, 1]), + Some(tvec![1, 1]), + c, + o, + ); + // an exported model arrives in OIHW; the metal rule must flip it + let op = Conv { pool_spec, kernel_fmt: KernelFormat::OIHW, group: 1, q_params: None }; + let input = ramp(dt, &[n, ih, iw, c], 1)?; + let w_oihw = ohwi_to_oihw(&ramp(dt, &[o, k, k, c], 2)?)?; + let bias = Tensor::zero_dt(dt, &[o])?; + let mut model = TypedModel::default(); + let i = model.add_source("i", dt.fact(&[n, ih, iw, c]))?; + let w = model.add_const("w", w_oihw.clone())?; + let b = model.add_const("b", bias.clone())?; + let out = model.wire_node("conv", op.clone(), &[i, w, b])?; + model.select_output_outlets(&out)?; + let cpu = model.clone().into_runnable()?.run(tvec![input.clone().into_tvalue()])?; + + let metal = MetalTransform::default().transform_into(model)?; + let ohwi = metal.nodes().iter().any(|n| { + n.op_as::() + .map(|c| c.op.kernel_fmt == KernelFormat::OHWI) + .unwrap_or(false) + }); + assert!(ohwi, "conv should have been moved to OHWI for the MLX kernel"); + let got = metal.into_runnable()?.run(tvec![input.into_tvalue()])?; + cpu[0] + .clone() + .into_tensor() + .close_enough(&got[0].clone().into_tensor(), Approximation::Approximate)?; + Ok(()) + } + + // How long each conv pipeline specialization takes to build. + // cargo test -p tract-metal bench_conv_pipeline_build -- --ignored --nocapture + #[test] + #[ignore] + fn bench_conv_pipeline_build() -> TractResult<()> { + use crate::{ConstantValues, Value}; + use std::time::Instant; + with_borrowed_metal_stream(|stream| { + let t = Instant::now(); + stream.load_library(LibraryName::MlxConv)?; + println!(" library compile {:8.1} ms", t.elapsed().as_secs_f64() * 1e3); + for (bm, bn) in [(32, 32), (64, 32), (32, 64), (64, 64)] { + for align in [true, false] { + for tname in ["float32", "float16"] { + let name = format!( + "implicit_gemm_conv_2d_general_{tname}_bm{bm}_bn{bn}_bk16_wm2_wn2" + ); + let consts = Some(ConstantValues::new(vec![(200, Value::Bool(align))])); + let t = Instant::now(); + stream.load_pipeline_with_constants(LibraryName::MlxConv, &name, consts)?; + println!( + " {tname} bm{bm} bn{bn} alC={align:<5} {:8.1} ms", + t.elapsed().as_secs_f64() * 1e3 + ); + } + } + } + Ok(()) + }) + } + + // Inception v3 conv shapes, to find which one the ported kernel handles badly. + // cargo test -p tract-metal bench_inception_shapes -- --ignored --nocapture + #[test] + #[ignore] + fn bench_inception_shapes() -> TractResult<()> { + use std::time::Instant; + let dt = DatumType::F32; + // (ih, iw, c, o, kh, kw, stride, valid) + for &(ih, iw, c, o, kh, kw, st, valid) in &[ + (299usize, 299usize, 3usize, 32usize, 3usize, 3usize, 2usize, true), + (149, 149, 32, 32, 3, 3, 1, true), + (147, 147, 32, 64, 3, 3, 1, false), + (73, 73, 80, 192, 3, 3, 1, true), + (35, 35, 48, 64, 5, 5, 1, false), + (35, 35, 64, 96, 3, 3, 1, false), + (17, 17, 128, 128, 1, 7, 1, false), + (17, 17, 128, 192, 7, 1, 1, false), + (8, 8, 384, 384, 1, 3, 1, false), + ] { + let pool_spec = PoolSpec::new( + DataFormat::NHWC, + tvec![kh, kw], + if valid { PaddingSpec::Valid } else { PaddingSpec::SameUpper }, + Some(tvec![1, 1]), + Some(tvec![st, st]), + c, + o, + ); + let op = Conv { + pool_spec: pool_spec.clone(), + kernel_fmt: KernelFormat::OHWI, + group: 1, + q_params: None, + }; + let op_dir = + Conv { pool_spec, kernel_fmt: KernelFormat::OIHW, group: 1, q_params: None }; + let input = ramp(dt, &[1, ih, iw, c], 1)?; + let w = ramp(dt, &[o, kh, kw, c], 2)?; + let w_dir = ohwi_to_oihw(&w)?; + let oh = if valid { (ih - kh) / st + 1 } else { (ih + st - 1) / st }; + let ow = if valid { (iw - kw) / st + 1 } else { (iw + st - 1) / st }; + let (d, m) = with_borrowed_metal_stream(|stream| { + let i = input.clone().into_device()?; + let wm = w.clone().into_device()?; + let wd = w_dir.clone().into_device()?; + let out = unsafe { DeviceTensor::uninitialized_dt(dt, &[1, oh, ow, o])? }; + let time = |f: &dyn Fn() -> TractResult<()>| -> TractResult { + f()?; + stream.wait_until_completed()?; + let t = Instant::now(); + for _ in 0..5 { + f()?; + } + stream.wait_until_completed()?; + Ok(t.elapsed().as_secs_f64() / 5.0) + }; + let d = time(&|| super::metal_conv_direct(stream, &op_dir, &i, &wd, None, &out))?; + let m = + time(&|| super::mlx_conv::dispatch_mlx_conv_2d(stream, &op, &i, &wm, &out))?; + Ok((d, m)) + })?; + println!( + " {ih}x{iw}x{c} -> {o}, {kh}x{kw}, s{st}{} direct {:9.3} ms mlx {:9.3} ms {:7.2}x", + if valid { " valid" } else { " same " }, + d * 1e3, + m * 1e3, + d / m + ); + } + Ok(()) + } + + // The kernel arrives in whatever layout the exporter used — TF gives HWIO, + // ONNX gives OIHW — and the ported kernel indexes OHWI. Each must reach it + // correctly; getting this wrong reads garbage loop bounds, not a crash. + #[test] + fn every_kernel_format_reaches_the_mlx_kernel() -> TractResult<()> { + use crate::MetalTransform; + use tract_core::transform::ModelTransform; + let dt = DatumType::F32; + let (n, ih, iw, c, o, kh, kw) = + (1usize, 12usize, 12usize, 16usize, 32usize, 3usize, 3usize); + let w_ohwi = ramp(dt, &[o, kh, kw, c], 2)?; + for fmt in [KernelFormat::OHWI, KernelFormat::OIHW, KernelFormat::HWIO] { + let weights = match fmt { + KernelFormat::OHWI => w_ohwi.clone(), + // [O,H,W,I] -> [O,I,H,W] + KernelFormat::OIHW => w_ohwi.clone().move_axis(3, 1)?, + // [O,H,W,I] -> [H,W,I,O] + KernelFormat::HWIO => w_ohwi.clone().move_axis(0, 3)?, + }; + let pool_spec = PoolSpec::new( + DataFormat::NHWC, + tvec![kh, kw], + PaddingSpec::SameUpper, + Some(tvec![1, 1]), + Some(tvec![1, 1]), + c, + o, + ); + let op = Conv { pool_spec, kernel_fmt: fmt, group: 1, q_params: None }; + let input = ramp(dt, &[n, ih, iw, c], 1)?; + let bias = Tensor::zero_dt(dt, &[o])?; + let mut model = TypedModel::default(); + let i = model.add_source("i", dt.fact(&[n, ih, iw, c]))?; + let w = model.add_const("w", weights)?; + let b = model.add_const("b", bias)?; + let out = model.wire_node("conv", op, &[i, w, b])?; + model.select_output_outlets(&out)?; + let cpu = model.clone().into_runnable()?.run(tvec![input.clone().into_tvalue()])?; + let metal = MetalTransform::default().transform_into(model)?; + let got = metal.into_runnable()?.run(tvec![input.clone().into_tvalue()])?; + cpu[0] + .clone() + .into_tensor() + .close_enough(&got[0].clone().into_tensor(), Approximation::Approximate) + .with_context(|| format!("kernel format {fmt:?}"))?; + } + Ok(()) + } +} diff --git a/metal/src/kernels/mod.rs b/metal/src/kernels/mod.rs index 8fce4884e9..d7a772bdde 100644 --- a/metal/src/kernels/mod.rs +++ b/metal/src/kernels/mod.rs @@ -32,7 +32,8 @@ const BASIC_MAT_MUL: &str = include_str!("matmul/basic/basic_mat_mul.metal"); const ARRAY_OPS: &str = include_str!("array/array_ops.metal"); const BIN_OPS: &str = include_str!("bin_ops.metal"); const NN_OPS: &str = concat!(include_str!("nn/nn_ops.metal"), include_str!("nn/pool.metal")); -const CONV_OPS: &str = include_str!("conv.metal"); +const CONV_OPS: &str = include_str!("conv/direct_conv.metal"); +const MLX_CONV: &str = include_str!("conv/mlx_conv.metal"); const ELEMENT_WISE_OPS: &str = include_str!("element_wise.metal"); const FFT_OPS: &str = include_str!("fft.metal"); const GDN_RECURRENT: &str = include_str!("gdn_recurrent.metal"); @@ -53,6 +54,7 @@ pub enum LibraryName { BinOps, ArrayOps, ConvOps, + MlxConv, NNOps, ElementWiseOps, Ggml, @@ -68,6 +70,7 @@ impl LibraryName { Self::ArrayOps => LibraryContent::Source(ARRAY_OPS), Self::BinOps => LibraryContent::Source(BIN_OPS), Self::ConvOps => LibraryContent::Source(CONV_OPS), + Self::MlxConv => LibraryContent::Source(MLX_CONV), Self::NNOps => LibraryContent::Source(NN_OPS), Self::ElementWiseOps => LibraryContent::Source(ELEMENT_WISE_OPS), Self::MlxGemm => LibraryContent::Source(MLX_GEMM), diff --git a/metal/src/transform.rs b/metal/src/transform.rs index 4dfd82fb0e..37a1371e41 100644 --- a/metal/src/transform.rs +++ b/metal/src/transform.rs @@ -10,6 +10,7 @@ use crate::{kernels, ops}; use tract_core::dyn_clone::clone_box; use tract_core::internal::translator::Translate; use tract_core::internal::*; +use tract_core::ops::cnn::KernelFormat; use tract_core::ops::cnn::conv::rewrite_kernel_conv_in_oihw; use tract_core::ops::cnn::{Conv, rewrite_conv_with_n_axis}; use tract_core::ops::einsum::prefix_matmul::{PrefixMatMul, rewrite_einsum_to_prefix_matmul}; @@ -110,6 +111,46 @@ fn cast_sdpa_mask_to_query_dt( Ok(Some(patch)) } +/// Metal-local kernel reordering: the shared rule puts every conv kernel in +/// `OIHW`, which the direct kernel indexes, but the ported MLX kernel wants +/// `OHWI`. Eligible convs are moved to `OHWI` instead; everything else falls +/// through to the shared rule. The kernel is normally a constant, so the +/// reorder folds. +fn rewrite_conv_kernel_metal( + ctx: &(), + model: &TypedModel, + node: &TypedNode, + name: &str, + conv: &Conv, +) -> TractResult> { + let in_facts = model.node_input_facts(node.id)?; + if crate::kernels::conv::mlx_conv::mlx_conv_eligible(conv, &in_facts) { + if conv.kernel_fmt == KernelFormat::OHWI { + return Ok(None); + } + // Reorder the constant here rather than wiring an axis move: the metal + // transform does not declutter afterwards, so a wired move would stay + // in the graph and run on every inference. + if let Some(kernel) = in_facts[1].konst.as_ref() { + let ohwi = match conv.kernel_fmt { + // [O,I,H,W] -> [O,H,W,I] + KernelFormat::OIHW => kernel.clone().into_tensor().move_axis(1, 3)?, + // [H,W,I,O] -> [O,H,W,I] + KernelFormat::HWIO => kernel.clone().into_tensor().move_axis(3, 0)?, + KernelFormat::OHWI => kernel.clone().into_tensor(), + }; + let mut patch = TypedModelPatch::default(); + let mut wire = patch.taps(model, &node.inputs)?; + wire[1] = patch.add_const(format!("{name}.kernel_ohwi"), ohwi)?; + let new = Conv { kernel_fmt: KernelFormat::OHWI, ..conv.clone() }; + let wire = patch.wire_node(name, new, &wire)?; + patch.shunt_outside(model, node.id.into(), wire[0])?; + return Ok(Some(patch)); + } + } + rewrite_kernel_conv_in_oihw(ctx, model, node, name, conv) +} + fn rewire_sdpa_metal(model: &mut TypedModel) -> TractResult<()> { Rewriter::default() .with_rule_for("cast-sdpa-mask-to-query-dt", cast_sdpa_mask_to_query_dt) @@ -190,7 +231,7 @@ impl MetalTransform { .rewrite(self, model)?; Rewriter::default() - .with_rule_for("rewrite_kernel_conv_in_oihw", rewrite_kernel_conv_in_oihw) + .with_rule_for("rewrite_conv_kernel_metal", rewrite_conv_kernel_metal) .with_rule_for("rewrite_conv_with_n_axis", rewrite_conv_with_n_axis) .with_rule_for("remove_rms_norm_cast", remove_rms_norm_cast) .with_rule_for("split_multi_axis_reduce", split_multi_axis_reduce)