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