[Mlir-commits] [mlir] b50ce4c - [MLIR][sparse] Add `soa` property to `sparse_tensor` Python bindings (#109135)

[llvmlistbot at llvm.org]

Author: Mateusz Sokół
Date: 2024-10-02T09:07:55-07:00
New Revision: b50ce4c81e71855bc01b9564d3bd239437847184

URL: https://github.com/llvm/llvm-project/commit/b50ce4c81e71855bc01b9564d3bd239437847184
DIFF: https://github.com/llvm/llvm-project/commit/b50ce4c81e71855bc01b9564d3bd239437847184.diff

LOG: [MLIR][sparse] Add `soa` property to `sparse_tensor` Python bindings (#109135)

Added: 
    

Modified: 
    mlir/include/mlir-c/Dialect/SparseTensor.h
    mlir/lib/Bindings/Python/DialectSparseTensor.cpp
    mlir/lib/CAPI/Dialect/SparseTensor.cpp
    mlir/test/Integration/Dialect/SparseTensor/python/test_output.py

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir-c/Dialect/SparseTensor.h b/mlir/include/mlir-c/Dialect/SparseTensor.h
index 125469f57c5f55..c816c1b58690ea 100644
--- 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
index 584981cfe99bf1..a730bf500be98c 100644
--- 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
index f2a0ab33c0224f..cf25b5263678fb 100644
--- 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
index 544273eb18835e..7d9aa37ba2890c 100644
--- 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)],