[Mlir-commits] [mlir] f767f09 - [mlir][sparse] sparse storage scheme type conversion
Aart Bik
llvmlistbot at llvm.org
Wed Aug 31 15:13:07 PDT 2022
Author: Aart Bik
Date: 2022-08-31T15:12:55-07:00
New Revision: f767f0925268981272e0b48cbc3709d761b5600e
URL: https://github.com/llvm/llvm-project/commit/f767f0925268981272e0b48cbc3709d761b5600e
DIFF: https://github.com/llvm/llvm-project/commit/f767f0925268981272e0b48cbc3709d761b5600e.diff
LOG: [mlir][sparse] sparse storage scheme type conversion
This builds a compound type for the buffers required for the sparse storage scheme defined by the MLIR sparse tensor types. The use of a tuple allows for a simple 1:1 type conversion. A subsequent pass can expand this tuple into its component with an isolated 1:N type conversion.
Reviewed By: Peiming
Differential Revision: https://reviews.llvm.org/D133050
Added:
Modified:
mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp
mlir/test/Dialect/SparseTensor/codegen.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp
index 86669260c970f..b905b442f0975 100644
--- a/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp
+++ b/mlir/lib/Dialect/SparseTensor/Transforms/SparseTensorCodegen.cpp
@@ -33,14 +33,72 @@ namespace {
// Helper methods.
//===----------------------------------------------------------------------===//
-/// Maps each sparse tensor type to the appropriate buffer.
-static Optional<Type> convertSparseTensorTypes(Type type) {
- if (getSparseTensorEncoding(type) != nullptr) {
- // TODO: this is just a dummy rule to get the ball rolling....
- RankedTensorType rTp = type.cast<RankedTensorType>();
- return MemRefType::get({ShapedType::kDynamicSize}, rTp.getElementType());
+/// Maps a sparse tensor type to the appropriate compounded buffers.
+static Optional<Type> convertSparseTensorType(Type type) {
+ 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 eltType = rType.getElementType();
+ //
+ // Sparse tensor storage for rank-dimensional tensor is organized as a
+ // single compound type with the following fields:
+ //
+ // struct {
+ // memref<rank x index> dimSize ; size in each dimension
+ // ; per-dimension d:
+ // ; if dense:
+ // <nothing>
+ // ; if compresed:
+ // memref<? x idx> indices-d ; indices for sparse dim d
+ // memref<? x ptr> pointers-d ; pointers for sparse dim d
+ // ; if singleton:
+ // memref<? x idx> indices-d ; indices for singleton dim d
+ // memref<? x eltType> values ; values
+ // };
+ //
+ // TODO: fill in the ? when statically known
+ //
+ // TODO: emit dimSizes when not needed (e.g. all-dense)
+ //
+ unsigned rank = rType.getShape().size();
+ SmallVector<Type, 8> fields;
+ fields.push_back(MemRefType::get({rank}, indexType));
+ for (unsigned r = 0; r < rank; r++) {
+ // Dimension level types apply in order to the reordered dimension.
+ // As a result, the compound type can be constructed directly in the given
+ // order. Clients of this type know what field is what from the sparse
+ // tensor type.
+ switch (enc.getDimLevelType()[r]) {
+ case SparseTensorEncodingAttr::DimLevelType::Dense:
+ break;
+ case SparseTensorEncodingAttr::DimLevelType::Compressed:
+ case SparseTensorEncodingAttr::DimLevelType::CompressedNu:
+ case SparseTensorEncodingAttr::DimLevelType::CompressedNo:
+ case SparseTensorEncodingAttr::DimLevelType::CompressedNuNo:
+ fields.push_back(MemRefType::get({ShapedType::kDynamicSize}, idxType));
+ fields.push_back(MemRefType::get({ShapedType::kDynamicSize}, ptrType));
+ break;
+ case SparseTensorEncodingAttr::DimLevelType::Singleton:
+ case SparseTensorEncodingAttr::DimLevelType::SingletonNu:
+ case SparseTensorEncodingAttr::DimLevelType::SingletonNo:
+ case SparseTensorEncodingAttr::DimLevelType::SingletonNuNo:
+ fields.push_back(MemRefType::get({ShapedType::kDynamicSize}, idxType));
+ break;
+ }
}
- return llvm::None;
+ fields.push_back(MemRefType::get({ShapedType::kDynamicSize}, eltType));
+ // Sparse tensor storage (temporarily) lives in a tuple. This allows a
+ // simple 1:1 type conversion during codegen. A subsequent pass uses
+ // a 1:N type conversion to expand the tuple into its fields.
+ return TupleType::get(context, fields);
}
//===----------------------------------------------------------------------===//
@@ -67,7 +125,7 @@ class SparseReturnConverter : public OpConversionPattern<func::ReturnOp> {
mlir::SparseTensorTypeToBufferConverter::SparseTensorTypeToBufferConverter() {
addConversion([](Type type) { return type; });
- addConversion(convertSparseTensorTypes);
+ addConversion(convertSparseTensorType);
}
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/Dialect/SparseTensor/codegen.mlir b/mlir/test/Dialect/SparseTensor/codegen.mlir
index a3cecaf367a40..f9a979e38d42b 100644
--- a/mlir/test/Dialect/SparseTensor/codegen.mlir
+++ b/mlir/test/Dialect/SparseTensor/codegen.mlir
@@ -1,14 +1,62 @@
// RUN: mlir-opt %s --sparse-tensor-codegen --canonicalize --cse | FileCheck %s
#SparseVector = #sparse_tensor.encoding<{
- dimLevelType = ["compressed"]
+ dimLevelType = [ "compressed" ],
+ indexBitWidth = 64,
+ pointerBitWidth = 32
}>
-// TODO: just a dummy memref rewriting to get the ball rolling....
+#Dense = #sparse_tensor.encoding<{
+ dimLevelType = [ "dense", "dense" ],
+ indexBitWidth = 64,
+ pointerBitWidth = 32
+}>
+
+#Row = #sparse_tensor.encoding<{
+ dimLevelType = [ "compressed", "dense" ],
+ indexBitWidth = 64,
+ pointerBitWidth = 32
+}>
+
+#CSR = #sparse_tensor.encoding<{
+ dimLevelType = [ "dense", "compressed" ],
+ indexBitWidth = 64,
+ pointerBitWidth = 32
+}>
+
+#DCSR = #sparse_tensor.encoding<{
+ dimLevelType = [ "compressed", "compressed" ],
+ indexBitWidth = 64,
+ pointerBitWidth = 32
+}>
// CHECK-LABEL: func @sparse_nop(
-// CHECK-SAME: %[[A:.*]]: memref<?xf64>) -> memref<?xf64> {
-// CHECK: return %[[A]] : memref<?xf64>
+// CHECK-SAME: %[[A:.*]]: tuple<memref<1xindex>, memref<?xi64>, memref<?xi32>, memref<?xf64>>) -> tuple<memref<1xindex>, memref<?xi64>, memref<?xi32>, memref<?xf64>>
+// CHECK: return %[[A]] : tuple<memref<1xindex>, memref<?xi64>, memref<?xi32>, memref<?xf64>>
func.func @sparse_nop(%arg0: tensor<?xf64, #SparseVector>) -> tensor<?xf64, #SparseVector> {
return %arg0 : tensor<?xf64, #SparseVector>
}
+
+// CHECK-LABEL: func @sparse_dense(
+// CHECK-SAME: %[[A:.*]]: tuple<memref<2xindex>, memref<?xf64>>)
+func.func @sparse_dense(%arg0: tensor<?x?xf64, #Dense>) {
+ return
+}
+
+// CHECK-LABEL: func @sparse_row(
+// CHECK-SAME: %[[A:.*]]: tuple<memref<2xindex>, memref<?xi64>, memref<?xi32>, memref<?xf64>>)
+func.func @sparse_row(%arg0: tensor<?x?xf64, #Row>) {
+ return
+}
+
+// CHECK-LABEL: func @sparse_csr(
+// CHECK-SAME: %[[A:.*]]: tuple<memref<2xindex>, memref<?xi64>, memref<?xi32>, memref<?xf64>>)
+func.func @sparse_csr(%arg0: tensor<?x?xf64, #CSR>) {
+ return
+}
+
+// CHECK-LABEL: func @sparse_dcsr(
+// CHECK-SAME: %[[A:.*]]: tuple<memref<2xindex>, memref<?xi64>, memref<?xi32>, memref<?xi64>, memref<?xi32>, memref<?xf64>>)
+func.func @sparse_dcsr(%arg0: tensor<?x?xf64, #DCSR>) {
+ return
+}
More information about the Mlir-commits
mailing list