853 lines · plain
1// RUN: mlir-opt %s -test-vector-sink-patterns -split-input-file | FileCheck %s2// RUN: mlir-opt -transform-preload-library='transform-library-paths=%p/vector-sink-transform.mlir' -transform-interpreter -split-input-file %s | FileCheck %s3 4//-----------------------------------------------------------------------------5// [Pattern: ReorderElementwiseOpsOnBroadcast]6//-----------------------------------------------------------------------------7 8// CHECK-LABEL: func.func @broadcast_scalar_with_bcast(9// CHECK-SAME: %[[ARG_0:.*]]: index, %[[ARG_1:.*]]: index) -> vector<1x4xindex> {10// CHECK: %[[ADD:.*]] = arith.addi %[[ARG_0]], %[[ARG_1]] : index11// CHECK: %[[BCAST:.*]] = vector.broadcast %[[ADD]] : index to vector<1x4xindex>12// CHECK: return %[[BCAST]] : vector<1x4xindex>13 14func.func @broadcast_scalar_with_bcast(%arg1: index, %arg2: index) -> vector<1x4xindex> {15 %0 = vector.broadcast %arg1 : index to vector<1x4xindex>16 %1 = vector.broadcast %arg2 : index to vector<1x4xindex>17 %2 = arith.addi %0, %1 : vector<1x4xindex>18 return %2 : vector<1x4xindex>19}20 21// CHECK-LABEL: func.func @broadcast_scalar_with_bcast_scalable(22// CHECK-SAME: %[[ARG_0:.*]]: index, %[[ARG_1:.*]]: index) -> vector<1x[4]xindex> {23// CHECK: %[[ADD:.*]] = arith.addi %[[ARG_0]], %[[ARG_1]] : index24// CHECK: %[[BCAST:.*]] = vector.broadcast %[[ADD]] : index to vector<1x[4]xindex>25// CHECK: return %[[BCAST]] : vector<1x[4]xindex>26 27func.func @broadcast_scalar_with_bcast_scalable(%arg1: index, %arg2: index) -> vector<1x[4]xindex> {28 %0 = vector.broadcast %arg1 : index to vector<1x[4]xindex>29 %1 = vector.broadcast %arg2 : index to vector<1x[4]xindex>30 %2 = arith.addi %0, %1 : vector<1x[4]xindex>31 return %2 : vector<1x[4]xindex>32}33 34// -----35 36// CHECK-LABEL: func.func @broadcast_scalar_with_bcast_and_splat(37// CHECK-SAME: %[[ARG1:.*]]: index,38// CHECK-SAME: %[[ARG2:.*]]: index) -> vector<1x4xindex> {39// CHECK: %[[ADD:.*]] = arith.addi %[[ARG1]], %[[ARG2]] : index40// CHECK: %[[BCAST:.*]] = vector.broadcast %[[ADD]] : index to vector<1x4xindex>41// CHECK: return %[[BCAST]] : vector<1x4xindex>42func.func @broadcast_scalar_with_bcast_and_splat(%arg1: index, %arg2: index) -> vector<1x4xindex> {43 %0 = vector.broadcast %arg1 : index to vector<1x4xindex>44 %1 = vector.broadcast %arg2 : index to vector<1x4xindex>45 %2 = arith.addi %0, %1 : vector<1x4xindex>46 return %2 : vector<1x4xindex>47}48 49// CHECK-LABEL: func.func @broadcast_scalar_with_bcast_and_splat_scalable(50// CHECK-SAME: %[[ARG1:.*]]: index,51// CHECK-SAME: %[[ARG2:.*]]: index) -> vector<1x[4]xindex> {52// CHECK: %[[ADD:.*]] = arith.addi %[[ARG1]], %[[ARG2]] : index53// CHECK: %[[BCAST:.*]] = vector.broadcast %[[ADD]] : index to vector<1x[4]xindex>54// CHECK: return %[[BCAST]] : vector<1x[4]xindex>55func.func @broadcast_scalar_with_bcast_and_splat_scalable(%arg1: index, %arg2: index) -> vector<1x[4]xindex> {56 %0 = vector.broadcast %arg1 : index to vector<1x[4]xindex>57 %1 = vector.broadcast %arg2 : index to vector<1x[4]xindex>58 %2 = arith.addi %0, %1 : vector<1x[4]xindex>59 return %2 : vector<1x[4]xindex>60}61 62// -----63 64// CHECK-LABEL: func.func @broadcast_vector(65// CHECK-SAME: %[[ARG_0:.*]]: vector<4xf32>,66// CHECK-SAME: %[[ARG_1:.*]]: vector<4xf32>) -> vector<3x4xf32> {67// CHECK: %[[ADDF:.*]] = arith.addf %[[ARG_0]], %[[ARG_1]] : vector<4xf32>68// CHECK: %[[BCAST:.*]] = vector.broadcast %[[ADDF]] : vector<4xf32> to vector<3x4xf32>69// CHECK: return %[[BCAST]] : vector<3x4xf32>70 71func.func @broadcast_vector(%arg1: vector<4xf32>, %arg2: vector<4xf32>) -> vector<3x4xf32> {72 %arg1_bcast = vector.broadcast %arg1 : vector<4xf32> to vector<3x4xf32>73 %arg2_bcast = vector.broadcast %arg2 : vector<4xf32> to vector<3x4xf32>74 %2 = arith.addf %arg1_bcast, %arg2_bcast : vector<3x4xf32>75 return %2 : vector<3x4xf32>76}77 78// CHECK-LABEL: func.func @broadcast_vector_scalable(79// CHECK-SAME: %[[ARG_0:.*]]: vector<[4]xf32>,80// CHECK-SAME: %[[ARG_1:.*]]: vector<[4]xf32>) -> vector<3x[4]xf32> {81// CHECK: %[[ADDF:.*]] = arith.addf %[[ARG_0]], %[[ARG_1]] : vector<[4]xf32>82// CHECK: %[[BCAST:.*]] = vector.broadcast %[[ADDF]] : vector<[4]xf32> to vector<3x[4]xf32>83// CHECK: return %[[BCAST]] : vector<3x[4]xf32>84 85func.func @broadcast_vector_scalable(%arg1: vector<[4]xf32>, %arg2: vector<[4]xf32>) -> vector<3x[4]xf32> {86 %arg1_bcast = vector.broadcast %arg1 : vector<[4]xf32> to vector<3x[4]xf32>87 %arg2_bcast = vector.broadcast %arg2 : vector<[4]xf32> to vector<3x[4]xf32>88 %2 = arith.addf %arg1_bcast, %arg2_bcast : vector<3x[4]xf32>89 return %2 : vector<3x[4]xf32>90}91 92// -----93 94// CHECK-LABEL: func.func @broadcast_scalar_and_vec(95// CHECK-SAME: %[[ARG1:.*]]: index,96// CHECK-SAME: %[[ARG2:.*]]: vector<4xindex>) -> vector<1x4xindex> {97// CHECK: %[[SPLAT:.*]] = vector.broadcast %[[ARG1]] : index to vector<1x4xindex>98// CHECK: %[[BCAST:.*]] = vector.broadcast %[[ARG2]] : vector<4xindex> to vector<1x4xindex>99// CHECK: %[[ADD:.*]] = arith.addi %[[SPLAT]], %[[BCAST]] : vector<1x4xindex>100// CHECK: return %[[ADD]] : vector<1x4xindex>101func.func @broadcast_scalar_and_vec(%arg1: index, %arg2: vector<4xindex>) -> vector<1x4xindex> {102 %0 = vector.broadcast %arg1 : index to vector<1x4xindex>103 %1 = vector.broadcast %arg2 : vector<4xindex> to vector<1x4xindex>104 %2 = arith.addi %0, %1 : vector<1x4xindex>105 return %2 : vector<1x4xindex>106}107 108// CHECK-LABEL: func.func @broadcast_scalar_and_vec_scalable(109// CHECK-SAME: %[[ARG1:.*]]: index,110// CHECK-SAME: %[[ARG2:.*]]: vector<[4]xindex>) -> vector<1x[4]xindex> {111// CHECK: %[[SPLAT:.*]] = vector.broadcast %[[ARG1]] : index to vector<1x[4]xindex>112// CHECK: %[[BCAST:.*]] = vector.broadcast %[[ARG2]] : vector<[4]xindex> to vector<1x[4]xindex>113// CHECK: %[[ADD:.*]] = arith.addi %[[SPLAT]], %[[BCAST]] : vector<1x[4]xindex>114// CHECK: return %[[ADD]] : vector<1x[4]xindex>115func.func @broadcast_scalar_and_vec_scalable(%arg1: index, %arg2: vector<[4]xindex>) -> vector<1x[4]xindex> {116 %0 = vector.broadcast %arg1 : index to vector<1x[4]xindex>117 %1 = vector.broadcast %arg2 : vector<[4]xindex> to vector<1x[4]xindex>118 %2 = arith.addi %0, %1 : vector<1x[4]xindex>119 return %2 : vector<1x[4]xindex>120}121 122// -----123 124// CHECK-LABEL: func.func @broadcast_vector_and_scalar(125// CHECK-SAME: %[[ARG_0:.*]]: i32,126// CHECK-SAME: %[[ARG_1:.*]]: vector<4xi32>) -> vector<4xi32> {127// CHECK: %[[BCAST:.*]] = vector.broadcast %[[ARG_0]] : i32 to vector<4xi32>128// CHECK: %[[ADD:.*]] = arith.addi %[[BCAST]], %[[ARG_1]] : vector<4xi32>129// CHECK: return %[[ADD]] : vector<4xi32>130 131func.func @broadcast_vector_and_scalar(%arg1: i32, %arg2: vector<4xi32>) -> vector<4xi32> {132 %arg1_bcast = vector.broadcast %arg1 : i32 to vector<4xi32>133 %2 = arith.addi %arg1_bcast, %arg2 : vector<4xi32>134 return %2 : vector<4xi32>135}136 137// CHECK-LABEL: func.func @broadcast_vector_and_scalar_scalable(138// CHECK-SAME: %[[ARG_0:.*]]: i32,139// CHECK-SAME: %[[ARG_1:.*]]: vector<[4]xi32>) -> vector<[4]xi32> {140// CHECK: %[[BCAST:.*]] = vector.broadcast %[[ARG_0]] : i32 to vector<[4]xi32>141// CHECK: %[[ADD:.*]] = arith.addi %[[BCAST]], %[[ARG_1]] : vector<[4]xi32>142// CHECK: return %[[ADD]] : vector<[4]xi32>143 144func.func @broadcast_vector_and_scalar_scalable(%arg1: i32, %arg2: vector<[4]xi32>) -> vector<[4]xi32> {145 %arg1_bcast = vector.broadcast %arg1 : i32 to vector<[4]xi32>146 %2 = arith.addi %arg1_bcast, %arg2 : vector<[4]xi32>147 return %2 : vector<[4]xi32>148}149 150// -----151 152#matmat_accesses = [153 affine_map<(i, j, k) -> (i, k)>,154 affine_map<(i, j, k) -> (k, j)>,155 affine_map<(i, j, k) -> (i, j)>156]157#matmat_trait = {158 indexing_maps = #matmat_accesses,159 iterator_types = ["parallel", "parallel", "reduction"]160}161 162// CHECK-LABEL: func.func @negative_not_elementwise163// CHECK-DAG: %[[F1:.*]] = arith.constant dense<1.000000e+00> : vector<2x2xf32>164// CHECK-DAG: %[[F2:.*]] = arith.constant dense<2.000000e+00> : vector<2x2xf32>165// CHECK-DAG: %[[F3:.*]] = arith.constant dense<3.000000e+00> : vector<2x2xf32>166// CHECK: %[[RES:.*]] = vector.contract {indexing_maps = [#map, #map1, #map2], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>} %[[F1]], %[[F2]], %[[F3]] : vector<2x2xf32>, vector<2x2xf32> into vector<2x2xf32>167func.func @negative_not_elementwise() -> vector<2x2xf32> {168 %f1 = arith.constant 1.0: f32169 %f2 = arith.constant 2.0: f32170 %f3 = arith.constant 3.0: f32171 172 %A = vector.broadcast %f1 : f32 to vector<2x2xf32>173 %B = vector.broadcast %f2 : f32 to vector<2x2xf32>174 %C = vector.broadcast %f3 : f32 to vector<2x2xf32>175 %res = vector.contract #matmat_trait %A, %B, %C176 : vector<2x2xf32>, vector<2x2xf32> into vector<2x2xf32>177 178 return %res : vector<2x2xf32>179}180 181// -----182 183// The source and the result for arith.cmp have different types184 185// CHECK-LABEL: func.func @source_and_result_mismatch(186// CHECK-SAME: %[[ARG0:.+]]: f32)187// CHECK: %[[COMPARE:.+]] = arith.cmpf uno, %[[ARG0]], %[[ARG0]]188// CHECK: %[[BROADCAST:.+]] = vector.broadcast %[[COMPARE]] : i1 to vector<1xi1>189// CHECK: return %[[BROADCAST]]190func.func @source_and_result_mismatch(%arg0 : f32) -> vector<1xi1> {191 %0 = vector.broadcast %arg0 : f32 to vector<1xf32>192 %1 = arith.cmpf uno, %0, %0 : vector<1xf32>193 return %1 : vector<1xi1>194}195 196// -----197 198// vector.fma only supports vectors - currently it's not possible to replace this with e.g.:199// %scalar_res = vector.fma %scalar_1, %scalar2200// %vec_res = vector.broadcast %scalar_res201//202// TODO: It should be possible to support this case203 204// CHECK-LABEL: func.func @negative_op_only_supports_vectors205 // CHECK: %[[BROADCAST:.+]] = vector.broadcast206 // CHECK: %[[RESULT:.+]] = vector.fma %[[BROADCAST]]207 // CHECK: return %[[RESULT]]208func.func @negative_op_only_supports_vectors(%arg0 : f32) -> vector<1xf32> {209 %0 = vector.broadcast %arg0 : f32 to vector<1xf32>210 %1 = vector.fma %0, %0, %0 : vector<1xf32>211 return %1 : vector<1xf32>212}213 214// -----215 216// CHECK-LABEL: func.func @broadcast_scalar_and_splat_const(217// CHECK-SAME: %[[ARG_0:.*]]: index) -> vector<1x4xindex> {218// CHECK: %[[NEW_CST:.*]] = arith.constant 2 : index219// CHECK: %[[ADD:.*]] = arith.addi %[[ARG_0]], %[[NEW_CST]] : index220// CHECK: %[[BCAST:.*]] = vector.broadcast %[[ADD]] : index to vector<1x4xindex>221// CHECK: return %[[BCAST]] : vector<1x4xindex>222 223func.func @broadcast_scalar_and_splat_const(%arg0: index) -> vector<1x4xindex> {224 %0 = vector.broadcast %arg0 : index to vector<1x4xindex>225 %cst = arith.constant dense<2> : vector<1x4xindex>226 %2 = arith.addi %0, %cst : vector<1x4xindex>227 return %2 : vector<1x4xindex>228}229 230// -----231 232// CHECK-LABEL: func.func @broadcast_scalar_and_splat_const_const_first(233// CHECK-SAME: %[[ARG_0:.*]]: index) -> vector<1x4xindex> {234// CHECK: %[[NEW_CST:.*]] = arith.constant 2 : index235// CHECK: %[[SUB:.*]] = arith.subi %[[NEW_CST]], %[[ARG_0]] : index236// CHECK: %[[BCAST:.*]] = vector.broadcast %[[SUB]] : index to vector<1x4xindex>237// CHECK: return %[[BCAST]] : vector<1x4xindex>238 239func.func @broadcast_scalar_and_splat_const_const_first(%arg0: index) -> vector<1x4xindex> {240 %0 = vector.broadcast %arg0 : index to vector<1x4xindex>241 %cst = arith.constant dense<2> : vector<1x4xindex>242 %2 = arith.subi %cst, %0 : vector<1x4xindex>243 return %2 : vector<1x4xindex>244}245 246// -----247 248// CHECK-LABEL: func.func @broadcast_vector_and_splat_const(249// CHECK-SAME: %[[ARG_0:.*]]: vector<4xf32>) -> vector<3x4xf32> {250// CHECK: %[[NEW_CST:.*]] = arith.constant dense<2.000000e+00> : vector<4xf32>251// CHECK: %[[ADD:.*]] = arith.mulf %[[ARG_0]], %[[NEW_CST]] : vector<4xf32>252// CHECK: %[[BCAST:.*]] = vector.broadcast %[[ADD]] : vector<4xf32> to vector<3x4xf32>253// CHECK: return %[[BCAST]] : vector<3x4xf32>254 255func.func @broadcast_vector_and_splat_const(%arg0: vector<4xf32>) -> vector<3x4xf32> {256 %0 = vector.broadcast %arg0 : vector<4xf32> to vector<3x4xf32>257 %cst = arith.constant dense<2.000000e+00> : vector<3x4xf32>258 %2 = arith.mulf %0, %cst : vector<3x4xf32>259 return %2 : vector<3x4xf32>260}261 262// -----263 264// CHECK-LABEL: func.func @negative_broadcast_with_non_splat_const(265// CHECK-SAME: %[[ARG_0:.*]]: index) -> vector<1x4xindex> {266// CHECK-DAG: %[[BCAST:.*]] = vector.broadcast %[[ARG_0]] : index to vector<1x4xindex>267// CHECK-DAG: %[[CST:.*]] = arith.constant dense<{{\[}}[0, 1, 2, 3]]> : vector<1x4xindex>268// CHECK: %[[ADD:.*]] = arith.addi %[[BCAST]], %[[CST]] : vector<1x4xindex>269// CHECK: return %[[ADD]] : vector<1x4xindex>270 271func.func @negative_broadcast_with_non_splat_const(%arg0: index) -> vector<1x4xindex> {272 %0 = vector.broadcast %arg0 : index to vector<1x4xindex>273 %cst = arith.constant dense<[[0, 1, 2, 3]]> : vector<1x4xindex>274 %2 = arith.addi %0, %cst : vector<1x4xindex>275 return %2 : vector<1x4xindex>276}277 278// -----279 280// CHECK-LABEL: func.func @broadcast_scalar_mixed_type(281// CHECK-SAME: %[[ARG_0:.*]]: f16) -> vector<1x4xf32> {282// CHECK: %[[EXTF:.*]] = arith.extf %[[ARG_0]] : f16 to f32283// CHECK: %[[BCAST:.*]] = vector.broadcast %[[EXTF]] : f32 to vector<1x4xf32>284// CHECK: return %[[BCAST]] : vector<1x4xf32>285 286func.func @broadcast_scalar_mixed_type(%arg0: f16) -> vector<1x4xf32> {287 %0 = vector.broadcast %arg0 : f16 to vector<1x4xf16>288 %1 = arith.extf %0 : vector<1x4xf16> to vector<1x4xf32>289 return %1 : vector<1x4xf32>290}291 292// -----293 294// CHECK-LABEL: func.func @broadcast_vector_mixed_type(295// CHECK-SAME: %[[ARG_0:.*]]: vector<4xf16>) -> vector<3x4xf32> {296// CHECK: %[[EXTF:.*]] = arith.extf %[[ARG_0]] : vector<4xf16> to vector<4xf32>297// CHECK: %[[BCAST:.*]] = vector.broadcast %[[EXTF]] : vector<4xf32> to vector<3x4xf32>298// CHECK: return %[[BCAST]] : vector<3x4xf32>299 300func.func @broadcast_vector_mixed_type(%arg0: vector<4xf16>) -> vector<3x4xf32> {301 %0 = vector.broadcast %arg0 : vector<4xf16> to vector<3x4xf16>302 %1 = arith.extf %0 : vector<3x4xf16> to vector<3x4xf32>303 return %1 : vector<3x4xf32>304}305 306// -----307 308// CHECK-LABEL: func.func @broadcast_scalar_and_splat_const_mixed_type(309// CHECK-SAME: %[[ARG_0:.*]]: f32) -> vector<1x4xf32> {310// CHECK: %[[NEW_CST:.*]] = arith.constant 3 : i32311// CHECK: %[[POW:.*]] = math.fpowi %[[ARG_0]], %[[NEW_CST]] : f32, i32312// CHECK: %[[BCAST:.*]] = vector.broadcast %[[POW]] : f32 to vector<1x4xf32>313// CHECK: return %[[BCAST]] : vector<1x4xf32>314 315func.func @broadcast_scalar_and_splat_const_mixed_type(%arg0: f32) -> vector<1x4xf32> {316 %0 = vector.broadcast %arg0 : f32 to vector<1x4xf32>317 %cst = arith.constant dense<3> : vector<1x4xi32>318 %2 = math.fpowi %0, %cst : vector<1x4xf32>, vector<1x4xi32>319 return %2 : vector<1x4xf32>320}321 322// -----323 324// CHECK-LABEL: func.func @broadcast_vector_and_splat_const_mixed_type(325// CHECK-SAME: %[[ARG_0:.*]]: vector<4xf32>) -> vector<3x4xf32> {326// CHECK: %[[NEW_CST:.*]] = arith.constant dense<3> : vector<4xi32>327// CHECK: %[[POW:.*]] = math.fpowi %[[ARG_0]], %[[NEW_CST]] : vector<4xf32>, vector<4xi32>328// CHECK: %[[BCAST:.*]] = vector.broadcast %[[POW]] : vector<4xf32> to vector<3x4xf32>329// CHECK: return %[[BCAST]] : vector<3x4xf32>330 331func.func @broadcast_vector_and_splat_const_mixed_type(%arg0: vector<4xf32>) -> vector<3x4xf32> {332 %0 = vector.broadcast %arg0 : vector<4xf32> to vector<3x4xf32>333 %cst = arith.constant dense<3> : vector<3x4xi32>334 %2 = math.fpowi %0, %cst : vector<3x4xf32>, vector<3x4xi32>335 return %2 : vector<3x4xf32>336}337 338//===----------------------------------------------------------------------===//339// [Pattern: ReorderCastOpsOnBroadcast]340//341// Reorder casting ops and vector ops. The casting ops have almost identical342// pattern, so only arith.extsi op is tested.343//===----------------------------------------------------------------------===//344 345// -----346 347func.func @broadcast_vector_extsi(%a : vector<4xi8>) -> vector<2x4xi32> {348 // CHECK: %[[EXT:.+]] = arith.extsi %{{.+}} : vector<4xi8> to vector<4xi32>349 // CHECK: vector.broadcast %[[EXT:.+]] : vector<4xi32> to vector<2x4xi32>350 %b = vector.broadcast %a : vector<4xi8> to vector<2x4xi8>351 %r = arith.extsi %b : vector<2x4xi8> to vector<2x4xi32>352 return %r : vector<2x4xi32>353}354 355// -----356 357func.func @broadcast_vector_extsi_scalable(%a : vector<[4]xi8>) -> vector<2x[4]xi32> {358 // CHECK: %[[EXT:.+]] = arith.extsi %{{.+}} : vector<[4]xi8> to vector<[4]xi32>359 // CHECK: vector.broadcast %[[EXT:.+]] : vector<[4]xi32> to vector<2x[4]xi32>360 %b = vector.broadcast %a : vector<[4]xi8> to vector<2x[4]xi8>361 %r = arith.extsi %b : vector<2x[4]xi8> to vector<2x[4]xi32>362 return %r : vector<2x[4]xi32>363}364 365// -----366 367func.func @broadcast_scalar_extsi(%a : i8) -> vector<2x4xi32> {368 // CHECK: %[[EXT:.+]] = arith.extsi %{{.+}} : i8 to i32369 // CHECK: vector.broadcast %[[EXT]] : i32 to vector<2x4xi32>370 %b = vector.broadcast %a : i8 to vector<2x4xi8>371 %r = arith.extsi %b : vector<2x4xi8> to vector<2x4xi32>372 return %r : vector<2x4xi32>373}374 375// -----376 377func.func @broadcast_scalar_extsi_scalable(%a : i8) -> vector<2x[4]xi32> {378 // CHECK: %[[EXT:.+]] = arith.extsi %{{.+}} : i8 to i32379 // CHECK: vector.broadcast %[[EXT]] : i32 to vector<2x[4]xi32>380 %b = vector.broadcast %a : i8 to vector<2x[4]xi8>381 %r = arith.extsi %b : vector<2x[4]xi8> to vector<2x[4]xi32>382 return %r : vector<2x[4]xi32>383}384 385//===----------------------------------------------------------------------===//386// [Pattern: ReorderElementwiseOpsOnTranspose]387//===----------------------------------------------------------------------===//388 389func.func @transpose_extsi(%a : vector<4x2xi8>) -> vector<2x4xi32> {390 // CHECK: %[[EXT:.+]] = arith.extsi %{{.+}} : vector<4x2xi8> to vector<4x2xi32>391 // CHECK: vector.transpose %[[EXT]], [1, 0] : vector<4x2xi32> to vector<2x4xi32>392 %b = vector.transpose %a, [1, 0]: vector<4x2xi8> to vector<2x4xi8>393 %r = arith.extsi %b : vector<2x4xi8> to vector<2x4xi32>394 return %r : vector<2x4xi32>395}396 397// -----398 399func.func @transpose_extsi_scalable(%a : vector<[4]x2xi8>) -> vector<2x[4]xi32> {400 // CHECK: %[[EXT:.+]] = arith.extsi %{{.+}} : vector<[4]x2xi8> to vector<[4]x2xi32>401 // CHECK: vector.transpose %[[EXT]], [1, 0] : vector<[4]x2xi32> to vector<2x[4]xi32>402 %b = vector.transpose %a, [1, 0]: vector<[4]x2xi8> to vector<2x[4]xi8>403 %r = arith.extsi %b : vector<2x[4]xi8> to vector<2x[4]xi32>404 return %r : vector<2x[4]xi32>405}406 407// -----408 409// CHECK-LABEL: func @transpose_elementwise_same_type410// CHECK-SAME: (%[[A:.+]]: vector<4x2xf32>, %[[B:.+]]: vector<4x2xf32>)411// CHECK: %[[ADD:.+]] = arith.addf %[[A]], %[[B]] : vector<4x2xf32>412// CHECK: %[[T:.+]] = vector.transpose %[[ADD]], [1, 0]413// CHECK: return %[[T]]414 415func.func @transpose_elementwise_same_type(%a : vector<4x2xf32>, %b : vector<4x2xf32>) -> vector<2x4xf32> {416 %at = vector.transpose %a, [1, 0]: vector<4x2xf32> to vector<2x4xf32>417 %bt = vector.transpose %b, [1, 0]: vector<4x2xf32> to vector<2x4xf32>418 %r = arith.addf %at, %bt : vector<2x4xf32>419 return %r : vector<2x4xf32>420}421 422// -----423 424// CHECK-LABEL: func @transpose_elementwise_same_type_scalable425// CHECK-SAME: (%[[A:.+]]: vector<[4]x2xf32>, %[[B:.+]]: vector<[4]x2xf32>)426// CHECK: %[[ADD:.+]] = arith.addf %[[A]], %[[B]] : vector<[4]x2xf32>427// CHECK: %[[T:.+]] = vector.transpose %[[ADD]], [1, 0]428// CHECK: return %[[T]]429 430func.func @transpose_elementwise_same_type_scalable(%a : vector<[4]x2xf32>, %b : vector<[4]x2xf32>) -> vector<2x[4]xf32> {431 %at = vector.transpose %a, [1, 0]: vector<[4]x2xf32> to vector<2x[4]xf32>432 %bt = vector.transpose %b, [1, 0]: vector<[4]x2xf32> to vector<2x[4]xf32>433 %r = arith.addf %at, %bt : vector<2x[4]xf32>434 return %r : vector<2x[4]xf32>435}436 437// -----438 439// CHECK-LABEL: func @transpose_elementwise_diff_operand_types440// CHECK-SAME: (%[[COND:.+]]: vector<4x2xi1>, %[[A:.+]]: vector<4x2xf32>, %[[B:.+]]: vector<4x2xf32>)441// CHECK: %[[S:.+]] = arith.select %[[COND]], %[[A]], %[[B]] : vector<4x2xi1>, vector<4x2xf32>442// CHECK: %[[T:.+]] = vector.transpose %[[S]], [1, 0] : vector<4x2xf32> to vector<2x4xf32>443// CHECK: return %[[T]]444func.func @transpose_elementwise_diff_operand_types(%cond: vector<4x2xi1>, %a : vector<4x2xf32>, %b : vector<4x2xf32>) -> vector<2x4xf32> {445 %condt = vector.transpose %cond, [1, 0]: vector<4x2xi1> to vector<2x4xi1>446 %at = vector.transpose %a, [1, 0]: vector<4x2xf32> to vector<2x4xf32>447 %bt = vector.transpose %b, [1, 0]: vector<4x2xf32> to vector<2x4xf32>448 %r = arith.select %condt, %at, %bt : vector<2x4xi1>, vector<2x4xf32>449 return %r : vector<2x4xf32>450}451 452// -----453 454// CHECK-LABEL: func @transpose_elementwise_diff_operand_types_scalable455// CHECK-SAME: (%[[COND:.+]]: vector<[4]x2xi1>, %[[A:.+]]: vector<[4]x2xf32>, %[[B:.+]]: vector<[4]x2xf32>)456// CHECK: %[[S:.+]] = arith.select %[[COND]], %[[A]], %[[B]] : vector<[4]x2xi1>, vector<[4]x2xf32>457// CHECK: %[[T:.+]] = vector.transpose %[[S]], [1, 0] : vector<[4]x2xf32> to vector<2x[4]xf32>458// CHECK: return %[[T]]459func.func @transpose_elementwise_diff_operand_types_scalable(%cond: vector<[4]x2xi1>, %a : vector<[4]x2xf32>, %b : vector<[4]x2xf32>) -> vector<2x[4]xf32> {460 %condt = vector.transpose %cond, [1, 0]: vector<[4]x2xi1> to vector<2x[4]xi1>461 %at = vector.transpose %a, [1, 0]: vector<[4]x2xf32> to vector<2x[4]xf32>462 %bt = vector.transpose %b, [1, 0]: vector<[4]x2xf32> to vector<2x[4]xf32>463 %r = arith.select %condt, %at, %bt : vector<2x[4]xi1>, vector<2x[4]xf32>464 return %r : vector<2x[4]xf32>465}466 467// -----468 469// CHECK-LABEL: func @transpose_elementwise_diff_operand_result_type470// CHECK-SAME: (%[[A:.+]]: vector<4x2xf32>, %[[B:.+]]: vector<4x2xf32>)471// CHECK: %[[CMP:.+]] = arith.cmpf olt, %[[A]], %[[B]] : vector<4x2xf32>472// CHECK: %[[T:.+]] = vector.transpose %[[CMP]], [1, 0] : vector<4x2xi1> to vector<2x4xi1>473// CHECK: return %[[T]]474func.func @transpose_elementwise_diff_operand_result_type(%a : vector<4x2xf32>, %b : vector<4x2xf32>) -> vector<2x4xi1> {475 %at = vector.transpose %a, [1, 0]: vector<4x2xf32> to vector<2x4xf32>476 %bt = vector.transpose %b, [1, 0]: vector<4x2xf32> to vector<2x4xf32>477 %r = arith.cmpf olt, %at, %bt : vector<2x4xf32>478 return %r : vector<2x4xi1>479}480 481// -----482 483// CHECK-LABEL: func @transpose_elementwise_diff_operand_result_type_scalable484// CHECK-SAME: (%[[A:.+]]: vector<[4]x2xf32>, %[[B:.+]]: vector<[4]x2xf32>)485// CHECK: %[[CMP:.+]] = arith.cmpf olt, %[[A]], %[[B]] : vector<[4]x2xf32>486// CHECK: %[[T:.+]] = vector.transpose %[[CMP]], [1, 0] : vector<[4]x2xi1> to vector<2x[4]xi1>487// CHECK: return %[[T]]488func.func @transpose_elementwise_diff_operand_result_type_scalable(%a : vector<[4]x2xf32>, %b : vector<[4]x2xf32>) -> vector<2x[4]xi1> {489 %at = vector.transpose %a, [1, 0]: vector<[4]x2xf32> to vector<2x[4]xf32>490 %bt = vector.transpose %b, [1, 0]: vector<[4]x2xf32> to vector<2x[4]xf32>491 %r = arith.cmpf olt, %at, %bt : vector<2x[4]xf32>492 return %r : vector<2x[4]xi1>493}494 495// -----496 497// CHECK-LABEL: func @transpose_elementwise_splat_constant498// CHECK-SAME: (%[[A:.+]]: vector<4x6x3x2xf32>)499// CHECK: %[[B:.+]] = arith.constant dense<5.000000e+00> : vector<4x6x3x2xf32>500// CHECK: %[[ADD:.+]] = arith.addf %[[A]], %[[B]] : vector<4x6x3x2xf32>501// CHECK: %[[T:.+]] = vector.transpose %[[ADD]], [1, 0, 3, 2] : vector<4x6x3x2xf32> to vector<6x4x2x3xf32>502// CHECK: return %[[T:.+]] : vector<6x4x2x3xf32>503 504func.func @transpose_elementwise_splat_constant(%a : vector<4x6x3x2xf32>) -> vector<6x4x2x3xf32> {505 %b = arith.constant dense<5.0> : vector<6x4x2x3xf32>506 %at = vector.transpose %a, [1, 0, 3, 2]: vector<4x6x3x2xf32> to vector<6x4x2x3xf32>507 %r = arith.addf %at, %b : vector<6x4x2x3xf32>508 return %r : vector<6x4x2x3xf32>509}510 511// -----512 513// CHECK-LABEL: func @transpose_elementwise_splat_constant_scalable514// CHECK-SAME: (%[[A:.+]]: vector<[4]x6x3x2xf32>)515// CHECK: %[[B:.+]] = arith.constant dense<5.000000e+00> : vector<[4]x6x3x2xf32>516// CHECK: %[[ADD:.+]] = arith.addf %[[A]], %[[B]] : vector<[4]x6x3x2xf32>517// CHECK: %[[T:.+]] = vector.transpose %[[ADD]], [1, 0, 3, 2] : vector<[4]x6x3x2xf32> to vector<6x[4]x2x3xf32>518// CHECK: return %[[T:.+]] : vector<6x[4]x2x3xf32>519 520func.func @transpose_elementwise_splat_constant_scalable(%a : vector<[4]x6x3x2xf32>) -> vector<6x[4]x2x3xf32> {521 %b = arith.constant dense<5.0> : vector<6x[4]x2x3xf32>522 %at = vector.transpose %a, [1, 0, 3, 2]: vector<[4]x6x3x2xf32> to vector<6x[4]x2x3xf32>523 %r = arith.addf %at, %b : vector<6x[4]x2x3xf32>524 return %r : vector<6x[4]x2x3xf32>525}526 527// -----528 529// CHECK-LABEL: func @transpose_elementwise_diff_map530// CHECK: vector.transpose531// CHECK: vector.transpose532// CHECK: arith.addf533func.func @transpose_elementwise_diff_map(%a : vector<4x6x3x2xf32>, %b: vector<6x2x4x3xf32>) -> vector<6x4x2x3xf32> {534 %at = vector.transpose %a, [1, 0, 3, 2]: vector<4x6x3x2xf32> to vector<6x4x2x3xf32>535 %bt = vector.transpose %b, [0, 2, 1, 3]: vector<6x2x4x3xf32> to vector<6x4x2x3xf32>536 %r = arith.addf %at, %bt : vector<6x4x2x3xf32>537 return %r : vector<6x4x2x3xf32>538}539 540// -----541 542// CHECK-LABEL: func @transpose_elementwise_diff_map_scalable543// CHECK: vector.transpose544// CHECK: vector.transpose545// CHECK: arith.addf546func.func @transpose_elementwise_diff_map_scalable(%a : vector<[4]x6x3x2xf32>, %b: vector<6x2x[4]x3xf32>) -> vector<6x[4]x2x3xf32> {547 %at = vector.transpose %a, [1, 0, 3, 2]: vector<[4]x6x3x2xf32> to vector<6x[4]x2x3xf32>548 %bt = vector.transpose %b, [0, 2, 1, 3]: vector<6x2x[4]x3xf32> to vector<6x[4]x2x3xf32>549 %r = arith.addf %at, %bt : vector<6x[4]x2x3xf32>550 return %r : vector<6x[4]x2x3xf32>551}552 553// -----554 555//-----------------------------------------------------------------------------556// [Pattern: ExtractOpFromElementwise]557//-----------------------------------------------------------------------------558 559// CHECK-LABEL: @extract_elementwise_scalar560// CHECK-SAME: (%[[ARG0:.*]]: vector<4xf32>, %[[ARG1:.*]]: vector<4xf32>)561func.func @extract_elementwise_scalar(%arg0: vector<4xf32>, %arg1: vector<4xf32>) -> f32 {562// CHECK: %[[EXT0:.*]] = vector.extract %[[ARG0]][1] : f32 from vector<4xf32>563// CHECK: %[[EXT1:.*]] = vector.extract %[[ARG1]][1] : f32 from vector<4xf32>564// CHECK: %[[RES:.*]] = arith.addf %[[EXT0]], %[[EXT1]] : f32565// CHECK: return %[[RES]] : f32566 %0 = arith.addf %arg0, %arg1 : vector<4xf32>567 %1 = vector.extract %0[1] : f32 from vector<4xf32>568 return %1 : f32569}570 571// CHECK-LABEL: @extract_elementwise_arg_res_different_types572// CHECK-SAME: (%[[ARG0:.*]]: vector<4xindex>)573func.func @extract_elementwise_arg_res_different_types(%arg0: vector<4xindex>) -> i64 {574// CHECK: %[[EXT:.*]] = vector.extract %[[ARG0]][1] : index from vector<4xindex>575// CHECK: %[[RES:.*]] = arith.index_cast %[[EXT]] : index to i64576// CHECK: return %[[RES]] : i64577 %0 = arith.index_cast %arg0: vector<4xindex> to vector<4xi64>578 %1 = vector.extract %0[1] : i64 from vector<4xi64>579 return %1 : i64580}581 582// CHECK-LABEL: @extract_elementwise_vec583// CHECK-SAME: (%[[ARG0:.*]]: vector<2x4xf32>, %[[ARG1:.*]]: vector<2x4xf32>)584func.func @extract_elementwise_vec(%arg0: vector<2x4xf32>, %arg1: vector<2x4xf32>) -> vector<4xf32> {585// CHECK: %[[EXT0:.*]] = vector.extract %[[ARG0]][1] : vector<4xf32> from vector<2x4xf32>586// CHECK: %[[EXT1:.*]] = vector.extract %[[ARG1]][1] : vector<4xf32> from vector<2x4xf32>587// CHECK: %[[RES:.*]] = arith.addf %[[EXT0]], %[[EXT1]] : vector<4xf32>588// CHECK: return %[[RES]] : vector<4xf32>589 %0 = arith.addf %arg0, %arg1 : vector<2x4xf32>590 %1 = vector.extract %0[1] : vector<4xf32> from vector<2x4xf32>591 return %1 : vector<4xf32>592}593 594// CHECK-LABEL: @negative_extract_elementwise_no_single_use595// CHECK-SAME: (%[[ARG0:.*]]: vector<4xf32>, %[[ARG1:.*]]: vector<4xf32>)596func.func @negative_extract_elementwise_no_single_use(%arg0: vector<4xf32>, %arg1: vector<4xf32>) -> (f32, vector<4xf32>) {597// Do not propagate extract, as elementwise has other uses.598// CHECK: %[[ELT:.*]] = arith.addf %[[ARG0]], %[[ARG1]] : vector<4xf32>599// CHECK: %[[EXT:.*]] = vector.extract %[[ELT]][1] : f32 from vector<4xf32>600// CHECK: return %[[EXT]], %[[ELT]] : f32, vector<4xf32>601 %0 = arith.addf %arg0, %arg1 : vector<4xf32>602 %1 = vector.extract %0[1] : f32 from vector<4xf32>603 return %1, %0 : f32, vector<4xf32>604}605 606// CHECK-LABEL: @negative_extract_elementwise_not_one_res607// CHECK-SAME: (%[[ARG0:.*]]: vector<4xi32>, %[[ARG1:.*]]: vector<4xi32>)608func.func @negative_extract_elementwise_not_one_res(%arg0: vector<4xi32>, %arg1: vector<4xi32>) -> i32 {609// Do not propagate extract, as elementwise has more than 1 result.610// CHECK: %[[LOW:.*]], %[[HIGH:.*]] = arith.mulsi_extended %[[ARG0]], %[[ARG1]] : vector<4xi32>611// CHECK: %[[EXT:.*]] = vector.extract %[[LOW]][1] : i32 from vector<4xi32>612// CHECK: return %[[EXT]] : i32613 %low, %hi = arith.mulsi_extended %arg0, %arg1 : vector<4xi32>614 %1 = vector.extract %low[1] : i32 from vector<4xi32>615 return %1 : i32616}617 618// CHECK-LABEL: @negative_extract_not_elementwise619// CHECK-SAME: (%[[ARG0:.*]]: vector<4xi64>)620func.func @negative_extract_not_elementwise(%arg0: vector<4xi64>) -> i64 {621// `test.increment` is not an elemewise op.622// CHECK: %[[INC:.*]] = test.increment %[[ARG0]] : vector<4xi64>623// CHECK: %[[RES:.*]] = vector.extract %[[INC]][1] : i64 from vector<4xi64>624// CHECK: return %[[RES]] : i64625 %0 = test.increment %arg0: vector<4xi64>626 %1 = vector.extract %0[1] : i64 from vector<4xi64>627 return %1 : i64628}629 630// CHECK-LABEL: @negative_extract_vec_fma631// CHECK-SAME: (%[[ARG0:.*]]: vector<4xf32>, %[[ARG1:.*]]: vector<4xf32>, %[[ARG2:.*]]: vector<4xf32>)632func.func @negative_extract_vec_fma(%arg0: vector<4xf32>, %arg1: vector<4xf32>, %arg2: vector<4xf32>) -> f32 {633// `vector.fma` doesn't suppport scalars.634// CHECK: %[[FMA:.*]] = vector.fma %[[ARG0]], %[[ARG1]], %[[ARG2]] : vector<4xf32>635// CHECK: %[[RES:.*]] = vector.extract %[[FMA]][1] : f32 from vector<4xf32>636// CHECK: return %[[RES]] : f32637 %0 = vector.fma %arg0, %arg1, %arg2: vector<4xf32>638 %1 = vector.extract %0[1] : f32 from vector<4xf32>639 return %1 : f32640}641 642// CHECK-LABEL: @negative_extract_dynamic_pos643func.func @negative_extract_dynamic_pos(%arg0: vector<4xf32>, %arg1 : vector<4xf32>, %idx : vector<4xindex>) -> f32 {644 // CHECK-NOT: vector.extract645 // CHECK: arith.addf %{{.*}}, %{{.*}} : vector<4xf32>646 // CHECK: vector.extract647 // CHECK: vector.extract648 %0 = arith.addf %arg0, %arg1 : vector<4xf32>649 %1 = vector.extract %idx[0] : index from vector<4xindex>650 %2 = vector.extract %0[%1] : f32 from vector<4xf32>651 return %2 : f32652}653 654//-----------------------------------------------------------------------------655// [Pattern: ExtractOpFromLoad]656//-----------------------------------------------------------------------------657 658// CHECK-LABEL: @extract_load_scalar659// CHECK-SAME: (%[[ARG0:.*]]: memref<?xf32>, %[[ARG1:.*]]: index)660func.func @extract_load_scalar(%arg0: memref<?xf32>, %arg1: index) -> f32 {661// CHECK: %[[RES:.*]] = memref.load %[[ARG0]][%[[ARG1]]] : memref<?xf32>662// CHECK: return %[[RES]] : f32663 %0 = vector.load %arg0[%arg1] : memref<?xf32>, vector<4xf32>664 %1 = vector.extract %0[0] : f32 from vector<4xf32>665 return %1 : f32666}667 668// CHECK-LABEL: @extract_load_index669// CHECK-SAME: (%[[ARG0:.*]]: memref<?xindex>, %[[ARG1:.*]]: index)670func.func @extract_load_index(%arg0: memref<?xindex>, %arg1: index) -> index {671// CHECK: %[[RES:.*]] = memref.load %[[ARG0]][%[[ARG1]]] : memref<?xindex>672// CHECK: return %[[RES]] : index673 %0 = vector.load %arg0[%arg1] : memref<?xindex>, vector<4xindex>674 %1 = vector.extract %0[0] : index from vector<4xindex>675 return %1 : index676}677 678// CHECK-LABEL: @extract_load_scalar_non_zero_off679// CHECK-SAME: (%[[ARG0:.*]]: memref<?xf32>, %[[ARG1:.*]]: index)680func.func @extract_load_scalar_non_zero_off(%arg0: memref<?xf32>, %arg1: index) -> f32 {681// CHECK: %[[C1:.*]] = arith.constant 1 : index682// CHECK: %[[OFF:.*]] = arith.addi %[[ARG1]], %[[C1]] overflow<nsw> : index683// CHECK: %[[RES:.*]] = memref.load %[[ARG0]][%[[OFF]]] : memref<?xf32>684// CHECK: return %[[RES]] : f32685 %0 = vector.load %arg0[%arg1] : memref<?xf32>, vector<4xf32>686 %1 = vector.extract %0[1] : f32 from vector<4xf32>687 return %1 : f32688}689 690// CHECK-LABEL: @extract_load_scalar_dyn_off691// CHECK-SAME: (%[[ARG0:.*]]: memref<?xf32>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index)692func.func @extract_load_scalar_dyn_off(%arg0: memref<?xf32>, %arg1: index, %arg2: index) -> f32 {693// CHECK: %[[OFF:.*]] = arith.addi %[[ARG1]], %[[ARG2]] overflow<nsw> : index694// CHECK: %[[RES:.*]] = memref.load %[[ARG0]][%[[OFF]]] : memref<?xf32>695// CHECK: return %[[RES]] : f32696 %0 = vector.load %arg0[%arg1] : memref<?xf32>, vector<4xf32>697 %1 = vector.extract %0[%arg2] : f32 from vector<4xf32>698 return %1 : f32699}700 701// CHECK-LABEL: @extract_load_vec_non_zero_off702// CHECK-SAME: (%[[ARG0:.*]]: memref<?x?xf32>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index)703func.func @extract_load_vec_non_zero_off(%arg0: memref<?x?xf32>, %arg1: index, %arg2: index) -> vector<4xf32> {704// CHECK: %[[C1:.*]] = arith.constant 1 : index705// CHECK: %[[OFF:.*]] = arith.addi %[[ARG1]], %[[C1]] overflow<nsw> : index706// CHECK: %[[RES:.*]] = vector.load %[[ARG0]][%[[OFF]], %[[ARG2]]] : memref<?x?xf32>, vector<4xf32>707// CHECK: return %[[RES]] : vector<4xf32>708 %0 = vector.load %arg0[%arg1, %arg2] : memref<?x?xf32>, vector<2x4xf32>709 %1 = vector.extract %0[1] : vector<4xf32> from vector<2x4xf32>710 return %1 : vector<4xf32>711}712 713// CHECK-LABEL: @extract_load_scalar_non_zero_off_2d_src_memref714// CHECK-SAME: (%[[ARG0:.*]]: memref<?x?xf32>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index)715func.func @extract_load_scalar_non_zero_off_2d_src_memref(%arg0: memref<?x?xf32>, %arg1: index, %arg2: index) -> f32 {716// CHECK: %[[C1:.*]] = arith.constant 1 : index717// CHECK: %[[OFF:.*]] = arith.addi %[[ARG2]], %[[C1]] overflow<nsw> : index718// CHECK: %[[RES:.*]] = memref.load %[[ARG0]][%[[ARG1]], %[[OFF]]] : memref<?x?xf32>719// CHECK: return %[[RES]] : f32720 %0 = vector.load %arg0[%arg1, %arg2] : memref<?x?xf32>, vector<4xf32>721 %1 = vector.extract %0[1] : f32 from vector<4xf32>722 return %1 : f32723}724 725// CHECK-LABEL: @extract_load_vec_high_rank726// CHECK-SAME: (%[[ARG0:.*]]: memref<?x?x?xf32>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index, %[[ARG3:.*]]: index)727func.func @extract_load_vec_high_rank(%arg0: memref<?x?x?xf32>, %arg1: index, %arg2: index, %arg3: index) -> vector<4xf32> {728// CHECK: %[[C1:.*]] = arith.constant 1 : index729// CHECK: %[[OFF:.*]] = arith.addi %[[ARG2]], %[[C1]] overflow<nsw> : index730// CHECK: %[[RES:.*]] = vector.load %[[ARG0]][%[[ARG1]], %[[OFF]], %[[ARG3]]] : memref<?x?x?xf32>, vector<4xf32>731// CHECK: return %[[RES]] : vector<4xf32>732 %0 = vector.load %arg0[%arg1, %arg2, %arg3] : memref<?x?x?xf32>, vector<2x4xf32>733 %1 = vector.extract %0[1] : vector<4xf32> from vector<2x4xf32>734 return %1 : vector<4xf32>735}736 737// CHECK-LABEL: @negative_extract_load_scalar_from_memref_of_vec738// CHECK-SAME: (%[[ARG0:.*]]: memref<?xvector<4xf32>>, %[[ARG1:.*]]: index)739func.func @negative_extract_load_scalar_from_memref_of_vec(%arg0: memref<?xvector<4xf32>>, %arg1: index) -> f32 {740// CHECK: %[[RES:.*]] = vector.load %[[ARG0]][%[[ARG1]]] : memref<?xvector<4xf32>>, vector<4xf32>741// CHECK: %[[EXT:.*]] = vector.extract %[[RES]][0] : f32 from vector<4xf32>742// CHECK: return %[[EXT]] : f32743 %0 = vector.load %arg0[%arg1] : memref<?xvector<4xf32>>, vector<4xf32>744 %1 = vector.extract %0[0] : f32 from vector<4xf32>745 return %1 : f32746}747 748// CHECK-LABEL: @negative_extract_load_scalar_from_memref_of_i1749// CHECK-SAME: (%[[ARG0:.*]]: memref<?xi1>, %[[ARG1:.*]]: index)750func.func @negative_extract_load_scalar_from_memref_of_i1(%arg0: memref<?xi1>, %arg1: index) -> i1 {751// Subbyte types are tricky, ignore them for now.752// CHECK: %[[RES:.*]] = vector.load %[[ARG0]][%[[ARG1]]] : memref<?xi1>, vector<8xi1>753// CHECK: %[[EXT:.*]] = vector.extract %[[RES]][0] : i1 from vector<8xi1>754// CHECK: return %[[EXT]] : i1755 %0 = vector.load %arg0[%arg1] : memref<?xi1>, vector<8xi1>756 %1 = vector.extract %0[0] : i1 from vector<8xi1>757 return %1 : i1758}759 760// CHECK-LABEL: @negative_extract_load_no_single_use761// CHECK-SAME: (%[[ARG0:.*]]: memref<?xf32>, %[[ARG1:.*]]: index)762func.func @negative_extract_load_no_single_use(%arg0: memref<?xf32>, %arg1: index) -> (f32, vector<4xf32>) {763// CHECK: %[[RES:.*]] = vector.load %[[ARG0]][%[[ARG1]]] : memref<?xf32>, vector<4xf32>764// CHECK: %[[EXT:.*]] = vector.extract %[[RES]][0] : f32 from vector<4xf32>765// CHECK: return %[[EXT]], %[[RES]] : f32, vector<4xf32>766 %0 = vector.load %arg0[%arg1] : memref<?xf32>, vector<4xf32>767 %1 = vector.extract %0[0] : f32 from vector<4xf32>768 return %1, %0 : f32, vector<4xf32>769}770 771// CHECK-LABEL: @negative_extract_load_scalable772// CHECK-SAME: (%[[ARG0:.*]]: memref<?xf32>, %[[ARG1:.*]]: index)773func.func @negative_extract_load_scalable(%arg0: memref<?xf32>, %arg1: index) -> f32 {774// CHECK: %[[RES:.*]] = vector.load %[[ARG0]][%[[ARG1]]] : memref<?xf32>, vector<[1]xf32>775// CHECK: %[[EXT:.*]] = vector.extract %[[RES]][0] : f32 from vector<[1]xf32>776// CHECK: return %[[EXT]] : f32777 %0 = vector.load %arg0[%arg1] : memref<?xf32>, vector<[1]xf32>778 %1 = vector.extract %0[0] : f32 from vector<[1]xf32>779 return %1 : f32780}781 782//-----------------------------------------------------------------------------783// [Pattern: StoreOpFromSplatOrBroadcast]784//-----------------------------------------------------------------------------785 786// CHECK-LABEL: @store_splat787// CHECK-SAME: (%[[ARG0:.*]]: memref<?xf32>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: f32)788func.func @store_splat(%arg0: memref<?xf32>, %arg1: index, %arg2: f32) {789// CHECK: memref.store %[[ARG2]], %[[ARG0]][%[[ARG1]]] : memref<?xf32>790 %0 = vector.broadcast %arg2 : f32 to vector<1xf32>791 vector.store %0, %arg0[%arg1] : memref<?xf32>, vector<1xf32>792 return793}794 795// CHECK-LABEL: @store_broadcast796// CHECK-SAME: (%[[ARG0:.*]]: memref<?xf32>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: f32)797func.func @store_broadcast(%arg0: memref<?xf32>, %arg1: index, %arg2: f32) {798// CHECK: memref.store %[[ARG2]], %[[ARG0]][%[[ARG1]]] : memref<?xf32>799 %0 = vector.broadcast %arg2 : f32 to vector<1xf32>800 vector.store %0, %arg0[%arg1] : memref<?xf32>, vector<1xf32>801 return802}803 804// CHECK-LABEL: @store_broadcast_1d_to_2d805// CHECK-SAME: (%[[ARG0:.*]]: memref<?x?xf32>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index, %[[ARG3:.*]]: vector<1xf32>)806func.func @store_broadcast_1d_to_2d(%arg0: memref<?x?xf32>, %arg1: index, %arg2: index, %arg3: vector<1xf32>) {807// CHECK: vector.store %[[ARG3]], %[[ARG0]][%[[ARG1]], %[[ARG2]]] : memref<?x?xf32>, vector<1xf32>808 %0 = vector.broadcast %arg3 : vector<1xf32> to vector<1x1xf32>809 vector.store %0, %arg0[%arg1, %arg2] : memref<?x?xf32>, vector<1x1xf32>810 return811}812 813// CHECK-LABEL: @negative_store_scalable814// CHECK-SAME: (%[[ARG0:.*]]: memref<?xf32>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: f32)815func.func @negative_store_scalable(%arg0: memref<?xf32>, %arg1: index, %arg2: f32) {816// CHECK: %[[RES:.*]] = vector.broadcast %[[ARG2]] : f32 to vector<[1]xf32>817// CHECK: vector.store %[[RES]], %[[ARG0]][%[[ARG1]]] : memref<?xf32>, vector<[1]xf32>818 %0 = vector.broadcast %arg2 : f32 to vector<[1]xf32>819 vector.store %0, %arg0[%arg1] : memref<?xf32>, vector<[1]xf32>820 return821}822 823// CHECK-LABEL: @negative_store_memref_of_vec824// CHECK-SAME: (%[[ARG0:.*]]: memref<?xvector<1xf32>>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: f32)825func.func @negative_store_memref_of_vec(%arg0: memref<?xvector<1xf32>>, %arg1: index, %arg2: f32) {826// CHECK: %[[RES:.*]] = vector.broadcast %[[ARG2]] : f32 to vector<1xf32>827// CHECK: vector.store %[[RES]], %[[ARG0]][%[[ARG1]]] : memref<?xvector<1xf32>>, vector<1xf32>828 %0 = vector.broadcast %arg2 : f32 to vector<1xf32>829 vector.store %0, %arg0[%arg1] : memref<?xvector<1xf32>>, vector<1xf32>830 return831}832 833// CHECK-LABEL: @negative_store_more_than_one_element834// CHECK-SAME: (%[[ARG0:.*]]: memref<?xf32>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: f32)835func.func @negative_store_more_than_one_element(%arg0: memref<?xf32>, %arg1: index, %arg2: f32) {836// CHECK: %[[RES:.*]] = vector.broadcast %[[ARG2]] : f32 to vector<4xf32>837// CHECK: vector.store %[[RES]], %[[ARG0]][%[[ARG1]]] : memref<?xf32>, vector<4xf32>838 %0 = vector.broadcast %arg2 : f32 to vector<4xf32>839 vector.store %0, %arg0[%arg1] : memref<?xf32>, vector<4xf32>840 return841}842 843// CHECK-LABEL: @negative_store_no_single_use844// CHECK-SAME: (%[[ARG0:.*]]: memref<?xf32>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: f32)845func.func @negative_store_no_single_use(%arg0: memref<?xf32>, %arg1: index, %arg2: f32) -> vector<1xf32> {846// CHECK: %[[RES:.*]] = vector.broadcast %[[ARG2]] : f32 to vector<1xf32>847// CHECK: vector.store %[[RES]], %[[ARG0]][%[[ARG1]]] : memref<?xf32>, vector<1xf32>848// CHECK: return %[[RES:.*]] : vector<1xf32>849 %0 = vector.broadcast %arg2 : f32 to vector<1xf32>850 vector.store %0, %arg0[%arg1] : memref<?xf32>, vector<1xf32>851 return %0 : vector<1xf32>852}853