[mlir][sparse] Enhancing sparse=>sparse conversion.

[wrengr]

Summary

Fixes: https://github.com/llvm/llvm-project/issues/51652

Depends On D122060

[wrengr]

rebase, and fixing build error

[wrengr]

redoing the diff

[wrengr]

re-renaming the old appendPointer to appendCurrentPointer (since I've realized the name "finalizeSegment" is better used elsewhere).

[wrengr]

Lots of reorganization, subsequent to all the recent changes in the parent revisions.

[wrengr]

rebase

[wrengr]

rebase

[wrengr]

Rebasing for D123166. Also removing a bunch of inline keywords, per MLIR style-guide.

[wrengr]

rebase

[wrengr]

Rebasing for the latest version of D122060 (2022-05-11T13:09)

This also introduces a minor change to the enumerator design. Namely, due to the splitting of SparseTensorEnumeratorBase<V> vs SparseTensorEnumerator<P,I,V>, I've had to add the new SparseTensorStorageBase::newEnumerator methods to allow constructing objects of the abstract base type without reintroducing the "<P,I> vs <Q,J> problem" for direct conversion.

[wrengr]

Arg! It looks like when I installed the github version of Arcanist to avoid the https://www.mail-archive.com/[e-mail]/msg1838310.html bug with Debian's version, the new version isn't automatically applying the formatting script :( So please bear with the overly-long lines for now

[wrengr]

(attempting to get the clang-formatter to run again)

[wrengr]

(Previous attempt didn't work, so fixing some things manually)

[wrengr]

Rerunning git-clang-format

[aartbik]

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

569

Yeah! All your abstraction work is paying off!

605

shall we make this an explicit null_ptr for clarity
(I am fine either way but perhaps making it explicit reads clearer)?

943

Fine for this revision, but would it make sense to move this into its own file in the longer run?
Just to split the storage vs. supporting classes a bit more. Probably applies to some of the other classes too by now.

1426

yeah, direct! Sweet!

mlir/test/Dialect/SparseTensor/conversion.mlir

5

period at end

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

2

period at end of comment

[wrengr]

Rebasing for D125596

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

static void insertScalarIntoDenseTensor(OpBuilder &builder, Location loc,
   for (unsigned i = 0; i < rank; i++) {
     Value idx = constantIndex(builder, loc, i);
     ivs.push_back(builder.create<memref::LoadOp>(loc, ind, idx));
   }
   Value elemV = builder.create<memref::LoadOp>(loc, elemPtr);
   builder.create<memref::StoreOp>(loc, elemV, tensor, ivs);
 }
 
+/// Determine if the runtime library supports direct conversion to the
+/// given target `dimTypes`.
+static bool canUseDirectConversion(
+    ArrayRef<SparseTensorEncodingAttr::DimLevelType> dimTypes) {
+  bool alreadyCompressed = false;
+  for (uint64_t rank = dimTypes.size(), r = 0; r < rank; r++) {
+    switch (dimTypes[r]) {
+    case SparseTensorEncodingAttr::DimLevelType::Compressed:
+      if (alreadyCompressed)
+        return false; // Multiple compressed dimensions not yet supported.
+      alreadyCompressed = true;
+      break;
+    case SparseTensorEncodingAttr::DimLevelType::Dense:
+      if (alreadyCompressed)
+        return false; // Dense after Compressed not yet supported.
+      break;
+    case SparseTensorEncodingAttr::DimLevelType::Singleton:
+      // Although Singleton isn't generally supported yet, the direct
+      // conversion method doesn't have any particular problems with
+      // singleton after compressed.
+      break;
+    }
+  }
+  return true;
+}
+
 //===----------------------------------------------------------------------===//
 // Conversion rules.
 //===----------------------------------------------------------------------===//
 
 /// Sparse conversion rule for returns.
 class SparseReturnConverter : public OpConversionPattern<func::ReturnOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
if (encDst && encSrc) {
       if (encDst == encSrc) {
         rewriter.replaceOp(op, adaptor.getOperands()); // hidden nop cast
         return success();
       }
       SmallVector<Value, 4> sizes;
       SmallVector<Value, 8> params;
       ShapedType stp = srcType.cast<ShapedType>();
       sizesFromPtr(rewriter, sizes, op, 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) {
+        newParams(rewriter, params, op, stp, encDst, Action::kSparseToSparse,
+                  sizes, src);
+        rewriter.replaceOp(op, genNewCall(rewriter, op, params));
+      } 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(),
             encSrc.getPointerBitWidth(), encSrc.getIndexBitWidth());
         newParams(rewriter, params, op, stp, enc, Action::kToCOO, sizes, src);
         Value coo = genNewCall(rewriter, op, params);
         params[3] = constantPointerTypeEncoding(rewriter, loc, encDst);
         params[4] = constantIndexTypeEncoding(rewriter, loc, encDst);
         params[6] = constantAction(rewriter, loc, Action::kFromCOO);
         params[7] = coo;
         Value dst = genNewCall(rewriter, op, params);
         genDelCOOCall(rewriter, op, stp.getElementType(), coo);
         rewriter.replaceOp(op, dst);
+      }
       return success();
     }
     if (!encDst && encSrc) {
       // This is sparse => dense conversion, which is handled as follows:
       //   dst = new Tensor(0);
       //   iter = src->toCOO();
       //   iter->startIterator();
       //   while (elem = iter->getNext()) {

mlir/lib/ExecutionEngine/SparseTensorUtils.cpp

 
 /// A version of `operator*` on `uint64_t` which checks for overflows.
 static inline uint64_t checkedMul(uint64_t lhs, uint64_t rhs) {
   assert((lhs == 0 || rhs <= std::numeric_limits<uint64_t>::max() / lhs) &&
          "Integer overflow");
   return lhs * rhs;
 }
 
+// TODO: adjust this so it can be used by `openSparseTensorCOO` too.
+// That version doesn't have the permutation, and the `sizes` are
+// a pointer/C-array rather than `std::vector`.
+//
+/// Asserts that the `sizes` (in target-order) under the `perm` (mapping
+/// semantic-order to target-order) are a refinement of the desired `shape`
+/// (in semantic-order).
+///
+/// Precondition: `perm` and `shape` must be valid for `rank`.
+static inline void
+assertPermutedSizesMatchShape(const std::vector<uint64_t> &sizes, uint64_t rank,
+                              const uint64_t *perm, const uint64_t *shape) {
+  assert(perm && shape);
+  assert(rank == sizes.size() && "Rank mismatch");
+  for (uint64_t r = 0; r < rank; r++)
+    assert((shape[r] == 0 || shape[r] == sizes[perm[r]]) &&
+           "Dimension size mismatch");
+}
+
 /// A sparse tensor element in coordinate scheme (value and indices).
 /// For example, a rank-1 vector element would look like
 ///   ({i}, a[i])
 /// and a rank-5 tensor element like
 ///   ({i,j,k,l,m}, a[i,j,k,l,m])
 /// We use pointer to a shared index pool rather than e.g. a direct
 /// vector since that (1) reduces the per-element memory footprint, and
 /// (2) centralizes the memory reservation and (re)allocation to one place.

 private:
   const std::vector<uint64_t> sizes; // per-dimension sizes
   std::vector<Element<V>> elements;  // all COO elements
   std::vector<uint64_t> indices;     // shared index pool
   bool iteratorLocked = false;
   unsigned iteratorPos = 0;
 };
 
