[mlir][sparse] export sparse tensor runtime support through header file

Exposing the C versions of the methods of the sparse runtime support lib
through header files will enable using the same methods in an MLIR program
as well as a C++ program, which will simplify future benchmarking comparisons
(e.g. comparing MLIR generated code with eigen for Matrix Market sparse matrices).

Reviewed By: penpornk

Differential Revision: https://reviews.llvm.org/D91316

diff --git a/mlir/include/mlir/ExecutionEngine/CRunnerUtils.h b/mlir/include/mlir/ExecutionEngine/CRunnerUtils.h
index 72df9ea..edcd8e0 100644 (file)
--- a/mlir/include/mlir/ExecutionEngine/CRunnerUtils.h
+++ b/mlir/include/mlir/ExecutionEngine/CRunnerUtils.h
@@ -209,5 +209,16 @@ extern "C" MLIR_CRUNNERUTILS_EXPORT void printClose();
 extern "C" MLIR_CRUNNERUTILS_EXPORT void printComma();
 extern "C" MLIR_CRUNNERUTILS_EXPORT void printNewline();
 
-#endif // EXECUTIONENGINE_CRUNNERUTILS_H_
+//===----------------------------------------------------------------------===//
+// Small runtime support for sparse tensors.
+//===----------------------------------------------------------------------===//
+extern "C" MLIR_CRUNNERUTILS_EXPORT void openMatrixC(char *filename,
+                                                     uint64_t *mdata,
+                                                     uint64_t *ndata,
+                                                     uint64_t *nnzdata);
+extern "C" MLIR_CRUNNERUTILS_EXPORT void
+readMatrixItemC(uint64_t *idata, uint64_t *jdata, double *ddata);
+extern "C" MLIR_CRUNNERUTILS_EXPORT void closeMatrix();
+extern "C" MLIR_CRUNNERUTILS_EXPORT char *getMatrix(uint64_t id);
 
+#endif // EXECUTIONENGINE_CRUNNERUTILS_H_

diff --git a/mlir/lib/ExecutionEngine/CRunnerUtils.cpp b/mlir/lib/ExecutionEngine/CRunnerUtils.cpp
index a6f2423..c3baa8d 100644 (file)
--- a/mlir/lib/ExecutionEngine/CRunnerUtils.cpp
+++ b/mlir/lib/ExecutionEngine/CRunnerUtils.cpp
@@ -32,4 +32,4 @@ extern "C" void printClose() { fputs(" )", stdout); }
 extern "C" void printComma() { fputs(", ", stdout); }
 extern "C" void printNewline() { fputc('\n', stdout); }
 
-#endif
+#endif // MLIR_CRUNNERUTILS_DEFINE_FUNCTIONS

diff --git a/mlir/lib/ExecutionEngine/SparseUtils.cpp b/mlir/lib/ExecutionEngine/SparseUtils.cpp
index 253a3f4..870cf4f 100644 (file)
--- a/mlir/lib/ExecutionEngine/SparseUtils.cpp
+++ b/mlir/lib/ExecutionEngine/SparseUtils.cpp
@@ -14,6 +14,10 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "mlir/ExecutionEngine/CRunnerUtils.h"
+
+#ifdef MLIR_CRUNNERUTILS_DEFINE_FUNCTIONS
+
 #include <cctype>
 #include <cinttypes>
 #include <cstdio>
@@ -112,22 +116,23 @@ static void readItem(FILE *file, char *name, uint64_t *i, uint64_t *j,
 // 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>:
+// Note that input parameters in the "MLIRized" version of a function mimic
+// the data 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) {
+extern "C" void openMatrixC(char *filename, uint64_t *mdata, uint64_t *ndata,
+                            uint64_t *nnzdata) {
   if (sparseFile != nullptr) {
     fprintf(stderr, "Other file still open %s vs. %s\n", sparseFilename,
             filename);
@@ -142,9 +147,16 @@ extern "C" void openMatrix(char *filename, uint64_t *mbase, uint64_t *mdata,
   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) {
+// "MLIRized" version.
+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) {
+  openMatrixC(filename, mdata, ndata, nnzdata);
+}
+
+extern "C" void readMatrixItemC(uint64_t *idata, uint64_t *jdata,
+                                double *ddata) {
   if (sparseFile == nullptr) {
     fprintf(stderr, "Cannot read item from unopened matrix\n");
     exit(1);
@@ -152,6 +164,13 @@ extern "C" void readMatrixItem(uint64_t *ibase, uint64_t *idata, int64_t ioff,
   readItem(sparseFile, sparseFilename, idata, jdata, ddata);
 }
 
+// "MLIRized" version.
+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) {
+  readMatrixItemC(idata, jdata, ddata);
+}
+
 extern "C" void closeMatrix() {
   if (sparseFile == nullptr) {
     fprintf(stderr, "Cannot close unopened matrix\n");
@@ -170,3 +189,5 @@ extern "C" char *getMatrix(uint64_t id) {
   char *env = getenv(var);
   return env;
 }
+
+#endif // MLIR_CRUNNERUTILS_DEFINE_FUNCTIONS