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

[llvmlistbot at llvm.org]

Author: Ingo Müller
Date: 2023-09-19T10:34:47+02:00
New Revision: 86ddbdd3e74797cc3ead26bfc9665d69dfbccdbf

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

LOG: [mlir][linalg][transform][python] Allow no args in MaskedVectorize. (#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.

Added: 
    

Modified: 
    mlir/python/mlir/dialects/_structured_transform_ops_ext.py
    mlir/test/python/dialects/transform_structured_ext.py

Removed: 
    


################################################################################
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):