307 lines · plain
1// RUN: mlir-opt %s -test-vector-transfer-flatten-patterns -split-input-file | FileCheck %s2 3///----------------------------------------------------------------------------------------4/// [Pattern: DropUnitDimFromElementwiseOps]5///----------------------------------------------------------------------------------------6 7func.func @fold_unit_dim_add_basic(%vec : vector<1x8xi32>) -> vector<1x8xi32> {8 %res = arith.addi %vec, %vec : vector<1x8xi32>9 return %res : vector<1x8xi32>10}11// CHECK-LABEL: func.func @fold_unit_dim_add_basic(12// CHECK-SAME: %[[VAL_0:.*]]: vector<1x8xi32>) -> vector<1x8xi32> {13// CHECK: %[[VAL_1:.*]] = vector.shape_cast %[[VAL_0]] : vector<1x8xi32> to vector<8xi32>14// CHECK: %[[VAL_2:.*]] = vector.shape_cast %[[VAL_0]] : vector<1x8xi32> to vector<8xi32>15// CHECK: %[[VAL_3:.*]] = arith.addi %[[VAL_1]], %[[VAL_2]] : vector<8xi32>16// CHECK: %[[VAL_4:.*]] = vector.shape_cast %[[VAL_3]] : vector<8xi32> to vector<1x8xi32>17// CHECK: return %[[VAL_4]] : vector<1x8xi32>18 19// -----20 21func.func @fold_unit_dim_add_leading_and_trailing(%vec : vector<1x8x1xi32>) -> vector<1x8x1xi32> {22 %res = arith.addi %vec, %vec : vector<1x8x1xi32>23 return %res : vector<1x8x1xi32>24}25// CHECK-LABEL: func.func @fold_unit_dim_add_leading_and_trailing(26// CHECK-SAME: %[[VAL_0:.*]]: vector<1x8x1xi32>) -> vector<1x8x1xi32> {27// CHECK: %[[VAL_1:.*]] = vector.shape_cast %[[VAL_0]] : vector<1x8x1xi32> to vector<8xi32>28// CHECK: %[[VAL_2:.*]] = vector.shape_cast %[[VAL_0]] : vector<1x8x1xi32> to vector<8xi32>29// CHECK: %[[VAL_3:.*]] = arith.addi %[[VAL_1]], %[[VAL_2]] : vector<8xi32>30// CHECK: %[[VAL_4:.*]] = vector.shape_cast %[[VAL_3]] : vector<8xi32> to vector<1x8x1xi32>31// CHECK: return %[[VAL_4]] : vector<1x8x1xi32>32 33// -----34 35func.func @fold_unit_dim_add(%vec_0 : vector<8x1xi32>,36 %vec_1 : vector<1x8xi32>) -> vector<8xi32> {37 %sc_vec_0 = vector.shape_cast %vec_0 : vector<8x1xi32> to vector<1x8xi32>38 %add = arith.addi %sc_vec_0, %vec_1 : vector<1x8xi32>39 %res = vector.shape_cast %add : vector<1x8xi32> to vector<8xi32>40 return %res : vector<8xi32>41}42 43// CHECK-LABEL: func.func @fold_unit_dim_add(44// CHECK-SAME: %[[VAL_0:.*]]: vector<8x1xi32>,45// CHECK-SAME: %[[VAL_1:.*]]: vector<1x8xi32>) -> vector<8xi32> {46// CHECK: %[[VAL_2:.*]] = vector.shape_cast %[[VAL_0]] : vector<8x1xi32> to vector<8xi32>47// CHECK: %[[VAL_3:.*]] = vector.shape_cast %[[VAL_1]] : vector<1x8xi32> to vector<8xi32>48// CHECK: %[[VAL_4:.*]] = arith.addi %[[VAL_2]], %[[VAL_3]] : vector<8xi32>49// CHECK: return %[[VAL_4]] : vector<8xi32>50 51// -----52 53func.func @fold_unit_dim_mulf(%vec_0 : vector<8x[2]x1xf32>,54 %vec_1 : vector<1x8x[2]xf32>) -> vector<8x[2]xf32> {55 %sc_vec_0 = vector.shape_cast %vec_0 : vector<8x[2]x1xf32> to vector<1x8x[2]xf32>56 %add = arith.mulf %sc_vec_0, %vec_1 : vector<1x8x[2]xf32>57 %res = vector.shape_cast %add : vector<1x8x[2]xf32> to vector<8x[2]xf32>58 return %res : vector<8x[2]xf32>59}60 61// CHECK-LABEL: func.func @fold_unit_dim_mulf(62// CHECK-SAME: %[[VAL_0:.*]]: vector<8x[2]x1xf32>,63// CHECK-SAME: %[[VAL_1:.*]]: vector<1x8x[2]xf32>) -> vector<8x[2]xf32> {64// CHECK: %[[VAL_2:.*]] = vector.shape_cast %[[VAL_0]] : vector<8x[2]x1xf32> to vector<8x[2]xf32>65// CHECK: %[[VAL_3:.*]] = vector.shape_cast %[[VAL_1]] : vector<1x8x[2]xf32> to vector<8x[2]xf32>66// CHECK: %[[VAL_4:.*]] = arith.mulf %[[VAL_2]], %[[VAL_3]] : vector<8x[2]xf32>67// CHECK: return %[[VAL_4]] : vector<8x[2]xf32>68 69// -----70 71func.func @fold_unit_dim_sitofp(%vec : vector<8x[2]x1xi8>) -> vector<8x[2]xf32> {72 %sc_vec_0 = vector.shape_cast %vec : vector<8x[2]x1xi8> to vector<1x8x[2]xi8>73 %add = arith.sitofp %sc_vec_0 : vector<1x8x[2]xi8> to vector<1x8x[2]xf32>74 %res = vector.shape_cast %add : vector<1x8x[2]xf32> to vector<8x[2]xf32>75 return %res : vector<8x[2]xf32>76}77 78// CHECK-LABEL: func.func @fold_unit_dim_sitofp(79// CHECK-SAME: %[[VAL_0:.*]]: vector<8x[2]x1xi8>) -> vector<8x[2]xf32> {80// CHECK: %[[VAL_1:.*]] = vector.shape_cast %[[VAL_0]] : vector<8x[2]x1xi8> to vector<8x[2]xi8>81// CHECK: %[[VAL_2:.*]] = arith.sitofp %[[VAL_1]] : vector<8x[2]xi8> to vector<8x[2]xf32>82// CHECK: return %[[VAL_2]] : vector<8x[2]xf32>83 84// -----85 86// All shape casts are folded away87 88func.func @fold_unit_dims_entirely(%vec_0 : vector<8xi32>,89 %vec_1 : vector<8xi32>,90 %vec_2 : vector<8xi32>) -> vector<8xi32> {91 %sc_vec_0 = vector.shape_cast %vec_0 : vector<8xi32> to vector<1x8xi32>92 %sc_vec_1 = vector.shape_cast %vec_1 : vector<8xi32> to vector<1x8xi32>93 %sc_vec_2 = vector.shape_cast %vec_2 : vector<8xi32> to vector<1x8xi32>94 %mul = arith.muli %sc_vec_0, %sc_vec_1 : vector<1x8xi32>95 %add = arith.addi %mul, %sc_vec_2 : vector<1x8xi32>96 %res = vector.shape_cast %add : vector<1x8xi32> to vector<8xi32>97 return %res : vector<8xi32>98}99 100// CHECK-LABEL: func.func @fold_unit_dims_entirely(101// CHECK-SAME: %[[VAL_0:.*]]: vector<8xi32>, %[[VAL_1:.*]]: vector<8xi32>,102// CHECK-SAME: %[[VAL_2:.*]]: vector<8xi32>) -> vector<8xi32> {103// CHECK: %[[VAL_3:.*]] = arith.muli %[[VAL_0]], %[[VAL_1]] : vector<8xi32>104// CHECK: %[[VAL_4:.*]] = arith.addi %[[VAL_3]], %[[VAL_2]] : vector<8xi32>105// CHECK: return %[[VAL_4]] : vector<8xi32>106 107// -----108 109func.func @fold_inner_unit_dim(%vec_0 : vector<8x1x3xf128>,110 %vec_1 : vector<1x8x3xf128>) -> vector<8x3xf128> {111 %sc_vec_1 = vector.shape_cast %vec_1 : vector<1x8x3xf128> to vector<8x1x3xf128>112 %mul = arith.mulf %vec_0, %sc_vec_1 : vector<8x1x3xf128>113 %res = vector.shape_cast %mul : vector<8x1x3xf128> to vector<8x3xf128>114 return %res : vector<8x3xf128>115}116 117// CHECK-LABEL: func.func @fold_inner_unit_dim(118// CHECK-SAME: %[[VAL_0:.*]]: vector<8x1x3xf128>,119// CHECK-SAME: %[[VAL_1:.*]]: vector<1x8x3xf128>) -> vector<8x3xf128> {120// CHECK: %[[VAL_2:.*]] = vector.shape_cast %[[VAL_0]] : vector<8x1x3xf128> to vector<8x3xf128>121// CHECK: %[[VAL_3:.*]] = vector.shape_cast %[[VAL_1]] : vector<1x8x3xf128> to vector<8x3xf128>122// CHECK: %[[VAL_4:.*]] = arith.mulf %[[VAL_2]], %[[VAL_3]] : vector<8x3xf128>123// CHECK: return %[[VAL_4]] : vector<8x3xf128>124 125// -----126 127func.func @fold_inner_unit_dim_scalable(%vec_0 : vector<8x1x[1]x3xf128>,128 %vec_1 : vector<1x8x[1]x3xf128>) -> vector<8x[1]x3xf128> {129 %sc_vec_1 = vector.shape_cast %vec_1 : vector<1x8x[1]x3xf128> to vector<8x1x[1]x3xf128>130 %mul = arith.mulf %vec_0, %sc_vec_1 : vector<8x1x[1]x3xf128>131 %res = vector.shape_cast %mul : vector<8x1x[1]x3xf128> to vector<8x[1]x3xf128>132 return %res : vector<8x[1]x3xf128>133}134 135// CHECK-LABEL: func.func @fold_inner_unit_dim_scalable(136// CHECK-SAME: %[[VAL_0:.*]]: vector<8x1x[1]x3xf128>,137// CHECK-SAME: %[[VAL_1:.*]]: vector<1x8x[1]x3xf128>) -> vector<8x[1]x3xf128> {138// CHECK: %[[VAL_2:.*]] = vector.shape_cast %[[VAL_0]] : vector<8x1x[1]x3xf128> to vector<8x[1]x3xf128>139// CHECK: %[[VAL_3:.*]] = vector.shape_cast %[[VAL_1]] : vector<1x8x[1]x3xf128> to vector<8x[1]x3xf128>140// CHECK: %[[VAL_4:.*]] = arith.mulf %[[VAL_2]], %[[VAL_3]] : vector<8x[1]x3xf128>141// CHECK: return %[[VAL_4]] : vector<8x[1]x3xf128>142 143// -----144 145func.func @fold_all_unit_dims(%vec: vector<1x1xf32>) -> vector<1xf32> {146 %0 = arith.mulf %vec, %vec : vector<1x1xf32>147 %res = vector.shape_cast %0 : vector<1x1xf32> to vector<1xf32>148 return %res : vector<1xf32>149}150 151// CHECK-LABEL: func.func @fold_all_unit_dims(152// CHECK-SAME: %[[VAL_0:.*]]: vector<1x1xf32>) -> vector<1xf32>153// CHECK: %[[VAL_1:.*]] = vector.shape_cast %[[VAL_0]] : vector<1x1xf32> to vector<1xf32>154// CHECK: %[[VAL_2:.*]] = vector.shape_cast %[[VAL_0]] : vector<1x1xf32> to vector<1xf32>155// CHECK: %[[VAL_3:.*]] = arith.mulf %[[VAL_1]], %[[VAL_2]] : vector<1xf32>156// CHECK: return %[[VAL_3]] : vector<1xf32>157 158///----------------------------------------------------------------------------------------159/// [Pattern: DropUnitDimsFromTransposeOp]160///----------------------------------------------------------------------------------------161 162func.func @transpose_with_internal_unit_dims(%vec: vector<1x1x4x[4]xf32>) -> vector<[4]x1x1x4xf32> {163 %res = vector.transpose %vec, [3, 0, 1, 2] : vector<1x1x4x[4]xf32> to vector<[4]x1x1x4xf32>164 return %res : vector<[4]x1x1x4xf32>165}166 167// CHECK-LABEL: func.func @transpose_with_internal_unit_dims(168// CHECK-SAME: %[[VEC:.*]]: vector<1x1x4x[4]xf32>)169// CHECK-NEXT: %[[DROP_DIMS:.*]] = vector.shape_cast %arg0 : vector<1x1x4x[4]xf32> to vector<4x[4]xf32>170// CHECK-NEXT: %[[TRANSPOSE:.*]] = vector.transpose %0, [1, 0] : vector<4x[4]xf32> to vector<[4]x4xf32>171// CHECK-NEXT: %[[RESTORE_DIMS:.*]] = vector.shape_cast %1 : vector<[4]x4xf32> to vector<[4]x1x1x4xf32>172// CHECK-NEXT: return %[[RESTORE_DIMS]] : vector<[4]x1x1x4xf32>173 174// -----175 176func.func @transpose_with_scalable_unit_dims(%vec: vector<[1]x1x2x4x1xf32>) -> vector<1x1x4x2x[1]xf32>177{178 %res = vector.transpose %vec, [4, 1, 3, 2, 0] : vector<[1]x1x2x4x1xf32> to vector<1x1x4x2x[1]xf32>179 return %res: vector<1x1x4x2x[1]xf32>180}181 182// CHECK-LABEL: func.func @transpose_with_scalable_unit_dims(183// CHECK-SAME: %[[VEC:.*]]: vector<[1]x1x2x4x1xf32>)184// CHECK-NEXT: %[[DROP_DIMS:.*]] = vector.shape_cast %[[VEC]] : vector<[1]x1x2x4x1xf32> to vector<[1]x2x4xf32>185// CHECK-NEXT: %[[TRANSPOSE:.*]] = vector.transpose %[[DROP_DIMS]], [2, 1, 0] : vector<[1]x2x4xf32> to vector<4x2x[1]xf32>186// CHECK-NEXT: %[[RESTORE_DIMS:.*]] = vector.shape_cast %[[TRANSPOSE]] : vector<4x2x[1]xf32> to vector<1x1x4x2x[1]xf32>187// CHECK-NEXT: return %[[RESTORE_DIMS]] : vector<1x1x4x2x[1]xf32>188 189// -----190 191func.func @transpose_with_all_unit_dims(%vec: vector<1x1x1xf32>) -> vector<1x1x1xf32> {192 %res = vector.transpose %vec, [0, 2, 1] : vector<1x1x1xf32> to vector<1x1x1xf32>193 return %res : vector<1x1x1xf32>194}195// The `vec` is returned because there are other flattening patterns that fold196// vector.shape_cast ops away.197// CHECK-LABEL: func.func @transpose_with_all_unit_dims198// CHECK-SAME: %[[VEC:.[a-zA-Z0-9]+]]199// CHECK-NEXT: return %[[VEC]]200 201// -----202 203func.func @negative_transpose_with_no_unit_dims(%vec: vector<4x2x3xf32>) -> vector<4x3x2xf32> {204 %res = vector.transpose %vec, [0, 2, 1] : vector<4x2x3xf32> to vector<4x3x2xf32>205 return %res : vector<4x3x2xf32>206}207 208// CHECK-LABEL: func.func @negative_transpose_with_no_unit_dims209// CHECK-NOT: vector.shape_cast210 211// -----212 213///----------------------------------------------------------------------------------------214/// [Pattern: DropUnitDimsFromScfForOp]215///----------------------------------------------------------------------------------------216 217func.func @scf_for_with_internal_unit_dims(%vec: vector<4x1x1x[4]xf32>) -> vector<4x1x1x[4]xf32> {218 %c0 = arith.constant 0 : index219 %c1 = arith.constant 1 : index220 %c4 = arith.constant 4 : index221 %res = scf.for %i = %c0 to %c4 step %c1 iter_args(%iter = %vec) -> vector<4x1x1x[4]xf32> {222 %s = math.sqrt %iter : vector<4x1x1x[4]xf32>223 scf.yield %s : vector<4x1x1x[4]xf32>224 }225 return %res : vector<4x1x1x[4]xf32>226}227 228// CHECK-LABEL: func.func @scf_for_with_internal_unit_dims229// CHECK-SAME: %[[VEC:[A-Za-z0-9]+]]: vector<4x1x1x[4]xf32>230// CHECK: %[[CAST:.+]] = vector.shape_cast %[[VEC]] : vector<4x1x1x[4]xf32> to vector<4x[4]xf32>231// CHECK: %[[LOOP:.+]] = scf.for {{.*}} iter_args(%[[ITER:.+]] = %[[CAST]])232// CHECK: %[[SQRT:.+]] = math.sqrt %[[ITER]] : vector<4x[4]xf32>233// CHECK: scf.yield %[[SQRT]]234// CHECK: %[[CASTBACK:.+]] = vector.shape_cast %[[LOOP]] : vector<4x[4]xf32> to vector<4x1x1x[4]xf32>235// CHECK: return %[[CASTBACK]]236 237// -----238 239func.func @scf_for_with_all_unit_dims(%vec: vector<1x1xf32>) -> vector<1x1xf32> {240 %c0 = arith.constant 0 : index241 %c1 = arith.constant 1 : index242 %c4 = arith.constant 4 : index243 %res = scf.for %i = %c0 to %c4 step %c1 iter_args(%iter = %vec) -> vector<1x1xf32> {244 %s = math.sqrt %iter : vector<1x1xf32>245 scf.yield %s : vector<1x1xf32>246 }247 return %res : vector<1x1xf32>248}249 250// CHECK-LABEL: func.func @scf_for_with_all_unit_dims251// CHECK-SAME: %[[VEC:[A-Za-z0-9]+]]: vector<1x1xf32>252// CHECK: %[[CAST:.+]] = vector.shape_cast %[[VEC]] : vector<1x1xf32> to vector<1xf32>253// CHECK: %[[LOOP:.+]] = scf.for {{.*}} iter_args(%[[ITER:.+]] = %[[CAST]])254// CHECK: %[[SQRT:.+]] = math.sqrt %[[ITER]] : vector<1xf32>255// CHECK: scf.yield %[[SQRT]]256// CHECK: %[[CASTBACK:.+]] = vector.shape_cast %[[LOOP]] : vector<1xf32> to vector<1x1xf32>257// CHECK: return %[[CASTBACK]]258 259// -----260 261func.func @scf_for_with_multiple_operands(%idx: index, %vec0: vector<1x4xf32>, %vec1: vector<1x4xf32>) -> vector<1x4xf32> {262 %c0 = arith.constant 0 : index263 %c1 = arith.constant 1 : index264 %c4 = arith.constant 4 : index265 %res:3 = scf.for %i = %c0 to %c4 step %c1266 iter_args(%id = %idx, %iter0 = %vec0, %iter1 = %vec1) -> (index, vector<1x4xf32>, vector<1x4xf32>) {267 %add = arith.addf %iter0, %iter1 : vector<1x4xf32>268 scf.yield %id, %add, %add : index, vector<1x4xf32>, vector<1x4xf32>269 }270 return %res#1 : vector<1x4xf32>271}272 273// CHECK-LABEL: func.func @scf_for_with_multiple_operands274// CHECK-SAME: %[[IDX:[A-Za-z0-9]+]]: index275// CHECK-SAME: %[[VEC0:[A-Za-z0-9]+]]: vector<1x4xf32>276// CHECK-SAME: %[[VEC1:[A-Za-z0-9]+]]: vector<1x4xf32>277// CHECK-DAG: %[[CAST0:.+]] = vector.shape_cast %[[VEC0]] : vector<1x4xf32> to vector<4xf32>278// CHECK-DAG: %[[CAST1:.+]] = vector.shape_cast %[[VEC1]] : vector<1x4xf32> to vector<4xf32>279// CHECK: %[[LOOP:.+]]:3 = scf.for280// CHECK-SAME: iter_args(%{{.*}} = %[[IDX]], %[[ITER0:.+]] = %[[CAST0]], %[[ITER1:.+]] = %[[CAST1]])281// CHECK: %[[ADD:.+]] = arith.addf %[[ITER0]], %[[ITER1]] : vector<4xf32>282// CHECK: scf.yield %{{.*}}, %[[ADD]], %[[ADD]]283// CHECK: %[[CASTBACK:.+]] = vector.shape_cast %[[LOOP]]#1 : vector<4xf32> to vector<1x4xf32>284// CHECK: return %[[CASTBACK]]285 286// -----287 288func.func @scf_for_with_scalable_unit_dims(%vec: vector<1x[1]xf32>) -> vector<1x[1]xf32> {289 %c0 = arith.constant 0 : index290 %c1 = arith.constant 1 : index291 %c4 = arith.constant 4 : index292 %res = scf.for %i = %c0 to %c4 step %c1 iter_args(%iter = %vec) -> vector<1x[1]xf32> {293 %s = math.sqrt %iter : vector<1x[1]xf32>294 scf.yield %s : vector<1x[1]xf32>295 }296 return %res : vector<1x[1]xf32>297}298 299// CHECK-LABEL: func.func @scf_for_with_scalable_unit_dims300// CHECK-SAME: %[[VEC:[A-Za-z0-9]+]]: vector<1x[1]xf32>301// CHECK: %[[CAST:.+]] = vector.shape_cast %[[VEC]] : vector<1x[1]xf32> to vector<[1]xf32>302// CHECK: %[[LOOP:.+]] = scf.for {{.*}} iter_args(%[[ITER:.+]] = %[[CAST]])303// CHECK: %[[SQRT:.+]] = math.sqrt %[[ITER]] : vector<[1]xf32>304// CHECK: scf.yield %[[SQRT]]305// CHECK: %[[CASTBACK:.+]] = vector.shape_cast %[[LOOP]] : vector<[1]xf32> to vector<1x[1]xf32>306// CHECK: return %[[CASTBACK]]307