diff --git a/include/tvm/relay/attrs/random.h b/include/tvm/relay/attrs/random.h index 8238f102dab8..46cab8831caf 100644 --- a/include/tvm/relay/attrs/random.h +++ b/include/tvm/relay/attrs/random.h @@ -37,6 +37,18 @@ struct ThreefryGenerateAttrs : public tvm::AttrsNode { } }; +struct UniformAttrs : public tvm::AttrsNode { + Array out_shape; + DataType out_dtype; + + TVM_DECLARE_ATTRS(UniformAttrs, "relay.attrs.UniformAttrs") { + TVM_ATTR_FIELD(out_shape).describe("Shape of random numbers to generate"); + TVM_ATTR_FIELD(out_dtype) + .set_default(NullValue()) + .describe("Data type of the generated numbers"); + } +}; + } // namespace relay } // namespace tvm #endif // TVM_RELAY_ATTRS_RANDOM_H_ diff --git a/python/tvm/relay/op/op_attrs.py b/python/tvm/relay/op/op_attrs.py index 4cc6e0f26b91..b96e8b7a66b5 100644 --- a/python/tvm/relay/op/op_attrs.py +++ b/python/tvm/relay/op/op_attrs.py @@ -562,3 +562,8 @@ class BatchToSpaceNDAttrs(Attrs): @tvm._ffi.register_object("relay.attrs.ThreefryGenerateAttrs") class ThreefryGenerateAttrs(Attrs): """Attributes used in ThreefryGenerateAttrs operators""" + + +@tvm._ffi.register_object("relay.attrs.UniformAttrs") +class UniformAttrs(Attrs): + """Attributes used in UniformAttrs operators""" diff --git a/python/tvm/relay/op/random/_kernel.py b/python/tvm/relay/op/random/_kernel.py index 8be3397008d5..e601a7073cff 100644 --- a/python/tvm/relay/op/random/_kernel.py +++ b/python/tvm/relay/op/random/_kernel.py @@ -27,3 +27,7 @@ register_pattern("random.threefry_generate", OpPattern.OPAQUE) register_strategy("random.threefry_split", strategy.threefry_split_strategy) register_pattern("random.threefry_split", OpPattern.OPAQUE) + +# Distribution +register_strategy("random.uniform", strategy.uniform_strategy) +register_pattern("random.uniform", OpPattern.OPAQUE) diff --git a/python/tvm/relay/op/random/kernel.py b/python/tvm/relay/op/random/kernel.py index 96634943128d..fc1248e85678 100644 --- a/python/tvm/relay/op/random/kernel.py +++ b/python/tvm/relay/op/random/kernel.py @@ -21,7 +21,7 @@ import sys import numpy as np -from ...expr import Constant +from ...expr import Constant, Expr, const from .... import nd from . import _make @@ -132,3 +132,55 @@ def foo(key): :py:func:`threefry_generate`. """ return _make.threefry_split(key) + + +def uniform(key, shape, dtype="float32", low=0.0, high=1.0): + """Draw samples from a uniform distribution. + + Samples are uniformly distributed over the half-open interval [low, high) + (includes low, but excludes high). In other words, any value within the + given interval is equally likely to be drawn by uniform. + + Example + ------- + + .. code-block:: python + + key = threefry_key(0) + key, random_values = uniform(key, (100,), low=0, high=10) + + 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.** + + shape : Sequence[int] + Desired outputs shape of random numbers. + + dtype : str + Desired outputs type of random numbers. + + low : float or relay.Expr, optional + Lower bound of the uniform distribution. + + high : float or relay.Expr, optional + Upper bound of the uniform distribution. + + Returns + ------- + new_key : relay.Expr + New random key to pass to future uses of random functions. + + random_values : relay.Expr + The generated uniform distributed random numbers. + """ + if not isinstance(low, Expr): + low = const(low, dtype=dtype) + if not isinstance(high, Expr): + high = const(high, dtype=dtype) + return _make.uniform(key, low, high, shape, dtype) diff --git a/python/tvm/relay/op/strategy/generic.py b/python/tvm/relay/op/strategy/generic.py index 7451b397265f..97d1a91f744f 100644 --- a/python/tvm/relay/op/strategy/generic.py +++ b/python/tvm/relay/op/strategy/generic.py @@ -1495,6 +1495,28 @@ def threefry_split_strategy(attrs, inputs, out_type, target): return strategy +# uniform +def wrap_compute_uniform(topi_compute): + """Wrap uniform topi compute""" + + def _compute_uniform(attrs, inputs, _): + return list(topi_compute(inputs[0], inputs[1], inputs[2], attrs.out_shape, attrs.out_dtype)) + + return _compute_uniform + + +@override_native_generic_func("uniform_strategy") +def uniform_strategy(attrs, inputs, out_type, target): + """uniform generic strategy""" + strategy = _op.OpStrategy() + strategy.add_implementation( + wrap_compute_uniform(topi.random.uniform), + wrap_topi_schedule(topi.generic.schedule_extern), + name="uniform.generic", + ) + return strategy + + def wrap_compute_scanop(topi_compute): """Wrap scanop style topi compute""" diff --git a/python/tvm/topi/random/kernel.py b/python/tvm/topi/random/kernel.py index b6c0b3fa5930..2831baba30cd 100644 --- a/python/tvm/topi/random/kernel.py +++ b/python/tvm/topi/random/kernel.py @@ -466,3 +466,69 @@ def gen_ir(out_ptr): out_ary = tvm.nd.array(np.ones((1,), "uint64"), device) tvm.build(s, [f], target=target)(out_ary) return out_ary.asnumpy()[0] == 0 + + +def uniform(gen, low, high, out_shape, out_dtype): + """Draw samples from a uniform distribution. + + Samples are uniformly distributed over the half-open interval [low, high) + (includes low, but excludes high). In other words, any value within the + given interval is equally likely to be drawn by uniform. + + Parameters + ---------- + gen : ThreefryKey + Generator state. Can be create with :py:func:`tvm.relay.threefry_key`. This should not be + reused in another function, otherwise random numbers will be repeated. + + low : Tensor[(), out_dtype] + Lower boundary of the output interval. All values generated will be + greater than or equal to low. + + high : Tensor[(), out_dtype] + Upper boundary of the output interval. All values generated will be + less than high. + + out_shape : Sequence[int] + Output shape of the random numbers. Product of all dimensions must be a multiple of 4. + + out_dtype : str + The output dtype. + + Returns + ------- + new_gen : ThreefryKey + New generator state that is distinct from `gen`. + + out : Tensor[out_shape, out_dtype] + Tensor of random numbers with shape `out_shape` and type `out_dtype`. + """ + new_gen, random_bits = threefry_generate(gen, out_shape) + assert out_dtype in ("float32", "float64"), ( + "Only support float32 or float64 for now, got %s" % out_dtype + ) + if out_dtype == "float32": + random_dtype = "uint32" + nbits = 32 + nfraction = 23 + elif out_dtype == "float64": + random_dtype = "uint64" + nbits = 64 + nfraction = 52 + nexp = nbits - nfraction - 1 + random_bits = random_bits.astype(random_dtype) + + fraction = tvm.topi.right_shift( + random_bits, tvm.tir.const(nbits - nfraction, dtype=random_dtype) + ) + exponent = tvm.topi.left_shift( + tvm.topi.full(out_shape, random_dtype, (1 << (nexp - 1)) - 1), + tvm.tir.const(nfraction, dtype=random_dtype), + ) + mantissa = tvm.topi.bitwise_or(fraction, exponent).astype(random_dtype) + standard_uniform_values = tvm.topi.reinterpret(mantissa, out_dtype) - tvm.tir.const( + 1, dtype=out_dtype + ) + uniform_values = tvm.topi.add(tvm.topi.multiply(standard_uniform_values, high - low), low) + + return new_gen, uniform_values diff --git a/src/relay/op/random/kernel.cc b/src/relay/op/random/kernel.cc index ec092a7e05f2..077a1ea11673 100644 --- a/src/relay/op/random/kernel.cc +++ b/src/relay/op/random/kernel.cc @@ -85,5 +85,52 @@ RELAY_REGISTER_OP("random.threefry_split") .add_argument("key", "Tensor", "Input Threefry key") .add_type_rel("ThreefrySplit", ThreefrySplitRel); +TVM_REGISTER_NODE_TYPE(UniformAttrs); + +bool UniformRel(const Array& types, int num_inputs, const Attrs& attrs, + const TypeReporter& reporter) { + const UniformAttrs* param = attrs.as(); + ICHECK_EQ(types.size(), 4) << "Uniform should have three inputs and one output"; + + std::vector oshape; + for (auto& x : param->out_shape) { + oshape.push_back(x); + } + DataType out_dtype = param->out_dtype; + // we are supporting float32 and float64 at the moment. + if (!(out_dtype.is_float() && (out_dtype.bits() == 32 || out_dtype.bits() == 64))) { + reporter->GetDiagCtx().EmitFatal(Diagnostic::Error(reporter->GetSpan()) + << "We only support generating uniform random value of " + << "type float32 or float64, got " << out_dtype << "."); + return false; + } + reporter->Assign(types[0], ThreefryKeyType()); + reporter->Assign(types[1], TensorType({}, out_dtype)); + reporter->Assign(types[2], TensorType({}, out_dtype)); + // generate returns the next key and an array of random values + reporter->Assign(types[3], TupleType({ThreefryKeyType(), TensorType(oshape, out_dtype)})); + return true; +} + +Expr MakeUniform(Expr key, Expr low, Expr high, Array out_shape, DataType out_dtype) { + auto attrs = make_object(); + attrs->out_shape = out_shape; + attrs->out_dtype = out_dtype; + static const Op& op = Op::Get("random.uniform"); + return Call(op, {key, low, high}, Attrs(attrs), {}); +} + +TVM_REGISTER_GLOBAL("relay.op.random._make.uniform").set_body_typed(MakeUniform); + +RELAY_REGISTER_OP("random.uniform") + .describe( + R"doc(Generate an array of random numbers under uniform distribution.)doc" TVM_ADD_FILELINE) + .set_num_inputs(3) + .set_attrs_type() + .add_argument("key", "Tensor", "Input Threefry key") + .add_argument("low", "Tensor", "Lower bound of the distribution") + .add_argument("high", "Tensor", "Higher bound of the distribution") + .add_type_rel("Uniform", UniformRel); + } // namespace relay } // namespace tvm diff --git a/tests/python/relay/test_prng.py b/tests/python/relay/test_prng.py index ba4fdc466ecc..f622bfca5625 100644 --- a/tests/python/relay/test_prng.py +++ b/tests/python/relay/test_prng.py @@ -103,6 +103,21 @@ def test_threefry_split_infer(): assert tvm.ir.structural_equal(f.ret_type, expected_type) +def test_uniform_infer(): + oshape = (12,) + odtypes = ["float32", "float64"] + for odtype in odtypes: + key_type = tvm.relay.TensorType([10], dtype="uint64") + gen_type = tvm.relay.TensorType(oshape, dtype=odtype) + expected_type = tvm.relay.TupleType([key_type, gen_type]) + + key = tvm.relay.random.threefry_key(1) + rand1 = tvm.relay.random.uniform(key, oshape, odtype) + f = tvm.relay.Function([], rand1) + f = run_infer_type(f) + assert tvm.ir.structural_equal(f.ret_type, expected_type) + + @pytest.mark.xfail(raises=tvm.error.TVMError) def test_threefry_generate_infer_fail(): # xfail: key size should be 10 diff --git a/tests/python/topi/python/test_topi_prng.py b/tests/python/topi/python/test_topi_prng.py index 4ad3a80c6a9e..9189fa792898 100644 --- a/tests/python/topi/python/test_topi_prng.py +++ b/tests/python/topi/python/test_topi_prng.py @@ -43,6 +43,23 @@ def threefry_generate(target, dev, gen, size): return out_gen.asnumpy(), rands.asnumpy() +def uniform(target, dev, gen, low, high, size, dtype): + gen_placeholder = tvm.te.placeholder(gen.shape, name="gen", dtype="uint64") + low_placeholder = tvm.te.placeholder(low.shape, name="low", dtype=dtype) + high_placeholder = tvm.te.placeholder(high.shape, name="high", dtype=dtype) + left_placeholder, right_placeholder = tvm.topi.random.uniform( + gen_placeholder, low_placeholder, high_placeholder, size, dtype + ) + s = tvm.topi.generic.schedule_extern([left_placeholder, right_placeholder]) + f = tvm.build( + s, [gen_placeholder, low_placeholder, high_placeholder, left_placeholder, right_placeholder] + ) + out_gen = tvm.nd.array(np.zeros(gen.shape, dtype="uint64")) + rands = tvm.nd.array(np.zeros(size, dtype=dtype)) + f(tvm.nd.array(gen), tvm.nd.array(low), tvm.nd.array(high), out_gen, rands) + return out_gen.asnumpy(), rands.asnumpy() + + @tvm.testing.parametrize_targets def test_threefry_split(target, dev): # test that results of split do not equal eachother or the input @@ -118,7 +135,24 @@ def test_threefry_wrapping(target, dev): ), f"{target} does not suppport wrapping unsigned integer arithmetic" +@tvm.testing.parametrize_targets +def test_uniform(target, dev): + gen = tvm.relay.random.threefry_key(0).data.asnumpy() + m = 1024 + n = 1024 + dtypes = ["float32", "float64"] + for dtype in dtypes: + low = np.array(5.0, dtype=dtype) + high = np.array(10.0, dtype=dtype) + new_gen, rands = uniform(target, dev, gen, low, high, (m, n), dtype) + assert (gen != new_gen).any() + assert abs(np.mean(rands) - 7.5) < 1e-1 + assert np.min(rands) >= 5.0 + assert np.max(rands) <= 10.0 + + if __name__ == "__main__": test_threefry_split(tvm.target.Target("llvm"), tvm.device("cpu")) test_threefry_generate(tvm.target.Target("llvm"), tvm.device("cpu")) test_threefry_wrapping(tvm.target.Target("llvm"), tvm.device("cpu")) + test_uniform(tvm.target.Target("llvm"), tvm.device("cpu"))