[Mlir-commits] [mlir] 2d3dcd4 - [mlir][linalg][transform][python] Add mix-in for BufferizeToAllocOp.

[Ingo Müller]

Author: Ingo Müller
Date: 2023-08-16T15:07:43Z
New Revision: 2d3dcd4aecb94ae862163fb9bed0879295f6725f

URL: https://github.com/llvm/llvm-project/commit/2d3dcd4aecb94ae862163fb9bed0879295f6725f
DIFF: https://github.com/llvm/llvm-project/commit/2d3dcd4aecb94ae862163fb9bed0879295f6725f.diff

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

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`.

Reviewed By: ftynse

Differential Revision: https://reviews.llvm.org/D158075

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 675d4237031d98..e34451af429c1e 100644
--- a/mlir/python/mlir/dialects/_structured_transform_ops_ext.py
+++ b/mlir/python/mlir/dialects/_structured_transform_ops_ext.py
@@ -84,6 +84,40 @@ def _get_int_int_array_attr(
     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."""
 

diff  --git a/mlir/test/python/dialects/transform_structured_ext.py b/mlir/test/python/dialects/transform_structured_ext.py
index f4b5e2e9e948aa..e9ad3f0e8fde46 100644
--- a/mlir/test/python/dialects/transform_structured_ext.py
+++ b/mlir/test/python/dialects/transform_structured_ext.py
@@ -18,6 +18,42 @@ def run(f):
     return f
 
 
+ at 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
+
+
+ at 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(