1// RUN: mlir-opt %s -test-linalg-transform-patterns=test-linalg-to-vector-patterns -split-input-file | FileCheck %s 2 3// ----- 4 5// CHECK-LABEL: contraction_dot 6func @contraction_dot(%A: memref<1584xf32>, %B: memref<1584xf32>, %C: memref<f32>) { 7 8// CHECK: arith.mulf %{{.*}}, %{{.*}} : vector<1584xf32> 9// CHECK: vector.multi_reduction <add>, %{{.*}} [0] : vector<1584xf32> to f32 10// CHECK: arith.addf %{{.*}}, %{{.*}} : f32 11 linalg.dot ins(%A, %B: memref<1584xf32>, memref<1584xf32>) 12 outs(%C: memref<f32>) 13 return 14} 15 16// ----- 17 18// CHECK-LABEL: contraction_matvec 19func @contraction_matvec(%A: memref<1584x1584xf32>, %B: memref<1584xf32>, %C: memref<1584xf32>) { 20 21// CHECK: arith.mulf %{{.*}}, %{{.*}} : vector<1584x1584xf32> 22// CHECK: vector.multi_reduction <add>, %{{.*}} [1] : vector<1584x1584xf32> to vector<1584xf32> 23// CHECK: arith.addf %{{.*}}, %{{.*}} : vector<1584xf32> 24 linalg.matvec ins(%A, %B: memref<1584x1584xf32>, memref<1584xf32>) 25 outs(%C: memref<1584xf32>) 26 return 27} 28 29// ----- 30 31// CHECK-LABEL: contraction_matmul 32func @contraction_matmul(%A: memref<1584x1584xf32>, %B: memref<1584x1584xf32>, %C: memref<1584x1584xf32>) { 33// CHECK: arith.mulf %{{.*}}, %{{.*}} : vector<1584x1584x1584xf32> 34// CHECK: vector.multi_reduction <add>, %{{.*}} [2] : vector<1584x1584x1584xf32> to vector<1584x1584xf32> 35// CHECK: arith.addf %{{.*}}, %{{.*}} : vector<1584x1584xf32> 36 linalg.matmul ins(%A, %B: memref<1584x1584xf32>, memref<1584x1584xf32>) 37 outs(%C: memref<1584x1584xf32>) 38 return 39} 40 41// ----- 42 43// CHECK-LABEL: contraction_batch_matmul 44func @contraction_batch_matmul(%A: memref<1584x1584x1584xf32>, %B: memref<1584x1584x1584xf32>, %C: memref<1584x1584x1584xf32>) { 45// CHECK: arith.mulf %{{.*}}, %{{.*}} : vector<1584x1584x1584x1584xf32> 46// CHECK: vector.multi_reduction <add>, %{{.*}} [3] : vector<1584x1584x1584x1584xf32> to vector<1584x1584x1584xf32> 47// CHECK: arith.addf %{{.*}}, %{{.*}} : vector<1584x1584x1584xf32> 48 linalg.batch_matmul 49 ins(%A, %B: memref<1584x1584x1584xf32>, memref<1584x1584x1584xf32>) 50 outs(%C: memref<1584x1584x1584xf32>) 51 return 52} 53 54// ----- 55 56#matmul_trait = { 57 args_in = 2, 58 args_out = 1, 59 indexing_maps = [ 60 affine_map<(m, n, k) -> (m, k)>, 61 affine_map<(m, n, k) -> (k, n)>, 62 affine_map<(m, n, k) -> (m, n)> 63 ], 64 iterator_types = ["parallel", "parallel", "reduction"] 65} 66 67// CHECK-LABEL: func @vectorization_test 68func @vectorization_test(%A: memref<8x16xf32>, %B: memref<16x32xf32>, 69 %C: memref<8x32xf32>) { 70 // CHECK: vector.transfer_read %{{.*}} : memref<8x16xf32>, vector<8x32x16xf32> 71 // CHECK: vector.transfer_read %{{.*}} : memref<16x32xf32>, vector<8x32x16xf32> 72 // CHECK: vector.transfer_read %{{.*}} : memref<8x32xf32>, vector<8x32xf32> 73 // CHECK: %[[MUL:.*]] = arith.mulf %{{.*}}, %{{.*}} : vector<8x32x16xf32> 74 // CHECK: %[[R:.*]] = vector.multi_reduction <add>, %[[MUL]] [2] : vector<8x32x16xf32> to vector<8x32xf32> 75 // CHECK: arith.addf %[[R]], %{{.*}} : vector<8x32xf32> 76 // CHECK: vector.transfer_write %{{.*}}, %{{.*}} : vector<8x32xf32>, memref<8x32xf32> 77 linalg.generic #matmul_trait 78 ins(%A, %B : memref<8x16xf32>, memref<16x32xf32>) 79 outs(%C : memref<8x32xf32>) { 80 ^bb(%a: f32, %b: f32, %c: f32) : 81 %d = arith.mulf %a, %b: f32 82 %e = arith.addf %c, %d: f32 83 linalg.yield %e : f32 84 } 85 return 86} 87 88// ----- 89 90#matmul_transpose_out_trait = { 91 args_in = 2, 92 args_out = 1, 93 indexing_maps = [ 94 affine_map<(m, n, k) -> (m, k)>, 95 affine_map<(m, n, k) -> (k, n)>, 96 affine_map<(m, n, k) -> (n, m)> 97 ], 98 iterator_types = ["parallel", "parallel", "reduction"] 99} 100 101// CHECK-LABEL: func @generic_output_transpose 102func @generic_output_transpose(%A: memref<8x16xf32>, %B: memref<16x32xf32>, 103 %C: memref<32x8xf32>) { 104 // CHECK: vector.transfer_read %{{.*}} : memref<8x16xf32>, vector<8x32x16xf32> 105 // CHECK: vector.transfer_read %{{.*}} : memref<16x32xf32>, vector<8x32x16xf32> 106 // CHECK: vector.transfer_read %{{.*}} : memref<32x8xf32>, vector<8x32xf32> 107 // CHECK: %[[MUL:.*]] = arith.mulf %{{.*}}, %{{.*}} : vector<8x32x16xf32> 108 // CHECK: %[[R:.*]] = vector.multi_reduction <add>, %[[MUL]] [2] : vector<8x32x16xf32> to vector<8x32xf32> 109 // CHECK: arith.addf %[[R]], %{{.*}} : vector<8x32xf32> 110 // CHECK: vector.transfer_write %{{.*}}, %{{.*}} : vector<8x32xf32>, memref<32x8xf32> 111 linalg.generic #matmul_transpose_out_trait 112 ins(%A, %B : memref<8x16xf32>, memref<16x32xf32>) 113 outs(%C : memref<32x8xf32>) { 114 ^bb(%a: f32, %b: f32, %c: f32) : 115 %d = arith.mulf %a, %b: f32 116 %e = arith.addf %c, %d: f32 117 linalg.yield %e : f32 118 } 119 return 120} 121 122// ----- 123 124#matmul_trait = { 125 args_in = 2, 126 args_out = 1, 127 indexing_maps = [ 128 affine_map<(m, n, k) -> (m, k)>, 129 affine_map<(m, n, k) -> (k, n)>, 130 affine_map<(m, n, k) -> (m, n)> 131 ], 132 iterator_types = ["parallel", "parallel", "reduction"] 133} 134 135// CHECK-LABEL: func @vectorization_test_integer 136func @vectorization_test_integer(%A: memref<8x16xi32>, %B: memref<16x32xi32>, 137 %C: memref<8x32xi32>) { 138 // CHECK: vector.transfer_read %{{.*}} : memref<8x16xi32>, vector<8x32x16xi32> 139 // CHECK: vector.transfer_read %{{.*}} : memref<16x32xi32>, vector<8x32x16xi32> 140 // CHECK: vector.transfer_read %{{.*}} : memref<8x32xi32>, vector<8x32xi32> 141 // CHECK: %[[MUL:.*]] = arith.muli %{{.*}}, %{{.*}} : vector<8x32x16xi32> 142 // CHECK: %[[R:.*]] = vector.multi_reduction <add>, %[[MUL]] [2] : vector<8x32x16xi32> to vector<8x32xi32> 143 // CHECK: arith.addi %[[R]], %{{.*}} : vector<8x32xi32> 144 145 // CHECK: vector.transfer_write %{{.*}}, %{{.*}} : vector<8x32xi32>, memref<8x32xi32> 146 linalg.generic #matmul_trait 147 ins(%A, %B : memref<8x16xi32>, memref<16x32xi32>) 148 outs(%C : memref<8x32xi32>) { 149 ^bb(%a: i32, %b: i32, %c: i32) : 150 %d = arith.muli %a, %b: i32 151 %e = arith.addi %c, %d: i32 152 linalg.yield %e : i32 153 } 154 return 155} 156 157// ----- 158 159// CHECK-LABEL: func @vectorization_test_2 160func @vectorization_test_2(%A: memref<8x16xf32>, %B: memref<16x32xf32>, 161 %C: memref<8x32xf32>) { 162 // CHECK: arith.mulf %{{.*}}, %{{.*}} : vector<8x32x16xf32> 163 // CHECK: vector.multi_reduction <add>, %{{.*}} [2] : vector<8x32x16xf32> to vector<8x32xf32> 164 // CHECK: arith.addf %{{.*}}, %{{.*}} : vector<8x32xf32> 165 linalg.matmul 166 ins(%A, %B: memref<8x16xf32>, memref<16x32xf32>) 167 outs(%C: memref<8x32xf32>) 168 return 169} 170 171// ----- 172 173// CHECK-LABEL: func @test_vectorize_scalar_input 174func @test_vectorize_scalar_input(%A : memref<8x16xf32>, %arg0 : f32) { 175 // CHECK: %[[V:.*]] = vector.broadcast {{.*}} : f32 to vector<8x16xf32> 176 // CHECK: vector.transfer_write %[[V]], {{.*}} : vector<8x16xf32>, memref<8x16xf32> 177 linalg.generic { 178 indexing_maps = [affine_map<(m, n) -> ()>, affine_map<(m, n) -> (m, n)>], 179 iterator_types = ["parallel", "parallel"]} 180 ins(%arg0 : f32) 181 outs(%A: memref<8x16xf32>) { 182 ^bb(%0: f32, %1: f32) : 183 linalg.yield %0 : f32 184 } 185 return 186} 187 188// ----- 189 190// CHECK-LABEL: func @test_vectorize_fill 191func @test_vectorize_fill(%A : memref<8x16xf32>, %arg0 : f32) { 192 // CHECK: %[[V:.*]] = vector.broadcast {{.*}} : f32 to vector<8x16xf32> 193 // CHECK: vector.transfer_write %[[V]], {{.*}} : vector<8x16xf32>, memref<8x16xf32> 194 linalg.fill ins(%arg0 : f32) outs(%A : memref<8x16xf32>) 195 return 196} 197 198// ----- 199 200// CHECK-LABEL: func @test_vectorize_fill 201func @test_vectorize_fill_scalar(%A : memref<f32>, %arg0 : f32) { 202 // CHECK-SAME: (%[[M:.*]]: memref<f32>, %[[val:.*]]: f32) 203 // CHECK: %[[VEC:.*]] = vector.broadcast %[[val]] : f32 to vector<f32> 204 // CHECK: vector.transfer_write %[[VEC]], %[[M]][] : vector<f32>, memref<f32> 205 linalg.fill ins(%arg0 : f32) outs(%A : memref<f32>) 206 return 207} 208 209// ----- 210 211// CHECK-LABEL: func @test_vectorize_copy 212func @test_vectorize_copy(%A : memref<8x16xf32>, %B : memref<8x16xf32>) { 213 // CHECK: %[[V:.*]] = vector.transfer_read {{.*}} : memref<8x16xf32>, vector<8x16xf32> 214 // CHECK: vector.transfer_write %[[V]], {{.*}} : vector<8x16xf32>, memref<8x16xf32> 215 memref.copy %A, %B : memref<8x16xf32> to memref<8x16xf32> 216 return 217} 218 219// ----- 220 221// CHECK-LABEL: func @test_vectorize_copy_scalar 222func @test_vectorize_copy_scalar(%A : memref<f32>, %B : memref<f32>) { 223 // CHECK-SAME: (%[[A:.*]]: memref<f32>, %[[B:.*]]: memref<f32>) 224 // CHECK: %[[V:.*]] = vector.transfer_read %[[A]][]{{.*}} : memref<f32>, vector<f32> 225 // CHECK: %[[val:.*]] = vector.extractelement %[[V]][] : vector<f32> 226 // CHECK: %[[VV:.*]] = vector.broadcast %[[val]] : f32 to vector<f32> 227 // CHECK: vector.transfer_write %[[VV]], %[[B]][] : vector<f32>, memref<f32> 228 memref.copy %A, %B : memref<f32> to memref<f32> 229 return 230} 231 232// ----- 233 234// CHECK-LABEL: func @test_vectorize_trailing_index 235 // CHECK-SAME: (%[[ARG0:.*]]: memref<1x2x4x8xindex>) 236func @test_vectorize_trailing_index(%arg0: memref<1x2x4x8xindex>) { 237 // CHECK-DAG: %[[CST0:.*]] = arith.constant dense<[0, 1, 2, 3, 4, 5, 6, 7]> : vector<8xindex> 238 // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index 239 linalg.generic { 240 indexing_maps = [ 241 affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>], 242 iterator_types = ["parallel", "parallel", "parallel", "parallel"]} 243 outs(%arg0: memref<1x2x4x8xindex>) { 244 ^bb0(%arg1: index): 245 // CHECK: %[[BCST:.*]] = vector.broadcast %[[CST0]] : vector<8xindex> to vector<1x2x4x8xindex> 246 // CHECK: vector.transfer_write %[[BCST]], %[[ARG0]][%[[C0]], %[[C0]], %[[C0]], %[[C0]]] {{.*}} : vector<1x2x4x8xindex>, memref<1x2x4x8xindex> 247 %0 = linalg.index 3 : index 248 linalg.yield %0 : index 249 } 250 return 251} 252 253// ----- 254 255// CHECK-LABEL: func @test_vectorize_inner_index 256 // CHECK-SAME: (%[[ARG0:.*]]: memref<1x2x4x8xindex>) 257func @test_vectorize_inner_index(%arg0: memref<1x2x4x8xindex>) { 258 // CHECK-DAG: %[[CST0:.*]] = arith.constant dense<[0, 1]> : vector<2xindex> 259 // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index 260 linalg.generic { 261 indexing_maps = [ 262 affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>], 263 iterator_types = ["parallel", "parallel", "parallel", "parallel"]} 264 outs(%arg0: memref<1x2x4x8xindex>) { 265 ^bb0(%arg1: index): 266 // CHECK: %[[BCST:.*]] = vector.broadcast %[[CST0]] : vector<2xindex> to vector<1x8x4x2xindex> 267 // CHECK: %[[TRAN:.*]] = vector.transpose %[[BCST]], [0, 3, 2, 1] : vector<1x8x4x2xindex> to vector<1x2x4x8xindex> 268 // CHECK: vector.transfer_write %[[TRAN]], %[[ARG0]][%[[C0]], %[[C0]], %[[C0]], %[[C0]]] {{.*}} : vector<1x2x4x8xindex>, memref<1x2x4x8xindex> 269 %0 = linalg.index 1 : index 270 linalg.yield %0 : index 271 } 272 return 273} 274 275// ----- 276 277// CHECK-LABEL: func @generic_vectorize 278 // CHECK-SAME: (%[[ARG0:.*]]: memref<4x256xf32>, %[[ARG1:.*]]: memref<4x256xf32>, 279 // CHECK-SAME: %[[ARG2:.*]]: memref<256xf32>, %[[ARG3:.*]]: f32) 280func @generic_vectorize(%arg0: memref<4x256xf32>, 281 %arg1: memref<4x256xf32>, 282 %arg2: memref<256xf32>, %i: f32) { 283 // CHECK-DAG: %[[CST0:.*]] = arith.constant dense<2.000000e+00> : vector<4x256xf32> 284 // CHECK-DAG: %[[CST1:.*]] = arith.constant dense<1.000000e+00> : vector<4x256xf32> 285 // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index 286 %c1_f32 = arith.constant 1.0 : f32 287 linalg.generic { 288 args_in = 0 : i64, 289 args_out = 10 : i64, 290 indexing_maps = [ 291 affine_map<(d0, d1) -> (d0, d1)>, 292 affine_map<(d0, d1) -> (d1)>, 293 affine_map<(d0, d1) -> (d0, d1)>, 294 affine_map<(d0, d1) -> (d0, d1)>, 295 affine_map<(d0, d1) -> (d0, d1)>, 296 affine_map<(d0, d1) -> (d0, d1)>, 297 affine_map<(d0, d1) -> (d0, d1)>, 298 affine_map<(d0, d1) -> (d0, d1)>, 299 affine_map<(d0, d1) -> (d0, d1)>, 300 affine_map<(d0, d1) -> (d0, d1)>, 301 affine_map<(d0, d1) -> (d0, d1)>, 302 affine_map<(d0, d1) -> (d0, d1)>], 303 iterator_types = ["parallel", "parallel"]} 304 ins(%arg1, %arg2: memref<4x256xf32>, memref<256xf32>) 305 outs( 306 %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0 : 307 memref<4x256xf32>, memref<4x256xf32>, memref<4x256xf32>, memref<4x256xf32>, 308 memref<4x256xf32>, memref<4x256xf32>, memref<4x256xf32>, memref<4x256xf32>, 309 memref<4x256xf32>, memref<4x256xf32>) { 310 ^bb0(%arg3 : f32, %arg4 : f32, %arg5: f32, %arg6: f32, %arg7: f32, %arg8: f32, 311 // CHECK: %[[V2:.*]] = vector.transfer_read %[[ARG1]][%[[C0]], %[[C0]]], {{.*}} : memref<4x256xf32>, vector<4x256xf32> 312 // CHECK: %[[V0:.*]] = vector.transfer_read %[[ARG2]][%[[C0]]], {{.*}} : memref<256xf32>, vector<4x256xf32> 313 // CHECK: %[[V3:.*]] = vector.transfer_read %[[ARG0]][%[[C0]], %[[C0]]], {{.*}} : memref<4x256xf32>, vector<4x256xf32> 314 // CHECK: %[[V1:.*]] = vector.transfer_read %[[ARG0]][%[[C0]], %[[C0]]], {{.*}} : memref<4x256xf32>, vector<4x256xf32> 315 %arg9 : f32, %arg10 : f32, %arg11 : f32, %arg12 : f32, %arg13 : f32, 316 %arg14 : f32): 317 // CHECK: %[[ADD:.*]] = arith.addf %[[V0]], %[[V1]] : vector<4x256xf32> 318 %6 = arith.addf %arg4, %arg6 : f32 319 // CHECK: %[[CMP:.*]] = arith.cmpf ogt, %[[V2]], %[[V1]] : vector<4x256xf32> 320 %7 = arith.cmpf ogt, %arg3, %arg6 : f32 321 // CHECK: %[[ARG3B:.*]] = vector.broadcast %[[ARG3]] : f32 to vector<4x256xf32> 322 %8 = arith.constant 2.0 : f32 323 // CHECK: %[[DIV:.*]] = arith.divf %[[V3]], %[[ARG3B]] : vector<4x256xf32> 324 %9 = arith.divf %arg5, %i : f32 325 // CHECK: %[[EXP:.*]] = math.exp2 %[[V3]] : vector<4x256xf32> 326 %10 = math.exp2 %arg5 : f32 327 // CHECK: %[[MUL:.*]] = arith.mulf %[[V3]], %[[CST0]] : vector<4x256xf32> 328 %11 = arith.mulf %arg5, %8 : f32 329 // CHECK: %[[RSQRT:.*]] = math.rsqrt %[[V3]] : vector<4x256xf32> 330 %12 = math.rsqrt %arg5 : f32 331 // CHECK: %[[SEL:.*]] = arith.select %[[CMP]], %[[V3]], %[[V1]] : vector<4x256xi1>, vector<4x256xf32> 332 %13 = arith.select %7, %arg5, %arg6 : f32 333 // CHECK: %[[SUB:.*]] = arith.subf %[[V3]], %[[V0]] : vector<4x256xf32> 334 %14 = arith.subf %arg5, %arg4 : f32 335 // CHECK: %[[TAN:.*]] = math.tanh %[[V3]] : vector<4x256xf32> 336 %15 = math.tanh %arg5 : f32 337 // CHECK: vector.transfer_write %[[ADD]], %[[ARG0]][%[[C0]], %[[C0]]] {{.*}} : vector<4x256xf32>, memref<4x256xf32> 338 // CHECK: vector.transfer_write %[[CST0]], %[[ARG0]][%[[C0]], %[[C0]]] {{.*}} : vector<4x256xf32>, memref<4x256xf32> 339 // CHECK: vector.transfer_write %[[CST1]], %[[ARG0]][%[[C0]], %[[C0]]] {{.*}} : vector<4x256xf32>, memref<4x256xf32> 340 // CHECK: vector.transfer_write %[[DIV]], %[[ARG0]][%[[C0]], %[[C0]]] {{.*}} : vector<4x256xf32>, memref<4x256xf32> 341 // CHECK: vector.transfer_write %[[EXP]], %[[ARG0]][%[[C0]], %[[C0]]] {{.*}} : vector<4x256xf32>, memref<4x256xf32> 342 // CHECK: vector.transfer_write %[[MUL]], %[[ARG0]][%[[C0]], %[[C0]]] {{.*}} : vector<4x256xf32>, memref<4x256xf32> 343 // CHECK: vector.transfer_write %[[RSQRT]], %[[ARG0]][%[[C0]], %[[C0]]] {{.*}} : vector<4x256xf32>, memref<4x256xf32> 344 // CHECK: vector.transfer_write %[[SEL]], %[[ARG0]][%[[C0]], %[[C0]]] {{.*}} : vector<4x256xf32>, memref<4x256xf32> 345 // CHECK: vector.transfer_write %[[SUB]], %[[ARG0]][%[[C0]], %[[C0]]] {{.*}} : vector<4x256xf32>, memref<4x256xf32> 346 // CHECK: vector.transfer_write %[[TAN]], %[[ARG0]][%[[C0]], %[[C0]]] {{.*}} : vector<4x256xf32>, memref<4x256xf32> 347 linalg.yield %6, %8, %c1_f32, %9, %10, %11, %12, %13, %14, %15 : f32, f32, 348 f32, f32, f32, f32, f32, f32, f32, f32 349 } 350 return 351} 352 353// ----- 354 355// CHECK-LABEL: func @generic_vectorize_tensor 356// CHECK-SAME: (%[[ARG0:.*]]: tensor<4x256xf32>, %[[ARG1:.*]]: tensor<4x256xf32>, 357// CHECK-SAME: %[[ARG2:.*]]: tensor<256xf32>, %[[ARG3:.*]]: f32) 358func @generic_vectorize_tensor(%arg0: tensor<4x256xf32>, 359 %arg1: tensor<4x256xf32>, %arg2: tensor<256xf32>, 360 %i: f32) -> (tensor<4x256xf32>, tensor<4x256xf32>, tensor<4x256xf32>, 361 tensor<4x256xf32>, tensor<4x256xf32>, tensor<4x256xf32>, tensor<4x256xf32>, 362 tensor<4x256xf32>, tensor<4x256xf32>, tensor<4x256xf32>) { 363 %c1_f32 = arith.constant 1.0 : f32 364 %r:10 = linalg.generic { 365 indexing_maps = [ 366 affine_map<(d0, d1) -> (d0, d1)>, 367 affine_map<(d0, d1) -> (d1)>, 368 affine_map<(d0, d1) -> (d0, d1)>, 369 affine_map<(d0, d1) -> (d0, d1)>, 370 affine_map<(d0, d1) -> (d0, d1)>, 371 affine_map<(d0, d1) -> (d0, d1)>, 372 affine_map<(d0, d1) -> (d0, d1)>, 373 affine_map<(d0, d1) -> (d0, d1)>, 374 affine_map<(d0, d1) -> (d0, d1)>, 375 affine_map<(d0, d1) -> (d0, d1)>, 376 affine_map<(d0, d1) -> (d0, d1)>, 377 affine_map<(d0, d1) -> (d0, d1)>], 378 iterator_types = ["parallel", "parallel"]} 379 ins(%arg1, %arg2: tensor<4x256xf32>, tensor<256xf32>) 380 outs( 381 %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0, %arg0 : 382 tensor<4x256xf32>, tensor<4x256xf32>, tensor<4x256xf32>, tensor<4x256xf32>, 383 tensor<4x256xf32>, tensor<4x256xf32>, tensor<4x256xf32>, tensor<4x256xf32>, 384 tensor<4x256xf32>, tensor<4x256xf32>) { 385 ^bb0(%arg3 : f32, %arg4 : f32, %arg5: f32, %arg6: f32, %arg7: f32, %arg8: f32, 386 %arg9 : f32, %arg10 : f32, %arg11 : f32, %arg12 : f32, %arg13 : f32, 387 %arg14 : f32): 388 // CHECK-DAG: %[[CST0:.*]] = arith.constant dense<2.000000e+00> : vector<4x256xf32> 389 // CHECK-DAG: %[[CST1:.*]] = arith.constant dense<1.000000e+00> : vector<4x256xf32> 390 // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index 391 // CHECK: %[[V2:.*]] = vector.transfer_read %[[ARG1]][%[[C0]], %[[C0]]], {{.*}} : tensor<4x256xf32>, vector<4x256xf32> 392 // CHECK: %[[V0:.*]] = vector.transfer_read %[[ARG2]][%[[C0]]], {{.*}} : tensor<256xf32>, vector<4x256xf32> 393 // CHECK: %[[V3:.*]] = vector.transfer_read %[[ARG0]][%[[C0]], %[[C0]]], {{.*}} : tensor<4x256xf32>, vector<4x256xf32> 394 // CHECK: %[[V1:.*]] = vector.transfer_read %[[ARG0]][%[[C0]], %[[C0]]], {{.*}} : tensor<4x256xf32>, vector<4x256xf32> 395 // CHECK: %[[ADD:.*]] = arith.addf %[[V0]], %[[V1]] : vector<4x256xf32> 396 %6 = arith.addf %arg4, %arg6 : f32 397 // CHECK: %[[CMP:.*]] = arith.cmpf ogt, %[[V2]], %[[V1]] : vector<4x256xf32> 398 %7 = arith.cmpf ogt, %arg3, %arg6 : f32 399 // CHECK: %[[ARG3B:.*]] = vector.broadcast %[[ARG3]] : f32 to vector<4x256xf32> 400 %8 = arith.constant 2.0 : f32 401 // CHECK: %[[DIV:.*]] = arith.divf %[[V3]], %[[ARG3B]] : vector<4x256xf32> 402 %9 = arith.divf %arg5, %i : f32 403 // CHECK: %[[EXP:.*]] = math.exp2 %[[V3]] : vector<4x256xf32> 404 %10 = math.exp2 %arg5 : f32 405 // CHECK: %[[MUL:.*]] = arith.mulf %[[V3]], %[[CST0]] : vector<4x256xf32> 406 %11 = arith.mulf %arg5, %8 : f32 407 // CHECK: %[[RSQRT:.*]] = math.rsqrt %[[V3]] : vector<4x256xf32> 408 %12 = math.rsqrt %arg5 : f32 409 // CHECK: %[[SEL:.*]] = arith.select %[[CMP]], %[[V3]], %[[V1]] : vector<4x256xi1>, vector<4x256xf32> 410 %13 = arith.select %7, %arg5, %arg6 : f32 411 // CHECK: %[[SUB:.*]] = arith.subf %[[V3]], %[[V0]] : vector<4x256xf32> 412 %14 = arith.subf %arg5, %arg4 : f32 413 // CHECK: %[[TAN:.*]] = math.tanh %[[V3]] : vector<4x256xf32> 414 %15 = math.tanh %arg5 : f32 415 // CHECK: %[[R0:.*]] = vector.transfer_write %[[ADD]], %[[ARG0]][%[[C0]], %[[C0]]] {{.*}} : vector<4x256xf32>, tensor<4x256xf32> 416 // CHECK: %[[R1:.*]] = vector.transfer_write %[[CST0]], %[[ARG0]][%[[C0]], %[[C0]]] {{.*}} : vector<4x256xf32>, tensor<4x256xf32> 417 // CHECK: %[[R2:.*]] = vector.transfer_write %[[CST1]], %[[ARG0]][%[[C0]], %[[C0]]] {{.*}} : vector<4x256xf32>, tensor<4x256xf32> 418 // CHECK: %[[R3:.*]] = vector.transfer_write %[[DIV]], %[[ARG0]][%[[C0]], %[[C0]]] {{.*}} : vector<4x256xf32>, tensor<4x256xf32> 419 // CHECK: %[[R4:.*]] = vector.transfer_write %[[EXP]], %[[ARG0]][%[[C0]], %[[C0]]] {{.*}} : vector<4x256xf32>, tensor<4x256xf32> 420 // CHECK: %[[R5:.*]] = vector.transfer_write %[[MUL]], %[[ARG0]][%[[C0]], %[[C0]]] {{.*}} : vector<4x256xf32>, tensor<4x256xf32> 421 // CHECK: %[[R6:.*]] = vector.transfer_write %[[RSQRT]], %[[ARG0]][%[[C0]], %[[C0]]] {{.*}} : vector<4x256xf32>, tensor<4x256xf32> 422 // CHECK: %[[R7:.*]] = vector.transfer_write %[[SEL]], %[[ARG0]][%[[C0]], %[[C0]]] {{.*}} : vector<4x256xf32>, tensor<4x256xf32> 423 // CHECK: %[[R8:.*]] = vector.transfer_write %[[SUB]], %[[ARG0]][%[[C0]], %[[C0]]] {{.*}} : vector<4x256xf32>, tensor<4x256xf32> 424 // CHECK: %[[R9:.*]] = vector.transfer_write %[[TAN]], %[[ARG0]][%[[C0]], %[[C0]]] {{.*}} : vector<4x256xf32>, tensor<4x256xf32> 425 linalg.yield %6, %8, %c1_f32, %9, %10, %11, %12, %13, %14, %15 : f32, f32, 426 f32, f32, f32, f32, f32, f32, f32, f32 427 } -> (tensor<4x256xf32>, tensor<4x256xf32>, tensor<4x256xf32>, 428 tensor<4x256xf32>, tensor<4x256xf32>, tensor<4x256xf32>, tensor<4x256xf32>, 429 tensor<4x256xf32>, tensor<4x256xf32>, tensor<4x256xf32>) 430 // CHECK: return %[[R0]], %[[R1]], %[[R2]], %[[R3]], %[[R4]], %[[R5]], %[[R6]], %[[R7]], %[[R8]], %[[R9]] : tensor<4x256xf32>, tensor<4x256xf32>, tensor<4x256xf32>, tensor<4x256xf32>, tensor<4x256xf32>, tensor<4x256xf32>, tensor<4x256xf32>, tensor<4x256xf32>, tensor<4x256xf32>, tensor<4x256xf32> 431 return %r#0, %r#1, %r#2, %r#3, %r#4, %r#5, %r#6, %r#7, %r#8, %r#9: 432 tensor<4x256xf32>, tensor<4x256xf32>, tensor<4x256xf32>, 433 tensor<4x256xf32>, tensor<4x256xf32>, tensor<4x256xf32>, tensor<4x256xf32>, 434 tensor<4x256xf32>, tensor<4x256xf32>, tensor<4x256xf32> 435} 436 437// ----- 438 439// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, 0, 0, d1)> 440// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0) -> (d0, 0, 0, 0)> 441// CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0) -> (0, 0, d0, 0)> 442// CHECK-DAG: #[[$MAP3:.*]] = affine_map<(d0, d1) -> (d1, 0, d0, 0)> 443// CHECK: func @generic_vectorize_broadcast_transpose 444// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index 445// CHECK-DAG: %[[CF:.*]] = arith.constant 0.000000e+00 : f32 446// CHECK: %[[V0:.*]] = vector.transfer_read %{{.*}}[%[[C0]], %[[C0]]], %[[CF]] {in_bounds = [true, true, true, true], permutation_map = #[[$MAP0]]} : memref<4x4xf32>, vector<4x4x4x4xf32> 447// CHECK: %[[V1:.*]] = vector.transfer_read %{{.*}}[%[[C0]]], %[[CF]] {in_bounds = [true, true, true, true], permutation_map = #[[$MAP1]]} : memref<4xf32>, vector<4x4x4x4xf32> 448// CHECK: %[[V2:.*]] = vector.transfer_read %{{.*}}[%[[C0]]], %[[CF]] {in_bounds = [true, true, true, true], permutation_map = #[[$MAP2]]} : memref<4xf32>, vector<4x4x4x4xf32> 449// CHECK: %[[V3:.*]] = vector.transfer_read %{{.*}}[%[[C0]], %[[C0]]], %[[CF]] {in_bounds = [true, true, true, true], permutation_map = #[[$MAP3]]} : memref<4x4xf32>, vector<4x4x4x4xf32> 450// CHECK: %[[SUB:.*]] = arith.subf %[[V0]], %[[V1]] : vector<4x4x4x4xf32> 451// CHECK: %[[ADD0:.*]] = arith.addf %[[V2]], %[[SUB]] : vector<4x4x4x4xf32> 452// CHECK: %[[ADD1:.*]] = arith.addf %[[V3]], %[[ADD0]] : vector<4x4x4x4xf32> 453// CHECK: vector.transfer_write %[[ADD1]], {{.*}} : vector<4x4x4x4xf32>, memref<4x4x4x4xf32> 454func @generic_vectorize_broadcast_transpose( 455 %A: memref<4xf32>, %B: memref<4x4xf32>, %C: memref<4x4x4x4xf32>) { 456 linalg.generic { 457 indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d3)>, 458 affine_map<(d0, d1, d2, d3) -> (d0)>, 459 affine_map<(d0, d1, d2, d3) -> (d2)>, 460 affine_map<(d0, d1, d2, d3) -> (d2, d0)>, 461 affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>], 462 iterator_types = ["parallel", "parallel", "parallel", "parallel"]} 463 ins(%B, %A, %A, %B: memref<4x4xf32>, memref<4xf32>, memref<4xf32>, memref<4x4xf32>) 464 outs(%C : memref<4x4x4x4xf32>) { 465 ^bb0(%arg0: f32, %arg1: f32, %arg2: f32, %arg3: f32, %arg4: f32): 466 %s = arith.subf %arg0, %arg1 : f32 467 %a = arith.addf %arg2, %s : f32 468 %b = arith.addf %arg3, %a : f32 469 linalg.yield %b : f32 470 } 471 return 472} 473 474// ----- 475 476// Test different input maps. 477#matmul_trait = { 478 indexing_maps = [ 479 affine_map<(d0, d1, d2, d3) -> (d1, d0)>, 480 affine_map<(d0, d1, d2, d3) -> (d3, d1)>, 481 affine_map<(d0, d1, d2, d3) -> (d3, d1, d0, d2)>, 482 affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)> 483 ], 484 iterator_types = ["parallel", "parallel", "parallel", "parallel"] 485} 486 487// CHECK-DAG: #[[MAP0:.*]] = affine_map<(d0, d1) -> (d1, d0, 0, 0)> 488// CHECK-DAG: #[[MAP1:.*]] = affine_map<(d0, d1) -> (0, d1, 0, d0)> 489// CHECK-DAG: #[[MAP2:.*]] = affine_map<(d0, d1, d2, d3) -> (d2, d1, d3, d0)> 490// CHECK: func @vectorization_transpose 491// CHECK: vector.transfer_read {{.*}}{in_bounds = [true, true, true, true], permutation_map = #[[MAP0]]} : memref<14x7xf32>, vector<7x14x8x16xf32> 492// CHECK: vector.transfer_read {{.*}}{in_bounds = [true, true, true, true], permutation_map = #[[MAP1]]} : memref<16x14xf32>, vector<7x14x8x16xf32> 493// CHECK: vector.transfer_read {{.*}}{in_bounds = [true, true, true, true], permutation_map = #[[MAP2]]} : memref<16x14x7x8xf32>, vector<7x14x8x16xf32> 494// CHECK: arith.addf {{.*}} : vector<7x14x8x16xf32> 495// CHECK: arith.addf {{.*}} : vector<7x14x8x16xf32> 496// CHECK: vector.transfer_write {{.*}} : vector<7x14x8x16xf32>, memref<7x14x8x16xf32> 497func @vectorization_transpose(%A: memref<14x7xf32>, %B: memref<16x14xf32>, 498 %C: memref<16x14x7x8xf32>, %D: memref<7x14x8x16xf32>) { 499 linalg.generic #matmul_trait 500 ins(%A, %B, %C : memref<14x7xf32>, memref<16x14xf32>, memref<16x14x7x8xf32>) 501 outs(%D : memref<7x14x8x16xf32>) { 502 ^bb(%a: f32, %b: f32, %c: f32, %d: f32) : 503 %e = arith.addf %a, %b: f32 504 %f = arith.addf %e, %c: f32 505 linalg.yield %f : f32 506 } 507 return 508} 509 510// ----- 511 512// CHECK-LABEL: func @matmul_tensors 513// CHECK-SAME: (%[[ARG0:.*]]: tensor<8x4xf32>, %[[ARG1:.*]]: tensor<4x12xf32>, 514// CHECK-SAME: %[[ARG2:.*]]: tensor<8x12xf32>) -> tensor<8x12xf32> 515func @matmul_tensors( 516 %arg0: tensor<8x4xf32>, %arg1: tensor<4x12xf32>, %arg2: tensor<8x12xf32>) 517 -> tensor<8x12xf32> { 518 // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index 519 // CHECK-DAG: %[[V0:.*]] = vector.transfer_read %[[ARG0]][%[[C0]], %[[C0]]], {{.*}} : tensor<8x4xf32>, vector<8x12x4xf32> 520 // CHECK-DAG: %[[V1:.*]] = vector.transfer_read %[[ARG1]][%[[C0]], %[[C0]]], {{.*}} : tensor<4x12xf32>, vector<8x12x4xf32> 521 // CHECK-DAG: %[[V2:.*]] = vector.transfer_read %[[ARG2]][%[[C0]], %[[C0]]], {{.*}} : tensor<8x12xf32>, vector<8x12xf32> 522 // 523 // linalg matmul lowers gets expanded to a 3D reduction, canonicalization later 524 // convert it to a 2D contract. 525 // CHECK: %[[MUL:.*]] = arith.mulf %[[V0]], %[[V1]] : vector<8x12x4xf32> 526 // CHECK: %[[R:.*]] = vector.multi_reduction <add>, %[[MUL]] [2] : vector<8x12x4xf32> to vector<8x12xf32> 527 // CHECK: %[[ADD:.*]] = arith.addf %[[R]], %[[V2]] : vector<8x12xf32> 528 // CHECK: %[[W:.*]] = vector.transfer_write %[[ADD]], %[[ARG2]][%[[C0]], %[[C0]]] {in_bounds = [true, true]} : vector<8x12xf32>, tensor<8x12xf32> 529 %0 = linalg.matmul ins(%arg0, %arg1: tensor<8x4xf32>, tensor<4x12xf32>) 530 outs(%arg2: tensor<8x12xf32>) 531 -> tensor<8x12xf32> 532 // CHECK: return %[[W]] : tensor<8x12xf32> 533 return %0 : tensor<8x12xf32> 534} 535 536// ----- 537 538// CHECK-LABEL: func @pad_static( 539// CHECK-SAME: %[[ARG0:.*]]: tensor<2x?x2xf32>, %[[PAD:.*]]: f32 540// CHECK-NOT: tensor.pad 541// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index 542// CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index 543// CHECK-DAG: %[[INIT:.*]] = linalg.init_tensor [2, 3, 4] : tensor<2x3x4xf32> 544// CHECK-DAG: %[[VEC:.*]] = vector.broadcast %[[PAD]] : f32 to vector<2x3x4xf32> 545// CHECK: %[[FILL:.*]] = vector.transfer_write %[[VEC]], %[[INIT]]{{.*}} : vector<2x3x4xf32>, tensor<2x3x4xf32> 546// CHECK: %[[READ:.*]] = vector.transfer_read %[[ARG0]][%[[C0]], %[[C0]], %[[C0]]], %[[PAD]] {in_bounds = [true, false, true]} : tensor<2x?x2xf32>, vector<2x3x2xf32> 547// CHECK: %[[RESULT:.*]] = vector.transfer_write %[[READ]], %[[FILL]][%[[C0]], %[[C0]], %[[C2]]] {in_bounds = [true, true, true]} : vector<2x3x2xf32>, tensor<2x3x4xf32> 548// CHECK: return %[[RESULT]] 549func @pad_static(%arg0: tensor<2x?x2xf32>, %pad_value: f32) -> tensor<2x3x4xf32> { 550 %0 = tensor.pad %arg0 low[0, 0, 2] high[0, 1, 0] { 551 ^bb0(%arg1: index, %arg2: index, %arg3: index): 552 tensor.yield %pad_value : f32 553 } : tensor<2x?x2xf32> to tensor<2x3x4xf32> 554 return %0 : tensor<2x3x4xf32> 555} 556 557// ----- 558 559// CHECK-LABEL: func @pad_static_source( 560// CHECK-SAME: %[[ARG0:.*]]: tensor<2x5x2xf32>, %[[PAD:.*]]: f32 561// CHECK-NOT: tensor.pad 562// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index 563// CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index 564// CHECK: %[[INIT:.*]] = linalg.init_tensor [2, 6, 4] : tensor<2x6x4xf32> 565// CHECK: %[[VEC:.*]] = vector.broadcast %[[PAD]] : f32 to vector<2x6x4xf32> 566// CHECK: %[[FILL:.*]] = vector.transfer_write %[[VEC]], %[[INIT]][%[[C0]], %[[C0]], %[[C0]]] {in_bounds = [true, true, true]} : vector<2x6x4xf32>, tensor<2x6x4xf32> 567// CHECK: %[[READ:.*]] = vector.transfer_read %[[ARG0]][%[[C0]], %[[C0]], %[[C0]]], %{{.*}} {in_bounds = [true, true, true]} : tensor<2x5x2xf32>, vector<2x5x2xf32> 568// CHECK: %[[WRITE:.*]] = vector.transfer_write %[[READ]], %[[FILL]][%[[C0]], %[[C0]], %[[C2]]] {in_bounds = [true, true, true]} : vector<2x5x2xf32>, tensor<2x6x4xf32> 569// CHECK: return %[[WRITE]] 570func @pad_static_source(%arg0: tensor<2x5x2xf32>, %pad_value: f32) -> tensor<2x6x4xf32> { 571 %0 = tensor.pad %arg0 low[0, 0, 2] high[0, 1, 0] { 572 ^bb0(%arg1: index, %arg2: index, %arg3: index): 573 tensor.yield %pad_value : f32 574 } : tensor<2x5x2xf32> to tensor<2x6x4xf32> 575 return %0 : tensor<2x6x4xf32> 576} 577 578// ----- 579 580// CHECK-LABEL: func @pad_static_dynamic( 581// CHECK-SAME: %[[SRC:.*]]: tensor<1x2x2x?xf32>, %[[LOW:.*]]: index, %[[HIGH:.*]]: index 582// CHECK-NOT: tensor.pad 583// CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index 584// CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index 585// CHECK-DAG: %[[C5:.*]] = arith.constant 5 : index 586// CHECK: %[[V0:.*]] = arith.addi %[[LOW]], %[[C2]] : index 587// CHECK: %[[V1:.*]] = arith.addi %[[V0]], %[[C3]] : index 588// CHECK: %[[V2:.*]] = arith.addi %[[HIGH]], %[[C5]] : index 589// CHECK: %[[DIM3:.*]] = tensor.dim %[[SRC]], %[[C3]] : tensor<1x2x2x?xf32> 590// CHECK: %[[V4:.*]] = arith.addi %[[DIM3]], %[[C3]] : index 591// CHECK: %[[V5:.*]] = arith.addi %[[V4]], %[[C2]] : index 592// CHECK: %[[INIT:.*]] = linalg.init_tensor [6, %[[V1]], %[[V2]], %[[V5]]] : tensor<6x?x?x?xf32> 593// CHECK: %[[FILL:.*]] = linalg.fill ins(%{{.*}} : f32) outs(%[[INIT]] : tensor<6x?x?x?xf32>) -> tensor<6x?x?x?xf32> 594// CHECK: %[[SRCDIM:.*]] = tensor.dim %[[SRC]], %[[C3]] : tensor<1x2x2x?xf32> 595// CHECK: %[[RESULT:.*]] = tensor.insert_slice %[[SRC]] into %[[FILL]][2, %[[LOW]], 3, 3] [1, 2, 2, %[[SRCDIM]]] [1, 1, 1, 1] : tensor<1x2x2x?xf32> into tensor<6x?x?x?xf32> 596// CHECK: return %[[RESULT]] 597func @pad_static_dynamic(%arg0: tensor<1x2x2x?xf32>, %low: index, %high: index, 598 %pad_value: f32) -> tensor<6x?x?x?xf32> { 599 %0 = tensor.pad %arg0 low[2, %low, 3, 3] high[3, 3, %high, 2] { 600 ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index): 601 tensor.yield %pad_value : f32 602 } : tensor<1x2x2x?xf32> to tensor<6x?x?x?xf32> 603 return %0 : tensor<6x?x?x?xf32> 604} 605 606// ----- 607 608// CHECK-LABEL: func @pad_and_transfer_read 609// CHECK-SAME: %[[ARG0:.*]]: tensor<5x6xf32> 610// CHECK-NOT: tensor.pad 611// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index 612// CHECK-DAG: %[[C5:.*]] = arith.constant 5.0 613// CHECK: %[[RESULT:.*]] = vector.transfer_read %[[ARG0]][%[[C0]], %[[C0]]], %[[C5]] : tensor<5x6xf32>, vector<7x9xf32> 614// CHECK: return %[[RESULT]] 615func @pad_and_transfer_read(%arg0: tensor<5x6xf32>) -> vector<7x9xf32> { 616 %c0 = arith.constant 0 : index 617 %c5 = arith.constant 5.0 : f32 618 %c6 = arith.constant 6.0 : f32 619 %0 = tensor.pad %arg0 low[0, 0] high[5, 7] { 620 ^bb0(%arg1: index, %arg2: index): 621 tensor.yield %c5 : f32 622 } : tensor<5x6xf32> to tensor<10x13xf32> 623 %1 = vector.transfer_read %0[%c0, %c0], %c6 624 : tensor<10x13xf32>, vector<7x9xf32> 625 return %1 : vector<7x9xf32> 626} 627 628// ----- 629 630func private @make_vector() -> vector<7x9xf32> 631 632// CHECK-LABEL: func @pad_and_transfer_write_static 633// CHECK-SAME: %[[ARG0:.*]]: tensor<5x6xf32> 634// CHECK-NOT: tensor.pad 635// CHECK: %[[C0:.*]] = arith.constant 0 : index 636// CHECK: %[[VEC0:.*]] = call @make_vector() : () -> vector<7x9xf32> 637// CHECK: %[[RESULT:.*]] = vector.transfer_write %[[VEC0]], %[[ARG0]][%[[C0]], %[[C0]]] : vector<7x9xf32>, tensor<5x6xf32> 638// CHECK: return %[[RESULT]] 639func @pad_and_transfer_write_static( 640 %arg0: tensor<5x6xf32>) -> tensor<5x6xf32> { 641 %c0 = arith.constant 0 : index 642 %c5 = arith.constant 5.0 : f32 643 %0 = tensor.pad %arg0 low[0, 0] high[5, 7] { 644 ^bb0(%arg2: index, %arg3: index): 645 tensor.yield %c5 : f32 646 } : tensor<5x6xf32> to tensor<10x13xf32> 647 %1 = call @make_vector() : () -> vector<7x9xf32> 648 %2 = vector.transfer_write %1, %0[%c0, %c0] 649 : vector<7x9xf32>, tensor<10x13xf32> 650 %3 = tensor.extract_slice %2[0, 0] [5, 6] [1, 1] : tensor<10x13xf32> to tensor<5x6xf32> 651 return %3 : tensor<5x6xf32> 652} 653 654// ----- 655 656func private @make_vector() -> vector<7x9xf32> 657 658// CHECK-LABEL: func @pad_and_transfer_write_dynamic_static 659// CHECK-SAME: %[[ARG0:.*]]: tensor<?x?xf32>, %[[SIZE:.*]]: index, %[[PADDING:.*]]: index 660// CHECK-NOT: tensor.pad 661// CHECK: %[[C0:.*]] = arith.constant 0 : index 662// CHECK: %[[SUB:.*]] = tensor.extract_slice %[[ARG0]][0, 0] [%[[SIZE]], 6] [1, 1] : tensor<?x?xf32> to tensor<?x6xf32> 663// CHECK: %[[VEC0:.*]] = call @make_vector() : () -> vector<7x9xf32> 664// CHECK: %[[RESULT:.*]] = vector.transfer_write %[[VEC0]], %[[SUB]][%[[C0]], %[[C0]]] : vector<7x9xf32>, tensor<?x6xf32> 665// CHECK: return %[[RESULT]] 666func @pad_and_transfer_write_dynamic_static( 667 %arg0: tensor<?x?xf32>, %size: index, %padding: index) -> tensor<?x6xf32> { 668 %c0 = arith.constant 0 : index 669 %c5 = arith.constant 5.0 : f32 670 %s = tensor.extract_slice %arg0[0, 0] [%size, 6] [1, 1] 671 : tensor<?x?xf32> to tensor<?x6xf32> 672 %0 = tensor.pad %s low[0, 0] high[%padding, 7] { 673 ^bb0(%arg2: index, %arg3: index): 674 tensor.yield %c5 : f32 675 } : tensor<?x6xf32> to tensor<?x13xf32> 676 %1 = call @make_vector() : () -> vector<7x9xf32> 677 %2 = vector.transfer_write %1, %0[%c0, %c0] 678 : vector<7x9xf32>, tensor<?x13xf32> 679 %3 = tensor.extract_slice %2[0, 0] [%size, 6] [1, 1] : tensor<?x13xf32> to tensor<?x6xf32> 680 return %3 : tensor<?x6xf32> 681} 682 683// ----- 684 685func private @make_vector() -> tensor<12x13xf32> 686 687// CHECK-LABEL: func @pad_and_insert_slice_source 688// CHECK-SAME: %[[ARG0:.*]]: tensor<5x6xf32> 689// CHECK-NOT: tensor.pad 690// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index 691// CHECK-DAG: %[[C5:.*]] = arith.constant 5.0 692// CHECK: %[[VEC0:.*]] = call @make_vector() : () -> tensor<12x13xf32> 693// CHECK: %[[READ:.*]] = vector.transfer_read %[[ARG0]][%[[C0]], %[[C0]]], %[[C5]] : tensor<5x6xf32>, vector<7x9xf32> 694// CHECK: %[[WRITE:.*]] = vector.transfer_write %[[READ]], %[[VEC0]][%[[C0]], %[[C0]]] {in_bounds = [true, true]} : vector<7x9xf32>, tensor<12x13xf32> 695// CHECK: return %[[WRITE]] 696func @pad_and_insert_slice_source( 697 %arg0: tensor<5x6xf32>) -> tensor<12x13xf32> { 698 %c0 = arith.constant 0 : index 699 %c5 = arith.constant 5.0 : f32 700 %0 = tensor.pad %arg0 low[0, 0] high[2, 3] { 701 ^bb0(%arg2: index, %arg3: index): 702 tensor.yield %c5 : f32 703 } : tensor<5x6xf32> to tensor<7x9xf32> 704 %1 = call @make_vector() : () -> tensor<12x13xf32> 705 %r = tensor.insert_slice %0 into %1[0, 0][7, 9][1, 1] : tensor<7x9xf32> into tensor<12x13xf32> 706 return %r : tensor<12x13xf32> 707} 708 709// ----- 710 711func private @make_vector() -> tensor<12x13xf32> 712 713// CHECK-LABEL: func @pad_and_insert_slice_dest 714// Check the insert slice is not rewritten if the padded result is used by the destination operand. 715// CHECK: %[[T1:.*]] = call @make_vector() : () -> tensor<12x13xf32> 716// CHECK: = tensor.insert_slice %[[T1]] into 717func @pad_and_insert_slice_dest( 718 %arg0: tensor<1x5x6xf32>) -> tensor<1x12x13xf32> { 719 %c5 = arith.constant 5.0 : f32 720 %0 = tensor.pad %arg0 low[0, 0, 0] high[0, 7, 7] { 721 ^bb0(%arg2: index, %arg3: index, %arg4: index): 722 tensor.yield %c5 : f32 723 } : tensor<1x5x6xf32> to tensor<1x12x13xf32> 724 %1 = call @make_vector() : () -> tensor<12x13xf32> 725 %r = tensor.insert_slice %1 into %0[0, 0, 0][1, 12, 13][1, 1, 1] : tensor<12x13xf32> into tensor<1x12x13xf32> 726 return %r : tensor<1x12x13xf32> 727} 728 729// ----- 730 731// CHECK-LABEL: func @pad_tensor_non_const_pad_value 732// CHECK-SAME: %[[ARG0:.*]]: tensor<5x6xf32> 733// CHECK-NOT: tensor.pad 734// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index 735// CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index 736// CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index 737// CHECK: %[[FILL:.*]] = tensor.generate 738// CHECK: %[[RES:.*]] = arith.mulf 739// CHECK: tensor.yield %[[RES]] : f32 740// CHECK: %[[READ:.*]] = vector.transfer_read %[[ARG0]][%[[C0]], %[[C0]]], %{{.*}} {in_bounds = [true, true]} : tensor<5x6xf32>, vector<5x6xf32> 741// CHECK: %[[WRITE:.*]] = vector.transfer_write %[[READ]], %[[FILL]][%[[C3]], %[[C4]]] {in_bounds = [true, true]} : vector<5x6xf32>, tensor<12x13xf32> 742// CHECK: return %[[WRITE]] 743func @pad_tensor_non_const_pad_value(%arg0: tensor<5x6xf32>) -> tensor<12x13xf32> { 744 %c0 = arith.constant 0 : index 745 %c5 = arith.constant 5.0 : f32 746 %0 = tensor.pad %arg0 low[3, 4] high[4, 3] { 747 ^bb0(%arg1: index, %arg2: index): 748 %i1 = arith.index_cast %arg1 : index to i32 749 %i2 = arith.index_cast %arg2 : index to i32 750 %f1 = arith.sitofp %i1 : i32 to f32 751 %f2 = arith.sitofp %i2 : i32 to f32 752 %m = arith.mulf %f1, %f2 : f32 753 tensor.yield %m : f32 754 } : tensor<5x6xf32> to tensor<12x13xf32> 755 return %0 : tensor<12x13xf32> 756} 757 758// ----- 759 760// CHECK-LABEL: func @sum_exp 761func @sum_exp(%input: tensor<4x16x8xf32>, %output: tensor<4x16xf32>) 762 -> tensor<4x16xf32> 763{ 764 // CHECK: vector.transfer_read {{.*}} : tensor<4x16x8xf32>, vector<4x16x8xf32> 765 // CHECK: vector.transfer_read {{.*}} {in_bounds = [true, true]} : tensor<4x16xf32>, vector<4x16xf32> 766 // CHECK: math.exp {{.*}} : vector<4x16x8xf32> 767 // CHECK: vector.multi_reduction <add>, %{{.*}} [2] : vector<4x16x8xf32> to vector<4x16xf32> 768 // CHECK: addf {{.*}} : vector<4x16xf32> 769 // CHECK: vector.transfer_write {{.*}} : vector<4x16xf32>, tensor<4x16xf32> 770 // CHECK: return {{.*}} : tensor<4x16xf32> 771 %0 = linalg.generic { 772 indexing_maps = [ 773 affine_map<(d0, d1, d2) -> (d0, d1, d2)>, 774 affine_map<(d0, d1, d2) -> (d0, d1)> 775 ], 776 iterator_types = ["parallel", "parallel", "reduction"] 777 } ins(%input : tensor<4x16x8xf32>) outs(%output : tensor<4x16xf32>) { 778 ^bb0(%arg0: f32, %arg1: f32): 779 %1 = math.exp %arg0 : f32 780 %2 = arith.addf %1, %arg1 : f32 781 linalg.yield %2 : f32 782 } -> tensor<4x16xf32> 783 return %0 : tensor<4x16xf32> 784} 785 786// ----- 787 788// CHECK-DAG: #[[$M1:.*]] = affine_map<(d0, d1) -> (d1, d0, 0, 0)> 789// CHECK-DAG: #[[$M2:.*]] = affine_map<(d0, d1) -> (0, 0, d1, d0)> 790// CHECK-DAG: #[[$M3:.*]] = affine_map<(d0, d1) -> (d1, d0)> 791 792// CHECK-LABEL: func @sum_exp_2 793func @sum_exp_2(%input: tensor<3x2xf32>, %input_2: tensor<5x4xf32>, %output: tensor<5x2xf32>) 794 -> tensor<5x2xf32> 795{ 796 // CHECK: vector.transfer_read {{.*}} {in_bounds = [true, true, true, true], permutation_map = #[[$M1]]} : tensor<3x2xf32>, vector<2x3x4x5xf32> 797 // CHECK: vector.transfer_read {{.*}} {in_bounds = [true, true, true, true], permutation_map = #[[$M2]]} : tensor<5x4xf32>, vector<2x3x4x5xf32> 798 // CHECK: vector.transfer_read {{.*}} {in_bounds = [true, true], permutation_map = #[[$M3]]} : tensor<5x2xf32>, vector<2x5xf32> 799 // CHECK: math.exp {{.*}} : vector<2x3x4x5xf32> 800 // CHECK: math.exp {{.*}} : vector<2x3x4x5xf32> 801 // CHECK: addf {{.*}} : vector<2x3x4x5xf32> 802 // CHECK: vector.multi_reduction <add>, {{.*}} [1, 2] : vector<2x3x4x5xf32> to vector<2x5xf32> 803 // CHECK: addf {{.*}} : vector<2x5xf32> 804 // CHECK: vector.transfer_write {{.*}} {in_bounds = [true, true], permutation_map = #[[$M3]]} : vector<2x5xf32>, tensor<5x2xf32> 805 // CHECK: return {{.*}} : tensor<5x2xf32> 806 %0 = linalg.generic { 807 indexing_maps = [ 808 affine_map<(d0, d1, d2, d3) -> (d1, d0)>, 809 affine_map<(d0, d1, d2, d3) -> (d3, d2)>, 810 affine_map<(d0, d1, d2, d3) -> (d3, d0)> 811 ], 812 iterator_types = ["parallel", "reduction", "reduction", "parallel"] 813 } ins(%input, %input_2 : tensor<3x2xf32>, tensor<5x4xf32>) outs(%output : tensor<5x2xf32>) { 814 ^bb0(%arg0: f32, %arg1: f32, %arg2: f32): 815 %1 = math.exp %arg0 : f32 816 %2 = math.exp %arg1 : f32 817 %3 = arith.addf %1, %2 : f32 818 %4 = arith.addf %3, %arg2 : f32 819 linalg.yield %4 : f32 820 } -> tensor<5x2xf32> 821 return %0 : tensor<5x2xf32> 822} 823 824// ----- 825 826// CHECK-LABEL: func @red_max_2d( 827func @red_max_2d(%arg0: tensor<4x4xf32>) -> tensor<4xf32> { 828 // CHECK: %[[CMINF:.+]] = arith.constant dense<-3.402820e+38> : vector<4xf32> 829 // CHECK: linalg.init_tensor [4] : tensor<4xf32> 830 // CHECK: vector.transfer_write {{.*}} : vector<4xf32>, tensor<4xf32> 831 // CHECK: %[[R:.+]] = vector.multi_reduction <maxf>, {{.*}} [1] : vector<4x4xf32> to vector<4xf32> 832 // CHECK: maxf %[[R]], %[[CMINF]] : vector<4xf32> 833 // CHECK: vector.transfer_write {{.*}} : vector<4xf32>, tensor<4xf32> 834 %ident = arith.constant -3.40282e+38 : f32 835 %init = linalg.init_tensor [4] : tensor<4xf32> 836 %fill = linalg.fill ins(%ident : f32) outs(%init : tensor<4xf32>) -> tensor<4xf32> 837 %red = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, 838 affine_map<(d0, d1) -> (d0)>], 839 iterator_types = ["parallel", "reduction"]} 840 ins(%arg0 : tensor<4x4xf32>) outs(%fill : tensor<4xf32>) { 841 ^bb0(%in0: f32, %out0: f32): 842 %max = arith.maxf %in0, %out0 : f32 843 linalg.yield %max : f32 844 } -> tensor<4xf32> 845 return %red : tensor<4xf32> 846} 847 848// ----- 849 850// CHECK-LABEL: func @red_min_2d( 851func @red_min_2d(%arg0: tensor<4x4xf32>) -> tensor<4xf32> { 852 // CHECK: %[[CMAXF:.+]] = arith.constant dense<3.402820e+38> : vector<4xf32> 853 // CHECK: linalg.init_tensor [4] : tensor<4xf32> 854 // CHECK: vector.transfer_write {{.*}} : vector<4xf32>, tensor<4xf32> 855 // CHECK: vector.transfer_read {{.*}} : tensor<4x4xf32>, vector<4x4xf32> 856 // CHECK: %[[R:.+]] = vector.multi_reduction <minf>, {{.*}} [1] : vector<4x4xf32> to vector<4xf32> 857 // CHECK: arith.minf %[[R]], %[[CMAXF]] : vector<4xf32> 858 // CHECK: vector.transfer_write {{.*}} : vector<4xf32>, tensor<4xf32> 859 %maxf32 = arith.constant 3.40282e+38 : f32 860 %init = linalg.init_tensor [4] : tensor<4xf32> 861 %fill = linalg.fill ins(%maxf32 : f32) outs(%init : tensor<4xf32>) -> tensor<4xf32> 862 %red = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, 863 affine_map<(d0, d1) -> (d0)>], 864 iterator_types = ["parallel", "reduction"]} 865 ins(%arg0 : tensor<4x4xf32>) outs(%fill : tensor<4xf32>) { 866 ^bb0(%in0: f32, %out0: f32): 867 %min = arith.minf %out0, %in0 : f32 868 linalg.yield %min : f32 869 } -> tensor<4xf32> 870 return %red : tensor<4xf32> 871} 872 873// ----- 874 875// CHECK-LABEL: func @red_mul_2d( 876func @red_mul_2d(%arg0: tensor<4x4xf32>) -> tensor<4xf32> { 877 // CHECK: linalg.init_tensor [4] : tensor<4xf32> 878 // CHECK: vector.transfer_write {{.*}} : vector<4xf32>, tensor<4xf32> 879 // CHECK: vector.transfer_read {{.*}} : tensor<4x4xf32>, vector<4x4xf32> 880 // CHECK: vector.multi_reduction <mul>, {{.*}} [1] : vector<4x4xf32> to vector<4xf32> 881 // CHECK: vector.transfer_write {{.*}} : vector<4xf32>, tensor<4xf32> 882 %ident = arith.constant 1.0 : f32 883 %init = linalg.init_tensor [4] : tensor<4xf32> 884 %fill = linalg.fill ins(%ident : f32) outs(%init : tensor<4xf32>) -> tensor<4xf32> 885 %red = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, 886 affine_map<(d0, d1) -> (d0)>], 887 iterator_types = ["parallel", "reduction"]} 888 ins(%arg0 : tensor<4x4xf32>) outs(%fill : tensor<4xf32>) { 889 ^bb0(%in0: f32, %out0: f32): 890 %mul = arith.mulf %in0, %out0 : f32 891 linalg.yield %mul : f32 892 } -> tensor<4xf32> 893 return %red : tensor<4xf32> 894} 895 896// ----- 897 898// CHECK-LABEL: func @red_or_2d( 899func @red_or_2d(%arg0: tensor<4x4xi1>) -> tensor<4xi1> { 900 // CHECK: linalg.init_tensor [4] : tensor<4xi1> 901 // CHECK: vector.transfer_write {{.*}} : vector<4xi1>, tensor<4xi1> 902 // CHECK: vector.transfer_read {{.*}} : tensor<4x4xi1>, vector<4x4xi1> 903 // CHECK: vector.multi_reduction <or>, {{.*}} [1] : vector<4x4xi1> to vector<4xi1> 904 // CHECK: vector.transfer_write {{.*}} : vector<4xi1>, tensor<4xi1> 905 %ident = arith.constant false 906 %init = linalg.init_tensor [4] : tensor<4xi1> 907 %fill = linalg.fill ins(%ident : i1) outs(%init : tensor<4xi1>) -> tensor<4xi1> 908 %red = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, 909 affine_map<(d0, d1) -> (d0)>], 910 iterator_types = ["parallel", "reduction"]} 911 ins(%arg0 : tensor<4x4xi1>) outs(%fill : tensor<4xi1>) { 912 ^bb0(%in0: i1, %out0: i1): 913 %or = arith.ori %in0, %out0 : i1 914 linalg.yield %or : i1 915 } -> tensor<4xi1> 916 return %red : tensor<4xi1> 917} 918 919// ----- 920 921// CHECK-LABEL: func @red_and_2d( 922func @red_and_2d(%arg0: tensor<4x4xi1>) -> tensor<4xi1> { 923 // CHECK: linalg.init_tensor [4] : tensor<4xi1> 924 // CHECK: vector.transfer_write {{.*}} : vector<4xi1>, tensor<4xi1> 925 // CHECK: vector.transfer_read {{.*}} : tensor<4x4xi1>, vector<4x4xi1> 926 // CHECK: vector.multi_reduction <and>, {{.*}} [1] : vector<4x4xi1> to vector<4xi1> 927 // CHECK: vector.transfer_write {{.*}} : vector<4xi1>, tensor<4xi1> 928 %ident = arith.constant true 929 %init = linalg.init_tensor [4] : tensor<4xi1> 930 %fill = linalg.fill ins(%ident : i1) outs(%init : tensor<4xi1>) -> tensor<4xi1> 931 %red = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, 932 affine_map<(d0, d1) -> (d0)>], 933 iterator_types = ["parallel", "reduction"]} 934 ins(%arg0 : tensor<4x4xi1>) outs(%fill : tensor<4xi1>) { 935 ^bb0(%in0: i1, %out0: i1): 936 %and = arith.andi %in0, %out0 : i1 937 linalg.yield %and : i1 938 } -> tensor<4xi1> 939 return %red : tensor<4xi1> 940} 941 942// ----- 943 944// CHECK-LABEL: func @red_xor_2d( 945func @red_xor_2d(%arg0: tensor<4x4xi1>) -> tensor<4xi1> { 946 // CHECK: linalg.init_tensor [4] : tensor<4xi1> 947 // CHECK: vector.transfer_write {{.*}} : vector<4xi1>, tensor<4xi1> 948 // CHECK: vector.transfer_read {{.*}} : tensor<4x4xi1>, vector<4x4xi1> 949 // CHECK: vector.multi_reduction <xor>, {{.*}} [1] : vector<4x4xi1> to vector<4xi1> 950 // CHECK: vector.transfer_write {{.*}} : vector<4xi1>, tensor<4xi1> 951 %ident = arith.constant false 952 %init = linalg.init_tensor [4] : tensor<4xi1> 953 %fill = linalg.fill ins(%ident : i1) outs(%init : tensor<4xi1>) -> tensor<4xi1> 954 %red = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, 955 affine_map<(d0, d1) -> (d0)>], 956 iterator_types = ["parallel", "reduction"]} 957 ins(%arg0 : tensor<4x4xi1>) outs(%fill : tensor<4xi1>) { 958 ^bb0(%in0: i1, %out0: i1): 959 %xor = arith.xori %in0, %out0 : i1 960 linalg.yield %xor : i1 961 } -> tensor<4xi1> 962 return %red : tensor<4xi1> 963} 964 965// ----- 966 967// CHECK-DAG: #[[$M5:.*]] = affine_map<(d0, d1) -> (d0, 0)> 968 969// CHECK-LABEL: func @explicit_broadcast( 970func @explicit_broadcast(%arg0: tensor<4x4xf32>, %arg1: tensor<4x1xf32>) -> tensor<4x4xf32> { 971 // CHECK: vector.transfer_read {{.*}} {in_bounds = [true, true]} : tensor<4x4xf32>, vector<4x4xf32> 972 // CHECK: vector.transfer_read {{.*}} {in_bounds = [true, true], permutation_map = #[[$M5]]} : tensor<4x1xf32>, vector<4x4xf32> 973 // CHECK: subf {{.*}} : vector<4x4xf32> 974 // CHECK: vector.transfer_write {{.*}} {in_bounds = [true, true]} : vector<4x4xf32>, tensor<4x4xf32> 975 %c0 = arith.constant 0.0 : f32 976 %init = linalg.init_tensor [4, 4] : tensor<4x4xf32> 977 %fill = linalg.fill ins(%c0 : f32) outs(%init : tensor<4x4xf32>) -> tensor<4x4xf32> 978 %red = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, 979 affine_map<(d0, d1) -> (d0, 0)>, 980 affine_map<(d0, d1) -> (d0, d1)>], 981 iterator_types = ["parallel", "parallel"]} 982 ins(%arg0, %arg1 : tensor<4x4xf32>, tensor<4x1xf32>) 983 outs(%fill : tensor<4x4xf32>) { 984 ^bb0(%arg7: f32, %arg8: f32, %arg9: f32): 985 %40 = arith.subf %arg7, %arg8 : f32 986 linalg.yield %40 : f32 987 } -> tensor<4x4xf32> 988 return %red : tensor<4x4xf32> 989} 990 991// ----- 992 993// CHECK-DAG: #[[$M6:.*]] = affine_map<(d0, d1) -> (d0, 0)> 994 995// CHECK-LABEL: func @fused_broadcast_red_2d 996func @fused_broadcast_red_2d(%arg0: tensor<4x4xf32>, %arg1: tensor<4x1xf32>) -> tensor<4xf32> { 997 // CHECK: vector.transfer_read {{.*}} {in_bounds = [true, true]} : tensor<4x4xf32>, vector<4x4xf32> 998 // CHECK: vector.transfer_read {{.*}} {in_bounds = [true, true], permutation_map = #[[$M6]]} : tensor<4x1xf32>, vector<4x4xf32> 999 // CHECK: subf {{.*}} : vector<4x4xf32> 1000 // CHECK: math.exp {{.*}} : vector<4x4xf32> 1001 // CHECK: vector.multi_reduction <add>, {{.*}} : vector<4x4xf32> to vector<4xf32> 1002 // CHECK: addf {{.*}} : vector<4xf32> 1003 // CHECK: vector.transfer_write {{.*}} {in_bounds = [true]} : vector<4xf32>, tensor<4xf32> 1004 %c0 = arith.constant 0.0 : f32 1005 %init = linalg.init_tensor [4] : tensor<4xf32> 1006 %fill = linalg.fill ins(%c0 : f32) outs(%init : tensor<4xf32>) -> tensor<4xf32> 1007 %red = linalg.generic {indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, 1008 affine_map<(d0, d1) -> (d0, 0)>, 1009 affine_map<(d0, d1) -> (d0)>], 1010 iterator_types = ["parallel", "reduction"]} 1011 ins(%arg0, %arg1 : tensor<4x4xf32>, tensor<4x1xf32>) 1012 outs(%fill : tensor<4xf32>) { 1013 ^bb0(%arg7: f32, %arg8: f32, %arg9: f32): 1014 %40 = arith.subf %arg7, %arg8 : f32 1015 %41 = math.exp %40 : f32 1016 %42 = arith.addf %41, %arg9 : f32 1017 linalg.yield %42 : f32 1018 } -> tensor<4xf32> 1019 return %red : tensor<4xf32> 1020} 1021 1022// ----- 1023 1024// CHECK-LABEL: func @reduce_1d( 1025// CHECK-SAME: %[[A:.*]]: tensor<32xf32> 1026func @reduce_1d(%arg0: tensor<32xf32>) -> tensor<f32> { 1027 // CHECK-DAG: %[[vF0:.*]] = arith.constant dense<0.000000e+00> : vector<f32> 1028 // CHECK-DAG: %[[F0:.*]] = arith.constant 0.000000e+00 : f32 1029 // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index 1030 %f0 = arith.constant 0.000000e+00 : f32 1031 1032 // CHECK: %[[init:.*]] = linalg.init_tensor [] : tensor<f32> 1033 %0 = linalg.init_tensor [] : tensor<f32> 1034 1035 // CHECK: %[[f:.*]] = vector.transfer_write %[[vF0]], %[[init]][] 1036 // CHECK-SAME: : vector<f32>, tensor<f32> 1037 %1 = linalg.fill ins(%f0 : f32) outs(%0 : tensor<f32>) -> tensor<f32> 1038 // CHECK: %[[r:.*]] = vector.transfer_read %[[A]][%[[C0]]] 1039 // CHECK-SAME: : tensor<32xf32>, vector<32xf32> 1040 // CHECK: %[[f0:.*]] = vector.extractelement %[[vF0]][] : vector<f32> 1041 // CHECK: %[[red:.*]] = vector.multi_reduction <add>, %[[r]] [0] 1042 // CHECK-SAME: : vector<32xf32> to f32 1043 // CHECK: %[[a:.*]] = arith.addf %[[red]], %[[f0]] : f32 1044 // CHECK: %[[red_v1:.*]] = vector.broadcast %[[a]] : f32 to vector<f32> 1045 // CHECK: %[[res:.*]] = vector.transfer_write %[[red_v1]], %[[f]][] 1046 // CHECK-SAME: : vector<f32>, tensor<f32> 1047 %2 = linalg.generic { 1048 indexing_maps = [affine_map<(d0) -> (d0)>, 1049 affine_map<(d0) -> ()>], 1050 iterator_types = ["reduction"]} 1051 ins(%arg0 : tensor<32xf32>) 1052 outs(%1 : tensor<f32>) { 1053 ^bb0(%a: f32, %b: f32): 1054 %3 = arith.addf %a, %b : f32 1055 linalg.yield %3 : f32 1056 } -> tensor<f32> 1057 1058 return %2 : tensor<f32> 1059} 1060