llvm-project/mlir/test/Dialect/Linalg/decompose-convolution.mlir
gysit c8f2139eb0 [mlir][linalg] Add decompose to CodegenStrategy.
Add the decompose patterns that lower higher dimensional convolutions to lower dimensional ones to CodegenStrategy and use CodegenStrategy to test the decompose patterns. Additionally, remove the assertion that checks the anchor op name is set in the CodegenStrategyTest pass. Removing the assertion allows us to simplify the pipelines used in the interchange and decompose tests.

Depends On D114797

Reviewed By: nicolasvasilache

Differential Revision: https://reviews.llvm.org/D114798
2021-11-30 15:48:29 +00:00

95 lines
5.5 KiB
MLIR

// RUN: mlir-opt -test-linalg-codegen-strategy="decompose" -split-input-file %s | FileCheck %s
// CHECK-LABEL: func @conv2d_nhwc_4x1x2x8_tensor
// CHECK-SAME: (%[[INPUT:.+]]: tensor<4x1x6x3xf32>, %[[FILTER:.+]]: tensor<1x2x3x8xf32>, %[[INIT:.+]]: tensor<4x1x2x8xf32>)
func @conv2d_nhwc_4x1x2x8_tensor(%input: tensor<4x1x6x3xf32>, %filter: tensor<1x2x3x8xf32>, %init: tensor<4x1x2x8xf32>) -> tensor<4x1x2x8xf32> {
%0 = linalg.conv_2d_nhwc_hwcf
{dilations = dense<[2, 3]> : tensor<2xi64>, strides = dense<[3, 2]> : tensor<2xi64>}
ins(%input, %filter : tensor<4x1x6x3xf32>, tensor<1x2x3x8xf32>)
outs(%init : tensor<4x1x2x8xf32>) -> tensor<4x1x2x8xf32>
return %0 : tensor<4x1x2x8xf32>
}
// CHECK: %[[INPUT_1D:.+]] = tensor.extract_slice %[[INPUT]]
// CHECK-SAME{LITERAL}: [0, 0, 0, 0] [4, 1, 6, 3] [1, 1, 1, 1] : tensor<4x1x6x3xf32> to tensor<4x6x3xf32>
// CHECK: %[[FILTER_1D:.+]] = tensor.extract_slice %[[FILTER]]
// CHECK-SAME{LITERAL}: [0, 0, 0, 0] [1, 2, 3, 8] [1, 1, 1, 1] : tensor<1x2x3x8xf32> to tensor<2x3x8xf32>
// CHECK: %[[INIT_1D:.+]] = tensor.extract_slice %[[INIT]]
// CHECK-SAME{LITERAL}: [0, 0, 0, 0] [4, 1, 2, 8] [1, 1, 1, 1] : tensor<4x1x2x8xf32> to tensor<4x2x8xf32>
// CHECK: %[[CONV_1D:.+]] = linalg.conv_1d_nwc_wcf
// CHECK-SAME: dilations = dense<3> : vector<1xi64>
// CHECK-SAME: strides = dense<2> : vector<1xi64>
// CHECK-SAME: ins(%[[INPUT_1D]], %[[FILTER_1D]] : tensor<4x6x3xf32>, tensor<2x3x8xf32>)
// CHECK-SAME: outs(%[[INIT_1D]] : tensor<4x2x8xf32>)
// CHECK: %[[CONV_2D:.+]] = tensor.insert_slice %[[CONV_1D]] into %[[INIT]]
// CHECK-SAME{LITERAL}: [0, 0, 0, 0] [4, 1, 2, 8] [1, 1, 1, 1] : tensor<4x2x8xf32> into tensor<4x1x2x8xf32>
// CHECK: return %[[CONV_2D]]
// -----
// CHECK-LABEL: func @conv2d_nhwc_qxqx1xq_tensor
// CHECK-SAME: (%[[INPUT:.+]]: tensor<?x?x1x?xf32>, %[[FILTER:.+]]: tensor<?x1x?x?xf32>, %[[INIT:.+]]: tensor<?x?x1x?xf32>)
func @conv2d_nhwc_qxqx1xq_tensor(%input: tensor<?x?x1x?xf32>, %filter: tensor<?x1x?x?xf32>, %init: tensor<?x?x1x?xf32>) -> tensor<?x?x1x?xf32> {
%0 = linalg.conv_2d_nhwc_hwcf
{dilations = dense<[2, 3]> : tensor<2xi64>, strides = dense<[3, 2]> : tensor<2xi64>}
ins(%input, %filter : tensor<?x?x1x?xf32>, tensor<?x1x?x?xf32>)
outs(%init : tensor<?x?x1x?xf32>) -> tensor<?x?x1x?xf32>
return %0 : tensor<?x?x1x?xf32>
}
// CHECK: %[[INPUT_1D:.+]] = tensor.extract_slice %[[INPUT]]
// CHECK-SAME: [0, 0, 0, 0] [%{{.*}}, %{{.*}}, 1, %{{.*}}] [1, 1, 1, 1] :
// CHECK-SAME: tensor<?x?x1x?xf32> to tensor<?x?x?xf32>
// CHECK: %[[FILTER_1D:.+]] = tensor.extract_slice %[[FILTER]]
// CHECK-SAME: [0, 0, 0, 0] [%{{.*}}, 1, %{{.*}}, %{{.*}}] [1, 1, 1, 1] :
// CHECK-SAME: tensor<?x1x?x?xf32> to tensor<?x?x?xf32>
// CHECK: %[[INIT_1D:.+]] = tensor.extract_slice %[[INIT]]
// CHECK-SAME: [0, 0, 0, 0] [%{{.*}}, %{{.*}}, 1, %{{.*}}] [1, 1, 1, 1] :
// CHECK-SAME: tensor<?x?x1x?xf32> to tensor<?x?x?xf32>
// CHECK: %[[CONV_1D:.+]] = linalg.conv_1d_nwc_wcf
// CHECK-SAME: dilations = dense<2> : vector<1xi64>
// CHECK-SAME: strides = dense<3> : vector<1xi64>
// CHECK-SAME: ins(%[[INPUT_1D]], %[[FILTER_1D]] : tensor<?x?x?xf32>, tensor<?x?x?xf32>)
// CHECK-SAME: outs(%[[INIT_1D]] : tensor<?x?x?xf32>)
// CHECK: %[[CONV_2D:.+]] = tensor.insert_slice %[[CONV_1D]] into %[[INIT]]
// CHECK-SAME: [0, 0, 0, 0] [%{{.*}}, %{{.*}}, 1, %{{.*}}] [1, 1, 1, 1] :
// CHECK-SAME: tensor<?x?x?xf32> into tensor<?x?x1x?xf32>
// CHECK: return %[[CONV_2D]]
// -----
// Do not convert convolution ops whose window dimensions are not ones.
// CHECK-LABEL: func @conv2d_nhwc_4x1x2x8_tensor
func @conv2d_nhwc_4x1x2x8_tensor(%input: tensor<4x3x5x3xf32>, %filter: tensor<2x2x3x8xf32>, %init: tensor<4x1x2x8xf32>) -> tensor<4x1x2x8xf32> {
// CHECK: linalg.conv_2d_nhwc_hwcf
%0 = linalg.conv_2d_nhwc_hwcf
{dilations = dense<[2, 3]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}
ins(%input, %filter : tensor<4x3x5x3xf32>, tensor<2x2x3x8xf32>)
outs(%init : tensor<4x1x2x8xf32>) -> tensor<4x1x2x8xf32>
return %0 : tensor<4x1x2x8xf32>
}
// -----
// CHECK-LABEL: func @depthwise_conv_2d_nhwc_hwc_tensor
func @depthwise_conv_2d_nhwc_hwc_tensor(%input: tensor<1x1x113x96xf32>, %filter: tensor<1x3x96xf32>, %out: tensor<1x1x56x96xf32>) -> tensor<1x1x56x96xf32> {
// CHECK: linalg.depthwise_conv_1d_nwc_wc
%0 = linalg.depthwise_conv_2d_nhwc_hwc {dilations = dense<1> : vector<2xi64>, strides = dense<2> : vector<2xi64>}
ins(%input, %filter: tensor<1x1x113x96xf32>, tensor<1x3x96xf32>)
outs(%out: tensor<1x1x56x96xf32>) -> tensor<1x1x56x96xf32>
return %0: tensor<1x1x56x96xf32>
}
// -----
// Do not convert convolution ops whose window dimensions are not ones.
// CHECK-LABEL: func @depthwise_conv_2d_nhwc_hwc_tensor
func @depthwise_conv_2d_nhwc_hwc_tensor(%input: tensor<1x113x113x96xf32>, %filter: tensor<3x3x96xf32>, %out: tensor<1x56x56x96xf32>) -> tensor<1x56x56x96xf32> {
// CHECK: linalg.depthwise_conv_2d_nhwc_hwc
%0 = linalg.depthwise_conv_2d_nhwc_hwc {dilations = dense<1> : vector<2xi64>, strides = dense<2> : vector<2xi64>}
ins(%input, %filter: tensor<1x113x113x96xf32>, tensor<3x3x96xf32>)
outs(%out: tensor<1x56x56x96xf32>) -> tensor<1x56x56x96xf32>
return %0: tensor<1x56x56x96xf32>
}