[mlir] [sparse] convenience runtime support to read Matrix Market format

[aartbik]

Summary

Setting up input data for benchmarks and integration tests can be tedious in
pure MLIR. With more sparse tensor work planned, this convenience library
simplifies reading sparse matrices in the popular Matrix Market Exchange
Format (see https://math.nist.gov/MatrixMarket). Note that this library
is *not* part of core MLIR. It is merely intended as a convenience library
for benchmarking and integration testing.

[penpornk]

I'm new to MLIR, so sorry in advance for noob MLIR comments!

mlir/integration_test/Sparse/CPU/matrix-market-example.mlir

28

To make this more generic, shouldn't a be allocated after we know m and n?

49

I think there should be a separate routine to convert sparse matrix to dense, i.e., one routine to read Matrix Market entries into 3 arrays, then another to convert these arrays to a dense matrix. If we materialize the dense matrix in the sparse matrix read routine, we wouldn't be able to read many large sparse matrices from real applications.

This would allow us to just read and store all entries in 3 arrays at once (instead of reading entry-by-entry): i[nnz], j[nnz], and val[nnz].
The official Matrix Market I/O C library can do that for us: See mm_read_unsymmetric_sparse.
This will avoid the need to check if a different file is already open as well.

70

Noob question, why isn't there a CHECK-SAME in front of ( 1, 0, 0, 1.4, 0 ) just like the four lines below?

mlir/integration_test/data/test.mtx

14

Great idea making the coordinates not sorted! This makes for a good test case for when we want to convert to CSR, etc. :)

mlir/lib/ExecutionEngine/SparseUtils.cpp

38

(If we are still going to have this routine)
Should we use the official library to read the header (mm_read_banner)? It covers all cases.

82

Should this be inline?

[aartbik]

Welcome to this repo, Penporn!
And thanks for the comments.

mlir/integration_test/Sparse/CPU/matrix-market-example.mlir

28

Yes, you are right. I was a bit lazy anticipating the size, but for illustration purposes, adapting to the size would be better. I changed this to read more like an example. Thanks!

49

Note that the purpose of this library is not to provide deep support for sparse computations, but merely something light weight that makes it a bit easier to set up tests and benchmarks.

I did it this way to avoid a difficult interaction between memory allocation at MLIR level and in the C library. By "coming back first" after reading the header, all allocation can be done in MLIR.

70

The -SAME part indicates that the CHECK continues at the same line as the previous CHECK.
The output for this program is:

5
5
9
( ( 1, 0, 0, 1.4, 0 ), ( 0, 2, 0, 0, 2.5 ), ( 0, 0, 3, 0, 0 ), ( 4.1, 0, 0, 4, 0 ), ( 0, 5.2, 0, 0, 5 ) )

mlir/lib/ExecutionEngine/SparseUtils.cpp

38

I strongly prefer our own very light-weight implementation over pulling in sources from elsewhere.
We can support the other formats as the need arises.

82

If you prefer, but given the static scope, the compiler could already decide to inline this based on its heuristics.

[aartbik]

modified example to illustrate usage better

[penpornk]

Thank you for the clarifications!

[aartbik]

Thank *you* for the review, I feel it indeed has made the example more clear!

[nicolasvasilache]

Very cool @aartbik , nice use of MLIR / C interop to avoid hardcoding stuff in IR prematurely.

mlir/integration_test/Sparse/CPU/matrix-market-example.mlir

58

This is a nice minimal first step to think about for adding sparse functionality to Linalg.

There are a few things going on that would be nice to separate:

1. indirection for sparse read accesses
2. dense writes (i.e. no extra indirection)
3. dynamic collapsing logic between k loop over nnz and "i, j loops".

Let's chat in our next VC.

[mehdi_amini]

mlir/lib/ExecutionEngine/SparseUtils.cpp

42

This whole API is forcing us into using files, while we ought to work with a buffer abstraction in general (i.e. we can get a file or a buffer from another source, but the API should abstract it).

Can you refactor this and use something like llvm::MemoryBuffer for example?

[aartbik]

mlir/lib/ExecutionEngine/SparseUtils.cpp

42

I like this idea (did not know about the built-in mem buffer support).
This will also make injecting test files a bit easier than currently is the case.
Thanks!

Coming up in next sparse utils enhancements....

