From 5bf93147d873ba0369cec4303733cc754b37ac1d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mateusz=20Sok=C3=B3=C5=82?= <mat646@gmail.com>
Date: Wed, 18 Sep 2024 12:59:23 +0000
Subject: [PATCH] Add `soa` property to `sparse_tensor` Python bindings

---
 mlir/include/mlir-c/Dialect/SparseTensor.h                    | 1 +
 mlir/lib/Bindings/Python/DialectSparseTensor.cpp              | 3 ++-
 mlir/lib/CAPI/Dialect/SparseTensor.cpp                        | 4 +++-
 .../Integration/Dialect/SparseTensor/python/test_output.py    | 4 ++++
 4 files changed, 10 insertions(+), 2 deletions(-)

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)],