From c46c3833ee72f756db88ce36b04dcdb336aa9f63 Mon Sep 17 00:00:00 2001 From: Cheng Date: Fri, 23 Jan 2026 08:38:26 +0900 Subject: [PATCH] Use cuda::std for math ops (#3041) --- mlx/backend/cuda/device/binary_ops.cuh | 29 ++++---- mlx/backend/cuda/device/fp16_math.cuh | 98 -------------------------- mlx/backend/cuda/device/unary_ops.cuh | 70 +++++++++--------- mlx/backend/cuda/reduce/reduce_ops.cuh | 12 ++-- 4 files changed, 55 insertions(+), 154 deletions(-) diff --git a/mlx/backend/cuda/device/binary_ops.cuh b/mlx/backend/cuda/device/binary_ops.cuh index 1ff22cff..b0b79628 100644 --- a/mlx/backend/cuda/device/binary_ops.cuh +++ b/mlx/backend/cuda/device/binary_ops.cuh @@ -19,7 +19,7 @@ struct FloorDivide { if constexpr (cuda::std::is_integral_v) { return x / y; } else { - return truncf(x / y); + return cuda::std::trunc(x / y); } } }; @@ -47,7 +47,7 @@ struct Remainder { } else if constexpr (is_complex_v) { return x % y; } else { - T r = fmod(x, y); + T r = cuda::std::fmod(x, y); if (r != 0 && (r < 0 != y < 0)) { r = r + y; } @@ -66,6 +66,7 @@ struct Equal { struct NaNEqual { template __device__ bool operator()(T x, T y) { + using cuda::std::isnan; if constexpr (is_complex_v) { return x == y || (isnan(x.real()) && isnan(y.real()) && isnan(x.imag()) && @@ -110,8 +111,8 @@ struct LogAddExp { template __device__ T operator()(T x, T y) { if constexpr (is_complex_v) { - if (isnan(x.real()) || isnan(x.imag()) || isnan(y.real()) || - isnan(y.imag())) { + if (cuda::std::isnan(x.real()) || cuda::std::isnan(x.imag()) || + cuda::std::isnan(y.real()) || cuda::std::isnan(y.imag())) { return { cuda::std::numeric_limits::quiet_NaN(), cuda::std::numeric_limits::quiet_NaN()}; @@ -120,7 +121,7 @@ struct LogAddExp { auto min = x.real() < y.real() ? x : y; auto min_real = min.real(); auto max_real = max.real(); - if (!isfinite(min_real) && (min_real == max_real)) { + if (!cuda::std::isfinite(min_real) && (min_real == max_real)) { if (min_real < 0) { return min; } else { @@ -130,7 +131,7 @@ struct LogAddExp { return Log1p{}(Exp{}(min - max)) + max; } } else { - if (isnan(x) || isnan(y)) { + if (cuda::std::isnan(x) || cuda::std::isnan(y)) { return cuda::std::numeric_limits::quiet_NaN(); } T maxval = max(x, y); @@ -138,7 +139,7 @@ struct LogAddExp { return (minval == -cuda::std::numeric_limits::infinity() || maxval == cuda::std::numeric_limits::infinity()) ? maxval - : T(float(maxval) + log1p(expf(minval - maxval))); + : T(maxval + cuda::std::log1p(cuda::std::exp(minval - maxval))); } }; }; @@ -149,12 +150,12 @@ struct Maximum { if constexpr (cuda::std::is_integral_v) { return max(x, y); } else if constexpr (is_complex_v) { - if (isnan(x.real()) || isnan(x.imag())) { + if (cuda::std::isnan(x.real()) || cuda::std::isnan(x.imag())) { return x; } return x > y ? x : y; } else { - if (isnan(x)) { + if (cuda::std::isnan(x)) { return x; } return x > y ? x : y; @@ -168,12 +169,12 @@ struct Minimum { if constexpr (cuda::std::is_integral_v) { return min(x, y); } else if constexpr (is_complex_v) { - if (isnan(x.real()) || isnan(x.imag())) { + if (cuda::std::isnan(x.real()) || cuda::std::isnan(x.imag())) { return x; } return x < y ? x : y; } else { - if (isnan(x)) { + if (cuda::std::isnan(x)) { return x; } return x < y ? x : y; @@ -219,9 +220,9 @@ struct Power { } return res; } else if constexpr (is_complex_v) { - return pow(base, exp); + return cuda::std::pow(base, exp); } else { - return powf(base, exp); + return cuda::std::pow(base, exp); } } }; @@ -285,7 +286,7 @@ struct RightShift { struct ArcTan2 { template __device__ T operator()(T y, T x) { - return atan2f(y, x); + return cuda::std::atan2(y, x); } }; diff --git a/mlx/backend/cuda/device/fp16_math.cuh b/mlx/backend/cuda/device/fp16_math.cuh index f6fa17bb..27998594 100644 --- a/mlx/backend/cuda/device/fp16_math.cuh +++ b/mlx/backend/cuda/device/fp16_math.cuh @@ -4,98 +4,14 @@ #include #include -#include #include namespace mlx::core::cu { -/////////////////////////////////////////////////////////////////////////////// -// Unary ops for half types. -/////////////////////////////////////////////////////////////////////////////// - -#if CUDART_VERSION < 12000 && __CUDA_ARCH__ < 800 -#define MLX_DEFINE_UNARY_OP(NAME, HALF_OP) \ - template \ - __forceinline__ __device__ auto NAME(T x) { \ - if constexpr (cuda::std::is_same_v) { \ - return HALF_OP(x); \ - } else { \ - return ::NAME(x); \ - } \ - } -#else -#define MLX_DEFINE_UNARY_OP(NAME, HALF_OP) \ - template \ - __forceinline__ __device__ auto NAME(T x) { \ - if constexpr (cuda::std::is_same_v) { \ - return HALF_OP(x); \ - } else if constexpr (cuda::std::is_same_v) { \ - return HALF_OP(x); \ - } else { \ - return ::NAME(x); \ - } \ - } -#endif - -#define MLX_DEFINE_UNARY_OP_FALLBCK(NAME) \ - template \ - __forceinline__ __device__ auto NAME(T x) { \ - if constexpr (cuda::std::is_same_v) { \ - return ::NAME(__half2float(x)); \ - } else if constexpr (cuda::std::is_same_v) { \ - return ::NAME(__bfloat162float(x)); \ - } else { \ - return ::NAME(x); \ - } \ - } - -MLX_DEFINE_UNARY_OP(abs, __habs) -MLX_DEFINE_UNARY_OP(ceil, hceil) -MLX_DEFINE_UNARY_OP(cos, hcos) -MLX_DEFINE_UNARY_OP(exp, hexp) -MLX_DEFINE_UNARY_OP(floor, hfloor) -MLX_DEFINE_UNARY_OP(isnan, __hisnan) -MLX_DEFINE_UNARY_OP(log, hlog) -MLX_DEFINE_UNARY_OP(log2, hlog2) -MLX_DEFINE_UNARY_OP(log10, hlog10) -MLX_DEFINE_UNARY_OP(rint, hrint) -MLX_DEFINE_UNARY_OP(rsqrt, hrsqrt) -MLX_DEFINE_UNARY_OP(sin, hsin) -MLX_DEFINE_UNARY_OP(sqrt, hsqrt) -MLX_DEFINE_UNARY_OP_FALLBCK(acos) -MLX_DEFINE_UNARY_OP_FALLBCK(acosh) -MLX_DEFINE_UNARY_OP_FALLBCK(asin) -MLX_DEFINE_UNARY_OP_FALLBCK(asinh) -MLX_DEFINE_UNARY_OP_FALLBCK(atan) -MLX_DEFINE_UNARY_OP_FALLBCK(atanh) -MLX_DEFINE_UNARY_OP_FALLBCK(cosh) -MLX_DEFINE_UNARY_OP_FALLBCK(log1p) -MLX_DEFINE_UNARY_OP_FALLBCK(sinh) -MLX_DEFINE_UNARY_OP_FALLBCK(tan) -#if __CUDA_ARCH__ >= 1280 -MLX_DEFINE_UNARY_OP(tanh, htanh) -#else -MLX_DEFINE_UNARY_OP_FALLBCK(tanh) -#endif - -#undef MLX_DEFINE_UNARY_OP -#undef MLX_DEFINE_UNARY_OP_FALLBCK - /////////////////////////////////////////////////////////////////////////////// // Binary ops for half types. /////////////////////////////////////////////////////////////////////////////// -#if CUDART_VERSION < 12000 && __CUDA_ARCH__ < 800 -#define MLX_DEFINE_BINARY_OP(NAME, HALF_OP) \ - template \ - __forceinline__ __device__ auto NAME(T x, T y) { \ - if constexpr (cuda::std::is_same_v) { \ - return HALF_OP(x, y); \ - } else { \ - return ::NAME(x, y); \ - } \ - } -#else #define MLX_DEFINE_BINARY_OP(NAME, HALF_OP) \ template \ __forceinline__ __device__ auto NAME(T x, T y) { \ @@ -107,26 +23,12 @@ MLX_DEFINE_UNARY_OP_FALLBCK(tanh) return ::NAME(x, y); \ } \ } -#endif MLX_DEFINE_BINARY_OP(max, __hmax) MLX_DEFINE_BINARY_OP(min, __hmin) #undef MLX_DEFINE_BINARY_OP -template -__forceinline__ __device__ T fmod(T x, T y) { - if constexpr (cuda::std::is_same_v) { - return __float2half(::fmod(__half2float(x), __half2float(y))); -#if CUDART_VERSION >= 12000 || __CUDA_ARCH__ >= 800 - } else if constexpr (cuda::std::is_same_v) { - return __float2bfloat16(::fmod(__bfloat162float(x), __bfloat162float(y))); -#endif - } else { - return ::fmod(x, y); - } -} - /////////////////////////////////////////////////////////////////////////////// // Additional C++ operator overrides between half types and native types. /////////////////////////////////////////////////////////////////////////////// diff --git a/mlx/backend/cuda/device/unary_ops.cuh b/mlx/backend/cuda/device/unary_ops.cuh index 61c587e2..8234cf8e 100644 --- a/mlx/backend/cuda/device/unary_ops.cuh +++ b/mlx/backend/cuda/device/unary_ops.cuh @@ -2,12 +2,12 @@ #pragma once -#include - #include "mlx/backend/cuda/device/fp16_math.cuh" #include "mlx/backend/cuda/device/utils.cuh" +#include #include +#include namespace mlx::core::cu { @@ -17,7 +17,7 @@ struct Abs { if constexpr (cuda::std::is_unsigned_v) { return x; } else { - return abs(x); + return cuda::std::abs(x); } } }; @@ -25,42 +25,42 @@ struct Abs { struct ArcCos { template __device__ T operator()(T x) { - return acos(x); + return cuda::std::acos(x); } }; struct ArcCosh { template __device__ T operator()(T x) { - return acosh(x); + return cuda::std::acosh(x); } }; struct ArcSin { template __device__ T operator()(T x) { - return asin(x); + return cuda::std::asin(x); } }; struct ArcSinh { template __device__ T operator()(T x) { - return asinh(x); + return cuda::std::asinh(x); } }; struct ArcTan { template __device__ T operator()(T x) { - return atan(x); + return cuda::std::atan(x); } }; struct ArcTanh { template __device__ T operator()(T x) { - return atanh(x); + return cuda::std::atanh(x); } }; @@ -77,9 +77,9 @@ struct Ceil { if constexpr (cuda::std::is_integral_v) { return x; } else if constexpr (is_complex_v) { - return T{ceil(x.real()), ceil(x.imag())}; + return T{cuda::std::ceil(x.real()), cuda::std::ceil(x.imag())}; } else { - return ceil(x); + return cuda::std::ceil(x); } } }; @@ -87,21 +87,21 @@ struct Ceil { struct Conjugate { template __device__ complex_t operator()(complex_t x) { - return conj(x); + return cuda::std::conj(x); } }; struct Cos { template __device__ T operator()(T x) { - return cos(x); + return cuda::std::cos(x); } }; struct Cosh { template __device__ T operator()(T x) { - return cosh(x); + return cuda::std::cosh(x); } }; @@ -134,20 +134,14 @@ struct ErfInv { struct Exp { template __device__ T operator()(T x) { - return exp(x); + return cuda::std::exp(x); } }; struct Expm1 { template __device__ T operator()(T x) { - if constexpr (cuda::std::is_same_v) { - return expm1(__half2float(x)); - } else if constexpr (cuda::std::is_same_v) { - return expm1(__bfloat162float(x)); - } else { - return expm1(x); - } + return cuda::std::expm1(x); } }; @@ -157,9 +151,9 @@ struct Floor { if constexpr (cuda::std::is_integral_v) { return x; } else if constexpr (is_complex_v) { - return T{floor(x.real()), floor(x.imag())}; + return T{cuda::std::floor(x.real()), cuda::std::floor(x.imag())}; } else { - return floor(x); + return cuda::std::floor(x); } } }; @@ -174,7 +168,7 @@ struct Imag { struct Log { template __device__ T operator()(T x) { - return log(x); + return cuda::std::log(x); } }; @@ -185,7 +179,7 @@ struct Log2 { auto y = Log{}(x); return {y.real() / CUDART_LN2_F, y.imag() / CUDART_LN2_F}; } else { - return log2(x); + return cuda::std::log2(x); } } }; @@ -193,7 +187,7 @@ struct Log2 { struct Log10 { template __device__ T operator()(T x) { - return log10(x); + return cuda::std::log10(x); } }; @@ -216,7 +210,7 @@ struct Log1p { return {logf(z0), theta}; } } else { - return log1p(z); + return cuda::std::log1p(z); } } }; @@ -249,9 +243,9 @@ struct Round { template __device__ T operator()(T x) { if constexpr (is_complex_v) { - return {rint(x.real()), rint(x.imag())}; + return {cuda::std::rint(x.real()), cuda::std::rint(x.imag())}; } else { - return rint(x); + return cuda::std::rint(x); } } }; @@ -259,7 +253,7 @@ struct Round { struct Sigmoid { template __device__ T operator()(T x) { - T y = 1 / (1 + exp(abs(x))); + T y = 1 / (1 + cuda::std::exp(cuda::std::abs(x))); return (x < 0) ? y : 1 - y; } }; @@ -286,14 +280,14 @@ struct Sign { struct Sin { template __device__ T operator()(T x) { - return sin(x); + return cuda::std::sin(x); } }; struct Sinh { template __device__ T operator()(T x) { - return sinh(x); + return cuda::std::sinh(x); } }; @@ -307,7 +301,7 @@ struct Square { struct Sqrt { template __device__ T operator()(T x) { - return sqrt(x); + return cuda::std::sqrt(x); } }; @@ -316,6 +310,10 @@ struct Rsqrt { __device__ T operator()(T x) { if constexpr (is_complex_v) { return 1.0f / Sqrt{}(x); + } else if constexpr (cuda::std::is_same_v) { + return rsqrt(__half2float(x)); + } else if constexpr (cuda::std::is_same_v) { + return rsqrt(__bfloat162float(x)); } else { return rsqrt(x); } @@ -325,14 +323,14 @@ struct Rsqrt { struct Tan { template __device__ T operator()(T x) { - return tan(x); + return cuda::std::tan(x); } }; struct Tanh { template __device__ T operator()(T x) { - return tanh(x); + return cuda::std::tanh(x); } }; diff --git a/mlx/backend/cuda/reduce/reduce_ops.cuh b/mlx/backend/cuda/reduce/reduce_ops.cuh index 7f8cad0c..6c6b1827 100644 --- a/mlx/backend/cuda/reduce/reduce_ops.cuh +++ b/mlx/backend/cuda/reduce/reduce_ops.cuh @@ -70,14 +70,14 @@ struct Min { template __device__ __forceinline__ T operator()(T a, T b) { if constexpr (is_complex_v) { - if (isnan(a.real()) || isnan(a.imag())) { + if (cuda::std::isnan(a.real()) || cuda::std::isnan(a.imag())) { return a; } - if (isnan(b.real()) || isnan(b.imag())) { + if (cuda::std::isnan(b.real()) || cuda::std::isnan(b.imag())) { return b; } } else if constexpr (!cuda::std::is_integral_v) { - if (isnan(a) || isnan(b)) { + if (cuda::std::isnan(a) || cuda::std::isnan(b)) { return cuda::std::numeric_limits::quiet_NaN(); } } @@ -94,14 +94,14 @@ struct Max { template __device__ __forceinline__ T operator()(T a, T b) { if constexpr (is_complex_v) { - if (isnan(a.real()) || isnan(a.imag())) { + if (cuda::std::isnan(a.real()) || cuda::std::isnan(a.imag())) { return a; } - if (isnan(b.real()) || isnan(b.imag())) { + if (cuda::std::isnan(b.real()) || cuda::std::isnan(b.imag())) { return b; } } else if constexpr (!cuda::std::is_integral_v) { - if (isnan(a) || isnan(b)) { + if (cuda::std::isnan(a) || cuda::std::isnan(b)) { return cuda::std::numeric_limits::quiet_NaN(); } }