[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

969–970

While here, can you update the comment to

actually stored size.

(singular, and value did not add much)

970

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

970

Per offline discussion, move this to codegen utils.

[bixia]

Applied clang-format.

[Peiming]

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

302

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

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

 /// region, the result type is a memref with unknown stride and offset.
 /// Otherwise, the result type is a memref without any specified layout.
 Value genToIndices(OpBuilder &builder, Location loc, Value tensor, uint64_t d,
                    uint64_t cooStart);
 
 /// Infers the result type and generates ToValuesOp.
 Value genToValues(OpBuilder &builder, Location loc, Value tensor);
 
+/// Generates code to retrieve the values size for the sparse tensor.
+Value genValMemSize(OpBuilder &builder, Location loc, Value tensor);
+
 } // namespace sparse_tensor
 } // namespace mlir
 
 #endif // MLIR_DIALECT_SPARSETENSOR_TRANSFORMS_CODEGENUTILS_H_

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

 //===- CodegenUtils.cpp - Utilities for generating MLIR -------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "CodegenUtils.h"
+#include "SparseTensorStorageLayout.h"
 
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Bufferization/IR/Bufferization.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/IR/Matchers.h"

 }
 
 Value sparse_tensor::genToValues(OpBuilder &builder, Location loc,
                                  Value tensor) {
   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);
 }
 No newline at end of file

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());
-    rewriter.replaceOp(op, desc.getValMemSize(rewriter, op.getLoc()));
+    // Query memSizes for the actually stored values.
+    rewriter.replaceOp(
+        op, genValMemSize(rewriter, op.getLoc(), adaptor.getTensor()));
     return success();
   }
 };
 
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // Public method for populating conversion rules.

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) {
+  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) {
+  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,