[mlir][sparse] factored out a new DimLvlMapping class

[wrengr]

Summary

The new class generalizes dim<->lvl conversion logic, so that it can be reused for both the runtime and codegen paths. This abstraction is also the first step towards supporting non-trivial non-permutations.

[Peiming]

mlir/lib/Dialect/SparseTensor/Transforms/DimLvlMapping.h

196

Do we really need the cache? It does not seem like a complicated operation anyway, plus it might make the class expensive to copy.

305

I think requires user to pass in a loc on every call is more flexible, maybe the location need to be updated between calls (maybe builder as well).

[Peiming]

mlir/lib/Dialect/SparseTensor/Transforms/DimLvlMapping.h

238

You could probably avoid the virtual function here using CTRP

[Peiming]

mlir/lib/Dialect/SparseTensor/Transforms/DimLvlMapping.h

193–194

Why not call enc.getxxx for these two? Should be a cheap direct field queries.

[wrengr]

mlir/lib/Dialect/SparseTensor/Transforms/DimLvlMapping.h

196

At the moment the cache is trivial, because we're only handling permutations; but that'll change once I add the code for handling non-permutations. Considering how aggressively everything else in MLIR memoizes stuff, I don't think doing this is out of place. If/when we start needing to pass these around, then we can split the class into class DimLvlMappingImpl {...} vs class DimLvlMapping { DimLvlMappingImpl impl; } like is done everywhere else in MLIR.

305

So far all our code is structured to reuse the same loc/builder everywhere, so it's a lot cleaner to just pass them into the ctor rather than passing them into every method call.

[wrengr]

mlir/lib/Dialect/SparseTensor/Transforms/DimLvlMapping.h

193–194

Take a look at the ctor. These fields are memoized to avoid repeating the conditionals there, rather than to avoid repeating the getter methods.

238

Yeah, we could use CRTP instead. Personally I prefer virtual methods since they give much better error messages when things go wrong, but I'm not sure how bad the performance cost is for this case.

[wrengr]

Updating NewCallParams::genBuffers to take a reference rather than copying the DimLvlMapping

[aartbik]

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

127

Why this change? Is it possible to have no internal Attribute implementation?

mlir/lib/Dialect/SparseTensor/Transforms/DimLvlMapping.h

70

This really is a wrapper around a SparseTensorEncodingAttr and provides a consistent terminology and look into these. As such, do you think SparseTensor/IR is a better place for this (with "public" header). The codegen class below clearly belongs here, but I am wondering if moving this to a general place will enable a lot more simplification in the long run.

Otherwise, it seems a very heavy abstraction for "just" codeegen....

[wrengr]

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

127

Yes, it's possible not to have the impl (e.g., when the Attribute was build via the default-ctor).

The operational justification for making this change is because calling env.hasIdDimOrdering() will call env.getDimOrdering() which will call env.getImpl()->dimOrdering — which means that calling hasIdDimOrdering on the null-attribute will end up calling operator-> on the nullptr, which will crash.

The semantic justification for making this change is because we interpret dense-tensors (aka tensors will null-STEA) in the same way as sparse-tensors with the default STEA. And since the default dimOrdering field is the identity, thus we interpret dense-tensors as having the identity ordering as well.

My original motivation for making the change was so that we can just say enc.hasIdDimOrdering() in the DimLvlMapping ctor, rather than needing to phrase the check as (!enc || enc.hasIdDimOrdering()). [Back when I started this CL the hasIdDimOrdering method wasn't around yet, so I just had a standalone function that performs the same implementation as I've given the method in this version of the CL.] However, since making the change I've noticed there are several other places in the code that use the latter expression, which could now also be cleaned up to simply call the method.

mlir/lib/Dialect/SparseTensor/Transforms/DimLvlMapping.h

70

It's more a wrapper around RankedTensorType, since it includes information about the shape as well as the encoding. But yes, I agree.

Fwiw, I've been working on a rewrite of this CL since the currently posted version doesn't quite work out for being easily reused by the codegen pass (with the way the codegen pass is currently structured). And in that rewrite I've been moving this class towards being more and more of just a wrapper around RankedTensorType / ShapedType+STEA. In particular I've removed the stuff for memoizing and querying the levels, since the only places that actually need that are places that need the full DimLvlBuilder versions instead.

Once I get that variation finalized, I'd be happy to move it to the IR folder and make it public. One of the stumbling blocks I've been having, though, is what exactly to rename it to. My initial unfiltered thought would be to name it SparseTensorType (because of it's original motivation); but that's not a good name since it's specifically designed to abstract over the differences between sparse- and dense-tensor types, whereas that name seems to suggest that it's specifically for sparse-tensors alone. I suppose I could go for something like SparseLikeTensorType or SparseTensorAdaptorType etc, but those are a mouthful and also all feel a bit off... thoughts?

[wrengr]

mlir/lib/Dialect/SparseTensor/Transforms/DimLvlMapping.h

70

Yeahno, after tossing it about a bit more, I'm thinking SparseTensorType might be the best name —since it does things like renaming the various methods to make clear whether they're talking about dimensions vs levels. I'll just be sure to add some documentation clarifying the fact that it can and should be used for both dense- and sparse-tensors per se.

[wrengr]

This patch has been redesigned and replaced by other patches (D143800, D143946, D143949, D144052, et seq.)

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

return SparseTensorEncodingAttr::get(
       getPointerBitWidth(), getIndexBitWidth());
 }
 
 bool SparseTensorEncodingAttr::isAllDense() const {
   return llvm::all_of(getDimLevelType(), isDenseDLT);
 }
 
 bool SparseTensorEncodingAttr::hasIdDimOrdering() const {
-  return !getDimOrdering() || getDimOrdering().isIdentity();
+  return !getImpl() || !getDimOrdering() || getDimOrdering().isIdentity();
 }
 
 std::optional<uint64_t>
 SparseTensorEncodingAttr::getStaticDimSliceOffset(unsigned dim) const {
   return getDimSlices()[dim].getStaticOffset();
 }
 
 std::optional<uint64_t>

