[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

275–276

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/SparseTensorCodegen.cpp

 /// Maps a sparse tensor type to the appropriate compounded buffers.
 static Optional<LogicalResult>
 convertSparseTensorType(Type type, SmallVectorImpl<Type> &fields) {
   auto enc = getSparseTensorEncoding(type);
   if (!enc)
     return llvm::None;
   // Construct the basic types.
   auto *context = type.getContext();
-  unsigned idxWidth = enc.getIndexBitWidth();
-  unsigned ptrWidth = enc.getPointerBitWidth();
   RankedTensorType rType = type.cast<RankedTensorType>();
   Type indexType = IndexType::get(context);
-  Type idxType = idxWidth ? IntegerType::get(context, idxWidth) : indexType;
-  Type ptrType = ptrWidth ? IntegerType::get(context, ptrWidth) : indexType;
+  Type idxType = enc.getIndexType();
+  Type ptrType = enc.getPointerType();
   Type eltType = rType.getElementType();
   //
   // Sparse tensor storage scheme for rank-dimensional tensor is organized
   // as a single compound type with the following fields. Note that every
   // memref with ? size actually behaves as a "vector", i.e. the stored
   // size is the capacity and the used size resides in the memSizes array.
   //
   // struct {
static void allocSchemeForRank(OpBuilder &builder, Location loc,
                                unsigned r0) {
   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.
-      unsigned ptrWidth = getSparseTensorEncoding(rtp).getPointerBitWidth();
-      Type indexType = builder.getIndexType();
-      Type ptrType = ptrWidth ? builder.getIntegerType(ptrWidth) : indexType;
+      Type ptrType = getSparseTensorEncoding(rtp).getPointerType();
       Value ptrZero = constantZero(builder, loc, ptrType);
       createPushback(builder, loc, fields, field, 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, rtp, fields, r);
-      linear = builder.create<arith::MulIOp>(loc, linear, size);
+    }
+    // Keep compounding the size, but nothing needs to be initialized
   }
   // Reached values array so prepare for an insertion.
   Value valZero = constantZero(builder, loc, rtp.getElementType());
   createPushback(builder, loc, fields, field, valZero, linear);
   assert(fields.size() == ++field);
 }
 
 /// Creates allocation operation.

 /// 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) {
   auto enc = getSparseTensorEncoding(type);
   assert(enc);
-  // Construct the basic types.
-  unsigned idxWidth = enc.getIndexBitWidth();
-  unsigned ptrWidth = enc.getPointerBitWidth();
   RankedTensorType rtp = type.cast<RankedTensorType>();
   Type indexType = builder.getIndexType();
-  Type idxType = idxWidth ? builder.getIntegerType(idxWidth) : indexType;
-  Type ptrType = ptrWidth ? builder.getIntegerType(ptrWidth) : indexType;
+  Type idxType = enc.getIndexType();
+  Type ptrType = enc.getPointerType();
   Type eltType = rtp.getElementType();
   auto shape = rtp.getShape();
   unsigned rank = shape.size();
   Value heuristic = constantIndex(builder, loc, 16);
   // Build original sizes.
   SmallVector<Value> sizes;
   for (unsigned r = 0, o = 0; r < rank; r++) {
     if (ShapedType::isDynamic(shape[r]))
static void genEndInsert(OpBuilder &builder, Location loc, RankedTensorType rtp,
   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 mz = constantIndex(builder, loc, getMemSizesIndex(field));
         Value hi = genLoad(builder, loc, fields[memSizesIdx], 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, fields[field], zero)};
         scf::ForOp loop = createFor(builder, loc, hi, inits, one);
         Value i = loop.getInductionVar();