[ExecEngine] Add scanMemref* API and an `input-dir` option in JIT runner

[avarmapml]

Summary

• This commit adds scanMemref* API for reading input data into memref during JIT execution.
• It also adds input-dir option as part of mlir-cpu-runner tool to let the user mention the absolute path to the input directory from which the scanMemref* API will extract the data from.

Signed-off-by: Abhishek Varma <[e-mail]>

[avarmapml]

Trivial changes to test case

[avarmapml]

Updated integer to string conversion logic.

• Build fails for Windows because to_string isn't available as a member of std.
• Implemented a trivial integer->string conversion to tackle the same.

[avarmapml]

Shifted execution test case within mlir/test/mlir-cpu-runner

[avarmapml]

Can someone review this patch?
I've addressed the Windows build failures resulting earlier.

@bondhugula

[avarmapml]

@ftynse @herhut Can you review this patch?

[bondhugula]

@herhut - would you be able to also review this?

mlir/include/mlir/ExecutionEngine/RunnerUtils.h

304

We should have these emitted to stderr?

304

Opening file or Opened file?
Nit: no space before colon.

309

Nit: Full stop at the end.

[avarmapml]

Trivial nit changes.

[mehdi_amini]

mlir/include/mlir/ExecutionEngine/RunnerUtils.h

38

Any reason you're not using the LLVM streams here?

44

That seems like quite a hack to pass state around to me, using an environment variable may an alternative?

[avarmapml]

Using environment variable and llvm::itostr

• Modified the code to incorporate using environment variable to pass the input directory path to the runtime utility.
• Using llvm::itostr instead of sstream.

mlir/include/mlir/ExecutionEngine/RunnerUtils.h

38

I was using sstream to get the int->string conversion go smoothly.
Wasn't aware about itostr as part of llvm/Support/StringExtras.h - but now using that instead.
Thanks for the nudge here!

44

Thanks for suggesting this! I've made the appropriate changes.

[mehdi_amini]

mlir/include/mlir/ExecutionEngine/RunnerUtils.h

44

Nit: constexpr char * instead of define.

296

LLVM should have enough file manipulation routine to avoid std::fstream here as well I think.

299

Left over debugging.

305

Can you move all the part of the implementation that is not dependent on T into a non-templated function that would be in the .cpp file? That would avoid exposing all of the file manipulation and the env var in the header.

mlir/lib/ExecutionEngine/JitRunner.cpp

389

I would hope we could remove all changes in this file: why do we even need a command line flag still?

[avarmapml]

Shifted opening file utility to RunnerUtils.cpp

• Shifted opening file utility based on environment variable to RunnerUtils.cpp.
• Handled other review comments.

[avarmapml]

Shifted the file opening utility and environment variable

mlir/include/mlir/ExecutionEngine/RunnerUtils.h

296

I tried exploring/implementing the code with both raw_fd_stream as well as llvm/Support/FileSystem.h but I believe they require the char * as the buffer input.
I guess using LLVM's file manipulation routines would thus involve some local manipulations perhaps.
std::fstream with that respect seems quite cleaner/ "friendly" to me.

Can you perhaps give me some pointers here perhaps for me to try?

299

I thought it'd be nice to let the user know the status of scanning each argument during runtime, so included that (the position of the statement itself should've been right after open though).

But I've removed it now in this revision and keeping statements pertaining just the beginning and completion of scanMemref*.

mlir/lib/ExecutionEngine/JitRunner.cpp

389

Well, if the user is indeed creating an environment variable prior to calling mlir-cpu-runner, then we don't require the flag as you rightly mentioned.

But I see two reasons for still requiring the flag and then internally creating the environment variable :-

1. Perhaps user can mess up naming the environment variable correctly?
2. Also, the test case I've included in this revision needs to know where the input files are. Not sure how we're to set environment variable in that case.

Keeping the flag makes the execution process similar to both the user and the lit tests.

Let me know your thoughts here, perhaps I'm missing something.

[mehdi_amini]

mlir/include/mlir/ExecutionEngine/RunnerUtils.h

296

All of LLVM tools (like mlir-opt, etc.) are using LLVM steams and they are still loading files: so it is definitely possible...
Quickly looking at mlir-opt implementation, the input file is passed to mlir::openInputFile which is implemented in mlir/lib/Support/FileUtilities.cpp.

There is a constructor for raw_fd_stream that is: raw_fd_stream(StringRef Filename, std::error_code &EC); : seems like this takes a FileName?

mlir/lib/ExecutionEngine/JitRunner.cpp

389

