1313 lines · plain
1// RUN: mlir-opt %s -convert-vector-to-arm-sme -split-input-file -allow-unregistered-dialect | FileCheck %s2 3//===----------------------------------------------------------------------===//4// vector.transfer_read5//===----------------------------------------------------------------------===//6 7// CHECK-LABEL: @transfer_read_2d_i88// CHECK: arm_sme.tile_load %{{.*}}[{{.*}}] : memref<?x?xi8>, vector<[16]x[16]xi8>9func.func @transfer_read_2d_i8(%src : memref<?x?xi8>) {10 %c0 = arith.constant 0 : index11 %pad = arith.constant 0 : i812 %0 = vector.transfer_read %src[%c0, %c0], %pad {in_bounds = [true, true]} : memref<?x?xi8>, vector<[16]x[16]xi8>13 "prevent.dce"(%0) : (vector<[16]x[16]xi8>) -> ()14 return15}16 17// -----18 19// CHECK-LABEL: @transfer_read_2d_i1620// CHECK: arm_sme.tile_load %{{.*}}[{{.*}}] : memref<?x?xi16>, vector<[8]x[8]xi16>21func.func @transfer_read_2d_i16(%src : memref<?x?xi16>) {22 %c0 = arith.constant 0 : index23 %pad = arith.constant 0 : i1624 %0 = vector.transfer_read %src[%c0, %c0], %pad {in_bounds = [true, true]} : memref<?x?xi16>, vector<[8]x[8]xi16>25 "prevent.dce"(%0) : (vector<[8]x[8]xi16>) -> ()26 return27}28 29// -----30 31// CHECK-LABEL: @transfer_read_2d_i3232// CHECK: arm_sme.tile_load %{{.*}}[{{.*}}] : memref<?x?xi32>, vector<[4]x[4]xi32>33func.func @transfer_read_2d_i32(%src : memref<?x?xi32>) {34 %c0 = arith.constant 0 : index35 %pad = arith.constant 0 : i3236 %0 = vector.transfer_read %src[%c0, %c0], %pad {in_bounds = [true, true]} : memref<?x?xi32>, vector<[4]x[4]xi32>37 "prevent.dce"(%0) : (vector<[4]x[4]xi32>) -> ()38 return39}40 41// -----42 43// CHECK-LABEL: @transfer_read_2d_i6444// CHECK: arm_sme.tile_load %{{.*}}[{{.*}}] : memref<?x?xi64>, vector<[2]x[2]xi64>45func.func @transfer_read_2d_i64(%src : memref<?x?xi64>) {46 %c0 = arith.constant 0 : index47 %pad = arith.constant 0 : i6448 %0 = vector.transfer_read %src[%c0, %c0], %pad {in_bounds = [true, true]} : memref<?x?xi64>, vector<[2]x[2]xi64>49 "prevent.dce"(%0) : (vector<[2]x[2]xi64>) -> ()50 return51}52 53// -----54 55// CHECK-LABEL: @transfer_read_2d_i12856// CHECK: arm_sme.tile_load %{{.*}}[{{.*}}] : memref<?x?xi128>, vector<[1]x[1]xi128>57func.func @transfer_read_2d_i128(%src : memref<?x?xi128>) {58 %c0 = arith.constant 0 : index59 %pad = arith.constant 0 : i12860 %0 = vector.transfer_read %src[%c0, %c0], %pad {in_bounds = [true, true]} : memref<?x?xi128>, vector<[1]x[1]xi128>61 "prevent.dce"(%0) : (vector<[1]x[1]xi128>) -> ()62 return63}64 65// -----66 67// CHECK-LABEL: @transfer_read_2d_f1668// CHECK: arm_sme.tile_load %{{.*}}[{{.*}}] : memref<?x?xf16>, vector<[8]x[8]xf16>69func.func @transfer_read_2d_f16(%src : memref<?x?xf16>) {70 %c0 = arith.constant 0 : index71 %pad = arith.constant 0.0 : f1672 %0 = vector.transfer_read %src[%c0, %c0], %pad {in_bounds = [true, true]} : memref<?x?xf16>, vector<[8]x[8]xf16>73 "prevent.dce"(%0) : (vector<[8]x[8]xf16>) -> ()74 return75}76 77// -----78 79// CHECK-LABEL: @transfer_read_2d_bf1680// CHECK: arm_sme.tile_load %{{.*}}[{{.*}}] : memref<?x?xbf16>, vector<[8]x[8]xbf16>81func.func @transfer_read_2d_bf16(%src : memref<?x?xbf16>) {82 %c0 = arith.constant 0 : index83 %pad = arith.constant 0.0 : bf1684 %0 = vector.transfer_read %src[%c0, %c0], %pad {in_bounds = [true, true]} : memref<?x?xbf16>, vector<[8]x[8]xbf16>85 "prevent.dce"(%0) : (vector<[8]x[8]xbf16>) -> ()86 return87}88 89// -----90 91// CHECK-LABEL: @transfer_read_2d_f3292// CHECK: arm_sme.tile_load %{{.*}}[{{.*}}] : memref<?x?xf32>, vector<[4]x[4]xf32>93func.func @transfer_read_2d_f32(%src : memref<?x?xf32>) {94 %c0 = arith.constant 0 : index95 %pad = arith.constant 0.0 : f3296 %0 = vector.transfer_read %src[%c0, %c0], %pad {in_bounds = [true, true]} : memref<?x?xf32>, vector<[4]x[4]xf32>97 "prevent.dce"(%0) : (vector<[4]x[4]xf32>) -> ()98 return99}100 101// -----102 103// CHECK-LABEL: @transfer_read_2d_f64104// CHECK: arm_sme.tile_load %{{.*}}[{{.*}}] : memref<?x?xf64>, vector<[2]x[2]xf64>105func.func @transfer_read_2d_f64(%src : memref<?x?xf64>) {106 %c0 = arith.constant 0 : index107 %pad = arith.constant 0.0 : f64108 %0 = vector.transfer_read %src[%c0, %c0], %pad {in_bounds = [true, true]} : memref<?x?xf64>, vector<[2]x[2]xf64>109 "prevent.dce"(%0) : (vector<[2]x[2]xf64>) -> ()110 return111}112 113// -----114 115// CHECK-LABEL: @transfer_read_2d_with_mask_i16116// CHECK: arm_sme.tile_load %{{.*}}[{{.*}}], {{.*}}, {{.*}} : memref<?x?xi16>, vector<[8]x[8]xi16>117func.func @transfer_read_2d_with_mask_i16(%src : memref<?x?xi16>, %mask : vector<[8]x[8]xi1>) {118 %c0 = arith.constant 0 : index119 %pad = arith.constant 0 : i16120 %0 = vector.transfer_read %src[%c0, %c0], %pad, %mask {in_bounds = [true, true]} : memref<?x?xi16>, vector<[8]x[8]xi16>121 "prevent.dce"(%0) : (vector<[8]x[8]xi16>) -> ()122 return123}124 125// -----126 127/// in-flight transpose128 129// CHECK-LABEL: @transfer_read_2d_transpose_i8130// CHECK: arm_sme.tile_load {{.*}} layout<vertical> : memref<?x?xi8>, vector<[16]x[16]xi8>131func.func @transfer_read_2d_transpose_i8(%src : memref<?x?xi8>) {132 %c0 = arith.constant 0 : index133 %pad = arith.constant 0 : i8134 %0 = vector.transfer_read %src[%c0, %c0], %pad {permutation_map = affine_map<(d0, d1) -> (d1, d0)>, in_bounds = [true, true]} : memref<?x?xi8>, vector<[16]x[16]xi8>135 "prevent.dce"(%0) : (vector<[16]x[16]xi8>) -> ()136 return137}138 139// -----140 141// CHECK-LABEL: @transfer_read_2d_transpose_with_mask_f32142// CHECK: arm_sme.tile_load {{.*}} layout<vertical> : memref<?x?xf32>, vector<[4]x[4]xf32>143func.func @transfer_read_2d_transpose_with_mask_f32(%src : memref<?x?xf32>, %mask : vector<[4]x[4]xi1>) {144 %c0 = arith.constant 0 : index145 %pad = arith.constant 0.0 : f32146 %0 = vector.transfer_read %src[%c0, %c0], %pad, %mask {permutation_map = affine_map<(d0, d1) -> (d1, d0)>, in_bounds = [true, true]} : memref<?x?xf32>, vector<[4]x[4]xf32>147 "prevent.dce"(%0) : (vector<[4]x[4]xf32>) -> ()148 return149}150 151// -----152 153// CHECK-LABEL: @fold_transpose_into_load154// CHECK-NOT: arm_sme.tile_store155// CHECK: arm_sme.tile_load {{.*}} layout<vertical> : memref<?x?xf32>, vector<[4]x[4]xf32>156// CHECK-NOT: arm_sme.tile_store157func.func @fold_transpose_into_load(%src : memref<?x?xf32>) {158 %c0 = arith.constant 0 : index159 %pad = arith.constant 0.0 : f32160 %0 = vector.transfer_read %src[%c0, %c0], %pad {in_bounds = [true, true]} : memref<?x?xf32>, vector<[4]x[4]xf32>161 %1 = vector.transpose %0, [1, 0] : vector<[4]x[4]xf32> to vector<[4]x[4]xf32>162 "prevent.dce"(%1) : (vector<[4]x[4]xf32>) -> ()163}164 165// -----166 167/// Transposes with more than a single use cannot be folded into load and will168/// instead be transposed via memory.169 170// CHECK-LABEL: @fold_transpose_into_load_multi_use171// CHECK: arm_sme.tile_load {{.*}} : memref<?x?xf32>, vector<[4]x[4]xf32>172// CHECK: arm_sme.tile_store {{.*}} : memref<?x?xf32>, vector<[4]x[4]xf32>173// CHECK: %[[TILE_TRANSPOSED_VIA_MEM:.*]] = arm_sme.tile_load {{.*}} layout<vertical> : memref<?x?xf32>, vector<[4]x[4]xf32>174// CHECK: "prevent.dce"(%[[TILE_TRANSPOSED_VIA_MEM]]) : (vector<[4]x[4]xf32>) -> ()175func.func @fold_transpose_into_load_multi_use(%src : memref<?x?xf32>) {176 %c0 = arith.constant 0 : index177 %pad = arith.constant 0.0 : f32178 %0 = vector.transfer_read %src[%c0, %c0], %pad {in_bounds = [true, true]} : memref<?x?xf32>, vector<[4]x[4]xf32>179 "test.some_use"(%0) : (vector<[4]x[4]xf32>) -> ()180 %1 = vector.transpose %0, [1, 0] : vector<[4]x[4]xf32> to vector<[4]x[4]xf32>181 "prevent.dce"(%1) : (vector<[4]x[4]xf32>) -> ()182}183 184// -----185 186//===----------------------------------------------------------------------===//187// vector.transfer_write188//===----------------------------------------------------------------------===//189 190// CHECK-LABEL: func.func @transfer_write_2d_i8(191// CHECK-SAME: %[[VECTOR:.*]]: vector<[16]x[16]xi8>,192// CHECK-SAME: %[[DEST:.*]]: memref<?x?xi8>) {193// CHECK: %[[C0:.*]] = arith.constant 0 : index194// CHECK: arm_sme.tile_store %[[VECTOR]], %[[DEST]]{{\[}}%[[C0]], %[[C0]]] : memref<?x?xi8>, vector<[16]x[16]xi8>195func.func @transfer_write_2d_i8(%vector : vector<[16]x[16]xi8>, %dest : memref<?x?xi8>) {196 %c0 = arith.constant 0 : index197 vector.transfer_write %vector, %dest[%c0, %c0] {in_bounds = [true, true]} : vector<[16]x[16]xi8>, memref<?x?xi8>198 return199}200 201// -----202 203// CHECK-LABEL: func.func @transfer_write_2d_i16(204// CHECK-SAME: %[[VECTOR:.*]]: vector<[8]x[8]xi16>,205// CHECK-SAME: %[[DEST:.*]]: memref<?x?xi16>) {206// CHECK: %[[C0:.*]] = arith.constant 0 : index207// CHECK: arm_sme.tile_store %[[VECTOR]], %[[DEST]]{{\[}}%[[C0]], %[[C0]]] : memref<?x?xi16>, vector<[8]x[8]xi16>208func.func @transfer_write_2d_i16(%vector : vector<[8]x[8]xi16>, %dest : memref<?x?xi16>) {209 %c0 = arith.constant 0 : index210 vector.transfer_write %vector, %dest[%c0, %c0] {in_bounds = [true, true]} : vector<[8]x[8]xi16>, memref<?x?xi16>211 return212}213 214// -----215 216// CHECK-LABEL: func.func @transfer_write_2d_i32(217// CHECK-SAME: %[[VECTOR:.*]]: vector<[4]x[4]xi32>,218// CHECK-SAME: %[[DEST:.*]]: memref<?x?xi32>) {219// CHECK: %[[C0:.*]] = arith.constant 0 : index220// CHECK: arm_sme.tile_store %[[VECTOR]], %[[DEST]]{{\[}}%[[C0]], %[[C0]]] : memref<?x?xi32>, vector<[4]x[4]xi32>221func.func @transfer_write_2d_i32(%vector : vector<[4]x[4]xi32>, %dest : memref<?x?xi32>) {222 %c0 = arith.constant 0 : index223 vector.transfer_write %vector, %dest[%c0, %c0] {in_bounds = [true, true]} : vector<[4]x[4]xi32>, memref<?x?xi32>224 return225}226 227// -----228 229// CHECK-LABEL: func.func @transfer_write_2d_i64(230// CHECK-SAME: %[[VECTOR:.*]]: vector<[2]x[2]xi64>,231// CHECK-SAME: %[[DEST:.*]]: memref<?x?xi64>) {232// CHECK: %[[C0:.*]] = arith.constant 0 : index233// CHECK: arm_sme.tile_store %[[VECTOR]], %[[DEST]]{{\[}}%[[C0]], %[[C0]]] : memref<?x?xi64>, vector<[2]x[2]xi64>234func.func @transfer_write_2d_i64(%vector : vector<[2]x[2]xi64>, %dest : memref<?x?xi64>) {235 %c0 = arith.constant 0 : index236 vector.transfer_write %vector, %dest[%c0, %c0] {in_bounds = [true, true]} : vector<[2]x[2]xi64>, memref<?x?xi64>237 return238}239 240// -----241 242// CHECK-LABEL: func.func @transfer_write_2d_f16(243// CHECK-SAME: %[[VECTOR:.*]]: vector<[8]x[8]xf16>,244// CHECK-SAME: %[[DEST:.*]]: memref<?x?xf16>) {245// CHECK: %[[C0:.*]] = arith.constant 0 : index246// CHECK: arm_sme.tile_store %[[VECTOR]], %[[DEST]]{{\[}}%[[C0]], %[[C0]]] : memref<?x?xf16>, vector<[8]x[8]xf16>247func.func @transfer_write_2d_f16(%vector : vector<[8]x[8]xf16>, %dest : memref<?x?xf16>) {248 %c0 = arith.constant 0 : index249 vector.transfer_write %vector, %dest[%c0, %c0] {in_bounds = [true, true]} : vector<[8]x[8]xf16>, memref<?x?xf16>250 return251}252 253// -----254 255// CHECK-LABEL: func.func @transfer_write_2d_bf16(256// CHECK-SAME: %[[VECTOR:.*]]: vector<[8]x[8]xbf16>,257// CHECK-SAME: %[[DEST:.*]]: memref<?x?xbf16>) {258// CHECK: %[[C0:.*]] = arith.constant 0 : index259// CHECK: arm_sme.tile_store %[[VECTOR]], %[[DEST]]{{\[}}%[[C0]], %[[C0]]] : memref<?x?xbf16>, vector<[8]x[8]xbf16>260func.func @transfer_write_2d_bf16(%vector : vector<[8]x[8]xbf16>, %dest : memref<?x?xbf16>) {261 %c0 = arith.constant 0 : index262 vector.transfer_write %vector, %dest[%c0, %c0] {in_bounds = [true, true]} : vector<[8]x[8]xbf16>, memref<?x?xbf16>263 return264}265 266// -----267 268// CHECK-LABEL: func.func @transfer_write_2d_f32(269// CHECK-SAME: %[[VECTOR:.*]]: vector<[4]x[4]xf32>,270// CHECK-SAME: %[[DEST:.*]]: memref<?x?xf32>) {271// CHECK: %[[C0:.*]] = arith.constant 0 : index272// CHECK: arm_sme.tile_store %[[VECTOR]], %[[DEST]]{{\[}}%[[C0]], %[[C0]]] : memref<?x?xf32>, vector<[4]x[4]xf32>273func.func @transfer_write_2d_f32(%vector : vector<[4]x[4]xf32>, %dest : memref<?x?xf32>) {274 %c0 = arith.constant 0 : index275 vector.transfer_write %vector, %dest[%c0, %c0] {in_bounds = [true, true]} : vector<[4]x[4]xf32>, memref<?x?xf32>276 return277}278 279// -----280 281// CHECK-LABEL: func.func @transfer_write_2d_f64(282// CHECK-SAME: %[[VECTOR:.*]]: vector<[2]x[2]xf64>,283// CHECK-SAME: %[[DEST:.*]]: memref<?x?xf64>) {284// CHECK: %[[C0:.*]] = arith.constant 0 : index285// CHECK: arm_sme.tile_store %[[VECTOR]], %[[DEST]]{{\[}}%[[C0]], %[[C0]]] : memref<?x?xf64>, vector<[2]x[2]xf64>286func.func @transfer_write_2d_f64(%vector : vector<[2]x[2]xf64>, %dest : memref<?x?xf64>) {287 %c0 = arith.constant 0 : index288 vector.transfer_write %vector, %dest[%c0, %c0] {in_bounds = [true, true]} : vector<[2]x[2]xf64>, memref<?x?xf64>289 return290}291 292// -----293 294// CHECK-LABEL: func.func @transfer_write_2d_with_mask_f64(295// CHECK-SAME: %[[VECTOR:.*]]: vector<[2]x[2]xf64>,296// CHECK-SAME: %[[DEST:.*]]: memref<?x?xf64>,297// CHECK-SAME: %[[MASK:.*]]: vector<[2]x[2]xi1>) {298// CHECK: %[[C0:.*]] = arith.constant 0 : index299// CHECK: arm_sme.tile_store %[[VECTOR]], %[[DEST]]{{\[}}%[[C0]], %[[C0]]], %[[MASK]] : memref<?x?xf64>, vector<[2]x[2]xf64>300func.func @transfer_write_2d_with_mask_f64(%vector : vector<[2]x[2]xf64>, %dest : memref<?x?xf64>, %mask : vector<[2]x[2]xi1>) {301 %c0 = arith.constant 0 : index302 vector.transfer_write %vector, %dest[%c0, %c0], %mask {in_bounds = [true, true]} : vector<[2]x[2]xf64>, memref<?x?xf64>303 return304}305 306// -----307 308/// in-flight transpose via vertical store.309 310// CHECK-LABEL: func.func @transfer_write_2d_transpose_i64(311// CHECK-SAME: %[[VECTOR:.*]]: vector<[2]x[2]xi64>,312// CHECK-SAME: %[[DEST:.*]]: memref<?x?xi64>) {313// CHECK: %[[C0:.*]] = arith.constant 0 : index314// CHECK: arm_sme.tile_store %[[VECTOR]], %[[DEST]]{{\[}}%[[C0]], %[[C0]]] layout<vertical> : memref<?x?xi64>, vector<[2]x[2]xi64>315func.func @transfer_write_2d_transpose_i64(%vector : vector<[2]x[2]xi64>, %dest : memref<?x?xi64>) {316 %c0 = arith.constant 0 : index317 vector.transfer_write %vector, %dest[%c0, %c0] {permutation_map = affine_map<(d0, d1) -> (d1, d0)>, in_bounds = [true, true]} : vector<[2]x[2]xi64>, memref<?x?xi64>318 return319}320 321// -----322 323/// in-flight transpose via vertical store with mask.324 325// CHECK-LABEL: func.func @transfer_write_2d_transpose_with_mask_bf16(326// CHECK-SAME: %[[VECTOR:.*]]: vector<[8]x[8]xbf16>,327// CHECK-SAME: %[[DEST:.*]]: memref<?x?xbf16>,328// CHECK-SAME: %[[MASK:.*]]: vector<[8]x[8]xi1>) {329// CHECK: %[[C0:.*]] = arith.constant 0 : index330// CHECK: arm_sme.tile_store %[[VECTOR]], %[[DEST]]{{\[}}%[[C0]], %[[C0]]], %[[MASK]] layout<vertical> : memref<?x?xbf16>, vector<[8]x[8]xbf16>331func.func @transfer_write_2d_transpose_with_mask_bf16(%vector : vector<[8]x[8]xbf16>, %dest : memref<?x?xbf16>, %mask : vector<[8]x[8]xi1>) {332 %c0 = arith.constant 0 : index333 vector.transfer_write %vector, %dest[%c0, %c0], %mask {permutation_map = affine_map<(d0, d1) -> (d1, d0)>, in_bounds = [true, true]} : vector<[8]x[8]xbf16>, memref<?x?xbf16>334 return335}336 337// -----338 339// CHECK-LABEL: func.func @transfer_write_slice(340// CHECK-SAME: %[[VECTOR:.*]]: vector<[4]x[4]xf32>,341// CHECK-SAME: %[[DEST:.*]]: memref<?x?xf32>,342// CHECK-SAME: %[[INDEX:.*]]: index) {343// CHECK: %[[C0:.*]] = arith.constant 0 : index344// CHECK: %[[MASK:.*]] = arith.constant dense<true> : vector<[4]xi1>345// CHECK: arm_sme.store_tile_slice %[[VECTOR]], %[[INDEX]], %[[MASK]], %[[DEST]][%[[INDEX]], %[[C0]]] : memref<?x?xf32>, vector<[4]xi1>, vector<[4]x[4]xf32>346func.func @transfer_write_slice(%vector: vector<[4]x[4]xf32>, %dest : memref<?x?xf32>, %slice_index: index) {347 %c0 = arith.constant 0 : index348 %slice = vector.extract %vector[%slice_index] : vector<[4]xf32> from vector<[4]x[4]xf32>349 vector.transfer_write %slice, %dest[%slice_index, %c0] { in_bounds = [true] }: vector<[4]xf32>, memref<?x?xf32>350 return351}352 353// -----354 355// CHECK-LABEL: func.func @transfer_write_slice_with_mask(356// CHECK-SAME: %[[VECTOR:.*]]: vector<[4]x[4]xf32>,357// CHECK-SAME: %[[DEST:.*]]: memref<?x?xf32>,358// CHECK-SAME: %[[MASK:.*]]: vector<[4]xi1>,359// CHECK-SAME: %[[INDEX:.*]]: index) {360// CHECK: %[[C0:.*]] = arith.constant 0 : index361// CHECK: arm_sme.store_tile_slice %[[VECTOR]], %[[INDEX]], %[[MASK]], %[[DEST]][%[[INDEX]], %[[C0]]] : memref<?x?xf32>, vector<[4]xi1>, vector<[4]x[4]xf32>362func.func @transfer_write_slice_with_mask(%vector: vector<[4]x[4]xf32>, %dest : memref<?x?xf32>, %mask: vector<[4]xi1>, %slice_index: index) {363 %c0 = arith.constant 0 : index364 %slice = vector.extract %vector[%slice_index] : vector<[4]xf32> from vector<[4]x[4]xf32>365 vector.transfer_write %slice, %dest[%slice_index, %c0], %mask { in_bounds = [true] }: vector<[4]xf32>, memref<?x?xf32>366 return367}368 369// -----370 371// CHECK-LABEL: func.func @transfer_write_vertical_slice372// CHECK: arm_sme.store_tile_slice {{.*}} layout<vertical>373func.func @transfer_write_vertical_slice(%vector: vector<[4]x[4]xf32>, %dest : memref<?x?xf32>, %slice_index: index) {374 %c0 = arith.constant 0 : index375 %slice = arm_sme.extract_tile_slice %vector[%slice_index] layout<vertical>376 : vector<[4]xf32> from vector<[4]x[4]xf32>377 vector.transfer_write %slice, %dest[%slice_index, %c0] { in_bounds = [true] }: vector<[4]xf32>, memref<?x?xf32>378 return379}380 381//===----------------------------------------------------------------------===//382// vector.broadcast383//===----------------------------------------------------------------------===//384 385// -----386 387// CHECK-LABEL: func.func @broadcast_vec2d_from_i32(388// CHECK-SAME: %[[SRC:.*]]: i32) {389// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index390// CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index391// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index392// CHECK: %[[SRC_1D:.*]] = vector.broadcast %[[SRC]] : i32 to vector<[4]xi32>393// CHECK: %[[INIT_TILE:.*]] = arm_sme.get_tile : vector<[4]x[4]xi32>394// CHECK: %[[VSCALE:.*]] = vector.vscale395// CHECK: %[[NUM_TILE_SLICES:.*]] = arith.muli %[[VSCALE]], %[[C4]] : index396// CHECK: %[[TILE:.*]] = scf.for %[[TILE_SLICE_INDEX:.*]] = %[[C0]] to %[[NUM_TILE_SLICES]] step %[[C1]] iter_args(%[[CURRENT_TILE:.*]] = %[[INIT_TILE]]) -> (vector<[4]x[4]xi32>) {397// CHECK: %[[NEW_TILE:.*]] = arm_sme.insert_tile_slice %[[SRC_1D]], %[[CURRENT_TILE]][%[[TILE_SLICE_INDEX]]] : vector<[4]xi32> into vector<[4]x[4]xi32>398// CHECK: scf.yield %[[NEW_TILE]] : vector<[4]x[4]xi32>399// CHECK: "prevent.dce"(%[[TILE]]) : (vector<[4]x[4]xi32>) -> ()400func.func @broadcast_vec2d_from_i32(%arg0: i32) {401 %0 = vector.broadcast %arg0 : i32 to vector<[4]x[4]xi32>402 "prevent.dce"(%0) : (vector<[4]x[4]xi32>) -> ()403 return404}405 406// -----407 408// CHECK-LABEL: func.func @broadcast_vec2d_from_vec0d(409// CHECK-SAME: %[[SRC:.*]]: vector<f32>) {410// CHECK: %[[SRC_1D:.*]] = vector.broadcast %[[SRC]] : vector<f32> to vector<[4]xf32>411// CHECK: scf.for412// CHECK: arm_sme.insert_tile_slice %[[SRC_1D]], {{.*}}413func.func @broadcast_vec2d_from_vec0d(%arg0: vector<f32>) {414 %0 = vector.broadcast %arg0 : vector<f32> to vector<[4]x[4]xf32>415 "prevent.dce"(%0) : (vector<[4]x[4]xf32>) -> ()416 return417}418 419// -----420 421// CHECK-LABEL: func.func @broadcast_vec2d_from_vec1d(422// CHECK-SAME: %[[SRC:.*]]: vector<[8]xi16>) {423// CHECK-NOT: vector.broadcast424// CHECK: scf.for425// CHECK: arm_sme.insert_tile_slice %[[SRC]], {{.*}}426func.func @broadcast_vec2d_from_vec1d(%arg0: vector<[8]xi16>) {427 %0 = vector.broadcast %arg0 : vector<[8]xi16> to vector<[8]x[8]xi16>428 "prevent.dce"(%0) : (vector<[8]x[8]xi16>) -> ()429 return430}431 432 433//===----------------------------------------------------------------------===//434// vector.transpose435//===----------------------------------------------------------------------===//436 437// -----438 439// CHECK-LABEL: func.func @transpose_i8(440// CHECK-SAME: %[[TILE:.*]]: vector<[16]x[16]xi8>)441// CHECK-DAG: %[[C16:.*]] = arith.constant 16 : index442// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index443// CHECK: %[[VSCALE:.*]] = vector.vscale444// CHECK: %[[MIN_TILE_SLICES:.*]] = arith.muli %[[VSCALE]], %[[C16]] : index445// CHECK: %[[NUM_TILE_SLICES:.*]] = memref.alloca(%[[MIN_TILE_SLICES]], %[[MIN_TILE_SLICES]]) : memref<?x?xi8>446// CHECK: arm_sme.tile_store %[[TILE]], %[[NUM_TILE_SLICES]]{{\[}}%[[C0]], %[[C0]]] : memref<?x?xi8>, vector<[16]x[16]xi8>447// CHECK: arm_sme.tile_load %[[NUM_TILE_SLICES]]{{\[}}%[[C0]], %[[C0]]] layout<vertical> : memref<?x?xi8>, vector<[16]x[16]xi8>448func.func @transpose_i8(%arg0: vector<[16]x[16]xi8>) {449 %0 = vector.transpose %arg0, [1, 0] : vector<[16]x[16]xi8> to vector<[16]x[16]xi8>450 "prevent.dce"(%0) : (vector<[16]x[16]xi8>) -> ()451 return452}453 454// -----455 456// CHECK-LABEL: @transpose_i16457// CHECK: arith.constant 8458// CHECK: arm_sme.tile_store {{.*}} : memref<?x?xi16>, vector<[8]x[8]xi16>459// CHECK: arm_sme.tile_load {{.*}} layout<vertical> : memref<?x?xi16>, vector<[8]x[8]xi16>460func.func @transpose_i16(%arg0: vector<[8]x[8]xi16>) {461 %0 = vector.transpose %arg0, [1, 0] : vector<[8]x[8]xi16> to vector<[8]x[8]xi16>462 "prevent.dce"(%0) : (vector<[8]x[8]xi16>) -> ()463 return464}465 466// -----467 468// CHECK-LABEL: @transpose_i32469// CHECK: arith.constant 4470// CHECK: arm_sme.tile_store {{.*}} : memref<?x?xi32>, vector<[4]x[4]xi32>471// CHECK: arm_sme.tile_load {{.*}} layout<vertical> : memref<?x?xi32>, vector<[4]x[4]xi32>472func.func @transpose_i32(%arg0: vector<[4]x[4]xi32>) {473 %0 = vector.transpose %arg0, [1, 0] : vector<[4]x[4]xi32> to vector<[4]x[4]xi32>474 "prevent.dce"(%0) : (vector<[4]x[4]xi32>) -> ()475 return476}477 478// -----479 480// CHECK-LABEL: @transpose_i64481// CHECK: arith.constant 2482// CHECK: arm_sme.tile_store {{.*}} : memref<?x?xi64>, vector<[2]x[2]xi64>483// CHECK: arm_sme.tile_load {{.*}} layout<vertical> : memref<?x?xi64>, vector<[2]x[2]xi64>484func.func @transpose_i64(%arg0: vector<[2]x[2]xi64>) {485 %0 = vector.transpose %arg0, [1, 0] : vector<[2]x[2]xi64> to vector<[2]x[2]xi64>486 "prevent.dce"(%0) : (vector<[2]x[2]xi64>) -> ()487 return488}489 490// -----491 492// CHECK-LABEL: @transpose_i128493// CHECK: %[[VSCALE:.*]] = vector.vscale494// CHECK: %[[NUM_TILE_SLICES:.*]] = memref.alloca(%[[VSCALE]], %[[VSCALE]]) : memref<?x?xi128>495// CHECK: arm_sme.tile_store {{.*}} : memref<?x?xi128>, vector<[1]x[1]xi128>496// CHECK: arm_sme.tile_load {{.*}} layout<vertical> : memref<?x?xi128>, vector<[1]x[1]xi128>497func.func @transpose_i128(%arg0: vector<[1]x[1]xi128>) {498 %0 = vector.transpose %arg0, [1, 0] : vector<[1]x[1]xi128> to vector<[1]x[1]xi128>499 "prevent.dce"(%0) : (vector<[1]x[1]xi128>) -> ()500 return501}502 503// -----504 505// CHECK-LABEL: @transpose_f16506// CHECK: arith.constant 8507// CHECK: arm_sme.tile_store {{.*}} : memref<?x?xf16>, vector<[8]x[8]xf16>508// CHECK: arm_sme.tile_load {{.*}} layout<vertical> : memref<?x?xf16>, vector<[8]x[8]xf16>509func.func @transpose_f16(%arg0: vector<[8]x[8]xf16>) {510 %0 = vector.transpose %arg0, [1, 0] : vector<[8]x[8]xf16> to vector<[8]x[8]xf16>511 "prevent.dce"(%0) : (vector<[8]x[8]xf16>) -> ()512 return513}514 515// -----516 517// CHECK-LABEL: @transpose_bf16518// CHECK: arith.constant 8519// CHECK: arm_sme.tile_store {{.*}} : memref<?x?xbf16>, vector<[8]x[8]xbf16>520// CHECK: arm_sme.tile_load {{.*}} layout<vertical> : memref<?x?xbf16>, vector<[8]x[8]xbf16>521func.func @transpose_bf16(%arg0: vector<[8]x[8]xbf16>) {522 %0 = vector.transpose %arg0, [1, 0] : vector<[8]x[8]xbf16> to vector<[8]x[8]xbf16>523 "prevent.dce"(%0) : (vector<[8]x[8]xbf16>) -> ()524 return525}526 527// -----528 529// CHECK-LABEL: @transpose_f32530// CHECK: arith.constant 4531// CHECK: arm_sme.tile_store {{.*}} : memref<?x?xf32>, vector<[4]x[4]xf32>532// CHECK: arm_sme.tile_load {{.*}} layout<vertical> : memref<?x?xf32>, vector<[4]x[4]xf32>533func.func @transpose_f32(%arg0: vector<[4]x[4]xf32>) {534 %0 = vector.transpose %arg0, [1, 0] : vector<[4]x[4]xf32> to vector<[4]x[4]xf32>535 "prevent.dce"(%0) : (vector<[4]x[4]xf32>) -> ()536 return537}538 539// -----540 541// CHECK-LABEL: @transpose_f64542// CHECK: arith.constant 2543// CHECK: arm_sme.tile_store {{.*}} : memref<?x?xf64>, vector<[2]x[2]xf64>544// CHECK: arm_sme.tile_load {{.*}} layout<vertical> : memref<?x?xf64>, vector<[2]x[2]xf64>545func.func @transpose_f64(%arg0: vector<[2]x[2]xf64>) {546 %0 = vector.transpose %arg0, [1, 0] : vector<[2]x[2]xf64> to vector<[2]x[2]xf64>547 "prevent.dce"(%0) : (vector<[2]x[2]xf64>) -> ()548 return549}550 551//===----------------------------------------------------------------------===//552// vector.outerproduct553//===----------------------------------------------------------------------===//554 555// -----556 557// CHECK-LABEL: @vector_outerproduct_masked_f16558// CHECK-SAME: (%[[LHS:.*]]: vector<[8]xf16>, %[[RHS:.*]]: vector<[8]xf16>, %[[ACC:.*]]: vector<[8]x[8]xf16>, %[[DIM0:.*]]: index, %[[DIM1:.*]]: index559func.func @vector_outerproduct_masked_f16(%lhs : vector<[8]xf16>, %rhs : vector<[8]xf16>, %acc : vector<[8]x[8]xf16>, %dim0 : index, %dim1 : index) {560 %mask = vector.create_mask %dim0, %dim1 : vector<[8]x[8]xi1>561 // CHECK: %[[LHS_MASK:.*]] = vector.create_mask %[[DIM0]] : vector<[8]xi1>562 // CHECK: %[[RHS_MASK:.*]] = vector.create_mask %[[DIM1]] : vector<[8]xi1>563 // CHECK: arm_sme.outerproduct %[[LHS]], %[[RHS]] acc(%[[ACC]]) masks(%[[LHS_MASK]], %[[RHS_MASK]]) : vector<[8]xf16>, vector<[8]xf16>564 %result = vector.mask %mask { vector.outerproduct %lhs, %rhs, %acc {kind = #vector.kind<add>} : vector<[8]xf16>, vector<[8]xf16> } : vector<[8]x[8]xi1> -> vector<[8]x[8]xf16>565 "prevent.dce"(%result) : (vector<[8]x[8]xf16>) -> ()566}567 568// -----569 570// CHECK-LABEL: @vector_outerproduct_masked_bf16571// CHECK-SAME: (%[[LHS:.*]]: vector<[8]xbf16>, %[[RHS:.*]]: vector<[8]xbf16>, %[[ACC:.*]]: vector<[8]x[8]xbf16>, %[[DIM0:.*]]: index, %[[DIM1:.*]]: index572func.func @vector_outerproduct_masked_bf16(%lhs : vector<[8]xbf16>, %rhs : vector<[8]xbf16>, %acc : vector<[8]x[8]xbf16>, %dim0 : index, %dim1 : index) {573 %mask = vector.create_mask %dim0, %dim1 : vector<[8]x[8]xi1>574 // CHECK: %[[LHS_MASK:.*]] = vector.create_mask %[[DIM0]] : vector<[8]xi1>575 // CHECK: %[[RHS_MASK:.*]] = vector.create_mask %[[DIM1]] : vector<[8]xi1>576 // CHECK: arm_sme.outerproduct %[[LHS]], %[[RHS]] acc(%[[ACC]]) masks(%[[LHS_MASK]], %[[RHS_MASK]]) : vector<[8]xbf16>, vector<[8]xbf16>577 %result = vector.mask %mask { vector.outerproduct %lhs, %rhs, %acc {kind = #vector.kind<add>} : vector<[8]xbf16>, vector<[8]xbf16> } : vector<[8]x[8]xi1> -> vector<[8]x[8]xbf16>578 "prevent.dce"(%result) : (vector<[8]x[8]xbf16>) -> ()579}580 581// -----582 583// CHECK-LABEL: @vector_outerproduct_masked_f32584// CHECK-SAME: (%[[LHS:.*]]: vector<[4]xf32>, %[[RHS:.*]]: vector<[4]xf32>, %[[ACC:.*]]: vector<[4]x[4]xf32>, %[[DIM0:.*]]: index, %[[DIM1:.*]]: index585func.func @vector_outerproduct_masked_f32(%lhs : vector<[4]xf32>, %rhs : vector<[4]xf32>, %acc : vector<[4]x[4]xf32>, %dim0 : index, %dim1 : index) {586 %mask = vector.create_mask %dim0, %dim1 : vector<[4]x[4]xi1>587 // CHECK: %[[LHS_MASK:.*]] = vector.create_mask %[[DIM0]] : vector<[4]xi1>588 // CHECK: %[[RHS_MASK:.*]] = vector.create_mask %[[DIM1]] : vector<[4]xi1>589 // CHECK: arm_sme.outerproduct %[[LHS]], %[[RHS]] acc(%[[ACC]]) masks(%[[LHS_MASK]], %[[RHS_MASK]]) : vector<[4]xf32>, vector<[4]xf32>590 %result = vector.mask %mask { vector.outerproduct %lhs, %rhs, %acc {kind = #vector.kind<add>} : vector<[4]xf32>, vector<[4]xf32> } : vector<[4]x[4]xi1> -> vector<[4]x[4]xf32>591 "prevent.dce"(%result) : (vector<[4]x[4]xf32>) -> ()592}593 594// -----595 596// CHECK-LABEL: @vector_outerproduct_masked_f64597// CHECK-SAME: (%[[LHS:.*]]: vector<[2]xf64>, %[[RHS:.*]]: vector<[2]xf64>, %[[ACC:.*]]: vector<[2]x[2]xf64>, %[[DIM0:.*]]: index, %[[DIM1:.*]]: index598func.func @vector_outerproduct_masked_f64(%lhs : vector<[2]xf64>, %rhs : vector<[2]xf64>, %acc : vector<[2]x[2]xf64>, %dim0 : index, %dim1 : index) {599 %mask = vector.create_mask %dim0, %dim1 : vector<[2]x[2]xi1>600 // CHECK: %[[LHS_MASK:.*]] = vector.create_mask %[[DIM0]] : vector<[2]xi1>601 // CHECK: %[[RHS_MASK:.*]] = vector.create_mask %[[DIM1]] : vector<[2]xi1>602 // CHECK: arm_sme.outerproduct %[[LHS]], %[[RHS]] acc(%[[ACC]]) masks(%[[LHS_MASK]], %[[RHS_MASK]]) : vector<[2]xf64>, vector<[2]xf64>603 %result = vector.mask %mask { vector.outerproduct %lhs, %rhs, %acc {kind = #vector.kind<add>} : vector<[2]xf64>, vector<[2]xf64> } : vector<[2]x[2]xi1> -> vector<[2]x[2]xf64>604 "prevent.dce"(%result) : (vector<[2]x[2]xf64>) -> ()605}606 607// -----608 609// CHECK-LABEL: @vector_outerproduct_f16610// CHECK-SAME: (%[[LHS:.*]]: vector<[8]xf16>, %[[RHS:.*]]: vector<[8]xf16>, %[[ACC:.*]]: vector<[8]x[8]xf16>611func.func @vector_outerproduct_f16(%lhs : vector<[8]xf16>, %rhs : vector<[8]xf16>, %acc : vector<[8]x[8]xf16>) {612 // CHECK: arm_sme.outerproduct %[[LHS]], %[[RHS]] acc(%[[ACC]]) : vector<[8]xf16>, vector<[8]xf16>613 %result = vector.outerproduct %lhs, %rhs, %acc {kind = #vector.kind<add>} : vector<[8]xf16>, vector<[8]xf16>614 "prevent.dce"(%result) : (vector<[8]x[8]xf16>) -> ()615}616 617// -----618 619// CHECK-LABEL: @vector_outerproduct_bf16620// CHECK-SAME: (%[[LHS:.*]]: vector<[8]xbf16>, %[[RHS:.*]]: vector<[8]xbf16>, %[[ACC:.*]]: vector<[8]x[8]xbf16>621func.func @vector_outerproduct_bf16(%lhs : vector<[8]xbf16>, %rhs : vector<[8]xbf16>, %acc : vector<[8]x[8]xbf16>) {622 // CHECK: arm_sme.outerproduct %[[LHS]], %[[RHS]] acc(%[[ACC]]) : vector<[8]xbf16>, vector<[8]xbf16>623 %result = vector.outerproduct %lhs, %rhs, %acc {kind = #vector.kind<add>} : vector<[8]xbf16>, vector<[8]xbf16>624 "prevent.dce"(%result) : (vector<[8]x[8]xbf16>) -> ()625}626 627// -----628 629// CHECK-LABEL: @vector_outerproduct_f32630// CHECK-SAME: (%[[LHS:.*]]: vector<[4]xf32>, %[[RHS:.*]]: vector<[4]xf32>, %[[ACC:.*]]: vector<[4]x[4]xf32>631func.func @vector_outerproduct_f32(%lhs : vector<[4]xf32>, %rhs : vector<[4]xf32>, %acc : vector<[4]x[4]xf32>) {632 // CHECK: arm_sme.outerproduct %[[LHS]], %[[RHS]] acc(%[[ACC]]) : vector<[4]xf32>, vector<[4]xf32>633 %result = vector.outerproduct %lhs, %rhs, %acc {kind = #vector.kind<add>} : vector<[4]xf32>, vector<[4]xf32>634 "prevent.dce"(%result) : (vector<[4]x[4]xf32>) -> ()635}636 637// -----638 639// CHECK-LABEL: @vector_outerproduct_f64640// CHECK-SAME: (%[[LHS:.*]]: vector<[2]xf64>, %[[RHS:.*]]: vector<[2]xf64>, %[[ACC:.*]]: vector<[2]x[2]xf64>641func.func @vector_outerproduct_f64(%lhs : vector<[2]xf64>, %rhs : vector<[2]xf64>, %acc : vector<[2]x[2]xf64>) {642 // CHECK: arm_sme.outerproduct %[[LHS]], %[[RHS]] acc(%[[ACC]]) : vector<[2]xf64>, vector<[2]xf64>643 %result = vector.outerproduct %lhs, %rhs, %acc {kind = #vector.kind<add>} : vector<[2]xf64>, vector<[2]xf64>644 "prevent.dce"(%result) : (vector<[2]x[2]xf64>) -> ()645}646 647//===----------------------------------------------------------------------===//648// vector.print649//===----------------------------------------------------------------------===//650 651// -----652 653func.func @vector_print_tile(%tile: vector<[4]x[4]xf32>)654{655 vector.print %tile : vector<[4]x[4]xf32>656 return657}658// CHECK-LABEL: func.func @vector_print_tile(659// CHECK-SAME: %[[TILE:.*]]: vector<[4]x[4]xf32>) {660// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index661// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index662// CHECK-DAG: %[[C4:.*]] = arith.constant 4 : index663// CHECK-DAG: %[[VSCALE:.*]] = vector.vscale664// CHECK-DAG: %[[NUM_TILE_SLICES:.*]] = arith.muli %[[VSCALE]], %[[C4]] : index665// CHECK-NEXT: scf.for %[[TILE_SLICE_INDEX:.*]] = %[[C0]] to %[[NUM_TILE_SLICES]] step %[[C1]] {666// CHECK-NEXT: %[[TILE_SLICE:.*]] = arm_sme.extract_tile_slice %[[TILE]][%[[TILE_SLICE_INDEX]]] : vector<[4]xf32> from vector<[4]x[4]xf32>667// CHECK-NEXT: vector.print %[[TILE_SLICE]] : vector<[4]xf32>668 669//===----------------------------------------------------------------------===//670// vector.load671//===----------------------------------------------------------------------===//672 673// -----674 675// CHECK-LABEL: @vector_load_i8_with_offset(676// CHECK-SAME: %[[MEMREF:.*]]: memref<?x?xi8>)677// CHECK: %[[C0:.*]] = arith.constant 0 : index678// CHECK: %[[C123:.*]] = arith.constant 123 : index679// CHECK: arm_sme.tile_load %[[MEMREF]][%[[C123]], %[[C0]]] : memref<?x?xi8>, vector<[16]x[16]xi8>680func.func @vector_load_i8_with_offset(%arg0 : memref<?x?xi8>) -> vector<[16]x[16]xi8> {681 %c0 = arith.constant 0 : index682 %c123 = arith.constant 123 : index683 %tile = vector.load %arg0[%c123, %c0] : memref<?x?xi8>, vector<[16]x[16]xi8>684 return %tile : vector<[16]x[16]xi8>685}686 687// -----688 689// CHECK-LABEL: @vector_load_i16(690// CHECK: arm_sme.tile_load {{.*}} : memref<?x?xi16>, vector<[8]x[8]xi16>691func.func @vector_load_i16(%arg0 : memref<?x?xi16>) -> vector<[8]x[8]xi16> {692 %c0 = arith.constant 0 : index693 %tile = vector.load %arg0[%c0, %c0] : memref<?x?xi16>, vector<[8]x[8]xi16>694 return %tile : vector<[8]x[8]xi16>695}696 697// -----698 699// CHECK-LABEL: @vector_load_i32(700// CHECK: arm_sme.tile_load {{.*}} : memref<?x?xi32>, vector<[4]x[4]xi32>701func.func @vector_load_i32(%arg0 : memref<?x?xi32>) -> vector<[4]x[4]xi32> {702 %c0 = arith.constant 0 : index703 %tile = vector.load %arg0[%c0, %c0] : memref<?x?xi32>, vector<[4]x[4]xi32>704 return %tile : vector<[4]x[4]xi32>705}706 707// -----708 709// CHECK-LABEL: @vector_load_i64(710// CHECK: arm_sme.tile_load {{.*}} : memref<?x?xi64>, vector<[2]x[2]xi64>711func.func @vector_load_i64(%arg0 : memref<?x?xi64>) -> vector<[2]x[2]xi64> {712 %c0 = arith.constant 0 : index713 %tile = vector.load %arg0[%c0, %c0] : memref<?x?xi64>, vector<[2]x[2]xi64>714 return %tile : vector<[2]x[2]xi64>715}716 717// -----718 719// CHECK-LABEL: @vector_load_f16(720// CHECK: arm_sme.tile_load {{.*}} : memref<?x?xf16>, vector<[8]x[8]xf16>721func.func @vector_load_f16(%arg0 : memref<?x?xf16>) -> vector<[8]x[8]xf16> {722 %c0 = arith.constant 0 : index723 %tile = vector.load %arg0[%c0, %c0] : memref<?x?xf16>, vector<[8]x[8]xf16>724 return %tile : vector<[8]x[8]xf16>725}726 727// -----728 729// CHECK-LABEL: @vector_load_bf16(730// CHECK: arm_sme.tile_load {{.*}} : memref<?x?xbf16>, vector<[8]x[8]xbf16>731func.func @vector_load_bf16(%arg0 : memref<?x?xbf16>) -> vector<[8]x[8]xbf16> {732 %c0 = arith.constant 0 : index733 %tile = vector.load %arg0[%c0, %c0] : memref<?x?xbf16>, vector<[8]x[8]xbf16>734 return %tile : vector<[8]x[8]xbf16>735}736 737// -----738 739// CHECK-LABEL: @vector_load_f32(740// CHECK: arm_sme.tile_load {{.*}} : memref<?x?xf32>, vector<[4]x[4]xf32>741func.func @vector_load_f32(%arg0 : memref<?x?xf32>) -> vector<[4]x[4]xf32> {742 %c0 = arith.constant 0 : index743 %tile = vector.load %arg0[%c0, %c0] : memref<?x?xf32>, vector<[4]x[4]xf32>744 return %tile : vector<[4]x[4]xf32>745}746 747// -----748 749// CHECK-LABEL: @vector_load_f64(750// CHECK: arm_sme.tile_load {{.*}} : memref<?x?xf64>, vector<[2]x[2]xf64>751func.func @vector_load_f64(%arg0 : memref<?x?xf64>) -> vector<[2]x[2]xf64> {752 %c0 = arith.constant 0 : index753 %tile = vector.load %arg0[%c0, %c0] : memref<?x?xf64>, vector<[2]x[2]xf64>754 return %tile : vector<[2]x[2]xf64>755}756 757// -----758 759// CHECK-LABEL: @vector_load_i128(760// CHECK: arm_sme.tile_load {{.*}} : memref<?x?xi128>, vector<[1]x[1]xi128>761func.func @vector_load_i128(%arg0 : memref<?x?xi128>) -> vector<[1]x[1]xi128> {762 %c0 = arith.constant 0 : index763 %tile = vector.load %arg0[%c0, %c0] : memref<?x?xi128>, vector<[1]x[1]xi128>764 return %tile : vector<[1]x[1]xi128>765}766 767 768//===----------------------------------------------------------------------===//769// vector.store770//===----------------------------------------------------------------------===//771 772// -----773 774// CHECK-LABEL: @vector_store_i8(775// CHECK-SAME: %[[MEMREF:.*]]: memref<?x?xi8>) {776// CHECK: %[[C0:.*]] = arith.constant 0 : index777// CHECK: %[[TILE:.*]] = arm_sme.get_tile : vector<[16]x[16]xi8>778// CHECK: arm_sme.tile_store %[[TILE]], %[[MEMREF]][%[[C0]], %[[C0]]] : memref<?x?xi8>, vector<[16]x[16]xi8>779func.func @vector_store_i8(%arg0 : memref<?x?xi8>) {780 %c0 = arith.constant 0 : index781 %tile = arm_sme.get_tile : vector<[16]x[16]xi8>782 vector.store %tile, %arg0[%c0, %c0] : memref<?x?xi8>, vector<[16]x[16]xi8>783 return784}785 786// -----787 788// CHECK-LABEL: @vector_store_i16789// CHECK: arm_sme.tile_store {{.*}} : memref<?x?xi16>, vector<[8]x[8]xi16>790func.func @vector_store_i16(%arg0 : memref<?x?xi16>) {791 %c0 = arith.constant 0 : index792 %tile = arm_sme.get_tile : vector<[8]x[8]xi16>793 vector.store %tile, %arg0[%c0, %c0] : memref<?x?xi16>, vector<[8]x[8]xi16>794 return795}796 797// -----798 799// CHECK-LABEL: @vector_store_i32800// CHECK: arm_sme.tile_store {{.*}} : memref<?x?xi32>, vector<[4]x[4]xi32>801func.func @vector_store_i32(%arg0 : memref<?x?xi32>) {802 %c0 = arith.constant 0 : index803 %tile = arm_sme.get_tile : vector<[4]x[4]xi32>804 vector.store %tile, %arg0[%c0, %c0] : memref<?x?xi32>, vector<[4]x[4]xi32>805 return806}807 808// -----809 810// CHECK-LABEL: @vector_store_i64811// CHECK: arm_sme.tile_store {{.*}} : memref<?x?xi64>, vector<[2]x[2]xi64>812func.func @vector_store_i64(%arg0 : memref<?x?xi64>) {813 %c0 = arith.constant 0 : index814 %tile = arm_sme.get_tile : vector<[2]x[2]xi64>815 vector.store %tile, %arg0[%c0, %c0] : memref<?x?xi64>, vector<[2]x[2]xi64>816 return817}818 819// -----820 821// CHECK-LABEL: @vector_store_f16822// CHECK: arm_sme.tile_store {{.*}} : memref<?x?xf16>, vector<[8]x[8]xf16>823func.func @vector_store_f16(%arg0 : memref<?x?xf16>) {824 %c0 = arith.constant 0 : index825 %tile = arm_sme.get_tile : vector<[8]x[8]xf16>826 vector.store %tile, %arg0[%c0, %c0] : memref<?x?xf16>, vector<[8]x[8]xf16>827 return828}829 830// -----831 832// CHECK-LABEL: @vector_store_bf16833// CHECK: arm_sme.tile_store {{.*}} : memref<?x?xbf16>, vector<[8]x[8]xbf16>834func.func @vector_store_bf16(%arg0 : memref<?x?xbf16>) {835 %c0 = arith.constant 0 : index836 %tile = arm_sme.get_tile : vector<[8]x[8]xbf16>837 vector.store %tile, %arg0[%c0, %c0] : memref<?x?xbf16>, vector<[8]x[8]xbf16>838 return839}840// -----841 842// CHECK-LABEL: @vector_store_f32843// CHECK: arm_sme.tile_store {{.*}} : memref<?x?xf32>, vector<[4]x[4]xf32>844func.func @vector_store_f32(%arg0 : memref<?x?xf32>) {845 %c0 = arith.constant 0 : index846 %tile = arm_sme.get_tile : vector<[4]x[4]xf32>847 vector.store %tile, %arg0[%c0, %c0] : memref<?x?xf32>, vector<[4]x[4]xf32>848 return849}850 851// -----852 853// CHECK-LABEL: @vector_store_f64854// CHECK: arm_sme.tile_store {{.*}} : memref<?x?xf64>, vector<[2]x[2]xf64>855func.func @vector_store_f64(%arg0 : memref<?x?xf64>) {856 %c0 = arith.constant 0 : index857 %tile = arm_sme.get_tile : vector<[2]x[2]xf64>858 vector.store %tile, %arg0[%c0, %c0] : memref<?x?xf64>, vector<[2]x[2]xf64>859 return860}861 862// -----863 864// CHECK-LABEL: @vector_store_i128865// CHECK: arm_sme.tile_store {{.*}} : memref<?x?xi128>, vector<[1]x[1]xi128>866func.func @vector_store_i128(%arg0 : memref<?x?xi128>) {867 %c0 = arith.constant 0 : index868 %tile = arm_sme.get_tile : vector<[1]x[1]xi128>869 vector.store %tile, %arg0[%c0, %c0] : memref<?x?xi128>, vector<[1]x[1]xi128>870 return871}872 873//===----------------------------------------------------------------------===//874// vector.insert875//===----------------------------------------------------------------------===//876 877// -----878 879// CHECK-LABEL: @vector_insert_slice_i32(880// CHECK-SAME: %[[SLICE:.*]]: vector<[4]xi32>,881// CHECK-SAME: %[[INDEX:.*]]: index)882func.func @vector_insert_slice_i32(%slice: vector<[4]xi32>, %row: index) -> vector<[4]x[4]xi32>{883 // CHECK-NEXT: %[[TILE:.*]] = arm_sme.get_tile : vector<[4]x[4]xi32>884 // CHECK-NEXT: arm_sme.insert_tile_slice %[[SLICE]], %[[TILE]][%[[INDEX]]] : vector<[4]xi32> into vector<[4]x[4]xi32>885 %tile = arm_sme.get_tile : vector<[4]x[4]xi32>886 %new_tile = vector.insert %slice, %tile[%row] : vector<[4]xi32> into vector<[4]x[4]xi32>887 return %new_tile : vector<[4]x[4]xi32>888}889 890// -----891 892// CHECK-LABEL: @vector_insert_slice_i8893func.func @vector_insert_slice_i8(%slice: vector<[16]xi8>, %row: index) -> vector<[16]x[16]xi8> {894 // CHECK: arm_sme.insert_tile_slice %{{.*}} : vector<[16]xi8> into vector<[16]x[16]xi8>895 %tile = arm_sme.get_tile : vector<[16]x[16]xi8>896 %new_tile = vector.insert %slice, %tile[%row] : vector<[16]xi8> into vector<[16]x[16]xi8>897 return %new_tile : vector<[16]x[16]xi8>898}899 900// -----901 902// CHECK-LABEL: @vector_insert_slice_i16903func.func @vector_insert_slice_i16(%slice: vector<[8]xi16>, %row: index) -> vector<[8]x[8]xi16> {904 // CHECK: arm_sme.insert_tile_slice %{{.*}} : vector<[8]xi16> into vector<[8]x[8]xi16>905 %tile = arm_sme.get_tile : vector<[8]x[8]xi16>906 %new_tile = vector.insert %slice, %tile[%row] : vector<[8]xi16> into vector<[8]x[8]xi16>907 return %new_tile : vector<[8]x[8]xi16>908}909 910// -----911 912// CHECK-LABEL: @vector_insert_slice_i64913func.func @vector_insert_slice_i64(%slice: vector<[2]xi64>, %row: index) -> vector<[2]x[2]xi64> {914 // CHECK: arm_sme.insert_tile_slice %{{.*}} : vector<[2]xi64> into vector<[2]x[2]xi64>915 %tile = arm_sme.get_tile : vector<[2]x[2]xi64>916 %new_tile = vector.insert %slice, %tile[%row] : vector<[2]xi64> into vector<[2]x[2]xi64>917 return %new_tile : vector<[2]x[2]xi64>918}919 920// -----921 922// CHECK-LABEL: @vector_insert_slice_i128923func.func @vector_insert_slice_i128(%slice: vector<[1]xi128>, %row: index) -> vector<[1]x[1]xi128> {924 // CHECK: arm_sme.insert_tile_slice %{{.*}} : vector<[1]xi128> into vector<[1]x[1]xi128>925 %tile = arm_sme.get_tile : vector<[1]x[1]xi128>926 %new_tile = vector.insert %slice, %tile[%row] : vector<[1]xi128> into vector<[1]x[1]xi128>927 return %new_tile : vector<[1]x[1]xi128>928}929 930// -----931 932// CHECK-LABEL: @vector_insert_slice_f16933func.func @vector_insert_slice_f16(%slice: vector<[8]xf16>, %row: index) -> vector<[8]x[8]xf16> {934 // CHECK: arm_sme.insert_tile_slice %{{.*}} : vector<[8]xf16> into vector<[8]x[8]xf16>935 %tile = arm_sme.get_tile : vector<[8]x[8]xf16>936 %new_tile = vector.insert %slice, %tile[%row] : vector<[8]xf16> into vector<[8]x[8]xf16>937 return %new_tile : vector<[8]x[8]xf16>938}939 940// -----941 942// CHECK-LABEL: @vector_insert_slice_bf16943func.func @vector_insert_slice_bf16(%slice: vector<[8]xbf16>, %row: index) -> vector<[8]x[8]xbf16> {944 // CHECK: arm_sme.insert_tile_slice %{{.*}} : vector<[8]xbf16> into vector<[8]x[8]xbf16>945 %tile = arm_sme.get_tile : vector<[8]x[8]xbf16>946 %new_tile = vector.insert %slice, %tile[%row] : vector<[8]xbf16> into vector<[8]x[8]xbf16>947 return %new_tile : vector<[8]x[8]xbf16>948}949 950// -----951 952// CHECK-LABEL: @vector_insert_slice_f32953func.func @vector_insert_slice_f32(%slice: vector<[4]xf32>, %row: index) -> vector<[4]x[4]xf32> {954 // CHECK: arm_sme.insert_tile_slice %{{.*}} : vector<[4]xf32> into vector<[4]x[4]xf32>955 %tile = arm_sme.get_tile : vector<[4]x[4]xf32>956 %new_tile = vector.insert %slice, %tile[%row] : vector<[4]xf32> into vector<[4]x[4]xf32>957 return %new_tile : vector<[4]x[4]xf32>958}959 960// -----961 962// CHECK-LABEL: @vector_insert_slice_f64963func.func @vector_insert_slice_f64(%slice: vector<[2]xf64>, %row: index) -> vector<[2]x[2]xf64> {964 // CHECK: arm_sme.insert_tile_slice %{{.*}} : vector<[2]xf64> into vector<[2]x[2]xf64>965 %tile = arm_sme.get_tile : vector<[2]x[2]xf64>966 %new_tile = vector.insert %slice, %tile[%row] : vector<[2]xf64> into vector<[2]x[2]xf64>967 return %new_tile : vector<[2]x[2]xf64>968}969 970// -----971 972// CHECK-LABEL: @vector_insert_element_i32(973// CHECK-SAME: %[[EL:.*]]: i32,974// CHECK-SAME: %[[ROW:.*]]: index,975// CHECK-SAME: %[[COL:.*]]: index)976func.func @vector_insert_element_i32(%el: i32, %row: index, %col: index) -> vector<[4]x[4]xi32> {977 // CHECK-NEXT: %[[TILE:.*]] = arm_sme.get_tile : vector<[4]x[4]xi32>978 // CHECK-NEXT: %[[SLICE:.*]] = arm_sme.extract_tile_slice %[[TILE]][%[[ROW]]] : vector<[4]xi32> from vector<[4]x[4]xi32>979 // CHECK-NEXT: %[[NEW_SLICE:.*]] = vector.insert %[[EL]], %[[SLICE]] [%[[COL]]] : i32 into vector<[4]xi32>980 // CHECK-NEXT: arm_sme.insert_tile_slice %[[NEW_SLICE]], %[[TILE]][%[[ROW]]] : vector<[4]xi32> into vector<[4]x[4]xi32>981 %tile = arm_sme.get_tile : vector<[4]x[4]xi32>982 %new_tile = vector.insert %el, %tile[%row, %col] : i32 into vector<[4]x[4]xi32>983 return %new_tile : vector<[4]x[4]xi32>984}985 986// -----987 988// CHECK-LABEL: @vector_insert_element_i8989func.func @vector_insert_element_i8(%el: i8, %row: index, %col: index) -> vector<[16]x[16]xi8> {990 // CHECK: %[[TILE:.*]] = arm_sme.get_tile : vector<[16]x[16]xi8>991 // CHECK: arm_sme.extract_tile_slice %[[TILE]]{{.*}} : vector<[16]xi8> from vector<[16]x[16]xi8>992 // CHECK: arm_sme.insert_tile_slice %{{.*}}, %[[TILE]][%{{.*}}] : vector<[16]xi8> into vector<[16]x[16]xi8>993 %tile = arm_sme.get_tile : vector<[16]x[16]xi8>994 %new_tile = vector.insert %el, %tile[%row, %col] : i8 into vector<[16]x[16]xi8>995 return %new_tile : vector<[16]x[16]xi8>996}997 998// -----999 1000// CHECK-LABEL: @vector_insert_element_i161001func.func @vector_insert_element_i16(%el: i16, %row: index, %col: index) -> vector<[8]x[8]xi16> {1002 // CHECK: %[[TILE:.*]] = arm_sme.get_tile : vector<[8]x[8]xi16>1003 // CHECK: arm_sme.extract_tile_slice %[[TILE]]{{.*}} : vector<[8]xi16> from vector<[8]x[8]xi16>1004 // CHECK: arm_sme.insert_tile_slice %{{.*}}, %[[TILE]][%{{.*}}] : vector<[8]xi16> into vector<[8]x[8]xi16>1005 %tile = arm_sme.get_tile : vector<[8]x[8]xi16>1006 %new_tile = vector.insert %el, %tile[%row, %col] : i16 into vector<[8]x[8]xi16>1007 return %new_tile : vector<[8]x[8]xi16>1008}1009 1010// -----1011 1012// CHECK-LABEL: @vector_insert_element_i641013func.func @vector_insert_element_i64(%el: i64, %row: index, %col: index) -> vector<[2]x[2]xi64> {1014 // CHECK: %[[TILE:.*]] = arm_sme.get_tile : vector<[2]x[2]xi64>1015 // CHECK: arm_sme.extract_tile_slice %[[TILE]]{{.*}} : vector<[2]xi64> from vector<[2]x[2]xi64>1016 // CHECK: arm_sme.insert_tile_slice %{{.*}}, %[[TILE]][%{{.*}}] : vector<[2]xi64> into vector<[2]x[2]xi64>1017 %tile = arm_sme.get_tile : vector<[2]x[2]xi64>1018 %new_tile = vector.insert %el, %tile[%row, %col] : i64 into vector<[2]x[2]xi64>1019 return %new_tile : vector<[2]x[2]xi64>1020}1021 1022// -----1023 1024// CHECK-LABEL: @vector_insert_element_i1281025func.func @vector_insert_element_i128(%el: i128, %row: index, %col: index) -> vector<[1]x[1]xi128> {1026 // CHECK: %[[TILE:.*]] = arm_sme.get_tile : vector<[1]x[1]xi128>1027 // CHECK: arm_sme.extract_tile_slice %[[TILE]]{{.*}} : vector<[1]xi128> from vector<[1]x[1]xi128>1028 // CHECK: arm_sme.insert_tile_slice %{{.*}}, %[[TILE]][%{{.*}}] : vector<[1]xi128> into vector<[1]x[1]xi128>1029 %tile = arm_sme.get_tile : vector<[1]x[1]xi128>1030 %new_tile = vector.insert %el, %tile[%row, %col] : i128 into vector<[1]x[1]xi128>1031 return %new_tile : vector<[1]x[1]xi128>1032}1033 1034// -----1035 1036// CHECK-LABEL: @vector_insert_element_f161037func.func @vector_insert_element_f16(%el: f16, %row: index, %col: index) -> vector<[8]x[8]xf16> {1038 // CHECK: %[[TILE:.*]] = arm_sme.get_tile : vector<[8]x[8]xf16>1039 // CHECK: arm_sme.extract_tile_slice %[[TILE]]{{.*}} : vector<[8]xf16> from vector<[8]x[8]xf16>1040 // CHECK: arm_sme.insert_tile_slice %{{.*}}, %[[TILE]][%{{.*}}] : vector<[8]xf16> into vector<[8]x[8]xf16>1041 %tile = arm_sme.get_tile : vector<[8]x[8]xf16>1042 %new_tile = vector.insert %el, %tile[%row, %col] : f16 into vector<[8]x[8]xf16>1043 return %new_tile : vector<[8]x[8]xf16>1044}1045 1046// -----1047 1048// CHECK-LABEL: @vector_insert_element_bf161049func.func @vector_insert_element_bf16(%el: bf16, %row: index, %col: index) -> vector<[8]x[8]xbf16> {1050 // CHECK: %[[TILE:.*]] = arm_sme.get_tile : vector<[8]x[8]xbf16>1051 // CHECK: arm_sme.extract_tile_slice %[[TILE]]{{.*}} : vector<[8]xbf16> from vector<[8]x[8]xbf16>1052 // CHECK: arm_sme.insert_tile_slice %{{.*}}, %[[TILE]][%{{.*}}] : vector<[8]xbf16> into vector<[8]x[8]xbf16>1053 %tile = arm_sme.get_tile : vector<[8]x[8]xbf16>1054 %new_tile = vector.insert %el, %tile[%row, %col] : bf16 into vector<[8]x[8]xbf16>1055 return %new_tile : vector<[8]x[8]xbf16>1056}1057 1058// -----1059 1060// CHECK-LABEL: @vector_insert_element_f321061func.func @vector_insert_element_f32(%el: f32, %row: index, %col: index) -> vector<[4]x[4]xf32> {1062 // CHECK: %[[TILE:.*]] = arm_sme.get_tile : vector<[4]x[4]xf32>1063 // CHECK: arm_sme.extract_tile_slice %[[TILE]]{{.*}} : vector<[4]xf32> from vector<[4]x[4]xf32>1064 // CHECK: arm_sme.insert_tile_slice %{{.*}}, %[[TILE]][%{{.*}}] : vector<[4]xf32> into vector<[4]x[4]xf32>1065 %tile = arm_sme.get_tile : vector<[4]x[4]xf32>1066 %new_tile = vector.insert %el, %tile[%row, %col] : f32 into vector<[4]x[4]xf32>1067 return %new_tile : vector<[4]x[4]xf32>1068}1069 1070// -----1071 1072// CHECK-LABEL: @vector_insert_element_f641073func.func @vector_insert_element_f64(%el: f64, %row: index, %col: index) -> vector<[2]x[2]xf64> {1074 // CHECK: %[[TILE:.*]] = arm_sme.get_tile : vector<[2]x[2]xf64>1075 // CHECK: arm_sme.extract_tile_slice %[[TILE]]{{.*}} : vector<[2]xf64> from vector<[2]x[2]xf64>1076 // CHECK: arm_sme.insert_tile_slice %{{.*}}, %[[TILE]][%{{.*}}] : vector<[2]xf64> into vector<[2]x[2]xf64>1077 %tile = arm_sme.get_tile : vector<[2]x[2]xf64>1078 %new_tile = vector.insert %el, %tile[%row, %col] : f64 into vector<[2]x[2]xf64>1079 return %new_tile : vector<[2]x[2]xf64>1080}1081 1082//===----------------------------------------------------------------------===//1083// vector.extract --> arm_sme.extract_tile_slice1084//===----------------------------------------------------------------------===//1085 1086// -----1087 1088// CHECK-LABEL: @vector_extract_slice_i32(1089// CHECK-SAME: %[[INDEX:.*]]: index)1090func.func @vector_extract_slice_i32(%row: index) -> vector<[4]xi32> {1091 // CHECK: %[[TILE:.*]] = arm_sme.get_tile : vector<[4]x[4]xi32>1092 // CHECK: arm_sme.extract_tile_slice %[[TILE]][%[[INDEX]]] : vector<[4]xi32> from vector<[4]x[4]xi32>1093 %tile = arm_sme.get_tile : vector<[4]x[4]xi32>1094 %slice = vector.extract %tile[%row] : vector<[4]xi32> from vector<[4]x[4]xi32>1095 return %slice : vector<[4]xi32>1096}1097 1098// -----1099 1100// CHECK-LABEL: @vector_extract_slice_i81101func.func @vector_extract_slice_i8(%row: index) -> vector<[16]xi8> {1102 // CHECK: arm_sme.extract_tile_slice {{.*}} : vector<[16]xi8> from vector<[16]x[16]xi8>1103 %tile = arm_sme.get_tile : vector<[16]x[16]xi8>1104 %slice = vector.extract %tile[%row] : vector<[16]xi8> from vector<[16]x[16]xi8>1105 return %slice : vector<[16]xi8>1106}1107 1108// -----1109 1110// CHECK-LABEL: @vector_extract_slice_i161111func.func @vector_extract_slice_i16(%row: index) -> vector<[8]xi16> {1112 // CHECK: arm_sme.extract_tile_slice {{.*}} : vector<[8]xi16> from vector<[8]x[8]xi16>1113 %tile = arm_sme.get_tile : vector<[8]x[8]xi16>1114 %slice = vector.extract %tile[%row] : vector<[8]xi16> from vector<[8]x[8]xi16>1115 return %slice : vector<[8]xi16>1116}1117 1118// -----1119 1120// CHECK-LABEL: @vector_extract_slice_i641121func.func @vector_extract_slice_i64(%row: index) -> vector<[2]xi64> {1122 // CHECK: arm_sme.extract_tile_slice {{.*}} : vector<[2]xi64> from vector<[2]x[2]xi64>1123 %tile = arm_sme.get_tile : vector<[2]x[2]xi64>1124 %slice = vector.extract %tile[%row] : vector<[2]xi64> from vector<[2]x[2]xi64>1125 return %slice : vector<[2]xi64>1126}1127 1128// -----1129 1130// CHECK-LABEL: @vector_extract_slice_i1281131func.func @vector_extract_slice_i128(%row: index) -> vector<[1]xi128> {1132 // CHECK: arm_sme.extract_tile_slice {{.*}} : vector<[1]xi128> from vector<[1]x[1]xi128>1133 %tile = arm_sme.get_tile : vector<[1]x[1]xi128>1134 %slice = vector.extract %tile[%row] : vector<[1]xi128> from vector<[1]x[1]xi128>1135 return %slice : vector<[1]xi128>1136}1137 1138// -----1139 1140// CHECK-LABEL: @vector_extract_slice_f161141func.func @vector_extract_slice_f16(%row: index) -> vector<[8]xf16> {1142 // CHECK: arm_sme.extract_tile_slice {{.*}} : vector<[8]xf16> from vector<[8]x[8]xf16>1143 %tile = arm_sme.get_tile : vector<[8]x[8]xf16>1144 %slice = vector.extract %tile[%row] : vector<[8]xf16> from vector<[8]x[8]xf16>1145 return %slice : vector<[8]xf16>1146}1147 1148// -----1149 1150// CHECK-LABEL: @vector_extract_slice_bf161151func.func @vector_extract_slice_bf16(%row: index) -> vector<[8]xbf16> {1152 // CHECK: arm_sme.extract_tile_slice {{.*}} : vector<[8]xbf16> from vector<[8]x[8]xbf16>1153 %tile = arm_sme.get_tile : vector<[8]x[8]xbf16>1154 %slice = vector.extract %tile[%row] : vector<[8]xbf16> from vector<[8]x[8]xbf16>1155 return %slice : vector<[8]xbf16>1156}1157 1158// -----1159 1160// CHECK-LABEL: @vector_extract_slice_f321161func.func @vector_extract_slice_f32(%row: index) -> vector<[4]xf32> {1162 // CHECK: arm_sme.extract_tile_slice {{.*}} : vector<[4]xf32> from vector<[4]x[4]xf32>1163 %tile = arm_sme.get_tile : vector<[4]x[4]xf32>1164 %slice = vector.extract %tile[%row] : vector<[4]xf32> from vector<[4]x[4]xf32>1165 return %slice : vector<[4]xf32>1166}1167 1168// -----1169 1170// CHECK-LABEL: @vector_extract_slice_f641171func.func @vector_extract_slice_f64(%row: index) -> vector<[2]xf64> {1172 // CHECK: arm_sme.extract_tile_slice {{.*}} : vector<[2]xf64> from vector<[2]x[2]xf64>1173 %tile = arm_sme.get_tile : vector<[2]x[2]xf64>1174 %slice = vector.extract %tile[%row] : vector<[2]xf64> from vector<[2]x[2]xf64>1175 return %slice : vector<[2]xf64>1176}1177 1178// -----1179 1180// CHECK-LABEL: @vector_extract_element(1181// CHECK-SAME: %[[ROW:.*]]: index,1182// CHECK-SAME: %[[COL:.*]]: index)1183func.func @vector_extract_element(%row: index, %col: index) -> i32 {1184 // CHECK-NEXT: %[[TILE:.*]] = arm_sme.get_tile : vector<[4]x[4]xi32>1185 // CHECK-NEXT: %[[SLICE:.*]] = arm_sme.extract_tile_slice %[[TILE]][%[[ROW]]] : vector<[4]xi32> from vector<[4]x[4]xi32>1186 // CHECK-NEXT: %[[EL:.*]] = vector.extract %[[SLICE]]{{\[}}%[[COL]]] : i32 from vector<[4]xi32>1187 %tile = arm_sme.get_tile : vector<[4]x[4]xi32>1188 %el = vector.extract %tile[%row, %col] : i32 from vector<[4]x[4]xi32>1189 return %el : i321190}1191 1192// -----1193 1194// CHECK-LABEL: @vector_extract_element_i81195func.func @vector_extract_element_i8(%row: index, %col: index) -> i8 {1196 // CHECK: %[[SLICE:.*]] = arm_sme.extract_tile_slice %{{.*}} : vector<[16]xi8> from vector<[16]x[16]xi8>1197 // CHECK-NEXT: %{{.*}} = vector.extract %[[SLICE]]{{\[}}%{{.*}}] : i8 from vector<[16]xi8>1198 %tile = arm_sme.get_tile : vector<[16]x[16]xi8>1199 %el = vector.extract %tile[%row, %col] : i8 from vector<[16]x[16]xi8>1200 return %el : i81201}1202 1203// -----1204 1205// CHECK-LABEL: @vector_extract_element_i161206func.func @vector_extract_element_i16(%row: index, %col: index) -> i16 {1207 // CHECK: %[[SLICE:.*]] = arm_sme.extract_tile_slice %{{.*}} : vector<[8]xi16> from vector<[8]x[8]xi16>1208 // CHECK-NEXT: %{{.*}} = vector.extract %[[SLICE]]{{\[}}%{{.*}}] : i16 from vector<[8]xi16>1209 %tile = arm_sme.get_tile : vector<[8]x[8]xi16>1210 %el = vector.extract %tile[%row, %col] : i16 from vector<[8]x[8]xi16>1211 return %el : i161212}1213 1214// -----1215 1216// CHECK-LABEL: @vector_extract_element_i641217func.func @vector_extract_element_i64(%row: index, %col: index) -> i64 {1218 // CHECK: %[[SLICE:.*]] = arm_sme.extract_tile_slice %{{.*}} : vector<[2]xi64> from vector<[2]x[2]xi64>1219 // CHECK-NEXT: %{{.*}} = vector.extract %[[SLICE]]{{\[}}%{{.*}}] : i64 from vector<[2]xi64>1220 %tile = arm_sme.get_tile : vector<[2]x[2]xi64>1221 %el = vector.extract %tile[%row, %col] : i64 from vector<[2]x[2]xi64>1222 return %el : i641223}1224 1225// -----1226 1227// CHECK-LABEL: @vector_extract_element_i1281228func.func @vector_extract_element_i128(%row: index, %col: index) -> i128 {1229 // CHECK: %[[SLICE:.*]] = arm_sme.extract_tile_slice %{{.*}} : vector<[1]xi128> from vector<[1]x[1]xi128>1230 // CHECK-NEXT: %{{.*}} = vector.extract %[[SLICE]]{{\[}}%{{.*}}] : i128 from vector<[1]xi128>1231 %tile = arm_sme.get_tile : vector<[1]x[1]xi128>1232 %el = vector.extract %tile[%row, %col] : i128 from vector<[1]x[1]xi128>1233 return %el : i1281234}1235 1236// -----1237 1238// CHECK-LABEL: @vector_extract_element_f161239func.func @vector_extract_element_f16(%row: index, %col: index) -> f16 {1240 // CHECK: %[[SLICE:.*]] = arm_sme.extract_tile_slice %{{.*}} : vector<[8]xf16> from vector<[8]x[8]xf16>1241 // CHECK-NEXT: %{{.*}} = vector.extract %[[SLICE]]{{\[}}%{{.*}}] : f16 from vector<[8]xf16>1242 %tile = arm_sme.get_tile : vector<[8]x[8]xf16>1243 %el = vector.extract %tile[%row, %col] : f16 from vector<[8]x[8]xf16>1244 return %el : f161245}1246 1247// -----1248 1249// CHECK-LABEL: @vector_extract_element_bf161250func.func @vector_extract_element_bf16(%row: index, %col: index) -> bf16 {1251 // CHECK: %[[SLICE:.*]] = arm_sme.extract_tile_slice %{{.*}} : vector<[8]xbf16> from vector<[8]x[8]xbf16>1252 // CHECK-NEXT: %{{.*}} = vector.extract %[[SLICE]]{{\[}}%{{.*}}] : bf16 from vector<[8]xbf16>1253 %tile = arm_sme.get_tile : vector<[8]x[8]xbf16>1254 %el = vector.extract %tile[%row, %col] : bf16 from vector<[8]x[8]xbf16>1255 return %el : bf161256}1257 1258// -----1259 1260// CHECK-LABEL: @vector_extract_element_f321261func.func @vector_extract_element_f32(%row: index, %col: index) -> f32 {1262 // CHECK: %[[SLICE:.*]] = arm_sme.extract_tile_slice %{{.*}} : vector<[4]xf32> from vector<[4]x[4]xf32>1263 // CHECK-NEXT: %{{.*}} = vector.extract %[[SLICE]]{{\[}}%{{.*}}] : f32 from vector<[4]xf32>1264 %tile = arm_sme.get_tile : vector<[4]x[4]xf32>1265 %el = vector.extract %tile[%row, %col] : f32 from vector<[4]x[4]xf32>1266 return %el : f321267}1268 1269// -----1270 1271// CHECK-LABEL: @vector_extract_element_f641272func.func @vector_extract_element_f64(%row: index, %col: index) -> f64 {1273 // CHECK: %[[SLICE:.*]] = arm_sme.extract_tile_slice %{{.*}} : vector<[2]xf64> from vector<[2]x[2]xf64>1274 // CHECK-NEXT: %{{.*}} = vector.extract %[[SLICE]]{{\[}}%{{.*}}] : f64 from vector<[2]xf64>1275 %tile = arm_sme.get_tile : vector<[2]x[2]xf64>1276 %el = vector.extract %tile[%row, %col] : f64 from vector<[2]x[2]xf64>1277 return %el : f641278}1279 1280//===----------------------------------------------------------------------===//1281// vector.extract --> arm_sve.psel1282//===----------------------------------------------------------------------===//1283 1284// -----1285 1286// CHECK-LABEL: @dynamic_vector_extract_mask_to_psel(1287// CHECK-SAME: %[[A:.*]]: index, %[[B:.*]]: index, %[[INDEX:.*]]: index)1288func.func @dynamic_vector_extract_mask_to_psel(%a: index, %b: index, %index: index) -> vector<[8]xi1>1289{1290 // CHECK: %[[MASK_ROWS:.*]] = vector.create_mask %[[A]] : vector<[4]xi1>1291 // CHECK: %[[MASK_COLS:.*]] = vector.create_mask %[[B]] : vector<[8]xi1>1292 // CHECK: arm_sve.psel %[[MASK_COLS]], %[[MASK_ROWS]][%[[INDEX]]] : vector<[8]xi1>, vector<[4]xi1>1293 %mask = vector.create_mask %a, %b : vector<[4]x[8]xi1>1294 %slice = vector.extract %mask[%index] : vector<[8]xi1> from vector<[4]x[8]xi1>1295 return %slice : vector<[8]xi1>1296}1297 1298// -----1299 1300// CHECK-LABEL: @vector_extract_mask_to_psel(1301// CHECK-SAME: %[[A:.*]]: index,1302// CHECK-SAME: %[[B:.*]]: index)1303func.func @vector_extract_mask_to_psel(%a: index, %b: index) -> vector<[2]xi1>1304{1305 // CHECK: %[[C1:.*]] = arith.constant 1 : index1306 // CHECK: %[[MASK_ROWS:.*]] = vector.create_mask %[[A]] : vector<[16]xi1>1307 // CHECK: %[[MASK_COLS:.*]] = vector.create_mask %[[B]] : vector<[2]xi1>1308 // CHECK: arm_sve.psel %[[MASK_COLS]], %[[MASK_ROWS]][%[[C1]]] : vector<[2]xi1>, vector<[16]xi1>1309 %mask = vector.create_mask %a, %b : vector<[16]x[2]xi1>1310 %slice = vector.extract %mask[1] : vector<[2]xi1> from vector<[16]x[2]xi1>1311 return %slice : vector<[2]xi1>1312}1313