[mlir][sparse] Cleaning up the dim/lvl distinction in SparseTensorConversion

[wrengr]

Summary

This change cleans up the conversion pass re the "dim"-vs-"lvl" and "sizes"-vs-"shape" distinctions of the runtime. A quick synopsis includes:

• Adds new SparseTensorStorageBase::getDimSize method, with sparseDimSize wrapper in SparseTensorRuntime.h, and genDimSizeCall generator in SparseTensorConversion.cpp
• Changes genLvlSizeCall to perform no logic, just generate the function call.
• Adds createOrFold{Dim,Lvl}Call functions to handle the logic of replacing gen{Dim,Lvl}SizeCall with constants whenever possible. The createOrFoldDimCall function replaces the old sizeFromPtrAtDim.
• Adds {get,fill}DimSizes functions for iterating createOrFoldDimCall across the whole type. These functions replace the old sizesFromPtr.
• Adds {get,fill}DimShape functions for lowering a ShapedType into constants. These functions replace the old sizesFromType.
• Changes the DimOp rewrite to do the right thing.
• Changes the ExpandOp rewrite to compute the proper expansion size.

Depends On D138365

[aartbik]

mlir/test/Dialect/SparseTensor/conversion.mlir

364

any reason for dropping the signature?
(and good catch on the brittle match using c1 ;-)

[wrengr]

mlir/test/Dialect/SparseTensor/conversion.mlir

364

Just hoping to avoid brittleness is all. The signature doesn't add any information here, and the different test files seem to vary in whether they include the signatures or not, so I went with the one that seemed more common. I can add it back in if you'd prefer

mlir/include/mlir/ExecutionEngine/SparseTensor/Storage.h

 //===----------------------------------------------------------------------===//
 // This forward decl is sufficient to split `SparseTensorStorageBase` into
 // its own header, but isn't sufficient for `SparseTensorStorage` to join it.
 template <typename V>
 class SparseTensorEnumeratorBase;
 
 // These macros ensure consistent error messages, without risk of incuring
 // an additional method call to do so.
+#define ASSERT_VALID_DIM(d)                                                    \
+  assert(d < getDimRank() && "Dimension index is out of bounds");
 #define ASSERT_VALID_LVL(l)                                                    \
   assert(l < getLvlRank() && "Level index is out of bounds");
 #define ASSERT_COMPRESSED_LVL(l)                                               \
   assert(isCompressedLvl(l) && "Level is not compressed");
 #define ASSERT_COMPRESSED_OR_SINGLETON_LVL(l)                                  \
   do {                                                                         \
     const DimLevelType dlt = getLvlType(l);                                    \
     (void)dlt;                                                                 \
public:
   uint64_t getDimRank() const { return dimSizes.size(); }
 
   /// Gets the number of storage-levels.
   uint64_t getLvlRank() const { return lvlSizes.size(); }
 
   /// Gets the tensor-dimension sizes array.
   const std::vector<uint64_t> &getDimSizes() const { return dimSizes; }
 
+  /// Safely looks up the size of the given tensor-dimension.
+  uint64_t getDimSize(uint64_t d) const {
+    ASSERT_VALID_DIM(d);
+    return dimSizes[d];
+  }
+
   /// Gets the storage-level sizes array.
   const std::vector<uint64_t> &getLvlSizes() const { return lvlSizes; }
 
   /// Safely looks up the size of the given storage-level.
   uint64_t getLvlSize(uint64_t l) const {
     ASSERT_VALID_LVL(l);
     return lvlSizes[l];
   }
private:
   std::vector<std::vector<I>> indices;
   std::vector<V> values;
   std::vector<uint64_t> lvlCursor; // cursor for lexicographic insertion.
 };
 
 #undef ASSERT_COMPRESSED_OR_SINGLETON_LVL
 #undef ASSERT_COMPRESSED_LVL
 #undef ASSERT_VALID_LVL
+#undef ASSERT_VALID_DIM
 
 //===----------------------------------------------------------------------===//
 /// A (higher-order) function object for enumerating the elements of some
 /// `SparseTensorStorage` under a permutation.  That is, the `forallElements`
 /// method encapsulates the loop-nest for enumerating the elements of
 /// the source tensor (in whatever order is best for the source tensor),
 /// and applies a permutation to the coordinates/indices before handing
 /// each element to the callback.  A single enumerator object can be

mlir/include/mlir/ExecutionEngine/SparseTensorRuntime.h

 // as well as by any external runtime that wants to interact with MLIR
 // compiler-generated code.
 //
 //===----------------------------------------------------------------------===//
 
 /// Tensor-storage method to get the size of the given level.
 MLIR_CRUNNERUTILS_EXPORT index_type sparseLvlSize(void *tensor, index_type l);
 
+/// Tensor-storage method to get the size of the given dimension.
+MLIR_CRUNNERUTILS_EXPORT index_type sparseDimSize(void *tensor, index_type d);
+
 /// Tensor-storage method to finalize lexicographic insertions.
 MLIR_CRUNNERUTILS_EXPORT void endInsert(void *tensor);
 
 /// Coordinate-scheme method to write to file in extended FROSTT format.
 #define DECL_OUTSPARSETENSOR(VNAME, V)                                         \
   MLIR_CRUNNERUTILS_EXPORT void outSparseTensor##VNAME(void *coo, void *dest,  \
                                                        bool sort);
 MLIR_SPARSETENSOR_FOREVERY_V(DECL_OUTSPARSETENSOR)

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

static func::CallOp replaceOpWithFuncCall(RewriterBase &rewriter, Operation *op,
                                           ValueRange operands,
                                           EmitCInterface emitCInterface) {
   auto fn = getFunc(op->getParentOfType<ModuleOp>(), name, resultType, operands,
                     emitCInterface);
   return rewriter.replaceOpWithNewOp<func::CallOp>(op, resultType, fn,
                                                    operands);
 }
 
