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
95 lines
5.5 KiB
MLIR
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>
|
|
}
|