triton-windows 3.3.0.post19__cp39-cp39-win_amd64.whl → 3.4.0.post20__cp39-cp39-win_amd64.whl

triton/language/math.py

@@ -1,5 +1,4 @@
 from . import core
-from . import semantic
 from functools import wraps
 from typing import List
 
@@ -85,107 +84,107 @@ def _add_math_3arg_docstr(name: str) -> core.Callable[[T], T]:
 @core.builtin
 @_check_dtype(dtypes=["int32", "int64", "uint32", "uint64"])
 @_add_math_2arg_docstr("most significant N bits of the 2N-bit product")
-def umulhi(x, y, _builder=None):
-    x = semantic.to_tensor(x, _builder)
-    y = semantic.to_tensor(y, _builder)
-    x, y = core.binary_op_type_legalization(x, y, _builder)
-    return core.tensor(_builder.create_umulhi(x.handle, y.handle), x.type)
+def umulhi(x, y, _semantic=None):
+    x = _semantic.to_tensor(x)
+    y = _semantic.to_tensor(y)
+    x, y = core.binary_op_type_legalization(x, y, _semantic)
+    return core.tensor(_semantic.builder.create_umulhi(x.handle, y.handle), x.type)
 
 
 @core.builtin
 @_check_dtype(dtypes=["fp32", "fp64"])
 @_add_math_1arg_docstr("exponential")
 @core._tensor_member_fn
-def exp(x, _builder=None):
-    x = semantic.to_tensor(x, _builder)
-    return core.tensor(_builder.create_exp(x.handle), x.type)
+def exp(x, _semantic=None):
+    x = _semantic.to_tensor(x)
+    return core.tensor(_semantic.builder.create_exp(x.handle), x.type)
 
 
 @core.builtin
 @_check_dtype(dtypes=["fp32", "fp64"])
 @_add_math_1arg_docstr("exponential (base 2)")
 @core._tensor_member_fn
-def exp2(x, _builder=None):
-    x = semantic.to_tensor(x, _builder)
-    return core.tensor(_builder.create_exp2(x.handle), x.type)
+def exp2(x, _semantic=None):
+    x = _semantic.to_tensor(x)
+    return core.tensor(_semantic.builder.create_exp2(x.handle), x.type)
 
 
 @core.builtin
 @_check_dtype(dtypes=["fp32", "fp64"])
 @_add_math_1arg_docstr("natural logarithm")
 @core._tensor_member_fn
-def log(x, _builder=None):
-    x = semantic.to_tensor(x, _builder)
-    return core.tensor(_builder.create_log(x.handle), x.type)
+def log(x, _semantic=None):
+    x = _semantic.to_tensor(x)
+    return core.tensor(_semantic.builder.create_log(x.handle), x.type)
 
 
 @core.builtin
 @_check_dtype(dtypes=["fp32", "fp64"])
 @_add_math_1arg_docstr("logarithm (base 2)")
 @core._tensor_member_fn
-def log2(x, _builder=None):
-    x = semantic.to_tensor(x, _builder)
-    return core.tensor(_builder.create_log2(x.handle), x.type)
+def log2(x, _semantic=None):
+    x = _semantic.to_tensor(x)
+    return core.tensor(_semantic.builder.create_log2(x.handle), x.type)
 
 
 @core.builtin
 @_check_dtype(dtypes=["fp32", "fp64"])
 @_add_math_1arg_docstr("cosine")
 @core._tensor_member_fn
-def cos(x, _builder=None):
-    x = semantic.to_tensor(x, _builder)
-    return core.tensor(_builder.create_cos(x.handle), x.type)
+def cos(x, _semantic=None):
+    x = _semantic.to_tensor(x)
+    return core.tensor(_semantic.builder.create_cos(x.handle), x.type)
 
 
 @core.builtin
 @_check_dtype(dtypes=["fp32", "fp64"])
 @_add_math_1arg_docstr("sine")
 @core._tensor_member_fn
-def sin(x, _builder=None):
-    x = semantic.to_tensor(x, _builder)
-    return core.tensor(_builder.create_sin(x.handle), x.type)
+def sin(x, _semantic=None):
+    x = _semantic.to_tensor(x)
+    return core.tensor(_semantic.builder.create_sin(x.handle), x.type)
 
 
 @core.builtin
 @_check_dtype(dtypes=["fp32", "fp64"])
 @_add_math_1arg_docstr("fast square root")
 @core._tensor_member_fn
