[mlir][sparse] add full dimension ordering support

[aartbik]

Summary

This revision completes the "dimension ordering" feature
of sparse tensor types that enables the programmer to
define a preferred order on dimension access (other than
the default left-to-right order). This enables e.g. selection
of column-major over row-major storage for sparse matrices,
but generalized to any rank, as in:

dimOrdering = affine_map<(i,j,k,l,m,n,o,p) -> (p,o,j,k,i,l,m,n)>

[bixia]

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

61

Is it true that values.size() == sz? Maybe we can remove sz as a parameter here.

146–147

Would it be better if we push this into the else branch of if(enc), and add perm.resize(sz) to the if branch?

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

22

not need?

[aartbik]

Thankx Bixia!

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

61

Fair point. Removed.

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

22

Yes, I added this so you could quickly test another value, but we indeed have the requirement of avoiding commented out code. Removed.

[aartbik]

addressed comments

mlir/include/mlir/ExecutionEngine/CRunnerUtils.h

 // Small runtime support library for timing execution and printing GFLOPS
 //===----------------------------------------------------------------------===//
 extern "C" MLIR_CRUNNERUTILS_EXPORT void print_flops(double flops);
 extern "C" MLIR_CRUNNERUTILS_EXPORT double rtclock();
 
 //===----------------------------------------------------------------------===//
 // Small runtime support library for sparse tensors.
 //===----------------------------------------------------------------------===//
-extern "C" MLIR_CRUNNERUTILS_EXPORT void *openTensorC(char *filename,
-                                                      uint64_t *idata);
+extern "C" MLIR_CRUNNERUTILS_EXPORT void *
+openTensorC(char *filename, uint64_t *idata, uint64_t *perm);
 extern "C" MLIR_CRUNNERUTILS_EXPORT void
 readTensorItemC(void *tensor, uint64_t *idata, double *ddata);
 extern "C" MLIR_CRUNNERUTILS_EXPORT void closeTensor(void *tensor);
 extern "C" MLIR_CRUNNERUTILS_EXPORT char *getTensorFilename(uint64_t id);
 
 #endif // EXECUTIONENGINE_CRUNNERUTILS_H_

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

case SparseTensorEncodingAttr::DimLevelType::Dense:
     return 0;
   case SparseTensorEncodingAttr::DimLevelType::Compressed:
     return 1;
   case SparseTensorEncodingAttr::DimLevelType::Singleton:
     return 2;
   }
 }
 
