[Relay][PRNG] Add uniform distribution generator wrt threefry PRNG #8041

zhuzilin · 2021-05-14T04:17:04Z

This PR adds a uniform distribution generator using the threefry PRNG introduced in #7083. We would need uniform to develop the training phase dropout as the following roadmap:

uniform -> bernoulli -> dropout

The algorithm used is basically the same as the one used in jax: using the random bits generated from threefry_generate as the fraction section of the float32 or float64. To be specific, I use the last 23 bits of the random bits for float32 and last 52 for float64. There is one different from the jax implementation. In jax, they used a bitcast to turn uint into float:

# jax implementation
def _uniform(key, shape, dtype, minval, maxval) -> jnp.ndarray:
  ...
  bits = _random_bits(key, nbits, shape)

  # The strategy here is to randomize only the mantissa bits with an exponent of
  # 1 (after applying the bias), then shift and scale to the desired range. The
  # bit-level transformation we use relies on Numpy and XLA having bit-for-bit
  # equivalent float representations, which might not be true on all platforms.
  float_bits = lax.bitwise_or(
      lax.shift_right_logical(bits, np.array(nbits - nmant, lax.dtype(bits))),
      np.array(1., dtype).view(_UINT_DTYPES[nbits]))
  floats = lax.bitcast_convert_type(float_bits, dtype) - np.array(1., dtype)
  return lax.max(
      minval,
      lax.reshape(floats * (maxval - minval) + minval, shape.positional))

However, as I haven't found the bitcast in te or topi, I use a divide to cast the type, which may be slower:

    def uniform_scalar(bits):
        bits = bits >> (nbits - nfraction)
        standard_uniform = bits.astype(out_dtype) / float(1 << nfraction)
        return standard_uniform

Thank you for your time on reviewing this PR. I may not be familiar enough with the tvm codebase at the moment, so I'm sorry for breaking any conventions in the community and I'd love to fix them :).

Gently ping @tqchen @altanh @tkonolige

tests/python/topi/python/test_topi_prng.py

FrozenGene · 2021-05-14T09:19:42Z

python/tvm/topi/random/kernel.py

+
+    Parameters
+    ----------
+    gen : Tensor[10, uint64]


What is the meaning of 10?

This is the ThreefryKeyType introduced in #7083. Please refer to:

tvm/src/relay/op/random/kernel.cc

Line 28 in c999a84

static TensorType ThreefryKeyType() { return TensorType({10}, tvm::DataType::UInt(64)); }

If so, let us add comment describe what is the meaning of 10.

You could probably say ThreefryKey instead of Tensor[10, uint64]

FrozenGene · 2021-05-14T09:20:29Z

python/tvm/topi/random/kernel.py

+        less than high.
+
+    out_shape : Sequence[int]
+        Output shape of the random numbers. Product of all dimensions must be a multiple of 4.


What is the reason of product must be a multiple of 4?

It's the property of the threefry key. Please refer to this comment: #7083 (comment)

Sorry, I just rethink about this problem. There should not be any restriction to the output shape... We could change the input restriction of the threefry_generate in other PR.

do you mind sending a PR for updating the threefry_generate output, or rather what approach do you have in mind? I tried to avoid this problem by truncating output buffer but this required an extra copy, wonder if you have something else

@altanh Sorry that I'm not familiar with the threefry algorithm. Is it possible to call _threefry twice in threefry_generate in the following form? something like:

out_array = irb.buffer_ptr(out_array_ptr) # deal with most of the array _threefry(irb, tmp, 0, tmp, 4, out_array, 0, out_len // 4) if out_len % 4 != 0: # generate remainders in a small tmp buffer tmp_array = irb.allocate(gen.dtype, 4, name="tmp", scope="global") # may need to update the tmp key in between # ... _threefry(irb, tmp, 0, tmp, 4, tmp_array, 0, out_len // 4) # only copy the tmp buffer for i in range(out_len // 4 * 4, out_len): out_array[i] = tmp_array[i%4]

In this way, we coud avoid copying the whole generated tensor.

Yeah, you could do that. Maybe submit it in a new PR?

@tkonolige Sure, I will submit one. Could you tell me what kind of update on key tmp we need before the second _threefry? I can only think of updating increment counter (tmp[7]).

You'll need to update the counter buffer to be equal to out_len

python/tvm/topi/random/kernel.py

FrozenGene · 2021-05-14T09:23:13Z

tests/python/relay/test_prng.py

@@ -103,6 +103,19 @@ def test_threefry_split_infer():
    assert tvm.ir.structural_equal(f.ret_type, expected_type)


+def test_uniform_infer():
+    oshape = (12,)
+    odtype = "float32"


Should cover more types. For example float64 you have implemented

tests/python/topi/python/test_topi_prng.py

python/tvm/relay/op/random/kernel.py

src/relay/op/random/kernel.cc