-def sqrt(x, _builder=None):
-    x = semantic.to_tensor(x, _builder)
-    return core.tensor(_builder.create_sqrt(x.handle), x.type)
+def sqrt(x, _semantic=None):
+    x = _semantic.to_tensor(x)
+    return core.tensor(_semantic.builder.create_sqrt(x.handle), x.type)
 
 
 @core.builtin
 @_check_dtype(dtypes=["fp32"])
 @_add_math_1arg_docstr("precise square root (rounding to nearest wrt the IEEE standard)")
 @core._tensor_member_fn
-def sqrt_rn(x, _builder=None):
-    x = semantic.to_tensor(x, _builder)
-    return core.tensor(_builder.create_precise_sqrt(x.handle), x.type)
+def sqrt_rn(x, _semantic=None):
+    x = _semantic.to_tensor(x)
+    return core.tensor(_semantic.builder.create_precise_sqrt(x.handle), x.type)
 
 
 @core.builtin
 @_check_dtype(dtypes=["fp32", "fp64"])
 @_add_math_1arg_docstr("inverse square root")
 @core._tensor_member_fn
-def rsqrt(x, _builder=None):
-    x = semantic.to_tensor(x, _builder)
-    return core.tensor(_builder.create_rsqrt(x.handle), x.type)
+def rsqrt(x, _semantic=None):
+    x = _semantic.to_tensor(x)
+    return core.tensor(_semantic.builder.create_rsqrt(x.handle), x.type)
 
 
 @core._tensor_member_fn
 @core.builtin
 @_add_math_1arg_docstr("absolute value")
-def abs(x, _builder=None):
-    x = semantic.to_tensor(x, _builder)
+def abs(x, _semantic=None):
+    x = _semantic.to_tensor(x)
     dtype = x.dtype
     if dtype.is_fp8e4b15():
-        mask = core.full(x.shape, 0x7F, core.int8, _builder=_builder)
-        return core.tensor(_builder.create_and(x.handle, mask.handle), x.type)
+        mask = core.full(x.shape, 0x7F, core.int8, _semantic=_semantic)
+        return core.tensor(_semantic.builder.create_and(x.handle, mask.handle), x.type)
     elif dtype.is_floating():
-        return core.tensor(_builder.create_fabs(x.handle), x.type)
+        return core.tensor(_semantic.builder.create_fabs(x.handle), x.type)
     elif dtype.is_int_signed():
-        return core.tensor(_builder.create_iabs(x.handle), x.type)
+        return core.tensor(_semantic.builder.create_iabs(x.handle), x.type)
     elif dtype.is_int_unsigned():
         return x  # no-op
     else:
@@ -194,57 +193,57 @@ def abs(x, _builder=None):
 
 @core.builtin
 @_add_math_2arg_docstr("fast division")
-def fdiv(x, y, ieee_rounding=False, _builder=None):
-    ieee_rounding = core._constexpr_to_value(ieee_rounding)
-    x = semantic.to_tensor(x, _builder)
-    y = semantic.to_tensor(y, _builder)
-    return semantic.fdiv(x, y, ieee_rounding, _builder)
+def fdiv(x, y, ieee_rounding=False, _semantic=None):
+    ieee_rounding = core._unwrap_if_constexpr(ieee_rounding)
+    x = _semantic.to_tensor(x)
+    y = _semantic.to_tensor(y)
+    return _semantic.fdiv(x, y, ieee_rounding)
 
 
 @core.builtin
 @_check_dtype(dtypes=["fp32"])
 @_add_math_2arg_docstr("precise division (rounding to nearest wrt the IEEE standard)")
-def div_rn(x, y, _builder=None):
-    x = semantic.to_tensor(x, _builder)
-    y = semantic.to_tensor(y, _builder)
-    x, y = core.binary_op_type_legalization(x, y, _builder)
-    return core.tensor(_builder.create_precise_divf(x.handle, y.handle), x.type)
+def div_rn(x, y, _semantic=None):
+    x = _semantic.to_tensor(x)
+    y = _semantic.to_tensor(y)
+    x, y = core.binary_op_type_legalization(x, y, _semantic)
+    return core.tensor(_semantic.builder.create_precise_divf(x.handle, y.handle), x.type)
 
 
 @core.builtin
 @_check_dtype(dtypes=["fp32", "fp64"])
 @_add_math_1arg_docstr("error function")
 @core._tensor_member_fn
