215 lines · plain
1// The following test examples of linalg convolution named ops lowered to linalg.generic and then2// lifted back up to named op.3// NOTE: Most tests in this file use dynamic shapes as the underlying transformations don't modify shapes. There's one exception that's added as a smoke test.4 5// RUN: mlir-opt %s -linalg-generalize-named-ops | mlir-opt --linalg-specialize-generic-ops | FileCheck %s --implicit-check-not=linalg.generic6 7// -----------------------------8// Convolution ops.9// -----------------------------10func.func @conv_1d(%in : tensor<?xf32>, %filter : tensor<?xf32>, %out : tensor<?xf32>) -> tensor<?xf32> {11 %0 = linalg.conv_1d12 ins(%in, %filter : tensor<?xf32>, tensor<?xf32>)13 outs(%out : tensor<?xf32>) -> tensor<?xf32>14 return %0 : tensor<?xf32>15}16// CHECK: @conv_1d17// CHECK: linalg.conv_1d18 19// -----20 21func.func @conv_1d_nwc_wcf(%input: tensor<?x?x?xf32>, %filter: tensor<?x?x?xf32>, %output: tensor<?x?x?xf32>) -> tensor<?x?x?xf32> {22 %0 = linalg.conv_1d_nwc_wcf23 {dilations = dense<3> : tensor<1xi64>, strides = dense<2> : tensor<1xi64>}24 ins (%input, %filter: tensor<?x?x?xf32>, tensor<?x?x?xf32>)25 outs (%output: tensor<?x?x?xf32>) -> tensor<?x?x?xf32>26 return %0 : tensor<?x?x?xf32>27}28// CHECK: @conv_1d_nwc_wcf29// CHECK: linalg.conv_1d_nwc_wcf30// CHECK-SAME: dilations = dense<3> : tensor<1xi64>, strides = dense<2> : tensor<1xi64>31 32// -----33 34func.func @conv_1d_ncw_fcw(%input: tensor<?x?x?xf32>, %filter: tensor<?x?x?xf32>, %output: tensor<?x?x?xf32>) -> tensor<?x?x?xf32> {35 %0 = linalg.conv_1d_ncw_fcw36 {dilations = dense<3> : tensor<1xi64>, strides = dense<2> : tensor<1xi64>}37 ins (%input, %filter: tensor<?x?x?xf32>, tensor<?x?x?xf32>)38 outs (%output: tensor<?x?x?xf32>) -> tensor<?x?x?xf32>39 return %0 : tensor<?x?x?xf32>40}41// CHECK: @conv_1d_ncw_fcw42// CHECK: linalg.conv_1d_ncw_fcw43// CHECK-SAME: dilations = dense<3> : tensor<1xi64>, strides = dense<2> : tensor<1xi64>44 45// -----46 47func.func @conv_2d(%in : tensor<?x?xf32>, %filter : tensor<?x?xf32>, %out : tensor<?x?xf32>) -> tensor<?x?xf32> {48 %0 = linalg.conv_2d49 ins(%in, %filter : tensor<?x?xf32>, tensor<?x?xf32>)50 outs(%out: tensor<?x?xf32>) -> tensor<?x?xf32>51 return %0 : tensor<?x?xf32>52}53// CHECK: @conv_2d54// CHECK: linalg.conv_2d55 56// -----57 58func.func @conv_3d(%in : tensor<?x?x?xf32>, %filter : tensor<?x?x?xf32>, %out : tensor<?x?x?xf32>) -> tensor<?x?x?xf32> {59 %0 = linalg.conv_3d60 ins(%in, %filter : tensor<?x?x?xf32>, tensor<?x?x?xf32>)61 outs(%out : tensor<?x?x?xf32>) -> tensor<?x?x?xf32>62 return %0 : tensor<?x?x?xf32>63}64// CHECK: @conv_3d65// CHECK: linalg.conv_3d66 67// -----68 69// -----------------------------70// Depthwise Convolution ops.71// -----------------------------72func.func @depthwise_conv_1d_ncw_cw(%input: tensor<?x?x?xf32>, %filter: tensor<?x?xf32>, %output: tensor<?x?x?xf32>) -> tensor<?x?x?xf32> {73 %0 = linalg.depthwise_conv_1d_ncw_cw74 {dilations = dense<3> : tensor<1xi64>, strides = dense<2> : tensor<1xi64>}75 ins (%input, %filter: tensor<?x?x?xf32>, tensor<?x?xf32>)76 outs (%output: tensor<?x?x?xf32>) -> tensor<?x?x?xf32>77 return %0 : tensor<?x?x?xf32>78}79// CHECK: @depthwise_conv_1d_ncw_cw80// CHECK: linalg.depthwise_conv_1d_ncw_cw81// CHECK-SAME: dilations = dense<3> : tensor<1xi64>, strides = dense<2> : tensor<1xi64>82 83// -----84 85func.func @depthwise_conv_1d_nwc_wc_static(%input: tensor<1x25x8xi8>, %filter: tensor<3x8xi8>, %output: tensor<1x10x8xi32>) -> tensor<1x10x8xi32> {86 %0 = linalg.depthwise_conv_1d_nwc_wc 87 {dilations = dense<3> : tensor<1xi64>, strides = dense<2> : tensor<1xi64>}88 ins (%input, %filter: tensor<1x25x8xi8>, tensor<3x8xi8>)89 outs (%output: tensor<1x10x8xi32>) -> tensor<1x10x8xi32>90 return %0 : tensor<1x10x8xi32>91}92// CHECK: @depthwise_conv_1d_nwc_wc_static93// CHECK: linalg.depthwise_conv_1d_nwc_wc94// CHECK-SAME: dilations = dense<3> : tensor<1xi64>, strides = dense<2> : tensor<1xi64>95 96// -----97 98func.func @depthwise_conv_1d_nwc_wcm(%input: tensor<?x?x?xf32>, %filter: tensor<?x?x?xf32>, %output: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> {99 %0 = linalg.depthwise_conv_1d_nwc_wcm100 {dilations = dense<1> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>}101 ins (%input, %filter: tensor<?x?x?xf32>, tensor<?x?x?xf32>)102 outs (%output: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32>103 return %0 : tensor<?x?x?x?xf32>104}105// CHECK: @depthwise_conv_1d_nwc_wcm106// CHECK: linalg.depthwise_conv_1d_nwc_wcm107// CHECK-SAME: dilations = dense<1> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>108 109// -----110 111func.func @depthwise_conv_2d_nchw_chw(%input: tensor<?x?x?x?xf16>, %filter: tensor<?x?x?xf16>, %output: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> {112 %0 = linalg.depthwise_conv_2d_nchw_chw113 {dilations = dense<[2,3]> : vector<2xi64>, strides = dense<[4,5]> : vector<2xi64>}114 ins (%input, %filter: tensor<?x?x?x?xf16>, tensor<?x?x?xf16>)115 outs (%output: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32>116 return %0 : tensor<?x?x?x?xf32>117}118// CHECK: @depthwise_conv_2d_nchw_chw119// CHECK: linalg.depthwise_conv_2d_nchw_chw120// CHECK-SAME: dilations = dense<[2, 3]> : tensor<2xi64>, strides = dense<[4, 5]> : tensor<2xi64>121 122// -----123 124func.func @depthwise_conv_3d_ndhwc_dhwcm(%input: tensor<?x?x?x?x?xf32>, %filter: tensor<?x?x?x?x?xf32>, %output: tensor<?x?x?x?x?x?xf32>) -> tensor<?x?x?x?x?x?xf32> {125 %0 = linalg.depthwise_conv_3d_ndhwc_dhwcm126 {dilations = dense<1> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>}127 ins (%input, %filter: tensor<?x?x?x?x?xf32>, tensor<?x?x?x?x?xf32>)128 outs (%output: tensor<?x?x?x?x?x?xf32>) -> tensor<?x?x?x?x?x?xf32>129 return %0 : tensor<?x?x?x?x?x?xf32>130}131// CHECK: @depthwise_conv_3d_ndhwc_dhwcm132// CHECK: linalg.depthwise_conv_3d_ndhwc_dhwcm133// CHECK-SAME: dilations = dense<1> : tensor<3xi64>, strides = dense<1> : tensor<3xi64>134 135// -----136 137// -----------------------------138// Pooling ops.139// -----------------------------140func.func @pooling_nhwc_max(%input: tensor<?x?x?x?xf32>, %filter: tensor<?x?xf32>, %output: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> {141 %0 = linalg.pooling_nhwc_max142 {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}143 ins (%input, %filter: tensor<?x?x?x?xf32>, tensor<?x?xf32>)144 outs (%output: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32>145 return %0 : tensor<?x?x?x?xf32>146}147// CHECK: @pooling_nhwc_max148// CHECK: linalg.pooling_nhwc_max149// CHECK-SAME: dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>150 151// -----152 153func.func @pooling_nhwc_min(%input: tensor<?x?x?x?xf32>, %filter: tensor<?x?xf32>, %output: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> {154 %0 = linalg.pooling_nhwc_min155 {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}156 ins (%input, %filter: tensor<?x?x?x?xf32>, tensor<?x?xf32>)157 outs (%output: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32>158 return %0 : tensor<?x?x?x?xf32>159}160// CHECK: @pooling_nhwc_min161// CHECK: linalg.pooling_nhwc_min162// CHECK-SAME: dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>163 164// -----165 166func.func @pooling_nhwc_sum(%input: tensor<?x?x?x?xf32>, %filter: tensor<?x?xf32>, %output: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> {167 %0 = linalg.pooling_nhwc_sum168 {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}169 ins (%input, %filter: tensor<?x?x?x?xf32>, tensor<?x?xf32>)170 outs (%output: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32>171 return %0 : tensor<?x?x?x?xf32>172}173// CHECK: @pooling_nhwc_sum174// CHECK: linalg.pooling_nhwc_sum175// CHECK-SAME: dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>176 177// -----178 179func.func @pooling_nhwc_max_unsigned(%input: tensor<?x?x?x?xi8>, %filter: tensor<?x?xi8>, %output: tensor<?x?x?x?xi32>) -> tensor<?x?x?x?xi32> {180 %0 = linalg.pooling_nhwc_max_unsigned181 {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}182 ins (%input, %filter: tensor<?x?x?x?xi8>, tensor<?x?xi8>)183 outs (%output: tensor<?x?x?x?xi32>) -> tensor<?x?x?x?xi32>184 return %0 : tensor<?x?x?x?xi32>185}186// CHECK: @pooling_nhwc_max_unsigned187// CHECK: linalg.pooling_nhwc_max_unsigned188// CHECK-SAME: dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>189 190// -----191 192func.func @pooling_nhwc_min_unsigned_integer(%input: tensor<?x?x?x?xi32>, %filter: tensor<?x?xi32>, %output: tensor<?x?x?x?xi32>) -> tensor<?x?x?x?xi32> {193 %0 = linalg.pooling_nhwc_min_unsigned194 {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}195 ins (%input, %filter: tensor<?x?x?x?xi32>, tensor<?x?xi32>)196 outs (%output: tensor<?x?x?x?xi32>) -> tensor<?x?x?x?xi32>197 return %0 : tensor<?x?x?x?xi32>198}199// CHECK: @pooling_nhwc_min_unsigned_integer200// CHECK: linalg.pooling_nhwc_min_unsigned201// CHECK-SAME: dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>202 203// -----204 205func.func @pooling_nhwc_min_unsigned_float(%input: tensor<?x?x?x?xf32>, %filter: tensor<?x?xf32>, %output: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> {206 %0 = linalg.pooling_nhwc_min_unsigned207 {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>}208 ins (%input, %filter: tensor<?x?x?x?xf32>, tensor<?x?xf32>)209 outs (%output: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32>210 return %0 : tensor<?x?x?x?xf32>211}212// CHECK: @pooling_nhwc_min_unsigned_float213// CHECK: linalg.pooling_nhwc_min214// CHECK-SAME: dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>215