Implementing cunumeric.random.BitGenerator

cunumeric/config.py

@@ -82,6 +82,7 @@ class CuNumericOpCode(IntEnum):
     BINARY_OP = _cunumeric.CUNUMERIC_BINARY_OP
     BINARY_RED = _cunumeric.CUNUMERIC_BINARY_RED
     BINCOUNT = _cunumeric.CUNUMERIC_BINCOUNT
+    BITGENERATOR = _cunumeric.CUNUMERIC_BITGENERATOR
     CHOOSE = _cunumeric.CUNUMERIC_CHOOSE
     CONTRACT = _cunumeric.CUNUMERIC_CONTRACT
     CONVERT = _cunumeric.CUNUMERIC_CONVERT

cunumeric/deferred.py

@@ -38,6 +38,8 @@
 from .thunk import NumPyThunk
 from .utils import get_arg_value_dtype
 
+# from defusedxml import NotSupportedError
+
 
 def _complex_field_dtype(dtype):
     if dtype == np.complex64:
@@ -1326,6 +1328,22 @@ def nonzero(self):
         task.execute()
         return results
 
+    def bitgenerator_random_raw(self, handle):
+        task = self.context.create_task(CuNumericOpCode.BITGENERATOR)
+        task.add_output(self.base)
+        task.add_scalar_arg(3, ty.int32)  # OP_RAND_RAW
+        task.add_scalar_arg(handle, ty.uint32)
+        # TODO: check if no function does this
+        totalsize = 1
+        for sz in self.shape:
+            totalsize = totalsize * sz
+        task.add_scalar_arg(totalsize, ty.uint64)
+        # strides
+        task.add_scalar_arg(self.compute_strides(self.shape), (ty.int64,))
+
+        task.add_broadcast(self.base, axes=tuple(range(1, self.ndim)))
+        task.execute()
+
     def random(self, gen_code, args=[]):
         task = self.context.create_task(CuNumericOpCode.RAND)
 

cunumeric/eager.py

@@ -564,6 +564,20 @@ def sort(self, rhs, argsort=False, axis=-1, kind="quicksort", order=None):
             else:
                 self.array = np.sort(rhs.array, axis, kind, order)
 
+    def bitgenerator_random_raw(self, handle):
+        if self.deferred is not None:
+            print("eager.py - bitgenerator_random_raw - deferred is not None")
+            self.deferred.bitgenerator_random_raw(handle)
+        else:
+            print("eager.py - bitgenerator_random_raw - deferred is None")
+            if self.array.size == 1:
+                self.array.fill(np.random.rand())
+            else:
+                a = np.random.randint(
+                    0, 2**32 - 1, *(self.array.shape), dtype=self.array.dtype
+                )
+                self.array[:] = a[:]
+
     def random_uniform(self):
         if self.deferred is not None:
             self.deferred.random_uniform()

cunumeric/random/BitGenerator.py

@@ -0,0 +1,87 @@
+# Copyright 2021-2022 NVIDIA Corporation
+#
+# Licensed 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.
+#
+
+import numpy as np
+from cunumeric.array import ndarray
+from cunumeric.runtime import runtime
+
+
+class BitGenerator:
+    # see bitgenerator_util.h
+    OP_CREATE = 1
+    OP_DESTROY = 2
+    OP_RAND_RAW = 3
+
+    # see bitgenerator_util.h
+    DEFAULT = 0
+    XORWOW = 1
+    MRG32K3A = 2
+    MTGP32 = 3
+    MT19937 = 4
+    PHILOX4_32_10 = 5
+
+    __slots__ = [
+        "handle",  # handle to the runtime id
+    ]
+
+    def __init__(self, seed=None, generatorType=DEFAULT):
+        if type(self) is BitGenerator:
+            raise NotImplementedError(
+                "BitGenerator is a base class and cannot be instantized"
+            )
+        self.handle = runtime.bitgenerator_create(generatorType)
+        if seed is not None:
+            runtime.bitgenerator_set_seed(self.handle, seed)
+
+    def __del__(self):
+        runtime.bitgenerator_destroy(self.handle)
+
+    # when output is false => skip ahead
+    def random_raw(self, shape=None, output=True):
+        if shape is None:
+            raise NotImplementedError("Empty shape not implemented")
+        if not isinstance(shape, tuple):
+            shape = (shape,)
+        if output:
+            res = ndarray(shape, dtype=np.dtype(np.uint32))
+            res._thunk.bitgenerator_random_raw(self.handle)
+            return res
+        else:
+            runtime.bitgenerator_random_raw(self.handle, shape)
+
+
+class XORWOW(BitGenerator):
+    def __init__(self, seed=None):
+        super().__init__(seed, BitGenerator.XORWOW)
+
+
+class MRG32k3a(BitGenerator):
+    def __init__(self, seed=None):
+        super().__init__(seed, BitGenerator.MRG32K3A)
+
+
+class MTGP32(BitGenerator):
+    def __init__(self, seed=None):
+        super().__init__(seed, BitGenerator.MTGP32)
+
+
+class MT19937(BitGenerator):
+    def __init__(self, seed=None):
+        super().__init__(seed, BitGenerator.MT19937)
+
+
+class PHILOX4_32_10(BitGenerator):
+    def __init__(self, seed=None):
+        super().__init__(seed, BitGenerator.PHILOX4_32_10)

