[Mlir-commits] [mlir] [mlir][linalg][transform][python] Allow args in MaskedVectorize. (PR #66541)

[Ingo Müller]

https://github.com/ingomueller-net created https://github.com/llvm/llvm-project/pull/66541

The mix-in of this op did not allow to pass in no argument. This special case is now handled correctly and covered by the tests.

>From bc912fad081fb04b44c5839290da864759d6a6f0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Ingo=20M=C3=BCller?= <ingomueller at google.com>
Date: Fri, 15 Sep 2023 19:28:02 +0000
Subject: [PATCH] [mlir][linalg][transform][python] Allow args in
 MaskedVectorize.

The mix-in of this op did not allow to pass in no argument. This special
case is now handled correctly and covered by the tests.
---
 .../mlir/dialects/_structured_transform_ops_ext.py     | 10 ++++++++--
 mlir/test/python/dialects/transform_structured_ext.py  | 10 ++++++++++
 2 files changed, 18 insertions(+), 2 deletions(-)

diff --git a/mlir/python/mlir/dialects/_structured_transform_ops_ext.py b/mlir/python/mlir/dialects/_structured_transform_ops_ext.py
index c5134b6e718f3b8..fd3dbca7c5a607a 100644
--- a/mlir/python/mlir/dialects/_structured_transform_ops_ext.py
+++ b/mlir/python/mlir/dialects/_structured_transform_ops_ext.py
@@ -366,7 +366,7 @@ class MaskedVectorizeOp:
     def __init__(
         self,
         target: Union[Operation, OpView, Value],
-        vector_sizes: Union[DynamicIndexList, ArrayAttr],
+        vector_sizes: Optional[Union[DynamicIndexList, ArrayAttr]] = None,
         *,
         vectorize_nd_extract: Optional[bool] = None,
         scalable_sizes: OptionalBoolList = None,
@@ -374,7 +374,13 @@ def __init__(
         loc=None,
         ip=None,
     ):
-        if scalable_sizes is None and static_vector_sizes is None:
+        if (
+            scalable_sizes is None
+            and static_vector_sizes is None
+            and vector_sizes is None
+        ):
+            dynamic_vector_sizes = []
+        elif scalable_sizes is None and static_vector_sizes is None:
             (
                 dynamic_vector_sizes,
                 static_vector_sizes,
diff --git a/mlir/test/python/dialects/transform_structured_ext.py b/mlir/test/python/dialects/transform_structured_ext.py
index 5d5ee945b66867e..69181160d5489b7 100644
--- a/mlir/test/python/dialects/transform_structured_ext.py
+++ b/mlir/test/python/dialects/transform_structured_ext.py
@@ -169,6 +169,16 @@ def testMatchOpNamesList(target):
     # CHECK-SAME: (!transform.any_op) -> !transform.any_op
 
 
+ at run
+ at create_sequence
+def testMaskedVectorizeNoArgs(target):
+    structured.MaskedVectorizeOp(target)
+    # CHECK-LABEL: TEST: testMaskedVectorizeNoArgs
+    # CHECK: transform.sequence
+    # CHECK: transform.structured.masked_vectorize
+    # CHECK-NOT:     vector_sizes
+
+
 @run
 @create_sequence
 def testMaskedVectorizeStatic(target):