[mlir][sparse] avoid using mutable descriptor when unnecessary (NFC)

[Peiming]

Summary

Use SparseTensorDescriptor whenever not calling setters, to avoid needing to create a temporal buffer for simple query purposes.

[Peiming]

fix typo

[Peiming]

clean up

[Peiming]

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

107

I removed the optional as it never returns std::nullopt.

[wrengr]

The CL description seems a bit misleading/backwards to me —since it places focus on *using* MutSparseTensorDescriptor, rather than focusing on not using it. So would be better to rephrase it to something like:

"""Use SparseTensorDescriptor whenever not calling setters, to avoid needing to create a temporal buffer for simple query purposes."""

(Though I'm thinking you mean "temporary" here rather than "temporal".)

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

205

I'm curious why the change here? To you anticipate needing to define ValueArrayRef to something other than ValueRange or SmallVectorImpl<Value>& ? If there's a valid reason to expect other uses, then I'm all for generalization; but if there'll only ever be those two instantiations, then it'd be better to stick with the bool templating since it explicitly ensures the invariant that those are the only two options.

(If you revert to using the bool parameter, I think you should keep the new documentation on the subclasses. Since that documentation is remains valid and helpful under both ways of structuring the code.)

259

The compiler should be smart enough to fix it, but using the helper function here means doing a redundant assertion.

[Peiming]

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

205

I think changing it to ValueArrayRef makes it easier for readers to understand what is actually used for storage the array.

One reason I used bool at the first place is that I used to implement both mutable/immutable descriptor in this class and using enable_if to only allows setters for mutable descriptor. But the situation is changed, and this class now only holds getters, so I feel like that it make less sense to create a SparseTensorDescriptorImp<true> when it actually provides no setters.

259

Yeah, but should be okay? I would like to reuse the code here (maybe we will, though unlikely, switch to some class that does not provide operator[]).

[Peiming]

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

205

I think changing it to ValueArrayRef makes it easier for readers to understand what is actually used for storage the array.

*for storing the array*

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

Value sparse_tensor::genToValues(OpBuilder &builder, Location loc,
   RankedTensorType srcTp = tensor.getType().cast<RankedTensorType>();
   Type valTp = get1DMemRefType(srcTp.getElementType(),
                                /*withLayout=*/false);
   return builder.create<ToValuesOp>(loc, valTp, tensor);
 }
 
 Value sparse_tensor::genValMemSize(OpBuilder &builder, Location loc,
                                    Value tensor) {
-  SmallVector<Value> fields;
-  auto desc = getMutDescriptorFromTensorTuple(tensor, fields);
-  return desc.getValMemSize(builder, loc);
-}
+  return getDescriptorFromTensorTuple(tensor).getValMemSize(builder, loc);
+}
 

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

static scf::ForOp createFor(OpBuilder &builder, Location loc, Value upper,
   scf::ForOp forOp = builder.create<scf::ForOp>(loc, lower, upper, one, fields);
   for (unsigned i = 0, e = fields.size(); i < e; i++)
     fields[i] = forOp.getRegionIterArg(i);
   builder.setInsertionPointToStart(forOp.getBody());
   return forOp;
 }
 
 /// Gets the dimension size for the given sparse tensor at the given
