1// RUN: mlir-opt -linalg-bufferize -canonicalize -cse -split-input-file %s | FileCheck %s 2 3#map0 = affine_map<(d0) -> (d0)> 4 5// In-depth checking of a basic case, this is testing 6// - bufferization.to_memref / bufferization.to_tensor materializations are 7// properly inserted 8// - payload is correctly carried over 9// - affine maps are correctly carried over 10// Later tests will not check all these details. 11 12// CHECK: #map = affine_map<(d0) -> (d0)> 13// CHECK-LABEL: func @basic( 14// CHECK-SAME: %[[TENSOR:.*]]: tensor<4xf32>) -> tensor<4xf32> { 15// CHECK: %[[MEMREF:.*]] = bufferization.to_memref %[[TENSOR]] : memref<4xf32> 16// CHECK: %[[RESULT_MEMREF:.*]] = memref.alloc() : memref<4xf32> 17// CHECK: linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel"]} 18// CHECK-SAME: ins(%[[MEMREF]] : memref<4xf32>) 19// CHECK-SAME: outs(%[[RESULT_MEMREF]] : memref<4xf32>) { 20// CHECK: ^bb0(%[[RESULT1:.*]]: f32, %[[UNUSED:.*]]: f32): 21// CHECK: %[[DIM1:.*]] = math.exp %[[RESULT1]] : f32 22// CHECK: linalg.yield %[[DIM1]] : f32 23// CHECK: } 24// CHECK: %[[RESULT:.*]] = bufferization.to_tensor %[[RESULT_MEMREF]] : memref<4xf32> 25// CHECK: return %[[RESULT]] : tensor<4xf32> 26func @basic(%arg0: tensor<4xf32>) -> tensor<4xf32> { 27 %0 = linalg.generic { 28 indexing_maps = [#map0, #map0], 29 iterator_types = ["parallel"] 30 } ins(%arg0 : tensor<4xf32>) 31 outs(%arg0 : tensor<4xf32>) { 32 ^bb0(%gen_arg1: f32, %out: f32): 33 %tmp1 = math.exp %gen_arg1 : f32 34 linalg.yield %tmp1 : f32 35 } -> tensor<4xf32> 36 return %0 : tensor<4xf32> 37} 38 39 40// ----- 41 42#map0 = affine_map<(d0) -> (d0)> 43 44// Same as above but with linalg.init_tensor op. 45 46// CHECK: #map = affine_map<(d0) -> (d0)> 47// CHECK-LABEL: func @init_tensor( 48// CHECK-SAME: %[[IN:.*]]: tensor<?xf32>, %[[SIZE:.*]]: index) 49// CHECK: %[[MEMREF:.*]] = bufferization.to_memref %[[IN]] : memref<?xf32> 50// CHECK: %[[OUT_BUF:.*]] = memref.alloc(%[[SIZE]]) : memref<?xf32> 51// CHECK: linalg.generic 52// CHECK-SAME: ins(%[[MEMREF]] : memref<?xf32>) 53// CHECK-SAME: outs(%[[OUT_BUF]] : memref<?xf32>) { 54func @init_tensor(%in : tensor<?xf32>, %size: index) -> tensor<?xf32> { 55 %init = linalg.init_tensor [%size] : tensor<?xf32> 56 %0 = linalg.generic { 57 indexing_maps = [#map0, #map0], 58 iterator_types = ["parallel"] 59 } ins(%in : tensor<?xf32>) 60 outs(%init : tensor<?xf32>) { 61 ^bb0(%gen_arg1: f32, %out: f32): 62 %tmp1 = math.exp %gen_arg1 : f32 63 linalg.yield %tmp1 : f32 64 } -> tensor<?xf32> 65 return %0 : tensor<?xf32> 66} 67 68 69// ----- 70 71#map0 = affine_map<(d0) -> (d0)> 72 73// CHECK-LABEL: func @multiple_results 74// CHECK: %[[RESULT0:.*]] = memref.alloc() : memref<4xf32> 75// CHECK: %[[RESULT1:.*]] = memref.alloc() : memref<4xf32> 76// CHECK: linalg.generic 77// CHECK-SAME: ins(%{{.*}} : memref<4xf32>) 78// CHECK-SAME: outs(%[[RESULT0]], %[[RESULT1]] : memref<4xf32>, memref<4xf32>) 79// CHECK-NEXT: ^bb0(%{{.*}}: f32, %{{.*}}: f32, %{{.*}}: f32): 80func @multiple_results(%arg0: tensor<4xf32>) -> (tensor<4xf32>, tensor<4xf32>) { 81 %0, %1 = linalg.generic { 82 indexing_maps = [#map0, #map0, #map0], 83 iterator_types = ["parallel"] 84 } ins(%arg0 : tensor<4xf32>) 85 outs (%arg0, %arg0 : tensor<4xf32>, tensor<4xf32>) { 86 ^bb0(%gen_arg1: f32, %out1: f32, %out2: f32): 87 %tmp1 = math.exp %gen_arg1 : f32 88 linalg.yield %tmp1, %tmp1 : f32, f32 89 } -> (tensor<4xf32>, tensor<4xf32>) 90 return %0, %1 : tensor<4xf32>, tensor<4xf32> 91} 92 93// ----- 94 95#map_2d = affine_map<(d0, d1) -> (d0, d1)> 96 97// Check that the allocs properly consider the different shapes of the output 98// operands. The permuted indexing maps translate to different output shapes. 99 100// CHECK-LABEL: func @dynamic_results( 101// CHECK-SAME: %[[ARG:.*]]: tensor<?x?xf32> 102// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index 103// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index 104// CHECK: %[[MEMREF_ARG:.*]] = bufferization.to_memref %[[ARG]] : memref<?x?xf32> 105// CHECK: %[[DIM0:.*]] = tensor.dim %[[ARG]], %[[C0]] : tensor<?x?xf32> 106// CHECK: %[[DIM1:.*]] = tensor.dim %[[ARG]], %[[C1]] : tensor<?x?xf32> 107// CHECK: %[[RESULT0:.*]] = memref.alloc(%[[DIM0]], %[[DIM1]]) : memref<?x?xf32> 108// CHECK: %[[RESULT1:.*]] = memref.alloc(%[[DIM0]], %[[DIM1]]) : memref<?x?xf32> 109// CHECK: linalg.generic 110// CHECK-SAME: ins(%[[MEMREF_ARG]] : memref<?x?xf32>) 111// CHECK-SAME: outs(%[[RESULT0]], %[[RESULT1]] : memref<?x?xf32>, memref<?x?xf32>) 112func @dynamic_results(%arg0: tensor<?x?xf32>) 113 -> (tensor<?x?xf32>, tensor<?x?xf32>) { 114 %0, %1 = linalg.generic { 115 indexing_maps = [#map_2d, #map_2d, #map_2d], 116 iterator_types = ["parallel", "parallel"] 117 } ins(%arg0 : tensor<?x?xf32>) 118 outs (%arg0, %arg0 : tensor<?x?xf32>, tensor<?x?xf32>) { 119 ^bb0(%gen_arg1: f32, %out1: f32, %out2: f32): 120 %tmp1 = math.exp %gen_arg1 : f32 121 linalg.yield %tmp1, %tmp1 : f32, f32 122 } -> (tensor<?x?xf32>, tensor<?x?xf32>) 123 return %0, %1 : tensor<?x?xf32>, tensor<?x?xf32> 124} 125 126// ----- 127 128#accesses = [ 129 affine_map<(i, j, k) -> (j, i, k)>, 130 affine_map<(i, j, k) -> (i, j)> 131] 132 133#trait = { 134 indexing_maps = #accesses, 135 iterator_types = ["parallel", "parallel", "reduction"] 136} 137 138// Check the bufferization of init tensors. 139 140// CHECK-LABEL: func @generic_with_init_tensor( 141// CHECK-SAME: %[[ARG0_TENSOR:.*]]: tensor<2x3x4xvector<3x4xi4>>, 142// CHECK-SAME: %[[ARG1_TENSOR:.*]]: tensor<3x2xf32>) -> tensor<3x2xf32> { 143// CHECK-DAG: %[[ARG0_MEMREF:.*]] = bufferization.to_memref %[[ARG0_TENSOR]] : memref<2x3x4xvector<3x4xi4>> 144// CHECK-DAG: %[[ARG1_MEMREF:.*]] = bufferization.to_memref %[[ARG1_TENSOR]] : memref<3x2xf32> 145// CHECK: %[[INIT_BUFFER:.*]] = memref.alloc() : memref<3x2xf32> 146// CHECK: linalg.copy(%[[ARG1_MEMREF]], %[[INIT_BUFFER]]) : memref<3x2xf32>, memref<3x2xf32> 147// CHECK: linalg.generic 148// CHECK-SAME: ins(%[[ARG0_MEMREF]] : memref<2x3x4xvector<3x4xi4>>) 149// CHECK-SAME: outs(%[[INIT_BUFFER]] : memref<3x2xf32>) { 150func @generic_with_init_tensor(%arg0: tensor<2x3x4xvector<3x4xi4>>, 151 %arg1: tensor<3x2xf32>) -> (tensor<3x2xf32>) { 152 153 %0 = linalg.generic #trait 154 ins(%arg0 : tensor<2x3x4xvector<3x4xi4>>) 155 outs(%arg1 : tensor<3x2xf32>) { 156 ^bb(%v0: vector<3x4xi4>, %v1: f32) : 157 linalg.yield %v1 : f32 158 } -> tensor<3x2xf32> 159 160 return %0 : tensor<3x2xf32> 161} 162 163// ----- 164 165// CHECK-DAG: #[[$MAP0:[0-9a-z]*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)> 166// CHECK-DAG: #[[$MAP1:[0-9a-z]*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1 * 2)> 167 168func private @make_index() -> index 169 170// CHECK-LABEL: func @bufferize_slice( 171// CHECK-SAME: %[[T:[0-9a-z]*]]: tensor<?x?xf32> 172func @bufferize_slice(%t : tensor<?x?xf32>) -> (tensor<2x3xf32>, tensor<2x?xf32>) { 173 // CHECK: %[[M:.*]] = bufferization.to_memref %[[T]] : memref<?x?xf32> 174 175 // CHECK: %[[IDX:.*]] = call @make_index() : () -> index 176 %i0 = call @make_index() : () -> index 177 178 // CHECK-NEXT: %[[A0:.*]] = memref.alloc() : memref<2x3xf32> 179 // CHECK-NEXT: %[[SM0:.*]] = memref.subview %[[M]][0, 0] [2, 3] [1, 1] 180 // CHECK-SAME: memref<?x?xf32> to memref<2x3xf32, #[[$MAP0]]> 181 // CHECK-NEXT: linalg.copy(%[[SM0]], %[[A0]]) : memref<2x3xf32, #[[$MAP0]]>, memref<2x3xf32> 182 // CHECK-NEXT: %[[RT0:.*]] = bufferization.to_tensor %[[A0]] : memref<2x3xf32> 183 %st0 = tensor.extract_slice %t[0, 0][2, 3][1, 1] : tensor<?x?xf32> to tensor<2x3xf32> 184 185 // CHECK-NEXT: %[[A1:.*]] = memref.alloc(%[[IDX]]) : memref<2x?xf32> 186 // CHECK-NEXT: %[[SM1:.*]] = memref.subview %[[M]][0, %[[IDX]]] [2, %[[IDX]]] [1, 2] 187 // CHECK-SAME: memref<?x?xf32> to memref<2x?xf32, #[[$MAP1]]> 188 // CHECK-NEXT: linalg.copy(%[[SM1]], %[[A1]]) : memref<2x?xf32, #[[$MAP1]]>, memref<2x?xf32> 189 // CHECK-NEXT: %[[RT1:.*]] = bufferization.to_tensor %[[A1]] : memref<2x?xf32> 190 %st1 = tensor.extract_slice %t[0, %i0][2, %i0][1, 2] : tensor<?x?xf32> to tensor<2x?xf32> 191 192 // CHECK-NEXT: return %[[RT0]], %[[RT1]] 193 return %st0, %st1 : tensor<2x3xf32>, tensor<2x?xf32> 194} 195 196// ----- 197 198// CHECK-DAG: #[[$MAP0:[0-9a-z]*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1)> 199// CHECK-DAG: #[[$MAP1:[0-9a-z]*]] = affine_map<(d0, d1)[s0, s1] -> (d0 * s1 + s0 + d1 * 2)> 200 201func private @make_index() -> index 202 203// CHECK-LABEL: func @bufferize_insert_slice( 204// CHECK-SAME: %[[T:[0-9a-z]*]]: tensor<?x?xf32> 205// CHECK-SAME: %[[ST0:[0-9a-z]*]]: tensor<2x3xf32> 206// CHECK-SAME: %[[ST1:[0-9a-z]*]]: tensor<2x?xf32> 207func @bufferize_insert_slice(%t : tensor<?x?xf32>, %st0 : tensor<2x3xf32>, %st1 : tensor<2x?xf32>) -> 208 (tensor<?x?xf32>, tensor<?x?xf32>) { 209 // CHECK-DAG: %[[M:.*]] = bufferization.to_memref %[[T]] : memref<?x?xf32> 210 // CHECK-DAG: %[[SM0:.*]] = bufferization.to_memref %[[ST0]] : memref<2x3xf32> 211 // CHECK-DAG: %[[SM1:.*]] = bufferization.to_memref %[[ST1]] : memref<2x?xf32> 212 213 %c0 = arith.constant 0 : index 214 %c1 = arith.constant 1 : index 215 // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index 216 // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index 217 %i0 = call @make_index() : () -> index 218 // CHECK: %[[IDX:.*]] = call @make_index() : () -> index 219 220 221 // CHECK-NEXT: %[[DIM0:.*]] = tensor.dim %[[T]], %[[C0]] : tensor<?x?xf32> 222 // CHECK-NEXT: %[[DIM1:.*]] = tensor.dim %[[T]], %[[C1]] : tensor<?x?xf32> 223 // CHECK-NEXT: %[[M_COPY0:.*]] = memref.alloc(%[[DIM0]], %[[DIM1]]) : memref<?x?xf32> 224 // CHECK-NEXT: linalg.copy(%[[M]], %[[M_COPY0]]) : memref<?x?xf32>, memref<?x?xf32> 225 // CHECK-NEXT: %[[SUBVIEW0:.*]] = memref.subview %[[M_COPY0]][0, 0] [2, 3] [1, 1] 226 // CHECK-SAME: memref<?x?xf32> to memref<2x3xf32, #[[$MAP0]]> 227 // CHECK-NEXT: linalg.copy(%[[SM0]], %[[SUBVIEW0]]) : memref<2x3xf32>, memref<2x3xf32, #[[$MAP0]]> 228 // CHECK-NEXT: %[[RT0:.*]] = bufferization.to_tensor %[[M_COPY0]] : memref<?x?xf32> 229 %t0 = tensor.insert_slice %st0 into %t[0, 0][2, 3][1, 1] : tensor<2x3xf32> into tensor<?x?xf32> 230 231 // CHECK-NEXT: %[[M_COPY1:.*]] = memref.alloc(%[[DIM0]], %[[DIM1]]) : memref<?x?xf32> 232 // CHECK-NEXT: linalg.copy(%[[M]], %[[M_COPY1]]) : memref<?x?xf32>, memref<?x?xf32> 233 // CHECK-NEXT: %[[SUBVIEW1:.*]] = memref.subview %[[M_COPY1]][0, %[[IDX]]] [2, %[[IDX]]] [1, 2] 234 // CHECK-SAME: memref<?x?xf32> to memref<2x?xf32, #[[$MAP1]]> 235 // CHECK-NEXT: linalg.copy(%[[SM1]], %[[SUBVIEW1]]) : memref<2x?xf32>, memref<2x?xf32, #[[$MAP1]]> 236 // CHECK-NEXT: %[[RT1:.*]] = bufferization.to_tensor %[[M_COPY1]] : memref<?x?xf32> 237 %t1 = tensor.insert_slice %st1 into %t[0, %i0][2, %i0][1, 2] : tensor<2x?xf32> into tensor<?x?xf32> 238 239 // CHECK: return %[[RT0]], %[[RT1]] 240 return %t0, %t1: tensor<?x?xf32>, tensor<?x?xf32> 241} 242 243// ----- 244 245// CHECK-LABEL: func @bufferize_fill( 246// CHECK-SAME: %[[IN:.*]]: tensor<?xf32> 247func @bufferize_fill(%arg0: tensor<?xf32>) -> tensor<?xf32> { 248 %c0 = arith.constant 0.0 : f32 249 // CHECK: %[[MEMREF:.*]] = bufferization.to_memref %[[IN]] : memref<?xf32> 250 // CHECK: linalg.fill(%cst, %[[MEMREF]]) : f32, memref<?xf32> 251 // CHECK: %[[TENSOR:.*]] = bufferization.to_tensor %[[MEMREF]] : memref<?xf32> 252 // CHECK: return %[[TENSOR]] 253 %0 = linalg.fill(%c0, %arg0) : f32, tensor<?xf32> -> tensor<?xf32> 254 return %0 : tensor<?xf32> 255} 256 257// ----- 258 259// CHECK-LABEL: func @bufferize_tensor_collapse_shape( 260// CHECK-SAME: %[[IN:.*]]: tensor<4x5xf32> 261func @bufferize_tensor_collapse_shape(%arg0: tensor<4x5xf32>) -> tensor<20xf32> { 262 %out = tensor.collapse_shape %arg0 [[0, 1]] : 263 tensor<4x5xf32> into tensor<20xf32> 264 return %out : tensor<20xf32> 265} 266// CHECK: %[[MEMREF:.*]] = bufferization.to_memref %[[IN]] : memref<4x5xf32> 267// CHECK: %[[RESHAPE:.*]] = memref.collapse_shape %[[MEMREF]] {{\[}}[0, 1]] 268// CHECK-SAME: : memref<4x5xf32> into memref<20xf32> 269// CHECK: %[[TENSOR:.*]] = bufferization.to_tensor %[[RESHAPE]] : memref<20xf32> 270// CHECK: return %[[TENSOR]] 271 272// ----- 273 274// CHECK-LABEL: func @pad_tensor_dynamic_shape( 275// CHECK-SAME: %[[IN:.*]]: tensor<4x?x2x?xf32>, 276// CHECK-SAME: %[[OFFSET:.*]]: index) -> tensor<4x?x?x?xf32> { 277func @pad_tensor_dynamic_shape(%arg0: tensor<4x?x2x?xf32>, %arg1: index) -> tensor<4x?x?x?xf32> { 278 %c0 = arith.constant 0 : index 279 %cst = arith.constant 0.0 : f32 280 %out = linalg.pad_tensor %arg0 low[%c0, %c0, %arg1, %c0] high[%c0, %c0, %c0, %arg1] { 281 ^bb0(%gen_arg1: index, %gen_arg2: index, %gen_arg3: index, %gen_arg4: index): 282 linalg.yield %cst : f32 283 } : tensor<4x?x2x?xf32> to tensor<4x?x?x?xf32> 284 return %out : tensor<4x?x?x?xf32> 285} 286 287// CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index 288// CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index 289// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index 290// CHECK-DAG: %[[CST:.*]] = arith.constant 0.000000e+00 : f32 291// CHECK: %[[IN_MEMREF:.*]] = bufferization.to_memref %[[IN]] : memref<4x?x2x?xf32> 292// CHECK: %[[DIM1:.*]] = tensor.dim %[[IN]], %[[C1]] : tensor<4x?x2x?xf32> 293// CHECK: %[[OUT_DIM2:.*]] = arith.addi %[[OFFSET]], %[[C2]] : index 294// CHECK: %[[DIM3:.*]] = tensor.dim %[[IN]], %[[C3]] : tensor<4x?x2x?xf32> 295// CHECK: %[[OUT_DIM3:.*]] = arith.addi %[[DIM3]], %[[OFFSET]] : index 296// CHECK: %[[FILLED:.*]] = memref.alloc(%[[DIM1]], %[[OUT_DIM2]], %[[OUT_DIM3]]) : memref<4x?x?x?xf32> 297// CHECK: linalg.fill(%[[CST]], %[[FILLED]]) : f32, memref<4x?x?x?xf32> 298// CHECK: %[[OUT:.*]] = memref.alloc(%[[DIM1]], %[[OUT_DIM2]], %[[OUT_DIM3]]) : memref<4x?x?x?xf32> 299// CHECK: linalg.copy(%[[FILLED]], %[[OUT]]) : memref<4x?x?x?xf32>, memref<4x?x?x?xf32> 300// CHECK: %[[INTERIOR:.*]] = memref.subview %[[OUT]][0, 0, %[[OFFSET]], 0] [4, %[[DIM1]], 2, %[[DIM3]]] [1, 1, 1, 1] : memref<4x?x?x?xf32> to memref<4x?x2x?xf32, #map> 301// CHECK: linalg.copy(%[[IN_MEMREF]], %[[INTERIOR]]) : memref<4x?x2x?xf32>, memref<4x?x2x?xf32, #map> 302// CHECK: %[[OUT_TENSOR:.*]] = bufferization.to_tensor %[[OUT]] : memref<4x?x?x?xf32> 303// CHECK: return %[[OUT_TENSOR]] : tensor<4x?x?x?xf32> 304// CHECK: } 305 306 307// ----- 308 309// CHECK-LABEL: func @vector_transfer 310func @vector_transfer(%in: tensor<4xf32>, %out: tensor<4xf32>) { 311 %c0 = arith.constant 0 : index 312 %cst = arith.constant 0.000000e+00 : f32 313 %read = vector.transfer_read %in[%c0], %cst {in_bounds = [true]} 314 : tensor<4xf32>, vector<4xf32> 315 %tanh = math.tanh %read : vector<4xf32> 316 %write = vector.transfer_write %tanh, %out[%c0] {in_bounds = [true]} 317 : vector<4xf32>, tensor<4xf32> 318 return 319 // CHECK: vector.transfer_read {{.*}} : memref<4xf32>, vector<4xf32> 320 // CHECK: vector.transfer_write {{.*}} : vector<4xf32>, memref<4xf32> 321} 322 323// ----- 324 325// CHECK-LABEL: func @bufferize_dot 326func @bufferize_dot(%in: tensor<4xf32>, %out: tensor<f32>) -> tensor<f32> { 327 %dot = linalg.dot ins(%in, %in : tensor<4xf32>, tensor<4xf32>) 328 outs(%out : tensor<f32>) -> tensor<f32> 329 return %dot : tensor<f32> 330 // CHECK: linalg.dot ins(%{{.*}}, %{{.*}} : memref<4xf32>, memref<4xf32>) 331 // CHECK-SAME: outs(%[[OUT:.*]] : memref<f32>) 332 // CHECK: %[[OUT_TENSOR:.*]] = bufferization.to_tensor %[[OUT]] : memref<f32> 333 // CHECK: return %[[OUT_TENSOR]] 334} 335