[mlir][sparse] replace support lib conversion with actual MLIR codegen

[aartbik]

Summary

Rationale:
Passing in a pointer to the memref data in order to implement the
dense to sparse conversion was a bit too low-level. This revision
improves upon that approach with a cleaner solution of generating
a loop nest in MLIR code itself that prepares the COO object before
passing it to our "swiss army knife" setup. This is much more
intuitive *and* now also allows for dynamic shapes.

[aartbik]

rebase with main (for smaller diff)

[bixia]

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

117

So, perm is an output to support getEltCall. Can we document this?
Also, google coding style prefers return values over output parameters. We may return a pair of Values here instead.

199

You mean add if to avoid making the call when val is zero?

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

112–121

This allows DimOp of non-sparse tensor types. Am I right? Can we add a comment for this

mlir/test/Dialect/SparseTensor/conversion.mlir

162

Shall we add test cases for non-trivial permutations and unknown dimension sizes?

[aartbik]

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

117

Output parameters are okay for llvm, but documenting this is of course always a good idea! Added this to the method doc.

199

Yes the != 0 is now inside the method, but we can avoid the call with an IF

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

112–121

Yes, but that was always the case already, i.e. passing existing ones through. In this particular case, I need to allow introducing them. But I added the comment as requested to clarify.

mlir/test/Dialect/SparseTensor/conversion.mlir

162

Yes, more tests are always good.

[aartbik]

addressed comments

[aartbik]

rebase

mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp

if (auto tp1 = op.source().getType().dyn_cast<RankedTensorType>()) {
     if (auto tp2 = op.dest().getType().dyn_cast<RankedTensorType>()) {
       assert(tp1.getRank() == tp2.getRank());
       auto shape1 = tp1.getShape();
       auto shape2 = tp2.getShape();
       for (unsigned d = 0, rank = tp1.getRank(); d < rank; d++) {
         if (shape1[d] != shape2[d])
           return op.emitError()
                  << "unexpected conversion mismatch in dimension " << d;
-        if (shape1[d] == MemRefType::kDynamicSize)
-          return op.emitError("unexpected dynamic size");
       }
       return success();
     }
   }
   return op.emitError("unexpected type in convert");
 }
 
 static LogicalResult verify(ToPointersOp op) {

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

 //===- SparseTensorLowering.cpp - Sparse tensor primitives conversion -----===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // Convert sparse tensor primitives to calls into a runtime support library.
 // Note that this is a current implementation choice to keep the conversion
 // simple. In principle, these primitives could also be converted to actual
 // elaborate IR code that implements the primitives on the selected sparse
 // tensor storage schemes.
 //
 //===----------------------------------------------------------------------===//
 
-#include "mlir/Dialect/LLVMIR/LLVMTypes.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 using namespace mlir;
 using namespace mlir::sparse_tensor;
moduleBuilder
         .create<FuncOp>(op->getLoc(), name,
                         FunctionType::get(context, operands.getTypes(), result))
         .setPrivate();
   }
   return SymbolRefAttr::get(context, name);
 }
 
 /// Generates a call into the "swiss army knife" method of the sparse runtime
