diff --git a/mlir/integration_test/CMakeLists.txt b/mlir/integration_test/CMakeLists.txt --- a/mlir/integration_test/CMakeLists.txt +++ b/mlir/integration_test/CMakeLists.txt @@ -28,3 +28,7 @@ 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/) diff --git a/mlir/integration_test/Sparse/CPU/lit.local.cfg b/mlir/integration_test/Sparse/CPU/lit.local.cfg new file mode 100644 --- /dev/null +++ b/mlir/integration_test/Sparse/CPU/lit.local.cfg @@ -0,0 +1,5 @@ +import sys + +# No JIT on win32. +if sys.platform == 'win32': + config.unsupported = True diff --git a/mlir/integration_test/Sparse/CPU/matrix-market-example.mlir b/mlir/integration_test/Sparse/CPU/matrix-market-example.mlir new file mode 100644 --- /dev/null +++ b/mlir/integration_test/Sparse/CPU/matrix-market-example.mlir @@ -0,0 +1,88 @@ +// 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, memref, memref, memref) -> () + func @readMatrixItem(memref, memref, memref) -> () + func @closeMatrix() -> () + func @getSparseMatrix(index) -> (!llvm.ptr) + + func @entry() { + %d0 = constant 0.0 : f64 + %c0 = constant 0 : index + %c1 = constant 1 : index + %c5 = constant 5 : index + %m = alloc() : memref + %n = alloc() : memref + %nnz = alloc() : memref + %i = alloc() : memref + %j = alloc() : memref + %d = alloc() : memref + %a = alloc() : memref<5x5xf64> + scf.for %ii = %c0 to %c5 step %c1 { + scf.for %jj = %c0 to %c5 step %c1 { + store %d0, %a[%ii, %jj] : memref<5x5xf64> + } + } + + // + // Read the sparse matrix. + // + %file = call @getSparseMatrix(%c0) : (index) -> (!llvm.ptr) + call @openMatrix(%file, %m, %n, %nnz) + : (!llvm.ptr, memref, + memref, memref) -> () + %u = load %nnz[] : memref + scf.for %k = %c0 to %u step %c1 { + call @readMatrixItem(%i, %j, %d) + : (memref, memref, memref) -> () + %idx = load %i[] : memref + %jdx = load %j[] : memref + %val = load %d[] : memref + store %val, %a[%idx, %jdx] : memref<5x5xf64> + } + call @closeMatrix() : () -> () + + // + // Test data. + // + %0 = load %m[] : memref + %1 = load %n[] : memref + %2 = load %nnz[] : memref + %A = vector.transfer_read %a[%c0, %c0], %d0 + : memref<5x5xf64>, vector<5x5xf64> + vector.print %0 : index + vector.print %1 : index + vector.print %2 : 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 + dealloc %n : memref + dealloc %nnz : memref + dealloc %i : memref + dealloc %j : memref + dealloc %d : memref + dealloc %a : memref<5x5xf64> + + return + } +} diff --git a/mlir/integration_test/data/test.mtx b/mlir/integration_test/data/test.mtx new file mode 100644 --- /dev/null +++ b/mlir/integration_test/data/test.mtx @@ -0,0 +1,15 @@ +%%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 diff --git a/mlir/lib/ExecutionEngine/CMakeLists.txt b/mlir/lib/ExecutionEngine/CMakeLists.txt --- a/mlir/lib/ExecutionEngine/CMakeLists.txt +++ b/mlir/lib/ExecutionEngine/CMakeLists.txt @@ -3,6 +3,7 @@ set(LLVM_OPTIONAL_SOURCES CRunnerUtils.cpp + SparseUtils.cpp ExecutionEngine.cpp RunnerUtils.cpp OptUtils.cpp @@ -70,6 +71,7 @@ add_mlir_library(mlir_c_runner_utils SHARED CRunnerUtils.cpp + SparseUtils.cpp EXCLUDE_FROM_LIBMLIR ) @@ -77,6 +79,7 @@ add_mlir_library(mlir_c_runner_utils_static CRunnerUtils.cpp + SparseUtils.cpp EXCLUDE_FROM_LIBMLIR ) diff --git a/mlir/lib/ExecutionEngine/SparseUtils.cpp b/mlir/lib/ExecutionEngine/SparseUtils.cpp new file mode 100644 --- /dev/null +++ b/mlir/lib/ExecutionEngine/SparseUtils.cpp @@ -0,0 +1,172 @@ +//===- 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 +#include +#include +#include +#include + +//===----------------------------------------------------------------------===// +// +// 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, +// memref, +// memref, +// memref) -> () +// .... prepare reading in m x n matrix A with nnz nonzero elements .... +// %u = load %nnz[] : memref +// scf.for %k = %c0 to %u step %c1 { +// call @readMatrixItem(%i, %j, %d) : (memref, +// memref, memref) -> () +// .... 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: +// 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; +}