[mlir][sparse] Breaking up openSparseTensor to better support non-permutations

[wrengr]

Summary

This commit updates how the SparseTensorConversion pass handles NewOp. It breaks up the underlying openSparseTensor function into two parts (SparseTensorReader::create and SparseTensorReader::readSparseTensor) so that the pass can inject code for constructing lvlSizes between those two parts. Migrating the construction of lvlSizes out of the runtime and into the pass is a necessary first step toward fully supporting non-permutations. (The alternative would be for the pass to generate a FuncOp for performing the construction and then passing that to the runtime; which doesn't seem to have any benefits over the design of this commit.) And since the pass now generates the code to call these two functions, this change also removes the Action::kFromFile value from the enum used by _mlir_ciface_newSparseTensor.

[wrengr]

Fixing a logic bug in SparseTensorReader::readCOO. This bug does not affect permutations (hence why not discovering it earlier), but would affect non-permutations.

[aartbik]

mlir/lib/ExecutionEngine/SparseTensorRuntime.cpp

656

agreed; being on the boundary of MLIR calling convention and nice C++ templating is really annoying sometimes ;-) it would be even nicer to be able to generate the lib "on demand" so that we only have the cases needed by the compiler (and also generate cases not covered yet)

[wrengr]

rebase

[wrengr]

rebase

[wrengr]

correcting rewriting_for_codegen.mlir post-rebase

mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h

 
 constexpr bool isComplexPrimaryType(PrimaryType valTy) {
   return PrimaryType::kC64 <= valTy && valTy <= PrimaryType::kC32;
 }
 
 /// The actions performed by @newSparseTensor.
 enum class Action : uint32_t {
   kEmpty = 0,
-  kFromFile = 1,
+  // newSparseTensor no longer handles `kFromFile=1`, so we leave this
+  // number reserved to help catch any code that still needs updating.
   kFromCOO = 2,
   kSparseToSparse = 3,
   kEmptyCOO = 4,
   kToCOO = 5,
   kToIterator = 6,
 };
 
 /// This enum defines all the sparse representations supportable by

mlir/include/mlir/ExecutionEngine/SparseTensor/File.h

public:
   explicit SparseTensorReader(const char *filename) : filename(filename) {
     assert(filename && "Received nullptr for filename");
   }
 
   // Disallows copying, to avoid duplicating the `file` pointer.
   SparseTensorReader(const SparseTensorReader &) = delete;
   SparseTensorReader &operator=(const SparseTensorReader &) = delete;
 
+  /// Factory method to allocate a new reader, open the file, read the
+  /// header, and validate that the actual contents of the file match
+  /// the expected `dimShape` and `valTp`.
+  static SparseTensorReader *create(const char *filename, uint64_t dimRank,
+                                    const uint64_t *dimShape,
+                                    PrimaryType valTp) {
+    SparseTensorReader *reader = new SparseTensorReader(filename);
+    reader->openFile();
+    reader->readHeader();
+    if (!reader->canReadAs(valTp))
+      MLIR_SPARSETENSOR_FATAL(
+          "Tensor element type %d not compatible with values in file %s\n",
+          static_cast<int>(valTp), filename);
+    reader->assertMatchesShape(dimRank, dimShape);
+    return reader;
+  }
+
   // This dtor tries to avoid leaking the `file`.  (Though it's better
   // to call `closeFile` explicitly when possible, since there are
   // circumstances where dtors are not called reliably.)
   ~SparseTensorReader() { closeFile(); }
 
   /// Opens the file for reading.
   void openFile();
 
public:
   /// valid after parsing the header.
   void assertMatchesShape(uint64_t rank, const uint64_t *shape) const;
 
   /// Reads a sparse tensor element from the next line in the input file and
   /// returns the value of the element. Stores the coordinates of the element
   /// to the `indices` array.
   template <typename V>
   V readCOOElement(uint64_t rank, uint64_t *indices) {
-    char *linePtr = readCOOIndices(rank, indices);
+    assert(rank == getRank() && "rank mismatch");
+    char *linePtr = readCOOIndices(indices);
     return detail::readCOOValue<V>(&linePtr, isPattern());
   }
 
+  /// Allocates a new COO object for `lvlSizes`, initializes it by reading
+  /// all the elements from the file and applying `dim2lvl` to their indices,
+  /// and then closes the file.
+  ///
+  /// Preconditions:
+  /// * `lvlSizes` must be valid for `lvlRank`.
+  /// * `dim2lvl` must be valid for `getRank()`.
+  /// * `dim2lvl` maps indices valid for `getDimSizes()` to indices
+  ///   valid for `lvlSizes`.
+  /// * the file's actual value type can be read as `V`.
+  ///
+  /// Asserts:
+  /// * `isValid()`
+  /// * `dim2lvl` is a permutation, and therefore also `lvlRank == getRank()`.
+  ///   (This requirement will be lifted once we functionalize `dim2lvl`.)
+  //
+  // NOTE: This method is factored out of `readSparseTensor` primarily to
+  // reduce code bloat (since the bulk of the code doesn't care about the
+  // `<P,I>` type template parameters).  But we leave it public since it's
+  // perfectly reasonable for clients to use.
+  template <typename V>
+  SparseTensorCOO<V> *readCOO(uint64_t lvlRank, const uint64_t *lvlSizes,
+                              const uint64_t *dim2lvl);
+
+  /// Allocates a new sparse-tensor storage object with the given encoding,
+  /// initializes it by reading all the elements from the file, and then
+  /// closes the file.  Preconditions/assertions are as per `readCOO`
+  /// and `SparseTensorStorage::newFromCOO`.
+  template <typename P, typename I, typename V>
+  SparseTensorStorage<P, I, V> *
+  readSparseTensor(uint64_t lvlRank, const uint64_t *lvlSizes,
+                   const DimLevelType *lvlTypes, const uint64_t *lvl2dim,
+                   const uint64_t *dim2lvl) {
+    auto *lvlCOO = readCOO<V>(lvlRank, lvlSizes, dim2lvl);
+    auto *tensor = SparseTensorStorage<P, I, V>::newFromCOO(
+        getRank(), getDimSizes(), lvlRank, lvlTypes, lvl2dim, *lvlCOO);
+    delete lvlCOO;
+    return tensor;
+  }
+
 private:
   /// Attempts to read a line from the file.  Is private because there's
   /// no reason for client code to call it.
   void readLine();
 
   /// Reads the next line of the input file and parses the coordinates
   /// into the `indices` argument.  Returns the position in the `line`
   /// buffer where the element's value should be parsed from.  This method
   /// has been factored out from `readCOOElement` to minimize code bloat
   /// for the generated library.
-  char *readCOOIndices(uint64_t rank, uint64_t *indices);
+  ///
+  /// Precondition: `indices` is valid for `getRank()`.
+  char *readCOOIndices(uint64_t *indices);
 
   /// Reads the MME header of a general sparse matrix of type real.
   void readMMEHeader();
 
   /// Reads the "extended" FROSTT header. Although not part of the
   /// documented format, we assume that the file starts with optional
   /// comments followed by two lines that define the rank, the number of
   /// nonzeros, and the dimensions sizes (one per rank) of the sparse tensor.
   void readExtFROSTTHeader();
 
   static constexpr int kColWidth = 1025;
   const char *const filename;
   FILE *file = nullptr;
   ValueKind valueKind_ = ValueKind::kInvalid;
   bool isSymmetric_ = false;
   uint64_t idata[512];
   char line[kColWidth];
 };
 
 //===----------------------------------------------------------------------===//
 
