[mlir][sparse] implement sparse tensor init operation

[aartbik]

Summary

Next step towards supporting sparse tensors outputs.
Also some minor refactoring of enum constants as well
as replacing tensor arguments with proper buffer arguments
(latter is required for more general sizes arguments for
the sparse_tensor.init operation, as well as more general
spares_tensor.convert operations later)

[aartbik]

rebased with main

[wrengr]

Overall I like the changes in SparseTensorConversion.cpp, but I'm a little worried about switching from APInt to attributes because of the comment about blowing up compiler memory usage in mlir/lib/Dialect/Linalg/Transforms/ElementwiseOpFusion.cpp on lines 1380–1382. Do we have anything set up to track performance metrics and prevent regressions?

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

32

Typo: ".." -> "."

148

This should be called genAlloca (or similar) since it generates an alloca rather than a malloc

179–182

Not sure if that'll allow the last statement to fit on one line, but it's a cleanup I've been meaning to make for a while. Feel free to choose a better name for dlt of course.

[wrengr]

mlir/lib/ExecutionEngine/SparseUtils.cpp

184

Shouldn't that just be "per-dimension"?

190–191

Throughout this file, I feel like the old d name makes more sense, since we're iterating through dimensions d1,d2,...dr up to the rank r.

194

Given that we keep calling this for all the loops, it should be floated out to the top of the function, to reduce boilerplate

267

Please rename this to rank, to avoid confusion with the sizes parameter and for consistency with the rest of the file

540

+1 :)

[wrengr]

conversion.mlir looks at least conceptually good to me, though I'm not familiar enough with FileCheck to sign off on it.

The other three files LGTM once the comments are addressed

[aartbik]

