xref: /llvm-project/mlir/test/Dialect/Linalg/swap-extract-slice-with-fill.mlir (revision c325e978b529e7e899a31cd2693fbc1a676afcf0)
1//RUN: mlir-opt -split-input-file -test-linalg-transform-patterns=test-swap-extract-slice-with-fill-pattern %s | FileCheck %s
2
3// CHECK-LABEL: func.func @swap_fill_insert_slice
4//  CHECK-SAME: (%[[INIT:.+]]: tensor<?x?x?xf32>, %[[OFFSET0:.+]]: index, %[[SIZE1:.+]]: index)
5//       CHECK:   %[[F0:.+]] = arith.constant 0.000000e+00 : f32
6//       CHECK:   %[[EXT:.+]] = tensor.extract_slice %[[INIT]][%[[OFFSET0]], 8, 4] [1, %[[SIZE1]], 6] [1, 3, 1]
7//       CHECK:   %[[FILL:.+]] = linalg.fill ins(%[[F0]] : f32) outs(%[[EXT]] : tensor<?x6xf32>) -> tensor<?x6xf32>
8//       CHECK:   return %[[FILL]]
9func.func @swap_fill_insert_slice(%init : tensor<?x?x?xf32>, %offset0: index, %size1: index) -> tensor<?x6xf32> {
10  %f0 = arith.constant 0.000000e+00 : f32
11  %0 = linalg.fill ins(%f0 : f32) outs(%init : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
12  %1 = tensor.extract_slice %0[%offset0, 8, 4] [1, %size1, 6] [1, 3, 1]
13    : tensor<?x?x?xf32> to tensor<?x6xf32>
14  return %1: tensor<?x6xf32>
15}
16
17// -----
18
19// CHECK-LABEL: func.func @dont_swap_fill_insert_slice_multi_user
20//       CHECK:   linalg.fill
21//       CHECK:   tensor.extract_slice
22func.func @dont_swap_fill_insert_slice_multi_user(%init : tensor<?x?x?xf32>, %offset0: index, %size1: index) -> (tensor<?x?x?xf32>, tensor<2x?x6xf32>) {
23  %f0 = arith.constant 0.000000e+00 : f32
24  %0 = linalg.fill ins(%f0 : f32) outs(%init : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>
25  %1 = tensor.extract_slice %0[%offset0, 8, 4] [2, %size1, 6] [1, 3, 1]
26    : tensor<?x?x?xf32> to tensor<2x?x6xf32>
27  return %0, %1: tensor<?x?x?xf32>, tensor<2x?x6xf32>
28}
29