[Mlir-commits] [mlir] ebc8466 - [mlir][sparse] Adding {pointer, index}OverheadTypeEncoding

[wren romano]

Author: wren romano
Date: 2022-03-23T12:04:47-07:00
New Revision: ebc8466481f938d799d0170e0caf1c7e64f52c53

URL: https://github.com/llvm/llvm-project/commit/ebc8466481f938d799d0170e0caf1c7e64f52c53
DIFF: https://github.com/llvm/llvm-project/commit/ebc8466481f938d799d0170e0caf1c7e64f52c53.diff

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

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

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

Reviewed By: aartbik

Differential Revision: https://reviews.llvm.org/D122056

Added: 
    

Modified: 
    mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.cpp
    mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.h

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.cpp
index ea9be3bddb547..38f15a1ee7c01 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.cpp
@@ -58,15 +58,24 @@ Type mlir::sparse_tensor::getOverheadType(Builder &builder, OverheadType ot) {
   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) {

diff  --git a/mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.h b/mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.h
index 9286cca808aa8..944605bab4b9f 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.h
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/CodegenUtils.h
@@ -38,6 +38,12 @@ 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);