[mlir][sparse] adding `SparseTensorType::get{Pointer,Index}Type` methods

[wrengr]

Summary

Depends On D143800

[aartbik]

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorType.h

237

Comment on pointer/index seems interchanged?

[wrengr]

correcting swapped documentiation

[wrengr]

git-clang-format

[wrengr]

rebase

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensor.h

 
 /// [deprecated] Convenience method to translate the given dimension to
 /// the corresponding level.  Requires: `0 <= d < dimRank`.
 DEPRECATED Level toStoredDim(SparseTensorEncodingAttr enc, Dimension d);
 DEPRECATED Level toStoredDim(RankedTensorType type, Dimension d);
 
 #undef DEPRECATED
 
+namespace detail {
+Type getIntegerOrIndexType(MLIRContext *ctx, unsigned bitwidth);
+} // namespace detail
+
 } // namespace sparse_tensor
 } // namespace mlir
 
 #endif // MLIR_DIALECT_SPARSETENSOR_IR_SPARSETENSOR_H_

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorType.h

public:
   /// Returns the index-overhead bitwidth, defaulting to zero.
   unsigned getIndexBitWidth() const { return enc ? enc.getIndexBitWidth() : 0; }
 
   /// Returns the pointer-overhead bitwidth, defaulting to zero.
   unsigned getPointerBitWidth() const {
     return enc ? enc.getPointerBitWidth() : 0;
   }
 
+  /// Returns the index-overhead MLIR type, defaulting to `IndexType`.
+  Type getIndexType() const {
+    return detail::getIntegerOrIndexType(getContext(), getIndexBitWidth());
+  }
+
+  /// Returns the pointer-overhead MLIR type, defaulting to `IndexType`.
+  Type getPointerType() const {
+    return detail::getIntegerOrIndexType(getContext(), getPointerBitWidth());
+  }
+
 private:
   // These two must be const, to ensure coherence of the memoized fields.
   const RankedTensorType rtp;
   const SparseTensorEncodingAttr enc;
   // Memoized to avoid frequent redundant conditionals.
   const Level lvlRank;
   const AffineMap dim2lvl;
 };

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

if ((offset == SparseTensorDimSliceAttr::kDynamic || offset >= 0) &&
       (size == SparseTensorDimSliceAttr::kDynamic || size > 0) &&
       (stride == SparseTensorDimSliceAttr::kDynamic || stride > 0)) {
     return success();
   }
   return emitError()
          << "expect positive value or ? for slice offset/size/stride";
 }
 