tkonolige · 2021-05-14T16:18:35Z

python/tvm/topi/random/kernel.py

+
+    standard_uniform_values = tvm.te.compute(out_shape, lambda *i: uniform_scalar(random_bits(*i)))
+
+    uniform_values = tvm.topi.add(tvm.topi.multiply(standard_uniform_values, high - low), low)


How well does this approach work when we have a large range (high - low)? It seems like we would be loosing a lot of potential randomness.

altanh · 2021-05-14T18:54:41Z

Thanks for this PR! I will be reading it soon, and just wanted to point you to a branch I worked on a while ago where I hacked a uniform op + dropout support: https://github.com/altanh/tvm/commits/prng (just in case it might be useful for you to check and compare).

However, as I haven't found the bitcast in te or topi, I use a divide to cast the type, which may be slower:

Perhaps this is the operation you're looking for? https://github.com/altanh/tvm/blob/2d9ac7710ab055d4f20e8b5a0a3580836723efac/python/tvm/topi/generic/algorithm.py#L465

Thanks!

zhuzilin · 2021-05-15T17:06:22Z

@FrozenGene @tkonolige @altanh Thank you for your reviews. I've updated this PR based on them.

I think there are 2 options: 1. Ref Conv2DRel. 2. You could restrict the type here and raise exception. I prefer option 1.

@FrozenGene Thank you for the clue! However, I haven't find how to restrict the dtype attributes in Conv2DRel or Conv2DAttrs... Should I add the type restriction to the UniformAttrs, or raise error in the MakeUniform and UniformRel?

How well does this approach work when we have a large range (high - low)? It seems like we would be loosing a lot of potential randomness.

@tkonolige As this approach is only using the fraction bits to represent float, there will be loss of randomness for all floats, at least (2^nexp-1) / 2^nexp of the float (nexp stands for the number of exponential digits). However, it's a little tricky to use all 64 digits of the random bit to represent a uniform distributed float number... Do you have any idea on that?

Perhaps this is the operation you're looking for? https://github.com/altanh/tvm/blob/2d9ac7710ab055d4f20e8b5a0a3580836723efac/python/tvm/topi/generic/algorithm.py#L465

@altanh Thank you for your references! The reinterpret is exactly what I was looking for. I've updated the algorithm and right now it is the same as the one used in jax.

FrozenGene · 2021-05-17T01:35:40Z

@FrozenGene Thank you for the clue! However, I haven't find how to restrict the dtype attributes in Conv2DRel or Conv2DAttrs... Should I add the type restriction to the UniformAttrs, or raise error in the MakeUniform and UniformRel?

Suggest UniformRel

zhuzilin · 2021-05-17T15:10:41Z

Suggest UniformRel

@FrozenGene Thank you. I've added the type restriction.

tkonolige

Looks pretty good to me. Just a couple of minor fixes.

python/tvm/relay/op/random/kernel.py

python/tvm/topi/random/kernel.py

altanh

Overall LGTM with some minor comments. I did want to request that we keep the output shape restriction in the documentation for now until a follow up PR is merged which relaxes it. Thanks for the work!

src/relay/op/random/kernel.cc

python/tvm/relay/op/strategy/generic.py

zhuzilin · 2021-05-18T08:24:47Z

@FrozenGene @altanh @tkonolige I've updated the PR upon the reviews. Could you take another look? Thank you~

tkonolige

Looks good to me. Just some small comments

python/tvm/topi/random/kernel.py

src/relay/op/random/kernel.cc

altanh

LGTM! I'm a bit uneasy about introducing a nondeterministic test based on averaging the random numbers but I imagine it will almost never fail. Also left a comment about comparing the min/max of the generated numbers - can we always guarantee <= or >= on the output or will there be some floating point inaccuracy cases where this might be violated?

tests/python/topi/python/test_topi_prng.py

zhuzilin · 2021-05-19T03:48:59Z

@FrozenGene Could you take another look of this PR? Thank you~

zhuzilin · 2021-05-21T02:36:18Z

@FrozenGene Could you have another look at this PR? Thank you!

FrozenGene · 2021-05-21T03:15:39Z

Thanks @zhuzilin @altanh @tkonolige merged now

…pache#8041) * Add uniform distribution generator wrt threefry PRNG * fix lint * remove the redundant print * modifications based on review * update docs * update uniform algorithm to use bit operations only * add type restrictions * minor fix upon review * update test and error information

Add uniform distribution generator wrt threefry PRNG

7baf726

zhuzilin force-pushed the prng_uniform branch 2 times, most recently from 2f831ea to 78bc8b7 Compare May 14, 2021 04:28

fix lint

fd68ff4

zhuzilin force-pushed the prng_uniform branch from 78bc8b7 to fd68ff4 Compare May 14, 2021 04:44

FrozenGene reviewed May 14, 2021

