307 lines · plain
1// RUN: mlir-opt %s -canonicalize="test-convergence" -split-input-file -allow-unregistered-dialect | FileCheck %s2 3// This file contains some tests of canonicalizations and foldings involving vector.transpose.4 5// +---------------------------------------------------------------------------6// Tests of FoldTransposeBroadcast: transpose(broadcast) -> broadcast7// +---------------------------------------------------------------------------8 9// CHECK-LABEL: func @transpose_scalar_broadcast110// CHECK-SAME: (%[[ARG:.+]]: vector<1xf32>)11// CHECK: %[[V:.+]] = vector.broadcast %[[ARG]] : vector<1xf32> to vector<1x8xf32>12// CHECK: return %[[V]] : vector<1x8xf32>13func.func @transpose_scalar_broadcast1(%value: vector<1xf32>) -> vector<1x8xf32> {14 %bcast = vector.broadcast %value : vector<1xf32> to vector<8x1xf32>15 %t = vector.transpose %bcast, [1, 0] : vector<8x1xf32> to vector<1x8xf32>16 return %t : vector<1x8xf32>17}18 19// -----20 21// CHECK-LABEL: func @transpose_scalar_broadcast222// CHECK-SAME: (%[[ARG:.+]]: f32)23// CHECK: %[[V:.+]] = vector.broadcast %[[ARG]] : f32 to vector<1x8xf32>24// CHECK: return %[[V]] : vector<1x8xf32>25func.func @transpose_scalar_broadcast2(%value: f32) -> vector<1x8xf32> {26 %bcast = vector.broadcast %value : f32 to vector<8x1xf32>27 %t = vector.transpose %bcast, [1, 0] : vector<8x1xf32> to vector<1x8xf32>28 return %t : vector<1x8xf32>29}30 31// -----32 33 34// CHECK-LABEL: broadcast_transpose_scalar_to_broadcast35// CHECK-SAME: %[[ARG:.*]]: i8) -> vector<2x3x4xi8> {36func.func @broadcast_transpose_scalar_to_broadcast(%arg0 : i8) -> vector<2x3x4xi8> {37// CHECK: %[[BC:.*]] = vector.broadcast %[[ARG]] : i8 to vector<2x3x4xi8>38 %0 = vector.broadcast %arg0 : i8 to vector<3x4x2xi8>39 %1 = vector.transpose %0, [2, 0, 1] : vector<3x4x2xi8> to vector<2x3x4xi8>40// CHECK: return %[[BC]] : vector<2x3x4xi8>41 return %1 : vector<2x3x4xi8>42}43 44// -----45 46// CHECK-LABEL: broadcast_transpose_ones_to_broadcast47// CHECK-SAME: %[[ARG:.*]]: vector<1x1x1xi8>) -> vector<2x3x4xi8> {48// CHECK: %[[RES:.*]] = vector.broadcast %[[ARG]] : vector<1x1x1xi8> to vector<2x3x4xi8>49// CHECK: return %[[RES]] : vector<2x3x4xi8>50func.func @broadcast_transpose_ones_to_broadcast(%arg0 : vector<1x1x1xi8>) -> vector<2x3x4xi8> {51 %0 = vector.broadcast %arg0 : vector<1x1x1xi8> to vector<3x4x2xi8>52 %1 = vector.transpose %0, [2, 0, 1] : vector<3x4x2xi8> to vector<2x3x4xi8>53 return %1 : vector<2x3x4xi8>54}55 56// -----57 58// CHECK-LABEL: broadcast_transpose_partial_ones_to_broadcast59// CHECK-SAME: %[[ARG:.*]]: vector<1xi8>) -> vector<8x1xi8> {60// CHECK: %[[RES:.*]] = vector.broadcast %[[ARG]] : vector<1xi8> to vector<8x1xi8>61// CHECK: return %[[RES]] : vector<8x1xi8>62func.func @broadcast_transpose_partial_ones_to_broadcast(%arg0 : vector<1xi8>) -> vector<8x1xi8> {63 %0 = vector.broadcast %arg0 : vector<1xi8> to vector<1x8xi8>64 %1 = vector.transpose %0, [1, 0] : vector<1x8xi8> to vector<8x1xi8>65 return %1 : vector<8x1xi8>66}67 68// -----69 70// CHECK-LABEL: broadcast_transpose_mixed_example71// CHECK-SAME: %[[ARG:.*]]: vector<4x1x1x7xi8>) -> vector<3x2x4x5x6x7xi8> {72// CHECK: %[[RES:.*]] = vector.broadcast %[[ARG]] : vector<4x1x1x7xi8> to vector<3x2x4x5x6x7xi8>73// CHECK: return %[[RES]] : vector<3x2x4x5x6x7xi8>74func.func @broadcast_transpose_mixed_example(%arg0 : vector<4x1x1x7xi8>) -> vector<3x2x4x5x6x7xi8> {75 %0 = vector.broadcast %arg0 : vector<4x1x1x7xi8> to vector<2x3x4x5x6x7xi8>76 %1 = vector.transpose %0, [1, 0, 2, 3, 4, 5] : vector<2x3x4x5x6x7xi8> to vector<3x2x4x5x6x7xi8>77 return %1 : vector<3x2x4x5x6x7xi8>78}79 80// -----81 82// CHECK-LABEL: broadcast_transpose_final_group83// CHECK-SAME: %[[ARG:.*]]: vector<4x7x1x1xi8>) -> vector<4x7x2x3xi8> {84// CHECK: %[[RES:.*]] = vector.broadcast %[[ARG]] : vector<4x7x1x1xi8> to vector<4x7x2x3xi8>85// CHECK: return %[[RES]] : vector<4x7x2x3xi8>86func.func @broadcast_transpose_final_group(%arg0 : vector<4x7x1x1xi8>) -> vector<4x7x2x3xi8> {87 %0 = vector.broadcast %arg0 : vector<4x7x1x1xi8> to vector<4x7x3x2xi8>88 %1 = vector.transpose %0, [0, 1, 3, 2] : vector<4x7x3x2xi8> to vector<4x7x2x3xi8>89 return %1 : vector<4x7x2x3xi8>90}91 92// -----93 94// CHECK-LABEL: negative_broadcast_transpose_square95// CHECK-SAME: %[[ARG:.*]]:96// CHECK: %[[BCT:.*]] = vector.broadcast %[[ARG]]97// CHECK: %[[TRP:.*]] = vector.transpose %[[BCT]], [1, 0]98// CHECK: return %[[TRP]] : vector<4x4xi8>99func.func @negative_broadcast_transpose_square(%arg0 : vector<4x1xi8>) -> vector<4x4xi8> {100 %0 = vector.broadcast %arg0 : vector<4x1xi8> to vector<4x4xi8>101 %1 = vector.transpose %0, [1, 0] : vector<4x4xi8> to vector<4x4xi8>102 return %1 : vector<4x4xi8>103}104 105// -----106 107// CHECK-LABEL: negative_broadcast_transpose_hypercube108// CHECK-SAME: %[[ARG:.*]]:109// CHECK: %[[BCT:.*]] = vector.broadcast %[[ARG]]110// CHECK: %[[TRP:.*]] = vector.transpose %[[BCT]], [1, 0, 3, 2]111// CHECK: return %[[TRP]] : vector<4x4x4x4xi8>112func.func @negative_broadcast_transpose_hypercube(%arg0 : vector<1x1x4xi8>) -> vector<4x4x4x4xi8> {113 %0 = vector.broadcast %arg0 : vector<1x1x4xi8> to vector<4x4x4x4xi8>114 %1 = vector.transpose %0, [1, 0, 3, 2] : vector<4x4x4x4xi8> to vector<4x4x4x4xi8>115 return %1 : vector<4x4x4x4xi8>116}117 118// -----119 120// CHECK-LABEL: negative_broadcast_transpose_102121// CHECK-SAME: %[[ARG:.*]]:122// CHECK: %[[BCT:.*]] = vector.broadcast %[[ARG]]123// CHECK: %[[TRP:.*]] = vector.transpose %[[BCT]], [1, 0, 2]124// CHECK: return %[[TRP]] : vector<3x3x3xi8>125func.func @negative_broadcast_transpose_102(%arg0 : vector<3x1x3xi8>) -> vector<3x3x3xi8> {126 %0 = vector.broadcast %arg0 : vector<3x1x3xi8> to vector<3x3x3xi8>127 %1 = vector.transpose %0, [1, 0, 2] : vector<3x3x3xi8> to vector<3x3x3xi8>128 return %1 : vector<3x3x3xi8>129}130 131// -----132 133// CHECK-LABEL: negative_broadcast_transpose_021134// CHECK-SAME: %[[ARG:.*]]:135// CHECK: %[[BCT:.*]] = vector.broadcast %[[ARG]]136// CHECK: %[[TRP:.*]] = vector.transpose %[[BCT]], [0, 2, 1]137// CHECK: return %[[TRP]] : vector<3x3x3xi8>138func.func @negative_broadcast_transpose_021(%arg0 : vector<3x1x3xi8>) -> vector<3x3x3xi8> {139 %0 = vector.broadcast %arg0 : vector<3x1x3xi8> to vector<3x3x3xi8>140 %1 = vector.transpose %0, [0, 2, 1] : vector<3x3x3xi8> to vector<3x3x3xi8>141 return %1 : vector<3x3x3xi8>142}143 144// -----145 146/// +--------------------------------------------------------------------------147/// Tests of ShapeCastOp::fold: shape_cast(transpose) -> shape_cast148/// +--------------------------------------------------------------------------149 150// In this test, the permutation maps the non-unit dimensions (1 and 2) as follows:151// 1 -> 0152// 2 -> 4153// Because 0 < 4, this permutation is order preserving and effectively a shape_cast.154// CHECK-LABEL: @shape_cast_of_transpose155// CHECK-SAME: %[[ARG:.*]]: vector<1x4x4x1x1xi8>) -> vector<4x4xi8> {156// CHECK: %[[SHAPE_CAST:.*]] = vector.shape_cast %[[ARG]] :157// CHECK-SAME: vector<1x4x4x1x1xi8> to vector<4x4xi8>158// CHECK: return %[[SHAPE_CAST]] : vector<4x4xi8>159func.func @shape_cast_of_transpose(%arg : vector<1x4x4x1x1xi8>) -> vector<4x4xi8> {160 %0 = vector.transpose %arg, [1, 0, 3, 4, 2]161 : vector<1x4x4x1x1xi8> to vector<4x1x1x1x4xi8>162 %1 = vector.shape_cast %0 : vector<4x1x1x1x4xi8> to vector<4x4xi8>163 return %1 : vector<4x4xi8>164}165 166// -----167 168// In this test, the permutation maps the non-unit dimensions (1 and 2) as follows:169// 1 -> 0170// 2 -> 4171// Because 0 < 4, this permutation is order preserving and effectively a shape_cast.172// (same as the example above, but one of the dims is scalable)173// CHECK-LABEL: @shape_cast_of_transpose_scalable174// CHECK-SAME: %[[ARG:.*]]: vector<1x[4]x4x1x1xi8>) -> vector<[4]x4xi8> {175// CHECK: %[[SHAPE_CAST:.*]] = vector.shape_cast %[[ARG]] :176// CHECK-SAME: vector<1x[4]x4x1x1xi8> to vector<[4]x4xi8>177// CHECK: return %[[SHAPE_CAST]] : vector<[4]x4xi8>178func.func @shape_cast_of_transpose_scalable(%arg : vector<1x[4]x4x1x1xi8>) -> vector<[4]x4xi8> {179 %0 = vector.transpose %arg, [1, 0, 3, 4, 2]180 : vector<1x[4]x4x1x1xi8> to vector<[4]x1x1x1x4xi8>181 %1 = vector.shape_cast %0 : vector<[4]x1x1x1x4xi8> to vector<[4]x4xi8>182 return %1 : vector<[4]x4xi8>183}184 185// -----186 187// In this test, the mapping of non-unit dimensions (1 and 2) is as follows:188// 1 -> 2189// 2 -> 1190// As this is not increasing (2 > 1), this transpose is not order191// preserving and cannot be treated as a shape_cast.192// CHECK-LABEL: @negative_shape_cast_of_transpose193// CHECK-SAME: %[[ARG:.*]]: vector<1x4x4x1xi8>) -> vector<4x4xi8> {194// CHECK: %[[TRANSPOSE:.*]] = vector.transpose %[[ARG]]195// CHECK: %[[SHAPE_CAST:.*]] = vector.shape_cast %[[TRANSPOSE]]196// CHECK: return %[[SHAPE_CAST]] : vector<4x4xi8>197func.func @negative_shape_cast_of_transpose(%arg : vector<1x4x4x1xi8>) -> vector<4x4xi8> {198 %0 = vector.transpose %arg, [0, 2, 1, 3]199 : vector<1x4x4x1xi8> to vector<1x4x4x1xi8>200 %1 = vector.shape_cast %0 : vector<1x4x4x1xi8> to vector<4x4xi8>201 return %1 : vector<4x4xi8>202}203 204// -----205 206/// +--------------------------------------------------------------------------207/// Tests of FoldTransposeShapeCast: transpose(shape_cast) -> shape_cast208/// +--------------------------------------------------------------------------209 210// A transpose that is 'order preserving' can be treated like a shape_cast. 211// CHECK-LABEL: @transpose_of_shape_cast212// CHECK-SAME: %[[ARG:.*]]: vector<2x3x1x1xi8>) -> vector<6x1x1xi8> {213// CHECK: %[[SHAPE_CAST:.*]] = vector.shape_cast %[[ARG]] :214// CHECK-SAME: vector<2x3x1x1xi8> to vector<6x1x1xi8>215// CHECK: return %[[SHAPE_CAST]] : vector<6x1x1xi8>216func.func @transpose_of_shape_cast(%arg : vector<2x3x1x1xi8>) -> vector<6x1x1xi8> {217 %0 = vector.shape_cast %arg : vector<2x3x1x1xi8> to vector<6x1x1xi8>218 %1 = vector.transpose %0, [0, 2, 1]219 : vector<6x1x1xi8> to vector<6x1x1xi8>220 return %1 : vector<6x1x1xi8>221}222 223// -----224 225// CHECK-LABEL: @transpose_of_shape_cast_scalable226// CHECK-SAME: %[[ARG:.*]]: vector<[2]x3x1x1xi8>) -> vector<[6]x1x1xi8> {227// CHECK: %[[SHAPE_CAST:.*]] = vector.shape_cast %[[ARG]] :228// CHECK-SAME: vector<[2]x3x1x1xi8> to vector<[6]x1x1xi8>229// CHECK: return %[[SHAPE_CAST]] : vector<[6]x1x1xi8>230func.func @transpose_of_shape_cast_scalable(%arg : vector<[2]x3x1x1xi8>) -> vector<[6]x1x1xi8> {231 %0 = vector.shape_cast %arg : vector<[2]x3x1x1xi8> to vector<[6]x1x1xi8>232 %1 = vector.transpose %0, [0, 2, 1]233 : vector<[6]x1x1xi8> to vector<[6]x1x1xi8>234 return %1 : vector<[6]x1x1xi8>235}236 237// -----238 239// Scalable 1 dimensions (i.e. [1]) should be treated as non-unit dimensions240// (hence no folding).241// CHECK-LABEL: @negative_transpose_of_shape_cast_scalable_unit242// CHECK: vector.shape_cast243// CHECK: vector.transpose244func.func @negative_transpose_of_shape_cast_scalable_unit(%arg : vector<[1]x4x1xi8>) -> vector<4x[1]xi8> {245 %0 = vector.shape_cast %arg : vector<[1]x4x1xi8> to vector<[1]x4xi8>246 %1 = vector.transpose %0, [1, 0] : vector<[1]x4xi8> to vector<4x[1]xi8>247 return %1 : vector<4x[1]xi8>248}249 250// -----251 252// Test of shape_cast (not) folding.253// CHECK-LABEL: @negative_transpose_of_shape_cast254// CHECK-SAME: %[[ARG:.*]]: vector<6xi8>) -> vector<2x3xi8> {255// CHECK: %[[SHAPE_CAST:.*]] = vector.shape_cast %[[ARG]] :256// CHECK: %[[TRANSPOSE:.*]] = vector.transpose %[[SHAPE_CAST]]257// CHECK: return %[[TRANSPOSE]] : vector<2x3xi8>258func.func @negative_transpose_of_shape_cast(%arg : vector<6xi8>) -> vector<2x3xi8> {259 %0 = vector.shape_cast %arg : vector<6xi8> to vector<3x2xi8>260 %1 = vector.transpose %0, [1, 0] : vector<3x2xi8> to vector<2x3xi8>261 return %1 : vector<2x3xi8>262}263 264// -----265 266// +-----------------------------------267// Tests of TransposeOp::fold268// +-----------------------------------269 270// CHECK-LABEL: transpose_1D_identity271// CHECK-SAME: [[ARG:%.*]]: vector<4xf32>272// CHECK-NEXT: return [[ARG]]273func.func @transpose_1D_identity(%arg : vector<4xf32>) -> vector<4xf32> {274 %0 = vector.transpose %arg, [0] : vector<4xf32> to vector<4xf32>275 return %0 : vector<4xf32>276}277 278// -----279 280// CHECK-LABEL: transpose_2D_identity281// CHECK-SAME: [[ARG:%.*]]: vector<4x3xf32>282// CHECK-NEXT: return [[ARG]]283func.func @transpose_2D_identity(%arg : vector<4x3xf32>) -> vector<4x3xf32> {284 %0 = vector.transpose %arg, [0, 1] : vector<4x3xf32> to vector<4x3xf32>285 return %0 : vector<4x3xf32>286}287 288// -----289 290// CHECK-LABEL: transpose_shape_and_order_preserving291// CHECK-SAME: [[ARG:%.*]]: vector<6x1x1x4xi8>292// CHECK-NEXT: return [[ARG]]293func.func @transpose_shape_and_order_preserving(%arg : vector<6x1x1x4xi8>) -> vector<6x1x1x4xi8> {294 %0 = vector.transpose %arg, [0, 2, 1, 3] : vector<6x1x1x4xi8> to vector<6x1x1x4xi8>295 return %0 : vector<6x1x1x4xi8>296}297 298// -----299 300// CHECK-LABEL: negative_transpose_fold301// CHECK: [[TRANSP:%.*]] = vector.transpose302// CHECK: return [[TRANSP]]303func.func @negative_transpose_fold(%arg : vector<2x2xi8>) -> vector<2x2xi8> {304 %0 = vector.transpose %arg, [1, 0] : vector<2x2xi8> to vector<2x2xi8>305 return %0 : vector<2x2xi8>306}307