dialects: (vector) Add for vector.transfer_read and vector.transfer_write operations

tests/dialects/test_vector.py

@@ -1,8 +1,11 @@
 import pytest
 
 from xdsl.dialects.builtin import (
+    AffineMapAttr,
+    ArrayAttr,
     IndexType,
     IntAttr,
+    IntegerAttr,
     MemRefType,
     VectorType,
     i1,
@@ -20,8 +23,11 @@
     MaskedstoreOp,
     PrintOp,
     StoreOp,
+    TransferReadOp,
+    TransferWriteOp,
 )
 from xdsl.ir import Attribute, OpResult
+from xdsl.ir.affine import AffineExpr, AffineMap
 from xdsl.utils.test_value import TestSSAValue
 
 
@@ -653,3 +659,64 @@ def test_vector_insert_element_0d_verify_empty_position():
         match="Expected position to be empty with 0-D vector.",
     ):
         insert_element.verify()
+
+
+def test_vector_transfer_write_construction():
+    x = AffineExpr.dimension(0)
+    vector_type = VectorType(IndexType(), [3])
+    memref_type = MemRefType(IndexType(), [3, 3])
+    # (x, y) -> x
+    permutation_map = AffineMapAttr(AffineMap(2, 0, (x,)))
+    in_bounds = ArrayAttr(
+        [IntegerAttr.from_bool(False) for _ in range(vector_type.get_num_dims())]
+    )
+
+    vector = TestSSAValue(vector_type)
+    source = TestSSAValue(memref_type)
+    index = TestSSAValue(IndexType())
+
+    transfer_write = TransferWriteOp(
+        vector,
+        source,
+        [index, index],
+        in_bounds,
+        permutation_map=permutation_map,
+    )
+
+    transfer_write.verify()
+
+    assert transfer_write.vector is vector
+    assert transfer_write.source is source
+    assert len(transfer_write.indices) == 2
+    assert transfer_write.indices[0] is index
+    assert transfer_write.permutation_map is permutation_map
+
+
+def test_vector_transfer_read_construction():
+    x = AffineExpr.dimension(0)
+    vector_type = VectorType(IndexType(), [3])
+    memref_type = MemRefType(IndexType(), [3, 3])
+    permutation_map = AffineMapAttr(AffineMap(2, 0, (x,)))
+    in_bounds = ArrayAttr(
+        [IntegerAttr.from_bool(False) for _ in range(vector_type.get_num_dims())]
+    )
+
+    source = TestSSAValue(memref_type)
+    index = TestSSAValue(IndexType())
+    padding = TestSSAValue(IndexType())
+
+    transfer_read = TransferReadOp(
+        source,
+        [index, index],
+        padding,
+        vector_type,
+        in_bounds,
+        permutation_map=permutation_map,
+    )
+
+    transfer_read.verify()
+
+    assert transfer_read.source is source
+    assert len(transfer_read.indices) == 2
+    assert transfer_read.indices[0] is index
+    assert transfer_read.permutation_map is permutation_map

tests/filecheck/dialects/vector/vector_ops.mlir

@@ -1,5 +1,5 @@
 // RUN: XDSL_ROUNDTRIP