-/// original dimension 'dim'. Returns std::nullopt if no sparse encoding is
-/// attached to the given tensor type.
-static std::optional<Value>
-sizeFromTensorAtDim(OpBuilder &builder, Location loc,
-                    const SparseTensorDescriptor &desc, unsigned dim) {
+/// original dimension 'dim'.
+static Value sizeFromTensorAtDim(OpBuilder &builder, Location loc,
+                                 SparseTensorDescriptor desc, unsigned dim) {
   RankedTensorType rtp = desc.getTensorType();
   // Access into static dimension can query original type directly.
   // Note that this is typically already done by DimOp's folding.
   auto shape = rtp.getShape();
   if (!ShapedType::isDynamic(shape[dim]))
     return constantIndex(builder, loc, shape[dim]);
 
   // Any other query can consult the dimSizes array at field DimSizesIdx,
   // accounting for the reordering applied to the sparse storage.
   return desc.getDimSize(builder, loc, toStoredDim(rtp, dim));
 }
 
-// Gets the dimension size at the given stored dimension 'd', either as a
+// Gets the dimension size at the given stored level 'lvl', either as a
 // constant for a static size, or otherwise dynamically through memSizes.
-Value sizeAtStoredDim(OpBuilder &builder, Location loc,
-                      MutSparseTensorDescriptor desc, unsigned d) {
-  RankedTensorType rtp = desc.getTensorType();
-  unsigned dim = toOrigDim(rtp, d);
-  auto shape = rtp.getShape();
-  if (!ShapedType::isDynamic(shape[dim]))
-    return constantIndex(builder, loc, shape[dim]);
-
-  return desc.getDimSize(builder, loc, d);
+static Value sizeFromTensorAtLvl(OpBuilder &builder, Location loc,
+                                 SparseTensorDescriptor desc, unsigned lvl) {
+  return sizeFromTensorAtDim(builder, loc, desc,
+                             toOrigDim(desc.getTensorType(), lvl));
 }
 
 static void createPushback(OpBuilder &builder, Location loc,
                            MutSparseTensorDescriptor desc,
                            SparseTensorFieldKind kind,
                            std::optional<unsigned> dim, Value value,
                            Value repeat = Value()) {
   Type etp = desc.getMemRefElementType(kind, dim);
for (unsigned r = r0; r < rank; r++) {
     }
     if (isSingletonDim(rtp, r)) {
       return; // nothing to do
     }
     // Keep compounding the size, but nothing needs to be initialized
     // at this level. We will eventually reach a compressed level or
     // otherwise the values array for the from-here "all-dense" case.
     assert(isDenseDim(rtp, r));
-    Value size = sizeAtStoredDim(builder, loc, desc, r);
+    Value size = sizeFromTensorAtLvl(builder, loc, desc, r);
     linear = builder.create<arith::MulIOp>(loc, linear, size);
   }
   // Reached values array so prepare for an insertion.
   Value valZero = constantZero(builder, loc, rtp.getElementType());
   createPushback(builder, loc, desc, SparseTensorFieldKind::ValMemRef,
                  std::nullopt, valZero, linear);
 }
 
if (isCompressedDim(rtp, d)) {
       //   <insert @ pos[d] at next dimension d + 1>
       createPushback(builder, loc, desc, SparseTensorFieldKind::IdxMemRef, d,
                      indices[d]);
     } else {
       assert(isDenseDim(rtp, d));
       // Construct the new position as:
       //   pos[d] = size * pos[d-1] + i[d]
       //   <insert @ pos[d] at next dimension d + 1>
-      Value size = sizeAtStoredDim(builder, loc, desc, d);
+      Value size = sizeFromTensorAtLvl(builder, loc, desc, d);
       Value mult = builder.create<arith::MulIOp>(loc, size, pos);
       pos = builder.create<arith::AddIOp>(loc, mult, indices[d]);
     }
   }
   // Reached the actual value append/insert.
   if (!isDenseDim(rtp, rank - 1))
     createPushback(builder, loc, desc, SparseTensorFieldKind::ValMemRef,
                    std::nullopt, value);
