1177 lines · plain
1// RUN: mlir-opt %s -linalg-fold-unit-extent-dims -split-input-file | FileCheck %s2// RUN: mlir-opt %s -linalg-fold-unit-extent-dims="use-rank-reducing-slices" -cse -split-input-file | FileCheck %s --check-prefix=CHECK-SLICES3 4#accesses = [5 affine_map<(i, j, k, l, m) -> (i, k, m)>,6 affine_map<(i, j, k, l, m) -> ()>,7 affine_map<(i, j, k, l, m) -> (i, k, j, l, m)>8]9 10#trait = {11 iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"],12 indexing_maps = #accesses,13 library_call = "some_external_func"14}15 16func.func @drop_one_trip_loops(%arg0 : tensor<?x1x?xf32>, %arg1 : f32, %shape: tensor<?x1x?x1x?xf32>) -> tensor<?x1x?x1x?xf32> {17 %0 = linalg.generic #trait18 ins(%arg0, %arg1 : tensor<?x1x?xf32>, f32)19 outs(%shape : tensor<?x1x?x1x?xf32>) {20 ^bb0(%arg2 : f32, %arg3 : f32, %arg4 : f32) :21 linalg.yield %arg3 : f3222 } -> tensor<?x1x?x1x?xf32>23 return %0 : tensor<?x1x?x1x?xf32>24}25// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2) -> (d0, d2)>26// CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1, d2) -> ()>27// CHECK-DAG: #[[$MAP3:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>28// CHECK-LABEL: func @drop_one_trip_loops29// CHECK: %[[C2:.*]] = arith.constant 2 : index30// CHECK: %[[C0:.*]] = arith.constant 0 : index31// CHECK: tensor.collapse_shape %{{.*}} {{\[\[}}0, 1], [2]]32// CHECK: tensor.collapse_shape %{{.*}} {{\[\[}}0, 1], [2, 3], [4]]33// CHECK: linalg.generic34// CHECK-SAME: indexing_maps = [#[[$MAP1]], #[[$MAP2]], #[[$MAP3]]]35// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"]36// CHECK: %[[DIM:.*]] = tensor.dim %{{.*}}, %[[C0]]37// CHECK: %[[DIM_1:.*]] = tensor.dim %{{.*}}, %[[C2]]38// CHECK: %[[DIM_2:.*]] = tensor.dim %{{.*}}, %[[C2]]39// CHECK: %[[EXPANDED:.*]] = tensor.expand_shape %{{.*}} {{\[\[}}0, 1], [2, 3], [4]] output_shape [%[[DIM]], 1, %[[DIM_1]], 1, %[[DIM_2]]] : tensor<?x?x?xf32> into tensor<?x1x?x1x?xf32>40 41// CHECK-SLICES-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2) -> (d0, d2)>42// CHECK-SLICES-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1, d2) -> ()>43// CHECK-SLICES-DAG: #[[$MAP3:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>44// CHECK-SLICES-LABEL: func @drop_one_trip_loops45// CHECK-SLICES: tensor.extract_slice %{{.*}}[0, 0, 0] [%{{.*}}, 1, %{{.*}}] [1, 1, 1] : tensor<?x1x?xf32> to tensor<?x?xf32>46// CHECK-SLICES: tensor.extract_slice %{{.*}}[0, 0, 0, 0, 0] [%{{.*}}, 1, %{{.*}}, 1, %{{.*}}] [1, 1, 1, 1, 1] : tensor<?x1x?x1x?xf32> to tensor<?x?x?xf32>47// CHECK-SLICES: linalg.generic48// CHECK-SLICES-SAME: indexing_maps = [#[[$MAP1]], #[[$MAP2]], #[[$MAP3]]]49// CHECK-SLICES-SAME: iterator_types = ["parallel", "parallel", "parallel"]50// CHECK-SLICES: tensor.insert_slice %{{.*}} into %{{.*}}[0, 0, 0, 0, 0] [%{{.*}}, 1, %{{.*}}, 1, %{{.*}}] [1, 1, 1, 1, 1] : tensor<?x?x?xf32> into tensor<?x1x?x1x?xf32>51 52 53// -----54 55#accesses = [56 affine_map<(i, j, k, l, m) -> (i, k, m)>,57 affine_map<(i, j, k, l, m) -> ()>,58 affine_map<(i, j, k, l, m) -> (i, k, j, l, m)>59]60 61#trait = {62 iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"],63 indexing_maps = #accesses,64 library_call = "some_external_func"65}66 67func.func @drop_one_trip_loops_all_ones(%arg0 : tensor<1x1x1xf32>, %arg1 : f32, %shape: tensor<1x1x?x1x1xf32>) -> tensor<1x1x?x1x1xf32> {68 %0 = linalg.generic #trait69 ins(%arg0, %arg1 : tensor<1x1x1xf32>, f32)70 outs(%shape : tensor<1x1x?x1x1xf32>) {71 ^bb0(%arg2 : f32, %arg3 : f32, %arg4 : f32) :72 linalg.yield %arg3 : f3273 } -> tensor<1x1x?x1x1xf32>74 return %0 : tensor<1x1x?x1x1xf32>75}76// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0) -> ()>77// CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0) -> (d0)>78// CHECK-LABEL: func @drop_one_trip_loops_all_ones79// CHECK: %[[C2:.*]] = arith.constant 2 : index80// CHECK: tensor.collapse_shape %{{.*}} []81// CHECK: tensor.collapse_shape %{{.*}} {{\[}}[0, 1, 2, 3, 4]]82// CHECK: linalg.generic83// CHECK-SAME: indexing_maps = [#[[$MAP1]], #[[$MAP1]], #[[$MAP2]]]84// CHECK-SAME: iterator_types = ["parallel"]85// CHECK: %[[DIM:.*]] = tensor.dim %{{.*}}, %[[C2]] : tensor<1x1x?x1x1xf32>86// CHECK: %[[EXPAND:.*]] = tensor.expand_shape %{{.*}} {{\[\[}}0, 1, 2, 3, 4]] output_shape [1, 1, %[[DIM]], 1, 1] : tensor<?xf32> into tensor<1x1x?x1x1xf32>87 88// -----89 90#accesses = [91 affine_map<(i, j, k, l, m) -> (i, k, m)>,92 affine_map<(i, j, k, l, m) -> (i, k, j, l, m)>93]94 95#trait = {96 iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"],97 indexing_maps = #accesses,98 library_call = "some_external_func"99}100 101func.func @drop_one_trip_loops_indexed102 (%arg0 : tensor<?x1x?xi32>, %shape: tensor<?x1x?x1x?xi32>) -> tensor<?x1x?x1x?xi32>103{104 %0 = linalg.generic #trait105 ins(%arg0 : tensor<?x1x?xi32>)106 outs(%shape: tensor<?x1x?x1x?xi32>) {107 ^bb0(%arg6 : i32, %arg7 : i32) :108 %idx0 = linalg.index 0 : index109 %idx1 = linalg.index 1 : index110 %idx2 = linalg.index 2 : index111 %idx3 = linalg.index 3 : index112 %idx4 = linalg.index 4 : index113 %1 = arith.addi %idx0, %idx1 : index114 %2 = arith.subi %1, %idx2 : index115 %3 = arith.subi %2, %idx3 : index116 %4 = arith.addi %3, %idx4 : index117 %5 = arith.index_cast %4 : index to i32118 %6 = arith.addi %5, %arg6 : i32119 linalg.yield %6 : i32120 } -> tensor<?x1x?x1x?xi32>121 return %0 : tensor<?x1x?x1x?xi32>122}123// The subtractions disappear the access map of the output tensor maps its unit124// dimensions 1 and 3 to the index dimensions 2 and 3.125// CHECK-LABEL: func @drop_one_trip_loops_indexed126// CHECK: linalg.generic127// CHECK: ^{{.+}}(128// CHECK-SAME: %[[ARG4:[a-zA-Z0-9]+]]: i32, %{{.*}}: i32)129// CHECK: %[[IDX0:.+]] = linalg.index 0 : index130// CHECK: %[[IDX1:.+]] = linalg.index 1 : index131// CHECK: %[[IDX2:.+]] = linalg.index 2 : index132// CHECK: %[[T3:.+]] = arith.addi %[[IDX0]], %[[IDX1]]133// CHECK: %[[T4:.+]] = arith.addi %[[T3]], %[[IDX2]]134// CHECK: %[[T5:.+]] = arith.index_cast %[[T4]] : index to i32135// CHECK: %[[T6:.+]] = arith.addi %[[T5]], %[[ARG4]] : i32136// CHECK: linalg.yield %[[T6]] : i32137 138// -----139 140#map0 = affine_map<(i, j) -> (i, j)>141#access = [#map0, #map0]142#trait = {143 iterator_types = ["parallel", "parallel"],144 indexing_maps = #access,145 library_call = "some_external_func"146}147 148func.func @drop_all_loops(%arg0 : tensor<1x1xf32>) -> tensor<1x1xf32>149{150 %0 = linalg.generic #trait151 ins(%arg0 : tensor<1x1xf32>)152 outs(%arg0 : tensor<1x1xf32>) {153 ^bb0(%arg1: f32, %arg2: f32) :154 linalg.yield %arg1 : f32155 } -> tensor<1x1xf32>156 return %0 : tensor<1x1xf32>157}158// CHECK: #[[$MAP0:.*]] = affine_map<() -> ()>159// CHECK-LABEL: func @drop_all_loops160// CHECK: tensor.collapse_shape %{{.*}} []161// CHECK: linalg.generic162// CHECK-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP0]]]163// CHECK-SAME: iterator_types = []164 165// -----166 167#map0 = affine_map<(i, j) -> (i, j)>168#access = [#map0, #map0]169#trait = {170 iterator_types = ["parallel", "parallel"],171 indexing_maps = #access,172 library_call = "some_external_func"173}174 175func.func @drop_all_loops_indexed176 (%arg0 : tensor<1x1xi32>) -> tensor<1x1xi32>{177 %0 = linalg.generic #trait178 ins(%arg0 : tensor<1x1xi32>)179 outs(%arg0 : tensor<1x1xi32>) {180 ^bb0(%arg3: i32, %arg4: i32) :181 %idx0 = linalg.index 0 : index182 %idx1 = linalg.index 1 : index183 %1 = arith.addi %idx0, %idx1 : index184 %2 = arith.index_cast %1 : index to i32185 %3 = arith.addi %2, %arg3 : i32186 linalg.yield %3 : i32187 } -> tensor<1x1xi32>188 return %0 : tensor<1x1xi32>189}190 191// CHECK-LABEL: func @drop_all_loops_indexed192// CHECK: linalg.generic193// CHECK: ^{{.+}}(%[[ARG1:.+]]: i32, %[[ARG2:.+]]: i32)194// CHECK: linalg.yield %[[ARG1]] : i32195 196// -----197 198#accesses = [199 affine_map<(d0) -> (0, d0)>,200 affine_map<(d0) -> (d0)>201]202 203#trait = {204 indexing_maps = #accesses,205 iterator_types = ["parallel"],206 library_call = "some_external_fn"207}208 209func.func @leading_dim_1_canonicalization(%arg0: tensor<1x5xf32>, %shape: tensor<5xf32>) -> tensor<5xf32> {210 %0 = linalg.generic #trait211 ins(%arg0 : tensor<1x5xf32>)212 outs(%shape : tensor<5xf32>) {213 ^bb0(%arg2: f32, %arg3: f32):214 linalg.yield %arg2 : f32215 } -> tensor<5xf32>216 return %0 : tensor<5xf32>217}218// CHECK: #[[$MAP1:.*]] = affine_map<(d0) -> (d0)>219 220// CHECK-LABEL: func @leading_dim_1_canonicalization221// CHECK: tensor.collapse_shape %{{.*}} {{\[}}[0, 1]]222// CHECK: linalg.generic223// CHECK-SAME: indexing_maps = [#[[$MAP1]], #[[$MAP1]]]224// CHECK-SAME: iterator_types = ["parallel"]225 226// -----227 228#accesses = [229 affine_map<(d0, d1) -> (0, d1)>,230 affine_map<(d0, d1) -> (d0, 0)>,231 affine_map<(d0, d1) -> (d0, d1)>232]233 234#trait = {235 indexing_maps = #accesses,236 iterator_types = ["parallel", "parallel"],237 library_call = "some_external_fn"238}239 240func.func @broadcast_test(%arg0 : tensor<5xf32>, %arg1 : tensor<5xf32>, %shape : tensor<5x5xf32>) -> tensor<5x5xf32>241{242 %0 = tensor.expand_shape %arg0 [[0, 1]] output_shape [1, 5] : tensor<5xf32> into tensor<1x5xf32>243 %1 = tensor.expand_shape %arg1 [[0, 1]] output_shape [5, 1] : tensor<5xf32> into tensor<5x1xf32>244 %2 = linalg.generic #trait245 ins(%0, %1 : tensor<1x5xf32>, tensor<5x1xf32>)246 outs(%shape : tensor<5x5xf32>) {247 ^bb0(%arg3: f32, %arg4: f32, %arg5: f32):248 %3 = arith.addf %arg3, %arg4 : f32249 linalg.yield %3 : f32250 } -> tensor<5x5xf32>251 return %2 : tensor<5x5xf32>252}253// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d1)>254// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d0)>255// CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1) -> (d0, d1)>256// CHECK-LABEL: func @broadcast_test257// CHECK-NOT: linalg.tensor_{{.*}}shape258// CHECK: linalg.generic259// CHECK-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]]]260// CHECK-SAME: iterator_types = ["parallel", "parallel"]261// CHECK-NOT: linalg.tensor_{{.*}}shape262 263// -----264 265#accesses = [266 affine_map<(d0, d1) -> (0, 0)>,267 affine_map<(d0, d1) -> (d0, d1)>268]269 270#trait = {271 indexing_maps = #accesses,272 iterator_types = ["parallel", "parallel"],273 library_call = "some_external_fn"274}275 276func.func @broadcast_scalar(%arg0 : tensor<1x1xf32>, %shape : tensor<?x?xf32>) -> tensor<?x?xf32>277{278 %0 = linalg.generic #trait279 ins(%arg0 : tensor<1x1xf32>)280 outs(%shape : tensor<?x?xf32>) {281 ^bb0(%arg2 : f32, %arg3 : f32):282 linalg.yield %arg2 : f32283 } -> tensor<?x?xf32>284 return %0 : tensor<?x?xf32>285}286// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1) -> ()>287// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d0, d1)>288// CHECK-LABEL: func @broadcast_scalar289// CHECK-SAME: %[[ARG0:.*]]: tensor<1x1xf32>290// CHECK: %[[A:.*]] = tensor.collapse_shape %[[ARG0]] []291// CHECK-SAME: tensor<1x1xf32> into tensor<f32>292// CHECK: linalg.generic293// CHECK-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP1]]]294// CHECK-SAME: iterator_types = ["parallel", "parallel"]295// CHECK-SAME: %[[A]]296 297// -----298 299#map0 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>300#map1 = affine_map<(d0, d1, d2) -> (d2)>301func.func @fold_unit_dim_tensor_reshape_op(%arg0 : tensor<5xf32>) -> tensor<2x5xf32>302{303 %1 = tensor.empty() : tensor<1x2x5xf32>304 %2 = linalg.generic {i64, indexing_maps = [#map1, #map0],305 iterator_types = ["parallel", "parallel", "parallel"]}306 ins(%arg0 : tensor<5xf32>) outs(%1 : tensor<1x2x5xf32>) {307 ^bb0(%arg1: f32, %arg2: f32):308 linalg.yield %arg1 : f32309 } -> tensor<1x2x5xf32>310 %3 = tensor.collapse_shape %2 [[0, 1], [2]]311 : tensor<1x2x5xf32> into tensor<2x5xf32>312 return %3 : tensor<2x5xf32>313}314// CHECK-LABEL: func @fold_unit_dim_tensor_reshape_op315// CHECK: %[[RESULT:.+]] = linalg.generic316// CHECK: return %[[RESULT]]317 318// -----319 320func.func @fold_unit_dim_for_empty_tensor(%input: tensor<1x1000xf32>) -> tensor<1xf32> {321 %cst = arith.constant 0.0 : f32322 %init = tensor.empty() : tensor<1xf32>323 %fill = linalg.fill ins(%cst : f32) outs(%init : tensor<1xf32>) -> tensor<1xf32>324 %add = linalg.generic {325 indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0)>],326 iterator_types = ["parallel", "reduction"]}327 ins(%input : tensor<1x1000xf32>)outs(%fill : tensor<1xf32>) {328 ^bb0(%arg1: f32, %arg2: f32):329 %1823 = arith.addf %arg1, %arg2 : f32330 linalg.yield %1823 : f32331 } -> tensor<1xf32>332 return %add : tensor<1xf32>333}334 335 336// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0) -> (d0)>337// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0) -> ()>338 339// CHECK: func @fold_unit_dim_for_empty_tensor340 341// CHECK: %[[INPUT_RESHAPE:.+]] = tensor.collapse_shape %{{.+}} {{\[}}[0, 1]] : tensor<1x1000xf32> into tensor<1000xf32>342// CHECK: %[[INIT:.+]] = tensor.empty() : tensor<f32>343// CHECK: %[[FILL:.+]] = linalg.fill ins(%cst : f32) outs(%[[INIT]] : tensor<f32>) -> tensor<f32>344// CHECK: %[[GENERIC:.+]] = linalg.generic345// CHECK-SAME: indexing_maps = [#[[MAP1]], #[[MAP2]]]346// CHECK-SAME: iterator_types = ["reduction"]347// CHECK-SAME: ins(%[[INPUT_RESHAPE]] : tensor<1000xf32>)348// CHECK-SAME: outs(%[[FILL]] : tensor<f32>)349// CHECK: %[[GENERIC_RESHAPE:.+]] = tensor.expand_shape %[[GENERIC]] [] output_shape [1] : tensor<f32> into tensor<1xf32>350// CHECK: return %[[GENERIC_RESHAPE:.+]] : tensor<1xf32>351 352 353// -----354 355func.func @fold_slice(356 %arg0 : tensor<1x?x?x1x?x1x1xf32>, %arg1 : tensor<1x?x?x?x?x1x1xf32>,357 %arg2 : index, %arg3 : index, %arg4 : index, %arg5 : index,358 %arg6 : index, %arg7 : index) -> (tensor<1x?x?x1x?x1x1xf32>, tensor<1x?x?x1x?x1x1xf32>) {359 %0 = tensor.extract_slice %arg0[0, %arg2, %arg3, 0, %arg4, 0, 0]360 [1, %arg5, %arg6, 1, %arg7, 1, 1] [1, 1, 1, 1, 1, 1, 1] :361 tensor<1x?x?x1x?x1x1xf32> to tensor<1x?x?x1x?x1x1xf32>362 %1 = tensor.extract_slice %arg1[%arg2, 0, %arg3, 0, 0, %arg4, 0]363 [1, %arg5, %arg6, 1, %arg7, 1, 1] [1, 1, 1, 1, 1, 1, 1] :364 tensor<1x?x?x?x?x1x1xf32> to tensor<1x?x?x1x?x1x1xf32>365 return %0, %1 : tensor<1x?x?x1x?x1x1xf32>, tensor<1x?x?x1x?x1x1xf32>366}367// CHECK: func @fold_slice368// CHECK-SAME: %[[ARG0:.+]]: tensor<1x?x?x1x?x1x1xf32>369// CHECK-SAME: %[[ARG1:.+]]: tensor<1x?x?x?x?x1x1xf32>370// CHECK: %[[SLICE1:.+]] = tensor.extract_slice %[[ARG0]]371// CHECK-SAME: to tensor<?x?x?xf32>372// CHECK: %[[RESULT1:.+]] = tensor.expand_shape %[[SLICE1]]373// CHECK-SAME: {{\[\[}}0, 1], [2], [3, 4, 5, 6]] output_shape [1, %arg5, %arg6, 1, %arg7, 1, 1] : tensor<?x?x?xf32> into tensor<1x?x?x1x?x1x1xf32>374// CHECK: %[[SLICE2:.+]] = tensor.extract_slice %[[ARG1]]375// CHECK-SAME: to tensor<?x?x?xf32>376// CHECK: %[[RESULT2:.+]] = tensor.expand_shape %[[SLICE2]]377// CHECK-SAME: {{\[\[}}0, 1], [2], [3, 4, 5, 6]] output_shape [1, %arg5, %arg6, 1, %arg7, 1, 1] : tensor<?x?x?xf32> into tensor<1x?x?x1x?x1x1xf32>378// CHECK: return %[[RESULT1]], %[[RESULT2]]379 380// -----381 382func.func @unit_dim_for_reduction(%arg0: tensor<1x?x1x?xf32>) -> tensor<1x?xf32> {383 %cst = arith.constant 1.000000e+00 : f32384 %c3 = arith.constant 3 : index385 %0 = tensor.dim %arg0, %c3 : tensor<1x?x1x?xf32>386 %1 = tensor.empty(%0) : tensor<1x?xf32>387 %2 = linalg.fill ins(%cst : f32) outs(%1 : tensor<1x?xf32>) -> tensor<1x?xf32>388 %3 = linalg.generic {389 indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>,390 affine_map<(d0, d1, d2, d3) -> (d0, d1)>],391 iterator_types = ["parallel", "parallel", "reduction", "reduction"]}392 ins(%arg0 : tensor<1x?x1x?xf32>)393 outs(%2 : tensor<1x?xf32>) {394 ^bb0(%arg1: f32, %arg2: f32):395 %4 = arith.addf %arg1, %arg2 : f32396 linalg.yield %4 : f32397 } -> tensor<1x?xf32>398 return %3 : tensor<1x?xf32>399}400// CHECK-DAG: #[[MAP:.+]] = affine_map<(d0, d1) -> (d0, d1)>401// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1) -> (d0)>402// CHECK: func @unit_dim_for_reduction403// CHECK-SAME: %[[ARG0:.+]]: tensor<1x?x1x?xf32>404// CHECK: %[[C1:.+]] = arith.constant 1 : index405// CHECK: %[[CST:.+]] = arith.constant 1.000000e+00 : f32406// CHECK: %[[C3:.+]] = arith.constant 3 : index407// CHECK: %[[DIM:.+]] = tensor.dim %arg0, %[[C3]] : tensor<1x?x1x?xf32>408// CHECK: %[[RESHAPE:.+]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0, 1, 2], [3]]409// CHECK: %[[INIT:.+]] = tensor.empty(%{{.+}}) : tensor<?xf32>410// CHECK: %[[FILL:.+]] = linalg.fill ins(%{{.+}}{{.*}}outs(%[[INIT]]411// CHECK: %[[RESULT:.+]] = linalg.generic412// CHECK-SAME: indexing_maps = [#[[MAP]], #[[MAP2]]]413// CHECK-SAME: iterator_types = ["parallel", "reduction"]414// CHECK-SAME: ins(%[[RESHAPE]] : tensor<?x?xf32>)415// CHECK-SAME: outs(%[[FILL]] : tensor<?xf32>)416// CHECK: %[[DIM_0:.*]] = tensor.dim %[[ARG0]], %[[C1]] : tensor<1x?x1x?xf32>417// CHECK: %[[EXPANDED:.*]] = tensor.expand_shape %[[GENERIC]] {{\[\[}}0, 1]] output_shape [1, %[[DIM_0]]] : tensor<?xf32> into tensor<1x?xf32>418// CHECK: return %[[EXPANDED]] : tensor<1x?xf32>419 420// -----421 422func.func @unit_dim_for_both_reduction(%arg0: tensor<1x?x1x1xf32>) -> tensor<1x1xf32> {423 %cst = arith.constant 1.000000e+00 : f32424 %c3 = arith.constant 3 : index425 %1 = tensor.empty() : tensor<1x1xf32>426 %2 = linalg.fill ins(%cst : f32) outs(%1 : tensor<1x1xf32>) -> tensor<1x1xf32>427 %3 = linalg.generic {428 indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>,429 affine_map<(d0, d1, d2, d3) -> (d0, d1)>],430 iterator_types = ["parallel", "parallel", "reduction", "reduction"]}431 ins(%arg0 : tensor<1x?x1x1xf32>)432 outs(%2 : tensor<1x1xf32>) {433 ^bb0(%arg1: f32, %arg2: f32):434 %4 = arith.addf %arg1, %arg2 : f32435 linalg.yield %4 : f32436 } -> tensor<1x1xf32>437 return %3 : tensor<1x1xf32>438}439// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0) -> (d0)>440// CHECK: func @unit_dim_for_both_reduction441// CHECK-SAME: %[[ARG0:.+]]: tensor<1x?x1x1xf32>442// CHECK-DAG: %[[RESHAPE:.+]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0, 1, 2, 3]443// CHECK: %[[INIT:.+]] = tensor.empty() : tensor<1xf32>444// CHECK: %[[FILL:.+]] = linalg.fill ins(%{{.+}}{{.*}}outs(%[[INIT]]445// CHECK: %[[INIT2:.+]] = tensor.empty() : tensor<1xf32>446// CHECK: %[[RESULT:.+]] = linalg.generic447// CHECK-SAME: indexing_maps = [#[[MAP2]], #[[MAP2]], #[[MAP2]]]448// CHECK-SAME: iterator_types = ["parallel"]449// CHECK-SAME: ins(%[[RESHAPE]], %[[FILL]] : tensor<?xf32>, tensor<1xf32>)450// CHECK-SAME: outs(%[[INIT2]] : tensor<1xf32>)451// CHECK: %[[RESULT_RESHAPE:.+]] = tensor.expand_shape %[[RESULT]] {{\[}}[0, 1]] output_shape [1, 1]452// CHECK: return %[[RESULT_RESHAPE]]453 454// -----455 456func.func @unit_dim_for_reduction_inner(%arg0: tensor<?x1x?x1xf32>) -> tensor<?x1xf32> {457 %cst = arith.constant 1.000000e+00 : f32458 %c2 = arith.constant 2 : index459 %0 = tensor.dim %arg0, %c2 : tensor<?x1x?x1xf32>460 %1 = tensor.empty(%0) : tensor<?x1xf32>461 %2 = linalg.fill ins(%cst : f32) outs(%1 : tensor<?x1xf32>) -> tensor<?x1xf32>462 %3 = linalg.generic {463 indexing_maps = [affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>,464 affine_map<(d0, d1, d2, d3) -> (d0, d1)>],465 iterator_types = ["parallel", "parallel", "reduction", "reduction"]}466 ins(%arg0 : tensor<?x1x?x1xf32>)467 outs(%2 : tensor<?x1xf32>) {468 ^bb0(%arg1: f32, %arg2: f32):469 %4 = arith.addf %arg1, %arg2 : f32470 linalg.yield %4 : f32471 } -> tensor<?x1xf32>472 return %3 : tensor<?x1xf32>473}474// CHECK-DAG: #[[MAP:.+]] = affine_map<(d0, d1) -> (d0, d1)>475// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1) -> (d0)>476// CHECK: func @unit_dim_for_reduction_inner477// CHECK-SAME: %[[ARG0:.+]]: tensor<?x1x?x1xf32>478// CHECK: %[[C0:.*]] = arith.constant 0 : index479// CHECK: %[[CST:.*]] = arith.constant 1.000000e+00 : f32480// CHECK: %[[C2:.*]] = arith.constant 2 : index481// CHECK: %[[DIM:.*]] = tensor.dim %arg0, %[[C2]] : tensor<?x1x?x1xf32>482// CHECK: %[[RESHAPE:.+]] = tensor.collapse_shape %[[ARG0]] {{\[}}[0, 1], [2, 3]]483// CHECK: %[[INIT:.+]] = tensor.empty(%{{.+}}) : tensor<?xf32>484// CHECK: %[[FILL:.+]] = linalg.fill ins(%{{.+}}{{.*}}outs(%[[INIT]]485// CHECK: %[[RESULT:.+]] = linalg.generic486// CHECK-SAME: indexing_maps = [#[[MAP]], #[[MAP2]]]487// CHECK-SAME: iterator_types = ["parallel", "reduction"]488// CHECK-SAME: ins(%[[RESHAPE]] : tensor<?x?xf32>)489// CHECK-SAME: outs(%[[FILL]] : tensor<?xf32>)490// CHECK: %[[DIM_0:.+]] = tensor.dim %[[ARG0]], %[[C0]] : tensor<?x1x?x1xf32>491// CHECK: %[[RESULT_RESHAPE:.+]] = tensor.expand_shape %[[RESULT]] {{\[}}[0, 1]] output_shape [%[[DIM_0]], 1] : tensor<?xf32> into tensor<?x1xf32>492// CHECK: return %[[RESULT_RESHAPE]]493 494// -----495 496func.func @slice_unit_dims(%arg0: tensor<1x3xf32>) -> tensor<1x1xf32> {497 %0 = tensor.extract_slice %arg0[0, 2] [1, 1] [1, 1] : tensor<1x3xf32> to tensor<1x1xf32>498 return %0 : tensor<1x1xf32>499}500// CHECK-LABEL: func @slice_unit_dims501// CHECK: %[[SLICE:.+]] = tensor.extract_slice502// CHECK-SAME: tensor<1x3xf32> to tensor<f32>503// CHECK: %[[RESULT:.+]] = tensor.expand_shape %[[SLICE]] [] output_shape [1, 1]504// CHECK: return %[[RESULT]]505 506// -----507 508func.func @rank_reduced_extract_slice(%arg0: tensor<1x1x3x1x3xf32>) -> tensor<1x3x3xf32> {509 %0 = tensor.extract_slice %arg0[0, 0, 0, 0, 0] [1, 1, 3, 1, 3] [1, 1, 1, 1, 1] : tensor<1x1x3x1x3xf32> to tensor<1x3x3xf32>510 return %0 : tensor<1x3x3xf32>511}512// CHECK-LABEL: func @rank_reduced_extract_slice513// CHECK: %[[SLICE:.+]] = tensor.extract_slice514// CHECK-SAME: tensor<1x1x3x1x3xf32> to tensor<3x3xf32>515// CHECK: %[[RESULT:.+]] = tensor.expand_shape %[[SLICE]] {{\[}}[0, 1], [2]] output_shape [1, 3, 3]516// CHECK: return %[[RESULT]]517 518// -----519 520func.func @insert_slice_unit_dims(%arg0: tensor<1x3xf32>, %arg1: tensor<1x1xf32>) -> tensor<1x3xf32> {521 %0 = tensor.insert_slice %arg1 into %arg0[0, 2] [1, 1] [1, 1] : tensor<1x1xf32> into tensor<1x3xf32>522 return %0 : tensor<1x3xf32>523}524// CHECK-LABEL: func @insert_slice_unit_dims525// CHECK: %[[RESHAPE:.+]] = tensor.collapse_shape %{{.+}} []526// CHECK: %[[RESULT:.+]] = tensor.insert_slice %[[RESHAPE]]527// CHECK-SAME: tensor<f32> into tensor<1x3xf32>528// CHECK: return %[[RESULT]]529 530// -----531 532func.func @rank_reduced_insert_slice(%arg0: tensor<1x1x3x1x3xf32>, %arg1: tensor<1x3x3xf32>) -> tensor<1x1x3x1x3xf32> {533 %0 = tensor.insert_slice %arg1 into %arg0[0, 0, 0, 0, 0] [1, 1, 3, 1, 3] [1, 1, 1, 1, 1] : tensor<1x3x3xf32> into tensor<1x1x3x1x3xf32>534 return %0 : tensor<1x1x3x1x3xf32>535}536// CHECK-LABEL: func @rank_reduced_insert_slice537// CHECK: %[[RESHAPE:.+]] = tensor.collapse_shape %{{.+}} {{\[}}[0, 1], [2]]538// CHECK: %[[RESULT:.+]] = tensor.insert_slice %[[RESHAPE]]539// CHECK-SAME: tensor<3x3xf32> into tensor<1x1x3x1x3xf32>540// CHECK: return %[[RESULT]]541 542// -----543 544#accesses = [545 affine_map<(i, j, k, l, m) -> (i, k, m)>,546 affine_map<(i, j, k, l, m) -> ()>,547 affine_map<(i, j, k, l, m) -> (i, k, j, l, m)>548]549 550#trait = {551 iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"],552 indexing_maps = #accesses,553 library_call = "some_external_func"554}555 556func.func @drop_one_trip_loops(%arg0 : memref<?x1x?xf32>, %arg1 : f32, %shape: memref<?x1x?x1x?xf32>) -> memref<?x1x?x1x?xf32> {557 linalg.generic #trait558 ins(%arg0, %arg1 : memref<?x1x?xf32>, f32)559 outs(%shape : memref<?x1x?x1x?xf32>) {560 ^bb0(%arg2 : f32, %arg3 : f32, %arg4 : f32) :561 linalg.yield %arg3 : f32562 }563 return %shape : memref<?x1x?x1x?xf32>564}565// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2) -> (d0, d2)>566// CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1, d2) -> ()>567// CHECK-DAG: #[[$MAP3:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>568// CHECK-LABEL: func @drop_one_trip_loops569// CHECK: memref.collapse_shape %{{.*}} {{\[}}[0, 1], [2]]570// CHECK: linalg.generic571// CHECK-SAME: indexing_maps = [#[[$MAP1]], #[[$MAP2]], #[[$MAP3]]]572// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"]573 574// -----575 576#accesses = [577 affine_map<(i, j, k, l, m) -> (i, k, m)>,578 affine_map<(i, j, k, l, m) -> (i, k, j, l, m)>579]580 581#trait = {582 iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"],583 indexing_maps = #accesses,584 library_call = "some_external_func"585}586 587func.func @drop_one_trip_loops_indexed588 (%arg0 : memref<?x1x?xi32>, %shape: memref<?x1x?x1x?xi32>) -> memref<?x1x?x1x?xi32>589{590 linalg.generic #trait591 ins(%arg0 : memref<?x1x?xi32>)592 outs(%shape: memref<?x1x?x1x?xi32>) {593 ^bb0(%arg6 : i32, %arg7 : i32) :594 %idx0 = linalg.index 0 : index595 %idx1 = linalg.index 1 : index596 %idx2 = linalg.index 2 : index597 %idx3 = linalg.index 3 : index598 %idx4 = linalg.index 4 : index599 %1 = arith.addi %idx0, %idx1 : index600 %2 = arith.subi %1, %idx2 : index601 %3 = arith.subi %2, %idx3 : index602 %4 = arith.addi %3, %idx4 : index603 %5 = arith.index_cast %4 : index to i32604 %6 = arith.addi %5, %arg6 : i32605 linalg.yield %6 : i32606 }607 return %shape : memref<?x1x?x1x?xi32>608}609// The subtractions disappear the access map of the output memref maps its unit610// dimensions 1 and 3 to the index dimensions 2 and 3.611// CHECK-LABEL: func @drop_one_trip_loops_indexed612// CHECK: linalg.generic613// CHECK: ^{{.+}}(614// CHECK-SAME: %[[ARG4:[a-zA-Z0-9]+]]: i32, %{{.*}}: i32)615// CHECK: %[[IDX0:.+]] = linalg.index 0 : index616// CHECK: %[[IDX1:.+]] = linalg.index 1 : index617// CHECK: %[[IDX2:.+]] = linalg.index 2 : index618// CHECK: %[[T3:.+]] = arith.addi %[[IDX0]], %[[IDX1]]619// CHECK: %[[T4:.+]] = arith.addi %[[T3]], %[[IDX2]]620// CHECK: %[[T5:.+]] = arith.index_cast %[[T4]] : index to i32621// CHECK: %[[T6:.+]] = arith.addi %[[T5]], %[[ARG4]] : i32622// CHECK: linalg.yield %[[T6]] : i32623 624// -----625 626#map0 = affine_map<(i, j) -> (i, j)>627#access = [#map0, #map0]628#trait = {629 iterator_types = ["parallel", "parallel"],630 indexing_maps = #access,631 library_call = "some_external_func"632}633 634func.func @drop_all_loops(%arg0 : memref<1x1xf32>) -> memref<1x1xf32>635{636 linalg.generic #trait637 ins(%arg0 : memref<1x1xf32>)638 outs(%arg0 : memref<1x1xf32>) {639 ^bb0(%arg1: f32, %arg2: f32) :640 linalg.yield %arg1 : f32641 }642 return %arg0 : memref<1x1xf32>643}644// CHECK: #[[$MAP0:.*]] = affine_map<() -> ()>645// CHECK-LABEL: func @drop_all_loops646// CHECK: memref.collapse_shape %{{.*}} []647// CHECK: linalg.generic648// CHECK-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP0]]]649// CHECK-SAME: iterator_types = []650 651// -----652 653#map0 = affine_map<(i, j) -> (i, j)>654#access = [#map0, #map0]655#trait = {656 iterator_types = ["parallel", "parallel"],657 indexing_maps = #access,658 library_call = "some_external_func"659}660 661func.func @drop_all_loops_indexed662 (%arg0 : memref<1x1xi32>) -> memref<1x1xi32>{663 linalg.generic #trait664 ins(%arg0 : memref<1x1xi32>)665 outs(%arg0 : memref<1x1xi32>) {666 ^bb0(%arg3: i32, %arg4: i32) :667 %idx0 = linalg.index 0 : index668 %idx1 = linalg.index 1 : index669 %1 = arith.addi %idx0, %idx1 : index670 %2 = arith.index_cast %1 : index to i32671 %3 = arith.addi %2, %arg3 : i32672 linalg.yield %3 : i32673 }674 return %arg0 : memref<1x1xi32>675}676 677// CHECK-LABEL: func @drop_all_loops_indexed678// CHECK: linalg.generic679// CHECK: ^{{.+}}(%[[ARG1:.+]]: i32, %[[ARG2:.+]]: i32)680// CHECK: linalg.yield %[[ARG1]] : i32681 682// -----683 684#accesses = [685 affine_map<(d0) -> (0, d0)>,686 affine_map<(d0) -> (d0)>687]688 689#trait = {690 indexing_maps = #accesses,691 iterator_types = ["parallel"],692 library_call = "some_external_fn"693}694 695func.func @leading_dim_1_canonicalization(%arg0: memref<1x5xf32>, %shape: memref<5xf32>) -> memref<5xf32> {696 linalg.generic #trait697 ins(%arg0 : memref<1x5xf32>)698 outs(%shape : memref<5xf32>) {699 ^bb0(%arg2: f32, %arg3: f32):700 linalg.yield %arg2 : f32701 }702 return %shape : memref<5xf32>703}704// CHECK: #[[$MAP1:.*]] = affine_map<(d0) -> (d0)>705 706// CHECK-LABEL: func @leading_dim_1_canonicalization707// CHECK: memref.collapse_shape %{{.*}} {{\[}}[0, 1]]708// CHECK: linalg.generic709// CHECK-SAME: indexing_maps = [#[[$MAP1]], #[[$MAP1]]]710// CHECK-SAME: iterator_types = ["parallel"]711 712// -----713 714#accesses = [715 affine_map<(d0, d1) -> (0, d1)>,716 affine_map<(d0, d1) -> (d0, 0)>,717 affine_map<(d0, d1) -> (d0, d1)>718]719 720#trait = {721 indexing_maps = #accesses,722 iterator_types = ["parallel", "parallel"],723 library_call = "some_external_fn"724}725 726func.func @broadcast_test(%arg0 : memref<5xf32>, %arg1 : memref<5xf32>, %shape : memref<5x5xf32>) -> memref<5x5xf32>727{728 %0 = memref.expand_shape %arg0 [[0, 1]] output_shape [1, 5] : memref<5xf32> into memref<1x5xf32>729 %1 = memref.expand_shape %arg1 [[0, 1]] output_shape [5, 1] : memref<5xf32> into memref<5x1xf32>730 linalg.generic #trait731 ins(%0, %1 : memref<1x5xf32>, memref<5x1xf32>)732 outs(%shape : memref<5x5xf32>) {733 ^bb0(%arg3: f32, %arg4: f32, %arg5: f32):734 %3 = arith.addf %arg3, %arg4 : f32735 linalg.yield %3 : f32736 }737 return %shape : memref<5x5xf32>738}739// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1) -> (d1)>740// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d0)>741// CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1) -> (d0, d1)>742// CHECK-LABEL: func @broadcast_test743// CHECK-NOT: linalg.memref_{{.*}}shape744// CHECK: linalg.generic745// CHECK-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]]]746// CHECK-SAME: iterator_types = ["parallel", "parallel"]747// CHECK-NOT: linalg.memref_{{.*}}shape748 749// -----750 751#accesses = [752 affine_map<(d0, d1) -> (0, 0)>,753 affine_map<(d0, d1) -> (d0, d1)>754]755 756#trait = {757 indexing_maps = #accesses,758 iterator_types = ["parallel", "parallel"],759 library_call = "some_external_fn"760}761 762func.func @broadcast_scalar(%arg0 : memref<1x1xf32>, %shape : memref<?x?xf32>) -> memref<?x?xf32>763{764 linalg.generic #trait765 ins(%arg0 : memref<1x1xf32>)766 outs(%shape : memref<?x?xf32>) {767 ^bb0(%arg2 : f32, %arg3 : f32):768 linalg.yield %arg2 : f32769 }770 return %shape : memref<?x?xf32>771}772// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1) -> ()>773// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1) -> (d0, d1)>774// CHECK-LABEL: func @broadcast_scalar775// CHECK-SAME: %[[ARG0:.*]]: memref<1x1xf32>776// CHECK: %[[A:.*]] = memref.collapse_shape %[[ARG0]] []777// CHECK-SAME: memref<1x1xf32> into memref<f32>778// CHECK: linalg.generic779// CHECK-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP1]]]780// CHECK-SAME: iterator_types = ["parallel", "parallel"]781// CHECK-SAME: %[[A]]782 783// -----784 785#map0 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>786#map1 = affine_map<(d0, d1, d2) -> (d2)>787func.func @fold_unit_dim_memref_reshape_op(%arg0 : memref<5xf32>) -> memref<2x5xf32>788{789 %1 = memref.alloc() : memref<1x2x5xf32>790 linalg.generic {i64, indexing_maps = [#map1, #map0],791 iterator_types = ["parallel", "parallel", "parallel"]}792 ins(%arg0 : memref<5xf32>) outs(%1 : memref<1x2x5xf32>) {793 ^bb0(%arg1: f32, %arg2: f32):794 linalg.yield %arg1 : f32795 }796 %3 = memref.collapse_shape %1 [[0, 1], [2]]797 : memref<1x2x5xf32> into memref<2x5xf32>798 return %3 : memref<2x5xf32>799}800// CHECK-LABEL: func @fold_unit_dim_memref_reshape_op801// CHECK: %[[ALLOC:.*]] = memref.alloc() : memref<1x2x5xf32>802// CHECK: %[[OUT:.*]] = memref.collapse_shape %[[ALLOC]]803// CHECK: linalg.generic804// CHECK-SAME: outs(%[[OUT:.*]] :805// CHECK: %[[RESULT:.*]] = memref.collapse_shape %[[ALLOC]]806// CHECK: return %[[RESULT]]807 808// -----809 810func.func @fold_unit_dim_for_init_memref(%input: memref<1x1000xf32>) -> memref<1xf32> {811 %cst = arith.constant 0.0 : f32812 %init = memref.alloc() : memref<1xf32>813 linalg.generic {814 indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>, affine_map<(d0, d1) -> (d0)>],815 iterator_types = ["parallel", "reduction"]}816 ins(%input : memref<1x1000xf32>)outs(%init : memref<1xf32>) {817 ^bb0(%arg1: f32, %arg2: f32):818 %1823 = arith.addf %arg1, %arg2 : f32819 linalg.yield %1823 : f32820 }821 return %init : memref<1xf32>822}823 824 825// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0) -> (d0)>826// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0) -> ()>827 828// CHECK: func @fold_unit_dim_for_init_memref829// CHECK: %[[INIT:.+]] = memref.alloc() : memref<1xf32>830// CHECK: %[[INPUT_RESHAPE:.+]] = memref.collapse_shape %{{.+}} {{\[}}[0, 1]] : memref<1x1000xf32> into memref<1000xf32>831// CHECK: %[[INIT_RESHAPE:.+]] = memref.collapse_shape %[[INIT]] [] : memref<1xf32> into memref<f32>832// CHECK: linalg.generic833// CHECK-SAME: indexing_maps = [#[[MAP1]], #[[MAP2]]]834// CHECK-SAME: iterator_types = ["reduction"]835// CHECK-SAME: ins(%[[INPUT_RESHAPE]] : memref<1000xf32>)836// CHECK-SAME: outs(%[[INIT_RESHAPE]] : memref<f32>)837// CHECK: return %[[INIT:.+]] : memref<1xf32>838 839 840// -----841// Test that nothing changes and no assertions are fired for memrefs with affine842// maps while still changing the other operations.843 844#accesses = [845 affine_map<(i, j, k, l, m) -> (i, k, m)>,846 affine_map<(i, j, k, l, m) -> ()>,847 affine_map<(i, j, k, l, m) -> (i, k, j, l, m)>848]849 850#trait = {851 iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel"],852 indexing_maps = #accesses,853 library_call = "some_external_func"854}855 856func.func @input_stays_same(%arg0 : memref<?x1x?xf32, strided<[?, 1, 1]>>, %arg1 : f32, %shape: memref<?x1x?x1x?xf32>) -> memref<?x1x?x1x?xf32> {857 linalg.generic #trait858 ins(%arg0, %arg1 : memref<?x1x?xf32, strided<[?, 1, 1]>>, f32)859 outs(%shape : memref<?x1x?x1x?xf32>) {860 ^bb0(%arg2 : f32, %arg3 : f32, %arg4 : f32) :861 linalg.yield %arg3 : f32862 }863 return %shape : memref<?x1x?x1x?xf32>864}865 866// CHECK-DAG: #[[MAP1:.*]] = affine_map<(d0, d1, d2) -> (d0, 0, d2)>867// CHECK-DAG: #[[MAP2:.*]] = affine_map<(d0, d1, d2) -> ()>868// CHECK-DAG: #[[MAP3:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>869// CHECK: func @input_stays_same(870// CHECK-SAME: %[[ARG0:.*]]: memref<?x1x?xf32, strided<[?, 1, 1]>>,871// CHECK-SAME: %[[ARG1:.*]]: f32, %[[ARG2:.*]]: memref<?x1x?x1x?xf32>)872// CHECK-SAME: -> memref<?x1x?x1x?xf32> {873// CHECK: %[[OUT:.*]] = memref.collapse_shape %[[ARG2]] {{\[}}[0, 1], [2, 3], [4]]874// CHECK-SAME: : memref<?x1x?x1x?xf32> into memref<?x?x?xf32>875// CHECK: linalg.generic876// CHECK-SAME: {indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]],877// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"]}878// CHECK-SAME: ins(%[[ARG0]], %[[ARG1]] : memref<?x1x?xf32, strided<[?, 1, 1]>>, f32)879// CHECK-SAME: outs(%[[OUT]] : memref<?x?x?xf32>) {880// CHECK: ^bb0(%{{.*}}: f32, %[[ARG:.*]]: f32, %{{.*}}: f32):881// CHECK: linalg.yield %[[ARG]] : f32882// CHECK: }883// CHECK: return %[[ARG2]] : memref<?x1x?x1x?xf32>884 885// -----886 887// Negative test for case with tensor encoding.888#matvec = {889 indexing_maps = [890 affine_map<(i,j) -> (i,j)>, // A891 affine_map<(i,j) -> (j)>, // b892 affine_map<(i,j) -> (i)> // x (out)893 ],894 iterator_types = ["parallel", "reduction"]895}896 897#CSR = #sparse_tensor.encoding<{ map = (d0, d1) -> (d0 : dense, d1 : compressed) }>898 899func.func @sparse_case(%arg0: tensor<8x8xf32, #CSR>, %arg1: tensor<8xf32>) -> tensor<8xf32> {900 %0 = tensor.empty() : tensor<8xf32>901 %1 = linalg.generic #matvec902 ins(%arg0, %arg1: tensor<8x8xf32, #CSR>, tensor<8xf32>)903 outs(%0: tensor<8xf32>) {904 ^bb(%a: f32, %b: f32, %x: f32):905 %m = arith.mulf %a, %b : f32906 %add = arith.addf %x, %m : f32907 linalg.yield %add : f32908 } -> tensor<8xf32>909 return %1: tensor<8xf32>910}911 912// CHECK-LABEL: func @sparse_case913// CHECK-NEXT: tensor.empty914// CHECK-NEXT: linalg.generic915 916// -----917 918func.func @parallel_insert_slice() -> tensor<4x2xf32> {919 %c2 = arith.constant 2 : index920 %c4 = arith.constant 4 : index921 %cst = arith.constant 0.000000e+00 : f32922 %0 = tensor.empty() : tensor<4x2xf32>923 %res = scf.forall (%arg0, %arg1) in (%c4, %c2) shared_outs(%o = %0) -> (tensor<4x2xf32>) {924 %1 = tensor.empty() : tensor<1x1xf32>925 %2 = linalg.fill ins(%cst : f32) outs(%1 : tensor<1x1xf32>) -> tensor<1x1xf32>926 // CHECK: scf.forall.in_parallel927 scf.forall.in_parallel {928 // CHECK: tensor.parallel_insert_slice %{{[0-9a-z]*}} into %{{[0-9a-z]*}}929 // CHECK-SAME: [%{{.*}}, %{{.*}}] [1, 1] [1, 1] : tensor<f32> into tensor<4x2xf32>930 tensor.parallel_insert_slice %2 into %o[%arg0, %arg1] [1, 1] [1, 1] :931 tensor<1x1xf32> into tensor<4x2xf32>932 }933 }934 return %res: tensor<4x2xf32>935}936 937// -----938 939#map0 = affine_map<(i, j) -> (i, j)>940#access = [#map0, #map0]941#trait = {942 iterator_types = ["parallel", "parallel"],943 indexing_maps = #access,944 library_call = "some_external_func"945}946 947func.func @drop_all_loops(%arg0 : memref<1x1xf32, 3>) -> memref<1x1xf32, 3>948{949 linalg.generic #trait950 ins(%arg0 : memref<1x1xf32, 3>)951 outs(%arg0 : memref<1x1xf32, 3>) {952 ^bb0(%arg1: f32, %arg2: f32) :953 linalg.yield %arg1 : f32954 }955 return %arg0 : memref<1x1xf32, 3>956}957 958// CHECK-LABEL: func @drop_all_loops959// CHECK: memref.collapse_shape960// CHECK-SAME: [] : memref<1x1xf32, 3> into memref<f32, 3>961// CHECK: linalg.generic{{.*}}memref<f32, 3>962 963// CHECK-SLICES-LABEL: func @drop_all_loops964// CHECK-SLICES: memref.subview %{{.*}}[0, 0] [1, 1] [1, 1] : memref<1x1xf32, 3> to memref<f32, strided<[]>, 3>965// CHECK-SLICES: linalg.generic{{.*}}memref<f32, strided<[]>, 3>966 967// -----968 969func.func @drop_unit_pad_dims(%arg0: tensor<1x1x3x1x1xf32>) -> tensor<1x2x3x1x3xf32>970{971 %c0 = arith.constant 0 : index972 %cst0 = arith.constant 0.0 : f32973 %0 = tensor.pad %arg0 low[0, 1, 0, %c0, 0] high[0, 0, 0, %c0, 2] {974 ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index):975 tensor.yield %cst0 : f32976 } : tensor<1x1x3x1x1xf32> to tensor<1x2x3x1x3xf32>977 return %0 : tensor<1x2x3x1x3xf32>978}979 980// CHECK-LABEL: func @drop_unit_pad_dims981// CHECK: %[[COLLAPSE:.+]] = tensor.collapse_shape982// CHECK-SAME: {{\[}}[0, 1], [2, 3], [4]{{\]}} : tensor<1x1x3x1x1xf32> into tensor<1x3x1xf32>983// CHECK: %[[PADDED:.+]] = tensor.pad %[[COLLAPSE]] low[1, 0, 0] high[0, 0, 2]984// CHECK: } : tensor<1x3x1xf32> to tensor<2x3x3xf32>985// CHECK: tensor.expand_shape %[[PADDED]]986// CHECK-SAME: {{\[}}[0, 1], [2, 3], [4]{{\]}} output_shape [1, 2, 3, 1, 3] : tensor<2x3x3xf32> into tensor<1x2x3x1x3xf32>987 988// CHECK-SLICES-LABEL: func @drop_unit_pad_dims989// CHECK-SLICES: %[[EXTRACT:.+]] = tensor.extract_slice990// CHECK-SLICES-SAME: [0, 0, 0, 0, 0] [1, 1, 3, 1, 1] [1, 1, 1, 1, 1] : tensor<1x1x3x1x1xf32> to tensor<1x3x1xf32>991// CHECK-SLICES: %[[PADDED:.+]] = tensor.pad %[[EXTRACT]] low[1, 0, 0] high[0, 0, 2]992// CHECK-SLICES: } : tensor<1x3x1xf32> to tensor<2x3x3xf32>993// CHECK-SLICES: tensor.insert_slice %[[PADDED]]994// CHECK-SLICES-SAME: [0, 0, 0, 0, 0] [1, 2, 3, 1, 3] [1, 1, 1, 1, 1] : tensor<2x3x3xf32> into tensor<1x2x3x1x3xf32>995 996// -----997 998func.func @drop_unit_pad_dynamic_dims(%arg0: tensor<1x?xf32>) -> tensor<1x?xf32>999{1000 %c0 = arith.constant 0 : index1001 %cst0 = arith.constant 0.0 : f321002 %0 = tensor.pad %arg0 low[0, 5] high[0, 6] {1003 ^bb0(%arg1: index, %arg2: index):1004 tensor.yield %cst0 : f321005 } : tensor<1x?xf32> to tensor<1x?xf32>1006 return %0 : tensor<1x?xf32>1007}1008 1009// CHECK-DAG: #[[$MAP1:.+]] = affine_map<()[s0] -> (s0 + 11)>1010// CHECK-LABEL: func @drop_unit_pad_dynamic_dims1011// CHECK: %[[C1:.*]] = arith.constant 1 : index1012// CHECK: %[[CST:.*]] = arith.constant 0.000000e+00 : f321013// CHECK: %[[COLLAPSE:.+]] = tensor.collapse_shape1014// CHECK-SAME: {{\[}}[0, 1]{{\]}} : tensor<1x?xf32> into tensor<?xf32>1015// CHECK: %[[PADDED:.+]] = tensor.pad %[[COLLAPSE]] low[5] high[6]1016// CHECK: } : tensor<?xf32> to tensor<?xf32>1017// CHECK: %[[DIM:.+]] = tensor.dim %{{.*}}, %[[C1]] : tensor<1x?xf32>1018// CHECK: %[[VAL_1:.+]] = affine.apply #[[$MAP1]]()[%[[DIM]]]1019// CHECK: %[[EXPANDED:.+]] = tensor.expand_shape %[[PADDED]] {{\[\[}}0, 1]] output_shape [1, %[[VAL_1]]] : tensor<?xf32> into tensor<1x?xf32>1020 1021// CHECK-SLICES: #[[$MAP:.+]] = affine_map<()[s0] -> (s0 + 11)>1022 1023// CHECK-SLICES-LABEL: func @drop_unit_pad_dynamic_dims1024// CHECK-SLICES-SAME: %[[ARG0:[A-Za-z0-9]+]]: tensor<1x?xf32>1025// CHECK-SLICES: %[[DIM:.+]] = tensor.dim %[[ARG0]], %c11026// CHECK-SLICES: %[[EXTRACT:.+]] = tensor.extract_slice1027// CHECK-SLICES-SAME: [0, 0] [1, %[[DIM]]] [1, 1] : tensor<1x?xf32> to tensor<?xf32>1028// CHECK-SLICES: %[[PADDED:.+]] = tensor.pad %[[EXTRACT]] low[5] high[6]1029// CHECK-SLICES: } : tensor<?xf32> to tensor<?xf32>1030// CHECK-SLICES: %[[PADDED_DIM:.+]] = affine.apply #[[$MAP]]()[%[[DIM]]]1031// CHECK-SLICES: %[[EMPTY:.+]] = tensor.empty(%[[PADDED_DIM]]) : tensor<1x?xf32>1032// CHECK-SLICES: tensor.insert_slice %[[PADDED]] into %[[EMPTY]]1033// CHECK-SLICES-SAME: [0, 0] [1, %[[PADDED_DIM]]] [1, 1] : tensor<?xf32> into tensor<1x?xf32>1034 1035// -----1036 1037func.func @do_not_drop_non_constant_padding(%arg0: tensor<1x1x3x1x1xf32>, %pad: f32) -> tensor<1x2x3x1x3xf32>1038{1039 %c0 = arith.constant 0 : index1040 %0 = tensor.pad %arg0 low[0, 1, 0, %c0, 0] high[0, 0, 0, %c0, 2] {1041 ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index):1042 %0 = arith.index_cast %arg3 : index to i641043 %1 = arith.sitofp %0 : i64 to f321044 %add = arith.addf %pad, %1 : f321045 tensor.yield %add : f321046 } : tensor<1x1x3x1x1xf32> to tensor<1x2x3x1x3xf32>1047 return %0 : tensor<1x2x3x1x3xf32>1048}1049 1050// CHECK-LABEL: func @do_not_drop_non_constant_padding1051// CHECK: tensor.pad %{{.*}} low[0, 1, 0, %c0, 0] high[0, 0, 0, %c0, 2]1052// CHECK: } : tensor<1x1x3x1x1xf32> to tensor<1x2x3x1x3xf32>1053 1054// CHECK-SLICES-LABEL: func @do_not_drop_non_constant_padding1055// CHECK-SLICES: tensor.pad %{{.*}} low[0, 1, 0, %c0, 0] high[0, 0, 0, %c0, 2]1056// CHECK-SLICES: } : tensor<1x1x3x1x1xf32> to tensor<1x2x3x1x3xf32>1057 1058// -----1059 1060func.func @drop_known_unit_constant_low_high(%arg0: tensor<1x383x128xf32>) -> tensor<1x384x128xf32> {1061 %c0 = arith.constant 0 : index1062 %c1 = arith.constant 1 : index1063 %cst = arith.constant 0.000000e+00 : f321064 %padded = tensor.pad %arg0 low[%c0, %c1, %c0] high[%c0, %c0, %c0] {1065 ^bb0(%arg1: index, %arg2: index, %arg3: index):1066 tensor.yield %cst : f321067 } : tensor<1x383x128xf32> to tensor<1x384x128xf32>1068 return %padded : tensor<1x384x128xf32>1069}1070// CHECK-LABEL: func @drop_known_unit_constant_low_high1071// CHECK: %[[COLLAPSE:.+]] = tensor.collapse_shape1072// CHECK-SAME: {{\[}}[0, 1], [2]] : tensor<1x383x128xf32> into tensor<383x128xf32>1073// CHECK: %[[PADDED:.+]] = tensor.pad %[[COLLAPSE]] low[1, 0] high[0, 0]1074// CHECK: } : tensor<383x128xf32> to tensor<384x128xf32>1075// CHECK: tensor.expand_shape %[[PADDED]]1076// CHECK-SAME: {{\[}}[0, 1], [2]] output_shape [1, 384, 128] : tensor<384x128xf32> into tensor<1x384x128xf32>1077 1078// -----1079 1080func.func @drop_unit_dim_mixed_static_dynamic(%arg0: tensor<1x?xf32>) -> tensor<1x16xf32> {1081 %c0 = arith.constant 0 : index1082 %c1 = arith.constant 1 : index1083 %cst = arith.constant 0.000000e+00 : f321084 %padded = tensor.pad %arg0 low[%c0, %c1] high[%c0, %c0] {1085 ^bb0(%arg1: index, %arg2: index):1086 tensor.yield %cst : f321087 } : tensor<1x?xf32> to tensor<1x16xf32>1088 return %padded : tensor<1x16xf32>1089}1090// CHECK-LABEL: func @drop_unit_dim_mixed_static_dynamic1091// CHECK: %[[CST:.*]] = arith.constant 0.000000e+00 : f321092// CHECK: %[[COLLAPSE:.+]] = tensor.collapse_shape %[[ARGS:.*]] : tensor<1x?xf32> into tensor<?xf32>1093// CHECK: %[[PADDED:.*]] = tensor.pad %[[COLLAPSE]] low[1] high[0] {1094// CHECK: ^bb0(%[[IDX:.*]]: index):1095// CHECK: tensor.yield %[[CST]] : f321096// CHECK: } : tensor<?xf32> to tensor<16xf32>1097// CHECK: %[[EXPANDED:.*]] = tensor.expand_shape %[[PADDED]] {{\[\[}}0, 1]] output_shape [1, 16] : tensor<16xf32> into tensor<1x16xf32>1098// CHECK: return %[[EXPANDED]] : tensor<1x16xf32>1099 1100// -----1101 1102#map = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1 + d4, d2 + d5, d6)>1103#map1 = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d4, d5, d6, d3)>1104#map2 = affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3)>1105module {1106 func.func @drop_unit_dim_corresponding_to_dynamic_dim(%arg0: tensor<1x?x?x1xf32>, %arg1: index) -> tensor<?x1x61x1xf32> {1107 %cst = arith.constant dense<1.000000e+00> : tensor<1x1x1x1xf32>1108 %0 = tensor.empty(%arg1) : tensor<?x1x61x1xf32>1109 %1 = linalg.generic {indexing_maps = [#map, #map1, #map2], iterator_types = ["parallel", "parallel", "parallel", "parallel", "reduction", "reduction", "reduction"]} ins(%arg0, %cst : tensor<1x?x?x1xf32>, tensor<1x1x1x1xf32>) outs(%0 : tensor<?x1x61x1xf32>) {1110 ^bb0(%in: f32, %in_0: f32, %out: f32):1111 %2 = arith.mulf %in, %in_0 : f321112 %3 = arith.addf %out, %2 : f321113 linalg.yield %3 : f321114 } -> tensor<?x1x61x1xf32>1115 return %1 : tensor<?x1x61x1xf32>1116 }1117}1118// CHECK: #[[$MAP1:.+]] = affine_map<(d0) -> (0, d0)>1119// CHECK: #[[$MAP2:.+]] = affine_map<(d0) -> ()>1120 1121// CHECK-LABEL: func @drop_unit_dim_corresponding_to_dynamic_dim1122// CHECK-SAME: %[[ARG0:.*]]: tensor<1x?x?x1xf32>,1123// CHECK-SAME: %[[ARG1:.*]]: index) -> tensor<?x1x61x1xf32> {1124// CHECK: %[[VAL_0:.*]] = arith.constant 0 : index1125// CHECK: %[[VAL_2:.*]] = arith.constant dense<1.000000e+00> : tensor<f32>1126// CHECK: %[[VAL_3:.*]] = tensor.collapse_shape %[[ARG0]] {{\[\[}}0, 1], [2, 3]] : tensor<1x?x?x1xf32> into tensor<?x?xf32>1127// CHECK: %[[VAL_4:.*]] = tensor.empty(%[[ARG1]]) : tensor<?x61xf32>1128// CHECK: %[[VAL_6:.*]] = tensor.empty(%[[ARG1]]) : tensor<?x61xf32>1129// CHECK: %[[VAL_7:.*]] = linalg.generic {indexing_maps = [#[[$MAP1]], #[[$MAP2]], #[[$MAP1]], #[[$MAP1]]], iterator_types = ["parallel"]} ins(%[[VAL_3]], %[[VAL_2]], %[[VAL_4]] : tensor<?x?xf32>, tensor<f32>, tensor<?x61xf32>) outs(%[[VAL_6]] : tensor<?x61xf32>) {1130// CHECK: ^bb0(%[[VAL_8:.*]]: f32, %[[VAL_9:.*]]: f32, %[[VAL_10:.*]]: f32, %[[VAL_11:.*]]: f32):1131// CHECK: %[[VAL_12:.*]] = arith.mulf %[[VAL_8]], %[[VAL_9]] : f321132// CHECK: %[[VAL_13:.*]] = arith.addf %[[VAL_10]], %[[VAL_12]] : f321133// CHECK: linalg.yield %[[VAL_13]] : f321134// CHECK: } -> tensor<?x61xf32>1135// CHECK: %[[VAL_14:.*]] = tensor.expand_shape %[[VAL_7]] {{\[\[}}0, 1], [2, 3]] output_shape {{\[}}%[[VAL_0]], 1, 61, 1] : tensor<?x61xf32> into tensor<?x1x61x1xf32>1136// CHECK: return %[[VAL_14]] : tensor<?x1x61x1xf32>1137// CHECK: }1138 1139// -----1140 1141func.func @no_fold_empty_tensor_dim_out_of_bounds(%arg0: tensor<1x?x10xf32>) -> tensor<1x?xf32> {1142 %cst = arith.constant 1.000000e+00 : f321143 %cst7 = arith.constant 7 : index1144 %dim = tensor.dim %arg0, %cst7 : tensor<1x?x10xf32>1145 %0 = tensor.empty(%dim) : tensor<1x?xf32>1146 %1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<1x?xf32>) -> tensor<1x?xf32>1147 return %1 : tensor<1x?xf32>1148}1149// CHECK-LABEL: func.func @no_fold_empty_tensor_dim_out_of_bounds1150// CHECK-SAME: %[[ARG0:.*]]: tensor<1x?x10xf32>) -> tensor<1x?xf32> {1151// CHECK: %[[CST:.*]] = arith.constant 1.000000e+00 : f321152// CHECK: %[[C7:.*]] = arith.constant 71153// CHECK: %[[DIM:.*]] = tensor.dim %[[ARG0]], %[[C7]] : tensor<1x?x10xf32>1154// CHECK: %[[VAL_0:.*]] = tensor.empty(%[[DIM]]) : tensor<1x?xf32>1155// CHECK: %[[VAL_1:.*]] = linalg.fill ins(%[[CST]] : f32) outs(%[[VAL_0]] : tensor<1x?xf32>) -> tensor<1x?xf32>1156// CHECK: return %[[VAL_1]] : tensor<1x?xf32>1157// CHECK: }1158 1159// -----1160 1161func.func @fold_empty_tensor_dim_op(%arg0: tensor<1x?x10xf32>) -> tensor<1x?xf32> {1162 %cst = arith.constant 1.000000e+00 : f321163 %cst2 = index.constant 21164 %dim10 = tensor.dim %arg0, %cst2 : tensor<1x?x10xf32>1165 %0 = tensor.empty(%dim10) : tensor<1x?xf32>1166 %1 = linalg.fill ins(%cst : f32) outs(%0 : tensor<1x?xf32>) -> tensor<1x?xf32>1167 return %1 : tensor<1x?xf32>1168}1169// CHECK-LABEL: func.func @fold_empty_tensor_dim_op1170// CHECK-SAME: %[[ARG0:.*]]: tensor<1x?x10xf32>) -> tensor<1x?xf32> {1171// CHECK: %[[CST:.*]] = arith.constant 1.000000e+00 : f321172// CHECK: %[[VAL_0:.*]] = tensor.empty() : tensor<1x10xf32>1173// CHECK: %[[VAL_1:.*]] = tensor.cast %[[VAL_0]] : tensor<1x10xf32> to tensor<1x?xf32>1174// CHECK: %[[VAL_2:.*]] = linalg.fill ins(%[[CST]] : f32) outs(%[[VAL_1]] : tensor<1x?xf32>) -> tensor<1x?xf32>1175// CHECK: return %[[VAL_2]] : tensor<1x?xf32>1176// CHECK: }1177