46 lines · plain
1// RUN: mlir-opt -split-input-file --test-linalg-transform-patterns="test-decompose-pad-tensor" %s | FileCheck %s2 3// CHECK-LABEL: func @generalize_pad_tensor_static_shape(4// CHECK-SAME: %[[IN:.*]]: tensor<1x28x28x1xf32>) -> tensor<1x32x32x1xf32> {5// CHECK: %[[C0:.*]] = arith.constant 0.000000e+00 : f326// CHECK: %[[INIT:.*]] = tensor.empty() : tensor<1x32x32x1xf32>7// CHECK: %[[FILL:.*]] = linalg.fill ins(%[[C0]] : f32) outs(%[[INIT]] : tensor<1x32x32x1xf32>) -> tensor<1x32x32x1xf32>8// CHECK: %[[PADDED:.*]] = tensor.insert_slice %[[IN]] into %[[FILL]][0, 2, 2, 0] [1, 28, 28, 1] [1, 1, 1, 1] : tensor<1x28x28x1xf32> into tensor<1x32x32x1xf32>9// CHECK: return %[[PADDED]] : tensor<1x32x32x1xf32>10func.func @generalize_pad_tensor_static_shape(%arg0: tensor<1x28x28x1xf32>) -> tensor<1x32x32x1xf32> {11 %cst = arith.constant 0.000000e+00 : f3212 %0 = tensor.pad %arg0 low[0, 2, 2, 0] high[0, 2, 2, 0] {13 ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):14 tensor.yield %cst : f3215 } : tensor<1x28x28x1xf32> to tensor<1x32x32x1xf32>16 return %0 : tensor<1x32x32x1xf32>17}18 19// CHECK-LABEL: func @generalize_pad_tensor_dynamic_shape(20// CHECK-SAME: %[[IN:.*]]: tensor<4x?x2x?xf32>,21// CHECK-SAME: %[[OFFSET:.*]]: index) -> tensor<4x?x?x?xf32> {22// CHECK-DAG: %[[CST:.*]] = arith.constant 0.000000e+00 : f3223// CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index24// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index25// CHECK-DAG: %[[C3:.*]] = arith.constant 3 : index26// CHECK: %[[DIM1:.*]] = tensor.dim %[[IN]], %[[C1]] : tensor<4x?x2x?xf32>27// CHECK: %[[OUT_DIM2:.*]] = arith.addi %[[OFFSET]], %[[C2]] : index28// CHECK: %[[DIM3:.*]] = tensor.dim %[[IN]], %[[C3]] : tensor<4x?x2x?xf32>29// CHECK: %[[OUT_DIM3:.*]] = arith.addi %[[DIM3]], %[[OFFSET]] : index30// CHECK: %[[INIT:.*]] = tensor.empty(%[[DIM1]], %[[OUT_DIM2]], %[[OUT_DIM3]]) : tensor<4x?x?x?xf32>31// CHECK: %[[FILL:.*]] = linalg.fill ins(%[[CST]] : f32) outs(%[[INIT]] : tensor<4x?x?x?xf32>) -> tensor<4x?x?x?xf32>32// CHECK: %[[DIM1_1:.*]] = tensor.dim %[[IN]], %[[C1]] : tensor<4x?x2x?xf32>33// CHECK: %[[DIM3_1:.*]] = tensor.dim %[[IN]], %[[C3]] : tensor<4x?x2x?xf32>34// CHECK: %[[PADDED:.*]] = tensor.insert_slice %[[IN]] into %[[FILL]][0, 0, %[[OFFSET]], 0] [4, %[[DIM1_1]], 2, %[[DIM3_1]]] [1, 1, 1, 1] : tensor<4x?x2x?xf32> into tensor<4x?x?x?xf32>35// CHECK: return %[[PADDED]] : tensor<4x?x?x?xf32>36// CHECK: }37func.func @generalize_pad_tensor_dynamic_shape(%arg0: tensor<4x?x2x?xf32>, %arg1: index) -> tensor<4x?x?x?xf32> {38 %c0 = arith.constant 0 : index39 %cst = arith.constant 0.0 : f3240 %out = tensor.pad %arg0 low[%c0, %c0, %arg1, %c0] high[%c0, %c0, %c0, %arg1] {41 ^bb0(%gen_arg1: index, %gen_arg2: index, %gen_arg3: index, %gen_arg4: index):42 tensor.yield %cst : f3243 } : tensor<4x?x2x?xf32> to tensor<4x?x?x?xf32>44 return %out : tensor<4x?x?x?xf32>45}46