static void genInsertionCallHelper(OpBuilder &builder,
   func::CallOp call = builder.create<func::CallOp>(loc, func, operands);
   for (size_t i = 0, e = desc.getNumFields(); i < e; i++) {
     desc.getFields()[i] = call.getResult(i);
   }
 }
 
 /// Generations insertion finalization code.
 static void genEndInsert(OpBuilder &builder, Location loc,
-                         MutSparseTensorDescriptor desc) {
+                         SparseTensorDescriptor desc) {
   RankedTensorType rtp = desc.getTensorType();
   unsigned rank = rtp.getShape().size();
   for (unsigned d = 0; d < rank; d++) {
     if (isCompressedDim(rtp, d)) {
       // Compressed dimensions need a pointer cleanup for all entries
       // that were not visited during the insertion pass.
       //
       // TODO: avoid cleanup and keep compressed scheme consistent at all
matchAndRewrite(tensor::DimOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     std::optional<int64_t> index = op.getConstantIndex();
     if (!index || !getSparseTensorEncoding(adaptor.getSource().getType()))
       return failure();
 
     auto desc = getDescriptorFromTensorTuple(adaptor.getSource());
     auto sz = sizeFromTensorAtDim(rewriter, op.getLoc(), desc, *index);
 
-    if (!sz)
-      return failure();
-
-    rewriter.replaceOp(op, *sz);
+    rewriter.replaceOp(op, sz);
     return success();
   }
 };
 
 /// Sparse codegen rule for trivial tensor casts.
 class SparseCastConverter : public OpConversionPattern<tensor::CastOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
public:
   matchAndRewrite(bufferization::DeallocTensorOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     auto enc = getSparseTensorEncoding(op.getTensor().getType());
     if (!enc)
       return failure();
 
     // Replace the sparse tensor deallocation with field deallocations.
     Location loc = op.getLoc();
-    SmallVector<Value> fields;
-    auto desc = getMutDescriptorFromTensorTuple(adaptor.getTensor(), fields);
+    auto desc = getDescriptorFromTensorTuple(adaptor.getTensor());
     for (auto input : desc.getMemRefFields())
       // Deallocate every buffer used to store the sparse tensor handler.
       rewriter.create<memref::DeallocOp>(loc, input);
 
     rewriter.eraseOp(op);
     return success();
   }
 };
 
 /// Sparse codegen rule for tensor rematerialization.
 class SparseTensorLoadConverter : public OpConversionPattern<LoadOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(LoadOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     // Prepare descriptor.
-    SmallVector<Value> fields;
-    auto desc = getMutDescriptorFromTensorTuple(adaptor.getTensor(), fields);
+    auto desc = getDescriptorFromTensorTuple(adaptor.getTensor());
     // Generate optional insertion finalization code.
     if (op.getHasInserts())
       genEndInsert(rewriter, op.getLoc(), desc);
     // Replace operation with resulting memrefs.
     rewriter.replaceOp(op, genTuple(rewriter, op.getLoc(), desc));
     return success();
   }
 };
matchAndRewrite(ExpandOp op, OpAdaptor adaptor,
     Type idxType = rewriter.getIndexType();
     // All initialization should be done on entry of the loop nest.
     rewriter.setInsertionPointAfter(op.getTensor().getDefiningOp());
     // Determine the size for access expansion (always the innermost stored
     // dimension size, translated back to original dimension). Note that we
     // recursively rewrite the new DimOp on the **original** tensor.
     unsigned innerDim = toOrigDim(srcType, srcType.getRank() - 1);
     auto sz = sizeFromTensorAtDim(rewriter, loc, desc, innerDim);
-    assert(sz); // This for sure is a sparse tensor
     // Generate a memref for `sz` elements of type `t`.
     auto genAlloc = [&](Type t) {
       auto memTp = MemRefType::get({ShapedType::kDynamic}, t);
-      return rewriter.create<memref::AllocOp>(loc, memTp, ValueRange{*sz});
+      return rewriter.create<memref::AllocOp>(loc, memTp, ValueRange{sz});
     };
     // Allocate temporary buffers for values/filled-switch and added.
     // 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(eltType);
     Value filled = genAlloc(boolType);
     Value added = genAlloc(idxType);
     Value zero = constantZero(rewriter, loc, idxType);
public:
   using OpAdaptor = typename ToIndicesBufferOp::Adaptor;
   using OpConversionPattern<ToIndicesBufferOp>::OpConversionPattern;
   LogicalResult
   matchAndRewrite(ToIndicesBufferOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     // Replace the requested pointer access with corresponding field.
     // The cast_op is inserted by type converter to intermix 1:N type
     // conversion.
-    SmallVector<Value> fields;
-    auto desc = getMutDescriptorFromTensorTuple(adaptor.getTensor(), fields);
+    auto desc = getDescriptorFromTensorTuple(adaptor.getTensor());
     rewriter.replaceOp(op, desc.getAOSMemRef());
 
     return success();
   }
 };
 
 /// Sparse codegen rule for value accesses.
 class SparseToValuesConverter : public OpConversionPattern<ToValuesOp> {

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

private:
   TypedValue<StorageSpecifierType> specifier;
 };
 
 /// A helper class around an array of values that corresponding to a sparse
 /// tensor, provides a set of meaningful APIs to query and update a particular
 /// field in a consistent way.
 /// Users should not make assumption on how a sparse tensor is laid out but
 /// instead relies on this class to access the right value for the right field.
