1// RUN: mlir-opt %s -allow-unregistered-dialect -test-linalg-transform-patterns=test-vector-transfer-forwarding-patterns | FileCheck %s
2
3// CHECK-LABEL: testAllocRead
4//  CHECK-SAME: %[[ARG0:[0-9a-zA-Z]*]]: memref
5//   CHECK-NOT: linalg.fill
6//   CHECK-NOT: linalg.copy
7//       CHECK: %[[ALLOC:.*]] = alloc
8//       CHECK: vector.transfer_read %[[ARG0]]
9//   CHECK-NOT: masked
10func @testAllocRead(%in: memref<? x f32>) -> vector<32 x f32> {
11  %c0 = constant 0: index
12  %f0 = constant 0.0: f32
13  %alloc = alloc() : memref<32 x f32>
14  %subview = subview %alloc[0][16][1] : memref<32 x f32> to memref<16 x f32>
15  linalg.copy(%in, %subview): memref<? x f32>, memref<16 x f32>
16  %0 = vector.transfer_read %alloc[%c0], %f0 {masked = [false]} : memref<32 x f32>, vector<32 x f32>
17  dealloc %alloc : memref<32 x f32>
18  return %0: vector<32 x f32>
19}
20
21// CHECK-LABEL: testAllocFillRead
22//  CHECK-SAME: %[[ARG0:[0-9a-zA-Z]*]]: memref
23//   CHECK-NOT: linalg.fill
24//   CHECK-NOT: linalg.copy
25//       CHECK: %[[ALLOC:.*]] = alloc
26//       CHECK: vector.transfer_read %[[ARG0]]
27//   CHECK-NOT: masked
28func @testAllocFillRead(%in: memref<? x f32>) -> vector<32 x f32> {
29  %c0 = constant 0: index
30  %f0 = constant 0.0: f32
31  %alloc = alloc() : memref<32 x f32>
32  linalg.fill(%alloc, %f0): memref<32 x f32>, f32
33  %subview = subview %alloc[0][16][1] : memref<32 x f32> to memref<16 x f32>
34  linalg.copy(%in, %subview): memref<? x f32>, memref<16 x f32>
35  %0 = vector.transfer_read %alloc[%c0], %f0 {masked = [false]} : memref<32 x f32>, vector<32 x f32>
36  dealloc %alloc : memref<32 x f32>
37  return %0: vector<32 x f32>
38}
39
40// CHECK-LABEL: testViewRead
41//  CHECK-SAME: %[[ARG0:[0-9a-zA-Z]*]]: memref
42//   CHECK-NOT: linalg.fill
43//   CHECK-NOT: linalg.copy
44//       CHECK: %[[ALLOC:.*]] = alloc
45//       CHECK: vector.transfer_read %[[ARG0]]
46//   CHECK-NOT: masked
47func @testViewRead(%in: memref<? x f32>) -> vector<32 x f32> {
48  %c0 = constant 0: index
49  %f0 = constant 0.0: f32
50  %alloc = alloc() : memref<128 x i8>
51  %view = view %alloc[%c0][] : memref<128 x i8> to memref<32 x f32>
52  %subview = subview %view[0][16][1] : memref<32 x f32> to memref<16 x f32>
53  linalg.copy(%in, %subview): memref<? x f32>, memref<16 x f32>
54  %0 = vector.transfer_read %view[%c0], %f0 {masked = [false]} : memref<32 x f32>, vector<32 x f32>
55  dealloc %alloc : memref<128 x i8>
56  return %0: vector<32 x f32>
57}
58
59// CHECK-LABEL: testViewFillRead
60//  CHECK-SAME: %[[ARG0:[0-9a-zA-Z]*]]: memref
61//   CHECK-NOT: linalg.fill
62//   CHECK-NOT: linalg.copy
63//       CHECK: %[[ALLOC:.*]] = alloc
64//       CHECK: vector.transfer_read %[[ARG0]]
65//   CHECK-NOT: masked
66func @testViewFillRead(%in: memref<? x f32>) -> vector<32 x f32> {
67  %c0 = constant 0: index
68  %f0 = constant 0.0: f32
69  %alloc = alloc() : memref<128 x i8>
70  %view = view %alloc[%c0][] : memref<128 x i8> to memref<32 x f32>
71  %subview = subview %view[0][16][1] : memref<32 x f32> to memref<16 x f32>
72  linalg.fill(%view, %f0): memref<32 x f32>, f32
73  linalg.copy(%in, %subview): memref<? x f32>, memref<16 x f32>
74  %0 = vector.transfer_read %view[%c0], %f0 {masked = [false]} : memref<32 x f32>, vector<32 x f32>
75  dealloc %alloc : memref<128 x i8>
76  return %0: vector<32 x f32>
77}
78
79// CHECK-LABEL: testAllocWrite
80//  CHECK-SAME: %[[ARG0:[0-9a-zA-Z]*]]: vector
81//  CHECK-SAME: %[[ARG1:[0-9a-zA-Z]*]]: memref
82//   CHECK-NOT: linalg.copy
83//       CHECK: %[[ALLOC:.*]] = alloc
84//       CHECK: vector.transfer_write %[[ARG0]], %[[ARG1]]
85//   CHECK-NOT: masked
86func @testAllocWrite(%vec: vector<32 x f32>, %out: memref<? x f32>) {
87  %c0 = constant 0: index
88  %f0 = constant 0.0: f32
89  %alloc = alloc() : memref<32 x f32>
90  %subview = subview %alloc[0][16][1] : memref<32 x f32> to memref<16 x f32>
91  vector.transfer_write %vec, %alloc[%c0] {masked = [false]} : vector<32 x f32>, memref<32 x f32>
92  linalg.copy(%subview, %out): memref<16 x f32>, memref<? x f32>
93  dealloc %alloc : memref<32 x f32>
94  return
95}
96
97// CHECK-LABEL: testViewWrite
98//  CHECK-SAME: %[[ARG0:[0-9a-zA-Z]*]]: vector
99//  CHECK-SAME: %[[ARG1:[0-9a-zA-Z]*]]: memref
100//   CHECK-NOT: linalg.copy
101//       CHECK: %[[ALLOC:.*]] = alloc
102//       CHECK: vector.transfer_write %[[ARG0]], %[[ARG1]]
103//   CHECK-NOT: masked
104func @testViewWrite(%vec: vector<32 x f32>, %out: memref<? x f32>) {
105  %c0 = constant 0: index
106  %f0 = constant 0.0: f32
107  %alloc = alloc() : memref<128 x i8>
108  %view = view %alloc[%c0][] : memref<128 x i8> to memref<32 x f32>
109  %subview = subview %view[0][16][1] : memref<32 x f32> to memref<16 x f32>
110  vector.transfer_write %vec, %view[%c0] {masked = [false]} : vector<32 x f32>, memref<32 x f32>
111  linalg.copy(%subview, %out): memref<16 x f32>, memref<? x f32>
112  dealloc %alloc : memref<128 x i8>
113  return
114}
115
116///===--------------------------------------------------------------------===///
117// Negative tests
118///===--------------------------------------------------------------------===///
119
120// This should fail the rewrite due to mismatching fill and transfer read value.
121// CHECK-LABEL: failAllocFillRead
122//  CHECK-SAME: %[[ARG0:[0-9a-zA-Z]*]]: memref
123//   CHECK-NOT: vector.transfer_read %[[ARG0]]
124//       CHECK: %[[ALLOC:.*]] = alloc
125//       CHECK: linalg.copy
126//       CHECK: vector.transfer_read %[[ALLOC]]
127func @failAllocFillRead(%in: memref<? x f32>) -> vector<32 x f32> {
128  %c0 = constant 0: index
129  %f0 = constant 0.0: f32
130  %f1 = constant 1.0: f32
131  %alloc = alloc() : memref<32 x f32>
132  linalg.fill(%alloc, %f0): memref<32 x f32>, f32
133  %subview = subview %alloc[0][16][1] : memref<32 x f32> to memref<16 x f32>
134  linalg.copy(%in, %subview): memref<? x f32>, memref<16 x f32>
135  "some_interleaved_use"(%subview) : (memref<16 x f32>) -> ()
136  %0 = vector.transfer_read %alloc[%c0], %f1: memref<32 x f32>, vector<32 x f32>
137  dealloc %alloc : memref<32 x f32>
138  return %0: vector<32 x f32>
139}
140
141// This should fail the rewrite due to some interleaved use.
142// CHECK-LABEL: failAllocWrite
143//  CHECK-SAME: %[[ARG0:[0-9a-zA-Z]*]]: vector
144//  CHECK-SAME: %[[ARG1:[0-9a-zA-Z]*]]: memref
145//   CHECK-NOT: vector.transfer_write %[[ARG0]], %[[ARG1]]
146//       CHECK: %[[ALLOC:.*]] = alloc
147//       CHECK: vector.transfer_write %[[ARG0]], %[[ALLOC]]
148//       CHECK: linalg.copy
149func @failAllocWrite(%vec: vector<32 x f32>, %out: memref<? x f32>) {
150  %c0 = constant 0: index
151  %f0 = constant 0.0: f32
152  %alloc = alloc() : memref<32 x f32>
153  %subview = subview %alloc[0][16][1] : memref<32 x f32> to memref<16 x f32>
154  vector.transfer_write %vec, %alloc[%c0] : vector<32 x f32>, memref<32 x f32>
155  "some_interleaved_use"(%subview) : (memref<16 x f32>) -> ()
156  linalg.copy(%subview, %out): memref<16 x f32>, memref<? x f32>
157  dealloc %alloc : memref<32 x f32>
158  return
159}
160