[mlir][linalg][transform][python] Extend mix-in for Vectorize

[martin-luecke]

Summary

Extends the existing mix-in for VectorizeOp with support for the missing unit attributes.

Also fixes the unintuitive implementation where
structured.VectorizeOp(target=target, vectorize_padding=False) still resulted in the creation of the UnitAttr vectorize_padding.

[ingomueller-net]

mlir/python/mlir/dialects/_structured_transform_ops_ext.py

793–794

That cast should not be necessary; you just tested if it's an instance of bool, no?

793–794

Is all of this really necessary? Doesn't the generated code handle Booleans? I one case, I see this:

if bool(bufferize_destination_only): attributes["bufferize_destination_only"] = _ods_ir.UnitAttr.get(
  _ods_get_default_loc_context(loc))

Doesn't this handle all cases as desired if you remove the whole 12 lines?

[martin-luecke]

Address review comments:

• simplify mix-in to rely more on the generated constructor

[martin-luecke]

Thank you for your comments @ingomueller-net.

I do not yet have commit privileges so if you are happy with the state I would appreciate it if you could commit this on my behalf.

[ingomueller-net]

OK, done! Unfortunately, I show up as an author now on git. I spend a bit of time trying to find out who was going to be the author and how to make sure it remains you but I didn't find anything. Turns out that the one who lands becomes the author...

mlir/python/mlir/dialects/_structured_transform_ops_ext.py

 
 class VectorizeOp:
     """Specialization for VectorizeOp class."""
 
     def __init__(
         self,
         target: Union[Operation, Value],
         *,
-        vectorize_padding: Union[bool, BoolAttr] = False,
+        disable_multi_reduction_to_contract_patterns: bool = False,
+        disable_transfer_permutation_map_lowering_patterns: bool = False,
+        vectorize_nd_extract: bool = False,
+        vectorize_padding: bool = False,
         loc=None,
         ip=None,
     ):
         pdl_operation_type = pdl.OperationType.get()
-        if isinstance(vectorize_padding, bool):
-            vectorize_padding = UnitAttr.get()
         super().__init__(
             pdl_operation_type,
             _get_op_result_or_value(target),
+            disable_multi_reduction_to_contract_patterns=disable_multi_reduction_to_contract_patterns,
+            disable_transfer_permutation_map_lowering_patterns=disable_transfer_permutation_map_lowering_patterns,
+            vectorize_nd_extract=vectorize_nd_extract,
             vectorize_padding=vectorize_padding,
             loc=loc,
             ip=ip,
         )

mlir/test/python/dialects/transform_structured_ext.py

with InsertionPoint(sequence.body):
         )
         transform.YieldOp()
     # CHECK-LABEL: TEST: testTileToForallMapping
     # CHECK: = transform.structured.tile_to_forall_op
     # CHECK-SAME: mapping = [#gpu.thread<y>, #gpu.thread<x>]
 
 
 @run
-def testVectorize():
+def testVectorizeAllAttrs():
     sequence = transform.SequenceOp(
         transform.FailurePropagationMode.Propagate, [], pdl.OperationType.get()
     )
     with InsertionPoint(sequence.body):
-        structured.VectorizeOp(sequence.bodyTarget, vectorize_padding=True)
+        structured.VectorizeOp(
+            sequence.bodyTarget,
+            disable_multi_reduction_to_contract_patterns=True,
+            disable_transfer_permutation_map_lowering_patterns=True,
+            vectorize_nd_extract=True,
+            vectorize_padding=True,
+        )
+        transform.YieldOp()
+    # CHECK-LABEL: TEST: testVectorizeAllAttrs
+    # CHECK: transform.sequence
+    # CHECK: = transform.structured.vectorize
+    # CHECK-SAME: disable_multi_reduction_to_contract_patterns
+    # CHECK-SAME: disable_transfer_permutation_map_lowering_patterns
+    # CHECK-SAME: vectorize_nd_extract
+    # CHECK-SAME: vectorize_padding
+
+
+@run
+def testVectorizeNoAttrs():
+    sequence = transform.SequenceOp(
+        transform.FailurePropagationMode.Propagate, [], pdl.OperationType.get()
+    )
+    with InsertionPoint(sequence.body):
+        structured.VectorizeOp(
+            sequence.bodyTarget,
+            disable_multi_reduction_to_contract_patterns=False,
+            disable_transfer_permutation_map_lowering_patterns=False,
+            vectorize_nd_extract=False,
+            vectorize_padding=False,
+        )
         transform.YieldOp()
-    # CHECK-LABEL: TEST: testVectorize
+    # CHECK-LABEL: TEST: testVectorizeNoAttrs
     # CHECK: transform.sequence
     # CHECK: = transform.structured.vectorize
-    # CHECK: {vectorize_padding}
+    # CHECK-NOT: disable_multi_reduction_to_contract_patterns
+    # CHECK-NOT: disable_transfer_permutation_map_lowering_patterns
+    # CHECK-NOT: vectorize_nd_extract
+    # CHECK-NOT: vectorize_padding
 
 
 @run
 def testMatchInterfaceEnum():
     names = ArrayAttr.get([StringAttr.get("test.dummy")])
     result_type = transform.AnyOpType.get()
     sequence = transform.SequenceOp(
         transform.FailurePropagationMode.Propagate, [], transform.AnyOpType.get()