[Mlir-commits] [mlir] [mlir][sparse] introduce SoA level property on singleton level. (PR #81942)
Peiming Liu
llvmlistbot at llvm.org
Thu Feb 15 16:11:05 PST 2024
https://github.com/PeimingLiu updated https://github.com/llvm/llvm-project/pull/81942
>From 08892f1356a17b3e0ff0c0ac48dc001cdad68c95 Mon Sep 17 00:00:00 2001
From: Peiming Liu <peiming at google.com>
Date: Thu, 15 Feb 2024 23:55:30 +0000
Subject: [PATCH 1/2] [mlir][sparse] introduce AoS level property to singleton
level.
---
.../mlir/Dialect/SparseTensor/IR/Enums.h | 12 ++++++++++--
.../SparseTensor/IR/SparseTensorAttrDefs.td | 17 ++++++++++++++---
.../SparseTensor/IR/Detail/LvlTypeParser.cpp | 6 ++++--
.../SparseTensor/IR/SparseTensorDialect.cpp | 19 +++++++++++++++++++
.../SparseTensor/invalid_encoding.mlir | 17 +++++++++++++++++
.../SparseTensor/roundtrip_encoding.mlir | 11 +++++++++++
6 files changed, 75 insertions(+), 7 deletions(-)
diff --git a/mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h b/mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h
index d03afed6f126ce..fa34bbe3c9d910 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h
+++ b/mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h
@@ -189,8 +189,9 @@ constexpr const char *toFormatString(LevelFormat lvlFmt) {
/// This enum defines all the nondefault properties for storage formats.
enum class LevelPropNonDefault : uint64_t {
- Nonunique = 0x0001,
- Nonordered = 0x0002,
+ Nonunique = 0x0001, // 0b001
+ Nonordered = 0x0002, // 0b010
+ SoA = 0x0004, // 0b100
};
/// Returns string representation of the given level properties.
@@ -200,6 +201,8 @@ constexpr const char *toPropString(LevelPropNonDefault lvlProp) {
return "nonunique";
case LevelPropNonDefault::Nonordered:
return "nonordered";
+ case LevelPropNonDefault::SoA:
+ return "soa";
}
return "";
}
@@ -334,6 +337,11 @@ struct LevelType {
propStr += ", ";
propStr += toPropString(LevelPropNonDefault::Nonordered);
}
+ if (isa<LevelPropNonDefault::SoA>()) {
+ if (!propStr.empty())
+ propStr += ", ";
+ propStr += toPropString(LevelPropNonDefault::SoA);
+ }
if (!propStr.empty())
lvlStr += ("(" + propStr + ")");
return lvlStr;
diff --git a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td
index 5b3b971f9a7f23..f0b832571e68ec 100644
--- a/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td
+++ b/mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td
@@ -157,11 +157,14 @@ def SparseTensorEncodingAttr : SparseTensor_Attr<"SparseTensorEncoding",
By default, each level-type has the property of being unique (no duplicate
coordinates at that level) and ordered (coordinates appear sorted at that
- level). The following properties can be added to a level-format to change
- this default behavior:
+ level). For singleton levels, the coordinates are fused with its parents in AoS
+ (array of structures) scheme. The following properties can be added to a level-format
+ to change this default behavior:
- **nonunique** : duplicate coordinates may appear at the level
- **nonordered** : coordinates may appear in arbribratry order
+ - **soa** : only applicable to singleton levels, fuses the singleton
+ level in SoA (structure of arrays) scheme.
In addition to the map, the following two fields are optional:
@@ -188,10 +191,18 @@ def SparseTensorEncodingAttr : SparseTensor_Attr<"SparseTensorEncoding",
}>
... tensor<?xf32, #SparseVector> ...
- // Sorted coordinate scheme.
+ // Sorted coordinate scheme (arranged in AoS format by default).
#SortedCOO = #sparse_tensor.encoding<{
map = (i, j) -> (i : compressed(nonunique), j : singleton)
}>
+ // coordinates = {x_crd, y_crd}[nnz]
+ ... tensor<?x?xf64, #SortedCOO> ...
+
+ // Sorted coordinate scheme (arranged in SoA format).
+ #SortedCOO = #sparse_tensor.encoding<{
+ map = (i, j) -> (i : compressed(nonunique), j : singleton(soa))
+ }>
+ // coordinates = {x_crd[nnz], y_crd[nnz]}
... tensor<?x?xf64, #SortedCOO> ...
// Batched sorted coordinate scheme, with high encoding.
diff --git a/mlir/lib/Dialect/SparseTensor/IR/Detail/LvlTypeParser.cpp b/mlir/lib/Dialect/SparseTensor/IR/Detail/LvlTypeParser.cpp
index 380cccc989ec6a..455e90baf0a715 100644
--- a/mlir/lib/Dialect/SparseTensor/IR/Detail/LvlTypeParser.cpp
+++ b/mlir/lib/Dialect/SparseTensor/IR/Detail/LvlTypeParser.cpp
@@ -87,10 +87,12 @@ ParseResult LvlTypeParser::parseProperty(AsmParser &parser,
auto loc = parser.getCurrentLocation();
ERROR_IF(failed(parser.parseOptionalKeyword(&strVal)),
"expected valid level property (e.g. nonordered, nonunique or high)")
- if (strVal.compare("nonunique") == 0) {
+ if (strVal.equals(toPropString(LevelPropNonDefault::Nonunique))) {
*properties |= static_cast<uint64_t>(LevelPropNonDefault::Nonunique);
- } else if (strVal.compare("nonordered") == 0) {
+ } else if (strVal.equals(toPropString(LevelPropNonDefault::Nonordered))) {
*properties |= static_cast<uint64_t>(LevelPropNonDefault::Nonordered);
+ } else if (strVal.equals(toPropString(LevelPropNonDefault::SoA))) {
+ *properties |= static_cast<uint64_t>(LevelPropNonDefault::SoA);
} else {
parser.emitError(loc, "unknown level property: ") << strVal;
return failure();
diff --git a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
index 6d02645d860e96..d78107cb53dcee 100644
--- a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
+++ b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
@@ -664,7 +664,26 @@ LogicalResult SparseTensorEncodingAttr::verify(
[](LevelType i) { return isSingletonLT(i); }))
return emitError() << "expected all singleton lvlTypes "
"following a singleton level";
+ // We can potentially support mixed SoA/AoS singleton levels.
+ if (!std::all_of(it, lvlTypes.end(), [it](LevelType i) {
+ return it->isa<LevelPropNonDefault::SoA>() ==
+ i.isa<LevelPropNonDefault::SoA>();
+ })) {
+ return emitError() << "expected all singleton lvlTypes stored in the "
+ "same memory layout (SoA vs AoS).";
+ }
}
+
+ // Property SoA can only be applied on singleton level.
+ auto soaLvls = llvm::make_filter_range(lvlTypes, [](LevelType lt) {
+ return lt.isa<LevelPropNonDefault::SoA>();
+ });
+ if (llvm::any_of(soaLvls, [](LevelType lt) {
+ return !lt.isa<LevelFormat::Singleton>();
+ })) {
+ return emitError() << "SoA is only applicable on singleton lvlTypes.";
+ }
+
// TODO: audit formats that actually are supported by backend.
if (auto it = std::find_if(lvlTypes.begin(), lvlTypes.end(), isNOutOfMLT);
it != std::end(lvlTypes)) {
diff --git a/mlir/test/Dialect/SparseTensor/invalid_encoding.mlir b/mlir/test/Dialect/SparseTensor/invalid_encoding.mlir
index a52a46b4fcfe7c..79fcbe661f1355 100644
--- a/mlir/test/Dialect/SparseTensor/invalid_encoding.mlir
+++ b/mlir/test/Dialect/SparseTensor/invalid_encoding.mlir
@@ -232,6 +232,23 @@ func.func private @too_many_lvl_decl(%arg0: tensor<?x?xf64, #TooManyLvlDecl>) {
return
}
+// -----
+
+// expected-error at +1{{expected all singleton lvlTypes stored in the same memory layout (SoA vs AoS).}}
+#COO_SoA = #sparse_tensor.encoding<{
+ map = (d0, d1, d2) -> (d0 : compressed(nonunique), d1 : singleton(soa, nonunique), d2 : singleton)
+}>
+func.func private @sparse_coo(tensor<?x?xf32, #COO_SoA>)
+
+// -----
+
+// expected-error at +1{{SoA is only applicable on singleton lvlTypes.}}
+#COO_SoA = #sparse_tensor.encoding<{
+ map = (d0, d1) -> (d0 : compressed(nonunique, soa), d1 : singleton(soa))
+}>
+func.func private @sparse_coo(tensor<?x?xf32, #COO_SoA>)
+
+
// -----
// expected-error at +2 {{use of undeclared identifier 'l1'}}
diff --git a/mlir/test/Dialect/SparseTensor/roundtrip_encoding.mlir b/mlir/test/Dialect/SparseTensor/roundtrip_encoding.mlir
index 64520638b253df..9d5118ceecc587 100644
--- a/mlir/test/Dialect/SparseTensor/roundtrip_encoding.mlir
+++ b/mlir/test/Dialect/SparseTensor/roundtrip_encoding.mlir
@@ -94,6 +94,17 @@ func.func private @sparse_sorted_coo(tensor<10x10xf64, #SortedCOO>)
// -----
+#COO_SoA = #sparse_tensor.encoding<{
+ map = (d0, d1) -> (d0 : compressed(nonunique), d1 : singleton(soa))
+}>
+
+// CHECK-DAG: #[[$COO_SoA:.*]] = #sparse_tensor.encoding<{ map = (d0, d1) -> (d0 : compressed(nonunique), d1 : singleton(soa)) }>
+// CHECK-LABEL: func private @sparse_coo(
+// CHECK-SAME: tensor<?x?xf32, #[[$COO_SoA]]>)
+func.func private @sparse_coo(tensor<?x?xf32, #COO_SoA>)
+
+// -----
+
#BSR = #sparse_tensor.encoding<{
map = ( i, j ) ->
( i floordiv 2 : dense,
>From 6f0a51308580c59cb946aabf4e322d74f1ae4787 Mon Sep 17 00:00:00 2001
From: Peiming Liu <peiming at google.com>
Date: Fri, 16 Feb 2024 00:10:52 +0000
Subject: [PATCH 2/2] address comments
---
mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp | 4 ++--
mlir/test/Dialect/SparseTensor/invalid_encoding.mlir | 1 -
2 files changed, 2 insertions(+), 3 deletions(-)
diff --git a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
index d78107cb53dcee..db359b4b7a5d09 100644
--- a/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
+++ b/mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
@@ -674,14 +674,14 @@ LogicalResult SparseTensorEncodingAttr::verify(
}
}
- // Property SoA can only be applied on singleton level.
+ // SoA property can only be applied on singleton level.
auto soaLvls = llvm::make_filter_range(lvlTypes, [](LevelType lt) {
return lt.isa<LevelPropNonDefault::SoA>();
});
if (llvm::any_of(soaLvls, [](LevelType lt) {
return !lt.isa<LevelFormat::Singleton>();
})) {
- return emitError() << "SoA is only applicable on singleton lvlTypes.";
+ return emitError() << "SoA is only applicable to singleton lvlTypes.";
}
// TODO: audit formats that actually are supported by backend.
diff --git a/mlir/test/Dialect/SparseTensor/invalid_encoding.mlir b/mlir/test/Dialect/SparseTensor/invalid_encoding.mlir
index 79fcbe661f1355..899a3071aeb890 100644
--- a/mlir/test/Dialect/SparseTensor/invalid_encoding.mlir
+++ b/mlir/test/Dialect/SparseTensor/invalid_encoding.mlir
@@ -248,7 +248,6 @@ func.func private @sparse_coo(tensor<?x?xf32, #COO_SoA>)
}>
func.func private @sparse_coo(tensor<?x?xf32, #COO_SoA>)
-
// -----
// expected-error at +2 {{use of undeclared identifier 'l1'}}
More information about the Mlir-commits
mailing list