forked from OSchip/llvm-project
64 lines
2.9 KiB
MLIR
64 lines
2.9 KiB
MLIR
// RUN: mlir-opt -split-input-file --test-linalg-transform-patterns="test-transform-pad-tensor" %s | FileCheck --check-prefix=CHECK %s
|
|
|
|
// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
|
|
// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2, d3) -> (d0 + 1, d1 + 1, d2 + 1, d3 + 2)>
|
|
// CHECK-LABEL: func @pad_tensor_with_memrefs
|
|
func.func @pad_tensor_with_memrefs(%arg0: memref<1x28x28x1xf32>) -> memref<2x31x31x3xf32> {
|
|
%cst = arith.constant 0.000000e+00 : f32
|
|
%0 = bufferization.to_tensor %arg0 : memref<1x28x28x1xf32>
|
|
%1 = tensor.pad %0 low[1, 1, 1, 2] high[0, 2, 2, 0] {
|
|
^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
|
|
tensor.yield %cst : f32
|
|
} : tensor<1x28x28x1xf32> to tensor<2x31x31x3xf32>
|
|
%2 = bufferization.to_memref %1 : memref<2x31x31x3xf32>
|
|
return %2 : memref<2x31x31x3xf32>
|
|
}
|
|
|
|
// CHECK: linalg.fill
|
|
// CHECK: linalg.generic
|
|
// CHECK-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP1]]]
|
|
|
|
// -----
|
|
|
|
// CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
|
|
// CHECK-DAG: #[[$MAP3:.*]] = affine_map<(d0, d1, d2) -> (d0 + 1, d1 + 2, d2 + 2)>
|
|
// CHECK-LABEL: func @pad_tensor_no_memrefs
|
|
func.func @pad_tensor_no_memrefs(%arg0: tensor<1x28x28xf32>) -> tensor<2x32x32xf32> {
|
|
%cst = arith.constant 0.000000e+00 : f32
|
|
%0 = tensor.pad %arg0 low[1, 2, 2] high[0, 2, 2] {
|
|
^bb0(%arg1: index, %arg2: index, %arg3: index):
|
|
tensor.yield %cst : f32
|
|
} : tensor<1x28x28xf32> to tensor<2x32x32xf32>
|
|
return %0 : tensor<2x32x32xf32>
|
|
}
|
|
|
|
// CHECK: linalg.fill
|
|
// CHECK: linalg.generic
|
|
// CHECK-SAME: indexing_maps = [#[[$MAP2]], #[[$MAP3]]]
|
|
|
|
// -----
|
|
|
|
// CHECK-DAG: #[[$MAP4:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
|
|
// CHECK-DAG: #[[$MAP5:.*]] = affine_map<(d0, d1, d2, d3) -> (d0, d1 + 2, d2 + 2, d3)>
|
|
// CHECK-LABEL: func @pad_tensor_detailed
|
|
func.func @pad_tensor_detailed(%arg0: tensor<1x28x28x1xf32>) -> tensor<1x32x32x1xf32> {
|
|
%cst = arith.constant 0.000000e+00 : f32
|
|
%0 = tensor.pad %arg0 low[0, 2, 2, 0] high[0, 2, 2, 0] {
|
|
^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
|
|
tensor.yield %cst : f32
|
|
} : tensor<1x28x28x1xf32> to tensor<1x32x32x1xf32>
|
|
return %0 : tensor<1x32x32x1xf32>
|
|
}
|
|
|
|
// CHECK: %[[ARG0:[a-zA-Z0-9_]+]]: tensor<1x28x28x1xf32>) -> tensor<1x32x32x1xf32>
|
|
// CHECK: %[[CTE:.+]] = arith.constant 0.000000e+00 : f32
|
|
// CHECK: %[[TMP:.+]] = linalg.init_tensor [1, 32, 32, 1] : tensor<1x32x32x1xf32>
|
|
// CHECK: %[[R1c:.+]] = linalg.fill
|
|
// CHECK: %[[R2c:.+]] = linalg.generic
|
|
// CHECK-SAME: indexing_maps = [#[[$MAP4]], #[[$MAP5]]]
|
|
// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel"]
|
|
// CHECK: ins(%arg0 : tensor<1x28x28x1xf32>) outs(%1 : tensor<1x32x32x1xf32>)
|
|
// CHECK: ^bb0(%[[VAL:.+]]: f32, %arg2: f32)
|
|
// CHECK: linalg.yield %[[VAL]] : f32
|
|
// CHECK: return %[[R2c:.+]]
|