[mlir][linalg][transform][python] Add mix-in for BufferizeToAllocOp.

[ingomueller-net]

Summary

Re-apply https://reviews.llvm.org/D157704.

The original patch broke the tests on Python 3.8 and got reverted by
0c4aad050c23254c3c612e860e1278961d161aef. This patch replaces the usage
of the vertical bar operator for type hints with Union.

[ingomueller-net]

The original patch has been reviewed and the update is trivial. Will self-accept when CI has passed.

mlir/python/mlir/dialects/_structured_transform_ops_ext.py

if isinstance(values, list):
                 [IntegerAttr.get(IntegerType.get_signless(64), v) for v in value]
             )
             for value in values
         ]
 
     return ArrayAttr.get(values)
 
 
+class BufferizeToAllocationOp:
+    """Specialization for BufferizeToAllocationOp class."""
+
+    def __init__(
+        self,
+        target: Union[Operation, OpView, Value],
+        *,
+        memory_space: Optional[Union[int, str, Attribute]] = None,
+        memcpy_op: Optional[str] = None,
+        alloc_op: Optional[str] = None,
+        bufferize_destination_only: Optional[bool] = None,
+        loc=None,
+        ip=None,
+    ):
+        # No other types are allowed, so hard-code those here.
+        allocated_buffer_type = transform.AnyValueType.get()
+        new_ops_type = transform.AnyOpType.get()
+
+        if isinstance(memory_space, int):
+            memory_space = str(memory_space)
+        if isinstance(memory_space, str):
+            memory_space = Attribute.parse(memory_space)
+
+        super().__init__(
+            allocated_buffer_type,
+            new_ops_type,
+            target,
+            memory_space=memory_space,
+            memcpy_op=memcpy_op,
+            alloc_op=alloc_op,
+            bufferize_destination_only=bufferize_destination_only,
+        )
+
+
 class DecomposeOp:
     """Specialization for DecomposeOp class."""
 
     def __init__(self, target: Union[Operation, Value], *, loc=None, ip=None):
         super().__init__(
             pdl.OperationType.get(), _get_op_result_or_value(target), loc=loc, ip=ip
         )
 

mlir/test/python/dialects/transform_structured_ext.py

with Context(), Location.unknown():
             print("\nTEST:", f.__name__)
             f()
         module.operation.verify()
         print(module)
     return f
 
 
 @run
+def testBufferizeToAllocationOpCompact():
+    sequence = transform.SequenceOp(
+        transform.FailurePropagationMode.PROPAGATE, [], pdl.OperationType.get()
+    )
+    with InsertionPoint(sequence.body):
+        structured.BufferizeToAllocationOp(sequence.bodyTarget)
+        transform.YieldOp()
+    # CHECK-LABEL: TEST: testBufferizeToAllocationOpCompact
+    # CHECK: transform.sequence
+    # CHECK: transform.structured.bufferize_to_allocation
+
+
+@run
+def testBufferizeToAllocationOpArgs():
+    sequence = transform.SequenceOp(
+        transform.FailurePropagationMode.PROPAGATE, [], pdl.OperationType.get()
+    )
+    with InsertionPoint(sequence.body):
+        structured.BufferizeToAllocationOp(
+            sequence.bodyTarget,
+            memory_space=3,
+            memcpy_op="memref.copy",
+            alloc_op="memref.alloca",
+            bufferize_destination_only=True,
+        )
+        transform.YieldOp()
+    # CHECK-LABEL: TEST: testBufferizeToAllocationOpArgs
+    # CHECK: transform.sequence
+    # CHECK: transform.structured.bufferize_to_allocation
+    # CHECK-SAME: alloc_op = "memref.alloca"
+    # CHECK-SAME: bufferize_destination_only
+    # CHECK-SAME: memcpy_op = "memref.copy"
+    # CHECK-SAME: memory_space = 3
+
+
+@run
 def testDecompose():
     sequence = transform.SequenceOp(
         transform.FailurePropagationMode.PROPAGATE, [], pdl.OperationType.get()
     )
     with InsertionPoint(sequence.body):
         structured.DecomposeOp(sequence.bodyTarget)
         transform.YieldOp()
     # CHECK-LABEL: TEST: testDecompose