-/// Generates call to lookup a level-size.
-static Value genLvlSizeCall(OpBuilder &builder, Location loc,
-                            SparseTensorEncodingAttr &enc, Value src,
+/// Generates call to lookup a level-size.  N.B., this only generates
+/// the raw function call, and therefore (intentionally) does not perform
+/// any dim<->lvl conversion or other logic.
+static Value genLvlSizeCall(OpBuilder &builder, Location loc, Value tensor,
                             uint64_t lvl) {
-  // Generate the call.
   StringRef name = "sparseLvlSize";
-  SmallVector<Value, 2> params{  // just two
-      src, constantIndex(builder, loc, toStoredDim(enc, lvl))};
+  SmallVector<Value, 2> params{tensor, constantIndex(builder, loc, lvl)};
   Type iTp = builder.getIndexType();
   return createFuncCall(builder, loc, name, iTp, params, EmitCInterface::Off)
       .getResult(0);
 }
 
-/// Compute the size from type (for static sizes) or from an already-converted
-/// opaque pointer source (for dynamic sizes) at the given dimension.
-//
-// FIXME: Need to rename this function to match `genLvlSizeCall` and hence
-// match the naming convention used in the runtime library.  However, it's
-// not entirely clear that all callsites of this function properly make the
-// "level"-vs-"dimension" distinction; so need to audit each callsite to
-// ensure this still does what they mean (possibly by having two separate
-// functions, one for levels and one for dimensions).  That also means
-// renaming `sizesFromPtr`, `sizesFromType`, etc, to make clear whether
-// they mean to be referring to level-sizes vs dimension-sizes.
-static Value sizeFromPtrAtDim(OpBuilder &builder, Location loc,
-                              SparseTensorEncodingAttr &enc, ShapedType stp,
-                              Value src, unsigned i) {
-  auto shape = stp.getShape();
-  if (shape[i] == ShapedType::kDynamic)
-    return genLvlSizeCall(builder, loc, enc, src, i);
-  return constantIndex(builder, loc, shape[i]);
+/// Generates call to lookup a dimension-size.  N.B., this only generates
+/// the raw function call, and therefore (intentionally) does not perform
+/// any dim<->lvl conversion or other logic.
+static Value genDimSizeCall(OpBuilder &builder, Location loc, Value tensor,
+                            uint64_t dim) {
+  StringRef name = "sparseDimSize";
+  SmallVector<Value, 2> params{tensor, constantIndex(builder, loc, dim)};
+  Type iTp = builder.getIndexType();
+  return createFuncCall(builder, loc, name, iTp, params, EmitCInterface::Off)
+      .getResult(0);
 }
 
-/// Populates given sizes array from type (for static sizes) and from
-/// an already-converted opaque pointer source (for dynamic sizes).
-static void sizesFromPtr(OpBuilder &builder, SmallVectorImpl<Value> &sizes,
-                         Location loc, SparseTensorEncodingAttr &enc,
-                         ShapedType stp, Value src) {
-  unsigned rank = stp.getRank();
-  sizes.reserve(rank);
-  for (unsigned i = 0; i < rank; i++)
-    sizes.push_back(sizeFromPtrAtDim(builder, loc, enc, stp, src, i));
+/// Looks up a level-size by returning a statically-computed constant
+/// (when possible), or by calling `genLvlSizeCall` (when dynamic).
+static Value createOrFoldLvlCall(OpBuilder &builder, Location loc,
+                                 SparseTensorEncodingAttr &enc, ShapedType stp,
+                                 Value tensor, unsigned lvl) {
+  // Only sparse tensors have "levels" to query.
+  assert(enc);
+  auto dimOrder = enc.getDimOrdering();
+  // TODO: The following implementation only handles permutations;
+  // we'll need to generalize this to handle arbitrary AffineExpr.
+  //
+  // There's no need to assert `isPermutation` here: because
+  // `getDimPosition` checks that the expr isa `AffineDimExpr`,
+  // which is all we care about (for supporting permutations).
+  unsigned dim = dimOrder ? dimOrder.getDimPosition(lvl) : lvl;
+  auto s = stp.getShape()[dim];
+  if (s != ShapedType::kDynamic)
+    return constantIndex(builder, loc, s);
+  // If we cannot statically compute the size from the shape, then we
+  // must dynamically query it.  (In principle we could also dynamically
+  // compute it, but since we already did so to construct the `tensor`
+  // in the first place, we might as well query rather than recompute.)
+  return genLvlSizeCall(builder, loc, tensor, lvl);
+}
+
+/// Looks up a dimension-size by returning a constant from the shape
+/// (for static sizes), or by calling `genDimSizeCall` (for dynamic sizes
+/// of sparse tensors) or `linalg::createOrFoldDimOp` (for dynamic sizes
+/// of dense tensors).
+static Value createOrFoldDimCall(OpBuilder &builder, Location loc,
+                                 SparseTensorEncodingAttr &enc, ShapedType stp,
+                                 Value tensor, unsigned dim) {
+  auto s = stp.getShape()[dim];
+  if (s != ShapedType::kDynamic)
+    return constantIndex(builder, loc, s);
+  if (enc)
+    return genDimSizeCall(builder, loc, tensor, dim);
+  return linalg::createOrFoldDimOp(builder, loc, tensor, dim);
 }
 
-/// Populates given sizes array from type.
-static void sizesFromType(OpBuilder &builder, SmallVectorImpl<Value> &sizes,
-                          Location loc, ShapedType stp) {
+/// Populates the array with the dimension-sizes of the given tensor.
+static void fillDimSizes(OpBuilder &builder, Location loc,
+                         SparseTensorEncodingAttr &enc, ShapedType stp,
+                         Value tensor, SmallVectorImpl<Value> &out) {
+  unsigned dimRank = stp.getRank();
+  out.reserve(dimRank);
+  for (unsigned d = 0; d < dimRank; d++)
+    out.push_back(createOrFoldDimCall(builder, loc, enc, stp, tensor, d));
+}
+
+/// Returns an array with the dimension-sizes of the given tensor.
+static SmallVector<Value> getDimSizes(OpBuilder &builder, Location loc,
+                                      SparseTensorEncodingAttr &enc,
+                                      ShapedType stp, Value tensor) {
+  SmallVector<Value> out;
+  fillDimSizes(builder, loc, enc, stp, tensor, out);
+  return out;
+}
+
+/// Populates the array with the dimension-shape of the given `ShapedType`,
+/// where dynamic sizes are represented by zero.
+static void fillDimShape(OpBuilder &builder, Location loc, ShapedType stp,
+                         SmallVectorImpl<Value> &out) {
   auto shape = stp.getShape();
-  unsigned rank = stp.getRank();
-  sizes.reserve(rank);
-  for (unsigned i = 0; i < rank; i++) {
-    uint64_t s = shape[i] == ShapedType::kDynamic ? 0 : shape[i];
-    sizes.push_back(constantIndex(builder, loc, s));
-  }
+  unsigned dimRank = stp.getRank();
+  out.reserve(dimRank);
+  for (unsigned d = 0; d < dimRank; d++) {
+    auto s = shape[d] == ShapedType::kDynamic ? 0 : shape[d];
+    out.push_back(constantIndex(builder, loc, s));
+  }
+
+/// Returns an array with the dimension-shape of the given `ShapedType`,
+/// where dynamic sizes are represented by zero.
+static SmallVector<Value> getDimShape(OpBuilder &builder, Location loc,
+                                      ShapedType stp) {
+  SmallVector<Value> out;
+  fillDimShape(builder, loc, stp, out);
+  return out;
 }
 
 /// Populates the given sizes array for concatenation from type (for static
 /// sizes) and from an already-converted opaque pointer source (for dynamic
 /// sizes).
 static void concatSizesFromInputs(OpBuilder &builder,
                                   SmallVectorImpl<Value> &sizes, Location loc,
                                   ShapedType dstTp, ValueRange srcs,
                                   unsigned dim) {
   auto dstShape = dstTp.getShape();
 
   auto srcTp = srcs[0].getType().cast<ShapedType>();
   auto srcEnc = getSparseTensorEncoding(srcTp);
   // We first fills the sizes from an input tensor, and then
   // compute the size of the concatenation dimension if necessary.
   if (srcEnc)
     // Reuses sizes from an arbitrary input tensor is fine.
-    sizesFromPtr(builder, sizes, loc, srcEnc, srcTp, srcs[0]);
+    fillDimSizes(builder, loc, srcEnc, srcTp, srcs[0], sizes);
   else
     sizesFromSrc(builder, sizes, loc, srcs[0]);
 
   // Sum up on the `dim` if the dimension is dynamic.
   if (dstShape[dim] != ShapedType::kDynamic) {
     // Faithfully take the static size.
     sizes[dim] = constantIndex(builder, loc, dstShape[dim]);
   } else {
     // Else, compute the shape dynamically.
     for (size_t i = 1, sz = srcs.size(); i < sz; i++) {
       auto srcTp = srcs[i].getType().cast<ShapedType>();
       auto encSrc = getSparseTensorEncoding(srcTp);
       Value srcSz =
-          encSrc ? sizeFromPtrAtDim(builder, loc, encSrc, srcTp, srcs[i], dim)
-                 : linalg::createOrFoldDimOp(builder, loc, srcs[i], dim);
+          createOrFoldDimCall(builder, loc, encSrc, srcTp, srcs[i], dim);
       // Sum up all the sizes.
       sizes[dim] = builder.create<arith::AddIOp>(loc, sizes[dim], srcSz);
     }
   }
 }
 
 /// Generates an uninitialized buffer of the given size and type,
 /// but returns it as type `memref<? x $tp>` (rather than as type
genSparse2SparseReshape(ReshapeOp op, typename ReshapeOp::Adaptor adaptor,
                         ConversionPatternRewriter &rewriter) {
   Location loc = op.getLoc();
   auto srcTp = op.getSrc().getType().template cast<RankedTensorType>();
   auto dstTp = op.getResult().getType().template cast<RankedTensorType>();
   auto encSrc = getSparseTensorEncoding(srcTp);
   auto encDst = getSparseTensorEncoding(dstTp);
   if (!encDst || !encSrc)
     return failure();
-
-  unsigned srcRank = srcTp.getRank();
-  unsigned dstRank = dstTp.getRank();
   Type elemTp = srcTp.getElementType();
   assert(elemTp == dstTp.getElementType() &&
          "reshape should not change element type");
   // Start an iterator over the source tensor (in original index order).
   auto noPerm = SparseTensorEncodingAttr::get(
       op->getContext(), encSrc.getDimLevelType(), AffineMap(), AffineMap(),
       encSrc.getPointerBitWidth(), encSrc.getIndexBitWidth());
-  SmallVector<Value> srcSizes;
-  sizesFromPtr(rewriter, srcSizes, loc, encSrc, srcTp, adaptor.getSrc());
+  SmallVector<Value> srcDimSizes =
+      getDimSizes(rewriter, loc, encSrc, srcTp, adaptor.getSrc());
   NewCallParams params(rewriter, loc);
-  Value iter = params.genBuffers(noPerm, srcSizes, srcTp)
+  Value iter = params.genBuffers(noPerm, srcDimSizes, srcTp)
                    .genNewCall(Action::kToIterator, adaptor.getSrc());
   // Start a new COO for the destination tensor.
-  SmallVector<Value> dstSizes;
-  if (dstTp.hasStaticShape()) {
-    sizesFromType(rewriter, dstSizes, loc, dstTp);
-  } else {
-    ArrayRef<int64_t> dstShape = dstTp.getShape();
-    genReshapeDstShape(loc, rewriter, dstSizes, srcSizes, dstShape,
-                       op.getReassociationIndices());
-  }
-  Value coo =
-      params.genBuffers(encDst, dstSizes, dstTp).genNewCall(Action::kEmptyCOO);
+  SmallVector<Value> dstDimSizes;
+  if (dstTp.hasStaticShape())
+    // Static "shapes" are in fact "sizes".
+    fillDimShape(rewriter, loc, dstTp, dstDimSizes);
+  else
+    genReshapeDstShape(loc, rewriter, dstDimSizes, srcDimSizes,
+                       dstTp.getShape(), op.getReassociationIndices());
+  Value coo = params.genBuffers(encDst, dstDimSizes, dstTp)
+                  .genNewCall(Action::kEmptyCOO);
   Value dstPerm = params.getDim2LvlMap();
   // Construct a while loop over the iterator.
-  Value srcIdx = genAlloca(rewriter, loc, srcRank, rewriter.getIndexType());
-  Value dstIdx = genAlloca(rewriter, loc, dstRank, rewriter.getIndexType());
+  Type iTp = rewriter.getIndexType();
+  Value srcIdx = genAlloca(rewriter, loc, srcTp.getRank(), iTp);
+  Value dstIdx = genAlloca(rewriter, loc, dstTp.getRank(), iTp);
   Value elemPtr = genAllocaScalar(rewriter, loc, elemTp);
   SmallVector<Value> noArgs;
   SmallVector<Type> noTypes;
   auto whileOp = rewriter.create<scf::WhileOp>(loc, noTypes, noArgs);
   Block *before = rewriter.createBlock(&whileOp.getBefore(), {}, noTypes);
   rewriter.setInsertionPointToEnd(before);
   Value cond = genGetNextCall(rewriter, loc, iter, srcIdx, elemPtr);
   rewriter.create<scf::ConditionOp>(loc, cond, before->getArguments());
   // Translate indices from source to target and insert. Note that we do
   // not need to store the value in elemPtr, as the value is still there.
   Block *after = rewriter.createBlock(&whileOp.getAfter(), {}, noTypes);
   rewriter.setInsertionPointToStart(after);
   translateIndices(loc, rewriter, op.getReassociationIndices(), dstTp, srcTp,
-                   dstIdx, srcIdx, dstSizes, srcSizes);
+                   dstIdx, srcIdx, dstDimSizes, srcDimSizes);
   genAddEltCall(rewriter, loc, elemTp, coo, elemPtr, dstIdx, dstPerm);
   rewriter.create<scf::YieldOp>(loc);
   // Final call to construct sparse tensor storage and free temporary resources.
   rewriter.setInsertionPointAfter(whileOp);
   Value dst = params.genNewCall(Action::kFromCOO, coo);
   genDelCOOCall(rewriter, loc, elemTp, coo);
   genDelIteratorCall(rewriter, loc, elemTp, iter);
   rewriter.replaceOp(op, dst);
static void genSparseCOOIterationLoop(
 
   unsigned rank = tensorTp.getRank();
   Type elemTp = tensorTp.getElementType();
 
   // Start an iterator over the tensor (in original index order).
   auto noPerm = SparseTensorEncodingAttr::get(
       rewriter.getContext(), enc.getDimLevelType(), AffineMap(), AffineMap(),
       enc.getPointerBitWidth(), enc.getIndexBitWidth());
-  SmallVector<Value> sizes;
-  sizesFromPtr(rewriter, sizes, loc, noPerm, tensorTp, t);
+  SmallVector<Value> dimSizes = getDimSizes(rewriter, loc, noPerm, tensorTp, t);
   Value iter = NewCallParams(rewriter, loc)
-                   .genBuffers(noPerm, sizes, tensorTp)
+                   .genBuffers(noPerm, dimSizes, tensorTp)
                    .genNewCall(Action::kToIterator, t);
 
   // Construct a while loop over the iterator.
   Value srcIdx = genAlloca(rewriter, loc, rank, rewriter.getIndexType());
   Value elemPtr = genAllocaScalar(rewriter, loc, elemTp);
   SmallVector<Value> noArgs;
   SmallVector<Type> noTypes;
   auto whileOp = rewriter.create<scf::WhileOp>(loc, noTypes, noArgs);
public:
   LogicalResult
   matchAndRewrite(func::ReturnOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     rewriter.replaceOpWithNewOp<func::ReturnOp>(op, adaptor.getOperands());
     return success();
   }
 };
 
-/// Sparse conversion rule for dimension accesses.
+/// Sparse conversion rule for accessing dimension-sizes.
 class SparseTensorToDimSizeConverter
     : public OpConversionPattern<tensor::DimOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(tensor::DimOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
-    // Only rewrite annotated DimOp with constant index.
-    auto enc = getSparseTensorEncoding(op.getSource().getType());
+    auto stp = op.getSource().getType().cast<ShapedType>();
+    // Only rewrite sparse DimOp.
+    auto enc = getSparseTensorEncoding(stp);
     if (!enc)
       return failure();
-    Optional<int64_t> index = op.getConstantIndex();
-    if (!index)
+    // Only rewrite DimOp with constant index.
+    Optional<int64_t> dim = op.getConstantIndex();
+    if (!dim)
       return failure();
     // Generate the call.
     Value src = adaptor.getOperands()[0];
-    int64_t idx = *index;
-    rewriter.replaceOp(op,
-                       genLvlSizeCall(rewriter, op->getLoc(), enc, src, idx));
+    rewriter.replaceOp(
+        op, createOrFoldDimCall(rewriter, op->getLoc(), enc, stp, src, *dim));
     return success();
   }
 };
 
 /// Sparse conversion rule for trivial tensor casts.
 class SparseCastConverter : public OpConversionPattern<tensor::CastOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
matchAndRewrite(NewOp op, OpAdaptor adaptor,
     auto stp = op.getType().cast<ShapedType>();
     auto enc = getSparseTensorEncoding(stp);
     if (!enc)
       return failure();
     const unsigned dimRank = stp.getRank();
     const unsigned lvlRank = enc.getDimLevelType().size();
     // Construct the dimShape.
     const auto dimShape = stp.getShape();
-    SmallVector<Value> dimShapeValues;
-    sizesFromType(rewriter, dimShapeValues, loc, stp);
+    SmallVector<Value> dimShapeValues = getDimShape(rewriter, loc, stp);
     Value dimShapeBuffer = genBuffer(rewriter, loc, dimShapeValues);
     // Allocate `SparseTensorReader` and perform all initial setup that
     // does not depend on lvlSizes (nor dim2lvl, lvl2dim, etc).
     Type opaqueTp = getOpaquePointerType(rewriter);
     Value valTp =
         constantPrimaryTypeEncoding(rewriter, loc, stp.getElementType());
     Value reader =
         createFuncCall(rewriter, loc, "createCheckedSparseTensorReader",
if (encDst && encSrc) {
       //   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.
       if (encDst == encSrc) {
         rewriter.replaceOp(op, adaptor.getOperands()); // hidden nop cast
         return success();
       }
-      SmallVector<Value> sizes;
       NewCallParams params(rewriter, loc);
       ShapedType stp = srcType.cast<ShapedType>();
-      sizesFromPtr(rewriter, sizes, loc, encSrc, stp, src);
+      SmallVector<Value> dimSizes =
+          getDimSizes(rewriter, loc, encSrc, stp, src);
       bool useDirectConversion;
       switch (options.sparseToSparseStrategy) {
       case SparseToSparseConversionStrategy::kViaCOO:
         useDirectConversion = false;
         break;
       case SparseToSparseConversionStrategy::kDirect:
         useDirectConversion = true;
         assert(canUseDirectConversion(encDst.getDimLevelType()) &&
                "Unsupported target for direct sparse-to-sparse conversion");
         break;
       case SparseToSparseConversionStrategy::kAuto:
         useDirectConversion = canUseDirectConversion(encDst.getDimLevelType());
         break;
       }
       if (useDirectConversion) {
-        rewriter.replaceOp(op, params.genBuffers(encDst, sizes, stp)
+        rewriter.replaceOp(op, params.genBuffers(encDst, dimSizes, stp)
                                    .genNewCall(Action::kSparseToSparse, src));
       } else { // use via-COO conversion.
         // Set up encoding with right mix of src and dst so that the two
         // method calls can share most parameters, while still providing
         // the correct sparsity information to either of them.
         auto enc = SparseTensorEncodingAttr::get(
             op->getContext(), encDst.getDimLevelType(), encDst.getDimOrdering(),
             encDst.getHigherOrdering(), encSrc.getPointerBitWidth(),
             encSrc.getIndexBitWidth());
         // TODO: This is the only place where `kToCOO` (or `kToIterator`)
         // is called with a non-identity permutation.  Is there any clean
         // way to push the permutation over to the `kFromCOO` side instead?
-        Value coo =
-            params.genBuffers(enc, sizes, stp).genNewCall(Action::kToCOO, src);
+        Value coo = params.genBuffers(enc, dimSizes, stp)
+                        .genNewCall(Action::kToCOO, src);
         Value dst = params.setTemplateTypes(encDst, stp)
                         .genNewCall(Action::kFromCOO, coo);
         genDelCOOCall(rewriter, loc, stp.getElementType(), coo);
         rewriter.replaceOp(op, dst);
       }
       return success();
     }
     if (!encDst && encSrc) {
if (!encDst && encSrc) {
       Type elemTp = dstTensorTp.getElementType();
       // Fabricate a no-permutation encoding for NewCallParams
       // The pointer/index types must be those of `src`.
       // The dimLevelTypes aren't actually used by Action::kToIterator.
       encDst = SparseTensorEncodingAttr::get(
           op->getContext(),
           SmallVector<DimLevelType>(rank, DimLevelType::Dense), AffineMap(),
           AffineMap(), encSrc.getPointerBitWidth(), encSrc.getIndexBitWidth());
-      SmallVector<Value> sizes;
-      sizesFromPtr(rewriter, sizes, loc, encSrc, srcTensorTp, src);
+      SmallVector<Value> dimSizes =
+          getDimSizes(rewriter, loc, encSrc, srcTensorTp, src);
       Value iter = NewCallParams(rewriter, loc)
-                       .genBuffers(encDst, sizes, dstTensorTp)
+                       .genBuffers(encDst, dimSizes, dstTensorTp)
                        .genNewCall(Action::kToIterator, src);
       Value ind = genAlloca(rewriter, loc, rank, rewriter.getIndexType());
       Value elemPtr = genAllocaScalar(rewriter, loc, elemTp);
       Block *insertionBlock = rewriter.getInsertionBlock();
       // TODO: Dense buffers should be allocated/deallocated via the callback
       // in BufferizationOptions.
-      Value dst = allocDenseTensor(rewriter, loc, dstTensorTp, sizes);
+      Value dst = allocDenseTensor(rewriter, loc, dstTensorTp, dimSizes);
       SmallVector<Value> noArgs;
       SmallVector<Type> noTypes;
       auto whileOp = rewriter.create<scf::WhileOp>(loc, noTypes, noArgs);
       Block *before = rewriter.createBlock(&whileOp.getBefore(), {}, noTypes);
       rewriter.setInsertionPointToEnd(before);
       Value cond = genGetNextCall(rewriter, loc, iter, ind, elemPtr);
       rewriter.create<scf::ConditionOp>(loc, cond, before->getArguments());
       Block *after = rewriter.createBlock(&whileOp.getAfter(), {}, noTypes);
matchAndRewrite(ExpandOp op, OpAdaptor adaptor,
     Location loc = op->getLoc();
     RankedTensorType srcType =
         op.getTensor().getType().cast<RankedTensorType>();
     Type eltType = srcType.getElementType();
     Type boolType = rewriter.getIntegerType(1);
     Type idxType = rewriter.getIndexType();
     // All initialization should be done on entry of the loop nest.
     rewriter.setInsertionPointAfter(op.getTensor().getDefiningOp());
-    // Determine the size for access expansion (always the innermost stored
-    // dimension size, translated back to original dimension).
-    auto enc = getSparseTensorEncoding(srcType);
-    unsigned innerDim = toOrigDim(srcType, srcType.getRank() - 1);
-    auto sz = sizeFromPtrAtDim(rewriter, loc, enc, srcType, adaptor.getTensor(),
-                               innerDim);
+    // Get the cardinality of valid coordinates for the innermost level.
+    auto srcEnc = getSparseTensorEncoding(srcType);
+    unsigned lvlRank =
+        srcEnc ? srcEnc.getDimLevelType().size() : srcType.getRank();
+    Value sz = createOrFoldLvlCall(rewriter, loc, srcEnc, srcType,
+                                   adaptor.getTensor(), lvlRank - 1);
     // Allocate temporary buffers for values, filled-switch, and indices.
     // We do not use stack buffers for this, since the expanded size may
     // be rather large (as it envelops a single expanded dense dimension).
     Value values = genAlloc(rewriter, loc, sz, eltType);
     Value filled = genAlloc(rewriter, loc, sz, boolType);
     Value indices = genAlloc(rewriter, loc, sz, idxType);
     Value zero = constantZero(rewriter, loc, idxType);
     // Reset the values/filled-switch to all-zero/false. Note that this
for (auto it : llvm::zip(op.getInputs(), adaptor.getInputs())) {
                 indVec[concatDim] = builder.create<arith::AddIOp>(
                     loc, indVec[concatDim], offset);
                 builder.create<memref::StoreOp>(loc, val, dst, indVec);
               }
             });
       }
       // Accumulate offset.
       // TODO: avoid calling sparseDimSize multiple times by caching the result!
-      Value curDim = encSrc ? sizeFromPtrAtDim(rewriter, loc, encSrc, srcTp,
-                                               adaptedOp, concatDim)
-                            : linalg::createOrFoldDimOp(rewriter, loc,
-                                                        adaptedOp, concatDim);
+      Value curDim = createOrFoldDimCall(rewriter, loc, encSrc, srcTp,
 
       offset = rewriter.create<arith::AddIOp>(loc, offset, curDim);
     }
     if (encDst) {
       // In sparse output case, the destination holds the COO.
       Value coo = dst;
       dst = params.genNewCall(Action::kFromCOO, coo);
       // Release resources.
public:
   LogicalResult
   matchAndRewrite(OutOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Location loc = op->getLoc();
     ShapedType srcType = op.getTensor().getType().cast<ShapedType>();
     // Convert to default permuted COO.
     Value src = adaptor.getOperands()[0];
     auto encSrc = getSparseTensorEncoding(srcType);
-    SmallVector<Value> sizes;
-    sizesFromPtr(rewriter, sizes, loc, encSrc, srcType, src);
+    SmallVector<Value> dimSizes =
+        getDimSizes(rewriter, loc, encSrc, srcType, src);
     auto enc = SparseTensorEncodingAttr::get(
         op->getContext(), encSrc.getDimLevelType(), AffineMap(), AffineMap(),
         encSrc.getPointerBitWidth(), encSrc.getIndexBitWidth());
     Value coo = NewCallParams(rewriter, loc)
-                    .genBuffers(enc, sizes, srcType)
+                    .genBuffers(enc, dimSizes, srcType)
                     .genNewCall(Action::kToCOO, src);
     // Then output the tensor to external file with indices in the externally
     // visible lexicographic index order. A sort is required if the source was
     // not in that order yet (note that the sort can be dropped altogether if
     // external format does not care about the order at all, but here we assume
     // it does).
     Value sort = constantI1(rewriter, loc,
                             encSrc.getDimOrdering() &&

mlir/lib/ExecutionEngine/SparseTensorRuntime.cpp

 
 //===----------------------------------------------------------------------===//
 //
 // Public functions which accept only C-style data structures to interact
 // with sparse tensors (which are only visible as opaque pointers externally).
 //
 //===----------------------------------------------------------------------===//
 
-index_type sparseLvlSize(void *tensor, index_type x) {
-  return static_cast<SparseTensorStorageBase *>(tensor)->getLvlSize(x);
+index_type sparseLvlSize(void *tensor, index_type l) {
+  return static_cast<SparseTensorStorageBase *>(tensor)->getLvlSize(l);
+}
+
+index_type sparseDimSize(void *tensor, index_type d) {
+  return static_cast<SparseTensorStorageBase *>(tensor)->getDimSize(d);
 }
 
 void endInsert(void *tensor) {
   return static_cast<SparseTensorStorageBase *>(tensor)->endInsert();
 }
 
 #define IMPL_OUTSPARSETENSOR(VNAME, V)                                         \
   void outSparseTensor##VNAME(void *coo, void *dest, bool sort) {              \

mlir/test/Dialect/SparseTensor/conversion.mlir

 //       CHECK: return %[[A]] : !llvm.ptr<i8>
 func.func @sparse_nop(%arg0: tensor<?xf64, #SparseVector>) -> tensor<?xf64, #SparseVector> {
   return %arg0 : tensor<?xf64, #SparseVector>
 }
 
 // CHECK-LABEL: func @sparse_dim1d(
 //  CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>)
 //       CHECK: %[[C:.*]] = arith.constant 0 : index
-//       CHECK: %[[D:.*]] = call @sparseLvlSize(%[[A]], %[[C]])
+//       CHECK: %[[D:.*]] = call @sparseDimSize(%[[A]], %[[C]])
 //       CHECK: return %[[D]] : index
 func.func @sparse_dim1d(%arg0: tensor<?xf64, #SparseVector>) -> index {
   %c = arith.constant 0 : index
   %0 = tensor.dim %arg0, %c : tensor<?xf64, #SparseVector>
   return %0 : index
 }
 
+// Querying the size of dimension 1 should do so; i.e., it should
+// not be permuted into a query for the size of level 2 (even though
+// dimension 1 is stored as level 2).
 // CHECK-LABEL: func @sparse_dim3d(
 //  CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>)
-//       CHECK: %[[C:.*]] = arith.constant 2 : index
-//       CHECK: %[[D:.*]] = call @sparseLvlSize(%[[A]], %[[C]])
+//       CHECK: %[[C:.*]] = arith.constant 1 : index
+//       CHECK: %[[D:.*]] = call @sparseDimSize(%[[A]], %[[C]])
 //       CHECK: return %[[D]] : index
 func.func @sparse_dim3d(%arg0: tensor<?x?x?xf64, #SparseTensor>) -> index {
-  // Querying for dimension 1 in the tensor type needs to be
-  // permuted into querying for dimension 2 in the stored sparse
-  // tensor scheme, since the latter honors the dimOrdering.
   %c = arith.constant 1 : index
   %0 = tensor.dim %arg0, %c : tensor<?x?x?xf64, #SparseTensor>
   return %0 : index
 }
 
+// Querying the size of a static dimension should be folded into a
+// constant (and we should be sure to get the size of dimension 1,
+// not dimension 2 nor level 1).
 // CHECK-LABEL: func @sparse_dim3d_const(
 //  CHECK-SAME: %[[A:.*]]: !llvm.ptr<i8>)
 //       CHECK: %[[C:.*]] = arith.constant 20 : index
 //       CHECK: return %[[C]] : index
 func.func @sparse_dim3d_const(%arg0: tensor<10x20x30xf64, #SparseTensor>) -> index {
-  // Querying for dimension 1 in the tensor type can be directly
-  // folded into the right value (even though it corresponds
-  // 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: %[[DimShape0:.*]] = memref.alloca() : memref<1xindex>
func.func @sparse_expansion2() -> memref<?xindex> {
   %values, %filled, %added, %count = sparse_tensor.expand %0
     : tensor<4x8xf64, #CSC> to memref<?xf64>, memref<?xi1>, memref<?xindex>
   return %added : memref<?xindex>
 }
 
 // CHECK-LABEL: func @sparse_expansion3(
 //       CHECK: %[[C1:.*]] = arith.constant 1 : index
 //       CHECK: %[[N:.*]] = call @newSparseTensor
-//       CHECK: %[[S:.*]] = call @sparseLvlSize(%[[N]], %c1) : (!llvm.ptr<i8>, index) -> index
+//       CHECK: %[[S:.*]] = call @sparseLvlSize(%[[N]], %[[C1]])
 //       CHECK: %[[A:.*]] = memref.alloc(%[[S]]) : memref<?xf64>
 //       CHECK: %[[B:.*]] = memref.alloc(%[[S]]) : memref<?xi1>
 //       CHECK: %[[C:.*]] = memref.alloc(%[[S]]) : memref<?xindex>
 //   CHECK-DAG: linalg.fill ins(%{{.*}} : f64) outs(%[[A]] : memref<?xf64>)
 //   CHECK-DAG: linalg.fill ins(%{{.*}} : i1) outs(%[[B]] : memref<?xi1>)
 //       CHECK: return %[[C]] : memref<?xindex>
 func.func @sparse_expansion3(%arg0: index, %arg1: index) -> memref<?xindex> {
   %0 = bufferization.alloc_tensor(%arg0, %arg1) : tensor<?x?xf64, #CSC>

mlir/test/Dialect/SparseTensor/convert_sparse2dense.mlir

 //   CHECK-DAG: %[[C0:.*]] = arith.constant 0 : i32
 //   CHECK-DAG: %[[C6:.*]] = arith.constant 6 : i32
 //
 //   CHECK-DAG: %[[LvlTypes:.*]] = memref.alloca() : memref<1xi8>
 //   CHECK-DAG: %[[LvlTypesP:.*]] = memref.cast %[[LvlTypes]] : memref<1xi8> to memref<?xi8>
 //   CHECK-DAG: memref.store %[[DenseDLT]], %[[LvlTypes]][%[[I0]]] : memref<1xi8>
 //   CHECK-DAG: %[[DimSizes:.*]] = memref.alloca() : memref<1xindex>
 //   CHECK-DAG: %[[DimSizesP:.*]] = memref.cast %[[DimSizes]] : memref<1xindex> to memref<?xindex>
-//   CHECK-DAG: %[[SizeI0:.*]] = call @sparseLvlSize(%[[Arg]], %[[I0]]) : (!llvm.ptr<i8>, index) -> index
+//   CHECK-DAG: %[[SizeI0:.*]] = call @sparseDimSize(%[[Arg]], %[[I0]]) : (!llvm.ptr<i8>, index) -> index
 //   CHECK-DAG: memref.store %[[SizeI0]], %[[DimSizes]][%[[I0]]] : memref<1xindex>
 //   CHECK-DAG: %[[LvlSizes:.*]] = memref.alloca() : memref<1xindex>
 //   CHECK-DAG: %[[LvlSizesP:.*]] = memref.cast %[[LvlSizes]] : memref<1xindex> to memref<?xindex>
 //   CHECK-DAG: %[[Iota:.*]] = memref.alloca() : memref<1xindex>
 //   CHECK-DAG: %[[IotaP:.*]] = memref.cast %[[Iota]] : memref<1xindex> to memref<?xindex>
 //   CHECK-DAG: memref.store %[[I0]], %[[Iota]][%[[I0]]] : memref<1xindex>
 //       CHECK: %[[Iter:.*]] = call @newSparseTensor(%[[DimSizesP]], %[[LvlSizesP]], %[[LvlTypesP]], %[[IotaP]], %[[IotaP]], %[[C0]], %[[C0]], %[[C6]], %[[C6]], %[[Arg]])
 //

 //   CHECK-DAG: %[[E0:.*]] = arith.constant 0.000000e+00 : f64
 //
 //   CHECK-DAG: %[[LvlTypes:.*]] = memref.alloca() : memref<2xi8>
 //   CHECK-DAG: %[[LvlTypesP:.*]] = memref.cast %[[LvlTypes]] : memref<2xi8> to memref<?xi8>
 //   CHECK-DAG: memref.store %[[DenseDLT]], %[[LvlTypes]][%[[I0]]] : memref<2xi8>
 //   CHECK-DAG: memref.store %[[DenseDLT]], %[[LvlTypes]][%[[I1]]] : memref<2xi8>
 //   CHECK-DAG: %[[DimSizes:.*]] = memref.alloca() : memref<2xindex>
 //   CHECK-DAG: %[[DimSizesP:.*]] = memref.cast %[[DimSizes]] : memref<2xindex> to memref<?xindex>
-//   CHECK-DAG: %[[SizeI0:.*]] = call @sparseLvlSize(%[[Arg]], %[[I0]]) : (!llvm.ptr<i8>, index) -> index
+//   CHECK-DAG: %[[SizeI0:.*]] = call @sparseDimSize(%[[Arg]], %[[I0]]) : (!llvm.ptr<i8>, index) -> index
 //   CHECK-DAG: memref.store %[[SizeI0]], %[[DimSizes]][%[[I0]]] : memref<2xindex>
 //   CHECK-DAG: memref.store %[[I4]], %[[DimSizes]][%[[I1]]] : memref<2xindex>
 //   CHECK-DAG: %[[LvlSizes:.*]] = memref.alloca() : memref<2xindex>
 //   CHECK-DAG: %[[LvlSizesP:.*]] = memref.cast %[[LvlSizes]] : memref<2xindex> to memref<?xindex>
 //   CHECK-DAG: %[[Iota:.*]] = memref.alloca() : memref<2xindex>
 //   CHECK-DAG: %[[IotaP:.*]] = memref.cast %[[Iota]] : memref<2xindex> to memref<?xindex>
 //   CHECK-DAG: memref.store %[[I0]], %[[Iota]][%[[I0]]] : memref<2xindex>
 //   CHECK-DAG: memref.store %[[I1]], %[[Iota]][%[[I1]]] : memref<2xindex>

 //   CHECK-DAG: %[[E0:.*]] = arith.constant 0.000000e+00 : f64
 //
 //   CHECK-DAG: %[[LvlTypes:.*]] = memref.alloca() : memref<2xi8>
 //   CHECK-DAG: %[[LvlTypesP:.*]] = memref.cast %[[LvlTypes]] : memref<2xi8> to memref<?xi8>
 //   CHECK-DAG: memref.store %[[DenseDLT]], %[[LvlTypes]][%[[I0]]] : memref<2xi8>
 //   CHECK-DAG: memref.store %[[DenseDLT]], %[[LvlTypes]][%[[I1]]] : memref<2xi8>
 //   CHECK-DAG: %[[DimSizes:.*]] = memref.alloca() : memref<2xindex>
 //   CHECK-DAG: %[[DimSizesP:.*]] = memref.cast %[[DimSizes]] : memref<2xindex> to memref<?xindex>
-//   CHECK-DAG: %[[SizeI1:.*]] = call @sparseLvlSize(%[[Arg]], %[[I1]]) : (!llvm.ptr<i8>, index) -> index
+//   CHECK-DAG: %[[SizeI1:.*]] = call @sparseDimSize(%[[Arg]], %[[I1]]) : (!llvm.ptr<i8>, index) -> index
 //   CHECK-DAG: memref.store %[[I2]], %[[DimSizes]][%[[I0]]] : memref<2xindex>
 //   CHECK-DAG: memref.store %[[SizeI1]], %[[DimSizes]][%[[I1]]] : memref<2xindex>
 //   CHECK-DAG: %[[LvlSizes:.*]] = memref.alloca() : memref<2xindex>
 //   CHECK-DAG: %[[LvlSizesP:.*]] = memref.cast %[[LvlSizes]] : memref<2xindex> to memref<?xindex>
 //   CHECK-DAG: %[[Iota:.*]] = memref.alloca() : memref<2xindex>
 //   CHECK-DAG: %[[IotaP:.*]] = memref.cast %[[Iota]] : memref<2xindex> to memref<?xindex>
 //   CHECK-DAG: memref.store %[[I0]], %[[Iota]][%[[I0]]] : memref<2xindex>
 //   CHECK-DAG: memref.store %[[I1]], %[[Iota]][%[[I1]]] : memref<2xindex>

 //   CHECK-DAG: %[[E0:.*]] = arith.constant 0.000000e+00 : f64
 //
 //   CHECK-DAG: %[[LvlTypes:.*]] = memref.alloca() : memref<2xi8>
 //   CHECK-DAG: %[[LvlTypesP:.*]] = memref.cast %[[LvlTypes]] : memref<2xi8> to memref<?xi8>
 //   CHECK-DAG: memref.store %[[DenseDLT]], %[[LvlTypes]][%[[I0]]] : memref<2xi8>
 //   CHECK-DAG: memref.store %[[DenseDLT]], %[[LvlTypes]][%[[I1]]] : memref<2xi8>
 //   CHECK-DAG: %[[DimSizes:.*]] = memref.alloca() : memref<2xindex>
 //   CHECK-DAG: %[[DimSizesP:.*]] = memref.cast %[[DimSizes]] : memref<2xindex> to memref<?xindex>
-//   CHECK-DAG: %[[SizeI0:.*]] = call @sparseLvlSize(%[[Arg]], %[[I0]]) : (!llvm.ptr<i8>, index) -> index
-//   CHECK-DAG: %[[SizeI1:.*]] = call @sparseLvlSize(%[[Arg]], %[[I1]]) : (!llvm.ptr<i8>, index) -> index
+//   CHECK-DAG: %[[SizeI0:.*]] = call @sparseDimSize(%[[Arg]], %[[I0]]) : (!llvm.ptr<i8>, index) -> index
+//   CHECK-DAG: %[[SizeI1:.*]] = call @sparseDimSize(%[[Arg]], %[[I1]]) : (!llvm.ptr<i8>, index) -> index
 //   CHECK-DAG: memref.store %[[SizeI0]], %[[DimSizes]][%[[I0]]] : memref<2xindex>
 //   CHECK-DAG: memref.store %[[SizeI1]], %[[DimSizes]][%[[I1]]] : memref<2xindex>
 //   CHECK-DAG: %[[LvlSizes:.*]] = memref.alloca() : memref<2xindex>
 //   CHECK-DAG: %[[LvlSizesP:.*]] = memref.cast %[[LvlSizes]] : memref<2xindex> to memref<?xindex>
 //   CHECK-DAG: %[[Iota:.*]] = memref.alloca() : memref<2xindex>
 //   CHECK-DAG: %[[IotaP:.*]] = memref.cast %[[Iota]] : memref<2xindex> to memref<?xindex>
 //   CHECK-DAG: memref.store %[[I0]], %[[Iota]][%[[I0]]] : memref<2xindex>
 //   CHECK-DAG: memref.store %[[I1]], %[[Iota]][%[[I1]]] : memref<2xindex>

mlir/test/Dialect/SparseTensor/sparse_expand.mlir

 // CHECK-SPARSE:   scf.for {{.*}} {
 // CHECK-SPARSE:   }
 // CHECK-SPARSE: }
 // CHECK-SPARSE: sparse_tensor.compress %[[A]], %[[B]], %[[C]], %[[COUNT]] into
 // CHECK-SPARSE: %[[RET:.*]] = sparse_tensor.load %{{.*}} hasInserts
 // CHECK-SPARSE: return %[[RET]]
 //
 // CHECK-CONVERT-LABEL: func @kernel(
-// CHECK-CONVERT: %[[C:.*]] = arith.constant 0 : index
-// CHECK-CONVERT: %{{.*}} = call @sparseLvlSize
-// CHECK-CONVERT: %[[N:.*]] = call @newSparseTensor
-// CHECK-CONVERT: %[[S:.*]] = call @sparseLvlSize(%[[N]], %[[C]])
+// CHECK-CONVERT-SAME: %[[A:.*]]: !llvm.ptr<i8>) -> !llvm.ptr<i8>
+// CHECK-CONVERT: %[[C0:.*]] = arith.constant 0 : index
+// CHECK-CONVERT: %[[N:.*]] = call @sparseDimSize(%[[A]], %[[C0]])
+// CHECK-CONVERT: %[[V:.*]] = call @newSparseTensor
+// CHECK-CONVERT: %[[S:.*]] = call @sparseLvlSize(%[[V]], %[[C0]])
 // CHECK-CONVERT: %[[A:.*]] = memref.alloc(%[[S]]) : memref<?xf64>
 // CHECK-CONVERT: %[[B:.*]] = memref.alloc(%[[S]]) : memref<?xi1>
 // CHECK-CONVERT: %[[C:.*]] = memref.alloc(%[[S]]) : memref<?xindex>
 // CHECK-CONVERT: linalg.fill ins(%{{.*}} : f64) outs(%[[A]] : memref<?xf64>)
 // CHECK-CONVERT: linalg.fill ins(%{{.*}} : i1) outs(%[[B]] : memref<?xi1>)
 // CHECK-CONVERT: scf.for {{.*}} {
 // CHECK-CONVERT:   scf.for {{.*}} {
 // CHECK-CONVERT:   }

mlir/test/Dialect/SparseTensor/sparse_perm_lower.mlir

 // CHECK-HIR:             scf.yield %[[VAL_15]] : f32
 // CHECK-HIR:           }
 // CHECK-HIR:           memref.store %[[VAL_12]], %[[VAL_10]][] : memref<f32>
 // CHECK-HIR:           %[[VAL_30:.*]] = bufferization.to_tensor %[[VAL_10]] : memref<f32>
 // CHECK-HIR:           return %[[VAL_30]] : tensor<f32>
 // CHECK-HIR:         }
 //
 // CHECK-MIR-LABEL:   func @sparse_dynamic_dims(
-// CHECK-MIR-SAME:      %[[VAL_0:.*]]: !llvm.ptr<i8>,
-// CHECK-MIR-SAME:      %[[VAL_1:.*]]: tensor<f32>) -> tensor<f32> {
-// CHECK-MIR-DAG:       %[[VAL_2:.*]] = arith.constant 2 : index
-// CHECK-MIR-DAG:       %[[VAL_3:.*]] = arith.constant 1 : index
-// CHECK-MIR-DAG:       %[[VAL_4:.*]] = arith.constant 0 : index
-// CHECK-MIR-DAG:       %[[VAL_5:.*]] = call @sparseLvlSize(%[[VAL_0]], %[[VAL_4]]) : (!llvm.ptr<i8>, index) -> index
-// CHECK-MIR-DAG:       %[[VAL_6:.*]] = call @sparseLvlSize(%[[VAL_0]], %[[VAL_3]]) : (!llvm.ptr<i8>, index) -> index
-// CHECK-MIR-DAG:       %[[VAL_7:.*]] = call @sparseLvlSize(%[[VAL_0]], %[[VAL_2]]) : (!llvm.ptr<i8>, index) -> index
-// CHECK-MIR-DAG:       %[[VAL_8:.*]] = call @sparseValuesF32(%[[VAL_0]]) : (!llvm.ptr<i8>) -> memref<?xf32>
-// CHECK-MIR-DAG:       %[[VAL_10:.*]] = bufferization.to_memref %[[VAL_1]] : memref<f32>
-// CHECK-MIR:           %[[VAL_11:.*]] = tensor.extract %[[VAL_1]][] : tensor<f32>
-// CHECK-MIR:           %[[VAL_12:.*]] = scf.for %[[VAL_13:.*]] = %[[VAL_4]] to %[[VAL_5]] step %[[VAL_3]] iter_args(%[[VAL_14:.*]] = %[[VAL_11]]) -> (f32) {
-// CHECK-MIR:             %[[VAL_15:.*]] = scf.for %[[VAL_16:.*]] = %[[VAL_4]] to %[[VAL_6]] step %[[VAL_3]] iter_args(%[[VAL_17:.*]] = %[[VAL_14]]) -> (f32) {
-// CHECK-MIR:               %[[VAL_18:.*]] = arith.muli %[[VAL_6]], %[[VAL_13]] : index
-// CHECK-MIR:               %[[VAL_19:.*]] = arith.addi %[[VAL_18]], %[[VAL_16]] : index
-// CHECK-MIR:               %[[VAL_20:.*]] = scf.for %[[VAL_21:.*]] = %[[VAL_4]] to %[[VAL_7]] step %[[VAL_3]] iter_args(%[[VAL_22:.*]] = %[[VAL_17]]) -> (f32) {
-// CHECK-MIR:                 %[[VAL_23:.*]] = arith.muli %[[VAL_7]], %[[VAL_19]] : index
-// CHECK-MIR:                 %[[VAL_24:.*]] = arith.addi %[[VAL_23]], %[[VAL_21]] : index
+// CHECK-MIR-SAME:      %[[ARGA:.*]]: !llvm.ptr<i8>,
+// CHECK-MIR-SAME:      %[[ARGX:.*]]: tensor<f32>) -> tensor<f32> {
+// CHECK-MIR-DAG:       %[[I0:.*]] = arith.constant 0 : index
+// CHECK-MIR-DAG:       %[[I1:.*]] = arith.constant 1 : index
+// CHECK-MIR-DAG:       %[[I2:.*]] = arith.constant 2 : index
+// CHECK-MIR-DAG:       %[[DimSize0:.*]] = call @sparseDimSize(%[[ARGA]], %[[I0]])
+// CHECK-MIR-DAG:       %[[DimSize1:.*]] = call @sparseDimSize(%[[ARGA]], %[[I1]])
+// CHECK-MIR-DAG:       %[[DimSize2:.*]] = call @sparseDimSize(%[[ARGA]], %[[I2]])
+// CHECK-MIR-DAG:       %[[VAL_8:.*]] = call @sparseValuesF32(%[[ARGA]]) : (!llvm.ptr<i8>) -> memref<?xf32>
+// CHECK-MIR-DAG:       %[[VAL_10:.*]] = bufferization.to_memref %[[ARGX]] : memref<f32>
+// CHECK-MIR:           %[[VAL_11:.*]] = tensor.extract %[[ARGX]][] : tensor<f32>
+// CHECK-MIR:           %[[VAL_12:.*]] = scf.for %[[D2:.*]] = %[[I0]] to %[[DimSize2]] step %[[I1]] iter_args(%[[VAL_14:.*]] = %[[VAL_11]]) -> (f32) {
+// CHECK-MIR:             %[[VAL_15:.*]] = scf.for %[[D0:.*]] = %[[I0]] to %[[DimSize0]] step %[[I1]] iter_args(%[[VAL_17:.*]] = %[[VAL_14]]) -> (f32) {
+// CHECK-MIR:               %[[VAL_18:.*]] = arith.muli %[[DimSize0]], %[[D2]] : index
+// CHECK-MIR:               %[[VAL_19:.*]] = arith.addi %[[VAL_18]], %[[D0]] : index
+// CHECK-MIR:               %[[VAL_20:.*]] = scf.for %[[D1:.*]] = %[[I0]] to %[[DimSize1]] step %[[I1]] iter_args(%[[VAL_22:.*]] = %[[VAL_17]]) -> (f32) {
+// CHECK-MIR:                 %[[VAL_23:.*]] = arith.muli %[[DimSize1]], %[[VAL_19]] : index
+// CHECK-MIR:                 %[[VAL_24:.*]] = arith.addi %[[VAL_23]], %[[D1]] : index
 // CHECK-MIR:                 %[[VAL_25:.*]] = memref.load %[[VAL_8]]{{\[}}%[[VAL_24]]] : memref<?xf32>
 // CHECK-MIR:                 %[[VAL_26:.*]] = arith.addf %[[VAL_22]], %[[VAL_25]] : f32
 // CHECK-MIR:                 scf.yield %[[VAL_26]] : f32
 // CHECK-MIR:               }
 // CHECK-MIR:               scf.yield %[[VAL_20]] : f32
 // CHECK-MIR:             }
 // CHECK-MIR:             scf.yield %[[VAL_15]] : f32
 // CHECK-MIR:           }