[mlir][sparse] introduce SoA level property on singleton level. (#81942)

Author: PeimingLiu

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;

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.

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();

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)) {

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@+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@+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@+2 {{use of undeclared identifier 'l1'}}
 #TooFewLvlDecl = #sparse_tensor.encoding<{
   map = {l0} (d0, d1) -> (l0 = d0 : dense, l1 = d1 : compressed)

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,