I thought of an environment variable to keep this entirely out of the JitRunner.cpp actually. I didn't quite get what is your problem with the test though?

Alternatively, we could scrap all of this and use an API instead: the IR call could pass in the directory where to load a memref (environment variables are a bit hacky anyway)

mlir/include/mlir/ExecutionEngine/RunnerUtils.h

 #endif // MLIR_RUNNERUTILS_EXPORT
 #else
 // Non-windows: use visibility attributes.
 #define MLIR_RUNNERUTILS_EXPORT __attribute__((visibility("default")))
 #endif // _WIN32
 
 #include <assert.h>
 #include <cmath>
+#include <fstream>
 #include <iostream>
+#include <sstream>
 
 #include "mlir/ExecutionEngine/CRunnerUtils.h"
 
+/// The file which contains the absolute path to the input directory from which
+/// scanMemref* API will extract each input file.
+#define INPUT_DIRECTORY_PATH_FILE "inputDirPathFile.dat"
+
 template <typename T, typename StreamType>
 void printMemRefMetaData(StreamType &os, const DynamicMemRefType<T> &v) {
   os << "base@ = " << reinterpret_cast<void *>(v.data) << " rank = " << v.rank
      << " offset = " << v.offset;
   auto print = [&](const int64_t *ptr) {
     if (v.rank == 0)
       return;
     os << ptr[0];

 }
 
 template <typename T>
 void printMemRef(UnrankedMemRefType<T> &m) {
   std::cout << "Unranked Memref ";
   printMemRef(DynamicMemRefType<T>(m));
 }
 
+/// MemRefDataScanner struct declares the scanning API template to read in
+/// input from a filestream into the memref passed.
+template <typename T>
+struct MemRefDataScanner {
+  static void scan(std::fstream &fs, T *base, int64_t rank, int64_t offset,
+                   const int64_t *sizes, const int64_t *strides);
+  static void scanFirst(std::fstream &fs, T *base, int64_t rank, int64_t offset,
+                        const int64_t *sizes, const int64_t *strides);
+  static void scanLast(std::fstream &fs, T *base, int64_t rank, int64_t offset,
+                       const int64_t *sizes, const int64_t *strides);
+};
+
+/// Helper function to scan the input file in a recursive structure that takes
+/// into account the shape of the memref.
+template <typename T>
+void MemRefDataScanner<T>::scanFirst(std::fstream &fs, T *base, int64_t rank,
+                                     int64_t offset, const int64_t *sizes,
+                                     const int64_t *strides) {
+  scan(fs, base, rank - 1, offset, sizes + 1, strides + 1);
+  // If single element, return early.
+  if (sizes[0] <= 1) {
+    return;
+  }
+}
+
+/// Helper function to scan the input file in a recursive structure that takes
+/// into account the shape of the memref.
+template <typename T>
+void MemRefDataScanner<T>::scan(std::fstream &fs, T *base, int64_t rank,
+                                int64_t offset, const int64_t *sizes,
+                                const int64_t *strides) {
+  if (rank == 0) {
+    fs >> base[offset];
+    return;
+  }
+  scanFirst(fs, base, rank, offset, sizes, strides);
+  for (unsigned i = 1; i + 1 < sizes[0]; ++i) {
+    scan(fs, base, rank - 1, offset + i * strides[0], sizes + 1, strides + 1);
+  }
+  if (sizes[0] <= 1)
+    return;
+  scanLast(fs, base, rank, offset, sizes, strides);
+}
+
+/// Helper function to scan the input file in a recursive structure that takes
+/// into account the shape of the memref.
+template <typename T>
+void MemRefDataScanner<T>::scanLast(std::fstream &fs, T *base, int64_t rank,
+                                    int64_t offset, const int64_t *sizes,
+                                    const int64_t *strides) {
+  scan(fs, base, rank - 1, offset + (sizes[0] - 1) * (*strides), sizes + 1,
+       strides + 1);
+}
+
+/// Opens the file in filestream and scans the file into the memref passed.
+template <typename T>
+static void scanMemRef(const DynamicMemRefType<T> &M, unsigned fileNum) {
+  std::ostringstream oss;
+  oss << fileNum;
+  std::string fileName = "/input" + oss.str() + ".dat";
+  std::fstream inputDirFs;
+  inputDirFs.open(INPUT_DIRECTORY_PATH_FILE, std::fstream::in);
+  std::string filePath = "";
+  inputDirFs >> filePath;
+  inputDirFs.close();
+  filePath = filePath + fileName;
+  std::cerr << "Opened file: " << filePath;
+  std::fstream fs;
+  fs.open(filePath, std::fstream::in);
+  MemRefDataScanner<T>::scan(fs, M.data, M.rank, M.offset, M.sizes, M.strides);
+  fs.close();
+  std::cerr << "\nMemref scanning complete.\n";
+}
+
+/// Casts the UnrankedMemRef into DynamicMemRef and calls the main scanMemRef
+/// function.
+template <typename T>
+static void scanMemRef(UnrankedMemRefType<T> &M, unsigned fileNum) {
+  std::cerr << "Scanning unranked memref:\n";
+  scanMemRef(DynamicMemRefType<T>(M), fileNum);
+}
+
 /// Verify the result of two computations are equivalent up to a small
 /// numerical error and return the number of errors.
 template <typename T>
 struct MemRefDataVerifier {
   /// Maximum number of errors printed by the verifier.
   static constexpr int printLimit = 10;
 
   /// Verify the relative difference of the values is smaller than epsilon.

 extern "C" MLIR_RUNNERUTILS_EXPORT void
 _mlir_ciface_printMemref3dF32(StridedMemRefType<float, 3> *m);
 extern "C" MLIR_RUNNERUTILS_EXPORT void
 _mlir_ciface_printMemref4dF32(StridedMemRefType<float, 4> *m);
 
 extern "C" MLIR_RUNNERUTILS_EXPORT void _mlir_ciface_printMemrefVector4x4xf32(
     StridedMemRefType<Vector2D<4, 4, float>, 2> *m);
 
+extern "C" MLIR_RUNNERUTILS_EXPORT void
+_mlir_ciface_scanMemrefI32(UnrankedMemRefType<int32_t> *M, unsigned fileNum);
+extern "C" MLIR_RUNNERUTILS_EXPORT void
+_mlir_ciface_scanMemrefI64(UnrankedMemRefType<int64_t> *M, unsigned fileNum);
+extern "C" MLIR_RUNNERUTILS_EXPORT void
+_mlir_ciface_scanMemrefF32(UnrankedMemRefType<float> *M, unsigned fileNum);
+extern "C" MLIR_RUNNERUTILS_EXPORT void
+_mlir_ciface_scanMemrefF64(UnrankedMemRefType<double> *M, unsigned fileNum);
+
+extern "C" MLIR_RUNNERUTILS_EXPORT void scanMemrefI32(int64_t rank, void *ptr,
+                                                      unsigned fileNum);
+extern "C" MLIR_RUNNERUTILS_EXPORT void scanMemrefI64(int64_t rank, void *ptr,
+                                                      unsigned fileNum);
+extern "C" MLIR_RUNNERUTILS_EXPORT void scanMemrefF32(int64_t rank, void *ptr,
+                                                      unsigned fileNum);
+extern "C" MLIR_RUNNERUTILS_EXPORT void scanMemrefF64(int64_t rank, void *ptr,
+                                                      unsigned fileNum);
+
 extern "C" MLIR_RUNNERUTILS_EXPORT int64_t _mlir_ciface_verifyMemRefI32(
     UnrankedMemRefType<int32_t> *actual, UnrankedMemRefType<int32_t> *expected);
 extern "C" MLIR_RUNNERUTILS_EXPORT int64_t _mlir_ciface_verifyMemRefF32(
     UnrankedMemRefType<float> *actual, UnrankedMemRefType<float> *expected);
 extern "C" MLIR_RUNNERUTILS_EXPORT int64_t _mlir_ciface_verifyMemRefF64(
     UnrankedMemRefType<double> *actual, UnrankedMemRefType<double> *expected);
 
 extern "C" MLIR_RUNNERUTILS_EXPORT int64_t verifyMemRefI32(int64_t rank,

mlir/lib/ExecutionEngine/JitRunner.cpp

 #include "mlir/Support/FileUtilities.h"
 
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/LegacyPassNameParser.h"
 #include "llvm/Support/CommandLine.h"
+#include "llvm/Support/Debug.h"
 #include "llvm/Support/FileUtilities.h"
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/StringSaver.h"
 #include "llvm/Support/ToolOutputFile.h"
 #include <cstdint>
+#include <fstream>
+#include <iostream>
 #include <numeric>
 #include <utility>
 
 using namespace mlir;
 using llvm::Error;
 
+/// Since `scan_memref` API will require the input directory path to extract
+/// input files, we will be storing the path obtained from command line argument
+/// into a file called `inputDirPathFile.dat`.
+#define INPUT_DIRECTORY_PATH_FILE "inputDirPathFile.dat"
+
 namespace {
 /// This options struct prevents the need for global static initializers, and
 /// is only initialized if the JITRunner is invoked.
 struct Options {
   llvm::cl::opt<std::string> inputFilename{llvm::cl::Positional,
                                            llvm::cl::desc("<input file>"),
                                            llvm::cl::init("-")};
   llvm::cl::opt<std::string> mainFuncName{
struct Options {
   llvm::cl::opt<bool> dumpObjectFile{
       "dump-object-file",
       llvm::cl::desc("Dump JITted-compiled object to file specified with "
                      "-object-filename (<input file>.o by default).")};
 
   llvm::cl::opt<std::string> objectFilename{
       "object-filename",
       llvm::cl::desc("Dump JITted-compiled object to file <input file>.o")};
+
+  llvm::cl::opt<std::string> inputDirectoryPath{
+      "input-dir", llvm::cl::desc("Input directory path"), llvm::cl::init(".")};
 };
 
 struct CompileAndExecuteConfig {
   /// LLVM module transformer that is passed to ExecutionEngine.
   std::function<llvm::Error(llvm::Module *)> transformer;
 
   /// A custom function that is passed to ExecutionEngine. It processes MLIR
   /// module and creates LLVM IR module.
int mlir::JitRunnerMain(int argc, char **argv, const DialectRegistry &registry,
   auto compileAndExecuteFn =
       StringSwitch<CompileAndExecuteFnT>(options.mainFuncType.getValue())
           .Case("i32", compileAndExecuteSingleReturnFunction<int32_t>)
           .Case("i64", compileAndExecuteSingleReturnFunction<int64_t>)
           .Case("f32", compileAndExecuteSingleReturnFunction<float>)
           .Case("void", compileAndExecuteVoidFunction)
           .Default(nullptr);
 
+  // Open `inputDirPathFile.dat` and store the input directory path obtained
+  // from the command line option `inputDirectoryPath`.
+  std::ofstream inputDirFs(INPUT_DIRECTORY_PATH_FILE);
+  inputDirFs << options.inputDirectoryPath.getValue();
+  inputDirFs.close();
   Error error = compileAndExecuteFn
                     ? compileAndExecuteFn(options, m.get(),
                                           options.mainFuncName.getValue(),
                                           compileAndExecuteConfig)
                     : makeStringError("unsupported function type");
 
   int exitCode = EXIT_SUCCESS;
   llvm::handleAllErrors(std::move(error),
                         [&exitCode](const llvm::ErrorInfoBase &info) {
                           llvm::errs() << "Error: ";
                           info.log(llvm::errs());
                           llvm::errs() << '\n';
                           exitCode = EXIT_FAILURE;
                         });
 
   return exitCode;
 }

mlir/lib/ExecutionEngine/RunnerUtils.cpp

 }
 extern "C" void _mlir_ciface_printMemref3dF32(StridedMemRefType<float, 3> *M) {
   impl::printMemRef(*M);
 }
 extern "C" void _mlir_ciface_printMemref4dF32(StridedMemRefType<float, 4> *M) {
   impl::printMemRef(*M);
 }
 
+/// scanMemref* APIs need an integer argument 'fileNum'. Scans a file
+/// identified by 'input<fileNum>.dat' and populates the memref passed with the
+/// file content.
+extern "C" void _mlir_ciface_scanMemrefI32(UnrankedMemRefType<int32_t> *M,
+                                           unsigned fileNum) {
+  impl::scanMemRef(*M, fileNum);
+}
+
+extern "C" void _mlir_ciface_scanMemrefI64(UnrankedMemRefType<int64_t> *M,
+                                           unsigned fileNum) {
+  impl::scanMemRef(*M, fileNum);
+}
+
+extern "C" void _mlir_ciface_scanMemrefF32(UnrankedMemRefType<float> *M,
+                                           unsigned fileNum) {
+  impl::scanMemRef(*M, fileNum);
+}
+
+extern "C" void _mlir_ciface_scanMemrefF64(UnrankedMemRefType<double> *M,
+                                           unsigned fileNum) {
+  impl::scanMemRef(*M, fileNum);
+}
+
+extern "C" void scanMemrefI32(int64_t rank, void *ptr, unsigned fileNum) {
+  UnrankedMemRefType<int32_t> descriptor = {rank, ptr};
+  _mlir_ciface_scanMemrefI32(&descriptor, fileNum);
+}
+
+extern "C" void scanMemrefI64(int64_t rank, void *ptr, unsigned fileNum) {
+  UnrankedMemRefType<int64_t> descriptor = {rank, ptr};
+  _mlir_ciface_scanMemrefI64(&descriptor, fileNum);
+}
+
+extern "C" void scanMemrefF32(int64_t rank, void *ptr, unsigned fileNum) {
+  UnrankedMemRefType<float> descriptor = {rank, ptr};
+  _mlir_ciface_scanMemrefF32(&descriptor, fileNum);
+}
+
+extern "C" void scanMemrefF64(int64_t rank, void *ptr, unsigned fileNum) {
+  UnrankedMemRefType<double> descriptor = {rank, ptr};
+  _mlir_ciface_scanMemrefF64(&descriptor, fileNum);
+}
+
 extern "C" int64_t
 _mlir_ciface_verifyMemRefI32(UnrankedMemRefType<int32_t> *actual,
                              UnrankedMemRefType<int32_t> *expected) {
   return impl::verifyMemRef(*actual, *expected);
 }
 
 extern "C" int64_t
 _mlir_ciface_verifyMemRefF32(UnrankedMemRefType<float> *actual,

mlir/test/mlir-cpu-runner/ExecutionTest/execution-test.mlir

+// RUN: mlir-opt -convert-linalg-to-loops -lower-affine -arith-expand \
+// RUN:          -canonicalize -convert-scf-to-openmp -convert-scf-to-cf \
+// RUN:          -canonicalize -convert-math-to-llvm -convert-vector-to-llvm \
+// RUN:          -convert-memref-to-llvm -convert-openmp-to-llvm -canonicalize \
+// RUN:          -convert-cf-to-llvm -reconcile-unrealized-casts -canonicalize %s | \
+// RUN: mlir-cpu-runner -O3 -e execution_test_driver -entry-point-result=void \
+// RUN:                 -input-dir=%mlir_src_root/test/mlir-cpu-runner/ExecutionTest \
+// RUN:                 -shared-libs=%mlir_runner_utils_dir/libmlir_runner_utils%shlibext,%mlir_runner_utils_dir/libmlir_c_runner_utils%shlibext | FileCheck %s
+
+// This test case demonstrates working of scanMemref* API in tandem with printMemref*.
+module  {
+  func.func private @printMemrefF32(memref<*xf32>)
+  func.func private @scanMemrefF32(memref<*xf32>, i32)
+  func.func @execution_test_driver() {
+    // Input arg0.
+    %0 = memref.alloc() : memref<10xf32>
+    %c0_i32 = arith.constant 0 : i32
+    %1 = memref.cast %0 : memref<10xf32> to memref<*xf32>
+    call @scanMemrefF32(%1, %c0_i32) : (memref<*xf32>, i32) -> ()
+
+    // Input arg1.
+    %2 = memref.alloc() : memref<10xf32>
+    %c1_i32 = arith.constant 1 : i32
+    %3 = memref.cast %2 : memref<10xf32> to memref<*xf32>
+    call @scanMemrefF32(%3, %c1_i32) : (memref<*xf32>, i32) -> ()
+
+    // Output arg2.
+    %4 = memref.alloc() : memref<10xf32>
+
+    // Call user-defined function `getSum`.
+    call @getElementWiseSum(%0, %2, %4) : (memref<10xf32>, memref<10xf32>, memref<10xf32>) -> ()
+
+    // Print Output arg2.
+    %5 = memref.cast %4 : memref<10xf32> to memref<*xf32>
+    call @printMemrefF32(%5) : (memref<*xf32>) -> ()
+
+    return
+  }
+  func.func @getElementWiseSum(%arg0: memref<10xf32>, %arg1: memref<10xf32>, %arg2: memref<10xf32>) {
+    affine.for %arg3 = 0 to 10 {
+      %0 = affine.load %arg0[%arg3] : memref<10xf32>
+      %1 = affine.load %arg1[%arg3] : memref<10xf32>
+      %2 = arith.addf %0, %1 : f32
+      affine.store %2, %arg2[%arg3] : memref<10xf32>
+    }
+    return
+  }
+}
+// %arg0 is initialized with input0.dat :   1    2    3    4    5    6    7    8    9   10
+// %arg1 is initialized with input1.dat :  11   12   13   14   15   16   17   18   19   20
+// Result for element-wise addition     :
+// CHECK:                                 [12,  14,  16,  18,  20,  22,  24,  26,  28,  30]

mlir/test/mlir-cpu-runner/ExecutionTest/input0.dat

[phone]

mlir/test/mlir-cpu-runner/ExecutionTest/input1.dat

[phone] 18 19 20