[mlir][sparse] Adding {pointer,index}OverheadTypeEncoding

[wrengr]

Summary

Work towards: https://github.com/llvm/llvm-project/issues/51652

The new functions fill the gap between overheadTypeEncoding and get{Pointer,Index}OverheadType.

[wrengr]

Since writing this differential I managed to adjust the implementation of enumerators in D122060 so that the subsequent D122061 can remain ignorant of the source tensor's <P,I> template parameters. Originally, the functions in this differential were used in D122061 in order to detect/assert that the source tensor's <P,I> match the target tensor's <P,I>.

Even though they're no longer being used anywhere, I think it's still useful to fill the API gap.

[aartbik]

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

42

here, and below, I don't feel very strongly about this, but using a parameter name even here seems more consistent with this file, so

const SparseTensorEncodingAttr &enc

[wrengr]

addressing style

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

 OverheadType overheadTypeEncoding(unsigned width);
 
 /// Converts an overhead storage type to its internal type-encoding.
 OverheadType overheadTypeEncoding(Type tp);
 
 /// Converts the internal type-encoding for overhead storage to an mlir::Type.
 Type getOverheadType(Builder &builder, OverheadType ot);
 
+/// Returns the OverheadType for pointer overhead storage.
+OverheadType pointerOverheadTypeEncoding(const SparseTensorEncodingAttr &enc);
+
+/// Returns the OverheadType for index overhead storage.
+OverheadType indexOverheadTypeEncoding(const SparseTensorEncodingAttr &enc);
+
 /// Returns the mlir::Type for pointer overhead storage.
 Type getPointerOverheadType(Builder &builder,
                             const SparseTensorEncodingAttr &enc);
 
 /// Returns the mlir::Type for index overhead storage.
 Type getIndexOverheadType(Builder &builder,
                           const SparseTensorEncodingAttr &enc);
 

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

Type mlir::sparse_tensor::getOverheadType(Builder &builder, OverheadType ot) {
   case OverheadType::kU16:
     return builder.getIntegerType(16);
   case OverheadType::kU8:
     return builder.getIntegerType(8);
   }
   llvm_unreachable("Unknown OverheadType");
 }
 
+OverheadType mlir::sparse_tensor::pointerOverheadTypeEncoding(
+    const SparseTensorEncodingAttr &enc) {
+  return overheadTypeEncoding(enc.getPointerBitWidth());
+}
+
+OverheadType mlir::sparse_tensor::indexOverheadTypeEncoding(
+    const SparseTensorEncodingAttr &enc) {
+  return overheadTypeEncoding(enc.getIndexBitWidth());
+}
+
 Type mlir::sparse_tensor::getPointerOverheadType(
     Builder &builder, const SparseTensorEncodingAttr &enc) {
-  return getOverheadType(builder,
-                         overheadTypeEncoding(enc.getPointerBitWidth()));
+  return getOverheadType(builder, pointerOverheadTypeEncoding(enc));
 }
 
 Type mlir::sparse_tensor::getIndexOverheadType(
     Builder &builder, const SparseTensorEncodingAttr &enc) {
-  return getOverheadType(builder, overheadTypeEncoding(enc.getIndexBitWidth()));
+  return getOverheadType(builder, indexOverheadTypeEncoding(enc));
 }
 
 StringRef mlir::sparse_tensor::overheadTypeFunctionSuffix(OverheadType ot) {
   switch (ot) {
   case OverheadType::kIndex:
     return "";
   case OverheadType::kU64:
     return "64";