[RELAY,TOPI] Threefry PRNG: splittable and stateless

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 argsort(data, axis=-1, is_ascend=1, dtype="int32"):

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, (1,))
+ 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/cuda.py

@@ -17,7 +17,6 @@
 """Definition of CUDA/GPU operator strategy."""
 # pylint: disable=invalid-name,unused-argument,wildcard-import,unused-wildcard-import
 from tvm import topi
-import tvm
 from tvm.auto_scheduler import is_auto_scheduler_enabled
 from tvm.te import SpecializedCondition
 from tvm.contrib import nvcc

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

@@ -1264,3 +1264,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

@@ -54,6 +54,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 *