[RELAY,TOPI] Threefry PRNG: splittable and stateless (apache#7083)

* [RELAY,TOPI] Threefry PRNG: splittable and stateless

* Fix sphinx?

* Lint fixes

* sphinx fixes round 2

* fix inputs for tests

* reorganize to random, fix uninitialized memory bug

* silence linter

* silence linter even further

* s

* strengthen Threefry key type checking, add tests

* replace static variable with function for Threefry key type

* lint fix

* Remove old todos, improve assert messages

* describe how random number is generated

* add tests for incorrect output size. also vary test sizes

Co-authored-by: Altan Haan <ahaan@octoml.ai>

include/tvm/relay/attrs/random.h

@@ -0,0 +1,42 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied.  See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+
+/*!
+ * \file tvm/relay/attrs/vision.h
+ * \brief Auxiliary attributes for random operators.
+ */
+#ifndef TVM_RELAY_ATTRS_RANDOM_H_
+#define TVM_RELAY_ATTRS_RANDOM_H_
+
+#include <tvm/ir/attrs.h>
+
+namespace tvm {
+namespace relay {
+
+struct ThreefryGenerateAttrs : public tvm::AttrsNode<ThreefryGenerateAttrs> {
+  Array<Integer> out_shape;
+
+  TVM_DECLARE_ATTRS(ThreefryGenerateAttrs, "relay.attrs.ThreefryGenerateAttrs") {
+    TVM_ATTR_FIELD(out_shape).describe("Shape of random numbers to generate");
+  }
+};
+
+}  // namespace relay
+}  // namespace tvm
+#endif  // TVM_RELAY_ATTRS_RANDOM_H_

python/tvm/relay/__init__.py

@@ -45,6 +45,7 @@
 from .op import vision
 from .op import contrib
 from .op import dyn
+from .op import random
 from .op.reduce import *
 from .op.tensor import *
 from .op.transform import *

python/tvm/relay/op/__init__.py

@@ -43,6 +43,7 @@
 from . import image
 from . import vision
 from . import op_attrs
+from . import random
 
 
 # operator registry

python/tvm/relay/op/algorithm.py

@@ -17,9 +17,9 @@
 """Classic algorithm operation"""
 from __future__ import absolute_import as _abs
 
+from ..expr import Constant, Expr, TupleWrapper
 from . import _make
 from .dyn import _make as _dyn_make
-from ..expr import TupleWrapper, Expr, Constant
 
 
 def sort(data, axis=-1, is_ascend=1):

python/tvm/relay/op/op_attrs.py

@@ -552,3 +552,8 @@ class SpaceToBatchNDAttrs(Attrs):
 @tvm._ffi.register_object("relay.attrs.BatchToSpaceNDAttrs")
 class BatchToSpaceNDAttrs(Attrs):
     """Attributes used in BatchToSpaceNDAttrs operators"""
+
+
+@tvm._ffi.register_object("relay.attrs.ThreefryGenerateAttrs")
+class ThreefryGenerateAttrs(Attrs):
+    """Attributes used in ThreefryGenerateAttrs operators"""

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

@@ -0,0 +1,20 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+# pylint: disable=wildcard-import
+"""PRNG related operators."""
+from .kernel import *
+from . import _kernel

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

@@ -0,0 +1,29 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Splittable and parallelizable PRNG kernels."""
+# pylint: disable=invalid-name,unused-argument
+from __future__ import absolute_import
+
+from .. import strategy
+from ..op import register_strategy, register_pattern, OpPattern
+
+
+# Threefry
+register_strategy("random.threefry_generate", strategy.threefry_generate_strategy)
+register_pattern("random.threefry_generate", OpPattern.OPAQUE)
+register_strategy("random.threefry_split", strategy.threefry_split_strategy)
+register_pattern("random.threefry_split", OpPattern.OPAQUE)

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

@@ -0,0 +1,20 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Constructor APIs"""
+import tvm._ffi
+
+tvm._ffi._init_api("relay.op.random._make", __name__)

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

@@ -0,0 +1,134 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+"""Splittable and parallelizable PRNG kernels."""
+# pylint: disable=invalid-name,unused-argument
+from __future__ import absolute_import
+
+import sys
+import numpy as np
+
+from ...expr import Constant
+from .... import nd
+from . import _make
+
+
+def threefry_key(seed):
+    """Create a new Threefry random number generator key.
+
+    Example
+    -------
+
+    .. code-block:: python
+
+        gen = threefry_key(0)
+        _, random_number = threefry_generate(gen, (4,))
+
+    Parameters
+    ----------
+    seed : int
+        Starting seed for the key
+
+    Returns
+    -------
+    key : relay.Expr
+        New key to pass to future uses of :py:func:`threefry_split` or
+        :py:func:`threefry_generate`.
+    """
+    s = np.frombuffer(seed.to_bytes(32, sys.byteorder), dtype="uint64")
+    a = np.concatenate((s, np.array([0, 0, 0, 0, 1 << 63, 0], dtype="uint64")))
+    return Constant(nd.array(a))
+
+
+def threefry_generate(key, shape):
+    """Generate an array of random bits (`uint64`) using the Threefry algorithm
+
+    Example
+    -------
+
+    .. code-block:: python
+
+        key = threefry_key(0)
+        new_key, random1 = threefry_generate(key, (4,))
+        _, random2 = threefry_generate(new_key, (4,))
+        # random1 and random2 are different random numbers
+
+    Parameters
+    ----------
+    key : relay.Expr
+        key that uniquely determines the random values. Multiple uses with the
+        same key will generate the same random values. This key should be
+        treated as an opaque pointer. You can create one from calling
+        :py:func:`threefry_key`, :py:func:`threefry_split`, or
+        :py:func:`threefry_generate`. **Do not use this key again after calling
+        this function.**
+
+    shape : Sequence[int]
+        Desired outputs shape of random numbers. **Currently the total
+        number of elements must be a multiple of 4.**
+
+    Returns
+    -------
+    new_key : relay.Expr
+        New key to pass to future uses of :py:func:`threefry_split` or
+        :py:func:`threefry_generate`.
+
+    random_array : relay.Expr
+        Array of random numbers. Has shape `shape`.
+    """
+    return _make.threefry_generate(key, shape)
+
+
+def threefry_split(key):
+    """Split an existing Threefry key into two new ones.
+
+    This is useful if you have to subsequent calls which each need their own
+    independent random number generation.
+
+    Example
+    -------
+
+    .. code-block:: python
+
+        def foo(key):
+            new_key, num = threefry_generate(key, (4,))
+            return num
+
+        key = threefry_key(0)
+        key1, key2 = threefry_split(key)
+        assert foo(key1) != foo(key2)
+
+    Parameters
+    ----------
+    key : relay.Expr
+        key that uniquely determines the random values. Multiple uses with the
+        same generator will generate the same random values. This generator should be
+        treated as an opaque pointer. You can create one from calling
+        :py:func:`threefry_key`, :py:func:`threefry_split`, or
+        :py:func:`threefry_generate`. **Do not use this generator again after calling
+        this function.**
+
+    Returns
+    -------
+    new_key_1 : relay.Expr
+        New key to pass to future uses of :py:func:`threefry_split` or
+        :py:func:`threefry_generate`.
+
+    new_key_2 : relay.Expr
+        New key to pass to future uses of :py:func:`threefry_split` or
+        :py:func:`threefry_generate`.
+    """
+    return _make.threefry_split(key)

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

@@ -1325,3 +1325,47 @@ def argwhere_strategy(attrs, inputs, out_type, target):
         name="argwhere.generic",
     )
     return strategy
