diff --git a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td --- a/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td +++ b/mlir/include/mlir/Dialect/Linalg/IR/LinalgStructuredOps.td @@ -505,6 +505,7 @@ >>; class GenericOpBase : LinalgStructuredBase_Op, SingleBlockImplicitTerminator<"YieldOp">]> { diff --git a/mlir/test/Dialect/Linalg/affine.mlir b/mlir/test/Dialect/Linalg/affine.mlir --- a/mlir/test/Dialect/Linalg/affine.mlir +++ b/mlir/test/Dialect/Linalg/affine.mlir @@ -9,6 +9,8 @@ // CHECK-DAG: #[[$clampMinMap:.*]] = affine_map<(d0) -> (d0, 0)> +// CHECK-DAG: #[[$tset:.*]] = affine_set<(d0, d1) : (d0 - d1 >= 0)> + func @matmul(%arg0: memref, %M: index, %N: index, %K: index) { %c0 = constant 0 : index %c1 = constant 1 : index @@ -149,3 +151,72 @@ // CHECK-NEXT: cmpf // CHECK-NEXT: select // CHECK-NEXT: affine.store + +//----------------------------------------------------------------------------// +// Generic ops to loops. +//----------------------------------------------------------------------------// +#transpose_accesses = [ + affine_map<(H,W)->(H,W)>, + affine_map<(H,W)->(W,H)> +] +#trans_trait = { + indexing_maps = #transpose_accesses, + iterator_types = ["parallel","parallel"] +} + +func @transpose(%in: memref, %out: memref){ + + // Transpose + linalg.generic #trans_trait + ins(%in : memref) + outs(%out : memref) { + ^bb0(%a: f32, %b:f32): + linalg.yield %a : f32 + } + return +} +// CHECK-LABEL: @transpose +// CHECK-NEXT: constant +// CHECK-NEXT: constant +// CHECK-NEXT: memref.dim +// CHECK-NEXT: memref.dim +// CHECK-NEXT: affine.for %[[idx0:.*]] = {{.*}} +// CHECK-NEXT: affine.for %[[idx1:.*]] = {{.*}} +// CHECK-NEXT: %[[v0:.*]] = affine.load %{{.*}}[%[[idx0]], %[[idx1]]] {{.*}} +// CHECK-NEXT: affine.store %[[v0]], %{{.*}}[%[[idx1]], %[[idx0]]] {{.*}} + +#set0 = affine_set<(d0,d1) : (d0-d1>=0)> +func @transpose_inplace(%out: memref){ + + linalg.indexed_generic #trans_trait + outs(%out, %out : memref, memref) { + ^bb0(%i: index, %j: index, %a: f32, %b:f32): + + // With the addtion of the AffineScope trait, linalg.*generic ops can have + // affine ops in its body + %r1, %r2 = affine.if #set0(%i,%j) -> (f32,f32) { + affine.yield %b,%a : f32, f32 + } else { + affine.yield %a,%b : f32, f32 + } + linalg.yield %r1, %r2 : f32 ,f32 + } + return +} + +// CHECK-LABEL: @transpose_inplace +// CHECK-NEXT: constant +// CHECK-NEXT: constant +// CHECK-NEXT: memref.dim +// CHECK-NEXT: memref.dim +// CHECK-NEXT: affine.for %[[idx0:.*]] = {{.*}} +// CHECK-NEXT: affine.for %[[idx1:.*]] = {{.*}} +// CHECK-NEXT: %[[v0:.*]] = affine.load %{{.*}}[%[[idx0]], %[[idx1]]] {{.*}} +// CHECK-NEXT: %[[v1:.*]] = affine.load %{{.*}}[%[[idx1]], %[[idx0]]] {{.*}} +// CHECK-NEXT: %[[ret:.*]]:2 = affine.if #[[$tset]](%[[idx0]], %[[idx1]]) {{.*}} +// CHECK-NEXT: affine.yield %[[v1]], %[[v0]] {{.*}} +// CHECK-NEXT: } else { +// CHECK-NEXT: affine.yield %[[v0]], %[[v1]] {{.*}} +// CHECK-NEXT: } +// CHECK-NEXT: affine.store %[[ret]]#0, %{{.*}}[%[[idx0]], %[[idx1]]] {{.*}} +// CHECK-NEXT: affine.store %[[ret]]#1, %{{.*}}[%[[idx1]], %[[idx0]]] {{.*}}