https://github.com/yinying-lisa-li updated
https://github.com/llvm/llvm-project/pull/77124
>From 1c774e6c6ae3c5c7be9291677651d20c8979c7f5 Mon Sep 17 00:00:00 2001
From: Yinying Li
Date: Fri, 5 Jan 2024 01:17:39 +
Subject: [PATCH 1/6] [mlir][sparse][CRunnerUtils] Add shuffle and shuffleFree
in CRunnerUtils to generate unique and random numbers
It's helpful for generating tensor with specified sparsity level.
---
.../mlir/ExecutionEngine/CRunnerUtils.h | 4 +
mlir/lib/ExecutionEngine/CRunnerUtils.cpp | 16 +++
.../SparseTensor/CPU/sparse_generate.mlir | 108 ++
3 files changed, 128 insertions(+)
create mode 100644
mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_generate.mlir
diff --git a/mlir/include/mlir/ExecutionEngine/CRunnerUtils.h
b/mlir/include/mlir/ExecutionEngine/CRunnerUtils.h
index 76b04145b482e4..747e5ca40ca6f6 100644
--- a/mlir/include/mlir/ExecutionEngine/CRunnerUtils.h
+++ b/mlir/include/mlir/ExecutionEngine/CRunnerUtils.h
@@ -486,6 +486,10 @@ extern "C" MLIR_CRUNNERUTILS_EXPORT void *rtsrand(uint64_t
s);
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 *);
+// Returns a pointer to an array of random numbers in the range of [0, s).
+extern "C" MLIR_CRUNNERUTILS_EXPORT void *shuffle(uint64_t s, void *g);
+// Deletes the array of random numbers.
+extern "C" MLIR_CRUNNERUTILS_EXPORT void shuffleFree(void *a);
//===--===//
// Runtime support library to allow the use of std::sort in MLIR program.
diff --git a/mlir/lib/ExecutionEngine/CRunnerUtils.cpp
b/mlir/lib/ExecutionEngine/CRunnerUtils.cpp
index e28e75eb110303..3a3261d1ad4e03 100644
--- a/mlir/lib/ExecutionEngine/CRunnerUtils.cpp
+++ b/mlir/lib/ExecutionEngine/CRunnerUtils.cpp
@@ -160,6 +160,22 @@ extern "C" void mlirAlignedFree(void *ptr) {
#endif
}
+/// Generates an array with unique and random numbers from 0 to s-1.
+extern "C" void *shuffle(uint64_t s, void *g) {
+ std::mt19937 *generator = static_cast(g);
+ uint64_t *output = new uint64_t[s];
+ std::vector arr(s);
+ std::iota(arr.begin(), arr.end(), 0);
+ std::shuffle(arr.begin(), arr.end(), *generator);
+ std::copy(arr.begin(), arr.end(), output);
+ return output;
+}
+
+extern "C" void shuffleFree(void *a) {
+ uint64_t *arr = static_cast(a);
+ delete[] arr;
+}
+
extern "C" void *rtsrand(uint64_t s) {
// Standard mersenne_twister_engine seeded with s.
return new std::mt19937(s);
diff --git
a/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_generate.mlir
b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_generate.mlir
new file mode 100644
index 00..250993d874b370
--- /dev/null
+++ b/mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_generate.mlir
@@ -0,0 +1,108 @@
+//--
+// WHEN CREATING A NEW TEST, PLEASE JUST COPY & PASTE WITHOUT EDITS.
+//
+// Set-up that's shared across all tests in this directory. In principle, this
+// config could be moved to lit.local.cfg. However, there are downstream users
that
+// do not use these LIT config files. Hence why this is kept inline.
+//
+// DEFINE: %{sparsifier_opts} = enable-runtime-library=true
+// DEFINE: %{sparsifier_opts_sve} = enable-arm-sve=true %{sparsifier_opts}
+// DEFINE: %{compile} = mlir-opt %s --sparsifier="%{sparsifier_opts}"
+// DEFINE: %{compile_sve} = mlir-opt %s --sparsifier="%{sparsifier_opts_sve}"
+// DEFINE: %{run_libs} = -shared-libs=%mlir_c_runner_utils,%mlir_runner_utils
+// DEFINE: %{run_opts} = -e entry -entry-point-result=void
+// DEFINE: %{run} = mlir-cpu-runner %{run_opts} %{run_libs}
+// DEFINE: %{run_sve} = %mcr_aarch64_cmd --march=aarch64 --mattr="+sve"
%{run_opts} %{run_libs}
+//
+// DEFINE: %{env} =
+//--
+
+// RUN: %{compile} | %{run} | FileCheck %s
+//
+// Do the same run, but now with direct IR generation.
+// REDEFINE: %{sparsifier_opts} = enable-runtime-library=false
+// RUN: %{compile} | %{run} | FileCheck %s
+//
+// Do the same run, but now with direct IR generation and vectorization.
+// REDEFINE: %{sparsifier_opts} = enable-runtime-library=false vl=2
reassociate-fp-reductions=true enable-index-optimizations=true
+// RUN: %{compile} | %{run} | FileCheck %s
+//
+// Do the same run, but now with direct IR generation and VLA vectorization.
+// RUN: %if mlir_arm_sve_tests %{ %{compile_sve} | %{run_sve} | FileCheck %s %}
+
+//
+// Integration test that generates a tensor with specified sparsity level.
+//
+
+!Generator = !llvm.ptr
+!Array = !llvm.ptr
+
+#SparseVector = #sparse_tensor.encoding<{
+ map = (d0) -> (d0 : compressed)
+}>
+
+module {
+ func.func private @rtsrand(index) -> (!Generator)
+ func.func private
