1// RUN: mlir-opt %s -linalg-tile-to-parallel-loops="linalg-tile-sizes=2,4,8" -split-input-file | FileCheck %s 2// RUN: mlir-opt %s -linalg-tile-to-parallel-loops="linalg-tile-sizes=2" -split-input-file | FileCheck %s -check-prefix=TILE1 3// RUN: mlir-opt %s -linalg-tile-to-parallel-loops="linalg-tile-sizes=2,4" -split-input-file | FileCheck %s -check-prefix=TILE2 4 5func @gemm(%arg0 : memref<?x?xf32>, 6 %arg1 : memref<?x?xf32>, 7 %arg2 : memref<?x?xf32>) 8{ 9 linalg.matmul ins(%arg0, %arg1: memref<?x?xf32>, memref<?x?xf32>) 10 outs(%arg2: memref<?x?xf32>) 11 return 12} 13// CHECK-LABEL: func @gemm 14// CHECK-DAG: %[[C2:.*]] = constant 2 : index 15// CHECK-DAG: %[[C4:.*]] = constant 4 : index 16// CHECK-DAG: %[[C8:.*]] = constant 8 : index 17// CHECK: scf.parallel (%[[ARG3:.*]], %[[ARG4:.*]]) = 18// CHECK-SAME: step (%[[C2]], %[[C4]]) 19// CHECK: scf.for %[[ARG5:.*]] = 20// CHECK-SAME: step %[[C8]] 21// CHECK: %[[SV1:.*]] = subview %{{.*}}[%[[ARG3]], %[[ARG5]]] 22// CHECK: %[[SV2:.*]] = subview %{{.*}}[%[[ARG5]], %[[ARG4]]] 23// CHECK: %[[SV3:.*]] = subview %{{.*}}[%[[ARG3]], %[[ARG4]]] 24// CHECK: linalg.matmul ins(%[[SV1]], %[[SV2]]{{.*}} outs(%[[SV3]] 25 26// TILE1-LABEL: func @gemm 27// TILE1-DAG: %[[C2:.*]] = constant 2 : index 28// TILE1: scf.parallel (%[[ARG3:.*]]) = 29// TILE1-SAME: step (%[[C2]]) 30// TILE1: %[[SV1:.*]] = subview %{{.*}}[%[[ARG3]], 0] 31// TILE1: %[[SV3:.*]] = subview %{{.*}}[%[[ARG3]], 0] 32// TILE1-NOT: subview 33// TILE1: linalg.matmul ins(%[[SV1]], %{{.*}} outs(%[[SV3]] 34 35// TILE2-LABEL: func @gemm 36// TILE2-DAG: %[[C2:.*]] = constant 2 : index 37// TILE2-DAG: %[[C4:.*]] = constant 4 : index 38// TILE2: scf.parallel (%[[ARG3:.*]], %[[ARG4:.*]]) = 39// TILE2-SAME: step (%[[C2]], %[[C4]]) 40// TILE2: %[[SV1:.*]] = subview %{{.*}}[%[[ARG3]], 0] 41// TILE2: %[[SV2:.*]] = subview %{{.*}}[0, %[[ARG4]]] 42// TILE2: %[[SV3:.*]] = subview %{{.*}}[%[[ARG3]], %[[ARG4]]] 43// TILE2: linalg.matmul ins(%[[SV1]], %[[SV2]]{{.*}} outs(%[[SV3]] 44 45// ----- 46 47#map0 = affine_map<(d0, d1, d2) -> (d0, d1, d2)> 48#map1 = affine_map<(d0, d1, d2) -> (d0, d2)> 49#map2 = affine_map<(d0, d1, d2) -> (d1)> 50#accesses = [#map0, #map1, #map2] 51#trait = { 52 args_in = 2 : i64, 53 args_out = 1 : i64, 54 iterator_types = ["reduction", "parallel", "reduction"], 55 indexing_maps = #accesses 56} 57 58func @reduction(%arg0 : memref<?x?x?xf32>, 59 %arg1 : memref<?x?xf32>, 60 %arg2 : memref<?xf32>) 61{ 62 linalg.generic #trait 63 ins(%arg0, %arg1 : memref<?x?x?xf32>, memref<?x?xf32>) 64 outs(%arg2 : memref<?xf32>) { 65 ^bb0(%arg3 : f32, %arg4 : f32, %arg5 : f32): 66 %0 = addf %arg3, %arg4 : f32 67 %1 = addf %0, %arg5 : f32 68 linalg.yield %1 : f32 69 } 70 return 71} 72 73// CHECK-LABEL: func @reduction 74// CHECK-DAG: %[[C2:.*]] = constant 2 : index 75// CHECK-DAG: %[[C4:.*]] = constant 4 : index 76// CHECK-DAG: %[[C8:.*]] = constant 8 : index 77// CHECK: scf.for %[[ARG3:.*]] = 78// CHECK-SAME: step %[[C2]] 79// CHECK: scf.parallel (%[[ARG4:.*]]) = 80// CHECK-SAME: step (%[[C4]]) 81// CHECK: scf.for %[[ARG5:.*]] = 82// CHECK-SAME: step %[[C8]] 83// CHECK: %[[SV1:.*]] = subview %{{.*}}[%[[ARG3]], %[[ARG4]], %[[ARG5]]] 84// CHECK: %[[SV2:.*]] = subview %{{.*}}[%[[ARG3]], %[[ARG5]]] 85// CHECK: %[[SV3:.*]] = subview %{{.*}}[%[[ARG4]]] 86// CHECK: linalg.generic 87// CHECK-SAME: ins(%[[SV1]], %[[SV2]] 88// CHECK-SAME: outs(%[[SV3]] 89 90// TILE1-LABEL: func @reduction 91// TILE1-DAG: %[[C2:.*]] = constant 2 : index 92// TILE1: scf.for %[[ARG3:.*]] = 93// TILE1-SAME: step %[[C2]] 94// TILE1: %[[SV1:.*]] = subview %{{.*}}[%[[ARG3]], 0, 0] 95// TILE1: %[[SV2:.*]] = subview %{{.*}}[%[[ARG3]], 0] 96// TILE1-NOT: subview 97// TILE1: linalg.generic 98// TILE1-SAME: ins(%[[SV1]], %[[SV2]] 99// TILE1-SAME: outs(%{{.*}} 100 101// TILE2-LABEL: func @reduction 102// TILE2-DAG: %[[C2:.*]] = constant 2 : index 103// TILE2-DAG: %[[C4:.*]] = constant 4 : index 104// TILE2: scf.for %[[ARG3:.*]] = 105// TILE2-SAME: step %[[C2]] 106// TILE2: scf.parallel (%[[ARG4:.*]]) = 107// TILE2-SAME: step (%[[C4]]) 108// TILE2: %[[SV1:.*]] = subview %{{.*}}[%[[ARG3]], %[[ARG4]], 0] 109// TILE2: %[[SV2:.*]] = subview %{{.*}}[%[[ARG3]], 0] 110// TILE2: %[[SV3:.*]] = subview %{{.*}}[%[[ARG4]]] 111// TILE2: linalg.generic 112// TILE2-SAME: ins(%[[SV1]], %[[SV2]] 113// TILE2-SAME: outs(%[[SV3]] 114