-/// support library for materializing sparse tensors into the computation.
-static void genNewCall(ConversionPatternRewriter &rewriter, Operation *op,
+/// support library for materializing sparse tensors into the computation. The
+/// method returns the call value and assigns the permutation to 'perm'.
+static Value genNewCall(ConversionPatternRewriter &rewriter, Operation *op,
                         SparseTensorEncodingAttr &enc, uint32_t action,
-                       Value ptr) {
+                        Value &perm, Value ptr = Value()) {
   Location loc = op->getLoc();
   ShapedType resType = op->getResult(0).getType().cast<ShapedType>();
   SmallVector<Value, 8> params;
   // Sparsity annotations in tensor constant form.
   SmallVector<APInt, 4> attrs;
   unsigned sz = enc.getDimLevelType().size();
   for (unsigned i = 0; i < sz; i++)
     attrs.push_back(
         APInt(8, getDimLevelTypeEncoding(enc.getDimLevelType()[i])));
   params.push_back(getTensor(rewriter, 8, loc, attrs));
   // Dimension sizes array of the enveloping *dense* tensor. Useful for either
   // verification of external data, or for construction of internal data.
   auto shape = resType.getShape();
   SmallVector<APInt, 4> sizes;
   for (unsigned i = 0; i < sz; i++) {
     uint64_t s = shape[i] == ShapedType::kDynamicSize ? 0 : shape[i];
     sizes.push_back(APInt(64, s));
   }
   params.push_back(getTensor(rewriter, 64, loc, sizes));
   // Dimension order permutation array. This is the "identity" permutation by
   // default, or otherwise the "reverse" permutation of a given ordering, so
   // that indices can be mapped quickly to the right position.
-  SmallVector<APInt, 4> perm(sz);
-  AffineMap p = enc.getDimOrdering();
-  if (p) {
-    assert(p.isPermutation() && p.getNumResults() == sz);
+  SmallVector<APInt, 4> rev(sz);
+  if (AffineMap p = enc.getDimOrdering()) {
     for (unsigned i = 0; i < sz; i++)
-      perm[p.getDimPosition(i)] = APInt(64, i);
+      rev[p.getDimPosition(i)] = APInt(64, i);
   } else {
     for (unsigned i = 0; i < sz; i++)
-      perm[i] = APInt(64, i);
+      rev[i] = APInt(64, i);
   }
-  params.push_back(getTensor(rewriter, 64, loc, perm));
+  perm = getTensor(rewriter, 64, loc, rev);
+  params.push_back(perm);
   // Secondary and primary types encoding.
   unsigned secPtr = getOverheadTypeEncoding(enc.getPointerBitWidth());
   unsigned secInd = getOverheadTypeEncoding(enc.getIndexBitWidth());
   unsigned primary = getPrimaryTypeEncoding(resType.getElementType());
   assert(primary);
   params.push_back(
       rewriter.create<ConstantOp>(loc, rewriter.getI64IntegerAttr(secPtr)));
   params.push_back(
       rewriter.create<ConstantOp>(loc, rewriter.getI64IntegerAttr(secInd)));
   params.push_back(
       rewriter.create<ConstantOp>(loc, rewriter.getI64IntegerAttr(primary)));
   // User action and pointer.
+  Type pTp = LLVM::LLVMPointerType::get(IntegerType::get(op->getContext(), 8));
+  if (!ptr)
+    ptr = rewriter.create<LLVM::NullOp>(loc, pTp);
   params.push_back(
       rewriter.create<ConstantOp>(loc, rewriter.getI32IntegerAttr(action)));
   params.push_back(ptr);
   // Generate the call to create new tensor.
-  Type ptrType =
-      LLVM::LLVMPointerType::get(IntegerType::get(op->getContext(), 8));
   StringRef name = "newSparseTensor";
-  rewriter.replaceOpWithNewOp<CallOp>(
-      op, ptrType, getFunc(op, name, ptrType, params), params);
+  auto call =
+      rewriter.create<CallOp>(loc, pTp, getFunc(op, name, pTp, params), params);
+  return call.getResult(0);
 }
 
-/// Generates a call that exposes the data pointer as a void pointer.
-// TODO: probing the data pointer directly is a bit raw; we should replace
-//       this with proper memref util calls once they become available.
-static bool genPtrCall(ConversionPatternRewriter &rewriter, Operation *op,
-                       Value val, Value &ptr) {
+/// Generates a call that adds one element to a coordinate scheme.
+static void genAddEltCall(ConversionPatternRewriter &rewriter, Operation *op,
+                          Value ptr, Value tensor, Value ind, Value perm,
+                          ValueRange ivs) {
   Location loc = op->getLoc();
-  ShapedType sType = op->getResult(0).getType().cast<ShapedType>();
-  Type eltType = sType.getElementType();
-  // Specialize name for the data type. Even though the final buffferized
-  // version only operates on pointers, different names are required to
-  // ensure type correctness for all intermediate states.
   StringRef name;
+  Type eltType = tensor.getType().cast<ShapedType>().getElementType();
   if (eltType.isF64())
-    name = "getPtrF64";
+    name = "addEltF64";
   else if (eltType.isF32())
-    name = "getPtrF32";
+    name = "addEltF32";
   else if (eltType.isInteger(64))
-    name = "getPtrI64";
+    name = "addEltI64";
   else if (eltType.isInteger(32))
-    name = "getPtrI32";
+    name = "addEltI32";
   else if (eltType.isInteger(16))
-    name = "getPtrI16";
+    name = "addEltI16";
   else if (eltType.isInteger(8))
-    name = "getPtrI8";
+    name = "addEltI8";
   else
-    return false;
-  auto memRefTp = MemRefType::get(sType.getShape(), eltType);
-  auto unrankedTp = UnrankedMemRefType::get(eltType, 0);
-  Value c = rewriter.create<memref::BufferCastOp>(loc, memRefTp, val);
-  Value d = rewriter.create<memref::CastOp>(loc, unrankedTp, c);
-  Type ptrType =
-      LLVM::LLVMPointerType::get(IntegerType::get(op->getContext(), 8));
-  auto call =
-      rewriter.create<CallOp>(loc, ptrType, getFunc(op, name, ptrType, d), d);
-  ptr = call.getResult(0);
-  return true;
+    llvm_unreachable("Unknown element type");
+  Value val = rewriter.create<tensor::ExtractOp>(loc, tensor, ivs);
+  // TODO: add if here?
+  unsigned i = 0;
+  for (auto iv : ivs) {
+    Value idx = rewriter.create<ConstantOp>(loc, rewriter.getIndexAttr(i++));
+    rewriter.create<memref::StoreOp>(loc, iv, ind, idx);
+  }
+  SmallVector<Value, 8> params;
+  params.push_back(ptr);
+  params.push_back(val);
+  params.push_back(ind);
+  params.push_back(perm);
+  Type pTp = LLVM::LLVMPointerType::get(IntegerType::get(op->getContext(), 8));
+  rewriter.create<CallOp>(loc, pTp, getFunc(op, name, pTp, params), params);
 }
 
 //===----------------------------------------------------------------------===//
 // Conversion rules.
 //===----------------------------------------------------------------------===//
 
 /// Sparse conversion rule for returns.
 class SparseReturnConverter : public OpConversionPattern<ReturnOp> {
class SparseTensorNewConverter : public OpConversionPattern<NewOp> {
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(NewOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     Type resType = op.getType();
     auto enc = getSparseTensorEncoding(resType);
     if (!enc)
       return failure();
-    genNewCall(rewriter, op, enc, 0, operands[0]);
+    Value perm;
+    rewriter.replaceOp(op, genNewCall(rewriter, op, enc, 0, perm, operands[0]));
     return success();
   }
 };
 
 /// Sparse conversion rule for the convert operator.
 class SparseTensorConvertConverter : public OpConversionPattern<ConvertOp> {
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(ConvertOp op, ArrayRef<Value> operands,
                   ConversionPatternRewriter &rewriter) const override {
     Type resType = op.getType();
     auto encDst = getSparseTensorEncoding(resType);
     auto encSrc = getSparseTensorEncoding(op.source().getType());
     // TODO: implement sparse => sparse
     //             and sparse => dense
     if (!encDst || encSrc)
       return failure();
-    // This is a dense => sparse conversion.
-    Value ptr;
-    if (!genPtrCall(rewriter, op, operands[0], ptr))
-      return failure();
-    genNewCall(rewriter, op, encDst, 1, ptr);
+    // This is a dense => sparse conversion, that is handled as follows:
+    //   t = newSparseCOO()
+    //   for i1 in dim1
+    //    ..
+    //     for ik in dimk
+    //       val = a[i1,..,ik]
+    //       if val != 0
+    //         t->add(val, [i1,..,ik], [p1,..,pk])
+    //   s = newSparseTensor(t)
+    // Note that the dense tensor traversal code is actually implemented
+    // using MLIR IR to avoid having to expose too much low-level
+    // memref traversal details to the runtime support library.
+    Location loc = op->getLoc();
+    ShapedType shape = resType.cast<ShapedType>();
+    auto memTp =
+        MemRefType::get({ShapedType::kDynamicSize}, rewriter.getIndexType());
+    Value perm;
+    Value ptr = genNewCall(rewriter, op, encDst, 2, perm);
+    Value tensor = operands[0];
+    Value arg = rewriter.create<ConstantOp>(
+        loc, rewriter.getIndexAttr(shape.getRank()));
+    Value ind = rewriter.create<memref::AllocaOp>(loc, memTp, ValueRange{arg});
+    SmallVector<Value> lo;
+    SmallVector<Value> hi;
+    SmallVector<Value> st;
+    Value zero = rewriter.create<ConstantOp>(loc, rewriter.getIndexAttr(0));
+    Value one = rewriter.create<ConstantOp>(loc, rewriter.getIndexAttr(1));
+    for (unsigned i = 0, rank = shape.getRank(); i < rank; i++) {
+      lo.push_back(zero);
+      hi.push_back(linalg::createOrFoldDimOp(rewriter, loc, tensor, i));
+      st.push_back(one);
+    }
+    scf::buildLoopNest(rewriter, op.getLoc(), lo, hi, st, {},
+                       [&](OpBuilder &builder, Location loc, ValueRange ivs,
+                           ValueRange args) -> scf::ValueVector {
+                         genAddEltCall(rewriter, op, ptr, tensor, ind, perm,
+                                       ivs);
+                         return {};
+                       });
+    rewriter.replaceOp(op, genNewCall(rewriter, op, encDst, 1, perm, ptr));
     return success();
   }
 };
 
 /// Sparse conversion rule for pointer accesses.
 class SparseTensorToPointersConverter
     : public OpConversionPattern<ToPointersOp> {
 public:

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

struct SparseTensorConversionPass
     : public SparseTensorConversionBase<SparseTensorConversionPass> {
   void runOnOperation() override {
     auto *ctx = &getContext();
     RewritePatternSet patterns(ctx);
     SparseTensorTypeConverter converter;
     ConversionTarget target(*ctx);
     target.addIllegalOp<NewOp, ConvertOp, ToPointersOp, ToIndicesOp, ToValuesOp,
                         ToTensorOp>();
+    // All dynamic rules below accept new function, call, return, and dimop
+    // operations as legal output of the rewriting provided that all sparse
+    // tensor types have been fully rewritten.
     target.addDynamicallyLegalOp<FuncOp>(
         [&](FuncOp op) { return converter.isSignatureLegal(op.getType()); });
     target.addDynamicallyLegalOp<CallOp>([&](CallOp op) {
       return converter.isSignatureLegal(op.getCalleeType());
     });
     target.addDynamicallyLegalOp<ReturnOp>(
         [&](ReturnOp op) { return converter.isLegal(op.getOperandTypes()); });
-    target.addLegalOp<ConstantOp, tensor::CastOp, memref::BufferCastOp,
-                      memref::CastOp>();
+    target.addDynamicallyLegalOp<tensor::DimOp>([&](tensor::DimOp op) {
+      return converter.isLegal(op.getOperandTypes());
+    });
+    // The following operations and dialects may be introduced by the
+    // rewriting rules, and are therefore marked as legal.
+    target.addLegalOp<ConstantOp, tensor::CastOp, tensor::ExtractOp>();
+    target.addLegalDialect<scf::SCFDialect, LLVM::LLVMDialect,
+                           memref::MemRefDialect>();
+    // Populate with rules and apply rewriting rules.
     populateFuncOpTypeConversionPattern(patterns, converter);
     populateCallOpTypeConversionPattern(patterns, converter);
     populateSparseTensorConversionPatterns(converter, patterns);
     if (failed(applyPartialConversion(getOperation(), target,
                                       std::move(patterns))))
       signalPassFailure();
   }
 };
 

mlir/lib/ExecutionEngine/SparseUtils.cpp

 //
 // Internal support for storing and reading sparse tensors.
 //
 // The following memory-resident sparse storage schemes are supported:
 //
 // (a) A coordinate scheme for temporarily storing and lexicographically
 //     sorting a sparse tensor by index.
 //
-// (b) A "one-size-fits-all" sparse storage scheme defined by per-rank
+// (b) A "one-size-fits-all" sparse tensor storage scheme defined by per-rank
 //     sparse/dense annnotations to be used by generated MLIR code.
 //
 // The following external formats are supported:
 //
 // (1) Matrix Market Exchange (MME): *.mtx
 //     https://math.nist.gov/MatrixMarket/formats.html
 //
 // (2) Formidable Repository of Open Sparse Tensors and Tools (FROSTT): *.tns

 /// A memory-resident sparse tensor in coordinate scheme (collection of
 /// elements). This data structure is used to read a sparse tensor from
 /// any external format into memory and sort the elements lexicographically
 /// by indices before passing it back to the client (most packed storage
 /// formats require the elements to appear in lexicographic index order).
 template <typename V>
 struct SparseTensor {
 public:
-  SparseTensor(const std::vector<uint64_t> &szs, uint64_t capacity = 0)
+  SparseTensor(const std::vector<uint64_t> &szs, uint64_t capacity)
       : sizes(szs), pos(0) {
     if (capacity)
       elements.reserve(capacity);
   }
   /// Adds element as indices and value.
   void add(const std::vector<uint64_t> &ind, V val) {
     assert(getRank() == ind.size());
     for (uint64_t r = 0, rank = getRank(); r < rank; r++)
       assert(ind[r] < sizes[r]); // within bounds
     elements.emplace_back(Element<V>(ind, val));
   }
   /// Sorts elements lexicographically by index.
   void sort() { std::sort(elements.begin(), elements.end(), lexOrder); }
   /// Primitive one-time iteration.
   const Element<V> &next() { return elements[pos++]; }
   /// Returns rank.
   uint64_t getRank() const { return sizes.size(); }
   /// Getter for sizes array.
   const std::vector<uint64_t> &getSizes() const { return sizes; }
   /// Getter for elements array.
   const std::vector<Element<V>> &getElements() const { return elements; }
 
+  /// Factory method.
+  static SparseTensor<V> *newSparseTensor(uint64_t size, uint64_t *sizes,
+                                          uint64_t *perm,
+                                          uint64_t capacity = 0) {
+    std::vector<uint64_t> indices(size);
+    for (uint64_t r = 0; r < size; r++)
+      indices[perm[r]] = sizes[r];
+    return new SparseTensor<V>(indices, capacity);
+  }
+
 private:
   /// Returns true if indices of e1 < indices of e2.
   static bool lexOrder(const Element<V> &e1, const Element<V> &e2) {
     assert(e1.indices.size() == e2.indices.size());
     for (uint64_t r = 0, rank = e1.indices.size(); r < rank; r++) {
       if (e1.indices[r] == e2.indices[r])
         continue;
       return e1.indices[r] < e2.indices[r];

 /// form of an imaginary SparseTensorType, until such a type becomes a
 /// first-class citizen of MLIR. In contrast to generating setup methods for
 /// each differently annotated sparse tensor, this method provides a convenient
 /// "one-size-fits-all" solution that simply takes an input tensor and
 /// annotations to implement all required setup in a general manner.
 template <typename P, typename I, typename V>
 class SparseTensorStorage : public SparseTensorStorageBase {
 public:
-  /// Constructs sparse tensor storage scheme following the given
-  /// per-rank dimension dense/sparse annotations.
+  /// Constructs a sparse tensor storage scheme from the given sparse
+  /// tensor in coordinate scheme following the given per-rank dimension
+  /// dense/sparse annotations.
   SparseTensorStorage(SparseTensor<V> *tensor, uint8_t *sparsity)
       : sizes(tensor->getSizes()), pointers(getRank()), indices(getRank()) {
     // Provide hints on capacity.
     // TODO: needs fine-tuning based on sparsity
     uint64_t nnz = tensor->getElements().size();
     values.reserve(nnz);
     for (uint64_t d = 0, s = 1, rank = getRank(); d < rank; d++) {
       s *= sizes[d];
public:
   void getPointers(std::vector<P> **out, uint64_t d) override {
     *out = &pointers[d];
   }
   void getIndices(std::vector<I> **out, uint64_t d) override {
     *out = &indices[d];
   }
   void getValues(std::vector<V> **out) override { *out = &values; }
 
-  // Factory method.
+  /// Factory method.
   static SparseTensorStorage<P, I, V> *newSparseTensor(SparseTensor<V> *t,
                                                        uint8_t *s) {
     t->sort(); // sort lexicographically
     SparseTensorStorage<P, I, V> *n = new SparseTensorStorage<P, I, V>(t, s);
     delete t;
     return n;
   }
 
 private:
-  /// Initializes sparse tensor storage scheme from a memory-resident
-  /// representation of an external sparse tensor. This method prepares
-  /// the pointers and indices arrays under the given per-rank dimension
-  /// dense/sparse annotations.
+  /// Initializes sparse tensor storage scheme from a memory-resident sparse
+  /// tensor in coordinate scheme. This method prepares the pointers and indices
+  /// arrays under the given per-rank dimension dense/sparse annotations.
   void traverse(SparseTensor<V> *tensor, uint8_t *sparsity, uint64_t lo,
                 uint64_t hi, uint64_t d) {
     const std::vector<Element<V>> &elements = tensor->getElements();
     // Once dimensions are exhausted, insert the numerical values.
     if (d == getRank()) {
       values.push_back(lo < hi ? elements[lo].value : 0);
       return;
     }
static SparseTensor<V> *openTensor(char *filename, uint64_t size,
   } else {
     fprintf(stderr, "Unknown format %s\n", filename);
     exit(1);
   }
   // Prepare sparse tensor object with per-rank dimension sizes
   // and the number of nonzeros as initial capacity.
   assert(size == idata[0] && "rank mismatch");
   uint64_t nnz = idata[1];
-  std::vector<uint64_t> indices(size);
-  for (uint64_t r = 0; r < size; r++) {
-    uint64_t sz = idata[2 + r];
-    assert((sizes[r] == 0 || sizes[r] == sz) && "dimension size mismatch");
-    indices[perm[r]] = sz;
-  }
-  SparseTensor<V> *tensor = new SparseTensor<V>(indices, nnz);
-  // Read all nonzero elements.
+  for (uint64_t r = 0; r < size; r++)
+    assert((sizes[r] == 0 || sizes[r] == idata[2 + r]) &&
+           "dimension size mismatch");
+  SparseTensor<V> *tensor =
+      SparseTensor<V>::newSparseTensor(size, idata + 2, perm, nnz);
+  //  Read all nonzero elements.
+  std::vector<uint64_t> indices(size);
   for (uint64_t k = 0; k < nnz; k++) {
     uint64_t idx = -1;
     for (uint64_t r = 0; r < size; r++) {
       if (fscanf(file, "%" PRIu64, &idx) != 1) {
         fprintf(stderr, "Cannot find next index in %s\n", filename);
         exit(1);
       }
       // Add 0-based index.
       indices[perm[r]] = idx - 1;
     }
     // The external formats always store the numerical values with the type
     // double, but we cast these values to the sparse tensor object type.
     double value;
     if (fscanf(file, "%lg\n", &value) != 1) {
       fprintf(stderr, "Cannot find next value in %s\n", filename);
       exit(1);
     }
     tensor->add(indices, value);
   }
   // Close the file and return tensor.
   fclose(file);
   return tensor;
 }
 
-/// Helper to copy a linearized dense tensor.
-template <typename V>
-static V *copyTensorTraverse(SparseTensor<V> *tensor,
-                             std::vector<uint64_t> &indices, uint64_t r,
-                             uint64_t rank, uint64_t *sizes, uint64_t *perm,
-                             V *data) {
-  for (uint64_t i = 0, sz = sizes[r]; i < sz; i++) {
-    indices[perm[r]] = i;
-    if (r + 1 == rank) {
-      V d = *data++;
-      if (d)
-        tensor->add(indices, d);
-    } else {
-      data =
-          copyTensorTraverse(tensor, indices, r + 1, rank, sizes, perm, data);
-    }
-  }
-  return data;
-}
-
-/// Copies the nonzeros of a linearized dense tensor into a memory-resident
-/// sparse tensor in coordinate scheme.
-template <typename V>
-static SparseTensor<V> *copyTensor(uint64_t size, uint64_t *sizes,
-                                   uint64_t *perm, V *data) {
-  std::vector<uint64_t> indices(size);
-  for (uint64_t r = 0; r < size; r++)
-    indices[perm[r]] = sizes[r];
-  SparseTensor<V> *tensor = new SparseTensor<V>(indices);
-  copyTensorTraverse<V>(tensor, indices, 0, size, sizes, perm, data);
-  return tensor;
-}
-
 } // anonymous namespace
 
 extern "C" {
 
 /// Helper method to read a sparse tensor filename from the environment,
 /// defined with the naming convention ${TENSOR0}, ${TENSOR1}, etc.
 char *getTensorFilename(uint64_t id) {
   char var[80];

 // by actual codegen in MLIR.
 //
 // Because we cannot use C++ templates with C linkage, some macro magic is used
 // to generate implementations for all required type combinations that can be
 // called from MLIR generated code.
 //
 //===----------------------------------------------------------------------===//
 
-struct UnrankedMemRef {
-  uint64_t rank;
-  void *descriptor;
-};
-
 #define TEMPLATE(NAME, TYPE)                                                   \
   struct NAME {                                                                \
     const TYPE *base;                                                          \
     const TYPE *data;                                                          \
     uint64_t off;                                                              \
     uint64_t sizes[1];                                                         \
     uint64_t strides[1];                                                       \
   }
 
 #define CASE(p, i, v, P, I, V)                                                 \
   if (ptrTp == (p) && indTp == (i) && valTp == (v)) {                          \
     SparseTensor<V> *tensor;                                                   \
     if (action == 0)                                                           \
       tensor = openTensor<V>(static_cast<char *>(ptr), asize, sizes, perm);    \
+    else if (action == 1)                                                      \
+      tensor = static_cast<SparseTensor<V> *>(ptr);                            \
     else                                                                       \
-      tensor = copyTensor<V>(asize, sizes, perm, static_cast<V *>(ptr));       \
+      return SparseTensor<V>::newSparseTensor(asize, sizes, perm);             \
     return SparseTensorStorage<P, I, V>::newSparseTensor(tensor, sparsity);    \
   }
 
 #define IMPL1(RET, NAME, TYPE, LIB)                                            \
   RET NAME(void *tensor) {                                                     \
     std::vector<TYPE> *v;                                                      \
     static_cast<SparseTensorStorageBase *>(tensor)->LIB(&v);                   \
     return {v->data(), v->data(), 0, {v->size()}, {1}};                        \
   }
 
 #define IMPL2(RET, NAME, TYPE, LIB)                                            \
   RET NAME(void *tensor, uint64_t d) {                                         \
     std::vector<TYPE> *v;                                                      \
     static_cast<SparseTensorStorageBase *>(tensor)->LIB(&v, d);                \
     return {v->data(), v->data(), 0, {v->size()}, {1}};                        \
   }
 
-#define PTR(NAME)                                                              \
-  const void *NAME(int64_t sz, UnrankedMemRef *m) { return m->descriptor; }
+#define IMPL3(NAME, TYPE)                                                      \
+  void *NAME(void *tensor, TYPE value, uint64_t *ibase, uint64_t *idata,       \
+             uint64_t ioff, uint64_t isize, uint64_t istride, uint64_t *pbase, \
+             uint64_t *pdata, uint64_t poff, uint64_t psize,                   \
+             uint64_t pstride) {                                               \
+    assert(istride == 1 && pstride == 1 && isize == psize);                    \
+    uint64_t *indx = idata + ioff;                                             \
+    if (!value)                                                                \
+      return tensor;                                                           \
+    uint64_t *perm = pdata + poff;                                             \
+    std::vector<uint64_t> indices(isize);                                      \
+    for (uint64_t r = 0; r < isize; r++)                                       \
+      indices[perm[r]] = indx[r];                                              \
+    static_cast<SparseTensor<TYPE> *>(tensor)->add(indices, value);            \
+    return tensor;                                                             \
+  }
 
 TEMPLATE(MemRef1DU64, uint64_t);
 TEMPLATE(MemRef1DU32, uint32_t);
 TEMPLATE(MemRef1DU16, uint16_t);
 TEMPLATE(MemRef1DU8, uint8_t);
 TEMPLATE(MemRef1DI64, int64_t);
 TEMPLATE(MemRef1DI32, int32_t);
 TEMPLATE(MemRef1DI16, int16_t);
enum PrimaryTypeEnum : uint64_t {
   kI64 = 3,
   kI32 = 4,
   kI16 = 5,
   kI8 = 6
 };
 
 /// Constructs a new sparse tensor. This is the "swiss army knife"
 /// method for materializing sparse tensors into the computation.
+///  action
+///    0 : ptr contains filename to read into storage
+///    1 : ptr contains coordinate scheme to assign to storage
+///    2 : returns coordinate scheme to fill (call back later with 1)
 void *newSparseTensor(uint8_t *abase, uint8_t *adata, uint64_t aoff,
                       uint64_t asize, uint64_t astride, uint64_t *sbase,
                       uint64_t *sdata, uint64_t soff, uint64_t ssize,
                       uint64_t sstride, uint64_t *pbase, uint64_t *pdata,
                       uint64_t poff, uint64_t psize, uint64_t pstride,
                       uint64_t ptrTp, uint64_t indTp, uint64_t valTp,
                       uint32_t action, void *ptr) {
   assert(astride == 1 && sstride == 1 && pstride == 1);
+  assert(asize == ssize && ssize == psize);
   uint8_t *sparsity = adata + aoff;
   uint64_t *sizes = sdata + soff;
   uint64_t *perm = pdata + poff;
 
   // 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);

 IMPL1(MemRef1DI16, sparseValuesI16, int16_t, getValues)
 IMPL1(MemRef1DI8, sparseValuesI8, int8_t, getValues)
 
 /// Releases sparse tensor storage.
 void delSparseTensor(void *tensor) {
   delete static_cast<SparseTensorStorageBase *>(tensor);
 }
 
-/// Helper to get pointer, one per value type.
-PTR(getPtrF64)
-PTR(getPtrF32)
-PTR(getPtrI64)
-PTR(getPtrI32)
-PTR(getPtrI16)
-PTR(getPtrI8)
+/// Helper to add value to coordinate scheme, one per value type.
+IMPL3(addEltF64, double)
+IMPL3(addEltF32, float)
+IMPL3(addEltI64, int64_t)
+IMPL3(addEltI32, int32_t)
+IMPL3(addEltI16, int16_t)
+IMPL3(addEltI8, int8_t)
 
 #undef TEMPLATE
 #undef CASE
 #undef IMPL1
 #undef IMPL2
+#undef IMPL3
 
 } // extern "C"
 
 #endif // MLIR_CRUNNERUTILS_DEFINE_FUNCTIONS

mlir/test/Dialect/SparseTensor/conversion.mlir

 //   CHECK-DAG: %[[Z:.*]] = tensor.cast %[[W]] : tensor<3xi64> to tensor<?xi64>
 //       CHECK: %[[T:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[A]])
 //       CHECK: return %[[T]] : !llvm.ptr<i8>
 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_convert(
+// CHECK-LABEL: func @sparse_convert_1d(
+//  CHECK-SAME: %[[A:.*]]: tensor<?xi32>) -> !llvm.ptr<i8>
+//   CHECK-DAG: %[[C0:.*]] = constant 0 : index
+//   CHECK-DAG: %[[C1:.*]] = constant 1 : index
+//   CHECK-DAG: %[[D0:.*]] = constant dense<0> : tensor<1xi64>
+//   CHECK-DAG: %[[D1:.*]] = constant dense<1> : tensor<1xi8>
+//   CHECK-DAG: %[[X:.*]] = tensor.cast %[[D1]] : tensor<1xi8> to tensor<?xi8>
+//   CHECK-DAG: %[[Y:.*]] = tensor.cast %[[D0]] : tensor<1xi64> to tensor<?xi64>
+//       CHECK: %[[C:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Y]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.}})
+//       CHECK: %[[M:.*]] = memref.alloca() : memref<1xindex>
+//       CHECK: %[[T:.*]] = memref.cast %[[M]] : memref<1xindex> to memref<?xindex>
+//       CHECK: %[[U:.*]] = tensor.dim %[[A]], %[[C0]] : tensor<?xi32>
+//       CHECK: scf.for %[[I:.*]] = %[[C0]] to %[[U]] step %[[C1]] {
+//       CHECK:   %[[E:.*]] = tensor.extract %[[A]][%[[I]]] : tensor<?xi32>
+//       CHECK:   memref.store %[[I]], %[[M]][%[[C0]]] : memref<1xindex>
+//       CHECK:   call @addEltI32(%[[C]], %[[E]], %[[T]], %[[Y]])
+//       CHECK: }
+//       CHECK: %[[T:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Y]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[C]])
+//       CHECK: return %[[T]] : !llvm.ptr<i8>
+func @sparse_convert_1d(%arg0: tensor<?xi32>) -> tensor<?xi32, #SparseVector> {
+  %0 = sparse_tensor.convert %arg0 : tensor<?xi32> to tensor<?xi32, #SparseVector>
+  return %0 : tensor<?xi32, #SparseVector>
+}
+
+// CHECK-LABEL: func @sparse_convert_2d(
 //  CHECK-SAME: %[[A:.*]]: tensor<2x4xf64>) -> !llvm.ptr<i8>
+//   CHECK-DAG: %[[C0:.*]] = constant 0 : index
+//   CHECK-DAG: %[[C1:.*]] = constant 1 : index
 //   CHECK-DAG: %[[U:.*]] = constant dense<[0, 1]> : tensor<2xi8>
 //   CHECK-DAG: %[[V:.*]] = constant dense<[2, 4]> : tensor<2xi64>
 //   CHECK-DAG: %[[W:.*]] = constant dense<[0, 1]> : tensor<2xi64>
-//       CHECK: %[[C:.*]] = memref.buffer_cast %arg0 : memref<2x4xf64>
-//       CHECK: %[[M:.*]] = memref.cast %[[C]] : memref<2x4xf64> to memref<*xf64>
-//       CHECK: %[[C:.*]] = call @getPtrF64(%[[M]]) : (memref<*xf64>) -> !llvm.ptr<i8>
 //   CHECK-DAG: %[[X:.*]] = tensor.cast %[[U]] : tensor<2xi8> to tensor<?xi8>
 //   CHECK-DAG: %[[Y:.*]] = tensor.cast %[[V]] : tensor<2xi64> to tensor<?xi64>
 //   CHECK-DAG: %[[Z:.*]] = tensor.cast %[[W]] : tensor<2xi64> to tensor<?xi64>
+//       CHECK: %[[C:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.}})
+//       CHECK: %[[M:.*]] = memref.alloca() : memref<2xindex>
+//       CHECK: %[[T:.*]] = memref.cast %[[M]] : memref<2xindex> to memref<?xindex>
+//       CHECK: scf.for %[[I:.*]] = %[[C0]] to %{{.*}} step %[[C1]] {
+//       CHECK:   scf.for %[[J:.*]] = %[[C0]] to %{{.*}} step %[[C1]] {
+//       CHECK:     %[[E:.*]] = tensor.extract %[[A]][%[[I]], %[[J]]] : tensor<2x4xf64>
+//       CHECK:     memref.store %[[I]], %[[M]][%[[C0]]] : memref<2xindex>
+//       CHECK:     memref.store %[[J]], %[[M]][%[[C1]]] : memref<2xindex>
+//       CHECK:     call @addEltF64(%[[C]], %[[E]], %[[T]], %[[Z]])
+//       CHECK:   }
+//       CHECK: }
 //       CHECK: %[[T:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[C]])
 //       CHECK: return %[[T]] : !llvm.ptr<i8>
-func @sparse_convert(%arg0: tensor<2x4xf64>) -> tensor<2x4xf64, #SparseMatrix> {
+func @sparse_convert_2d(%arg0: tensor<2x4xf64>) -> tensor<2x4xf64, #SparseMatrix> {
   %0 = sparse_tensor.convert %arg0 : tensor<2x4xf64> to tensor<2x4xf64, #SparseMatrix>
   return %0 : tensor<2x4xf64, #SparseMatrix>
 }
 
+// CHECK-LABEL: func @sparse_convert_3d(
+//  CHECK-SAME: %[[A:.*]]: tensor<?x?x?xf64>) -> !llvm.ptr<i8>
+//   CHECK-DAG: %[[C0:.*]] = constant 0 : index
+//   CHECK-DAG: %[[C1:.*]] = constant 1 : index
+//   CHECK-DAG: %[[C2:.*]] = constant 2 : index
+//   CHECK-DAG: %[[U:.*]] = constant dense<[0, 1, 1]> : tensor<3xi8>
+//   CHECK-DAG: %[[V:.*]] = constant dense<0> : tensor<3xi64>
+//   CHECK-DAG: %[[W:.*]] = constant dense<[1, 2, 0]> : tensor<3xi64>
+//   CHECK-DAG: %[[X:.*]] = tensor.cast %[[U]] : tensor<3xi8> to tensor<?xi8>
+//   CHECK-DAG: %[[Y:.*]] = tensor.cast %[[V]] : tensor<3xi64> to tensor<?xi64>
+//   CHECK-DAG: %[[Z:.*]] = tensor.cast %[[W]] : tensor<3xi64> to tensor<?xi64>
+//       CHECK: %[[C:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.}})
+//       CHECK: %[[M:.*]] = memref.alloca() : memref<3xindex>
+//       CHECK: %[[T:.*]] = memref.cast %[[M]] : memref<3xindex> to memref<?xindex>
+//       CHECK: %[[U1:.*]] = tensor.dim %[[A]], %[[C0]] : tensor<?x?x?xf64>
+//       CHECK: %[[U2:.*]] = tensor.dim %[[A]], %[[C1]] : tensor<?x?x?xf64>
+//       CHECK: %[[U3:.*]] = tensor.dim %[[A]], %[[C2]] : tensor<?x?x?xf64>
+//       CHECK: scf.for %[[I:.*]] = %[[C0]] to %[[U1]] step %[[C1]] {
+//       CHECK:   scf.for %[[J:.*]] = %[[C0]] to %[[U2]] step %[[C1]] {
+//       CHECK:     scf.for %[[K:.*]] = %[[C0]] to %[[U3]] step %[[C1]] {
+//       CHECK:       %[[E:.*]] = tensor.extract %[[A]][%[[I]], %[[J]], %[[K]]] : tensor<?x?x?xf64>
+//       CHECK:       memref.store %[[I]], %[[M]][%[[C0]]] : memref<3xindex>
+//       CHECK:       memref.store %[[J]], %[[M]][%[[C1]]] : memref<3xindex>
+//       CHECK:       memref.store %[[K]], %[[M]][%[[C2]]] : memref<3xindex>
+//       CHECK:       call @addEltF64(%[[C]], %[[E]], %[[T]], %[[Z]])
+//       CHECK:     }
+//       CHECK:   }
+//       CHECK: }
+//       CHECK: %[[T:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[C]])
+//       CHECK: return %[[T]] : !llvm.ptr<i8>
+func @sparse_convert_3d(%arg0: tensor<?x?x?xf64>) -> tensor<?x?x?xf64, #SparseTensor> {
+  %0 = sparse_tensor.convert %arg0 : tensor<?x?x?xf64> to tensor<?x?x?xf64, #SparseTensor>
+  return %0 : tensor<?x?x?xf64, #SparseTensor>
+}
+
 // CHECK-LABEL: func @sparse_pointers(
 //  CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>)
 //       CHECK: %[[C:.*]] = constant 0 : index
 //       CHECK: %[[T:.*]] = call @sparsePointers(%[[A]], %[[C]]) : (!llvm.ptr<i8>, index) -> memref<?xindex>
 //       CHECK: return %[[T]] : memref<?xindex>
 func @sparse_pointers(%arg0: tensor<128xf64, #SparseVector>) -> memref<?xindex> {
   %c = constant 0 : index
   %0 = sparse_tensor.pointers %arg0, %c : tensor<128xf64, #SparseVector> to memref<?xindex>