1059 lines · plain
1// RUN: mlir-opt %s -linalg-fuse-elementwise-ops -split-input-file | FileCheck %s2 3// CHECK-DAG: [[$MAP0:#[a-zA-Z0-9_]*]] = affine_map<(d0, d1) -> (d0, d1)>4#map0 = affine_map<(d0, d1) -> (d0, d1)>5 6// CHECK-LABEL: @add_mul_fusion7func.func @add_mul_fusion(%arg0: tensor<?x?xf32>, %arg1 : tensor<?x?xf32>, %arg2 : tensor<?x?xf32>) -> tensor<?x?xf32>8{9 %c0 = arith.constant 0 : index10 %c1 = arith.constant 1 : index11 %0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>12 %1 = tensor.dim %arg0, %c1 : tensor<?x?xf32>13 %2 = tensor.empty(%0, %1) : tensor<?x?xf32>14 %3 = linalg.generic {indexing_maps = [#map0, #map0, #map0], iterator_types = ["parallel", "parallel"]}15 ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)16 outs(%2 : tensor<?x?xf32>) {17 ^bb0(%arg3: f32, %arg4: f32, %arg5: f32):18 %4 = arith.addf %arg3, %arg4 : f3219 linalg.yield %4 : f3220 } -> tensor<?x?xf32>21 // CHECK: linalg.generic {22 // CHECK-SAME: indexing_maps = {{\[}}[[$MAP0]], [[$MAP0]], [[$MAP0]], [[$MAP0]]{{\]}}23 %4 = linalg.generic {indexing_maps = [#map0, #map0, #map0], iterator_types = ["parallel", "parallel"]}24 ins(%3, %arg2 : tensor<?x?xf32>, tensor<?x?xf32>)25 outs(%2 : tensor<?x?xf32>) {26 // CHECK: ^{{[a-zA-Z0-9_]*}}27 // CHECK-SAME: [[ARG0:%[a-zA-Z0-9_]*]]28 // CHECK-SAME: [[ARG1:%[a-zA-Z0-9_]*]]29 // CHECK-SAME: [[ARG2:%[a-zA-Z0-9_]*]]30 ^bb0(%arg5: f32, %arg6: f32, %arg7: f32):31 // CHECK: [[T1:%[a-zA-Z0-9_]*]] = arith.addf [[ARG0]], [[ARG1]]32 // CHECK-NOT: linalg.yield33 // CHECK: arith.mulf [[T1]], [[ARG2]]34 // CHECK: linalg.yield35 %5 = arith.mulf %arg5, %arg6 : f3236 linalg.yield %5 : f3237 } -> tensor<?x?xf32>38 return %4 : tensor<?x?xf32>39}40 41// -----42 43// CHECK-DAG: [[$MAP0:#[a-zA-Z0-9_]*]] = affine_map<(d0, d1) -> (d0, d1)>44// CHECK-DAG: [[$MAP1:#[a-zA-Z0-9_]*]] = affine_map<(d0, d1) -> ()>45#map0 = affine_map<(d0, d1) -> (d0, d1)>46#map1 = affine_map<(d0, d1) -> ()>47 48// CHECK-LABEL: @scalar_add_mul_fusion49func.func @scalar_add_mul_fusion(%arg0: tensor<?x?xf32>, %arg1 : f32, %arg2 : f32) -> tensor<?x?xf32>50{51 %c0 = arith.constant 0 : index52 %c1 = arith.constant 1 : index53 %0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>54 %1 = tensor.dim %arg0, %c1 : tensor<?x?xf32>55 %2 = tensor.empty(%0, %1) : tensor<?x?xf32>56 %3 = linalg.generic {indexing_maps = [#map0, #map1, #map0], iterator_types = ["parallel", "parallel"]}57 ins(%arg0, %arg1 : tensor<?x?xf32>, f32)58 outs(%2 : tensor<?x?xf32>) {59 ^bb0(%arg3: f32, %arg4: f32, %arg5: f32):60 %4 = arith.addf %arg3, %arg4 : f3261 linalg.yield %4 : f3262 } -> tensor<?x?xf32>63 // CHECK: linalg.generic {64 // CHECK-SAME: indexing_maps = {{\[}}[[$MAP0]], [[$MAP1]], [[$MAP1]], [[$MAP0]]{{\]}}65 %4 = linalg.generic {indexing_maps = [#map0, #map1, #map0], iterator_types = ["parallel", "parallel"]}66 ins(%3, %arg2 : tensor<?x?xf32>, f32)67 outs(%2 : tensor<?x?xf32>) {68 // CHECK: ^{{[a-zA-Z0-9_]*}}69 // CHECK-SAME: [[ARG3:%[a-zA-Z0-9_]*]]70 // CHECK-SAME: [[ARG4:%[a-zA-Z0-9_]*]]71 // CHECK-SAME: [[ARG5:%[a-zA-Z0-9_]*]]72 ^bb0(%arg5: f32, %arg6: f32, %arg7: f32):73 // CHECK: [[T1:%[a-zA-Z0-9_]*]] = arith.addf [[ARG3]], [[ARG4]]74 // CHECK-NOT: linalg.yield75 // CHECK: arith.mulf [[T1]], [[ARG5]]76 // CHECK: linalg.yield77 %5 = arith.mulf %arg5, %arg6 : f3278 linalg.yield %5 : f3279 } -> tensor<?x?xf32>80 return %4 : tensor<?x?xf32>81}82 83// -----84 85// CHECK-DAG: [[$MAP0:#[a-zA-Z0-9_]*]] = affine_map<(d0, d1) -> (d0, d1)>86// CHECK-DAG: [[$MAP1:#[a-zA-Z0-9_]*]] = affine_map<(d0, d1) -> (d1, d0)>87#map0 = affine_map<(d0, d1) -> (d0, d1)>88#map1 = affine_map<(d0, d1) -> (d1, d0)>89 90// CHECK-LABEL: @transpose_add_mul_fusion91func.func @transpose_add_mul_fusion(%arg0: tensor<?x?xf32>, %arg1 : tensor<?x?xf32>, %arg2 : tensor<?x?xf32>) -> tensor<?x?xf32>92{93 %c0 = arith.constant 0 : index94 %c1 = arith.constant 1 : index95 %0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>96 %1 = tensor.dim %arg0, %c1 : tensor<?x?xf32>97 %2 = tensor.empty(%0, %1) : tensor<?x?xf32>98 %3 = linalg.generic {indexing_maps = [#map0, #map1, #map0], iterator_types = ["parallel", "parallel"]}99 ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)100 outs(%2 : tensor<?x?xf32>) {101 ^bb0(%arg3: f32, %arg4: f32, %arg5: f32):102 %4 = arith.addf %arg3, %arg4 : f32103 linalg.yield %4 : f32104 } -> tensor<?x?xf32>105 // CHECK: linalg.generic {106 // CHECK-SAME: indexing_maps = {{\[}}[[$MAP0]], [[$MAP1]], [[$MAP0]], [[$MAP0]]{{\]}}107 %4 = linalg.generic {indexing_maps = [#map0, #map0, #map0], iterator_types = ["parallel", "parallel"]}108 ins(%3, %arg2 : tensor<?x?xf32>, tensor<?x?xf32>)109 outs(%2 : tensor<?x?xf32>) {110 ^bb0(%arg5: f32, %arg6: f32, %arg7: f32):111 %5 = arith.mulf %arg5, %arg6 : f32112 linalg.yield %5 : f32113 } -> tensor<?x?xf32>114 return %4 : tensor<?x?xf32>115}116 117// -----118 119// CHECK-DAG: [[$MAP0:#[a-zA-Z0-9_]*]] = affine_map<(d0, d1) -> (d0, d1)>120// CHECK-DAG: [[$MAP1:#[a-zA-Z0-9_]*]] = affine_map<(d0, d1) -> (d1, d0)>121#map0 = affine_map<(d0, d1) -> (d0, d1)>122#map1 = affine_map<(d0, d1) -> (d1, d0)>123 124// CHECK-LABEL: @add_transpose_mul_fusion125func.func @add_transpose_mul_fusion(%arg0: tensor<?x?xf32>, %arg1 : tensor<?x?xf32>, %arg2 : tensor<?x?xf32>) -> tensor<?x?xf32>126{127 %c0 = arith.constant 0 : index128 %c1 = arith.constant 1 : index129 %0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>130 %1 = tensor.dim %arg0, %c1 : tensor<?x?xf32>131 %2 = tensor.empty(%0, %1) : tensor<?x?xf32>132 %3 = linalg.generic {indexing_maps = [#map0, #map1, #map0], iterator_types = ["parallel", "parallel"]}133 ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)134 outs(%2 : tensor<?x?xf32>) {135 ^bb0(%arg3: f32, %arg4: f32, %arg5: f32):136 %4 = arith.addf %arg3, %arg4 : f32137 linalg.yield %4 : f32138 } -> tensor<?x?xf32>139 // CHECK: linalg.generic {140 // CHECK-SAME: indexing_maps = {{\[}}[[$MAP1]], [[$MAP0]], [[$MAP0]], [[$MAP0]]{{\]}}141 %4 = linalg.generic {indexing_maps = [#map1, #map0, #map0], iterator_types = ["parallel", "parallel"]}142 ins(%3, %arg2 : tensor<?x?xf32>, tensor<?x?xf32>)143 outs(%2 : tensor<?x?xf32>){144 ^bb0(%arg5: f32, %arg6: f32, %arg7: f32):145 %5 = arith.mulf %arg5, %arg6 : f32146 linalg.yield %5 : f32147 } -> tensor<?x?xf32>148 return %4 : tensor<?x?xf32>149}150 151// -----152 153// CHECK-DAG: [[$MAP0:#[a-zA-Z0-9_]*]] = affine_map<(d0, d1) -> (d0, d1)>154// CHECK-DAG: [[$MAP1:#[a-zA-Z0-9_]*]] = affine_map<(d0, d1) -> (d0)>155#map0 = affine_map<(d0, d1) -> (d0, d1)>156#map1 = affine_map<(d0, d1) -> (d0)>157#map2 = affine_map<(d0) -> (d0)>158 159// CHECK-LABEL: @add_broadcast_mul_fusion160func.func @add_broadcast_mul_fusion(%arg0: tensor<?xf32>, %arg1 : tensor<?xf32>, %arg2 : tensor<?x?xf32>) -> tensor<?x?xf32>161{162 %c0 = arith.constant 0 : index163 %c1 = arith.constant 1 : index164 %0 = tensor.dim %arg0, %c0 : tensor<?xf32>165 %1 = tensor.empty(%0) : tensor<?xf32>166 %2 = linalg.generic {indexing_maps = [#map2, #map2, #map2], iterator_types = ["parallel"]}167 ins(%arg0, %arg1 : tensor<?xf32>, tensor<?xf32>)168 outs(%1 : tensor<?xf32>) {169 ^bb0(%arg3: f32, %arg4: f32, %arg5: f32):170 %3 = arith.addf %arg3, %arg4 : f32171 linalg.yield %3 : f32172 } -> tensor<?xf32>173 // CHECK: linalg.generic {174 // CHECK-SAME: indexing_maps = {{\[}}[[$MAP1]], [[$MAP1]], [[$MAP0]], [[$MAP0]]175 %3 = tensor.dim %arg2, %c1 : tensor<?x?xf32>176 %4 = tensor.empty(%0, %3) : tensor<?x?xf32>177 %5 = linalg.generic {indexing_maps = [#map1, #map0, #map0], iterator_types = ["parallel", "parallel"]}178 ins(%2, %arg2 : tensor<?xf32>, tensor<?x?xf32>)179 outs(%4 : tensor<?x?xf32>){180 ^bb0(%arg5: f32, %arg6: f32, %arg7: f32):181 %6 = arith.mulf %arg5, %arg6 : f32182 linalg.yield %6 : f32183 } -> tensor<?x?xf32>184 return %5 : tensor<?x?xf32>185}186 187// -----188 189// CHECK: #[[$MAP0:.*]] = affine_map<() -> ()>190#map0 = affine_map<() -> ()>191 192// CHECK-LABEL: @add_mul_scalar_fusion193func.func @add_mul_scalar_fusion(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<f32>194{195 %0 = tensor.empty() : tensor<f32>196 %1 = linalg.generic {indexing_maps = [#map0, #map0, #map0], iterator_types = []}197 ins(%arg0, %arg1 : tensor<f32>, tensor<f32>)198 outs(%0 : tensor<f32>) {199 ^bb0(%arg3: f32, %arg4: f32, %arg5: f32):200 %2 = arith.addf %arg3, %arg4 : f32201 linalg.yield %2 : f32202 } -> tensor<f32>203 // CHECK: linalg.generic {204 // CHECK: arith.addf205 // CHECK: arith.mulf206 %2 = linalg.generic {indexing_maps = [#map0, #map0, #map0], iterator_types = []}207 ins(%1, %arg2 : tensor<f32>, tensor<f32>)208 outs(%0 : tensor<f32>) {209 ^bb0(%arg3: f32, %arg4: f32, %arg5: f32):210 %3 = arith.mulf %arg3, %arg4 : f32211 linalg.yield %3 : f32212 } -> tensor<f32>213 214 return %2 : tensor<f32>215}216 217// -----218 219#map0 = affine_map<(d0, d1, d2) -> (d0)>220#map1 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>221func.func @generic_op_constant_fusion(%arg0 : tensor<5x?x?xf32>) -> tensor<5x?x?xf32>222{223 %c0 = arith.constant 0 : index224 %c1 = arith.constant 1 : index225 %c2 = arith.constant 2 : index226 %cst = arith.constant dense<42.0> : tensor<5xf32>227 %0 = tensor.dim %arg0, %c1 : tensor<5x?x?xf32>228 %1 = tensor.dim %arg0, %c2 : tensor<5x?x?xf32>229 %2 = tensor.empty(%0, %1) : tensor<5x?x?xf32>230 %3 = linalg.generic {231 indexing_maps = [#map0, #map1, #map1],232 iterator_types = ["parallel", "parallel", "parallel"]}233 ins(%cst, %arg0 : tensor<5xf32>, tensor<5x?x?xf32>)234 outs(%2 : tensor<5x?x?xf32>) {235 ^bb0(%arg1: f32, %arg2: f32, %arg3: f32):236 %4 = arith.mulf %arg1, %arg2 : f32237 linalg.yield %4 : f32238 } -> tensor<5x?x?xf32>239 return %3 : tensor<5x?x?xf32>240}241// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>242// CHECK-LABEL: func @generic_op_constant_fusion243// CHECK: %[[CST:.*]] = arith.constant {{.*}} : f32244// CHECK: linalg.generic245// CHECK: ^{{.+}}(%[[ARG1:[a-zA-Z0-9_]+]]: f32, %{{.+}}: f32):246// CHECK: arith.mulf %[[ARG1]], %[[CST]]247 248// -----249 250#map0 = affine_map<(d0, d1, d2) -> ()>251#map1 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>252func.func @generic_op_zero_dim_constant_fusion(%arg0 : tensor<5x?x?xf32>)253 -> tensor<5x?x?xf32>254{255 %c0 = arith.constant 0 : index256 %c1 = arith.constant 1 : index257 %c2 = arith.constant 2 : index258 %cst = arith.constant dense<42.0> : tensor<f32>259 %0 = tensor.dim %arg0, %c1 : tensor<5x?x?xf32>260 %1 = tensor.dim %arg0, %c2 : tensor<5x?x?xf32>261 %2 = tensor.empty(%0, %1) : tensor<5x?x?xf32>262 %3 = linalg.generic {263 indexing_maps = [#map0, #map1, #map1],264 iterator_types = ["parallel", "parallel", "parallel"]}265 ins(%cst, %arg0 : tensor<f32>, tensor<5x?x?xf32>)266 outs(%2 : tensor<5x?x?xf32>) {267 ^bb0(%arg1: f32, %arg2: f32, %arg3: f32):268 %4 = arith.mulf %arg1, %arg2 : f32269 linalg.yield %4 : f32270 } -> tensor<5x?x?xf32>271 return %3 : tensor<5x?x?xf32>272}273// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>274// CHECK-LABEL: func @generic_op_zero_dim_constant_fusion275// CHECK: %[[CST:.*]] = arith.constant {{.*}} : f32276// CHECK: linalg.generic277// CHECK: ^{{.*}}(%[[ARG1:[a-zA-Z0-9_]*]]: f32, %{{.*}}: f32)278// CHECK: arith.mulf %[[ARG1]], %[[CST]]279 280// -----281 282#map0 = affine_map<(d0, d1) -> (d0, d1)>283func.func @producer_indexed_consumer_fusion(%arg0: tensor<?x?xi32>,284 %arg1: tensor<?x?xi32>) -> tensor<?x?xi32> {285 %c0 = arith.constant 0 : index286 %c1 = arith.constant 1 : index287 %0 = tensor.dim %arg0, %c0 : tensor<?x?xi32>288 %1 = tensor.dim %arg0, %c1 : tensor<?x?xi32>289 %2 = tensor.empty(%0, %1) : tensor<?x?xi32>290 %3 = linalg.generic {291 indexing_maps = [#map0, #map0, #map0],292 iterator_types = ["parallel", "parallel"] }293 ins(%arg0, %arg1 : tensor<?x?xi32>, tensor<?x?xi32>)294 outs(%2 : tensor<?x?xi32>) {295 ^bb0(%arg2: i32, %arg3: i32, %arg4: i32):296 %10 = arith.addi %arg2, %arg3 : i32297 linalg.yield %10 : i32298 } -> tensor<?x?xi32>299 %4 = linalg.generic {300 indexing_maps = [#map0, #map0],301 iterator_types = ["parallel", "parallel"] }302 ins(%3 : tensor<?x?xi32>)303 outs(%2 : tensor<?x?xi32>) {304 ^bb0(%arg2: i32, %arg3: i32):305 %idx0 = linalg.index 0 : index306 %idx1 = linalg.index 1 : index307 %5 = arith.index_cast %idx0 : index to i32308 %6 = arith.index_cast %idx1 : index to i32309 %7 = arith.addi %arg2, %5 : i32310 %8 = arith.subi %7, %6 : i32311 linalg.yield %8 : i32312 } -> tensor<?x?xi32>313 return %4 : tensor<?x?xi32>314}315// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>316// CHECK-LABEL: func @producer_indexed_consumer_fusion317// CHECK: linalg.generic318// CHECK-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP0]], #[[$MAP0]]]319// CHECK: ^{{[a-zA-Z0-9_]*}}320// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]*]]: i32321// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]*]]: i32322// CHECK: %[[VAL1:.+]] = arith.addi %[[ARG0]], %[[ARG1]] : i32323// CHECK: %[[IDX0:.+]] = linalg.index 0 : index324// CHECK: %[[IDX1:.+]] = linalg.index 1 : index325// CHECK: %[[ADD_OPERAND:.+]] = arith.index_cast %[[IDX0]] : index to i32326// CHECK: %[[SUB_OPERAND:.+]] = arith.index_cast %[[IDX1]] : index to i32327// CHECK: %[[VAL2:.+]] = arith.addi %[[VAL1]], %[[ADD_OPERAND]] : i32328// CHECK: %[[VAL3:.+]] = arith.subi %[[VAL2]], %[[SUB_OPERAND]] : i32329// CHECK: linalg.yield %[[VAL3]] : i32330// CHECK-NOT: linalg.generic331 332// -----333 334#map0 = affine_map<(d0, d1) -> (d0, d1)>335func.func @indexed_producer_consumer_fusion(%arg0: tensor<?x?xi32>) -> tensor<?x?xi32> {336 %c0 = arith.constant 0 : index337 %c1 = arith.constant 1 : index338 %0 = tensor.dim %arg0, %c0 : tensor<?x?xi32>339 %1 = tensor.dim %arg0, %c1 : tensor<?x?xi32>340 %2 = tensor.empty(%0, %1) : tensor<?x?xi32>341 %3 = linalg.generic {342 indexing_maps = [#map0, #map0],343 iterator_types = ["parallel", "parallel"] }344 ins(%arg0 : tensor<?x?xi32>)345 outs(%2 : tensor<?x?xi32>) {346 ^bb0(%arg4: i32, %arg5: i32):347 %idx0 = linalg.index 0 : index348 %idx1 = linalg.index 1 : index349 %4 = arith.index_cast %idx0 : index to i32350 %5 = arith.index_cast %idx1 : index to i32351 %6 = arith.addi %arg4, %4 : i32352 %7 = arith.subi %6, %5 : i32353 linalg.yield %7 : i32354 } -> tensor<?x?xi32>355 %4 = linalg.generic {356 indexing_maps = [#map0, #map0, #map0],357 iterator_types = ["parallel", "parallel"] }358 ins(%3, %arg0 : tensor<?x?xi32>, tensor<?x?xi32>)359 outs(%2 : tensor<?x?xi32>) {360 ^bb0(%arg2: i32, %arg3: i32, %arg4: i32):361 %10 = arith.addi %arg2, %arg3 : i32362 linalg.yield %10 : i32363 } -> tensor<?x?xi32>364 return %4 : tensor<?x?xi32>365}366// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>367// CHECK-LABEL: func @indexed_producer_consumer_fusion368// CHECK: linalg.generic369// CHECK-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP0]]]370// CHECK: ^{{[a-zA-Z0-9_]*}}371// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]*]]: i32372// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]*]]: i32373// CHECK: %[[IDX0:.+]] = linalg.index 0 : index374// CHECK: %[[IDX1:.+]] = linalg.index 1 : index375// CHECK: %[[ADD_OPERAND:.+]] = arith.index_cast %[[IDX0]] : index to i32376// CHECK: %[[SUB_OPERAND:.+]] = arith.index_cast %[[IDX1]] : index to i32377// CHECK: %[[VAL1:.+]] = arith.addi %[[ARG0]], %[[ADD_OPERAND]] : i32378// CHECK: %[[VAL2:.+]] = arith.subi %[[VAL1]], %[[SUB_OPERAND]] : i32379// CHECK: %[[VAL3:.+]] = arith.addi %[[VAL2]], %[[ARG0]] : i32380// CHECK: linalg.yield %[[VAL3]] : i32381// CHECK-NOT: linalg.generic382 383// -----384 385// The indices of the first generic op are swapped after fusion.386#map0 = affine_map<(d0, d1) -> (d1, d0)>387#map1 = affine_map<(d0, d1) -> (d0, d1)>388func.func @indexed_producer_indexed_consumer_fusion(%arg0: tensor<?x?xi32>)389 -> tensor<?x?xi32> {390 %c0 = arith.constant 0 : index391 %c1 = arith.constant 1 : index392 %0 = tensor.dim %arg0, %c0 : tensor<?x?xi32>393 %1 = tensor.dim %arg0, %c1 : tensor<?x?xi32>394 %2 = tensor.empty(%0, %1) : tensor<?x?xi32>395 %3 = linalg.generic {396 indexing_maps = [#map0, #map0],397 iterator_types = ["parallel", "parallel"] }398 ins(%arg0 : tensor<?x?xi32>)399 outs(%2 : tensor<?x?xi32>) {400 ^bb0(%arg2: i32, %arg3: i32):401 %idx0 = linalg.index 0 : index402 %idx1 = linalg.index 1 : index403 %4 = arith.index_cast %idx0 : index to i32404 %5 = arith.index_cast %idx1 : index to i32405 %6 = arith.addi %arg2, %4 : i32406 %7 = arith.subi %5, %6 : i32407 linalg.yield %7 : i32408 } -> tensor<?x?xi32>409 %4= linalg.generic {410 indexing_maps = [#map1, #map1],411 iterator_types = ["parallel", "parallel"] }412 ins(%3 : tensor<?x?xi32>)413 outs(%2 : tensor<?x?xi32>) {414 ^bb0(%arg2: i32, %arg3: i32):415 %idx0 = linalg.index 0 : index416 %idx1 = linalg.index 1 : index417 %5 = arith.index_cast %idx0 : index to i32418 %6 = arith.index_cast %idx1 : index to i32419 %7 = arith.addi %arg2, %5 : i32420 %8 = arith.subi %7, %6 : i32421 linalg.yield %8 : i32422 } -> tensor<?x?xi32>423 return %4 : tensor<?x?xi32>424}425// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>426// CHECK-LABEL: func @indexed_producer_indexed_consumer_fusion427// CHECK: linalg.generic428// CHECK-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP0]]]429// CHECK: ^{{[a-zA-Z0-9_]*}}430// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]*]]: i32431// CHECK: %[[IDX0:.+]] = linalg.index 0 : index432// CHECK: %[[IDX1:.+]] = linalg.index 1 : index433// CHECK: %[[ADD_OPERAND1:.+]] = arith.index_cast %[[IDX1]] : index to i32434// CHECK: %[[SUB_OPERAND1:.+]] = arith.index_cast %[[IDX0]] : index to i32435// CHECK: %[[VAL1:.+]] = arith.addi %[[ARG0]], %[[ADD_OPERAND1]] : i32436// CHECK: %[[VAL2:.+]] = arith.subi %[[SUB_OPERAND1]], %[[VAL1]] : i32437// CHECK: %[[IDX2:.+]] = linalg.index 0 : index438// CHECK: %[[IDX3:.+]] = linalg.index 1 : index439// CHECK: %[[ADD_OPERAND2:.+]] = arith.index_cast %[[IDX2]] : index to i32440// CHECK: %[[SUB_OPERAND2:.+]] = arith.index_cast %[[IDX3]] : index to i32441// CHECK: %[[VAL3:.+]] = arith.addi %[[VAL2]], %[[ADD_OPERAND2]] : i32442// CHECK: %[[VAL4:.+]] = arith.subi %[[VAL3]], %[[SUB_OPERAND2]] : i32443// CHECK: linalg.yield %[[VAL4]] : i32444// CHECK-NOT: linalg.generic445 446// -----447 448#map1 = affine_map<(d0) -> (d0)>449#map2 = affine_map<(d0, d1) -> (d0, d1)>450#map3 = affine_map<(d0, d1) -> (d1)>451func.func @one_dim_indexed_producer_consumer_fusion(%arg0 : tensor<?xi32>,452 %arg1 : tensor<?x?xi32>) -> tensor<?x?xi32> {453 %c0 = arith.constant 0 : index454 %c1 = arith.constant 1 : index455 %d0 = tensor.dim %arg0, %c0 : tensor<?xi32>456 %0 = tensor.empty(%d0) : tensor<?xi32>457 %1 = linalg.generic458 {indexing_maps = [#map1, #map1],459 iterator_types = ["parallel"]}460 ins(%arg0 : tensor<?xi32>) outs(%0 : tensor<?xi32>) {461 ^bb0(%arg2 : i32, %arg3 : i32):462 %2 = linalg.index 0 : index463 %3 = arith.index_cast %2 : index to i32464 %4 = arith.addi %arg2, %3 : i32465 linalg.yield %4 : i32466 } -> tensor<?xi32>467 %2 = tensor.dim %arg1, %c0 : tensor<?x?xi32>468 %3 = tensor.dim %arg1, %c1 : tensor<?x?xi32>469 %4 = tensor.empty(%2, %3) : tensor<?x?xi32>470 %5 = linalg.generic471 {indexing_maps = [#map2, #map3, #map2],472 iterator_types = ["parallel", "parallel"]}473 ins(%arg1, %1 : tensor<?x?xi32>, tensor<?xi32>)474 outs(%4 : tensor<?x?xi32>) {475 ^bb0(%arg2 : i32, %arg3 : i32, %arg4: i32):476 %6 = arith.addi %arg2, %arg3 : i32477 linalg.yield %6 : i32478 } -> tensor<?x?xi32>479 return %5 : tensor<?x?xi32>480}481// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d0, d1)>482// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d1)>483// CHECK-LABEL: func @one_dim_indexed_producer_consumer_fusion484// CHECK: linalg.generic485// CHECK-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP0]]]486// CHECK: ^{{[a-zA-Z0-9_]*}}487// CHECK-SAME: (%[[ARG0:[a-zA-Z0-9_]*]]: i32, %[[ARG1:[a-zA-Z0-9_]*]]: i32488// CHECK: %[[IDX1:.+]] = linalg.index 1 : index489// CHECK: %[[VAL1:.+]] = arith.index_cast %[[IDX1]] : index to i32490// CHECK: %[[VAL2:.+]] = arith.addi %[[ARG1]], %[[VAL1]] : i32491// CHECK: %[[VAL3:.+]] = arith.addi %[[ARG0]], %[[VAL2]] : i32492// CHECK: linalg.yield %[[VAL3]] : i32493// CHECK-NOT: linalg.generic494 495// -----496 497func.func @scalar_generic_fusion498 (%arg0: tensor<5x1x1xf32>, %arg1 : tensor<i32>) -> tensor<10xf32>499{500 %c0 = arith.constant 0 : index501 %cst = arith.constant dense<1.000000e+00> : tensor<10xf32>502 %0 = tensor.empty() : tensor<f32>503 %1 = linalg.generic504 {indexing_maps = [affine_map<() -> ()>, affine_map<() -> ()>],505 iterator_types = []}506 ins(%arg1 : tensor<i32>) outs(%0 : tensor<f32>) {507 ^bb0(%arg2: i32, %arg3: f32):508 %3 = arith.index_cast %arg2 : i32 to index509 %4 = tensor.extract %arg0[%3, %c0, %c0] : tensor<5x1x1xf32>510 linalg.yield %4 : f32511 } -> tensor<f32>512 %2 = tensor.empty() : tensor<10xf32>513 %3 = linalg.generic514 {indexing_maps = [affine_map<(d0) -> ()>, affine_map<(d0) -> (d0)>,515 affine_map<(d0) -> (d0)>],516 iterator_types = ["parallel"]}517 ins(%1, %cst : tensor<f32>, tensor<10xf32>) outs(%2 : tensor<10xf32>) {518 ^bb0(%arg2: f32, %arg3: f32, %arg4: f32):519 %4 = arith.mulf %arg2, %arg3 : f32520 linalg.yield %4 : f32521 } -> tensor<10xf32>522 return %3 : tensor<10xf32>523}524// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0) -> ()>525// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0) -> (d0)>526// CHECK: func @scalar_generic_fusion527// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<5x1x1xf32>528// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<i32>529// CHECK: %[[T0:.+]] = linalg.generic530// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]]]531// CHECK-SAME: iterator_types = ["parallel"]532// CHECK-SAME: ins(%[[ARG1]] : tensor<i32>)533// CHECK: tensor.extract %[[ARG0]]534// CHECK: linalg.yield535// CHECK: return %[[T0]]536 537// -----538 539func.func @constant_fusion(%arg0 : tensor<4xf32>) -> (tensor<4xf32>) {540 %cst = arith.constant dense<1.0> : tensor<4xf32>541 %1 = tensor.empty() : tensor<4xf32>542 %2 = linalg.generic543 {indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>,544 affine_map<(d0) -> (d0)>],545 iterator_types = ["parallel"]}546 ins (%arg0, %cst : tensor<4xf32>, tensor<4xf32>)547 outs (%1 : tensor<4xf32>) {548 ^bb0(%arg1: f32, %arg2: f32, %arg3: f32):549 %3 = arith.addf %arg1, %arg2 : f32550 linalg.yield %3 : f32551 } -> tensor<4xf32>552 return %2 : tensor<4xf32>553}554 555// CHECK-DAG: #[[MAP:.+]] = affine_map<(d0) -> (d0)>556// CHECK: func @constant_fusion(%[[ARG0:.+]]: tensor<4xf32>)557// CHECK-DAG: %[[CST:.+]] = arith.constant 1.000000e+00 : f32558// CHECK-DAG: %[[T0:.+]] = tensor.empty() : tensor<4xf32>559// CHECK: %[[T1:.+]] = linalg.generic560// CHECK-SAME: indexing_maps = [#[[MAP]], #[[MAP]]]561// CHECK-SAME: ins(%[[ARG0]] : tensor<4xf32>)562// CHECK-SAME: outs(%[[T0]] : tensor<4xf32>)563// CHECK: ^{{.+}}(564// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: f32, %[[ARG2:[a-zA-Z0-9_]+]]: f32)565// CHECK: %[[T2:.+]] = arith.addf %[[ARG1]], %[[CST]]566// CHECK: linalg.yield %[[T2]]567// CHECK: return %[[T1]]568 569// -----570 571#map0 = affine_map<(d0, d1) -> (d0, d1)>572#map1 = affine_map<(d0) -> (0, d0)>573#map2 = affine_map<(d0) -> (0)>574func.func @consumer_with_reduction(%arg0: tensor<1x10xf32>,575 %arg1: tensor<1x10xf32>,576 %arg2: tensor<1xf32>) -> tensor<1xf32> {577 %init = tensor.empty() : tensor<1x10xf32>578 %0 = linalg.generic579 {indexing_maps = [#map0, #map0, #map0],580 iterator_types = ["parallel", "parallel"]}581 ins(%arg0, %arg1 : tensor<1x10xf32>, tensor<1x10xf32>)582 outs(%init : tensor<1x10xf32>) {583 ^bb0(%arg3: f32, %arg4: f32, %arg5: f32):584 %2 = arith.addf %arg3, %arg4 : f32585 linalg.yield %2 : f32586 } -> tensor<1x10xf32>587 %1 = linalg.generic588 {indexing_maps = [#map1, #map2],589 iterator_types = ["reduction"]}590 ins(%0 : tensor<1x10xf32>)591 outs(%arg2 : tensor<1xf32>) {592 ^bb0(%arg3: f32, %arg4: f32):593 %2 = arith.addf %arg3, %arg4 : f32594 linalg.yield %2 : f32595 } -> tensor<1xf32>596 return %1 : tensor<1xf32>597}598// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0) -> (0, d0)>599// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0) -> (0)>600// CHECK: func @consumer_with_reduction(%[[ARG0:.+]]: tensor<1x10xf32>, %[[ARG1:.+]]: tensor<1x10xf32>, %[[ARG2:.+]]: tensor<1xf32>)601// CHECK: %[[RES:.+]] = linalg.generic602// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP0]], #[[MAP1]]]603// CHECK-SAME: iterator_types = ["reduction"]604// CHECK-SAME: ins(%[[ARG0]], %[[ARG1]] : tensor<1x10xf32>, tensor<1x10xf32>)605// CHECK: ^{{.+}}(%[[T0:.+]]: f32, %[[T1:.+]]: f32, %[[T2:.+]]: f32)606// CHECK: %[[T3:.+]] = arith.addf %[[T0]], %[[T1]] : f32607// CHECK: %[[T4:.+]] = arith.addf %[[T3]], %[[T2]] : f32608// CHECK: linalg.yield %[[T4]]609// CHECK: return %[[RES]]610 611// -----612 613// CHECK-LABEL: func @sigmoid_dynamic_dim(614// CHECK: %[[RES:.*]] = linalg.generic615// CHECK-NOT: linalg.generic616// CHECK: return %[[RES]]617func.func @sigmoid_dynamic_dim(%0: tensor<?x1xf32>) -> tensor<?x1xf32> {618 %cp5 = arith.constant 5.000000e-01 : f32619 %c0 = arith.constant 0 : index620 %shape = shape.shape_of %0 : tensor<?x1xf32> -> tensor<?xindex>621 %extend = shape.to_extent_tensor %shape : tensor<?xindex> -> tensor<2xindex>622 %extracted = tensor.extract %extend[%c0] : tensor<2xindex>623 %init0 = tensor.empty(%extracted) : tensor<?x1xf32>624 %1 = linalg.generic {indexing_maps = [625 affine_map<(d0, d1) -> (d0, d1)>],626 iterator_types = ["parallel", "parallel"]627 }628 outs(%init0 : tensor<?x1xf32>) {629 ^bb0(%a: f32):630 linalg.yield %cp5 : f32631 } -> tensor<?x1xf32>632 %d0 = tensor.dim %0, %c0 : tensor<?x1xf32>633 %init1 = tensor.empty(%d0) : tensor<?x1xf32>634 %2 = linalg.generic {indexing_maps = [635 affine_map<(d0, d1) -> (d0, d1)>,636 affine_map<(d0, d1) -> (d0, d1)>,637 affine_map<(d0, d1) -> (d0, d1)>],638 iterator_types = ["parallel", "parallel"]639 }640 ins(%0, %1 : tensor<?x1xf32>, tensor<?x1xf32>)641 outs(%init1 : tensor<?x1xf32>) {642 ^bb0(%a: f32, %b: f32, %c: f32):643 %m = arith.mulf %a, %b : f32644 linalg.yield %m : f32645 } -> tensor<?x1xf32>646 return %2 : tensor<?x1xf32>647}648 649// -----650 651func.func private @compute1(%a: f64) -> f64652func.func private @compute2(%a: f64, %b: i32) -> i32653 654// CHECK-LABEL: func @generic_index_op2(655func.func @generic_index_op2(%arg0: tensor<1x8xf64>, %arg1: tensor<1x8xi32>) -> tensor<1x8xi32> {656 %0 = linalg.generic {657 indexing_maps = [affine_map<(i, j) -> (i, j)>],658 iterator_types = ["parallel", "parallel"]}659 outs(%arg0 : tensor<1x8xf64>) {660 ^bb0(%a: f64):661 %r = func.call @compute1(%a) : (f64) -> f64662 linalg.yield %r : f64663 } -> tensor<1x8xf64>664 665 // CHECK-NEXT: %[[R:.*]]:2 = linalg.generic666 // CHECK: bb0(%[[BBA:[0-9a-zA-Z_]*]]: f64, %[[BBB:[0-9a-zA-Z_]*]]: i32):667 // CHECK-NEXT: %[[A:.*]] = func.call @compute1(%[[BBA]]) : (f64) -> f64668 // CHECK-NEXT: %[[B:.*]] = func.call @compute2(%[[A]], %[[BBB]]) : (f64, i32) -> i32669 // CHECK-NEXT: linalg.yield %[[A]], %[[B]] : f64, i32670 // CHECK-NEXT: } -> (tensor<1x8xf64>, tensor<1x8xi32>)671 %1 = linalg.generic {672 indexing_maps = [affine_map<(i, j) -> (i, j)>, affine_map<(i, j) -> (i, j)>],673 iterator_types = ["parallel", "parallel"]}674 ins(%0 : tensor<1x8xf64>)675 outs(%arg1 : tensor<1x8xi32>) {676 ^bb0(%a: f64, %b: i32):677 %r = func.call @compute2(%a, %b) : (f64, i32) -> i32678 linalg.yield %r : i32679 } -> tensor<1x8xi32>680 681 // CHECK-NEXT: return %[[R]]#1 : tensor<1x8xi32>682 return %1 : tensor<1x8xi32>683}684 685// -----686 687// CHECK-LABEL: func @no_fuse_constant_with_reduction688func.func @no_fuse_constant_with_reduction() -> tensor<3xf32>689{690 // CHECK: %[[CONST:.+]] = arith.constant {{.+}} : tensor<3x2xf32>691 // CHECK: %[[RESULT:.+]] = linalg.generic692 // CHECK-SAME: ins(%[[CONST]] : tensor<3x2xf32>)693 // CHECK: return %[[RESULT]]694 %three = arith.constant dense<3.0> : tensor<3x2xf32>695 %init = tensor.empty() : tensor<3xf32>696 %result = linalg.generic {697 indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,698 affine_map<(d0, d1) -> (d0)>],699 iterator_types = ["parallel", "reduction"]}700 ins(%three : tensor<3x2xf32>) outs(%init : tensor<3xf32>) {701 ^bb0(%arg0 : f32, %arg1 : f32):702 %0 = arith.addf %arg0, %arg1 : f32703 linalg.yield %0 : f32704 } -> tensor<3xf32>705 return %result : tensor<3xf32>706}707 708// -----709 710#map = affine_map<(d0, d1) -> (d0, d1)>711#trait = {712 indexing_maps = [#map, #map],713 iterator_types = ["parallel", "parallel"]714}715func.func @break_outs_dependency(%arg0 : tensor<?x?xf32>) -> tensor<?x?xf32>716{717 %0 = linalg.generic #trait ins(%arg0 : tensor<?x?xf32>) outs(%arg0 : tensor<?x?xf32>) {718 ^bb0(%arg1 : f32, %arg2 : f32) :719 %1 = arith.addf %arg1, %arg1 : f32720 linalg.yield %1 : f32721 } -> tensor<?x?xf32>722 %2 = linalg.generic #trait ins(%0 : tensor<?x?xf32>) outs(%0 : tensor<?x?xf32>) {723 ^bb0(%arg1 : f32, %arg2 : f32) :724 %3 = arith.mulf %arg1, %arg1 : f32725 linalg.yield %3 : f32726 } -> tensor<?x?xf32>727 return %2 : tensor<?x?xf32>728}729// CHECK: func @break_outs_dependency(730// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?xf32>)731// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index732// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index733// CHECK-DAG: %[[D0:.+]] = tensor.dim %[[ARG0]], %[[C0]]734// CHECK-DAG: %[[D1:.+]] = tensor.dim %[[ARG0]], %[[C1]]735// CHECK-DAG: %[[INIT:.+]] = tensor.empty(%[[D0]], %[[D1]])736// CHECK: %[[GENERIC1:.+]] = linalg.generic737// CHECK-SAME: outs(%[[INIT]] : tensor<?x?xf32>)738// CHECK-DAG: %[[D0:.+]] = tensor.dim %[[GENERIC1]], %[[C0]]739// CHECK-DAG: %[[D1:.+]] = tensor.dim %[[GENERIC1]], %[[C1]]740// CHECK-DAG: %[[INIT:.+]] = tensor.empty(%[[D0]], %[[D1]])741// CHECK: %[[RESULT:.+]] = linalg.generic742// CHECK-SAME: outs(%[[INIT]] : tensor<?x?xf32>)743 744// -----745 746func.func @fuse_scalar_constant(%arg0 : tensor<?x?xf32>) -> (tensor<?x?xf32>, tensor<?x?xi32>) {747 %cst = arith.constant 4.0 : f32748 %c42 = arith.constant 42 : i32749 %c0 = arith.constant 0 : index750 %c1 = arith.constant 1 : index751 %d0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>752 %d1 = tensor.dim %arg0, %c1 : tensor<?x?xf32>753 %0 = tensor.empty(%d0, %d1) : tensor<?x?xf32>754 %1 = tensor.empty(%d0, %d1) : tensor<?x?xi32>755 %2:2 = linalg.generic {756 indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>,757 affine_map<(d0, d1) -> ()>,758 affine_map<(d0, d1) -> ()>,759 affine_map<(d0, d1) -> (d0, d1)>,760 affine_map<(d0, d1) -> (d0, d1)>],761 iterator_types = ["parallel", "parallel"]}762 ins(%arg0, %cst, %c42 : tensor<?x?xf32>, f32, i32)763 outs(%0, %1 : tensor<?x?xf32>, tensor<?x?xi32>) {764 ^bb0(%arg1 : f32, %arg2 : f32, %arg3 : i32, %arg4 : f32, %arg5 : i32) :765 %3 = arith.addf %arg1, %arg2 : f32766 linalg.yield %3, %arg3 : f32, i32767 } -> (tensor<?x?xf32>, tensor<?x?xi32>)768 return %2#0, %2#1 : tensor<?x?xf32>, tensor<?x?xi32>769}770// CHECK-LABEL: func @fuse_scalar_constant771// CHECK-DAG: %[[CST:.+]] = arith.constant 4.000000e+00 : f32772// CHECK-DAG: %[[C42:.+]] = arith.constant 42 : i32773// CHECK: linalg.generic774// CHECK-SAME: ins(%{{.+}} : tensor<?x?xf32>)775// CHECK: %[[YIELD:.+]] = arith.addf %{{.+}}, %[[CST]] : f32776// CHECK: linalg.yield %[[YIELD]], %[[C42]] : f32, i32777 778// -----779 780// Fusing the broadcast into a reduction would require to insert extra knowledge781// about the size of the reduction dimension. As long, as this is not782// implemented, we check that two linalg operations remain.783// TODO: Support this case in element-wise fusion.784 785#map0 = affine_map<(d0, d1) -> ()>786#map1 = affine_map<(d0, d1) -> (d0, d1)>787#map2 = affine_map<(d0, d1) -> (d1, d0)>788#map3 = affine_map<(d0, d1) -> (d0)>789 790// CHECK-LABEL: @no_fusion_missing_reduction_shape791// CHECK: linalg.generic792// CHECK: linalg.generic793func.func @no_fusion_missing_reduction_shape(%arg0: tensor<f32>, %arg1: index) -> tensor<?xf32> {794 %cst = arith.constant 0xFF800000 : f32795 %4 = tensor.empty(%arg1, %arg1) : tensor<?x?xf32>796 %5 = linalg.generic {797 indexing_maps = [#map0, #map1],798 iterator_types = ["parallel", "parallel"]799 } ins(%arg0 : tensor<f32>) outs(%4 : tensor<?x?xf32>) {800 ^bb0(%arg2: f32, %arg3: f32):801 linalg.yield %arg2 : f32802 } -> tensor<?x?xf32>803 %6 = tensor.empty(%arg1) : tensor<?xf32>804 %7 = linalg.fill ins(%cst : f32) outs(%6 : tensor<?xf32>) -> tensor<?xf32>805 %8 = linalg.generic {806 indexing_maps = [#map2, #map3],807 iterator_types = ["parallel", "reduction"]808 } ins(%5 : tensor<?x?xf32>) outs(%7 : tensor<?xf32>) {809 ^bb0(%arg2: f32, %arg3: f32):810 %9 = arith.maximumf %arg2, %arg3 : f32811 linalg.yield %9 : f32812 } -> tensor<?xf32>813 return %8 : tensor<?xf32>814}815 816// -----817 818func.func @fusion_different_axes(%arg0 : tensor<5000xi64>, %arg1 : tensor<5000xi32>) -> tensor<5000xi32> {819 %c1_i32 = arith.constant 1 : i32820 %0 = linalg.generic {821 indexing_maps = [affine_map<(d0) -> (d0)>],822 iterator_types = ["parallel"]}823 outs(%arg0 : tensor<5000xi64>) {824 ^bb0(%arg3: i64): // no predecessors825 %22 = linalg.index 0 : index826 %23 = arith.index_cast %22 : index to i64827 linalg.yield %23 : i64828 } -> tensor<5000xi64>829 %1 = tensor.empty() : tensor<5000xi32>830 %2 = linalg.generic {831 indexing_maps = [affine_map<(d0, d1) -> (d0)>, affine_map<(d0, d1) -> (d1)>],832 iterator_types = ["parallel", "parallel"]}833 ins(%0 : tensor<5000xi64>) outs(%1 : tensor<5000xi32>) {834 ^bb0(%arg3: i64, %arg5: i32): // no predecessors835 %22 = arith.index_cast %arg3 : i64 to index836 %23 = tensor.extract %arg1[%22] : tensor<5000xi32>837 linalg.yield %23 : i32838 } -> tensor<5000xi32>839 return %2 : tensor<5000xi32>840}841// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1) -> (d0)>842// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1) -> (d1)>843// CHECK: func @fusion_different_axes(844// CHECK-SAME: %[[ARG0:.+]]: tensor<5000xi64>845// CHECK-SAME: %[[ARG1:.+]]: tensor<5000xi32>846// CHECK-DAG: %[[INIT0:.+]] = tensor.empty() : tensor<5000xi64>847// CHECK-DAG: %[[INIT1:.+]] = tensor.empty() : tensor<5000xi32>848// CHECK: %[[RESULT:.+]]:2 = linalg.generic849// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]]]850// CHECK-SAME: outs(%[[INIT0]], %[[INIT1]] :851// CHECK-NEXT: ^bb0(852// CHECK-SAME: %[[B0:.+]]: i64853// CHECK-SAME: %[[B1:.+]]: i32854// CHECK-DAG: %[[T0:.+]] = linalg.index 0855// CHECK-DAG: %[[CAST1:.+]] = arith.index_cast %[[T0]] : index to i64856// CHECK-DAG: %[[CAST2:.+]] = arith.index_cast %[[CAST1]] : i64 to index857// CHECK: %[[EXTRACT:.+]] = tensor.extract %[[ARG1]][%[[CAST2]]]858// CHECK: linalg.yield %[[CAST1]], %[[EXTRACT]]859// CHECK: return %[[RESULT]]#1860 861// -----862 863func.func @fusion_different_axes_indexed(%arg0: tensor<2x2xi32>) -> tensor<2xi32> {864 %0 = tensor.empty() : tensor<2x2xi32>865 %1 = linalg.generic {866 indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0, d1)>],867 iterator_types = ["parallel", "parallel"]}868 ins(%arg0 : tensor<2x2xi32>) outs(%0 : tensor<2x2xi32>) {869 ^bb0(%in: i32, %out: i32):870 %2 = linalg.index 1 : index871 %3 = arith.index_cast %2 : index to i32872 linalg.yield %3 : i32873 } -> tensor<2x2xi32>874 %4 = tensor.empty() : tensor<2xi32>875 %5 = linalg.generic {876 indexing_maps = [affine_map<(d0) -> (d0, 1)>, affine_map<(d0) -> (d0)>],877 iterator_types = ["parallel"]}878 ins(%1 : tensor<2x2xi32>) outs(%4 : tensor<2xi32>) {879 ^bb0(%in: i32, %out: i32):880 linalg.yield %in : i32881 } -> tensor<2xi32>882 return %5 : tensor<2xi32>883}884 885// CHECK-DAG: #[[MAP:.+]] = affine_map<(d0) -> (d0)>886// CHECK: func @fusion_different_axes_indexed(887// CHECK-SAME: %[[ARG0:.+]]: tensor<2x2xi32>888// CHECK-DAG: %[[CST:.+]] = arith.constant 1 : i32889// CHECK-DAG: %[[INIT:.+]] = tensor.empty() : tensor<2xi32>890// CHECK: %[[RESULT:.+]] = linalg.generic891// CHECK-SAME: indexing_maps = [#[[MAP]]]892// CHECK-SAME: outs(%[[INIT]] :893// CHECK-NEXT: ^bb0(894// CHECK-SAME: %[[B0:.+]]: i32895// CHECK: linalg.yield %[[CST]] : i32896// CHECK: return %[[RESULT]]897 898// -----899 900// CHECK-LABEL: func @fold_fill_generic_basic901// CHECK-SAME: (%[[ARG0:.*]]: tensor<?xf32>) -> tensor<?xf32> {902// CHECK-NOT: linalg.fill903// CHECK: %[[GENERIC_OP:.*]] = linalg.generic904// CHECK-SAME: ins(%[[ARG0]] : tensor<?xf32>)905// CHECK-SAME: outs({{.*}} : tensor<?xf32>) {906#map0 = affine_map<(d0) -> (d0)>907func.func @fold_fill_generic_basic(%arg0: tensor<?xf32>) -> (tensor<?xf32>) {908 %c0 = arith.constant 0 : index909 %cst = arith.constant 7.0 : f32910 %0 = tensor.dim %arg0, %c0 : tensor<?xf32>911 %1 = tensor.empty(%0) : tensor<?xf32>912 %2 = linalg.fill ins(%cst : f32) outs(%1 : tensor<?xf32>) -> tensor<?xf32>913 %3 = tensor.empty(%0) : tensor<?xf32>914 %4 = linalg.generic {indexing_maps = [#map0, #map0, #map0], iterator_types=["parallel"]} ins(%arg0, %2 : tensor<?xf32>, tensor<?xf32>) outs (%3:tensor<?xf32>) {915 ^bb0(%arg1: f32, %arg2: f32, %arg3: f32):916 %5 = arith.addf %arg1, %arg2 : f32917 linalg.yield %5 : f32918 } -> tensor<?xf32>919 return %4 : tensor<?xf32>920}921 922// -----923 924// CHECK-LABEL: func @fold_fill_generic_mixedaccess925// CHECK-NOT: linalg.fill926// CHECK: %[[GENERIC_OP:.*]] = linalg.generic927// CHECK-NOT: ins928// CHECK-SAME: outs({{.*}} : tensor<?x?xf32>) {929#map0 = affine_map<(d0, d1) -> (d0, d1)>930#map1 = affine_map<(d0, d1) -> (d1, d0)>931func.func @fold_fill_generic_mixedaccess(%arg0: tensor<?x?xf32>) -> (tensor<?x?xf32>) {932 %c0 = arith.constant 0 : index933 %c1 = arith.constant 0 : index934 %cst1 = arith.constant 7.0 : f32935 %cst2 = arith.constant 6.0 : f32936 %0 = tensor.dim %arg0, %c0 : tensor<?x?xf32>937 %1 = tensor.dim %arg0, %c1 : tensor<?x?xf32>938 %2 = tensor.empty(%0, %1) : tensor<?x?xf32>939 %3 = linalg.fill ins(%cst1 : f32) outs(%2 : tensor<?x?xf32>) -> tensor<?x?xf32>940 %4 = tensor.empty(%1, %0) : tensor<?x?xf32>941 %5 = linalg.fill ins(%cst2 : f32) outs(%4 : tensor<?x?xf32>) -> tensor<?x?xf32>942 %6 = tensor.empty(%0, %1) : tensor<?x?xf32>943 %7 = linalg.generic {indexing_maps = [#map0, #map1, #map0], iterator_types=["parallel","parallel"]} ins(%3, %5 : tensor<?x?xf32>, tensor<?x?xf32>) outs (%6:tensor<?x?xf32>) {944 ^bb0(%arg1: f32, %arg2: f32, %arg3: f32):945 %8 = arith.divf %arg1, %arg2 : f32946 linalg.yield %8 : f32947 } -> tensor<?x?xf32>948 return %7 : tensor<?x?xf32>949}950 951// -----952 953#map = affine_map<() -> ()>954module {955 func.func @fuse_multi_result_producer(%arg0: tensor<f32>, %arg1: tensor<f32>, %arg2: tensor<f32>, %arg3: tensor<f32>, %arg4: tensor<f32>) -> tensor<f32> {956 %0 = tensor.empty() : tensor<f32>957 %1 = tensor.empty() : tensor<f32>958 %2:2 = linalg.generic {959 indexing_maps = [#map, #map, #map, #map, #map], iterator_types = []}960 ins(%arg0, %arg1, %arg1 : tensor<f32>, tensor<f32>, tensor<f32>) outs(%0, %1 : tensor<f32>, tensor<f32>) {961 ^bb0(%arg5: f32, %arg6: f32, %arg7: f32, %arg8: f32, %arg9: f32):962 %4 = arith.addf %arg5, %arg6 : f32963 %5 = arith.addf %4, %arg7 : f32964 linalg.yield %4, %5 : f32, f32965 } -> (tensor<f32>, tensor<f32>)966 %3 = linalg.generic {967 indexing_maps = [#map, #map, #map], iterator_types = []}968 ins(%2#1, %arg1 : tensor<f32>, tensor<f32>) outs(%arg4 : tensor<f32>) {969 ^bb0(%arg5: f32, %arg6: f32, %arg7: f32):970 %4 = arith.addf %arg5, %arg6 : f32971 %5 = arith.addf %4, %arg6 : f32972 linalg.yield %5 : f32973 } -> tensor<f32>974 return %3 : tensor<f32>975 }976}977// CHECK-LABEL: func.func @fuse_multi_result_producer978// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<f32>979// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: tensor<f32>980// CHECK: %[[INIT:.+]] = tensor.empty981// CHECK: %[[GENERIC:.+]] = linalg.generic982// CHECK-SAME: ins(%[[ARG0]], %[[ARG1]] :983// CHECK-SAME: outs(%[[INIT]] :984// CHECK-NEXT: ^bb0985// CHECK-SAME: %[[B0:[a-zA-Z0-9_]+]]: f32986// CHECK-SAME: %[[B1:[a-zA-Z0-9_]+]]: f32987// CHECK-DAG: %[[T0:.+]] = arith.addf %[[B0]], %[[B1]]988// CHECK-DAG: %[[T1:.+]] = arith.addf %[[T0]], %[[B1]]989// CHECK-DAG: %[[T2:.+]] = arith.addf %[[T1]], %[[B1]]990// CHECK-DAG: %[[T3:.+]] = arith.addf %[[T2]], %[[B1]]991// CHECK: linalg.yield %[[T3]] : f32992// CHECK: return %[[GENERIC]]993 994// -----995 996// In this test we expect the first two linalg.generic operations to be fused into one, but the third one (the matmul) to remain separate. 997// The reason is that when the pattern is applied the 1st time, the fusion of the first two operations produces a fused operation with 998// an additional result and ana dditional output indexing map that is not a permutation / not invertible. 999// The fused op will still produce also the original result (and its output indexing map), which is preserved because the new indexing map 1000// is not invertible. Thus the fused op will have 2 results, but only the 2nd one will be used by the following matmul op as an input argument.1001// When trying to apply the fusion pattern again, the matmul op won't be fused because the operand to fuse was not produced with an invertible indexing map.1002 1003#map = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>1004#map1 = affine_map<(d0, d1, d2, d3) -> (d0 * 4 + d1 * 2 + d2 + d3, 0, 0, 0)>1005#map2 = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d6)>1006#map3 = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d3, d4, d5, d6)>1007#map4 = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3, d4, d5)>1008module {1009 func.func @fuse_only_as_long_as_result_map_is_permutation(%arg0: tensor<1x1x2x1xf32>, %arg1: tensor<1x1x2x1xf32>) -> tensor<1x1x2x4xf32> {1010 %c2 = arith.constant 2 : index1011 %c1 = arith.constant 1 : index1012 %cst = arith.constant 0.000000e+00 : f321013 %c0 = arith.constant 0 : index1014 %0 = tensor.empty() : tensor<1x2x2x1xf32>1015 %1 = linalg.generic {indexing_maps = [#map], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} outs(%0 : tensor<1x2x2x1xf32>) {1016 ^bb0(%out: f32):1017 %6 = linalg.index 1 : index1018 %7 = linalg.index 2 : index1019 %8 = arith.cmpi ult, %6, %c1 : index1020 %9 = arith.cmpi ult, %7, %c2 : index1021 %10 = arith.andi %8, %9 : i11022 %11 = scf.if %10 -> (f32) {1023 %extracted = tensor.extract %arg1[%c0, %6, %7, %c0] : tensor<1x1x2x1xf32>1024 scf.yield %extracted : f321025 } else {1026 scf.yield %cst : f321027 }1028 linalg.yield %11 : f321029 } -> tensor<1x2x2x1xf32>1030 %2 = tensor.empty() : tensor<4x1x1x1xf32>1031 %3 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%1 : tensor<1x2x2x1xf32>) outs(%2 : tensor<4x1x1x1xf32>) {1032 ^bb0(%in: f32, %out: f32):1033 linalg.yield %in : f321034 } -> tensor<4x1x1x1xf32>1035 %4 = tensor.empty() : tensor<1x1x2x4xf32>1036 %expanded = tensor.expand_shape %4 [[0], [1], [2], [3, 4, 5]] output_shape [1, 1, 2, 4, 1, 1] : tensor<1x1x2x4xf32> into tensor<1x1x2x4x1x1xf32>1037 %5 = linalg.generic {indexing_maps = [#map2, #map3, #map4], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel", "parallel", "reduction"]} ins(%arg0, %3 : tensor<1x1x2x1xf32>, tensor<4x1x1x1xf32>) outs(%expanded : tensor<1x1x2x4x1x1xf32>) {1038 ^bb0(%in: f32, %in_0: f32, %out: f32):1039 %6 = arith.mulf %in, %in_0 : f321040 %7 = arith.addf %6, %out : f321041 linalg.yield %7 : f321042 } -> tensor<1x1x2x4x1x1xf32>1043 %collapsed = tensor.collapse_shape %5 [[0], [1], [2], [3, 4, 5]] : tensor<1x1x2x4x1x1xf32> into tensor<1x1x2x4xf32>1044 return %collapsed : tensor<1x1x2x4xf32>1045 }1046}1047// CHECK-DAG: #[[MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>1048// CHECK-DAG: #[[MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d0 * 4 + d1 * 2 + d2 + d3, 0, 0, 0)>1049// CHECK: func.func @fuse_only_as_long_as_result_map_is_permutation1050// CHECK-SAME: (%[[ARG0:.*]]: tensor<1x1x2x1xf32>, %[[ARG1:.*]]: tensor<1x1x2x1xf32>) -> tensor<1x1x2x4xf32> {1051// CHECK-DAG: %[[OUT0:.+]] = tensor.empty() : tensor<1x2x2x1xf32>1052// CHECK-DAG: %[[OUT1:.+]] = tensor.empty() : tensor<4x1x1x1xf32>1053// CHECK: %[[FUSED:.+]]:2 = linalg.generic {indexing_maps = [#[[MAP0]], #[[MAP1]]], iterator_types = ["parallel", "parallel", "parallel", "parallel"]}1054// CHECK-SAME: outs(%[[OUT0]], %[[OUT1]] : tensor<1x2x2x1xf32>, tensor<4x1x1x1xf32>)1055// CHECK-NOT: linalg.generic1056// CHECK: tensor.expand_shape1057// CHECK: linalg.generic {{.*}}, iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel", "parallel", "reduction"]}1058// CHECK-SAME: ins(%[[ARG0]], %[[FUSED]]#1 : tensor<1x1x2x1xf32>, tensor<4x1x1x1xf32>)1059