View reviewed changes

tests/python/topi/python/test_topi_prng.py Outdated Show resolved Hide resolved

FrozenGene reviewed May 14, 2021

View reviewed changes

tests/python/topi/python/test_topi_prng.py Outdated Show resolved Hide resolved

remove the redundant print

135da62

zhuzilin requested a review from FrozenGene May 14, 2021 05:04

FrozenGene reviewed May 14, 2021

View reviewed changes

tqchen assigned FrozenGene May 14, 2021

tqchen added the status: need review label May 14, 2021

tkonolige reviewed May 14, 2021

View reviewed changes

zhuzilin added 2 commits May 16, 2021 00:09

modifications based on review

cbaec73

update docs

024e176

update uniform algorithm to use bit operations only

fd1368b

zhuzilin requested a review from FrozenGene May 16, 2021 02:33

add type restrictions

bf577e4

tkonolige approved these changes May 17, 2021

View reviewed changes

python/tvm/relay/op/random/kernel.py Outdated Show resolved Hide resolved

python/tvm/topi/random/kernel.py Show resolved Hide resolved

altanh reviewed May 17, 2021

View reviewed changes

src/relay/op/random/kernel.cc Outdated Show resolved Hide resolved

python/tvm/relay/op/strategy/generic.py Show resolved Hide resolved

zhuzilin force-pushed the prng_uniform branch from 7e4e0bc to 49383a7 Compare May 18, 2021 09:24

minor fix upon review

aa372d4

zhuzilin force-pushed the prng_uniform branch from 49383a7 to aa372d4 Compare May 18, 2021 09:26

tkonolige approved these changes May 18, 2021

View reviewed changes

python/tvm/topi/random/kernel.py Outdated Show resolved Hide resolved

src/relay/op/random/kernel.cc Outdated Show resolved Hide resolved

altanh approved these changes May 18, 2021

View reviewed changes

tests/python/topi/python/test_topi_prng.py Outdated Show resolved Hide resolved

tests/python/topi/python/test_topi_prng.py Outdated Show resolved Hide resolved

update test and error information

c707ce0

FrozenGene approved these changes May 21, 2021

View reviewed changes

FrozenGene merged commit e438a73 into apache:main May 21, 2021

FrozenGene added status: accepted and removed status: need review labels May 21, 2021

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Relay][PRNG] Add uniform distribution generator wrt threefry PRNG #8041

[Relay][PRNG] Add uniform distribution generator wrt threefry PRNG #8041

zhuzilin commented May 14, 2021

FrozenGene May 14, 2021

zhuzilin May 14, 2021

FrozenGene May 14, 2021

tkonolige May 14, 2021

zhuzilin May 15, 2021

FrozenGene May 14, 2021

zhuzilin May 14, 2021

zhuzilin May 15, 2021 •

edited

Loading

altanh May 17, 2021

zhuzilin May 18, 2021

tkonolige May 18, 2021

zhuzilin May 19, 2021

tkonolige May 19, 2021

FrozenGene May 14, 2021

tkonolige May 14, 2021

altanh commented May 14, 2021

zhuzilin commented May 15, 2021 •

edited

Loading

FrozenGene commented May 17, 2021 •

edited

Loading

zhuzilin commented May 17, 2021

tkonolige left a comment

altanh left a comment

zhuzilin commented May 18, 2021 •

edited

Loading

tkonolige left a comment

altanh left a comment

zhuzilin commented May 19, 2021

zhuzilin commented May 21, 2021

FrozenGene commented May 21, 2021


		standard_uniform_values = tvm.te.compute(out_shape, lambda i: uniform_scalar(random_bits(i)))

		uniform_values = tvm.topi.add(tvm.topi.multiply(standard_uniform_values, high - low), low)

[Relay][PRNG] Add uniform distribution generator wrt threefry PRNG #8041

[Relay][PRNG] Add uniform distribution generator wrt threefry PRNG #8041

Conversation

zhuzilin commented May 14, 2021

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

zhuzilin May 15, 2021 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

altanh commented May 14, 2021

zhuzilin commented May 15, 2021 • edited Loading

FrozenGene commented May 17, 2021 • edited Loading

zhuzilin commented May 17, 2021

tkonolige left a comment

Choose a reason for hiding this comment

altanh left a comment

Choose a reason for hiding this comment

zhuzilin commented May 18, 2021 • edited Loading

tkonolige left a comment

Choose a reason for hiding this comment

altanh left a comment

Choose a reason for hiding this comment

zhuzilin commented May 19, 2021

zhuzilin commented May 21, 2021

FrozenGene commented May 21, 2021

zhuzilin May 15, 2021 •

edited

Loading

zhuzilin commented May 15, 2021 •

edited

Loading

FrozenGene commented May 17, 2021 •

edited

Loading

zhuzilin commented May 18, 2021 •

edited

Loading