[mlir][sparse] add asserts on reading in tensor data
[lldb.git] / mlir / lib / ExecutionEngine / SparseUtils.cpp
1 //===- SparseUtils.cpp - Sparse Utils for MLIR execution ------------------===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 //
9 // This file implements a light-weight runtime support library that is useful
10 // for sparse tensor manipulations. The functionality provided in this library
11 // is meant to simplify benchmarking, testing, and debugging MLIR code that
12 // operates on sparse tensors. The provided functionality is **not** part
13 // of core MLIR, however.
14 //
15 //===----------------------------------------------------------------------===//
16
17 #include "mlir/ExecutionEngine/CRunnerUtils.h"
18
19 #ifdef MLIR_CRUNNERUTILS_DEFINE_FUNCTIONS
20
21 #include <algorithm>
22 #include <cassert>
23 #include <cctype>
24 #include <cinttypes>
25 #include <cstdio>
26 #include <cstdlib>
27 #include <cstring>
28 #include <vector>
29
30 //===----------------------------------------------------------------------===//
31 //
32 // Internal support for reading sparse tensors in one of the following
33 // external file formats:
34 //
35 // (1) Matrix Market Exchange (MME): *.mtx
36 //     https://math.nist.gov/MatrixMarket/formats.html
37 //
38 // (2) Formidable Repository of Open Sparse Tensors and Tools (FROSTT): *.tns
39 //     http://frostt.io/tensors/file-formats.html
40 //
41 //===----------------------------------------------------------------------===//
42
43 namespace {
44
45 /// A sparse tensor element in coordinate scheme (value and indices).
46 /// For example, a rank-1 vector element would look like
47 ///   ({i}, a[i])
48 /// and a rank-5 tensor element like
49 ///   ({i,j,k,l,m}, a[i,j,k,l,m])
50 struct Element {
51   Element(const std::vector<uint64_t> &ind, double val)
52       : indices(ind), value(val){};
53   std::vector<uint64_t> indices;
54   double value;
55 };
56
57 /// A memory-resident sparse tensor in coordinate scheme (collection of
58 /// elements). This data structure is used to read a sparse tensor from
59 /// external file format into memory and sort the elements lexicographically
60 /// by indices before passing it back to the client (most packed storage
61 /// formats require the elements to appear in lexicographic index order).
62 struct SparseTensor {
63 public:
64   SparseTensor(const std::vector<uint64_t> &szs, uint64_t capacity)
65       : sizes(szs), pos(0) {
66     elements.reserve(capacity);
67   }
68   // Add element as indices and value.
69   void add(const std::vector<uint64_t> &ind, double val) {
70     assert(sizes.size() == ind.size());
71     for (int64_t r = 0, rank = sizes.size(); r < rank; r++)
72       assert(ind[r] < sizes[r]); // within bounds
73     elements.emplace_back(Element(ind, val));
74   }
75   // Sort elements lexicographically by index.
76   void sort() { std::sort(elements.begin(), elements.end(), lexOrder); }
77   // Primitive one-time iteration.
78   const Element &next() { return elements[pos++]; }
79
80 private:
81   // Returns true if indices of e1 < indices of e2.
82   static bool lexOrder(const Element &e1, const Element &e2) {
83     assert(e1.indices.size() == e2.indices.size());
84     for (int64_t r = 0, rank = e1.indices.size(); r < rank; r++) {
85       if (e1.indices[r] == e2.indices[r])
86         continue;
87       return e1.indices[r] < e2.indices[r];
88     }
89     return false;
90   }
91
92   std::vector<uint64_t> sizes; // per-rank dimension sizes
93   std::vector<Element> elements;
94   uint64_t pos;
95 };
96
97 /// Helper to convert string to lower case.
98 static char *toLower(char *token) {
99   for (char *c = token; *c; c++)
100     *c = tolower(*c);
101   return token;
102 }
103
104 /// Read the MME header of a general sparse matrix of type real.
105 static void readMMEHeader(FILE *file, char *name, uint64_t *idata) {
106   char line[1025];
107   char header[64];
108   char object[64];
109   char format[64];
110   char field[64];
111   char symmetry[64];
112   // Read header line.
113   if (fscanf(file, "%63s %63s %63s %63s %63s\n", header, object, format, field,
114              symmetry) != 5) {
115     fprintf(stderr, "Corrupt header in %s\n", name);
116     exit(1);
117   }
118   // Make sure this is a general sparse matrix.
119   if (strcmp(toLower(header), "%%matrixmarket") ||
120       strcmp(toLower(object), "matrix") ||
121       strcmp(toLower(format), "coordinate") || strcmp(toLower(field), "real") ||
122       strcmp(toLower(symmetry), "general")) {
123     fprintf(stderr,
124             "Cannot find a general sparse matrix with type real in %s\n", name);
125     exit(1);
126   }
127   // Skip comments.
128   while (1) {
129     if (!fgets(line, 1025, file)) {
130       fprintf(stderr, "Cannot find data in %s\n", name);
131       exit(1);
132     }
133     if (line[0] != '%')
134       break;
135   }
136   // Next line contains M N NNZ.
137   idata[0] = 2; // rank
138   if (sscanf(line, "%" PRIu64 "%" PRIu64 "%" PRIu64 "\n", idata + 2, idata + 3,
139              idata + 1) != 3) {
140     fprintf(stderr, "Cannot find size in %s\n", name);
141     exit(1);
142   }
143 }
144
145 /// Read the "extended" FROSTT header. Although not part of the documented
146 /// format, we assume that the file starts with optional comments followed
147 /// by two lines that define the rank, the number of nonzeros, and the
148 /// dimensions sizes (one per rank) of the sparse tensor.
149 static void readExtFROSTTHeader(FILE *file, char *name, uint64_t *idata) {
150   char line[1025];
151   // Skip comments.
152   while (1) {
153     if (!fgets(line, 1025, file)) {
154       fprintf(stderr, "Cannot find data in %s\n", name);
155       exit(1);
156     }
157     if (line[0] != '#')
158       break;
159   }
160   // Next line contains RANK and NNZ.
161   if (sscanf(line, "%" PRIu64 "%" PRIu64 "\n", idata, idata + 1) != 2) {
162     fprintf(stderr, "Cannot find metadata in %s\n", name);
163     exit(1);
164   }
165   // Followed by a line with the dimension sizes (one per rank).
166   for (uint64_t r = 0; r < idata[0]; r++) {
167     if (fscanf(file, "%" PRIu64, idata + 2 + r) != 1) {
168       fprintf(stderr, "Cannot find dimension size %s\n", name);
169       exit(1);
170     }
171   }
172 }
173
174 } // anonymous namespace
175
176 //===----------------------------------------------------------------------===//
177 //
178 // Public API of the sparse runtime support library that enables MLIR code
179 // to read a sparse tensor from an external format (MME for FROSTT).
180 //
181 // For example, a sparse matrix in MME can be read as follows.
182 //
183 //   %tensor = call @openTensor(%fileName, %idata)
184 //     : (!llvm.ptr<i8>, memref<?xindex>) -> (!llvm.ptr<i8>)
185 //   %rank = load %idata[%c0] : memref<?xindex>    # always 2 for MME
186 //   %nnz  = load %idata[%c1] : memref<?xindex>
187 //   %m    = load %idata[%c2] : memref<?xindex>
188 //   %n    = load %idata[%c3] : memref<?xindex>
189 //   .. prepare reading in m x n sparse tensor A with nnz nonzero elements ..
190 //   scf.for %k = %c0 to %nnz step %c1 {
191 //     call @readTensorItem(%tensor, %idata, %ddata)
192 //       : (!llvm.ptr<i8>, memref<?xindex>, memref<?xf64>) -> ()
193 //     %i = load %idata[%c0] : memref<?xindex>
194 //     %j = load %idata[%c1] : memref<?xindex>
195 //     %d = load %ddata[%c0] : memref<?xf64>
196 //     .. process next nonzero element A[i][j] = d
197 //        where the elements appear in lexicographic order ..
198 //   }
199 //   call @closeTensor(%tensor) : (!llvm.ptr<i8>) -> ()
200 //
201 //
202 // Note that input parameters in the "MLIRized" version of a function mimic
203 // the data layout of a MemRef<?xT>:
204 //
205 //   struct MemRef {
206 //     T *base;
207 //     T *data;
208 //     int64_t off;
209 //     int64_t sizes[1];
210 //     int64_t strides[1];
211 //   }
212 //
213 //===----------------------------------------------------------------------===//
214
215 /// Reads in a sparse tensor with the given filename. The call yields a
216 /// pointer to an opaque memory-resident sparse tensor object that is only
217 /// understood by other methods in the sparse runtime support library. An
218 /// array parameter is used to pass the rank, the number of nonzero elements,
219 /// and the dimension sizes (one per rank).
220 extern "C" void *openTensorC(char *filename, uint64_t *idata) {
221   // Open the file.
222   FILE *file = fopen(filename, "r");
223   if (!file) {
224     fprintf(stderr, "Cannot find %s\n", filename);
225     exit(1);
226   }
227   // Perform some file format dependent set up.
228   if (strstr(filename, ".mtx")) {
229     readMMEHeader(file, filename, idata);
230   } else if (strstr(filename, ".tns")) {
231     readExtFROSTTHeader(file, filename, idata);
232   } else {
233     fprintf(stderr, "Unknown format %s\n", filename);
234     exit(1);
235   }
236   // Prepare sparse tensor object with per-rank dimension sizes
237   // and the number of nonzeros as initial capacity.
238   uint64_t rank = idata[0];
239   uint64_t nnz = idata[1];
240   std::vector<uint64_t> indices(rank);
241   for (uint64_t r = 0; r < rank; r++)
242     indices[r] = idata[2 + r];
243   SparseTensor *tensor = new SparseTensor(indices, nnz);
244   // Read all nonzero elements.
245   for (uint64_t k = 0; k < nnz; k++) {
246     for (uint64_t r = 0; r < rank; r++) {
247       if (fscanf(file, "%" PRIu64, &indices[r]) != 1) {
248         fprintf(stderr, "Cannot find next index in %s\n", filename);
249         exit(1);
250       }
251       indices[r]--; // 0-based index
252     }
253     double value;
254     if (fscanf(file, "%lg\n", &value) != 1) {
255       fprintf(stderr, "Cannot find next value in %s\n", filename);
256       exit(1);
257     }
258     tensor->add(indices, value);
259   }
260   // Close the file and return sorted tensor.
261   fclose(file);
262   tensor->sort(); // sort lexicographically
263   return tensor;
264 }
265
266 /// "MLIRized" version.
267 extern "C" void *openTensor(char *filename, uint64_t *ibase, uint64_t *idata,
268                             uint64_t ioff, uint64_t isize, uint64_t istride) {
269   assert(istride == 1);
270   return openTensorC(filename, idata + ioff);
271 }
272
273 /// Yields the next element from the given opaque sparse tensor object.
274 extern "C" void readTensorItemC(void *tensor, uint64_t *idata, double *ddata) {
275   const Element &e = static_cast<SparseTensor *>(tensor)->next();
276   for (uint64_t r = 0, rank = e.indices.size(); r < rank; r++)
277     idata[r] = e.indices[r];
278   ddata[0] = e.value;
279 }
280
281 /// "MLIRized" version.
282 extern "C" void readTensorItem(void *tensor, uint64_t *ibase, uint64_t *idata,
283                                uint64_t ioff, uint64_t isize, uint64_t istride,
284                                double *dbase, double *ddata, uint64_t doff,
285                                uint64_t dsize, uint64_t dstride) {
286   assert(istride == 1 && dstride == 1);
287   readTensorItemC(tensor, idata + ioff, ddata + doff);
288 }
289
290 /// Closes the given opaque sparse tensor object, releasing its memory
291 /// resources. After this call, the opague object cannot be used anymore.
292 extern "C" void closeTensor(void *tensor) {
293   delete static_cast<SparseTensor *>(tensor);
294 }
295
296 /// Helper method to read a sparse tensor filename from the environment,
297 /// defined with the naming convention ${TENSOR0}, ${TENSOR1}, etc.
298 extern "C" char *getTensorFilename(uint64_t id) {
299   char var[80];
300   sprintf(var, "TENSOR%" PRIu64, id);
301   char *env = getenv(var);
302   return env;
303 }
304
305 #endif // MLIR_CRUNNERUTILS_DEFINE_FUNCTIONS