gaspersic/triton
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

[FRONTEND][RANDOM] Make sure offset dtype is always uint32 before calling uint32_to_uniform_float (#427)

Browse Source
This commit is contained in:
Philippe Tillet
2022-01-11 11:08:49 -08:00
committed by GitHub
parent bf32205edc
commit bbc78f6516
1 changed files with 3 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

4
python/triton/language/random.py
View File

@@ -76,7 +76,7 @@ def uint32_to_uniform_float(x):
""" """
Numerically stable function to convert a random uint32 into a random float uniformly sampled in [0, 1). Numerically stable function to convert a random uint32 into a random float uniformly sampled in [0, 1).
""" """
two_to_the_minus_32 = 2.328306e-10 two_to_the_minus_32: tl.constexpr = 2.328306e-10
return x * two_to_the_minus_32 return x * two_to_the_minus_32
@@ -89,6 +89,7 @@ def rand(seed, offset, n_rounds: tl.constexpr = N_ROUNDS_DEFAULT):
:param seed: The seed for generating random numbers. :param seed: The seed for generating random numbers.
:param offsets: The offsets to generate random numbers for. :param offsets: The offsets to generate random numbers for.
""" """
offset = offset.to(tl.uint32, bitcast=True)
source = randint(seed, offset, n_rounds) source = randint(seed, offset, n_rounds)
return uint32_to_uniform_float(source) return uint32_to_uniform_float(source)
@@ -102,6 +103,7 @@ def rand4x(seed, offsets, n_rounds: tl.constexpr = N_ROUNDS_DEFAULT):
:param seed: The seed for generating random numbers. :param seed: The seed for generating random numbers.
:param offsets: The offsets to generate random numbers for. :param offsets: The offsets to generate random numbers for.
""" """
offsets = offsets.to(tl.uint32, bitcast=True)
i1, i2, i3, i4 = randint4x(seed, offsets, n_rounds) i1, i2, i3, i4 = randint4x(seed, offsets, n_rounds)
u1 = uint32_to_uniform_float(i1) u1 = uint32_to_uniform_float(i1)
u2 = uint32_to_uniform_float(i2) u2 = uint32_to_uniform_float(i2)