cunumeric/random/__init__.py

@@ -16,6 +16,7 @@
 import numpy.random as _nprandom
 from cunumeric.random.random import *
 from cunumeric.coverage import clone_module
+from cunumeric.random.BitGenerator import *
 
 clone_module(_nprandom, globals())
 

cunumeric/runtime.py

@@ -61,6 +61,7 @@ def __init__(self, legate_context):
         self.legate_context = legate_context
         self.legate_runtime = get_legate_runtime()
         self.current_random_epoch = 0
+        self.current_random_bitgenid = 0
         self.destroyed = False
         self.api_calls = []
 
@@ -218,6 +219,62 @@ def create_scalar(self, array: memoryview, dtype, shape=None, wrap=False):
             result = future
         return result
 
+    def bitgenerator_create(self, generatorType):
+        task = self.legate_context.create_task(
+            CuNumericOpCode.BITGENERATOR,
+            manual=True,
+            launch_domain=Rect(lo=(0,), hi=(self.num_procs,)),
+        )
+        self.current_random_bitgenid = self.current_random_bitgenid + 1
+        task.add_scalar_arg(1, ty.int32)  # OP_CREATE
+        task.add_scalar_arg(self.current_random_bitgenid, ty.uint32)
+        task.add_scalar_arg(generatorType, ty.uint64)
+        task.execute()
+        self.legate_runtime.issue_execution_fence()
+        return self.current_random_bitgenid
+
+    def bitgenerator_destroy(self, handle):
+        self.legate_runtime.issue_execution_fence()
+        task = self.legate_context.create_task(
+            CuNumericOpCode.BITGENERATOR,
+            manual=True,
+            launch_domain=Rect(lo=(0,), hi=(self.num_procs,)),
+        )
+        task.add_scalar_arg(2, ty.int32)  # OP_DESTROY
+        task.add_scalar_arg(handle, ty.uint32)
+        task.add_scalar_arg(0, ty.uint64)
+        task.execute()
+
+    def bitgenerator_set_seed(self, handle, seed):
+        if not isinstance(seed, int):
+            raise NotImplementedError("Non integer seed is not implemented")
+        task = self.legate_context.create_task(
+            CuNumericOpCode.BITGENERATOR,
+            manual=True,
+            launch_domain=Rect(lo=(0,), hi=(self.num_procs,)),
+        )
+        task.add_scalar_arg(4, ty.int32)  # OP_SET_SEED
+        task.add_scalar_arg(handle, ty.uint32)
+        task.add_scalar_arg(seed, ty.uint64)
+        task.execute()
+
+    def bitgenerator_random_raw(self, handle, size):
+        # here, no output: we discard generated numbers... - just a skipahead
+        task = self.legate_context.create_task(
+            CuNumericOpCode.BITGENERATOR,
+            manual=True,
+            launch_domain=Rect(lo=(0,), hi=(self.num_procs,)),
+        )
+        task.add_scalar_arg(3, ty.int32)  # OP_RAND_RAW
+        task.add_scalar_arg(handle, ty.uint32)
+        gencount = 1
+        for sz in size:
+            gencount = gencount * sz
+        task.add_scalar_arg(gencount, ty.uint64)  # size of the output
+        task.execute()
+        # for consistent random ordering
+        self.legate_runtime.issue_execution_fence()
+
     def set_next_random_epoch(self, epoch):
         self.current_random_epoch = epoch
 

cunumeric/thunk.py

@@ -189,6 +189,10 @@ def bincount(self, rhs, weights=None):
     def nonzero(self):
         ...
 
+    @abstractmethod
+    def bitgenerator_random_raw(self, bitgen):
+        ...
+
     @abstractmethod
     def random_uniform(self):
         ...

src/Makefile

