936 lines · plain
1// RUN: mlir-opt %s -test-linalg-elementwise-fusion-patterns=fuse-with-reshape-by-collapsing -split-input-file | FileCheck %s2// RUN: mlir-opt %s -test-linalg-elementwise-fusion-patterns=fuse-with-reshape-by-collapsing-control -split-input-file | FileCheck %s --check-prefix=CONTROL3 4// Static problem sizes. Checks all aspects of fusion by collapsing. Rest of the5// tests only check a subset of conditions.6#map0 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d0, d1, d2, d3, d4, d5, d6, d7)>7#map1 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d0, d1, d2)>8#map2 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d3, d4, d5, d6)>9#map3 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d0, d1, d2, d3, d4, d5, d6, d7)>10#map4 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d1, d2, d0, d7, d3, d4, d5, d6)>11func.func @fuse_by_collapsing(%arg0 : tensor<2x12x5x336x9xi32>,12 %arg1 : tensor<2x3x4xi32>, %arg2 : tensor<5x6x7x8xi32>) -> (tensor<2x3x4x5x6x7x8x9xi32>, tensor<3x4x2x9x5x6x7x8xi32>) {13 %expand = tensor.expand_shape %arg0 [[0], [1, 2], [3], [4, 5, 6], [7]] output_shape [2, 3, 4, 5, 6, 7, 8, 9] : tensor<2x12x5x336x9xi32> into tensor<2x3x4x5x6x7x8x9xi32>14 %init_0 = tensor.empty() : tensor<2x3x4x5x6x7x8x9xi32>15 %init_1 = tensor.empty() : tensor<3x4x2x9x5x6x7x8xi32>16 %generic:2 = linalg.generic {17 indexing_maps = [#map0, #map1, #map2, #map3, #map4],18 iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel", "parallel", "parallel", "parallel"]}19 ins(%expand, %arg1, %arg2 : tensor<2x3x4x5x6x7x8x9xi32>, tensor<2x3x4xi32>, tensor<5x6x7x8xi32>)20 outs(%init_0, %init_1 : tensor<2x3x4x5x6x7x8x9xi32>, tensor<3x4x2x9x5x6x7x8xi32>) {21 ^bb0(%b0 : i32, %b1 : i32, %b2 : i32, %b3 : i32, %b4 : i32):22 %t0 = arith.addi %b0, %b1 : i3223 %t1 = arith.addi %t0, %b2 : i3224 linalg.yield %t1, %t1 : i32, i3225 } -> (tensor<2x3x4x5x6x7x8x9xi32>, tensor<3x4x2x9x5x6x7x8xi32>)26 return %generic#0, %generic#1 : tensor<2x3x4x5x6x7x8x9xi32>, tensor<3x4x2x9x5x6x7x8xi32>27}28// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>29// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1)>30// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d2, d3)>31// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d1, d0, d4, d2, d3)>32// CHECK: func @fuse_by_collapsing(33// CHECK-SAME: %[[ARG0:.+]]: tensor<2x12x5x336x9xi32>34// CHECK-SAME: %[[ARG1:.+]]: tensor<2x3x4xi32>35// CHECK-SAME: %[[ARG2:.+]]: tensor<5x6x7x8xi32>36// CHECK-DAG: %[[INIT0:.+]] = tensor.empty() : tensor<2x3x4x5x6x7x8x9xi32>37// CHECK-DAG: %[[INIT1:.+]] = tensor.empty() : tensor<3x4x2x9x5x6x7x8xi32>38// CHECK-DAG: %[[ARG1_RESHAPE:.+]] = tensor.collapse_shape %[[ARG1]] {{\[}}[0], [1, 2]{{\]}}39// CHECK-DAG: %[[ARG2_RESHAPE:.+]] = tensor.collapse_shape %[[ARG2]] {{\[}}[0], [1, 2, 3]{{\]}}40// CHECK-DAG: %[[INIT0_RESHAPE:.+]] = tensor.collapse_shape %[[INIT0]] {{\[}}[0], [1, 2], [3], [4, 5, 6], [7]{{\]}}41// CHECK-DAG: %[[INIT1_RESHAPE:.+]] = tensor.collapse_shape %[[INIT1]] {{\[}}[0, 1], [2], [3], [4], [5, 6, 7]{{\]}}42// CHECK: %[[COLLAPSED_OP:.+]]:2 = linalg.generic43// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]], #[[MAP0]], #[[MAP3]]]44// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]45// CHECK-SAME: ins(%[[ARG0]], %[[ARG1_RESHAPE]], %[[ARG2_RESHAPE]] :46// CHECK-SAME: outs(%[[INIT0_RESHAPE]], %[[INIT1_RESHAPE]] :47// CHECK: %[[RESULT0_RESHAPE:.+]] = tensor.expand_shape %[[COLLAPSED_OP]]#0 {{\[}}[0], [1, 2], [3], [4, 5, 6], [7]{{\]}} output_shape [2, 3, 4, 5, 6, 7, 8, 9]48// CHECK: %[[RESULT1_RESHAPE:.+]] = tensor.expand_shape %[[COLLAPSED_OP]]#1 {{\[}}[0, 1], [2], [3], [4], [5, 6, 7]{{\]}} output_shape [3, 4, 2, 9, 5, 6, 7, 8]49// CHECK: return %[[RESULT0_RESHAPE]], %[[RESULT1_RESHAPE]]50 51// CONTROL: func @fuse_by_collapsing(52// CONTROL-SAME: %[[ARG0:.+]]: tensor<2x12x5x336x9xi32>53// CONTROL-SAME: %[[ARG1:.+]]: tensor<2x3x4xi32>54// CONTROL-SAME: %[[ARG2:.+]]: tensor<5x6x7x8xi32>55// CONTROL: %[[EXPAND:.+]] = tensor.expand_shape %[[ARG0]]56// CONTROL: %[[GENERIC:.+]]:2 = linalg.generic57// CONTROL-SAME: ins(%[[EXPAND]],58// CONTROL: return %[[GENERIC]]#0, %[[GENERIC]]#159 60// -----61 62#map0 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d0, d1, d2, d3, d4, d5, d6, d7)>63#map1 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d0, d1, d2)>64#map2 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d3, d4, d5, d6)>65#map3 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d0, d1, d2, d3, d4, d5, d6, d7)>66func.func @fuse_by_collapsing_indexing_op(%arg0 : tensor<2x12x5x336x9xi32>,67 %arg1 : tensor<2x3x4xi32>, %arg2 : tensor<5x6x7x8xi32>) -> tensor<2x3x4x5x6x7x8x9xi32> {68 %expand = tensor.expand_shape %arg0 [[0], [1, 2], [3], [4, 5, 6], [7]] output_shape [2, 3, 4, 5, 6, 7, 8, 9] : tensor<2x12x5x336x9xi32> into tensor<2x3x4x5x6x7x8x9xi32>69 %init = tensor.empty() : tensor<2x3x4x5x6x7x8x9xi32>70 %generic = linalg.generic {71 indexing_maps = [#map0, #map1, #map2, #map3],72 iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel", "parallel", "parallel", "parallel"]}73 ins(%expand, %arg1, %arg2 : tensor<2x3x4x5x6x7x8x9xi32>, tensor<2x3x4xi32>, tensor<5x6x7x8xi32>)74 outs(%init : tensor<2x3x4x5x6x7x8x9xi32>) {75 ^bb0(%b0 : i32, %b1 : i32, %b2 : i32, %b3 : i32):76 %iv0 = linalg.index 0: index77 %iv1 = linalg.index 1: index78 %t0 = arith.addi %iv0, %iv1 : index79 %iv2 = linalg.index 2 : index80 %t1 = arith.addi %t0, %iv2 : index81 %iv3 = linalg.index 3 : index82 %t2 = arith.addi %t1, %iv3 : index83 %iv4 = linalg.index 4 : index84 %t3 = arith.addi %t2, %iv4 : index85 %iv5 = linalg.index 5 : index86 %t4 = arith.addi %t3, %iv5 : index87 %iv6 = linalg.index 6 : index88 %t5 = arith.addi %t4, %iv6 : index89 %iv7 = linalg.index 7 : index90 %t6 = arith.addi %t5, %iv7 : index91 %yield = arith.index_cast %t6 : index to i3292 linalg.yield %yield : i3293 } -> tensor<2x3x4x5x6x7x8x9xi32>94 return %generic : tensor<2x3x4x5x6x7x8x9xi32>95}96// CHECK-LABEL: func @fuse_by_collapsing_indexing_op(97// CHECK-DAG: %[[C4:.+]] = arith.constant 4 : index98// CHECK-DAG: %[[C8:.+]] = arith.constant 8 : index99// CHECK-DAG: %[[C7:.+]] = arith.constant 7 : index100// CHECK: %[[IV0:.+]] = linalg.index 0101// CHECK: %[[IV1:.+]] = linalg.index 1102// CHECK: %[[REM_IV1:.+]] = arith.remsi %[[IV1]], %[[C4]]103// CHECK: %[[DIV_IV1:.+]] = arith.divsi %[[IV1]], %[[C4]]104// CHECK: %[[IV2:.+]] = linalg.index 2105// CHECK: %[[IV3:.+]] = linalg.index 3106// CHECK: %[[REM1_IV3:.+]] = arith.remsi %[[IV3]], %[[C8]]107// CHECK: %[[DIV1_IV3:.+]] = arith.divsi %[[IV3]], %[[C8]]108// CHECK: %[[REM2_IV3:.+]] = arith.remsi %[[DIV1_IV3]], %[[C7]]109// CHECK: %[[DIV2_IV3:.+]] = arith.divsi %[[DIV1_IV3]], %[[C7]]110// CHECK: %[[IV4:.+]] = linalg.index 4111// CHECK: %[[T0:.+]] = arith.addi %[[IV0]], %[[DIV_IV1]]112// CHECK: %[[T1:.+]] = arith.addi %[[T0]], %[[REM_IV1]]113// CHECK: %[[T2:.+]] = arith.addi %[[T1]], %[[IV2]]114// CHECK: %[[T3:.+]] = arith.addi %[[T2]], %[[DIV2_IV3]]115// CHECK: %[[T4:.+]] = arith.addi %[[T3]], %[[REM2_IV3]]116// CHECK: %[[T5:.+]] = arith.addi %[[T4]], %[[REM1_IV3]]117// CHECK: %[[T6:.+]] = arith.addi %[[T5]], %[[IV4]]118// CHECK: %[[YIELD:.+]] = arith.index_cast %[[T6]]119// CHECK: linalg.yield %[[YIELD]]120 121// -----122 123#map0 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d7, d5, d6, d0, d1, d2, d3, d4)>124#map1 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d5, d6, d0)>125#map2 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d4, d1, d2, d3)>126#map3 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d0, d1, d2, d3, d4, d5, d6, d7)>127func.func @fuse_by_collapsing_change_reshape_order(%arg0 : tensor<9x56x2x60x6xi32>,128 %arg1 : tensor<7x8x2xi32>, %arg2 : tensor<6x3x4x5xi32>) -> tensor<2x3x4x5x6x7x8x9xi32> {129 %expand = tensor.expand_shape %arg0 [[0], [1, 2], [3], [4, 5, 6], [7]] output_shape [9, 7, 8, 2, 3, 4, 5, 6] : tensor<9x56x2x60x6xi32> into tensor<9x7x8x2x3x4x5x6xi32>130 %init = tensor.empty() : tensor<2x3x4x5x6x7x8x9xi32>131 %generic = linalg.generic {132 indexing_maps = [#map0, #map1, #map2, #map3],133 iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel", "parallel", "parallel", "parallel"]}134 ins(%expand, %arg1, %arg2 : tensor<9x7x8x2x3x4x5x6xi32>, tensor<7x8x2xi32>, tensor<6x3x4x5xi32>)135 outs(%init : tensor<2x3x4x5x6x7x8x9xi32>) {136 ^bb0(%b0 : i32, %b1 : i32, %b2 : i32, %b3 : i32):137 %t0 = arith.addi %b0, %b1 : i32138 %t1 = arith.addi %t0, %b2 : i32139 linalg.yield %t1 : i32140 } -> tensor<2x3x4x5x6x7x8x9xi32>141 return %generic : tensor<2x3x4x5x6x7x8x9xi32>142}143// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d4, d3, d0, d1, d2)>144// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d3, d0)>145// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d2, d1)>146// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>147// CHECK: func @fuse_by_collapsing_change_reshape_order(148// CHECK-SAME: %[[ARG0:.+]]: tensor<9x56x2x60x6xi32>149// CHECK-SAME: %[[ARG1:.+]]: tensor<7x8x2xi32>150// CHECK-SAME: %[[ARG2:.+]]: tensor<6x3x4x5xi32>151// CHECK-DAG: %[[INIT:.+]] = tensor.empty()152// CHECK-DAG: %[[ARG1_RESHAPE:.+]] = tensor.collapse_shape %[[ARG1]] {{\[}}[0, 1], [2]{{\]}}153// CHECK-DAG: %[[ARG2_RESHAPE:.+]] = tensor.collapse_shape %[[ARG2]] {{\[}}[0], [1, 2, 3]{{\]}}154// CHECK-DAG: %[[INIT_RESHAPE:.+]] = tensor.collapse_shape %[[INIT]] {{\[}}[0], [1, 2, 3], [4], [5, 6], [7]{{\]}}155// CHECK: %[[COLLAPSED_OP:.+]] = linalg.generic156// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]], #[[MAP3]]]157// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]158// CHECK-SAME: ins(%[[ARG0]], %[[ARG1_RESHAPE]], %[[ARG2_RESHAPE]] :159// CHECK-SAME: outs(%[[INIT_RESHAPE]] :160// CHECK: %[[RESULT_RESHAPE:.+]] = tensor.expand_shape %[[COLLAPSED_OP]] {{\[}}[0], [1, 2, 3], [4], [5, 6], [7]{{\]}} output_shape [2, 3, 4, 5, 6, 7, 8, 9]161// CHECK: return %[[RESULT_RESHAPE]]162 163// -----164 165// Dynamic case. Only checks things not covered by `fuse_by_collapsing` test above.166#map0 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d7, d5, d6, d0, d1, d2, d3, d4)>167#map1 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d5, d6, d0)>168#map2 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d4, d1, d2, d3)>169#map3 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d0, d1, d2, d3, d4, d5, d6, d7)>170func.func @fuse_by_collapsing_dynamic(%arg0 : tensor<?x?x?x?x?xi32>,171 %arg1 : tensor<?x?x?xi32>, %arg2 : tensor<?x?x?x?xi32>, %sz0: index, %sz1: index, %sz2: index, %sz3: index, %sz4: index) -> tensor<?x3x?x5x?x7x?x?xi32> {172 %c0 = arith.constant 0 : index173 %c1 = arith.constant 1 : index174 %c2 = arith.constant 2 : index175 %expand = tensor.expand_shape %arg0 [[0], [1, 2], [3], [4, 5, 6], [7]] output_shape [%sz0, 7, %sz1, %sz2, 3, %sz3, 5, %sz4]176 : tensor<?x?x?x?x?xi32> into tensor<?x7x?x?x3x?x5x?xi32>177 %d0 = tensor.dim %arg1, %c2 : tensor<?x?x?xi32>178 %d2 = tensor.dim %arg2, %c2 : tensor<?x?x?x?xi32>179 %d4 = tensor.dim %arg2, %c0 : tensor<?x?x?x?xi32>180 %d6 = tensor.dim %arg1, %c1 : tensor<?x?x?xi32>181 %d7 = tensor.dim %arg0, %c0 : tensor<?x?x?x?x?xi32>182 %init = tensor.empty(%d0, %d2, %d4, %d6, %d7) : tensor<?x3x?x5x?x7x?x?xi32>183 %generic = linalg.generic {184 indexing_maps = [#map0, #map1, #map2, #map3],185 iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel", "parallel", "parallel", "parallel"]}186 ins(%expand, %arg1, %arg2 : tensor<?x7x?x?x3x?x5x?xi32>, tensor<?x?x?xi32>, tensor<?x?x?x?xi32>)187 outs(%init : tensor<?x3x?x5x?x7x?x?xi32>) {188 ^bb0(%b0 : i32, %b1 : i32, %b2 : i32, %b3 : i32):189 %iv0 = linalg.index 0: index190 %iv1 = linalg.index 1: index191 %t0 = arith.addi %iv0, %iv1 : index192 %iv2 = linalg.index 2 : index193 %t1 = arith.addi %t0, %iv2 : index194 %iv3 = linalg.index 3 : index195 %t2 = arith.addi %t1, %iv3 : index196 %iv4 = linalg.index 4 : index197 %t3 = arith.addi %t2, %iv4 : index198 %iv5 = linalg.index 5 : index199 %t4 = arith.addi %t3, %iv5 : index200 %iv6 = linalg.index 6 : index201 %t5 = arith.addi %t4, %iv6 : index202 %iv7 = linalg.index 7 : index203 %t6 = arith.addi %t5, %iv7 : index204 %yield = arith.index_cast %t6 : index to i32205 linalg.yield %yield : i32206 } -> tensor<?x3x?x5x?x7x?x?xi32>207 return %generic : tensor<?x3x?x5x?x7x?x?xi32>208}209// CHECK: func @fuse_by_collapsing_dynamic210// CHECK-SAME: (%[[ARG0:.+]]: tensor<?x?x?x?x?xi32>, %[[SZ0:.+]]: index, %[[SZ1:.+]]: index, %[[SZ2:.+]]: index, %[[SZ3:.+]]: index, %[[SZ4:.+]]: index)211// CHECK-DAG: %[[C2:.+]] = arith.constant 2 : index212// CHECK-DAG: %[[C5:.+]] = arith.constant 5 : index213// CHECK: %[[EXPAND:.+]] = tensor.expand_shape %[[ARG0]]214// CHECK-DAG: %[[D0:.+]] = tensor.dim %[[EXPAND]], %[[C2]]215// CHECK-DAG: %[[D1:.+]] = tensor.dim %[[EXPAND]], %[[C5]]216// CHECK: linalg.generic217// CHECK: %[[IV0:.+]] = linalg.index 1218// CHECK: %[[REM1_IV0:.+]] = arith.remsi %[[IV0]], %[[C5]]219// CHECK: %[[DIV1_IV0:.+]] = arith.divsi %[[IV0]], %[[C5]]220// CHECK: %[[REM2_IV0:.+]] = arith.remsi %[[DIV1_IV0]], %[[D1]]221// CHECK: %[[DIV2_IV0:.+]] = arith.divsi %[[DIV1_IV0]], %[[D1]]222// CHECK: %[[IV1:.+]] = linalg.index 3223// CHECK: %[[REM1_IV1:.+]] = arith.remsi %[[IV1]], %[[D0]]224// CHECK: %[[DIV1_IV1:.+]] = arith.divsi %[[IV1]], %[[D0]]225 226// -----227 228#map0 = affine_map<(d0, d1) -> (d0, d1)>229func.func @fuse_by_collapsing_dynamic_2(%arg0 : tensor<?xf32>, %sz0: index, %sz1: index) -> tensor<?x?xf32> {230 %0 = tensor.expand_shape %arg0 [[0, 1]] output_shape [%sz0, %sz1] : tensor<?xf32> into tensor<?x?xf32>231 %init = tensor.empty(%sz1, %sz0) : tensor<?x?xf32>232 %1 = linalg.generic {233 indexing_maps = [#map0, #map0],234 iterator_types = ["parallel", "parallel"]}235 ins(%0 : tensor<?x?xf32>) 236 outs(%init : tensor<?x?xf32>) {237 ^bb0(%b0 : f32, %b1 : f32):238 %out = arith.negf %b0 : f32239 linalg.yield %out : f32240 } -> tensor<?x?xf32>241 return %1 : tensor<?x?xf32>242}243 244// CHECK-LABEL: func @fuse_by_collapsing_dynamic_2245// CHECK-SAME: %[[ARG0:.+]]: tensor<?xf32>246// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index247// CHECK-DAG: %[[C1:.+]] = arith.constant 1 : index248// CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[ARG0]]249// CHECK-DAG: %[[DIM0:.+]] = tensor.dim %[[EXPANDED]], %[[C0]]250// CHECK-DAG: %[[DIM1:.+]] = tensor.dim %[[EXPANDED]], %[[C1]]251// CHECK: %[[OUT:.+]] = linalg.generic252// CHECK-SAME: ins(%[[ARG0]] : tensor<?xf32>)253// CHECK-SAME: outs(%{{.*}} : tensor<?xf32>)254// CHECK: %[[EXPANDED_1:.+]] = tensor.expand_shape %[[OUT]]255// CHECK-SAME: output_shape [%[[DIM0]], %[[DIM1]]]256// CHECK: return %[[EXPANDED_1]]257 258// -----259 260#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>261#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d3)>262func.func @fuse_reductions(%arg0 : tensor<2x?x5xf32>, %arg1 : tensor<2x5xf32>, %sz0: index) -> tensor<2x5xf32> {263 %0 = tensor.expand_shape %arg0 [[0], [1, 2], [3]] output_shape [2, 6, %sz0, 5] : tensor<2x?x5xf32> into tensor<2x6x?x5xf32>264 %1 = linalg.generic {265 indexing_maps = [#map0, #map1],266 iterator_types = ["parallel", "reduction", "reduction", "parallel"]}267 ins(%0 : tensor<2x6x?x5xf32>) outs(%arg1 : tensor<2x5xf32>) {268 ^bb0(%b0 : f32, %b1 : f32):269 %2 = arith.addf %b0, %b1 : f32270 linalg.yield %2 : f32271 } -> tensor<2x5xf32>272 return %1 : tensor<2x5xf32>273}274// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>275// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d0, d2)>276// CHECK: func @fuse_reductions(277// CHECK-SAME: %[[ARG0:.+]]: tensor<2x?x5xf32>278// CHECK-SAME: %[[ARG1:.+]]: tensor<2x5xf32>279// CHECK-SAME: %[[SZ0:.+]]: index) -> tensor<2x5xf32>280// CHECK: %[[GENERIC:.+]] = linalg.generic281// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]]]282// CHECK-SAME: iterator_types = ["parallel", "reduction", "parallel"]283// CHECK-SAME: ins(%[[ARG0]] : tensor<2x?x5xf32>)284// CHECK-SAME: outs(%[[ARG1]] : tensor<2x5xf32>)285 286// -----287 288// Test no fusion because the folded dimensions are not all preserved.289#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>290#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d1)>291func.func @no_fuse_unpreserved_folding(%arg0 : tensor<2x12x5xf32>, %arg1 : tensor<2x3xf32>) -> tensor<2x3x4x5xf32> {292 %0 = tensor.expand_shape %arg0 [[0], [1, 2], [3]] output_shape [2, 3, 4, 5] : tensor<2x12x5xf32> into tensor<2x3x4x5xf32>293 %init = tensor.empty(): tensor<2x3x4x5xf32>294 %1 = linalg.generic {295 indexing_maps = [#map0, #map1, #map0],296 iterator_types = ["parallel", "parallel", "parallel", "parallel"]}297 ins(%0, %arg1 : tensor<2x3x4x5xf32>, tensor<2x3xf32>) outs(%init : tensor<2x3x4x5xf32>) {298 ^bb0(%b0 : f32, %b1 : f32, %b2 : f32):299 %2 = arith.addf %b0, %b1 : f32300 linalg.yield %2 : f32301 } -> tensor<2x3x4x5xf32>302 return %1 : tensor<2x3x4x5xf32>303}304// CHECK: func @no_fuse_unpreserved_folding305// CHECK-SAME: %[[ARG0:.+]]: tensor<2x12x5xf32>306// CHECK-SAME: %[[ARG1:.+]]: tensor<2x3xf32>307// CHECK: %[[RESHAPE:.+]] = tensor.expand_shape %[[ARG0]]308// CHECK: %[[GENERIC:.+]] = linalg.generic309// CHECK-SAME: ins(%[[RESHAPE]], %[[ARG1]] :310// CHECK: return %[[GENERIC]]311 312// -----313 314// Test no fusion because the folded dimensions are not all preserved.315#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>316#map1 = affine_map<(d0, d1, d2, d3) -> (d0)>317#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d2, d1, d3)>318func.func @no_fuse_unpreserved_folding_transpose(%arg0 : tensor<2x12x5xf32>, %arg1 : tensor<2xf32>) -> tensor<2x4x3x5xf32> {319 %0 = tensor.expand_shape %arg0 [[0], [1, 2], [3]] output_shape [2, 3, 4, 5] : tensor<2x12x5xf32> into tensor<2x3x4x5xf32>320 %init = tensor.empty() : tensor<2x4x3x5xf32>321 %1 = linalg.generic {322 indexing_maps = [#map0, #map1, #map2],323 iterator_types = ["parallel", "parallel", "parallel", "parallel"]}324 ins(%0, %arg1 : tensor<2x3x4x5xf32>, tensor<2xf32>) outs(%init : tensor<2x4x3x5xf32>) {325 ^bb0(%b0 : f32, %b1 : f32, %b2 : f32):326 %2 = arith.addf %b0, %b1 : f32327 linalg.yield %2 : f32328 } -> tensor<2x4x3x5xf32>329 return %1 : tensor<2x4x3x5xf32>330}331// CHECK: func @no_fuse_unpreserved_folding_transpose332// CHECK-SAME: %[[ARG0:.+]]: tensor<2x12x5xf32>333// CHECK-SAME: %[[ARG1:.+]]: tensor<2xf32>334// CHECK: %[[RESHAPE:.+]] = tensor.expand_shape %[[ARG0]]335// CHECK: %[[GENERIC:.+]] = linalg.generic336// CHECK-SAME: ins(%[[RESHAPE]], %[[ARG1]] :337// CHECK: return %[[GENERIC]]338 339// -----340 341// Test no fusion because the iterator types of folded dims are not preserved.342#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>343#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d1)>344#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d3)>345func.func @no_fuse_mismatched_iterator_types(%arg0 : tensor<2x12x5xf32>, %arg1 : tensor<2x3xf32>) -> tensor<2x5xf32> {346 %0 = tensor.expand_shape %arg0 [[0], [1, 2], [3]] output_shape [2, 3, 4, 5] : tensor<2x12x5xf32> into tensor<2x3x4x5xf32>347 %init = tensor.empty() : tensor<2x5xf32>348 %1 = linalg.generic {349 indexing_maps = [#map0, #map1, #map2],350 iterator_types = ["parallel", "reduction", "parallel", "parallel"]}351 ins(%0, %arg1 : tensor<2x3x4x5xf32>, tensor<2x3xf32>) outs(%init : tensor<2x5xf32>) {352 ^bb0(%b0 : f32, %b1 : f32, %b2 : f32):353 %2 = arith.addf %b0, %b1 : f32354 linalg.yield %2 : f32355 } -> tensor<2x5xf32>356 return %1 : tensor<2x5xf32>357}358// CHECK: func @no_fuse_mismatched_iterator_types359// CHECK-SAME: %[[ARG0:.+]]: tensor<2x12x5xf32>360// CHECK-SAME: %[[ARG1:.+]]: tensor<2x3xf32>361// CHECK: %[[RESHAPE:.+]] = tensor.expand_shape %[[ARG0]]362// CHECK: %[[GENERIC:.+]] = linalg.generic363// CHECK-SAME: ins(%[[RESHAPE]], %[[ARG1]] :364// CHECK: return %[[GENERIC]]365 366// -----367 368// Test control of fusion using control function369// Test no fusion because the folded dimensions are not all preserved.370#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1)>371#map1 = affine_map<(d0, d1, d2, d3) -> (d2, d3)>372#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>373func.func @control_fusion(%arg0 : tensor<6xf32>, %arg1 : tensor<20xf32>) -> tensor<2x3x4x5xf32> {374 %0 = tensor.expand_shape %arg0 [[0, 1]] output_shape [2, 3] : tensor<6xf32> into tensor<2x3xf32>375 %1 = tensor.expand_shape %arg1 [[0, 1]] output_shape [4, 5] : tensor<20xf32> into tensor<4x5xf32>376 %init = tensor.empty() : tensor<2x3x4x5xf32>377 %2 = linalg.generic {378 indexing_maps = [#map0, #map1, #map2],379 iterator_types = ["parallel", "parallel", "parallel", "parallel"]}380 ins(%0, %1 : tensor<2x3xf32>, tensor<4x5xf32>) outs(%init : tensor<2x3x4x5xf32>) {381 ^bb0(%b0 : f32, %b1 : f32, %b2 : f32):382 %3 = arith.addf %b0, %b1 : f32383 linalg.yield %3 : f32384 } -> tensor<2x3x4x5xf32>385 return %2 : tensor<2x3x4x5xf32>386}387// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1) -> (d0)>388// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1) -> (d1)>389// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1) -> (d0, d1)>390// CHECK: func @control_fusion(391// CHECK-SAME: %[[ARG0:.+]]: tensor<6xf32>392// CHECK-SAME: %[[ARG1:.+]]: tensor<20xf32>393// CHECK: %[[GENERIC:.+]] = linalg.generic394// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]]]395// CHECK-SAME: iterator_types = ["parallel", "parallel"]396// CHECK-SAME: ins(%[[ARG0]], %[[ARG1]] :397// CHECK-SAME: outs(%{{.+}}: tensor<6x20xf32>)398// CHECK: %[[RESHAPE1:.+]] = tensor.expand_shape %[[GENERIC]] {{\[}}[0], [1, 2]{{\]}} output_shape [6, 4, 5]399// CHECK: %[[RESHAPE2:.+]] = tensor.expand_shape %[[RESHAPE1]] {{\[}}[0, 1], [2], [3]{{\]}} output_shape [2, 3, 4, 5]400// CHECK: return %[[RESHAPE2]]401 402// CONTROL-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>403// CONTROL-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d2)>404// CONTROL-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>405// CONTROL: func @control_fusion(406// CONTROL-SAME: %[[ARG0:.+]]: tensor<6xf32>407// CONTROL-SAME: %[[ARG1:.+]]: tensor<20xf32>408// CONTROL: %[[EXPAND:.+]] = tensor.expand_shape %[[ARG0]]409// CONTROL: %[[INIT:.+]] = tensor.empty()410// CONTROL: %[[INIT_RESHAPE:.+]] = tensor.collapse_shape %[[INIT]] {{\[}}[0], [1], [2, 3]{{\]}}411// CONTROL: %[[GENERIC:.+]] = linalg.generic412// CONTROL-SAME: ins(%[[EXPAND]], %[[ARG1]] :413// CONTROL-SAME: outs(%[[INIT_RESHAPE]] :414// CONTROL: %[[RESULT:.+]] = tensor.expand_shape %[[GENERIC]] {{\[}}[0], [1], [2, 3]{{\]}} output_shape [2, 3, 4, 5]415 416// -----417 418// Corner case that isnt handled currently.419#map = affine_map<(d0) -> (d0)>420func.func @zero_D_test(%arg0: tensor<f32>) -> tensor<1xf32> {421 %0 = tensor.expand_shape %arg0 [] output_shape [1] : tensor<f32> into tensor<1xf32>422 %init = tensor.empty() : tensor<1xf32>423 %1 = linalg.generic {424 indexing_maps = [#map, #map],425 iterator_types = ["parallel"]}426 ins(%0: tensor<1xf32>) outs(%init : tensor<1xf32>) {427 ^bb0(%b0 : f32, %b1 : f32):428 linalg.yield %b0: f32429 } -> tensor<1xf32>430 return %1 : tensor<1xf32>431}432// CHECK: func @zero_D_test433// CHECK-SAME: %[[ARG0:.+]]: tensor<f32>434// CHECK: %[[EXPAND:.+]] = tensor.expand_shape %[[ARG0]]435// CHECK: %[[GENERIC:.+]] = linalg.generic436// CHECK-SAME: ins(%[[EXPAND]] :437// CHECK: return %[[GENERIC]]438 439// -----440 441#map0 = affine_map<(d0, d1, d2, d3) -> (d1, d0, d2, d3)>442#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>443func.func @fuse_only_one_reassociation(%arg0 : tensor<?x?xf32>, %arg1 : tensor<4x?x?x8xf32>, %sz0: index, %sz1: index) -> tensor<4x?x?x8xf32> {444 %0 = tensor.expand_shape %arg0 [[0, 1], [2, 3]] output_shape [%sz0, 4, %sz1, 8] : tensor<?x?xf32> into tensor<?x4x?x8xf32>445 %1 = linalg.generic {446 indexing_maps = [#map0, #map1, #map1],447 iterator_types = ["parallel", "parallel", "parallel", "parallel"]}448 ins(%0, %arg1 : tensor<?x4x?x8xf32>, tensor<4x?x?x8xf32>)449 outs(%arg1 : tensor<4x?x?x8xf32>) {450 ^bb0(%b0: f32, %b1 : f32, %b2 : f32):451 %2 = arith.addf %b0, %b1 : f32452 linalg.yield %2 : f32453 } -> tensor<4x?x?x8xf32>454 return %1 : tensor<4x?x?x8xf32>455}456// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2) -> (d1, d0, d2)>457// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>458// CHECK: func @fuse_only_one_reassociation459// CHECK-SAME: (%[[ARG0:.+]]: tensor<?x?xf32>, %[[ARG1:.+]]: tensor<4x?x?x8xf32>, %[[SZ0:.+]]: index, %[[SZ1:.+]]: index)460// CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index461// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index462// CHECK-DAG: %[[EXPAND_ARG0:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0, 1], [2, 3]{{\]}} output_shape [%[[SZ0]], 4, %[[SZ1]], 8]463// CHECK-DAG: %[[DIM:.+]] = tensor.dim %[[EXPAND_ARG0]], %[[C0]] : tensor<?x4x?x8xf32>464// CHECK-DAG: %[[DIM_2:.+]] = tensor.dim %[[EXPAND_ARG0]], %[[C2]] : tensor<?x4x?x8xf32>465// CHECK-DAG: %[[COLLAPSE_ARG0:.+]] = tensor.collapse_shape %[[EXPAND_ARG0]] {{\[}}[0], [1], [2, 3]{{\]}}466// CHECK-DAG: %[[COLLAPSE_ARG1_0:.+]] = tensor.collapse_shape %[[ARG1]] {{\[}}[0], [1], [2, 3]{{\]}}467// CHECK-DAG: %[[COLLAPSE_ARG1_1:.+]] = tensor.collapse_shape %[[ARG1]] {{\[}}[0], [1], [2, 3]{{\]}}468// CHECK: %[[GENERIC:.+]] = linalg.generic469// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP1]]]470// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"]471// CHECK-SAME: ins(%[[COLLAPSE_ARG0]], %[[COLLAPSE_ARG1_0]] :472// CHECK-SAME: outs(%[[COLLAPSE_ARG1_1]] :473// CHECK: %[[EXPANDED_3:.+]] = tensor.expand_shape %[[GENERIC]] {{\[\[}}0], [1], [2, 3]] output_shape [4, %[[DIM]], %[[DIM_2]], 8] : tensor<4x?x?xf32> into tensor<4x?x?x8xf32>474// CHECK: return %[[EXPANDED_3]]475 476// -----477 478#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3, d1)>479#map1 = affine_map<(d0, d1, d2, d3) -> (d3, d1, d0, d2)>480func.func @fold_non_consecutive_dims(%arg0 : tensor<?x?xi32>, %sz0: index, %sz1: index) -> tensor<?x8x?x4xi32> {481 %c0 = arith.constant 0 : index482 %c2 = arith.constant 2 : index483 %0 = tensor.expand_shape %arg0 [[0, 1], [2, 3]] output_shape [%sz0, 4, %sz1, 8] : tensor<?x?xi32> into tensor<?x4x?x8xi32>484 %d0 = tensor.dim %0, %c0 : tensor<?x4x?x8xi32>485 %d1 = tensor.dim %0, %c2 : tensor<?x4x?x8xi32>486 %init = tensor.empty(%d1, %d0) : tensor<?x8x?x4xi32>487 %1 = linalg.generic {488 indexing_maps = [#map0, #map1],489 iterator_types = ["parallel", "parallel", "parallel", "parallel"]}490 ins(%0 : tensor<?x4x?x8xi32>) outs(%init : tensor<?x8x?x4xi32>) {491 ^bb0(%b0 : i32, %b1 : i32):492 %2 = linalg.index 0 : index493 %3 = linalg.index 1 : index494 %4 = linalg.index 2 : index495 %5 = linalg.index 3 : index496 %6 = arith.addi %2, %3 : index497 %7 = arith.addi %6, %4 : index498 %8 = arith.addi %7, %5 : index499 %9 = arith.index_cast %8 : index to i32500 linalg.yield %9: i32501 } -> tensor<?x8x?x4xi32>502 return %1 : tensor<?x8x?x4xi32>503}504// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1) -> (d0, d1)>505// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1) -> (d1, d0)>506// CHECK: func @fold_non_consecutive_dims(507// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?xi32>, %[[SZ0:.+]]: index, %[[SZ1:.+]]: index)508// CHECK-DAG: %[[C4:.+]] = arith.constant 4 : index509// CHECK-DAG: %[[C8:.+]] = arith.constant 8 : index510// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index511// CHECK-DAG: %[[C2:.+]] = arith.constant 2 : index512// CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[ARG0]] {{\[\[}}0, 1], [2, 3]] output_shape [%[[SZ0]], 4, %[[SZ1]], 8] : tensor<?x?xi32> into tensor<?x4x?x8xi32>513// CHECK-DAG: %[[DIM:.+]] = tensor.dim %[[EXPANDED]], %[[C0]]514// CHECK-DAG: %[[DIM_0:.+]] = tensor.dim %[[EXPANDED]], %[[C2]]515// CHECK: %[[INIT:.+]] = tensor.empty(%[[DIM_0]], %[[DIM]])516// CHECK-DAG: %[[DIM_1:.+]] = tensor.dim %[[EXPANDED]], %[[C0]]517// CHECK-DAG: %[[DIM_2:.+]] = tensor.dim %[[EXPANDED]], %[[C2]]518// CHECK: %[[COLLAPSE_INIT:.+]] = tensor.collapse_shape %[[INIT]] {{\[}}[0, 1], [2, 3]{{\]}}519// CHECK: %[[GENERIC:.+]] = linalg.generic520// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]]]521// CHECK-SAME: iterator_types = ["parallel", "parallel"]522// CHECK-SAME: ins(%[[ARG0]] :523// CHECK-SAME: outs(%[[COLLAPSE_INIT]] :524// CHECK-NEXT: ^bb{{[0-9]}}525// CHECK: %[[ID0:.+]] = linalg.index 0526// CHECK-DAG: %[[T0:.+]] = arith.remsi %[[ID0]], %[[C4]]527// CHECK-DAG: %[[T1:.+]] = arith.divsi %[[ID0]], %[[C4]]528// CHECK: %[[ID1:.+]] = linalg.index 1529// CHECK-DAG: %[[T2:.+]] = arith.remsi %[[ID1]], %[[C8]]530// CHECK-DAG: %[[T3:.+]] = arith.divsi %[[ID1]], %[[C8]]531// CHECK-DAG: %[[T4:.+]] = arith.addi %[[T1]], %[[T2]]532// CHECK-DAG: %[[T5:.+]] = arith.addi %[[T4]], %[[T0]]533// CHECK-DAG: %[[T6:.+]] = arith.addi %[[T5]], %[[T3]]534// CHECK-DAG: %[[T7:.+]] = arith.index_cast %[[T6]]535// CHECK: linalg.yield %[[T7]]536// CHECK: %[[EXPANDED_3:.+]] = tensor.expand_shape %[[GENERIC]] {{\[\[}}0, 1], [2, 3]] output_shape [%[[DIM_2]], 8, %[[DIM_1]], 4] : tensor<?x?xi32> into tensor<?x8x?x4xi32>537// CHECK: return %[[EXPANDED_3]]538 539// -----540 541// None of the folded iteration space dims are contiguous reduction dimensions.542// So no change in the code.543#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3, d1)>544#map1 = affine_map<(d0, d1, d2, d3) -> ()>545func.func @no_fold_non_consecutive_reduction_dims(%arg0 : tensor<?x?xi32>, %sz0: index, %sz1: index) -> tensor<i32> {546 %c0 = arith.constant 0 : index547 %c2 = arith.constant 2 : index548 %0 = tensor.expand_shape %arg0 [[0, 1], [2, 3]] output_shape [%sz0, 4, %sz1, 8] : tensor<?x?xi32> into tensor<?x4x?x8xi32>549 %init = tensor.empty() : tensor<i32>550 %1 = linalg.generic {551 indexing_maps = [#map0, #map1],552 iterator_types = ["reduction", "reduction", "reduction", "reduction"]}553 ins(%0 : tensor<?x4x?x8xi32>) outs(%init : tensor<i32>) {554 ^bb0(%b0 : i32, %b1 : i32):555 %2 = linalg.index 0 : index556 %3 = linalg.index 1 : index557 %4 = linalg.index 2 : index558 %5 = linalg.index 3 : index559 %6 = arith.addi %2, %3 : index560 %7 = arith.addi %6, %4 : index561 %8 = arith.addi %7, %5 : index562 %9 = arith.index_cast %8 : index to i32563 linalg.yield %9: i32564 } -> tensor<i32>565 return %1 : tensor<i32>566}567// CHECK: func @no_fold_non_consecutive_reduction_dims(568// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?xi32>, %[[SZ0:.+]]: index, %[[SZ1:.+]]: index)569// CHECK: %[[EXPAND_ARG0:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0, 1], [2, 3]{{\]}} output_shape [%[[SZ0]], 4, %[[SZ1]], 8]570// CHECK: %[[GENERIC:.+]] = linalg.generic571// CHECK-SAME: ins(%[[EXPAND_ARG0]] :572// CHECK: return %[[GENERIC]]573 574// -----575 576func.func @fuse_by_collapsing_pad(%arg0 : tensor<2x12x5x336x9xi32>) -> tensor<8x3x4x17x6x7x8x14xi32> {577 %expand = tensor.expand_shape %arg0 [[0], [1, 2], [3], [4, 5, 6], [7]] output_shape [2, 3, 4, 5, 6, 7, 8, 9] : tensor<2x12x5x336x9xi32> into tensor<2x3x4x5x6x7x8x9xi32>578 %cst = arith.constant 0 : i32579 %padded_0 = tensor.pad %expand low[1, 0, 0, 8, 0, 0, 0, 3] high[5, 0, 0, 4, 0, 0, 0, 2] {580 ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index,581 %arg5: index, %arg6: index, %arg7: index, %arg8: index):582 tensor.yield %cst : i32583 } : tensor<2x3x4x5x6x7x8x9xi32> to tensor<8x3x4x17x6x7x8x14xi32>584 return %padded_0 : tensor<8x3x4x17x6x7x8x14xi32>585}586// CHECK: func @fuse_by_collapsing_pad(587// CHECK-SAME: %[[ARG0:.+]]: tensor<2x12x5x336x9xi32>)588// CHECK: %[[PAD:.+]] = tensor.pad %[[ARG0]]589// CHECK-SAME: low[1, 0, 8, 0, 3] high[5, 0, 4, 0, 2]590// CHECK: tensor<2x12x5x336x9xi32> to tensor<8x12x17x336x14xi32>591// CHECK: %[[EXPAND:.+]] = tensor.expand_shape %[[PAD]] {{\[}}[0], [1, 2], [3], [4, 5, 6], [7]]592// CHECK-SAME: output_shape [8, 3, 4, 17, 6, 7, 8, 14] : tensor<8x12x17x336x14xi32> into tensor<8x3x4x17x6x7x8x14xi32>593// CHECK: return %[[EXPAND]]594 595// -----596 597func.func @no_fuse_by_collapsing_pad_non_constant_padding(%arg0 : tensor<2x12xi32>) -> tensor<8x3x4xi32> {598 %expand = tensor.expand_shape %arg0 [[0], [1, 2]] output_shape [2, 3, 4] : tensor<2x12xi32> into tensor<2x3x4xi32>599 %cst = arith.constant 0 : i32600 %padded_0 = tensor.pad %expand low[1, 0, 0] high[5, 0, 0] {601 ^bb0(%arg1: index, %arg2: index, %arg3: index):602 %pad_val = arith.index_cast %arg1 : index to i32603 tensor.yield %pad_val : i32604 } : tensor<2x3x4xi32> to tensor<8x3x4xi32>605 return %padded_0 : tensor<8x3x4xi32>606}607// CHECK: func @no_fuse_by_collapsing_pad_non_constant_padding(608// CHECK-SAME: %[[ARG0:.+]]: tensor<2x12xi32>)609// CHECK: %[[EXPAND:.+]] = tensor.expand_shape %[[ARG0]]610// CHECK: %[[PAD:.+]] = tensor.pad %[[EXPAND]]611// CHECK: return %[[PAD]]612 613// -----614 615func.func @no_fuse_by_collapsing_pad(%arg0 : tensor<2x12x5x336x9xi32>) -> tensor<8x5x4x17x6x7x8x14xi32> {616 %expand = tensor.expand_shape %arg0 [[0], [1, 2], [3], [4, 5, 6], [7]] output_shape [2, 3, 4, 5, 6, 7, 8, 9] : tensor<2x12x5x336x9xi32> into tensor<2x3x4x5x6x7x8x9xi32>617 %cst = arith.constant 0 : i32618 %padded_0 = tensor.pad %expand low[1, 2, 0, 8, 0, 0, 0, 3] high[5, 0, 0, 4, 0, 0, 0, 2] {619 ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index,620 %arg5: index, %arg6: index, %arg7: index, %arg8: index):621 tensor.yield %cst : i32622 } : tensor<2x3x4x5x6x7x8x9xi32> to tensor<8x5x4x17x6x7x8x14xi32>623 return %padded_0 : tensor<8x5x4x17x6x7x8x14xi32>624}625// CHECK: func @no_fuse_by_collapsing_pad(626// CHECK-SAME: %[[ARG0:.+]]: tensor<2x12x5x336x9xi32>)627// CHECK: %[[EXPAND_ARG0:.+]] = tensor.expand_shape %[[ARG0]] {{\[}}[0], [1, 2], [3], [4, 5, 6], [7]]628// CHECK-SAME: output_shape [2, 3, 4, 5, 6, 7, 8, 9] : tensor<2x12x5x336x9xi32> into tensor<2x3x4x5x6x7x8x9xi32>629// CHECK: %[[PAD:.+]] = tensor.pad %[[EXPAND_ARG0]]630// CHECK-SAME: low[1, 2, 0, 8, 0, 0, 0, 3] high[5, 0, 0, 4, 0, 0, 0, 2]631// CHECK: tensor<2x3x4x5x6x7x8x9xi32> to tensor<8x5x4x17x6x7x8x14xi32>632// CHECK: return %[[PAD]]633 634// -----635 636func.func @fuse_by_collapsing_dynamic_pad(%arg0 : tensor<?x?x?x?xf32>,637 %s0 : index, %s1 : index, %s2 : index, %s3 : index, %s4 : index, %s5 : index,638 %l0 : index, %l1 : index, %h0 : index, %h1 : index) -> tensor<?x?x?x?x?x?xf32> {639 %expand = tensor.expand_shape %arg0 [[0], [1, 2], [3], [4, 5]] output_shape [%s0, %s1, %s2, %s3, %s4, %s5] : tensor<?x?x?x?xf32> into tensor<?x?x?x?x?x?xf32>640 %cst = arith.constant 0.0 : f32641 %padded_0 = tensor.pad %expand low[%l0, 0, 0, %l1, 0, 0] high[%h0, 0, 0, %h1, 0, 0] {642 ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: index):643 tensor.yield %cst : f32644 } : tensor<?x?x?x?x?x?xf32> to tensor<?x?x?x?x?x?xf32>645 return %padded_0 : tensor<?x?x?x?x?x?xf32>646}647// CHECK-DAG: #[[MAP:.+]] = affine_map<()[s0, s1, s2] -> (s0 + s1 + s2)>648// CHECK: func @fuse_by_collapsing_dynamic_pad(649// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?x?x?xf32>650// CHECK-SAME: %[[S0:.+]]: index, %[[S1:.+]]: index, %[[S2:.+]]: index, %[[S3:.+]]: index, %[[S4:.+]]: index, %[[S5:.+]]: index, %[[L0:.+]]: index, %[[L1:.+]]: index, %[[H0:.+]]: index, %[[H1:.+]]: index651// CHECK: %[[PAD_SIZE0:.+]] = affine.apply #[[MAP]]()[%[[L0]], %[[H0]], %[[S0]]]652// CHECK: %[[PAD_SIZE1:.+]] = affine.apply #[[MAP]]()[%[[L1]], %[[H1]], %[[S3]]]653// CHECK: %[[PAD:.+]] = tensor.pad %[[ARG0]]654// CHECK-SAME: low[%[[L0]], 0, %[[L1]], 0] high[%[[H0]], 0, %[[H1]], 0]655// CHECK: tensor<?x?x?x?xf32> to tensor<?x?x?x?xf32>656// CHECK: %[[EXPAND:.+]] = tensor.expand_shape %[[PAD]] {{\[}}[0], [1, 2], [3], [4, 5]]657// CHECK-SAME: output_shape [%[[PAD_SIZE0]], %[[S1]], %[[S2]], %[[PAD_SIZE1]], %[[S4]], %[[S5]]] : tensor<?x?x?x?xf32> into tensor<?x?x?x?x?x?xf32>658// CHECK: return %[[EXPAND]]659 660// -----661 662func.func @collapse_shape_with_producer_pad(%arg0: tensor<2x3x4x5x6x7x8x9xi32>) -> tensor<8x12x17x336x14xi32> {663 %cst = arith.constant 0 : i32664 %padded = tensor.pad %arg0 low[1, 0, 0, 8, 0, 0, 0, 3] high[5, 0, 0, 4, 0, 0, 0, 2] {665 ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index,666 %arg5: index, %arg6: index, %arg7: index, %arg8: index):667 tensor.yield %cst : i32668 } : tensor<2x3x4x5x6x7x8x9xi32> to tensor<8x3x4x17x6x7x8x14xi32>669 %collapsed = tensor.collapse_shape %padded [[0], [1, 2], [3], [4, 5, 6], [7]]670 : tensor<8x3x4x17x6x7x8x14xi32> into tensor<8x12x17x336x14xi32>671 return %collapsed : tensor<8x12x17x336x14xi32>672}673// CHECK: func @collapse_shape_with_producer_pad674// CHECK-SAME: %[[ARG0:.+]]: tensor<2x3x4x5x6x7x8x9xi32>675// CHECK: %[[COLLAPSE:.+]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0], [1, 2], [3], [4, 5, 6], [7]]676// CHECK: %[[PAD:.+]] = tensor.pad %[[COLLAPSE]] low[1, 0, 8, 0, 3] high[5, 0, 4, 0, 2]677// CHECK: return %[[PAD]]678 679// -----680 681func.func @collapse_shape_with_producer_pad_dynamic(%arg0: tensor<?x?x?x?x?x?xf32>,682 %l0 : index, %l1 : index, %h0 : index, %h1 : index) -> tensor<?x?x?x?xf32> {683 %cst = arith.constant 0.0 : f32684 %padded = tensor.pad %arg0 low[%l0, 0, 0, %l1, 0, 0] high[%h0, 0, 0, %h1, 0, 0] {685 ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index, %arg6: index):686 tensor.yield %cst : f32687 } : tensor<?x?x?x?x?x?xf32> to tensor<?x?x?x?x?x?xf32>688 %collapsed = tensor.collapse_shape %padded [[0], [1, 2], [3], [4, 5]]689 : tensor<?x?x?x?x?x?xf32> into tensor<?x?x?x?xf32>690 return %collapsed : tensor<?x?x?x?xf32>691}692// CHECK: func @collapse_shape_with_producer_pad_dynamic693// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?x?x?x?x?xf32>694// CHECK-SAME: %[[L0:.+]]: index, %[[L1:.+]]: index, %[[H0:.+]]: index, %[[H1:.+]]: index695// CHECK: %[[COLLAPSE:.+]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0], [1, 2], [3], [4, 5]]696// CHECK: %[[PAD:.+]] = tensor.pad %[[COLLAPSE]] low[%[[L0]], 0, %[[L1]], 0] high[%[[H0]], 0, %[[H1]], 0]697// CHECK: return %[[PAD]]698 699// -----700 701func.func @collapse_shape_with_producer_pad_non_constant_padding(%arg0 : tensor<2x3x4xi32>) -> tensor<8x12xi32> {702 %cst = arith.constant 0 : i32703 %padded_0 = tensor.pad %arg0 low[1, 0, 0] high[5, 0, 0] {704 ^bb0(%arg1: index, %arg2: index, %arg3: index):705 %pad_val = arith.index_cast %arg1 : index to i32706 tensor.yield %pad_val : i32707 } : tensor<2x3x4xi32> to tensor<8x3x4xi32>708 %collapsed = tensor.collapse_shape %padded_0 [[0], [1, 2]] : tensor<8x3x4xi32> into tensor<8x12xi32>709 return %collapsed : tensor<8x12xi32>710}711// CHECK: func @collapse_shape_with_producer_pad_non_constant_padding(712// CHECK-SAME: %[[ARG0:.+]]: tensor<2x3x4xi32>)713// CHECK: %[[PAD:.+]] = tensor.pad %[[ARG0]]714// CHECK: %[[COLLAPSED:.+]] = tensor.collapse_shape %[[PAD]]715// CHECK: return %[[COLLAPSED]]716 717// -----718// Static problem sizes. Checks all aspects of fusion by collapsing with bubbling up collapse shapes.719#map0 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d0, d1, d2, d3, d4, d5, d6, d7)>720#map1 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d0, d1, d2)>721#map2 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d3, d4, d5, d6)>722#map3 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d0, d1, d2, d3, d4, d5, d6, d7)>723#map4 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d1, d2, d0, d7, d3, d4, d5, d6)>724func.func @fuse_by_collapsing_bubblecollapse(%arg0 : tensor<2x3x4x5x6x7x8x9xi32>,725 %arg1 : tensor<2x3x4xi32>, %arg2 : tensor<5x6x7x8xi32>) -> (tensor<2x12x5x336x9xi32>, tensor<12x2x9x5x336xi32>) {726 %init_0 = tensor.empty() : tensor<2x3x4x5x6x7x8x9xi32>727 %init_1 = tensor.empty() : tensor<3x4x2x9x5x6x7x8xi32>728 %generic:2 = linalg.generic {729 indexing_maps = [#map0, #map1, #map2, #map3, #map4],730 iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel", "parallel", "parallel", "parallel"]}731 ins(%arg0, %arg1, %arg2 : tensor<2x3x4x5x6x7x8x9xi32>, tensor<2x3x4xi32>, tensor<5x6x7x8xi32>)732 outs(%init_0, %init_1 : tensor<2x3x4x5x6x7x8x9xi32>, tensor<3x4x2x9x5x6x7x8xi32>) {733 ^bb0(%b0 : i32, %b1 : i32, %b2 : i32, %b3 : i32, %b4 : i32):734 %t0 = arith.addi %b0, %b1 : i32735 %t1 = arith.addi %t0, %b2 : i32736 linalg.yield %t1, %t1 : i32, i32737 } -> (tensor<2x3x4x5x6x7x8x9xi32>, tensor<3x4x2x9x5x6x7x8xi32>)738 %collapse_1 = tensor.collapse_shape %generic#0 [[0], [1, 2], [3], [4, 5, 6], [7]] : tensor<2x3x4x5x6x7x8x9xi32> into tensor<2x12x5x336x9xi32>739 %collapse_2 = tensor.collapse_shape %generic#1 [[0, 1], [2], [3], [4], [5, 6, 7]] : tensor<3x4x2x9x5x6x7x8xi32> into tensor<12x2x9x5x336xi32>740 return %collapse_1, %collapse_2 : tensor<2x12x5x336x9xi32>, tensor<12x2x9x5x336xi32>741}742// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>743// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1)>744// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d2, d3)>745// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d1, d0, d4, d2, d3)>746// CHECK: func @fuse_by_collapsing_bubblecollapse(747// CHECK-SAME: %[[ARG0:.+]]: tensor<2x3x4x5x6x7x8x9xi32>748// CHECK-SAME: %[[ARG1:.+]]: tensor<2x3x4xi32>749// CHECK-SAME: %[[ARG2:.+]]: tensor<5x6x7x8xi32>750// CHECK-DAG: %[[INIT0:.+]] = tensor.empty() : tensor<2x3x4x5x6x7x8x9xi32>751// CHECK-DAG: %[[INIT1:.+]] = tensor.empty() : tensor<3x4x2x9x5x6x7x8xi32>752// CHECK-DAG: %[[ARG0_RESHAPE:.+]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0], [1, 2], [3], [4, 5, 6], [7]{{\]}}753// CHECK-DAG: %[[ARG1_RESHAPE:.+]] = tensor.collapse_shape %[[ARG1]] {{\[}}[0], [1, 2]{{\]}}754// CHECK-DAG: %[[ARG2_RESHAPE:.+]] = tensor.collapse_shape %[[ARG2]] {{\[}}[0], [1, 2, 3]{{\]}}755// CHECK-DAG: %[[INIT0_RESHAPE:.+]] = tensor.collapse_shape %[[INIT0]] {{\[}}[0], [1, 2], [3], [4, 5, 6], [7]{{\]}}756// CHECK-DAG: %[[INIT1_RESHAPE:.+]] = tensor.collapse_shape %[[INIT1]] {{\[}}[0, 1], [2], [3], [4], [5, 6, 7]{{\]}}757// CHECK: %[[COLLAPSED_OP:.+]]:2 = linalg.generic758// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]], #[[MAP0]], #[[MAP3]]]759// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]760// CHECK-SAME: ins(%[[ARG0_RESHAPE]], %[[ARG1_RESHAPE]], %[[ARG2_RESHAPE]] :761// CHECK-SAME: outs(%[[INIT0_RESHAPE]], %[[INIT1_RESHAPE]] :762// CHECK: return %[[COLLAPSED_OP]]#0, %[[COLLAPSED_OP]]#1763 764// -----765 766#map0 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d0, d1, d2, d3, d4, d5, d6, d7)>767#map1 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d0, d1, d2)>768#map2 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d3, d4, d5, d6)>769#map3 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d0, d1, d2, d3, d4, d5, d6, d7)>770func.func @fuse_by_collapsing_indexing_op_bubblecollapse(%arg0 : tensor<2x3x4x5x6x7x8x9xi32>,771 %arg1 : tensor<2x3x4xi32>, %arg2 : tensor<5x6x7x8xi32>) -> tensor<2x12x5x336x9xi32> {772 %init = tensor.empty() : tensor<2x3x4x5x6x7x8x9xi32>773 %generic = linalg.generic {774 indexing_maps = [#map0, #map1, #map2, #map3],775 iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel", "parallel", "parallel", "parallel"]}776 ins(%arg0, %arg1, %arg2 : tensor<2x3x4x5x6x7x8x9xi32>, tensor<2x3x4xi32>, tensor<5x6x7x8xi32>)777 outs(%init : tensor<2x3x4x5x6x7x8x9xi32>) {778 ^bb0(%b0 : i32, %b1 : i32, %b2 : i32, %b3 : i32):779 %iv0 = linalg.index 0: index780 %iv1 = linalg.index 1: index781 %t0 = arith.addi %iv0, %iv1 : index782 %iv2 = linalg.index 2 : index783 %t1 = arith.addi %t0, %iv2 : index784 %iv3 = linalg.index 3 : index785 %t2 = arith.addi %t1, %iv3 : index786 %iv4 = linalg.index 4 : index787 %t3 = arith.addi %t2, %iv4 : index788 %iv5 = linalg.index 5 : index789 %t4 = arith.addi %t3, %iv5 : index790 %iv6 = linalg.index 6 : index791 %t5 = arith.addi %t4, %iv6 : index792 %iv7 = linalg.index 7 : index793 %t6 = arith.addi %t5, %iv7 : index794 %yield = arith.index_cast %t6 : index to i32795 linalg.yield %yield : i32796 } -> tensor<2x3x4x5x6x7x8x9xi32>797 %collapse = tensor.collapse_shape %generic [[0], [1, 2], [3], [4, 5, 6], [7]] : tensor<2x3x4x5x6x7x8x9xi32> into tensor<2x12x5x336x9xi32>798 return %collapse : tensor<2x12x5x336x9xi32>799}800// CHECK-LABEL: func @fuse_by_collapsing_indexing_op_bubblecollapse(801// CHECK-DAG: %[[C4:.+]] = arith.constant 4 : index802// CHECK-DAG: %[[C8:.+]] = arith.constant 8 : index803// CHECK-DAG: %[[C7:.+]] = arith.constant 7 : index804// CHECK: %[[IV0:.+]] = linalg.index 0805// CHECK: %[[IV1:.+]] = linalg.index 1806// CHECK: %[[REM_IV1:.+]] = arith.remsi %[[IV1]], %[[C4]]807// CHECK: %[[DIV_IV1:.+]] = arith.divsi %[[IV1]], %[[C4]]808// CHECK: %[[IV2:.+]] = linalg.index 2809// CHECK: %[[IV3:.+]] = linalg.index 3810// CHECK: %[[REM1_IV3:.+]] = arith.remsi %[[IV3]], %[[C8]]811// CHECK: %[[DIV1_IV3:.+]] = arith.divsi %[[IV3]], %[[C8]]812// CHECK: %[[REM2_IV3:.+]] = arith.remsi %[[DIV1_IV3]], %[[C7]]813// CHECK: %[[DIV2_IV3:.+]] = arith.divsi %[[DIV1_IV3]], %[[C7]]814// CHECK: %[[IV4:.+]] = linalg.index 4815// CHECK: %[[T0:.+]] = arith.addi %[[IV0]], %[[DIV_IV1]]816// CHECK: %[[T1:.+]] = arith.addi %[[T0]], %[[REM_IV1]]817// CHECK: %[[T2:.+]] = arith.addi %[[T1]], %[[IV2]]818// CHECK: %[[T3:.+]] = arith.addi %[[T2]], %[[DIV2_IV3]]819// CHECK: %[[T4:.+]] = arith.addi %[[T3]], %[[REM2_IV3]]820// CHECK: %[[T5:.+]] = arith.addi %[[T4]], %[[REM1_IV3]]821// CHECK: %[[T6:.+]] = arith.addi %[[T5]], %[[IV4]]822// CHECK: %[[YIELD:.+]] = arith.index_cast %[[T6]]823// CHECK: linalg.yield %[[YIELD]]824 825// -----826 827#map0 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d7, d5, d6, d0, d1, d2, d3, d4)>828#map1 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d5, d6, d0)>829#map2 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d4, d1, d2, d3)>830#map3 = affine_map<(d0, d1, d2, d3, d4, d5, d6, d7) -> (d0, d1, d2, d3, d4, d5, d6, d7)>831func.func @fuse_by_collapsing_change_reshape_order_bubblecollapse(%arg0 : tensor<9x7x8x2x3x4x5x6xi32>,832 %arg1 : tensor<7x8x2xi32>, %arg2 : tensor<6x3x4x5xi32>) -> tensor<2x60x6x56x9xi32> {833 %init = tensor.empty() : tensor<2x3x4x5x6x7x8x9xi32>834 %generic = linalg.generic {835 indexing_maps = [#map0, #map1, #map2, #map3],836 iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel", "parallel", "parallel", "parallel"]}837 ins(%arg0, %arg1, %arg2 : tensor<9x7x8x2x3x4x5x6xi32>, tensor<7x8x2xi32>, tensor<6x3x4x5xi32>)838 outs(%init : tensor<2x3x4x5x6x7x8x9xi32>) {839 ^bb0(%b0 : i32, %b1 : i32, %b2 : i32, %b3 : i32):840 %t0 = arith.addi %b0, %b1 : i32841 %t1 = arith.addi %t0, %b2 : i32842 linalg.yield %t1 : i32843 } -> tensor<2x3x4x5x6x7x8x9xi32>844 %collapse = tensor.collapse_shape %generic [[0], [1, 2, 3], [4], [5, 6], [7]] : tensor<2x3x4x5x6x7x8x9xi32> into tensor<2x60x6x56x9xi32>845 return %collapse : tensor<2x60x6x56x9xi32>846}847 848// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d4, d3, d0, d1, d2)>849// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d3, d0)>850// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d2, d1)>851// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>852// CHECK: func @fuse_by_collapsing_change_reshape_order_bubblecollapse(853// CHECK-SAME: %[[ARG0:.+]]: tensor<9x7x8x2x3x4x5x6xi32>854// CHECK-SAME: %[[ARG1:.+]]: tensor<7x8x2xi32>855// CHECK-SAME: %[[ARG2:.+]]: tensor<6x3x4x5xi32>856// CHECK-DAG: %[[INIT:.+]] = tensor.empty()857// CHECK-DAG: %[[ARG0_RESHAPE:.+]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0], [1, 2], [3], [4, 5, 6], [7]{{\]}}858// CHECK-DAG: %[[ARG1_RESHAPE:.+]] = tensor.collapse_shape %[[ARG1]] {{\[}}[0, 1], [2]{{\]}}859// CHECK-DAG: %[[ARG2_RESHAPE:.+]] = tensor.collapse_shape %[[ARG2]] {{\[}}[0], [1, 2, 3]{{\]}}860// CHECK-DAG: %[[INIT_RESHAPE:.+]] = tensor.collapse_shape %[[INIT]] {{\[}}[0], [1, 2, 3], [4], [5, 6], [7]{{\]}}861// CHECK: %[[COLLAPSED_OP:.+]] = linalg.generic862// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]], #[[MAP2]], #[[MAP3]]]863// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]864// CHECK-SAME: ins(%[[ARG0_RESHAPE]], %[[ARG1_RESHAPE]], %[[ARG2_RESHAPE]] :865// CHECK-SAME: outs(%[[INIT_RESHAPE]] :866// CHECK: return %[[COLLAPSED_OP]]867 868// CONTROL: func @fuse_by_collapsing_change_reshape_order_bubblecollapse(869// CONTROL-SAME: %[[ARG0:.+]]: tensor<9x7x8x2x3x4x5x6xi32>870// CONTROL-SAME: %[[ARG1:.+]]: tensor<7x8x2xi32>871// CONTROL-SAME: %[[ARG2:.+]]: tensor<6x3x4x5xi32>872// CONTROL: %[[GENERIC:.+]] = linalg.generic873// CONTROL-SAME: ins(%[[ARG0]],874// CONTROL: %[[COLLAPSE:.+]] = tensor.collapse_shape %[[GENERIC]]875// CONTROL: return %[[COLLAPSE]]876 877// -----878 879// Check that new ops are inserted at `%0` because `%0` is also used by `tensor.dim`.880#map0 = affine_map<(d0, d1) -> (d0, d1)>881func.func @fuse_by_collapsing_correct_insertion(%arg0 : tensor<?x?xf32>,882 %sz0: index, %sz1: index) -> (tensor<?xf32>, index) {883 %c0 = arith.constant 0 : index884 %init = tensor.empty(%sz1, %sz0) : tensor<?x?xf32>885 %0 = linalg.generic {886 indexing_maps = [#map0, #map0],887 iterator_types = ["parallel", "parallel"]}888 ins(%arg0 : tensor<?x?xf32>)889 outs(%init : tensor<?x?xf32>) {890 ^bb0(%b0 : f32, %b1 : f32):891 %out = arith.negf %b0 : f32892 linalg.yield %out : f32893 } -> tensor<?x?xf32>894 %dim = tensor.dim %0, %c0 : tensor<?x?xf32>895 %1 = tensor.collapse_shape %0 [[0, 1]] : tensor<?x?xf32> into tensor<?xf32>896 return %1, %dim : tensor<?xf32>, index897}898 899// CHECK-LABEL: func @fuse_by_collapsing_correct_insertion900// CHECK-SAME: %[[ARG0:.+]]: tensor<?x?xf32>901// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index902// CHECK: %[[COLLAPSE:.+]] = tensor.collapse_shape %[[ARG0]]903// CHECK: %[[OUT:.+]] = linalg.generic904// CHECK-SAME: ins(%[[COLLAPSE]] : tensor<?xf32>)905// CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[OUT]]906// CHECK: %[[DIM:.+]] = tensor.dim %[[EXPANDED]], %[[C0]]907// CHECK: return %[[OUT]], %[[DIM]]908 909// -----910 911#map = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2, d3, d4)>912#map1 = affine_map<(d0, d1, d2, d3, d4) -> (d0, d3, d4, d1, d2)>913func.func @partial_fuse_by_collapsing(%arg0: tensor<4x?x32x128x192xf16>, %arg1: tensor<4x128x192x?x32xf32>) -> tensor<512x192x?xf32> {914 %0 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"]} ins(%arg0 : tensor<4x?x32x128x192xf16>) outs(%arg1 : tensor<4x128x192x?x32xf32>) {915 ^bb0(%in: f16, %out: f32):916 linalg.yield %out : f32917 } -> tensor<4x128x192x?x32xf32>918 %collapsed = tensor.collapse_shape %0 [[0, 1], [2], [3, 4]] : tensor<4x128x192x?x32xf32> into tensor<512x192x?xf32>919 return %collapsed : tensor<512x192x?xf32>920}921// CHECK-LABEL: func @partial_fuse_by_collapsing922// CHECK-SAME: %[[ARG0:.+]]: tensor<4x?x32x128x192xf16>923// CHECK-SAME: %[[ARG1:.+]]: tensor<4x128x192x?x32xf32>924// CHECK-DAG: %[[COLLAPSED0:.+]] = tensor.collapse_shape %[[ARG0]]925// CHECK-SAME: tensor<4x?x32x128x192xf16> into tensor<4x?x128x192xf16>926// CHECK-DAG: %[[COLLAPSED1:.+]] = tensor.collapse_shape %[[ARG1]]927// CHECK-SAME: tensor<4x128x192x?x32xf32> into tensor<4x128x192x?xf32>928// CHECK: %[[GENERIC:.+]] = linalg.generic929// CHECK-SAME: ins(%[[COLLAPSED0]]930// CHECK-SAME: outs(%[[COLLAPSED1]]931// CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[GENERIC]]932// CHECK-SAME: tensor<4x128x192x?xf32> into tensor<4x128x192x?x32xf32>933// CHECK: %[[COLLAPSED:.+]] = tensor.collapse_shape %[[EXPANDED]]934// CHECK-SAME: tensor<4x128x192x?x32xf32> into tensor<512x192x?xf32>935// CHECK: return %[[COLLAPSED]] : tensor<512x192x?xf32>936