1// RUN: mlir-opt %s -linalg-fusion-for-tensor-ops -split-input-file | FileCheck %s 2 3#map0 = affine_map<(d0, d1, d2) -> (d2, d0, d1)> 4#map1 = affine_map<(d0, d1, d2) -> (d1, d2, d0)> 5func @generic_op_reshape_producer_fusion(%arg0 : tensor<?x?x?x?xf32>, 6 %arg1 : tensor<?x?x?xf32>) -> 7 tensor<?x?x?xf32> 8{ 9 %0 = linalg.tensor_reshape %arg0 [affine_map<(i, j, k, l) -> (i)>, 10 affine_map<(i, j, k, l) -> (j, k)>, 11 affine_map<(i, j, k, l) -> (l)>] : 12 tensor<?x?x?x?xf32> into tensor<?x?x?xf32> 13 %1 = linalg.generic { 14 indexing_maps = [#map0, #map1, #map1], 15 iterator_types = ["parallel", "parallel", "parallel"]} 16 ins(%0, %arg1 : tensor<?x?x?xf32>, tensor<?x?x?xf32>) { 17 ^bb0(%arg3: f32, %arg4: f32): // no predecessors 18 %1 = mulf %arg3, %arg4 : f32 19 linalg.yield %1 : f32 20 } -> tensor<?x?x?xf32> 21 return %1 : tensor<?x?x?xf32> 22} 23 24// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3) -> (d0)> 25// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d1)> 26// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d2, d3)> 27// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3) -> (d3, d0, d1, d2)> 28// CHECK-DAG: #[[MAP4:.+]] = affine_map<(d0, d1, d2, d3) -> (d2, d3, d0, d1)> 29// CHECK: func @generic_op_reshape_producer_fusion 30// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?x?x?xf32> 31// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x?x?xf32> 32// CHECK: %[[T0:.+]] = linalg.tensor_reshape %[[ARG1]] 33// CHECK-SAME: [#[[MAP0]], #[[MAP1]], #[[MAP2]]] 34// CHECK-SAME: tensor<?x?x?xf32> into tensor<?x?x?x?xf32> 35// CHECK: %[[T1:.+]] = linalg.generic 36// CHECK-SAME: indexing_maps = [#[[MAP3]], #[[MAP4]], #[[MAP4]]] 37// CHECK-SAME: ["parallel", "parallel", "parallel", "parallel"] 38// CHECK-SAME: ins(%[[ARG0]], %[[T0]] : tensor<?x?x?x?xf32>, tensor<?x?x?x?xf32>) 39// CHECK: %[[T2:.+]] = linalg.tensor_reshape 40// CHECK-SAME: [#[[MAP0]], #[[MAP1]], #[[MAP2]]] 41// CHECK-SAME: tensor<?x?x?x?xf32> into tensor<?x?x?xf32> 42// CHECK: return %[[T2]] 43 44// ----- 45 46#map0 = affine_map<(d0, d1) -> (d0, d1)> 47func @generic_op_reshape_consumer_fusion(%arg0 : tensor<?x?xf32>, 48 %arg1 : tensor<?x?xf32>) -> 49 tensor<?x?x4x5xf32> 50{ 51 %0 = linalg.generic { 52 indexing_maps = [#map0, #map0, #map0], 53 iterator_types = ["parallel", "parallel"]} 54 ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>) { 55 ^bb0(%arg3: f32, %arg4: f32): // no predecessors 56 %1 = mulf %arg3, %arg4 : f32 57 linalg.yield %1 : f32 58 } -> tensor<?x?xf32> 59 %1 = linalg.tensor_reshape %0 [affine_map<(i, j, k, l) -> (i)>, 60 affine_map<(i, j, k, l) -> (j, k, l)>] : 61 tensor<?x?xf32> into tensor<?x?x4x5xf32> 62 return %1 : tensor<?x?x4x5xf32> 63} 64 65// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3) -> (d0)> 66// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d1, d2, d3)> 67// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)> 68// CHECK: func @generic_op_reshape_consumer_fusion 69// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?xf32> 70// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x?xf32> 71// CHECK: %[[T0:.+]] = linalg.tensor_reshape %[[ARG0]] 72// CHECK-SAME: [#[[MAP0]], #[[MAP1]]] 73// CHECK-SAME: tensor<?x?xf32> into tensor<?x?x4x5xf32> 74// CHECK: %[[T1:.+]] = linalg.tensor_reshape %[[ARG1]] 75// CHECK-SAME: [#[[MAP0]], #[[MAP1]]] 76// CHECK-SAME: tensor<?x?xf32> into tensor<?x?x4x5xf32> 77// CHECK: %[[T2:.+]] = linalg.generic 78// CHECK-SAME: indexing_maps = [#[[MAP2]], #[[MAP2]], #[[MAP2]]] 79// CHECK-SAME: ["parallel", "parallel", "parallel", "parallel"] 80// CHECK-SAME: ins(%[[T0]], %[[T1]] : tensor<?x?x4x5xf32>, tensor<?x?x4x5xf32>) 81// CHECK: return %[[T2]] : tensor<?x?x4x5xf32> 82 83 84// ----- 85 86func @reshape_as_consumer_permutation 87 (%a : tensor<?x?x?xf32>, %b : tensor<?x?xf32>) 88 -> tensor<?x?x?x?x?x?xf32> { 89 %c = linalg.generic { 90 indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d0, d2)>, 91 affine_map<(d0, d1, d2) -> (d1, d2)>, 92 affine_map<(d0, d1, d2) -> (d0, d2, d1)>], 93 iterator_types = ["parallel", "parallel", "parallel"]} 94 ins(%a, %b : tensor<?x?x?xf32>, tensor<?x?xf32>) { 95 ^bb0(%arg0 : f32, %arg1: f32): 96 %1 = addf %arg0, %arg1 : f32 97 linalg.yield %1 : f32 98 } -> tensor<?x?x?xf32> 99 %d = linalg.tensor_reshape %c 100 [affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1)>, 101 affine_map<(d0, d1, d2, d3, d4, d5) -> (d2)>, 102 affine_map<(d0, d1, d2, d3, d4, d5) -> (d3, d4, d5)>] 103 : tensor<?x?x?xf32> into tensor<?x?x?x?x?x?xf32> 104 return %d : tensor<?x?x?x?x?x?xf32> 105} 106// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2)> 107// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d3, d4)> 108// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d5)> 109// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)> 110// CHECK-DAG: #[[MAP4:.+]] = affine_map<(d0, d1, d2, d3) -> (d3)> 111// CHECK-DAG: #[[MAP5:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d2, d3, d4, d0, d1, d5)> 112// CHECK-DAG: #[[MAP6:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d2, d3, d4, d5)> 113// CHECK-DAG: #[[MAP7:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d5, d2, d3, d4)> 114// CHECK: func @reshape_as_consumer_permutation 115// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?x?xf32> 116// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x?xf32> 117// CHECK: %[[T0:.+]] = linalg.tensor_reshape %[[ARG0]] 118// CHECK-SAME: [#[[MAP0]], #[[MAP1]], #[[MAP2]]] 119// CHECK-SAME: tensor<?x?x?xf32> into tensor<?x?x?x?x?x?xf32> 120// CHECK: %[[T1:.+]] = linalg.tensor_reshape %[[ARG1]] 121// CHECK-SAME: [#[[MAP3]], #[[MAP4]]] 122// CHECK-SAME: tensor<?x?xf32> into tensor<?x?x?x?xf32> 123// CHECK: %[[T2:.+]] = linalg.generic 124// CHECK-SAME: indexing_maps = [#[[MAP5]], #[[MAP6]], #[[MAP7]]] 125// CHECK-SAME: ["parallel", "parallel", "parallel", "parallel", "parallel", "parallel"] 126// CHECK-SAME: ins(%[[T0]], %[[T1]] : tensor<?x?x?x?x?x?xf32>, tensor<?x?x?x?xf32>) 127// CHECK: return %[[T2]] : tensor<?x?x?x?x?x?xf32> 128 129// ----- 130 131#map0 = affine_map<(d0, d1) -> (d0, d1)> 132#map1 = affine_map<(d0, d1, d2) -> (d0, d1)> 133#map2 = affine_map<(d0, d1, d2) -> (d2)> 134 135func @generic_op_reshape_consumer_static(%arg0: tensor<264x4xf32>) 136 -> tensor<8x33x4xf32> { 137 %cst = constant dense<2.000000e+00> : tensor<264x4xf32> 138 %0 = linalg.generic { 139 indexing_maps = [#map0, #map0, #map0], 140 iterator_types = ["parallel", "parallel"]} 141 ins(%arg0, %cst : tensor<264x4xf32>, tensor<264x4xf32>) { 142 ^bb0(%arg1: f32, %arg2: f32): // no predecessors 143 %2 = mulf %arg1, %arg2 : f32 144 linalg.yield %2 : f32 145 } -> tensor<264x4xf32> 146 %1 = linalg.tensor_reshape %0 [#map1, #map2] : 147 tensor<264x4xf32> into tensor<8x33x4xf32> 148 return %1 : tensor<8x33x4xf32> 149} 150 151// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)> 152// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d2)> 153// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)> 154// CHECK: func @generic_op_reshape_consumer_static 155// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<264x4xf32> 156// CHECK: %[[T0:.+]] = linalg.tensor_reshape %[[ARG0]] 157// CHECK-SAME: [#[[MAP0]], #[[MAP1]]] 158// CHECK-SAME: tensor<264x4xf32> into tensor<8x33x4xf32> 159// CHECK: %[[T1:.+]] = linalg.generic 160// CHECK-SAME: indexing_maps = [#[[MAP2]], #[[MAP2]]] 161// CHECK-SAME: ["parallel", "parallel", "parallel"] 162// CHECK-SAME: ins(%[[T0]] : tensor<8x33x4xf32>) 163// CHECK: return %[[T1]] : tensor<8x33x4xf32> 164 165// ----- 166 167func @scalar_reshape(%arg0 : tensor<1x10xf32>, %arg1 : tensor<1xf32>) 168 -> tensor<1x10xf32> { 169 %0 = linalg.tensor_reshape %arg1 [] : tensor<1xf32> into tensor<f32> 170 %1 = linalg.generic 171 {indexing_maps = [affine_map<(d0) -> ()>, affine_map<(d0) -> (d0)>], 172 iterator_types = ["parallel"]} ins(%0 : tensor<f32>) { 173 ^bb0(%arg2: f32): // no predecessors 174 linalg.yield %arg2 : f32 175 } -> tensor<10xf32> 176 %2 = linalg.tensor_reshape %1 [affine_map<(d0, d1) -> (d0, d1)>] 177 : tensor<10xf32> into tensor<1x10xf32> 178 return %2 : tensor<1x10xf32> 179} 180 181// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1) -> ()> 182// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1) -> (d0, d1)> 183// CHECK: func @scalar_reshape 184// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<1x10xf32> 185// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: tensor<1xf32> 186// CHECK: %[[T0:.+]] = linalg.tensor_reshape %[[ARG1]] [] 187// CHECK-SAME: tensor<1xf32> into tensor<f32> 188// CHECK: %[[T1:.+]] = linalg.generic 189// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]]] 190// CHECK-SAME: iterator_types = ["parallel", "parallel"] 191// CHECK-SAME: ins(%[[T0]] : tensor<f32>) 192// CHECK: return %[[T1]] : tensor<1x10xf32> 193 194// ----- 195 196#map0 = affine_map<(d0, d1, d2) -> (d2, d0, d1)> 197#map1 = affine_map<(d0, d1, d2) -> (d1, d2, d0)> 198func @indexed_generic_op_reshape_producer_fusion(%arg0 : tensor<?x?x4x?xi32>, 199 %arg1 : tensor<?x?x?xi32>) -> 200 tensor<?x?x?xi32> 201{ 202 %0 = linalg.tensor_reshape %arg0 [affine_map<(i, j, k, l) -> (i)>, 203 affine_map<(i, j, k, l) -> (j, k)>, 204 affine_map<(i, j, k, l) -> (l)>] : 205 tensor<?x?x4x?xi32> into tensor<?x?x?xi32> 206 %1 = linalg.indexed_generic { 207 indexing_maps = [#map0, #map1, #map1], 208 iterator_types = ["parallel", "parallel", "parallel"]} 209 ins(%0, %arg1 : tensor<?x?x?xi32>, tensor<?x?x?xi32>) { 210 ^bb0(%arg3 : index, %arg4 : index, %arg5 : index, %arg6: i32, %arg7: i32): 211 %1 = muli %arg6, %arg7 : i32 212 %2 = index_cast %arg3 : index to i32 213 %3 = addi %1, %2 : i32 214 %4 = index_cast %arg4 : index to i32 215 %5 = addi %3, %4 : i32 216 %6 = index_cast %arg5 : index to i32 217 %7 = addi %5, %6 : i32 218 linalg.yield %7 : i32 219 } -> tensor<?x?x?xi32> 220 return %1 : tensor<?x?x?xi32> 221} 222 223// The generic op version of the test check for the op structure. Only 224// checking the op body here. 225// CHECK: #[[MAP:.+]] = affine_map<(d0, d1) -> (d0 * 4 + d1)> 226// CHECK: func @indexed_generic_op_reshape_producer_fusion 227// CHECK: linalg.indexed_generic 228// CHECK: ^{{.*}}( 229// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: index, %[[ARG3:[a-zA-Z0-9]+]]: index, 230// CHECK-SAME: %[[ARG4:[a-zA-Z0-9]+]]: index, %[[ARG5:[a-zA-Z0-9]+]]: index, 231// CHECK-SAME: %[[ARG6:[a-zA-Z0-9]+]]: i32, %[[ARG7:[a-zA-Z0-9]+]]: i32) 232// CHECK: %[[T3:.+]] = affine.apply #[[MAP]](%[[ARG2]], %[[ARG3]]) 233// CHECK: %[[T4:.+]] = muli %[[ARG6]], %[[ARG7]] 234// CHECK: %[[T5:.+]] = index_cast %[[T3]] 235// CHECK: %[[T6:.+]] = addi %[[T4]], %[[T5]] 236// CHECK: %[[T7:.+]] = index_cast %[[ARG4]] 237// CHECK: %[[T8:.+]] = addi %[[T6]], %[[T7]] 238// CHECK: %[[T9:.+]] = index_cast %[[ARG5]] 239// CHECK: %[[T10:.+]] = addi %[[T8]], %[[T9]] 240// CHECK: linalg.yield %[[T10]] 241 242// ----- 243 244#map0 = affine_map<(d0, d1) -> (d0, d1)> 245func @indexed_generic_op_reshape_consumer_fusion(%arg0 : tensor<?x?xi32>, 246 %arg1 : tensor<?x?xi32>) -> 247 tensor<?x?x4x5xi32> 248{ 249 %0 = linalg.indexed_generic { 250 indexing_maps = [#map0, #map0, #map0], 251 iterator_types = ["parallel", "parallel"]} 252 ins(%arg0, %arg1 : tensor<?x?xi32>, tensor<?x?xi32>) { 253 ^bb0(%arg3 : index, %arg4 : index, %arg5: i32, %arg6: i32): // no predecessors 254 %1 = muli %arg5, %arg6 : i32 255 %2 = index_cast %arg3 : index to i32 256 %3 = addi %1, %2 : i32 257 %4 = index_cast %arg4 : index to i32 258 %5 = addi %3, %4 : i32 259 linalg.yield %5 : i32 260 } -> tensor<?x?xi32> 261 %1 = linalg.tensor_reshape %0 [affine_map<(i, j, k, l) -> (i)>, 262 affine_map<(i, j, k, l) -> (j, k, l)>] : 263 tensor<?x?xi32> into tensor<?x?x4x5xi32> 264 return %1 : tensor<?x?x4x5xi32> 265} 266// The generic op version of the test check for the op structure. Only 267// checking the op body here. 268// CHECK: #[[MAP:.+]] = affine_map<(d0, d1, d2) -> (d0 * 20 + d1 * 5 + d2)> 269// CHECK: func @indexed_generic_op_reshape_consumer_fusion 270// CHECK: linalg.indexed_generic 271// CHECK: ^{{.*}}( 272// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: index, %[[ARG3:[a-zA-Z0-9]+]]: index, 273// CHECK-SAME: %[[ARG4:[a-zA-Z0-9]+]]: index, %[[ARG5:[a-zA-Z0-9]+]]: index, 274// CHECK-SAME: %[[ARG6:[a-zA-Z0-9]+]]: i32, %[[ARG7:[a-zA-Z0-9]+]]: i32) 275// CHECK: %[[T3:.+]] = affine.apply #[[MAP]](%[[ARG3]], %[[ARG4]], %[[ARG5]]) 276// CHECK: %[[T4:.+]] = muli %[[ARG6]], %[[ARG7]] 277// CHECK: %[[T5:.+]] = index_cast %[[ARG2]] 278// CHECK: %[[T6:.+]] = addi %[[T4]], %[[T5]] 279// CHECK: %[[T7:.+]] = index_cast %[[T3]] 280// CHECK: %[[T8:.+]] = addi %[[T6]], %[[T7]] 281// CHECK: linalg.yield %[[T8]] 282 283// ----- 284 285func @reshape_as_consumer_permutation 286 (%a : tensor<210x6x4xi32>, %b : tensor<210x4xi32>) 287 -> tensor<2x3x4x5x6x7xi32> { 288 %c = linalg.indexed_generic { 289 indexing_maps = [affine_map<(d0, d1, d2) -> (d1, d0, d2)>, 290 affine_map<(d0, d1, d2) -> (d1, d2)>, 291 affine_map<(d0, d1, d2) -> (d0, d2, d1)>], 292 iterator_types = ["parallel", "parallel", "parallel"]} 293 ins(%a, %b : tensor<210x6x4xi32>, tensor<210x4xi32>) { 294 ^bb0(%arg0 : index, %arg1 : index, %arg2 : index, %arg3 : i32, %arg4: i32): 295 %1 = addi %arg3, %arg4 : i32 296 %2 = index_cast %arg0 : index to i32 297 %3 = addi %1, %2 : i32 298 %4 = index_cast %arg1 : index to i32 299 %5 = addi %3, %4 : i32 300 %6 = index_cast %arg2 : index to i32 301 %7 = addi %5, %6 : i32 302 linalg.yield %7 : i32 303 } -> tensor<6x4x210xi32> 304 %d = linalg.tensor_reshape %c 305 [affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1)>, 306 affine_map<(d0, d1, d2, d3, d4, d5) -> (d2)>, 307 affine_map<(d0, d1, d2, d3, d4, d5) -> (d3, d4, d5)>] 308 : tensor<6x4x210xi32> into tensor<2x3x4x5x6x7xi32> 309 return %d : tensor<2x3x4x5x6x7xi32> 310} 311 312 313// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2)> 314// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d3, d4)> 315// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d5)> 316// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)> 317// CHECK-DAG: #[[MAP4:.+]] = affine_map<(d0, d1, d2, d3) -> (d3)> 318// CHECK-DAG: #[[MAP5:.+]] = affine_map<(d0, d1) -> (d0 * 3 + d1)> 319// CHECK-DAG: #[[MAP6:.+]] = affine_map<(d0, d1, d2) -> (d0 * 42 + d1 * 7 + d2)> 320// CHECK-DAG: #[[MAP7:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d2, d3, d4, d0, d1, d5)> 321// CHECK-DAG: #[[MAP8:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d2, d3, d4, d5)> 322// CHECK-DAG: #[[MAP9:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d5, d2, d3, d4)> 323// CHECK: func @reshape_as_consumer_permutation 324// CHECK-SAME: %[[ARG0:.+]]: tensor<210x6x4xi32> 325// CHECK-SAME: %[[ARG1:.+]]: tensor<210x4xi32> 326// CHECK-DAG: %[[T0:.+]] = linalg.tensor_reshape %[[ARG0]] 327// CHECK-SAME: [#[[MAP0]], #[[MAP1]], #[[MAP2]]] 328// CHECK-DAG: %[[T1:.+]] = linalg.tensor_reshape %[[ARG1]] 329// CHECK-SAME: [#[[MAP3]], #[[MAP4]]] 330// CHECK: %[[T2:.+]] = linalg.indexed_generic 331// CHECK-SAME: indexing_maps = [#[[MAP7]], #[[MAP8]], #[[MAP9]]] 332// CHECK-SAME: ins(%[[T0]], %[[T1]] : tensor<{{.+}}>, tensor<{{.+}}>) 333// CHECK: ^{{.+}}( 334// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: index, %[[ARG3:[a-zA-Z0-9]+]]: index, 335// CHECK-SAME: %[[ARG4:[a-zA-Z0-9]+]]: index, %[[ARG5:[a-zA-Z0-9]+]]: index, 336// CHECK-SAME: %[[ARG6:[a-zA-Z0-9]+]]: index, %[[ARG7:[a-zA-Z0-9]+]]: index, 337// CHECK-SAME: %[[ARG8:[a-zA-Z0-9]+]]: i32, %[[ARG9:[a-zA-Z0-9]+]]: i32) 338// CHECK-DAG: %[[T3:.+]] = affine.apply #[[MAP5]](%[[ARG2]], %[[ARG3]]) 339// CHECK-DAG: %[[T4:.+]] = affine.apply #[[MAP6]](%[[ARG4]], %[[ARG5]], %[[ARG6]]) 340// CHECK-DAG: %[[T5:.+]] = addi %[[ARG8]], %[[ARG9]] 341// CHECK: %[[T6:.+]] = index_cast %[[T3]] 342// CHECK: %[[T7:.+]] = addi %[[T5]], %[[T6]] 343// CHECK: %[[T8:.+]] = index_cast %[[T4]] 344// CHECK: %[[T9:.+]] = addi %[[T7]], %[[T8]] 345// CHECK: %[[T10:.+]] = index_cast %[[ARG7]] 346// CHECK: %[[T11:.+]] = addi %[[T9]], %[[T10]] 347 348// ----- 349 350func @reshape_as_producer_projected_permutation 351 (%arg0 : tensor<33x8x?xi32>) -> tensor<264x?x4xi32> { 352 %0 = linalg.tensor_reshape %arg0 [affine_map<(d0, d1, d2) -> (d0, d1)>, 353 affine_map<(d0, d1, d2) -> (d2)>] 354 : tensor<33x8x?xi32> into tensor<264x?xi32> 355 %1 = linalg.indexed_generic 356 {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1)>, 357 affine_map<(d0, d1, d2) -> (d0, d1, d2)>], 358 iterator_types = ["parallel", "parallel", "parallel"]} ins(%0 : tensor<264x?xi32>) { 359 ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: i32): // no predecessors 360 %2 = index_cast %arg1 : index to i32 361 %3 = addi %arg4, %2 : i32 362 %4 = index_cast %arg2 : index to i32 363 %5 = addi %3, %4 : i32 364 %6 = index_cast %arg3 : index to i32 365 %7 = addi %5, %6 : i32 366 linalg.yield %7 : i32 367 } -> tensor<264x?x4xi32> 368 return %1 : tensor<264x?x4xi32> 369} 370 371// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)> 372// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)> 373// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1) -> (d0 * 8 + d1)> 374// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1)> 375// CHECK-DAG: #[[MAP4:.+]] = affine_map<(d0, d1, d2, d3) -> (d2)> 376// CHECK-DAG: #[[MAP5:.+]] = affine_map<(d0, d1, d2, d3) -> (d3)> 377// CHECK: @reshape_as_producer_projected_permutation 378// CHECK-SAME: %[[ARG0:.+]]: tensor<33x8x?xi32> 379// CHECK: %[[RES:.+]] = linalg.indexed_generic 380// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]]] 381// CHECK-SAME: ins(%[[ARG0]] : tensor<33x8x?xi32>) 382// CHECK: ^{{.+}}( 383// CHECK-SAME: %[[ARG1:[a-zA-Z0-9]+]]: index, 384// CHECK-SAME: %[[ARG2:[a-zA-Z0-9]+]]: index, 385// CHECK-SAME: %[[ARG3:[a-zA-Z0-9]+]]: index, 386// CHECK-SAME: %[[ARG4:[a-zA-Z0-9]+]]: index, 387// CHECK-SAME: %[[ARG5:[a-zA-Z0-9]+]]: i32) 388// CHECK: %[[T0:.+]] = affine.apply #[[MAP2]](%[[ARG1]], %[[ARG2]]) 389// CHECK: %[[T1:.+]] = index_cast %[[T0]] : index to i32 390// CHECK: %[[T2:.+]] = addi %[[ARG5]], %[[T1]] : i32 391// CHECK: %[[T3:.+]] = index_cast %[[ARG3]] : index to i32 392// CHECK: %[[T4:.+]] = addi %[[T2]], %[[T3]] : i32 393// CHECK: %[[T5:.+]] = index_cast %[[ARG4]] : index to i32 394// CHECK: %[[T6:.+]] = addi %[[T4]], %[[T5]] : i32 395// CHECK: linalg.yield %[[T6]] : i32 396// CHECK: %[[RES2:.+]] = linalg.tensor_reshape %[[RES]] 397// CHECK-SAME: [#[[MAP3]], #[[MAP4]], #[[MAP5]]] 398// CHECK-SAME: : tensor<33x8x?x4xi32> into tensor<264x?x4xi32> 399// CHECK: return %[[RES2]] : tensor<264x?x4xi32> 400 401// ----- 402 403#map0 = affine_map<(d0, d1) -> (d0, d1)> 404#map1 = affine_map<(d0, d1) -> (d1, d0)> 405func @generic_op_reshape_consumer_fusion_projected(%arg0 : tensor<?x?xf32>, 406 %arg1 : tensor<?x?xf32>) -> 407 tensor<?x?x4x5xf32> 408{ 409 %0 = linalg.generic { 410 indexing_maps = [#map0, #map0, #map1], 411 iterator_types = ["parallel", "parallel"]} 412 ins(%arg0, %arg1 : tensor<?x?xf32>, tensor<?x?xf32>) { 413 ^bb0(%arg3: f32, %arg4: f32): // no predecessors 414 %1 = mulf %arg3, %arg4 : f32 415 linalg.yield %1 : f32 416 } -> tensor<?x?xf32> 417 %1 = linalg.tensor_reshape %0 [affine_map<(i, j, k, l) -> (i)>, 418 affine_map<(i, j, k, l) -> (j, k, l)>] : 419 tensor<?x?xf32> into tensor<?x?x4x5xf32> 420 return %1 : tensor<?x?x4x5xf32> 421} 422 423// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)> 424// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d3)> 425// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)> 426// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3) -> (d3, d0, d1, d2)> 427// CHECK: func @generic_op_reshape_consumer_fusion_projected 428// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<?x?xf32> 429// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: tensor<?x?xf32> 430// CHECK: %[[T0:.+]] = linalg.tensor_reshape %[[ARG0]] 431// CHECK-SAME: [#[[MAP0]], #[[MAP1]]] 432// CHECK-SAME: tensor<?x?xf32> into tensor<?x4x5x?xf32> 433// CHECK: %[[T1:.+]] = linalg.tensor_reshape %[[ARG1]] 434// CHECK-SAME: [#[[MAP0]], #[[MAP1]]] 435// CHECK-SAME: tensor<?x?xf32> into tensor<?x4x5x?xf32> 436// CHECK: %[[T2:.+]] = linalg.generic 437// CHECK-SAME: indexing_maps = [#[[MAP2]], #[[MAP2]], #[[MAP3]]] 438// CHECK-SAME: ["parallel", "parallel", "parallel", "parallel"] 439// CHECK-SAME: ins(%[[T0]], %[[T1]] : tensor<?x4x5x?xf32>, tensor<?x4x5x?xf32>) 440// CHECK: return %[[T2]] : tensor<?x?x4x5xf32> 441