758 lines · plain
1// RUN: mlir-opt %s | mlir-opt | FileCheck %s2// RUN: mlir-opt %s --mlir-print-op-generic | mlir-opt | FileCheck %s3 4// TODO: Re-enable LLVM lowering test.5//6// Test that we can lower all the way to LLVM without crashing, don't check results here.7// DISABLED: mlir-opt %s -o=/dev/null 2>&18 9func.func @views(%arg0: index) {10 %c0 = arith.constant 0 : index11 %0 = arith.muli %arg0, %arg0 : index12 %1 = memref.alloc (%0) : memref<?xi8>13 %3 = memref.view %1[%c0][%arg0, %arg0] : memref<?xi8> to memref<?x?xf32>14 %4 = memref.view %1[%c0][%arg0, %arg0] : memref<?xi8> to memref<?x?xvector<4x4xf32>>15 memref.dealloc %1 : memref<?xi8>16 return17}18// CHECK-LABEL: func @views19// CHECK: arith.muli %{{.*}}, %{{.*}} : index20// CHECK-NEXT: memref.alloc(%{{.*}}) : memref<?xi8>21// CHECK-NEXT: memref.view %{{.*}}[%{{.*}}][%{{.*}}] :22// CHECK-SAME: memref<?xi8> to memref<?x?xf32>23// CHECK-NEXT: memref.view %{{.*}}[%{{.*}}][%{{.*}}] :24// CHECK-SAME: memref<?xi8> to memref<?x?xvector<4x4xf32>>25// CHECK-NEXT: memref.dealloc %{{.*}} : memref<?xi8>26 27// -----28 29func.func @ops(%arg0: memref<?x?xf32, strided<[?, 1], offset: ?>>,30 %arg1: memref<?xf32, strided<[1], offset: ?>>,31 %arg2: memref<?xf32, strided<[1], offset: ?>>,32 %arg3: memref<f32>) {33 linalg.matmul ins(%arg0, %arg0 : memref<?x?xf32, strided<[?, 1], offset: ?>>,34 memref<?x?xf32, strided<[?, 1], offset: ?>>)35 outs(%arg0 : memref<?x?xf32, strided<[?, 1], offset: ?>>)36 linalg.matvec ins(%arg0, %arg1: memref<?x?xf32, strided<[?, 1], offset: ?>>,37 memref<?xf32, strided<[1], offset: ?>>)38 outs(%arg2: memref<?xf32, strided<[1], offset: ?>>)39 linalg.dot ins(%arg1, %arg2: memref<?xf32, strided<[1], offset: ?>>,40 memref<?xf32, strided<[1], offset: ?>>)41 outs(%arg3: memref<f32>)42 return43}44// CHECK-LABEL: func @ops(%45// CHECK: linalg.matmul46// CHECK-SAME: ins(%{{.*}}, %{{.*}} : memref<?x?xf32, strided<[?, 1], offset: ?>>,47// CHECK-SAME: memref<?x?xf32, strided<[?, 1], offset: ?>>)48// CHECK-SAME: outs(%{{.*}} : memref<?x?xf32, strided<[?, 1], offset: ?>>)49// CHECK: linalg.matvec50// CHECK-SAME: ins(%{{.*}}, %{{.*}}: memref<?x?xf32, strided<[?, 1], offset: ?>>,51// CHECK-SAME: memref<?xf32, strided<[1], offset: ?>>)52// CHECK-SAME: outs(%{{.*}}: memref<?xf32, strided<[1], offset: ?>>)53// CHECK: linalg.dot54// CHECK-SAME: ins(%{{.*}}, %{{.*}}: memref<?xf32, strided<[1], offset: ?>>,55// CHECK-SAME: memref<?xf32, strided<[1], offset: ?>>)56// CHECK-SAME: outs(%{{.*}}: memref<f32>)57 58// -----59 60func.func @fill_view(%arg0: memref<?xf32, strided<[1], offset: ?>>, %arg1: f32) {61 linalg.fill ins(%arg1 : f32) outs(%arg0 : memref<?xf32, strided<[1], offset: ?>>)62 return63}64// CHECK-LABEL: func @fill_view(65// CHECK: %{{.*}}: memref<?xf32, strided<[1], offset: ?>>, %{{.*}}: f32) {66// CHECK: linalg.fill ins(%{{.*}} : f32) outs(%{{.*}} : memref<?xf32, strided<[1], offset: ?>>)67 68// -----69 70func.func @memref_transpose(%arg0: memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>) {71 %0 = memref.transpose %arg0 (i, j, k) -> (k, j, i) : memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>> to memref<?x?x?xf32, strided<[1, ?, ?], offset: ?>>72 return73}74// CHECK-LABEL: func @memref_transpose75// CHECK: memref.transpose %{{.*}} ([[i:.*]], [[j:.*]], [[k:.*]]) -> ([[k]], [[j]], [[i]]) :76// CHECK-SAME: memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>> to memref<?x?x?xf32, strided<[1, ?, ?], offset: ?>>77 78// -----79 80 81func.func @fill_view3(%arg0: memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>, %arg1: f32) {82 linalg.fill ins(%arg1 : f32) outs(%arg0 : memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>)83 return84}85// CHECK-LABEL: func @fill_view3(86// CHECK: %{{.*}}: memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>, %{{.*}}: f32) {87// CHECK: linalg.fill ins(%{{.*}} : f32) outs(%{{.*}} : memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>)88 89// -----90 91#accesses_0 = [92 affine_map<(i, j, k) -> (j, i)>,93 affine_map<(i, j, k) -> ()>,94 affine_map<(i, j, k) -> (i, k, i + j)>95]96 97#trait_0 = {98 indexing_maps = #accesses_0,99 iterator_types = ["parallel", "parallel", "parallel"],100 library_call = "some_external_function_name_1"101}102 103func.func @generic(%arg0: memref<?x?xvector<3x4xi4>, strided<[?, 1], offset: ?>>,104 %arg1: memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>) {105 %cst = arith.constant 0.0 : f32106 linalg.generic #trait_0107 ins(%arg0, %cst : memref<?x?xvector<3x4xi4>, strided<[?, 1], offset: ?>>, f32)108 outs(%arg1 : memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>)109 attrs = {foo = 1} {110 ^bb(%0: vector<3x4xi4>, %1: f32, %2: f32) :111 linalg.yield %1 : f32112 }113 return114}115// CHECK-LABEL: func @generic116// CHECK: linalg.generic {117// CHECK-SAME: indexing_maps = [#{{[0-9a-z]*}}, #{{[0-9a-z]*}}, #{{[0-9a-z]*}}],118// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"],119// CHECK-SAME: library_call = "some_external_function_name_1"}120// CHECK-SAME: ins({{.*}}, {{.*}} : memref<?x?xvector<3x4xi4>, strided<[?, 1], offset: ?>>, f32)121// CHECK-SAME: outs({{.*}} : memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>)122// CHECK-SAME: {foo = 1 : i64}123 124// -----125 126#map0 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>127func.func @generic_without_inputs(%arg0 : memref<?x?x?xf32>) {128 linalg.generic {indexing_maps = [#map0],129 iterator_types = ["parallel", "parallel", "parallel"]}130 outs(%arg0 : memref<?x?x?xf32>) {131 ^bb0(%arg3: f32):132 %cst = arith.constant 0.000000e+00 : f32133 linalg.yield %cst : f32134 }135 return136}137 138// CHECK-LABEL: func @generic_without_inputs139// CHECK: linalg.generic140// CHECK-NOT: ins141 142// -----143 144#accesses_1 = [145 affine_map<(i, j, k) -> (j, i)>,146 affine_map<(i, j, k) -> (i, k, i + j)>,147 affine_map<(i, j, k) -> (i, k, i + j)>148]149 150#trait_1 = {151 indexing_maps = #accesses_1,152 iterator_types = ["parallel", "parallel", "parallel"],153 library_call = "some_external_function_name_1"154}155 156func.func @generic_with_tensor_input_and_output(157 %arg0: tensor<?x?xvector<3x4xi4>>, %arg1: tensor<?x?x?xf32>)158 -> (tensor<?x?x?xf32>) {159 %0 = linalg.generic #trait_1160 ins(%arg0, %arg1 : tensor<?x?xvector<3x4xi4>>, tensor<?x?x?xf32>)161 outs(%arg1 : tensor<?x?x?xf32>)162 attrs = {foo = 1} {163 ^bb(%0: vector<3x4xi4>, %1: f32, %2: f32) :164 %f0 = arith.constant 0.0 : f32165 linalg.yield %f0 : f32166 } -> tensor<?x?x?xf32>167 return %0 : tensor<?x?x?xf32>168}169// CHECK-LABEL: func @generic_with_tensor_input_and_output170// CHECK: linalg.generic {171// CHECK-SAME: indexing_maps = [#{{.*}}, #{{.*}}], iterator_types = ["parallel", "parallel", "parallel"],172// CHECK-SAME: library_call = "some_external_function_name_1"}173// CHECK-SAME: ins({{.*}} : tensor<?x?xvector<3x4xi4>>, tensor<?x?x?xf32>)174// CHECK-SAME: outs({{.*}} : tensor<?x?x?xf32>)175// CHECK-SAME: {foo = 1 : i64}176// CHECK: -> tensor<?x?x?xf32>177// CHECK: return {{.*}} : tensor<?x?x?xf32>178 179// -----180 181func.func @generic_with_multiple_tensor_outputs(182 %arg0: tensor<?xi32>, %arg1: tensor<?xi32>, %arg2: i32)183 -> (tensor<i32>, tensor<i32>) {184 %c0 = arith.constant 0 : index185 %0 = tensor.empty() : tensor<i32>186 %1 = linalg.fill ins(%arg2 : i32) outs(%0 : tensor<i32>) -> tensor<i32>187 %2 = tensor.empty() : tensor<i32>188 %3 = linalg.fill ins(%arg2 : i32) outs(%2 : tensor<i32>) -> tensor<i32>189 %4:2 = linalg.generic {190 indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>, affine_map<(d0) -> ()>, affine_map<(d0) -> ()>],191 iterator_types = ["reduction"]}192 ins(%arg0, %arg1 : tensor<?xi32>, tensor<?xi32>)193 outs(%1, %3 : tensor<i32>, tensor<i32>) {194 ^bb0(%arg3: i32, %arg4: i32, %arg5: i32, %arg6: i32):195 %5 = arith.cmpi sge, %arg3, %arg5 : i32196 %6 = arith.select %5, %arg3, %arg5 : i32197 %7 = arith.cmpi eq, %arg3, %arg5 : i32198 %8 = arith.cmpi slt, %arg4, %arg6 : i32199 %9 = arith.select %8, %arg4, %arg6 : i32200 %10 = arith.select %5, %arg4, %arg6 : i32201 %11 = arith.select %7, %9, %10 : i32202 linalg.yield %6, %11 : i32, i32203 } -> (tensor<i32>, tensor<i32>)204 return %4#0, %4#1 : tensor<i32>, tensor<i32>205}206// CHECK-LABEL: func @generic_with_multiple_tensor_outputs207// CHECK: %{{.*}} = linalg.generic {208// CHECK-SAME: ins({{.*}} : tensor<?xi32>, tensor<?xi32>)209// CHECK-SAME: outs({{.*}} : tensor<i32>, tensor<i32>)210// CHECK: } -> (tensor<i32>, tensor<i32>)211 212// -----213 214#broadcast_access = [215 affine_map<(i, j) -> ()>,216 affine_map<(i, j) -> (i, j)>217]218 219#trait_broadcast = {220 indexing_maps = #broadcast_access,221 iterator_types = ["parallel", "parallel"],222 library_call = "some_broadcast_external_fn"223}224 225func.func @generic_op_zero_rank(%arg0: tensor<f32>, %arg1 : tensor<3x4xf32>) -> (tensor<3x4xf32>)226{227 %0 = linalg.generic #trait_broadcast228 ins(%arg0 : tensor<f32>)229 outs(%arg1 : tensor<3x4xf32>) {230 ^bb(%a: f32, %b: f32) :231 linalg.yield %a : f32232 } -> tensor<3x4xf32>233 return %0 : tensor<3x4xf32>234}235 236// -----237 238 239#accesses_3 = [240 affine_map<(i, j, k) -> (j, i)>,241 affine_map<(i, j, k) -> (i, k, i + j)>242]243 244#trait_3 = {245 indexing_maps = #accesses_3,246 iterator_types = ["parallel", "parallel", "parallel"],247 library_call = "some_external_function_name_2"248}249 250func.func @generic_region(%arg0: memref<?x?xvector<3x4xi4>, strided<[?, 1], offset: ?>>,251 %arg1: memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>) {252 linalg.generic #trait_3253 ins(%arg0 : memref<?x?xvector<3x4xi4>, strided<[?, 1], offset: ?>>)254 outs(%arg1 : memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>)255 attrs = {foo = 1} {256 ^bb(%a: vector<3x4xi4>, %b: f32) :257 %0 = linalg.index 0 : index258 %1 = linalg.index 1 : index259 %2 = linalg.index 2 : index260 linalg.yield %b : f32261 }262 return263}264// CHECK-LABEL: func @generic_region265// CHECK: linalg.generic {266// CHECK-SAME: indexing_maps = [#{{[0-9a-z]*}}, #{{[0-9a-z]*}}],267// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"],268// CHECK-SAME: library_call = "some_external_function_name_2"269// CHECK-SAME: ins({{.*}} : memref<?x?xvector<3x4xi4>, strided<[?, 1], offset: ?>>)270// CHECK-SAME: outs({{.*}} : memref<?x?x?xf32, strided<[?, ?, 1], offset: ?>>)271// CHECK-SAME: attrs = {foo = 1 : i64} {272// CHECK: ^{{.*}}(%{{.*}}: vector<3x4xi4>, %{{.*}}: f32):273// CHECK: %{{.*}} = linalg.index 0 : index274// CHECK: %{{.*}} = linalg.index 1 : index275// CHECK: %{{.*}} = linalg.index 2 : index276// CHECK: linalg.yield %{{.*}} : f32277 278// -----279 280#accessA = affine_map<(batch, m, n, k) -> (batch, m, k)>281#accessB = affine_map<(batch, m, n, k) -> (batch, k, n)>282#accessC = affine_map<(batch, m, n, k) -> (batch, m, n)>283func.func @named_ops(%a3: memref<?x?x?xf32>, %b3: memref<?x?x?xf32>, %c3: memref<?x?x?xf32>,284 %ta3: tensor<?x?x?xf32>, %tb3: tensor<?x?x?xf32>, %tc3: tensor<?x?x?xf32>)285 -> (tensor<?x?x?xf32>, tensor<?x?x?xf32>)286{287 linalg.batch_matmul ins(%a3, %b3: memref<?x?x?xf32>, memref<?x?x?xf32>)288 outs(%c3: memref<?x?x?xf32>)289 linalg.contract290 indexing_maps = [#accessA, #accessB, #accessC]291 ins(%a3, %b3: memref<?x?x?xf32>, memref<?x?x?xf32>)292 outs(%c3: memref<?x?x?xf32>)293 %res1 = linalg.batch_matmul294 ins(%ta3, %tb3: tensor<?x?x?xf32>, tensor<?x?x?xf32>)295 outs(%tc3: tensor<?x?x?xf32>)296 -> tensor<?x?x?xf32>297 %res2 = linalg.contract298 indexing_maps = [#accessA, #accessB, #accessC]299 ins(%ta3, %tb3: tensor<?x?x?xf32>, tensor<?x?x?xf32>)300 outs(%tc3: tensor<?x?x?xf32>) -> tensor<?x?x?xf32>301 return %res1, %res2 : tensor<?x?x?xf32>, tensor<?x?x?xf32>302}303// CHECK-LABEL: func @named_ops304// CHECK: linalg.batch_matmul305// CHECK: linalg.contract306// CHECK: linalg.batch_matmul307// CHECK: linalg.contract308 309// -----310 311func.func @fill_tensor(%arg0 : index, %arg1 : index, %arg2 : f32) -> tensor<?x?xf32> {312 %0 = tensor.empty(%arg0, %arg1) : tensor<?x?xf32>313 %1 = linalg.fill ins(%arg2 : f32) outs(%0 : tensor<?x?xf32>) -> tensor<?x?xf32>314 return %1 : tensor<?x?xf32>315}316// CHECK: %{{.+}} = linalg.fill ins(%{{.+}} : f32) outs(%{{.+}} : tensor<?x?xf32>) -> tensor<?x?xf32>317 318// -----319 320func.func @mixed_parallel_reduced_results(%arg0 : tensor<?x?x?xf32>,321 %arg1 : tensor<?x?xf32>, %arg2 : tensor<?x?x?xf32>, %arg3 : tensor<?x?xf32>) ->322 (tensor<?x?x?xf32>, tensor<?x?xf32>) {323 %0:2 = linalg.generic {324 indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1)>,325 affine_map<(d0, d1, d2) -> (d0, d1, d2)>, affine_map<(d0, d1, d2) -> (d0, d1)>],326 iterator_types = ["parallel", "parallel", "reduction"]}327 ins(%arg0, %arg1 : tensor<?x?x?xf32>, tensor<?x?xf32>)328 outs(%arg2, %arg3 : tensor<?x?x?xf32>, tensor<?x?xf32>) {329 ^bb0(%b0 : f32, %b1 : f32, %b2 : f32, %b3 : f32):330 %1 = arith.mulf %b0, %b1 : f32331 %2 = arith.addf %1, %b3 : f32332 linalg.yield %1, %2 : f32, f32333 } -> (tensor<?x?x?xf32>, tensor<?x?xf32>)334 return %0#0, %0#1 : tensor<?x?x?xf32>, tensor<?x?xf32>335}336// CHECK-LABEL: func @mixed_parallel_reduced_results337// CHECK: linalg.generic338 339// -----340 341func.func @map_no_inputs(%init: tensor<64xf32>) -> tensor<64xf32> {342 %add = linalg.map343 outs(%init:tensor<64xf32>)344 (%out: f32) {345 %0 = arith.constant 0.0: f32346 linalg.yield %0: f32347 }348 func.return %add : tensor<64xf32>349}350// CHECK-LABEL: func @map_no_inputs351// CHECK: linalg.map outs352// CHECK-NEXT: (%[[OUT:.*]]: f32) {353// CHECK-NEXT: arith.constant354// CHECK-NEXT: linalg.yield355// CHECK-NEXT: }356 357// -----358 359func.func @map_binary(%lhs: tensor<64xf32>, %rhs: tensor<64xf32>,360 %init: tensor<64xf32>) -> tensor<64xf32> {361 %add = linalg.map362 ins(%lhs, %rhs: tensor<64xf32>, tensor<64xf32>)363 outs(%init:tensor<64xf32>)364 (%lhs_elem: f32, %rhs_elem: f32, %out: f32) {365 %0 = arith.addf %lhs_elem, %rhs_elem: f32366 linalg.yield %0: f32367 }368 func.return %add : tensor<64xf32>369}370// CHECK-LABEL: func @map_binary371// CHECK: linalg.map { arith.addf } ins372// CHECK-SAME: outs373 374// -----375 376func.func @map_binary_memref(%lhs: memref<64xf32>, %rhs: memref<64xf32>,377 %init: memref<64xf32>) {378 linalg.map379 ins(%lhs, %rhs: memref<64xf32>, memref<64xf32>)380 outs(%init:memref<64xf32>)381 (%lhs_elem: f32, %rhs_elem: f32, %out: f32) {382 %0 = arith.addf %lhs_elem, %rhs_elem: f32383 linalg.yield %0: f32384 }385 func.return386}387// CHECK-LABEL: func @map_binary_memref388// CHECK: linalg.map389 390// -----391 392func.func @map_unary(%input: tensor<64xf32>, %init: tensor<64xf32>) -> tensor<64xf32> {393 %abs = linalg.map394 ins(%input:tensor<64xf32>)395 outs(%init:tensor<64xf32>)396 (%input_elem: f32, %out: f32) {397 %0 = math.absf %input_elem: f32398 linalg.yield %0: f32399 }400 func.return %abs : tensor<64xf32>401}402// CHECK-LABEL: func @map_unary403// CHECK: linalg.map404 405// -----406 407func.func @map_unary_memref(%input: memref<64xf32>, %init: memref<64xf32>) {408 linalg.map409 ins(%input:memref<64xf32>)410 outs(%init:memref<64xf32>)411 (%input_elem: f32, %out: f32) {412 %0 = math.absf %input_elem: f32413 linalg.yield %0: f32414 }415 func.return416}417// CHECK-LABEL: func @map_unary_memref418// CHECK: linalg.map419 420// -----421 422func.func @reduce(%input: tensor<16x32x64xf32>,423 %init: tensor<16x64xf32>) -> tensor<16x64xf32> {424 %reduce = linalg.reduce425 ins(%input:tensor<16x32x64xf32>)426 outs(%init:tensor<16x64xf32>)427 dimensions = [1]428 (%in: f32, %out: f32) {429 %0 = arith.addf %out, %in: f32430 linalg.yield %0: f32431 }432 func.return %reduce : tensor<16x64xf32>433}434// CHECK-LABEL: func @reduce435// CHECK: linalg.reduce { arith.addf } ins436// CHECK-SAME: outs437// CHECK-SAME: dimensions = [1]438 439 440// -----441 442 443func.func @reduce_not_short_form_compatible(%input: tensor<16x32x64xf32>,444 %init: tensor<16x64xf32>) -> tensor<16x64xf32> {445 %reduce = linalg.reduce446 ins(%input:tensor<16x32x64xf32>)447 outs(%init:tensor<16x64xf32>)448 dimensions = [1]449 (%in1: f32, %in2: f32) {450 %0 = arith.addf %in1, %in2: f32451 linalg.yield %in1: f32452 }453 func.return %reduce : tensor<16x64xf32>454}455 456// CHECK-LABEL: func @reduce_not_short_form_compatible457// CHECK-SAME: %[[INPUT:.*]]: tensor<16x32x64xf32>458// CHECK-SAME: %[[INIT:.*]]: tensor<16x64xf32>459// CHECK-NOT: linalg.reduce { arith.addf } ins(%[[INPUT]] : tensor<16x32x64xf32>460// CHECK: linalg.reduce ins(%[[INPUT]] : tensor<16x32x64xf32>) outs(%[[INIT]] : tensor<16x64xf32>) 461// CHECK-SAME: dimensions = [1]462// CHECK: (%[[IN1:.*]]: f32, %[[IN2:.*]]: f32) {463// CHECK-NEXT: %[[ADD_RESULT:.*]] = arith.addf %[[IN1]], %[[IN2]] : f32464// CHECK-NEXT: linalg.yield %[[IN1]] : f32465// CHECK-NEXT: }466 467// -----468 469func.func @reduce_memref(%input: memref<16x32x64xf32>,470 %init: memref<16x64xf32>) {471 linalg.reduce472 ins(%input:memref<16x32x64xf32>)473 outs(%init:memref<16x64xf32>)474 dimensions = [1]475 (%in: f32, %out: f32) {476 %0 = arith.addf %out, %in: f32477 linalg.yield %0: f32478 }479 func.return480}481// CHECK-LABEL: func @reduce482// CHECK: linalg.reduce { arith.addf } ins483// CHECK-SAME: outs484// CHECK-SAME: dimensions = [1]485 486// -----487 488func.func @variadic_reduce(%input1: tensor<16x32x64xf32>,489 %init1: tensor<16x64xf32>, %input2: tensor<16x32x64xi64>,490 %init2: tensor<16x64xi64>) -> (tensor<16x64xf32>, tensor<16x64xi64>) {491 %reduce, %reduce2 = linalg.reduce492 ins(%input1, %input2 : tensor<16x32x64xf32>, tensor<16x32x64xi64>)493 outs(%init1, %init2 : tensor<16x64xf32>, tensor<16x64xi64>)494 dimensions = [1]495 (%in1: f32, %in2: i64, %out1: f32, %out2: i64) {496 %0 = arith.addf %in1, %out1: f32497 %1 = arith.addi %in2, %out2: i64498 linalg.yield %0, %1: f32, i64499 }500 func.return %reduce, %reduce2 : tensor<16x64xf32>, tensor<16x64xi64>501}502// CHECK-LABEL: func @variadic_reduce503// CHECK: linalg.reduce504// CHECK-NOT: { arith.addf505 506// -----507 508func.func @variadic_reduce_memref(%input1: memref<16x32x64xf32>,509 %init1: memref<16x64xf32>, %input2: memref<16x32x64xi64>,510 %init2: memref<16x64xi64>) {511 linalg.reduce512 ins(%input1, %input2 : memref<16x32x64xf32>, memref<16x32x64xi64>)513 outs(%init1, %init2 : memref<16x64xf32>, memref<16x64xi64>)514 dimensions = [1]515 (%in1: f32, %in2: i64, %out1: f32, %out2: i64) {516 %0 = arith.addf %in1, %out1: f32517 %1 = arith.addi %in2, %out2: i64518 linalg.yield %0, %1: f32, i64519 }520 func.return521}522// CHECK-LABEL: func @variadic_reduce_memref523// CHECK: linalg.reduce524// CHECK-NOT: { arith.addf525 526// -----527 528func.func @transpose(%input: tensor<16x32x64xf32>,529 %init: tensor<32x64x16xf32>) -> tensor<32x64x16xf32> {530 %transpose = linalg.transpose531 ins(%input:tensor<16x32x64xf32>)532 outs(%init:tensor<32x64x16xf32>)533 permutation = [1, 2, 0]534 func.return %transpose : tensor<32x64x16xf32>535}536// CHECK-LABEL: func @transpose537// CHECK: linalg.transpose ins538// CHECK-SAME: outs539// CHECK-SAME: permutation540 541// -----542 543func.func @transpose_memref(%input: memref<16x32x64xf32>,544 %init: memref<32x64x16xf32>) {545 linalg.transpose546 ins(%input:memref<16x32x64xf32>)547 outs(%init:memref<32x64x16xf32>)548 permutation = [1, 2, 0]549 func.return550}551// CHECK-LABEL: func @transpose_memref552 553// -----554 555func.func @broadcast_static_sizes(%input: tensor<8x32xf32>,556 %init: tensor<8x16x32xf32>) -> tensor<8x16x32xf32> {557 %bcast = linalg.broadcast558 ins(%input:tensor<8x32xf32>)559 outs(%init:tensor<8x16x32xf32>)560 dimensions = [1]561 func.return %bcast : tensor<8x16x32xf32>562}563// CHECK-LABEL: func @broadcast_static_sizes564// CHECK: linalg.broadcast ins565// CHECK-SAME: outs566// CHECK-SAME: dimensions567 568// -----569 570func.func @broadcast_with_dynamic_sizes(571 %input: tensor<8x?xf32>, %init: tensor<8x16x?xf32>)572 -> tensor<8x16x?xf32> {573 %bcast = linalg.broadcast574 ins(%input:tensor<8x?xf32>)575 outs(%init:tensor<8x16x?xf32>)576 dimensions = [1]577 func.return %bcast : tensor<8x16x?xf32>578}579// CHECK-LABEL: func @broadcast_with_dynamic_sizes580// CHECK: linalg.broadcast ins581// CHECK-SAME: outs582// CHECK-SAME: dimensions583 584// -----585 586func.func @broadcast_memref(%input: memref<8x32xf32>,587 %init: memref<8x16x32xf32>) {588 linalg.broadcast589 ins(%input:memref<8x32xf32>)590 outs(%init:memref<8x16x32xf32>)591 dimensions = [1]592 func.return593}594 595// CHECK-LABEL: func @broadcast_memref596// CHECK: linalg.broadcast ins597// CHECK-SAME: outs598// CHECK-SAME: dimensions599 600// -----601 602func.func @map_arith_with_attr(%lhs: tensor<64xf32>, %rhs: tensor<64xf32>,603 %init: tensor<64xf32>) -> tensor<64xf32> {604 %add = linalg.map605 ins(%lhs, %rhs: tensor<64xf32>, tensor<64xf32>)606 outs(%init:tensor<64xf32>)607 (%lhs_elem: f32, %rhs_elem: f32, %out: f32) {608 %0 = arith.addf %lhs_elem, %rhs_elem fastmath<fast> : f32609 linalg.yield %0: f32610 }611 func.return %add : tensor<64xf32>612}613 614// CHECK-LABEL: func @map_arith_with_attr615// CHECK-NEXT: %[[MAPPED:.*]] = linalg.map616// CHECK-SAME: { arith.addf {fastmath = #arith.fastmath<fast>} }617// CHECK-SAME: ins618// CHECK-SAME: outs619// CHECK-NEXT: return %[[MAPPED]] : tensor<64xf32>620 621// -----622 623func.func @map_not_short_form_compatible(%lhs: tensor<1x32xf32>, %rhs: tensor<1x32xf32>, %init: tensor<1x32xf32>) -> tensor<1x32xf32> {624 %mapped = linalg.map ins(%lhs, %rhs : tensor<1x32xf32>, tensor<1x32xf32>) outs(%init : tensor<1x32xf32>)625 (%in_1: f32, %in_2: f32, %out: f32) {626 %1 = arith.maximumf %in_1, %in_2 : f32627 linalg.yield %in_1 : f32628 }629 func.return %mapped : tensor<1x32xf32>630}631 632// CHECK-LABEL: func @map_not_short_form_compatible633// CHECK-SAME: %[[LHS:.*]]: tensor<1x32xf32>, %[[RHS:.*]]: tensor<1x32xf32>, %[[INIT:.*]]: tensor<1x32xf32>634// CHECK-NOT: linalg.map { arith.maximumf } ins(%[[LHS]] : tensor<1x32xf32>635// CHECK: linalg.map ins(%[[LHS]], %[[RHS]] : tensor<1x32xf32>, tensor<1x32xf32>) 636// CHECK-SAME: outs(%[[INIT]] : tensor<1x32xf32>)637// CHECK-NEXT: (%[[IN1:.*]]: f32, %[[IN2:.*]]: f32, %[[OUT:.*]]: f32) {638// CHECK-NEXT: %[[MAX_RESULT:.*]] = arith.maximumf %[[IN1]], %[[IN2]] : f32639// CHECK-NEXT: linalg.yield %[[IN1]] : f32640// CHECK-NEXT: }641 642// -----643 644func.func @reduce_arith_with_attr(%input: tensor<16x32x64xf32>,645 %init: tensor<16x64xf32>) -> tensor<16x64xf32> {646 %reduce = linalg.reduce647 ins(%input:tensor<16x32x64xf32>)648 outs(%init:tensor<16x64xf32>)649 dimensions = [1]650 (%in: f32, %out: f32) {651 %0 = arith.addf %out, %in fastmath<fast> : f32652 linalg.yield %0: f32653 }654 func.return %reduce : tensor<16x64xf32>655}656// CHECK-LABEL: func @reduce_arith_with_attr657// CHECK-NEXT: %[[REDUCED:.*]] = linalg.reduce658// CHECK-SAME: { arith.addf {fastmath = #arith.fastmath<fast>} }659// CHECK-SAME: ins660// CHECK-SAME: outs661// CHECK-SAME: dimensions = [1]662// CHECK-NEXT: return %[[REDUCED]] : tensor<16x64xf32>663 664// -----665 666func.func @softmax(%arg0: tensor<2x16x32xf32>) -> tensor<2x16x32xf32> {667 %0 = tensor.empty() : tensor<2x16x32xf32>668 %1 = linalg.softmax dimension(2) ins(%arg0 : tensor<2x16x32xf32>) outs(%0: tensor<2x16x32xf32>) -> tensor<2x16x32xf32>669 return %1 : tensor<2x16x32xf32>670}671// CHECK: func.func @softmax(%[[ARG0:[a-zA-Z0-9_]+]]: tensor<2x16x32xf32>) -> tensor<2x16x32xf32> {672// CHECK: %[[D0:.+]] = tensor.empty() : tensor<2x16x32xf32>673// CHECK: %[[D1:.+]] = linalg.softmax dimension(2) ins(%[[ARG0]] : tensor<2x16x32xf32>) outs(%[[D0]] :674// CHECK-SAME: tensor<2x16x32xf32>) -> tensor<2x16x32xf32>675// CHECK: return %[[D1]] : tensor<2x16x32xf32>676// CHECK: }677 678// -----679 680func.func @winograd(%arg0: tensor<2x6x6x5xf32>, %arg1: tensor<2x3x3x5xf32>, %arg2: tensor<1xf32>, %arg3: tensor<2x4x4x2xf32>) -> tensor<2x4x4x2xf32> {681 %0 = tensor.empty() : tensor<6x6x5x2xf32>682 %1 = linalg.winograd_filter_transform fmr(F_4_3) ins(%arg1 : tensor<2x3x3x5xf32>) outs(%0 : tensor<6x6x5x2xf32>) -> tensor<6x6x5x2xf32>683 %2 = tensor.empty() : tensor<6x6x1x1x2x5xf32>684 %3 = linalg.winograd_input_transform fmr(F_4_3) ins(%arg0 : tensor<2x6x6x5xf32>) outs(%2 : tensor<6x6x1x1x2x5xf32>) -> tensor<6x6x1x1x2x5xf32>685 %collapsed = tensor.collapse_shape %1 [[0, 1], [2], [3]] : tensor<6x6x5x2xf32> into tensor<36x5x2xf32>686 %collapsed_0 = tensor.collapse_shape %3 [[0, 1], [2, 3, 4], [5]] : tensor<6x6x1x1x2x5xf32> into tensor<36x2x5xf32>687 %4 = tensor.empty() : tensor<36x2x2xf32>688 %5 = linalg.batch_matmul ins(%collapsed_0, %collapsed : tensor<36x2x5xf32>, tensor<36x5x2xf32>) outs(%4 : tensor<36x2x2xf32>) -> tensor<36x2x2xf32>689 %expanded = tensor.expand_shape %5 [[0, 1], [2, 3, 4], [5]] output_shape [6, 6, 1, 1, 2, 2] : tensor<36x2x2xf32> into tensor<6x6x1x1x2x2xf32>690 %6 = linalg.winograd_output_transform fmr(F_4_3) ins(%expanded : tensor<6x6x1x1x2x2xf32>) outs(%arg3 : tensor<2x4x4x2xf32>) -> tensor<2x4x4x2xf32>691 return %6 : tensor<2x4x4x2xf32>692}693 694// CHECK-LABEL: func @winograd695// CHECK: linalg.winograd_filter_transform fmr(F_4_3)696// CHECK: linalg.winograd_input_transform fmr(F_4_3)697// CHECK: linalg.winograd_output_transform fmr(F_4_3)698 699// -----700 701func.func @winograd_filter_dyn(%arg0: tensor<?x3x3x?xf32>, %arg1: tensor<6x6x?x?xf32>) -> tensor<6x6x?x?xf32> {702 %0 = linalg.winograd_filter_transform fmr(F_4_3) ins(%arg0 : tensor<?x3x3x?xf32>) outs(%arg1 : tensor<6x6x?x?xf32>) -> tensor<6x6x?x?xf32>703 return %0 : tensor<6x6x?x?xf32>704}705 706// CHECK-LABEL: func @winograd_filter_dyn707// CHECK: linalg.winograd_filter_transform fmr(F_4_3) ins(%arg0 : tensor<?x3x3x?xf32>) outs(%arg1 : tensor<6x6x?x?xf32>) -> tensor<6x6x?x?xf32>708 709// -----710 711func.func @winograd_input_dyn(%arg0: tensor<?x?x?x?xf32>, %arg1: tensor<6x6x?x?x?x?xf32>) -> tensor<6x6x?x?x?x?xf32> {712 %0 = linalg.winograd_input_transform fmr(F_4_3) ins(%arg0 : tensor<?x?x?x?xf32>) outs(%arg1 : tensor<6x6x?x?x?x?xf32>) -> tensor<6x6x?x?x?x?xf32>713 return %0 : tensor<6x6x?x?x?x?xf32>714}715 716// CHECK-LABEL: func @winograd_input_dyn717// CHECK: linalg.winograd_input_transform fmr(F_4_3) ins(%arg0 : tensor<?x?x?x?xf32>) outs(%arg1 : tensor<6x6x?x?x?x?xf32>) -> tensor<6x6x?x?x?x?xf32>718 719// -----720 721func.func @winograd_output_dyn(%arg0: tensor<6x6x?x?x?x?xf32>, %arg1: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> {722 %0 = linalg.winograd_output_transform fmr(F_4_3) ins(%arg0 : tensor<6x6x?x?x?x?xf32>) outs(%arg1 : tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32>723 return %0 : tensor<?x?x?x?xf32>724}725 726// CHECK-LABEL: func @winograd_output_dyn727// CHECK: linalg.winograd_output_transform fmr(F_4_3) ins(%arg0 : tensor<6x6x?x?x?x?xf32>) outs(%arg1 : tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32>728 729// -----730 731func.func @conv2d_channel_first_q(%img: tensor<100x3x224x224xi32>, %filt: tensor<64x3x5x5xi32>, %a: i32, %b: i32) -> tensor<100x64x220x220xi32> {732 %init = arith.constant dense<0> : tensor<100x64x220x220xi32>733 %1 = linalg.conv_2d_nchw_fchw_q {dilations = dense<1> : tensor<2xi64>,734 strides = dense<1> : tensor<2xi64>}735 ins(%img, %filt, %a, %b : tensor<100x3x224x224xi32>, tensor<64x3x5x5xi32>, i32, i32)736 outs(%init : tensor<100x64x220x220xi32>) -> tensor<100x64x220x220xi32>737 return %1 : tensor<100x64x220x220xi32>738}739 740// CHECK-LABEL: func @conv2d_channel_first_q(741// CHECK: %[[arg0:[a-zA-z0-9]*]]: tensor<100x3x224x224xi32>, %[[arg1:[a-zA-z0-9]*]]: tensor<64x3x5x5xi32>, %[[arg2:[a-zA-z0-9]*]]: i32, %[[arg3:[a-zA-z0-9]*]]: i32)742// CHECK: linalg.conv_2d_nchw_fchw_q {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} ins(%[[arg0]], %[[arg1]], %[[arg2]], %[[arg3]] : tensor<100x3x224x224xi32>, tensor<64x3x5x5xi32>, i32, i32) outs(%{{.*}} : tensor<100x64x220x220xi32>) -> tensor<100x64x220x220xi32>743 744// -----745 746func.func @conv2d_channel_first_q_promote(%img: tensor<100x3x224x224xi8>, %filt: tensor<64x3x5x5xi8>, %a: i8, %b: i8) -> tensor<100x64x220x220xi32> {747 %init = arith.constant dense<0> : tensor<100x64x220x220xi32>748 %1 = linalg.conv_2d_nchw_fchw_q {dilations = dense<1> : tensor<2xi64>,749 strides = dense<1> : tensor<2xi64>}750 ins(%img, %filt, %a, %b : tensor<100x3x224x224xi8>, tensor<64x3x5x5xi8>, i8, i8)751 outs(%init : tensor<100x64x220x220xi32>) -> tensor<100x64x220x220xi32>752 return %1 : tensor<100x64x220x220xi32>753}754 755// CHECK-LABEL: func @conv2d_channel_first_q_promote(756// CHECK: %[[arg0:[a-zA-z0-9]*]]: tensor<100x3x224x224xi8>, %[[arg1:[a-zA-z0-9]*]]: tensor<64x3x5x5xi8>, %[[arg2:[a-zA-z0-9]*]]: i8, %[[arg3:[a-zA-z0-9]*]]: i8)757// CHECK: linalg.conv_2d_nchw_fchw_q {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} ins(%[[arg0]], %[[arg1]], %[[arg2]], %[[arg3]] : tensor<100x3x224x224xi8>, tensor<64x3x5x5xi8>, i8, i8) outs(%{{.*}} : tensor<100x64x220x220xi32>) -> tensor<100x64x220x220xi32>758