addressed Wreng's comments, thanks!

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

 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 using namespace mlir;
 using namespace mlir::sparse_tensor;
 
 namespace {
 
+/// New tensor storage action times.. Keep these values consistent with
+/// the sparse runtime support library.
+enum Action : uint32_t {
+  kEmpty = 0,
+  kFromFile = 1,
+  kFromCOO = 2,
+  kEmptyCOO = 3,
+  kToCOO = 4
+};
+
 //===----------------------------------------------------------------------===//
 // Helper methods.
 //===----------------------------------------------------------------------===//
 
 /// Returns internal type encoding for primary storage. Keep these
 /// values consistent with the sparse runtime support library.
 static unsigned getPrimaryTypeEncoding(Type tp) {
   if (tp.isF64())

 }
 
 /// Generates a constant of `i32` type.
 inline static Value constantI32(ConversionPatternRewriter &rewriter,
                                 Location loc, int32_t i) {
   return rewriter.create<arith::ConstantIntOp>(loc, i, 32);
 }
 
-/// Returns integers of given width and values as a constant tensor.
-/// We cast the static shape into a dynamic shape to ensure that the
-/// method signature remains uniform across different tensor dimensions.
-static Value getTensor(ConversionPatternRewriter &rewriter, unsigned width,
-                       Location loc, ArrayRef<APInt> values) {
-  Type etp = rewriter.getIntegerType(width);
-  unsigned sz = values.size();
-  RankedTensorType tt1 = RankedTensorType::get({sz}, etp);
-  RankedTensorType tt2 = RankedTensorType::get({ShapedType::kDynamicSize}, etp);
-  auto elts = rewriter.create<arith::ConstantOp>(
-      loc, DenseElementsAttr::get(tt1, values));
-  return rewriter.create<tensor::CastOp>(loc, tt2, elts);
+/// Generates a constant of `i8` type.
+inline static Value constantI8(ConversionPatternRewriter &rewriter,
+                               Location loc, int8_t i) {
+  return rewriter.create<arith::ConstantIntOp>(loc, i, 8);
 }
 
 /// Returns a function reference (first hit also inserts into module). Sets
 /// the "_emit_c_interface" on the function declaration when requested,
 /// so that LLVM lowering generates a wrapper function that takes care
 /// of ABI complications with passing in and returning MemRefs to C functions.
 static FlatSymbolRefAttr getFunc(Operation *op, StringRef name,
                                  TypeRange resultType, ValueRange operands,
func = moduleBuilder.create<FuncOp>(
         FunctionType::get(context, operands.getTypes(), resultType));
     func.setPrivate();
     if (emitCInterface)
       func->setAttr("llvm.emit_c_interface", UnitAttr::get(context));
   }
   return result;
 }
 
+/// Generates a temporary buffer of the given size and type.
+static Value genAlloc(ConversionPatternRewriter &rewriter, Location loc,
+                      unsigned sz, Type tp) {
+  auto memTp = MemRefType::get({ShapedType::kDynamicSize}, tp);
+  Value a = constantIndex(rewriter, loc, sz);
+  return rewriter.create<memref::AllocaOp>(loc, memTp, ValueRange{a});
+}
+
+/// Fills a temporary buffer of the given type with arguments.
+static Value genBuffer(ConversionPatternRewriter &rewriter, Location loc,
+                       ArrayRef<Value> values) {
+  unsigned sz = values.size();
+  assert(sz >= 1);
+  Value buffer = genAlloc(rewriter, loc, sz, values[0].getType());
+  for (unsigned i = 0; i < sz; i++) {
+    Value idx = constantIndex(rewriter, loc, i);
+    rewriter.create<memref::StoreOp>(loc, values[i], buffer, idx);
+  }
+  return buffer;
+}
+
 /// Generates a call into the "swiss army knife" method of the sparse runtime
 /// 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 &perm, Value ptr = Value()) {
+                        Value &perm, ValueRange szs, 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;
+  SmallVector<Value, 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));
+    attrs.push_back(constantI8(
+        rewriter, loc, getDimLevelTypeEncoding(enc.getDimLevelType()[i])));
+  params.push_back(genBuffer(rewriter, 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;
+  SmallVector<Value, 4> sizes;
+  if (szs.size() > 0) {
+    for (Value s : szs)
+      sizes.push_back(
+          rewriter.create<arith::IndexCastOp>(loc, s, rewriter.getI64Type()));
+  } else {
     for (unsigned i = 0; i < sz; i++) {
       uint64_t s = shape[i] == ShapedType::kDynamicSize ? 0 : shape[i];
-    sizes.push_back(APInt(64, s));
+      sizes.push_back(constantI64(rewriter, loc, s));
     }
-  params.push_back(getTensor(rewriter, 64, loc, sizes));
+  }
+  params.push_back(genBuffer(rewriter, 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> rev(sz);
+  SmallVector<Value, 4> rev(sz);
   if (AffineMap p = enc.getDimOrdering()) {
     for (unsigned i = 0; i < sz; i++)
-      rev[p.getDimPosition(i)] = APInt(64, i);
+      rev[p.getDimPosition(i)] = constantI64(rewriter, loc, i);
   } else {
     for (unsigned i = 0; i < sz; i++)
-      rev[i] = APInt(64, i);
+      rev[i] = constantI64(rewriter, loc, i);
   }
-  perm = getTensor(rewriter, 64, loc, rev);
+  perm = genBuffer(rewriter, 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(constantI64(rewriter, loc, secPtr));
   params.push_back(constantI64(rewriter, loc, secInd));
Value val = rewriter.create<tensor::ExtractOp>(loc, indices,
                                                    ValueRange{ivs[0], idx});
     val =
         rewriter.create<arith::IndexCastOp>(loc, val, rewriter.getIndexType());
     rewriter.create<memref::StoreOp>(loc, val, ind, idx);
   }
   return rewriter.create<tensor::ExtractOp>(loc, values, ivs[0]);
 }
 
-/// Generates code to stack-allocate a `memref<?xindex>` where the `?`
-/// is the given `rank`.  This array is intended to serve as a reusable
-/// buffer for storing the indices of a single tensor element, to avoid
-/// allocation in the body of loops.
-static Value allocaIndices(ConversionPatternRewriter &rewriter, Location loc,
-                           int64_t rank) {
-  auto indexTp = rewriter.getIndexType();
-  auto memTp = MemRefType::get({ShapedType::kDynamicSize}, indexTp);
-  Value arg = constantIndex(rewriter, loc, rank);
-  return rewriter.create<memref::AllocaOp>(loc, memTp, ValueRange{arg});
-}
-
 //===----------------------------------------------------------------------===//
 // Conversion rules.
 //===----------------------------------------------------------------------===//
 
 /// Sparse conversion rule for returns.
 class SparseReturnConverter : public OpConversionPattern<ReturnOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
class SparseTensorNewConverter : public OpConversionPattern<NewOp> {
   LogicalResult
   matchAndRewrite(NewOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Type resType = op.getType();
     auto enc = getSparseTensorEncoding(resType);
     if (!enc)
       return failure();
     Value perm;
+    rewriter.replaceOp(op, genNewCall(rewriter, op, enc, kFromFile, perm, {},
+                                      adaptor.getOperands()[0]));
+    return success();
+  }
+};
+
+/// Sparse conversion rule for the init operator.
+class SparseTensorInitConverter : public OpConversionPattern<InitOp> {
+  using OpConversionPattern::OpConversionPattern;
+  LogicalResult
+  matchAndRewrite(InitOp op, OpAdaptor adaptor,
+                  ConversionPatternRewriter &rewriter) const override {
+    Type resType = op.getType();
+    auto enc = getSparseTensorEncoding(resType);
+    if (!enc)
+      return failure();
+    Value perm;
     rewriter.replaceOp(
-        op, genNewCall(rewriter, op, enc, 0, perm, adaptor.getOperands()[0]));
+        op, genNewCall(rewriter, op, enc, kEmpty, perm, adaptor.getOperands()));
     return success();
   }
 };
 
 /// Sparse conversion rule for the convert operator.
 class SparseTensorConvertConverter : public OpConversionPattern<ConvertOp> {
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(ConvertOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Type resType = op.getType();
     auto encDst = getSparseTensorEncoding(resType);
     auto encSrc = getSparseTensorEncoding(op.source().getType());
     auto src = adaptor.getOperands()[0];
     if (encDst && encSrc) {
       // This is a sparse => sparse conversion, which is handled as follows:
       //   t = src->toCOO();         ; src to COO in dst order
       //   dst = newSparseTensor(t)
       // Using the coordinate scheme as an intermediate does not always
       // yield the fastest conversion but avoids the need for a full
       // O(N^2) conversion matrix.
       Value perm;
-      Value coo = genNewCall(rewriter, op, encDst, 3, perm, src);
-      rewriter.replaceOp(op, genNewCall(rewriter, op, encDst, 1, perm, coo));
+      Value coo = genNewCall(rewriter, op, encDst, kToCOO, perm, {}, src);
+      rewriter.replaceOp(
+          op, genNewCall(rewriter, op, encDst, kFromCOO, perm, {}, coo));
       return success();
     }
     if (!encDst || encSrc) {
       // TODO: sparse => dense
       return failure();
     }
     // This is a dense => sparse conversion or a sparse constant in COO =>
     // sparse conversion, which is handled as follows:
matchAndRewrite(ConvertOp op, OpAdaptor adaptor,
     // using MLIR IR to avoid having to expose too much low-level
     // memref traversal details to the runtime support library.
     // Also note that the code below only generates the "new" ops and
     // the loop-nest per se; whereas the entire body of the innermost
     // loop is generated by genAddElt().
     Location loc = op->getLoc();
     ShapedType shape = resType.cast<ShapedType>();
     Value perm;
-    Value ptr = genNewCall(rewriter, op, encDst, 2, perm);
-    Value ind = allocaIndices(rewriter, loc, shape.getRank());
+    Value ptr = genNewCall(rewriter, op, encDst, kEmptyCOO, perm, {});
+    Value ind =
+        genAlloc(rewriter, loc, shape.getRank(), rewriter.getIndexType());
     SmallVector<Value> lo;
     SmallVector<Value> hi;
     SmallVector<Value> st;
     Value zero = constantIndex(rewriter, loc, 0);
     Value one = constantIndex(rewriter, loc, 1);
     auto indicesValues = genSplitSparseConstant(rewriter, loc, src);
     bool isCOOConstant = indicesValues.hasValue();
     Value indices;
scf::buildLoopNest(
           if (isCOOConstant)
             val = genIndexAndValueForSparse(rewriter, loc, indices, values, ind,
                                             ivs, rank);
           else
             val = genIndexAndValueForDense(rewriter, loc, src, ind, ivs);
           genAddEltCall(rewriter, op, eltType, ptr, val, ind, perm);
           return {};
         });
-    rewriter.replaceOp(op, genNewCall(rewriter, op, encDst, 1, perm, ptr));
+    rewriter.replaceOp(
+        op, genNewCall(rewriter, op, encDst, kFromCOO, perm, {}, ptr));
     return success();
   }
 };
 
 /// Sparse conversion rule for the release operator.
 class SparseTensorReleaseConverter : public OpConversionPattern<ReleaseOp> {
 public:
   using OpConversionPattern::OpConversionPattern;

 // Public method for populating conversion rules.
 //===----------------------------------------------------------------------===//
 
 /// Populates the given patterns list with conversion rules required for
 /// the sparsification of linear algebra operations.
 void mlir::populateSparseTensorConversionPatterns(TypeConverter &typeConverter,
                                                   RewritePatternSet &patterns) {
   patterns.add<SparseReturnConverter, SparseTensorToDimSizeConverter,
-               SparseTensorNewConverter, SparseTensorConvertConverter,
-               SparseTensorReleaseConverter, SparseTensorToPointersConverter,
-               SparseTensorToIndicesConverter, SparseTensorToValuesConverter,
-               SparseTensorToTensorConverter>(typeConverter,
-                                              patterns.getContext());
+               SparseTensorNewConverter, SparseTensorInitConverter,
+               SparseTensorConvertConverter, SparseTensorReleaseConverter,
+               SparseTensorToPointersConverter, SparseTensorToIndicesConverter,
+               SparseTensorToValuesConverter, SparseTensorToTensorConverter>(
+      typeConverter, patterns.getContext());
 }

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>();
+    target.addIllegalOp<ConvertOp, NewOp, ToIndicesOp, ToPointersOp, ToTensorOp,
+                        ToValuesOp>();
     // 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.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<arith::ConstantOp, ConstantOp, arith::IndexCastOp,
-                      tensor::CastOp, tensor::ExtractOp, arith::CmpFOp,
-                      arith::CmpIOp>();
-    target.addLegalDialect<scf::SCFDialect, LLVM::LLVMDialect,
-                           memref::MemRefDialect>();
+    target.addLegalOp<arith::CmpFOp, arith::CmpIOp, arith::ConstantOp,
+                      arith::IndexCastOp, tensor::CastOp, tensor::ExtractOp>();
+    target.addLegalDialect<LLVM::LLVMDialect, memref::MemRefDialect,
+                           scf::SCFDialect>();
     // 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

 /// annotations on dense/sparse. This data structure provides a bufferized
 /// form of a sparse tensor type. 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 a sparse tensor storage scheme from the given sparse
-  /// tensor in coordinate scheme following the given per-rank dimension
-  /// dense/sparse annotations.
-  SparseTensorStorage(SparseTensorCOO<V> *tensor, const uint8_t *sparsity,
-                      const uint64_t *perm)
-      : sizes(tensor->getSizes()), rev(getRank()), pointers(getRank()),
-        indices(getRank()) {
+  /// Constructs a sparse tensor storage scheme with the given dimensions,
+  /// permutation, and per-rank dimension dense/sparse annotations, using
+  /// the coordinate scheme tensor for the initial contents if provided.
+  SparseTensorStorage(const std::vector<uint64_t> &szs, const uint64_t *perm,
+                      const uint8_t *sparsity, SparseTensorCOO<V> *tensor)
+      : sizes(szs), rev(getRank()), pointers(getRank()), indices(getRank()) {
     // Store "reverse" permutation.
-    for (uint64_t d = 0, rank = getRank(); d < rank; d++)
-      rev[perm[d]] = d;
-    // Provide hints on capacity.
+    for (uint64_t r = 0, rank = getRank(); r < rank; r++)
+      rev[perm[r]] = r;
+    // Provide hints on capacity of pointers and indices.
     // 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];
-      if (sparsity[d] == kCompressed) {
-        pointers[d].reserve(s + 1);
-        indices[d].reserve(s);
+    for (uint64_t r = 0, s = 1, rank = getRank(); r < rank; r++) {
+      s *= sizes[r];
+      if (sparsity[r] == kCompressed) {
+        pointers[r].reserve(s + 1);
+        indices[r].reserve(s);
         s = 1;
       } else {
-        assert(sparsity[d] == kDense && "singleton not yet supported");
+        assert(sparsity[r] == kDense && "singleton not yet supported");
       }
     }
     // Prepare sparse pointer structures for all dimensions.
-    for (uint64_t d = 0, rank = getRank(); d < rank; d++)
-      if (sparsity[d] == kCompressed)
-        pointers[d].push_back(0);
-    // Then setup the tensor.
+    for (uint64_t r = 0, rank = getRank(); r < rank; r++)
+      if (sparsity[r] == kCompressed)
+        pointers[r].push_back(0);
+    // Then assign contents from coordinate scheme tensor if provided.
+    if (tensor) {
+      uint64_t nnz = tensor->getElements().size();
+      values.reserve(nnz);
       fromCOO(tensor, sparsity, 0, nnz, 0);
     }
+  }
 
   virtual ~SparseTensorStorage() {}
 
   /// Get the rank of the tensor.
   uint64_t getRank() const { return sizes.size(); }
 
   /// Get the size in the given dimension of the tensor.
   uint64_t getDimSize(uint64_t d) override {
SparseTensorCOO<V> *toCOO(const uint64_t *perm) {
     for (uint64_t r = 0; r < size; r++)
       reord[r] = perm[rev[r]];
     std::vector<uint64_t> idx(size);
     toCOO(tensor, reord, idx, 0, 0);
     assert(tensor->getElements().size() == values.size());
     return tensor;
   }
 
-  /// Factory method. Expects a coordinate scheme that respects the same
-  /// permutation as is desired for the new sparse storage scheme.
-  static SparseTensorStorage<P, I, V> *newSparseTensor(SparseTensorCOO<V> *t,
-                                                       const uint8_t *sparsity,
-                                                       const uint64_t *perm) {
-    t->sort(); // sort lexicographically
-    SparseTensorStorage<P, I, V> *n =
-        new SparseTensorStorage<P, I, V>(t, sparsity, perm);
-    delete t;
+  /// Factory method. Constructs a sparse tensor storage scheme with the given
+  /// dimensions, permutation, and per-rank dimension dense/sparse annotations,
+  /// using the coordinate scheme tensor for the initial contents if provided.
+  /// In the latter case, the coordinate scheme must respect the same
+  /// permutation as is desired for the new sparse tensor storage.
+  static SparseTensorStorage<P, I, V> *
+  newSparseTensor(uint64_t size, const uint64_t *sizes, const uint64_t *perm,
+                  const uint8_t *sparsity, SparseTensorCOO<V> *tensor) {
+    SparseTensorStorage<P, I, V> *n = nullptr;
+    if (tensor) {
+      assert(tensor->getRank() == size);
+      for (uint64_t r = 0; r < size; r++)
+        assert(sizes[r] == 0 || sizes[r] == tensor->getSizes()[perm[r]]);
+      tensor->sort(); // sort lexicographically
+      n = new SparseTensorStorage<P, I, V>(tensor->getSizes(), perm, sparsity,
+                                           tensor);
+      delete tensor;
+    } else {
+      std::vector<uint64_t> permsz(size);
+      for (uint64_t r = 0; r < size; r++)
+        permsz[perm[r]] = sizes[r];
+      n = new SparseTensorStorage<P, I, V>(permsz, perm, sparsity, tensor);
+    }
     return n;
   }
 
 private:
   /// 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 fromCOO(SparseTensorCOO<V> *tensor, const uint8_t *sparsity, uint64_t lo,
enum PrimaryTypeEnum : uint64_t {
   kF64 = 1,
   kF32 = 2,
   kI64 = 3,
   kI32 = 4,
   kI16 = 5,
   kI8 = 6
 };
 
-enum Action : uint32_t { kFromFile = 0, kFromCOO = 1, kNewCOO = 2, kToCOO = 3 };
+enum Action : uint32_t {
+  kEmpty = 0,
+  kFromFile = 1,
+  kFromCOO = 2,
+  kEmptyCOO = 3,
+  kToCOO = 4
+};
 
 #define CASE(p, i, v, P, I, V)                                                 \
   if (ptrTp == (p) && indTp == (i) && valTp == (v)) {                          \
     SparseTensorCOO<V> *tensor = nullptr;                                      \
     if (action == kFromFile)                                                   \
       tensor =                                                                 \
           openSparseTensorCOO<V>(static_cast<char *>(ptr), size, sizes, perm); \
     else if (action == kFromCOO)                                               \
       tensor = static_cast<SparseTensorCOO<V> *>(ptr);                         \
-    else if (action == kNewCOO)                                                \
+    else if (action == kEmptyCOO)                                              \
       return SparseTensorCOO<V>::newSparseTensorCOO(size, sizes, perm);        \
     else if (action == kToCOO)                                                 \
       return static_cast<SparseTensorStorage<P, I, V> *>(ptr)->toCOO(perm);    \
     else                                                                       \
-      assert(0);                                                               \
-    return SparseTensorStorage<P, I, V>::newSparseTensor(tensor, sparsity,     \
-                                                         perm);                \
+      assert(action == kEmpty);                                                \
+    return SparseTensorStorage<P, I, V>::newSparseTensor(size, sizes, perm,    \
+                                                         sparsity, tensor);    \
   }
 
 #define IMPL1(NAME, TYPE, LIB)                                                 \
   void _mlir_ciface_##NAME(StridedMemRefType<TYPE, 1> *ref, void *tensor) {    \
     assert(ref);                                                               \
     assert(tensor);                                                            \
     std::vector<TYPE> *v;                                                      \
     static_cast<SparseTensorStorageBase *>(tensor)->LIB(&v);                   \
void *_mlir_ciface_##NAME(void *tensor, TYPE value,                          \
     static_cast<SparseTensorCOO<TYPE> *>(tensor)->add(indices, value);         \
     return tensor;                                                             \
   }
 
 /// Constructs a new sparse tensor. This is the "swiss army knife"
 /// method for materializing sparse tensors into the computation.
 ///
 /// action:
-/// kFromFile = ptr contains filename to read into storage
-/// kFromCOO = ptr contains coordinate scheme to assign to new storage
-/// kNewCOO = returns empty coordinate scheme to fill and use with kFromCOO
+/// kEmpty = returns empty storage to fill later
+/// kFromFile = returns storage, where ptr contains filename to read
+/// kFromCOO = returns storage, where ptr contains coordinate scheme to assign
+/// kEmptyCOO = returns empty coordinate scheme to fill and use with kFromCOO
 /// kToCOO = returns coordinate scheme from storage in ptr to use with kFromCOO
 void *
 _mlir_ciface_newSparseTensor(StridedMemRefType<uint8_t, 1> *aref, // NOLINT
                              StridedMemRefType<uint64_t, 1> *sref,
                              StridedMemRefType<uint64_t, 1> *pref,
                              uint64_t ptrTp, uint64_t indTp, uint64_t valTp,
                              uint32_t action, void *ptr) {
   assert(aref && sref && pref);
void *convertToMLIRSparseTensor(uint64_t rank, uint64_t nse, uint64_t *shape,
   for (uint64_t i = 0, base = 0; i < nse; i++) {
     for (uint64_t j = 0; j < rank; j++)
       idx[j] = indices[base + j];
     tensor->add(idx, values[i]);
     base += rank;
   }
   // Return sparse tensor storage format as opaque pointer.
   return SparseTensorStorage<uint64_t, uint64_t, double>::newSparseTensor(
-      tensor, sparse.data(), perm.data());
+      rank, shape, perm.data(), sparse.data(), tensor);
 }
 
 } // extern "C"
 
 #endif // MLIR_CRUNNERUTILS_DEFINE_FUNCTIONS

mlir/test/Dialect/SparseTensor/conversion.mlir

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: %[[U:.*]] = arith.constant dense<1> : tensor<1xi8>
-//   CHECK-DAG: %[[V:.*]] = arith.constant dense<128> : tensor<1xi64>
-//   CHECK-DAG: %[[W:.*]] = arith.constant dense<0> : tensor<1xi64>
-//   CHECK-DAG: %[[X:.*]] = tensor.cast %[[U]] : tensor<1xi8> to tensor<?xi8>
-//   CHECK-DAG: %[[Y:.*]] = tensor.cast %[[V]] : tensor<1xi64> to tensor<?xi64>
-//   CHECK-DAG: %[[Z:.*]] = tensor.cast %[[W]] : tensor<1xi64> to tensor<?xi64>
+//   CHECK-DAG: %[[P:.*]] = memref.alloca() : memref<1xi8>
+//   CHECK-DAG: %[[Q:.*]] = memref.alloca() : memref<1xi64>
+//   CHECK-DAG: %[[R:.*]] = memref.alloca() : memref<1xi64>
+//   CHECK-DAG: %[[X:.*]] = memref.cast %[[P]] : memref<1xi8> to memref<?xi8>
+//   CHECK-DAG: %[[Y:.*]] = memref.cast %[[Q]] : memref<1xi64> to memref<?xi64>
+//   CHECK-DAG: %[[Z:.*]] = memref.cast %[[R]] : memref<1xi64> to memref<?xi64>
 //       CHECK: %[[T:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[A]])
 //       CHECK: return %[[T]] : !llvm.ptr<i8>
 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: %[[U:.*]] = arith.constant dense<[0, 1]> : tensor<2xi8>
-//   CHECK-DAG: %[[V:.*]] = arith.constant dense<0> : tensor<2xi64>
-//   CHECK-DAG: %[[W:.*]] = arith.constant dense<[0, 1]> : tensor<2xi64>
-//   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-DAG: %[[P:.*]] = memref.alloca() : memref<2xi8>
+//   CHECK-DAG: %[[Q:.*]] = memref.alloca() : memref<2xi64>
+//   CHECK-DAG: %[[R:.*]] = memref.alloca() : memref<2xi64>
+//   CHECK-DAG: %[[X:.*]] = memref.cast %[[P]] : memref<2xi8> to memref<?xi8>
+//   CHECK-DAG: %[[Y:.*]] = memref.cast %[[Q]] : memref<2xi64> to memref<?xi64>
+//   CHECK-DAG: %[[Z:.*]] = memref.cast %[[R]] : memref<2xi64> to memref<?xi64>
 //       CHECK: %[[T:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[A]])
 //       CHECK: return %[[T]] : !llvm.ptr<i8>
 func @sparse_new2d(%arg0: !llvm.ptr<i8>) -> tensor<?x?xf32, #SparseMatrix> {
   %0 = sparse_tensor.new %arg0 : !llvm.ptr<i8> to tensor<?x?xf32, #SparseMatrix>
   return %0 : tensor<?x?xf32, #SparseMatrix>
 }
 
 // CHECK-LABEL: func @sparse_new3d(
 //  CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) -> !llvm.ptr<i8>
-//   CHECK-DAG: %[[U:.*]] = arith.constant dense<[0, 1, 1]> : tensor<3xi8>
-//   CHECK-DAG: %[[V:.*]] = arith.constant dense<0> : tensor<3xi64>
-//   CHECK-DAG: %[[W:.*]] = arith.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-DAG: %[[P:.*]] = memref.alloca() : memref<3xi8>
+//   CHECK-DAG: %[[Q:.*]] = memref.alloca() : memref<3xi64>
+//   CHECK-DAG: %[[R:.*]] = memref.alloca() : memref<3xi64>
+//   CHECK-DAG: %[[X:.*]] = memref.cast %[[P]] : memref<3xi8> to memref<?xi8>
+//   CHECK-DAG: %[[Y:.*]] = memref.cast %[[Q]] : memref<3xi64> to memref<?xi64>
+//   CHECK-DAG: %[[Z:.*]] = memref.cast %[[R]] : memref<3xi64> to memref<?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_init(
+//  CHECK-SAME: %[[I:.*]]: index,
+//  CHECK-SAME: %[[J:.*]]: index) -> !llvm.ptr<i8>
+//   CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+//   CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+//   CHECK-DAG: %[[P:.*]] = memref.alloca() : memref<2xi8>
+//   CHECK-DAG: %[[Q:.*]] = memref.alloca() : memref<2xi64>
+//   CHECK-DAG: %[[R:.*]] = memref.alloca() : memref<2xi64>
+//   CHECK-DAG: %[[X:.*]] = memref.cast %[[P]] : memref<2xi8> to memref<?xi8>
+//   CHECK-DAG: %[[Y:.*]] = memref.cast %[[Q]] : memref<2xi64> to memref<?xi64>
+//   CHECK-DAG: %[[Z:.*]] = memref.cast %[[R]] : memref<2xi64> to memref<?xi64>
+//   CHECK-DAG: %[[II:.*]] = arith.index_cast %[[I]] : index to i64
+//   CHECK-DAG: %[[JJ:.*]] = arith.index_cast %[[J]] : index to i64
+//   CHECK-DAG: memref.store %[[II]], %[[Q]][%[[C0]]] : memref<2xi64>
+//   CHECK-DAG: memref.store %[[JJ]], %[[Q]][%[[C1]]] : memref<2xi64>
+//       CHECK: %[[A:.*]] = llvm.mlir.null : !llvm.ptr<i8>
+//       CHECK: %[[T:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[A]])
+//       CHECK: return %[[T]] : !llvm.ptr<i8>
+func @sparse_init(%arg0: index, %arg1: index) -> tensor<?x?xf64, #SparseMatrix> {
+  %0 = sparse_tensor.init [%arg0, %arg1] : tensor<?x?xf64, #SparseMatrix>
+  return %0 : tensor<?x?xf64, #SparseMatrix>
+}
+
 // CHECK-LABEL: func @sparse_release(
 //  CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>)
 //       CHECK: call @delSparseTensor(%[[A]]) : (!llvm.ptr<i8>) -> ()
 //       CHECK: return
 func @sparse_release(%arg0: tensor<128xf64, #SparseVector>) {
   sparse_tensor.release %arg0 : tensor<128xf64, #SparseVector>
   return
 }
 
 // CHECK-LABEL: func @sparse_nop_convert(
 //  CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) -> !llvm.ptr<i8>
 //       CHECK: return %[[A]] : !llvm.ptr<i8>
 func @sparse_nop_convert(%arg0: tensor<64xf32, #SparseVector>) -> tensor<64xf32, #SparseVector> {
   %0 = sparse_tensor.convert %arg0 : tensor<64xf32, #SparseVector> to tensor<64xf32, #SparseVector>
   return %0 : tensor<64xf32, #SparseVector>
 }
 
 // CHECK-LABEL: func @sparse_convert_1d(
 //  CHECK-SAME: %[[A:.*]]: tensor<?xi32>) -> !llvm.ptr<i8>
 //   CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
 //   CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
-//   CHECK-DAG: %[[D0:.*]] = arith.constant dense<0> : tensor<1xi64>
-//   CHECK-DAG: %[[D1:.*]] = arith.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-DAG: %[[P:.*]] = memref.alloca() : memref<1xi8>
+//   CHECK-DAG: %[[Q:.*]] = memref.alloca() : memref<1xi64>
+//   CHECK-DAG: %[[R:.*]] = memref.alloca() : memref<1xi64>
+//   CHECK-DAG: %[[X:.*]] = memref.cast %[[P]] : memref<1xi8> to memref<?xi8>
+//   CHECK-DAG: %[[Y:.*]] = memref.cast %[[Q]] : memref<1xi64> to memref<?xi64>
+//   CHECK-DAG: %[[Z:.*]] = memref.cast %[[R]] : memref<1xi64> to memref<?xi64>
+//       CHECK: %[[C:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.}})
 //       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:   call @addEltI32(%[[C]], %[[E]], %[[T]], %[[Z]])
 //       CHECK: }
-//       CHECK: %[[T:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Y]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[C]])
+//       CHECK: %[[T:.*]] = call @newSparseTensor(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[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_1d_ss(
 //  CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>)
 //       CHECK: %[[C:.*]] = call @newSparseTensor(%{{.}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[A]])
 //       CHECK: %[[T:.*]] = call @newSparseTensor(%{{.}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[C]])
 //       CHECK: return %[[T]] : !llvm.ptr<i8>
 func @sparse_convert_1d_ss(%arg0: tensor<?xf32, #SparseVector64>) -> tensor<?xf32, #SparseVector32> {
   %0 = sparse_tensor.convert %arg0 : tensor<?xf32, #SparseVector64> to tensor<?xf32, #SparseVector32>
   return %0 : tensor<?xf32, #SparseVector32>
 }
 
 // CHECK-LABEL: func @sparse_convert_2d(
 //  CHECK-SAME: %[[A:.*]]: tensor<2x4xf64>) -> !llvm.ptr<i8>
 //   CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
 //   CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
-//   CHECK-DAG: %[[U:.*]] = arith.constant dense<[0, 1]> : tensor<2xi8>
-//   CHECK-DAG: %[[V:.*]] = arith.constant dense<[2, 4]> : tensor<2xi64>
-//   CHECK-DAG: %[[W:.*]] = arith.constant dense<[0, 1]> : tensor<2xi64>
-//   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-DAG: %[[P:.*]] = memref.alloca() : memref<2xi8>
+//   CHECK-DAG: %[[Q:.*]] = memref.alloca() : memref<2xi64>
+//   CHECK-DAG: %[[R:.*]] = memref.alloca() : memref<2xi64>
+//   CHECK-DAG: %[[X:.*]] = memref.cast %[[P]] : memref<2xi8> to memref<?xi8>
+//   CHECK-DAG: %[[Y:.*]] = memref.cast %[[Q]] : memref<2xi64> to memref<?xi64>
+//   CHECK-DAG: %[[Z:.*]] = memref.cast %[[R]] : memref<2xi64> to memref<?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: %[[T:.*]] = call @newSparseTensor(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[C]])
 //       CHECK: return %[[T]] : !llvm.ptr<i8>
 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>
 }
 
-#CSR = #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ] }>
-
-// CHECK-LABEL:   func @entry() -> !llvm.ptr<i8> {
-// CHECK:           %[[C1:.*]] = arith.constant 1 : i32
-// CHECK:           %[[Offset:.*]] = arith.constant dense<[0, 1]> : tensor<2xi64>
-// CHECK:           %[[Dims:.*]] = arith.constant dense<[8, 7]> : tensor<2xi64>
-// CHECK:           %[[Base:.*]] = arith.constant dense<[0, 1]> : tensor<2xi8>
-// CHECK:           %[[I2:.*]] = arith.constant 2 : index
-// CHECK:           %[[SparseV:.*]] = arith.constant dense<[1.000000e+00, 5.000000e+00]> : tensor<2xf32>
-// CHECK:           %[[SparseI:.*]] = arith.constant dense<{{\[\[}}0, 0], [1, 6]]> : tensor<2x2xi64>
-// CHECK:           %[[I1:.*]] = arith.constant 1 : index
-// CHECK:           %[[I0:.*]] = arith.constant 0 : index
-// CHECK:           %[[C2:.*]] = arith.constant 2 : i32
-// CHECK:           %[[BaseD:.*]] = tensor.cast %[[Base]] : tensor<2xi8> to tensor<?xi8>
-// CHECK:           %[[DimsD:.*]] = tensor.cast %[[Dims]] : tensor<2xi64> to tensor<?xi64>
-// CHECK:           %[[OffsetD:.*]] = tensor.cast %[[Offset]] : tensor<2xi64> to tensor<?xi64>
-// CHECK:           %[[TCOO:.*]] = call @newSparseTensor(%[[BaseD]], %[[DimsD]], %[[OffsetD]], %{{.*}}, %{{.*}}, %{{.*}}, %[[C2]], %{{.}})
-// CHECK:           %[[Index:.*]] = memref.alloca() : memref<2xindex>
-// CHECK:           %[[IndexD:.*]] = memref.cast %[[Index]] : memref<2xindex> to memref<?xindex>
-// CHECK:           scf.for %[[IV:.*]] = %[[I0]] to %[[I2]] step %[[I1]] {
-// CHECK:             %[[VAL0:.*]] = tensor.extract %[[SparseI]]{{\[}}%[[IV]], %[[I0]]] : tensor<2x2xi64>
-// CHECK:             %[[VAL1:.*]] = arith.index_cast %[[VAL0]] : i64 to index
-// CHECK:             memref.store %[[VAL1]], %[[Index]]{{\[}}%[[I0]]] : memref<2xindex>
-// CHECK:             %[[VAL2:.*]] = tensor.extract %[[SparseI]]{{\[}}%[[IV]], %[[I1]]] : tensor<2x2xi64>
-// CHECK:             %[[VAL3:.*]] = arith.index_cast %[[VAL2]] : i64 to index
-// CHECK:             memref.store %[[VAL3]], %[[Index]]{{\[}}%[[I1]]] : memref<2xindex>
-// CHECK:             %[[VAL4:.*]] = tensor.extract %[[SparseV]]{{\[}}%[[IV]]] : tensor<2xf32>
-// CHECK:             call @addEltF32(%[[TCOO]], %[[VAL4]], %[[IndexD]], %[[OffsetD]])
-// CHECK:           }
-// CHECK:           %[[T:.*]] = call @newSparseTensor(%[[BaseD]], %[[DimsD]], %[[OffsetD]], %{{.*}}, %{{.*}}, %[[C1]], %{{.*}})
-// CHECK:           return %[[T]] : !llvm.ptr<i8>
-func @entry() -> tensor<8x7xf32, #CSR>{
+// CHECK-LABEL: func @sparse_constant() -> !llvm.ptr<i8> {
+//   CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
+//   CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
+//   CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index
+//   CHECK-DAG: %[[P:.*]] = memref.alloca() : memref<2xi8>
+//   CHECK-DAG: %[[Q:.*]] = memref.alloca() : memref<2xi64>
+//   CHECK-DAG: %[[R:.*]] = memref.alloca() : memref<2xi64>
+//   CHECK-DAG: %[[X:.*]] = memref.cast %[[P]] : memref<2xi8> to memref<?xi8>
+//   CHECK-DAG: %[[Y:.*]] = memref.cast %[[Q]] : memref<2xi64> to memref<?xi64>
+//   CHECK-DAG: %[[Z:.*]] = memref.cast %[[R]] : memref<2xi64> to memref<?xi64>
+//       CHECK: %[[C:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.}})
+//       CHECK: %[[M:.*]] = memref.alloca() : memref<2xindex>
+//       CHECK: %[[N:.*]] = memref.cast %[[M]] : memref<2xindex> to memref<?xindex>
+//       CHECK: scf.for %[[I:.*]] = %[[C0]] to %[[C2]] step %[[C1]] {
+//       CHECK:   memref.store %{{.*}}, %[[M]][%[[C0]]] : memref<2xindex>
+//       CHECK:   memref.store %{{.*}}, %[[M]][%[[C1]]] : memref<2xindex>
+//       CHECK:   %[[V:.*]] = tensor.extract %{{.*}}[%[[I]]] : tensor<2xf32>
+//       CHECK:   call @addEltF32(%{{.*}}, %[[V]], %[[N]], %{{.*}})
+//       CHECK: }
+//       CHECK: %[[T:.*]] = call @newSparseTensor(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[C]])
+//       CHECK: return %[[T]] : !llvm.ptr<i8>
+func @sparse_constant() -> tensor<8x7xf32, #SparseMatrix>{
   // Initialize a tensor.
   %0 = arith.constant sparse<[[0, 0], [1, 6]], [1.0, 5.0]> : tensor<8x7xf32>
   // Convert the tensor to a sparse tensor.
-  %1 = sparse_tensor.convert %0 : tensor<8x7xf32> to tensor<8x7xf32, #CSR>
-  return %1 : tensor<8x7xf32, #CSR>
+  %1 = sparse_tensor.convert %0 : tensor<8x7xf32> to tensor<8x7xf32, #SparseMatrix>
+  return %1 : tensor<8x7xf32, #SparseMatrix>
 }
 
 // CHECK-LABEL: func @sparse_convert_3d(
 //  CHECK-SAME: %[[A:.*]]: tensor<?x?x?xf64>) -> !llvm.ptr<i8>
 //   CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index
 //   CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index
 //   CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index
-//   CHECK-DAG: %[[U:.*]] = arith.constant dense<[0, 1, 1]> : tensor<3xi8>
-//   CHECK-DAG: %[[V:.*]] = arith.constant dense<0> : tensor<3xi64>
-//   CHECK-DAG: %[[W:.*]] = arith.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-DAG: %[[P:.*]] = memref.alloca() : memref<3xi8>
+//   CHECK-DAG: %[[Q:.*]] = memref.alloca() : memref<3xi64>
+//   CHECK-DAG: %[[R:.*]] = memref.alloca() : memref<3xi64>
+//   CHECK-DAG: %[[X:.*]] = memref.cast %[[P]] : memref<3xi8> to memref<?xi8>
+//   CHECK-DAG: %[[Y:.*]] = memref.cast %[[Q]] : memref<3xi64> to memref<?xi64>
+//   CHECK-DAG: %[[Z:.*]] = memref.cast %[[R]] : memref<3xi64> to memref<?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: %[[N:.*]] = 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:       call @addEltF64(%[[C]], %[[E]], %[[N]], %[[Z]])
 //       CHECK:     }
 //       CHECK:   }
 //       CHECK: }
-//       CHECK: %[[T:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[C]])
+//       CHECK: %[[T:.*]] = call @newSparseTensor(%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[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>)