93 lines · plain
1// RUN: mlir-opt %s -allow-unregistered-dialect -transform-interpreter -canonicalize --split-input-file -verify-diagnostics| FileCheck %s2 3func.func @conv2d(%arg0: tensor<2x10x10x5xf32>, %arg1: tensor<2x3x3x5xf32>, %arg2: tensor<1xf32>, %arg3: tensor<2x8x8x2xf32>) -> tensor<2x8x8x2xf32> {4 %0 = linalg.conv_2d_nhwc_fhwc {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} ins(%arg0, %arg1 : tensor<2x10x10x5xf32>, tensor<2x3x3x5xf32>) outs(%arg3 : tensor<2x8x8x2xf32>) -> tensor<2x8x8x2xf32>5 return %0 : tensor<2x8x8x2xf32>6}7 8module attributes {transform.with_named_sequence} {9 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {10 %0 = transform.structured.match ops{["linalg.conv_2d_nhwc_fhwc"]} in %arg1 : (!transform.any_op) -> !transform.any_op11 %1 = transform.structured.winograd_conv2d %0 { fmr = 1: i32 } : (!transform.any_op) -> (!transform.any_op)12 transform.yield13 }14}15 16// CHECK-LABEL: func.func @conv2d17// CHECK: linalg.winograd_filter_transform fmr(F_4_3)18// CHECK: linalg.winograd_input_transform fmr(F_4_3)19// CHECK: linalg.batch_matmul20// CHECK: linalg.winograd_output_transform fmr(F_4_3)21 22// -----23 24func.func @conv2d_unaligned(%arg0: tensor<2x11x11x5xf32>, %arg1: tensor<2x3x3x5xf32>, %arg2: tensor<1xf32>, %arg3: tensor<2x9x9x2xf32>) -> tensor<2x9x9x2xf32> {25 %0 = linalg.conv_2d_nhwc_fhwc {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} ins(%arg0, %arg1 : tensor<2x11x11x5xf32>, tensor<2x3x3x5xf32>) outs(%arg3 : tensor<2x9x9x2xf32>) -> tensor<2x9x9x2xf32>26 return %0 : tensor<2x9x9x2xf32>27}28 29module attributes {transform.with_named_sequence} {30 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {31 %0 = transform.structured.match ops{["linalg.conv_2d_nhwc_fhwc"]} in %arg1 : (!transform.any_op) -> !transform.any_op32 %1 = transform.structured.winograd_conv2d %0 { fmr = 1: i32 } : (!transform.any_op) -> (!transform.any_op)33 transform.yield34 }35}36 37// CHECK-LABEL: func.func @conv2d_unaligned38// CHECK: linalg.winograd_filter_transform fmr(F_4_3)39// CHECK: tensor.pad40// CHECK-SAME: low[0, 0, 0, 0] high[0, 3, 3, 0]41// CHECK: linalg.winograd_input_transform fmr(F_4_3)42// CHECK: tensor.pad43// CHECK-SAME: low[0, 0, 0, 0] high[0, 3, 3, 0]44// CHECK: linalg.winograd_output_transform fmr(F_4_3)45 46// -----47 48func.func @conv2d_unsupported(%arg0: tensor<2x10x10x5xf32>, %arg1: tensor<3x3x5x2xf32>, %arg2: tensor<1xf32>, %arg3: tensor<2x8x8x2xf32>) -> tensor<2x8x8x2xf32> {49 %0 = linalg.conv_2d_nhwc_hwcf {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} ins(%arg0, %arg1 : tensor<2x10x10x5xf32>, tensor<3x3x5x2xf32>) outs(%arg3 : tensor<2x8x8x2xf32>) -> tensor<2x8x8x2xf32>50 return %0 : tensor<2x8x8x2xf32>51}52 53module attributes {transform.with_named_sequence} {54 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {55 %0 = transform.structured.match ops{["linalg.conv_2d_nhwc_hwcf"]} in %arg1 : (!transform.any_op) -> !transform.any_op56 // expected-error @+1 {{this operation is not supported to convert to Winograd Conv2D}}57 %1 = transform.structured.winograd_conv2d %0 { fmr = 1: i32 } : (!transform.any_op) -> (!transform.any_op)58 transform.yield59 }60}61 62// -----63 64func.func @conv2d_unsupported_type(%arg0: memref<2x10x10x5xf32>, %arg1: memref<2x3x3x5xf32>, %arg2: memref<2x8x8x2xf32>) {65 linalg.conv_2d_nhwc_fhwc {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} ins(%arg0, %arg1 : memref<2x10x10x5xf32>, memref<2x3x3x5xf32>) outs(%arg2 : memref<2x8x8x2xf32>)66 return67}68 69module attributes {transform.with_named_sequence} {70 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {71 %0 = transform.structured.match ops{["linalg.conv_2d_nhwc_fhwc"]} in %arg1 : (!transform.any_op) -> !transform.any_op72 // expected-error @+1 {{apply Winograd Conv2D failed}}73 %1 = transform.structured.winograd_conv2d %0 { fmr = 1: i32 } : (!transform.any_op) -> (!transform.any_op)74 transform.yield75 }76}77 78// -----79 80func.func @conv2d(%arg0: tensor<2x?x?x5xf32>, %arg1: tensor<2x3x3x5xf32>, %arg2: tensor<1xf32>, %arg3: tensor<2x?x?x2xf32>) -> tensor<2x?x?x2xf32> {81 %0 = linalg.conv_2d_nhwc_fhwc {dilations = dense<1> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} ins(%arg0, %arg1 : tensor<2x?x?x5xf32>, tensor<2x3x3x5xf32>) outs(%arg3 : tensor<2x?x?x2xf32>) -> tensor<2x?x?x2xf32>82 return %0 : tensor<2x?x?x2xf32>83}84 85module attributes {transform.with_named_sequence} {86 transform.named_sequence @__transform_main(%arg1: !transform.any_op {transform.readonly}) {87 %0 = transform.structured.match ops{["linalg.conv_2d_nhwc_fhwc"]} in %arg1 : (!transform.any_op) -> !transform.any_op88 // expected-error @+1 {{apply Winograd Conv2D failed}}89 %1 = transform.structured.winograd_conv2d %0 { fmr = 1: i32 } : (!transform.any_op) -> (!transform.any_op)90 transform.yield91 }92}93