kolandr/gridbot
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Initial import: grid-bot — grid trading bot for BTC-USDT on Cifra Markets

Browse Source
This commit is contained in:
Kolp
2026-09-24 13:22:23 +07:00
commit 642cc11a9f
18968 changed files with 5683248 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
kronos-venv/lib/python3.12/site-packages/torch/include/c10/metal/atomic.h
+289
View File
@@ -0,0 +1,289 @@
#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
#pragma once
#include <metal_atomic>
namespace c10 {
namespace metal {
// Atomic operations helper
template <typename T>
struct AtomicType {};
template <typename T>
using AtomicType_t = typename AtomicType<T>::type;
template <typename AT, typename T>
static inline void atomic_binary_op_helper(
device ::metal::atomic<AT>* data,
long offset,
T value,
T (*op)(T, T)) {
auto ptr = data + offset;
auto old = ::metal::atomic_load_explicit(ptr, ::metal::memory_order_relaxed);
T val;
do {
val = op(old, value);
} while (!::metal::atomic_compare_exchange_weak_explicit(
ptr,
&old,
val,
::metal::memory_order_relaxed,
::metal::memory_order_relaxed));
}
template <>
struct AtomicType<float> {
using type = ::metal::atomic<float>;
static inline void atomic_add(device type* data, long offset, float value) {
::metal::atomic_fetch_add_explicit(
data + offset, value, ::metal::memory_order_relaxed);
}
static inline void atomic_binary_op(
device type* data,
long offset,
float value,
float (*op)(float, float)) {
atomic_binary_op_helper(data, offset, value, op);
}
};
template <>
struct AtomicType<int> {
using type = ::metal::atomic<int>;
static inline void atomic_add(device type* data, long offset, int value) {
::metal::atomic_fetch_add_explicit(
data + offset, value, ::metal::memory_order_relaxed);
}
static inline void atomic_binary_op(
device type* data,
long offset,
int value,
int (*op)(int, int)) {
atomic_binary_op_helper(data, offset, value, op);
}
};
// As of Metal3.2 atomic operations are not supported on half-precision floats,
// so they must be simulated Using atomic compare and exchange over 32-bit
// atomic type
template <typename T>
static inline void atomic_add_helper(
device ::metal::atomic<uint>* data,
long offset,
T value) {
// atomic<uint> requires 4-byte alignment; fix up misaligned pointers
auto addr = reinterpret_cast<ulong>(data);
auto misalign = (addr % alignof(::metal::atomic<uint>)) / sizeof(T);
data = reinterpret_cast<device ::metal::atomic<uint>*>(
reinterpret_cast<device char*>(data) - misalign * sizeof(T));
offset += misalign;
constexpr auto elem_per_enum = sizeof(uint) / sizeof(T);
auto ptr = data + (offset / elem_per_enum);
auto old = ::metal::atomic_load_explicit(ptr, ::metal::memory_order_relaxed);
union {
uint i;
T t[elem_per_enum];
} val;
do {
val.i = old;
val.t[offset & (elem_per_enum - 1)] += value;
} while (!::metal::atomic_compare_exchange_weak_explicit(
ptr,
&old,
val.i,
::metal::memory_order_relaxed,
::metal::memory_order_relaxed));
}
template <typename T>
static inline void atomic_binary_op_helper(
device ::metal::atomic<uint>* data,
long offset,
T value,
T (*Op)(T, T)) {
// atomic<uint> requires 4-byte alignment; fix up misaligned pointers
auto addr = reinterpret_cast<ulong>(data);
auto misalign = (addr % alignof(::metal::atomic<uint>)) / sizeof(T);
data = reinterpret_cast<device ::metal::atomic<uint>*>(
reinterpret_cast<device char*>(data) - misalign * sizeof(T));
offset += misalign;
constexpr auto elem_per_enum = sizeof(uint) / sizeof(T);
auto ptr = data + (offset / elem_per_enum);
auto old = ::metal::atomic_load_explicit(ptr, ::metal::memory_order_relaxed);
union {
uint i;
T t[elem_per_enum];
} val;
do {
val.i = old;
val.t[offset & (elem_per_enum - 1)] =
Op(val.t[offset & (elem_per_enum - 1)], value);
} while (!::metal::atomic_compare_exchange_weak_explicit(
ptr,
&old,
val.i,
::metal::memory_order_relaxed,
::metal::memory_order_relaxed));
}
template <>
struct AtomicType<half> {
using type = ::metal::atomic<uint>;
static inline void atomic_add(device type* data, long offset, half value) {
atomic_add_helper(data, offset, value);
}
static inline void atomic_binary_op(
device type* data,
long offset,
half value,
half (*op)(half, half)) {
atomic_binary_op_helper(data, offset, value, op);
}
};
template <>
struct AtomicType<short> {
using type = ::metal::atomic<uint>;
static inline void atomic_add(device type* data, long offset, short value) {
atomic_add_helper(data, offset, value);
}
static inline void atomic_binary_op(
device type* data,
long offset,
short value,
short (*op)(short, short)) {
atomic_binary_op_helper(data, offset, value, op);
}
};
template <>
struct AtomicType<char> {
using type = ::metal::atomic<uint>;
static inline void atomic_add(device type* data, long offset, char value) {
atomic_add_helper(data, offset, value);
}
static inline void atomic_binary_op(
device type* data,
long offset,
char value,
char (*op)(char, char)) {
atomic_binary_op_helper(data, offset, value, op);
}
};
template <>
struct AtomicType<uchar> {
using type = ::metal::atomic<uint>;
static inline void atomic_add(device type* data, long offset, char value) {
atomic_add_helper(data, offset, value);
}
static inline void atomic_binary_op(
device type* data,
long offset,
uchar value,
uchar (*op)(uchar, uchar)) {
atomic_binary_op_helper(data, offset, value, op);
}
};
template <>
struct AtomicType<bfloat> {
using type = ::metal::atomic<uint>;
static inline void atomic_add(device type* data, long offset, bfloat value) {
atomic_add_helper<bfloat>(data, offset, value);
}
static inline void atomic_binary_op(
device type* data,
long offset,
bfloat value,
bfloat (*op)(bfloat, bfloat)) {
atomic_binary_op_helper(data, offset, value, op);
}
};
// Metal supports atomic_store_explicit for bools, but
// sizeof(::metal::atomic_bool) is 4 Therefore it could not be used to
// atomically modify unaligned memory, so fall back to compare and exchange
// trick As accumulation over booleans are just or operation, do nothing if
// value is false
template <>
struct AtomicType<bool> {
using type = ::metal::atomic<uint>;
static inline void atomic_add(device type* data, long offset, bool value) {
if (!value) {
return;
}
auto ptr = data + (offset >> 2);
auto old =
::metal::atomic_load_explicit(ptr, ::metal::memory_order_relaxed);
union {
uint i;
bool t[4];
} val;
do {
val.i = old;
val.t[offset & 3] = true;
} while (!::metal::atomic_compare_exchange_weak_explicit(
ptr,
&old,
val.i,
::metal::memory_order_relaxed,
::metal::memory_order_relaxed));
}
};
// ComplexHalf atomic op
template <>
struct AtomicType<half2> {
using type = ::metal::atomic<uint>;
static inline void atomic_add(device type* data, long offset, half2 value) {
auto ptr = data + offset;
auto old =
::metal::atomic_load_explicit(ptr, ::metal::memory_order_relaxed);
while (!::metal::atomic_compare_exchange_weak_explicit(
ptr,
&old,
as_type<uint>(as_type<half2>(old) + value),
::metal::memory_order_relaxed,
::metal::memory_order_relaxed))
;
}
};
// There are no atomic 64-bit add in Metal yet, but templates below implements a
// consistent add I.e. if multiple threads are modify the same 64-bit value,
// results stored at the address will eventually be equal to its original value
// plus sum of all operands
template <>
struct AtomicType<long> {
using type = ::metal::atomic<uint>;
static inline void atomic_add(device type* data, long offset, long value) {
const auto value_bits = as_type<ulong>(value);
const uint low = static_cast<uint>(value_bits);
uint high = static_cast<uint>(value_bits >> 32);
auto ptr = data + (offset << 1);
auto old_low =
atomic_fetch_add_explicit(ptr, low, ::metal::memory_order_relaxed);
high += (old_low + low < old_low) ? 1 : 0;
atomic_fetch_add_explicit(ptr + 1, high, ::metal::memory_order_relaxed);
}
};
// ComplexFloat atomic op, which again is not really atomic, but eventually
// consistent
template <>
struct AtomicType<float2> {
using type = ::metal::atomic<float>;
static inline void atomic_add(device type* data, long offset, float2 value) {
auto ptr = data + (offset << 1);
atomic_fetch_add_explicit(ptr + 0, value.x, ::metal::memory_order_relaxed);
atomic_fetch_add_explicit(ptr + 1, value.y, ::metal::memory_order_relaxed);
}
};
} // namespace metal
} // namespace c10
#else
#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
#endif // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
kronos-venv/lib/python3.12/site-packages/torch/include/c10/metal/common.h
+53
View File
@@ -0,0 +1,53 @@
#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
#pragma once
// Set of global constants that could be shareable between CPU and Metal code
#ifdef __METAL__
#include <metal_array>
#define C10_METAL_CONSTEXPR constant constexpr
#else
#include <array>
#define C10_METAL_CONSTEXPR constexpr
#endif
#define C10_METAL_ALL_TYPES_FUNCTOR(_) \
_(Byte, 0) \
_(Char, 1) \
_(Short, 2) \
_(Int, 3) \
_(Long, 4) \
_(Half, 5) \
_(Float, 6) \
_(ComplexHalf, 8) \
_(ComplexFloat, 9) \
_(Bool, 11) \
_(BFloat16, 15) \
_(UInt16, 27) \
_(UInt32, 28) \
_(UInt64, 29)
namespace c10 {
namespace metal {
C10_METAL_CONSTEXPR unsigned max_ndim = 16;
C10_METAL_CONSTEXPR unsigned simdgroup_size = 32;
#ifdef __METAL__
template <typename T, unsigned N>
using array = ::metal::array<T, N>;
#else
template <typename T, unsigned N>
using array = std::array<T, N>;
#endif
enum class ScalarType {
#define _DEFINE_ENUM_VAL_(_v, _n) _v = _n,
C10_METAL_ALL_TYPES_FUNCTOR(_DEFINE_ENUM_VAL_)
#undef _DEFINE_ENUM_VAL_
};
} // namespace metal
} // namespace c10
#else
#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
#endif // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
kronos-venv/lib/python3.12/site-packages/torch/include/c10/metal/error.h
+116
View File
@@ -0,0 +1,116 @@
#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
#pragma once
#include <c10/metal/common.h>
namespace c10 {
namespace metal {
C10_METAL_CONSTEXPR unsigned error_message_count = 30;
struct ErrorMessage {
char file[128];
char func[128];
char message[250];
unsigned int line;
};
struct ErrorMessages {
#ifdef __METAL__
::metal::atomic<unsigned int> count;
#else
unsigned int count;
#endif
ErrorMessage msg[error_message_count];
};
#ifdef __METAL__
namespace detail {
static uint strncpy(device char* dst, constant const char* src, unsigned len) {
uint i = 0;
while (src[i] != 0 && i < len - 1) {
dst[i] = src[i];
i++;
}
dst[i] = 0;
return i;
}
inline uint print_arg(
device char* ptr,
unsigned len,
constant const char* arg) {
return strncpy(ptr, arg, len);
}
// Returns number length as string in base10
static inline uint base10_length(long num) {
uint rc = 1;
if (num < 0) {
num = -num;
rc += 1;
}
while (num > 9) {
num /= 10;
rc++;
}
return rc;
}
// Converts signed integer to string
inline uint print_arg(device char* ptr, unsigned len, long arg) {
const auto arg_len = base10_length(arg);
if (arg_len >= len)
return 0;
if (arg < 0) {
ptr[0] = '-';
arg = -arg;
}
uint idx = 1;
do {
ptr[arg_len - idx] = '0' + (arg % 10);
arg /= 10;
idx++;
} while (arg > 0);
ptr[arg_len] = 0;
return arg_len;
}
template <typename T>
inline void print_args(device char* ptr, unsigned len, T arg) {
print_arg(ptr, len, arg);
}
template <typename T, typename... Args>
inline void print_args(device char* ptr, unsigned len, T arg, Args... args) {
const auto rc = print_arg(ptr, len, arg);
print_args(ptr + rc, len - rc, args...);
}
} // namespace detail
template <typename... Args>
static void report_error(
device ErrorMessages* msgs,
constant const char* file,
int line,
constant const char* func,
Args... args) {
const auto idx =
atomic_fetch_add_explicit(&msgs->count, 1, ::metal::memory_order_relaxed);
if (idx >= error_message_count) {
return;
}
device auto* msg = &msgs->msg[idx];
detail::strncpy(msg->file, file, 128);
detail::strncpy(msg->func, func, 128);
detail::print_args(msg->message, 250, args...);
msg->line = line;
}
#define TORCH_REPORT_ERROR(buf, ...) \
::c10::metal::report_error(buf, __FILE__, __LINE__, __func__, __VA_ARGS__)
#endif
} // namespace metal
} // namespace c10
#else
#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
#endif // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
kronos-venv/lib/python3.12/site-packages/torch/include/c10/metal/expm1f.h
+102
View File
@@ -0,0 +1,102 @@
#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
// Copy-and-pasted from:
// https://github.com/ml-explore/mlx/blob/99c33d011d63174f50cea37c3eede002958be6d3/mlx/backend/metal/kernels/expm1f.h
#pragma once
#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.
namespace c10 {
namespace metal {
/* 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)
*/
inline 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 = ::metal::fma(1.442695f, a, 12582912.f); // 0x1.715476p0, 0x1.8p23
j = j - 12582912.0f; // 0x1.8p23
i = (int)j;
f = ::metal::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 = ::metal::fma(r, f, 1.39309070e-3f); // 0x1.6d30bcp-10
r = ::metal::fma(r, f, 8.33343994e-3f); // 0x1.1111f6p-7
r = ::metal::fma(r, f, 4.16668020e-2f); // 0x1.55559ep-5
r = ::metal::fma(r, f, 1.66666716e-1f); // 0x1.55555cp-3
r = ::metal::fma(r, f, 4.99999970e-1f); // 0x1.fffffep-2
u = (j == 1) ? (f + 0.5f) : f;
v = ::metal::fma(r, s, u);
s = 0.5f * b;
t = ::metal::ldexp(s, i);
y = t - s;
x = (t - y) - s; // double-float canonicalization of difference
r = ::metal::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 */
inline float expm1f(float a) {
float r;
r = expm1f_scaled_unchecked(a, 1.0f);
/* handle severe overflow and underflow */
if (::metal::abs(a - 1.0f) > 88.0f) {
r = ::metal::pow(2, a);
r = ::metal::fma(r, r, -1.0f);
}
return r;
}
} // namespace metal
} // namespace c10
#else
#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
#endif // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
kronos-venv/lib/python3.12/site-packages/torch/include/c10/metal/igamma.h
+749
View File
@@ -0,0 +1,749 @@
#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
#pragma once
#include <c10/metal/utils.h>
#include <metal_math>
#include <metal_stdlib>
using namespace c10::metal;
using namespace metal;
namespace c10 {
namespace metal {
template <typename T>
inline float log_gamma(const T);
inline float expm1f(float a);
template <typename T>
float erfc(T x);
} // namespace metal
} // namespace c10
namespace {
template <typename T>
inline float lgamma(const T a) {
return log_gamma(a);
}
inline float expm1(float a) {
return expm1f(a);
}
// NOTE: The following code was ported directly from the CUDA implementation in
// `aten/src/ATen/native/cuda/IGammaKernel.cu`
/*
* This implementation of the regularized incomplete gamma functions and
* their helper functions are derived from the implementation of SciPy's
* gammainc, Cephes's igam and igamc, and Boost's Lanczos approximations.
* See NOTICE for the licenses.
*/
// regularized lower & upper incomplete gamma
template <typename scalar_t>
scalar_t ratevl(
scalar_t x,
const scalar_t num[],
int64_t M,
const scalar_t denom[],
int64_t N) {
// evaluating rational function, i.e., the ratio of two polynomials
// the coefficients for numerator are given by `num` while coeffs for
// denumerator are given by `denom`
using accscalar_t = opmath_t<scalar_t>;
int64_t i, dir;
accscalar_t y, num_ans, denom_ans;
accscalar_t absx = ::fabs(x);
thread const accscalar_t* p;
if (absx > 1) {
/* Evaluate as a polynomial in 1/x. */
dir = -1;
p = num + M;
y = 1 / x;
} else {
dir = 1;
p = num;
y = x;
}
/* Evaluate the numerator */
num_ans = *p;
p += dir;
for (i = 1; i <= M; i++) {
num_ans = num_ans * y + *p;
p += dir;
}
/* Evaluate the denominator */
if (absx > 1) {
p = denom + N;
} else {
p = denom;
}
denom_ans = *p;
p += dir;
for (i = 1; i <= N; i++) {
denom_ans = denom_ans * y + *p;
p += dir;
}
if (absx > 1) {
i = N - M;
return ::pow(x, static_cast<accscalar_t>(i)) * num_ans / denom_ans;
} else {
return num_ans / denom_ans;
}
}
template <typename scalar_t>
scalar_t lanczos_sum_expg_scaled(scalar_t x) {
// lanczos approximation
using accscalar_t = opmath_t<scalar_t>;
const accscalar_t lanczos_sum_expg_scaled_num[13] = {
0.006061842346248906525783753964555936883222,
0.5098416655656676188125178644804694509993,
19.51992788247617482847860966235652136208,
449.9445569063168119446858607650988409623,
6955.999602515376140356310115515198987526,
75999.29304014542649875303443598909137092,
601859.6171681098786670226533699352302507,
3481712.15498064590882071018964774556468,
14605578.08768506808414169982791359218571,
43338889.32467613834773723740590533316085,
86363131.28813859145546927288977868422342,
103794043.1163445451906271053616070238554,
56906521.91347156388090791033559122686859};
const accscalar_t lanczos_sum_expg_scaled_denom[13] = {
1.,
66.,
1925.,
32670.,
357423.,
2637558.,
13339535.,
45995730.,
105258076.,
150917976.,
120543840.,
39916800.,
0};
return ratevl(
static_cast<accscalar_t>(x),
lanczos_sum_expg_scaled_num,
sizeof(lanczos_sum_expg_scaled_num) /
sizeof(lanczos_sum_expg_scaled_num[0]) -
1,
lanczos_sum_expg_scaled_denom,
sizeof(lanczos_sum_expg_scaled_denom) /
sizeof(lanczos_sum_expg_scaled_denom[0]) -
1);
}
template <typename scalar_t>
scalar_t _igam_helper_fac(scalar_t a, scalar_t x) {
// compute x^a * exp(-a) / gamma(a)
// corrected from (15) and (16) in [igam2] by replacing exp(x - a) with
// exp(a - x).
using accscalar_t = opmath_t<scalar_t>;
accscalar_t ax, fac, res, num, numfac;
const accscalar_t MAXLOG = 88.72283905206835;
const accscalar_t EXP1 = 2.718281828459045;
const accscalar_t lanczos_g = 6.024680040776729583740234375;
if (::fabs(a - x) > 0.4 * ::fabs(a)) {
ax = a * ::log(x) - x - ::lgamma(a);
if (ax < -MAXLOG) {
return 0.0;
}
return ::exp(ax);
}
fac = a + lanczos_g - 0.5;
res = ::sqrt(fac / EXP1) / lanczos_sum_expg_scaled(a);
if ((a < 200) && (x < 200)) {
res *= ::exp(a - x) * ::pow(x / fac, a);
} else {
num = x - a - lanczos_g + 0.5;
numfac = num / fac;
res *= ::exp(a * (::log1p(numfac) - numfac) + x * (0.5 - lanczos_g) / fac);
}
return res;
}
template <typename scalar_t>
scalar_t _igam_helper_series(scalar_t a, scalar_t x) {
// Compute igam using DLMF 8.11.4. [igam1]
using accscalar_t = opmath_t<scalar_t>;
const accscalar_t MACHEP = 5.9604644775390625E-8;
const int MAXITER = 2000;
int i;
accscalar_t ans, ax, c, r;
ax = _igam_helper_fac(a, x);
if (ax == 0.0) {
return 0.0;
}
/* power series */
r = a;
c = 1.0;
ans = 1.0;
for (i = 0; i < MAXITER; i++) {
r += 1.0;
c *= x / r;
ans += c;
if (c <= MACHEP * ans) {
break;
}
}
return (ans * ax / a);
}
template <typename scalar_t>
scalar_t _igamc_helper_series(scalar_t a, scalar_t x) {
// Compute igamc using DLMF 8.7.3 [igam1]. This is related to the series in
// _igam_helper_series but extra care is taken to avoid cancellation.
using accscalar_t = opmath_t<scalar_t>;
int n;
accscalar_t fac = 1;
accscalar_t sum = 0;
accscalar_t term, logx;
const int MAXITER = 2000;
const accscalar_t MACHEP = 5.9604644775390625E-8;
for (n = 1; n < MAXITER; n++) {
fac *= -x / n;
term = fac / (a + n);
sum += term;
if (::fabs(term) <= MACHEP * ::fabs(sum)) {
break;
}
}
logx = ::log(x);
term = -::expm1(a * logx - ::lgamma(1 + a));
return term - ::exp(a * logx - ::lgamma(a)) * sum;
}
template <typename scalar_t>
scalar_t _igam_helper_asymptotic_series(scalar_t a, scalar_t x, bool igam) {
// Compute igam/igamc using DLMF 8.12.3/8.12.4 [igam1]
using accscalar_t = opmath_t<scalar_t>;
const accscalar_t d[25][25] = {
{-3.3333333333333333e-1, 8.3333333333333333e-2,
-1.4814814814814815e-2, 1.1574074074074074e-3,
3.527336860670194e-4, -1.7875514403292181e-4,
3.9192631785224378e-5, -2.1854485106799922e-6,
-1.85406221071516e-6, 8.296711340953086e-7,
-1.7665952736826079e-7, 6.7078535434014986e-9,
1.0261809784240308e-8, -4.3820360184533532e-9,
9.1476995822367902e-10, -2.551419399494625e-11,
-5.8307721325504251e-11, 2.4361948020667416e-11,
-5.0276692801141756e-12, 1.1004392031956135e-13,
3.3717632624009854e-13, -1.3923887224181621e-13,
2.8534893807047443e-14, -5.1391118342425726e-16,
-1.9752288294349443e-15},
{-1.8518518518518519e-3, -3.4722222222222222e-3, 2.6455026455026455e-3,
-9.9022633744855967e-4, 2.0576131687242798e-4, -4.0187757201646091e-7,
-1.8098550334489978e-5, 7.6491609160811101e-6, -1.6120900894563446e-6,
4.6471278028074343e-9, 1.378633446915721e-7, -5.752545603517705e-8,
1.1951628599778147e-8, -1.7543241719747648e-11, -1.0091543710600413e-9,
4.1627929918425826e-10, -8.5639070264929806e-11, 6.0672151016047586e-14,
7.1624989648114854e-12, -2.9331866437714371e-12, 5.9966963656836887e-13,
-2.1671786527323314e-16, -4.9783399723692616e-14, 2.0291628823713425e-14,
-4.13125571381061e-15},
{4.1335978835978836e-3, -2.6813271604938272e-3, 7.7160493827160494e-4,
2.0093878600823045e-6, -1.0736653226365161e-4, 5.2923448829120125e-5,
-1.2760635188618728e-5, 3.4235787340961381e-8, 1.3721957309062933e-6,
-6.298992138380055e-7, 1.4280614206064242e-7, -2.0477098421990866e-10,
-1.4092529910867521e-8, 6.228974084922022e-9, -1.3670488396617113e-9,
9.4283561590146782e-13, 1.2872252400089318e-10, -5.5645956134363321e-11,
1.1975935546366981e-11, -4.1689782251838635e-15, -1.0940640427884594e-12,
4.6622399463901357e-13, -9.905105763906906e-14, 1.8931876768373515e-17,
8.8592218725911273e-15},
{6.4943415637860082e-4, 2.2947209362139918e-4, -4.6918949439525571e-4,
2.6772063206283885e-4, -7.5618016718839764e-5, -2.3965051138672967e-7,
1.1082654115347302e-5, -5.6749528269915966e-6, 1.4230900732435884e-6,
-2.7861080291528142e-11, -1.6958404091930277e-7, 8.0994649053880824e-8,
-1.9111168485973654e-8, 2.3928620439808118e-12, 2.0620131815488798e-9,
-9.4604966618551322e-10, 2.1541049775774908e-10, -1.388823336813903e-14,
-2.1894761681963939e-11, 9.7909989511716851e-12, -2.1782191880180962e-12,
6.2088195734079014e-17, 2.126978363279737e-13, -9.3446887915174333e-14,
2.0453671226782849e-14},
{-8.618882909167117e-4, 7.8403922172006663e-4,
-2.9907248030319018e-4, -1.4638452578843418e-6,
6.6414982154651222e-5, -3.9683650471794347e-5,
1.1375726970678419e-5, 2.5074972262375328e-10,
-1.6954149536558306e-6, 8.9075075322053097e-7,
-2.2929348340008049e-7, 2.956794137544049e-11,
2.8865829742708784e-8, -1.4189739437803219e-8,
3.4463580499464897e-9, -2.3024517174528067e-13,
-3.9409233028046405e-10, 1.8602338968504502e-10,
-4.356323005056618e-11, 1.2786001016296231e-15,
4.6792750266579195e-12, -2.1492464706134829e-12,
4.9088156148096522e-13, -6.3385914848915603e-18,
-5.0453320690800944e-14},
{-3.3679855336635815e-4, -6.9728137583658578e-5, 2.7727532449593921e-4,
-1.9932570516188848e-4, 6.7977804779372078e-5, 1.419062920643967e-7,
-1.3594048189768693e-5, 8.0184702563342015e-6, -2.2914811765080952e-6,
-3.252473551298454e-10, 3.4652846491085265e-7, -1.8447187191171343e-7,
4.8240967037894181e-8, -1.7989466721743515e-14, -6.3061945000135234e-9,
3.1624176287745679e-9, -7.8409242536974293e-10, 5.1926791652540407e-15,
9.3589442423067836e-11, -4.5134262161632782e-11, 1.0799129993116827e-11,
-3.661886712685252e-17, -1.210902069055155e-12, 5.6807435849905643e-13,
-1.3249659916340829e-13},
{5.3130793646399222e-4, -5.9216643735369388e-4, 2.7087820967180448e-4,
7.9023532326603279e-7, -8.1539693675619688e-5, 5.6116827531062497e-5,
-1.8329116582843376e-5, -3.0796134506033048e-9, 3.4651553688036091e-6,
-2.0291327396058604e-6, 5.7887928631490037e-7, 2.338630673826657e-13,
-8.8286007463304835e-8, 4.7435958880408128e-8, -1.2545415020710382e-8,
8.6496488580102925e-14, 1.6846058979264063e-9, -8.5754928235775947e-10,
2.1598224929232125e-10, -7.6132305204761539e-16, -2.6639822008536144e-11,
1.3065700536611057e-11, -3.1799163902367977e-12, 4.7109761213674315e-18,
3.6902800842763467e-13},
{3.4436760689237767e-4, 5.1717909082605922e-5,
-3.3493161081142236e-4, 2.812695154763237e-4,
-1.0976582244684731e-4, -1.2741009095484485e-7,
2.7744451511563644e-5, -1.8263488805711333e-5,
5.7876949497350524e-6, 4.9387589339362704e-10,
-1.0595367014026043e-6, 6.1667143761104075e-7,
-1.7562973359060462e-7, -1.2974473287015439e-12,
2.695423606288966e-8, -1.4578352908731271e-8,
3.887645959386175e-9, -3.8810022510194121e-17,
-5.3279941738772867e-10, 2.7437977643314845e-10,
-6.9957960920705679e-11, 2.5899863874868481e-17,
8.8566890996696381e-12, -4.403168815871311e-12,
1.0865561947091654e-12},
{-6.5262391859530942e-4, 8.3949872067208728e-4, -4.3829709854172101e-4,
-6.969091458420552e-7, 1.6644846642067548e-4, -1.2783517679769219e-4,
4.6299532636913043e-5, 4.5579098679227077e-9, -1.0595271125805195e-5,
6.7833429048651666e-6, -2.1075476666258804e-6, -1.7213731432817145e-11,
3.7735877416110979e-7, -2.1867506700122867e-7, 6.2202288040189269e-8,
6.5977038267330006e-16, -9.5903864974256858e-9, 5.2132144922808078e-9,
-1.3991589583935709e-9, 5.382058999060575e-16, 1.9484714275467745e-10,
-1.0127287556389682e-10, 2.6077347197254926e-11, -5.0904186999932993e-18,
-3.3721464474854592e-12},
{-5.9676129019274625e-4, -7.2048954160200106e-5,
6.7823088376673284e-4, -6.4014752602627585e-4,
2.7750107634328704e-4, 1.8197008380465151e-7,
-8.4795071170685032e-5, 6.105192082501531e-5,
-2.1073920183404862e-5, -8.8585890141255994e-10,
4.5284535953805377e-6, -2.8427815022504408e-6,
8.7082341778646412e-7, 3.6886101871706965e-12,
-1.5344695190702061e-7, 8.862466778790695e-8,
-2.5184812301826817e-8, -1.0225912098215092e-14,
3.8969470758154777e-9, -2.1267304792235635e-9,
5.7370135528051385e-10, -1.887749850169741e-19,
-8.0931538694657866e-11, 4.2382723283449199e-11,
-1.1002224534207726e-11},
{1.3324454494800656e-3, -1.9144384985654775e-3, 1.1089369134596637e-3,
9.932404122642299e-7, -5.0874501293093199e-4, 4.2735056665392884e-4,
-1.6858853767910799e-4, -8.1301893922784998e-9, 4.5284402370562147e-5,
-3.127053674781734e-5, 1.044986828530338e-5, 4.8435226265680926e-11,
-2.1482565873456258e-6, 1.329369701097492e-6, -4.0295693092101029e-7,
-1.7567877666323291e-13, 7.0145043163668257e-8, -4.040787734999483e-8,
1.1474026743371963e-8, 3.9642746853563325e-18, -1.7804938269892714e-9,
9.7480262548731646e-10, -2.6405338676507616e-10, 5.794875163403742e-18,
3.7647749553543836e-11},
{1.579727660730835e-3, 1.6251626278391582e-4, -2.0633421035543276e-3,
2.1389686185689098e-3, -1.0108559391263003e-3, -3.9912705529919201e-7,
3.6235025084764691e-4, -2.8143901463712154e-4, 1.0449513336495887e-4,
2.1211418491830297e-9, -2.5779417251947842e-5, 1.7281818956040463e-5,
-5.6413773872904282e-6, -1.1024320105776174e-11, 1.1223224418895175e-6,
-6.8693396379526735e-7, 2.0653236975414887e-7, 4.6714772409838506e-14,
-3.5609886164949055e-8, 2.0470855345905963e-8, -5.8091738633283358e-9,
-1.332821287582869e-16, 9.0354604391335133e-10, -4.9598782517330834e-10,
1.3481607129399749e-10},
{-4.0725121195140166e-3, 6.4033628338080698e-3, -4.0410161081676618e-3,
-2.183732802866233e-6, 2.1740441801254639e-3, -1.9700440518418892e-3,
8.3595469747962458e-4, 1.9445447567109655e-8, -2.5779387120421696e-4,
1.9009987368139304e-4, -6.7696499937438965e-5, -1.4440629666426572e-10,
1.5712512518742269e-5, -1.0304008744776893e-5, 3.304517767401387e-6,
7.9829760242325709e-13, -6.4097794149313004e-7, 3.8894624761300056e-7,
-1.1618347644948869e-7, -2.816808630596451e-15, 1.9878012911297093e-8,
-1.1407719956357511e-8, 3.2355857064185555e-9, 4.1759468293455945e-20,
-5.0423112718105824e-10},
{-5.9475779383993003e-3, -5.4016476789260452e-4, 8.7910413550767898e-3,
-9.8576315587856125e-3, 5.0134695031021538e-3, 1.2807521786221875e-6,
-2.0626019342754683e-3, 1.7109128573523058e-3, -6.7695312714133799e-4,
-6.9011545676562133e-9, 1.8855128143995902e-4, -1.3395215663491969e-4,
4.6263183033528039e-5, 4.0034230613321351e-11, -1.0255652921494033e-5,
6.612086372797651e-6, -2.0913022027253008e-6, -2.0951775649603837e-13,
3.9756029041993247e-7, -2.3956211978815887e-7, 7.1182883382145864e-8,
8.925574873053455e-16, -1.2101547235064676e-8, 6.9350618248334386e-9,
-1.9661464453856102e-9},
{1.7402027787522711e-2, -2.9527880945699121e-2, 2.0045875571402799e-2,
7.0289515966903407e-6, -1.2375421071343148e-2, 1.1976293444235254e-2,
-5.4156038466518525e-3, -6.3290893396418616e-8, 1.8855118129005065e-3,
-1.473473274825001e-3, 5.5515810097708387e-4, 5.2406834412550662e-10,
-1.4357913535784836e-4, 9.9181293224943297e-5, -3.3460834749478311e-5,
-3.5755837291098993e-12, 7.1560851960630076e-6, -4.5516802628155526e-6,
1.4236576649271475e-6, 1.8803149082089664e-14, -2.6623403898929211e-7,
1.5950642189595716e-7, -4.7187514673841102e-8, -6.5107872958755177e-17,
7.9795091026746235e-9},
{3.0249124160905891e-2, 2.4817436002649977e-3, -4.9939134373457022e-2,
5.9915643009307869e-2, -3.2483207601623391e-2, -5.7212968652103441e-6,
1.5085251778569354e-2, -1.3261324005088445e-2, 5.5515262632426148e-3,
3.0263182257030016e-8, -1.7229548406756723e-3, 1.2893570099929637e-3,
-4.6845138348319876e-4, -1.830259937893045e-10, 1.1449739014822654e-4,
-7.7378565221244477e-5, 2.5625836246985201e-5, 1.0766165333192814e-12,
-5.3246809282422621e-6, 3.349634863064464e-6, -1.0381253128684018e-6,
-5.608909920621128e-15, 1.9150821930676591e-7, -1.1418365800203486e-7,
3.3654425209171788e-8},
{-9.9051020880159045e-2, 1.7954011706123486e-1, -1.2989606383463778e-1,
-3.1478872752284357e-5, 9.0510635276848131e-2, -9.2828824411184397e-2,
4.4412112839877808e-2, 2.7779236316835888e-7, -1.7229543805449697e-2,
1.4182925050891573e-2, -5.6214161633747336e-3, -2.39598509186381e-9,
1.6029634366079908e-3, -1.1606784674435773e-3, 4.1001337768153873e-4,
1.8365800754090661e-11, -9.5844256563655903e-5, 6.3643062337764708e-5,
-2.076250624489065e-5, -1.1806020912804483e-13, 4.2131808239120649e-6,
-2.6262241337012467e-6, 8.0770620494930662e-7, 6.0125912123632725e-16,
-1.4729737374018841e-7},
{-1.9994542198219728e-1, -1.5056113040026424e-2, 3.6470239469348489e-1,
-4.6435192311733545e-1, 2.6640934719197893e-1, 3.4038266027147191e-5,
-1.3784338709329624e-1, 1.276467178337056e-1, -5.6213828755200985e-2,
-1.753150885483011e-7, 1.9235592956768113e-2, -1.5088821281095315e-2,
5.7401854451350123e-3, 1.0622382710310225e-9, -1.5335082692563998e-3,
1.0819320643228214e-3, -3.7372510193945659e-4, -6.6170909729031985e-12,
8.4263617380909628e-5, -5.5150706827483479e-5, 1.7769536448348069e-5,
3.8827923210205533e-14, -3.53513697488768e-6, 2.1865832130045269e-6,
-6.6812849447625594e-7},
{7.2438608504029431e-1, -1.3918010932653375, 1.0654143352413968,
1.876173868950258e-4, -8.2705501176152696e-1, 8.9352433347828414e-1,
-4.4971003995291339e-1, -1.6107401567546652e-6, 1.9235590165271091e-1,
-1.6597702160042609e-1, 6.8882222681814333e-2, 1.3910091724608687e-8,
-2.146911561508663e-2, 1.6228980898865892e-2, -5.9796016172584256e-3,
-1.1287469112826745e-10, 1.5167451119784857e-3, -1.0478634293553899e-3,
3.5539072889126421e-4, 8.1704322111801517e-13, -7.7773013442452395e-5,
5.0291413897007722e-5, -1.6035083867000518e-5, 1.2469354315487605e-14,
3.1369106244517615e-6},
{1.6668949727276811, 1.165462765994632e-1, -3.3288393225018906,
4.4692325482864037, -2.6977693045875807, -2.600667859891061e-4,
1.5389017615694539, -1.4937962361134612, 6.8881964633233148e-1,
1.3077482004552385e-6, -2.5762963325596288e-1, 2.1097676102125449e-1,
-8.3714408359219882e-2, -7.7920428881354753e-9, 2.4267923064833599e-2,
-1.7813678334552311e-2, 6.3970330388900056e-3, 4.9430807090480523e-11,
-1.5554602758465635e-3, 1.0561196919903214e-3, -3.5277184460472902e-4,
9.3002334645022459e-14, 7.5285855026557172e-5, -4.8186515569156351e-5,
1.5227271505597605e-5},
{-6.6188298861372935, 1.3397985455142589e+1, -1.0789350606845146e+1,
-1.4352254537875018e-3, 9.2333694596189809, -1.0456552819547769e+1,
5.5105526029033471, 1.2024439690716742e-5, -2.5762961164755816,
2.3207442745387179, -1.0045728797216284, -1.0207833290021914e-7,
3.3975092171169466e-1, -2.6720517450757468e-1, 1.0235252851562706e-1,
8.4329730484871625e-10, -2.7998284958442595e-2, 2.0066274144976813e-2,
-7.0554368915086242e-3, 1.9402238183698188e-12, 1.6562888105449611e-3,
-1.1082898580743683e-3, 3.654545161310169e-4, -5.1290032026971794e-11,
-7.6340103696869031e-5},
{-1.7112706061976095e+1, -1.1208044642899116, 3.7131966511885444e+1,
-5.2298271025348962e+1, 3.3058589696624618e+1, 2.4791298976200222e-3,
-2.061089403411526e+1, 2.088672775145582e+1, -1.0045703956517752e+1,
-1.2238783449063012e-5, 4.0770134274221141, -3.473667358470195,
1.4329352617312006, 7.1359914411879712e-8, -4.4797257159115612e-1,
3.4112666080644461e-1, -1.2699786326594923e-1, -2.8953677269081528e-10,
3.3125776278259863e-2, -2.3274087021036101e-2, 8.0399993503648882e-3,
-1.177805216235265e-9, -1.8321624891071668e-3, 1.2108282933588665e-3,
-3.9479941246822517e-4},
{7.389033153567425e+1, -1.5680141270402273e+2, 1.322177542759164e+2,
1.3692876877324546e-2, -1.2366496885920151e+2, 1.4620689391062729e+2,
-8.0365587724865346e+1, -1.1259851148881298e-4, 4.0770132196179938e+1,
-3.8210340013273034e+1, 1.719522294277362e+1, 9.3519707955168356e-7,
-6.2716159907747034, 5.1168999071852637, -2.0319658112299095,
-4.9507215582761543e-9, 5.9626397294332597e-1, -4.4220765337238094e-1,
1.6079998700166273e-1, -2.4733786203223402e-8, -4.0307574759979762e-2,
2.7849050747097869e-2, -9.4751858992054221e-3, 6.419922235909132e-6,
2.1250180774699461e-3},
{2.1216837098382522e+2, 1.3107863022633868e+1, -4.9698285932871748e+2,
7.3121595266969204e+2, -4.8213821720890847e+2, -2.8817248692894889e-2,
3.2616720302947102e+2, -3.4389340280087117e+2, 1.7195193870816232e+2,
1.4038077378096158e-4, -7.52594195897599e+1, 6.651969984520934e+1,
-2.8447519748152462e+1, -7.613702615875391e-7, 9.5402237105304373,
-7.5175301113311376, 2.8943997568871961, -4.6612194999538201e-7,
-8.0615149598794088e-1, 5.8483006570631029e-1, -2.0845408972964956e-1,
1.4765818959305817e-4, 5.1000433863753019e-2, -3.3066252141883665e-2,
1.5109265210467774e-2},
{-9.8959643098322368e+2, 2.1925555360905233e+3, -1.9283586782723356e+3,
-1.5925738122215253e-1, 1.9569985945919857e+3, -2.4072514765081556e+3,
1.3756149959336496e+3, 1.2920735237496668e-3, -7.525941715948055e+2,
7.3171668742208716e+2, -3.4137023466220065e+2, -9.9857390260608043e-6,
1.3356313181291573e+2, -1.1276295161252794e+2, 4.6310396098204458e+1,
-7.9237387133614756e-6, -1.4510726927018646e+1, 1.1111771248100563e+1,
-4.1690817945270892, 3.1008219800117808e-3, 1.1220095449981468,
-7.6052379926149916e-1, 3.6262236505085254e-1, 2.216867741940747e-1,
4.8683443692930507e-1}};
int k, n, sgn;
int maxpow = 0;
const accscalar_t MACHEP = 5.9604644775390625E-8;
accscalar_t lambda = x / a;
accscalar_t sigma = (x - a) / a;
accscalar_t eta, res, ck, ckterm, term, absterm;
accscalar_t absoldterm = INFINITY;
accscalar_t etapow[25] = {1};
accscalar_t sum = 0;
accscalar_t afac = 1;
if (igam) {
sgn = -1;
} else {
sgn = 1;
}
if (lambda > 1) {
eta = ::sqrt(-2 * (::log1p(sigma) - sigma));
} else if (lambda < 1) {
eta = -::sqrt(-2 * (::log1p(sigma) - sigma));
} else {
eta = 0;
}
res = 0.5 * ::erfc(sgn * eta * ::sqrt(a / 2));
for (k = 0; k < 25; k++) {
ck = d[k][0];
for (n = 1; n < 25; n++) {
if (n > maxpow) {
etapow[n] = eta * etapow[n - 1];
maxpow += 1;
}
ckterm = d[k][n] * etapow[n];
ck += ckterm;
if (::fabs(ckterm) < MACHEP * ::fabs(ck)) {
break;
}
}
term = ck * afac;
absterm = ::fabs(term);
if (absterm > absoldterm) {
break;
}
sum += term;
if (absterm < MACHEP * ::fabs(sum)) {
break;
}
absoldterm = absterm;
afac /= a;
}
res += sgn * ::exp(-0.5 * a * eta * eta) * sum / ::sqrt(2 * 3.1415926535 * a);
return res;
}
template <typename scalar_t>
scalar_t _igamc_helper_continued_fraction(scalar_t a, scalar_t x) {
// Compute igamc using DLMF 8.9.2. [igam1]
using accscalar_t = opmath_t<scalar_t>;
int i;
accscalar_t ans, ax, c, yc, r, t, y, z;
accscalar_t pk, pkm1, pkm2, qk, qkm1, qkm2;
const int MAXITER = 2000;
const accscalar_t MACHEP = 5.9604644775390625E-8;
const accscalar_t BIG = 16777216.;
const accscalar_t BIGINV = 5.9604644775390625E-8;
ax = _igam_helper_fac(a, x);
if (ax == 0.0) {
return 0.0;
}
/* continued fraction */
y = 1.0 - a;
z = x + y + 1.0;
c = 0.0;
pkm2 = 1.0;
qkm2 = x;
pkm1 = x + 1.0;
qkm1 = z * x;
ans = pkm1 / qkm1;
for (i = 0; i < MAXITER; i++) {
c += 1.0;
y += 1.0;
z += 2.0;
yc = y * c;
pk = pkm1 * z - pkm2 * yc;
qk = qkm1 * z - qkm2 * yc;
if (qk != 0) {
r = pk / qk;
t = ::fabs((ans - r) / r);
ans = r;
} else {
t = 1.0;
}
pkm2 = pkm1;
pkm1 = pk;
qkm2 = qkm1;
qkm1 = qk;
if (::fabs(pk) > BIG) {
pkm2 *= BIGINV;
pkm1 *= BIGINV;
qkm2 *= BIGINV;
qkm1 *= BIGINV;
}
if (t <= MACHEP) {
break;
}
}
return ans * ax;
}
template <typename scalar_t>
scalar_t calc_igammac(scalar_t a, scalar_t x) {
/* the calculation of the regularized upper incomplete gamma function
* is done differently based on the values of a and x:
* - if x and/or a is at the boundary of defined region, then assign the
* result at the boundary
* - if a is large and a ~ x, then using Uniform Asymptotic Expansions for
* Large Parameter (see DLMF 8.12.4 [igam1])
* - if x > 1.1 and x < a, using the subtraction from the regularized lower
* incomplete gamma
* - otherwise, calculate the series from [igam2] eq (5)
*/
using accscalar_t = opmath_t<scalar_t>;
accscalar_t absxma_a;
const accscalar_t SMALL = 20.0;
const accscalar_t LARGE = 200.0;
const accscalar_t SMALLRATIO = 0.3;
const accscalar_t LARGERATIO = 4.5;
if ((x < 0) || (a < 0)) {
// out of defined-region of the function
return NAN;
} else if (a == 0) {
if (x > 0) {
return 0.0;
} else {
return NAN;
}
} else if (x == 0) {
return 1.0;
} else if (isinf(a)) {
if (isinf(x)) {
return NAN;
}
return 1.0;
} else if (isinf(x)) {
return 0.0;
}
absxma_a = ::fabs(x - a) / a;
if ((a > SMALL) && (a < LARGE) && (absxma_a < SMALLRATIO)) {
return _igam_helper_asymptotic_series(a, x, 0);
} else if ((a > LARGE) && (absxma_a < LARGERATIO / ::sqrt(a))) {
return _igam_helper_asymptotic_series(a, x, 0);
}
if (x > 1.1) {
if (x < a) {
return 1.0 - _igam_helper_series(a, x);
} else {
return _igamc_helper_continued_fraction(a, x);
}
} else if (x <= 0.5) {
if (-0.4 / ::log(x) < a) {
return 1.0 - _igam_helper_series(a, x);
} else {
return _igamc_helper_series(a, x);
}
} else {
if (x * 1.1 < a) {
return 1.0 - _igam_helper_series(a, x);
} else {
return _igamc_helper_series(a, x);
}
}
}
template <typename scalar_t>
scalar_t calc_igamma(scalar_t a, scalar_t x) {
/* the calculation of the regularized lower incomplete gamma function
* is done differently based on the values of a and x:
* - if x and/or a is at the boundary of defined region, then assign the
* result at the boundary
* - if a is large and a ~ x, then using Uniform Asymptotic Expansions for
* Large Parameter (see DLMF 8.12.3 [igam1])
* - if x > 1 and x > a, using the subtraction from the regularized upper
* incomplete gamma
* - otherwise, calculate the series from [igam2] eq (4)
*/
using accscalar_t = opmath_t<scalar_t>;
accscalar_t absxma_a;
const accscalar_t SMALL = 20.0;
const accscalar_t LARGE = 200.0;
const accscalar_t SMALLRATIO = 0.3;
const accscalar_t LARGERATIO = 4.5;
// boundary values following SciPy
if ((x < 0) || (a < 0)) {
// out of defined-region of the function
return NAN;
} else if (a == 0) {
if (x > 0) {
return 1.0;
} else {
return NAN;
}
} else if (x == 0) {
return 0.0; // zero integration limit
} else if (isinf(a)) {
if (isinf(x)) {
return NAN;
}
return 0.0;
} else if (isinf(x)) {
return 1.0;
}
/* Asymptotic regime where a ~ x. */
absxma_a = ::fabs(x - a) / a;
if ((a > SMALL) && (a < LARGE) && (absxma_a < SMALLRATIO)) {
return _igam_helper_asymptotic_series(a, x, 1);
} else if ((a > LARGE) && (absxma_a < LARGERATIO / ::sqrt(a))) {
return _igam_helper_asymptotic_series(a, x, 1);
}
if ((x > 1.0) && (x > a)) {
return 1.0 - calc_igammac(a, x);
}
return _igam_helper_series(a, x);
}
} // namespace
// end of regularized lower & upper incomplete gamma
namespace c10 {
namespace metal {
template <typename T>
inline T igamma(T a, T b) {
return calc_igamma(a, b);
}
template <typename T>
inline T igammac(T a, T b) {
return calc_igammac(a, b);
}
} // namespace metal
} // namespace c10
#else
#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
#endif // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
kronos-venv/lib/python3.12/site-packages/torch/include/c10/metal/indexing.h
+1085
View File
File diff suppressed because it is too large Load Diff
kronos-venv/lib/python3.12/site-packages/torch/include/c10/metal/random.h
+83
View File
@@ -0,0 +1,83 @@
#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
// Philox Counter based RNG implementation for Metal
// Borrowed from aten/src/ATen/core/PhiloxRNGEngine.h
// Which in turn borrowed from
// http://www.thesalmons.org/john/random123/papers/random123sc11.pdf
#pragma once
#include <metal_stdlib>
namespace c10 {
namespace metal {
namespace detail {
constexpr float uint32_to_uniform_float(uint32_t value) {
// maximum value such that `MAX_INT * scale < 1.0` (with float rounding)
constexpr float scale = 4.6566127342e-10;
return static_cast<float>(value & 0x7FFFFFFF) * scale;
}
inline uint2 splitlong(ulong v) {
return uint2(v >> 32, v & 0xffffffff);
}
} // namespace detail
namespace philox4 {
uint2 mulhilo(uint a, uint b) {
auto rc = static_cast<ulong>(a) * b;
return detail::splitlong(rc);
}
uint4 single_round(uint4 ctr, uint2 key) {
constexpr uint kPhiloxSA = 0xD2511F53;
constexpr uint kPhiloxSB = 0xCD9E8D57;
auto rc0 = mulhilo(kPhiloxSA, ctr.x);
auto rc1 = mulhilo(kPhiloxSB, ctr.z);
return uint4(rc1.x ^ ctr.y ^ key.x, rc1.y, rc0.x ^ ctr.w ^ key.y, rc0.y);
}
uint4 multiple_rounds(uint4 ctr, uint2 key, uint rounds) {
constexpr uint2 kPhilox10 = {0x9E3779B9, 0xBB67AE85};
for (uint round = 0; round < rounds - 1; ++round) {
ctr = single_round(ctr, key);
key += kPhilox10;
}
return ctr;
}
uint4 rand(long seed, long index) {
uint4 ctr = 0;
ctr.zw = detail::splitlong(index);
return multiple_rounds(ctr, detail::splitlong(seed), 10);
}
} // namespace philox4
float randn(long seed, long index) {
auto value = philox4::rand(seed, index);
float u1 = 1.0 - detail::uint32_to_uniform_float(value.x);
float u2 = 1.0 - detail::uint32_to_uniform_float(value.y);
return ::metal::sqrt(-2.0 * ::metal::log(u1)) *
::metal::cos(2.0 * M_PI_F * u2);
}
float rand(long seed, long index) {
auto value = philox4::rand(seed, index);
return detail::uint32_to_uniform_float(value.x);
}
long randint64(long seed, long index, long low, long high) {
auto range = high - low;
auto value = philox4::rand(seed, index);
// TODO: Implement better algorithm for large ranges
return low +
static_cast<long>(detail::uint32_to_uniform_float(value.x) * range);
}
} // namespace metal
} // namespace c10
#else
#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
#endif // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
kronos-venv/lib/python3.12/site-packages/torch/include/c10/metal/reduction_utils.h
+364
View File
@@ -0,0 +1,364 @@
#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
#pragma once
#include <c10/metal/utils.h>
#include <metal_compute>
namespace c10 {
namespace metal {
namespace detail {
template <typename T>
struct simd_type {
using t = T;
};
// Helper that allows one to run simd ops over bfl16 by upcasting them to fp32
template <typename T>
using simd_type_t = typename simd_type<T>::t;
template <>
struct simd_type<bfloat> {
using t = float;
};
} // namespace detail
template <typename T>
inline ::metal::enable_if_t<!::metal::is_same_v<T, long>, T> simd_sum(T val) {
return T(::metal::simd_sum(detail::simd_type_t<T>(val)));
}
template <typename T>
inline ::metal::enable_if_t<!::metal::is_same_v<T, long>, T> simd_prod(T val) {
return T(::metal::simd_product(detail::simd_type_t<T>(val)));
}
// Extend simd_broadcast to 64-bit integral types using int2 trick
template <
typename T,
::metal::enable_if_t<::metal::is_integral_v<T> && sizeof(T) == 8, bool> =
true>
inline T simd_broadcast(T val, ushort lane_id) {
return as_type<T>(::metal::simd_broadcast(as_type<int2>(val), lane_id));
}
template <
typename T,
::metal::enable_if_t<!::metal::is_integral_v<T> || sizeof(T) != 8, bool> =
true>
inline T simd_broadcast(T val, ushort lane_id) {
return ::metal::simd_broadcast(val, lane_id);
}
// Floating simd_min/max with nan propagation
template <
typename T,
::metal::enable_if_t<::metal::is_floating_point_v<T>, bool> = true>
inline T simd_max(T val) {
if (::metal::simd_any(::metal::isnan(val))) {
return ::metal::numeric_limits<T>::quiet_NaN();
}
return T(::metal::simd_max(detail::simd_type_t<T>(val)));
}
template <
typename T,
::metal::enable_if_t<::metal::is_floating_point_v<T>, bool> = true>
inline T simd_min(T val) {
if (::metal::simd_any(::metal::isnan(val))) {
return ::metal::numeric_limits<T>::quiet_NaN();
}
return T(::metal::simd_min(detail::simd_type_t<T>(val)));
}
template <
typename T,
::metal::enable_if_t<::metal::is_integral_v<T> && sizeof(T) != 8, bool> =
true>
inline T simd_max(T val) {
return ::metal::simd_max(val);
}
template <
typename T,
::metal::enable_if_t<::metal::is_integral_v<T> && sizeof(T) != 8, bool> =
true>
inline T simd_min(T val) {
return ::metal::simd_min(val);
}
// Metal does not support SIMD reductions over 64-bit types, but it could be
// implement using simd_shuffle_down, that yields result in log2(simdgroup_size)
// iterations Use fill variant, as shuffle down returns garbage if inactive
// thread is referenced (on M1/M2, works fine on M4) and broadcast result to all
// threads in the end. Implementation heavily borrows from
// https://github.com/ml-explore/mlx/blob/86389bf9707f46101af45d90510e8e97c8a90b93/mlx/backend/metal/kernels/reduction/ops.h#L16
template <typename T>
inline ::metal::enable_if_t<::metal::is_same_v<T, long>, T> simd_sum(T val) {
for (ushort i = simdgroup_size / 2; i > 0; i /= 2) {
val += as_type<T>(
::metal::simd_shuffle_and_fill_down(as_type<int2>(val), int2(0), i));
}
return simd_broadcast(val, 0);
}
template <typename T>
inline ::metal::enable_if_t<::metal::is_same_v<T, long>, T> simd_prod(T val) {
for (ushort i = simdgroup_size / 2; i > 0; i /= 2) {
val *= as_type<T>(
::metal::simd_shuffle_and_fill_down(as_type<int2>(val), int2(0), i));
}
return simd_broadcast(val, 0);
}
template <typename T>
inline ::metal::enable_if_t<::metal::is_same_v<T, long>, T> simd_max(T val) {
for (ushort i = simdgroup_size / 2; i > 0; i /= 2) {
val = ::metal::max(
val,
as_type<T>(::metal::simd_shuffle_and_fill_down(
as_type<int2>(val), int2(0), i)));
}
return simd_broadcast(val, 0);
}
template <typename T>
inline ::metal::enable_if_t<::metal::is_same_v<T, long>, T> simd_min(T val) {
for (ushort i = simdgroup_size / 2; i > 0; i /= 2) {
val = ::metal::min(
val,
as_type<T>(::metal::simd_shuffle_and_fill_down(
as_type<int2>(val), int2(0), i)));
}
return simd_broadcast(val, 0);
}
// argmin/argmax helpers using simd_ballot
template <
typename T,
::metal::enable_if_t<::metal::is_integral_v<T>, bool> = true>
inline ::c10::metal::pair<T, ushort> simd_argmin(T val) {
const auto rc = simd_min(val);
const auto vote = ::metal::simd_ballot(val == rc);
return {rc, static_cast<ushort>(::metal::ctz(static_cast<ulong>(vote)))};
}
template <
typename T,
::metal::enable_if_t<::metal::is_floating_point_v<T>, bool> = true>
inline ::c10::metal::pair<T, ushort> simd_argmin(T val) {
const auto rc = simd_min(val);
const auto vote = ::metal::simd_ballot(val == rc || ::metal::isnan(val));
return {rc, static_cast<ushort>(::metal::ctz(static_cast<ulong>(vote)))};
}
template <
typename T,
::metal::enable_if_t<::metal::is_integral_v<T>, bool> = true>
inline ::c10::metal::pair<T, ushort> simd_argmax(T val) {
const auto rc = simd_max(val);
const auto vote = ::metal::simd_ballot(val == rc);
return {rc, static_cast<ushort>(::metal::ctz(static_cast<ulong>(vote)))};
}
template <
typename T,
::metal::enable_if_t<::metal::is_floating_point_v<T>, bool> = true>
inline ::c10::metal::pair<T, ushort> simd_argmax(T val) {
const auto rc = simd_max(val);
const auto vote = ::metal::simd_ballot(val == rc || ::metal::isnan(val));
return {rc, static_cast<ushort>(::metal::ctz(static_cast<ulong>(vote)))};
}
template <typename ARG_T, typename IDX_T>
inline c10::metal::pair<ARG_T, IDX_T> simd_argmin(ARG_T val, IDX_T idx_val) {
auto rc = simd_argmin(val);
return {rc.first, simd_broadcast(idx_val, rc.second)};
}
template <typename ARG_T, typename IDX_T>
inline c10::metal::pair<ARG_T, IDX_T> simd_argmax(ARG_T val, IDX_T idx_val) {
auto rc = simd_argmax(val);
return {rc.first, simd_broadcast(idx_val, rc.second)};
}
// Below algorithms are written with hardcoded assumption that simdgroup is 32
// and threadgroup_max is 1024, i.e. reduction can be done in two stages max
template <typename T>
opmath_t<T> threadgroup_sum(
threadgroup opmath_t<T>* data,
T val,
unsigned idx,
unsigned size) {
auto rc = simd_sum(static_cast<opmath_t<T>>(val));
if (idx % simdgroup_size == 0) {
data[idx / simdgroup_size] = rc;
}
if (size > simdgroup_size) {
::metal::threadgroup_barrier(::metal::mem_flags::mem_threadgroup);
if (idx < ((size + simdgroup_size - 1) / simdgroup_size)) {
auto rc1 = simd_sum(data[idx]);
if (idx == 0) {
data[0] = rc1;
}
}
}
::metal::threadgroup_barrier(::metal::mem_flags::mem_threadgroup);
return data[0];
}
template <typename T>
opmath_t<T> threadgroup_prod(
threadgroup opmath_t<T>* data,
T val,
unsigned idx,
unsigned size) {
auto rc = simd_prod(static_cast<opmath_t<T>>(val));
if (idx % simdgroup_size == 0) {
data[idx / simdgroup_size] = rc;
}
if (size > simdgroup_size) {
::metal::threadgroup_barrier(::metal::mem_flags::mem_threadgroup);
if (idx < ((size + simdgroup_size - 1) / simdgroup_size)) {
auto rc1 = simd_prod(data[idx]);
if (idx == 0) {
data[0] = rc1;
}
}
}
::metal::threadgroup_barrier(::metal::mem_flags::mem_threadgroup);
return data[0];
}
template <typename T>
T threadgroup_max(threadgroup T* data, T val, unsigned idx, unsigned size) {
auto rc = simd_max(val);
if (idx % simdgroup_size == 0) {
data[idx / simdgroup_size] = rc;
}
if (size > simdgroup_size) {
::metal::threadgroup_barrier(::metal::mem_flags::mem_threadgroup);
if (idx < ((size + simdgroup_size - 1) / simdgroup_size)) {
auto rc1 = simd_max(data[idx]);
if (idx == 0) {
data[0] = rc1;
}
}
}
::metal::threadgroup_barrier(::metal::mem_flags::mem_threadgroup);
return data[0];
}
template <typename T>
T threadgroup_min(threadgroup T* data, T val, unsigned idx, unsigned size) {
auto rc = simd_min(val);
if (idx % simdgroup_size == 0) {
data[idx / simdgroup_size] = rc;
}
if (size > simdgroup_size) {
::metal::threadgroup_barrier(::metal::mem_flags::mem_threadgroup);
if (idx < ((size + simdgroup_size - 1) / simdgroup_size)) {
auto rc1 = simd_min(data[idx]);
if (idx == 0) {
data[0] = rc1;
}
}
}
::metal::threadgroup_barrier(::metal::mem_flags::mem_threadgroup);
return data[0];
}
template <typename T>
float3 threadgroup_welford_reduce(threadgroup T* data, unsigned size) {
::metal::threadgroup_barrier(::metal::mem_flags::mem_threadgroup);
float m = data[0];
float m2 = 0;
for (unsigned idx = 1; idx < size; ++idx) {
float delta = data[idx] - m;
m += delta / (idx + 1);
m2 += delta * (data[idx] - m);
}
return float3(m, m2, size);
}
// Each vec3type is tuple of mean, m2 and weight
template <typename T>
float3 welford_combine(T a, T b) {
float delta = b.x - a.x;
float new_weight = a.z + b.z;
auto w2_over_w = new_weight != 0 ? b.z / new_weight : 0.0;
return float3(
a.x + delta * w2_over_w,
a.y + b.y + delta * delta * a.z * w2_over_w,
new_weight);
}
template <typename T>
float3 threadgroup_welford_combine(threadgroup T* data, unsigned size) {
::metal::threadgroup_barrier(::metal::mem_flags::mem_threadgroup);
float3 rc = data[0];
for (unsigned idx = 1; idx < size; ++idx) {
rc = welford_combine(rc, data[idx]);
}
return rc;
}
template <typename ARG_T, typename IDX_T>
IDX_T threadgroup_argmax(
threadgroup ARG_T* arg_data,
threadgroup IDX_T* idx_data,
ARG_T val,
IDX_T idx_val,
unsigned idx,
unsigned size) {
auto rc = simd_argmax(val, idx_val);
if (size <= simdgroup_size) {
return rc.second;
}
if (idx % simdgroup_size == 0) {
arg_data[idx / simdgroup_size] = rc.first;
idx_data[idx / simdgroup_size] = rc.second;
}
::metal::threadgroup_barrier(::metal::mem_flags::mem_threadgroup);
if (idx < ((size + simdgroup_size - 1) / simdgroup_size)) {
auto rc1 = simd_argmax(arg_data[idx], idx_data[idx]);
if (idx == 0) {
idx_data[0] = rc1.second;
}
}
::metal::threadgroup_barrier(::metal::mem_flags::mem_threadgroup);
return idx_data[0];
}
template <typename ARG_T, typename IDX_T>
IDX_T threadgroup_argmin(
threadgroup ARG_T* arg_data,
threadgroup IDX_T* idx_data,
ARG_T val,
IDX_T idx_val,
unsigned idx,
unsigned size) {
auto rc = simd_argmin(val, idx_val);
if (size <= simdgroup_size) {
return rc.second;
}
if (idx % simdgroup_size == 0) {
arg_data[idx / simdgroup_size] = rc.first;
idx_data[idx / simdgroup_size] = rc.second;
}
::metal::threadgroup_barrier(::metal::mem_flags::mem_threadgroup);
if (idx < ((size + simdgroup_size - 1) / simdgroup_size)) {
auto rc1 = simd_argmin(arg_data[idx], idx_data[idx]);
if (idx == 0) {
idx_data[0] = rc1.second;
}
}
::metal::threadgroup_barrier(::metal::mem_flags::mem_threadgroup);
return idx_data[0];
}
} // namespace metal
} // namespace c10
#else
#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
#endif // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
kronos-venv/lib/python3.12/site-packages/torch/include/c10/metal/special_math.h
+3086
View File
File diff suppressed because it is too large Load Diff
kronos-venv/lib/python3.12/site-packages/torch/include/c10/metal/utils.h
+528
View File
@@ -0,0 +1,528 @@
#if !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)
// Metal helper functions
#pragma once
#include <c10/metal/common.h>
#include <metal_stdlib>
namespace c10 {
namespace metal {
namespace detail {
template <typename T>
struct vectypes {};
template <>
struct vectypes<float> {
using type4 = float4;
using type3 = float3;
using type2 = float2;
};
template <>
struct vectypes<half> {
using type4 = half4;
using type3 = half3;
using type2 = half2;
};
template <>
struct vectypes<bfloat> {
using type4 = bfloat4;
using type3 = bfloat3;
using type2 = bfloat2;
};
template <>
struct vectypes<short> {
using type4 = short4;
using type3 = short3;
using type2 = short2;
};
template <>
struct vectypes<int> {
using type4 = int4;
using type3 = int3;
using type2 = int2;
};
template <>
struct vectypes<long> {
using type4 = short4;
using type3 = short3;
using type2 = short2;
};
template <typename T>
struct OpMathType {
using type = T;
};
template <>
struct OpMathType<half> {
using type = float;
};
template <>
struct OpMathType<short> {
using type = int;
};
template <>
struct OpMathType<char> {
using type = int;
};
template <>
struct OpMathType<uchar> {
using type = int;
};
template <>
struct OpMathType<bfloat> {
using type = float;
};
// Type promotion structure for higher precision accumulation
template <typename T>
struct AccumulationType {
using type = T;
};
// Specialization for half - promote to float for accumulation
template <>
struct AccumulationType<half> {
using type = float;
};
// Specialization for bfloat - promote to float for accumulation
template <>
struct AccumulationType<bfloat> {
using type = float;
};
} // namespace detail
template <typename T>
::metal::enable_if_t<::metal::is_floating_point_v<T>, T> max(T a, T b) {
return ::metal::isunordered(a, b) ? NAN : ::metal::max(a, b);
}
template <typename T, typename U>
::metal::enable_if_t<::metal::is_integral_v<T>&& ::metal::is_integral_v<U>, T>
max(T a, U b) {
return ::metal::max(a, static_cast<T>(b));
}
template <typename T>
::metal::enable_if_t<::metal::is_floating_point_v<T>, T> min(T a, T b) {
return ::metal::isunordered(a, b) ? NAN : ::metal::min(a, b);
}
template <typename T, typename U>
::metal::enable_if_t<::metal::is_integral_v<T>&& ::metal::is_integral_v<U>, T>
min(T a, U b) {
return ::metal::min(a, static_cast<T>(b));
}
template <>
inline bfloat min(bfloat a, bfloat b) {
return bfloat(
::metal::isunordered(a, b) ? NAN : ::metal::min(float(a), float(b)));
}
template <>
inline bfloat max(bfloat a, bfloat b) {
return bfloat(
::metal::isunordered(a, b) ? NAN : ::metal::max(float(a), float(b)));
}
template <typename T>
using vec2type_t = typename detail::vectypes<T>::type2;
template <typename T>
using vec4type_t = typename detail::vectypes<T>::type4;
template <typename T>
using opmath_t = typename detail::OpMathType<T>::type;
template <typename T>
using accum_t = typename detail::AccumulationType<T>::type;
// TODO: Move it to type_traits header may be
template <typename F, typename... Args>
using result_of = decltype(::metal::declval<F>()(::metal::declval<Args>()...));
template <typename T>
constexpr constant bool is_complex_v =
::metal::is_same_v<T, float2> || ::metal::is_same_v<T, half2>;
template <typename T>
constexpr constant bool is_scalar_floating_point_v =
::metal::is_floating_point_v<T> && ::metal::is_scalar_v<T>;
template <typename T>
constexpr constant bool is_scalar_integral_v =
::metal::is_integral_v<T> && ::metal::is_scalar_v<T>;
template <typename U, typename V>
using common_dtype = decltype(U(0) + V(0));
// floor_divide
template <
typename T,
typename U,
::metal::enable_if_t<
is_scalar_integral_v<T> && is_scalar_integral_v<U>,
bool> = true>
inline common_dtype<T, U> floor_divide(T x, U y) {
const auto quot = x / y;
return (x < 0) == (y < 0) ? quot : (x % y != 0) ? quot - 1 : quot;
}
template <
typename T,
typename U,
::metal::enable_if_t<
is_scalar_floating_point_v<T> && is_scalar_floating_point_v<U>,
bool> = true>
inline common_dtype<T, U> floor_divide(T x, U y) {
return ::metal::floor(x / y);
}
// Workaround for Metal compiler bug: the compiler produces wrong results
// when optimizing fused (x / A) % B expressions for integral types.
template <
typename T,
typename U,
::metal::enable_if_t<
is_scalar_integral_v<T> && is_scalar_integral_v<U>,
bool> = true>
inline common_dtype<T, U> safe_mod(volatile T x, U y) {
return x % y;
}
// fmod
template <
typename T,
typename U,
::metal::enable_if_t<
is_scalar_integral_v<T> && is_scalar_integral_v<U>,
bool> = true>
inline common_dtype<T, U> fmod(T x, U y) {
return x % y;
}
template <
typename T,
typename U,
::metal::enable_if_t<
is_scalar_floating_point_v<T> && is_scalar_floating_point_v<U>,
bool> = true>
inline common_dtype<T, U> fmod(T x, U y) {
return ::metal::fmod(x, y);
}
// cast_to primitives
// - No-op if types as the same
template <
typename T,
typename U,
::metal::enable_if_t<::metal::is_same_v<U, T>, bool> = true>
inline T cast_to(const U from) {
return from;
}
// - Simple cast between scalar and complex dtypes
template <
typename T,
typename U,
::metal::enable_if_t<
!::metal::is_same_v<U, T> && (is_complex_v<T> == is_complex_v<U>),
bool> = true>
inline T cast_to(const U from) {
return static_cast<T>(from);
}
// - Scalar to complex
template <
typename T,
typename U,
::metal::enable_if_t<is_complex_v<T> && !is_complex_v<U>, bool> = true>
inline T cast_to(const U from) {
return T(float(from), 0.0);
}
// - Complex to scalar (should not really be used, but exists for compliteness)
template <
typename T,
typename U,
::metal::enable_if_t<!is_complex_v<T> && is_complex_v<U>, bool> = true>
inline T cast_to(const U from) {
return static_cast<T>(from.x);
}
// Generalizable math operators (used for both scalar and complex)
template <
typename T,
typename U,
::metal::enable_if_t<!is_complex_v<T>, bool> = true>
inline common_dtype<T, U> mul(const T x, const U y) {
return x * y;
}
template <
typename T,
typename U,
::metal::enable_if_t<is_complex_v<T> && is_complex_v<U>, bool> = true>
inline common_dtype<T, U> mul(const T x, const U y) {
return T(x.x * y.x - x.y * y.y, x.x * y.y + x.y * y.x);
}
template <
typename T,
typename U,
::metal::enable_if_t<!is_complex_v<T>, bool> = true>
inline common_dtype<T, U> div(const T x, const U y) {
return x / y;
}
template <
typename T,
typename U,
::metal::enable_if_t<is_complex_v<T> && is_complex_v<U>, bool> = true>
inline common_dtype<T, U> div(const T x, const U y) {
return T(::metal::dot(x, y), x.y * y.x - x.x * y.y) / ::metal::dot(y, y);
}
// Remainder operator
template <
typename T,
typename U,
::metal::enable_if_t<
is_scalar_floating_point_v<T> || is_scalar_floating_point_v<U>,
bool> = true>
inline float remainder(const T x, const U y) {
const auto x_f = static_cast<float>(x);
const auto y_f = static_cast<float>(y);
return x_f - y_f * floor_divide(x_f, y_f);
}
template <
typename T,
typename U,
::metal::enable_if_t<
is_scalar_integral_v<T> && is_scalar_integral_v<U>,
bool> = true>
inline common_dtype<T, U> remainder(const T x, const U y) {
auto rc = x % y;
return rc == 0 || (x ^ y) > 0 ? rc : rc + y;
}
// Based on aten/src/ATen/native/Pow.h
template <
typename T,
::metal::enable_if_t<is_scalar_integral_v<T>, bool> = true>
inline T powi_impl(T a, T b) {
T result = 1;
while (b) {
if (b & 1) {
result *= a;
}
b /= 2;
a *= a;
}
return result;
}
template <
typename T,
typename U,
::metal::enable_if_t<
is_scalar_floating_point_v<T> || is_scalar_floating_point_v<U>,
bool> = true>
inline float pow(T a, U b) {
return ::metal::precise::pow(static_cast<float>(a), static_cast<float>(b));
}
// Complex pow - use polar form: a = r*e^(i*theta)
// a^b = exp(b * log(a)) = exp(b * (log(r) + i*theta))
template <
typename T,
typename U,
::metal::enable_if_t<is_complex_v<T> && is_complex_v<U>, bool> = true>
inline float2 pow(T a, U b) {
// Convert a to polar form
// Use explicit computation instead of length() due to numerical issues
const auto r = ::metal::precise::sqrt(a.x * a.x + a.y * a.y);
// Special case: if r is 0, return 0
if (r == 0.0) {
return float2(0.0, 0.0);
}
const auto theta = ::metal::precise::atan2(a.y, a.x);
const auto log_r = ::metal::precise::log(r);
// Calculate a^b = r^b * e^(i*theta*b)
// new_r = exp(b.x * log(r) - b.y * theta)
// new_theta = b.x * theta + b.y * log(r)
const auto new_r = ::metal::precise::exp(b.x * log_r - b.y * theta);
const auto new_theta = b.x * theta + b.y * log_r;
return float2(
new_r * ::metal::precise::cos(new_theta),
new_r * ::metal::precise::sin(new_theta));
}
// Integral pow - unsigned types
template <
typename T,
typename U,
::metal::enable_if_t<
is_scalar_integral_v<T> && !::metal::is_signed_v<T>,
bool> = true>
inline T pow(T a, U b) {
return powi_impl(a, T(b));
}
// Integral pow - signed types
template <
typename T,
typename U,
::metal::enable_if_t<
is_scalar_integral_v<T>&& ::metal::is_signed_v<T>,
bool> = true>
inline T pow(T a, U b) {
if (b < 0) {
if (a == 1) {
return 1;
} else if (a == -1) {
auto negative = (-b) % static_cast<T>(2);
return negative ? -1 : 1;
} else {
return 0;
}
}
return powi_impl(a, T(b));
}
// Based on algorithm described in
// https://docs.oracle.com/cd/E19957-01/806-3568/ncg_goldberg.html#1202
inline float log1p(float x) {
const auto xp1 = 1.0f + x;
// First two elements of Taylor series for log(1+x) in Horner's form are:
// log(1+x) = x * (1 - x * (.5 ...)), but if 1 + x == x, then it's just x
if (xp1 == 1.0f) {
return x;
}
auto rc = ::metal::precise::log(xp1);
if (x > -.5 && x < .5) {
// Order of operations is important here for higher precision
rc *= x / (xp1 - 1.0f);
}
return rc;
}
// The function is ported from mlx
inline float2 log1p(float2 in) {
float x = in.x;
float y = in.y;
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};
}
}
template <typename T1, typename T2 = T1>
struct pair {
T1 first;
T2 second;
};
template <typename T>
inline T conj(T a) {
return a;
}
template <>
inline half2 conj(half2 a) {
return half2(a.x, -a.y);
}
template <>
inline float2 conj(float2 a) {
return float2(a.x, -a.y);
}
// The following implementation of hypot provides better numerical stability
// than the naive implementation. It is based on:
// https://github.com/pearu/functional_algorithms/blob/7dbbfd7db225b1c202e0e364fc435423ccf52dbe/functional_algorithms/algorithms.py#L168
//
// This implementation changes the naive formula for the hypotenuse of a right
// triangle, `h = sqrt(a^2 + b^2)`, into three alternate forms to be used in
// different cases. The reason why the naive formula is unstable is because of
// the square terms. If `a` or `b` are very large or very small floating point
// numbers, then their squares will resolve to inf or 0.
//
// Assume `a >= b >= 0`. We can first change the formula to:
// `h = a sqrt(1 + (b / a)^2)`
// `h = a sqrt(1 + r)`
// where `r = (b / a)^2`. Since `a >= b >= 0`, then `1 >= r >= 0`.
//
// Case 1: `a == b`
// The formula simplifies to `h = a sqrt(2)`.
//
// Case 2: `1 >> r > 0`
// Due to floating point error, `sqrt(1 + r)` resolves to 1. So we use the
// binomial approximation `sqrt(1 + r) ≈ 1 + r / 2`, and the formula becomes
// `h ≈ a + a r / 2`.
//
// Case 3: All other cases.
// Use `h = a sqrt(1 + r)`.
inline float hypot(float a_, float b_) {
auto a = max(a_, b_);
auto b = min(a_, b_);
auto b_over_a = c10::metal::div(b, a);
auto r = c10::metal::mul(b_over_a, b_over_a);
auto sqrt_1_plus_r = ::metal::precise::sqrt(1 + r);
auto h1 = M_SQRT2_F * a;
auto h2 = a + a * r / 2;
auto h3 = a * sqrt_1_plus_r;
bool is_h1 = (a == b);
bool is_h2 = ((sqrt_1_plus_r == 1) && (r > 0));
return ::metal::select(::metal::select(h3, h2, is_h2), h1, is_h1);
}
#define INSTANTIATE_FOR_ALL_TYPES(MACRO) \
MACRO(float); \
MACRO(half); \
MACRO(bfloat); \
MACRO(float2); \
MACRO(long); \
MACRO(char); \
MACRO(uchar); \
MACRO(short); \
MACRO(int);
#define INSTANTIATE_FOR_FLOAT_TYPES(MACRO) \
MACRO(float); \
MACRO(half); \
MACRO(bfloat);
} // namespace metal
} // namespace c10
#else
#error "This file should not be included when either TORCH_STABLE_ONLY or TORCH_TARGET_VERSION is defined."
#endif // !defined(TORCH_STABLE_ONLY) && !defined(TORCH_TARGET_VERSION)