[RANDOM] Init contrib.random Library (#684)

* [RANDOM] Init contrib.random library

* [RANDOM] Add uniform

* [RANDOM] Fix lint

* [RANDOM] Add comments and tests

* [RANDOM] Fix lint

Makefile

@@ -131,6 +131,7 @@ ifeq ($(USE_GRAPH_RUNTIME), 1)
 endif
 
 include make/contrib/cblas.mk
+include make/contrib/random.mk
 include make/contrib/nnpack.mk
 include make/contrib/cudnn.mk
 include make/contrib/mps.mk

make/config.mk

@@ -62,6 +62,9 @@ USE_GRAPH_RUNTIME = 1
 # Whether use BLAS, choices: openblas, atlas, blas, apple
 USE_BLAS = none
 
+# Whether use contrib.random in runtime
+USE_RANDOM = 0
+
 # Whether use NNPack
 USE_NNPACK = 0
 # NNPACK_PATH = none

make/contrib/random.mk

+RANDOM_CONTRIB_SRC = $(wildcard src/contrib/random/*.cc)
+RANDOM_CONTRIB_OBJ = $(patsubst src/%.cc, build/%.o, $(RANDOM_CONTRIB_SRC))
+
+ifeq ($(USE_RANDOM), 1)
+	RUNTIME_DEP += $(RANDOM_CONTRIB_OBJ)
+endif

python/tvm/api.py

@@ -371,6 +371,8 @@ def extern(shape, inputs, fcompute, name="extern", dtype=None, tag=""):
             raise ValueError("Cannot infer output type, please provide dtype argument")
         infered_type = types.pop()
         dtype = [infered_type for _ in shape]
+    if isinstance(dtype, str):
+        dtype = [dtype]
 
     for shp, dt in zip(shape, dtype):
         output_placeholders.append(decl_buffer(shp, dt, name))

python/tvm/contrib/random.py

+"""External function interface to random library."""
+from __future__ import absolute_import as _abs
+
+from .. import api as _api
+from .. import intrin as _intrin
+from .._ffi.function import _init_api
+
+
+def randint(low, high, size, dtype='int32'):
+    """Return random integers from low (inclusive) to high (exclusive).
+    Return random integers from the "discrete uniform" distribution of the
+    specified dtype in the "half-open" interval [low, high).
+
+    Parameters
+    ----------
+    low : int
+        Lowest (signed) integer to be drawn from the distribution
+    high : int
+        One above the largest (signed) integer to be drawn from the distribution
+
+    Returns
+    -------
+    out : Tensor
+        A tensor with specified size and dtype
+    """
+    assert 'int' in dtype, "the type of randint output must be int or uint"
+    return _api.extern(size, [], lambda ins, outs: _intrin.call_packed(
+        "tvm.contrib.random.randint", int(low), int(high), outs[0]), dtype=dtype)
+
+
+def uniform(low, high, size):
+    """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
+    ----------
+    low : float
+        Lower boundary of the output interval. All values generated will be
+        greater than or equal to low.
+    high : float
+        Upper boundary of the output interval. All values generated will be
+        less than high.
+    size : tuple of ints
+        Output shape. If the given shape is, e.g., (m, n, k), then m * n * k
+        samples are drawn.
+
+    Returns
+    -------
+    out : Tensor
+        A tensor with specified size and dtype.
+    """
+    return _api.extern(size, [], lambda ins, outs: _intrin.call_packed(
+        "tvm.contrib.random.uniform", float(low), float(high), outs[0]), dtype='float32')
+
+_init_api("tvm.contrib.random")

src/contrib/random/random.cc

+/*!
+ *  Copyright (c) 2017 by Contributors
+ * \file External random functions for tensor.
+ */
+#include <tvm/runtime/registry.h>
+#include <tvm/runtime/util.h>
+#include <dmlc/logging.h>
+#include <dmlc/thread_local.h>
+#include <algorithm>
+#include <random>
+#include <ctime>
+
+#define DLPACK_INTEGER_TYPE_SWITCH(type, DType, ...)    \
+  if (type.code == kDLInt && type.bits == 32) {         \
+    typedef int32_t DType;                              \
+    {__VA_ARGS__}                                       \
+  } else if (type.code == kDLInt && type.bits == 16) {  \
+    typedef int16_t DType;                              \
+    {__VA_ARGS__}                                       \
+  } else if (type.code == kDLInt && type.bits == 8) {   \
+    typedef int8_t DType;                               \
+    {__VA_ARGS__}                                       \
+  } else if (type.code == kDLUInt && type.bits == 32) { \
+    typedef uint32_t DType;                             \
+    {__VA_ARGS__}                                       \
+  } else if (type.code == kDLUInt && type.bits == 16) { \
+    typedef uint16_t DType;                             \
+    {__VA_ARGS__}                                       \
+  } else if (type.code == kDLUInt && type.bits == 8) {  \
+    typedef uint8_t DType;                              \
+    {__VA_ARGS__}                                       \
+  } else {                                              \
+    LOG(FATAL) << "unknown data type";                  \
+  }
+
+namespace tvm {
+namespace contrib {
+
+using namespace runtime;
+
+class RandomEngine {
+ public:
+  RandomEngine() {
+    this->Seed(time(0));
+  }
+  explicit RandomEngine(int seed) {
+    this->Seed(seed);
+  }
+
+  ~RandomEngine() {}
+
+  inline void Seed(int seed) {
+    rnd_engine_.seed(seed);
+    this->rseed_ = static_cast<unsigned>(seed);
+  }
+
+  inline unsigned GetSeed() const {
+    return rseed_;
+  }
+
+  inline unsigned GetRandInt() {
+    return rnd_engine_();
+  }
+
+  void SampleUniform(DLTensor* data, float low, float high) {
+    CHECK_GT(high, low) << "high must be bigger than low";
+    CHECK(data->strides == nullptr);
+
+    DLDataType dtype = data->dtype;
+    int64_t size = 1;
+    for (int i = 0; i < data->ndim; ++i) {
+      size *= data->shape[i];
+    }
+
+    CHECK(dtype.code == kDLFloat && dtype.bits == 32 && dtype.lanes == 1);
+
+    if (data->ctx.device_type == kDLCPU) {
+      std::uniform_real_distribution<float> uniform_dist(low, high);
+      std::generate_n(static_cast<float*>(data->data), size, [&] () {
+        return uniform_dist(rnd_engine_);
+      });
+    } else {
+      LOG(FATAL) << "Do not support random.randint on this device yet";
+    }
+  }
+
+ private:
+  std::mt19937 rnd_engine_;
+  unsigned rseed_;
+};
+
+struct RandomThreadLocalEntry {
+  RandomEngine random_engine;
+  static RandomThreadLocalEntry* ThreadLocal();
+};
+
+typedef dmlc::ThreadLocalStore<RandomThreadLocalEntry> RandomThreadLocalStore;
+
+RandomThreadLocalEntry* RandomThreadLocalEntry::ThreadLocal() {
+  return RandomThreadLocalStore::Get();
+}
+
+
+TVM_REGISTER_GLOBAL("tvm.contrib.random.randint")
+.set_body([](TVMArgs args, TVMRetValue *ret) {
+    RandomThreadLocalEntry *entry = RandomThreadLocalEntry::ThreadLocal();
+    int64_t low = args[0];
+    int64_t high = args[1];
+    DLTensor* out = args[2];
+    CHECK_GT(high, low) << "high must be bigger than low";
+    CHECK(out->strides == nullptr);
+
+    DLDataType dtype = out->dtype;
+    int64_t size = 1;
+    for (int i = 0; i < out->ndim; ++i) {
+      size *= out->shape[i];
+    }
+
+    DLPACK_INTEGER_TYPE_SWITCH(dtype, DType, {
+      int64_t numeric_low = std::numeric_limits<DType>::min();
+      int64_t numeric_high = std::numeric_limits<DType>::max();
+      numeric_high += 1;  // exclusive upper bound
+      low = std::max(low, numeric_low);
+      high = std::min(high, numeric_high);
+
+      if (out->ctx.device_type == kDLCPU) {
+          // file the data with random byte
+          std::generate_n(static_cast<DType*>(out->data), size, [&] () {
+            unsigned rint = entry->random_engine.GetRandInt();
+            return low + rint % (high - low);
+          });
+      } else {
+        LOG(FATAL) << "Do not support random.randint on this device yet";
+      }
+    })
+  });
+
+TVM_REGISTER_GLOBAL("tvm.contrib.random.uniform")
+.set_body([](TVMArgs args, TVMRetValue *ret) {
+    RandomThreadLocalEntry *entry = RandomThreadLocalEntry::ThreadLocal();
+    double low = args[0];
+    double high = args[1];
+    DLTensor* out = args[2];
+    entry->random_engine.SampleUniform(out, low, high);
+  });
+
+
+}  // namespace contrib
+}  // namespace tvm

tests/python/contrib/test_random.py

+import tvm
+import numpy as np
+from tvm.contrib import random
+
+def test_randint():
+    m = 1024
+    n = 1024
+    A = random.randint(-127, 128, size=(m, n), dtype='int32')
+    s = tvm.create_schedule(A.op)
+
+    def verify(target="llvm"):
+        if not tvm.module.enabled(target):
+            print("skip because %s is not enabled..." % target)
+            return
+        if not tvm.get_global_func("tvm.contrib.random.randint", True):
+            print("skip because extern function is not avalable")
+            return
+        ctx = tvm.cpu(0)
+        f = tvm.build(s, [A], target)
+        a = tvm.nd.array(np.zeros((m, n), dtype=A.dtype), ctx)
+        f(a)
+        na = a.asnumpy()
+        assert abs(np.mean(na)) < 0.2
+        assert np.min(na) == -127
+        assert np.max(na) == 127
+    verify()
+
+
+def test_uniform():
+    m = 1024
+    n = 1024
+    A = random.uniform(0, 1, size=(m, n))
+    s = tvm.create_schedule(A.op)
+
+    def verify(target="llvm"):
+        if not tvm.module.enabled(target):
+            print("skip because %s is not enabled..." % target)
+            return
+        if not tvm.get_global_func("tvm.contrib.random.uniform", True):
+            print("skip because extern function is not avalable")
+            return
+        ctx = tvm.cpu(0)
+        f = tvm.build(s, [A], target)
+        a = tvm.nd.array(np.zeros((m, n), dtype=A.dtype), ctx)
+        f(a)
+        na = a.asnumpy()
+        assert abs(np.mean(na) - 0.5) < 1e-2
+        assert abs(np.min(na) - 0.0) < 1e-3
+        assert abs(np.max(na) - 1.0) < 1e-3
+    verify()
+
+
+if __name__ == "__main__":
+    test_randint()
+    test_uniform()