mlir/integration_test/CMakeLists.txt

 
 # Let target "check-mlir" imply the target "check-mlir-integration" as well,
 # provided that MLIR_INCLUDE_INTEGRATION_TESTS=ON during the initial setup.
 add_dependencies(check-mlir check-mlir-integration)
 
 add_lit_testsuites(MLIR_INTEGRATION ${CMAKE_CURRENT_SOURCE_DIR}
   DEPENDS ${MLIR_INTEGRATION_TEST_DEPENDS}
   )
+
+# Copy test data over.
+file(COPY ${CMAKE_CURRENT_SOURCE_DIR}/data/test.mtx
+        DESTINATION ${MLIR_INTEGRATION_TEST_DIR}/data/)

mlir/integration_test/Sparse/CPU/lit.local.cfg

+import sys
+
+# No JIT on win32.
+if sys.platform == 'win32':
+    config.unsupported = True

mlir/integration_test/Sparse/CPU/matrix-market-example.mlir

+// RUN: mlir-opt %s \
+// RUN:  -convert-scf-to-std -convert-vector-to-scf \
+// RUN:  -convert-linalg-to-llvm -convert-vector-to-llvm | \
+// RUN: SPARSE_MATRIX0="%mlir_integration_test_dir/data/test.mtx" \
+// RUN: mlir-cpu-runner \
+// RUN:  -e entry -entry-point-result=void  \
+// RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
+// RUN: FileCheck %s
+
+module {
+  func @openMatrix(!llvm.ptr<i8>, memref<index>, memref<index>, memref<index>) -> ()
+  func @readMatrixItem(memref<index>, memref<index>, memref<f64>) -> ()
+  func @closeMatrix() -> ()
+  func @getSparseMatrix(index) -> (!llvm.ptr<i8>)
+
+  func @entry() {
+    %d0  = constant 0.0 : f64
+    %c0  = constant 0 : index
+    %c1  = constant 1 : index
+    %c5  = constant 5 : index
+    %m   = alloc() : memref<index>
+    %n   = alloc() : memref<index>
+    %nnz = alloc() : memref<index>
+    %i   = alloc() : memref<index>
+    %j   = alloc() : memref<index>
+    %d   = alloc() : memref<f64>
+
+    //
+    // Read the header of a sparse matrix. This yields the
+    // size (m x n) and number of nonzero elements (nnz).
+    //
+    %file = call @getSparseMatrix(%c0) : (index) -> (!llvm.ptr<i8>)
+    call @openMatrix(%file, %m, %n, %nnz)
+        : (!llvm.ptr<i8>, memref<index>,
+                    memref<index>, memref<index>) -> ()
+    %M = load %m[]   : memref<index>
+    %N = load %n[]   : memref<index>
+    %Z = load %nnz[] : memref<index>
+
+    //
+    // At this point, code should prepare a proper sparse storage
+    // scheme for an m x n matrix with nnz nonzero elements. For
+    // simplicity, however, here we simply set up a dense matrix.
+    //
+    %a = alloc(%M, %N) : memref<?x?xf64>
+    scf.for %ii = %c0 to %M step %c1 {
+      scf.for %jj = %c0 to %N step %c1 {
+        store %d0, %a[%ii, %jj] : memref<?x?xf64>
+      }
+    }
+
+    //
+    // Now we are ready to read in the nonzero elements of the
+    // sparse matrix and insert these into a sparse storage
+    // scheme. In this example, we simply insert them in the
+    // dense matrix.
+    //
+    scf.for %k = %c0 to %Z step %c1 {
+      call @readMatrixItem(%i, %j, %d)
+          : (memref<index>, memref<index>, memref<f64>) -> ()
+      %I = load %i[] : memref<index>
+      %J = load %j[] : memref<index>
+      %D = load %d[] : memref<f64>
+      store %D, %a[%I, %J] : memref<?x?xf64>
+    }
+    call @closeMatrix() : () -> ()
+
+    //
+    // Verify that the results are as expected.
+    //
+    %A = vector.transfer_read %a[%c0, %c0], %d0 : memref<?x?xf64>, vector<5x5xf64>
+    vector.print %M : index
+    vector.print %N : index
+    vector.print %Z : index
+    vector.print %A : vector<5x5xf64>
+    //
+    // CHECK: 5
+    // CHECK: 5
+    // CHECK: 9
+    //
+    // CHECK: ( ( 1, 0, 0, 1.4, 0 ),
+    // CHECK-SAME: ( 0, 2, 0, 0, 2.5 ),
+    // CHECK-SAME: ( 0, 0, 3, 0, 0 ),
+    // CHECK-SAME: ( 4.1, 0, 0, 4, 0 ),
+    // CHECK-SAME: ( 0, 5.2, 0, 0, 5 ) )
+
+    //
+    // Free.
+    //
+    dealloc %m   : memref<index>
+    dealloc %n   : memref<index>
+    dealloc %nnz : memref<index>
+    dealloc %i   : memref<index>
+    dealloc %j   : memref<index>
+    dealloc %d   : memref<f64>
+    dealloc %a   : memref<?x?xf64>
+
+    return
+  }
+}

