[Mlir-commits] [mlir] 3032c07 - [mlir][crunner] Add support for random number generation.

[llvmlistbot at llvm.org]

Author: bixia1
Date: 2022-12-06T08:54:00-08:00
New Revision: 3032c07d3acb049d0be6ffd0ee5cae9f85771d84

URL: https://github.com/llvm/llvm-project/commit/3032c07d3acb049d0be6ffd0ee5cae9f85771d84
DIFF: https://github.com/llvm/llvm-project/commit/3032c07d3acb049d0be6ffd0ee5cae9f85771d84.diff

LOG: [mlir][crunner] Add support for random number generation.

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D139374

Added: 
    

Modified: 
    mlir/include/mlir/ExecutionEngine/CRunnerUtils.h
    mlir/lib/ExecutionEngine/CRunnerUtils.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/ExecutionEngine/CRunnerUtils.h b/mlir/include/mlir/ExecutionEngine/CRunnerUtils.h
index f76e8c0eb0db4..e145aaf04c22e 100644
--- a/mlir/include/mlir/ExecutionEngine/CRunnerUtils.h
+++ b/mlir/include/mlir/ExecutionEngine/CRunnerUtils.h
@@ -476,4 +476,14 @@ extern "C" MLIR_CRUNNERUTILS_EXPORT void printNewline();
 extern "C" MLIR_CRUNNERUTILS_EXPORT void printFlops(double flops);
 extern "C" MLIR_CRUNNERUTILS_EXPORT double rtclock();
 
+//===----------------------------------------------------------------------===//
+// Runtime support library for random number generation.
+//===----------------------------------------------------------------------===//
+// Uses a seed to initialize a random generator and returns the generator.
+extern "C" MLIR_CRUNNERUTILS_EXPORT void *rtsrand(uint64_t s);
+// Returns a random number in the range of [0, m).
+extern "C" MLIR_CRUNNERUTILS_EXPORT uint64_t rtrand(void *, uint64_t m);
+// Deletes the random number generator.
+extern "C" MLIR_CRUNNERUTILS_EXPORT void rtdrand(void *);
+
 #endif // MLIR_EXECUTIONENGINE_CRUNNERUTILS_H

diff  --git a/mlir/lib/ExecutionEngine/CRunnerUtils.cpp b/mlir/lib/ExecutionEngine/CRunnerUtils.cpp
index e5e0d6f94f147..d9486f95b627c 100644
--- a/mlir/lib/ExecutionEngine/CRunnerUtils.cpp
+++ b/mlir/lib/ExecutionEngine/CRunnerUtils.cpp
@@ -29,6 +29,7 @@
 #include <cinttypes>
 #include <cstdio>
 #include <cstdlib>
+#include <random>
 #include <string.h>
 
 #ifdef MLIR_CRUNNERUTILS_DEFINE_FUNCTIONS
@@ -148,4 +149,20 @@ extern "C" void mlirAlignedFree(void *ptr) {
 #endif
 }
 
+extern "C" void *rtsrand(uint64_t s) {
+  // Standard mersenne_twister_engine seeded with s.
+  return new std::mt19937(s);
+}
+
+extern "C" uint64_t rtrand(void *g, uint64_t m) {
+  std::mt19937 *generator = static_cast<std::mt19937 *>(g);
+  std::uniform_int_distribution<uint64_t> distrib(0, m);
+  return distrib(*generator);
+}
+
+extern "C" void rtdrand(void *g) {
+  std::mt19937 *generator = static_cast<std::mt19937 *>(g);
+  delete generator;
+}
+
 #endif // MLIR_CRUNNERUTILS_DEFINE_FUNCTIONS