[Mlir-commits] [mlir] 088c7ce - [mlir][sparse] introduce SoA level property on singleton level. (#81942)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Feb 15 16:41:14 PST 2024
Author: Peiming Liu
Date: 2024-02-15T16:41:10-08:00
New Revision: 088c7ce429197a809f8943e5699a7fd0650fe9f7
URL: https://github.com/llvm/llvm-project/commit/088c7ce429197a809f8943e5699a7fd0650fe9f7
DIFF: https://github.com/llvm/llvm-project/commit/088c7ce429197a809f8943e5699a7fd0650fe9f7.diff
LOG: [mlir][sparse] introduce SoA level property on singleton level. (#81942)
Added:
Modified:
mlir/include/mlir/Dialect/SparseTensor/IR/Enums.h
mlir/include/mlir/Dialect/SparseTensor/IR/SparseTensorAttrDefs.td
mlir/lib/Dialect/SparseTensor/IR/Detail/LvlTypeParser.cpp
mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp
mlir/test/Dialect/SparseTensor/invalid_encoding.mlir
mlir/test/Dialect/SparseTensor/roundtrip_encoding.mlir
Removed:
################################################################################
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..db359b4b7a5d09 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).";
+ }
}
+
+ // 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 to 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..9ed3cee2591475 100644
--- a/mlir/test/Dialect/SparseTensor/invalid_encoding.mlir
+++ b/mlir/test/Dialect/SparseTensor/invalid_encoding.mlir
@@ -234,6 +234,22 @@ func.func private @too_many_lvl_decl(%arg0: tensor<?x?xf64, #TooManyLvlDecl>) {
// -----
+// 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 to 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'}}
#TooFewLvlDecl = #sparse_tensor.encoding<{
map = {l0} (d0, d1) -> (l0 = d0 : dense, l1 = d1 : compressed)
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,
More information about the Mlir-commits
mailing list