dialects: Implement and test memref.CopyOp.

tests/dialects/test_memref.py

@@ -20,6 +20,7 @@
 from xdsl.dialects.memref import (
     Alloc,
     Alloca,
+    CopyOp,
     Dealloc,
     MemRefType,
     Load,
@@ -393,3 +394,37 @@ def test_memref_dma_wait():
         wrong_tag = TestSSAValue(wrong_tag_type)
 
         DmaWaitOp.get(wrong_tag, [index], num_elements).verify()
+
+
+def test_memref_copy():
+    typ = MemRefType.from_element_type_and_shape(i32, [4])
+    typ3 = MemRefType.from_element_type_and_shape(i32, [3])
+    typ64 = MemRefType.from_element_type_and_shape(i64, [4])
+    source = TestSSAValue(typ)
+    destination = TestSSAValue(typ)
+
+    copy = CopyOp(source, destination)
+
+    copy.verify()
+    assert isinstance(copy, CopyOp)
+    assert copy.source.typ == typ
+    assert copy.destination.typ == typ
+
+    destination = TestSSAValue(typ3)
+
+    copy = CopyOp(source, destination)
+
+    with pytest.raises(
+        VerifyException, match="Expected source and destination to have the same shape."
+    ):
+        copy.verify()
+
+    destination = TestSSAValue(typ64)
+
+    copy = CopyOp(source, destination)
+
+    with pytest.raises(
+        VerifyException,
+        match="Expected source and destination to have the same element type.",
+    ):
+        copy.verify()

tests/filecheck/mlir-conversion/with-mlir/dialects/memref/memref_ops_mlir_conversion.mlir

@@ -14,6 +14,7 @@
     %6 = "memref.subview"(%5) {"operand_segment_sizes" = array<i32: 1, 0, 0, 0>, "static_offsets" = array<i64: 0, 0>, "static_sizes" = array<i64: 1, 1>, "static_strides" = array<i64: 1, 1>} : (memref<10x2xindex>) -> memref<1x1xindex>
     %7 = "memref.cast"(%5) : (memref<10x2xindex>) -> memref<?x?xindex>
     %no_align = "memref.alloca"() {i64, "operand_segment_sizes" = array<i32: 0, 0>} : () -> memref<1xindex>
+    "memref.copy"(%no_align, %2) : (memref<1xindex>, memref<1xindex>) -> ()
     "memref.dealloc"(%no_align) : (memref<1xindex>) -> ()
     "memref.dealloc"(%2) : (memref<1xindex>) -> ()
     "memref.dealloc"(%5) : (memref<10x2xindex>) -> ()
@@ -49,6 +50,7 @@
 // CHECK-NEXT: %6 = "memref.subview"(%5) {"operand_segment_sizes" = array<i32: 1, 0, 0, 0>, "static_offsets" = array<i64: 0, 0>, "static_sizes" = array<i64: 1, 1>, "static_strides" = array<i64: 1, 1>} : (memref<10x2xindex>) -> memref<1x1xindex>
 // CHECK-NEXT: %7 = "memref.cast"(%5) : (memref<10x2xindex>) -> memref<?x?xindex>
 // CHECK-NEXT: %8 = "memref.alloca"() {"i64", "operand_segment_sizes" = array<i32: 0, 0>} : () -> memref<1xindex>
+// CHECK-NEXT: "memref.copy"(%8, %2) : (memref<1xindex>, memref<1xindex>) -> ()
 // CHECK-NEXT: "memref.dealloc"(%8) : (memref<1xindex>) -> ()
 // CHECK-NEXT: "memref.dealloc"(%2) : (memref<1xindex>) -> ()
 // CHECK-NEXT: "memref.dealloc"(%5) : (memref<10x2xindex>) -> ()

xdsl/dialects/memref.py

@@ -30,6 +30,7 @@
     ContainerType,
 )
 from xdsl.ir import (
+    Attribute,
     TypeAttribute,
     Operation,
     SSAValue,
@@ -593,12 +594,35 @@ def verify_(self) -> None:
             raise VerifyException("Expected tag to be a memref of i32")
 
 
+@irdl_op_definition
+class CopyOp(IRDLOperation):
+    name = "memref.copy"
+    source: Operand = operand_def(MemRefType)
+    destination: Operand = operand_def(MemRefType)
+
+    def __init__(self, source: SSAValue | Operation, destination: SSAValue | Operation):
+        super().__init__([source, destination])
+
+    def verify_(self) -> None:
+        source = cast(MemRefType[Attribute], self.source.typ)
+        destination = cast(MemRefType[Attribute], self.destination.typ)
+        if source.get_shape() != destination.get_shape():
+            raise VerifyException(
+                f"Expected source and destination to have the same shape."
+            )
+        if source.get_element_type() != destination.get_element_type():
+            raise VerifyException(
+                f"Expected source and destination to have the same element type."
+            )
+
+
 MemRef = Dialect(
     [
         Load,
         Store,
         Alloc,
         Alloca,
+        CopyOp,
         Dealloc,
         GetGlobal,
         Global,