L-electron-Rare/mlx
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Use cuda::std for math ops (#3041)

Browse Source
This commit is contained in:
Cheng
2026-01-23 08:38:26 +09:00
committed by GitHub
parent faea3e6d34
commit c46c3833ee
4 changed files with 55 additions and 154 deletions
Download Patch File Download Diff File Expand all files Collapse all files
mlx/backend/cuda/device/binary_ops.cuh
+15 -14
View File
@@ -19,7 +19,7 @@ struct FloorDivide {
if constexpr (cuda::std::is_integral_v<T>) {
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<T>) {
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 <typename T>
__device__ bool operator()(T x, T y) {
using cuda::std::isnan;
if constexpr (is_complex_v<T>) {
return x == y ||
(isnan(x.real()) && isnan(y.real()) && isnan(x.imag()) &&
@@ -110,8 +111,8 @@ struct LogAddExp {
template <typename T>
__device__ T operator()(T x, T y) {
if constexpr (is_complex_v<T>) {
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<float>::quiet_NaN(),
cuda::std::numeric_limits<float>::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<T>::quiet_NaN();
}
T maxval = max(x, y);
@@ -138,7 +139,7 @@ struct LogAddExp {
return (minval == -cuda::std::numeric_limits<T>::infinity() ||
maxval == cuda::std::numeric_limits<T>::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<T>) {
return max(x, y);
} else if constexpr (is_complex_v<T>) {
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<T>) {
return min(x, y);
} else if constexpr (is_complex_v<T>) {
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<T>) {
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 <typename T>
__device__ T operator()(T y, T x) {
return atan2f(y, x);
return cuda::std::atan2(y, x);
}
};
mlx/backend/cuda/device/fp16_math.cuh
-98
View File
@@ -4,98 +4,14 @@
#include <cuda_bf16.h>
#include <cuda_fp16.h>
#include <cuda/std/limits>
#include <cuda/std/type_traits>
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 <typename T> \
__forceinline__ __device__ auto NAME(T x) { \
if constexpr (cuda::std::is_same_v<T, __half>) { \
return HALF_OP(x); \
} else { \
return ::NAME(x); \
} \
}
#else
#define MLX_DEFINE_UNARY_OP(NAME, HALF_OP) \
template <typename T> \
__forceinline__ __device__ auto NAME(T x) { \
if constexpr (cuda::std::is_same_v<T, __half>) { \
return HALF_OP(x); \
} else if constexpr (cuda::std::is_same_v<T, __nv_bfloat16>) { \
return HALF_OP(x); \
} else { \
return ::NAME(x); \
} \
}
#endif
#define MLX_DEFINE_UNARY_OP_FALLBCK(NAME) \
template <typename T> \
__forceinline__ __device__ auto NAME(T x) { \
if constexpr (cuda::std::is_same_v<T, __half>) { \
return ::NAME(__half2float(x)); \
} else if constexpr (cuda::std::is_same_v<T, __nv_bfloat16>) { \
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 <typename T> \
__forceinline__ __device__ auto NAME(T x, T y) { \
if constexpr (cuda::std::is_same_v<T, __half>) { \
return HALF_OP(x, y); \
} else { \
return ::NAME(x, y); \
} \
}
#else
#define MLX_DEFINE_BINARY_OP(NAME, HALF_OP) \
template <typename T> \
__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 <typename T>
__forceinline__ __device__ T fmod(T x, T y) {
if constexpr (cuda::std::is_same_v<T, __half>) {
return __float2half(::fmod(__half2float(x), __half2float(y)));
#if CUDART_VERSION >= 12000 || __CUDA_ARCH__ >= 800
} else if constexpr (cuda::std::is_same_v<T, __nv_bfloat16>) {
return __float2bfloat16(::fmod(__bfloat162float(x), __bfloat162float(y)));
#endif
} else {
return ::fmod(x, y);
}
}
///////////////////////////////////////////////////////////////////////////////
// Additional C++ operator overrides between half types and native types.
///////////////////////////////////////////////////////////////////////////////
mlx/backend/cuda/device/unary_ops.cuh
+34 -36
View File
@@ -2,12 +2,12 @@
#pragma once
#include <cuda_fp8.h>
#include "mlx/backend/cuda/device/fp16_math.cuh"
#include "mlx/backend/cuda/device/utils.cuh"
#include <cuda_fp8.h>
#include <math_constants.h>
#include <cuda/std/cmath>
namespace mlx::core::cu {
@@ -17,7 +17,7 @@ struct Abs {
if constexpr (cuda::std::is_unsigned_v<T>) {
return x;
} else {
return abs(x);
return cuda::std::abs(x);
}
}
};
@@ -25,42 +25,42 @@ struct Abs {
struct ArcCos {
template <typename T>
__device__ T operator()(T x) {
return acos(x);
return cuda::std::acos(x);
}
};
struct ArcCosh {
template <typename T>
__device__ T operator()(T x) {
return acosh(x);
return cuda::std::acosh(x);
}
};
struct ArcSin {
template <typename T>
__device__ T operator()(T x) {
return asin(x);
return cuda::std::asin(x);
}
};
struct ArcSinh {
template <typename T>
__device__ T operator()(T x) {
return asinh(x);
return cuda::std::asinh(x);
}
};
struct ArcTan {
template <typename T>
__device__ T operator()(T x) {
return atan(x);
return cuda::std::atan(x);
}
};
struct ArcTanh {
template <typename T>
__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<T>) {
return x;
} else if constexpr (is_complex_v<T>) {
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 <typename T>
__device__ complex_t<T> operator()(complex_t<T> x) {
return conj(x);
return cuda::std::conj(x);
}
};
struct Cos {
template <typename T>
__device__ T operator()(T x) {
return cos(x);
return cuda::std::cos(x);
}
};
struct Cosh {
template <typename T>
__device__ T operator()(T x) {
return cosh(x);
return cuda::std::cosh(x);
}
};
@@ -134,20 +134,14 @@ struct ErfInv {
struct Exp {
template <typename T>
__device__ T operator()(T x) {
return exp(x);
return cuda::std::exp(x);
}
};
struct Expm1 {
template <typename T>
__device__ T operator()(T x) {
if constexpr (cuda::std::is_same_v<T, __half>) {
return expm1(__half2float(x));
} else if constexpr (cuda::std::is_same_v<T, __nv_bfloat16>) {
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<T>) {
return x;
} else if constexpr (is_complex_v<T>) {
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 <typename T>
__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 <typename T>
__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 <typename T>
__device__ T operator()(T x) {
if constexpr (is_complex_v<T>) {
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 <typename T>
__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 <typename T>
__device__ T operator()(T x) {
return sin(x);
return cuda::std::sin(x);
}
};
struct Sinh {
template <typename T>
__device__ T operator()(T x) {
return sinh(x);
return cuda::std::sinh(x);
}
};
@@ -307,7 +301,7 @@ struct Square {
struct Sqrt {
template <typename T>
__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<T>) {
return 1.0f / Sqrt{}(x);
} else if constexpr (cuda::std::is_same_v<T, __half>) {
return rsqrt(__half2float(x));
} else if constexpr (cuda::std::is_same_v<T, __nv_bfloat16>) {
return rsqrt(__bfloat162float(x));
} else {
return rsqrt(x);
}
@@ -325,14 +323,14 @@ struct Rsqrt {
struct Tan {
template <typename T>
__device__ T operator()(T x) {
return tan(x);
return cuda::std::tan(x);
}
};
struct Tanh {
template <typename T>
__device__ T operator()(T x) {
return tanh(x);
return cuda::std::tanh(x);
}
};
mlx/backend/cuda/reduce/reduce_ops.cuh
+6 -6
View File
@@ -70,14 +70,14 @@ struct Min {
template <typename T>
__device__ __forceinline__ T operator()(T a, T b) {
if constexpr (is_complex_v<T>) {
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<T>) {
if (isnan(a) || isnan(b)) {
if (cuda::std::isnan(a) || cuda::std::isnan(b)) {
return cuda::std::numeric_limits<float>::quiet_NaN();
}
}
@@ -94,14 +94,14 @@ struct Max {
template <typename T>
__device__ __forceinline__ T operator()(T a, T b) {
if constexpr (is_complex_v<T>) {
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<T>) {
if (isnan(a) || isnan(b)) {
if (cuda::std::isnan(a) || cuda::std::isnan(b)) {
return cuda::std::numeric_limits<float>::quiet_NaN();
}
}