hugo/DS4Server
1
0
Fork 0
Code Issues 4 Pull Requests Actions Packages Projects Releases Wiki Activity
Files
48c2f751b4cf643d2bc1ca12a7f4a3b155e7d120
DS4Server/metal/mtplx-qsa-compiled-header.metal
T

2487 lines
82 KiB
Metal
Raw Blame History

// Copyright © 2025 Apple Inc.
// Auto generated source for mlx/backend/metal/kernels/utils.h
///////////////////////////////////////////////////////////////////////////////
// Contents from "mlx/backend/metal/kernels/bf16.h"
///////////////////////////////////////////////////////////////////////////////
#line 1 "mlx/backend/metal/kernels/bf16.h"
// Copyright © 2023 Apple Inc.
#include <metal_stdlib>
using namespace metal;
typedef bfloat bfloat16_t;
inline uint16_t bfloat16_to_uint16(const bfloat16_t x) {
return as_type<uint16_t>(x);
}
inline bfloat16_t uint16_to_bfloat16(const uint16_t x) {
return as_type<bfloat16_t>(x);
}
///////////////////////////////////////////////////////////////////////////////
// Contents from "mlx/backend/metal/kernels/bf16_math.h"
///////////////////////////////////////////////////////////////////////////////
#line 1 "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<bfloat>(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<otype>(__metal_fabs(static_cast<ctype>(x), mfast)); \
} \
METAL_FUNC otype acos(itype x) { \
return static_cast<otype>(__metal_acos(static_cast<ctype>(x), mfast)); \
} \
METAL_FUNC otype acosh(itype x) { \
return static_cast<otype>(__metal_acosh(static_cast<ctype>(x), mfast)); \
} \
METAL_FUNC otype asin(itype x) { \
return static_cast<otype>(__metal_asin(static_cast<ctype>(x), mfast)); \
} \
METAL_FUNC otype asinh(itype x) { \
return static_cast<otype>(__metal_asinh(static_cast<ctype>(x), mfast)); \
} \
METAL_FUNC otype atan(itype y_over_x) { \
return static_cast<otype>( \
__metal_atan(static_cast<ctype>(y_over_x), mfast)); \
} \
METAL_FUNC otype atan2(itype y, itype x) { \
return static_cast<otype>( \
__metal_atan2(static_cast<ctype>(y), static_cast<ctype>(x), mfast)); \
} \
METAL_FUNC otype atanh(itype x) { \
return static_cast<otype>(__metal_atanh(static_cast<ctype>(x), mfast)); \
} \
METAL_FUNC otype ceil(itype x) { \
return static_cast<otype>(__metal_ceil(static_cast<ctype>(x), mfast)); \
} \
METAL_FUNC otype cos(itype x) { \
return static_cast<otype>(__metal_cos(static_cast<ctype>(x), mfast)); \
} \
METAL_FUNC otype cosh(itype x) { \
return static_cast<otype>(__metal_cosh(static_cast<ctype>(x), mfast)); \
} \
METAL_FUNC otype cospi(itype x) { \
return static_cast<otype>(__metal_cospi(static_cast<ctype>(x), mfast)); \
} \
METAL_FUNC otype divide(itype x, itype y) { \
return static_cast<otype>( \
__metal_divide(static_cast<ctype>(x), static_cast<ctype>(y), mfast)); \
} \
METAL_FUNC otype exp(itype x) { \
return static_cast<otype>(__metal_exp(static_cast<ctype>(x), mfast)); \
} \
METAL_FUNC otype exp10(itype x) { \
return static_cast<otype>(__metal_exp10(static_cast<ctype>(x), mfast)); \
} \
METAL_FUNC otype exp2(itype x) { \
return static_cast<otype>(__metal_exp2(static_cast<ctype>(x), mfast)); \
} \
METAL_FUNC otype fabs(itype x) { \
return static_cast<otype>(__metal_fabs(static_cast<ctype>(x), mfast)); \
} \
METAL_FUNC otype fdim(itype x, itype y) { \
ctype t = static_cast<ctype>(x - y); \
return static_cast<otype>(select(t, ctype(0), t < ctype(0) || x == y)); \
} \
METAL_FUNC otype floor(itype x) { \
return static_cast<otype>(__metal_floor(static_cast<ctype>(x), mfast)); \
} \
METAL_FUNC otype fma(itype x, itype y, itype z) { \
return static_cast<otype>(__metal_fma( \
static_cast<ctype>(x), static_cast<ctype>(y), static_cast<ctype>(z))); \
} \
METAL_FUNC otype fmax(itype x, itype y) { \
return static_cast<otype>( \
__metal_fmax(static_cast<ctype>(x), static_cast<ctype>(y), mfast)); \
} \
METAL_FUNC otype fmax3(itype x, itype y, itype z) { \
return static_cast<otype>(__metal_fmax3( \
static_cast<ctype>(x), \
static_cast<ctype>(y), \
static_cast<ctype>(z), \
mfast)); \
} \
METAL_FUNC otype fmedian3(itype x, itype y, itype z) { \
return static_cast<otype>(__metal_fmedian3( \
static_cast<ctype>(x), \
static_cast<ctype>(y), \
static_cast<ctype>(z), \
mfast)); \
} \
METAL_FUNC otype fmin(itype x, itype y) { \
return static_cast<otype>( \
__metal_fmin(static_cast<ctype>(x), static_cast<ctype>(y), mfast)); \
} \
METAL_FUNC otype fmin3(itype x, itype y, itype z) { \
return static_cast<otype>(__metal_fmin3( \
static_cast<ctype>(x), \
static_cast<ctype>(y), \
static_cast<ctype>(z), \
mfast)); \
} \
METAL_FUNC otype fmod(itype x, itype y) { \
return static_cast<otype>( \
__metal_fmod(static_cast<ctype>(x), static_cast<ctype>(y), mfast)); \
} \
METAL_FUNC otype fract(itype x) { \
return static_cast<otype>(__metal_fract(static_cast<ctype>(x), mfast)); \
} \
METAL_FUNC otype frexp(itype x, thread int& exp) { \
return static_cast<otype>(__metal_frexp(static_cast<ctype>(x), &exp)); \
} \
METAL_FUNC otype ldexp(itype x, int k) { \
return static_cast<otype>(__metal_ldexp(static_cast<ctype>(x), k, mfast)); \
} \
METAL_FUNC otype log(itype x) { \
return static_cast<otype>(__metal_log(static_cast<ctype>(x), mfast)); \
} \
METAL_FUNC otype log10(itype x) { \
return static_cast<otype>(__metal_log10(static_cast<ctype>(x), mfast)); \
} \
METAL_FUNC otype log2(itype x) { \
return static_cast<otype>(__metal_log2(static_cast<ctype>(x), mfast)); \
} \
METAL_FUNC otype max(itype x, itype y) { \
return static_cast<otype>( \
__metal_fmax(static_cast<ctype>(x), static_cast<ctype>(y), mfast)); \
} \
METAL_FUNC otype max3(itype x, itype y, itype z) { \
return static_cast<otype>(__metal_fmax3( \
static_cast<ctype>(x), \
static_cast<ctype>(y), \
static_cast<ctype>(z), \
mfast)); \
} \
METAL_FUNC otype median3(itype x, itype y, itype z) { \
return static_cast<otype>(__metal_fmedian3( \
static_cast<ctype>(x), \
static_cast<ctype>(y), \
static_cast<ctype>(z), \
mfast)); \
} \
METAL_FUNC otype min(itype x, itype y) { \
return static_cast<otype>( \
__metal_fmin(static_cast<ctype>(x), static_cast<ctype>(y), mfast)); \
} \
METAL_FUNC otype min3(itype x, itype y, itype z) { \
return static_cast<otype>(__metal_fmin3( \
static_cast<ctype>(x), \
static_cast<ctype>(y), \
static_cast<ctype>(z), \
mfast)); \
} \
METAL_FUNC otype nextafter(itype x, itype y) { \
return static_cast<otype>( \
__metal_nextafter(static_cast<ctype>(x), static_cast<ctype>(y))); \
} \
METAL_FUNC otype pow(itype x, itype y) { \
return static_cast<otype>( \
__metal_pow(static_cast<ctype>(x), static_cast<ctype>(y), mfast)); \
} \
METAL_FUNC otype powr(itype x, itype y) { \
return static_cast<otype>( \
__metal_powr(static_cast<ctype>(x), static_cast<ctype>(y), mfast)); \
} \
METAL_FUNC otype rint(itype x) { \
return static_cast<otype>(__metal_rint(static_cast<ctype>(x), mfast)); \
} \
METAL_FUNC otype round(itype x) { \
return static_cast<otype>(__metal_round(static_cast<ctype>(x), mfast)); \
} \
METAL_FUNC otype rsqrt(itype x) { \
return static_cast<otype>(__metal_rsqrt(static_cast<ctype>(x), mfast)); \
} \
METAL_FUNC otype sin(itype x) { \
return static_cast<otype>(__metal_sin(static_cast<ctype>(x), mfast)); \
} \
METAL_FUNC otype sinh(itype x) { \
return static_cast<otype>(__metal_sinh(static_cast<ctype>(x), mfast)); \
} \
METAL_FUNC otype sinpi(itype x) { \
return static_cast<otype>(__metal_sinpi(static_cast<ctype>(x), mfast)); \
} \
METAL_FUNC otype sqrt(itype x) { \
return static_cast<otype>(__metal_sqrt(static_cast<ctype>(x), mfast)); \
} \
METAL_FUNC otype tan(itype x) { \
return static_cast<otype>(__metal_tan(static_cast<ctype>(x), mfast)); \
} \
METAL_FUNC otype tanh(itype x) { \
return static_cast<otype>(__metal_tanh(static_cast<ctype>(x), mfast)); \
} \
METAL_FUNC otype tanpi(itype x) { \
return static_cast<otype>(__metal_tanpi(static_cast<ctype>(x), mfast)); \
} \
METAL_FUNC otype trunc(itype x) { \
return static_cast<otype>(__metal_trunc(static_cast<ctype>(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<otype>(__metal_simd_max(static_cast<ctype>(data))); \
} \
\
METAL_FUNC otype simd_min(itype data) { \
return static_cast<otype>(__metal_simd_min(static_cast<ctype>(data))); \
} \
\
METAL_FUNC otype simd_prefix_exclusive_product(itype data) { \
return static_cast<otype>( \
__metal_simd_prefix_exclusive_product(static_cast<ctype>(data))); \
} \
\
METAL_FUNC otype simd_prefix_exclusive_sum(itype data) { \
return static_cast<otype>( \
__metal_simd_prefix_exclusive_sum(static_cast<ctype>(data))); \
} \
\
METAL_FUNC otype simd_prefix_inclusive_product(itype data) { \
return static_cast<otype>( \
__metal_simd_prefix_inclusive_product(static_cast<ctype>(data))); \
} \
\
METAL_FUNC otype simd_prefix_inclusive_sum(itype data) { \
return static_cast<otype>( \
__metal_simd_prefix_inclusive_sum(static_cast<ctype>(data))); \
} \
\
METAL_FUNC otype simd_product(itype data) { \
return static_cast<otype>(__metal_simd_product(static_cast<ctype>(data))); \
} \
\
METAL_FUNC otype simd_sum(itype data) { \
return static_cast<otype>(__metal_simd_sum(static_cast<ctype>(data))); \
} \
\
METAL_FUNC otype simd_xor(itype data) { \
return static_cast<otype>(__metal_simd_xor(static_cast<ctype>(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
///////////////////////////////////////////////////////////////////////////////
// Contents from "mlx/backend/metal/kernels/complex.h"
///////////////////////////////////////////////////////////////////////////////
#line 1 "mlx/backend/metal/kernels/complex.h"
// Copyright © 2023 Apple Inc.
#include <metal_stdlib>
using namespace metal;
template <typename T>
struct complex_t;
template <typename T>
static constexpr constant bool is_complex_v = false;
template <typename T>
static constexpr constant bool is_complex_v<complex_t<T>> = true;
// Metal accepts explicit bfloat casts that is_convertible_v reports as false.
template <typename From, typename To>
static constexpr constant bool is_lane_convertible_v =
is_convertible_v<From, To> ||
(is_same_v<To, bfloat16_t> && is_convertible_v<From, float>) ||
(is_same_v<From, bfloat16_t> && is_convertible_v<float, To>);
template <typename T>
struct complex_t {
using value_type = T;
T real;
T imag;
// Constructors
constexpr complex_t(T real, T imag) thread : real(real), imag(imag) {};
constexpr complex_t() thread : real(0), imag(0) {};
constexpr complex_t() threadgroup : real(0), imag(0) {};
// Conversions from scalar types
template <
typename U,
typename = typename enable_if<
!is_complex_v<U> && is_lane_convertible_v<U, T>>::type>
constexpr complex_t(U x) thread : real(static_cast<T>(x)),
imag(static_cast<T>(0)) {}
template <
typename U,
typename = typename enable_if<
!is_complex_v<U> && is_lane_convertible_v<U, T>>::type>
constexpr complex_t(U x) threadgroup : real(static_cast<T>(x)),
imag(static_cast<T>(0)) {}
template <
typename U,
typename = typename enable_if<
!is_complex_v<U> && is_lane_convertible_v<U, T>>::type>
constexpr complex_t(U x) device : real(static_cast<T>(x)),
imag(static_cast<T>(0)) {}
template <
typename U,
typename = typename enable_if<
!is_complex_v<U> && is_lane_convertible_v<U, T>>::type>
constexpr complex_t(U x) constant : real(static_cast<T>(x)),
imag(static_cast<T>(0)) {}
// Conversions between complex types
template <
typename U,
typename = typename enable_if<
!is_same_v<U, T> && is_lane_convertible_v<U, T>>::type>
constexpr complex_t(complex_t<U> x) thread : real(static_cast<T>(x.real)),
imag(static_cast<T>(x.imag)) {}
template <
typename U,
typename = typename enable_if<
!is_same_v<U, T> && is_lane_convertible_v<U, T>>::type>
constexpr complex_t(complex_t<U> x) threadgroup
: real(static_cast<T>(x.real)),
imag(static_cast<T>(x.imag)) {}
template <
typename U,
typename = typename enable_if<
!is_same_v<U, T> && is_lane_convertible_v<U, T>>::type>
constexpr complex_t(complex_t<U> x) device : real(static_cast<T>(x.real)),
imag(static_cast<T>(x.imag)) {}
template <
typename U,
typename = typename enable_if<
!is_same_v<U, T> && is_lane_convertible_v<U, T>>::type>
constexpr complex_t(complex_t<U> x) constant : real(static_cast<T>(x.real)),
imag(static_cast<T>(x.imag)) {}
// Conversions to and from two-lane vectors (the FFT lane representation)
constexpr complex_t(vec<T, 2> v) thread : real(v.x), imag(v.y) {};
constexpr complex_t(vec<T, 2> v) threadgroup : real(v.x), imag(v.y) {};
constexpr complex_t(vec<T, 2> v) device : real(v.x), imag(v.y) {};
constexpr complex_t(vec<T, 2> v) constant : real(v.x), imag(v.y) {};
constexpr operator vec<T, 2>() const thread {
return vec<T, 2>(real, imag);
}
constexpr operator vec<T, 2>() const threadgroup {
return vec<T, 2>(real, imag);
}
constexpr operator vec<T, 2>() const device {
return vec<T, 2>(real, imag);
}
constexpr operator vec<T, 2>() const constant {
return vec<T, 2>(real, imag);
}
// Conversions to scalar types
template <
typename U,
typename = typename enable_if<
!is_complex_v<U> && is_lane_convertible_v<T, U>>::type>
constexpr operator U() const thread {
return static_cast<U>(real);
}
template <
typename U,
typename = typename enable_if<
!is_complex_v<U> && is_lane_convertible_v<T, U>>::type>
constexpr operator U() const threadgroup {
return static_cast<U>(real);
}
template <
typename U,
typename = typename enable_if<
!is_complex_v<U> && is_lane_convertible_v<T, U>>::type>
constexpr operator U() const device {
return static_cast<U>(real);
}
template <
typename U,
typename = typename enable_if<
!is_complex_v<U> && is_lane_convertible_v<T, U>>::type>
constexpr operator U() const constant {
return static_cast<U>(real);
}
};
using complex32_t = complex_t<half>;
using complex64_t = complex_t<float>;
static_assert(sizeof(complex32_t) == 2 * sizeof(half));
static_assert(sizeof(complex64_t) == 2 * sizeof(float));
static_assert(sizeof(complex_t<bfloat16_t>) == 2 * sizeof(bfloat16_t));
template <typename T>
constexpr complex_t<T> operator-(complex_t<T> x) {
return {-x.real, -x.imag};
}
template <typename T>
constexpr bool operator>=(complex_t<T> a, complex_t<T> b) {
return (a.real > b.real) || (a.real == b.real && a.imag >= b.imag);
}
template <typename T>
constexpr bool operator>(complex_t<T> a, complex_t<T> b) {
return (a.real > b.real) || (a.real == b.real && a.imag > b.imag);
}
template <typename T>
constexpr bool operator<=(complex_t<T> a, complex_t<T> b) {
return operator>=(b, a);
}
template <typename T>
constexpr bool operator<(complex_t<T> a, complex_t<T> b) {
return operator>(b, a);
}
template <typename T>
constexpr bool operator==(complex_t<T> a, complex_t<T> b) {
return a.real == b.real && a.imag == b.imag;
}
template <typename T>
constexpr complex_t<T> operator+(complex_t<T> a, complex_t<T> b) {
return {a.real + b.real, a.imag + b.imag};
}
template <typename T>
constexpr thread complex_t<T>& operator+=(
thread complex_t<T>& a,
complex_t<T> b) {
a.real += b.real;
a.imag += b.imag;
return a;
}
template <typename T>
constexpr threadgroup complex_t<T>& operator+=(
threadgroup complex_t<T>& a,
complex_t<T> b) {
a.real += b.real;
a.imag += b.imag;
return a;
}
template <typename T>
constexpr device complex_t<T>& operator+=(
device complex_t<T>& a,
complex_t<T> b) {
a.real += b.real;
a.imag += b.imag;
return a;
}
template <
typename T,
typename U,
enable_if_t<!is_complex_v<U> && is_lane_convertible_v<U, T>, bool> = true>
constexpr complex_t<T> operator+(U a, complex_t<T> b) {
return {static_cast<T>(a) + b.real, b.imag};
}
template <
typename T,
typename U,
enable_if_t<!is_complex_v<U> && is_lane_convertible_v<U, T>, bool> = true>
constexpr complex_t<T> operator+(complex_t<T> a, U b) {
return {a.real + static_cast<T>(b), a.imag};
}
template <typename T>
constexpr complex_t<T> operator-(complex_t<T> a, complex_t<T> b) {
return {a.real - b.real, a.imag - b.imag};
}
template <
typename T,
typename U,
enable_if_t<!is_complex_v<U> && is_lane_convertible_v<U, T>, bool> = true>
constexpr complex_t<T> operator-(U a, complex_t<T> b) {
return {static_cast<T>(a) - b.real, -b.imag};
}
template <
typename T,
typename U,
enable_if_t<!is_complex_v<U> && is_lane_convertible_v<U, T>, bool> = true>
constexpr complex_t<T> operator-(complex_t<T> a, U b) {
return {a.real - static_cast<T>(b), a.imag};
}
template <typename T>
constexpr complex_t<T> operator*(complex_t<T> a, complex_t<T> b) {
return {a.real * b.real - a.imag * b.imag, a.real * b.imag + a.imag * b.real};
}
template <typename T>
constexpr complex_t<T> operator/(complex_t<T> a, complex_t<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};
}
template <
typename T,
typename U,
enable_if_t<!is_complex_v<U> && is_lane_convertible_v<U, T>, bool> = true>
constexpr complex_t<T> operator/(U a, complex_t<T> b) {
auto scalar = static_cast<T>(a);
auto denom = b.real * b.real + b.imag * b.imag;
auto x = scalar * b.real;
auto y = -scalar * b.imag;
return {x / denom, y / denom};
}
template <typename T>
constexpr complex_t<T> operator%(complex_t<T> a, complex_t<T> b) {
auto real = a.real - (b.real * static_cast<int64_t>(a.real / b.real));
auto imag = a.imag - (b.imag * static_cast<int64_t>(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};
}
static_assert(
(complex_t<half>{1.0h, 2.0h} * complex_t<half>{3.0h, 4.0h}).real == -5.0h);
static_assert(
(complex_t<bfloat16_t>{bfloat16_t(1.0f), bfloat16_t(2.0f)} *
complex_t<bfloat16_t>{bfloat16_t(3.0f), bfloat16_t(4.0f)})
.real == bfloat16_t(-5.0f));
///////////////////////////////////////////////////////////////////////////////
// Contents from "mlx/backend/metal/kernels/defines.h"
///////////////////////////////////////////////////////////////////////////////
#line 1 "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<a, b>
#define instantiate_kernel(name, func, ...) \
template [[host_name( \
name)]] [[kernel]] decltype(func<__VA_ARGS__>) func<__VA_ARGS__>;
///////////////////////////////////////////////////////////////////////////////
// Contents from "mlx/backend/metal/kernels/logging.h"
///////////////////////////////////////////////////////////////////////////////
#line 1 "mlx/backend/metal/kernels/logging.h"
// Copyright © 2025 Apple Inc.
#if defined(__METAL_VERSION__) && (__METAL_VERSION__ >= 320)
#include <metal_logging>
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 <typename... Args>
void log_debug(constant char*, Args...) const thread {}
template <typename... Args>
void log_debug(constant char*, Args...) const constant {}
};
} // namespace mlx
#endif
///////////////////////////////////////////////////////////////////////////////
// Contents from "mlx/backend/metal/kernels/utils.h"
///////////////////////////////////////////////////////////////////////////////
#line 1 "mlx/backend/metal/kernels/utils.h"
// Copyright © 2023-2024 Apple Inc.
#include <metal_math>
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 <typename U>
struct WorkPerThread {
static_assert(sizeof(U) <= 8, "Type too large");
static constexpr int constant n = 8 / sizeof(U);
};
///////////////////////////////////////////////////////////////////////////////
// Type limits utils
///////////////////////////////////////////////////////////////////////////////
template <typename U>
struct Limits {
static const constant U max = metal::numeric_limits<U>::max();
static const constant U min = metal::numeric_limits<U>::min();
static const constant U finite_max = metal::numeric_limits<U>::max();
static const constant U finite_min = metal::numeric_limits<U>::min();
};
#define instantiate_default_limit(type) \
template <> \
struct Limits<type> { \
static constexpr constant type max = metal::numeric_limits<type>::max(); \
static constexpr constant type min = metal::numeric_limits<type>::min(); \
static constexpr constant type finite_max = \
metal::numeric_limits<type>::max(); \
static constexpr constant type finite_min = \
metal::numeric_limits<type>::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<type> { \
static constexpr constant type max = \
metal::numeric_limits<type>::infinity(); \
static constexpr constant type min = \
-metal::numeric_limits<type>::infinity(); \
static constexpr constant type finite_max = \
metal::numeric_limits<type>::max(); \
static constexpr constant type finite_min = \
-metal::numeric_limits<type>::max(); \
};
instantiate_float_limit(half);
instantiate_float_limit(float);
instantiate_float_limit(bfloat16_t);
template <>
struct Limits<bool> {
static constexpr constant bool max = true;
static constexpr constant bool min = false;
};
template <typename T>
struct Limits<complex_t<T>> {
inline static constexpr constant complex_t<T> max = complex_t<T>(
metal::numeric_limits<T>::infinity(),
metal::numeric_limits<T>::infinity());
inline static constexpr constant complex_t<T> min = complex_t<T>(
-metal::numeric_limits<T>::infinity(),
-metal::numeric_limits<T>::infinity());
};
///////////////////////////////////////////////////////////////////////////////
// Indexing utils
///////////////////////////////////////////////////////////////////////////////
#define MLX_MTL_PRAGMA_UNROLL _Pragma("clang loop unroll(full)")
///////////////////////////////////////////////////////////////////////////////
// Single Array with generic dims
template <typename IdxT = int64_t>
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 <typename IdxT = int64_t>
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 <typename IdxT = int64_t>
METAL_FUNC IdxT elem_to_loc_1(uint elem, constant const int64_t& stride) {
return elem * IdxT(stride);
}
template <typename IdxT = int64_t>
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 <typename IdxT = int64_t>
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 <typename IdxT = int64_t>
METAL_FUNC vec<IdxT, 2> 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<IdxT, 2> 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 <typename IdxT = int64_t>
METAL_FUNC vec<IdxT, 3> 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<IdxT, 3> 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 <int DIM, typename OffsetT = size_t, bool General = true>
struct LoopedElemToLoc {
int dim;
LoopedElemToLoc<DIM - 1, OffsetT, General> inner_looper;
OffsetT offset{0};
int index{0};
LoopedElemToLoc(int dim) thread : dim(dim), inner_looper(dim - 1) {}
void next(const constant int* shape, const constant int64_t* strides) thread {
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)
thread {
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() thread {
return offset;
}
};
template <typename OffsetT>
struct LoopedElemToLoc<1, OffsetT, true> {
int dim;
OffsetT offset{0};
uint index{0};
LoopedElemToLoc(int dim) thread : dim(dim) {}
void next(const constant int* shape, const constant int64_t* strides) thread {
index++;
if (dim > 1) {
offset = elem_to_loc<OffsetT>(index, shape, strides, dim);
} else {
offset += OffsetT(strides[0]);
}
}
void next(int n, const constant int* shape, const constant int64_t* strides)
thread {
index += n;
if (dim > 1) {
offset = elem_to_loc<OffsetT>(index, shape, strides, dim);
} else {
offset = index * OffsetT(strides[0]);
}
}
OffsetT location() thread {
return offset;
}
};
template <typename OffsetT>
struct LoopedElemToLoc<1, OffsetT, false> {
OffsetT offset{0};
LoopedElemToLoc(int) thread {}
void next(const constant int*, const constant int64_t* strides) thread {
offset += OffsetT(strides[0]);
}
void next(int n, const constant int*, const constant int64_t* strides)
thread {
offset += n * OffsetT(strides[0]);
}
OffsetT location() thread {
return offset;
}
};
///////////////////////////////////////////////////////////////////////////////
// Calculation utils
///////////////////////////////////////////////////////////////////////////////
/** Compute ceil((float)N/(float)M) */
template <typename T, typename U>
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<float>::max) {
return Limits<float>::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<float>(x);
if (xp1 == Limits<float>::max) {
return Limits<bfloat16_t>::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<uint64_t>(
metal::simd_shuffle_down(as_type<uint2>(data), delta));
}
inline int64_t simd_shuffle_down(int64_t data, uint16_t delta) {
return as_type<int64_t>(
metal::simd_shuffle_down(as_type<uint2>(data), delta));
}
inline bool simd_shuffle_down(bool data, uint16_t delta) {
return simd_shuffle_down(static_cast<uint32_t>(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<uint64_t>(metal::simd_shuffle_up(as_type<uint2>(data), delta));
}
inline int64_t simd_shuffle_up(int64_t data, uint16_t delta) {
return as_type<int64_t>(metal::simd_shuffle_up(as_type<uint2>(data), delta));
}
inline bool simd_shuffle_up(bool data, uint16_t delta) {
return simd_shuffle_up(static_cast<uint32_t>(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<uint64_t>(metal::simd_shuffle_and_fill_up(
as_type<uint2>(data), as_type<uint2>(filling), delta));
}
inline int64_t
simd_shuffle_and_fill_up(int64_t data, int64_t filling, uint16_t delta) {
return as_type<int64_t>(metal::simd_shuffle_and_fill_up(
as_type<uint2>(data), as_type<uint2>(filling), delta));
}
inline bool simd_shuffle_and_fill_up(bool data, bool filling, uint16_t delta) {
return simd_shuffle_and_fill_up(
static_cast<uint32_t>(data), static_cast<uint32_t>(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<uint64_t>(metal::simd_shuffle(as_type<uint2>(data), lane));
}
inline int64_t simd_shuffle(int64_t data, uint16_t lane) {
return as_type<int64_t>(metal::simd_shuffle(as_type<uint2>(data), lane));
}
inline bool simd_shuffle(bool data, uint16_t lane) {
return simd_shuffle(static_cast<uint32_t>(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 <bool condition, typename T, typename U>
struct ConditionalType {
using type = U;
};
template <typename T, typename U>
struct ConditionalType<true, T, U> {
using type = T;
};
///////////////////////////////////////////////////////////////////////////////
// Type casting utils
///////////////////////////////////////////////////////////////////////////////
template <typename U, typename T>
inline U cast_to(T val) {
return static_cast<U>(val);
}
template <>
inline bool cast_to<bool, float>(float val) {
return (as_type<uint32_t>(val) & 0x7FFFFFFF) != 0;
}
template <>
inline bool cast_to<bool, bfloat16_t>(bfloat16_t val) {
return (as_type<uint16_t>(val) & 0x7FFF) != 0;
}
template <>
inline bool cast_to<bool, complex64_t>(complex64_t val) {
return cast_to<bool, float>(val.real) || cast_to<bool, float>(val.imag);
}
///////////////////////////////////////////////////////////////////////////////
// Copyright © 2025 Apple Inc.
// Auto generated source for mlx/backend/metal/kernels/unary_ops.h
///////////////////////////////////////////////////////////////////////////////
// Contents from "mlx/backend/metal/kernels/cexpf.h"
///////////////////////////////////////////////////////////////////////////////
#line 1 "mlx/backend/metal/kernels/cexpf.h"
// Copyright © 2025 Apple Inc.
// Copyright © 2008-2013 NVIDIA Corporation
// Copyright © 2013 Filipe RNC Maia
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// Forked from
// https://github.com/NVIDIA/cccl/blob/main/thrust/thrust/detail/complex/cexpf.h
// TODO: We should use thrust::exp but the thrust header in old CUDA versions
// can not be used in JIT.
#include <metal_math>
using ieee_float_shape_type = union {
float value;
uint32_t word;
};
inline void get_float_word(thread uint32_t& i, float d) {
ieee_float_shape_type gf_u;
gf_u.value = (d);
(i) = gf_u.word;
}
inline void get_float_word(thread int32_t& i, float d) {
ieee_float_shape_type gf_u;
gf_u.value = (d);
(i) = gf_u.word;
}
inline void set_float_word(thread float& d, uint32_t i) {
ieee_float_shape_type sf_u;
sf_u.word = (i);
(d) = sf_u.value;
}
inline float frexp_expf(float x, thread int* expt) {
const uint32_t k = 235;
const float kln2 = 162.88958740F;
float exp_x;
uint32_t hx;
exp_x = metal::exp(x - kln2);
get_float_word(hx, exp_x);
*expt = (hx >> 23) - (0x7f + 127) + k;
set_float_word(exp_x, (hx & 0x7fffff) | ((0x7f + 127) << 23));
return exp_x;
}
inline complex64_t ldexp_cexpf(complex64_t z, int expt) {
float x, y, exp_x, scale1, scale2;
int ex_expt, half_expt;
x = z.real;
y = z.imag;
exp_x = frexp_expf(x, &ex_expt);
expt += ex_expt;
half_expt = expt / 2;
set_float_word(scale1, (0x7f + half_expt) << 23);
half_expt = expt - half_expt;
set_float_word(scale2, (0x7f + half_expt) << 23);
return complex64_t{
metal::cos(y) * exp_x * scale1 * scale2,
metal::sin(y) * exp_x * scale1 * scale2};
}
inline complex64_t cexpf(const thread complex64_t& z) {
float x, y, exp_x;
uint32_t hx, hy;
const uint32_t exp_ovfl = 0x42b17218, cexp_ovfl = 0x43400074;
x = z.real;
y = z.imag;
get_float_word(hy, y);
hy &= 0x7fffffff;
/* cexp(x + I 0) = exp(x) + I 0 */
if (hy == 0) {
return complex64_t{metal::exp(x), y};
}
get_float_word(hx, x);
/* cexp(0 + I y) = cos(y) + I sin(y) */
if ((hx & 0x7fffffff) == 0) {
return complex64_t{metal::cos(y), metal::sin(y)};
}
if (hy >= 0x7f800000) {
if ((hx & 0x7fffffff) != 0x7f800000) {
/* cexp(finite|NaN +- I Inf|NaN) = NaN + I NaN */
return complex64_t{y - y, y - y};
} else if (hx & 0x80000000) {
/* cexp(-Inf +- I Inf|NaN) = 0 + I 0 */
return complex64_t{0.0, 0.0};
} else {
/* cexp(+Inf +- I Inf|NaN) = Inf + I NaN */
return complex64_t{x, y - y};
}
}
if (hx >= exp_ovfl && hx <= cexp_ovfl) {
/*
* x is between 88.7 and 192, so we must scale to avoid
* overflow in expf(x).
*/
return ldexp_cexpf(z, 0);
} else {
/*
* Cases covered here:
* - x < exp_ovfl and exp(x) won't overflow (common case)
* - x > cexp_ovfl, so exp(x) * s overflows for all s > 0
* - x = +-Inf (generated by exp())
* - x = NaN (spurious inexact exception from y)
*/
exp_x = metal::exp(x);
return complex64_t{exp_x * metal::cos(y), exp_x * metal::sin(y)};
}
}
///////////////////////////////////////////////////////////////////////////////
// Contents from "mlx/backend/metal/kernels/expm1f.h"
///////////////////////////////////////////////////////////////////////////////
#line 1 "mlx/backend/metal/kernels/expm1f.h"
// Copyright © 2023 Apple Inc.
#include <metal_math>
// Original license copied below:
// Copyright (c) 2015-2023 Norbert Juffa
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without
// modification, are permitted provided that the following conditions
// are met:
//
// 1. Redistributions of source code must retain the above copyright
// notice, this list of conditions and the following disclaimer.
//
// 2. Redistributions in binary form must reproduce the above copyright
// notice, this list of conditions and the following disclaimer in the
// documentation and/or other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
// "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
// LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
// A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
// HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
// SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
// LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
// DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
// THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
/* Compute exponential base e minus 1. Maximum ulp error = 0.997458
i = rint(a/log(2)), f = a-i*log(2). Then expm1(a) = 2**i * (expm1(f)+1) - 1.
Compute r = expm1(f). Then expm1(a)= 2 * (0.5 * 2**i * r + 0.5 * 2**i - 0.5).
With t = 0.5*2**i, expm1(a) = 2*(r * t + t-0.5). However, for best accuracy,
when i == 1, expm1(a)= 2*(r + 0.5), and when i == 0, expm1(a) = r.
NOTE: Scale factor b is only applied if i < 0 or i > 1 (should be power of 2)
*/
float expm1f_scaled_unchecked(float a, float b) {
float f, j, r, s, t, u, v, x, y;
int i;
// exp(a) = 2**i * exp(f); i = rintf (a / log(2))
j = fma(1.442695f, a, 12582912.f); // 0x1.715476p0, 0x1.8p23
j = j - 12582912.0f; // 0x1.8p23
i = (int)j;
f = fma(j, -6.93145752e-1f, a);
// approximate r = exp(f)-1 on interval [-log(2)/2, +log(2)/2]
s = f * f;
if (a == 0.0f)
s = a; // ensure -0 is passed through
// err = 0.997458 ulp1 = 11081805
r = 1.97350979e-4f; // 0x1.9de000p-13
r = fma(r, f, 1.39309070e-3f); // 0x1.6d30bcp-10
r = fma(r, f, 8.33343994e-3f); // 0x1.1111f6p-7
r = fma(r, f, 4.16668020e-2f); // 0x1.55559ep-5
r = fma(r, f, 1.66666716e-1f); // 0x1.55555cp-3
r = fma(r, f, 4.99999970e-1f); // 0x1.fffffep-2
u = (j == 1) ? (f + 0.5f) : f;
v = fma(r, s, u);
s = 0.5f * b;
t = ldexp(s, i);
y = t - s;
x = (t - y) - s; // double-float canonicalization of difference
r = fma(v, t, x) + y;
r = r + r;
if (j == 0)
r = v;
if (j == 1)
r = v + v;
return r;
}
/* Compute exponential base e minus 1. max ulp err = 0.99746 */
float expm1f(float a) {
float r;
r = expm1f_scaled_unchecked(a, 1.0f);
/* handle severe overflow and underflow */
if (abs(a - 1.0f) > 88.0f) {
r = pow(2, a);
r = fma(r, r, -1.0f);
}
return r;
}
///////////////////////////////////////////////////////////////////////////////
// Contents from "mlx/backend/metal/kernels/erf.h"
///////////////////////////////////////////////////////////////////////////////
#line 1 "mlx/backend/metal/kernels/erf.h"
// Copyright © 2023 Apple Inc.
#include <metal_math>
/*
* Approximation to the error function.
* Based on code from:
* https://stackoverflow.com/questions/35148198/efficient-faithfully-rounded-implementation-of-error-function-erff#answer-35148199
*/
float erf(float a) {
float r, s, t, u;
t = metal::abs(a);
s = a * a;
if (t > 0.927734375f) {
// maximum error 0.99527 ulp
r = metal::fma(
-1.72853470e-5f, t, 3.83197126e-4f); // -0x1.220000p-16,0x1.91cfb2p-12
u = metal::fma(
-3.88396438e-3f, t, 2.42546219e-2f); // -0x1.fd1438p-9, 0x1.8d6342p-6
r = metal::fma(r, s, u);
r = metal::fma(r, t, -1.06777877e-1f); // -0x1.b55cb8p-4
r = metal::fma(r, t, -6.34846687e-1f); // -0x1.450aa0p-1
r = metal::fma(r, t, -1.28717512e-1f); // -0x1.079d0cp-3
r = metal::fma(r, t, -t);
r = -expm1f(r);
r = metal::copysign(r, a);
} else {
// maximum error 0.98929 ulp
r = -5.96761703e-4f; // -0x1.38e000p-11
r = metal::fma(r, s, 4.99119423e-3f); // 0x1.471a58p-8
r = metal::fma(r, s, -2.67681349e-2f); // -0x1.b691b2p-6
r = metal::fma(r, s, 1.12819925e-1f); // 0x1.ce1c44p-4
r = metal::fma(r, s, -3.76125336e-1f); // -0x1.812700p-2
r = metal::fma(r, s, 1.28379166e-1f); // 0x1.06eba8p-3
r = metal::fma(r, a, a);
}
return r;
}
float erfinv(float a) {
auto t = metal::fma(a, 0.0f - a, 1.0f);
t = metal::log(t);
float p;
if (metal::abs(t) > 6.125f) { // maximum ulp error = 2.35793
p = 3.03697567e-10f; // 0x1.4deb44p-32
p = metal::fma(p, t, 2.93243101e-8f); // 0x1.f7c9aep-26
p = metal::fma(p, t, 1.22150334e-6f); // 0x1.47e512p-20
p = metal::fma(p, t, 2.84108955e-5f); // 0x1.dca7dep-16
p = metal::fma(p, t, 3.93552968e-4f); // 0x1.9cab92p-12
p = metal::fma(p, t, 3.02698812e-3f); // 0x1.8cc0dep-9
p = metal::fma(p, t, 4.83185798e-3f); // 0x1.3ca920p-8
p = metal::fma(p, t, -2.64646143e-1f); // -0x1.0eff66p-2
p = metal::fma(p, t, 8.40016484e-1f); // 0x1.ae16a4p-1
} else { // maximum ulp error = 2.35002
p = 5.43877832e-9f; // 0x1.75c000p-28
p = metal::fma(p, t, 1.43285448e-7f); // 0x1.33b402p-23
p = metal::fma(p, t, 1.22774793e-6f); // 0x1.499232p-20
p = metal::fma(p, t, 1.12963626e-7f); // 0x1.e52cd2p-24
p = metal::fma(p, t, -5.61530760e-5f); // -0x1.d70bd0p-15
p = metal::fma(p, t, -1.47697632e-4f); // -0x1.35be90p-13
p = metal::fma(p, t, 2.31468678e-3f); // 0x1.2f6400p-9
p = metal::fma(p, t, 1.15392581e-2f); // 0x1.7a1e50p-7
p = metal::fma(p, t, -2.32015476e-1f); // -0x1.db2aeep-3
p = metal::fma(p, t, 8.86226892e-1f); // 0x1.c5bf88p-1
}
return a * p;
}
///////////////////////////////////////////////////////////////////////////////
// Contents from "mlx/backend/metal/kernels/fp8.h"
///////////////////////////////////////////////////////////////////////////////
#line 1 "mlx/backend/metal/kernels/fp8.h"
struct fp8_e4m3 {
template <typename T>
fp8_e4m3(T f) thread {
// From PyTorch
// https://github.com/pytorch/pytorch/blob/e3643e1e0e923f0fc063dfab6f45c956d568919d/c10/util/Float8_e4m3fn.h#L148
uint32_t fp8_max = 543 << 21;
uint32_t denorm_mask = 141 << 23;
uint32_t f_bits = as_type<uint32_t>(static_cast<float>(f));
uint32_t sign = f_bits & 0x80000000;
f_bits ^= sign;
if (f_bits >= fp8_max) {
// Default behavior saturates to min/max
bits = 0x7E;
} else {
if (f_bits < (121 << 23)) {
f_bits = as_type<uint32_t>(
as_type<float>(f_bits) + as_type<float>(denorm_mask));
bits = static_cast<uint8_t>(f_bits - denorm_mask);
} else {
// resulting mantissa is odd
uint8_t mant_odd = (f_bits >> 20) & 1;
f_bits += ((uint32_t)(7 - 127) << 23) + 0x7FFFF;
f_bits += mant_odd;
bits = static_cast<uint8_t>(f_bits >> 20);
}
}
bits |= static_cast<uint8_t>(sign >> 24);
}
operator float16_t() thread {
uint16_t v = (bits & 127) << 7;
half converted = as_type<half>(v);
converted *= 256.0;
auto sign = bits & 128;
return (sign ? -converted : converted);
}
operator bfloat16_t() thread {
return static_cast<bfloat16_t>(this->operator float16_t());
}
operator float() thread {
return static_cast<float>(this->operator float16_t());
}
uint8_t bits;
};
struct fp8_e8m0 {
fp8_e8m0(float x) thread {
if (!metal::isfinite(x)) {
bits = 0xFF;
return;
}
if (x < 0.0f) {
bits = 0x00;
return;
}
float le = metal::log2(x);
int n = int(metal::round(le));
n = n < -127 ? -127 : n;
n = n > 127 ? 127 : n;
bits = static_cast<uint8_t>(n + 127);
}
operator bfloat16_t() thread {
uint16_t out = (bits == 0 ? 0x40 : (static_cast<uint16_t>(bits) << 7));
return as_type<bfloat16_t>(out);
}
operator float() thread {
uint32_t out = (bits == 0 ? 0x400000 : (static_cast<uint16_t>(bits) << 23));
return as_type<float>(out);
}
uint8_t bits;
};
// Smallest E8M0 >= x. Scales are amax/max_element, so rounding one down
// leaves the block's largest elements outside the element range, where they
// saturate. Matches the CUDA backend, which rounds up via cutlass ue8m0.
inline float mx_scale_round_up(float x) {
fp8_e8m0 s(x);
if (s.bits < 0xFE && float(s) < x) {
s.bits += 1;
}
return float(s);
}
///////////////////////////////////////////////////////////////////////////////
// Contents from "mlx/backend/metal/kernels/unary_ops.h"
///////////////////////////////////////////////////////////////////////////////
#line 1 "mlx/backend/metal/kernels/unary_ops.h"
// Copyright © 2023-2024 Apple Inc.
#include <metal_integer>
#include <metal_math>
namespace {
constant float inf = metal::numeric_limits<float>::infinity();
}
struct Abs {
template <typename T>
T operator()(T x) thread {
return metal::abs(x);
};
uint8_t operator()(uint8_t x) thread {
return x;
};
uint16_t operator()(uint16_t x) thread {
return x;
};
uint32_t operator()(uint32_t x) thread {
return x;
};
uint64_t operator()(uint64_t x) thread {
return x;
};
bool operator()(bool x) thread {
return x;
};
complex64_t operator()(complex64_t x) thread {
return {metal::precise::sqrt(x.real * x.real + x.imag * x.imag), 0};
};
};
struct ArcCos {
template <typename T>
T operator()(T x) thread {
return metal::precise::acos(x);
};
complex64_t operator()(complex64_t x) thread;
};
struct ArcCosh {
template <typename T>
T operator()(T x) thread {
return metal::precise::acosh(x);
};
};
struct ArcSin {
template <typename T>
T operator()(T x) thread {
return metal::precise::asin(x);
};
complex64_t operator()(complex64_t x) thread;
};
struct ArcSinh {
template <typename T>
T operator()(T x) thread {
return metal::precise::asinh(x);
};
};
struct ArcTan {
template <typename T>
T operator()(T x) thread {
return metal::precise::atan(x);
};
complex64_t operator()(complex64_t x) thread;
};
struct ArcTanh {
template <typename T>
T operator()(T x) thread {
return metal::precise::atanh(x);
};
};
struct BitwiseInvert {
template <typename T>
T operator()(T x) thread {
return ~x;
};
};
struct Ceil {
template <typename T>
T operator()(T x) thread {
return metal::ceil(x);
};
int8_t operator()(int8_t x) thread {
return x;
};
int16_t operator()(int16_t x) thread {
return x;
};
int32_t operator()(int32_t x) thread {
return x;
};
int64_t operator()(int64_t x) thread {
return x;
};
uint8_t operator()(uint8_t x) thread {
return x;
};
uint16_t operator()(uint16_t x) thread {
return x;
};
uint32_t operator()(uint32_t x) thread {
return x;
};
uint64_t operator()(uint64_t x) thread {
return x;
};
bool operator()(bool x) thread {
return x;
};
};
struct Cos {
template <typename T>
T operator()(T x) thread {
return metal::precise::cos(x);
};
complex64_t operator()(complex64_t x) thread {
return {
metal::precise::cos(x.real) * metal::precise::cosh(x.imag),
-metal::precise::sin(x.real) * metal::precise::sinh(x.imag)};
};
};
struct Cosh {
template <typename T>
T operator()(T x) thread {
return metal::precise::cosh(x);
};
complex64_t operator()(complex64_t x) thread {
return {
metal::precise::cosh(x.real) * metal::precise::cos(x.imag),
metal::precise::sinh(x.real) * metal::precise::sin(x.imag)};
};
};
struct Conjugate {
complex64_t operator()(complex64_t x) thread {
return complex64_t{x.real, -x.imag};
}
};
struct Erf {
template <typename T>
T operator()(T x) thread {
return static_cast<T>(erf(static_cast<float>(x)));
};
};
struct ErfInv {
template <typename T>
T operator()(T x) thread {
return static_cast<T>(erfinv(static_cast<float>(x)));
};
};
struct Exp {
template <typename T>
T operator()(T x) thread {
return metal::precise::exp(x);
};
complex64_t operator()(complex64_t x) thread {
return cexpf(x);
}
};
struct Expm1 {
template <typename T>
T operator()(T x) thread {
return static_cast<T>(expm1f(static_cast<float>(x)));
};
};
struct Floor {
template <typename T>
T operator()(T x) thread {
return metal::floor(x);
};
int8_t operator()(int8_t x) thread {
return x;
};
int16_t operator()(int16_t x) thread {
return x;
};
int32_t operator()(int32_t x) thread {
return x;
};
int64_t operator()(int64_t x) thread {
return x;
};
uint8_t operator()(uint8_t x) thread {
return x;
};
uint16_t operator()(uint16_t x) thread {
return x;
};
uint32_t operator()(uint32_t x) thread {
return x;
};
uint64_t operator()(uint64_t x) thread {
return x;
};
bool operator()(bool x) thread {
return x;
};
};
struct Imag {
float operator()(complex64_t x) thread {
return x.imag;
};
};
struct Log {
template <typename T>
T operator()(T x) thread {
return metal::precise::log(x);
};
complex64_t operator()(complex64_t x) thread {
auto r = metal::precise::log(Abs{}(x).real);
auto i = metal::precise::atan2(x.imag, x.real);
return {r, i};
};
};
struct Log2 {
template <typename T>
T operator()(T x) thread {
return metal::precise::log2(x);
};
complex64_t operator()(complex64_t x) thread {
auto y = Log{}(x);
return {y.real / M_LN2_F, y.imag / M_LN2_F};
};
};
struct Log10 {
template <typename T>
T operator()(T x) thread {
return metal::precise::log10(x);
};
complex64_t operator()(complex64_t x) thread {
auto y = Log{}(x);
return {y.real / M_LN10_F, y.imag / M_LN10_F};
};
};
struct Log1p {
template <typename T>
T operator()(T x) thread {
return log1p(x);
};
};
struct LogicalNot {
template <typename T>
T operator()(T x) thread {
return !x;
};
};
struct Negative {
template <typename T>
T operator()(T x) thread {
return -x;
};
};
struct Real {
float operator()(complex64_t x) thread {
return x.real;
};
};
struct Round {
template <typename T>
T operator()(T x) thread {
return metal::rint(x);
};
complex64_t operator()(complex64_t x) thread {
return {metal::rint(x.real), metal::rint(x.imag)};
};
};
struct Sigmoid {
template <typename T>
T operator()(T x) thread {
auto y = 1 / (1 + metal::exp(metal::abs(x)));
return (x < 0) ? y : 1 - y;
}
};
struct Sign {
template <typename T>
T operator()(T x) thread {
return (x > T(0)) - (x < T(0));
};
uint32_t operator()(uint32_t x) thread {
return x != 0;
};
complex64_t operator()(complex64_t x) thread {
if (x == complex64_t(0)) {
return x;
}
return x /
(complex64_t)metal::precise::sqrt(x.real * x.real + x.imag * x.imag);
};
};
struct Sin {
template <typename T>
T operator()(T x) thread {
return metal::precise::sin(x);
};
complex64_t operator()(complex64_t x) thread {
return {
metal::precise::sin(x.real) * metal::precise::cosh(x.imag),
metal::precise::cos(x.real) * metal::precise::sinh(x.imag)};
};
};
struct Sinh {
template <typename T>
T operator()(T x) thread {
return metal::precise::sinh(x);
};
complex64_t operator()(complex64_t x) thread {
return {
metal::precise::sinh(x.real) * metal::precise::cos(x.imag),
metal::precise::cosh(x.real) * metal::precise::sin(x.imag)};
};
};
struct Square {
template <typename T>
T operator()(T x) thread {
return x * x;
};
};
struct Sqrt {
template <typename T>
T operator()(T x) thread {
return metal::precise::sqrt(x);
};
complex64_t operator()(complex64_t x) thread {
if (x.real == 0.0 && x.imag == 0.0) {
return {0.0, 0.0};
}
auto r = Abs{}(x).real;
auto a = metal::precise::sqrt((r + x.real) / 2.0);
auto b_abs = metal::precise::sqrt((r - x.real) / 2.0);
auto b = metal::copysign(b_abs, x.imag);
return {a, b};
}
};
struct Rsqrt {
template <typename T>
T operator()(T x) thread {
return metal::precise::rsqrt(x);
};
complex64_t operator()(complex64_t x) thread {
return 1.0 / Sqrt{}(x);
}
};
struct Tan {
template <typename T>
T operator()(T x) thread {
return metal::precise::tan(x);
};
complex64_t operator()(complex64_t x) thread {
float tan_a = metal::precise::tan(x.real);
float tanh_b = metal::precise::tanh(x.imag);
float t1 = tan_a * tanh_b;
float denom = 1. + t1 * t1;
return {(tan_a - tanh_b * t1) / denom, (tanh_b + tan_a * t1) / denom};
};
};
struct Tanh {
template <typename T>
T operator()(T x) thread {
return metal::precise::tanh(x);
};
complex64_t operator()(complex64_t x) thread {
float tanh_a = metal::precise::tanh(x.real);
float tan_b = metal::precise::tan(x.imag);
float t1 = tanh_a * tan_b;
float denom = 1. + t1 * t1;
return {(tanh_a + tan_b * t1) / denom, (tan_b - tanh_a * t1) / denom};
};
};
complex64_t ArcCos::operator()(complex64_t x) thread {
auto i = complex64_t{0.0, 1.0};
auto y = Log{}(x + i * Sqrt{}(1.0 - x * x));
return {y.imag, -y.real};
};
complex64_t ArcSin::operator()(complex64_t x) thread {
auto i = complex64_t{0.0, 1.0};
auto y = Log{}(i * x + Sqrt{}(1.0 - x * x));
return {y.imag, -y.real};
};
complex64_t ArcTan::operator()(complex64_t x) thread {
auto i = complex64_t{0.0, 1.0};
auto ix = i * x;
return (1.0 / complex64_t{0.0, 2.0}) * Log{}((1.0 + ix) / (1.0 - ix));
};
struct ToFP8 {
template <typename T>
uint8_t operator()(T f) thread {
return fp8_e4m3(f).bits;
}
};
struct FromFP8 {
float operator()(uint8_t x) thread {
return float(*(thread fp8_e4m3*)(&x));
}
};
///////////////////////////////////////////////////////////////////////////////
// Copyright © 2025 Apple Inc.
// Auto generated source for mlx/backend/metal/kernels/binary_ops.h
///////////////////////////////////////////////////////////////////////////////
// Contents from "mlx/backend/metal/kernels/binary_ops.h"
///////////////////////////////////////////////////////////////////////////////
#line 1 "mlx/backend/metal/kernels/binary_ops.h"
// Copyright © 2023-2024 Apple Inc.
#include <metal_integer>
#include <metal_math>
constant mlx::os_log logger("mlx", "binary_ops");
struct Add {
template <typename T>
T operator()(T x, T y) thread {
return x + y;
}
};
struct FloorDivide {
template <typename T>
metal::enable_if_t<metal::is_integral_v<T> & !metal::is_signed_v<T>, T>
operator()(T x, T y) thread {
return x / y;
}
template <typename T>
metal::enable_if_t<metal::is_integral_v<T> & metal::is_signed_v<T>, T>
operator()(T x, T y) thread {
auto q = x / y;
if (x % y != 0 && (x < 0) != (y < 0)) {
q -= 1;
}
return q;
}
template <typename T>
metal::enable_if_t<!metal::is_integral_v<T>, T> operator()(T x, T y) thread {
return floor(x / y);
}
template <>
complex64_t operator()(complex64_t x, complex64_t y) thread {
// Complex is not supported, simply make compiler happy.
return x / y;
}
};
struct Divide {
template <typename T>
T operator()(T x, T y) thread {
return x / y;
}
};
struct Remainder {
template <typename T>
metal::enable_if_t<metal::is_integral_v<T> & !metal::is_signed_v<T>, T>
operator()(T x, T y) thread {
return x % y;
}
template <typename T>
metal::enable_if_t<metal::is_integral_v<T> & metal::is_signed_v<T>, T>
operator()(T x, T y) thread {
auto r = x % y;
if (r != 0 && (r < 0 != y < 0)) {
r += y;
}
return r;
}
template <typename T>
metal::enable_if_t<!metal::is_integral_v<T>, T> operator()(T x, T y) thread {
T r = fmod(x, y);
if (r != 0 && (r < 0 != y < 0)) {
r += y;
}
return r;
}
template <>
complex64_t operator()(complex64_t x, complex64_t y) thread {
return x % y;
}
};
struct Equal {
template <typename T>
bool operator()(T x, T y) thread {
return x == y;
}
};
struct NaNEqual {
template <typename T>
bool operator()(T x, T y) thread {
return x == y || (metal::isnan(x) && metal::isnan(y));
}
template <>
bool operator()(complex64_t x, complex64_t y) thread {
return x == y ||
(metal::isnan(x.real) && metal::isnan(y.real) && metal::isnan(x.imag) &&
metal::isnan(y.imag)) ||
(x.real == y.real && metal::isnan(x.imag) && metal::isnan(y.imag)) ||
(metal::isnan(x.real) && metal::isnan(y.real) && x.imag == y.imag);
}
};
struct Greater {
template <typename T>
bool operator()(T x, T y) thread {
return x > y;
}
};
struct GreaterEqual {
template <typename T>
bool operator()(T x, T y) thread {
return x >= y;
}
};
struct Less {
template <typename T>
bool operator()(T x, T y) thread {
return x < y;
}
};
struct LessEqual {
template <typename T>
bool operator()(T x, T y) thread {
return x <= y;
}
};
struct LogAddExp {
template <typename T>
T operator()(T x, T y) thread {
if (metal::isnan(x) || metal::isnan(y)) {
return metal::numeric_limits<T>::quiet_NaN();
}
constexpr T inf = metal::numeric_limits<T>::infinity();
T maxval = metal::max(x, y);
T minval = metal::min(x, y);
return (minval == -inf || maxval == inf)
? maxval
: (maxval + log1p(metal::exp(minval - maxval)));
};
complex64_t operator()(complex64_t x, complex64_t y) thread {
if (metal::isnan(x.real) || metal::isnan(x.imag) || metal::isnan(y.real) ||
metal::isnan(y.imag)) {
return metal::numeric_limits<float>::quiet_NaN();
}
constexpr float inf = metal::numeric_limits<float>::infinity();
complex64_t maxval = x > y ? x : y;
complex64_t minval = x < y ? x : y;
if (minval.real == -inf || maxval.real == inf)
return maxval;
float m = metal::exp(minval.real - maxval.real);
complex64_t dexp{
m * metal::cos(minval.imag - maxval.imag),
m * metal::sin(minval.imag - maxval.imag),
};
return maxval + log1p(dexp);
}
};
struct Maximum {
template <typename T>
metal::enable_if_t<metal::is_integral_v<T>, T> operator()(T x, T y) thread {
return metal::max(x, y);
}
template <typename T>
metal::enable_if_t<!metal::is_integral_v<T>, T> operator()(T x, T y) thread {
if (metal::isnan(x)) {
return x;
}
return x > y ? x : y;
}
template <>
complex64_t operator()(complex64_t x, complex64_t y) thread {
if (metal::isnan(x.real) || metal::isnan(x.imag)) {
return x;
}
return x > y ? x : y;
}
};
struct Minimum {
template <typename T>
metal::enable_if_t<metal::is_integral_v<T>, T> operator()(T x, T y) thread {
return metal::min(x, y);
}
template <typename T>
metal::enable_if_t<!metal::is_integral_v<T>, T> operator()(T x, T y) thread {
if (metal::isnan(x)) {
return x;
}
return x < y ? x : y;
}
template <>
complex64_t operator()(complex64_t x, complex64_t y) thread {
if (metal::isnan(x.real) || metal::isnan(x.imag)) {
return x;
}
return x < y ? x : y;
}
};
struct Multiply {
template <typename T>
T operator()(T x, T y) thread {
return x * y;
}
};
struct NotEqual {
template <typename T>
bool operator()(T x, T y) thread {
return x != y;
}
template <>
bool operator()(complex64_t x, complex64_t y) thread {
return x.real != y.real || x.imag != y.imag;
}
};
struct Power {
template <typename T>
metal::enable_if_t<!metal::is_integral_v<T>, T> operator()(T base, T exp)
thread {
return metal::pow(base, exp);
}
template <typename T>
metal::enable_if_t<metal::is_integral_v<T>, T> operator()(T base, T exp)
thread {
T res = 1;
// Undefined to raise integer to negative power
if (exp < 0) {
logger.log_debug(
"int pow exp<0 (base=%ld exp=%ld)", (long)base, (long)exp);
return 0;
}
while (exp) {
if (exp & 1) {
res *= base;
}
exp >>= 1;
base *= base;
}
return res;
}
template <>
complex64_t operator()(complex64_t x, complex64_t y) thread {
if (x.real == 0 && x.imag == 0) {
if (metal::isnan(y.real) || metal::isnan(y.imag)) {
auto nan = metal::numeric_limits<float>::quiet_NaN();
return {nan, nan};
}
return {0.0, 0.0};
}
auto x_theta = metal::atan2(x.imag, x.real);
auto x_ln_r = 0.5 * metal::log(x.real * x.real + x.imag * x.imag);
auto mag = metal::exp(y.real * x_ln_r - y.imag * x_theta);
auto phase = y.imag * x_ln_r + y.real * x_theta;
return {mag * metal::cos(phase), mag * metal::sin(phase)};
}
};
struct Subtract {
template <typename T>
T operator()(T x, T y) thread {
return x - y;
}
};
struct LogicalAnd {
template <typename T>
T operator()(T x, T y) thread {
return x && y;
};
};
struct LogicalOr {
template <typename T>
T operator()(T x, T y) thread {
return x || y;
};
};
struct BitwiseAnd {
template <typename T>
T operator()(T x, T y) thread {
return x & y;
};
};
struct BitwiseOr {
template <typename T>
T operator()(T x, T y) thread {
return x | y;
};
};
struct BitwiseXor {
template <typename T>
T operator()(T x, T y) thread {
return x ^ y;
};
};
struct LeftShift {
template <typename T>
T operator()(T x, T y) thread {
return x << y;
};
};
struct RightShift {
template <typename T>
T operator()(T x, T y) thread {
return x >> y;
};
};
struct ArcTan2 {
template <typename T>
T operator()(T y, T x) thread {
return metal::precise::atan2(y, x);
}
};
struct DivMod {
template <typename T>
metal::array<T, 2> operator()(T x, T y) thread {
return {FloorDivide{}(x, y), Remainder{}(x, y)};
};
};
///////////////////////////////////////////////////////////////////////////////
// Copyright © 2025 Apple Inc.
// Auto generated source for mlx/backend/metal/kernels/ternary_ops.h
///////////////////////////////////////////////////////////////////////////////
// Contents from "mlx/backend/metal/kernels/ternary_ops.h"
///////////////////////////////////////////////////////////////////////////////
#line 1 "mlx/backend/metal/kernels/ternary_ops.h"
// Copyright © 2023-2024 Apple Inc.
struct Select {
template <typename T>
T operator()(bool condition, T x, T y) thread {
return condition ? x : y;
}
};
///////////////////////////////////////////////////////////////////////////////
Reference in New Issue View Git Blame Copy Permalink