[mlir][sparse] Move some member functions from SparseTensorDescriptorImpl to MutSparseTensorDescriptor.

[bixia]

Summary

This is to prepare for implementing AOS optimization.

[aartbik]

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

970

While here, can you update the comment to

actually stored size.

(singular, and value did not add much)

972

It is a bit strange that we need the mutable class now to query a read-only field.
As discussed offline, perhaps these classes mean different things now.

Also, even in the "new" class, can we have a version that does not need passing in the dummy fields?

[bixia]

Fix comments, add codegen util for generating code to access values memref size.

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

972

Per offline discussion, move this to codegen utils.

[bixia]

Applied clang-format.

[Peiming]

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

301

This template parameter is no longer needed, the sole purpose of it was to enable SFINAE

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

 /// Sparse codegen rule for number of entries operator.
 class SparseNumberOfEntriesConverter
     : public OpConversionPattern<NumberOfEntriesOp> {
 public:
   using OpConversionPattern::OpConversionPattern;
   LogicalResult
   matchAndRewrite(NumberOfEntriesOp op, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     // Query memSizes for the actually stored values size.
-    auto desc = getDescriptorFromTensorTuple(adaptor.getTensor());
+    SmallVector<Value> fields;
+    auto desc = getMutDescriptorFromTensorTuple(adaptor.getTensor(), fields);
     rewriter.replaceOp(op, desc.getValMemSize(rewriter, op.getLoc()));
     return success();
   }
 };
 
 } // namespace
 
 //===----------------------------------------------------------------------===//

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

 
 /// 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>
 class SparseTensorDescriptorImpl {
-private:
+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.
Value getSpecifierField(OpBuilder &builder, Location loc,
     SparseTensorSpecifier md(fields.back());
     return md.getSpecifierField(builder, loc, kind, dim);
   }
 
   Value getDimSize(OpBuilder &builder, Location loc, unsigned dim) const {
     return getSpecifierField(builder, loc, StorageSpecifierKind::DimSize, dim);
   }
 
-  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 {
-    return getSpecifierField(builder, loc, StorageSpecifierKind::IdxMemSize,
-                             dim);
-  }
-
-  Value getValMemSize(OpBuilder &builder, Location loc) const {
-    return getSpecifierField(builder, loc, StorageSpecifierKind::ValMemSize,
-                             std::nullopt);
-  }
-
   Value getPtrMemRef(unsigned ptrDim) const {
     return getMemRefField(SparseTensorFieldKind::PtrMemRef, ptrDim);
   }
 
   Value getIdxMemRef(unsigned idxDim) const {
     return getMemRefField(SparseTensorFieldKind::IdxMemRef, idxDim);
   }
 
Value getMemRefField(unsigned fidx) const {
     return fields[fidx];
   }
 
   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);
+  }
+
+  Type getMemRefElementType(SparseTensorFieldKind kind,
+                            Optional<unsigned> dim) const {
+    return getMemRefField(kind, dim)
+        .getType()
+        .template cast<MemRefType>()
+        .getElementType();
+  }
+
+  RankedTensorType getTensorType() const { return rType; }
+  ValueArrayRef getFields() const { return fields; }
+
+protected:
+  RankedTensorType rType;
+  ValueArrayRef fields;
+};
+
+class MutSparseTensorDescriptor : public SparseTensorDescriptorImpl<true> {
+public:
+  MutSparseTensorDescriptor(Type tp, ValueArrayRef buffers)
+      : SparseTensorDescriptorImpl<true>(tp, buffers) {}
+
+  ///
+  /// Getters: get the value for required field.
+  ///
+
+  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 {
+    return getSpecifierField(builder, loc, StorageSpecifierKind::IdxMemSize,
+                             dim);
+  }
+
+  Value getValMemSize(OpBuilder &builder, Location loc) const {
+    return getSpecifierField(builder, loc, StorageSpecifierKind::ValMemSize,
+                             std::nullopt);
+  }
+
   ///
   /// Setters: update the value for required field (only enabled for
   /// MutSparseTensorDescriptor).
   ///
 
   template <typename T = Value>
-  void setMemRefField(SparseTensorFieldKind kind, Optional<unsigned> dim,
-                      std::enable_if_t<mut, T> v) {
+  void setMemRefField(SparseTensorFieldKind kind, Optional<unsigned> dim, T v) {
     fields[getMemRefFieldIndex(kind, dim)] = v;
   }
 
-  template <typename T = Value>
-  void setMemRefField(unsigned fidx, std::enable_if_t<mut, T> v) {
+  template <typename T = Value> void setMemRefField(unsigned fidx, T v) {
     assert(fidx < fields.size() - 1);
     fields[fidx] = v;
   }
 
-  template <typename T = Value>
-  void setField(unsigned fidx, std::enable_if_t<mut, T> v) {
+  template <typename T = Value> void setField(unsigned fidx, T v) {
     assert(fidx < fields.size());
     fields[fidx] = v;
   }
 
   template <typename T = Value>
   void setSpecifierField(OpBuilder &builder, Location loc,
                          StorageSpecifierKind kind, Optional<unsigned> dim,
-                         std::enable_if_t<mut, T> v) {
+                         T v) {
     SparseTensorSpecifier md(fields.back());
     md.setSpecifierField(builder, loc, v, kind, dim);
     fields.back() = md;
   }
 
   template <typename T = Value>
-  void setDimSize(OpBuilder &builder, Location loc, unsigned dim,
-                  std::enable_if_t<mut, T> v) {
+  void setDimSize(OpBuilder &builder, Location loc, unsigned dim, T v) {
     setSpecifierField(builder, loc, StorageSpecifierKind::DimSize, dim, v);
   }
-
-  ValueRange getMemRefFields() const {
-    ValueRange ret = fields;
-    // drop the last metadata fields
-    return ret.slice(0, fields.size() - 1);
-  }
-
-  Type getMemRefElementType(SparseTensorFieldKind kind,
-                            Optional<unsigned> dim) const {
-    return getMemRefField(kind, dim)
-        .getType()
-        .template cast<MemRefType>()
-        .getElementType();
-  }
-
-  RankedTensorType getTensorType() const { return rType; }
-  ValueArrayRef getFields() const { return fields; }
-
-private:
-  RankedTensorType rType;
-  ValueArrayRef fields;
 };
 
 using SparseTensorDescriptor = SparseTensorDescriptorImpl<false>;
-using MutSparseTensorDescriptor = SparseTensorDescriptorImpl<true>;
 
 /// 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,