-static Type getIntegerOrIndexType(MLIRContext *ctx, unsigned bitwidth) {
+Type mlir::sparse_tensor::detail::getIntegerOrIndexType(MLIRContext *ctx,
+                                                        unsigned bitwidth) {
   if (bitwidth)
     return IntegerType::get(ctx, bitwidth);
   return IndexType::get(ctx);
 }
 
 Type SparseTensorEncodingAttr::getPointerType() const {
-  return getIntegerOrIndexType(getContext(), getPointerBitWidth());
+  return detail::getIntegerOrIndexType(getContext(), getPointerBitWidth());
 }
 
 Type SparseTensorEncodingAttr::getIndexType() const {
-  return getIntegerOrIndexType(getContext(), getIndexBitWidth());
+  return detail::getIntegerOrIndexType(getContext(), getIndexBitWidth());
 }
 
 SparseTensorEncodingAttr SparseTensorEncodingAttr::withoutOrdering() const {
   return SparseTensorEncodingAttr::get(
       getContext(), getDimLevelType(), AffineMap(), AffineMap(),
       getPointerBitWidth(), getIndexBitWidth());
 }
 

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

static void allocSchemeForRank(OpBuilder &builder, Location loc,
   Value linear = constantIndex(builder, loc, 1);
   const Level lvlRank = stt.getLvlRank();
   for (Level l = startLvl; l < lvlRank; l++) {
     const auto dlt = stt.getLvlType(l);
     if (isCompressedDLT(dlt)) {
       // Append linear x pointers, initialized to zero. Since each compressed
       // dimension initially already has a single zero entry, this maintains
       // the desired "linear + 1" length property at all times.
-      Type ptrType = stt.getEncoding().getPointerType();
+      Type ptrType = stt.getPointerType();
       Value ptrZero = constantZero(builder, loc, ptrType);
       createPushback(builder, loc, desc, SparseTensorFieldKind::PtrMemRef, l,
                      ptrZero, linear);
       return;
     }
     if (isSingletonDLT(dlt)) {
       return; // nothing to do
     }
foreachFieldAndTypeInSparseTensor(
       });
 
   MutSparseTensorDescriptor desc(stt, fields);
 
   // Initialize the storage scheme to an empty tensor. Initialized memSizes
   // to all zeros, sets the dimSizes to known values and gives all pointer
   // fields an initial zero entry, so that it is easier to maintain the
   // "linear + 1" length property.
-  Value ptrZero =
-      constantZero(builder, loc, stt.getEncoding().getPointerType());
+  Value ptrZero = constantZero(builder, loc, stt.getPointerType());
   for (Level lvlRank = stt.getLvlRank(), l = 0; l < lvlRank; l++) {
     // Fills dim sizes array.
     // FIXME: this method seems to set *level* sizes, but the name is confusing
     // FIXME: `toOrigDim` is deprecated.
     desc.setDimSize(builder, loc, l, dimSizes[toOrigDim(stt, l)]);
     // Pushes a leading zero to pointers memref.
     if (stt.isCompressedLvl(l))
       createPushback(builder, loc, desc, SparseTensorFieldKind::PtrMemRef, l,
for (Level l = 0; l < lvlRank; l++) {
     if (isCompressedDLT(dlt)) {
       // 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
       // times?
       //
       if (l > 0) {
-        Type ptrType = stt.getEncoding().getPointerType();
+        Type ptrType = stt.getPointerType();
         Value ptrMemRef = desc.getPtrMemRef(l);
         Value hi = desc.getPtrMemSize(builder, loc, l);
         Value zero = constantIndex(builder, loc, 0);
         Value one = constantIndex(builder, loc, 1);
         // Vector of only one, but needed by createFor's prototype.
         SmallVector<Value, 1> inits{genLoad(builder, loc, ptrMemRef, zero)};
         scf::ForOp loop = createFor(builder, loc, hi, inits, one);
         Value i = loop.getInductionVar();

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

 #undef RETURN_ON_FALSE
 }
 
 void sparse_tensor::foreachFieldAndTypeInSparseTensor(
     SparseTensorType stt,
     llvm::function_ref<bool(Type, FieldIndex, SparseTensorFieldKind, Level,
                             DimLevelType)>
         callback) {
-  const auto enc = stt.getEncoding();
-  assert(enc);
+  assert(stt.hasEncoding());
   // Construct the basic types.
-  Type idxType = enc.getIndexType();
-  Type ptrType = enc.getPointerType();
+  Type idxType = stt.getIndexType();
+  Type ptrType = stt.getPointerType();
   Type eltType = stt.getElementType();
 
-  Type metaDataType = StorageSpecifierType::get(enc);
+  Type metaDataType = StorageSpecifierType::get(stt.getEncoding());
   // memref<? x ptr>  pointers
   Type ptrMemType = MemRefType::get({ShapedType::kDynamic}, ptrType);
   // memref<? x idx>  indices
   Type idxMemType = MemRefType::get({ShapedType::kDynamic}, idxType);
   // memref<? x eltType> values
   Type valMemType = MemRefType::get({ShapedType::kDynamic}, eltType);
 
   foreachFieldInSparseTensor(
-      enc,
+      stt.getEncoding(),
       [metaDataType, ptrMemType, idxMemType, valMemType,
        callback](FieldIndex fieldIdx, SparseTensorFieldKind fieldKind,
                  Level lvl, DimLevelType dlt) -> bool {
         switch (fieldKind) {
         case SparseTensorFieldKind::StorageSpec:
           return callback(metaDataType, fieldIdx, fieldKind, lvl, dlt);
         case SparseTensorFieldKind::PtrMemRef:
           return callback(ptrMemType, fieldIdx, fieldKind, lvl, dlt);