+// Forward.
+template <typename V>
+class SparseTensorEnumeratorBase;
+
 /// Abstract base class for `SparseTensorStorage<P,I,V>`.  This class
 /// takes responsibility for all the `<P,I,V>`-independent aspects
 /// of the tensor (e.g., shape, sparsity, permutation).  In addition,
 /// we use function overloading to implement "partial" method
 /// specialization, which the C-API relies on to catch type errors
 /// arising from our use of opaque pointers.
 class SparseTensorStorageBase {
 public:
public:
   const std::vector<DimLevelType> &getDimTypes() const { return dimTypes; }
 
   /// Safely check if the (storage-order) dimension uses compressed storage.
   bool isCompressedDim(uint64_t d) const {
     assert(d < getRank());
     return (dimTypes[d] == DimLevelType::kCompressed);
   }
 
+  /// Allocate a new enumerator.
+  virtual void newEnumerator(SparseTensorEnumeratorBase<double> **, uint64_t,
+                             const uint64_t *) const {
+    fatal("enumf64");
+  }
+  virtual void newEnumerator(SparseTensorEnumeratorBase<float> **, uint64_t,
+                             const uint64_t *) const {
+    fatal("enumf32");
+  }
+  virtual void newEnumerator(SparseTensorEnumeratorBase<int64_t> **, uint64_t,
+                             const uint64_t *) const {
+    fatal("enumi64");
+  }
+  virtual void newEnumerator(SparseTensorEnumeratorBase<int32_t> **, uint64_t,
+                             const uint64_t *) const {
+    fatal("enumi32");
+  }
+  virtual void newEnumerator(SparseTensorEnumeratorBase<int16_t> **, uint64_t,
+                             const uint64_t *) const {
+    fatal("enumi16");
+  }
+  virtual void newEnumerator(SparseTensorEnumeratorBase<int8_t> **, uint64_t,
+                             const uint64_t *) const {
+    fatal("enumi8");
+  }
+  virtual void newEnumerator(SparseTensorEnumeratorBase<complex64> **, uint64_t,
+                             const uint64_t *) const {
+    fatal("enumc64");
+  }
+  virtual void newEnumerator(SparseTensorEnumeratorBase<complex32> **, uint64_t,
+                             const uint64_t *) const {
+    fatal("enumc32");
+  }
+
   /// Overhead storage.
   virtual void getPointers(std::vector<uint64_t> **, uint64_t) { fatal("p64"); }
   virtual void getPointers(std::vector<uint32_t> **, uint64_t) { fatal("p32"); }
   virtual void getPointers(std::vector<uint16_t> **, uint64_t) { fatal("p16"); }
   virtual void getPointers(std::vector<uint8_t> **, uint64_t) { fatal("p8"); }
   virtual void getIndices(std::vector<uint64_t> **, uint64_t) { fatal("i64"); }
   virtual void getIndices(std::vector<uint32_t> **, uint64_t) { fatal("i32"); }
   virtual void getIndices(std::vector<uint16_t> **, uint64_t) { fatal("i16"); }

 /// A memory-resident sparse tensor using a storage scheme based on
 /// per-dimension sparse/dense annotations. This data structure provides a
 /// bufferized form of a sparse tensor type. In contrast to generating setup
 /// methods for each differently annotated sparse tensor, this method provides
 /// a convenient "one-size-fits-all" solution that simply takes an input tensor
 /// and annotations to implement all required setup in a general manner.
 template <typename P, typename I, typename V>
 class SparseTensorStorage : public SparseTensorStorageBase {
+  /// Private constructor to share code between the other constructors.
+  /// Beware that the object is not necessarily guaranteed to be in a
+  /// valid state after this constructor alone; e.g., `isCompressedDim(d)`
+  /// doesn't entail `!(pointers[d].empty())`.
+  ///
+  /// Precondition: `perm` and `sparsity` must be valid for `szs.size()`.
+  SparseTensorStorage(const std::vector<uint64_t> &szs, const uint64_t *perm,
+                      const DimLevelType *sparsity)
+      : SparseTensorStorageBase(szs, perm, sparsity), pointers(getRank()),
+        indices(getRank()), idx(getRank()) {}
+
 public:
   /// Constructs a sparse tensor storage scheme with the given dimensions,
   /// permutation, and per-dimension dense/sparse annotations, using
   /// the coordinate scheme tensor for the initial contents if provided.
   ///
   /// Precondition: `perm` and `sparsity` must be valid for `szs.size()`.
   SparseTensorStorage(const std::vector<uint64_t> &szs, const uint64_t *perm,
-                      const DimLevelType *sparsity,
-                      SparseTensorCOO<V> *coo = nullptr)
-      : SparseTensorStorageBase(szs, perm, sparsity), pointers(getRank()),
-        indices(getRank()), idx(getRank()) {
+                      const DimLevelType *sparsity, SparseTensorCOO<V> *coo)
+      : SparseTensorStorage(szs, perm, sparsity) {
     // Provide hints on capacity of pointers and indices.
     // TODO: needs much fine-tuning based on actual sparsity; currently
     //       we reserve pointer/index space based on all previous dense
     //       dimensions, which works well up to first sparse dim; but
     //       we should really use nnz and dense/sparse distribution.
     bool allDense = true;
     uint64_t sz = 1;
     for (uint64_t r = 0, rank = getRank(); r < rank; r++) {
if (coo) {
       uint64_t nnz = elements.size();
       values.reserve(nnz);
       fromCOO(elements, 0, nnz, 0);
     } else if (allDense) {
       values.resize(sz, 0);
     }
   }
 
+  /// Constructs a sparse tensor storage scheme with the given dimensions,
+  /// permutation, and per-dimension dense/sparse annotations, using
+  /// the given sparse tensor for the initial contents.
+  ///
+  /// Preconditions:
+  /// * `perm` and `sparsity` must be valid for `szs.size()`.
+  /// * The `tensor` must have the same value type `V`.
+  SparseTensorStorage(const std::vector<uint64_t> &szs, const uint64_t *perm,
+                      const DimLevelType *sparsity,
+                      const SparseTensorStorageBase &tensor);
+
   ~SparseTensorStorage() override = default;
 
   /// Partially specialize these getter methods based on template types.
   void getPointers(std::vector<P> **out, uint64_t d) override {
     assert(d < getRank());
     *out = &pointers[d];
   }
   void getIndices(std::vector<I> **out, uint64_t d) override {
public:
   /// Finalizes lexicographic insertions.
   void endInsert() override {
     if (values.empty())
       finalizeSegment(0);
     else
       endPath(0);
   }
 
+  void newEnumerator(SparseTensorEnumeratorBase<V> **out, uint64_t rank,
+                     const uint64_t *perm) const override {
+    *out = new SparseTensorEnumerator<P, I, V>(*this, rank, perm);
+  }
+
   /// Returns this sparse tensor storage scheme as a new memory-resident
   /// sparse tensor in coordinate scheme with the given dimension order.
   ///
   /// Precondition: `perm` must be valid for `getRank()`.
   SparseTensorCOO<V> *toCOO(const uint64_t *perm) const {
-    SparseTensorEnumerator<P, I, V> enumerator(*this, getRank(), perm);
+    SparseTensorEnumeratorBase<V> *enumerator;
+    newEnumerator(&enumerator, getRank(), perm);
     SparseTensorCOO<V> *coo =
-        new SparseTensorCOO<V>(enumerator.permutedSizes(), values.size());
-    enumerator.forallElements([&coo](const std::vector<uint64_t> &ind, V val) {
+        new SparseTensorCOO<V>(enumerator->permutedSizes(), values.size());
+    enumerator->forallElements([&coo](const std::vector<uint64_t> &ind, V val) {
       coo->add(ind, val);
     });
     // TODO: This assertion assumes there are no stored zeros,
     // or if there are then that we don't filter them out.
     // Cf., <https://github.com/llvm/llvm-project/issues/54179>
     assert(coo->getElements().size() == values.size());
+    delete enumerator;
     return coo;
   }
 
   /// Factory method. Constructs a sparse tensor storage scheme with the given
   /// dimensions, permutation, and per-dimension dense/sparse annotations,
   /// using the coordinate scheme tensor for the initial contents if provided.
   /// In the latter case, the coordinate scheme must respect the same
   /// permutation as is desired for the new sparse tensor storage.
   ///
   /// Precondition: `shape`, `perm`, and `sparsity` must be valid for `rank`.
   static SparseTensorStorage<P, I, V> *
   newSparseTensor(uint64_t rank, const uint64_t *shape, const uint64_t *perm,
                   const DimLevelType *sparsity, SparseTensorCOO<V> *coo) {
     SparseTensorStorage<P, I, V> *n = nullptr;
     if (coo) {
-      assert(coo->getRank() == rank && "Tensor rank mismatch");
       const auto &coosz = coo->getSizes();
-      for (uint64_t r = 0; r < rank; r++)
-        assert(shape[r] == 0 || shape[r] == coosz[perm[r]]);
+      assertPermutedSizesMatchShape(coosz, rank, perm, shape);
       n = new SparseTensorStorage<P, I, V>(coosz, perm, sparsity, coo);
     } else {
       std::vector<uint64_t> permsz(rank);
       for (uint64_t r = 0; r < rank; r++) {
         assert(shape[r] > 0 && "Dimension size zero has trivial storage");
         permsz[perm[r]] = shape[r];
       }
-      n = new SparseTensorStorage<P, I, V>(permsz, perm, sparsity);
+      // We pass the null `coo` to ensure we select the intended constructor.
+      n = new SparseTensorStorage<P, I, V>(permsz, perm, sparsity, coo);
     }
     return n;
   }
 
+  /// Factory method. Constructs a sparse tensor storage scheme with
+  /// the given dimensions, permutation, and per-dimension dense/sparse
+  /// annotations, using the sparse tensor for the initial contents.
+  ///
+  /// Preconditions:
+  /// * `shape`, `perm`, and `sparsity` must be valid for `rank`.
+  /// * The `tensor` must have the same value type `V`.
+  static SparseTensorStorage<P, I, V> *
+  newSparseTensor(uint64_t rank, const uint64_t *shape, const uint64_t *perm,
+                  const DimLevelType *sparsity,
+                  const SparseTensorStorageBase *source) {
+    assert(source && "Got nullptr for source");
+    SparseTensorEnumeratorBase<V> *enumerator;
+    source->newEnumerator(&enumerator, rank, perm);
+    const auto &permsz = enumerator->permutedSizes();
+    assertPermutedSizesMatchShape(permsz, rank, perm, shape);
+    auto *tensor =
+        new SparseTensorStorage<P, I, V>(permsz, perm, sparsity, *source);
+    delete enumerator;
+    return tensor;
+  }
+
 private:
   /// Appends an arbitrary new position to `pointers[d]`.  This method
   /// checks that `pos` is representable in the `P` type; however, it
   /// does not check that `pos` is semantically valid (i.e., larger than
   /// the previous position and smaller than `indices[d].capacity()`).
   void appendPointer(uint64_t d, uint64_t pos, uint64_t count = 1) {
     assert(isCompressedDim(d));
     assert(pos <= std::numeric_limits<P>::max() &&
if (isCompressedDim(d)) {
         return; // Short-circuit, since it'll be a nop.
       if (d + 1 == getRank())
         values.insert(values.end(), i - full, 0);
       else
         finalizeSegment(d + 1, 0, i - full);
     }
   }
 
+  /// Writes the given coordinate to `indices[d][pos]`.  This method
+  /// checks that `i` is representable in the `I` type; however, it
+  /// does not check that `i` is semantically valid (i.e., in bounds
+  /// for `sizes[d]` and not elsewhere occurring in the same segment).
+  void writeIndex(uint64_t d, uint64_t pos, uint64_t i) {
+    assert(isCompressedDim(d));
+    // Subscript assignment to `std::vector` requires that the `pos`-th
+    // entry has been initialized; thus we must be sure to check `size()`
+    // here, instead of `capacity()` as would be ideal.
+    assert(pos < indices[d].size() && "Index position is out of bounds");
+    assert(i <= std::numeric_limits<I>::max() &&
+           "Index value is too large for the I-type");
+    indices[d][pos] = static_cast<I>(i);
+  }
+
+  /// Computes the assembled-size associated with the `d`-th dimension,
+  /// given the assembled-size associated with the `(d-1)`-th dimension.
+  /// "Assembled-sizes" correspond to the (nominal) sizes of overhead
+  /// storage, as opposed to "dimension-sizes" which are the cardinality
+  /// of coordinates for that dimension.
+  ///
+  /// Precondition: the `pointers[d]` array must be fully initialized
+  /// before calling this method.
+  uint64_t assembledSize(uint64_t parentSz, uint64_t d) const {
+    if (isCompressedDim(d))
+      return pointers[d][parentSz];
+    // else if dense:
+    return parentSz * getDimSizes()[d];
+  }
+
   /// Initializes sparse tensor storage scheme from a memory-resident sparse
   /// tensor in coordinate scheme. This method prepares the pointers and
   /// indices arrays under the given per-dimension dense/sparse annotations.
   ///
   /// Preconditions:
   /// (1) the `elements` must be lexicographically sorted.
   /// (2) the indices of every element are valid for `sizes` (equal rank
   ///     and pointwise less-than).
if (d == Base::getRank()) {
       for (uint64_t i = 0; i < sz; i++) {
         cursor_reord_d = i;
         forallElements(yield, pstart + i, d + 1);
       }
     }
   }
 };
 
+/// Statistics regarding the number of nonzero subtensors in
+/// a source tensor, for direct sparse=>sparse conversion a la
+/// <https://arxiv.org/abs/2001.02609>.
+///
+/// N.B., this class stores references to the parameters passed to
+/// the constructor; thus, objects of this class must not outlive
+/// those parameters.
+class SparseTensorNNZ {
+public:
+  /// Allocate the statistics structure for the desired sizes and
+  /// sparsity (in the target tensor's storage-order).  This constructor
+  /// does not actually populate the statistics, however; for that see
+  /// `initialize`.
+  ///
+  /// Precondition: `szs` must not contain zeros.
+  SparseTensorNNZ(const std::vector<uint64_t> &szs,
+                  const std::vector<DimLevelType> &sparsity)
+      : dimSizes(szs), dimTypes(sparsity), nnz(getRank()) {
+    assert(dimSizes.size() == dimTypes.size() && "Rank mismatch");
+    bool uncompressed = true;
+    uint64_t sz = 1; // the product of all `dimSizes` strictly less than `r`.
+    for (uint64_t rank = getRank(), r = 0; r < rank; r++) {
+      switch (dimTypes[r]) {
+      case DimLevelType::kCompressed:
+        assert(uncompressed &&
+               "Multiple compressed layers not currently supported");
+        uncompressed = false;
+        nnz[r].resize(sz, 0); // Both allocate and zero-initialize.
+        break;
+      case DimLevelType::kDense:
+        assert(uncompressed &&
+               "Dense after compressed not currently supported");
+        break;
+      case DimLevelType::kSingleton:
+        // Singleton after Compressed causes no problems for allocating
+        // `nnz` nor for the yieldPos loop.  This remains true even
+        // when adding support for multiple compressed dimensions or
+        // for dense-after-compressed.
+        break;
+      }
+      sz = checkedMul(sz, dimSizes[r]);
+    }
+  }
+
+  // We disallow copying to help avoid leaking the stored references.
+  SparseTensorNNZ(const SparseTensorNNZ &) = delete;
+  SparseTensorNNZ &operator=(const SparseTensorNNZ &) = delete;
+
+  /// Returns the rank of the target tensor.
+  uint64_t getRank() const { return dimSizes.size(); }
+
+  /// Enumerate the source tensor to fill in the statistics.  The
+  /// enumerator should already incorporate the permutation (from
+  /// semantic-order to the target storage-order).
+  template <typename V>
+  void initialize(SparseTensorEnumeratorBase<V> &enumerator) {
+    assert(enumerator.getRank() == getRank() && "Tensor rank mismatch");
+    assert(enumerator.permutedSizes() == dimSizes && "Tensor size mismatch");
+    enumerator.forallElements(
+        [this](const std::vector<uint64_t> &ind, V) { add(ind); });
+  }
+
+  /// The type of callback functions which receive an nnz-statistic.
+  using NNZConsumer = const std::function<void(uint64_t)> &;
+
+  /// Lexicographically enumerates all indicies for dimensions strictly
+  /// less than `stopDim`, and passes their nnz statistic to the callback.
+  /// Since our use-case only requires the statistic not the coordinates
+  /// themselves, we do not bother to construct those coordinates.
+  void forallIndices(uint64_t stopDim, NNZConsumer yield) const {
+    assert(stopDim < getRank() && "Stopping-dimension is out of bounds");
+    assert(dimTypes[stopDim] == DimLevelType::kCompressed &&
+           "Cannot look up non-compressed dimensions");
+    forallIndices(yield, stopDim, 0, 0);
+  }
+
+private:
+  /// Adds a new element (i.e., increment its statistics).  We use
+  /// a method rather than inlining into the lambda in `initialize`,
+  /// to avoid spurious templating over `V`.  And this method is private
+  /// to avoid needing to re-assert validity of `ind` (which is guaranteed
+  /// by `forallElements`).
+  void add(const std::vector<uint64_t> &ind) {
+    uint64_t parentPos = 0;
+    for (uint64_t rank = getRank(), r = 0; r < rank; r++) {
+      if (dimTypes[r] == DimLevelType::kCompressed)
+        nnz[r][parentPos]++;
+      parentPos = parentPos * dimSizes[r] + ind[r];
+    }
+  }
+
+  /// Recursive component of the public `forallIndices`.
+  void forallIndices(NNZConsumer yield, uint64_t stopDim, uint64_t parentPos,
+                     uint64_t d) const {
+    assert(d <= stopDim);
+    if (d == stopDim) {
+      assert(parentPos < nnz[d].size() && "Cursor is out of range");
+      yield(nnz[d][parentPos]);
+    } else {
+      const uint64_t sz = dimSizes[d];
+      const uint64_t pstart = parentPos * sz;
+      for (uint64_t i = 0; i < sz; i++)
+        forallIndices(yield, stopDim, pstart + i, d + 1);
+    }
+  }
+
+  // All of these are in the target storage-order.
+  const std::vector<uint64_t> &dimSizes;
+  const std::vector<DimLevelType> &dimTypes;
+  std::vector<std::vector<uint64_t>> nnz;
+};
+
+template <typename P, typename I, typename V>
+SparseTensorStorage<P, I, V>::SparseTensorStorage(
+    const std::vector<uint64_t> &szs, const uint64_t *perm,
+    const DimLevelType *sparsity, const SparseTensorStorageBase &tensor)
+    : SparseTensorStorage(szs, perm, sparsity) {
+  SparseTensorEnumeratorBase<V> *enumerator;
+  tensor.newEnumerator(&enumerator, getRank(), perm);
+  {
+    // Initialize the statistics structure.
+    SparseTensorNNZ nnz(getDimSizes(), getDimTypes());
+    nnz.initialize(*enumerator);
+    // Initialize "pointers" overhead (and allocate "indices", "values").
+    uint64_t parentSz = 1; // assembled-size (not dimension-size) of `r-1`.
+    for (uint64_t rank = getRank(), r = 0; r < rank; r++) {
+      if (isCompressedDim(r)) {
+        pointers[r].reserve(parentSz + 1);
+        pointers[r].push_back(0);
+        uint64_t currentPos = 0;
+        nnz.forallIndices(r, [this, &currentPos, r](uint64_t n) {
+          currentPos += n;
+          appendPointer(r, currentPos);
+        });
+        assert(pointers[r].size() == parentSz + 1 &&
+               "Final pointers size doesn't match allocated size");
+        // That assertion entails `assembledSize(parentSz, r)`
+        // is now in a valid state.  That is, `pointers[r][parentSz]`
+        // equals the present value of `currentPos`, which is the
+        // correct assembled-size for `indices[r]`.
+      }
+      // Update assembled-size for the next iteration.
+      parentSz = assembledSize(parentSz, r);
+      // Ideally we need only `indices[r].reserve(parentSz)`, however
+      // the `std::vector` implementation forces us to initialize it too.
+      // That is, in the yieldPos loop we need random-access assignment
+      // to `indices[r]`; however, `std::vector`'s subscript-assignment
+      // only allows assigning to already-initialized positions.
+      if (isCompressedDim(r))
+        indices[r].resize(parentSz, 0);
+    }
+    values.resize(parentSz, 0); // Both allocate and zero-initialize.
+  }
+  // The yieldPos loop
+  enumerator->forallElements([this](const std::vector<uint64_t> &ind, V val) {
+    uint64_t parentSz = 1, parentPos = 0;
+    for (uint64_t rank = getRank(), r = 0; r < rank; r++) {
+      if (isCompressedDim(r)) {
+        // If `parentPos == parentSz` then it's valid as an array-lookup;
+        // however, it's semantically invalid here since that entry
+        // does not represent a segment of `indices[r]`.  Moreover, that
+        // entry must be immutable for `assembledSize` to remain valid.
+        assert(parentPos < parentSz && "Pointers position is out of bounds");
+        const uint64_t currentPos = pointers[r][parentPos];
+        // This increment won't overflow the `P` type, since it can't
+        // exceed the original value of `pointers[r][parentPos+1]`
+        // which was already verified to be within bounds for `P`
+        // when it was written to the array.
+        pointers[r][parentPos]++;
+        writeIndex(r, currentPos, ind[r]);
+        parentPos = currentPos;
+      } else { // Dense dimension.
+        parentPos = parentPos * getDimSizes()[r] + ind[r];
+      }
+      parentSz = assembledSize(parentSz, r);
+    }
+    assert(parentPos < values.size() && "Value position is out of bounds");
+    values[parentPos] = val;
+  });
+  // No longer need the enumerator, so we'll delete it ASAP.
+  delete enumerator;
+  // The finalizeYieldPos loop
+  for (uint64_t parentSz = 1, rank = getRank(), r = 0; r < rank; r++) {
+    if (isCompressedDim(r)) {
+      assert(parentSz == pointers[r].size() - 1 &&
+             "Actual pointers size doesn't match the expected size");
+      // Can't check all of them, but at least we can check the last one.
+      assert(pointers[r][parentSz - 1] == pointers[r][parentSz] &&
+             "Pointers got corrupted");
+      // TODO: optimize this by using `memmove` or similar.
+      for (uint64_t n = 0; n < parentSz; n++) {
+        const uint64_t parentPos = parentSz - n;
+        pointers[r][parentPos] = pointers[r][parentPos - 1];
+      }
+      pointers[r][0] = 0;
+    }
+    parentSz = assembledSize(parentSz, r);
+  }
+}
+
 /// Helper to convert string to lower case.
 static char *toLower(char *token) {
   for (char *c = token; *c; c++)
     *c = tolower(*c);
   return token;
 }
 
 /// Read the MME header of a general sparse matrix of type real.
if (action <= Action::kFromCOO) {                                          \
         assert(action == Action::kEmpty);                                      \
       }                                                                        \
       auto *tensor = SparseTensorStorage<P, I, V>::newSparseTensor(            \
           rank, shape, perm, sparsity, coo);                                   \
       if (action == Action::kFromFile)                                         \
         delete coo;                                                            \
       return tensor;                                                           \
     }                                                                          \
+    if (action == Action::kSparseToSparse) {                                   \
+      auto *tensor = static_cast<SparseTensorStorageBase *>(ptr);              \
+      return SparseTensorStorage<P, I, V>::newSparseTensor(rank, shape, perm,  \
+                                                           sparsity, tensor);  \
+    }                                                                          \
     if (action == Action::kEmptyCOO)                                           \
       return SparseTensorCOO<V>::newSparseTensorCOO(rank, shape, perm);        \
     coo = static_cast<SparseTensorStorage<P, I, V> *>(ptr)->toCOO(perm);       \
     if (action == Action::kToIterator) {                                       \
       coo->startIterator();                                                    \
     } else {                                                                   \
       assert(action == Action::kToCOO);                                        \
     }                                                                          \

mlir/test/Dialect/SparseTensor/conversion.mlir

-// RUN: mlir-opt %s --sparse-tensor-conversion --canonicalize --cse | FileCheck %s
+// First use with `kViaCOO` for sparse2sparse conversion (the old way).
+// RUN: mlir-opt %s --sparse-tensor-conversion="s2s-strategy=1" \
+// RUN:    --canonicalize --cse | FileCheck %s
+//
+// Now again with `kAuto` (the new default).
+// RUN: mlir-opt %s --sparse-tensor-conversion="s2s-strategy=0" \
+// RUN:    --canonicalize --cse | FileCheck %s -check-prefix=CHECKAUTO
 
 #SparseVector = #sparse_tensor.encoding<{
   dimLevelType = ["compressed"]
 }>
 
 #SparseVector64 = #sparse_tensor.encoding<{
   dimLevelType = ["compressed"],
   pointerBitWidth = 64,

 //   CHECK-DAG: %[[R:.*]] = memref.alloca() : memref<1xindex>
 //   CHECK-DAG: %[[X:.*]] = memref.cast %[[P]] : memref<1xi8> to memref<?xi8>
 //   CHECK-DAG: %[[Y:.*]] = memref.cast %[[Q]] : memref<1xindex> to memref<?xindex>
 //   CHECK-DAG: %[[Z:.*]] = memref.cast %[[R]] : memref<1xindex> to memref<?xindex>
 //       CHECK: %[[C:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %[[ToCOO]], %[[A]])
 //       CHECK: %[[T:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %[[FromCOO]], %[[C]])
 //       CHECK: call @delSparseTensorCOOF32(%[[C]])
 //       CHECK: return %[[T]] : !llvm.ptr<i8>
+// CHECKAUTO-LABEL: func @sparse_convert_1d_ss(
+//  CHECKAUTO-SAME: %[[A:.*]]: !llvm.ptr<i8>)
+//   CHECKAUTO-DAG: %[[SparseToSparse:.*]] = arith.constant 3 : i32
+//   CHECKAUTO-DAG: %[[P:.*]] = memref.alloca() : memref<1xi8>
+//   CHECKAUTO-DAG: %[[Q:.*]] = memref.alloca() : memref<1xindex>
+//   CHECKAUTO-DAG: %[[R:.*]] = memref.alloca() : memref<1xindex>
+//   CHECKAUTO-DAG: %[[X:.*]] = memref.cast %[[P]] : memref<1xi8> to memref<?xi8>
+//   CHECKAUTO-DAG: %[[Y:.*]] = memref.cast %[[Q]] : memref<1xindex> to memref<?xindex>
+//   CHECKAUTO-DAG: %[[Z:.*]] = memref.cast %[[R]] : memref<1xindex> to memref<?xindex>
+//       CHECKAUTO: %[[T:.*]] = call @newSparseTensor(%[[X]], %[[Y]], %[[Z]], %{{.*}}, %{{.*}}, %{{.*}}, %[[SparseToSparse]], %[[A]])
+//       CHECKAUTO: return %[[T]] : !llvm.ptr<i8>
 func.func @sparse_convert_1d_ss(%arg0: tensor<?xf32, #SparseVector64>) -> tensor<?xf32, #SparseVector32> {
   %0 = sparse_tensor.convert %arg0 : tensor<?xf32, #SparseVector64> to tensor<?xf32, #SparseVector32>
   return %0 : tensor<?xf32, #SparseVector32>
 }
 
 // CHECK-LABEL: func @sparse_convert_2d(
 //  CHECK-SAME: %[[A:.*]]: tensor<2x4xf64>) -> !llvm.ptr<i8>
 //   CHECK-DAG: %[[EmptyCOO:.*]] = arith.constant 4 : i32

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

+// Force this file to use the kDirect method for sparse2sparse.
+// RUN: mlir-opt %s --sparse-compiler="s2s-strategy=2" | \
+// RUN: mlir-cpu-runner -e entry -entry-point-result=void \
+// RUN:  -shared-libs=%mlir_integration_test_dir/libmlir_c_runner_utils%shlibext | \
+// RUN: FileCheck %s
+
+#Tensor1 = #sparse_tensor.encoding<{
+  dimLevelType = [ "dense", "dense", "compressed" ]
+}>
+
+// NOTE: dense after compressed is not currently supported for the target
+// of direct-sparse2sparse conversion.  (It's fine for the source though.)
+#Tensor2 = #sparse_tensor.encoding<{
+  dimLevelType = [ "dense", "compressed", "dense" ]
+}>
+
+#Tensor3 = #sparse_tensor.encoding<{
+  dimLevelType = [ "dense", "dense", "compressed" ],
+  dimOrdering = affine_map<(i,j,k) -> (i,k,j)>
+}>
+
+module {
+  //
+  // Utilities for output and releasing memory.
+  //
+  func.func @dump(%arg0: tensor<2x3x4xf64>) {
+    %c0 = arith.constant 0 : index
+    %d0 = arith.constant -1.0 : f64
+    %0 = vector.transfer_read %arg0[%c0, %c0, %c0], %d0: tensor<2x3x4xf64>, vector<2x3x4xf64>
+    vector.print %0 : vector<2x3x4xf64>
+    return
+  }
+  func.func @dumpAndRelease_234(%arg0: tensor<2x3x4xf64>) {
+    call @dump(%arg0) : (tensor<2x3x4xf64>) -> ()
+    %1 = bufferization.to_memref %arg0 : memref<2x3x4xf64>
+    memref.dealloc %1 : memref<2x3x4xf64>
+    return
+  }
+
+  //
+  // Main driver.
+  //
+  func.func @entry() {
+    //
+    // Initialize a 3-dim dense tensor.
+    //
+    %src = arith.constant dense<[
+       [  [  1.0,  2.0,  3.0,  4.0 ],
+          [  5.0,  6.0,  7.0,  8.0 ],
+          [  9.0, 10.0, 11.0, 12.0 ] ],
+       [  [ 13.0, 14.0, 15.0, 16.0 ],
+          [ 17.0, 18.0, 19.0, 20.0 ],
+          [ 21.0, 22.0, 23.0, 24.0 ] ]
+    ]> : tensor<2x3x4xf64>
+
+    //
+    // Convert dense tensor directly to various sparse tensors.
+    //
+    %s1 = sparse_tensor.convert %src : tensor<2x3x4xf64> to tensor<2x3x4xf64, #Tensor1>
+    %s2 = sparse_tensor.convert %src : tensor<2x3x4xf64> to tensor<2x3x4xf64, #Tensor2>
+    %s3 = sparse_tensor.convert %src : tensor<2x3x4xf64> to tensor<2x3x4xf64, #Tensor3>
+
+    //
+    // Convert sparse tensor directly to another sparse format.
+    //
+    %t13 = sparse_tensor.convert %s1 : tensor<2x3x4xf64, #Tensor1> to tensor<2x3x4xf64, #Tensor3>
+    %t21 = sparse_tensor.convert %s2 : tensor<2x3x4xf64, #Tensor2> to tensor<2x3x4xf64, #Tensor1>
+    %t23 = sparse_tensor.convert %s2 : tensor<2x3x4xf64, #Tensor2> to tensor<2x3x4xf64, #Tensor3>
+    %t31 = sparse_tensor.convert %s3 : tensor<2x3x4xf64, #Tensor3> to tensor<2x3x4xf64, #Tensor1>
+
+    //
+    // Convert sparse tensor back to dense.
+    //
+    %d13 = sparse_tensor.convert %t13 : tensor<2x3x4xf64, #Tensor3> to tensor<2x3x4xf64>
+    %d21 = sparse_tensor.convert %t21 : tensor<2x3x4xf64, #Tensor1> to tensor<2x3x4xf64>
+    %d23 = sparse_tensor.convert %t23 : tensor<2x3x4xf64, #Tensor3> to tensor<2x3x4xf64>
+    %d31 = sparse_tensor.convert %t31 : tensor<2x3x4xf64, #Tensor1> to tensor<2x3x4xf64>
+
+    //
+    // Check round-trip equality.  And release dense tensors.
+    //
+    // CHECK-COUNT-5: ( ( ( 1, 2, 3, 4 ), ( 5, 6, 7, 8 ), ( 9, 10, 11, 12 ) ), ( ( 13, 14, 15, 16 ), ( 17, 18, 19, 20 ), ( 21, 22, 23, 24 ) ) )
+    call @dump(%src) : (tensor<2x3x4xf64>) -> ()
+    call @dumpAndRelease_234(%d13) : (tensor<2x3x4xf64>) -> ()
+    call @dumpAndRelease_234(%d21) : (tensor<2x3x4xf64>) -> ()
+    call @dumpAndRelease_234(%d23) : (tensor<2x3x4xf64>) -> ()
+    call @dumpAndRelease_234(%d31) : (tensor<2x3x4xf64>) -> ()
+
+    //
+    // Release sparse tensors.
+    //
+    sparse_tensor.release %t13 : tensor<2x3x4xf64, #Tensor3>
+    sparse_tensor.release %t21 : tensor<2x3x4xf64, #Tensor1>
+    sparse_tensor.release %t23 : tensor<2x3x4xf64, #Tensor3>
+    sparse_tensor.release %t31 : tensor<2x3x4xf64, #Tensor1>
+    sparse_tensor.release %s1 : tensor<2x3x4xf64, #Tensor1>
+    sparse_tensor.release %s2 : tensor<2x3x4xf64, #Tensor2>
+    sparse_tensor.release %s3 : tensor<2x3x4xf64, #Tensor3>
+
+    return
+  }
+}

mlir/test/Integration/Dialect/SparseTensor/python/test_stress.py

def main():
 
   # CHECK-LABEL: TEST: test_stress
   print("\nTEST: test_stress")
   with ir.Context() as ctx, ir.Location.unknown():
     par = 0
     vec = 0
     vl = 1
     e = False
+    # Disable direct sparse2sparse conversion, because it doubles the time!
+    # TODO: While direct s2s is far too slow for per-commit testing,
+    # we should have some framework ensure that we run this test with
+    # `s2s=0` on a regular basis, to ensure that it does continue to work.
+    s2s = 1
     sparsification_options = (
         f'parallelization-strategy={par} '
         f'vectorization-strategy={vec} '
         f'vl={vl} '
-        f'enable-simd-index32={e}')
+        f'enable-simd-index32={e} '
+        f's2s-strategy={s2s}')
     compiler = sparse_compiler.SparseCompiler(
         options=sparsification_options, opt_level=0, shared_libs=[support_lib])
     f64 = ir.F64Type.get()
     # Be careful about increasing this because
     #     len(types) = 1 + 2^rank * rank! * len(bitwidths)^2
     shape = range(2, 6)
     rank = len(shape)
     # All combinations.