[mlir][linalg][transform][python] Allow no args in MaskedVectorize.

[ingomueller-net]

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.

[llvmbot]

@llvm/pr-subscribers-mlir

Changes

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.

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

2 Files Affected:

• (modified) mlir/python/mlir/dialects/_structured_transform_ops_ext.py (+8-2)
• (modified) mlir/test/python/dialects/transform_structured_ext.py (+10)

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 +@run +@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):