[mlir][sparse] add getPointerType/getIndexType to SparseTensorEncodingAttr.

[Peiming]

Summary

add new interfaces to SparseTensorEncodingAttr to construct the pointer/index types based on pointer/index bitwidth.

[wrengr]

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

47

You should also pass in the signedness semantics here. The default is Signless but we really do mean Unsigned.

53

ditto.

[wrengr]

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td

162

Should this use "///" so that the comment is picked up by doxygen? (I don't know what the style guide is re extraClassDeclaration)

165

ditto

[Peiming]

rebase against main

[Peiming]

change comment style.

[Peiming]

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td

162

You are right, I checked other td file to confirm.

[wrengr]

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

46

Any reason why this floated out rather than inlined at the use site? If you instead floated out auto ctx = getContext(); you might be able to fit everything on one line, if that's the concern

47

After some offline discussion, I rescind this suggestion. Let's handle that in a separate CL instead.

52

ditto

[aartbik]

Still LGTM, if you fix Wren's doc suggestion.

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

47

Discussed internally, keep for now i32/i64 in this refactoring, and investigate ui32/ui64 independently.

[aartbik]

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

213–214

Should this go to next line?

[Peiming]

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

46

Yes, because IndexType and IntegerType are incompatible to be used by ?: operator

[Peiming]

fix comments.

mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td

let parameters = (
     // A mapping between the original and higher-ordering index space.
     "AffineMap":$higherOrdering,
     // The required bit width for pointer storage.
     "unsigned":$pointerBitWidth,
     // The required bit width for index storage.
     "unsigned":$indexBitWidth
   );
 
+  let extraClassDeclaration = [{
+    // Returns the type for pointer storage based on pointerBitWidth
+    Type getPointerType() const;
+
+    // Returns the type for index storage based on indexBitWidth
+    Type getIndexType() const;
+  }];
+
   let genVerifyDecl = 1;
   let hasCustomAssemblyFormat = 1;
 }
 
 def IsSparseTensorPred
   : CPred<"!!::mlir::sparse_tensor::getSparseTensorEncoding($_self)">;
 
 // The following four follow the same idiom as `TensorOf`, `AnyTensor`,

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

static bool acceptBitWidth(unsigned bitWidth) {
   case 32:
   case 64:
     return true;
   default:
     return false;
   }
 }
 
