1// RUN: mlir-opt %s -test-vector-transfer-full-partial-split | FileCheck %s
2// RUN: mlir-opt %s -test-vector-transfer-full-partial-split=use-linalg-copy | FileCheck %s --check-prefix=LINALG
3
4// CHECK-DAG: #[[$map_p4:.*]] = affine_map<()[s0] -> (s0 + 4)>
5// CHECK-DAG: #[[$map_p8:.*]] = affine_map<()[s0] -> (s0 + 8)>
6// CHECK-DAG: #[[$map_2d_stride_1:.*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)>
7
8// LINALG-DAG: #[[$map_p4:.*]] = affine_map<()[s0] -> (s0 + 4)>
9// LINALG-DAG: #[[$map_p8:.*]] = affine_map<()[s0] -> (s0 + 8)>
10// LINALG-DAG: #[[$map_2d_stride_1:.*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)>
11// LINALG-DAG: #[[$map_2d_dynamic:.*]] = affine_map<(d0, d1)[s0, s1, s2] -> (d0 * s1 + s0 + d1 * s2)>
12// LINALG-DAG: #[[$bounds_map_4:.*]] = affine_map<(d0, d1, d2) -> (d0 - d1, 4)>
13// LINALG-DAG: #[[$bounds_map_8:.*]] = affine_map<(d0, d1, d2) -> (d0 - d1, 8)>
14
15// CHECK-LABEL: split_vector_transfer_read_2d(
16//  CHECK-SAME: %[[A:[a-zA-Z0-9]*]]: memref
17//  CHECK-SAME: %[[i:[a-zA-Z0-9]*]]: index
18//  CHECK-SAME: %[[j:[a-zA-Z0-9]*]]: index
19
20// LINALG-LABEL: split_vector_transfer_read_2d(
21//  LINALG-SAME: %[[A:[a-zA-Z0-9]*]]: memref
22//  LINALG-SAME: %[[i:[a-zA-Z0-9]*]]: index
23//  LINALG-SAME: %[[j:[a-zA-Z0-9]*]]: index
24func @split_vector_transfer_read_2d(%A: memref<?x8xf32>, %i: index, %j: index) -> vector<4x8xf32> {
25  %c0 = constant 0 : index
26  %f0 = constant 0.0 : f32
27
28  //  CHECK-DAG: %[[c0:.*]] = constant 0 : index
29  //  CHECK-DAG: %[[c8:.*]] = constant 8 : index
30  //  CHECK-DAG: %[[cst:.*]] = constant 0.000000e+00 : f32
31  // alloca for boundary full tile
32  //      CHECK: %[[alloc:.*]] = alloca() {alignment = 32 : i64} : memref<4x8xf32>
33  // %i + 4 <= dim(%A, 0)
34  //      CHECK: %[[idx0:.*]] = affine.apply #[[$map_p4]]()[%[[i]]]
35  //      CHECK: %[[d0:.*]] = dim %[[A]], %[[c0]] : memref<?x8xf32>
36  //      CHECK: %[[cmp0:.*]] = cmpi "sle", %[[idx0]], %[[d0]] : index
37  // %j + 8 <= dim(%A, 1)
38  //      CHECK: %[[idx1:.*]] = affine.apply #[[$map_p8]]()[%[[j]]]
39  //      CHECK: %[[cmp1:.*]] = cmpi "sle", %[[idx1]], %[[c8]] : index
40  // are both conds true
41  //      CHECK: %[[cond:.*]] = and %[[cmp0]], %[[cmp1]] : i1
42  //      CHECK: %[[ifres:.*]]:3 = scf.if %[[cond]] -> (memref<?x8xf32>, index, index) {
43  //               inBounds, just yield %A
44  //      CHECK:   scf.yield %[[A]], %[[i]], %[[j]] : memref<?x8xf32>, index, index
45  //      CHECK: } else {
46  //               slow path, fill tmp alloc and yield a memref_casted version of it
47  //      CHECK:   %[[slow:.*]] = vector.transfer_read %[[A]][%[[i]], %[[j]]], %cst :
48  // CHECK-SAME:     memref<?x8xf32>, vector<4x8xf32>
49  //      CHECK:   %[[cast_alloc:.*]] = vector.type_cast %[[alloc]] :
50  // CHECK-SAME:     memref<4x8xf32> to memref<vector<4x8xf32>>
51  //      CHECK:   store %[[slow]], %[[cast_alloc]][] : memref<vector<4x8xf32>>
52  //      CHECK:   %[[yielded:.*]] = memref_cast %[[alloc]] :
53  // CHECK-SAME:     memref<4x8xf32> to memref<?x8xf32>
54  //      CHECK:   scf.yield %[[yielded]], %[[c0]], %[[c0]] :
55  // CHECK-SAME:     memref<?x8xf32>, index, index
56  //      CHECK: }
57  //      CHECK: %[[res:.*]] = vector.transfer_read %[[ifres]]#0[%[[ifres]]#1, %[[ifres]]#2], %[[cst]]
58  // CHECK_SAME:   {masked = [false, false]} : memref<?x8xf32>, vector<4x8xf32>
59
60  //  LINALG-DAG: %[[c0:.*]] = constant 0 : index
61  //  LINALG-DAG: %[[c1:.*]] = constant 1 : index
62  //  LINALG-DAG: %[[c4:.*]] = constant 4 : index
63  //  LINALG-DAG: %[[c8:.*]] = constant 8 : index
64  //  LINALG-DAG: %[[cst:.*]] = constant 0.000000e+00 : f32
65  // alloca for boundary full tile
66  //      LINALG: %[[alloc:.*]] = alloca() {alignment = 32 : i64} : memref<4x8xf32>
67  // %i + 4 <= dim(%A, 0)
68  //      LINALG: %[[idx0:.*]] = affine.apply #[[$map_p4]]()[%[[i]]]
69  //      LINALG: %[[d0:.*]] = dim %[[A]], %[[c0]] : memref<?x8xf32>
70  //      LINALG: %[[cmp0:.*]] = cmpi "sle", %[[idx0]], %[[d0]] : index
71  // %j + 8 <= dim(%A, 1)
72  //      LINALG: %[[idx1:.*]] = affine.apply #[[$map_p8]]()[%[[j]]]
73  //      LINALG: %[[cmp1:.*]] = cmpi "sle", %[[idx1]], %[[c8]] : index
74  // are both conds true
75  //      LINALG: %[[cond:.*]] = and %[[cmp0]], %[[cmp1]] : i1
76  //      LINALG: %[[ifres:.*]]:3 = scf.if %[[cond]] -> (memref<?x8xf32>, index, index) {
77  //               inBounds, just yield %A
78  //      LINALG:   scf.yield %[[A]], %[[i]], %[[j]] : memref<?x8xf32>, index, index
79  //      LINALG: } else {
80  //               slow path, fill tmp alloc and yield a memref_casted version of it
81  //      LINALG:   linalg.fill(%[[alloc]], %[[cst]]) : memref<4x8xf32>, f32
82  //      LINALG:   %[[d0:.*]] = dim %[[A]], %[[c0]] : memref<?x8xf32>
83  //      LINALG:   %[[sv0:.*]] = affine.min #[[$bounds_map_4]](%[[d0]], %[[i]], %[[c4]])
84  //      LINALG:   %[[sv1:.*]] = affine.min #[[$bounds_map_8]](%[[c8]], %[[j]], %[[c8]])
85  //      LINALG:   %[[sv:.*]] = subview %[[A]][%[[i]], %[[j]]] [%[[sv0]], %[[sv1]]] [%[[c1]], %[[c1]]]
86  // LINALG-SAME:     memref<?x8xf32> to memref<?x?xf32, #[[$map_2d_dynamic]]>
87  //      LINALG:   linalg.copy(%[[sv]], %[[alloc]]) : memref<?x?xf32, #[[$map_2d_dynamic]]>, memref<4x8xf32>
88  //      LINALG:   %[[yielded:.*]] = memref_cast %[[alloc]] :
89  // LINALG-SAME:     memref<4x8xf32> to memref<?x8xf32>
90  //      LINALG:   scf.yield %[[yielded]], %[[c0]], %[[c0]] :
91  // LINALG-SAME:     memref<?x8xf32>, index, index
92  //      LINALG: }
93  //      LINALG: %[[res:.*]] = vector.transfer_read %[[ifres]]#0[%[[ifres]]#1, %[[ifres]]#2], %[[cst]]
94  // LINALG_SAME:   {masked = [false, false]} : memref<?x8xf32>, vector<4x8xf32>
95  %1 = vector.transfer_read %A[%i, %j], %f0 : memref<?x8xf32>, vector<4x8xf32>
96
97  // LINALG: return %[[res]] : vector<4x8xf32>
98  return %1: vector<4x8xf32>
99}
100
101// CHECK-LABEL: split_vector_transfer_read_strided_2d(
102//  CHECK-SAME: %[[A:[a-zA-Z0-9]*]]: memref
103//  CHECK-SAME: %[[i:[a-zA-Z0-9]*]]: index
104//  CHECK-SAME: %[[j:[a-zA-Z0-9]*]]: index
105
106// LINALG-LABEL: split_vector_transfer_read_strided_2d(
107//  LINALG-SAME: %[[A:[a-zA-Z0-9]*]]: memref
108//  LINALG-SAME: %[[i:[a-zA-Z0-9]*]]: index
109//  LINALG-SAME: %[[j:[a-zA-Z0-9]*]]: index
110func @split_vector_transfer_read_strided_2d(
111    %A: memref<7x8xf32, offset:?, strides:[?, 1]>,
112    %i: index, %j: index) -> vector<4x8xf32> {
113  %c0 = constant 0 : index
114  %f0 = constant 0.0 : f32
115
116  //  CHECK-DAG: %[[c0:.*]] = constant 0 : index
117  //  CHECK-DAG: %[[c7:.*]] = constant 7 : index
118  //  CHECK-DAG: %[[c8:.*]] = constant 8 : index
119  //  CHECK-DAG: %[[cst:.*]] = constant 0.000000e+00 : f32
120  // alloca for boundary full tile
121  //      CHECK: %[[alloc:.*]] = alloca() {alignment = 32 : i64} : memref<4x8xf32>
122  // %i + 4 <= dim(%A, 0)
123  //      CHECK: %[[idx0:.*]] = affine.apply #[[$map_p4]]()[%[[i]]]
124  //      CHECK: %[[cmp0:.*]] = cmpi "sle", %[[idx0]], %[[c7]] : index
125  // %j + 8 <= dim(%A, 1)
126  //      CHECK: %[[idx1:.*]] = affine.apply #[[$map_p8]]()[%[[j]]]
127  //      CHECK: %[[cmp1:.*]] = cmpi "sle", %[[idx1]], %[[c8]] : index
128  // are both conds true
129  //      CHECK: %[[cond:.*]] = and %[[cmp0]], %[[cmp1]] : i1
130  //      CHECK: %[[ifres:.*]]:3 = scf.if %[[cond]] -> (memref<?x8xf32, #[[$map_2d_stride_1]]>, index, index) {
131  //               inBounds but not cast-compatible: yield a memref_casted form of %A
132  //      CHECK:   %[[casted:.*]] = memref_cast %arg0 :
133  // CHECK-SAME:     memref<7x8xf32, #[[$map_2d_stride_1]]> to memref<?x8xf32, #[[$map_2d_stride_1]]>
134  //      CHECK:   scf.yield %[[casted]], %[[i]], %[[j]] :
135  // CHECK-SAME:     memref<?x8xf32, #[[$map_2d_stride_1]]>, index, index
136  //      CHECK: } else {
137  //               slow path, fill tmp alloc and yield a memref_casted version of it
138  //      CHECK:   %[[slow:.*]] = vector.transfer_read %[[A]][%[[i]], %[[j]]], %cst :
139  // CHECK-SAME:     memref<7x8xf32, #[[$map_2d_stride_1]]>, vector<4x8xf32>
140  //      CHECK:   %[[cast_alloc:.*]] = vector.type_cast %[[alloc]] :
141  // CHECK-SAME:     memref<4x8xf32> to memref<vector<4x8xf32>>
142  //      CHECK:   store %[[slow]], %[[cast_alloc]][] :
143  // CHECK-SAME:     memref<vector<4x8xf32>>
144  //      CHECK:   %[[yielded:.*]] = memref_cast %[[alloc]] :
145  // CHECK-SAME:     memref<4x8xf32> to memref<?x8xf32, #[[$map_2d_stride_1]]>
146  //      CHECK:   scf.yield %[[yielded]], %[[c0]], %[[c0]] :
147  // CHECK-SAME:     memref<?x8xf32, #[[$map_2d_stride_1]]>, index, index
148  //      CHECK: }
149  //      CHECK: %[[res:.*]] = vector.transfer_read {{.*}} {masked = [false, false]} :
150  // CHECK-SAME:   memref<?x8xf32, #[[$map_2d_stride_1]]>, vector<4x8xf32>
151
152  //  LINALG-DAG: %[[c0:.*]] = constant 0 : index
153  //  LINALG-DAG: %[[c1:.*]] = constant 1 : index
154  //  LINALG-DAG: %[[c4:.*]] = constant 4 : index
155  //  LINALG-DAG: %[[c7:.*]] = constant 7 : index
156  //  LINALG-DAG: %[[c8:.*]] = constant 8 : index
157  //  LINALG-DAG: %[[cst:.*]] = constant 0.000000e+00 : f32
158  // alloca for boundary full tile
159  //      LINALG: %[[alloc:.*]] = alloca() {alignment = 32 : i64} : memref<4x8xf32>
160  // %i + 4 <= dim(%A, 0)
161  //      LINALG: %[[idx0:.*]] = affine.apply #[[$map_p4]]()[%[[i]]]
162  //      LINALG: %[[cmp0:.*]] = cmpi "sle", %[[idx0]], %[[c7]] : index
163  // %j + 8 <= dim(%A, 1)
164  //      LINALG: %[[idx1:.*]] = affine.apply #[[$map_p8]]()[%[[j]]]
165  //      LINALG: %[[cmp1:.*]] = cmpi "sle", %[[idx1]], %[[c8]] : index
166  // are both conds true
167  //      LINALG: %[[cond:.*]] = and %[[cmp0]], %[[cmp1]] : i1
168  //      LINALG: %[[ifres:.*]]:3 = scf.if %[[cond]] -> (memref<?x8xf32, #[[$map_2d_stride_1]]>, index, index) {
169  //               inBounds but not cast-compatible: yield a memref_casted form of %A
170  //      LINALG:   %[[casted:.*]] = memref_cast %arg0 :
171  // LINALG-SAME:     memref<7x8xf32, #[[$map_2d_stride_1]]> to memref<?x8xf32, #[[$map_2d_stride_1]]>
172  //      LINALG:   scf.yield %[[casted]], %[[i]], %[[j]] :
173  // LINALG-SAME:     memref<?x8xf32, #[[$map_2d_stride_1]]>, index, index
174  //      LINALG: } else {
175  //               slow path, fill tmp alloc and yield a memref_casted version of it
176  //      LINALG:   linalg.fill(%[[alloc]], %[[cst]]) : memref<4x8xf32>, f32
177  //      LINALG:   %[[sv0:.*]] = affine.min #[[$bounds_map_4]](%[[c7]], %[[i]], %[[c4]])
178  //      LINALG:   %[[sv1:.*]] = affine.min #[[$bounds_map_8]](%[[c8]], %[[j]], %[[c8]])
179  //      LINALG:   %[[sv:.*]] = subview %[[A]][%[[i]], %[[j]]] [%[[sv0]], %[[sv1]]] [%[[c1]], %[[c1]]]
180  // LINALG-SAME:     memref<7x8xf32, #[[$map_2d_stride_1]]> to memref<?x?xf32, #[[$map_2d_dynamic]]>
181  //      LINALG:   linalg.copy(%[[sv]], %[[alloc]]) : memref<?x?xf32, #[[$map_2d_dynamic]]>, memref<4x8xf32>
182  //      LINALG:   %[[yielded:.*]] = memref_cast %[[alloc]] :
183  // LINALG-SAME:     memref<4x8xf32> to memref<?x8xf32, #[[$map_2d_stride_1]]>
184  //      LINALG:   scf.yield %[[yielded]], %[[c0]], %[[c0]] :
185  // LINALG-SAME:     memref<?x8xf32, #[[$map_2d_stride_1]]>, index, index
186  //      LINALG: }
187  //      LINALG: %[[res:.*]] = vector.transfer_read {{.*}} {masked = [false, false]} :
188  // LINALG-SAME:   memref<?x8xf32, #[[$map_2d_stride_1]]>, vector<4x8xf32>
189  %1 = vector.transfer_read %A[%i, %j], %f0 :
190    memref<7x8xf32, offset:?, strides:[?, 1]>, vector<4x8xf32>
191
192  // CHECK: return %[[res]] : vector<4x8xf32>
193  return %1 : vector<4x8xf32>
194}
195