@@ -50,7 +50,7 @@ LD_FLAGS += -L$(OPENBLAS_PATH)/lib -l$(OPENBLAS_LIBNAME) -Wl,-rpath,$(OPENBLAS_P
 LD_FLAGS += -L$(TBLIS_PATH)/lib -ltblis -Wl,-rpath,$(TBLIS_PATH)/lib
 ifeq ($(strip $(USE_CUDA)),1)
 DEVICE_LD_FLAGS += -lcufft_static
-LD_FLAGS += -lcublas -lcusolver -lcufft_static -lculibos
+LD_FLAGS += -lcublas -lcusolver -lcufft_static -lculibos -lcurand
 LD_FLAGS += -L$(CUTENSOR_PATH)/lib -lcutensor -Wl,-rpath,$(CUTENSOR_PATH)/lib
 LD_FLAGS += -L$(NCCL_PATH)/lib -lnccl -Wl,-rpath,$(NCCL_PATH)/lib
 endif

src/cunumeric.mk

@@ -43,6 +43,7 @@ GEN_CPU_SRC += cunumeric/ternary/where.cc               \
 							 cunumeric/matrix/trsm.cc                 \
 							 cunumeric/matrix/util.cc                 \
 							 cunumeric/random/rand.cc                 \
+							 cunumeric/random/bitgenerator.cc         \
 							 cunumeric/search/nonzero.cc              \
 							 cunumeric/set/unique.cc                  \
 							 cunumeric/set/unique_reduce.cc           \
@@ -80,6 +81,7 @@ GEN_CPU_SRC += cunumeric/ternary/where_omp.cc          \
 							 cunumeric/matrix/trsm_omp.cc            \
 							 cunumeric/matrix/util_omp.cc            \
 							 cunumeric/random/rand_omp.cc            \
+							 cunumeric/random/bitgenerator_omp.cc    \
 							 cunumeric/search/nonzero_omp.cc         \
 							 cunumeric/set/unique_omp.cc             \
 							 cunumeric/sort/sort_omp.cc              \
@@ -119,6 +121,7 @@ GEN_GPU_SRC += cunumeric/ternary/where.cu               \
 							 cunumeric/matrix/trilu.cu                \
 							 cunumeric/matrix/trsm.cu                 \
 							 cunumeric/random/rand.cu                 \
+							 cunumeric/random/bitgenerator.cu         \
 							 cunumeric/search/nonzero.cu              \
 							 cunumeric/set/unique.cu                  \
 							 cunumeric/sort/sort.cu                   \

src/cunumeric/cunumeric_c.h

@@ -27,6 +27,7 @@ enum CuNumericOpCode {
   CUNUMERIC_BINARY_OP,
   CUNUMERIC_BINARY_RED,
   CUNUMERIC_BINCOUNT,
+  CUNUMERIC_BITGENERATOR,
   CUNUMERIC_CHOOSE,
   CUNUMERIC_CONTRACT,
   CUNUMERIC_CONVERT,

src/cunumeric/mapper.cc

@@ -160,6 +160,16 @@ std::vector<StoreMapping> CuNumericMapper::store_mappings(
       mappings.back().policy.exact = true;
       return std::move(mappings);
     }
+    case CUNUMERIC_BITGENERATOR: {
+      std::vector<StoreMapping> mappings;
+      auto& outputs = task.outputs();
+      for (auto& output : outputs) {
+        mappings.push_back(StoreMapping::default_mapping(output, options.front()));
+        mappings.back().policy.ordering.c_order();
+        mappings.back().policy.exact = true;
+      }
+      return std::move(mappings);
+    }
     default: {
       return {};
     }

src/cunumeric/random/bitgenerator.cc

@@ -0,0 +1,78 @@
+/* Copyright 2021-2022 NVIDIA Corporation
+ *
+ * Licensed 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.
+ *
+ */
+
+#include "cunumeric/random/bitgenerator.h"
+#include "cunumeric/random/bitgenerator_template.inl"
+#include "cunumeric/random/bitgenerator_util.h"
+
+#include "cunumeric/random/curand_help.h"
+
+#include "cunumeric/random/bitgenerator_curand.inl"
+
+namespace cunumeric {
+
+using namespace Legion;
+using namespace legate;
+
+template <>
+struct CURANDGeneratorBuilder<VariantKind::CPU> {
+  static CURANDGenerator* build(BitGeneratorType gentype)
+  {
+    curandGenerator_t gen;
+    CHECK_CURAND(::curandCreateGeneratorHost(&gen, get_curandRngType(gentype)));
+    CURANDGenerator* cugenptr = new CURANDGenerator();
+    CURANDGenerator& cugen    = *cugenptr;
+    cugen.gen                 = gen;
+    cugen.offset              = 0;
+    cugen.type                = get_curandRngType(gentype);
+    cugen.supports_skipahead  = supportsSkipAhead(cugen.type);
+    cugen.dev_buffer_size     = cugen.DEFAULT_DEV_BUFFER_SIZE;
+    cugen.dev_buffer          = (uint32_t*)::malloc(cugen.dev_buffer_size * sizeof(uint32_t));
+    return cugenptr;
+  }
+
+  static void destroy(CURANDGenerator* cugenptr)
+  {
+    // wait for rand jobs and clean-up resources
+    std::lock_guard<std::mutex> guard(cugenptr->lock);
+    ::free(cugenptr->dev_buffer);
+    CHECK_CURAND(::curandDestroyGenerator(cugenptr->gen));
+  }
+};
+
+template <>
+std::map<Legion::Processor, std::unique_ptr<generatormap<VariantKind::CPU>>>
+  BitGeneratorImplBody<VariantKind::CPU>::m_generators = {};
+
+template <>
+std::mutex BitGeneratorImplBody<VariantKind::CPU>::lock_generators = {};
+
+/*static*/ void BitGeneratorTask::cpu_variant(TaskContext& context)
+{
+  bitgenerator_template<VariantKind::CPU>(context);
+}
+
+namespace  // unnamed
+{
+
+static void __attribute__((constructor)) register_tasks(void)
+{
+  BitGeneratorTask::register_variants();
+}
+
+}  // namespace
+
+}  // namespace cunumeric