diff --git a/mlir/lib/ExecutionEngine/SparseUtils.cpp b/mlir/lib/ExecutionEngine/SparseUtils.cpp --- a/mlir/lib/ExecutionEngine/SparseUtils.cpp +++ b/mlir/lib/ExecutionEngine/SparseUtils.cpp @@ -48,9 +48,9 @@ /// and a rank-5 tensor element like /// ({i,j,k,l,m}, a[i,j,k,l,m]) struct Element { - Element(const std::vector &ind, double val) + Element(const std::vector &ind, double val) : indices(ind), value(val){}; - std::vector indices; + std::vector indices; double value; }; @@ -61,9 +61,15 @@ /// formats require the elements to appear in lexicographic index order). struct SparseTensor { public: - SparseTensor(int64_t capacity) : pos(0) { elements.reserve(capacity); } + SparseTensor(const std::vector &szs, uint64_t capacity) + : sizes(szs), pos(0) { + elements.reserve(capacity); + } // Add element as indices and value. - void add(const std::vector &ind, double val) { + void add(const std::vector &ind, double val) { + assert(sizes.size() == ind.size()); + for (int64_t r = 0, rank = sizes.size(); r < rank; r++) + assert(ind[r] < sizes[r]); // within bounds elements.emplace_back(Element(ind, val)); } // Sort elements lexicographically by index. @@ -82,6 +88,8 @@ } return false; } + + std::vector sizes; // per-rank dimension sizes std::vector elements; uint64_t pos; }; @@ -225,20 +233,24 @@ fprintf(stderr, "Unknown format %s\n", filename); exit(1); } - // Read all nonzero elements. + // Prepare sparse tensor object with per-rank dimension sizes + // and the number of nonzeros as initial capacity. uint64_t rank = idata[0]; uint64_t nnz = idata[1]; - SparseTensor *tensor = new SparseTensor(nnz); - std::vector indices(rank); - double value; + std::vector indices(rank); + for (uint64_t r = 0; r < rank; r++) + indices[r] = idata[2 + r]; + SparseTensor *tensor = new SparseTensor(indices, nnz); + // Read all nonzero elements. for (uint64_t k = 0; k < nnz; k++) { for (uint64_t r = 0; r < rank; r++) { - if (fscanf(file, "%" PRId64, &indices[r]) != 1) { + if (fscanf(file, "%" PRIu64, &indices[r]) != 1) { fprintf(stderr, "Cannot find next index in %s\n", filename); exit(1); } indices[r]--; // 0-based index } + double value; if (fscanf(file, "%lg\n", &value) != 1) { fprintf(stderr, "Cannot find next value in %s\n", filename); exit(1);