mlir/lib/Dialect/SparseTensor/Transforms/CMakeLists.txt

 add_mlir_dialect_library(MLIRSparseTensorTransforms
   BufferizableOpInterfaceImpl.cpp
   CodegenEnv.cpp
   CodegenUtils.cpp
+  DimLvlMapping.cpp
   LoopEmitter.cpp
   SparseBufferRewriting.cpp
   SparseStorageSpecifierToLLVM.cpp
   SparseTensorCodegen.cpp
   SparseTensorConversion.cpp
   SparseTensorPasses.cpp
   SparseTensorRewriting.cpp
   SparseTensorStorageLayout.cpp
add_mlir_dialect_library(MLIRSparseTensorTransforms
   MLIRSCFUtils
   MLIRSparseTensorDialect
   MLIRSparseTensorEnums
   MLIRSparseTensorUtils
   MLIRTensorDialect
   MLIRTransforms
   MLIRVectorDialect
 )
+
+# To make sure we adhere to the style guide:
+# <https://llvm.org/docs/CodingStandards.html#provide-a-virtual-method-anchor-for-classes-in-headers>
+check_cxx_compiler_flag(-Wweak-vtables
+  COMPILER_SUPPORTS_WARNING_WEAK_VTABLES)
+if(COMPILER_SUPPORTS_WARNING_WEAK_VTABLES)
+  target_compile_options(MLIRSparseTensorTransforms PUBLIC
+    "-Wweak-vtables")
+endif()

mlir/lib/Dialect/SparseTensor/Transforms/DimLvlMapping.h

+//===- DimLvlMapping.h - Utilities for dimension<->level --------*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file defines classes for managing all aspects of conversion
+// between dimensions and levels.  The classes are designed to be used
+// by both SparseTensorConversionPass and SparseTensorCodegenPass.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_DIMLVLMAPPING_H_
+#define MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_DIMLVLMAPPING_H_
+
+#include "CodegenUtils.h"
+
+#include "mlir/Dialect/Linalg/Utils/Utils.h"
+#include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
+
+namespace mlir {
+namespace sparse_tensor {
+
+//===----------------------------------------------------------------------===//
+//
+// Type aliases to help code be more self-documenting.  Unfortunately
+// these are not type-checked, so they only provide documentation rather
+// than doing anything to prevent mixups.
+//
+//===----------------------------------------------------------------------===//
+
+/// The type of dimension identifiers, and dimension-ranks.  We use the
+/// same type for both identifiers and ranks because the latter are used
+/// mainly for ordering-comparisons against the former (just like how the
+/// one-past-the-end iterators are used).
+using Dimension = uint64_t;
+
+/// The type of level identifiers, and level-ranks.  We use the same
+/// type for both identifiers and ranks because the latter are used
+/// mainly for ordering-comparisons against the former (just like how
+/// the one-past-the-end iterators are used).
+using Level = uint64_t;
+
+/// The type for individual components of a compile-time shape.  We avoid
+/// calling this "size" because we use the term "sizes" to indicate the
+/// actual run-time sizes, whereas this type also allows the value
+/// `ShapedType::kDynamic`.
+using Ship = int64_t;
+
+//===----------------------------------------------------------------------===//
+//
+// Helper functions.
+//
+//===----------------------------------------------------------------------===//
+
+/// Constructs a `constantIndex` when the potentially-dynamic size is
+/// not in fact dynamic.  When it is dynamic, returns the nullptr instead.
+inline Value constantSize(OpBuilder &builder, Location loc, Ship s) {
+  if (ShapedType::isDynamic(s))
+    return nullptr;
+  assert(s > 0);
+  return constantIndex(builder, loc, s);
+}
+
+//===----------------------------------------------------------------------===//
+/// A class for encapsulating and managing non-codegen aspects of
+/// converting between dimensions and levels.
+///
+/// This class memoizes the level-sizes to avoid redundant computation.
+/// We mark the methods `getLvlShape`, `getLvlShip`, etc as "const"
+/// to indicate intent and ensure a uniform API (with respect to the
+/// dimension variants of those methods).  However, since these methods
+/// may update the memo table, they are only logically-const not
+/// immutably-const.  Therefore, this class is not threadsafe (for now).
+class DimLvlMapping {
+public:
+  DimLvlMapping(RankedTensorType rtp)
+      : DimLvlMapping(rtp, getSparseTensorEncoding(rtp)) {}
+
+  // We memoize `lvlRank` and `dim2lvl` to avoid redundant checks later on.
+  DimLvlMapping(ShapedType stp, SparseTensorEncodingAttr enc)
+      : stp(stp), enc(enc),
+        lvlRank(isSparse() ? enc.getDimLevelType().size() : getDimRank()),
+        dim2lvl(enc.hasIdDimOrdering() ? AffineMap() : enc.getDimOrdering()),
+        lvlShape(isIdentity() ? 0 : lvlRank, kEmptyMemo) {
+    assert((!isIdentity() || getDimRank() == lvlRank) && "Rank mismatch");
+  }
+
+  /// Constructs a new `DimLvlMapping` for `getEncoding().withoutOrdering()`.
+  DimLvlMapping withoutOrdering() const {
+    return DimLvlMapping(stp, enc.withoutOrdering());
+  }
+
+  MLIRContext *getContext() const { return stp.getContext(); }
+
+  ShapedType getShapedType() const { return stp; }
+
+  Type getElementType() const { return stp.getElementType(); }
+
+  /// Returns the encoding (or nullptr for dense tensors).
+  SparseTensorEncodingAttr getEncoding() const { return enc; }
+
+  /// Returns true for tensors which have an encoding, and false for
+  /// those which do not.  Therefore tensors with an all-dense encoding
+  /// return true.
+  bool isSparse() const { return static_cast<bool>(enc); }
+
+  /// Returns true for tensors which do not have an encoding, and false
+  /// for tensors which do.  Therefore tensors with an all-dense encoding
+  /// return false.
+  bool isDense() const { return !enc; }
+
+  /// Returns true if the dimToLvl mapping is the identity.
+  bool isIdentity() const { return !dim2lvl; }
+
+  /// Returns the dimToLvl mapping (or nullptr for the identity).
+  AffineMap getDimToLvlMap() const { return dim2lvl; }
+
+  /// Returns the dimToLvl mapping, where the identity map is expanded out
+  /// into a full `AffineMap`.  This method is provided as a convenience,
+  /// but for most purposes other methods (`isIdentity`, `getDimToLvlMap`,
+  /// `getLvlShape`, etc) will be more helpful.
+  AffineMap getExpandedDimToLvlMap() const {
+    return dim2lvl
+               ? dim2lvl
+               : AffineMap::getMultiDimIdentityMap(getDimRank(), getContext());
+  }
+
+  /// Returns the dimension-rank.
+  Dimension getDimRank() const { return stp.getRank(); }
+
+  /// Returns the level-rank.
+  Level getLvlRank() const { return lvlRank; }
+
+  /// Returns the dimension-shape.
+  ArrayRef<Ship> getDimShape() const { return stp.getShape(); }
+
+  /// Returns the level-shape.
+  ArrayRef<Ship> getLvlShape() const;
+
+  /// Returns the requested dimension-ship.
+  Ship getDimShip(Dimension d) const {
+    assert(d < getDimRank());
+    return getDimShape()[d];
+  }
+
+  /// Returns the requested level-ship.
+  Ship getLvlShip(Level l) const {
+    if (isIdentity())
+      return getDimShip(l);
+    assert(l < lvlRank); // Ensure precondition of `setLvlShip`.
+    setLvlShip(l);
+    return lvlShape[l];
+  }
+
+  /// Returns true if the dimension-shape is static.
+  bool hasStaticDimShape() const { return stp.hasStaticShape(); }
+
+  /// Returns true if the level-shape is static.
+  bool hasStaticLvlShape() const {
+    return !ShapedType::isDynamicShape(getLvlShape());
+  }
+
+  /// Returns true if the dimension-ship is dynamic.
+  bool isDynamicDim(Dimension d) const {
+    return ShapedType::isDynamic(getDimShip(d));
+  }
+
+  /// Returns true if the level-ship is dynamic.
+  bool isDynamicLvl(Level l) const {
+    return ShapedType::isDynamic(getLvlShip(l));
+  }
+
+private:
+  /// Computes the requested level-ship, without performing any assertions
+  /// or checks that `getLvlShape` would want to hoist out of the loop.
+  ///
+  /// Preconditions:
+  /// * `dim2lvl` is non-null
+  /// * `l` is valid for `lvlRank`.
+  void setLvlShip(Level l) const;
+
+  // The value indicating level-ships which have not yet been memoized.
+  // Since we disallow zero as a level-size, that means we can safely
+  // use zero to indicate the lack of a memo.
+  static constexpr Ship kEmptyMemo = 0;
+
+  const ShapedType stp;
+  const SparseTensorEncodingAttr enc;
+  // Memoized to avoid frequent redundant conditionals.
+  const Level lvlRank;
+  const AffineMap dim2lvl;
+  // Memoized to avoid redundant computation.
+  mutable SmallVector<Ship> lvlShape;
+};
+
+//===----------------------------------------------------------------------===//
+/// Abstract base class for performing codegen associated with `DimLvlMapping`.
+class DimLvlBuilder : public DimLvlMapping {
+  /// Out-of-line virtual method to ensure we avoid weak-vtables:
+  /// <https://llvm.org/docs/CodingStandards.html#provide-a-virtual-method-anchor-for-classes-in-headers>
+  virtual void anchor();
+
+public:
+  DimLvlBuilder(OpBuilder &builder, Location loc, ShapedType stp,
+                SparseTensorEncodingAttr enc)
+      : DimLvlMapping(stp, enc), builder(builder), loc(loc) {}
+
+  DimLvlBuilder(OpBuilder &builder, Location loc, RankedTensorType rtp)
+      : DimLvlMapping(rtp), builder(builder), loc(loc) {}
+
+  DimLvlBuilder(OpBuilder &builder, Location loc, DimLvlMapping dlm)
+      : DimLvlMapping(dlm), builder(builder), loc(loc) {}
+
+protected:
+  // Since this class is virtual, we must disallow public copying in
+  // order to avoid "slicing".  Since this class has data members,
+  // that means making copying protected.
+  // <https://github.com/isocpp/CppCoreGuidelines/blob/master/CppCoreGuidelines.md#Rc-copy-virtual>
+  DimLvlBuilder(const DimLvlBuilder &) = default;
+  // Copy-assignment would be implicitly deleted (because the base class
+  // has const fields), so we explicitly delete it for clarity.
+  DimLvlBuilder &operator=(const DimLvlBuilder &) = delete;
+
+public:
+  virtual ~DimLvlBuilder() = default;
+
+  /// Generates code to lookup a dimension-size.  This should only
+  /// generate the lookup itself, and not perform any dim<->lvl
+  /// conversion or other logic.
+  virtual Value lookupDimSizeImpl(Value tensor, Dimension d) const = 0;
+
+  /// Generates code to lookup a level-size.  This should only
+  /// generate the lookup itself, and not perform any dim<->lvl
+  /// conversion or other logic.
+  virtual Value lookupLvlSizeImpl(Value tensor, Level l) const = 0;
+
+  /// Looks up a dimension-size by returning a constant from the shape
+  /// (for static sizes), or by calling `lookupDimSizeImpl` (for dynamic
+  /// sizes of sparse tensors) or `linalg::createOrFoldDimOp` (for dynamic
+  /// sizes of dense tensors).
+  Value lookupDimSize(Value tensor, Dimension d) const {
+    const auto sz = constantSize(builder, loc, getDimShip(d));
+    return sz           ? sz
+           : isSparse() ? lookupDimSizeImpl(tensor, d)
+                        : linalg::createOrFoldDimOp(builder, loc, tensor, d);
+  }
+
+  /// Looks up a level-size by returning a statically-computed constant
+  /// (when possible), or by calling `lookupLvlSizeImpl` (when dynamic).
+  Value lookupLvlSize(Value tensor, Level l) const {
+    const auto sz = constantSize(builder, loc, getLvlShip(l));
+    return sz ? sz : lookupLvlSizeImpl(tensor, l);
+  }
+
+  /// Populates the array with the dimension-shape of the `ShapedType`,
+  /// where dynamic sizes are represented by zero.
+  void reflectDimShape(SmallVectorImpl<Value> &out) const;
+
+  /// Returns an array with the dimension-shape of the `ShapedType`,
+  /// where dynamic sizes are represented by zero.
+  SmallVector<Value> reflectDimShape() const {
+    SmallVector<Value> out;
+    reflectDimShape(out);
+    return out;
+  }
+
+  /// Populates the array with the dimension-sizes of the given tensor.
+  void lookupDimSizes(Value tensor, SmallVectorImpl<Value> &out) const;
+
+  /// Returns an array with the dimension-sizes of the given tensor.
+  SmallVector<Value> lookupDimSizes(Value tensor) const {
+    SmallVector<Value> out;
+    lookupDimSizes(tensor, out);
+    return out;
+  }
+
+  /// Populates the array with the level-sizes of the given tensor.
+  void lookupLvlSizes(Value tensor, SmallVectorImpl<Value> &out) const;
+
+  /// Returns an array with the level-sizes of the given tensor.
+  SmallVector<Value> lookupLvlSizes(Value tensor) const {
+    SmallVector<Value> out;
+    lookupLvlSizes(tensor, out);
+    return out;
+  }
+
+  /// Generates code for computing level-sizes from the given
+  /// dimension-sizes, and populates the array with the results.
+  void computeLvlSizes(ValueRange dimSizes, SmallVectorImpl<Value> &out) const;
+
+  /// Returns an array with the level-sizes for the given dimension-sizes.
+  SmallVector<Value> computeLvlSizes(ValueRange dimSizes) const {
+    // TODO: when `isIdentity`, we'd like to just return `dimSizes`
+    // directly rather than copying.
+    SmallVector<Value> out;
+    computeLvlSizes(dimSizes, out);
+    return out;
+  }
+
+protected:
+  OpBuilder &builder;
+  Location loc;
+};
+
+} // namespace sparse_tensor
+} // namespace mlir
+
+#endif // MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_DIMLVLMAPPING_H_

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

+//===- DimLvlMapping.cpp - Utilities for dimension<->level ----------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file defines classes for managing all aspects of conversion
+// between dimensions and levels.  The classes are designed to be used
+// by both SparseTensorConversionPass and SparseTensorCodegenPass.
+//
+//===----------------------------------------------------------------------===//
+
+#include "DimLvlMapping.h"
+#include "CodegenUtils.h"
+
+using namespace mlir;
+using namespace mlir::sparse_tensor;
+
+//===----------------------------------------------------------------------===//
+
+ArrayRef<Ship> DimLvlMapping::getLvlShape() const {
+  if (isIdentity())
+    return getDimShape();
+  for (Level l = 0; l < lvlRank; ++l)
+    setLvlShip(l);
+  return lvlShape;
+}
+
+// TODO: Figure out how to make `setLvlShip` and `computeLvlSizes` share logic.
+void DimLvlMapping::setLvlShip(Level l) const {
+  if (lvlShape[l] != kEmptyMemo)
+    return; // Already been set.
+  // TODO: The following implementation only handles permutations;
+  // we 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).
+  lvlShape[l] = getDimShip(dim2lvl.getDimPosition(l));
+  // Ensure that the computation doesn't result in a value that
+  // coincides with the empty memo.
+  assert(lvlShape[l] != kEmptyMemo);
+}
+
+//===----------------------------------------------------------------------===//
+/// Out-of-line virtual method to ensure we avoid weak-vtables:
+/// <https://llvm.org/docs/CodingStandards.html#provide-a-virtual-method-anchor-for-classes-in-headers>
+void DimLvlBuilder::anchor() {}
+
+void DimLvlBuilder::reflectDimShape(SmallVectorImpl<Value> &out) const {
+  out.clear();
+  out.reserve(getDimRank());
+  const auto zero = constantIndex(builder, loc, 0);
+  for (auto s : getDimShape()) {
+    const auto sz = constantSize(builder, loc, s);
+    out.push_back(sz ? sz : zero);
+  }
+}
+
+void DimLvlBuilder::lookupDimSizes(Value tensor,
+                                   SmallVectorImpl<Value> &out) const {
+  const Dimension dimRank = getDimRank();
+  out.clear();
+  out.reserve(dimRank);
+  for (Dimension d = 0; d < dimRank; ++d)
+    out.push_back(lookupDimSize(tensor, d));
+}
+
+void DimLvlBuilder::lookupLvlSizes(Value tensor,
+                                   SmallVectorImpl<Value> &out) const {
+  out.clear();
+  out.reserve(getLvlRank());
+  // We inline `lookupLvlSize` to avoid the redundant checks of its
+  // underlying call to `getLvlShip`.
+  for (const auto &l : llvm::enumerate(getLvlShape())) {
+    const auto sz = constantSize(builder, loc, l.value());
+    out.push_back(sz ? sz : lookupLvlSizeImpl(tensor, l.index()));
+  }
+}
+
+// TODO: Figure out how to make `setLvlShip` and `computeLvlSizes` share logic.
+void DimLvlBuilder::computeLvlSizes(ValueRange dimSizes,
+                                    SmallVectorImpl<Value> &out) const {
+  out.clear();
+  out.reserve(getLvlRank());
+  if (isIdentity()) {
+    for (auto sz : dimSizes)
+      out.push_back(sz);
+    return;
+  }
+  const auto dim2lvl = getDimToLvlMap();
+  // TODO: The following implementation only handles permutations;
+  // we need to generalize this to handle arbitrary AffineExpr.
+  //
+  // There's no need to assert `isPermutation` before calling `getDimPosition`,
+  // because the latter already checks that the expr isa `AffineDimExpr`
+  // which is all we care about (for supporting permutations).
+  for (const auto &l : llvm::enumerate(getLvlShape())) {
+    const auto sz = constantSize(builder, loc, l.value());
+    out.push_back(sz ? sz : dimSizes[dim2lvl.getDimPosition(l.index())]);
+  }
+}

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

 // together with runtime support library keeps the conversion relatively
 // simple, but at the expense of IR opacity, which obscures opportunities
 // for subsequent optimization of the IR. An alternative is provided by
 // the SparseTensorCodegen pass.
 //
 //===----------------------------------------------------------------------===//
 
 #include "CodegenUtils.h"
+#include "DimLvlMapping.h"
 
 #include "mlir/Dialect/Bufferization/IR/BufferizableOpInterface.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/SparseTensor/IR/Enums.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 #include "mlir/Dialect/SparseTensor/Transforms/Passes.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Transforms/DialectConversion.h"
 
 using namespace mlir;
 using namespace mlir::sparse_tensor;
 
 namespace {
 
 //===----------------------------------------------------------------------===//
 // Helper methods.
 //===----------------------------------------------------------------------===//
 
+template <typename T>
+static RankedTensorType getRankedTensorType(T t) {
+  return t.getType().template cast<RankedTensorType>();
+}
+
 /// Maps each sparse tensor type to an opaque pointer.
 static std::optional<Type> convertSparseTensorTypes(Type type) {
   if (getSparseTensorEncoding(type) != nullptr)
     return LLVM::LLVMPointerType::get(IntegerType::get(type.getContext(), 8));
   return std::nullopt;
 }
 
 /// Replaces the `op` with  a `CallOp` to the function reference returned
 /// by `getFunc()`.
 static func::CallOp replaceOpWithFuncCall(RewriterBase &rewriter, Operation *op,
                                           StringRef name, TypeRange resultType,
                                           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.  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) {
+/// Specialize `DimLvlBuilder` for generating calls into the runtime library.
+class RuntimeDimLvlBuilder final : public DimLvlBuilder {
+public:
+  using DimLvlBuilder::DimLvlBuilder;
+
+  ~RuntimeDimLvlBuilder() final = default;
+
+  Value lookupLvlSizeImpl(Value tensor, Level lvl) const final {
     StringRef name = "sparseLvlSize";
     SmallVector<Value, 2> params{tensor, constantIndex(builder, loc, lvl)};
     Type iTp = builder.getIndexType();
     return createFuncCall(builder, loc, name, iTp, params, EmitCInterface::Off)
         .getResult(0);
   }
 
-/// 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) {
+  Value lookupDimSizeImpl(Value tensor, Dimension dim) const final {
     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);
   }
-
-/// 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 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 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,
+static void concatDimSizesFromInputs(OpBuilder &builder, 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.
-    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 =
-          createOrFoldDimCall(builder, loc, encSrc, srcTp, srcs[i], dim);
-      // Sum up all the sizes.
-      sizes[dim] = builder.create<arith::AddIOp>(loc, sizes[dim], srcSz);
-    }
+                                     Dimension dim,
+                                     SmallVectorImpl<Value> &sizes) {
+  assert(dim < dstTp.getRank());
+  // First, fill the sizes from an arbitrary source tensor.
+  RuntimeDimLvlBuilder(builder, loc, getRankedTensorType(srcs[0]))
+      .lookupDimSizes(srcs[0], sizes);
+
+  // Now, compute the size of the concatenation dimension.
+  // If the `dstTp` has a static size, then faithfully take it.
+  if (const auto sz = constantSize(builder, loc, dstTp.getShape()[dim])) {
+    sizes[dim] = sz;
+    return;
+  }
+  // Otherwise, dynamically sum up the actual sizes of the sources.
+  for (const auto &src : llvm::drop_begin(srcs)) {
+    sizes[dim] = builder.create<arith::AddIOp>(
+        loc, sizes[dim],
+        RuntimeDimLvlBuilder(builder, loc, getRankedTensorType(src))
+            .lookupDimSize(src, dim));
   }
 }
 
 /// Generates an uninitialized buffer of the given size and type,
 /// but returns it as type `memref<? x $tp>` (rather than as type
 /// `memref<$sz x $tp>`). Unlike temporary buffers on the stack,
 /// this buffer must be explicitly deallocated by client.
 static Value genAlloc(RewriterBase &rewriter, Location loc, Value sz, Type tp) {
public:
   NewCallParams(OpBuilder &builder, Location loc)
       : builder(builder), loc(loc), pTp(getOpaquePointerType(builder)) {}
 
   /// Initializes all static parameters (i.e., those which indicate
   /// type-level information such as the encoding and sizes), generating
   /// MLIR buffers as needed, and returning `this` for method chaining.
   /// This method does not set the action and pointer arguments, since
   /// those are handled by `genNewCall` instead.
-  NewCallParams &genBuffers(SparseTensorEncodingAttr enc, ValueRange sizes,
-                            ShapedType stp);
+  NewCallParams &genBuffers(DimLvlMapping dlm, ValueRange dimSizes);
 
   /// (Re)sets the C++ template type parameters, and returns `this`
   /// for method chaining.  This is already done as part of `genBuffers`,
   /// but is factored out so that it can also be called independently
   /// whenever subsequent `genNewCall` calls want to reuse the same
   /// buffers but different type parameters.
   //
   // TODO: This is only ever used by sparse2sparse-viaCOO `ConvertOp`;
private:
   Location loc;
   Type pTp;
   Value params[kNumParams];
 };
 
 // TODO: see the note at `_mlir_ciface_newSparseTensor` about how
 // the meaning of the various arguments (e.g., "sizes" vs "shapes")
 // is inconsistent between the different actions.
-NewCallParams &NewCallParams::genBuffers(SparseTensorEncodingAttr enc,
-                                         ValueRange dimSizes, ShapedType stp) {
-  const unsigned lvlRank = enc.getDimLevelType().size();
-  const unsigned dimRank = stp.getRank();
+NewCallParams &NewCallParams::genBuffers(DimLvlMapping dlm,
+                                         ValueRange dimSizes) {
+  const unsigned lvlRank = dlm.getLvlRank();
+  const unsigned dimRank = dlm.getDimRank();
   // Sparsity annotations.
-  params[kParamLvlTypes] = genLvlTypesBuffer(builder, loc, enc);
+  params[kParamLvlTypes] = genLvlTypesBuffer(builder, loc, dlm.getEncoding());
   // Dimension-sizes array of the enveloping tensor.  Useful for either
   // verification of external data, or for construction of internal data.
   assert(dimSizes.size() == dimRank && "Dimension-rank mismatch");
   params[kParamDimSizes] = allocaBuffer(builder, loc, dimSizes);
   // The level-sizes array must be passed as well, since for arbitrary
   // dim2lvl mappings it cannot be trivially reconstructed at runtime.
-  // For now however, since we're still assuming permutations, we will
-  // initialize this parameter alongside the `dim2lvl` and `lvl2dim`
-  // parameters below.  We preinitialize `lvlSizes` for code symmetry.
-  SmallVector<Value> lvlSizes(lvlRank);
+  SmallVector<Value> lvlSizes =
+      RuntimeDimLvlBuilder(builder, loc, dlm).computeLvlSizes(dimSizes);
+  params[kParamLvlSizes] = allocaBuffer(builder, loc, lvlSizes);
   // The dimension-to-level mapping and its inverse.  We must preinitialize
   // `dim2lvl` so that the true branch below can perform random-access
   // `operator[]` assignment.  We preinitialize `lvl2dim` for code symmetry.
   SmallVector<Value> dim2lvl(dimRank);
   SmallVector<Value> lvl2dim(lvlRank);
-  auto dimOrder = enc.getDimOrdering();
-  if (dimOrder) {
+  if (!dlm.isIdentity()) {
+    const auto dimOrder = dlm.getEncoding().getDimOrdering();
     assert(dimOrder.isPermutation());
     for (unsigned l = 0; l < lvlRank; l++) {
       // The `d`th source variable occurs in the `l`th result position.
       uint64_t d = dimOrder.getDimPosition(l);
       dim2lvl[d] = constantIndex(builder, loc, l);
       lvl2dim[l] = constantIndex(builder, loc, d);
-      lvlSizes[l] = dimSizes[d];
     }
   } else {
     assert(dimRank == lvlRank && "Rank mismatch");
-    for (unsigned i = 0; i < lvlRank; i++) {
+    for (unsigned i = 0; i < lvlRank; i++)
       dim2lvl[i] = lvl2dim[i] = constantIndex(builder, loc, i);
-      lvlSizes[i] = dimSizes[i];
-    }
   }
-  params[kParamLvlSizes] = allocaBuffer(builder, loc, lvlSizes);
   params[kParamLvl2Dim] = allocaBuffer(builder, loc, lvl2dim);
-  params[kParamDim2Lvl] =
-      dimOrder ? allocaBuffer(builder, loc, dim2lvl) : params[kParamLvl2Dim];
+  params[kParamDim2Lvl] = dlm.isIdentity()
+                              ? params[kParamLvl2Dim]
+                              : allocaBuffer(builder, loc, dim2lvl);
   // Secondary and primary types encoding.
-  setTemplateTypes(enc, stp);
+  setTemplateTypes(dlm.getEncoding(), dlm.getShapedType());
   // Finally, make note that initialization is complete.
   assert(isInitialized() && "Initialization failed");
   // And return `this` for method chaining.
   return *this;
 }
 
 /// Generates a call to obtain the values array.
 static Value genValuesCall(OpBuilder &builder, Location loc, ShapedType tp,

 ///     coo->add(reshape(elem.indices), elem.value)
 ///   }
 ///   s = newSparseTensor(coo)
 template <typename ReshapeOp>
 static LogicalResult
 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)
+  const auto srcTp = getRankedTensorType(op.getSrc());
+  const auto dstTp = getRankedTensorType(op.getResult());
+  RuntimeDimLvlBuilder srcDLM(rewriter, loc, srcTp);
+  RuntimeDimLvlBuilder dstDLM(rewriter, loc, dstTp);
+  if (srcDLM.isDense() || dstDLM.isDense())
     return failure();
-  Type elemTp = srcTp.getElementType();
-  assert(elemTp == dstTp.getElementType() &&
+  const Type elemTp = srcDLM.getElementType();
+  assert(elemTp == dstDLM.getElementType() &&
          "reshape should not change element type");
   // Start an iterator over the source tensor (in original index order).
-  const auto noPerm = encSrc.withoutOrdering();
-  SmallVector<Value> srcDimSizes =
-      getDimSizes(rewriter, loc, encSrc, srcTp, adaptor.getSrc());
+  const SmallVector<Value> srcDimSizes =
+      srcDLM.lookupDimSizes(adaptor.getSrc());
   NewCallParams params(rewriter, loc);
-  Value iter = params.genBuffers(noPerm, srcDimSizes, srcTp)
+  const Value iter = params.genBuffers(srcDLM.withoutOrdering(), srcDimSizes)
                          .genNewCall(Action::kToIterator, adaptor.getSrc());
   // Start a new COO for the destination tensor.
   SmallVector<Value> dstDimSizes;
-  if (dstTp.hasStaticShape())
+  if (dstDLM.hasStaticDimShape())
     // Static "shapes" are in fact "sizes".
-    fillDimShape(rewriter, loc, dstTp, dstDimSizes);
+    dstDLM.reflectDimShape(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();
+                       dstDLM.getDimShape(), op.getReassociationIndices());
+  const Value coo =
+      params.genBuffers(dstDLM, dstDimSizes).genNewCall(Action::kEmptyCOO);
+  const Value dstPerm = params.getDim2LvlMap();
   // Construct a while loop over the iterator.
-  Type iTp = rewriter.getIndexType();
-  Value srcIdx = genAlloca(rewriter, loc, srcTp.getRank(), iTp);
-  Value dstIdx = genAlloca(rewriter, loc, dstTp.getRank(), iTp);
+  const Type iTp = rewriter.getIndexType();
+  Value srcIdx = genAlloca(rewriter, loc, srcDLM.getDimRank(), iTp);
+  Value dstIdx = genAlloca(rewriter, loc, dstDLM.getDimRank(), 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());

 //   }
 // TODO: It can be used by other operators (ReshapeOp, ConvertOP) conversion to
 // reduce code repetition!
 // TODO: rename to `genSparseIterationLoop`?
 static void genSparseCOOIterationLoop(
     ConversionPatternRewriter &rewriter, Location loc, Value t,
     RankedTensorType tensorTp,
     function_ref<void(OpBuilder &, Location, Value, Value)> bodyBuilder) {
-  auto enc = getSparseTensorEncoding(tensorTp);
-  assert(enc && "Generating Sparse Tensor COO Loop on a Dense Tensor!");
-
-  unsigned rank = tensorTp.getRank();
-  Type elemTp = tensorTp.getElementType();
+  RuntimeDimLvlBuilder dlm(rewriter, loc, tensorTp);
+  assert(dlm.isSparse() &&
+         "Generating Sparse Tensor COO Loop on a Dense Tensor!");
 
   // Start an iterator over the tensor (in original index order).
-  const auto noPerm = enc.withoutOrdering();
-  SmallVector<Value> dimSizes = getDimSizes(rewriter, loc, noPerm, tensorTp, t);
   Value iter = NewCallParams(rewriter, loc)
-                   .genBuffers(noPerm, dimSizes, tensorTp)
+                   .genBuffers(dlm.withoutOrdering(), dlm.lookupDimSizes(t))
                    .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);
+  Value srcIdx =
+      genAlloca(rewriter, loc, dlm.getDimRank(), rewriter.getIndexType());
+  Value elemPtr = genAllocaScalar(rewriter, loc, dlm.getElementType());
   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());
   Block *after = rewriter.createBlock(&whileOp.getAfter(), {}, noTypes);
   rewriter.setInsertionPointToStart(after);
 
-  bool hasDenseDim = llvm::any_of(
-      enc.getDimLevelType(), [](DimLevelType dlt) { return isDenseDLT(dlt); });
+  bool hasDenseDim =
+      llvm::any_of(dlm.getEncoding().getDimLevelType(), isDenseDLT);
   if (hasDenseDim) {
     Value elemV = rewriter.create<memref::LoadOp>(loc, elemPtr);
     Value isZero = genIsNonzero(rewriter, loc, elemV);
     scf::IfOp ifOp = rewriter.create<scf::IfOp>(loc, isZero, /*else*/ false);
     rewriter.setInsertionPointToStart(&ifOp.getThenRegion().front());
   }
   // Callback here to build loop body.
   bodyBuilder(rewriter, loc, srcIdx, elemPtr);
 
   // Exit the scope from the IfOp.
   if (hasDenseDim)
     rewriter.setInsertionPointToEnd(after);
 
   rewriter.create<scf::YieldOp>(loc);
   // Finish generating loop.
   rewriter.setInsertionPointAfter(whileOp);
 
   // Free memory for iterator.
-  genDelIteratorCall(rewriter, loc, elemTp, iter);
+  genDelIteratorCall(rewriter, loc, dlm.getElementType(), iter);
 }
 
 // Generate loop that iterates over a dense tensor.
 //   for i1 in dim1
 //    ..
 //     for ik in dimk
 //       val = a[i1,..,ik]
 //       if val != 0

 /// 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 {
-    auto stp = op.getSource().getType().cast<ShapedType>();
+    const auto rtp = getRankedTensorType(op.getSource());
+    RuntimeDimLvlBuilder dlm(rewriter, op->getLoc(), rtp);
     // Only rewrite sparse DimOp.
-    auto enc = getSparseTensorEncoding(stp);
-    if (!enc)
+    if (dlm.isDense())
       return failure();
     // Only rewrite DimOp with constant index.
     std::optional<int64_t> dim = op.getConstantIndex();
     if (!dim)
       return failure();
     // Generate the call.
     Value src = adaptor.getOperands()[0];
-    rewriter.replaceOp(
-        op, createOrFoldDimCall(rewriter, op->getLoc(), enc, stp, src, *dim));
+    rewriter.replaceOp(op, dlm.lookupDimSize(src, *dim));
     return success();
   }
 };
 
 /// Sparse conversion rule for trivial tensor casts.
 class SparseCastConverter : public OpConversionPattern<tensor::CastOp> {
 public:
   using OpConversionPattern::OpConversionPattern;

 /// Sparse conversion rule for the new operator.
 class SparseTensorNewConverter : public OpConversionPattern<NewOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(NewOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Location loc = op.getLoc();
-    auto stp = op.getType().cast<ShapedType>();
-    auto enc = getSparseTensorEncoding(stp);
-    if (!enc)
+    const auto rtp = getRankedTensorType(op);
+    const RuntimeDimLvlBuilder dlm(rewriter, loc, rtp);
+    const auto enc = dlm.getEncoding();
+    if (dlm.isDense())
       return failure();
-    const unsigned dimRank = stp.getRank();
-    const unsigned lvlRank = enc.getDimLevelType().size();
+    const unsigned dimRank = dlm.getDimRank();
+    const unsigned lvlRank = dlm.getLvlRank();
     // Construct the dimShape.
-    const auto dimShape = stp.getShape();
-    SmallVector<Value> dimShapeValues = getDimShape(rewriter, loc, stp);
+    SmallVector<Value> dimShapeValues = dlm.reflectDimShape();
     Value dimShapeBuffer = allocaBuffer(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());
+        constantPrimaryTypeEncoding(rewriter, loc, rtp.getElementType());
     Value reader =
         createFuncCall(rewriter, loc, "createCheckedSparseTensorReader",
                        opaqueTp,
                        {adaptor.getOperands()[0], dimShapeBuffer, valTp},
                        EmitCInterface::On)
             .getResult(0);
     // Construct the lvlSizes.  If the dimShape is static, then it's
     // identical to dimSizes: so we can compute lvlSizes entirely at
     // compile-time.  If dimShape is dynamic, then we'll need to generate
     // code for computing lvlSizes from the `reader`'s actual dimSizes.
     //
     // TODO: For now we're still assuming `dim2lvl` is a permutation.
     // But since we're computing lvlSizes here (rather than in the runtime),
     // we can easily generalize that simply by adjusting this code.
     //
     // FIXME: reduce redundancy vs `NewCallParams::genBuffers`.
     Value dimSizesBuffer;
-    if (!stp.hasStaticShape()) {
+    if (!dlm.hasStaticDimShape()) {
       Type indexTp = rewriter.getIndexType();
       auto memTp = MemRefType::get({ShapedType::kDynamic}, indexTp);
       dimSizesBuffer =
           createFuncCall(rewriter, loc, "getSparseTensorReaderDimSizes", memTp,
                          reader, EmitCInterface::On)
               .getResult(0);
     }
     Value lvlSizesBuffer;
     Value lvl2dimBuffer;
     Value dim2lvlBuffer;
     if (auto dimOrder = enc.getDimOrdering()) {
       assert(dimOrder.isPermutation() && "Got non-permutation");
+      const auto dimShape = dlm.getDimShape();
       // We preinitialize `dim2lvlValues` since we need random-access writing.
       // And we preinitialize the others for stylistic consistency.
       SmallVector<Value> lvlSizeValues(lvlRank);
       SmallVector<Value> lvl2dimValues(lvlRank);
       SmallVector<Value> dim2lvlValues(dimRank);
       for (unsigned l = 0; l < lvlRank; l++) {
         // The `d`th source variable occurs in the `l`th result position.
         uint64_t d = dimOrder.getDimPosition(l);
         Value lvl = constantIndex(rewriter, loc, l);
         Value dim = constantIndex(rewriter, loc, d);
         dim2lvlValues[d] = lvl;
         lvl2dimValues[l] = dim;
         lvlSizeValues[l] =
-            (dimShape[d] == ShapedType::kDynamic)
+            ShapedType::isDynamic(dimShape[d])
                 ? rewriter.create<memref::LoadOp>(loc, dimSizesBuffer, dim)
                 : dimShapeValues[d];
       }
       lvlSizesBuffer = allocaBuffer(rewriter, loc, lvlSizeValues);
       lvl2dimBuffer = allocaBuffer(rewriter, loc, lvl2dimValues);
       dim2lvlBuffer = allocaBuffer(rewriter, loc, dim2lvlValues);
     } else {
       assert(dimRank == lvlRank && "Rank mismatch");
public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(bufferization::AllocTensorOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     if (op.getCopy())
       return rewriter.notifyMatchFailure(op,
                                          "sparse tensor copy not implemented");
     Location loc = op.getLoc();
-    RankedTensorType resType = op.getType();
-    auto enc = getSparseTensorEncoding(resType);
-    if (!enc)
+    const DimLvlMapping dlm(op.getType());
+    if (dlm.isDense())
       return failure();
     // Gather all dimension sizes as SSA values.
-    SmallVector<Value> sizes;
+    const Dimension dimRank = dlm.getDimRank();
+    SmallVector<Value> dimSizes;
+    dimSizes.reserve(dimRank);
     unsigned int operandCtr = 0;
-    for (int64_t i = 0; i < resType.getRank(); ++i) {
-      if (resType.isDynamicDim(i)) {
-        sizes.push_back(adaptor.getOperands()[operandCtr++]);
-      } else {
-        sizes.push_back(
-            rewriter.create<arith::ConstantIndexOp>(loc, op.getStaticSize(i)));
-      }
+    for (Dimension d = 0; d < dimRank; ++d) {
+      dimSizes.push_back(
+          dlm.isDynamicDim(d)
+              ? adaptor.getOperands()[operandCtr++]
+              : constantIndex(rewriter, loc, op.getStaticSize(d)));
     }
     // Generate the call to construct empty tensor. The sizes are
     // explicitly defined by the arguments to the alloc operator.
-    rewriter.replaceOp(op,
-                       NewCallParams(rewriter, loc)
-                           .genBuffers(enc, sizes, resType.cast<ShapedType>())
+    rewriter.replaceOp(op, NewCallParams(rewriter, loc)
+                               .genBuffers(dlm, dimSizes)
                                .genNewCall(Action::kEmpty));
     return success();
   }
 };
 
 /// Sparse conversion rule for the convert operator.
 class SparseTensorConvertConverter : public OpConversionPattern<ConvertOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   SparseTensorConvertConverter(MLIRContext *context,
                                SparseTensorConversionOptions o)
       : OpConversionPattern<ConvertOp>(context), options(o) {}
   SparseTensorConvertConverter(TypeConverter &typeConv, MLIRContext *context,
                                SparseTensorConversionOptions o)
       : OpConversionPattern<ConvertOp>(typeConv, context), options(o) {}
 
   LogicalResult
   matchAndRewrite(ConvertOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
-    Location loc = op->getLoc();
-    Type resType = op.getType();
-    Type srcType = op.getSource().getType();
-    auto encDst = getSparseTensorEncoding(resType);
-    auto encSrc = getSparseTensorEncoding(srcType);
-    Value src = adaptor.getOperands()[0];
-    if (encDst && encSrc) {
+    const Location loc = op->getLoc();
+    const auto srcTp = getRankedTensorType(op.getSource());
+    const auto dstTp = getRankedTensorType(op);
+    RuntimeDimLvlBuilder srcDLM(rewriter, loc, srcTp);
+    DimLvlMapping dstDLM(dstTp);
+    if (srcDLM.isDense() && dstDLM.isDense())
+      return failure();
+
+    const Type elemTp = srcDLM.getElementType();
+    const Value src = adaptor.getOperands()[0];
+    if (srcDLM.isSparse() && dstDLM.isSparse()) {
+      const auto dstEnc = dstDLM.getEncoding();
+      const auto srcEnc = srcDLM.getEncoding();
       // This is a sparse => sparse conversion, which is handled as follows:
       //   t = src->toCOO();         ; src to COO in dst order
       //   dst = newSparseTensor(t)
       // Using the coordinate scheme as an intermediate does not always
       // yield the fastest conversion but avoids the need for a full
       // O(N^2) conversion matrix.
-      if (encDst == encSrc) {
+      if (dstEnc == srcEnc) {
         rewriter.replaceOp(op, adaptor.getOperands()); // hidden nop cast
         return success();
       }
       NewCallParams params(rewriter, loc);
-      ShapedType stp = srcType.cast<ShapedType>();
-      SmallVector<Value> dimSizes =
-          getDimSizes(rewriter, loc, encSrc, stp, src);
+      const SmallVector<Value> dimSizes = srcDLM.lookupDimSizes(src);
       bool useDirectConversion;
       switch (options.sparseToSparseStrategy) {
       case SparseToSparseConversionStrategy::kViaCOO:
         useDirectConversion = false;
         break;
       case SparseToSparseConversionStrategy::kDirect:
         useDirectConversion = true;
-        assert(canUseDirectConversion(encDst.getDimLevelType()) &&
+        assert(canUseDirectConversion(dstEnc.getDimLevelType()) &&
                "Unsupported target for direct sparse-to-sparse conversion");
         break;
       case SparseToSparseConversionStrategy::kAuto:
-        useDirectConversion = canUseDirectConversion(encDst.getDimLevelType());
+        useDirectConversion = canUseDirectConversion(dstEnc.getDimLevelType());
         break;
       }
       if (useDirectConversion) {
-        rewriter.replaceOp(op, params.genBuffers(encDst, dimSizes, stp)
-                                   .genNewCall(Action::kSparseToSparse, src));
+        rewriter.replaceOp(
+            op, params.genBuffers(DimLvlMapping(srcTp, dstEnc), dimSizes)
       } 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());
+        const auto mixedEnc = SparseTensorEncodingAttr::get(
+            op->getContext(), dstEnc.getDimLevelType(), dstEnc.getDimOrdering(),
+            dstEnc.getHigherOrdering(), srcEnc.getPointerBitWidth(),
+            srcEnc.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, dimSizes, stp)
+        Value coo = params.genBuffers(DimLvlMapping(srcTp, mixedEnc), dimSizes)
                         .genNewCall(Action::kToCOO, src);
-        Value dst = params.setTemplateTypes(encDst, stp)
+        Value dst = params.setTemplateTypes(dstEnc, srcTp)
                         .genNewCall(Action::kFromCOO, coo);
-        genDelCOOCall(rewriter, loc, stp.getElementType(), coo);
+        genDelCOOCall(rewriter, loc, elemTp, coo);
         rewriter.replaceOp(op, dst);
       }
       return success();
     }
-    if (!encDst && encSrc) {
+    if (srcDLM.isSparse() && dstDLM.isDense()) {
+      const auto srcEnc = srcDLM.getEncoding();
       // This is sparse => dense conversion, which is handled as follows:
       //   dst = new Tensor(0);
       //   iter = new SparseTensorIterator(src);
       //   while (elem = iter->getNext()) {
       //     dst[elem.indices] = elem.value;
       //   }
       //   delete iter;
-      RankedTensorType dstTensorTp = resType.cast<RankedTensorType>();
-      RankedTensorType srcTensorTp = srcType.cast<RankedTensorType>();
-      unsigned rank = dstTensorTp.getRank();
-      Type elemTp = dstTensorTp.getElementType();
+      const unsigned dstDimRank = dstDLM.getDimRank();
       // 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(
+      const auto dstEnc = SparseTensorEncodingAttr::get(
           op->getContext(),
-          SmallVector<DimLevelType>(rank, DimLevelType::Dense), AffineMap(),
-          AffineMap(), encSrc.getPointerBitWidth(), encSrc.getIndexBitWidth());
-      SmallVector<Value> dimSizes =
-          getDimSizes(rewriter, loc, encSrc, srcTensorTp, src);
+          SmallVector<DimLevelType>(dstDimRank, DimLevelType::Dense),
+          AffineMap(), AffineMap(), srcEnc.getPointerBitWidth(),
+          srcEnc.getIndexBitWidth());
+      SmallVector<Value> dimSizes = srcDLM.lookupDimSizes(src);
       Value iter = NewCallParams(rewriter, loc)
-                       .genBuffers(encDst, dimSizes, dstTensorTp)
+                       .genBuffers(DimLvlMapping(dstTp, dstEnc), dimSizes)
                        .genNewCall(Action::kToIterator, src);
-      Value ind = genAlloca(rewriter, loc, rank, rewriter.getIndexType());
+      Value ind = genAlloca(rewriter, loc, dstDimRank, 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, dimSizes);
+      Value dst = allocDenseTensor(rewriter, loc, dstTp, 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);
       rewriter.setInsertionPointToStart(after);
-      SmallVector<Value> ivs = loadIndices(rewriter, loc, rank, ind);
+      SmallVector<Value> ivs = loadIndices(rewriter, loc, dstDimRank, ind);
       insertScalarIntoDenseTensor(rewriter, loc, elemPtr, dst, ivs);
       rewriter.create<scf::YieldOp>(loc);
       rewriter.setInsertionPointAfter(whileOp);
       genDelIteratorCall(rewriter, loc, elemTp, iter);
-      rewriter.replaceOpWithNewOp<bufferization::ToTensorOp>(op, resType, dst);
+      rewriter.replaceOpWithNewOp<bufferization::ToTensorOp>(op, dstTp, dst);
       // Deallocate the buffer.
       if (bufferization::allocationDoesNotEscape(op->getOpResult(0))) {
         rewriter.setInsertionPoint(insertionBlock->getTerminator());
         deallocDenseTensor(rewriter, loc, dst);
       }
       return success();
     }
-    if (!encDst && !encSrc) {
-      // dense => dense
-      return failure();
-    }
+    assert(srcDLM.isDense() && dstDLM.isSparse());
     // This is a dense => sparse conversion or a sparse constant in COO =>
     // sparse conversion, which is handled as follows:
     //   t = newSparseCOO()
     //   ...code to fill the COO tensor t...
     //   s = newSparseTensor(t)
     //
     // To fill the COO tensor from a dense tensor:
     //   for i1 in dim1
matchAndRewrite(ConvertOp op, OpAdaptor adaptor,
     //     t->add(val, [i1,..,ik], [p1,..,pk])
     //
     // Note that the dense tensor traversal code is actually implemented
     // using MLIR IR to avoid having to expose too much low-level
     // memref traversal details to the runtime support library.
     // Also note that the code below only generates the "new" ops and
     // the loop-nest per se; whereas the entire body of the innermost
     // loop is generated by genAddElt().
-    ShapedType stp = resType.cast<ShapedType>();
-    unsigned rank = stp.getRank();
+    const unsigned dstDimRank = dstDLM.getDimRank();
     SmallVector<Value> sizes;
     sizesFromSrc(rewriter, sizes, loc, src);
     NewCallParams params(rewriter, loc);
-    Value coo =
-        params.genBuffers(encDst, sizes, stp).genNewCall(Action::kEmptyCOO);
-    Value ind = genAlloca(rewriter, loc, rank, rewriter.getIndexType());
-    Value perm = params.getDim2LvlMap();
-    Type eltType = stp.getElementType();
-    Value elemPtr = genAllocaScalar(rewriter, loc, eltType);
+    Value coo = params.genBuffers(dstDLM, sizes).genNewCall(Action::kEmptyCOO);
+    Value ind = genAlloca(rewriter, loc, dstDimRank, rewriter.getIndexType());
+    const Value perm = params.getDim2LvlMap();
+    Value elemPtr = genAllocaScalar(rewriter, loc, elemTp);
     genDenseTensorOrSparseConstantIterLoop(
-        rewriter, loc, src, rank,
-        [&](OpBuilder &builder, Location loc, Value val, ValueRange indices) {
-          for (unsigned i = 0; i < rank; i++) {
-            Value idx = constantIndex(builder, loc, i);
-            builder.create<memref::StoreOp>(loc, indices[i], ind, idx);
+        rewriter, loc, src, dstDimRank,
+        [&](OpBuilder &builder, Location loc, Value val, ValueRange dimInd) {
+          // TODO: rewrite this to use `storeIndices`
+          for (unsigned d = 0; d < dstDimRank; d++) {
+            Value dim = constantIndex(builder, loc, d);
+            builder.create<memref::StoreOp>(loc, dimInd[d], ind, dim);
           }
           builder.create<memref::StoreOp>(loc, val, elemPtr);
-          genAddEltCall(builder, loc, eltType, coo, elemPtr, ind, perm);
+          genAddEltCall(builder, loc, elemTp, coo, elemPtr, ind, perm);
         });
     // Final call to construct sparse tensor storage.
     Value dst = params.genNewCall(Action::kFromCOO, coo);
-    genDelCOOCall(rewriter, loc, eltType, coo);
+    genDelCOOCall(rewriter, loc, elemTp, coo);
     rewriter.replaceOp(op, dst);
     return success();
   }
 
 private:
   /// Options to control sparse code generation.
   SparseTensorConversionOptions options;
 };
public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(InsertOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     // Note that the current regime only allows for strict lexicographic
     // index order. All values are passed by reference through stack
     // allocated memrefs.
     Location loc = op->getLoc();
-    auto tp = op.getTensor().getType().cast<RankedTensorType>();
+    const auto tp = getRankedTensorType(op.getTensor());
     auto elemTp = tp.getElementType();
     unsigned rank = tp.getRank();
     auto mref = genAlloca(rewriter, loc, rank, rewriter.getIndexType());
     auto vref = genAllocaScalar(rewriter, loc, elemTp);
     for (unsigned i = 0; i < rank; i++)
       rewriter.create<memref::StoreOp>(loc, adaptor.getIndices()[i], mref,
                                        constantIndex(rewriter, loc, i));
     rewriter.create<memref::StoreOp>(loc, adaptor.getValue(), vref);
     SmallString<12> name{"lexInsert", primaryTypeFunctionSuffix(elemTp)};
     createFuncCall(rewriter, loc, name, {}, {adaptor.getTensor(), mref, vref},
                    EmitCInterface::On);
     rewriter.replaceOp(op, adaptor.getTensor());
     return success();
   }
 };
 
 /// Sparse conversion rule for the expand operator.
 class SparseTensorExpandConverter : public OpConversionPattern<ExpandOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(ExpandOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Location loc = op->getLoc();
-    RankedTensorType srcType =
-        op.getTensor().getType().cast<RankedTensorType>();
+    const auto srcType = getRankedTensorType(op.getTensor());
     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());
     // 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);
+    RuntimeDimLvlBuilder srcDLM(rewriter, loc, srcType);
+    Value sz =
+        srcDLM.lookupLvlSize(adaptor.getTensor(), srcDLM.getLvlRank() - 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
matchAndRewrite(CompressOp op, OpAdaptor adaptor,
     // all-zero/false by only iterating over the set elements, so the
     // complexity remains proportional to the sparsity of the expanded
     // access pattern.
     Value values = adaptor.getValues();
     Value filled = adaptor.getFilled();
     Value added = adaptor.getAdded();
     Value count = adaptor.getCount();
     Value tensor = adaptor.getTensor();
-    auto tp = op.getTensor().getType().cast<RankedTensorType>();
+    const auto tp = getRankedTensorType(op.getTensor());
     Type elemTp = tp.getElementType();
     unsigned rank = tp.getRank();
     auto mref = genAlloca(rewriter, loc, rank, rewriter.getIndexType());
     for (unsigned i = 0; i < rank - 1; i++)
       rewriter.create<memref::StoreOp>(loc, adaptor.getIndices()[i], mref,
                                        constantIndex(rewriter, loc, i));
     SmallString<12> name{"expInsert", primaryTypeFunctionSuffix(elemTp)};
     createFuncCall(rewriter, loc, name, {},
matchAndRewrite(ConcatenateOp op, OpAdaptor adaptor,
     //    a = malloc(shapeOf(a)) or newSparseAllDense(shapeOf(a))
     //    for i, j, k // dense input
     //      a[ adjustForOffset(i,j,k) ] = b[i,j,k]
     //
     //    for elem in sparse_input
     //      a[ adjustForOffset(elem.indices) ] = elem.value
     //    return a
     Location loc = op.getLoc();
-    auto dstTp = op.getType().cast<RankedTensorType>();
+    const auto dstTp = getRankedTensorType(op);
     auto encDst = getSparseTensorEncoding(dstTp);
     Type elemTp = dstTp.getElementType();
     uint64_t concatDim = op.getDimension().getZExtValue();
     unsigned rank = dstTp.getRank();
 
     Value dst;     // destination tensor
     Value dstPerm; // destination tensor permutation (if sparse out)
     // A pointer to the value being inserted (if dense => sparse)
     Value elemPtr;
     // Memory that holds the COO for destination tensor (if sparse out)
     Value dstIdx;
     // The offset applied to the dimenstion to be concated (starting from 0)
     Value offset = constantIndex(rewriter, loc, 0);
 
     SmallVector<Value> sizes;
     NewCallParams params(rewriter, loc);
-    concatSizesFromInputs(rewriter, sizes, loc, dstTp, op.getInputs(),
-                          concatDim);
+    concatDimSizesFromInputs(rewriter, loc, dstTp, op.getInputs(), concatDim,
+                             sizes);
 
     bool allDense = false;
     Value dstTensor;
     if (encDst) {
       allDense = encDst.isAllDense();
       // Start a new COO or an initialized annotated all dense sparse tensor.
-      dst = params.genBuffers(encDst, sizes, dstTp)
+      RuntimeDimLvlBuilder dstDLM(rewriter, loc, dstTp, encDst);
+      dst = params.genBuffers(dstDLM, sizes)
                 .genNewCall(allDense ? Action::kEmpty : Action::kEmptyCOO);
       dstIdx = genAlloca(rewriter, loc, rank, rewriter.getIndexType());
       if (allDense) {
         dstTensor = dst;
         // Get the values buffer for the sparse tensor and reshape it to the
         // corresponding dense tensor shape.
         dst = genValuesCall(rewriter, loc,
                             MemRefType::get({ShapedType::kDynamic}, elemTp),
auto dimIdx2LvlIdx = [&](ValueRange dIdx) -> SmallVector<Value> {
       SmallVector<Value> lIdx;
       for (unsigned i = 0; i < dIdx.size(); i++)
         lIdx.push_back(dIdx[toOrigDim(encDst, i)]);
       return lIdx;
     };
     for (auto it : llvm::zip(op.getInputs(), adaptor.getInputs())) {
       Value orignalOp = std::get<0>(it); // Input (with encoding) from Op
       Value adaptedOp = std::get<1>(it); // Input (type converted) from adaptor
-      RankedTensorType srcTp = orignalOp.getType().cast<RankedTensorType>();
-      auto encSrc = getSparseTensorEncoding(srcTp);
-      if (encSrc) {
+      const auto srcTp = getRankedTensorType(orignalOp);
+      RuntimeDimLvlBuilder srcDLM(rewriter, loc, srcTp);
+      if (srcDLM.isSparse()) {
         genSparseCOOIterationLoop(
             rewriter, loc, adaptedOp, srcTp,
             [&](OpBuilder &builder, Location loc, Value idx,
                 Value elemPtr) -> void {
               SmallVector<Value> dimInd =
                   loadIndices(builder, loc, rank, idx, concatDim, offset);
               if (encDst && !allDense) {
                 // Case: sparse => sparse, except for annotated all dense.
for (auto it : llvm::zip(op.getInputs(), adaptor.getInputs())) {
                 if (allDense)
                   lvlInd = dimIdx2LvlIdx(lvlInd);
                 builder.create<memref::StoreOp>(loc, val, dst, lvlInd);
               }
             });
       }
       // Accumulate offset.
       // TODO: avoid calling sparseDimSize multiple times by caching the result!
-      Value curDim = createOrFoldDimCall(rewriter, loc, encSrc, srcTp,
-                                         adaptedOp, concatDim);
+      Value curDim = srcDLM.lookupDimSize(adaptedOp, concatDim);
 
       offset = rewriter.create<arith::AddIOp>(loc, offset, curDim);
     }
     if (encDst) {
       if (!allDense) {
         // In sparse output case, the destination holds the COO.
         Value coo = dst;
         dst = params.genNewCall(Action::kFromCOO, coo);

 
 /// Sparse conversion rule for the output operator.
 class SparseTensorOutConverter : public OpConversionPattern<OutOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(OutOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
-    Location loc = op->getLoc();
-    ShapedType srcType = op.getTensor().getType().cast<ShapedType>();
+    const Location loc = op->getLoc();
+    const auto srcTp = getRankedTensorType(op.getTensor());
+    RuntimeDimLvlBuilder srcDLM(rewriter, loc, srcTp);
     // Convert to default permuted COO.
-    Value src = adaptor.getOperands()[0];
-    auto encSrc = getSparseTensorEncoding(srcType);
-    SmallVector<Value> dimSizes =
-        getDimSizes(rewriter, loc, encSrc, srcType, src);
-    const auto enc = encSrc.withoutOrdering();
-    Value coo = NewCallParams(rewriter, loc)
-                    .genBuffers(enc, dimSizes, srcType)
-                    .genNewCall(Action::kToCOO, src);
+    Value coo;
+    {
+      const Value src = adaptor.getOperands()[0];
+      const SmallVector<Value> dimSizes = srcDLM.lookupDimSizes(src);
+      coo = NewCallParams(rewriter, loc)
+                .genBuffers(srcDLM.withoutOrdering(), dimSizes)
+    }
     // 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() &&
-                                !encSrc.getDimOrdering().isIdentity());
+    const Value sort = constantI1(rewriter, loc, !srcDLM.isIdentity());
     SmallVector<Value, 3> outParams{coo, adaptor.getOperands()[1], sort};
-    Type eltType = srcType.getElementType();
-    SmallString<18> name{"outSparseTensor", primaryTypeFunctionSuffix(eltType)};
+    const Type elemTp = srcDLM.getElementType();
+    SmallString<18> name{"outSparseTensor", primaryTypeFunctionSuffix(elemTp)};
     createFuncCall(rewriter, loc, name, {}, outParams, EmitCInterface::Off);
-    genDelCOOCall(rewriter, loc, eltType, coo);
+    genDelCOOCall(rewriter, loc, elemTp, coo);
     rewriter.eraseOp(op);
     return success();
   }
 };
 
 } // namespace
 
 //===----------------------------------------------------------------------===//