-/// Reads a sparse tensor with the given filename into a memory-resident
-/// sparse tensor.
-///
-/// Preconditions:
-/// * `dimShape` and `dim2lvl` must be valid for `dimRank`.
-/// * `lvlTypes` and `lvl2dim` must be valid for `lvlRank`.
-/// * `dim2lvl` is the inverse of `lvl2dim`.
-///
-/// Asserts:
-/// * the file's actual value type can be read as `valTp`.
-/// * the file's actual dimension-sizes match the expected `dimShape`.
-/// * `dim2lvl` is a permutation, and therefore also `dimRank == lvlRank`.
-//
-// TODO: As currently written, this function uses `dim2lvl` in two
-// places: first, to construct the level-sizes from the file's actual
-// dimension-sizes; and second, to map the file's dimension-indices into
-// level-indices.  The latter can easily generalize to arbitrary mappings,
-// however the former cannot.  Thus, once we functionalize the mappings,
-// this function will need both the sizes-to-sizes and indices-to-indices
-// variants of the `dim2lvl` mapping.  For the `lvl2dim` direction we only
-// need the indices-to-indices variant, for handing off to `newFromCOO`.
-template <typename P, typename I, typename V>
-inline SparseTensorStorage<P, I, V> *
-openSparseTensor(uint64_t dimRank, const uint64_t *dimShape, uint64_t lvlRank,
-                 const DimLevelType *lvlTypes, const uint64_t *lvl2dim,
-                 const uint64_t *dim2lvl, const char *filename,
-                 PrimaryType valTp) {
-  // Read the file's header and check the file's actual element type and
-  // dimension-sizes against the expected element type and dimension-shape.
-  SparseTensorReader stfile(filename);
-  stfile.openFile();
-  stfile.readHeader();
-  if (!stfile.canReadAs(valTp))
-    MLIR_SPARSETENSOR_FATAL(
-        "Tensor element type %d not compatible with values in file %s\n",
-        static_cast<int>(valTp), filename);
-  stfile.assertMatchesShape(dimRank, dimShape);
-  const uint64_t *dimSizes = stfile.getDimSizes();
-  // Construct the level-sizes from the file's dimension-sizes
-  // TODO: This doesn't generalize to arbitrary mappings. (See above.)
-  assert(dimRank == lvlRank && "Rank mismatch");
+template <typename V>
+SparseTensorCOO<V> *SparseTensorReader::readCOO(uint64_t lvlRank,
+                                                const uint64_t *lvlSizes,
+                                                const uint64_t *dim2lvl) {
+  assert(isValid() && "Attempt to readCOO() before readHeader()");
+  // Construct a `PermutationRef` for the `pushforward` below.
+  // TODO: This specific implementation does not generalize to arbitrary
+  // mappings, but once we functionalize the `dim2lvl` argument we can
+  // simply use that function instead.
+  const uint64_t dimRank = getRank();
+  assert(lvlRank == dimRank && "Rank mismatch");
   detail::PermutationRef d2l(dimRank, dim2lvl);
-  std::vector<uint64_t> lvlSizes = d2l.pushforward(dimRank, dimSizes);
   // Prepare a COO object with the number of nonzeros as initial capacity.
-  uint64_t nnz = stfile.getNNZ();
-  auto *lvlCOO = new SparseTensorCOO<V>(lvlSizes, nnz);
+  const uint64_t nnz = getNNZ();
+  auto *lvlCOO = new SparseTensorCOO<V>(lvlRank, lvlSizes, nnz);
   // Read all nonzero elements.
   std::vector<uint64_t> dimInd(dimRank);
   std::vector<uint64_t> lvlInd(lvlRank);
+  // Do some manual LICM, to avoid assertions in the for-loop.
+  const bool addSymmetric = (isSymmetric() && dimRank == 2);
+  const bool isPattern_ = isPattern();
   for (uint64_t k = 0; k < nnz; ++k) {
-    const V value = stfile.readCOOElement<V>(dimRank, dimInd.data());
+    // We inline `readCOOElement` here in order to avoid redundant
+    // assertions, since they're guaranteed by the call to `isValid()`
+    // and the construction of `dimInd` above.
+    char *linePtr = readCOOIndices(dimInd.data());
+    const V value = detail::readCOOValue<V>(&linePtr, isPattern_);
     d2l.pushforward(dimRank, dimInd.data(), lvlInd.data());
     // TODO: <https://github.com/llvm/llvm-project/issues/54179>
     lvlCOO->add(lvlInd, value);
     // We currently chose to deal with symmetric matrices by fully
     // constructing them.  In the future, we may want to make symmetry
     // implicit for storage reasons.
-    if (stfile.isSymmetric() && lvlInd[0] != lvlInd[1])
+    if (addSymmetric && lvlInd[0] != lvlInd[1])
       lvlCOO->add({lvlInd[1], lvlInd[0]}, value);
   }
-  // Close the file, convert the COO to SparseTensorStorage, and return.
-  stfile.closeFile();
-  auto *tensor = SparseTensorStorage<P, I, V>::newFromCOO(
-      dimRank, dimSizes, lvlRank, lvlTypes, lvl2dim, *lvlCOO);
-  delete lvlCOO;
-  return tensor;
+  // Close the file and return the COO.
+  closeFile();
+  return lvlCOO;
 }
 
 /// Writes the sparse tensor to `filename` in extended FROSTT format.
 template <typename V>
 inline void writeExtFROSTT(const SparseTensorCOO<V> &coo,
                            const char *filename) {
   assert(filename && "Got nullptr for filename");
   auto &dimSizes = coo.getDimSizes();

mlir/include/mlir/ExecutionEngine/SparseTensorRuntime.h

MLIR_CRUNNERUTILS_EXPORT void convertFromMLIRSparseTensor##VNAME(            \
       void *tensor, uint64_t *pRank, uint64_t *pNse, uint64_t **pShape,        \
       V **pValues, uint64_t **pIndices);
 MLIR_SPARSETENSOR_FOREVERY_V(DECL_CONVERTFROMMLIRSPARSETENSOR)
 #undef DECL_CONVERTFROMMLIRSPARSETENSOR
 
 /// Creates a SparseTensorReader for reading a sparse tensor from a file with
 /// the given file name. This opens the file and read the header meta data based
 /// of the sparse tensor format derived from the suffix of the file name.
+//
+// FIXME: update `SparseTensorCodegenPass` to use
+// `_mlir_ciface_createCheckedSparseTensorReader` instead.
 MLIR_CRUNNERUTILS_EXPORT void *createSparseTensorReader(char *filename);
 
+/// Constructs a new SparseTensorReader object, opens the file, reads the
+/// header, and validates that the actual contents of the file match
+/// the expected `dimShapeRef` and `valTp`.
+MLIR_CRUNNERUTILS_EXPORT void *_mlir_ciface_createCheckedSparseTensorReader(
+    char *filename, StridedMemRefType<index_type, 1> *dimShapeRef,
+    PrimaryType valTp);
+
+/// Constructs a new sparse-tensor storage object with the given encoding,
+/// initializes it by reading all the elements from the file, and then
+/// closes the file.
+MLIR_CRUNNERUTILS_EXPORT void *_mlir_ciface_newSparseTensorFromReader(
+    void *p, StridedMemRefType<index_type, 1> *lvlSizesRef,
+    StridedMemRefType<DimLevelType, 1> *lvlTypesRef,
+    StridedMemRefType<index_type, 1> *lvl2dimRef,
+    StridedMemRefType<index_type, 1> *dim2lvlRef, OverheadType ptrTp,
+    OverheadType indTp, PrimaryType valTp);
+
 /// Returns the rank of the sparse tensor being read.
 MLIR_CRUNNERUTILS_EXPORT index_type getSparseTensorReaderRank(void *p);
 
 /// Returns the is_symmetric bit for the sparse tensor being read.
 MLIR_CRUNNERUTILS_EXPORT bool getSparseTensorReaderIsSymmetric(void *p);
 
 /// Returns the number of non-zero values for the sparse tensor being read.
 MLIR_CRUNNERUTILS_EXPORT index_type getSparseTensorReaderNNZ(void *p);
 
 /// Returns the size of a dimension for the sparse tensor being read.
 MLIR_CRUNNERUTILS_EXPORT index_type getSparseTensorReaderDimSize(void *p,
                                                                  index_type d);
 
-/// Returns all dimension sizes for the sparse tensor being read.
-MLIR_CRUNNERUTILS_EXPORT void _mlir_ciface_getSparseTensorReaderDimSizes(
+/// SparseTensorReader method to copy the dimension-sizes into the
+/// provided memref.
+//
+// FIXME: update `SparseTensorCodegenPass` to use
+// `_mlir_ciface_getSparseTensorReaderDimSizes` instead.
+MLIR_CRUNNERUTILS_EXPORT void _mlir_ciface_copySparseTensorReaderDimSizes(
     void *p, StridedMemRefType<index_type, 1> *dref);
 
+/// SparseTensorReader method to obtain direct access to the
+/// dimension-sizes array.
+MLIR_CRUNNERUTILS_EXPORT void _mlir_ciface_getSparseTensorReaderDimSizes(
+    StridedMemRefType<index_type, 1> *out, void *p);
+
 /// Releases the SparseTensorReader. This also closes the file associated with
 /// the reader.
 MLIR_CRUNNERUTILS_EXPORT void delSparseTensorReader(void *p);
 
 /// Returns the next element for the sparse tensor being read.
 #define DECL_GETNEXT(VNAME, V)                                                 \
   MLIR_CRUNNERUTILS_EXPORT void _mlir_ciface_getSparseTensorReaderNext##VNAME( \
       void *p, StridedMemRefType<index_type, 1> *iref,                         \

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorConversion.cpp

static Value sizeFromPtrAtDim(OpBuilder &builder, Location loc,
   return constantIndex(builder, loc, shape[i]);
 }
 
 /// Populates given sizes array from type (for static sizes) and from
 /// an already-converted opaque pointer source (for dynamic sizes).
 static void sizesFromPtr(OpBuilder &builder, SmallVectorImpl<Value> &sizes,
                          Location loc, SparseTensorEncodingAttr &enc,
                          ShapedType stp, Value src) {
-  for (unsigned i = 0, rank = stp.getRank(); i < rank; i++)
+  unsigned rank = stp.getRank();
+  sizes.reserve(rank);
+  for (unsigned i = 0; i < rank; i++)
     sizes.push_back(sizeFromPtrAtDim(builder, loc, enc, stp, src, i));
 }
 
 /// Populates given sizes array from type.
 static void sizesFromType(OpBuilder &builder, SmallVectorImpl<Value> &sizes,
                           Location loc, ShapedType stp) {
   auto shape = stp.getShape();
-  for (unsigned i = 0, rank = stp.getRank(); i < rank; i++) {
+  unsigned rank = stp.getRank();
+  sizes.reserve(rank);
+  for (unsigned i = 0; i < rank; i++) {
     uint64_t s = shape[i] == ShapedType::kDynamicSize ? 0 : shape[i];
     sizes.push_back(constantIndex(builder, loc, s));
   }
 }
 
 /// Populates the given sizes array for concatenation from type (for static
 /// sizes) and from an already-converted opaque pointer source (for dynamic
 /// sizes).
static Value genBuffer(OpBuilder &builder, Location loc, ValueRange values) {
   Value buffer = genAlloca(builder, loc, sz, values[0].getType());
   for (unsigned i = 0; i < sz; i++) {
     Value idx = constantIndex(builder, loc, i);
     builder.create<memref::StoreOp>(loc, values[i], buffer, idx);
   }
   return buffer;
 }
 
+/// Generates a temporary buffer for the level-types of the given encoding.
+static Value genLvlTypesBuffer(OpBuilder &builder, Location loc,
+                               SparseTensorEncodingAttr enc) {
+  SmallVector<Value> lvlTypes;
+  auto dlts = enc.getDimLevelType();
+  lvlTypes.reserve(dlts.size());
+  for (auto dlt : dlts)
+    lvlTypes.push_back(constantDimLevelTypeEncoding(builder, loc, dlt));
+  return genBuffer(builder, loc, lvlTypes);
+}
+
 /// This class abstracts over the API of `_mlir_ciface_newSparseTensor`:
 /// the "swiss army knife" method of the sparse runtime support library
 /// for materializing sparse tensors into the computation.  This abstraction
 /// reduces the need to make modifications to client code whenever that
 /// API changes.
 class NewCallParams final {
 public:
   /// Allocates the `ValueRange` for the `func::CallOp` parameters,

 // TODO: see the note at `_mlir_ciface_newSparseTensor` about how
 // the meaning of the various arguments (e.g., "sizes" vs "shapes")
 // is inconsistent between the different actions.
 NewCallParams &NewCallParams::genBuffers(SparseTensorEncodingAttr enc,
                                          ValueRange dimSizes, ShapedType stp) {
   const unsigned lvlRank = enc.getDimLevelType().size();
   const unsigned dimRank = stp.getRank();
   // Sparsity annotations.
-  SmallVector<Value> lvlTypes;
-  for (auto dlt : enc.getDimLevelType())
-    lvlTypes.push_back(constantDimLevelTypeEncoding(builder, loc, dlt));
-  assert(lvlTypes.size() == lvlRank && "Level-rank mismatch");
-  params[kParamLvlTypes] = genBuffer(builder, loc, lvlTypes);
+  params[kParamLvlTypes] = genLvlTypesBuffer(builder, loc, enc);
   // Dimension-sizes array of the enveloping tensor.  Useful for either
   // verification of external data, or for construction of internal data.
   assert(dimSizes.size() == dimRank && "Dimension-rank mismatch");
   params[kParamDimSizes] = genBuffer(builder, loc, dimSizes);
   // The level-sizes array must be passed as well, since for arbitrary
   // dim2lvl mappings it cannot be trivially reconstructed at runtime.
   // For now however, since we're still assuming permutations, we will
   // initialize this parameter alongside the `dim2lvl` and `lvl2dim`

 /// Sparse conversion rule for the new operator.
 class SparseTensorNewConverter : public OpConversionPattern<NewOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(NewOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Location loc = op.getLoc();
-    Type resType = op.getType();
-    auto enc = getSparseTensorEncoding(resType);
+    auto stp = op.getType().cast<ShapedType>();
+    auto enc = getSparseTensorEncoding(stp);
     if (!enc)
       return failure();
-    // Generate the call to construct tensor from ptr. The sizes are
-    // inferred from the result type of the new operator.
-    SmallVector<Value> sizes;
-    ShapedType stp = resType.cast<ShapedType>();
-    sizesFromType(rewriter, sizes, loc, stp);
-    Value ptr = adaptor.getOperands()[0];
-    rewriter.replaceOp(op, NewCallParams(rewriter, loc)
-                               .genBuffers(enc, sizes, stp)
-                               .genNewCall(Action::kFromFile, ptr));
+    const unsigned dimRank = stp.getRank();
+    const unsigned lvlRank = enc.getDimLevelType().size();
+    // Construct the dimShape.
+    const auto dimShape = stp.getShape();
+    SmallVector<Value> dimShapeValues;
+    sizesFromType(rewriter, dimShapeValues, loc, stp);
+    Value dimShapeBuffer = genBuffer(rewriter, loc, dimShapeValues);
+    // Allocate `SparseTensorReader` and perform all initial setup that
+    // does not depend on lvlSizes (nor dim2lvl, lvl2dim, etc).
+    Type opaqueTp = getOpaquePointerType(rewriter);
+    Value valTp =
+        constantPrimaryTypeEncoding(rewriter, loc, stp.getElementType());
+    Value reader =
+        createFuncCall(rewriter, loc, "createCheckedSparseTensorReader",
+                       opaqueTp,
+                       {adaptor.getOperands()[0], dimShapeBuffer, valTp},
+                       EmitCInterface::On)
+            .getResult(0);
+    // Construct the lvlSizes.  If the dimShape is static, then it's
+    // identical to dimSizes: so we can compute lvlSizes entirely at
+    // compile-time.  If dimShape is dynamic, then we'll need to generate
+    // code for computing lvlSizes from the `reader`'s actual dimSizes.
+    //
+    // TODO: For now we're still assuming `dim2lvl` is a permutation.
+    // But since we're computing lvlSizes here (rather than in the runtime),
+    // we can easily generalize that simply by adjusting this code.
+    //
+    // FIXME: reduce redundancy vs `NewCallParams::genBuffers`.
+    Value dimSizesBuffer;
+    if (!stp.hasStaticShape()) {
+      Type indexTp = rewriter.getIndexType();
+      auto memTp = MemRefType::get({ShapedType::kDynamicSize}, indexTp);
+      dimSizesBuffer =
+          createFuncCall(rewriter, loc, "getSparseTensorReaderDimSizes", memTp,
+                         reader, EmitCInterface::On)
+              .getResult(0);
+    }
+    Value lvlSizesBuffer;
+    Value lvl2dimBuffer;
+    Value dim2lvlBuffer;
+    if (auto dimOrder = enc.getDimOrdering()) {
+      assert(dimOrder.isPermutation() && "Got non-permutation");
+      // We preinitialize `dim2lvlValues` since we need random-access writing.
+      // And we preinitialize the others for stylistic consistency.
+      SmallVector<Value> lvlSizeValues(lvlRank);
+      SmallVector<Value> lvl2dimValues(lvlRank);
+      SmallVector<Value> dim2lvlValues(dimRank);
+      for (unsigned l = 0; l < lvlRank; l++) {
+        // The `d`th source variable occurs in the `l`th result position.
+        uint64_t d = dimOrder.getDimPosition(l);
+        Value lvl = constantIndex(rewriter, loc, l);
+        Value dim = constantIndex(rewriter, loc, d);
+        dim2lvlValues[d] = lvl;
+        lvl2dimValues[l] = dim;
+        lvlSizeValues[l] =
+            (dimShape[d] == ShapedType::kDynamicSize)
+                ? rewriter.create<memref::LoadOp>(loc, dimSizesBuffer, dim)
+                : dimShapeValues[d];
+      }
+      lvlSizesBuffer = genBuffer(rewriter, loc, lvlSizeValues);
+      lvl2dimBuffer = genBuffer(rewriter, loc, lvl2dimValues);
+      dim2lvlBuffer = genBuffer(rewriter, loc, dim2lvlValues);
+    } else {
+      assert(dimRank == lvlRank && "Rank mismatch");
+      SmallVector<Value> iotaValues;
+      iotaValues.reserve(lvlRank);
+      for (unsigned i = 0; i < lvlRank; i++)
+        iotaValues.push_back(constantIndex(rewriter, loc, i));
+      lvlSizesBuffer = dimSizesBuffer ? dimSizesBuffer : dimShapeBuffer;
+      dim2lvlBuffer = lvl2dimBuffer = genBuffer(rewriter, loc, iotaValues);
+    }
+    // Use the `reader` to parse the file.
+    SmallVector<Value, 8> params{
+        reader,
+        lvlSizesBuffer,
+        genLvlTypesBuffer(rewriter, loc, enc),
+        lvl2dimBuffer,
+        dim2lvlBuffer,
+        constantPointerTypeEncoding(rewriter, loc, enc),
+        constantIndexTypeEncoding(rewriter, loc, enc),
+        valTp};
+    Value tensor = createFuncCall(rewriter, loc, "newSparseTensorFromReader",
+                                  opaqueTp, params, EmitCInterface::On)
+                       .getResult(0);
+    // Free the memory for `reader`.
+    createFuncCall(rewriter, loc, "delSparseTensorReader", {}, {reader},
+                   EmitCInterface::Off);
+    rewriter.replaceOp(op, tensor);
     return success();
   }
 };
 
 /// Sparse conversion rule for the alloc operator.
 class SparseTensorAllocConverter
     : public OpConversionPattern<bufferization::AllocTensorOp> {
 public:

mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorRewriting.cpp

 /// Sparse rewriting rule for the new operator.
 struct NewRewriter : public OpRewritePattern<NewOp> {
   using OpRewritePattern::OpRewritePattern;
   LogicalResult matchAndRewrite(NewOp op,
                                 PatternRewriter &rewriter) const override {
     Location loc = op.getLoc();
     auto dstTp = op.getResult().getType().template cast<RankedTensorType>();
     SparseTensorEncodingAttr encDst = getSparseTensorEncoding(dstTp);
-    if (!encDst) {
+    if (!encDst)
       return failure();
-    }
 
     // Create a sparse tensor reader.
     Value fileName = op.getSource();
     Type opaqueTp = getOpaquePointerType(rewriter);
     Value reader = createFuncCall(rewriter, loc, "createSparseTensorReader",
                                   {opaqueTp}, {fileName}, EmitCInterface::Off)
                        .getResult(0);
 
     // Allocate a temporary buffer for storing dimension sizes and indices.
     Type indexTp = rewriter.getIndexType();
     uint64_t rank = dstTp.getRank();
     Value dimSizes = genAlloca(rewriter, loc, rank, indexTp);
 
     // If the result tensor has dynamic dimensions, get the dynamic sizes from
     // the sparse tensor reader.
     SmallVector<Value> dynSizesArray;
     if (!dstTp.hasStaticShape()) {
-      createFuncCall(rewriter, loc, "getSparseTensorReaderDimSizes", {},
+      createFuncCall(rewriter, loc, "copySparseTensorReaderDimSizes", {},
                      {reader, dimSizes}, EmitCInterface::On)
           .getResult(0);
       ArrayRef<int64_t> dstShape = dstTp.getShape();
       for (auto &d : llvm::enumerate(dstShape)) {
         if (d.value() == ShapedType::kDynamicSize) {
           dynSizesArray.push_back(rewriter.create<memref::LoadOp>(
               loc, dimSizes, constantIndex(rewriter, loc, d.index())));
         }
if (rank == 2) {
       symmetric = Value();
     }
     Type eltTp = dstTp.getElementType();
     Value value = genAllocaScalar(rewriter, loc, eltTp);
     scf::ForOp forOp = rewriter.create<scf::ForOp>(loc, c0, nnz, c1,
                                                    ArrayRef<Value>(cooBuffer));
     rewriter.setInsertionPointToStart(forOp.getBody());
 
-    SmallString<18> getNextFuncName{"getSparseTensorReaderNext",
+    SmallString<29> getNextFuncName{"getSparseTensorReaderNext",
                                     primaryTypeFunctionSuffix(eltTp)};
     Value indices = dimSizes; // Reuse the indices memref to store indices.
     createFuncCall(rewriter, loc, getNextFuncName, {}, {reader, indices, value},
                    EmitCInterface::On);
     SmallVector<Value> indicesArray;
     for (uint64_t i = 0; i < rank; i++) {
       indicesArray.push_back(rewriter.create<memref::LoadOp>(
           loc, indices, constantIndex(rewriter, loc, i)));
Value writer =
                        {op.getDest()}, EmitCInterface::Off)
             .getResult(0);
     Value rankValue = constantIndex(rewriter, loc, rank);
     createFuncCall(rewriter, loc, "outSparseTensorWriterMetaData", {},
                    {writer, rankValue, nnz, dimSizes}, EmitCInterface::On);
 
     Value indices = dimSizes; // Reuse the dimSizes buffer for indices.
     Type eltTp = srcTp.getElementType();
-    SmallString<18> outNextFuncName{"outSparseTensorWriterNext",
+    SmallString<29> outNextFuncName{"outSparseTensorWriterNext",
                                     primaryTypeFunctionSuffix(eltTp)};
     Value value = genAllocaScalar(rewriter, loc, eltTp);
     ModuleOp module = op->getParentOfType<ModuleOp>();
     // For each element in the source tensor, output the element.
     rewriter.create<ForeachOp>(
         loc, src, llvm::None,
         [&](OpBuilder &builder, Location loc, ValueRange args, Value v,
             ValueRange reduc) {

mlir/lib/ExecutionEngine/SparseTensor/File.cpp

 }
 
 /// Attempts to read a line from the file.
 void SparseTensorReader::readLine() {
   if (!fgets(line, kColWidth, file))
     MLIR_SPARSETENSOR_FATAL("Cannot read next line of %s\n", filename);
 }
 
-char *SparseTensorReader::readCOOIndices(uint64_t rank, uint64_t *indices) {
-  assert(rank == getRank() && "Rank mismatch");
+char *SparseTensorReader::readCOOIndices(uint64_t *indices) {
   readLine();
   // Local variable for tracking the parser's position in the `line` buffer.
   char *linePtr = line;
-  for (uint64_t r = 0; r < rank; ++r) {
+  for (uint64_t rank = getRank(), r = 0; r < rank; ++r) {
     // Parse the 1-based index.
     uint64_t idx = strtoul(linePtr, &linePtr, 10);
     // Store the 0-based index.
     indices[r] = idx - 1;
   }
   return linePtr;
 }
 

mlir/lib/ExecutionEngine/SparseTensorRuntime.cpp

#define MEMREF_GET_USIZE(MEMREF)                                               \
   detail::checkOverflowCast<uint64_t>((MEMREF)->sizes[0])
 
 #define ASSERT_USIZE_EQ(MEMREF, SZ)                                            \
   assert(detail::safelyEQ(MEMREF_GET_USIZE(MEMREF), (SZ)) &&                   \
          "Memref size mismatch")
 
 #define MEMREF_GET_PAYLOAD(MEMREF) ((MEMREF)->data + (MEMREF)->offset)
 
-// We make this a function rather than a macro mainly for type safety
-// reasons.  This function does not modify the vector, but it cannot
-// be marked `const` because it is stored into the non-`const` memref.
-template <typename T>
-static void vectorToMemref(std::vector<T> &v, StridedMemRefType<T, 1> &ref) {
-  ref.basePtr = ref.data = v.data();
+/// Initializes the memref with the provided size and data pointer.  This
+/// is designed for functions which want to "return" a memref that aliases
+/// into memory owned by some other object (e.g., `SparseTensorStorage`),
+/// without doing any actual copying.  (The "return" is in scarequotes
+/// because the `_mlir_ciface_` calling convention migrates any returned
+/// memrefs into an out-parameter passed before all the other function
+/// parameters.)
+///
+/// We make this a function rather than a macro mainly for type safety
+/// reasons.  This function does not modify the data pointer, but it
+/// cannot be marked `const` because it is stored into the (necessarily)
+/// non-`const` memref.  This function is templated over the `DataSizeT`
+/// to work around signedness warnings due to many data types having
+/// varying signedness across different platforms.  The templating allows
+/// this function to ensure that it does the right thing and never
+/// introduces errors due to implicit conversions.
+template <typename DataSizeT, typename T>
+static inline void aliasIntoMemref(DataSizeT size, T *data,
+                                   StridedMemRefType<T, 1> &ref) {
+  ref.basePtr = ref.data = data;
   ref.offset = 0;
-  using SizeT = typename std::remove_reference_t<decltype(ref.sizes[0])>;
-  ref.sizes[0] = detail::checkOverflowCast<SizeT>(v.size());
+  using MemrefSizeT = typename std::remove_reference_t<decltype(ref.sizes[0])>;
+  ref.sizes[0] = detail::checkOverflowCast<MemrefSizeT>(size);
   ref.strides[0] = 1;
 }
 
 } // anonymous namespace
 
 extern "C" {
 
 //===----------------------------------------------------------------------===//
 //
 // Public functions which operate on MLIR buffers (memrefs) to interact
 // with sparse tensors (which are only visible as opaque pointers externally).
 //
 //===----------------------------------------------------------------------===//
 
 #define CASE(p, i, v, P, I, V)                                                 \
   if (ptrTp == (p) && indTp == (i) && valTp == (v)) {                          \
     switch (action) {                                                          \
     case Action::kEmpty:                                                       \
       return SparseTensorStorage<P, I, V>::newEmpty(                           \
           dimRank, dimSizes, lvlRank, lvlSizes, lvlTypes, lvl2dim);            \
-    case Action::kFromFile: {                                                  \
-      char *filename = static_cast<char *>(ptr);                               \
-      return openSparseTensor<P, I, V>(dimRank, dimSizes, lvlRank, lvlTypes,   \
-                                       lvl2dim, dim2lvl, filename, v);         \
-    }                                                                          \
     case Action::kFromCOO: {                                                   \
       assert(ptr && "Received nullptr for SparseTensorCOO object");            \
       auto &coo = *static_cast<SparseTensorCOO<V> *>(ptr);                     \
       return SparseTensorStorage<P, I, V>::newFromCOO(                         \
           dimRank, dimSizes, lvlRank, lvlTypes, lvl2dim, coo);                 \
     }                                                                          \
     case Action::kSparseToSparse: {                                            \
       assert(ptr && "Received nullptr for SparseTensorStorage object");        \

 
 #define IMPL_SPARSEVALUES(VNAME, V)                                            \
   void _mlir_ciface_sparseValues##VNAME(StridedMemRefType<V, 1> *ref,          \
                                         void *tensor) {                        \
     assert(ref &&tensor);                                                      \
     std::vector<V> *v;                                                         \
     static_cast<SparseTensorStorageBase *>(tensor)->getValues(&v);             \
     assert(v);                                                                 \
-    vectorToMemref(*v, *ref);                                                  \
+    aliasIntoMemref(v->size(), v->data(), *ref);                               \
   }
 MLIR_SPARSETENSOR_FOREVERY_V(IMPL_SPARSEVALUES)
 #undef IMPL_SPARSEVALUES
 
 #define IMPL_GETOVERHEAD(NAME, TYPE, LIB)                                      \
   void _mlir_ciface_##NAME(StridedMemRefType<TYPE, 1> *ref, void *tensor,      \
                            index_type d) {                                     \
     assert(ref &&tensor);                                                      \
     std::vector<TYPE> *v;                                                      \
     static_cast<SparseTensorStorageBase *>(tensor)->LIB(&v, d);                \
     assert(v);                                                                 \
-    vectorToMemref(*v, *ref);                                                  \
+    aliasIntoMemref(v->size(), v->data(), *ref);                               \
   }
 #define IMPL_SPARSEPOINTERS(PNAME, P)                                          \
   IMPL_GETOVERHEAD(sparsePointers##PNAME, P, getPointers)
 MLIR_SPARSETENSOR_FOREVERY_O(IMPL_SPARSEPOINTERS)
 #undef IMPL_SPARSEPOINTERS
 
 #define IMPL_SPARSEINDICES(INAME, I)                                           \
   IMPL_GETOVERHEAD(sparseIndices##INAME, I, getIndices)
void _mlir_ciface_expInsert##VNAME(                                          \
     bool *filled = MEMREF_GET_PAYLOAD(fref);                                   \
     index_type *added = MEMREF_GET_PAYLOAD(aref);                              \
     static_cast<SparseTensorStorageBase *>(tensor)->expInsert(                 \
         cursor, values, filled, added, count);                                 \
   }
 MLIR_SPARSETENSOR_FOREVERY_V(IMPL_EXPINSERT)
 #undef IMPL_EXPINSERT
 
-void _mlir_ciface_getSparseTensorReaderDimSizes(
+void *_mlir_ciface_createCheckedSparseTensorReader(
+    char *filename, StridedMemRefType<index_type, 1> *dimShapeRef,
+    PrimaryType valTp) {
+  ASSERT_NO_STRIDE(dimShapeRef);
+  const uint64_t dimRank = MEMREF_GET_USIZE(dimShapeRef);
+  const index_type *dimShape = MEMREF_GET_PAYLOAD(dimShapeRef);
+  auto *reader = SparseTensorReader::create(filename, dimRank, dimShape, valTp);
+  return static_cast<void *>(reader);
+}
+
+// FIXME: update `SparseTensorCodegenPass` to use
+// `_mlir_ciface_getSparseTensorReaderDimSizes` instead.
+void _mlir_ciface_copySparseTensorReaderDimSizes(
     void *p, StridedMemRefType<index_type, 1> *dref) {
   assert(p);
+  SparseTensorReader &reader = *static_cast<SparseTensorReader *>(p);
   ASSERT_NO_STRIDE(dref);
+  const uint64_t dimRank = MEMREF_GET_USIZE(dref);
+  ASSERT_USIZE_EQ(dref, reader.getRank());
   index_type *dimSizes = MEMREF_GET_PAYLOAD(dref);
-  SparseTensorReader &file = *static_cast<SparseTensorReader *>(p);
-  const index_type *sizes = file.getDimSizes();
-  index_type rank = file.getRank();
-  for (index_type r = 0; r < rank; ++r)
-    dimSizes[r] = sizes[r];
+  const index_type *fileSizes = reader.getDimSizes();
+  for (uint64_t d = 0; d < dimRank; ++d)
+    dimSizes[d] = fileSizes[d];
+}
+
+void _mlir_ciface_getSparseTensorReaderDimSizes(
+    StridedMemRefType<index_type, 1> *out, void *p) {
+  assert(out && p);
+  SparseTensorReader &reader = *static_cast<SparseTensorReader *>(p);
+  auto *dimSizes = const_cast<uint64_t *>(reader.getDimSizes());
+  aliasIntoMemref(reader.getRank(), dimSizes, *out);
 }
 
 #define IMPL_GETNEXT(VNAME, V)                                                 \
   void _mlir_ciface_getSparseTensorReaderNext##VNAME(                          \
       void *p, StridedMemRefType<index_type, 1> *iref,                         \
       StridedMemRefType<V, 0> *vref) {                                         \
     assert(p &&vref);                                                          \
+    auto &reader = *static_cast<SparseTensorReader *>(p);                      \
     ASSERT_NO_STRIDE(iref);                                                    \
+    const uint64_t rank = MEMREF_GET_USIZE(iref);                              \
     index_type *indices = MEMREF_GET_PAYLOAD(iref);                            \
-    SparseTensorReader *stfile = static_cast<SparseTensorReader *>(p);         \
-    index_type rank = stfile->getRank();                                       \
     V *value = MEMREF_GET_PAYLOAD(vref);                                       \
-    *value = stfile->readCOOElement<V>(rank, indices);                         \
+    *value = reader.readCOOElement<V>(rank, indices);                          \
   }
 MLIR_SPARSETENSOR_FOREVERY_V(IMPL_GETNEXT)
 #undef IMPL_GETNEXT
 
+void *_mlir_ciface_newSparseTensorFromReader(
+    void *p, StridedMemRefType<index_type, 1> *lvlSizesRef,
+    StridedMemRefType<DimLevelType, 1> *lvlTypesRef,
+    StridedMemRefType<index_type, 1> *lvl2dimRef,
+    StridedMemRefType<index_type, 1> *dim2lvlRef, OverheadType ptrTp,
+    OverheadType indTp, PrimaryType valTp) {
+  assert(p);
+  SparseTensorReader &reader = *static_cast<SparseTensorReader *>(p);
+  ASSERT_NO_STRIDE(lvlSizesRef);
+  ASSERT_NO_STRIDE(lvlTypesRef);
+  ASSERT_NO_STRIDE(lvl2dimRef);
+  ASSERT_NO_STRIDE(dim2lvlRef);
+  const uint64_t dimRank = reader.getRank();
+  const uint64_t lvlRank = MEMREF_GET_USIZE(lvlSizesRef);
+  ASSERT_USIZE_EQ(lvlTypesRef, lvlRank);
+  ASSERT_USIZE_EQ(lvl2dimRef, lvlRank);
+  ASSERT_USIZE_EQ(dim2lvlRef, dimRank);
+  const index_type *lvlSizes = MEMREF_GET_PAYLOAD(lvlSizesRef);
+  const DimLevelType *lvlTypes = MEMREF_GET_PAYLOAD(lvlTypesRef);
+  const index_type *lvl2dim = MEMREF_GET_PAYLOAD(lvl2dimRef);
+  const index_type *dim2lvl = MEMREF_GET_PAYLOAD(dim2lvlRef);
+  //
+  // FIXME(wrengr): Really need to define a separate x-macro for handling
+  // all this. (Or ideally some better, entirely-different approach)
+#define CASE(p, i, v, P, I, V)                                                 \
+  if (ptrTp == OverheadType::p && indTp == OverheadType::i &&                  \
+      valTp == PrimaryType::v)                                                 \
+    return static_cast<void *>(reader.readSparseTensor<P, I, V>(               \
+        lvlRank, lvlSizes, lvlTypes, lvl2dim, dim2lvl));
+#define CASE_SECSAME(p, v, P, V) CASE(p, p, v, P, P, V)
+  // Rewrite kIndex to kU64, to avoid introducing a bunch of new cases.
+  // This is safe because of the static_assert above.
+  if (ptrTp == OverheadType::kIndex)
+    ptrTp = OverheadType::kU64;
+  if (indTp == OverheadType::kIndex)
+    indTp = OverheadType::kU64;
+  // Double matrices with all combinations of overhead storage.
+  CASE(kU64, kU64, kF64, uint64_t, uint64_t, double);
+  CASE(kU64, kU32, kF64, uint64_t, uint32_t, double);
+  CASE(kU64, kU16, kF64, uint64_t, uint16_t, double);
+  CASE(kU64, kU8, kF64, uint64_t, uint8_t, double);
+  CASE(kU32, kU64, kF64, uint32_t, uint64_t, double);
+  CASE(kU32, kU32, kF64, uint32_t, uint32_t, double);
+  CASE(kU32, kU16, kF64, uint32_t, uint16_t, double);
+  CASE(kU32, kU8, kF64, uint32_t, uint8_t, double);
+  CASE(kU16, kU64, kF64, uint16_t, uint64_t, double);
+  CASE(kU16, kU32, kF64, uint16_t, uint32_t, double);
+  CASE(kU16, kU16, kF64, uint16_t, uint16_t, double);
+  CASE(kU16, kU8, kF64, uint16_t, uint8_t, double);
+  CASE(kU8, kU64, kF64, uint8_t, uint64_t, double);
+  CASE(kU8, kU32, kF64, uint8_t, uint32_t, double);
+  CASE(kU8, kU16, kF64, uint8_t, uint16_t, double);
+  CASE(kU8, kU8, kF64, uint8_t, uint8_t, double);
+  // Float matrices with all combinations of overhead storage.
+  CASE(kU64, kU64, kF32, uint64_t, uint64_t, float);
+  CASE(kU64, kU32, kF32, uint64_t, uint32_t, float);
+  CASE(kU64, kU16, kF32, uint64_t, uint16_t, float);
+  CASE(kU64, kU8, kF32, uint64_t, uint8_t, float);
+  CASE(kU32, kU64, kF32, uint32_t, uint64_t, float);
+  CASE(kU32, kU32, kF32, uint32_t, uint32_t, float);
+  CASE(kU32, kU16, kF32, uint32_t, uint16_t, float);
+  CASE(kU32, kU8, kF32, uint32_t, uint8_t, float);
+  CASE(kU16, kU64, kF32, uint16_t, uint64_t, float);
+  CASE(kU16, kU32, kF32, uint16_t, uint32_t, float);
+  CASE(kU16, kU16, kF32, uint16_t, uint16_t, float);
+  CASE(kU16, kU8, kF32, uint16_t, uint8_t, float);
+  CASE(kU8, kU64, kF32, uint8_t, uint64_t, float);
+  CASE(kU8, kU32, kF32, uint8_t, uint32_t, float);
+  CASE(kU8, kU16, kF32, uint8_t, uint16_t, float);
+  CASE(kU8, kU8, kF32, uint8_t, uint8_t, float);
+  // Two-byte floats with both overheads of the same type.
+  CASE_SECSAME(kU64, kF16, uint64_t, f16);
+  CASE_SECSAME(kU64, kBF16, uint64_t, bf16);
+  CASE_SECSAME(kU32, kF16, uint32_t, f16);
+  CASE_SECSAME(kU32, kBF16, uint32_t, bf16);
+  CASE_SECSAME(kU16, kF16, uint16_t, f16);
+  CASE_SECSAME(kU16, kBF16, uint16_t, bf16);
+  CASE_SECSAME(kU8, kF16, uint8_t, f16);
+  CASE_SECSAME(kU8, kBF16, uint8_t, bf16);
+  // Integral matrices with both overheads of the same type.
+  CASE_SECSAME(kU64, kI64, uint64_t, int64_t);
+  CASE_SECSAME(kU64, kI32, uint64_t, int32_t);
+  CASE_SECSAME(kU64, kI16, uint64_t, int16_t);
+  CASE_SECSAME(kU64, kI8, uint64_t, int8_t);
+  CASE_SECSAME(kU32, kI64, uint32_t, int64_t);
+  CASE_SECSAME(kU32, kI32, uint32_t, int32_t);
+  CASE_SECSAME(kU32, kI16, uint32_t, int16_t);
+  CASE_SECSAME(kU32, kI8, uint32_t, int8_t);
+  CASE_SECSAME(kU16, kI64, uint16_t, int64_t);
+  CASE_SECSAME(kU16, kI32, uint16_t, int32_t);
+  CASE_SECSAME(kU16, kI16, uint16_t, int16_t);
+  CASE_SECSAME(kU16, kI8, uint16_t, int8_t);
+  CASE_SECSAME(kU8, kI64, uint8_t, int64_t);
+  CASE_SECSAME(kU8, kI32, uint8_t, int32_t);
+  CASE_SECSAME(kU8, kI16, uint8_t, int16_t);
+  CASE_SECSAME(kU8, kI8, uint8_t, int8_t);
+  // Complex matrices with wide overhead.
+  CASE_SECSAME(kU64, kC64, uint64_t, complex64);
+  CASE_SECSAME(kU64, kC32, uint64_t, complex32);
+
+  // Unsupported case (add above if needed).
+  // TODO: better pretty-printing of enum values!
+  MLIR_SPARSETENSOR_FATAL(
+      "unsupported combination of types: <P=%d, I=%d, V=%d>\n",
+      static_cast<int>(ptrTp), static_cast<int>(indTp),
+      static_cast<int>(valTp));
+#undef CASE_SECSAME
+#undef CASE
+}
+
 void _mlir_ciface_outSparseTensorWriterMetaData(
     void *p, index_type rank, index_type nnz,
     StridedMemRefType<index_type, 1> *dref) {
   assert(p);
   ASSERT_NO_STRIDE(dref);
   assert(rank != 0);
   index_type *dimSizes = MEMREF_GET_PAYLOAD(dref);
   SparseTensorWriter &file = *static_cast<SparseTensorWriter *>(p);
char *getTensorFilename(index_type id) {
   char *env = getenv(var);
   if (!env)
     MLIR_SPARSETENSOR_FATAL("Environment variable %s is not set\n", var);
   return env;
 }
 
 void readSparseTensorShape(char *filename, std::vector<uint64_t> *out) {
   assert(out && "Received nullptr for out-parameter");
-  SparseTensorReader stfile(filename);
-  stfile.openFile();
-  stfile.readHeader();
-  stfile.closeFile();
-  const uint64_t rank = stfile.getRank();
-  const uint64_t *dimSizes = stfile.getDimSizes();
-  out->reserve(rank);
-  out->assign(dimSizes, dimSizes + rank);
+  SparseTensorReader reader(filename);
+  reader.openFile();
+  reader.readHeader();
+  reader.closeFile();
+  const uint64_t dimRank = reader.getRank();
+  const uint64_t *dimSizes = reader.getDimSizes();
+  out->reserve(dimRank);
+  out->assign(dimSizes, dimSizes + dimRank);
 }
 
 // We can't use `static_cast` here because `DimLevelType` is an enum-class.
 #define IMPL_CONVERTTOMLIRSPARSETENSOR(VNAME, V)                               \
   void *convertToMLIRSparseTensor##VNAME(                                      \
       uint64_t rank, uint64_t nse, uint64_t *shape, V *values,                 \
       uint64_t *indices, uint64_t *perm, uint8_t *sparse) {                    \
     return toMLIRSparseTensor<V>(rank, nse, shape, values, indices, perm,      \
                                  reinterpret_cast<DimLevelType *>(sparse));    \
   }
 MLIR_SPARSETENSOR_FOREVERY_V(IMPL_CONVERTTOMLIRSPARSETENSOR)
 #undef IMPL_CONVERTTOMLIRSPARSETENSOR
 
 #define IMPL_CONVERTFROMMLIRSPARSETENSOR(VNAME, V)                             \
   void convertFromMLIRSparseTensor##VNAME(void *tensor, uint64_t *pRank,       \
                                           uint64_t *pNse, uint64_t **pShape,   \
                                           V **pValues, uint64_t **pIndices) {  \
     fromMLIRSparseTensor<V>(                                                   \
         static_cast<SparseTensorStorage<uint64_t, uint64_t, V> *>(tensor),     \
         pRank, pNse, pShape, pValues, pIndices);                               \
   }
 MLIR_SPARSETENSOR_FOREVERY_V(IMPL_CONVERTFROMMLIRSPARSETENSOR)
 #undef IMPL_CONVERTFROMMLIRSPARSETENSOR
 
+// FIXME: update `SparseTensorCodegenPass` to use
+// `_mlir_ciface_createCheckedSparseTensorReader` instead.
 void *createSparseTensorReader(char *filename) {
-  SparseTensorReader *stfile = new SparseTensorReader(filename);
-  stfile->openFile();
-  stfile->readHeader();
-  return static_cast<void *>(stfile);
+  SparseTensorReader *reader = new SparseTensorReader(filename);
+  reader->openFile();
+  reader->readHeader();
+  return static_cast<void *>(reader);
 }
 
 index_type getSparseTensorReaderRank(void *p) {
   return static_cast<SparseTensorReader *>(p)->getRank();
 }
 
 bool getSparseTensorReaderIsSymmetric(void *p) {
   return static_cast<SparseTensorReader *>(p)->isSymmetric();

mlir/test/Dialect/SparseTensor/conversion.mlir

func.func @sparse_dim3d_const(%arg0: tensor<10x20x30xf64, #SparseTensor>) -> index {
   // to dimension 2 in the stored sparse tensor scheme).
   %c = arith.constant 1 : index
   %0 = tensor.dim %arg0, %c : tensor<10x20x30xf64, #SparseTensor>
   return %0 : index
 }
 
 // CHECK-LABEL: func @sparse_new1d(
 //  CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) -> !llvm.ptr<i8>
-//   CHECK-DAG: %[[FromFile:.*]] = arith.constant 1 : i32
-//   CHECK-DAG: %[[DimSizes0:.*]] = memref.alloca() : memref<1xindex>
-//   CHECK-DAG: %[[LvlSizes0:.*]] = memref.alloca() : memref<1xindex>
-//   CHECK-DAG: %[[LvlTypes0:.*]] = memref.alloca() : memref<1xi8>
+//   CHECK-DAG: %[[DimShape0:.*]] = memref.alloca() : memref<1xindex>
+//   CHECK-DAG: %[[DimShape:.*]] = memref.cast %[[DimShape0]] : memref<1xindex> to memref<?xindex>
+//       CHECK: %[[Reader:.*]] = call @createCheckedSparseTensorReader(%[[A]], %[[DimShape]], %{{.*}})
 //   CHECK-DAG: %[[Iota0:.*]] = memref.alloca() : memref<1xindex>
-//   CHECK-DAG: %[[DimSizes:.*]] = memref.cast %[[DimSizes0]] : memref<1xindex> to memref<?xindex>
-//   CHECK-DAG: %[[LvlSizes:.*]] = memref.cast %[[LvlSizes0]] : memref<1xindex> to memref<?xindex>
-//   CHECK-DAG: %[[LvlTypes:.*]] = memref.cast %[[LvlTypes0]] : memref<1xi8> to memref<?xi8>
 //   CHECK-DAG: %[[Iota:.*]] = memref.cast %[[Iota0]] : memref<1xindex> to memref<?xindex>
-//       CHECK: %[[T:.*]] = call @newSparseTensor(%[[DimSizes]], %[[LvlSizes]], %[[LvlTypes]], %[[Iota]], %[[Iota]], %{{.*}}, %{{.*}}, %{{.*}}, %[[FromFile]], %[[A]])
+//   CHECK-DAG: %[[LvlTypes0:.*]] = memref.alloca() : memref<1xi8>
+//   CHECK-DAG: %[[LvlTypes:.*]] = memref.cast %[[LvlTypes0]] : memref<1xi8> to memref<?xi8>
+//       CHECK: %[[T:.*]] = call @newSparseTensorFromReader(%[[Reader]], %[[DimShape]], %[[LvlTypes]], %[[Iota]], %[[Iota]], %{{.*}}, %{{.*}}, %{{.*}})
+//       CHECK: call @delSparseTensorReader(%[[Reader]])
 //       CHECK: return %[[T]] : !llvm.ptr<i8>
 func.func @sparse_new1d(%arg0: !llvm.ptr<i8>) -> tensor<128xf64, #SparseVector> {
   %0 = sparse_tensor.new %arg0 : !llvm.ptr<i8> to tensor<128xf64, #SparseVector>
   return %0 : tensor<128xf64, #SparseVector>
 }
 
 // CHECK-LABEL: func @sparse_new2d(
 //  CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) -> !llvm.ptr<i8>
-//   CHECK-DAG: %[[FromFile:.*]] = arith.constant 1 : i32
-//   CHECK-DAG: %[[DimSizes0:.*]] = memref.alloca() : memref<2xindex>
-//   CHECK-DAG: %[[LvlSizes0:.*]] = memref.alloca() : memref<2xindex>
-//   CHECK-DAG: %[[LvlTypes0:.*]] = memref.alloca() : memref<2xi8>
+//   CHECK-DAG: %[[DimShape0:.*]] = memref.alloca() : memref<2xindex>
+//   CHECK-DAG: %[[DimShape:.*]] = memref.cast %[[DimShape0]] : memref<2xindex> to memref<?xindex>
+//       CHECK: %[[Reader:.*]] = call @createCheckedSparseTensorReader(%[[A]], %[[DimShape]], %{{.*}})
+//       CHECK: %[[DimSizes:.*]] = call @getSparseTensorReaderDimSizes(%[[Reader]])
 //   CHECK-DAG: %[[Iota0:.*]] = memref.alloca() : memref<2xindex>
-//   CHECK-DAG: %[[DimSizes:.*]] = memref.cast %[[DimSizes0]] : memref<2xindex> to memref<?xindex>
-//   CHECK-DAG: %[[LvlSizes:.*]] = memref.cast %[[LvlSizes0]] : memref<2xindex> to memref<?xindex>
-//   CHECK-DAG: %[[LvlTypes:.*]] = memref.cast %[[LvlTypes0]] : memref<2xi8> to memref<?xi8>
 //   CHECK-DAG: %[[Iota:.*]] = memref.cast %[[Iota0]] : memref<2xindex> to memref<?xindex>
-//       CHECK: %[[T:.*]] = call @newSparseTensor(%[[DimSizes]], %[[LvlSizes]], %[[LvlTypes]], %[[Iota]], %[[Iota]], %{{.*}}, %{{.*}}, %{{.*}}, %[[FromFile]], %[[A]])
+//   CHECK-DAG: %[[LvlTypes0:.*]] = memref.alloca() : memref<2xi8>
+//   CHECK-DAG: %[[LvlTypes:.*]] = memref.cast %[[LvlTypes0]] : memref<2xi8> to memref<?xi8>
+//       CHECK: %[[T:.*]] = call @newSparseTensorFromReader(%[[Reader]], %[[DimSizes]], %[[LvlTypes]], %[[Iota]], %[[Iota]], %{{.*}}, %{{.*}}, %{{.*}})
+//       CHECK: call @delSparseTensorReader(%[[Reader]])
 //       CHECK: return %[[T]] : !llvm.ptr<i8>
 func.func @sparse_new2d(%arg0: !llvm.ptr<i8>) -> tensor<?x?xf32, #CSR> {
   %0 = sparse_tensor.new %arg0 : !llvm.ptr<i8> to tensor<?x?xf32, #CSR>
   return %0 : tensor<?x?xf32, #CSR>
 }
 
 // CHECK-LABEL: func @sparse_new3d(
 //  CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) -> !llvm.ptr<i8>
-//   CHECK-DAG: %[[FromFile:.*]] = arith.constant 1 : i32
-//   CHECK-DAG: %[[DimSizes0:.*]] = memref.alloca() : memref<3xindex>
+//   CHECK-DAG: %[[DimShape0:.*]] = memref.alloca() : memref<3xindex>
+//   CHECK-DAG: %[[DimShape:.*]] = memref.cast %[[DimShape0]] : memref<3xindex> to memref<?xindex>
+//       CHECK: %[[Reader:.*]] = call @createCheckedSparseTensorReader(%[[A]], %[[DimShape]], %{{.*}})
+//       CHECK: %[[DimSizes:.*]] = call @getSparseTensorReaderDimSizes(%[[Reader]])
 //   CHECK-DAG: %[[LvlSizes0:.*]] = memref.alloca() : memref<3xindex>
-//   CHECK-DAG: %[[LvlTypes0:.*]] = memref.alloca() : memref<3xi8>
-//   CHECK-DAG: %[[Lvl2Dim0:.*]] = memref.alloca() : memref<3xindex>
-//   CHECK-DAG: %[[Dim2Lvl0:.*]] = memref.alloca() : memref<3xindex>
-//   CHECK-DAG: %[[DimSizes:.*]] = memref.cast %[[DimSizes0]] : memref<3xindex> to memref<?xindex>
 //   CHECK-DAG: %[[LvlSizes:.*]] = memref.cast %[[LvlSizes0]] : memref<3xindex> to memref<?xindex>
-//   CHECK-DAG: %[[LvlTypes:.*]] = memref.cast %[[LvlTypes0]] : memref<3xi8> to memref<?xi8>
+//   CHECK-DAG: %[[Lvl2Dim0:.*]] = memref.alloca() : memref<3xindex>
 //   CHECK-DAG: %[[Lvl2Dim:.*]] = memref.cast %[[Lvl2Dim0]] : memref<3xindex> to memref<?xindex>
+//   CHECK-DAG: %[[Dim2Lvl0:.*]] = memref.alloca() : memref<3xindex>
 //   CHECK-DAG: %[[Dim2Lvl:.*]] = memref.cast %[[Dim2Lvl0]] : memref<3xindex> to memref<?xindex>
-//       CHECK: %[[T:.*]] = call @newSparseTensor(%[[DimSizes]], %[[LvlSizes]], %[[LvlTypes]], %[[Lvl2Dim]], %[[Dim2Lvl]], %{{.*}}, %{{.*}}, %{{.*}}, %[[FromFile]], %[[A]])
+//   CHECK-DAG: %[[LvlTypes0:.*]] = memref.alloca() : memref<3xi8>
+//   CHECK-DAG: %[[LvlTypes:.*]] = memref.cast %[[LvlTypes0]] : memref<3xi8> to memref<?xi8>
+//       CHECK: %[[T:.*]] = call @newSparseTensorFromReader(%[[Reader]], %[[LvlSizes]], %[[LvlTypes]], %[[Lvl2Dim]], %[[Dim2Lvl]], %{{.*}}, %{{.*}}, %{{.*}})
+//       CHECK: call @delSparseTensorReader(%[[Reader]])
 //       CHECK: return %[[T]] : !llvm.ptr<i8>
 func.func @sparse_new3d(%arg0: !llvm.ptr<i8>) -> tensor<?x?x?xf32, #SparseTensor> {
   %0 = sparse_tensor.new %arg0 : !llvm.ptr<i8> to tensor<?x?x?xf32, #SparseTensor>
   return %0 : tensor<?x?x?xf32, #SparseTensor>
 }
 
 // CHECK-LABEL: func @sparse_init(
 //  CHECK-SAME: %[[I:.*]]: index,

mlir/test/Dialect/SparseTensor/rewriting_for_codegen.mlir

 // RUN: mlir-opt %s -post-sparsification-rewrite="enable-runtime-library=false enable-convert=false" | \
 // RUN: FileCheck %s
 
 #CSR = #sparse_tensor.encoding<{
   dimLevelType = ["dense", "compressed"]
 }>
 
 // CHECK-LABEL:   func.func @sparse_new(
 // CHECK-SAME:    %[[A:.*]]: !llvm.ptr<i8>) -> tensor<?x?xf32, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>> {
 // CHECK-DAG:     %[[C2:.*]] = arith.constant 2 : index
 // CHECK-DAG:     %[[C0:.*]] = arith.constant 0 : index
 // CHECK-DAG:     %[[C1:.*]] = arith.constant 1 : index
 // CHECK:         %[[R:.*]] = call @createSparseTensorReader(%[[A]])
 // CHECK:         %[[DS:.*]] = memref.alloca(%[[C2]]) : memref<?xindex>
-// CHECK:         call @getSparseTensorReaderDimSizes(%[[R]], %[[DS]])
+// CHECK:         call @copySparseTensorReaderDimSizes(%[[R]], %[[DS]])
 // CHECK:         %[[D0:.*]] = memref.load %[[DS]]{{\[}}%[[C0]]]
 // CHECK:         %[[D1:.*]] = memref.load %[[DS]]{{\[}}%[[C1]]]
 // CHECK:         %[[T:.*]] = bufferization.alloc_tensor(%[[D0]], %[[D1]])
 // CHECK:         %[[N:.*]] = call @getSparseTensorReaderNNZ(%[[R]])
 // CHECK:         %[[S:.*]] = call @getSparseTensorReaderIsSymmetric(%[[R]])
 // CHECK:         %[[VB:.*]] = memref.alloca()
 // CHECK:         %[[T2:.*]] = scf.for %{{.*}} = %[[C0]] to %[[N]] step %[[C1]] iter_args(%[[A2:.*]] = %[[T]])
 // CHECK:           func.call @getSparseTensorReaderNextF32(%[[R]], %[[DS]], %[[VB]])

mlir/test/Integration/Dialect/SparseTensor/CPU/sparse_file_io.mlir

module {
 
   func.func private @getTensorFilename(index) -> (!Filename)
 
   func.func private @createSparseTensorReader(!Filename) -> (!TensorReader)
   func.func private @delSparseTensorReader(!TensorReader) -> ()
   func.func private @getSparseTensorReaderRank(!TensorReader) -> (index)
   func.func private @getSparseTensorReaderNNZ(!TensorReader) -> (index)
   func.func private @getSparseTensorReaderIsSymmetric(!TensorReader) -> (i1)
-  func.func private @getSparseTensorReaderDimSizes(!TensorReader,
+  func.func private @copySparseTensorReaderDimSizes(!TensorReader,
     memref<?xindex>) -> () attributes { llvm.emit_c_interface }
   func.func private @getSparseTensorReaderNextF32(!TensorReader,
     memref<?xindex>, memref<f32>) -> () attributes { llvm.emit_c_interface }
 
   func.func private @createSparseTensorWriter(!Filename) -> (!TensorWriter)
   func.func private @delSparseTensorWriter(!TensorWriter)
   func.func private @outSparseTensorWriterMetaData(!TensorWriter, index, index,
     memref<?xindex>) -> () attributes { llvm.emit_c_interface }
func.func @readTensorFileAndDump(%fileName: !Filename) {
     %rank = call @getSparseTensorReaderRank(%tensor) : (!TensorReader) -> index
     vector.print %rank : index
     %nnz = call @getSparseTensorReaderNNZ(%tensor) : (!TensorReader) -> index
     vector.print %nnz : index
     %symmetric = call @getSparseTensorReaderIsSymmetric(%tensor)
       : (!TensorReader) -> i1
     vector.print %symmetric : i1
     %dimSizes = memref.alloc(%rank) : memref<?xindex>
-    func.call @getSparseTensorReaderDimSizes(%tensor, %dimSizes)
+    func.call @copySparseTensorReaderDimSizes(%tensor, %dimSizes)
       : (!TensorReader, memref<?xindex>) -> ()
     call @dumpi(%dimSizes) : (memref<?xindex>) -> ()
     %x0s, %x1s, %vs = call @readTensorFile(%tensor)
       : (!TensorReader) -> (memref<?xindex>, memref<?xindex>, memref<?xf32>)
 
     call @dumpi(%x0s) : (memref<?xindex>) -> ()
     call @dumpi(%x1s) : (memref<?xindex>) -> ()
     call @dumpf(%vs) : (memref<?xf32>) -> ()
func.func @createTensorFileFrom(%fileName0: !Filename, %fileName1: !Filename) {
     %tensor0 = call @createSparseTensorReader(%fileName0)
       : (!Filename) -> (!TensorReader)
     %tensor1 = call @createSparseTensorWriter(%fileName1)
       : (!Filename) -> (!TensorWriter)
 
     %rank = call @getSparseTensorReaderRank(%tensor0) : (!TensorReader) -> index
     %nnz = call @getSparseTensorReaderNNZ(%tensor0) : (!TensorReader) -> index
     %dimSizes = memref.alloc(%rank) : memref<?xindex>
-    func.call @getSparseTensorReaderDimSizes(%tensor0,%dimSizes)
+    func.call @copySparseTensorReaderDimSizes(%tensor0, %dimSizes)
       : (!TensorReader, memref<?xindex>) -> ()
     call @outSparseTensorWriterMetaData(%tensor1, %rank, %nnz, %dimSizes)
       : (!TensorWriter, index, index, memref<?xindex>) -> ()
 
     //TODO: handle isSymmetric.
     // Assume rank == 2.
     %indices = memref.alloc(%rank) : memref<?xindex>
     %value = memref.alloca() : memref<f32>