559 lines · plain
1// RUN: mlir-opt %s -transform-interpreter -split-input-file -verify-diagnostics | FileCheck %s2 3// Check that the im2col patterns are properly connected with the4// transform dialect.5 6// Non static shapes are not supported.7// Check that we emit an error.8// TODO: Hook up the rewriter errors in transform dialect.9func.func @conv_non_static(%arg0: tensor<?x?x?x?xf32>, %arg1: tensor<3x3x4x16xf32>, %arg2: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32> {10 // expected-note@below {{when applied to this op}}11 %0 = linalg.conv_2d_nhwc_hwcf12 {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64> }13 ins(%arg0, %arg1: tensor<?x?x?x?xf32>, tensor<3x3x4x16xf32>)14 outs(%arg2: tensor<?x?x?x?xf32>) -> tensor<?x?x?x?xf32>15 return %0 : tensor<?x?x?x?xf32>16}17 18module attributes {transform.with_named_sequence} {19 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {20 %0 = transform.structured.match ops{["linalg.conv_2d_nhwc_hwcf"]} in %arg1 : (!transform.any_op) -> !transform.any_op21 // expected-error@below {{failed to apply}}22 %1:2 = transform.structured.convert_conv2d_to_img2col %0 : (!transform.any_op) -> (!transform.any_op, !transform.any_op)23 transform.yield24 }25}26 27// -----28 29// Memref semantics is not supported.30// Check that we emit an error.31func.func @negative_conv_memref(%arg0: memref<1x16x16x4xf32>, %arg1: memref<16x3x3x4xf32>, %arg2: memref<1x14x14x16xf32>) {32 // expected-note@below {{when applied to this op}}33 linalg.conv_2d_nhwc_fhwc {dilations = dense<1> : memref<2xi64>, strides = dense<1> : memref<2xi64> }34 ins(%arg0, %arg1: memref<1x16x16x4xf32>, memref<16x3x3x4xf32>) outs(%arg2: memref<1x14x14x16xf32>)35 return36}37 38module attributes {transform.with_named_sequence} {39 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {40 %0 = transform.structured.match ops{["linalg.conv_2d_nhwc_fhwc"]} in %arg1 : (!transform.any_op) -> !transform.any_op41 // expected-error@below {{failed to apply}}42 %img2col_tensor_producer, %transformed = transform.structured.convert_conv2d_to_img2col %0 : (!transform.any_op) -> (!transform.any_op, !transform.any_op)43 transform.yield44 }45}46 47// -----48 49// Check that we get the proper handles for the img2col tensor producer50// and the final instruction.51 52// CHECK: IR printer: tensor_producer53// CHECK-NEXT: %[[COL_TENSOR:.+]] = linalg.generic54// CHECK-SAME: affine_map<(d0, d1, d2) -> (d0, d1, d2)>]55// CHECK: ^bb0(%[[OUT_DATA:.+]]: f32)56 57// CHECK: IR printer: transformed58// CHECK: tensor.expand_shape %{{[^ ]*}} {{\[}}[0], [1, 2], [3]] output_shape [1, 14, 14, 16] : tensor<1x196x16xf32> into tensor<1x14x14x16xf32>59 60// Im2col maps61// CHECK-DAG: #[[MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1 floordiv 14 + d2 floordiv 12, d1 mod 14 + (d2 mod 12) floordiv 4, d2 mod 4)>62// CHECK-DAG: #[[MAPI2C:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>63// Matmul maps64// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>65// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d3, d2)>66// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>67 68// CHECK: @conv_1643313669// CHECK-SAME: %[[INPUT:.+]]: tensor<1x16x16x4xf32>70// CHECK-SAME: %[[FILTER:.+]]: tensor<3x3x4x16xf32>71// CHECK-SAME: %[[OUTPUT:.+]]: tensor<1x14x14x16xf32>72// CHECK-DAG: %[[COLLAPSED_FILTER:.+]] = tensor.collapse_shape %[[FILTER]] {{\[}}[0, 1, 2], [3]] : tensor<3x3x4x16xf32> into tensor<36x16xf32>73// CHECK-DAG: %[[COLLAPSED_OUT:.+]] = tensor.collapse_shape %[[OUTPUT]] {{\[}}[0], [1, 2], [3]] : tensor<1x14x14x16xf32> into tensor<1x196x16xf32>74// CHECK: %[[INIT_COL_TENSOR:.+]] = tensor.empty() : tensor<1x196x36xf32>75 76// CHECK: %[[COL_TENSOR:.+]] = linalg.generic77// CHECK-SAME: indexing_maps = [#[[MAP]], #[[MAPI2C]]]78// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"]79// CHECK-SAME: ins(%[[INPUT]] : tensor<1x16x16x4xf32>)80// CHECK-SAME: outs(%[[INIT_COL_TENSOR]] : tensor<1x196x36xf32>)81// CHECK: ^bb0(%[[IN:.+]]: f32, %out: f32):82// CHECK: linalg.yield %[[IN]] : f3283// CHECK: } -> tensor<1x196x36xf32>84 85// CHECK: %[[MATMUL_RESULT:.+]] = linalg.generic86// CHECK-SAME: #[[MAP1]]87// CHECK-SAME: #[[MAP2]]88// CHECK-SAME: #[[MAP3]]89// CHECK-SAME: ins(%[[COL_TENSOR]], %[[COLLAPSED_FILTER]] : tensor<1x196x36xf32>, tensor<36x16xf32>)90// CHECK-SAME: outs(%[[COLLAPSED_OUT]] : tensor<1x196x16xf32>)91// CHECK: ^bb0(%[[ARG0:.+]]: f32, %[[ARG1:.+]]: f32, %[[ARG2:.+]]: f32)92// CHECK: %[[MUL:.+]] = arith.mulf %[[ARG0]], %[[ARG1]] : f3293// CHECK: %[[ADD:.+]] = arith.addf %[[MUL]], %[[ARG2]] : f3294// CHECK: linalg.yield %[[ADD]] : f3295// CHECK: } -> tensor<1x196x16xf32>96// CHECK: %[[RESULT:.+]] = tensor.expand_shape %[[MATMUL_RESULT]] {{\[}}[0], [1, 2], [3]] output_shape [1, 14, 14, 16] : tensor<1x196x16xf32> into tensor<1x14x14x16xf32>97// CHECK: return %[[RESULT]]98 99func.func @conv_16433136(%arg0: tensor<1x16x16x4xf32>, %arg1: tensor<3x3x4x16xf32>, %arg2: tensor<1x14x14x16xf32>) -> tensor<1x14x14x16xf32> {100 %0 = linalg.conv_2d_nhwc_hwcf101 {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64> }102 ins(%arg0, %arg1: tensor<1x16x16x4xf32>, tensor<3x3x4x16xf32>)103 outs(%arg2: tensor<1x14x14x16xf32>) -> tensor<1x14x14x16xf32>104 return %0 : tensor<1x14x14x16xf32>105}106 107module attributes {transform.with_named_sequence} {108 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {109 %0 = transform.structured.match ops{["linalg.conv_2d_nhwc_hwcf"]} in %arg1 : (!transform.any_op) -> !transform.any_op110 %img2col_tensor_producer, %transformed = transform.structured.convert_conv2d_to_img2col %0 : (!transform.any_op) -> (!transform.any_op, !transform.any_op)111 transform.print %img2col_tensor_producer {name = "tensor_producer"}: !transform.any_op112 transform.print %transformed {name = "transformed"}: !transform.any_op113 transform.yield114 }115}116 117// -----118 119// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3, d1)>120// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>121// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2) -> (d1, d2, d0)>122// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>123// CHECK-DAG: #[[MAP4:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2 + d4, d3 + d5)>124// CHECK-DAG: #[[MAP5:.+]] = affine_map<(d0, d1, d2, d3, d4, d5) -> (d0, d1, d2, d3, d4, d5)>125// CHECK-DAG: #[[MAP6:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d1, d2)>126// CHECK: @depthwise_conv_hwc_114x16x3127// CHECK-SAME: %[[INPUT:.+]]: tensor<1x114x114x16xf32>128// CHECK-SAME: %[[FILTER:.+]]: tensor<3x3x16xf32>129// CHECK-SAME: %[[OUTPUT:.+]]: tensor<1x112x112x16xf32>130// CHECK: %[[INPUT_T_INIT:.+]] = tensor.empty() : tensor<1x16x114x114xf32>131// CHECK: %[[INPUT_T:.+]] = linalg.generic132// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]]]133// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel"]134// CHECK-SAME: ins(%[[INPUT]] : tensor<1x114x114x16xf32>) outs(%[[INPUT_T_INIT]] : tensor<1x16x114x114xf32>) {135// CHECK-NEXT: ^bb0(%[[ARG3:.+]]: f32, %[[ARG4:.+]]: f32):136// CHECK-NEXT: linalg.yield %[[ARG3]] : f32137// CHECK-NEXT: } -> tensor<1x16x114x114xf32>138// CHECK: %[[FILTER_T_INIT:.+]] = tensor.empty() : tensor<16x3x3xf32>139// CHECK: %[[FILTER_T:.+]] = linalg.generic140// CHECK-SAME: indexing_maps = [#[[MAP2]], #[[MAP3]]141// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"]142// CHECK-SAME: ins(%[[FILTER]] : tensor<3x3x16xf32>) outs(%[[FILTER_T_INIT]] : tensor<16x3x3xf32>) {143// CHECK-NEXT: ^bb0(%{{.*}}: f32, %{{.*}}: f32):144// CHECK: linalg.yield145// CHECK: } -> tensor<16x3x3xf32>146// CHECK: %[[INIT_OUTPUT_TENSOR:.+]] = tensor.empty() : tensor<1x16x112x112xf32>147// CHECK: %[[OUTPUT_T:.+]] = linalg.generic148// CHECK-SAME: indexing_maps = [#[[MAP0]], #[[MAP1]]]149// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel"]150// CHECK-SAME: ins(%[[OUTPUT]] : tensor<1x112x112x16xf32>) outs(%[[INIT_OUTPUT_TENSOR]] : tensor<1x16x112x112xf32>) {151// CHECK-NEXT: ^bb0(%{{.*}}: f32, %{{.*}}: f32):152// CHECK-NEXT: linalg.yield153// CHECK-NEXT: } -> tensor<1x16x112x112xf32>154// CHECK: %[[INIT_COL_TENSOR:.+]] = tensor.empty() : tensor<1x16x112x112x3x3xf32>155// CHECK: %[[COL_TENSOR:.+]] = linalg.generic156// CHECK-SAME: indexing_maps = [#[[MAP4]], #[[MAP5]]]157// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel", "parallel", "parallel"]158// CHECK-SAME: ins(%[[INPUT_T]] : tensor<1x16x114x114xf32>) outs(%[[INIT_COL_TENSOR]] : tensor<1x16x112x112x3x3xf32>) {159// CHECK-NEXT: ^bb0(%{{.*}}: f32, %{{.*}}: f32):160// CHECK-NEXT: linalg.yield161// CHECK-NEXT: } -> tensor<1x16x112x112x3x3xf32>162// CHECK: %[[COL_TENSOR_R:.+]] = tensor.collapse_shape %[[COL_TENSOR]]163// CHECK-SAME: tensor<1x16x112x112x3x3xf32> into tensor<16x12544x9xf32>164// CHECK: %[[FILTER_T_R:.+]] = tensor.collapse_shape %[[FILTER_T]]165// CHECK-SAME: tensor<16x3x3xf32> into tensor<16x9xf32>166// CHECK: %[[OUTPUT_T_R:.+]] = tensor.collapse_shape %[[OUTPUT_T]]167// CHECK-SAME: tensor<1x16x112x112xf32> into tensor<16x12544xf32>168// CHECK: %[[BMV_RESULT:.+]] = linalg.batch_matvec ins(%[[COL_TENSOR_R]], %[[FILTER_T_R]] : tensor<16x12544x9xf32>, tensor<16x9xf32>) outs(%[[OUTPUT_T_R]] : tensor<16x12544xf32>) -> tensor<16x12544xf32>169// CHECK: %[[RESULT_R:.+]] = tensor.expand_shape %[[BMV_RESULT]]170// CHECK-SAME: tensor<16x12544xf32> into tensor<1x16x112x112xf32>171// CHECK: %[[RESULT_INIT:.+]] = tensor.empty() : tensor<1x112x112x16xf32>172// CHECK: %[[RESULT:.+]] = linalg.generic173// CHECK-SAME: indexing_maps = [#[[MAP6]], #[[MAP1]]]174// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "parallel"]175// CHECK-SAME: ins(%[[RESULT_R]] : tensor<1x16x112x112xf32>) outs(%[[RESULT_INIT]] : tensor<1x112x112x16xf32>) {176// CHECK-NEXT: ^bb0(%{{.*}}: f32, %{{.*}}: f32):177// CHECK-NEXT: linalg.yield178// CHECK-NEXT: } -> tensor<1x112x112x16xf32>179// CHECK: return %[[RESULT]] : tensor<1x112x112x16xf32>180func.func @depthwise_conv_hwc_114x16x3(%input: tensor<1x114x114x16xf32>, %filter: tensor<3x3x16xf32>, %output: tensor<1x112x112x16xf32>) -> tensor<1x112x112x16xf32> {181 %0 = linalg.depthwise_conv_2d_nhwc_hwc {182 dilations = dense<1> : tensor<2xi64>,183 strides = dense<1> : tensor<2xi64>184 } ins(%input, %filter : tensor<1x114x114x16xf32>, tensor<3x3x16xf32>) outs(%output : tensor<1x112x112x16xf32>) -> tensor<1x112x112x16xf32>185 return %0 : tensor<1x112x112x16xf32>186}187 188module attributes {transform.with_named_sequence} {189 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {190 %0 = transform.structured.match ops{["linalg.depthwise_conv_2d_nhwc_hwc"]} in %arg1 : (!transform.any_op) -> !transform.any_op191 %1:2 = transform.structured.convert_conv2d_to_img2col %0 : (!transform.any_op) -> (!transform.any_op, !transform.any_op)192 transform.yield193 }194}195 196// -----197 198// Im2col maps199// CHECK-DAG: #[[MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1 floordiv 14 + d2 floordiv 12, d1 mod 14 + (d2 mod 12) floordiv 4, d2 mod 4)>200// CHECK-DAG: #[[MAPI2C:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>201 202// CHECK-DAG: #[[LHSMAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>203// CHECK-DAG: #[[RHSMAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d3, d2)>204// CHECK-DAG: #[[RESMAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>205 206// CHECK: func.func @batch_nhwc_conv207// CHECK-SAME: (%[[INPUT:.+]]: tensor<8x16x16x4xf32>, %[[FILTER:.+]]: tensor<3x3x4x16xf32>, %[[INIT:.+]]: tensor<8x14x14x16xf32>)208// CHECK-DAG: %[[CS_FILTER:.+]] = tensor.collapse_shape %[[FILTER]] {{\[}}[0, 1, 2], [3]] : tensor<3x3x4x16xf32> into tensor<36x16xf32>209// CHECK-DAG: %[[CS_RESULT:.+]] = tensor.collapse_shape %[[INIT]] {{\[}}[0], [1, 2], [3]] : tensor<8x14x14x16xf32> into tensor<8x196x16xf32>210// CHECK: %[[IT:.+]] = tensor.empty() : tensor<8x196x36xf32>211// CHECK: %[[IMG2COL:.+]] = linalg.generic212// CHECK-SAME: indexing_maps = [#[[MAP]], #[[MAPI2C]]]213// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"]214// CHECK-SAME: ins(%[[INPUT]] : tensor<8x16x16x4xf32>)215// CHECK-SAME: outs(%[[IT]] : tensor<8x196x36xf32>)216// CHECK: ^bb0(%[[IN:.+]]: f32, %out: f32):217// CHECK: linalg.yield %[[IN]] : f32218// CHECK: } -> tensor<8x196x36xf32>219// CHECK: %[[MATMUL:.+]] = linalg.generic220// CHECK-SAME: indexing_maps = [#[[LHSMAP]], #[[RHSMAP]], #[[RESMAP]]],221// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel", "reduction"]222// CHECK-SAME: ins(%[[IMG2COL]], %[[CS_FILTER]] : tensor<8x196x36xf32>, tensor<36x16xf32>)223// CHECK-SAME: outs(%[[CS_RESULT]] : tensor<8x196x16xf32>)224// CHECK: ^bb0(%[[ARG0:.+]]: f32, %[[ARG1:.+]]: f32, %[[ARG2:.+]]: f32):225// CHECK: %[[MUL:.+]] = arith.mulf %[[ARG0]], %[[ARG1]] : f32226// CHECK: %[[ADD:.+]] = arith.addf %[[MUL]], %[[ARG2]] : f32227// CHECK: linalg.yield %[[ADD]] : f32228// CHECK: } -> tensor<8x196x16xf32>229// CHECK: %[[CS_FINAL:.+]] = tensor.expand_shape %[[MATMUL]] {{\[}}[0], [1, 2], [3]] output_shape [8, 14, 14, 16] : tensor<8x196x16xf32> into tensor<8x14x14x16xf32>230// CHECK: return %[[CS_FINAL]]231func.func @batch_nhwc_conv(%arg0: tensor<8x16x16x4xf32>, %arg1: tensor<3x3x4x16xf32>, %arg2: tensor<8x14x14x16xf32>) -> tensor<8x14x14x16xf32> {232 %0 = linalg.conv_2d_nhwc_hwcf233 {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64> }234 ins(%arg0, %arg1: tensor<8x16x16x4xf32>, tensor<3x3x4x16xf32>)235 outs(%arg2: tensor<8x14x14x16xf32>) -> tensor<8x14x14x16xf32>236 return %0 : tensor<8x14x14x16xf32>237}238 239module attributes {transform.with_named_sequence} {240 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {241 %0 = transform.structured.match ops{["linalg.conv_2d_nhwc_hwcf"]} in %arg1 : (!transform.any_op) -> !transform.any_op242 %1:2 = transform.structured.convert_conv2d_to_img2col %0 : (!transform.any_op) -> (!transform.any_op, !transform.any_op)243 transform.yield244 }245}246 247// -----248 249// Im2col maps250// CHECK-DAG: #[[MAP:.+]] = affine_map<(d0, d1, d2) -> (d0, d1 floordiv 9, d2 floordiv 14 + (d1 mod 9) floordiv 3, d2 mod 14 + d1 mod 3)>251// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>252 253// CHECK-DAG: #[[LHSMAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d1, d3)>254// CHECK-DAG: #[[RHSMAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d3, d2)>255// CHECK-DAG: #[[RESMAP:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>256 257// CHECK: func.func @batch_nchw_conv258// CHECK-SAME: (%[[INPUT:.+]]: tensor<8x4x16x16xf32>, %[[FILTER:.+]]: tensor<16x4x3x3xf32>, %[[INIT:.+]]: tensor<8x16x14x14xf32>)259// CHECK-DAG: %[[CS_FILTER:.+]] = tensor.collapse_shape %[[FILTER]] {{\[}}[0], [1, 2, 3]] : tensor<16x4x3x3xf32> into tensor<16x36xf32>260// CHECK-DAG: %[[CS_RESULT:.+]] = tensor.collapse_shape %[[INIT]] {{\[}}[0], [1], [2, 3]] : tensor<8x16x14x14xf32> into tensor<8x16x196xf32>261// CHECK: %[[IT:.+]] = tensor.empty() : tensor<8x36x196xf32>262// CHECK: %[[IMG2COL:.+]] = linalg.generic263// CHECK-SAME: indexing_maps = [#[[MAP]], #[[MAP1]]]264// CHECK-SAME: iterator_types = ["parallel", "parallel", "parallel"]265// CHECK-SAME: ins(%[[INPUT]] : tensor<8x4x16x16xf32>)266// CHECK-SAME: outs(%[[IT]] : tensor<8x36x196xf32>)267// CHECK: ^bb0(%[[IN:.+]]: f32, %out: f32):268// CHECK: linalg.yield %[[IN]] : f32269// CHECK: } -> tensor<8x16x196xf32>270// CHECK: %[[CS_FINAL:.+]] = tensor.expand_shape %[[MATMUL]] {{\[}}[0], [1], [2, 3]] output_shape [8, 16, 14, 14] : tensor<8x16x196xf32> into tensor<8x16x14x14xf32>271// CHECK: return %[[CS_FINAL]]272func.func @batch_nchw_conv(%arg0: tensor<8x4x16x16xf32>, %arg1: tensor<16x4x3x3xf32>, %arg2: tensor<8x16x14x14xf32>) -> tensor<8x16x14x14xf32> {273 %0 = linalg.conv_2d_nchw_fchw274 {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64> }275 ins(%arg0, %arg1: tensor<8x4x16x16xf32>, tensor<16x4x3x3xf32>)276 outs(%arg2: tensor<8x16x14x14xf32>) -> tensor<8x16x14x14xf32>277 return %0 : tensor<8x16x14x14xf32>278}279 280module attributes {transform.with_named_sequence} {281 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {282 %0 = transform.structured.match ops{["linalg.conv_2d_nchw_fchw"]} in %arg1 : (!transform.any_op) -> !transform.any_op283 %1:2 = transform.structured.convert_conv2d_to_img2col %0 : (!transform.any_op) -> (!transform.any_op, !transform.any_op)284 transform.yield285 }286}287 288// -----289 290// Check that the encoding on the filter (weights) tensor is propagated when applying the transform. 291 292// CHECK: func.func @batch_nchw_conv_with_filter_encoding(%[[INPUT:.+]]: tensor<8x4x16x16xf32>, %[[FILTER:.*]]: tensor<16x4x3x3xf32, 42 : i32>, %[[OUTPUT:.*]]: tensor<8x16x14x14xf32>)293// CHECK-DAG: %[[COLLAPSED_FILTER:.+]] = tensor.collapse_shape %[[FILTER]]294 // CHECK-SAME{LITERAL}: [[0], [1, 2, 3]] : tensor<16x4x3x3xf32, 42 : i32> into tensor<16x36xf32, 42 : i32>295// CHECK: %[[COL_TENSOR:.+]] = linalg.generic {{.*}} ins(%[[INPUT]] : tensor<8x4x16x16xf32>)296// CHECK: %[[MATMUL_RESULT:.+]] = linalg.generic {{.*}} ins(%[[COLLAPSED_FILTER]], %[[COL_TENSOR]] : tensor<16x36xf32, 42 : i32>, tensor<8x36x196xf32>)297func.func @batch_nchw_conv_with_filter_encoding(%arg0: tensor<8x4x16x16xf32>, %arg1: tensor<16x4x3x3xf32, 42 : i32>, %arg2: tensor<8x16x14x14xf32>) -> tensor<8x16x14x14xf32> {298 %0 = linalg.conv_2d_nchw_fchw299 {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64> }300 ins(%arg0, %arg1: tensor<8x4x16x16xf32>, tensor<16x4x3x3xf32, 42 : i32>)301 outs(%arg2: tensor<8x16x14x14xf32>) -> tensor<8x16x14x14xf32>302 return %0 : tensor<8x16x14x14xf32>303}304 305module attributes {transform.with_named_sequence} {306 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {307 %0 = transform.structured.match ops{["linalg.conv_2d_nchw_fchw"]} in %arg1 : (!transform.any_op) -> !transform.any_op308 %1:2 = transform.structured.convert_conv2d_to_img2col %0 : (!transform.any_op) -> (!transform.any_op, !transform.any_op)309 transform.yield310 }311}312 313// -----314 315// CHECK: IR printer: tensor_producer316// CHECK-NEXT: %[[COL_TENSOR:.+]] = linalg.generic317// CHECK-SAME: affine_map<(d0, d1, d2) -> (d0, d1 floordiv 14 + d2 floordiv 12, d1 mod 14 + (d2 mod 12) floordiv 4, d2 mod 4)>318// CHECK-SAME: affine_map<(d0, d1, d2) -> (d0, d1, d2)>]319// CHECK: ^bb0(%[[IN_DATA:.+]]: f32, %[[OUT_DATA:.+]]: f32)320// CHECK: linalg.yield %[[IN_DATA]] : f32321 322// CHECK: IR printer: transformed323// CHECK: tensor.expand_shape %{{[^ ]*}} {{\[}}[0], [1, 2], [3]] output_shape [1, 14, 14, 16] : tensor<1x196x16xf32> into tensor<1x14x14x16xf32>324 325// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2) -> (d0, d1 floordiv 14 + d2 floordiv 12, d1 mod 14 + (d2 mod 12) floordiv 4, d2 mod 4)>326// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>327// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>328// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3) -> (d2, d3)>329// CHECK-DAG: #[[MAP4:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>330// CHECK: @conv_2d_nhwc_fhwc331// CHECK-SAME: %[[INPUT:.+]]: tensor<1x16x16x4xf32>332// CHECK-SAME: %[[FILTER:.+]]: tensor<16x3x3x4xf32>333// CHECK-SAME: %[[OUTPUT:.+]]: tensor<1x14x14x16xf32>334// CHECK-DAG: %[[COLLAPSED_FILTER:.+]] = tensor.collapse_shape %[[FILTER]] {{\[}}[0], [1, 2, 3]] : tensor<16x3x3x4xf32> into tensor<16x36xf32>335// CHECK-DAG: %[[COLLAPSED_OUT:.+]] = tensor.collapse_shape %[[OUTPUT]] {{\[}}[0], [1, 2], [3]] : tensor<1x14x14x16xf32> into tensor<1x196x16xf32>336// CHECK: %[[INIT_COL_TENSOR:.+]] = tensor.empty() : tensor<1x196x36xf32>337// CHECK: %[[COL_TENSOR:.+]] = linalg.generic338// CHECK-SAME: [#[[MAP0]], #[[MAP1]]], {{.*}} ins(%[[INPUT]] : tensor<1x16x16x4xf32>) outs(%[[INIT_COL_TENSOR]] : tensor<1x196x36xf32>)339// CHECK: ^bb0(%[[OUT_DATA:.+]]: f32)340// CHECK: linalg.yield %{{.+}} : f32341// CHECK: %[[MATMUL_RESULT:.+]] = linalg.generic342// CHECK-SAME: #[[MAP2]]343// CHECK-SAME: #[[MAP3]]344// CHECK-SAME: #[[MAP4]]345// CHECK-SAME: ins(%[[COL_TENSOR]], %[[COLLAPSED_FILTER]] : tensor<1x196x36xf32>, tensor<16x36xf32>)346// CHECK-SAME: outs(%[[COLLAPSED_OUT]] : tensor<1x196x16xf32>)347// CHECK: ^bb0(%[[ARG0:.+]]: f32, %[[ARG1:.+]]: f32, %[[ARG2:.+]]: f32)348// CHECK: %[[MUL:.+]] = arith.mulf %[[ARG0]], %[[ARG1]] : f32349// CHECK: %[[ADD:.+]] = arith.addf %[[MUL]], %[[ARG2]] : f32350// CHECK: linalg.yield %[[ADD]] : f32351// CHECK: } -> tensor<1x196x16xf32>352// CHECK: %[[RESULT:.+]] = tensor.expand_shape %[[MATMUL_RESULT]] {{\[}}[0], [1, 2], [3]] output_shape [1, 14, 14, 16] : tensor<1x196x16xf32> into tensor<1x14x14x16xf32>353// CHECK: return %[[RESULT]]354 355func.func @conv_2d_nhwc_fhwc(%arg0: tensor<1x16x16x4xf32>, %arg1: tensor<16x3x3x4xf32>, %arg2: tensor<1x14x14x16xf32>) -> tensor<1x14x14x16xf32> {356 %0 = linalg.conv_2d_nhwc_fhwc357 {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64> }358 ins(%arg0, %arg1: tensor<1x16x16x4xf32>, tensor<16x3x3x4xf32>)359 outs(%arg2: tensor<1x14x14x16xf32>) -> tensor<1x14x14x16xf32>360 return %0 : tensor<1x14x14x16xf32>361}362 363module attributes {transform.with_named_sequence} {364 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {365 %0 = transform.structured.match ops{["linalg.conv_2d_nhwc_fhwc"]} in %arg1 : (!transform.any_op) -> !transform.any_op366 %img2col_tensor_producer, %transformed = transform.structured.convert_conv2d_to_img2col %0 : (!transform.any_op) -> (!transform.any_op, !transform.any_op)367 transform.print %img2col_tensor_producer {name = "tensor_producer"}: !transform.any_op368 transform.print %transformed {name = "transformed"}: !transform.any_op369 transform.yield370 }371}372 373// -----374 375// Check that the encoding on the filter (weights) tensor is propagated when applying the transform. 376 377// CHECK: func.func @conv_2d_nhwc_fhwc_with_filter_encoding(%[[INPUT:.+]]: tensor<1x16x16x4xf32>, %[[FILTER:.*]]: tensor<16x3x3x4xf32, 42 : i32>, %[[OUTPUT:.*]]: tensor<1x14x14x16xf32>)378// CHECK-DAG: %[[COLLAPSED_FILTER:.+]] = tensor.collapse_shape %[[FILTER]]379 // CHECK-SAME{LITERAL}: [[0], [1, 2, 3]] : tensor<16x3x3x4xf32, 42 : i32> into tensor<16x36xf32, 42 : i32>380// CHECK: %[[COL_TENSOR:.+]] = linalg.generic {{.*}} ins(%[[INPUT]] : tensor<1x16x16x4xf32>)381// CHECK: %[[MATMUL_RESULT:.+]] = linalg.generic {{.*}} ins(%[[COL_TENSOR]], %[[COLLAPSED_FILTER]] : tensor<1x196x36xf32>, tensor<16x36xf32, 42 : i32>)382func.func @conv_2d_nhwc_fhwc_with_filter_encoding(%input: tensor<1x16x16x4xf32>, %filter: tensor<16x3x3x4xf32, 42 : i32>, %out: tensor<1x14x14x16xf32>) -> tensor<1x14x14x16xf32> {383 %0 = linalg.conv_2d_nhwc_fhwc384 { dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64> }385 ins(%input, %filter: tensor<1x16x16x4xf32>, tensor<16x3x3x4xf32, 42 : i32>)386 outs(%out: tensor<1x14x14x16xf32>) -> tensor<1x14x14x16xf32>387 return %0 : tensor<1x14x14x16xf32>388}389 390module attributes {transform.with_named_sequence} {391 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {392 %0 = transform.structured.match ops{["linalg.conv_2d_nhwc_fhwc"]} in %arg1 : (!transform.any_op) -> !transform.any_op393 %1:2 = transform.structured.convert_conv2d_to_img2col %0 : (!transform.any_op) -> (!transform.any_op, !transform.any_op)394 transform.yield395 }396}397 398// -----399 400// Check for signed extend when the input type is smaller than the accumulator type.401 402// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>403// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>404// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d3, d2)>405// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>406// CHECK: @conv_integer_extend407// CHECK: %[[MATMUL_RESULT:.+]] = linalg.generic {indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]]408// CHECK-SAME: ins(%{{.*}}, %{{.*}} : tensor<1x196x36xi8>, tensor<36x16xi8>)409// CHECK-SAME: outs(%[[COLLAPSED_OUT]] : tensor<1x196x16xi32>)410// CHECK: ^bb0(%[[ARG0:.+]]: i8, %[[ARG1:.+]]: i8, %[[ARG2:.+]]: i32)411// CHECK: %[[EXT0:.+]] = arith.extsi %[[ARG0]] : i8 to i32412// CHECK: %[[EXT1:.+]] = arith.extsi %[[ARG1]] : i8 to i32413// CHECK: %[[MUL:.+]] = arith.muli %[[EXT0]], %[[EXT1]] : i32414// CHECK: %[[ADD:.+]] = arith.addi %[[MUL]], %[[ARG2]] : i32415// CHECK: linalg.yield %[[ADD]] : i32416// CHECK: } -> tensor<1x196x16xi32>417// CHECK: %[[RESULT:.+]] = tensor.expand_shape %[[MATMUL_RESULT]] {{\[}}[0], [1, 2], [3]] output_shape [1, 14, 14, 16] : tensor<1x196x16xi32> into tensor<1x14x14x16xi32>418// CHECK: return %[[RESULT]]419 420func.func @conv_integer_extend(%arg0: tensor<1x16x16x4xi8>, %arg1: tensor<3x3x4x16xi8>, %arg2: tensor<1x14x14x16xi32>) -> tensor<1x14x14x16xi32> {421 %0 = linalg.conv_2d_nhwc_hwcf422 {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64> }423 ins(%arg0, %arg1: tensor<1x16x16x4xi8>, tensor<3x3x4x16xi8>)424 outs(%arg2: tensor<1x14x14x16xi32>) -> tensor<1x14x14x16xi32>425 return %0 : tensor<1x14x14x16xi32>426}427 428module attributes {transform.with_named_sequence} {429 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {430 %0 = transform.structured.match ops{["linalg.conv_2d_nhwc_hwcf"]} in %arg1 : (!transform.any_op) -> !transform.any_op431 %img2col_tensor_producer, %transformed = transform.structured.convert_conv2d_to_img2col %0 : (!transform.any_op) -> (!transform.any_op, !transform.any_op)432 transform.print %img2col_tensor_producer {name = "tensor_producer"}: !transform.any_op433 transform.print %transformed {name = "transformed"}: !transform.any_op434 transform.yield435 }436}437 438// -----439 440// Check for compatible complex case.441 442// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>443// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>444// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d3, d2)>445// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>446// CHECK: @conv_complex447// CHECK: %[[MATMUL_RESULT:.+]] = linalg.generic {indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]]448// CHECK-SAME: ins(%{{.*}}, %{{.*}} : tensor<1x196x36xcomplex<f32>>, tensor<36x16xcomplex<f32>>)449// CHECK-SAME: outs(%[[COLLAPSED_OUT]] : tensor<1x196x16xcomplex<f32>>)450// CHECK: ^bb0(%[[ARG0:.+]]: complex<f32>, %[[ARG1:.+]]: complex<f32>, %[[ARG2:.+]]: complex<f32>)451// CHECK: %[[MUL:.+]] = complex.mul %[[ARG0]], %[[ARG1]] : complex<f32>452// CHECK: %[[ADD:.+]] = complex.add %[[MUL]], %[[ARG2]] : complex<f32>453// CHECK: linalg.yield %[[ADD]] : complex<f32>454// CHECK: } -> tensor<1x196x16xcomplex<f32>>455// CHECK: %[[RESULT:.+]] = tensor.expand_shape %[[MATMUL_RESULT]] {{\[}}[0], [1, 2], [3]] output_shape [1, 14, 14, 16] : tensor<1x196x16xcomplex<f32>> into tensor<1x14x14x16xcomplex<f32>>456// CHECK: return %[[RESULT]]457 458func.func @conv_complex(%arg0: tensor<1x16x16x4xcomplex<f32>>, %arg1: tensor<3x3x4x16xcomplex<f32>>, %arg2: tensor<1x14x14x16xcomplex<f32>>) -> tensor<1x14x14x16xcomplex<f32>> {459 %0 = linalg.conv_2d_nhwc_hwcf460 {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64> }461 ins(%arg0, %arg1: tensor<1x16x16x4xcomplex<f32>>, tensor<3x3x4x16xcomplex<f32>>)462 outs(%arg2: tensor<1x14x14x16xcomplex<f32>>) -> tensor<1x14x14x16xcomplex<f32>>463 return %0 : tensor<1x14x14x16xcomplex<f32>>464}465 466module attributes {transform.with_named_sequence} {467 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {468 %0 = transform.structured.match ops{["linalg.conv_2d_nhwc_hwcf"]} in %arg1 : (!transform.any_op) -> !transform.any_op469 %img2col_tensor_producer, %transformed = transform.structured.convert_conv2d_to_img2col %0 : (!transform.any_op) -> (!transform.any_op, !transform.any_op)470 transform.print %img2col_tensor_producer {name = "tensor_producer"}: !transform.any_op471 transform.print %transformed {name = "transformed"}: !transform.any_op472 transform.yield473 }474}475 476// -----477 478// Check for compatible complex extended case.479 480// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>481// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>482// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d3, d2)>483// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>484// CHECK: @conv_complex_extended485// CHECK: %[[MATMUL_RESULT:.+]] = linalg.generic {indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]]486// CHECK-SAME: ins(%{{.*}}, %{{.*}} : tensor<1x196x36xcomplex<f32>>, tensor<36x16xcomplex<f16>>)487// CHECK-SAME: outs(%[[COLLAPSED_OUT]] : tensor<1x196x16xcomplex<f32>>)488// CHECK: ^bb0(%[[ARG0:.+]]: complex<f32>, %[[ARG1:.+]]: complex<f16>, %[[ARG2:.+]]: complex<f32>)489// CHECK: %[[REAL:.+]] = complex.re %[[ARG1]] : complex<f16>490// CHECK: %[[IMAG:.+]] = complex.im %[[ARG1]] : complex<f16>491// CHECK: %[[REEXT:.+]] = arith.extf %[[REAL]] : f16 to f32492// CHECK: %[[IMEXT:.+]] = arith.extf %[[IMAG]] : f16 to f32493// CHECK: %[[COMPLEX:.+]] = complex.create %[[REEXT]], %[[IMEXT]] : complex<f32>494// CHECK: %[[MUL:.+]] = complex.mul %[[ARG0]], %[[COMPLEX]] : complex<f32>495// CHECK: %[[ADD:.+]] = complex.add %[[MUL]], %[[ARG2]] : complex<f32>496// CHECK: linalg.yield %[[ADD]] : complex<f32>497// CHECK: } -> tensor<1x196x16xcomplex<f32>>498// CHECK: %[[RESULT:.+]] = tensor.expand_shape %[[MATMUL_RESULT]] {{\[}}[0], [1, 2], [3]] output_shape [1, 14, 14, 16] : tensor<1x196x16xcomplex<f32>> into tensor<1x14x14x16xcomplex<f32>>499// CHECK: return %[[RESULT]]500 501func.func @conv_complex_extended(%arg0: tensor<1x16x16x4xcomplex<f32>>, %arg1: tensor<3x3x4x16xcomplex<f16>>, %arg2: tensor<1x14x14x16xcomplex<f32>>) -> tensor<1x14x14x16xcomplex<f32>> {502 %0 = linalg.conv_2d_nhwc_hwcf503 {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64> }504 ins(%arg0, %arg1: tensor<1x16x16x4xcomplex<f32>>, tensor<3x3x4x16xcomplex<f16>>)505 outs(%arg2: tensor<1x14x14x16xcomplex<f32>>) -> tensor<1x14x14x16xcomplex<f32>>506 return %0 : tensor<1x14x14x16xcomplex<f32>>507}508 509module attributes {transform.with_named_sequence} {510 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {511 %0 = transform.structured.match ops{["linalg.conv_2d_nhwc_hwcf"]} in %arg1 : (!transform.any_op) -> !transform.any_op512 %img2col_tensor_producer, %transformed = transform.structured.convert_conv2d_to_img2col %0 : (!transform.any_op) -> (!transform.any_op, !transform.any_op)513 transform.print %img2col_tensor_producer {name = "tensor_producer"}: !transform.any_op514 transform.print %transformed {name = "transformed"}: !transform.any_op515 transform.yield516 }517}518 519// -----520 521// Check for compatible complex extended case.522 523// CHECK-DAG: #[[MAP0:.+]] = affine_map<(d0, d1, d2) -> (d0, d1, d2)>524// CHECK-DAG: #[[MAP1:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>525// CHECK-DAG: #[[MAP2:.+]] = affine_map<(d0, d1, d2, d3) -> (d3, d2)>526// CHECK-DAG: #[[MAP3:.+]] = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2)>527// CHECK: @conv_complex_f16_extended528// CHECK: %[[MATMUL_RESULT:.+]] = linalg.generic {indexing_maps = [#[[MAP1]], #[[MAP2]], #[[MAP3]]]529// CHECK-SAME: ins(%{{.*}}, %{{.*}} : tensor<1x196x36xcomplex<f32>>, tensor<36x16xf16>)530// CHECK-SAME: outs(%[[COLLAPSED_OUT]] : tensor<1x196x16xcomplex<f32>>)531// CHECK: ^bb0(%[[ARG0:.+]]: complex<f32>, %[[ARG1:.+]]: f16, %[[ARG2:.+]]: complex<f32>)532// CHECK: %[[EXT:.+]] = arith.extf %[[ARG1]] : f16 to f32533// CHECK: %[[ZERO:.+]] = arith.constant 0.000000e+00 : f32534// CHECK: %[[COMPLEX:.+]] = complex.create %[[EXT]], %[[ZERO]]535// CHECK: %[[MUL:.+]] = complex.mul %[[ARG0]], %[[COMPLEX]] : complex<f32>536// CHECK: %[[ADD:.+]] = complex.add %[[MUL]], %[[ARG2]] : complex<f32>537// CHECK: linalg.yield %[[ADD]] : complex<f32>538// CHECK: } -> tensor<1x196x16xcomplex<f32>>539// CHECK: %[[RESULT:.+]] = tensor.expand_shape %[[MATMUL_RESULT]] {{\[}}[0], [1, 2], [3]] output_shape [1, 14, 14, 16] : tensor<1x196x16xcomplex<f32>> into tensor<1x14x14x16xcomplex<f32>>540// CHECK: return %[[RESULT]]541 542func.func @conv_complex_f16_extended(%arg0: tensor<1x16x16x4xcomplex<f32>>, %arg1: tensor<3x3x4x16xf16>, %arg2: tensor<1x14x14x16xcomplex<f32>>) -> tensor<1x14x14x16xcomplex<f32>> {543 %0 = linalg.conv_2d_nhwc_hwcf544 {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64> }545 ins(%arg0, %arg1: tensor<1x16x16x4xcomplex<f32>>, tensor<3x3x4x16xf16>)546 outs(%arg2: tensor<1x14x14x16xcomplex<f32>>) -> tensor<1x14x14x16xcomplex<f32>>547 return %0 : tensor<1x14x14x16xcomplex<f32>>548}549 550module attributes {transform.with_named_sequence} {551 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {552 %0 = transform.structured.match ops{["linalg.conv_2d_nhwc_hwcf"]} in %arg1 : (!transform.any_op) -> !transform.any_op553 %img2col_tensor_producer, %transformed = transform.structured.convert_conv2d_to_img2col %0 : (!transform.any_op) -> (!transform.any_op, !transform.any_op)554 transform.print %img2col_tensor_producer {name = "tensor_producer"}: !transform.any_op555 transform.print %transformed {name = "transformed"}: !transform.any_op556 transform.yield557 }558}559