+Type SparseTensorEncodingAttr::getPointerType() const {
+  unsigned ptrWidth = getPointerBitWidth();
+  Type indexType = IndexType::get(getContext());
+  return ptrWidth ? IntegerType::get(getContext(), ptrWidth) : indexType;
+}
+
+Type SparseTensorEncodingAttr::getIndexType() const {
+  unsigned idxWidth = getIndexBitWidth();
+  Type indexType = IndexType::get(getContext());
+  return idxWidth ? IntegerType::get(getContext(), idxWidth) : indexType;
+}
+
 Attribute SparseTensorEncodingAttr::parse(AsmParser &parser, Type type) {
   if (failed(parser.parseLess()))
     return {};
   // Parse the data as a dictionary.
   DictionaryAttr dict;
   if (failed(parser.parseAttribute(dict)))
     return {};
   if (failed(parser.parseGreater()))

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

public:
 
   RankedTensorType getTensorType() const { return rType; }
   Storage getFields() const { return fields; }
 
   Type getElementType(unsigned fidx) const {
     return fields[fidx].getType().template cast<MemRefType>().getElementType();
   }
 
-  // TODO: a better places for these functions should be in
-  // SparseTensorEncodingAttr.
-  Type getPtrElementType() const {
-    auto *ctx = rType.getContext();
-    unsigned ptrWidth = getSparseTensorEncoding(rType).getPointerBitWidth();
-    Type indexType = IndexType::get(ctx);
-    return ptrWidth ? IntegerType::get(ctx, ptrWidth) : indexType;
-  }
-
-  Type getIdxElementType() const {
-    auto *ctx = rType.getContext();
-    unsigned idxWidth = getSparseTensorEncoding(rType).getIndexBitWidth();
-    Type indexType = IndexType::get(ctx);
-    return idxWidth ? IntegerType::get(ctx, idxWidth) : indexType;
-  }
-
 private:
   unsigned getFieldIndex(unsigned dim, SparseTensorFieldKind kind) const {
     unsigned fieldIdx = -1u;
     foreachFieldInSparseTensor(
         getSparseTensorEncoding(rType),
         [dim, kind, &fieldIdx](unsigned fIdx, SparseTensorFieldKind fKind,
                                unsigned fDim, DimLevelType dlt) -> bool {
           if (fDim == dim && kind == fKind) {

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

foreachFieldAndTypeInSparseTensor(
       });
   return success();
 }
 
 /// Generates code that allocates a sparse storage scheme for given rank.
 static void allocSchemeForRank(OpBuilder &builder, Location loc,
                                MutSparseTensorDescriptor desc, unsigned r0) {
   RankedTensorType rtp = desc.getTensorType();
+  auto enc = getSparseTensorEncoding(rtp);
   unsigned rank = rtp.getShape().size();
   Value linear = constantIndex(builder, loc, 1);
   for (unsigned r = r0; r < rank; r++) {
     if (isCompressedDim(rtp, r)) {
       // 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 = desc.getPtrElementType();
+      Type ptrType = enc.getPointerType();
       Value ptrZero = constantZero(builder, loc, ptrType);
       unsigned fidx = desc.getPtrMemRefIndex(r);
       createPushback(builder, loc, desc, fidx, ptrZero, linear);
       return;
     }
     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);
     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, desc.getValMemRefIndex(), valZero, linear);

 ///
 /// TODO: for efficiency, we will need heuristis to make educated guesses
 ///       on the required capacities (see heuristic variable).
 ///
 static void createAllocFields(OpBuilder &builder, Location loc, Type type,
                               ValueRange dynSizes, bool enableInit,
                               SmallVectorImpl<Value> &fields) {
   RankedTensorType rtp = type.cast<RankedTensorType>();
+  auto enc = getSparseTensorEncoding(rtp);
   Value heuristic = constantIndex(builder, loc, 16);
 
   foreachFieldAndTypeInSparseTensor(
       rtp,
-      [&](Type fType, unsigned fIdx, SparseTensorFieldKind fKind,
+      [&builder, &fields, loc, heuristic,
+       enableInit](Type fType, unsigned fIdx, SparseTensorFieldKind fKind,
                    unsigned /*dim*/, DimLevelType /*dlt*/) -> bool {
         assert(fields.size() == fIdx);
         auto memRefTp = fType.cast<MemRefType>();
         Value field;
         switch (fKind) {
         case SparseTensorFieldKind::DimSizes:
         case SparseTensorFieldKind::MemSizes:
           field = builder.create<memref::AllocOp>(loc, memRefTp);
           break;
static void createAllocFields(OpBuilder &builder, Location loc, Type type,
   // 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.
   builder.create<linalg::FillOp>(
       loc, constantZero(builder, loc, builder.getIndexType()),
       desc.getMemSizesMemRef()); // zero memSizes
 
-  Value ptrZero = constantZero(builder, loc, desc.getPtrElementType());
+  Value ptrZero = constantZero(builder, loc, enc.getPointerType());
   for (unsigned r = 0; r < rank; r++) {
     unsigned ro = toOrigDim(rtp, r);
     // Fills dim sizes array.
     genStore(builder, loc, sizes[ro], desc.getDimSizesMemRef(),
              constantIndex(builder, loc, r));
 
     // Pushes a leading zero to pointers memref.
     if (isCompressedDim(rtp, r)) {
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
       // times?
       //
       if (d > 0) {
-        unsigned ptrWidth = getSparseTensorEncoding(rtp).getPointerBitWidth();
-        Type indexType = builder.getIndexType();
-        Type ptrType = ptrWidth ? builder.getIntegerType(ptrWidth) : indexType;
+        Type ptrType = getSparseTensorEncoding(rtp).getPointerType();
         Value ptrMemRef = desc.getPtrMemRef(d);
         Value mz = constantIndex(builder, loc, desc.getPtrMemSizesIndex(d));
         Value hi = genLoad(builder, loc, desc.getMemSizesMemRef(), mz);
         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);