Add utility functions to initialize 1-D memrefs.

Currently supported type: f16, bf16, f32.
intel · Aug 7, 2023 · 73ab8ba · 73ab8ba
1 parent d08b6f3
commit 73ab8ba
Show file tree

Hide file tree

Showing 2 changed files with 41 additions and 0 deletions.
diff --git a/include/imex/ExecutionEngine/ImexRunnerUtils.h b/include/imex/ExecutionEngine/ImexRunnerUtils.h
@@ -33,6 +33,24 @@
 #include "mlir/ExecutionEngine/Float16bits.h"
 #include "mlir/ExecutionEngine/RunnerUtils.h"
 
+template <typename T, int N> struct MemRefDescriptor {
+  T *allocated;
+  T *aligned;
+  int64_t offset;
+  int64_t sizes[N];
+  int64_t strides[N];
+};
+
+extern "C" IMEX_RUNNERUTILS_EXPORT void
+_mlir_ciface_fillResource1DBF16(MemRefDescriptor<bf16, 1> *ptr, // NOLINT
+                                float value);
+extern "C" IMEX_RUNNERUTILS_EXPORT void
+_mlir_ciface_fillResource1DF16(MemRefDescriptor<f16, 1> *ptr, // NOLINT
+                               float value);
+extern "C" IMEX_RUNNERUTILS_EXPORT void
+_mlir_ciface_fillResource1DF32(MemRefDescriptor<float, 1> *ptr, // NOLINT
+                               float value);
+
 extern "C" IMEX_RUNNERUTILS_EXPORT void
 _mlir_ciface_printMemrefBF16(UnrankedMemRefType<bf16> *m);
 extern "C" IMEX_RUNNERUTILS_EXPORT void

diff --git a/lib/ExecutionEngine/ImexRunnerUtils.cpp b/lib/ExecutionEngine/ImexRunnerUtils.cpp
@@ -21,6 +21,29 @@
 
 // NOLINTBEGIN(*-identifier-naming)
 
+/// Fills the given 1D bf16 memref with the given float value.
+extern "C" void
+_mlir_ciface_fillResource1DBF16(MemRefDescriptor<bf16, 1> *ptr, // NOLINT
+                                float value) {
+  bf16 bf16_val(value);
+  std::fill_n(ptr->allocated, ptr->sizes[0], bf16_val);
+}
+
+/// Fills the given 1D f16 memref with the given float value.
+extern "C" void
+_mlir_ciface_fillResource1DF16(MemRefDescriptor<f16, 1> *ptr, // NOLINT
+                               float value) {
+  f16 f16_val(value);
+  std::fill_n(ptr->allocated, ptr->sizes[0], f16_val);
+}
+
+/// Fills the given 1D float (f32) memref with the given float value.
+extern "C" void
+_mlir_ciface_fillResource1DF32(MemRefDescriptor<float, 1> *ptr, // NOLINT
+                               float value) {
+  std::fill_n(ptr->allocated, ptr->sizes[0], value);
+}
+
 extern "C" void _mlir_ciface_printMemrefBF16(UnrankedMemRefType<bf16> *M) {
   impl::printMemRef(*M);
 }