triton-windows 3.1.0.post17__cp39-cp39-win_amd64.whl

triton/language/extra/libdevice.py

@@ -0,0 +1,1213 @@
+from .cuda import libdevice as cuda_libdevice
+from .hip import libdevice as hip_libdevice
+from triton.language import core
+from functools import wraps
+from typing import TypeVar
+
+T = TypeVar('T')
+
+
+def dispatch(fn: T) -> T:
+    """Dispatch a function to a correct implementation."""
+    assert callable(fn)
+
+    @wraps(fn)
+    def wrapper(*args, **kwargs):
+        _backend = kwargs["_builder"].options.backend_name
+        if _backend == 'cuda':
+            _curr_libdevice_module = cuda_libdevice
+        elif _backend == 'hip':
+            _curr_libdevice_module = hip_libdevice
+        else:
+            raise RuntimeError('unknown backend')
+
+        try:
+            _impl = getattr(_curr_libdevice_module, fn.__name__)
+        except AttributeError:
+            raise RuntimeError(f'`{_backend}` does not provide support for `{fn.__name__}` extra function')
+
+        return _impl(*args, **kwargs)
+
+    return wrapper
+
+
+@core.extern
+@dispatch
+def clz(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def popc(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def byte_perm(arg0, arg1, arg2, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def mulhi(arg0, arg1, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def mul24(arg0, arg1, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def brev(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def sad(arg0, arg1, arg2, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def abs(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def floor(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def rcp64h(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def rsqrt(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def ceil(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def trunc(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def exp2(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def saturatef(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def fma_rn(arg0, arg1, arg2, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def fma_rz(arg0, arg1, arg2, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def fma_rd(arg0, arg1, arg2, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def fma_ru(arg0, arg1, arg2, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def fast_dividef(arg0, arg1, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def div_rn(arg0, arg1, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def div_rz(arg0, arg1, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def div_rd(arg0, arg1, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def div_ru(arg0, arg1, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def rcp_rn(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def rcp_rz(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def rcp_rd(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def rcp_ru(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def sqrt_rn(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def sqrt_rz(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def sqrt_rd(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def sqrt_ru(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def sqrt(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def add_rn(arg0, arg1, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def add_rz(arg0, arg1, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def add_rd(arg0, arg1, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def add_ru(arg0, arg1, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def mul_rn(arg0, arg1, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def mul_rz(arg0, arg1, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def mul_rd(arg0, arg1, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def mul_ru(arg0, arg1, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def double2float_rn(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def double2float_rz(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def double2float_rd(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def double2float_ru(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def double2int_rn(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def double2int_rz(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def double2int_rd(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def double2int_ru(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def double2uint_rn(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def double2uint_rz(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def double2uint_rd(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def double2uint_ru(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def int2double_rn(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def uint2double_rn(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def float2int_rn(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def float2int_rz(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def float2int_rd(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def float2int_ru(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def float2uint_rn(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def float2uint_rz(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def float2uint_rd(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def float2uint_ru(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def int2float_rn(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def int2float_rz(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def int2float_rd(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def int2float_ru(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def uint2float_rn(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def uint2float_rz(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def uint2float_rd(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def uint2float_ru(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def hiloint2double(arg0, arg1, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def double2loint(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def double2hiint(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def float2ll_rn(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def float2ll_rz(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def float2ll_rd(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def float2ll_ru(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def float2ull_rn(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def float2ull_rz(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def float2ull_rd(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def float2ull_ru(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def double2ll_rn(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def double2ll_rz(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def double2ll_rd(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def double2ll_ru(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def double2ull_rn(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def double2ull_rz(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def double2ull_rd(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def double2ull_ru(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def ll2float_rn(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def ll2float_rz(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def ll2float_rd(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def ll2float_ru(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def ull2float_rn(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def ull2float_rz(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def ull2float_rd(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def ull2float_ru(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def ll2double_rn(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def ll2double_rz(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def ll2double_rd(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def ll2double_ru(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def ull2double_rn(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def ull2double_rz(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def ull2double_rd(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def ull2double_ru(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def int_as_float(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def float_as_int(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def uint_as_float(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def float_as_uint(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def longlong_as_double(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def double_as_longlong(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def fast_sinf(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def fast_cosf(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def fast_log2f(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def fast_logf(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def fast_expf(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def fast_tanf(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def fast_exp10f(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def fast_log10f(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def fast_powf(arg0, arg1, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def hadd(arg0, arg1, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def rhadd(arg0, arg1, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def sub_rn(arg0, arg1, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def sub_rz(arg0, arg1, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def sub_rd(arg0, arg1, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def sub_ru(arg0, arg1, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def rsqrt_rn(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def ffs(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def rint(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def llrint(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def nearbyint(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def isnan(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def signbit(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def copysign(arg0, arg1, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def finitef(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def isinf(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def nextafter(arg0, arg1, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def sin(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def cos(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def sinpi(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def cospi(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def tan(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def log2(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def exp(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def exp10(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def cosh(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def sinh(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def tanh(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def atan2(arg0, arg1, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def atan(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def asin(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def acos(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def log(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def log10(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def log1p(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def acosh(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def asinh(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def atanh(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def expm1(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def hypot(arg0, arg1, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def rhypot(arg0, arg1, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def norm3d(arg0, arg1, arg2, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def rnorm3d(arg0, arg1, arg2, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def norm4d(arg0, arg1, arg2, arg3, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def rnorm4d(arg0, arg1, arg2, arg3, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def cbrt(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def rcbrt(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def j0(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def j1(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def y0(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def y1(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def yn(arg0, arg1, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def jn(arg0, arg1, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def cyl_bessel_i0(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def cyl_bessel_i1(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def erf(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def erfinv(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def erfc(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def erfcx(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def erfcinv(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def normcdfinv(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def normcdf(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def lgamma(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def ldexp(arg0, arg1, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def scalbn(arg0, arg1, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def fmod(arg0, arg1, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def remainder(arg0, arg1, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def fma(arg0, arg1, arg2, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def pow(arg0, arg1, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def tgamma(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def round(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def llround(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def fdim(arg0, arg1, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def ilogb(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def logb(arg0, _builder=None):
+    ...
+
+
+@core.extern
+@dispatch
+def isfinited(arg0, _builder=None):
+    ...