971 lines · plain
1// RUN: mlir-opt %s -test-linalg-elementwise-fusion-patterns=fuse-with-reshape-by-expansion -split-input-file | FileCheck %s2 3#map0 = affine_map<(d0, d1, d2) -> (d2, d0, d1)>4#map1 = affine_map<(d0, d1, d2) -> (d1, d2, d0)>5#map2 = affine_map<(d0, d1, d2) -> ()>6func.func @generic_op_reshape_producer_fusion(%arg0 : tensor<?x?x4x?xf32>,7 %arg1 : tensor<?x?x?xf32>,8 %arg2 : f32) ->9 tensor<?x?x?xf32>10{11 %0 = tensor.collapse_shape %arg0 [[0], [1, 2], [3]] :12 tensor<?x?x4x?xf32> into tensor<?x?x?xf32>13 %1 = linalg.generic {14 indexing_maps = [#map0, #map1, #map2, #map1],15 iterator_types = ["parallel", "parallel", "parallel"]}16 ins(%0, %arg1, %arg2 : tensor<?x?x?xf32>, tensor<?x?x?xf32>, f32)17 outs(%arg1 : tensor<?x?x?xf32>) {18 ^bb0(%arg3: f32, %arg4: f32, %arg5: f32, %s: f32):19 %1 = arith.mulf %arg3, %arg4 : f3220 %2 = arith.addf %1, %arg5 : f3221 linalg.yield %2 : f3222 } -> tensor<?x?x?xf32>23 return %1 : tensor<?x?x?xf32>24}25 26// CHECK-DAG: #[[MAP5:.+]] = affine_map<(d0, d1, d2, d3) -> (d3, d0, d1, d2)>27// CHECK-DAG: #[[MAP6:.+]] = affine_map<(d0, d1, d2, d3) -> (d2, d3, d0, d1)>28// CHECK-DAG: #[[MAP7:.+]] = affine_map<(d0, d1, d2, d3) -> ()>29// CHECK: func @generic_op_reshape_producer_fusion30// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?x4x?xf32>31// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x?x?xf32>32// CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]+]]: f3233// CHECK: %[[C3:.+]] = arith.constant 3 : index34// CHECK: %[[C1:.+]] = arith.constant 1 : index35// CHECK: %[[C0:.+]] = arith.constant 0 : index36// CHECK: %[[DIM:.+]] = tensor.dim %[[ARG0]], %[[C0]] : tensor<?x?x4x?xf32>37// CHECK: %[[DIM_0:.+]] = tensor.dim %[[ARG0]], %[[C1]] : tensor<?x?x4x?xf32>38// CHECK: %[[DIM_1:.+]] = tensor.dim %[[ARG0]], %[[C3]] : tensor<?x?x4x?xf32>39// CHECK: %[[T1:.+]] = tensor.expand_shape %[[ARG1]] {{\[\[}}0], [1], [2, 3]] output_shape [%[[DIM_1]], %[[DIM]], %[[DIM_0]], 4] : tensor<?x?x?xf32> into tensor<?x?x?x4xf32>40// CHECK: %[[T2:.+]] = tensor.expand_shape %[[ARG1]] {{\[\[}}0], [1], [2, 3]] output_shape [%[[DIM_1]], %[[DIM]], %[[DIM_0]], 4] : tensor<?x?x?xf32> into tensor<?x?x?x4xf32>41// CHECK: %[[T3:.+]] = linalg.generic42// CHECK-SAME: indexing_maps = [#[[MAP5]], #[[MAP6]], #[[MAP7]], #[[MAP6]]]43// CHECK-SAME: ["parallel", "parallel", "parallel", "parallel"]44// CHECK-SAME: ins(%[[ARG0]], %[[T1]], %[[ARG2]] : tensor<?x?x4x?xf32>, tensor<?x?x?x4xf32>, f32)45// CHECK-SAME: outs(%[[T2]] : tensor<?x?x?x4xf32>)46// CHECK: %[[T4:.+]] = tensor.collapse_shape %[[T3]]47// CHECK-SAME: [0], [1], [2, 3]48// CHECK-SAME: tensor<?x?x?x4xf32> into tensor<?x?x?xf32>49// CHECK: return %[[T4]]50 51// -----52 53#map0 = affine_map<(d0, d1) -> (d0, d1)>54#map1 = affine_map<(d0, d1) -> ()>55func.func @generic_op_reshape_consumer_fusion(%arg0 : tensor<?x?xf32>,56 %arg1 : tensor<?x?xf32>,57 %arg2 : f32,58 %sz0: index,59 %sz1: index) ->60 tensor<?x4x?x5xf32>61{62 %0 = linalg.generic {63 indexing_maps = [#map0, #map0, #map1, #map0],64 iterator_types = ["parallel", "parallel"]}65 ins(%arg0, %arg1, %arg2 : tensor<?x?xf32>, tensor<?x?xf32>, f32)66 outs(%arg0 : tensor<?x?xf32>) {67 ^bb0(%arg3: f32, %arg4: f32, %arg5: f32, %s: f32):68 %1 = arith.mulf %arg3, %arg4 : f3269 %2 = arith.addf %1, %arg5 : f3270 linalg.yield %2 : f3271 } -> tensor<?x?xf32>72 %1 = tensor.expand_shape %0 [[0], [1, 2, 3]] output_shape [%sz0, 4, %sz1, 5] :73 tensor<?x?xf32> into tensor<?x4x?x5xf32>74 return %1 : tensor<?x4x?x5xf32>75}76 77// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>78// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3) -> ()>79 80// CHECK: func @generic_op_reshape_consumer_fusion81// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?xf32>82// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x?xf32>83// CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]+]]: f3284// CHECK-SAME: %[[SZ0:.+]]: index, %[[SZ1:.+]]: index85// CHECK: %[[T0:.+]] = tensor.expand_shape %[[ARG0]] {{\[\[}}0], [1, 2, 3]] output_shape [%[[SZ0]], 4, %[[SZ1]], 5] : tensor<?x?xf32> into tensor<?x4x?x5xf32>86// CHECK: %[[T1:.+]] = tensor.expand_shape %[[ARG1]] {{\[\[}}0], [1, 2, 3]] output_shape [%[[SZ0]], 4, %[[SZ1]], 5] : tensor<?x?xf32> into tensor<?x4x?x5xf32>87// CHECK: %[[T2:.+]] = tensor.expand_shape %[[ARG0]] {{\[\[}}0], [1, 2, 3]] output_shape [%[[SZ0]], 4, %[[SZ1]], 5] : tensor<?x?xf32> into tensor<?x4x?x5xf32>88// CHECK: %[[T3:.+]] = linalg.generic89// CHECK-SAME: indexing_maps = [#[[MAP2]], #[[MAP2]], #[[MAP3]], #[[MAP2]]]90// CHECK-SAME: ["parallel", "parallel", "parallel", "parallel"]91// CHECK-SAME: ins(%[[T0]], %[[T1]], %[[ARG2]] : tensor<?x4x?x5xf32>, tensor<?x4x?x5xf32>, f32)92// CHECK-SAME: outs(%[[T2]] : tensor<?x4x?x5xf32>)93// CHECK: return %[[T3]] : tensor<?x4x?x5xf32>94 95 96// -----97 98func.func @reshape_as_consumer_permutation99 (%a : tensor<?x?x?xf32>, %b : tensor<?x?xf32>, %sz0: index, %sz1: index, %sz2: index)100 -> tensor<?x2x?x3x4x?xf32> {101 %c = linalg.generic {102 indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d0, d2)>,103 affine_map<(d0, d1, d2) -> (d1, d2)>,104 affine_map<(d0, d1, d2) -> (d0, d2, d1)>],105 iterator_types = ["parallel", "parallel", "parallel"]}106 ins(%a, %b : tensor<?x?x?xf32>, tensor<?x?xf32>)107 outs(%a : tensor<?x?x?xf32>) {108 ^bb0(%arg0 : f32, %arg1: f32, %s: f32):109 %1 = arith.addf %arg0, %arg1 : f32110 linalg.yield %1 : f32111 } -> tensor<?x?x?xf32>112 %d = tensor.expand_shape %c [[0, 1], [2], [3, 4, 5]] output_shape [%sz0, 2, %sz1, 3, 4, %sz2] : tensor<?x?x?xf32> into tensor<?x2x?x3x4x?xf32>113 return %d : tensor<?x2x?x3x4x?xf32>114}115// CHECK-DAG: #[[MAP8:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d2, d3, d4, d0, d1, d5)>116// CHECK-DAG: #[[MAP9:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d2, d3, d4, d5)>117// CHECK-DAG: #[[MAP10:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d5, d2, d3, d4)>118// CHECK: func @reshape_as_consumer_permutation119// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?x?xf32>120// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x?xf32>121// CHECK-SAME: %[[SZ0:.+]]: index, %[[SZ1:.+]]: index, %[[SZ2:.+]]: index122// CHECK: %[[T0:.+]] = tensor.expand_shape %[[ARG0]] {{\[\[}}0, 1, 2], [3, 4], [5]] output_shape [3, 4, %[[SZ2]], %[[SZ0]], 2, %[[SZ1]]] : tensor<?x?x?xf32> into tensor<3x4x?x?x2x?xf32>123// CHECK: %[[T1:.+]] = tensor.expand_shape %[[ARG1]] {{\[\[}}0, 1, 2], [3]] output_shape [3, 4, %[[SZ2]], %[[SZ1]]] : tensor<?x?xf32> into tensor<3x4x?x?xf32>124// CHECK: %[[T2:.+]] = tensor.expand_shape %[[ARG0]] {{\[\[}}0, 1], [2], [3, 4, 5]] output_shape [%[[SZ0]], 2, %[[SZ1]], 3, 4, %[[SZ2]]] : tensor<?x?x?xf32> into tensor<?x2x?x3x4x?xf32>125// CHECK: %[[T3:.+]] = linalg.generic126// CHECK-SAME: indexing_maps = [#[[MAP8]], #[[MAP9]], #[[MAP10]]]127// CHECK-SAME: ["parallel", "parallel", "parallel", "parallel", "parallel", "parallel"]128// CHECK-SAME: ins(%[[T0]], %[[T1]] : tensor<3x4x?x?x2x?xf32>, tensor<3x4x?x?xf32>)129// CHECK-SAME: outs(%[[T2]] : tensor<?x2x?x3x4x?xf32>)130// CHECK: return %[[T3]] : tensor<?x2x?x3x4x?xf32>131 132// -----133 134#map0 = affine_map<(d0, d1) -> (d0, d1)>135#map1 = affine_map<(d0, d1, d2) -> (d0, d1)>136#map2 = affine_map<(d0, d1, d2) -> (d2)>137 138func.func @generic_op_reshape_consumer_static(%arg0: tensor<264x4xf32>)139 -> tensor<8x33x4xf32> {140 %cst = arith.constant dense<2.000000e+00> : tensor<264x4xf32>141 %0 = tensor.empty() : tensor<264x4xf32>142 %1 = linalg.generic {143 indexing_maps = [#map0, #map0, #map0],144 iterator_types = ["parallel", "parallel"]}145 ins(%arg0, %cst : tensor<264x4xf32>, tensor<264x4xf32>)146 outs(%0 : tensor<264x4xf32>) {147 ^bb0(%arg1: f32, %arg2: f32, %s: f32):148 %2 = arith.mulf %arg1, %arg2 : f32149 linalg.yield %2 : f32150 } -> tensor<264x4xf32>151 %2 = tensor.expand_shape %1 [[0, 1], [2]] output_shape [8, 33, 4] :152 tensor<264x4xf32> into tensor<8x33x4xf32>153 return %2 : tensor<8x33x4xf32>154}155 156// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>157// CHECK: func @generic_op_reshape_consumer_static158// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<264x4xf32>159// CHECK-DAG: %[[CST:.+]] = arith.constant160// CHECK-SAME: : tensor<8x33x4xf32>161// CHECK-DAG: %[[INIT:.+]] = tensor.empty()162// CHECK: %[[T0:.+]] = tensor.expand_shape %[[ARG0]] {{\[\[}}0, 1], [2]] output_shape [8, 33, 4] : tensor<264x4xf32> into tensor<8x33x4xf32>163// CHECK: %[[T1:.+]] = tensor.expand_shape %[[INIT]] {{\[\[}}0, 1], [2]] output_shape [8, 33, 4] : tensor<264x4xf32> into tensor<8x33x4xf32>164// CHECK: %[[T2:.+]] = linalg.generic165// CHECK-SAME: indexing_maps = [#[[MAP2]], #[[MAP2]], #[[MAP2]]]166// CHECK-SAME: ["parallel", "parallel", "parallel"]167// CHECK-SAME: ins(%[[T0]], %[[CST]] :168// CHECK-SAME: outs(%[[T1]] : tensor<8x33x4xf32>)169// CHECK: return %[[T2]] : tensor<8x33x4xf32>170 171// -----172 173func.func @reshape_as_consumer_transpose174 (%a : tensor<4x210x6xf32>)175 -> tensor<2x3x4x5x6x7xf32> {176 %b = tensor.empty() : tensor<6x4x210xf32>177 %c = linalg.transpose178 ins(%a : tensor<4x210x6xf32>)179 outs(%b : tensor<6x4x210xf32>) permutation = [2, 0, 1]180 %d = tensor.expand_shape %c [[0, 1], [2], [3, 4, 5]] output_shape [2, 3, 4, 5, 6, 7] : tensor<6x4x210xf32> into tensor<2x3x4x5x6x7xf32>181 return %d : tensor<2x3x4x5x6x7xf32>182}183// CHECK: func @reshape_as_consumer_transpose184// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<4x210x6xf32>185// CHECK-DAG: %[[INIT:.+]] = tensor.empty()186// CHECK-DAG: %[[T0:.+]] = tensor.expand_shape %[[ARG0]] {{\[\[}}0], [1, 2, 3], [4, 5]] output_shape [4, 5, 6, 7, 2, 3] : tensor<4x210x6xf32> into tensor<4x5x6x7x2x3xf32>187// CHECK-DAG: %[[T1:.+]] = tensor.expand_shape %[[INIT]] {{\[\[}}0, 1], [2], [3, 4, 5]] output_shape [2, 3, 4, 5, 6, 7] : tensor<6x4x210xf32> into tensor<2x3x4x5x6x7xf32188// CHECK: %[[T2:.+]] = linalg.transpose ins(%[[T0]] : tensor<4x5x6x7x2x3xf32>)189// CHECK-SAME: outs(%[[T1]] : tensor<2x3x4x5x6x7xf32>)190// CHECK-SAME: permutation = [4, 5, 0, 1, 2, 3]191// CHECK: return %[[T2]] : tensor<2x3x4x5x6x7xf32>192 193 194// -----195 196#map0 = affine_map<(d0, d1, d2) -> (d2, d0, d1)>197#map1 = affine_map<(d0, d1, d2) -> (d1, d2, d0)>198func.func @indexed_consumer_reshape_producer_fusion(%arg0 : tensor<?x?x4x?xi32>,199 %arg1 : tensor<?x?x?xi32>) ->200 tensor<?x?x?xi32>201{202 %0 = tensor.collapse_shape %arg0 [[0], [1, 2], [3]]:203 tensor<?x?x4x?xi32> into tensor<?x?x?xi32>204 %1 = linalg.generic {205 indexing_maps = [#map0, #map1, #map1],206 iterator_types = ["parallel", "parallel", "parallel"]}207 ins(%0, %arg1 : tensor<?x?x?xi32>, tensor<?x?x?xi32>)208 outs(%0 : tensor<?x?x?xi32>) {209 ^bb0(%arg3: i32, %arg4: i32, %s: i32):210 %idx0 = linalg.index 0 : index211 %idx1 = linalg.index 1 : index212 %idx2 = linalg.index 2 : index213 %1 = arith.muli %arg3, %arg4 : i32214 %2 = arith.index_cast %idx0 : index to i32215 %3 = arith.addi %1, %2 : i32216 %4 = arith.index_cast %idx1 : index to i32217 %5 = arith.addi %3, %4 : i32218 %6 = arith.index_cast %idx2 : index to i32219 %7 = arith.addi %5, %6 : i32220 linalg.yield %7 : i32221 } -> tensor<?x?x?xi32>222 return %1 : tensor<?x?x?xi32>223}224 225// Only check the body in the indexed version of the test.226// CHECK: #[[MAP:.+]] = affine_map<()[s0, s1] -> (s0 + s1 * 4)>227// CHECK: func @indexed_consumer_reshape_producer_fusion228// CHECK: linalg.generic229// CHECK: ^{{.*}}(230// CHECK-SAME: %[[ARG3:[a-zA-Z0-9_]+]]: i32, %[[ARG4:[a-zA-Z0-9_]+]]: i32,231// CHECK-SAME: %[[ARG8:[a-zA-Z0-9_]+]]: i32)232// CHECK-DAG: %[[IDX0:.+]] = linalg.index 0 : index233// CHECK-DAG: %[[IDX1:.+]] = linalg.index 1 : index234// CHECK-DAG: %[[IDX2:.+]] = linalg.index 2 : index235// CHECK-DAG: %[[IDX3:.+]] = linalg.index 3 : index236// CHECK-DAG: %[[T3:.+]] = affine.apply #[[MAP]]()[%[[IDX1]], %[[IDX0]]]237// CHECK: %[[T4:.+]] = arith.muli %[[ARG3]], %[[ARG4]]238// CHECK: %[[T5:.+]] = arith.index_cast %[[T3]]239// CHECK: %[[T6:.+]] = arith.addi %[[T4]], %[[T5]]240// CHECK: %[[T7:.+]] = arith.index_cast %[[IDX2]]241// CHECK: %[[T8:.+]] = arith.addi %[[T6]], %[[T7]]242// CHECK: %[[T9:.+]] = arith.index_cast %[[IDX3]]243// CHECK: %[[T10:.+]] = arith.addi %[[T8]], %[[T9]]244// CHECK: linalg.yield %[[T10]]245 246// -----247 248#map0 = affine_map<(d0, d1) -> (d0, d1)>249func.func @indexed_producer_reshape_consumer_fusion(%arg0 : tensor<?x?xi32>,250 %arg1 : tensor<?x?xi32>, 251 %sz0: index, %sz1: index) ->252 tensor<?x?x4x5xi32>253{254 %0 = linalg.generic {255 indexing_maps = [#map0, #map0, #map0],256 iterator_types = ["parallel", "parallel"]}257 ins(%arg0, %arg1 : tensor<?x?xi32>, tensor<?x?xi32>)258 outs(%arg0 : tensor<?x?xi32>) {259 ^bb0(%arg3: i32, %arg4: i32, %s: i32):260 %idx0 = linalg.index 0 : index261 %idx1 = linalg.index 1 : index262 %1 = arith.muli %arg3, %arg4 : i32263 %2 = arith.index_cast %idx0 : index to i32264 %3 = arith.addi %1, %2 : i32265 %4 = arith.index_cast %idx1 : index to i32266 %5 = arith.addi %3, %4 : i32267 linalg.yield %5 : i32268 } -> tensor<?x?xi32>269 %1 = tensor.expand_shape %0 [[0], [1, 2, 3]] output_shape [%sz0, %sz1, 4, 5] :270 tensor<?x?xi32> into tensor<?x?x4x5xi32>271 return %1 : tensor<?x?x4x5xi32>272}273 274// Only check the body in the indexed version of the test.275// CHECK-DAG: #[[MAP1:.+]] = affine_map<()[s0, s1, s2] -> (s0 * 5 + s1 * 20 + s2)>276// CHECK: func @indexed_producer_reshape_consumer_fusion277// CHECK: linalg.generic278// CHECK: ^{{.*}}(279// CHECK-SAME: %[[ARG3:[a-zA-Z0-9_]+]]: i32, %[[ARG4:[a-zA-Z0-9_]+]]: i32,280// CHECK-SAME: %[[ARG5:[a-zA-Z0-9_]+]]: i32)281// CHECK-DAG: %[[IDX0:.+]] = linalg.index 0 : index282// CHECK-DAG: %[[IDX1:.+]] = linalg.index 1 : index283// CHECK-DAG: %[[IDX2:.+]] = linalg.index 2 : index284// CHECK-DAG: %[[IDX3:.+]] = linalg.index 3 : index285// CHECK: %[[T1:.+]] = affine.apply #[[MAP1]]()[%[[IDX2]], %[[IDX1]], %[[IDX3]]]286// CHECK: %[[T4:.+]] = arith.muli %[[ARG3]], %[[ARG4]]287// CHECK: %[[T5:.+]] = arith.index_cast %[[IDX0]]288// CHECK: %[[T6:.+]] = arith.addi %[[T4]], %[[T5]]289// CHECK: %[[T7:.+]] = arith.index_cast %[[T1]]290// CHECK: %[[T8:.+]] = arith.addi %[[T6]], %[[T7]]291// CHECK: linalg.yield %[[T8]]292 293// -----294 295func.func @reshape_as_consumer_permutation296 (%a : tensor<210x6x4xi32>, %b : tensor<210x4xi32>)297 -> tensor<2x3x4x5x6x7xi32> {298 %shape = tensor.empty() : tensor<6x4x210xi32>299 %c = linalg.generic {300 indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d0, d2)>,301 affine_map<(d0, d1, d2) -> (d1, d2)>,302 affine_map<(d0, d1, d2) -> (d0, d2, d1)>],303 iterator_types = ["parallel", "parallel", "parallel"]}304 ins(%a, %b : tensor<210x6x4xi32>, tensor<210x4xi32>)305 outs(%shape : tensor<6x4x210xi32>) {306 ^bb0(%arg3 : i32, %arg4: i32, %s: i32):307 %idx0 = linalg.index 0 : index308 %idx1 = linalg.index 1 : index309 %idx2 = linalg.index 2 : index310 %1 = arith.addi %arg3, %arg4 : i32311 %2 = arith.index_cast %idx0 : index to i32312 %3 = arith.addi %1, %2 : i32313 %4 = arith.index_cast %idx1 : index to i32314 %5 = arith.addi %3, %4 : i32315 %6 = arith.index_cast %idx2 : index to i32316 %7 = arith.addi %5, %6 : i32317 linalg.yield %7 : i32318 } -> tensor<6x4x210xi32>319 %d = tensor.expand_shape %c [[0, 1], [2], [3, 4, 5]] output_shape [2, 3, 4, 5, 6, 7] : tensor<6x4x210xi32> into tensor<2x3x4x5x6x7xi32>320 return %d : tensor<2x3x4x5x6x7xi32>321}322 323// -----324 325// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d2, d3, d4, d0, d1, d5)>326// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d2, d3, d4, d5)>327// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d5, d2, d3, d4)>328// CHECK-DAG: #[[MAP3:.+]] = affine_map<()[s0, s1] -> (s0 + s1 * 3)>329// CHECK-DAG: #[[MAP4:.+]] = affine_map<()[s0, s1, s2] -> (s0 * 7 + s1 * 42 + s2)>330// CHECK: func @reshape_as_consumer_permutation331// CHECK-SAME: %[[ARG0:.+]]: tensor<210x6x4xi32>332// CHECK-SAME: %[[ARG1:.+]]: tensor<210x4xi32>333// CHECK-DAG: %[[INIT:.+]] = tensor.empty()334// CHECK: %[[T1:.+]] = tensor.expand_shape %[[ARG0]] {{\[\[}}0, 1, 2], [3, 4], [5]] output_shape [5, 6, 7, 2, 3, 4] : tensor<210x6x4xi32> into tensor<5x6x7x2x3x4xi32>335// CHECK: %[[T2:.+]] = tensor.expand_shape %[[ARG1]] {{\[\[}}0, 1, 2], [3]] output_shape [5, 6, 7, 4] : tensor<210x4xi32> into tensor<5x6x7x4xi32>336// CHECK: %[[T3:.+]] = tensor.expand_shape %[[INIT]] {{\[\[}}0, 1], [2], [3, 4, 5]] output_shape [2, 3, 4, 5, 6, 7] : tensor<6x4x210xi32> into tensor<2x3x4x5x6x7xi32>337// CHECK: %[[T4:.+]] = linalg.generic338// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]]339// CHECK-SAME: ins(%[[T1]], %[[T2]] : tensor<5x6x7x2x3x4xi32>, tensor<5x6x7x4xi32>)340// CHECK-SAME: outs(%[[T3]] : tensor<2x3x4x5x6x7xi32>)341// CHECK: ^{{.+}}(342// CHECK-SAME: %[[ARG8:[a-zA-Z0-9_]+]]: i32, %[[ARG9:[a-zA-Z0-9_]+]]: i32,343// CHECK-SAME: %[[ARG10:[a-zA-Z0-9_]+]]: i32)344// CHECK-DAG: %[[IDX0:.+]] = linalg.index 0 : index345// CHECK-DAG: %[[IDX1:.+]] = linalg.index 1 : index346// CHECK-DAG: %[[IDX2:.+]] = linalg.index 2 : index347// CHECK-DAG: %[[IDX3:.+]] = linalg.index 3 : index348// CHECK-DAG: %[[IDX4:.+]] = linalg.index 4 : index349// CHECK-DAG: %[[IDX5:.+]] = linalg.index 5 : index350// CHECK-DAG: %[[T5:.+]] = affine.apply #[[MAP3]]()[%[[IDX1]], %[[IDX0]]]351// CHECK-DAG: %[[T6:.+]] = affine.apply #[[MAP4]]()[%[[IDX3]], %[[IDX2]], %[[IDX4]]]352// CHECK-DAG: %[[T8:.+]] = arith.addi %[[ARG8]], %[[ARG9]]353// CHECK: %[[T9:.+]] = arith.index_cast %[[T5]]354// CHECK: %[[T10:.+]] = arith.addi %[[T8]], %[[T9]]355// CHECK: %[[T11:.+]] = arith.index_cast %[[T6]]356// CHECK: %[[T12:.+]] = arith.addi %[[T10]], %[[T11]]357// CHECK: %[[T13:.+]] = arith.index_cast %[[IDX5]]358// CHECK: %[[T14:.+]] = arith.addi %[[T12]], %[[T13]]359 360// -----361 362func.func @reshape_as_producer_projected_permutation(363 %arg0 : tensor<33x8x?xi32>, %shape : tensor<264x?x4xi32>) -> tensor<264x?x4xi32>364{365 %0 = tensor.collapse_shape %arg0 [[0, 1], [2]]366 : tensor<33x8x?xi32> into tensor<264x?xi32>367 %1 = linalg.generic368 {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1)>,369 affine_map<(d0, d1, d2) -> (d0, d1, d2)>],370 iterator_types = ["parallel", "parallel", "parallel"]}371 ins(%0 : tensor<264x?xi32>)372 outs(%shape : tensor<264x?x4xi32>) {373 ^bb0(%arg1: i32, %s: i32):374 %idx0 = linalg.index 0 : index375 %idx1 = linalg.index 1 : index376 %idx2 = linalg.index 2 : index377 %2 = arith.index_cast %idx0 : index to i32378 %3 = arith.addi %arg1, %2 : i32379 %4 = arith.index_cast %idx1 : index to i32380 %5 = arith.addi %3, %4 : i32381 %6 = arith.index_cast %idx2 : index to i32382 %7 = arith.addi %5, %6 : i32383 linalg.yield %7 : i32384 } -> tensor<264x?x4xi32>385 return %1 : tensor<264x?x4xi32>386}387 388// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>389// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>390// CHECK-DAG: #[[MAP2:.+]] = affine_map<()[s0, s1] -> (s0 + s1 * 8)>391// CHECK: @reshape_as_producer_projected_permutation392// CHECK-SAME: %[[ARG0:.+]]: tensor<33x8x?xi32>393// CHECK: %[[RES:.+]] = linalg.generic394// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]]]395// CHECK-SAME: ins(%[[ARG0]] : tensor<33x8x?xi32>)396// CHECK: ^{{.+}}(397// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: i32,398// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: i32)399// CHECK-DAG: %[[IDX0:.+]] = linalg.index 0 : index400// CHECK-DAG: %[[IDX1:.+]] = linalg.index 1 : index401// CHECK-DAG: %[[IDX2:.+]] = linalg.index 2 : index402// CHECK-DAG: %[[IDX3:.+]] = linalg.index 3 : index403// CHECK-DAG: %[[T0:.+]] = affine.apply #[[MAP2]]()[%[[IDX1]], %[[IDX0]]]404// CHECK: %[[T1:.+]] = arith.index_cast %[[T0]] : index to i32405// CHECK: %[[T2:.+]] = arith.addi %[[ARG1]], %[[T1]] : i32406// CHECK: %[[T3:.+]] = arith.index_cast %[[IDX2]] : index to i32407// CHECK: %[[T4:.+]] = arith.addi %[[T2]], %[[T3]] : i32408// CHECK: %[[T5:.+]] = arith.index_cast %[[IDX3]] : index to i32409// CHECK: %[[T6:.+]] = arith.addi %[[T4]], %[[T5]] : i32410// CHECK: linalg.yield %[[T6]] : i32411// CHECK: %[[RES2:.+]] = tensor.collapse_shape %[[RES]]412// CHECK-SAME: [0, 1], [2], [3]413// CHECK-SAME: : tensor<33x8x?x4xi32> into tensor<264x?x4xi32>414// CHECK: return %[[RES2]] : tensor<264x?x4xi32>415 416// -----417 418#map0 = affine_map<(d0, d1) -> (d0, d1)>419#map1 = affine_map<(d0, d1) -> (d1, d0)>420func.func @generic_op_reshape_consumer_fusion_projected(%arg0 : tensor<?x?xf32>,421 %arg1 : tensor<?x?xf32>,422 %sz0: index, %sz1: index) ->423 tensor<?x?x4x5xf32>424{425 %0 = linalg.generic {426 indexing_maps = [#map0, #map0, #map1],427 iterator_types = ["parallel", "parallel"]}428 ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)429 outs(%arg0 : tensor<?x?xf32>) {430 ^bb0(%arg3: f32, %arg4: f32, %s: f32):431 %1 = arith.mulf %arg3, %arg4 : f32432 linalg.yield %1 : f32433 } -> tensor<?x?xf32>434 %1 = tensor.expand_shape %0 [[0], [1, 2, 3]] output_shape [%sz0, %sz1, 4, 5] :435 tensor<?x?xf32> into tensor<?x?x4x5xf32>436 return %1 : tensor<?x?x4x5xf32>437}438 439// CHECK-DAG: #[[MAP4:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>440// CHECK-DAG: #[[MAP5:.+]] = affine_map<(d0, d1, d2, d3) -> (d3, d0, d1, d2)>441// CHECK: func @generic_op_reshape_consumer_fusion_projected442// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?xf32>443// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x?xf32>444// CHECK-SAME: %[[SZ0:.+]]: index, %[[SZ1:.+]]: index445// CHECK: %[[T0:.+]] = tensor.expand_shape %[[ARG0]] {{\[\[}}0, 1, 2], [3]] output_shape [%[[SZ1]], 4, 5, %[[SZ0]]] : tensor<?x?xf32> into tensor<?x4x5x?xf32>446// CHECK: %[[T1:.+]] = tensor.expand_shape %[[ARG1]] {{\[\[}}0, 1, 2], [3]] output_shape [%[[SZ1]], 4, 5, %[[SZ0]]] : tensor<?x?xf32> into tensor<?x4x5x?xf32>447// CHECK: %[[T2:.+]] = tensor.expand_shape %[[ARG0]] {{\[\[}}0], [1, 2, 3]] output_shape [%[[SZ0]], %[[SZ1]], 4, 5] : tensor<?x?xf32> into tensor<?x?x4x5xf32>448// CHECK: %[[T3:.+]] = linalg.generic449// CHECK-SAME: indexing_maps = [#[[MAP4]], #[[MAP4]], #[[MAP5]]]450// CHECK-SAME: ["parallel", "parallel", "parallel", "parallel"]451// CHECK-SAME: ins(%[[T0]], %[[T1]] : tensor<?x4x5x?xf32>, tensor<?x4x5x?xf32>)452// CHECK-SAME: outs(%[[T2]] : tensor<?x?x4x5xf32>)453// CHECK: return %[[T3]] : tensor<?x?x4x5xf32>454 455// -----456 457func.func @fuse_collapse_reduction(%arg0: tensor<10x10x20xf32>) -> tensor<100xf32> {458 %c0 = arith.constant 0 : index459 %c_0 = arith.constant 0.0 : f32460 %0 = tensor.collapse_shape %arg0 [[0, 1], [2]] : tensor<10x10x20xf32> into tensor<100x20xf32>461 %2 = tensor.empty() : tensor<100xf32>462 %3 = linalg.fill ins(%c_0 : f32) outs(%2 : tensor<100xf32>) -> tensor<100xf32>463 %4 = linalg.generic {464 indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0)>],465 iterator_types = ["parallel", "reduction"]}466 ins(%0 : tensor<100x20xf32>) outs(%3 : tensor<100xf32>) {467 ^bb0(%arg1 : f32, %arg2: f32):468 %4 = arith.addf %arg1, %arg2 : f32469 linalg.yield %4 : f32470 } -> tensor<100xf32>471 return %4 : tensor<100xf32>472}473 474// CHECK: func @fuse_collapse_reduction475// CHECK-SAME: %[[ARG0:.+]]: tensor<10x10x20xf32>476// CHECK: %[[GENERIC:.+]] = linalg.generic477// CHECK-SAME: ins(%[[ARG0]] : tensor<10x10x20xf32>)478// CHECK: %[[COLLAPSE:.+]] = tensor.collapse_shape %[[GENERIC]]479// CHECK: return %[[COLLAPSE]]480 481// -----482 483func.func @fuse_dynamic_dims(%arg0: tensor<?x?xf32>) -> tensor<?xf32> {484 %c0 = arith.constant 0 : index485 %0 = tensor.collapse_shape %arg0 [[0, 1]] : tensor<?x?xf32> into tensor<?xf32>486 %1 = tensor.dim %0, %c0 : tensor<?xf32>487 %2 = tensor.empty(%1) : tensor<?xf32>488 %3 = linalg.generic {489 indexing_maps = [affine_map<(d0) -> (d0)>, affine_map<(d0) -> (d0)>],490 iterator_types = ["parallel"]}491 ins(%0 : tensor<?xf32>) outs(%2 : tensor<?xf32>) {492 ^bb0(%arg1 : f32, %arg2: f32):493 %4 = arith.addf %arg1, %arg1 : f32494 linalg.yield %4 : f32495 } -> tensor<?xf32>496 return %3 : tensor<?xf32>497}498 499// CHECK: func @fuse_dynamic_dims500// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?xf32>501// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index502// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index503// CHECK: %[[RESHAPE:.+]] = tensor.collapse_shape %[[ARG0]]504// CHECK: %[[EMPTY:.+]] = tensor.empty505// CHECK-DAG: %[[D0:.+]] = tensor.dim %[[ARG0]], %[[C0]]506// CHECK-DAG: %[[D1:.+]] = tensor.dim %[[ARG0]], %[[C1]]507// CHECK: %[[EXPAND_SHAPE:.+]] = tensor.expand_shape %[[EMPTY]] {{\[}}[0, 1]{{\]}}508// CHECK-SAME: output_shape [%[[D0]], %[[D1]]]509// CHECK: %[[GENERIC:.+]] = linalg.generic510// CHECK-SAME: ins(%[[ARG0]] :511// CHECK-SAME: outs(%[[EXPAND_SHAPE]] :512// CHECK: %[[COLLAPSE:.+]] = tensor.collapse_shape %[[GENERIC]] {{\[}}[0, 1]{{\]}}513// CHECK: return %[[COLLAPSE]]514 515// -----516 517func.func @reshape_as_consumer_permutation_with_multiple_results518 (%a : tensor<?x?x?xf32>, %b : tensor<?x?xf32>, %sz0: index, 519 %sz1: index, %sz2: index, %sz3: index, %sz4: index)520 -> (tensor<?x2x?x3x4x?xf32>, tensor<?x?x2x3x4x?xf32>) {521 %c:2 = linalg.generic {522 indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d0, d2)>,523 affine_map<(d0, d1, d2) -> (d1, d2)>,524 affine_map<(d0, d1, d2) -> (d0, d2, d1)>,525 affine_map<(d0, d1, d2) -> (d2, d0, d1)>],526 iterator_types = ["parallel", "parallel", "parallel"]}527 ins(%a, %b : tensor<?x?x?xf32>, tensor<?x?xf32>)528 outs(%a, %a : tensor<?x?x?xf32>, tensor<?x?x?xf32>) {529 ^bb0(%arg0 : f32, %arg1: f32, %s: f32, %t : f32):530 %1 = arith.addf %arg0, %arg1 : f32531 linalg.yield %1, %1 : f32, f32532 } -> (tensor<?x?x?xf32>, tensor<?x?x?xf32>)533 %d = tensor.expand_shape %c#0 [[0, 1], [2], [3, 4, 5]] output_shape [%sz0, 2, %sz1, 3, 4, %sz2] : tensor<?x?x?xf32> into tensor<?x2x?x3x4x?xf32>534 %e = tensor.expand_shape %c#1 [[0], [1, 2], [3, 4, 5]] output_shape [%sz3, %sz4, 2, 3, 4, %sz2] : tensor<?x?x?xf32> into tensor<?x?x2x3x4x?xf32>535 return %d, %e : tensor<?x2x?x3x4x?xf32>, tensor<?x?x2x3x4x?xf32>536}537// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d2, d3, d4, d0, d1, d5)>538// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d2, d3, d4, d5)>539// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d5, d2, d3, d4)>540// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d5, d0, d1, d2, d3, d4)>541// CHECK: func @reshape_as_consumer_permutation_with_multiple_results542// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<?x?x?xf32>543// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<?x?xf32>544// CHECK-SAME: %[[SZ0:.+]]: index, %[[SZ1:.+]]: index, %[[SZ2:.+]]: index, %[[SZ3:.+]]: index, %[[SZ4:.+]]: index545// CHECK: %[[RESHAPE0:.+]] = tensor.expand_shape %[[ARG0]] {{\[\[}}0, 1, 2], [3, 4], [5]] output_shape [3, 4, %[[SZ2]], %[[SZ4]], 2, %[[SZ3]]] : tensor<?x?x?xf32> into tensor<3x4x?x?x2x?xf32>546// CHECK: %[[RESHAPE1:.+]] = tensor.expand_shape %[[ARG1]] {{\[\[}}0, 1, 2], [3]] output_shape [3, 4, %[[SZ2]], %[[SZ3]]] : tensor<?x?xf32> into tensor<3x4x?x?xf32>547// CHECK: %[[RESHAPE2:.+]] = tensor.expand_shape %[[ARG0]] {{\[\[}}0, 1], [2], [3, 4, 5]] output_shape [%[[SZ4]], 2, %[[SZ3]], 3, 4, %[[SZ2]]] : tensor<?x?x?xf32> into tensor<?x2x?x3x4x?xf32>548// CHECK: %[[RESHAPE3:.+]] = tensor.expand_shape %[[ARG0]] {{\[\[}}0], [1, 2], [3, 4, 5]] output_shape [%[[SZ3]], %[[SZ4]], 2, 3, 4, %[[SZ2]]] : tensor<?x?x?xf32> into tensor<?x?x2x3x4x?xf32>549// CHECK: %[[GENERIC:.+]]:2 = linalg.generic550// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]], #[[MAP3]]]551// CHECK-SAME: ins(%[[RESHAPE0]], %[[RESHAPE1]] :552// CHECK-SAME: outs(%[[RESHAPE2]], %[[RESHAPE3]] :553// CHECK: return %[[GENERIC]]#0, %[[GENERIC]]#1554 555// -----556 557#map0 = affine_map<(d0, d1) -> (d1)>558#map1 = affine_map<(d0, d1) -> (d0, d1)>559module {560 func.func @multi_result_op_expansion(%arg0: tensor<512xf32>, %arg1: tensor<512xf32>,561 %arg2: tensor<512xf32>, %arg3: tensor<200x512xf32>) -> tensor<25x8x1x512xf32> {562 %0:2 = linalg.generic {563 indexing_maps = [#map0, #map0, #map0, #map1],564 iterator_types = ["parallel", "parallel"]}565 ins(%arg0, %arg1 : tensor<512xf32>, tensor<512xf32>)566 outs(%arg2, %arg3 : tensor<512xf32>, tensor<200x512xf32>) {567 ^bb0(%arg4: f32, %arg5: f32, %arg6: f32, %arg7: f32):568 %2 = arith.addf %arg4, %arg5 : f32569 linalg.yield %2, %2 : f32, f32570 } -> (tensor<512xf32>, tensor<200x512xf32>)571 %1 = tensor.expand_shape %0#1 [[0, 1, 2], [3]] output_shape [25, 8, 1, 512] : tensor<200x512xf32> into tensor<25x8x1x512xf32>572 return %1 : tensor<25x8x1x512xf32>573 }574}575// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3) -> (d3)>576// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>577// CHECK: func.func @multi_result_op_expansion(578// CHECK-SAME: %[[ARG0:[a-zA-Z0-9]+]]: tensor<512xf32>579// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: tensor<512xf32>580// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: tensor<512xf32>581// CHECK-SAME: %[[ARG3:[a-zA-Z0-9]+]]: tensor<200x512xf32>582// CHECK: %[[OUTS:.+]] = tensor.expand_shape %[[ARG3]] {{\[\[}}0, 1, 2], [3]] output_shape [25, 8, 1, 512] : tensor<200x512xf32> into tensor<25x8x1x512xf32>583// CHECK: %[[GENERIC:.+]]:2 = linalg.generic584// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP0]], #[[MAP0]], #[[MAP1]]]585// CHECK-SAME: ins(%[[ARG0]], %[[ARG1]] :586// CHECK-SAME: outs(%[[ARG2]], %[[OUTS]] :587// CHECK: return %[[GENERIC]]#1588 589// -----590 591#map0 = affine_map<(d0, d1, d2) -> (d0, d2)>592#map1 = affine_map<(d0, d1, d2) -> (d1, d2)>593#map2 = affine_map<(d0, d1, d2) -> (d0, d1)>594func.func @generic_op_reshape_consumer_fusion_reduction(%arg0 : tensor<?x?xf32>,595 %arg1 : tensor<?x?xf32>,596 %arg2 : tensor<?x?xf32>,597 %sz0: index,598 %sz1: index) ->599 tensor<?x?x4x5xf32>600{601 %0 = linalg.generic {602 indexing_maps = [#map0, #map1, #map2],603 iterator_types = ["parallel", "parallel", "reduction"]}604 ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)605 outs(%arg2 : tensor<?x?xf32>) {606 ^bb0(%arg3: f32, %arg4: f32, %s: f32):607 %1 = arith.mulf %arg3, %arg4 : f32608 linalg.yield %1 : f32609 } -> tensor<?x?xf32>610 %1 = tensor.expand_shape %0 [[0], [1, 2, 3]] output_shape [%sz0, %sz1, 4, 5] :611 tensor<?x?xf32> into tensor<?x?x4x5xf32>612 return %1 : tensor<?x?x4x5xf32>613}614 615// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d4)>616// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d1, d2, d3, d4)>617// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3)>618// CHECK: func @generic_op_reshape_consumer_fusion_reduction619// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?xf32>620// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x?xf32>621// CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]+]]: tensor<?x?xf32>622// CHECK-SAME: %[[SZ0:.+]]: index, %[[SZ1:.+]]: index623// CHECK: %[[C1:.+]] = arith.constant 1 : index624// CHECK: %[[DIM:.+]] = tensor.dim %[[ARG0]], %[[C1]] : tensor<?x?xf32>625// CHECK: %[[T1:.+]] = tensor.expand_shape %[[ARG1]] {{\[\[}}0, 1, 2], [3]] output_shape [%[[SZ1]], 4, 5, %[[DIM]]] : tensor<?x?xf32> into tensor<?x4x5x?xf32>626// CHECK: %[[T2:.+]] = tensor.expand_shape %[[ARG2]] {{\[\[}}0], [1, 2, 3]] output_shape [%[[SZ0]], %[[SZ1]], 4, 5] : tensor<?x?xf32> into tensor<?x?x4x5xf32>627// CHECK: %[[T3:.+]] = linalg.generic628// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]]629// CHECK-SAME: ["parallel", "parallel", "parallel", "parallel", "reduction"]630// CHECK-SAME: ins(%[[ARG0]], %[[T1]] : tensor<?x?xf32>, tensor<?x4x5x?xf32>)631// CHECK-SAME: outs(%[[T2]] : tensor<?x?x4x5xf32>)632// CHECK: return %[[T3]] : tensor<?x?x4x5xf32>633 634// -----635 636#map0 = affine_map<(d0, d1, d2) -> (d2, d0)>637#map1 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>638#map2 = affine_map<(d0, d1, d2) -> (d0, d2)>639func.func @generic_op_reshape_producer_fusion_with_reduction(%arg0 : tensor<?x7x?x8xf32>,640 %arg1 : tensor<?x4x?xf32>,641 %arg2 : tensor<?x?xf32>) ->642 tensor<?x?xf32>643{644 %0 = tensor.collapse_shape %arg0 [[0, 1], [2, 3]] :645 tensor<?x7x?x8xf32> into tensor<?x?xf32>646 %1 = linalg.generic {647 indexing_maps = [#map0, #map1, #map2],648 iterator_types = ["parallel", "reduction", "parallel"]}649 ins(%0, %arg1 : tensor<?x?xf32>, tensor<?x4x?xf32>)650 outs(%arg2 : tensor<?x?xf32>) {651 ^bb0(%arg3: f32, %arg4: f32, %arg5: f32):652 %1 = arith.mulf %arg3, %arg4 : f32653 %2 = arith.addf %1, %arg5 : f32654 linalg.yield %2 : f32655 } -> tensor<?x?xf32>656 return %1 : tensor<?x?xf32>657}658 659// CHECK-DAG: #[[$MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d3, d4, d0, d1)>660// CHECK-DAG: #[[$MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>661// CHECK-DAG: #[[$MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d3, d4)>662// CHECK: func @generic_op_reshape_producer_fusion_with_reduction663// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x7x?x8xf32>664// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x4x?xf32>665// CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]+]]: tensor<?x?xf32>666// CHECK: %[[C2:.+]] = arith.constant 2 : index667// CHECK: %[[C0:.+]] = arith.constant 0 : index668// CHECK: %[[DIM:.+]] = tensor.dim %[[ARG0]], %[[C0]] : tensor<?x7x?x8xf32>669// CHECK: %[[DIM_0:.+]] = tensor.dim %[[ARG0]], %[[C2]] : tensor<?x7x?x8xf32>670// CHECK: %[[T1:.+]] = tensor.expand_shape %[[ARG1]] {{\[\[}}0, 1], [2], [3, 4]] output_shape [%[[DIM_0]], 8, 4, %[[DIM]], 7] : tensor<?x4x?xf32> into tensor<?x8x4x?x7xf32>671// CHECK: %[[T2:.+]] = tensor.expand_shape %[[ARG2]] {{\[\[}}0, 1], [2, 3]] output_shape [%[[DIM_0]], 8, %[[DIM]], 7] : tensor<?x?xf32> into tensor<?x8x?x7xf32>672// CHECK: %[[T3:.+]] = linalg.generic673// CHECK-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]]]674// CHECK-SAME: ["parallel", "parallel", "reduction", "parallel", "parallel"]675// CHECK-SAME: ins(%[[ARG0]], %[[T1]] : tensor<?x7x?x8xf32>, tensor<?x8x4x?x7xf32>)676// CHECK-SAME: outs(%[[T2]] : tensor<?x8x?x7xf32>)677// CHECK: %[[T4:.+]] = tensor.collapse_shape %[[T3]]678// CHECK-SAME: [0, 1], [2, 3]679// CHECK-SAME: tensor<?x8x?x7xf32> into tensor<?x?xf32>680// CHECK: return %[[T4]]681 682// -----683 684func.func @linalg_add_reshape_consumer_fusion(%arg0 : tensor<?x?xf32>,685 %arg1 : tensor<?x?xf32>,686 %arg2 : tensor<?x?xf32>,687 %sz0: index,688 %sz1: index) ->689 tensor<?x?x4x5xf32>690{691 %0 = linalg.add ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)692 outs(%arg2 : tensor<?x?xf32>) -> tensor<?x?xf32>693 %1 = tensor.expand_shape %0 [[0], [1, 2, 3]] output_shape [%sz0, %sz1, 4, 5] :694 tensor<?x?xf32> into tensor<?x?x4x5xf32>695 return %1 : tensor<?x?x4x5xf32>696}697 698// CHECK-DAG: #[[MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>699// CHECK: func @linalg_add_reshape_consumer_fusion700// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?xf32>701// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x?xf32>702// CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]+]]: tensor<?x?xf32>703// CHECK-SAME: %[[SZ0:.+]]: index, %[[SZ1:.+]]: index704// CHECK: %[[T1:.+]] = tensor.expand_shape %[[ARG0]] {{\[\[}}0], [1, 2, 3]] output_shape [%[[SZ0]], %[[SZ1]], 4, 5] : tensor<?x?xf32> into tensor<?x?x4x5xf32>705// CHECK: %[[T2:.+]] = tensor.expand_shape %[[ARG1]] {{\[\[}}0], [1, 2, 3]] output_shape [%[[SZ0]], %[[SZ1]], 4, 5] : tensor<?x?xf32> into tensor<?x?x4x5xf32>706// CHECK: %[[T3:.+]] = tensor.expand_shape %[[ARG2]] {{\[\[}}0], [1, 2, 3]] output_shape [%[[SZ0]], %[[SZ1]], 4, 5] : tensor<?x?xf32> into tensor<?x?x4x5xf32>707// CHECK: %[[T4:.+]] = linalg.generic708// CHECK-SAME: indexing_maps = [#[[MAP]], #[[MAP]], #[[MAP]]]709// CHECK-SAME: ["parallel", "parallel", "parallel", "parallel"]710// CHECK-SAME: ins(%[[T1]], %[[T2]] : tensor<?x?x4x5xf32>, tensor<?x?x4x5xf32>)711// CHECK-SAME: outs(%[[T3]] : tensor<?x?x4x5xf32>)712// CHECK: return %[[T4]] : tensor<?x?x4x5xf32>713 714// -----715 716func.func @linalg_add_reshape_producer_fusion(%arg0 : tensor<?x7x?x8xf32>,717 %arg1 : tensor<?x?xf32>,718 %arg2 : tensor<?x?xf32>) ->719 tensor<?x?xf32>720{721 %0 = tensor.collapse_shape %arg0 [[0, 1], [2, 3]] :722 tensor<?x7x?x8xf32> into tensor<?x?xf32>723 %1 = linalg.add ins(%0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>)724 outs(%arg2 : tensor<?x?xf32>) -> tensor<?x?xf32>725 return %1 : tensor<?x?xf32>726}727 728// CHECK-DAG: #[[$MAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>729// CHECK: func @linalg_add_reshape_producer_fusion730// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x7x?x8xf32>731// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x?xf32>732// CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]+]]: tensor<?x?xf32>733// CHECK: %[[C2:.+]] = arith.constant 2 : index734// CHECK: %[[C0:.+]] = arith.constant 0 : index735// CHECK: %[[DIM:.+]] = tensor.dim %[[ARG0]], %[[C0]] : tensor<?x7x?x8xf32>736// CHECK: %[[DIM_0:.+]] = tensor.dim %[[ARG0]], %[[C2]] : tensor<?x7x?x8xf32>737// CHECK: %[[T1:.+]] = tensor.expand_shape %[[ARG1]] {{\[\[}}0, 1], [2, 3]] output_shape [%[[DIM]], 7, %[[DIM_0]], 8] : tensor<?x?xf32> into tensor<?x7x?x8xf32>738// CHECK: %[[T2:.+]] = tensor.expand_shape %[[ARG2]] {{\[\[}}0, 1], [2, 3]] output_shape [%[[DIM]], 7, %[[DIM_0]], 8] : tensor<?x?xf32> into tensor<?x7x?x8xf32>739// CHECK: %[[T3:.+]] = linalg.generic740// CHECK-SAME: indexing_maps = [#[[$MAP]], #[[$MAP]], #[[$MAP]]]741// CHECK-SAME: ["parallel", "parallel", "parallel", "parallel"]742// CHECK-SAME: ins(%[[ARG0]], %[[T1]] : tensor<?x7x?x8xf32>, tensor<?x7x?x8xf32>)743// CHECK-SAME: outs(%[[T2]] : tensor<?x7x?x8xf32>)744// CHECK: %[[T4:.+]] = tensor.collapse_shape %[[T3]]745// CHECK-SAME: [0, 1], [2, 3]746// CHECK-SAME: tensor<?x7x?x8xf32> into tensor<?x?xf32>747// CHECK: return %[[T4]]748 749// -----750 751func.func @linalg_copy_reshape_producer_fusion(%arg0 : tensor<?x7x?x8xf32>,752 %arg1 : tensor<?x?xf32>) ->753 tensor<?x?xf32>754{755 %0 = tensor.collapse_shape %arg0 [[0, 1], [2, 3]] :756 tensor<?x7x?x8xf32> into tensor<?x?xf32>757 %1 = linalg.copy ins(%0 : tensor<?x?xf32>)758 outs(%arg1 : tensor<?x?xf32>) -> tensor<?x?xf32>759 return %1 : tensor<?x?xf32>760}761 762// CHECK: func @linalg_copy_reshape_producer_fusion763// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x7x?x8xf32>764// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x?xf32>765// CHECK-DAG: %[[C2:.+]] = arith.constant 2 : index766// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index767// CHECK-DAG: %[[DIM:.+]] = tensor.dim %[[ARG0]], %[[C0]]768// CHECK-DAG: %[[DIM_0:.+]] = tensor.dim %[[ARG0]], %[[C2]]769// CHECK: %[[T1:.+]] = tensor.expand_shape %[[ARG1]] {{\[\[}}0, 1], [2, 3]] output_shape [%[[DIM]], 7, %[[DIM_0]], 8] : tensor<?x?xf32> into tensor<?x7x?x8xf32>770// CHECK: %[[T2:.+]] = linalg.copy771// CHECK-SAME: ins(%[[ARG0]] : tensor<?x7x?x8xf32>)772// CHECK-SAME: outs(%[[T1]] : tensor<?x7x?x8xf32>)773// CHECK: %[[T3:.+]] = tensor.collapse_shape %[[T2]]774// CHECK-SAME: [0, 1], [2, 3]775// CHECK-SAME: tensor<?x7x?x8xf32> into tensor<?x?xf32>776// CHECK: return %[[T3]]777 778// -----779 780func.func @reshape_as_producer_transpose781 (%a : tensor<4x5x6x7x2x3xf32>)782 -> tensor<6x4x210xf32> {783 %b = tensor.empty() : tensor<6x4x210xf32>784 %c = tensor.collapse_shape %a [[0], [1, 2, 3], [4, 5]] :785 tensor<4x5x6x7x2x3xf32> into tensor<4x210x6xf32>786 %d = linalg.transpose787 ins(%c : tensor<4x210x6xf32>)788 outs(%b : tensor<6x4x210xf32>) permutation = [2, 0, 1]789 return %d : tensor<6x4x210xf32>790}791 792// CHECK: func @reshape_as_producer_transpose793// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<4x5x6x7x2x3xf32>794// CHECK-DAG: %[[INIT:.+]] = tensor.empty()795// CHECK-DAG: %[[T0:.+]] = tensor.expand_shape %[[INIT]] {{\[\[}}0, 1], [2], [3, 4, 5]] output_shape [2, 3, 4, 5, 6, 7] : tensor<6x4x210xf32> into tensor<2x3x4x5x6x7xf32>796// CHECK: %[[T1:.+]] = linalg.transpose ins(%[[ARG0]] : tensor<4x5x6x7x2x3xf32>)797// CHECK-SAME: outs(%[[T0]] : tensor<2x3x4x5x6x7xf32>)798// CHECK-SAME: permutation = [4, 5, 0, 1, 2, 3]799// CHECK: %[[T2:.+]] = tensor.collapse_shape %[[T1]] {{\[\[}}0, 1], [2], [3, 4, 5]] : tensor<2x3x4x5x6x7xf32> into tensor<6x4x210xf32>800// CHECK: return %[[T2]] : tensor<6x4x210xf32>801 802 803// -----804 805func.func @fuse_by_expanding_pad(%arg0 : tensor<2x3x4x5x6x7x8x9xi32>) -> tensor<8x12x17x336x14xi32> {806 %collapse = tensor.collapse_shape %arg0 [[0], [1, 2], [3], [4, 5, 6], [7]] : tensor<2x3x4x5x6x7x8x9xi32> into tensor<2x12x5x336x9xi32>807 %cst = arith.constant 0 : i32808 %padded_0 = tensor.pad %collapse low[1, 0, 8, 0, 3] high[5, 0, 4, 0, 2] {809 ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index):810 tensor.yield %cst : i32811 } : tensor<2x12x5x336x9xi32> to tensor<8x12x17x336x14xi32>812 return %padded_0 : tensor<8x12x17x336x14xi32>813}814// CHECK: func @fuse_by_expanding_pad(815// CHECK-SAME: %[[ARG0:.+]]: tensor<2x3x4x5x6x7x8x9xi32>)816// CHECK: %[[PAD:.+]] = tensor.pad %[[ARG0]]817// CHECK-SAME: low[1, 0, 0, 8, 0, 0, 0, 3] high[5, 0, 0, 4, 0, 0, 0, 2]818// CHECK: tensor<2x3x4x5x6x7x8x9xi32> to tensor<8x3x4x17x6x7x8x14xi32>819// CHECK: %[[COLLAPSE:.+]] = tensor.collapse_shape %[[PAD]] {{\[}}[0], [1, 2], [3], [4, 5, 6], [7]]820// CHECK-SAME: : tensor<8x3x4x17x6x7x8x14xi32> into tensor<8x12x17x336x14xi32>821// CHECK: return %[[COLLAPSE]]822 823// -----824 825func.func @no_fuse_by_expanding_pad_non_constant_padding(%arg0 : tensor<2x3x4xi32>) -> tensor<8x12xi32> {826 %collapse = tensor.collapse_shape %arg0 [[0], [1, 2]] : tensor<2x3x4xi32> into tensor<2x12xi32>827 %padded_0 = tensor.pad %collapse low[1, 0] high[5, 0] {828 ^bb0(%arg1: index, %arg2: index):829 %pad_val = arith.index_cast %arg1 : index to i32830 tensor.yield %pad_val : i32831 } : tensor<2x12xi32> to tensor<8x12xi32>832 return %padded_0 : tensor<8x12xi32> 833}834// CHECK: func @no_fuse_by_expanding_pad_non_constant_padding(835// CHECK-SAME: %[[ARG0:.+]]: tensor<2x3x4xi32>)836// CHECK: %[[COLLAPSE:.+]] = tensor.collapse_shape %[[ARG0]]837// CHECK: %[[PAD:.+]] = tensor.pad %[[COLLAPSE]]838// CHECK: return %[[PAD]]839 840// -----841 842func.func @no_fuse_by_expanding_pad(%arg0 : tensor<2x3x4x5x6x7x8x9xi32>) -> tensor<8x12x17x339x14xi32> {843 %collapse = tensor.collapse_shape %arg0 [[0], [1, 2], [3], [4, 5, 6], [7]] : tensor<2x3x4x5x6x7x8x9xi32> into tensor<2x12x5x336x9xi32>844 %cst = arith.constant 0 : i32845 %padded_0 = tensor.pad %collapse low[1, 0, 8, 0, 3] high[5, 0, 4, 3, 2] {846 ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index):847 tensor.yield %cst : i32848 } : tensor<2x12x5x336x9xi32> to tensor<8x12x17x339x14xi32>849 return %padded_0 : tensor<8x12x17x339x14xi32>850}851// CHECK: func @no_fuse_by_expanding_pad(852// CHECK-SAME: %[[ARG0:.+]]: tensor<2x3x4x5x6x7x8x9xi32>)853// CHECK: %[[COLLAPSE:.+]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0], [1, 2], [3], [4, 5, 6], [7]]854// CHECK-SAME: : tensor<2x3x4x5x6x7x8x9xi32> into tensor<2x12x5x336x9xi32>855// CHECK: %[[PAD:.+]] = tensor.pad %[[COLLAPSE]]856// CHECK-SAME: low[1, 0, 8, 0, 3] high[5, 0, 4, 3, 2]857// CHECK: tensor<2x12x5x336x9xi32> to tensor<8x12x17x339x14xi32>858// CHECK: return %[[PAD]]859 860// -----861 862func.func @fuse_by_expanding_dynamic_pad(%arg0 : tensor<?x?x?x?x?x?xi32>, %l0: index, %l1: index, %h0: index, %h1: index) -> tensor<?x?x?x?xi32> {863 %collapse = tensor.collapse_shape %arg0 [[0], [1, 2], [3], [4, 5]] : tensor<?x?x?x?x?x?xi32> into tensor<?x?x?x?xi32>864 %cst = arith.constant 0 : i32865 %padded_0 = tensor.pad %collapse low[%l0, 0, %l1, 0] high[%h0, 0, %h1, 0] {866 ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):867 tensor.yield %cst : i32868 } : tensor<?x?x?x?xi32> to tensor<?x?x?x?xi32>869 return %padded_0 : tensor<?x?x?x?xi32>870}871// CHECK: func @fuse_by_expanding_dynamic_pad(872// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?x?x?x?x?xi32>873// CHECK-SAME: %[[L0:.+]]: index, %[[L1:.+]]: index, %[[H0:.+]]: index, %[[H1:.+]]: index874// CHECK: %[[PAD:.+]] = tensor.pad %[[ARG0]]875// CHECK-SAME: low[%[[L0]], 0, 0, %[[L1]], 0, 0] high[%[[H0]], 0, 0, %[[H1]], 0, 0]876// CHECK: tensor<?x?x?x?x?x?xi32> to tensor<?x?x?x?x?x?xi32>877// CHECK: %[[COLLAPSE:.+]] = tensor.collapse_shape %[[PAD]] {{\[}}[0], [1, 2], [3], [4, 5]]878// CHECK-SAME: : tensor<?x?x?x?x?x?xi32> into tensor<?x?x?x?xi32>879// CHECK: return %[[COLLAPSE]]880 881// -----882 883func.func @expand_shape_with_producer_pad(%arg0: tensor<2x12x5x336x9xi32>) -> tensor<8x3x4x17x6x7x8x14xi32> {884 %cst = arith.constant 0 : i32885 %padded = tensor.pad %arg0 low[1, 0, 8, 0, 3] high[5, 0, 4, 0, 2] {886 ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index):887 tensor.yield %cst : i32888 } : tensor<2x12x5x336x9xi32> to tensor<8x12x17x336x14xi32>889 %expanded = tensor.expand_shape %padded [[0], [1, 2], [3], [4, 5, 6], [7]] output_shape [8, 3, 4, 17, 6, 7, 8, 14]890 : tensor<8x12x17x336x14xi32> into tensor<8x3x4x17x6x7x8x14xi32>891 return %expanded : tensor<8x3x4x17x6x7x8x14xi32>892}893// CHECK: func @expand_shape_with_producer_pad894// CHECK-SAME: %[[ARG0:.+]]: tensor<2x12x5x336x9xi32>895// CHECK: %[[EXPAND:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0], [1, 2], [3], [4, 5, 6], [7]] output_shape [2, 3, 4, 5, 6, 7, 8, 9]896// CHECK: %[[PAD:.+]] = tensor.pad %[[EXPAND]] low[1, 0, 0, 8, 0, 0, 0, 3] high[5, 0, 0, 4, 0, 0, 0, 2]897// CHECK: return %[[PAD]]898 899// -----900 901func.func @expand_shape_with_producer_pad_dynamic(%arg0: tensor<?x?x?x?xf32>,902 %s0: index, %s1: index, %s2: index, %s3: index, %s4: index, %s5: index,903 %l0: index, %l1: index, %h0: index, %h1: index) -> tensor<?x?x?x?x?x?xf32> {904 %cst = arith.constant 0.0 : f32905 %padded = tensor.pad %arg0 low[%l0, 0, %l1, 0] high[%h0, 0, %h1, 0] {906 ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):907 tensor.yield %cst : f32908 } : tensor<?x?x?x?xf32> to tensor<?x?x?x?xf32>909 %expanded = tensor.expand_shape %padded [[0], [1, 2], [3], [4, 5]] output_shape [%s0, %s1, %s2, %s3, %s4, %s5]910 : tensor<?x?x?x?xf32> into tensor<?x?x?x?x?x?xf32>911 return %expanded : tensor<?x?x?x?x?x?xf32>912}913// CHECK: func @expand_shape_with_producer_pad_dynamic914// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?x?x?xf32>915// CHECK-SAME: %[[S0:.+]]: index, %[[S1:.+]]: index, %[[S2:.+]]: index, %[[S3:.+]]: index, %[[S4:.+]]: index, %[[S5:.+]]: index, %[[L0:.+]]: index, %[[L1:.+]]: index, %[[H0:.+]]: index, %[[H1:.+]]: index916// CHECK: %[[DIM0:.+]] = tensor.dim %[[ARG0]], %[[C0:.+]] : tensor<?x?x?x?xf32>917// CHECK: %[[DIM2:.+]] = tensor.dim %[[ARG0]], %[[C2:.+]] : tensor<?x?x?x?xf32>918// CHECK: %[[EXPAND:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0], [1, 2], [3], [4, 5]] output_shape [%[[DIM0]], %[[S1]], %[[S2]], %[[DIM2]], %[[S4]], %[[S5]]]919// CHECK: %[[PAD:.+]] = tensor.pad %[[EXPAND]] low[%[[L0]], 0, 0, %[[L1]], 0, 0] high[%[[H0]], 0, 0, %[[H1]], 0, 0]920// CHECK: return %[[PAD]]921 922// -----923 924func.func @expand_shape_with_producer_pad_non_constant_padding(%arg0 : tensor<2x12xi32>) -> tensor<8x3x4xi32> {925 %padded_0 = tensor.pad %arg0 low[1, 0] high[5, 0] {926 ^bb0(%arg1: index, %arg2: index):927 %pad_val = arith.index_cast %arg1 : index to i32928 tensor.yield %pad_val : i32929 } : tensor<2x12xi32> to tensor<8x12xi32>930 %expand = tensor.expand_shape %padded_0 [[0], [1, 2]] output_shape [8, 3, 4] : tensor<8x12xi32> into tensor<8x3x4xi32>931 return %expand : tensor<8x3x4xi32> 932}933// CHECK: func @expand_shape_with_producer_pad_non_constant_padding(934// CHECK-SAME: %[[ARG0:.+]]: tensor<2x12xi32>)935// CHECK: %[[PAD:.+]] = tensor.pad %[[ARG0]]936// CHECK: %[[EXPAND:.+]] = tensor.expand_shape %[[PAD]]937// CHECK: return %[[EXPAND]]938 939// -----940 941func.func @move_operand_deps(%arg0 : tensor<?x128xf16>,942 %arg1 : tensor<4x?x32x128xf16>, %empty : tensor<4x?x32x128xf16>) -> tensor<4x?x32x8x16xf16> {943 %c0 = arith.constant 0 : index944 %0 = linalg.generic {945 indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d1, d3)>,946 affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>],947 iterator_types = ["parallel", "parallel", "parallel", "parallel"]}948 ins(%arg0 : tensor<?x128xf16>)949 outs(%empty : tensor<4x?x32x128xf16>) {950 ^bb0(%b0: f16, %b1 : f16) :951 %iv0 = linalg.index 0 : index952 %iv1 = linalg.index 1 : index953 %iv2 = linalg.index 2 : index954 %iv3 = linalg.index 3 : index955 %1 = tensor.extract %arg1[%iv0, %iv1, %iv2, %iv3] : tensor<4x?x32x128xf16>956 %2 = arith.addf %1, %b0 : f16957 linalg.yield %2 : f16958 } -> tensor<4x?x32x128xf16>959 %1 = tensor.dim %arg0, %c0 : tensor<?x128xf16>960 %2 = tensor.expand_shape %0 [[0], [1], [2], [3, 4]] output_shape [4, %1, 32, 8, 16]961 : tensor<4x?x32x128xf16> into tensor<4x?x32x8x16xf16>962 func.return %2 : tensor<4x?x32x8x16xf16>963}964// CHECK: func @move_operand_deps(965// CHECK-SAME: %[[ARG0:.+]]: tensor<?x128xf16>966// CHECK-DAG: %[[MOVED_OP:.+]] = tensor.dim %[[ARG0]]967// CHECK-DAG: %[[EXPANDED:.+]] = tensor.expand_shape %[[ARG0]]968// CHECK: %[[GENERIC:.+]] = linalg.generic969// CHECK-SAME: ins(%[[EXPANDED]] :970// CHECK: return %[[GENERIC]]971