[MLIR][sparse] Add soa property to sparse_tensor Python bindings (#…

…109135)
llvm · Oct 2, 2024 · b50ce4c · b50ce4c
1 parent 504585d
commit b50ce4c
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 2 deletions.
diff --git a/mlir/include/mlir-c/Dialect/SparseTensor.h b/mlir/include/mlir-c/Dialect/SparseTensor.h
@@ -39,6 +39,7 @@ enum MlirSparseTensorLevelFormat {
 enum MlirSparseTensorLevelPropertyNondefault {
   MLIR_SPARSE_PROPERTY_NON_UNIQUE = 0x0001,
   MLIR_SPARSE_PROPERTY_NON_ORDERED = 0x0002,
+  MLIR_SPARSE_PROPERTY_SOA = 0x0004,
 };
 
 //===----------------------------------------------------------------------===//

diff --git a/mlir/lib/Bindings/Python/DialectSparseTensor.cpp b/mlir/lib/Bindings/Python/DialectSparseTensor.cpp
@@ -33,7 +33,8 @@ static void populateDialectSparseTensorSubmodule(const py::module &m) {
   py::enum_<MlirSparseTensorLevelPropertyNondefault>(m, "LevelProperty",
                                                      py::module_local())
       .value("non_ordered", MLIR_SPARSE_PROPERTY_NON_ORDERED)
-      .value("non_unique", MLIR_SPARSE_PROPERTY_NON_UNIQUE);
+      .value("non_unique", MLIR_SPARSE_PROPERTY_NON_UNIQUE)
+      .value("soa", MLIR_SPARSE_PROPERTY_SOA);
 
   mlir_attribute_subclass(m, "EncodingAttr",
                           mlirAttributeIsASparseTensorEncodingAttr)

diff --git a/mlir/lib/CAPI/Dialect/SparseTensor.cpp b/mlir/lib/CAPI/Dialect/SparseTensor.cpp
@@ -36,7 +36,9 @@ static_assert(
 static_assert(static_cast<int>(MLIR_SPARSE_PROPERTY_NON_ORDERED) ==
                       static_cast<int>(LevelPropNonDefault::Nonordered) &&
                   static_cast<int>(MLIR_SPARSE_PROPERTY_NON_UNIQUE) ==
-                      static_cast<int>(LevelPropNonDefault::Nonunique),
+                      static_cast<int>(LevelPropNonDefault::Nonunique) &&
+                  static_cast<int>(MLIR_SPARSE_PROPERTY_SOA) ==
+                      static_cast<int>(LevelPropNonDefault::SoA),
               "MlirSparseTensorLevelProperty (C-API) and "
               "LevelPropertyNondefault (C++) mismatch");
 

diff --git a/mlir/test/Integration/Dialect/SparseTensor/python/test_output.py b/mlir/test/Integration/Dialect/SparseTensor/python/test_output.py
@@ -129,6 +129,10 @@ def main():
         prop = st.LevelProperty
         levels = [
             [builder(fmt.compressed, [prop.non_unique]), builder(fmt.singleton)],
+            [
+                builder(fmt.compressed, [prop.non_unique]),
+                builder(fmt.singleton, [prop.soa]),
+            ],
             [builder(fmt.dense), builder(fmt.compressed)],
             [builder(fmt.dense), builder(fmt.loose_compressed)],
             [builder(fmt.compressed), builder(fmt.compressed)],