1// RUN: mlir-opt --test-transform-dialect-interpreter --split-input-file %s | FileCheck %s
2
3// CHECK-LABEL: @conv_2d_nhwc_hwcf
4// CHECK-SAME: %[[ARG0:.+]]: tensor<?x1x?x?xf32>,
5// CHECK-SAME: %[[ARG1:.+]]: tensor<1x?x?x?xf32>
6// CHECK-SAME: %[[ARG2:.+]]: tensor<?x1x?x?xf32>
7func.func @conv_2d_nhwc_hwcf(%input: tensor<?x1x?x?xf32>, %filter: tensor<1x?x?x?xf32>, %init: tensor<?x1x?x?xf32>) -> tensor<?x1x?x?xf32> {
8  // CHECK: %[[SLICE0:.+]] = tensor.extract_slice %[[ARG0]]
9  // CHECK: %[[SLICE1:.+]] = tensor.extract_slice %[[ARG1]]
10  // CHECK: %[[SLICE2:.+]] = tensor.extract_slice %[[ARG2]]
11  // CHECK: %[[SLICERES:.+]] = linalg.conv_1d_nwc_wcf
12  // CHECK: %[[RES:.+]] = tensor.insert_slice %[[SLICERES]] into %[[ARG2]]
13  %0 = linalg.conv_2d_nhwc_hwcf {dilations = dense<1> : tensor<2xi64>,
14                                 strides = dense<1> : tensor<2xi64>}
15     ins (%input, %filter: tensor<?x1x?x?xf32>, tensor<1x?x?x?xf32>)
16    outs (%init: tensor<?x1x?x?xf32>) -> tensor<?x1x?x?xf32>
17  // CHECK: return %[[RES]]
18  return %0 : tensor<?x1x?x?xf32>
19}
20
21// CHECK-LABEL: @depthwise_conv_2d_nhwc_hwc
22// CHECK-SAME: %[[ARG0:.+]]: tensor<1x1x113x96xf32>
23// CHECK-SAME: %[[ARG1:.+]]: tensor<1x3x96xf32>
24func.func @depthwise_conv_2d_nhwc_hwc(%input: tensor<1x1x113x96xf32>, %filter: tensor<1x3x96xf32>) -> tensor<1x1x56x96xf32> {
25  // CHECK: %[[RES:.+]] = linalg.init_tensor
26  %init = linalg.init_tensor [1, 1, 56, 96] : tensor<1x1x56x96xf32>
27  // CHECK: %[[SLICE0:.+]] = tensor.extract_slice %[[ARG0]]
28  // CHECK: %[[SLICE1:.+]] = tensor.extract_slice %[[ARG1]]
29  // CHECK: %[[SLICERES:.+]] = tensor.extract_slice %[[RES]]
30  // CHECK: %[[OPRES:.+]] = linalg.depthwise_conv_1d_nwc_wc
31  // CHECK-SAME: ins(%[[SLICE0]], %[[SLICE1]]
32  // CHECK-SAME: outs(%[[SLICERES]]
33  // CHECK: %[[INSERTED:.+]] = tensor.insert_slice %[[OPRES]] into %[[RES]]
34  %0 = linalg.depthwise_conv_2d_nhwc_hwc {dilations = dense<1> : vector<2xi64>, strides = dense<2> : vector<2xi64>}
35         ins(%input, %filter: tensor<1x1x113x96xf32>, tensor<1x3x96xf32>)
36         outs(%init: tensor<1x1x56x96xf32>) -> tensor<1x1x56x96xf32>
37  // CHECK: %[[INSERTED]]
38  return %0: tensor<1x1x56x96xf32>
39}
40
41transform.with_pdl_patterns {
42^bb0(%arg0: !pdl.operation):
43  transform.sequence %arg0 {
44  ^bb1(%arg1: !pdl.operation):
45    %0 = transform.structured.match interface{LinalgOp} in %arg1
46    %1 = transform.structured.decompose %0
47  }
48}
49