-def erf(x, _builder=None):
-    x = semantic.to_tensor(x, _builder)
-    return core.tensor(_builder.create_erf(x.handle), x.type)
+def erf(x, _semantic=None):
+    x = _semantic.to_tensor(x)
+    return core.tensor(_semantic.builder.create_erf(x.handle), x.type)
 
 
 @core.builtin
 @_check_dtype(dtypes=["fp32", "fp64"])
 @_add_math_1arg_docstr("floor")
 @core._tensor_member_fn
-def floor(x, _builder=None):
-    x = semantic.to_tensor(x, _builder)
-    return core.tensor(_builder.create_floor(x.handle), x.type)
+def floor(x, _semantic=None):
+    x = _semantic.to_tensor(x)
+    return core.tensor(_semantic.builder.create_floor(x.handle), x.type)
 
 
 @core.builtin
 @_check_dtype(dtypes=["fp32", "fp64"])
 @_add_math_1arg_docstr("ceil")
 @core._tensor_member_fn
-def ceil(x, _builder=None):
-    x = semantic.to_tensor(x, _builder)
-    return core.tensor(_builder.create_ceil(x.handle), x.type)
+def ceil(x, _semantic=None):
+    x = _semantic.to_tensor(x)
+    return core.tensor(_semantic.builder.create_ceil(x.handle), x.type)
 
 
 @core.builtin
 @_add_math_3arg_docstr("fused multiply-add")
-def fma(x, y, z, _builder=None):
-    x = semantic.to_tensor(x, _builder)
-    y = semantic.to_tensor(y, _builder)
-    z = semantic.to_tensor(z, _builder)
-    x, y = core.binary_op_type_legalization(x, y, _builder)
-    z, x = core.binary_op_type_legalization(z, x, _builder)
-    z, y = core.binary_op_type_legalization(z, y, _builder)
-    return core.tensor(_builder.create_fma(x.handle, y.handle, z.handle), x.type)
+def fma(x, y, z, _semantic=None):
+    x = _semantic.to_tensor(x)
+    y = _semantic.to_tensor(y)
+    z = _semantic.to_tensor(z)
+    x, y = core.binary_op_type_legalization(x, y, _semantic)
+    z, x = core.binary_op_type_legalization(z, x, _semantic)
+    z, y = core.binary_op_type_legalization(z, y, _semantic)
+    return core.tensor(_semantic.builder.create_fma(x.handle, y.handle, z.handle), x.type)

triton/language/random.py

@@ -51,6 +51,7 @@ def philox(seed, c0, c1, c2, c3, n_rounds: tl.constexpr = N_ROUNDS_DEFAULT):
     c1 = tl.to_tensor(c1)
     c2 = tl.to_tensor(c2)
     c3 = tl.to_tensor(c3)
+
     if tl.constexpr(c0.dtype.primitive_bitwidth) == 32:
         int_dtype = tl.uint32
         seed_hi = ((seed >> 32) & 0xffffffff).to(tl.uint32)
@@ -60,6 +61,7 @@ def philox(seed, c0, c1, c2, c3, n_rounds: tl.constexpr = N_ROUNDS_DEFAULT):
         int_dtype = tl.uint64
         seed_hi = tl.full((1, ), 0, dtype=int_dtype)
         seed_lo = seed
+
     c0 = c0.to(int_dtype, bitcast=True)
     c1 = c1.to(int_dtype, bitcast=True)
     c2 = c2.to(int_dtype, bitcast=True)
@@ -96,8 +98,16 @@ def randint4x(seed, offset, n_rounds: tl.constexpr = N_ROUNDS_DEFAULT):
     :param offsets: The offsets to generate random numbers for.
     """
     # _0 = tl.zeros(offset.shape, offset.dtype)
-    _0 = offset * 0
-    return philox(seed, offset, _0, _0, _0, n_rounds)
+
+    offset_lo = offset.to(tl.uint32)
+    _0 = offset_lo * 0
+
+    if tl.constexpr(offset.dtype.primitive_bitwidth) > 32:
+        offset_hi = (offset >> 32).to(tl.uint32)
+    else:
+        offset_hi = _0
+
+    return philox(seed, offset_lo, offset_hi, _0, _0, n_rounds)
 
 
 # -------------------