(dialects): add linalg.tranpose

tests/filecheck/mlir-conversion/with-mlir/dialects/linalg/invalid_ops.mlir

@@ -1,4 +1,4 @@
-// RUN: xdsl-opt --verify-diagnostics %s | filecheck %s
+// RUN: xdsl-opt --verify-diagnostics --split-input-file %s | filecheck %s
 
 
 builtin.module {
@@ -9,4 +9,38 @@ builtin.module {
 
 }
 
+// -----
+
+builtin.module {
+  %0, %1 = "test.op"() : () -> (tensor<16x64xf32>, tensor<64x16x1xf32>)
+
+  // CHECK: Operation does not verify: Input rank (2) does not match output rank (3)
+   %res_transpose = "linalg.transpose"(%0, %1) {"permutation" = array<i64: 1, 0>} : (tensor<16x64xf32>, tensor<64x16x1xf32>) -> tensor<64x16x1xf32>
+
+}
+
+// -----
+
+builtin.module {
+  %0, %1 = "test.op"() : () -> (tensor<16x64xf32>, tensor<64x16xf32>)
+
+  // CHECK: Operation does not verify: Input rank (2) does not match size of permutation (3)
+   %res_transpose = "linalg.transpose"(%0, %1) {"permutation" = array<i64: 1, 2, 3>} : (tensor<16x64xf32>, tensor<64x16xf32>) -> tensor<64x16xf32>
+
+}
+
+// -----
+
+builtin.module {
+  %0, %1 = "test.op"() : () -> (tensor<16x32x64xf32>, tensor<32x64x16xf32>)
+
+  // CHECK:  Operation does not verify: dim(result, 1) = 64 doesn't match dim(input, permutation[1]) = 32
+  %res_transpose = "linalg.transpose"(%0, %1) {"permutation" = array<i64: 1, 1, 2>} : (tensor<16x32x64xf32>, tensor<32x64x16xf32>) -> tensor<32x64x16xf32>
+
+}
+
+
+
+
+
 

tests/filecheck/mlir-conversion/with-mlir/dialects/linalg/ops.mlir

@@ -23,6 +23,10 @@ linalg.generic {indexing_maps = [affine_map<(d0, d1) -> ()>, affine_map<(d0, d1)
 
 %mul = linalg.mul ins(%2, %2 : tensor<2x3xf32>, tensor<2x3xf32>) outs(%3 : tensor<2x3xf32>) -> tensor<2x3xf32>
 
+%5, %6 = "test.op"() : () -> (tensor<16x64xf32>, tensor<64x16xf32>)
+
+%transposed = linalg.transpose ins(%5 : tensor<16x64xf32>) outs(%6 : tensor<64x16xf32>) permutation = [1, 0]
+
 %sum_2 = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>], iterator_types = ["parallel", "parallel"]} ins(%2, %2 : tensor<2x3xf32>, tensor<2x3xf32>) outs(%2 : tensor<2x3xf32>) {
 ^bb0(%in: f32, %in_0: f32, %out: f32):
     %acc = arith.addf %in, %in_0 : f32
@@ -47,6 +51,8 @@ linalg.generic {indexing_maps = [affine_map<(d0, d1) -> ()>, affine_map<(d0, d1)
 // CHECK-NEXT:    %cst = arith.constant 0.000000e+00 : f32
 // CHECK-NEXT:    %3 = linalg.fill ins(%cst : f32) outs(%1#0 : tensor<2x3xf32>) -> tensor<2x3xf32>
 // CHECK-NEXT:    %4 = linalg.mul  ins(%1#0, %1#0 : tensor<2x3xf32>, tensor<2x3xf32>) outs(%1#1 : tensor<2x3xf32>) -> tensor<2x3xf32>
+// CHECK-NEXT:    %5:2 = "test.op"() : () -> (tensor<16x64xf32>, tensor<64x16xf32>)
+// CHECK-NEXT:    %transposed  = linalg.transpose ins(%5#0 : tensor<16x64xf32>) outs(%5#1 : tensor<64x16xf32>) permutation = [1, 0]
 // CHECK-NEXT:    %{{.*}} = linalg.generic {indexing_maps = [#map1, #map1, #map1], iterator_types = ["parallel", "parallel"]} ins(%1#0, %1#0 : tensor<2x3xf32>, tensor<2x3xf32>) outs(%1#0 : tensor<2x3xf32>) {
 // CHECK-NEXT:    ^bb0(%in: f32, %in_0: f32, %out: f32):
 // CHECK-NEXT:      %{{.*}} = arith.addf %in, %in_0 : f32

xdsl/dialects/linalg.py

@@ -12,9 +12,13 @@
     AnyShapedType,
     AnyTensorType,
     ArrayAttr,
+    DenseArrayBase,
     IntegerType,
+    MemRefType,
     ShapedType,
     StringAttr,
+    TensorType,
+    i64,
 )
 from xdsl.dialects.utils import (
     AbstractYieldOperation,
@@ -27,8 +31,10 @@
     ParsePropInAttrDict,
     VarOperand,
     VarOpResult,
+    attr_def,
     irdl_attr_definition,
     irdl_op_definition,
+    operand_def,
     opt_attr_def,
     prop_def,
     region_def,
@@ -490,6 +496,111 @@ def __init__(
         )
 
 
+@irdl_op_definition
+class TransposeOp(IRDLOperation):
+    """
+    Transpose operator
+
+    See https://mlir.llvm.org/docs/Dialects/Linalg/#linalgtranspose-linalgtransposeop
+    """
+
+    name = "linalg.transpose"
+
+    input = operand_def(MemRefType | AnyTensorType)
+    init = operand_def(MemRefType | AnyTensorType)
+    result = var_result_def(AnyTensorType)
+
+    permutation = attr_def(DenseArrayBase)
+
+    def __init__(
+        self,
+        input: SSAValue,
+        init: SSAValue,
+        permutation: Attribute,
+        result: Attribute | None = None,
+    ):
+        super().__init__(
+            attributes={
+                "permutation": permutation,
+            },
+            operands=(input, init),
+            result_types=(result,),
+        )
+
+    def verify_(self) -> None:
+
+        assert isinstance(input_type := self.input.type, TensorType | MemRefType)
+        assert isinstance(init_type := self.init.type, TensorType | MemRefType)
+
+        input_shape = input_type.get_shape()
+        init_shape = init_type.get_shape()
+
+        if (input_rank := len(input_shape)) != (init_rank := len(init_shape)):
+            raise VerifyException(
+                f"Input rank ({input_rank}) does not match output rank ({init_rank})"
+            )
+        if (input_rank := len(input_shape)) != (
+            permutation_size := len(self.permutation.data)
+        ):
+            raise VerifyException(
+                f"Input rank ({input_rank}) does not match size of permutation ({permutation_size})"
+            )
+
+        permutation_shape = cast(list[int], self.permutation.as_tuple())
+
+        for i in range(len(input_shape)):
+            input_dimension = input_shape[permutation_shape[i]]
+            init_dimension = init_shape[i]
+
+            if input_dimension != init_dimension:
+                raise VerifyException(
+                    f"dim(result, {i}) = {init_dimension} "
+                    f"doesn't match dim(input, permutation[{i}]) = {input_dimension}"
+                )
+
+    def print(self, printer: Printer):
+        printer.print_string(" ins(")
+        printer.print(self.input)
+        printer.print_string(":")
+        printer.print(self.input.type)
+        printer.print_string(")")
+        printer.print_string(" outs(")
+        printer.print(self.init)
+        printer.print_string(":")
+        printer.print(self.init.type)
+        printer.print_string(") ")
+        printer.print_string("permutation")
+        printer.print_string(" = ")
+        printer.print(list(self.permutation.as_tuple()))
+
+    @classmethod
+    def parse(cls, parser: Parser) -> Self:
+        parser.parse_characters("ins")
+        parser.parse_punctuation("(")
+        input = parser.parse_operand()
+        parser.parse_punctuation(":")
+        parser.parse_type()
+        parser.parse_punctuation(")")
+        parser.parse_characters("outs")
+        parser.parse_punctuation("(")
+        init = parser.parse_operand()
+        parser.parse_punctuation(":")
+        result = parser.parse_type()
+        parser.parse_punctuation(")")
+        parser.parse_keyword("permutation")
+        parser.parse_punctuation("=")
+        permutation = parser.parse_comma_separated_list(
+            parser.Delimiter.SQUARE, parser.parse_integer
+        )
+        transpose = cls(
+            input,
+            init,
+            DenseArrayBase.create_dense_int_or_index(i64, permutation),
+            result,
+        )
+        return transpose
+
+
 Linalg = Dialect(
     "linalg",
     [
@@ -498,6 +609,7 @@ def __init__(
         AddOp,
         FillOp,
         MulOp,
+        TransposeOp,
     ],
     [
         IteratorTypeAttr,