|
150
|
1 // RUN: mlir-opt %s -linalg-tile-to-parallel-loops -linalg-tile-sizes=2 | FileCheck %s -check-prefix=TILE-2 --dump-input-on-failure
|
|
|
2 // RUN: mlir-opt %s -linalg-tile-to-parallel-loops -linalg-tile-sizes=0,2 | FileCheck %s -check-prefix=TILE-02 --dump-input-on-failure
|
|
|
3 // RUN: mlir-opt %s -linalg-tile-to-parallel-loops -linalg-tile-sizes=0,0,2 | FileCheck %s -check-prefix=TILE-002 --dump-input-on-failure
|
|
|
4 // RUN: mlir-opt %s -linalg-tile-to-parallel-loops -linalg-tile-sizes=2,3,4 | FileCheck %s -check-prefix=TILE-234 --dump-input-on-failure
|
|
|
5
|
|
|
6 #id_2d = affine_map<(i, j) -> (i, j)>
|
|
|
7 #pointwise_2d_trait = {
|
|
|
8 args_in = 2,
|
|
|
9 args_out = 1,
|
|
|
10 indexing_maps = [#id_2d, #id_2d, #id_2d],
|
|
|
11 iterator_types = ["parallel", "parallel"]
|
|
|
12 }
|
|
|
13
|
|
|
14 func @sum(%lhs: memref<?x?xf32, offset: ?, strides: [?, 1]>,
|
|
|
15 %rhs: memref<?x?xf32, offset: ?, strides: [?, 1]>,
|
|
|
16 %sum: memref<?x?xf32, offset: ?, strides: [?, 1]>) {
|
|
|
17 linalg.generic #pointwise_2d_trait %lhs, %rhs, %sum {
|
|
|
18 ^bb0(%lhs_in: f32, %rhs_in: f32, %sum_out: f32):
|
|
|
19 %result = addf %lhs_in, %rhs_in : f32
|
|
|
20 linalg.yield %result : f32
|
|
|
21 }: memref<?x?xf32, offset: ?, strides: [?, 1]>,
|
|
|
22 memref<?x?xf32, offset: ?, strides: [?, 1]>,
|
|
|
23 memref<?x?xf32, offset: ?, strides: [?, 1]>
|
|
|
24 return
|
|
|
25 }
|
|
|
26 // TILE-2-LABEL: func @sum(
|
|
|
27 // TILE-2-SAME: [[LHS:%.*]]: {{.*}}, [[RHS:%.*]]: {{.*}}, [[SUM:%.*]]: {{.*}}) {
|
|
|
28 // TILE-2-DAG: [[C0:%.*]] = constant 0 : index
|
|
|
29 // TILE-2-DAG: [[C1:%.*]] = constant 1 : index
|
|
|
30 // TILE-2-DAG: [[C2:%.*]] = constant 2 : index
|
|
|
31 // TILE-2: [[LHS_ROWS:%.*]] = dim [[LHS]], 0
|
|
|
32 // TILE-2: loop.parallel ([[I:%.*]]) = ([[C0]]) to ([[LHS_ROWS]]) step ([[C2]]) {
|
|
|
33 // TILE-2-NO: loop.parallel
|
|
|
34 // TILE-2: [[LHS_SUBVIEW:%.*]] = std.subview [[LHS]]
|
|
|
35 // TILE-2: [[RHS_SUBVIEW:%.*]] = std.subview [[RHS]]
|
|
|
36 // TILE-2: [[SUM_SUBVIEW:%.*]] = std.subview [[SUM]]
|
|
|
37 // TILE-2: linalg.generic {{.*}} [[LHS_SUBVIEW]], [[RHS_SUBVIEW]], [[SUM_SUBVIEW]] {
|
|
|
38
|
|
|
39 // TILE-02-LABEL: func @sum(
|
|
|
40 // TILE-02-SAME: [[LHS:%.*]]: {{.*}}, [[RHS:%.*]]: {{.*}}, [[SUM:%.*]]: {{.*}}) {
|
|
|
41 // TILE-02-DAG: [[C0:%.*]] = constant 0 : index
|
|
|
42 // TILE-02-DAG: [[C1:%.*]] = constant 1 : index
|
|
|
43 // TILE-02-DAG: [[C2:%.*]] = constant 2 : index
|
|
|
44 // TILE-02: [[LHS_COLS:%.*]] = dim [[LHS]], 1
|
|
|
45 // TILE-02: loop.parallel ([[I:%.*]]) = ([[C0]]) to ([[LHS_COLS]]) step ([[C2]]) {
|
|
|
46 // TILE-02-NO: loop.parallel
|
|
|
47 // TILE-02: [[LHS_SUBVIEW:%.*]] = std.subview [[LHS]]
|
|
|
48 // TILE-02: [[RHS_SUBVIEW:%.*]] = std.subview [[RHS]]
|
|
|
49 // TILE-02: [[SUM_SUBVIEW:%.*]] = std.subview [[SUM]]
|
|
|
50 // TILE-02: linalg.generic {{.*}} [[LHS_SUBVIEW]], [[RHS_SUBVIEW]], [[SUM_SUBVIEW]] {
|
|
|
51
|
|
|
52 // TILE-002-LABEL: func @sum(
|
|
|
53 // TILE-002-SAME: [[LHS:%.*]]: {{.*}}, [[RHS:%.*]]: {{.*}}, [[SUM:%.*]]: {{.*}}) {
|
|
|
54 // TILE-002-NO: loop.parallel
|
|
|
55 // TILE-002: linalg.generic {{.*}} [[LHS]], [[RHS]], [[SUM]] {
|
|
|
56
|
|
|
57 // TILE-234-LABEL: func @sum(
|
|
|
58 // TILE-234-SAME: [[LHS:%.*]]: {{.*}}, [[RHS:%.*]]: {{.*}}, [[SUM:%.*]]: {{.*}}) {
|
|
|
59 // TILE-234-DAG: [[C0:%.*]] = constant 0 : index
|
|
|
60 // TILE-234-DAG: [[C1:%.*]] = constant 1 : index
|
|
|
61 // TILE-234-DAG: [[C2:%.*]] = constant 2 : index
|
|
|
62 // TILE-234-DAG: [[C3:%.*]] = constant 3 : index
|
|
|
63 // TILE-234: [[LHS_ROWS:%.*]] = dim [[LHS]], 0
|
|
|
64 // TILE-234: [[LHS_COLS:%.*]] = dim [[LHS]], 1
|
|
|
65 // TILE-234: loop.parallel ([[I:%.*]], [[J:%.*]]) = ([[C0]], [[C0]]) to ([[LHS_ROWS]], [[LHS_COLS]]) step ([[C2]], [[C3]]) {
|
|
|
66 // TILE-234-NO: loop.parallel
|
|
|
67 // TILE-234: [[LHS_SUBVIEW:%.*]] = std.subview [[LHS]]
|
|
|
68 // TILE-234: [[RHS_SUBVIEW:%.*]] = std.subview [[RHS]]
|
|
|
69 // TILE-234: [[SUM_SUBVIEW:%.*]] = std.subview [[SUM]]
|
|
|
70 // TILE-234: linalg.generic {{.*}} [[LHS_SUBVIEW]], [[RHS_SUBVIEW]], [[SUM_SUBVIEW]] {
|