mlir/integration_test/data/test.mtx

+%%MatrixMarket matrix coordinate real general
+%
+% This is a test sparse matrix in Matrix Market Exchange Format.
+% see https://math.nist.gov/MatrixMarket
+%
+5 5 9
+1 1 1.0
+1 4 1.4
+2 2 2.0
+2 5 2.5
+3 3 3.0
+4 4 4.0
+4 1 4.1
+5 5 5.0
+5 2 5.2

mlir/lib/ExecutionEngine/CMakeLists.txt

 # Exclude these from libMLIR.so because the JIT infrastructure
 # is a big dependency which most don't need.
 
 set(LLVM_OPTIONAL_SOURCES
   CRunnerUtils.cpp
+  SparseUtils.cpp
   ExecutionEngine.cpp
   RunnerUtils.cpp
   OptUtils.cpp
   JitRunner.cpp
   )
 
 add_mlir_library(MLIRExecutionEngine
   ExecutionEngine.cpp
add_mlir_library(MLIRJitRunner
   MLIRTransforms
   MLIRStandardToLLVM
   MLIRSupport
 )
 
 add_mlir_library(mlir_c_runner_utils
   SHARED
   CRunnerUtils.cpp
+  SparseUtils.cpp
 
   EXCLUDE_FROM_LIBMLIR
   )
 set_property(TARGET mlir_c_runner_utils PROPERTY CXX_STANDARD 11)
 
 add_mlir_library(mlir_c_runner_utils_static
   CRunnerUtils.cpp
+  SparseUtils.cpp
 
   EXCLUDE_FROM_LIBMLIR
   )
 set_property(TARGET mlir_c_runner_utils_static PROPERTY CXX_STANDARD 11)
 target_compile_definitions(mlir_c_runner_utils PRIVATE mlir_c_runner_utils_EXPORTS)
 
 add_mlir_library(mlir_runner_utils
   SHARED
   RunnerUtils.cpp
 
   EXCLUDE_FROM_LIBMLIR
 
   LINK_LIBS PUBLIC
   mlir_c_runner_utils_static
 )
 target_compile_definitions(mlir_runner_utils PRIVATE mlir_runner_utils_EXPORTS)

mlir/lib/ExecutionEngine/SparseUtils.cpp

+//===- SparseUtils.cpp - Sparse Utils for MLIR execution ------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements a light-weight runtime library that is useful for
+// sparse tensor manipulations. The functionality provided in this library
+// is meant to simplify benchmarking, testing, and debugging MLIR code that
+// operates on sparse tensors. The provided functionality is **not** part
+// of core MLIR, however.
+//
+//===----------------------------------------------------------------------===//
+
+#include <cctype>
+#include <cinttypes>
+#include <cstdio>
+#include <cstdlib>
+#include <cstring>
+
+//===----------------------------------------------------------------------===//
+//
+// Internal support for reading matrices in the Matrix Market Exchange Format.
+// See https://math.nist.gov/MatrixMarket for details on this format.
+//
+//===----------------------------------------------------------------------===//
+
+// Helper to convert string to lower case.
+static char *toLower(char *token) {
+  for (char *c = token; *c; c++)
+    *c = tolower(*c);
+  return token;
+}
+
+// Read the header of a general sparse matrix of type real.
+//
+// TODO: support other formats as well?
+//
+static void readHeader(FILE *file, char *name, uint64_t *m, uint64_t *n,
+                       uint64_t *nnz) {
+  char line[1025];
+  char header[64];
+  char object[64];
+  char format[64];
+  char field[64];
+  char symmetry[64];
+  // Read header line.
+  if (fscanf(file, "%63s %63s %63s %63s %63s\n", header, object, format, field,
+             symmetry) != 5) {
+    fprintf(stderr, "Corrupt header in %s\n", name);
+    exit(1);
+  }
+  // Make sure this is a general sparse matrix.
+  if (strcmp(toLower(header), "%%matrixmarket") ||
+      strcmp(toLower(object), "matrix") ||
+      strcmp(toLower(format), "coordinate") || strcmp(toLower(field), "real") ||
+      strcmp(toLower(symmetry), "general")) {
+    fprintf(stderr,
+            "Cannot find a general sparse matrix with type real in %s\n", name);
+    exit(1);
+  }
+  // Skip comments.
+  while (1) {
+    if (!fgets(line, 1025, file)) {
+      fprintf(stderr, "Cannot find data in %s\n", name);
+      exit(1);
+    }
+    if (line[0] != '%')
+      break;
+  }
+  // Next line contains M N NNZ.
+  if (sscanf(line, "%" PRIu64 "%" PRIu64 "%" PRIu64, m, n, nnz) != 3) {
+    fprintf(stderr, "Cannot find size in %s\n", name);
+    exit(1);
+  }
+}
+
+// Read next data item.
+static void readItem(FILE *file, char *name, uint64_t *i, uint64_t *j,
+                     double *d) {
+  if (fscanf(file, "%" PRIu64 " %" PRIu64 " %lg\n", i, j, d) != 3) {
+    fprintf(stderr, "Cannot find next data item in %s\n", name);
+    exit(1);
+  }
+  // Translate 1-based to 0-based.
+  *i = *i - 1;
+  *j = *j - 1;
+}
+
+//===----------------------------------------------------------------------===//
+//
+// Public API of the sparse runtime library.
+//
+// Enables MLIR code to read a matrix in Matrix Market Exchange Format
+// as follows:
+//
+//   call @openMatrix("A.mtx", %m, %n, %nnz) : (!llvm.ptr<i8>,
+//                                              memref<index>,
+//                                              memref<index>,
+//                                              memref<index>) -> ()
+//   .... prepare reading in m x n matrix A with nnz nonzero elements ....
+//   %u = load %nnz[] : memref<index>
+//   scf.for %k = %c0 to %u step %c1 {
+//     call @readMatrixItem(%i, %j, %d) : (memref<index>,
+//                                         memref<index>, memref<f64>) -> ()
+//     .... process next nonzero element A[i][j] = d ....
+//   }
+//   call @closeMatrix() : () -> ()
+//
+// The implementation is *not* thread-safe. Also, only *one* matrix file can
+// be open at the time. A matrix file must be closed before reading in a next.
+//
+// Note that input parameters mimic the layout of a MemRef<T>:
+//   struct MemRef {
+//     T *base;
+//     T *data;
+//     int64_t off;
+//   }
+//===----------------------------------------------------------------------===//
+
+// Currently open matrix. This is *not* thread-safe or re-entrant.
+static FILE *sparseFile = nullptr;
+static char *sparseFilename = nullptr;
+
+extern "C" void openMatrix(char *filename, uint64_t *mbase, uint64_t *mdata,
+                           int64_t moff, uint64_t *nbase, uint64_t *ndata,
+                           int64_t noff, uint64_t *nnzbase, uint64_t *nnzdata,
+                           int64_t nnzoff) {
+  if (sparseFile != nullptr) {
+    fprintf(stderr, "Other file still open %s vs. %s\n", sparseFilename,
+            filename);
+    exit(1);
+  }
+  sparseFile = fopen(filename, "r");
+  if (!sparseFile) {
+    fprintf(stderr, "Cannot find %s\n", filename);
+    exit(1);
+  }
+  sparseFilename = filename;
+  readHeader(sparseFile, filename, mdata, ndata, nnzdata);
+}
+
+extern "C" void readMatrixItem(uint64_t *ibase, uint64_t *idata, int64_t ioff,
+                               uint64_t *jbase, uint64_t *jdata, int64_t joff,
+                               double *dbase, double *ddata, int64_t doff) {
+  if (sparseFile == nullptr) {
+    fprintf(stderr, "Cannot read item from unopened matrix\n");
+    exit(1);
+  }
+  readItem(sparseFile, sparseFilename, idata, jdata, ddata);
+}
+
+extern "C" void closeMatrix() {
+  if (sparseFile == nullptr) {
+    fprintf(stderr, "Cannot close unopened matrix\n");
+    exit(1);
+  }
+  fclose(sparseFile);
+  sparseFile = nullptr;
+  sparseFilename = nullptr;
+}
+
+// Helper method to read sparse matrix filenames from the environment, defined
+// with the naming convention ${SPARSE_MATRIX0}, ${SPARSE_MATRIX1}, etc.
+extern "C" char *getSparseMatrix(uint64_t id) {
+  char var[80];
+  sprintf(var, "SPARSE_MATRIX%lu", id);
+  char *env = getenv(var);
+  return env;
+}