[mlir][python] Fix how the mlir variadic Python accessor `_ods_equally_sized_accessor` is used (#101132) #106003

kasper0406 · 2024-08-25T19:23:18Z

As reported in #101132, this fixes two bugs:

When accessing variadic operands inside an operation, it must be accessed as self.operation.operands instead of operation.operands
The implementation of the equally_sized_accessor function is doing wrong arithmetics when calculating the resulting index and group sizes.

I have added a test for the equally_sized_accessor function, which did not have a test previously.

github-actions · 2024-08-25T19:23:37Z

Thank you for submitting a Pull Request (PR) to the LLVM Project!

This PR will be automatically labeled and the relevant teams will be notified.

If you wish to, you can add reviewers by using the "Reviewers" section on this page.

If this is not working for you, it is probably because you do not have write permissions for the repository. In which case you can instead tag reviewers by name in a comment by using @ followed by their GitHub username.

If you have received no comments on your PR for a week, you can request a review by "ping"ing the PR by adding a comment “Ping”. The common courtesy "ping" rate is once a week. Please remember that you are asking for valuable time from other developers.

If you have further questions, they may be answered by the LLVM GitHub User Guide.

You can also ask questions in a comment on this PR, on the LLVM Discord or on the forums.

llvmbot · 2024-08-25T19:24:05Z

@llvm/pr-subscribers-mlir-core

@llvm/pr-subscribers-mlir

Author: Kasper Nielsen (kasper0406)

Changes

As reported in #101132, this fixes two bugs:

When accessing variadic operands inside an operation, it must be accessed as self.operation.operands instead of operation.operands
The implementation of the equally_sized_accessor function is doing wrong arithmetics when calculating the resulting index and group sizes.

The above code had no tests previously. I have added a test for the equally_sized_accessor function.

Full diff: https://github.com/llvm/llvm-project/pull/106003.diff

4 Files Affected:

(modified) mlir/python/mlir/dialects/_ods_common.py (+1-1)
(modified) mlir/test/mlir-tblgen/op-python-bindings.td (+6-6)
(modified) mlir/test/python/dialects/ods_helpers.py (+67)
(modified) mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp (+1-1)

diff --git a/mlir/python/mlir/dialects/_ods_common.py b/mlir/python/mlir/dialects/_ods_common.py
index 1e7e8244ed4420..0b56a376d23813 100644
--- a/mlir/python/mlir/dialects/_ods_common.py
+++ b/mlir/python/mlir/dialects/_ods_common.py
@@ -65,7 +65,7 @@ def equally_sized_accessor(
           group.
     """
 
-    total_variadic_length = len(elements) - n_variadic + 1
+    total_variadic_length = len(elements) - n_preceding_simple
     # This should be enforced by the C++-side trait verifier.
     assert total_variadic_length % n_variadic == 0
 
diff --git a/mlir/test/mlir-tblgen/op-python-bindings.td b/mlir/test/mlir-tblgen/op-python-bindings.td
index 9f202ba08608c6..d0642968854fe5 100644
--- a/mlir/test/mlir-tblgen/op-python-bindings.td
+++ b/mlir/test/mlir-tblgen/op-python-bindings.td
@@ -480,17 +480,17 @@ def SameVariadicOperandSizeOp : TestOp<"same_variadic_operand",
                                        [SameVariadicOperandSize]> {
   // CHECK: @builtins.property
   // CHECK: def variadic1(self):
-  // CHECK:   start, pg = _ods_equally_sized_accessor(operation.operands, 2, 0, 0)
+  // CHECK:   start, pg = _ods_equally_sized_accessor(self.operation.operands, 2, 0, 0)
   // CHECK:   return self.operation.operands[start:start + pg]
   //
   // CHECK: @builtins.property
   // CHECK: def non_variadic(self):
-  // CHECK:   start, pg = _ods_equally_sized_accessor(operation.operands, 2, 0, 1)
+  // CHECK:   start, pg = _ods_equally_sized_accessor(self.operation.operands, 2, 0, 1)
   // CHECK:   return self.operation.operands[start]
   //
   // CHECK: @builtins.property
   // CHECK: def variadic2(self):
-  // CHECK:   start, pg = _ods_equally_sized_accessor(operation.operands, 2, 1, 1)
+  // CHECK:   start, pg = _ods_equally_sized_accessor(self.operation.operands, 2, 1, 1)
   // CHECK:   return self.operation.operands[start:start + pg]
   let arguments = (ins Variadic<AnyType>:$variadic1, AnyType:$non_variadic,
                    Variadic<AnyType>:$variadic2);
@@ -506,17 +506,17 @@ def SameVariadicResultSizeOp : TestOp<"same_variadic_result",
                                       [SameVariadicResultSize]> {
   // CHECK: @builtins.property
   // CHECK: def variadic1(self):
-  // CHECK:   start, pg = _ods_equally_sized_accessor(operation.results, 2, 0, 0)
+  // CHECK:   start, pg = _ods_equally_sized_accessor(self.operation.results, 2, 0, 0)
   // CHECK:   return self.operation.results[start:start + pg]
   //
   // CHECK: @builtins.property
   // CHECK: def non_variadic(self):
-  // CHECK:   start, pg = _ods_equally_sized_accessor(operation.results, 2, 0, 1)
+  // CHECK:   start, pg = _ods_equally_sized_accessor(self.operation.results, 2, 0, 1)
   // CHECK:   return self.operation.results[start]
   //
   // CHECK: @builtins.property
   // CHECK: def variadic2(self):
-  // CHECK:   start, pg = _ods_equally_sized_accessor(operation.results, 2, 1, 1)
+  // CHECK:   start, pg = _ods_equally_sized_accessor(self.operation.results, 2, 1, 1)
   // CHECK:   return self.operation.results[start:start + pg]
   let results = (outs Variadic<AnyType>:$variadic1, AnyType:$non_variadic,
                  Variadic<AnyType>:$variadic2);
diff --git a/mlir/test/python/dialects/ods_helpers.py b/mlir/test/python/dialects/ods_helpers.py
index 0d2a18e0eb0af2..cb0d0528a6a866 100644
--- a/mlir/test/python/dialects/ods_helpers.py
+++ b/mlir/test/python/dialects/ods_helpers.py
@@ -3,6 +3,7 @@
 import gc
 
 from mlir.ir import *
+from mlir.dialects._ods_common import equally_sized_accessor
 
 
 def run(f):
@@ -208,3 +209,69 @@ class TestOp(OpView):
 
 
 run(testOdsBuildDefaultCastError)
+
+
+def testOdsEquallySizedAccessor():
+    class TestOpMultiResultSegments(OpView):
+        OPERATION_NAME = "custom.test_op"
+        _ODS_REGIONS = (1, True)
+
+    with Context() as ctx, Location.unknown():
+        ctx.allow_unregistered_dialects = True
+        m = Module.create()
+        with InsertionPoint(m.body):
+            v = add_dummy_value()
+            ts = [IntegerType.get_signless(i * 8) for i in range(4)]
+
+            op = TestOpMultiResultSegments.build_generic(
+                results=[ts[0], ts[1], ts[2], ts[3]], operands=[v]
+            )
+            start, pg = equally_sized_accessor(op.results, 3, 1, 0)
+            # CHECK: start: 1, pg: 1
+            print(f"start: {start}, pg: {pg}")
+            # CHECK: i8
+            print(op.results[start].type)
+
+            start, pg = equally_sized_accessor(op.results, 3, 1, 1)
+            # CHECK: start: 2, pg: 1
+            print(f"start: {start}, pg: {pg}")
+            # CHECK: i16
+            print(op.results[start].type)
+
+
+run(testOdsEquallySizedAccessor)
+
+
+def testOdsEquallySizedAccessorMultipleSegments():
+    class TestOpMultiResultSegments(OpView):
+        OPERATION_NAME = "custom.test_op"
+        _ODS_REGIONS = (1, True)
+        _ODS_RESULT_SEGMENTS = [0, -1, -1]
+
+    def types(lst):
+        return [e.type for e in lst]
+
+    with Context() as ctx, Location.unknown():
+        ctx.allow_unregistered_dialects = True
+        m = Module.create()
+        with InsertionPoint(m.body):
+            v = add_dummy_value()
+            ts = [IntegerType.get_signless(i * 8) for i in range(7)]
+
+            op = TestOpMultiResultSegments.build_generic(
+                results=[ts[0], [ts[1], ts[2], ts[3]], [ts[4], ts[5], ts[6]]], operands=[v]
+            )
+            start, pg = equally_sized_accessor(op.results, 2, 1, 0)
+            # CHECK: start: 1, pg: 3
+            print(f"start: {start}, pg: {pg}")
+            # CHECK: [IntegerType(i8), IntegerType(i16), IntegerType(i24)]
+            print(types(op.results[start:start + pg]))
+
+            start, pg = equally_sized_accessor(op.results, 2, 1, 1)
+            # CHECK: start: 4, pg: 3
+            print(f"start: {start}, pg: {pg}")
+            # CHECK: [IntegerType(i32), IntegerType(i40), IntegerType(i48)]
+            print(types(op.results[start:start + pg]))
+
+
+run(testOdsEquallySizedAccessorMultipleSegments)
diff --git a/mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp b/mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp
index 052020acdcb764..97f2cc6c3f5763 100644
--- a/mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp
+++ b/mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp
@@ -145,7 +145,7 @@ constexpr const char *opOneVariadicTemplate = R"Py(
 constexpr const char *opVariadicEqualPrefixTemplate = R"Py(
   @builtins.property
   def {0}(self):
-    start, pg = _ods_equally_sized_accessor(operation.{1}s, {2}, {3}, {4}))Py";
+    start, pg = _ods_equally_sized_accessor(self.operation.{1}s, {2}, {3}, {4}))Py";
 
 /// Second part of the template for equally-sized case, accessing a single
 /// element:

github-actions · 2024-08-25T22:34:25Z

✅ With the latest revision this PR passed the Python code formatter.

mlir/test/python/dialects/ods_helpers.py

github-actions · 2024-08-30T12:58:04Z

✅ With the latest revision this PR passed the C/C++ code formatter.

makslevental · 2024-08-30T13:55:24Z

mlir/test/python/python_test_ops.td

@@ -227,4 +227,16 @@ def OptionalOperandOp : TestOp<"optional_operand_op"> {
  let results = (outs I32:$result);
 }

+def SameVariadicResultSizeOp : TestOp<"same_variadic_result",
+                                      [SameVariadicResultSize]> {
+  let results = (outs Variadic<AnyType>:$variadic1, AnyType:$non_variadic,


Sorry to be harassing you but can you add tests that permute the non-variadic - so like maybe

non-var, var, var
var, var, non-var

basically I'm too lazy to think really hard about the arithmetic 😅

No worries. I’ve added those tests now 😊

makslevental · 2024-08-30T22:03:46Z

mlir/tools/mlir-tblgen/OpPythonBindingGen.cpp

-  auto getNumVariableLengthOperands = [](const Operator &oper) {
-    return oper.getNumVariableLengthOperands();
-  };
-  emitElementAccessors(op, os, "operand", getNumVariableLengthOperands,
-                       getNumOperands, getOperand);


lol someone must've really liked functional programming to do it like this...

makslevental

Excellent work (and I did check the arithmetic lol). Thanks for the thoroughness. Let me know if you need me to merge.

kasper0406 · 2024-08-31T05:05:05Z

I have no power to merge it, so would appreciate if you can hit the button!

github-actions · 2024-08-31T07:17:51Z

@kasper0406 Congratulations on having your first Pull Request (PR) merged into the LLVM Project!

Your changes will be combined with recent changes from other authors, then tested by our build bots. If there is a problem with a build, you may receive a report in an email or a comment on this PR.

Please check whether problems have been caused by your change specifically, as the builds can include changes from many authors. It is not uncommon for your change to be included in a build that fails due to someone else's changes, or infrastructure issues.

How to do this, and the rest of the post-merge process, is covered in detail here.

If your change does cause a problem, it may be reverted, or you can revert it yourself. This is a normal part of LLVM development. You can fix your changes and open a new PR to merge them again.

If you don't get any reports, no action is required from you. Your changes are working as expected, well done!

kasper0406 added 2 commits August 21, 2024 16:55

Fix MLIR Python bindings when trying to access varadic elements

c35573e

Add and fix tests

c1e68b4

kasper0406 requested review from ftynse, makslevental and stellaraccident as code owners August 25, 2024 19:23

llvmbot added mlir:core MLIR Core Infrastructure mlir:python MLIR Python bindings mlir labels Aug 25, 2024

kasper0406 added 3 commits August 26, 2024 08:42

Fix code style

6be3f4c

Another coding style check

81fd53c

Do not skip tests

9db5bb6

joker-eph changed the title ~~Fix how the mlir variadic Python acessor _ods_equally_sized_accessor is used (#101132)~~ Fix how the mlir variadic Python accessor _ods_equally_sized_accessor is used (#101132) Aug 26, 2024

kasper0406 changed the title ~~Fix how the mlir variadic Python accessor _ods_equally_sized_accessor is used (#101132)~~ [mlir][python] Fix how the mlir variadic Python accessor _ods_equally_sized_accessor is used (#101132) Aug 27, 2024

jpienaar reviewed Aug 29, 2024

View reviewed changes

mlir/test/python/dialects/ods_helpers.py Outdated Show resolved Hide resolved

Rename pg to elements_per_group

b861ec0

makslevental reviewed Aug 29, 2024

View reviewed changes

mlir/test/python/dialects/ods_helpers.py Show resolved Hide resolved

Additional fixes and better tests

816dea1

Fix code style

f553697

makslevental reviewed Aug 30, 2024

View reviewed changes

Added tests for different variadic group permutations and sizes

1f7db9b

makslevental reviewed Aug 30, 2024

View reviewed changes

makslevental approved these changes Aug 30, 2024

View reviewed changes

makslevental merged commit 3766ba4 into llvm:main Aug 31, 2024
8 checks passed

kasper0406 deleted the kn/fix-mlir-python-bindings branch August 31, 2024 07:30

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[mlir][python] Fix how the mlir variadic Python accessor `_ods_equally_sized_accessor` is used (#101132) #106003

[mlir][python] Fix how the mlir variadic Python accessor `_ods_equally_sized_accessor` is used (#101132) #106003

Uh oh!

kasper0406 commented Aug 25, 2024 •

edited

Loading

Uh oh!

github-actions bot commented Aug 25, 2024

Uh oh!

llvmbot commented Aug 25, 2024 •

edited

Loading

Uh oh!

github-actions bot commented Aug 25, 2024 •

edited

Loading

Uh oh!

Uh oh!

Uh oh!

github-actions bot commented Aug 30, 2024 •

edited

Loading

Uh oh!

makslevental Aug 30, 2024

Uh oh!

kasper0406 Aug 30, 2024

Uh oh!

makslevental Aug 30, 2024

Uh oh!

makslevental left a comment

Uh oh!

kasper0406 commented Aug 31, 2024

Uh oh!

Uh oh!

github-actions bot commented Aug 31, 2024

Uh oh!

Uh oh!

[mlir][python] Fix how the mlir variadic Python accessor _ods_equally_sized_accessor is used (#101132) #106003

[mlir][python] Fix how the mlir variadic Python accessor _ods_equally_sized_accessor is used (#101132) #106003

Uh oh!

Conversation

kasper0406 commented Aug 25, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions bot commented Aug 25, 2024

Uh oh!

llvmbot commented Aug 25, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

github-actions bot commented Aug 25, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

github-actions bot commented Aug 30, 2024 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

makslevental Aug 30, 2024

Choose a reason for hiding this comment

Uh oh!

kasper0406 Aug 30, 2024

Choose a reason for hiding this comment

Uh oh!

makslevental Aug 30, 2024

Choose a reason for hiding this comment

Uh oh!

makslevental left a comment

Choose a reason for hiding this comment

Uh oh!

kasper0406 commented Aug 31, 2024

Uh oh!

Uh oh!

github-actions bot commented Aug 31, 2024

Uh oh!

Uh oh!

[mlir][python] Fix how the mlir variadic Python accessor `_ods_equally_sized_accessor` is used (#101132) #106003

[mlir][python] Fix how the mlir variadic Python accessor `_ods_equally_sized_accessor` is used (#101132) #106003

kasper0406 commented Aug 25, 2024 •

edited

Loading

llvmbot commented Aug 25, 2024 •

edited

Loading

github-actions bot commented Aug 25, 2024 •

edited

Loading

github-actions bot commented Aug 30, 2024 •

edited

Loading