mlir/test/Dialect/Linalg/lower-pad-tensor.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 @pad_tensor_with_memrefs(%arg0: memref<1x28x28x1xf32>) -> memref<2x31x31x3xf32> {
  %cst = arith.constant 0.000000e+00 : f32
  %0 = bufferization.to_tensor %arg0 : memref<1x28x28x1xf32>
  %1 = linalg.pad_tensor %0 low[1, 1, 1, 2] high[0, 2, 2, 0]  {
  ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):  // no predecessors
    linalg.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 @pad_tensor_no_memrefs(%arg0: tensor<1x28x28xf32>) -> tensor<2x32x32xf32> {
  %cst = arith.constant 0.000000e+00 : f32
  %0 = linalg.pad_tensor %arg0 low[1, 2, 2] high[0, 2, 2]  {
  ^bb0(%arg1: index, %arg2: index, %arg3: index):  // no predecessors
    linalg.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 @pad_tensor_detailed(%arg0: tensor<1x28x28x1xf32>) -> tensor<1x32x32x1xf32> {
  %cst = arith.constant 0.000000e+00 : f32
  %0 = linalg.pad_tensor %arg0 low[0, 2, 2, 0] high[0, 2, 2, 0]  {
  ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):  // no predecessors
    linalg.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:.+]]