-
+#map = affine_map<(d0, d1) -> (d0)>
 builtin.module {
   func.func private @vector_test(%0 : memref<4x4xindex>, %1 : vector<1xi1>, %2 : index) {
     %3 = "vector.load"(%0, %2, %2) : (memref<4x4xindex>, index, index) -> vector<2xindex>
@@ -10,6 +10,9 @@ builtin.module {
     "vector.maskedstore"(%0, %2, %2, %1, %6) : (memref<4x4xindex>, index, index, vector<1xi1>, vector<1xindex>) -> ()
     "vector.print"(%6) : (vector<1xindex>) -> ()
     %7 = "vector.create_mask"(%2) : (index) -> vector<2xi1>
+    %8 = "vector.transfer_read"(%0, %2, %2, %2) <{"in_bounds" = [true], "operandSegmentSizes" = array<i32: 1, 2, 1, 0>, "permutation_map" = #map}> : (memref<4x4xindex>, index, index, index) -> vector<4xindex>
+    "vector.transfer_write"(%8, %0, %2, %2) <{"in_bounds" = [true], "operandSegmentSizes" = array<i32: 1, 1, 2, 0>, "permutation_map" = #map}> :  (vector<4xindex>, memref<4x4xindex>, index, index) -> ()
+
     func.return
   }
 }
@@ -25,6 +28,8 @@ builtin.module {
 // CHECK-NEXT:     "vector.maskedstore"(%0, %2, %2, %1, %6) : (memref<4x4xindex>, index, index, vector<1xi1>, vector<1xindex>) -> ()
 // CHECK-NEXT:     "vector.print"(%6) : (vector<1xindex>) -> ()
 // CHECK-NEXT:     %7 = "vector.create_mask"(%2) : (index) -> vector<2xi1>
+// CHECK-NEXT:     %8 = "vector.transfer_read"(%0, %2, %2, %2) <{"in_bounds" = [true], "operandSegmentSizes" = array<i32: 1, 2, 1, 0>, "permutation_map" = affine_map<(d0, d1) -> (d0)>}> : (memref<4x4xindex>, index, index, index) -> vector<4xindex>
+// CHECK-NEXT:     "vector.transfer_write"(%8, %0, %2, %2) <{"in_bounds" = [true], "operandSegmentSizes" = array<i32: 1, 1, 2, 0>, "permutation_map" = affine_map<(d0, d1) -> (d0)>}> : (vector<4xindex>, memref<4x4xindex>, index, index) -> ()
 // CHECK-NEXT:     func.return
 // CHECK-NEXT:   }
 // CHECK-NEXT: }

tests/filecheck/dialects/vector/vector_transfer_read_verify.mlir

@@ -0,0 +1,43 @@
+// RUN: xdsl-opt --split-input-file --verify-diagnostics %s | filecheck %s
+
+%source, %index, %padding = "test.op"() : () -> (vector<4x3xf32>, index, f32)
+"vector.transfer_read"(%source, %index, %index, %padding) <{in_bounds=[true], operandSegmentSizes = array<i32: 1, 3, 1, 0>, permutation_map = affine_map<() -> (0)>}> : (vector<4x3xf32>, index, index, f32) -> vector<1x1x2x3xf32>
+// CHECK: Expected tensor or memref type, got vector<4x3xf32>
+
+// -----
+
+%source, %index, %padding = "test.op"() : () -> (memref<?x?xf32>, index, f32)
+"vector.transfer_read"(%source, %index, %index, %index, %padding) <{in_bounds=[true], operandSegmentSizes = array<i32: 1, 3, 1, 0>, permutation_map = affine_map<() -> (0)>}> : (memref<?x?xf32>, index, index, index, f32) -> vector<128xf32>
+// CHECK: Expected an index for each memref/tensor dimension
+
+// -----
+
+%source, %index, %padding = "test.op"() : () -> (memref<?x?xf32>, index, f32)
+"vector.transfer_read"(%source, %index, %index, %padding) <{in_bounds=[true], operandSegmentSizes = array<i32: 1, 2, 1, 0>, permutation_map = affine_map<(d0) -> (d0)>}> : (memref<?x?xf32>, index, index, f32) -> vector<128xf32>
+// CHECK: requires a permutation_map with input dims of the same rank as the source type
+
+// -----
+
+%source, %index, %padding = "test.op"() : () -> (memref<?x?xf32>, index, f32)
+"vector.transfer_read"(%source, %index, %index, %padding) <{in_bounds=[true, true], operandSegmentSizes = array<i32: 1, 2, 1, 0>, permutation_map = affine_map<(d0, d1) -> (d0, d1)>}> : (memref<?x?xf32>, index, index, f32) -> vector<128xf32>
+// CHECK: requires a permutation_map with result dims of the same rank as the vector type
+
+// -----
+
+%source, %index, %padding = "test.op"() : () -> (memref<?x?xf32>, index, f32)
+"vector.transfer_read"(%source, %index, %index, %padding) <{in_bounds=[true, true], operandSegmentSizes = array<i32: 1, 2, 1, 0>, permutation_map = affine_map<(d0, d1) -> (d0 + d1)>}> : (memref<?x?xf32>, index, index, f32) -> vector<128xf32>
+// CHECK: requires a projected permutation_map (at most one dim or the zero constant can appear in each result
+
+// -----
+
+%source, %index, %padding = "test.op"() : () -> (memref<?x?xf32>, index, f32)
+"vector.transfer_read"(%source, %index, %index, %padding) <{in_bounds=[true], operandSegmentSizes = array<i32: 1, 2, 1, 0>, permutation_map = affine_map<(d0, d1) -> (d0 + 1)>}> : (memref<?x?xf32>, index, index, f32) -> vector<128xf32>
+// CHECK: requires a projected permutation_map (at most one dim or the zero constant can appear in each result)
+
+// -----
+
+%source, %index, %padding = "test.op"() : () -> (memref<?x?x?xf32>, index, f32)
+"vector.transfer_read"(%source, %index, %index, %index, %padding) <{in_bounds=[true, true], operandSegmentSizes = array<i32: 1, 3, 1, 0>, permutation_map = affine_map<(d0, d1, d2) -> (d0, d0)>}> : (memref<?x?x?xf32>, index, index, index, f32) -> vector<3x7xf32>
+// CHECK: requires a permutation_map that is a permutation (found one dim used more than once)
+
+// TODO transfer other tests from mlir/test/Dialect/Vector/invalid.mlir once verification for vector element types is implemented

tests/filecheck/dialects/vector/vector_transfer_write_verify.mlir

@@ -0,0 +1,49 @@
+// RUN: xdsl-opt --split-input-file --verify-diagnostics %s | filecheck %s
+
+%source, %index = "test.op"() : () -> (vector<4x3xf32>, index)
+"vector.transfer_write"(%source, %source, %index, %index) <{in_bounds=[true], operandSegmentSizes = array<i32: 1, 1, 2, 0>, permutation_map = affine_map<() -> (0)>}> : (vector<4x3xf32>, vector<4x3xf32>, index, index) -> ()
+// CHECK: Expected tensor or memref type
+
+// -----
+
+%source, %vector, %index = "test.op"() : () -> (memref<?x?xf32>, vector<128xf32>, index)
+"vector.transfer_write"(%vector, %source, %index, %index, %index) <{in_bounds=[true], operandSegmentSizes = array<i32: 1, 1, 3, 0>, permutation_map = affine_map<() -> (0)>}> : (vector<128xf32>, memref<?x?xf32>, index, index, index) -> ()
+// CHECK: Expected an index for each memref/tensor dimension
+
+// -----
+
+%source, %vector, %index = "test.op"() : () -> (memref<?x?xf32>, vector<128xf32>, index)
+"vector.transfer_write"(%vector, %source, %index, %index) <{in_bounds=[true], operandSegmentSizes = array<i32: 1, 1, 2, 0>, permutation_map = affine_map<(d0) -> (d0)>}> : (vector<128xf32>, memref<?x?xf32>, index, index) -> ()
+// CHECK: requires a permutation_map with input dims of the same rank as the source type
+
+// -----
+
+%source, %vector, %index = "test.op"() : () -> (memref<?x?xf32>, vector<128xf32>, index)
+"vector.transfer_write"(%vector, %source, %index, %index) <{in_bounds=[true, true], operandSegmentSizes = array<i32: 1, 1, 2, 0>, permutation_map = affine_map<(d0, d1) -> (d0, d1)>}> : (vector<128xf32>, memref<?x?xf32>, index, index) -> ()
+// CHECK: requires a permutation_map with result dims of the same rank as the vector type
+
+// -----
+
+%source, %vector, %index = "test.op"() : () -> (memref<?x?xf32>, vector<128xf32>, index)
+"vector.transfer_write"(%vector, %source, %index, %index) <{in_bounds=[true, true], operandSegmentSizes = array<i32: 1, 1, 2, 0>, permutation_map = affine_map<(d0, d1) -> (d0 + d1)>}> : (vector<128xf32>, memref<?x?xf32>, index, index) -> ()
+// CHECK: requires a projected permutation_map (at most one dim or the zero constant can appear in each result)
+
+// -----
+
+%source, %vector, %index = "test.op"() : () -> (memref<?x?xf32>, vector<128xf32>, index)
+"vector.transfer_write"(%vector, %source, %index, %index) <{in_bounds=[true], operandSegmentSizes = array<i32: 1, 1, 2, 0>, permutation_map = affine_map<(d0, d1) -> (d0 + 1)>}> : (vector<128xf32>, memref<?x?xf32>, index, index) -> ()
+// CHECK: requires a projected permutation_map (at most one dim or the zero constant can appear in each result)
+
+// -----
+
+%source, %vector, %index = "test.op"() : () -> (memref<?x?x?xf32>, vector<3x7xf32>, index)
+"vector.transfer_write"(%vector, %source, %index, %index, %index) <{in_bounds=[true, true], operandSegmentSizes = array<i32: 1, 1, 3, 0>, permutation_map = affine_map<(d0, d1, d2) -> (d0, d0)>}> : (vector<3x7xf32>, memref<?x?x?xf32>, index, index, index) -> ()
+// CHECK: requires a permutation_map that is a permutation (found one dim used more than once)
+
+// -----
+
+%source, %vector, %index = "test.op"() : () -> (memref<?xf32>, vector<7xf32>, index)
+"vector.transfer_write"(%vector, %source, %index) <{in_bounds=[true], operandSegmentSizes = array<i32: 1, 1, 1, 0>, permutation_map = affine_map<(d0) -> (0)>}> : (vector<7xf32>, memref<?xf32>, index) -> ()
+// CHECK: requires a permutation_map that is a permutation (found one dim used more than once)
+
+// TODO transfer other tests from mlir/test/Dialect/Vector/invalid.mlir once verification for vector element types is implemented