+
+
+# threefry_generate
+def wrap_compute_threefry_generate(topi_compute):
+    """Wrap threefry_generate topi compute"""
+
+    def _compute_threefry_generate(attrs, inputs, _):
+        return topi_compute(inputs[0], attrs.out_shape)
+
+    return _compute_threefry_generate
+
+
+@override_native_generic_func("threefry_generate_strategy")
+def threefry_generate_strategy(attrs, inputs, out_type, target):
+    """threefry_generate generic strategy"""
+    strategy = _op.OpStrategy()
+    strategy.add_implementation(
+        wrap_compute_threefry_generate(topi.random.threefry_generate),
+        wrap_topi_schedule(topi.generic.schedule_extern),
+        name="threefry_generate.generic",
+    )
+    return strategy
+
+
+# threefry_split
+def wrap_compute_threefry_split(topi_compute):
+    """Wrap threefry_split topi compute"""
+
+    def _compute_threefry_split(attrs, inputs, _):
+        return topi_compute(inputs[0])
+
+    return _compute_threefry_split
+
+
+@override_native_generic_func("threefry_split_strategy")
+def threefry_split_strategy(attrs, inputs, out_type, target):
+    """threefry_split generic strategy"""
+    strategy = _op.OpStrategy()
+    strategy.add_implementation(
+        wrap_compute_threefry_split(topi.random.threefry_split),
+        wrap_topi_schedule(topi.generic.schedule_extern),
+        name="threefry_split.generic",
+    )
+    return strategy

python/tvm/topi/__init__.py

@@ -55,6 +55,7 @@
 from . import image
 from . import sparse
 from . import hls
+from . import random
 
 # error reporting
 from .utils import InvalidShapeError

python/tvm/topi/random/__init__.py

@@ -0,0 +1,22 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#   http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied.  See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+# pylint: disable=wildcard-import
+"""Pseudorandom generator kernels and operators."""
+from __future__ import absolute_import
+
+from .kernel import *