+/// 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 accross 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<ConstantOp>(loc, DenseElementsAttr::get(tt1, values));
+  return rewriter.create<tensor::CastOp>(loc, tt2, elts);
+}
+
 /// Returns function reference (first hit also inserts into module).
 static FlatSymbolRefAttr getFunc(Operation *op, StringRef name, Type result,
                                  ValueRange operands) {
   MLIRContext *context = op->getContext();
   auto module = op->getParentOfType<ModuleOp>();
   auto func = module.lookupSymbol<FuncOp>(name);
   if (!func) {
     OpBuilder moduleBuilder(module.getBodyRegion());
matchAndRewrite(NewOp op, ArrayRef<Value> operands,
     MLIRContext *context = op->getContext();
     SmallVector<Value, 5> params;
     // Sparse encoding.
     auto enc = getSparseTensorEncoding(resType);
     if (!enc)
       return failure();
     // User pointer.
     params.push_back(operands[0]);
-    // Sparsity annotations in tensor constant form. Note that we cast
-    // the static shape into a dynamic shape to ensure that the method
-    // signature remains uniform accross different tensor dimensions.
+    // 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])));
-    Type etp = rewriter.getIntegerType(8);
-    RankedTensorType tt1 = RankedTensorType::get({sz}, etp);
-    RankedTensorType tt2 =
-        RankedTensorType::get({ShapedType::kDynamicSize}, etp);
-    auto elts =
-        rewriter.create<ConstantOp>(loc, DenseElementsAttr::get(tt1, attrs));
-    params.push_back(rewriter.create<tensor::CastOp>(loc, tt2, elts));
-    // Seconary and primary types encoding.
+    params.push_back(getTensor(rewriter, 8, loc, attrs));
+    // 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);
+      for (unsigned i = 0; i < sz; i++)
+        perm[p.getDimPosition(i)] = APInt(64, i);
+    } else {
+      for (unsigned i = 0; i < sz; i++)
+        perm[i] = APInt(64, i);
+    }
+    params.push_back(getTensor(rewriter, 64, loc, perm));
+    // Secondary and primary types encoding.
     unsigned secPtr = getOverheadTypeEncoding(enc.getPointerBitWidth());
     unsigned secInd = getOverheadTypeEncoding(enc.getIndexBitWidth());
     unsigned primary;
     if (eltType.isF64())
       primary = 1;
     else if (eltType.isF32())
       primary = 2;
     else if (eltType.isInteger(32))

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

struct CodeGen {
   Value redVal;
   // Current vector length and mask.
   unsigned curVecLength;
   Value curVecMask;
 };
 
 } // namespace
 
+// Helper method to apply dimension ordering permutation.
+static unsigned perm(SparseTensorEncodingAttr &enc, unsigned d) {
+  if (enc) {
+    auto order = enc.getDimOrdering();
+    if (order) {
+      assert(order.isPermutation());
+      return order.getDimPosition(d);
+    }
+  }
+  return d;
+}
+
 // Helper method to translate dim level type to internal representation.
 static Dim toDim(SparseTensorEncodingAttr &enc, unsigned d) {
   if (enc) {
     SparseTensorEncodingAttr::DimLevelType tp = enc.getDimLevelType()[d];
     if (tp == SparseTensorEncodingAttr::DimLevelType::Compressed)
       return Dim::kSparse;
     if (tp == SparseTensorEncodingAttr::DimLevelType::Singleton)
       return Dim::kSingle;
   }
   return Dim::kDense;
 }
 
 /// Helper method to inspect sparse encodings in the tensor types.
 /// Fills the per-dimension sparsity information for all tensors.
 static bool findSparseAnnotations(Merger &merger, linalg::GenericOp op) {
   bool annotated = false;
   unsigned numTensors = op.getNumShapedOperands();
   unsigned lhs = numTensors - 1;
   for (unsigned t = 0; t < numTensors; t++) {
     auto map = op.getIndexingMap(t);
-    unsigned rank = op.getShapedType(t).getRank();
+    if (!map.isProjectedPermutation())
+      return false;
     auto enc = getSparseTensorEncoding(op.getShapedType(t));
     if (enc) {
       annotated = true;
-      if (enc.getDimOrdering() && !enc.getDimOrdering().isIdentity())
-        return false; // TODO: handle permutations
       if (t == lhs)
         return false; // TODO: handle sparse outputs
     }
-    for (unsigned d = 0; d < rank; d++) {
-      unsigned idx = map.getDimPosition(d);
+    assert(map.getNumResults() == op.getShapedType(t).getRank());
+    for (unsigned d = 0, rank = map.getNumResults(); d < rank; d++) {
+      unsigned idx = map.getDimPosition(perm(enc, d));
       merger.setDim(t, idx, toDim(enc, d));
     }
   }
   return annotated;
 }
 
 /// A DFS helper to compute a topological sort. Note that recursion is
 /// bounded by the number of implicit loops, which is always small.
static bool computeIterationGraph(Merger &merger, linalg::GenericOp op,
   // for the implicit loop indices i_0 .. i_n-1.
   unsigned n = op.getNumLoops();
   std::vector<std::vector<bool>> adjM(n, std::vector<bool>(n, false));
 
   // Iterate over the indexing maps of every tensor in the tensor expression.
   unsigned numTensors = op.getNumShapedOperands();
   for (unsigned t = 0; t < numTensors; t++) {
     auto map = op.getIndexingMap(t);
+    auto enc = getSparseTensorEncoding(op.getShapedType(t));
     assert(map.getNumDims() == n);
     // Skip dense tensor constraints when sparse only is requested.
-    if (sparseOnly && !getSparseTensorEncoding(op.getShapedType(t)))
+    if (sparseOnly && !enc)
       continue;
-    // At the moment, we take the index variables in the tensor access
-    // expression in the order in which they appear (conceptually a
-    // "row-major" layout of every tensor). So, a tensor access A_ijk
-    // forces the ordering i < j < k on the loop indices.
-    // TODO: support affine map to define alternative dimension orders.
-    for (unsigned d = 1, e = map.getNumResults(); d < e; d++) {
-      unsigned f = map.getDimPosition(d - 1);
-      unsigned t = map.getDimPosition(d);
+    // Each tensor expression and optional dimension ordering (row-major
+    // by default) puts an ordering constraint on the loop indices. For
+    // example, the tensor expresion A_ijk forces the ordering i < j < k
+    // on the loop indices if no explicit dimension ordering is given.
+    for (unsigned d = 1, rank = map.getNumResults(); d < rank; d++) {
+      unsigned f = map.getDimPosition(perm(enc, d - 1));
+      unsigned t = map.getDimPosition(perm(enc, d));
       adjM[f][t] = true;
     }
   }
 
   // Topologically sort the iteration graph to determine loop order.
   // Report failure for a cyclic iteration graph.
   topSort.clear();
   topSort.reserve(n);
   std::vector<unsigned> visit(n, 0);
   for (unsigned i = 0; i < n; i++)
     if (visit[i] == 0)
       if (!topSortDFS(i, visit, topSort, adjM))
         return false; // cycle!
   std::reverse(std::begin(topSort), std::end(topSort));
   return true;
 }
 
 /// Traverses the SSA tree (possibly a DAG) to build a tensor expression.
 /// This simplifies constructing (sub)expressions during iteration lattice
 /// building (compared to using the SSA representation everywhere).
 static Optional<unsigned> buildTensorExp(Merger &merger, linalg::GenericOp op,
                                          Value val) {
   if (auto arg = val.dyn_cast<BlockArgument>()) {
     unsigned argN = arg.getArgNumber();
-    if (arg.getOwner()->getParentOp() == op) {
-      // Any parameter of the generic op is considered a tensor,
-      // indexed by the implicit loop bounds.
-      auto map = op.getIndexingMap(argN);
-      if (map.isProjectedPermutation())
-        return merger.addExp(Kind::kTensor, argN);
-      // Cannot handle (yet).
-      return None;
-    }
+    if (arg.getOwner()->getParentOp() == op)
     // Any parameter of a higher op is invariant.
     return merger.addExp(Kind::kInvariant, val);
   }
   Operation *def = val.getDefiningOp();
   if (def->getBlock() != &op.region().front()) {
     // Something defined outside is invariant.
     return merger.addExp(Kind::kInvariant, val);
   } else if (def->getNumOperands() == 2) {
static void genBuffers(Merger &merger, CodeGen &codegen,
   for (unsigned t = 0; t < numTensors; t++) {
     Value tensor = t < numInputs ? op.getInput(t) : op.getOutput(0);
     auto tensorType = op.getShapedType(t);
     auto shape = tensorType.getShape();
     auto map = op.getIndexingMap(t);
     auto enc = getSparseTensorEncoding(tensorType);
     // Scan all dimensions of current tensor.
     args.clear();
-    for (unsigned d = 0, rank = shape.size(); d < rank; d++) {
-      unsigned i = map.getDimPosition(d);
+    for (unsigned d = 0, rank = map.getNumResults(); d < rank; d++) {
+      unsigned idx = map.getDimPosition(perm(enc, d));
       // Handle sparse storage schemes.
-      if (merger.isDim(t, i, Dim::kSparse)) {
+      if (merger.isDim(t, idx, Dim::kSparse)) {
         auto dynShape = {ShapedType::kDynamicSize};
         auto ptrTp = MemRefType::get(
             dynShape, genIntType(rewriter, enc.getPointerBitWidth()));
         auto indTp = MemRefType::get(
             dynShape, genIntType(rewriter, enc.getIndexBitWidth()));
         Value dim = rewriter.create<ConstantIndexOp>(loc, d);
         // Generate sparse primitives to obtains pointer and indices.
-        codegen.pointers[t][i] =
+        codegen.pointers[t][idx] =
             rewriter.create<ToPointersOp>(loc, ptrTp, tensor, dim);
-        codegen.indices[t][i] =
+        codegen.indices[t][idx] =
             rewriter.create<ToIndicesOp>(loc, indTp, tensor, dim);
       }
       // Find lower and upper bound in current dimension.
       Value up;
       if (shape[d] == MemRefType::kDynamicSize) {
         up = rewriter.create<memref::DimOp>(loc, tensor, d);
         args.push_back(up);
       } else {
         up = rewriter.create<ConstantIndexOp>(loc, shape[d]);
       }
-      codegen.sizes[i] = codegen.highs[t][i] = up;
+      codegen.sizes[idx] = codegen.highs[t][idx] = up;
     }
     // Perform the required bufferization. All dense inputs materialize
     // from the input tensor. The dense output tensor needs special
     // handling. Sparse inputs use a sparse primitive to obtain the values.
     if (!enc) {
       auto denseTp = MemRefType::get(shape, tensorType.getElementType());
       if (t < numInputs)
         codegen.buffers[t] =
if (codegen.curVecLength > 1 && !val.getType().isa<VectorType>())
       return genVectorInvariantValue(codegen, rewriter, val);
     return val;
   }
   // Actual load.
   SmallVector<Value, 4> args;
   unsigned tensor = merger.exp(exp).e0;
   auto map = op.getIndexingMap(tensor);
   auto enc = getSparseTensorEncoding(op.getShapedType(tensor));
-  for (unsigned i = 0, m = map.getNumResults(); i < m; ++i) {
-    unsigned idx = map.getDimPosition(i);
+  for (unsigned d = 0, rank = map.getNumResults(); d < rank; d++) {
+    unsigned idx = map.getDimPosition(perm(enc, d));
     args.push_back(codegen.loops[idx]); // universal dense index
     if (enc) {
       args.clear();
       args.push_back(codegen.pidxs[tensor][idx]); // position index
     }
   }
   Location loc = op.getLoc();
   Value ptr = codegen.buffers[tensor];
if (codegen.curVecLength > 1)
       rhs = rewriter.create<SelectOp>(loc, codegen.curVecMask, rhs,
                                       codegen.redVal);
     codegen.redVal = rhs;
     return;
   }
   // Actual store.
   SmallVector<Value, 4> args;
   auto map = op.getIndexingMap(tensor);
-  for (unsigned i = 0, m = map.getNumResults(); i < m; ++i) {
-    unsigned idx = map.getDimPosition(i);
+  assert(!getSparseTensorEncoding(op.getShapedType(tensor)));
+  for (unsigned d = 0, rank = map.getNumResults(); d < rank; d++) {
+    unsigned idx = map.getDimPosition(d);
     args.push_back(codegen.loops[idx]); // universal dense index
   }
   Value ptr = codegen.buffers[tensor];
   if (codegen.curVecLength > 1)
     genVectorStore(codegen, rewriter, rhs, ptr, args);
   else
     rewriter.create<memref::StoreOp>(loc, rhs, ptr, args);
 }

 static void genInvariants(Merger &merger, CodeGen &codegen,
                           PatternRewriter &rewriter, linalg::GenericOp op,
                           unsigned exp, unsigned ldx, bool hoist) {
   if (merger.exp(exp).kind == Kind::kTensor) {
     // Inspect tensor indices.
     bool atLevel = ldx == -1u;
     unsigned tensor = merger.exp(exp).e0;
     auto map = op.getIndexingMap(tensor);
-    for (unsigned i = 0, m = map.getNumResults(); i < m; ++i) {
-      unsigned idx = map.getDimPosition(i);
+    auto enc = getSparseTensorEncoding(op.getShapedType(tensor));
+    for (unsigned d = 0, rank = map.getNumResults(); d < rank; d++) {
+      unsigned idx = map.getDimPosition(perm(enc, d));
       if (!codegen.loops[idx])
         return; // still in play
       else if (idx == ldx)
         atLevel = true;
     }
     // All exhausted at this level (atLevel denotes exactly at this level).
     unsigned lhs = op.getNumShapedOperands() - 1;
     if (lhs == tensor) {

 /// For now, we reject vectorization of such cases.
 /// TODO: implement strided load/stores on dense arrays
 static bool denseUnitStrides(Merger &merger, linalg::GenericOp op,
                              unsigned idx) {
   unsigned numTensors = op.getNumShapedOperands();
   for (unsigned t = 0; t < numTensors; t++) {
     if (!getSparseTensorEncoding(op.getShapedType(t))) {
       auto map = op.getIndexingMap(t);
-      unsigned r = map.getNumResults();
-      for (unsigned i = 0; i < r; i++) {
-        if (map.getDimPosition(i) == idx && i != r - 1)
+      for (unsigned d = 0, rank = map.getNumResults(); d < rank; d++) {
+        if (map.getDimPosition(d) == idx && d != rank - 1)
           return false;
       }
     }
   }
   return true;
 }
 
 /// Generates a for-loop on a single index.

mlir/lib/ExecutionEngine/SparseUtils.cpp

private:
   std::vector<uint64_t> sizes; // per-rank dimension sizes
   std::vector<std::vector<P>> pointers;
   std::vector<std::vector<I>> indices;
   std::vector<V> values;
 };
 
 /// Templated reader.
 template <typename P, typename I, typename V>
-void *newSparseTensor(char *filename, uint8_t *sparsity, uint64_t size) {
+void *newSparseTensor(char *filename, uint8_t *sparsity, uint64_t *perm,
+                      uint64_t size) {
   uint64_t idata[64];
-  SparseTensor *t = static_cast<SparseTensor *>(openTensorC(filename, idata));
+  SparseTensor *t =
+      static_cast<SparseTensor *>(openTensorC(filename, idata, perm));
   assert(size == t->getRank()); // sparsity array must match rank
   SparseTensorStorageBase *tensor =
       new SparseTensorStorage<P, I, V>(t, sparsity);
   delete t;
   return tensor;
 }
 
 /// Helper to convert string to lower case.

 
 extern "C" {
 
 /// Reads in a sparse tensor with the given filename. The call yields a
 /// pointer to an opaque memory-resident sparse tensor object that is only
 /// understood by other methods in the sparse runtime support library. An
 /// array parameter is used to pass the rank, the number of nonzero elements,
 /// and the dimension sizes (one per rank).
-void *openTensorC(char *filename, uint64_t *idata) {
+void *openTensorC(char *filename, uint64_t *idata, uint64_t *perm) {
   // Open the file.
   FILE *file = fopen(filename, "r");
   if (!file) {
     fprintf(stderr, "Cannot find %s\n", filename);
     exit(1);
   }
   // Perform some file format dependent set up.
   if (strstr(filename, ".mtx")) {
     readMMEHeader(file, filename, idata);
   } else if (strstr(filename, ".tns")) {
     readExtFROSTTHeader(file, filename, idata);
   } 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.
   uint64_t rank = idata[0];
   uint64_t nnz = idata[1];
   std::vector<uint64_t> indices(rank);
   for (uint64_t r = 0; r < rank; r++)
+    if (perm)
+      indices[perm[r]] = idata[2 + r];
+    else
       indices[r] = idata[2 + r];
   SparseTensor *tensor = new SparseTensor(indices, nnz);
   // Read all nonzero elements.
   for (uint64_t k = 0; k < nnz; k++) {
+    uint64_t idx = -1;
     for (uint64_t r = 0; r < rank; r++) {
-      if (fscanf(file, "%" PRIu64, &indices[r]) != 1) {
+      if (fscanf(file, "%" PRIu64, &idx) != 1) {
         fprintf(stderr, "Cannot find next index in %s\n", filename);
         exit(1);
       }
-      indices[r]--; // 0-based index
+      // Add 0-based index.
+      if (perm)
+        indices[perm[r]] = idx - 1;
+      else
+        indices[r] = idx - 1;
     }
     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 sorted tensor.
   fclose(file);
   tensor->sort(); // sort lexicographically
   return tensor;
 }
 
 /// "MLIRized" version.
 void *openTensor(char *filename, uint64_t *ibase, uint64_t *idata,
                  uint64_t ioff, uint64_t isize, uint64_t istride) {
   assert(istride == 1);
-  return openTensorC(filename, idata + ioff);
+  return openTensorC(filename, idata + ioff, nullptr);
 }
 
 /// Yields the next element from the given opaque sparse tensor object.
 void readTensorItemC(void *tensor, uint64_t *idata, double *ddata) {
   const Element &e = static_cast<SparseTensor *>(tensor)->next();
   for (uint64_t r = 0, rank = e.indices.size(); r < rank; r++)
     idata[r] = e.indices[r];
   ddata[0] = e.value;
struct NAME {                                                                \
     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))                            \
-  return newSparseTensor<P, I, V>(filename, sparsity, asize)
+  return newSparseTensor<P, I, V>(filename, sparsity, perm, asize)
 
 #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}};                        \
   }
 
enum PrimaryTypeEnum : uint64_t {
   kF32 = 2,
   kI32 = 3,
   kI16 = 4,
   kI8 = 5
 };
 
 void *newSparseTensor(char *filename, uint8_t *abase, uint8_t *adata,
                       uint64_t aoff, uint64_t asize, uint64_t astride,
-                      uint64_t ptrTp, uint64_t indTp, uint64_t valTp) {
-  assert(astride == 1);
+                      uint64_t *pbase, uint64_t *pdata, uint64_t poff,
+                      uint64_t psize, uint64_t pstride, uint64_t ptrTp,
+                      uint64_t indTp, uint64_t valTp) {
+  assert(astride == 1 && pstride == 1);
   uint8_t *sparsity = adata + aoff;
+  uint64_t *perm = pdata + poff;
 
   // The most common cases: 64-bit or 32-bit overhead, double/float values.
   CASE(kU64, kU64, kF64, uint64_t, uint64_t, double);
   CASE(kU64, kU64, kF32, uint64_t, uint64_t, float);
   CASE(kU64, kU32, kF64, uint64_t, uint32_t, double);
   CASE(kU64, kU32, kF32, uint64_t, uint32_t, float);
   CASE(kU32, kU64, kF64, uint32_t, uint64_t, double);
   CASE(kU32, kU64, kF32, uint32_t, uint64_t, float);

mlir/test/Dialect/SparseTensor/conversion.mlir

#SparseVector32 = #sparse_tensor.encoding<{
   pointerBitWidth = 32,
   indexBitWidth = 32
 }>
 
 #SparseMatrix = #sparse_tensor.encoding<{
   dimLevelType = ["dense", "compressed"]
 }>
 
+#SparseTensor = #sparse_tensor.encoding<{
+  dimLevelType = ["dense", "compressed", "compressed"],
+  dimOrdering = affine_map<(i,j,k) -> (k,i,j)>
+}>
+
 // CHECK-LABEL: func @sparse_dim(
 //  CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>)
 //       CHECK: %[[C:.*]] = constant 0 : index
 //       CHECK: %[[D:.*]] = call @sparseDimSize(%[[A]], %[[C]])
 //       CHECK: return %[[D]] : index
 func @sparse_dim(%arg0: tensor<?xf64, #SparseVector>) -> index {
   %c = constant 0 : index
   %0 = memref.dim %arg0, %c : tensor<?xf64, #SparseVector>
   return %0 : index
 }
 
 // CHECK-LABEL: func @sparse_new1d(
 //  CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>) -> !llvm.ptr<i8>
 //       CHECK: %[[D:.*]] = constant dense<1> : tensor<1xi8>
 //       CHECK: %[[C:.*]] = tensor.cast %[[D]] : tensor<1xi8> to tensor<?xi8>
-//       CHECK: %[[T:.*]] = call @newSparseTensor(%[[A]], %[[C]], %{{.*}}, %{{.*}}, %{{.*}}) : (!llvm.ptr<i8>, tensor<?xi8>, i64, i64, i64) -> !llvm.ptr<i8>
+//       CHECK: %[[P:.*]] = constant dense<0> : tensor<1xi64>
+//       CHECK: %[[Q:.*]] = tensor.cast %[[P]] : tensor<1xi64> to tensor<?xi64>
+//       CHECK: %[[T:.*]] = call @newSparseTensor(%[[A]], %[[C]], %[[Q]], %{{.*}}, %{{.*}}, %{{.*}}) : (!llvm.ptr<i8>, tensor<?xi8>, tensor<?xi64>, i64, i64, i64) -> !llvm.ptr<i8>
 //       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: %[[D:.*]] = constant dense<[0, 1]> : tensor<2xi8>
 //       CHECK: %[[C:.*]] = tensor.cast %[[D]] : tensor<2xi8> to tensor<?xi8>
-//       CHECK: %[[T:.*]] = call @newSparseTensor(%[[A]], %[[C]], %{{.*}}, %{{.*}}, %{{.*}}) : (!llvm.ptr<i8>, tensor<?xi8>, i64, i64, i64) -> !llvm.ptr<i8>
+//       CHECK: %[[P:.*]] = constant dense<[0, 1]> : tensor<2xi64>
+//       CHECK: %[[Q:.*]] = tensor.cast %[[P]] : tensor<2xi64> to tensor<?xi64>
+//       CHECK: %[[T:.*]] = call @newSparseTensor(%[[A]], %[[C]], %[[Q]], %{{.*}}, %{{.*}}, %{{.*}}) : (!llvm.ptr<i8>, tensor<?xi8>, tensor<?xi64>, i64, i64, i64) -> !llvm.ptr<i8>
 //       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: %[[D:.*]] = constant dense<[0, 1, 1]> : tensor<3xi8>
+//       CHECK: %[[C:.*]] = tensor.cast %[[D]] : tensor<3xi8> to tensor<?xi8>
+//       CHECK: %[[P:.*]] = constant dense<[1, 2, 0]> : tensor<3xi64>
+//       CHECK: %[[Q:.*]] = tensor.cast %[[P]] : tensor<3xi64> to tensor<?xi64>
+//       CHECK: %[[T:.*]] = call @newSparseTensor(%[[A]], %[[C]], %[[Q]], %{{.*}}, %{{.*}}, %{{.*}}) : (!llvm.ptr<i8>, tensor<?xi8>, tensor<?xi64>, i64, i64, i64) -> !llvm.ptr<i8>
+//       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_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>

mlir/test/Dialect/SparseTensor/sparse_lower.mlir

indexing_maps = [
     affine_map<(i,j) -> (j)>,    // b
     affine_map<(i,j) -> (i)>     // x (out)
   ],
   iterator_types = ["parallel","reduction"],
   doc = "x(i) += A(i,j) * b(j)"
 }
 
 // CHECK-HIR-LABEL:   func @matvec(
-// CHECK-HIR-SAME:                 %[[VAL_0:.*]]: tensor<64x64xf64, #sparse_tensor.encoding<{{.*}}>>,
+// CHECK-HIR-SAME:                 %[[VAL_0:.*]]: tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ], pointerBitWidth = 0, indexBitWidth = 0 }>>,
 // CHECK-HIR-SAME:                 %[[VAL_1:.*]]: tensor<64xf64>,
-// CHECK-HIR-SAME:                 %[[VAL_2:.*]]: tensor<64xf64>) -> tensor<64xf64> {
-// CHECK-HIR:           %[[VAL_3:.*]] = constant 64 : index
+// CHECK-HIR-SAME:                 %[[VAL_2:.*]]: tensor<32xf64>) -> tensor<32xf64> {
+// CHECK-HIR:           %[[VAL_3:.*]] = constant 32 : index
 // CHECK-HIR:           %[[VAL_4:.*]] = constant 0 : index
 // CHECK-HIR:           %[[VAL_5:.*]] = constant 1 : index
-// CHECK-HIR:           %[[VAL_7:.*]] = sparse_tensor.pointers %[[VAL_0]], %[[VAL_5]] : tensor<64x64xf64, #sparse_tensor.encoding<{{.*}}>> to memref<?xindex>
-// CHECK-HIR:           %[[VAL_8:.*]] = sparse_tensor.indices %[[VAL_0]], %[[VAL_5]] : tensor<64x64xf64, #sparse_tensor.encoding<{{.*}}>> to memref<?xindex>
-// CHECK-HIR:           %[[VAL_9:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<64x64xf64, #sparse_tensor.encoding<{{.*}}>> to memref<?xf64>
-// CHECK-HIR:           %[[VAL_10:.*]] = memref.buffer_cast %[[VAL_1]] : memref<64xf64>
-// CHECK-HIR:           %[[VAL_11:.*]] = memref.buffer_cast %[[VAL_2]] : memref<64xf64>
-// CHECK-HIR:           %[[VAL_12:.*]] = memref.alloc() : memref<64xf64>
-// CHECK-HIR:           linalg.copy(%[[VAL_11]], %[[VAL_12]]) : memref<64xf64>, memref<64xf64>
-// CHECK-HIR:           scf.for %[[VAL_13:.*]] = %[[VAL_4]] to %[[VAL_3]] step %[[VAL_5]] {
-// CHECK-HIR:             %[[VAL_14:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_13]]] : memref<?xindex>
-// CHECK-HIR:             %[[VAL_15:.*]] = addi %[[VAL_13]], %[[VAL_5]] : index
-// CHECK-HIR:             %[[VAL_16:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_15]]] : memref<?xindex>
-// CHECK-HIR:             %[[VAL_17:.*]] = memref.load %[[VAL_12]]{{\[}}%[[VAL_13]]] : memref<64xf64>
-// CHECK-HIR:             %[[VAL_18:.*]] = scf.for %[[VAL_19:.*]] = %[[VAL_14]] to %[[VAL_16]] step %[[VAL_5]] iter_args(%[[VAL_20:.*]] = %[[VAL_17]]) -> (f64) {
-// CHECK-HIR:               %[[VAL_21:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_19]]] : memref<?xindex>
-// CHECK-HIR:               %[[VAL_22:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_19]]] : memref<?xf64>
-// CHECK-HIR:               %[[VAL_23:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_21]]] : memref<64xf64>
-// CHECK-HIR:               %[[VAL_24:.*]] = mulf %[[VAL_22]], %[[VAL_23]] : f64
-// CHECK-HIR:               %[[VAL_25:.*]] = addf %[[VAL_20]], %[[VAL_24]] : f64
-// CHECK-HIR:               scf.yield %[[VAL_25]] : f64
+// CHECK-HIR:           %[[VAL_6:.*]] = sparse_tensor.pointers %[[VAL_0]], %[[VAL_5]] : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ], pointerBitWidth = 0, indexBitWidth = 0 }>> to memref<?xindex>
+// CHECK-HIR:           %[[VAL_7:.*]] = sparse_tensor.indices %[[VAL_0]], %[[VAL_5]] : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ], pointerBitWidth = 0, indexBitWidth = 0 }>> to memref<?xindex>
+// CHECK-HIR:           %[[VAL_8:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ], pointerBitWidth = 0, indexBitWidth = 0 }>> to memref<?xf64>
+// CHECK-HIR:           %[[VAL_9:.*]] = memref.buffer_cast %[[VAL_1]] : memref<64xf64>
+// CHECK-HIR:           %[[VAL_10:.*]] = memref.buffer_cast %[[VAL_2]] : memref<32xf64>
+// CHECK-HIR:           %[[VAL_11:.*]] = memref.alloc() : memref<32xf64>
+// CHECK-HIR:           linalg.copy(%[[VAL_10]], %[[VAL_11]]) : memref<32xf64>, memref<32xf64>
+// CHECK-HIR:           scf.for %[[VAL_12:.*]] = %[[VAL_4]] to %[[VAL_3]] step %[[VAL_5]] {
+// CHECK-HIR:             %[[VAL_13:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_12]]] : memref<?xindex>
+// CHECK-HIR:             %[[VAL_14:.*]] = addi %[[VAL_12]], %[[VAL_5]] : index
+// CHECK-HIR:             %[[VAL_15:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_14]]] : memref<?xindex>
+// CHECK-HIR:             %[[VAL_16:.*]] = memref.load %[[VAL_11]]{{\[}}%[[VAL_12]]] : memref<32xf64>
+// CHECK-HIR:             %[[VAL_17:.*]] = scf.for %[[VAL_18:.*]] = %[[VAL_13]] to %[[VAL_15]] step %[[VAL_5]] iter_args(%[[VAL_19:.*]] = %[[VAL_16]]) -> (f64) {
+// CHECK-HIR:               %[[VAL_20:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_18]]] : memref<?xindex>
+// CHECK-HIR:               %[[VAL_21:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_18]]] : memref<?xf64>
+// CHECK-HIR:               %[[VAL_22:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_20]]] : memref<64xf64>
+// CHECK-HIR:               %[[VAL_23:.*]] = mulf %[[VAL_21]], %[[VAL_22]] : f64
+// CHECK-HIR:               %[[VAL_24:.*]] = addf %[[VAL_19]], %[[VAL_23]] : f64
+// CHECK-HIR:               scf.yield %[[VAL_24]] : f64
 // CHECK-HIR:             }
-// CHECK-HIR:             store %[[VAL_26:.*]], %[[VAL_12]]{{\[}}%[[VAL_13]]] : memref<64xf64>
+// CHECK-HIR:             memref.store %[[VAL_25:.*]], %[[VAL_11]]{{\[}}%[[VAL_12]]] : memref<32xf64>
 // CHECK-HIR:           }
-// CHECK-HIR:           %[[VAL_27:.*]] = memref.tensor_load %[[VAL_12]] : memref<64xf64>
-// CHECK-HIR:           return %[[VAL_27]] : tensor<64xf64>
+// CHECK-HIR:           %[[VAL_26:.*]] = memref.tensor_load %[[VAL_11]] : memref<32xf64>
+// CHECK-HIR:           return %[[VAL_26]] : tensor<32xf64>
 // CHECK-HIR:         }
 
 // CHECK-MIR-LABEL:   func @matvec(
 // CHECK-MIR-SAME:                 %[[VAL_0:.*]]: !llvm.ptr<i8>,
 // CHECK-MIR-SAME:                 %[[VAL_1:.*]]: tensor<64xf64>,
-// CHECK-MIR-SAME:                 %[[VAL_2:.*]]: tensor<64xf64>) -> tensor<64xf64> {
-// CHECK-MIR:           %[[VAL_3:.*]] = constant 64 : index
+// CHECK-MIR-SAME:                 %[[VAL_2:.*]]: tensor<32xf64>) -> tensor<32xf64> {
+// CHECK-MIR:           %[[VAL_3:.*]] = constant 32 : index
 // CHECK-MIR:           %[[VAL_4:.*]] = constant 0 : index
 // CHECK-MIR:           %[[VAL_5:.*]] = constant 1 : index
 // CHECK-MIR:           %[[VAL_6:.*]] = call @sparsePointers(%[[VAL_0]], %[[VAL_5]]) : (!llvm.ptr<i8>, index) -> memref<?xindex>
 // CHECK-MIR:           %[[VAL_7:.*]] = call @sparseIndices(%[[VAL_0]], %[[VAL_5]]) : (!llvm.ptr<i8>, index) -> memref<?xindex>
 // CHECK-MIR:           %[[VAL_8:.*]] = call @sparseValuesF64(%[[VAL_0]]) : (!llvm.ptr<i8>) -> memref<?xf64>
 // CHECK-MIR:           %[[VAL_9:.*]] = memref.buffer_cast %[[VAL_1]] : memref<64xf64>
-// CHECK-MIR:           %[[VAL_10:.*]] = memref.buffer_cast %[[VAL_2]] : memref<64xf64>
-// CHECK-MIR:           %[[VAL_11:.*]] = memref.alloc() : memref<64xf64>
+// CHECK-MIR:           %[[VAL_10:.*]] = memref.buffer_cast %[[VAL_2]] : memref<32xf64>
+// CHECK-MIR:           %[[VAL_11:.*]] = memref.alloc() : memref<32xf64>
 // CHECK-MIR:           scf.for %[[VAL_12:.*]] = %[[VAL_4]] to %[[VAL_3]] step %[[VAL_5]] {
-// CHECK-MIR:             %[[VAL_13:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_12]]] : memref<64xf64>
-// CHECK-MIR:             store %[[VAL_13]], %[[VAL_11]]{{\[}}%[[VAL_12]]] : memref<64xf64>
+// CHECK-MIR:             %[[VAL_13:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_12]]] : memref<32xf64>
+// CHECK-MIR:             memref.store %[[VAL_13]], %[[VAL_11]]{{\[}}%[[VAL_12]]] : memref<32xf64>
 // CHECK-MIR:           }
 // CHECK-MIR:           scf.for %[[VAL_14:.*]] = %[[VAL_4]] to %[[VAL_3]] step %[[VAL_5]] {
 // CHECK-MIR:             %[[VAL_15:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_14]]] : memref<?xindex>
 // CHECK-MIR:             %[[VAL_16:.*]] = addi %[[VAL_14]], %[[VAL_5]] : index
 // CHECK-MIR:             %[[VAL_17:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_16]]] : memref<?xindex>
-// CHECK-MIR:             %[[VAL_18:.*]] = memref.load %[[VAL_11]]{{\[}}%[[VAL_14]]] : memref<64xf64>
+// CHECK-MIR:             %[[VAL_18:.*]] = memref.load %[[VAL_11]]{{\[}}%[[VAL_14]]] : memref<32xf64>
 // CHECK-MIR:             %[[VAL_19:.*]] = scf.for %[[VAL_20:.*]] = %[[VAL_15]] to %[[VAL_17]] step %[[VAL_5]] iter_args(%[[VAL_21:.*]] = %[[VAL_18]]) -> (f64) {
 // CHECK-MIR:               %[[VAL_22:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_20]]] : memref<?xindex>
 // CHECK-MIR:               %[[VAL_23:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_20]]] : memref<?xf64>
 // CHECK-MIR:               %[[VAL_24:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_22]]] : memref<64xf64>
 // CHECK-MIR:               %[[VAL_25:.*]] = mulf %[[VAL_23]], %[[VAL_24]] : f64
 // CHECK-MIR:               %[[VAL_26:.*]] = addf %[[VAL_21]], %[[VAL_25]] : f64
 // CHECK-MIR:               scf.yield %[[VAL_26]] : f64
 // CHECK-MIR:             }
-// CHECK-MIR:             store %[[VAL_27:.*]], %[[VAL_11]]{{\[}}%[[VAL_14]]] : memref<64xf64>
+// CHECK-MIR:             memref.store %[[VAL_27:.*]], %[[VAL_11]]{{\[}}%[[VAL_14]]] : memref<32xf64>
 // CHECK-MIR:           }
-// CHECK-MIR:           %[[VAL_28:.*]] = memref.tensor_load %[[VAL_11]] : memref<64xf64>
-// CHECK-MIR:           return %[[VAL_28]] : tensor<64xf64>
+// CHECK-MIR:           %[[VAL_28:.*]] = memref.tensor_load %[[VAL_11]] : memref<32xf64>
+// CHECK-MIR:           return %[[VAL_28]] : tensor<32xf64>
 // CHECK-MIR:         }
 
 // CHECK-LIR-LABEL:   func @matvec(
 // CHECK-LIR-SAME:                 %[[VAL_0:.*]]: !llvm.ptr<i8>,
 // CHECK-LIR-SAME:                 %[[VAL_1:.*]]: memref<64xf64>,
-// CHECK-LIR-SAME:                 %[[VAL_2:.*]]: memref<64xf64>) -> memref<64xf64> {
-// CHECK-LIR:           %[[VAL_3:.*]] = constant 64 : index
+// CHECK-LIR-SAME:                 %[[VAL_2:.*]]: memref<32xf64>) -> memref<32xf64> {
+// CHECK-LIR:           %[[VAL_3:.*]] = constant 32 : index
 // CHECK-LIR:           %[[VAL_4:.*]] = constant 0 : index
 // CHECK-LIR:           %[[VAL_5:.*]] = constant 1 : index
 // CHECK-LIR:           %[[VAL_6:.*]] = call @sparsePointers(%[[VAL_0]], %[[VAL_5]]) : (!llvm.ptr<i8>, index) -> memref<?xindex>
 // CHECK-LIR:           %[[VAL_7:.*]] = call @sparseIndices(%[[VAL_0]], %[[VAL_5]]) : (!llvm.ptr<i8>, index) -> memref<?xindex>
 // CHECK-LIR:           %[[VAL_8:.*]] = call @sparseValuesF64(%[[VAL_0]]) : (!llvm.ptr<i8>) -> memref<?xf64>
-// CHECK-LIR:           %[[VAL_9:.*]] = memref.alloc() : memref<64xf64>
+// CHECK-LIR:           %[[VAL_9:.*]] = memref.alloc() : memref<32xf64>
 // CHECK-LIR:           scf.for %[[VAL_10:.*]] = %[[VAL_4]] to %[[VAL_3]] step %[[VAL_5]] {
-// CHECK-LIR:             %[[VAL_11:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_10]]] : memref<64xf64>
-// CHECK-LIR:             store %[[VAL_11]], %[[VAL_9]]{{\[}}%[[VAL_10]]] : memref<64xf64>
+// CHECK-LIR:             %[[VAL_11:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_10]]] : memref<32xf64>
+// CHECK-LIR:             memref.store %[[VAL_11]], %[[VAL_9]]{{\[}}%[[VAL_10]]] : memref<32xf64>
 // CHECK-LIR:           }
 // CHECK-LIR:           scf.for %[[VAL_12:.*]] = %[[VAL_4]] to %[[VAL_3]] step %[[VAL_5]] {
 // CHECK-LIR:             %[[VAL_13:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_12]]] : memref<?xindex>
 // CHECK-LIR:             %[[VAL_14:.*]] = addi %[[VAL_12]], %[[VAL_5]] : index
 // CHECK-LIR:             %[[VAL_15:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_14]]] : memref<?xindex>
-// CHECK-LIR:             %[[VAL_16:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_12]]] : memref<64xf64>
+// CHECK-LIR:             %[[VAL_16:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_12]]] : memref<32xf64>
 // CHECK-LIR:             %[[VAL_17:.*]] = scf.for %[[VAL_18:.*]] = %[[VAL_13]] to %[[VAL_15]] step %[[VAL_5]] iter_args(%[[VAL_19:.*]] = %[[VAL_16]]) -> (f64) {
 // CHECK-LIR:               %[[VAL_20:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_18]]] : memref<?xindex>
 // CHECK-LIR:               %[[VAL_21:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_18]]] : memref<?xf64>
 // CHECK-LIR:               %[[VAL_22:.*]] = memref.load %[[VAL_1]]{{\[}}%[[VAL_20]]] : memref<64xf64>
 // CHECK-LIR:               %[[VAL_23:.*]] = mulf %[[VAL_21]], %[[VAL_22]] : f64
 // CHECK-LIR:               %[[VAL_24:.*]] = addf %[[VAL_19]], %[[VAL_23]] : f64
 // CHECK-LIR:               scf.yield %[[VAL_24]] : f64
 // CHECK-LIR:             }
-// CHECK-LIR:             store %[[VAL_25:.*]], %[[VAL_9]]{{\[}}%[[VAL_12]]] : memref<64xf64>
+// CHECK-LIR:             memref.store %[[VAL_25:.*]], %[[VAL_9]]{{\[}}%[[VAL_12]]] : memref<32xf64>
 // CHECK-LIR:           }
-// CHECK-LIR:           return %[[VAL_9]] : memref<64xf64>
+// CHECK-LIR:           return %[[VAL_9]] : memref<32xf64>
 // CHECK-LIR:         }
 
-func @matvec(%arga: tensor<64x64xf64, #CSR>,
+func @matvec(%arga: tensor<32x64xf64, #CSR>,
              %argb: tensor<64xf64>,
-       %argx: tensor<64xf64>) -> tensor<64xf64> {
+             %argx: tensor<32xf64>) -> tensor<32xf64> {
   %0 = linalg.generic #trait_matvec
-      ins(%arga, %argb : tensor<64x64xf64, #CSR>, tensor<64xf64>)
-      outs(%argx: tensor<64xf64>) {
+      ins(%arga, %argb : tensor<32x64xf64, #CSR>, tensor<64xf64>)
+      outs(%argx: tensor<32xf64>) {
     ^bb(%A: f64, %b: f64, %x: f64):
       %0 = mulf %A, %b : f64
       %1 = addf %x, %0 : f64
       linalg.yield %1 : f64
-  } -> tensor<64xf64>
-  return %0 : tensor<64xf64>
+  } -> tensor<32xf64>
+  return %0 : tensor<32xf64>
 }

mlir/test/Dialect/SparseTensor/sparse_lower_col.mlir

+// RUN: mlir-opt %s -sparsification | FileCheck %s --check-prefix=CHECK-HIR
+//
+// RUN: mlir-opt %s -sparsification --sparse-tensor-conversion                 \
+// RUN: --convert-linalg-to-loops | FileCheck %s --check-prefix=CHECK-MIR
+//
+// RUN: mlir-opt %s -sparsification --sparse-tensor-conversion                 \
+// RUN: --convert-linalg-to-loops --func-bufferize --tensor-constant-bufferize \
+// RUN: --tensor-bufferize --finalizing-bufferize |                            \
+// RUN: FileCheck %s --check-prefix=CHECK-LIR
+
+#CSC = #sparse_tensor.encoding<{
+  dimLevelType = [ "dense", "compressed" ],
+  dimOrdering = affine_map<(i,j) -> (j,i)>
+}>
+
+#trait_matvec = {
+  indexing_maps = [
+    affine_map<(i,j) -> (i,j)>,  // A
+    affine_map<(i,j) -> (j)>,    // b
+    affine_map<(i,j) -> (i)>     // x (out)
+  ],
+  iterator_types = ["parallel","reduction"],
+  doc = "x(i) += A(i,j) * b(j)"
+}
+
+// CHECK-HIR-LABEL:   func @matvec(
+// CHECK-HIR-SAME:                 %[[VAL_0:.*]]: tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ], dimOrdering = affine_map<(d0, d1) -> (d1, d0)>, pointerBitWidth = 0, indexBitWidth = 0 }>>,
+// CHECK-HIR-SAME:                 %[[VAL_1:.*]]: tensor<64xf64>,
+// CHECK-HIR-SAME:                 %[[VAL_2:.*]]: tensor<32xf64>) -> tensor<32xf64> {
+// CHECK-HIR:           %[[VAL_3:.*]] = constant 64 : index
+// CHECK-HIR:           %[[VAL_4:.*]] = constant 0 : index
+// CHECK-HIR:           %[[VAL_5:.*]] = constant 1 : index
+// CHECK-HIR:           %[[VAL_6:.*]] = sparse_tensor.pointers %[[VAL_0]], %[[VAL_5]] : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ], dimOrdering = affine_map<(d0, d1) -> (d1, d0)>, pointerBitWidth = 0, indexBitWidth = 0 }>> to memref<?xindex>
+// CHECK-HIR:           %[[VAL_7:.*]] = sparse_tensor.indices %[[VAL_0]], %[[VAL_5]] : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ], dimOrdering = affine_map<(d0, d1) -> (d1, d0)>, pointerBitWidth = 0, indexBitWidth = 0 }>> to memref<?xindex>
+// CHECK-HIR:           %[[VAL_8:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ], dimOrdering = affine_map<(d0, d1) -> (d1, d0)>, pointerBitWidth = 0, indexBitWidth = 0 }>> to memref<?xf64>
+// CHECK-HIR:           %[[VAL_9:.*]] = memref.buffer_cast %[[VAL_1]] : memref<64xf64>
+// CHECK-HIR:           %[[VAL_10:.*]] = memref.buffer_cast %[[VAL_2]] : memref<32xf64>
+// CHECK-HIR:           %[[VAL_11:.*]] = memref.alloc() : memref<32xf64>
+// CHECK-HIR:           linalg.copy(%[[VAL_10]], %[[VAL_11]]) : memref<32xf64>, memref<32xf64>
+// CHECK-HIR:           scf.for %[[VAL_12:.*]] = %[[VAL_4]] to %[[VAL_3]] step %[[VAL_5]] {
+// CHECK-HIR:             %[[VAL_13:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_12]]] : memref<64xf64>
+// CHECK-HIR:             %[[VAL_14:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_12]]] : memref<?xindex>
+// CHECK-HIR:             %[[VAL_15:.*]] = addi %[[VAL_12]], %[[VAL_5]] : index
+// CHECK-HIR:             %[[VAL_16:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_15]]] : memref<?xindex>
+// CHECK-HIR:             scf.for %[[VAL_17:.*]] = %[[VAL_14]] to %[[VAL_16]] step %[[VAL_5]] {
+// CHECK-HIR:               %[[VAL_18:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_17]]] : memref<?xindex>
+// CHECK-HIR:               %[[VAL_19:.*]] = memref.load %[[VAL_11]]{{\[}}%[[VAL_18]]] : memref<32xf64>
+// CHECK-HIR:               %[[VAL_20:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_17]]] : memref<?xf64>
+// CHECK-HIR:               %[[VAL_21:.*]] = mulf %[[VAL_20]], %[[VAL_13]] : f64
+// CHECK-HIR:               %[[VAL_22:.*]] = addf %[[VAL_19]], %[[VAL_21]] : f64
+// CHECK-HIR:               memref.store %[[VAL_22]], %[[VAL_11]]{{\[}}%[[VAL_18]]] : memref<32xf64>
+// CHECK-HIR:             }
+// CHECK-HIR:           }
+// CHECK-HIR:           %[[VAL_23:.*]] = memref.tensor_load %[[VAL_11]] : memref<32xf64>
+// CHECK-HIR:           return %[[VAL_23]] : tensor<32xf64>
+// CHECK-HIR:         }
+
+// CHECK-MIR-LABEL:   func @matvec(
+// CHECK-MIR-SAME:                 %[[VAL_0:.*]]: !llvm.ptr<i8>,
+// CHECK-MIR-SAME:                 %[[VAL_1:.*]]: tensor<64xf64>,
+// CHECK-MIR-SAME:                 %[[VAL_2:.*]]: tensor<32xf64>) -> tensor<32xf64> {
+// CHECK-MIR:           %[[VAL_3:.*]] = constant 64 : index
+// CHECK-MIR:           %[[VAL_4:.*]] = constant 32 : index
+// CHECK-MIR:           %[[VAL_5:.*]] = constant 0 : index
+// CHECK-MIR:           %[[VAL_6:.*]] = constant 1 : index
+// CHECK-MIR:           %[[VAL_7:.*]] = call @sparsePointers(%[[VAL_0]], %[[VAL_6]]) : (!llvm.ptr<i8>, index) -> memref<?xindex>
+// CHECK-MIR:           %[[VAL_8:.*]] = call @sparseIndices(%[[VAL_0]], %[[VAL_6]]) : (!llvm.ptr<i8>, index) -> memref<?xindex>
+// CHECK-MIR:           %[[VAL_9:.*]] = call @sparseValuesF64(%[[VAL_0]]) : (!llvm.ptr<i8>) -> memref<?xf64>
+// CHECK-MIR:           %[[VAL_10:.*]] = memref.buffer_cast %[[VAL_1]] : memref<64xf64>
+// CHECK-MIR:           %[[VAL_11:.*]] = memref.buffer_cast %[[VAL_2]] : memref<32xf64>
+// CHECK-MIR:           %[[VAL_12:.*]] = memref.alloc() : memref<32xf64>
+// CHECK-MIR:           scf.for %[[VAL_13:.*]] = %[[VAL_5]] to %[[VAL_4]] step %[[VAL_6]] {
+// CHECK-MIR:             %[[VAL_14:.*]] = memref.load %[[VAL_11]]{{\[}}%[[VAL_13]]] : memref<32xf64>
+// CHECK-MIR:             memref.store %[[VAL_14]], %[[VAL_12]]{{\[}}%[[VAL_13]]] : memref<32xf64>
+// CHECK-MIR:           }
+// CHECK-MIR:           scf.for %[[VAL_15:.*]] = %[[VAL_5]] to %[[VAL_3]] step %[[VAL_6]] {
+// CHECK-MIR:             %[[VAL_16:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_15]]] : memref<64xf64>
+// CHECK-MIR:             %[[VAL_17:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_15]]] : memref<?xindex>
+// CHECK-MIR:             %[[VAL_18:.*]] = addi %[[VAL_15]], %[[VAL_6]] : index
+// CHECK-MIR:             %[[VAL_19:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_18]]] : memref<?xindex>
+// CHECK-MIR:             scf.for %[[VAL_20:.*]] = %[[VAL_17]] to %[[VAL_19]] step %[[VAL_6]] {
+// CHECK-MIR:               %[[VAL_21:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_20]]] : memref<?xindex>
+// CHECK-MIR:               %[[VAL_22:.*]] = memref.load %[[VAL_12]]{{\[}}%[[VAL_21]]] : memref<32xf64>
+// CHECK-MIR:               %[[VAL_23:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_20]]] : memref<?xf64>
+// CHECK-MIR:               %[[VAL_24:.*]] = mulf %[[VAL_23]], %[[VAL_16]] : f64
+// CHECK-MIR:               %[[VAL_25:.*]] = addf %[[VAL_22]], %[[VAL_24]] : f64
+// CHECK-MIR:               memref.store %[[VAL_25]], %[[VAL_12]]{{\[}}%[[VAL_21]]] : memref<32xf64>
+// CHECK-MIR:             }
+// CHECK-MIR:           }
+// CHECK-MIR:           %[[VAL_26:.*]] = memref.tensor_load %[[VAL_12]] : memref<32xf64>
+// CHECK-MIR:           return %[[VAL_26]] : tensor<32xf64>
+// CHECK-MIR:         }
+
+// CHECK-LIR-LABEL:   func @matvec(
+// CHECK-LIR-SAME:                 %[[VAL_0:.*]]: !llvm.ptr<i8>,
+// CHECK-LIR-SAME:                 %[[VAL_1:.*]]: memref<64xf64>,
+// CHECK-LIR-SAME:                 %[[VAL_2:.*]]: memref<32xf64>) -> memref<32xf64> {
+// CHECK-LIR:           %[[VAL_3:.*]] = constant 64 : index
+// CHECK-LIR:           %[[VAL_4:.*]] = constant 32 : index
+// CHECK-LIR:           %[[VAL_5:.*]] = constant 0 : index
+// CHECK-LIR:           %[[VAL_6:.*]] = constant 1 : index
+// CHECK-LIR:           %[[VAL_7:.*]] = call @sparsePointers(%[[VAL_0]], %[[VAL_6]]) : (!llvm.ptr<i8>, index) -> memref<?xindex>
+// CHECK-LIR:           %[[VAL_8:.*]] = call @sparseIndices(%[[VAL_0]], %[[VAL_6]]) : (!llvm.ptr<i8>, index) -> memref<?xindex>
+// CHECK-LIR:           %[[VAL_9:.*]] = call @sparseValuesF64(%[[VAL_0]]) : (!llvm.ptr<i8>) -> memref<?xf64>
+// CHECK-LIR:           %[[VAL_10:.*]] = memref.alloc() : memref<32xf64>
+// CHECK-LIR:           scf.for %[[VAL_11:.*]] = %[[VAL_5]] to %[[VAL_4]] step %[[VAL_6]] {
+// CHECK-LIR:             %[[VAL_12:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_11]]] : memref<32xf64>
+// CHECK-LIR:             memref.store %[[VAL_12]], %[[VAL_10]]{{\[}}%[[VAL_11]]] : memref<32xf64>
+// CHECK-LIR:           }
+// CHECK-LIR:           scf.for %[[VAL_13:.*]] = %[[VAL_5]] to %[[VAL_3]] step %[[VAL_6]] {
+// CHECK-LIR:             %[[VAL_14:.*]] = memref.load %[[VAL_1]]{{\[}}%[[VAL_13]]] : memref<64xf64>
+// CHECK-LIR:             %[[VAL_15:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_13]]] : memref<?xindex>
+// CHECK-LIR:             %[[VAL_16:.*]] = addi %[[VAL_13]], %[[VAL_6]] : index
+// CHECK-LIR:             %[[VAL_17:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_16]]] : memref<?xindex>
+// CHECK-LIR:             scf.for %[[VAL_18:.*]] = %[[VAL_15]] to %[[VAL_17]] step %[[VAL_6]] {
+// CHECK-LIR:               %[[VAL_19:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_18]]] : memref<?xindex>
+// CHECK-LIR:               %[[VAL_20:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_19]]] : memref<32xf64>
+// CHECK-LIR:               %[[VAL_21:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_18]]] : memref<?xf64>
+// CHECK-LIR:               %[[VAL_22:.*]] = mulf %[[VAL_21]], %[[VAL_14]] : f64
+// CHECK-LIR:               %[[VAL_23:.*]] = addf %[[VAL_20]], %[[VAL_22]] : f64
+// CHECK-LIR:               memref.store %[[VAL_23]], %[[VAL_10]]{{\[}}%[[VAL_19]]] : memref<32xf64>
+// CHECK-LIR:             }
+// CHECK-LIR:           }
+// CHECK-LIR:           return %[[VAL_10]] : memref<32xf64>
+// CHECK-LIR:         }
+
+func @matvec(%arga: tensor<32x64xf64, #CSC>,
+             %argb: tensor<64xf64>,
+             %argx: tensor<32xf64>) -> tensor<32xf64> {
+  %0 = linalg.generic #trait_matvec
+      ins(%arga, %argb : tensor<32x64xf64, #CSC>, tensor<64xf64>)
+      outs(%argx: tensor<32xf64>) {
+    ^bb(%A: f64, %b: f64, %x: f64):
+      %0 = mulf %A, %b : f64
+      %1 = addf %x, %0 : f64
+      linalg.yield %1 : f64
+  } -> tensor<32xf64>
+  return %0 : tensor<32xf64>
+}

mlir/test/Dialect/SparseTensor/sparse_lower_inplace.mlir

indexing_maps = [
     affine_map<(i,j) -> (j)>,    // b
     affine_map<(i,j) -> (i)>     // x (out)
   ],
   iterator_types = ["parallel","reduction"],
   doc = "x(i) += A(i,j) * b(j)"
 }
 
 // CHECK-HIR-LABEL:   func @matvec(
-// CHECK-HIR-SAME:                 %[[VAL_0:.*]]: tensor<64x64xf64, #sparse_tensor.encoding<{{.*}}>>,
+// CHECK-HIR-SAME:                 %[[VAL_0:.*]]: tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ], pointerBitWidth = 0, indexBitWidth = 0 }>>,
 // CHECK-HIR-SAME:                 %[[VAL_1:.*]]: tensor<64xf64>,
-// CHECK-HIR-SAME:                 %[[VAL_2:.*]]: tensor<64xf64> {linalg.inplaceable = true}) -> tensor<64xf64> {
-// CHECK-HIR:           %[[VAL_3:.*]] = constant 64 : index
+// CHECK-HIR-SAME:                 %[[VAL_2:.*]]: tensor<32xf64> {linalg.inplaceable = true}) -> tensor<32xf64> {
+// CHECK-HIR:           %[[VAL_3:.*]] = constant 32 : index
 // CHECK-HIR:           %[[VAL_4:.*]] = constant 0 : index
 // CHECK-HIR:           %[[VAL_5:.*]] = constant 1 : index
-// CHECK-HIR:           %[[VAL_6:.*]] = sparse_tensor.pointers %[[VAL_0]], %[[VAL_5]] : tensor<64x64xf64, #sparse_tensor.encoding<{{.*}}>> to memref<?xindex>
-// CHECK-HIR:           %[[VAL_7:.*]] = sparse_tensor.indices %[[VAL_0]], %[[VAL_5]] : tensor<64x64xf64, #sparse_tensor.encoding<{{.*}}>> to memref<?xindex>
-// CHECK-HIR:           %[[VAL_8:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<64x64xf64, #sparse_tensor.encoding<{{.*}}>> to memref<?xf64>
+// CHECK-HIR:           %[[VAL_6:.*]] = sparse_tensor.pointers %[[VAL_0]], %[[VAL_5]] : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ], pointerBitWidth = 0, indexBitWidth = 0 }>> to memref<?xindex>
+// CHECK-HIR:           %[[VAL_7:.*]] = sparse_tensor.indices %[[VAL_0]], %[[VAL_5]] : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ], pointerBitWidth = 0, indexBitWidth = 0 }>> to memref<?xindex>
+// CHECK-HIR:           %[[VAL_8:.*]] = sparse_tensor.values %[[VAL_0]] : tensor<32x64xf64, #sparse_tensor.encoding<{ dimLevelType = [ "dense", "compressed" ], pointerBitWidth = 0, indexBitWidth = 0 }>> to memref<?xf64>
 // CHECK-HIR:           %[[VAL_9:.*]] = memref.buffer_cast %[[VAL_1]] : memref<64xf64>
-// CHECK-HIR:           %[[VAL_10:.*]] = memref.buffer_cast %[[VAL_2]] : memref<64xf64>
+// CHECK-HIR:           %[[VAL_10:.*]] = memref.buffer_cast %[[VAL_2]] : memref<32xf64>
 // CHECK-HIR:           scf.for %[[VAL_11:.*]] = %[[VAL_4]] to %[[VAL_3]] step %[[VAL_5]] {
 // CHECK-HIR:             %[[VAL_12:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_11]]] : memref<?xindex>
 // CHECK-HIR:             %[[VAL_13:.*]] = addi %[[VAL_11]], %[[VAL_5]] : index
 // CHECK-HIR:             %[[VAL_14:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_13]]] : memref<?xindex>
-// CHECK-HIR:             %[[VAL_15:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_11]]] : memref<64xf64>
+// CHECK-HIR:             %[[VAL_15:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_11]]] : memref<32xf64>
 // CHECK-HIR:             %[[VAL_16:.*]] = scf.for %[[VAL_17:.*]] = %[[VAL_12]] to %[[VAL_14]] step %[[VAL_5]] iter_args(%[[VAL_18:.*]] = %[[VAL_15]]) -> (f64) {
 // CHECK-HIR:               %[[VAL_19:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_17]]] : memref<?xindex>
 // CHECK-HIR:               %[[VAL_20:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_17]]] : memref<?xf64>
 // CHECK-HIR:               %[[VAL_21:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_19]]] : memref<64xf64>
 // CHECK-HIR:               %[[VAL_22:.*]] = mulf %[[VAL_20]], %[[VAL_21]] : f64
 // CHECK-HIR:               %[[VAL_23:.*]] = addf %[[VAL_18]], %[[VAL_22]] : f64
 // CHECK-HIR:               scf.yield %[[VAL_23]] : f64
 // CHECK-HIR:             }
-// CHECK-HIR:             memref.store %[[VAL_24:.*]], %[[VAL_10]]{{\[}}%[[VAL_11]]] : memref<64xf64>
+// CHECK-HIR:             memref.store %[[VAL_24:.*]], %[[VAL_10]]{{\[}}%[[VAL_11]]] : memref<32xf64>
 // CHECK-HIR:           }
-// CHECK-HIR:           %[[VAL_25:.*]] = memref.tensor_load %[[VAL_10]] : memref<64xf64>
-// CHECK-HIR:           return %[[VAL_25]] : tensor<64xf64>
+// CHECK-HIR:           %[[VAL_25:.*]] = memref.tensor_load %[[VAL_10]] : memref<32xf64>
+// CHECK-HIR:           return %[[VAL_25]] : tensor<32xf64>
 // CHECK-HIR:         }
 
 // CHECK-MIR-LABEL:   func @matvec(
 // CHECK-MIR-SAME:                 %[[VAL_0:.*]]: !llvm.ptr<i8>,
 // CHECK-MIR-SAME:                 %[[VAL_1:.*]]: tensor<64xf64>,
-// CHECK-MIR-SAME:                 %[[VAL_2:.*]]: tensor<64xf64> {linalg.inplaceable = true}) -> tensor<64xf64> {
-// CHECK-MIR:           %[[VAL_3:.*]] = constant 64 : index
+// CHECK-MIR-SAME:                 %[[VAL_2:.*]]: tensor<32xf64> {linalg.inplaceable = true}) -> tensor<32xf64> {
+// CHECK-MIR:           %[[VAL_3:.*]] = constant 32 : index
 // CHECK-MIR:           %[[VAL_4:.*]] = constant 0 : index
 // CHECK-MIR:           %[[VAL_5:.*]] = constant 1 : index
 // CHECK-MIR:           %[[VAL_6:.*]] = call @sparsePointers(%[[VAL_0]], %[[VAL_5]]) : (!llvm.ptr<i8>, index) -> memref<?xindex>
 // CHECK-MIR:           %[[VAL_7:.*]] = call @sparseIndices(%[[VAL_0]], %[[VAL_5]]) : (!llvm.ptr<i8>, index) -> memref<?xindex>
 // CHECK-MIR:           %[[VAL_8:.*]] = call @sparseValuesF64(%[[VAL_0]]) : (!llvm.ptr<i8>) -> memref<?xf64>
 // CHECK-MIR:           %[[VAL_9:.*]] = memref.buffer_cast %[[VAL_1]] : memref<64xf64>
-// CHECK-MIR:           %[[VAL_10:.*]] = memref.buffer_cast %[[VAL_2]] : memref<64xf64>
+// CHECK-MIR:           %[[VAL_10:.*]] = memref.buffer_cast %[[VAL_2]] : memref<32xf64>
 // CHECK-MIR:           scf.for %[[VAL_11:.*]] = %[[VAL_4]] to %[[VAL_3]] step %[[VAL_5]] {
 // CHECK-MIR:             %[[VAL_12:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_11]]] : memref<?xindex>
 // CHECK-MIR:             %[[VAL_13:.*]] = addi %[[VAL_11]], %[[VAL_5]] : index
 // CHECK-MIR:             %[[VAL_14:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_13]]] : memref<?xindex>
-// CHECK-MIR:             %[[VAL_15:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_11]]] : memref<64xf64>
+// CHECK-MIR:             %[[VAL_15:.*]] = memref.load %[[VAL_10]]{{\[}}%[[VAL_11]]] : memref<32xf64>
 // CHECK-MIR:             %[[VAL_16:.*]] = scf.for %[[VAL_17:.*]] = %[[VAL_12]] to %[[VAL_14]] step %[[VAL_5]] iter_args(%[[VAL_18:.*]] = %[[VAL_15]]) -> (f64) {
 // CHECK-MIR:               %[[VAL_19:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_17]]] : memref<?xindex>
 // CHECK-MIR:               %[[VAL_20:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_17]]] : memref<?xf64>
 // CHECK-MIR:               %[[VAL_21:.*]] = memref.load %[[VAL_9]]{{\[}}%[[VAL_19]]] : memref<64xf64>
 // CHECK-MIR:               %[[VAL_22:.*]] = mulf %[[VAL_20]], %[[VAL_21]] : f64
 // CHECK-MIR:               %[[VAL_23:.*]] = addf %[[VAL_18]], %[[VAL_22]] : f64
 // CHECK-MIR:               scf.yield %[[VAL_23]] : f64
 // CHECK-MIR:             }
-// CHECK-MIR:             memref.store %[[VAL_24:.*]], %[[VAL_10]]{{\[}}%[[VAL_11]]] : memref<64xf64>
+// CHECK-MIR:             memref.store %[[VAL_24:.*]], %[[VAL_10]]{{\[}}%[[VAL_11]]] : memref<32xf64>
 // CHECK-MIR:           }
-// CHECK-MIR:           %[[VAL_25:.*]] = memref.tensor_load %[[VAL_10]] : memref<64xf64>
-// CHECK-MIR:           return %[[VAL_25]] : tensor<64xf64>
+// CHECK-MIR:           %[[VAL_25:.*]] = memref.tensor_load %[[VAL_10]] : memref<32xf64>
+// CHECK-MIR:           return %[[VAL_25]] : tensor<32xf64>
 // CHECK-MIR:         }
 
 // CHECK-LIR-LABEL:   func @matvec(
 // CHECK-LIR-SAME:                 %[[VAL_0:.*]]: !llvm.ptr<i8>,
 // CHECK-LIR-SAME:                 %[[VAL_1:.*]]: memref<64xf64>,
-// CHECK-LIR-SAME:                 %[[VAL_2:.*]]: memref<64xf64> {linalg.inplaceable = true}) -> memref<64xf64> {
-// CHECK-LIR:           %[[VAL_3:.*]] = constant 64 : index
+// CHECK-LIR-SAME:                 %[[VAL_2:.*]]: memref<32xf64> {linalg.inplaceable = true}) -> memref<32xf64> {
+// CHECK-LIR:           %[[VAL_3:.*]] = constant 32 : index
 // CHECK-LIR:           %[[VAL_4:.*]] = constant 0 : index
 // CHECK-LIR:           %[[VAL_5:.*]] = constant 1 : index
 // CHECK-LIR:           %[[VAL_6:.*]] = call @sparsePointers(%[[VAL_0]], %[[VAL_5]]) : (!llvm.ptr<i8>, index) -> memref<?xindex>
 // CHECK-LIR:           %[[VAL_7:.*]] = call @sparseIndices(%[[VAL_0]], %[[VAL_5]]) : (!llvm.ptr<i8>, index) -> memref<?xindex>
 // CHECK-LIR:           %[[VAL_8:.*]] = call @sparseValuesF64(%[[VAL_0]]) : (!llvm.ptr<i8>) -> memref<?xf64>
 // CHECK-LIR:           scf.for %[[VAL_9:.*]] = %[[VAL_4]] to %[[VAL_3]] step %[[VAL_5]] {
 // CHECK-LIR:             %[[VAL_10:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_9]]] : memref<?xindex>
 // CHECK-LIR:             %[[VAL_11:.*]] = addi %[[VAL_9]], %[[VAL_5]] : index
 // CHECK-LIR:             %[[VAL_12:.*]] = memref.load %[[VAL_6]]{{\[}}%[[VAL_11]]] : memref<?xindex>
-// CHECK-LIR:             %[[VAL_13:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_9]]] : memref<64xf64>
+// CHECK-LIR:             %[[VAL_13:.*]] = memref.load %[[VAL_2]]{{\[}}%[[VAL_9]]] : memref<32xf64>
 // CHECK-LIR:             %[[VAL_14:.*]] = scf.for %[[VAL_15:.*]] = %[[VAL_10]] to %[[VAL_12]] step %[[VAL_5]] iter_args(%[[VAL_16:.*]] = %[[VAL_13]]) -> (f64) {
 // CHECK-LIR:               %[[VAL_17:.*]] = memref.load %[[VAL_7]]{{\[}}%[[VAL_15]]] : memref<?xindex>
 // CHECK-LIR:               %[[VAL_18:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_15]]] : memref<?xf64>
 // CHECK-LIR:               %[[VAL_19:.*]] = memref.load %[[VAL_1]]{{\[}}%[[VAL_17]]] : memref<64xf64>
 // CHECK-LIR:               %[[VAL_20:.*]] = mulf %[[VAL_18]], %[[VAL_19]] : f64
 // CHECK-LIR:               %[[VAL_21:.*]] = addf %[[VAL_16]], %[[VAL_20]] : f64
 // CHECK-LIR:               scf.yield %[[VAL_21]] : f64
 // CHECK-LIR:             }
-// CHECK-LIR:             memref.store %[[VAL_22:.*]], %[[VAL_2]]{{\[}}%[[VAL_9]]] : memref<64xf64>
+// CHECK-LIR:             memref.store %[[VAL_22:.*]], %[[VAL_2]]{{\[}}%[[VAL_9]]] : memref<32xf64>
 // CHECK-LIR:           }
-// CHECK-LIR:           return %[[VAL_2]] : memref<64xf64>
+// CHECK-LIR:           return %[[VAL_2]] : memref<32xf64>
 // CHECK-LIR:         }
 
-func @matvec(%arga: tensor<64x64xf64, #CSR>,
+func @matvec(%arga: tensor<32x64xf64, #CSR>,
              %argb: tensor<64xf64>,
-       %argx: tensor<64xf64> {linalg.inplaceable = true}) -> tensor<64xf64> {
+       %argx: tensor<32xf64> {linalg.inplaceable = true}) -> tensor<32xf64> {
   %0 = linalg.generic #trait_matvec
-      ins(%arga, %argb : tensor<64x64xf64, #CSR>, tensor<64xf64>)
-      outs(%argx: tensor<64xf64>) {
+      ins(%arga, %argb : tensor<32x64xf64, #CSR>, tensor<64xf64>)
+      outs(%argx: tensor<32xf64>) {
     ^bb(%A: f64, %b: f64, %x: f64):
       %0 = mulf %A, %b : f64
       %1 = addf %x, %0 : f64
       linalg.yield %1 : f64
-  } -> tensor<64xf64>
-  return %0 : tensor<64xf64>
+  } -> tensor<32xf64>
+  return %0 : tensor<32xf64>
 }

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

+// RUN: mlir-opt %s \
+// RUN:   --sparsification --sparse-tensor-conversion \
+// RUN:   --convert-linalg-to-loops --convert-vector-to-scf --convert-scf-to-std \
+// RUN:   --func-bufferize --tensor-constant-bufferize --tensor-bufferize \
+// RUN:   --std-bufferize --finalizing-bufferize  \
+// RUN:   --convert-vector-to-llvm --convert-std-to-llvm | \
+// RUN: TENSOR0="%mlir_integration_test_dir/data/test.tns" \
+// RUN: mlir-cpu-runner \
+// RUN:  -e entry -entry-point-result=void  \
+// RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
+// RUN: FileCheck %s
+
+!Filename = type !llvm.ptr<i8>
+
+#SparseTensor = #sparse_tensor.encoding<{
+  dimLevelType = [ "compressed", "compressed", "compressed", "compressed",
+                   "compressed", "compressed", "compressed", "compressed" ],
+  // Note that any dimOrdering permutation should give the same results
+  // since, even though it impacts the sparse storage scheme layout,
+  // it should not change the semantics.
+  dimOrdering = affine_map<(i,j,k,l,m,n,o,p) -> (p,o,j,k,i,l,m,n)>
+}>
+
+#trait_flatten = {
+  indexing_maps = [
+    affine_map<(i,j,k,l,m,n,o,p) -> (i,j,k,l,m,n,o,p)>, // A
+    affine_map<(i,j,k,l,m,n,o,p) -> (i,j)>              // X (out)
+  ],
+  iterator_types = [ "parallel",  "parallel",  "reduction", "reduction",
+                     "reduction", "reduction", "reduction", "reduction" ],
+  doc = "X(i,j) += A(i,j,k,l,m,n,o,p)"
+}
+
+//
+// Integration test that lowers a kernel annotated as sparse to
+// actual sparse code, initializes a matching sparse storage scheme
+// from file, and runs the resulting code with the JIT compiler.
+//
+module {
+  //
+  // A kernel that flattens a rank 8 tensor into a dense matrix.
+  //
+  func @kernel_flatten(%arga: tensor<7x3x3x3x3x3x5x3xf64, #SparseTensor>,
+                       %argx: tensor<7x3xf64>) -> tensor<7x3xf64> {
+    %0 = linalg.generic #trait_flatten
+      ins(%arga: tensor<7x3x3x3x3x3x5x3xf64, #SparseTensor>)
+      outs(%argx: tensor<7x3xf64>) {
+      ^bb(%a: f64, %x: f64):
+        %0 = addf %x, %a : f64
+        linalg.yield %0 : f64
+    } -> tensor<7x3xf64>
+    return %0 : tensor<7x3xf64>
+  }
+
+  func private @getTensorFilename(index) -> (!Filename)
+
+  //
+  // Main driver that reads tensor from file and calls the sparse kernel.
+  //
+  func @entry() {
+    %d0 = constant 0.0 : f64
+    %c0 = constant 0 : index
+    %c1 = constant 1 : index
+    %c3 = constant 3 : index
+    %c7 = constant 7 : index
+
+    // Setup matrix memory that is initialized to zero.
+    %xdata = memref.alloc() : memref<7x3xf64>
+    scf.for %i = %c0 to %c7 step %c1 {
+      scf.for %j = %c0 to %c3 step %c1 {
+        memref.store %d0, %xdata[%i, %j] : memref<7x3xf64>
+      }
+    }
+    %x = memref.tensor_load %xdata : memref<7x3xf64>
+
+    // Read the sparse tensor from file, construct sparse storage.
+    %fileName = call @getTensorFilename(%c0) : (index) -> (!Filename)
+    %a = sparse_tensor.new %fileName : !llvm.ptr<i8> to tensor<7x3x3x3x3x3x5x3xf64, #SparseTensor>
+
+    // Call the kernel.
+    %0 = call @kernel_flatten(%a, %x)
+      : (tensor<7x3x3x3x3x3x5x3xf64, #SparseTensor>, tensor<7x3xf64>) -> tensor<7x3xf64>
+
+    // Print the result for verification.
+    //
+    // CHECK: ( 6.25, 0, 0 )
+    // CHECK: ( 4.224, 6.21, 0 )
+    // CHECK: ( 0, 0, 15.455 )
+    // CHECK: ( 0, 0, 0 )
+    // CHECK: ( 0, 0, 0 )
+    // CHECK: ( 0, 0, 0 )
+    // CHECK: ( 7, 0, 0 )
+    //
+    %r = memref.buffer_cast %0 : memref<7x3xf64>
+    scf.for %i = %c0 to %c7 step %c1 {
+      %v = vector.transfer_read %r[%i, %c0], %d0: memref<7x3xf64>, vector<3xf64>
+      vector.print %v : vector<3xf64>
+    }
+
+    // Release the resources.
+    memref.dealloc %xdata : memref<7x3xf64>
+
+    return
+  }
+}