[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:41 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 cada69459046dfce690e7713fe406c4966d2e2ff 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     | 3 +--
 2 files changed, 3 insertions(+), 4 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..9ed3cee2591475 100644
--- a/mlir/test/Dialect/SparseTensor/invalid_encoding.mlir
+++ b/mlir/test/Dialect/SparseTensor/invalid_encoding.mlir
@@ -242,13 +242,12 @@ func.func private @sparse_coo(tensor<?x?xf32, #COO_SoA>)
 
 // -----
 
-// expected-error at +1{{SoA is only applicable on singleton lvlTypes.}}
+// 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'}}



More information about the Mlir-commits mailing list