-template <bool mut>
+template <typename ValueArrayRef>
 class SparseTensorDescriptorImpl {
 protected:
-  // Uses ValueRange for immuatable descriptors; uses SmallVectorImpl<Value> &
-  // for mutable descriptors.
-  // Using SmallVector for mutable descriptor allows users to reuse it as a tmp
-  // buffers to append value for some special cases, though users should be
-  // responsible to restore the buffer to legal states after their use. It is
-  // probably not a clean way, but it is the most efficient way to avoid copying
-  // the fields into another SmallVector. If a more clear way is wanted, we
-  // should change it to MutableArrayRef instead.
-  using ValueArrayRef = typename std::conditional<mut, SmallVectorImpl<Value> &,
-                                                  ValueRange>::type;
-
   SparseTensorDescriptorImpl(Type tp, ValueArrayRef fields)
       : rType(tp.cast<RankedTensorType>()), fields(fields) {
     assert(getSparseTensorEncoding(tp) &&
            getNumFieldsFromEncoding(getSparseTensorEncoding(tp)) ==
                fields.size());
     // We should make sure the class is trivially copyable (and should be small
     // enough) such that we can pass it by value.
-    static_assert(
-        std::is_trivially_copyable_v<SparseTensorDescriptorImpl<mut>>);
+    static_assert(std::is_trivially_copyable_v<
+                  SparseTensorDescriptorImpl<ValueArrayRef>>);
   }
 
 public:
   unsigned getMemRefFieldIndex(SparseTensorFieldKind kind,
                                std::optional<unsigned> dim) const {
     // Delegates to storage layout.
     StorageLayout layout(getSparseTensorEncoding(rType));
     return layout.getMemRefFieldIndex(kind, dim);
public:
   }
 
   Value getValMemRef() const {
     return getMemRefField(SparseTensorFieldKind::ValMemRef, std::nullopt);
   }
 
   Value getMemRefField(SparseTensorFieldKind kind,
                        std::optional<unsigned> dim) const {
-    return fields[getMemRefFieldIndex(kind, dim)];
+    return getField(getMemRefFieldIndex(kind, dim));
   }
 
   Value getMemRefField(unsigned fidx) const {
     assert(fidx < fields.size() - 1);
-    return fields[fidx];
+    return getField(fidx);
   }
 
   Value getPtrMemSize(OpBuilder &builder, Location loc, unsigned dim) const {
     return getSpecifierField(builder, loc, StorageSpecifierKind::PtrMemSize,
                              dim);
   }
 
   Value getIdxMemSize(OpBuilder &builder, Location loc, unsigned dim) const {
public:
   Type getMemRefElementType(SparseTensorFieldKind kind,
                             std::optional<unsigned> dim) const {
     return getMemRefField(kind, dim)
         .getType()
         .template cast<MemRefType>()
         .getElementType();
   }
 
+  Value getField(unsigned fidx) const {
+    assert(fidx < fields.size());
+    return fields[fidx];
+  }
+
+  ValueRange getMemRefFields() const {
+    ValueRange ret = fields;
+    // Drop the last metadata fields.
+    return ret.slice(0, fields.size() - 1);
+  }
+
+  std::pair<unsigned, unsigned>
+  getIdxMemRefIndexAndStride(unsigned idxDim) const {
+    StorageLayout layout(getSparseTensorEncoding(rType));
+    return layout.getFieldIndexAndStride(SparseTensorFieldKind::IdxMemRef,
+                                         idxDim);
+  }
+
+  Value getAOSMemRef() const {
+    auto enc = getSparseTensorEncoding(rType);
+    unsigned cooStart = getCOOStart(enc);
+    assert(cooStart < enc.getDimLevelType().size());
+    return getMemRefField(SparseTensorFieldKind::IdxMemRef, cooStart);
+  }
+
   RankedTensorType getTensorType() const { return rType; }
   ValueArrayRef getFields() const { return fields; }
 
 protected:
   RankedTensorType rType;
   ValueArrayRef fields;
 };
 
-class MutSparseTensorDescriptor : public SparseTensorDescriptorImpl<true> {
+/// Uses ValueRange for immuatable descriptors;
+class SparseTensorDescriptor : public SparseTensorDescriptorImpl<ValueRange> {
 public:
-  MutSparseTensorDescriptor(Type tp, ValueArrayRef buffers)
-      : SparseTensorDescriptorImpl<true>(tp, buffers) {}
+  SparseTensorDescriptor(Type tp, ValueRange buffers)
+      : SparseTensorDescriptorImpl<ValueRange>(tp, buffers) {}
 
-  Value getField(unsigned fidx) const {
-    assert(fidx < fields.size());
-    return fields[fidx];
-  }
+  Value getIdxMemRefOrView(OpBuilder &builder, Location loc,
+                           unsigned idxDim) const;
+};
 
-  ValueRange getMemRefFields() const {
-    ValueRange ret = fields;
-    // Drop the last metadata fields.
-    return ret.slice(0, fields.size() - 1);
+/// Uses SmallVectorImpl<Value> & for mutable descriptors.
+/// Using SmallVector for mutable descriptor allows users to reuse it as a
+/// tmp buffers to append value for some special cases, though users should
+/// be responsible to restore the buffer to legal states after their use. It
+/// is probably not a clean way, but it is the most efficient way to avoid
+/// copying the fields into another SmallVector. If a more clear way is
+/// wanted, we should change it to MutableArrayRef instead.
+class MutSparseTensorDescriptor
+    : public SparseTensorDescriptorImpl<SmallVectorImpl<Value> &> {
+public:
+  MutSparseTensorDescriptor(Type tp, SmallVectorImpl<Value> &buffers)
+      : SparseTensorDescriptorImpl<SmallVectorImpl<Value> &>(tp, buffers) {}
+
+  // Allow implicit type conversion from mutable descriptors to immutable ones
+  // (but not vice versa).
+  /*implicit*/ operator SparseTensorDescriptor() const {
+    return SparseTensorDescriptor(rType, fields);
   }
 
   ///
-  /// Setters: update the value for required field (only enabled for
-  /// MutSparseTensorDescriptor).
+  /// Adds additional setters for mutable descriptor, update the value for
+  /// required field.
   ///
 
   void setMemRefField(SparseTensorFieldKind kind, std::optional<unsigned> dim,
                       Value v) {
     fields[getMemRefFieldIndex(kind, dim)] = v;
   }
 
   void setMemRefField(unsigned fidx, Value v) {
void setSpecifierField(OpBuilder &builder, Location loc,
     SparseTensorSpecifier md(fields.back());
     md.setSpecifierField(builder, loc, v, kind, dim);
     fields.back() = md;
   }
 
   void setDimSize(OpBuilder &builder, Location loc, unsigned dim, Value v) {
     setSpecifierField(builder, loc, StorageSpecifierKind::DimSize, dim, v);
   }
-
-  std::pair<unsigned, unsigned>
-  getIdxMemRefIndexAndStride(unsigned idxDim) const {
-    StorageLayout layout(getSparseTensorEncoding(rType));
-    return layout.getFieldIndexAndStride(SparseTensorFieldKind::IdxMemRef,
-                                         idxDim);
-  }
-
-  Value getAOSMemRef() const {
-    auto enc = getSparseTensorEncoding(rType);
-    unsigned cooStart = getCOOStart(enc);
-    assert(cooStart < enc.getDimLevelType().size());
-    return getMemRefField(SparseTensorFieldKind::IdxMemRef, cooStart);
-  }
-};
-
-class SparseTensorDescriptor : public SparseTensorDescriptorImpl<false> {
-public:
-  SparseTensorDescriptor(Type tp, ValueArrayRef buffers)
-      : SparseTensorDescriptorImpl<false>(tp, buffers) {}
-
-  Value getIdxMemRefOrView(OpBuilder &builder, Location loc,
-                           unsigned idxDim) const;
 };
 
 /// Returns the "tuple" value of the adapted tensor.
 inline UnrealizedConversionCastOp getTuple(Value tensor) {
   return llvm::cast<UnrealizedConversionCastOp>(tensor.getDefiningOp());
 }
 
 /// Packs the given values as a "tuple" value.
 inline Value genTuple(OpBuilder &builder, Location loc, Type tp,
                       ValueRange values) {
   return builder.create<UnrealizedConversionCastOp>(loc, TypeRange(tp), values)
       .getResult(0);
 }
 
 inline Value genTuple(OpBuilder &builder, Location loc,
-                      MutSparseTensorDescriptor desc) {
+                      SparseTensorDescriptor desc) {
   return genTuple(builder, loc, desc.getTensorType(), desc.getFields());
 }
 
 inline SparseTensorDescriptor getDescriptorFromTensorTuple(Value tensor) {
   auto tuple = getTuple(tensor);
   return SparseTensorDescriptor(tuple.getResultTypes()[0], tuple.getInputs());
 }