1// RUN: mlir-opt %s -pass-pipeline='func.func(canonicalize)' -split-input-file -allow-unregistered-dialect | FileCheck %s 2 3// ----- 4 5// CHECK-LABEL: create_vector_mask_to_constant_mask 6func.func @create_vector_mask_to_constant_mask() -> (vector<4x3xi1>) { 7 %c2 = arith.constant 2 : index 8 %c3 = arith.constant 3 : index 9 // CHECK: vector.constant_mask [3, 2] : vector<4x3xi1> 10 %0 = vector.create_mask %c3, %c2 : vector<4x3xi1> 11 return %0 : vector<4x3xi1> 12} 13 14// ----- 15 16// CHECK-LABEL: create_scalable_vector_mask_to_constant_mask 17func.func @create_scalable_vector_mask_to_constant_mask() -> (vector<[8]xi1>) { 18 %c-1 = arith.constant -1 : index 19 // CHECK: vector.constant_mask [0] : vector<[8]xi1> 20 %0 = vector.create_mask %c-1 : vector<[8]xi1> 21 return %0 : vector<[8]xi1> 22} 23 24// ----- 25 26// CHECK-LABEL: create_vector_mask_to_constant_mask_truncation 27func.func @create_vector_mask_to_constant_mask_truncation() -> (vector<4x3xi1>) { 28 %c2 = arith.constant 2 : index 29 %c5 = arith.constant 5 : index 30 // CHECK: vector.constant_mask [4, 2] : vector<4x3xi1> 31 %0 = vector.create_mask %c5, %c2 : vector<4x3xi1> 32 return %0 : vector<4x3xi1> 33} 34 35// ----- 36 37// CHECK-LABEL: create_vector_mask_to_constant_mask_truncation_neg 38func.func @create_vector_mask_to_constant_mask_truncation_neg() -> (vector<4x3xi1>) { 39 %cneg2 = arith.constant -2 : index 40 %c5 = arith.constant 5 : index 41 // CHECK: vector.constant_mask [0, 0] : vector<4x3xi1> 42 %0 = vector.create_mask %c5, %cneg2 : vector<4x3xi1> 43 return %0 : vector<4x3xi1> 44} 45 46// ----- 47 48// CHECK-LABEL: create_vector_mask_to_constant_mask_truncation_zero 49func.func @create_vector_mask_to_constant_mask_truncation_zero() -> (vector<4x3xi1>) { 50 %c2 = arith.constant 2 : index 51 %c0 = arith.constant 0 : index 52 // CHECK: vector.constant_mask [0, 0] : vector<4x3xi1> 53 %0 = vector.create_mask %c0, %c2 : vector<4x3xi1> 54 return %0 : vector<4x3xi1> 55} 56 57// ----- 58 59func.func @extract_strided_slice_of_constant_mask() -> (vector<2x2xi1>) { 60 %0 = vector.constant_mask [2, 2] : vector<4x3xi1> 61 %1 = vector.extract_strided_slice %0 62 {offsets = [0, 0], sizes = [2, 2], strides = [1, 1]} 63 : vector<4x3xi1> to vector<2x2xi1> 64 // CHECK: vector.constant_mask [2, 2] : vector<2x2xi1> 65 return %1 : vector<2x2xi1> 66} 67 68// ----- 69 70func.func @extract_strided_slice_of_constant_mask() -> (vector<2x2xi1>) { 71 %0 = vector.constant_mask [2, 2] : vector<4x3xi1> 72 %1 = vector.extract_strided_slice %0 73 {offsets = [1, 0], sizes = [2, 2], strides = [1, 1]} 74 : vector<4x3xi1> to vector<2x2xi1> 75 // CHECK: vector.constant_mask [1, 2] : vector<2x2xi1> 76 return %1 : vector<2x2xi1> 77} 78 79// ----- 80 81func.func @extract_strided_slice_of_constant_mask() -> (vector<2x2xi1>) { 82 %0 = vector.constant_mask [2, 2] : vector<4x3xi1> 83 %1 = vector.extract_strided_slice %0 84 {offsets = [0, 1], sizes = [2, 2], strides = [1, 1]} 85 : vector<4x3xi1> to vector<2x2xi1> 86 // CHECK: vector.constant_mask [2, 1] : vector<2x2xi1> 87 return %1 : vector<2x2xi1> 88} 89 90// ----- 91 92func.func @extract_strided_slice_of_constant_mask() -> (vector<2x2xi1>) { 93 %0 = vector.constant_mask [2, 2] : vector<4x3xi1> 94 %1 = vector.extract_strided_slice %0 95 {offsets = [2, 0], sizes = [2, 2], strides = [1, 1]} 96 : vector<4x3xi1> to vector<2x2xi1> 97 // CHECK: vector.constant_mask [0, 0] : vector<2x2xi1> 98 return %1 : vector<2x2xi1> 99} 100 101// ----- 102 103func.func @extract_strided_slice_of_constant_mask() -> (vector<2x1xi1>) { 104 %0 = vector.constant_mask [2, 2] : vector<4x3xi1> 105 %1 = vector.extract_strided_slice %0 106 {offsets = [0, 2], sizes = [2, 1], strides = [1, 1]} 107 : vector<4x3xi1> to vector<2x1xi1> 108 // CHECK: vector.constant_mask [0, 0] : vector<2x1xi1> 109 return %1 : vector<2x1xi1> 110} 111 112// ----- 113 114func.func @extract_strided_slice_of_constant_mask() -> (vector<2x1xi1>) { 115 %0 = vector.constant_mask [2, 2] : vector<4x3xi1> 116 %1 = vector.extract_strided_slice %0 117 {offsets = [0, 1], sizes = [2, 1], strides = [1, 1]} 118 : vector<4x3xi1> to vector<2x1xi1> 119 // CHECK: vector.constant_mask [2, 1] : vector<2x1xi1> 120 return %1 : vector<2x1xi1> 121} 122 123// ----- 124 125func.func @extract_strided_slice_of_constant_mask() -> (vector<2x1xi1>) { 126 %0 = vector.constant_mask [2, 2] : vector<4x3xi1> 127 %1 = vector.extract_strided_slice %0 128 {offsets = [1, 1], sizes = [2, 1], strides = [1, 1]} 129 : vector<4x3xi1> to vector<2x1xi1> 130 // CHECK: vector.constant_mask [1, 1] : vector<2x1xi1> 131 return %1 : vector<2x1xi1> 132} 133 134// ----- 135 136// CHECK-LABEL: extract_strided_fold 137// CHECK-SAME: (%[[ARG:.*]]: vector<4x3xi1>) 138// CHECK-NEXT: return %[[ARG]] : vector<4x3xi1> 139func.func @extract_strided_fold(%arg : vector<4x3xi1>) -> (vector<4x3xi1>) { 140 %0 = vector.extract_strided_slice %arg 141 {offsets = [0, 0], sizes = [4, 3], strides = [1, 1]} 142 : vector<4x3xi1> to vector<4x3xi1> 143 return %0 : vector<4x3xi1> 144} 145 146// ----- 147 148// CHECK-LABEL: extract_strided_fold_insert 149// CHECK-SAME: (%[[ARG:.*]]: vector<4x4xf32> 150// CHECK-NEXT: return %[[ARG]] : vector<4x4xf32> 151func.func @extract_strided_fold_insert(%a: vector<4x4xf32>, %b: vector<8x16xf32>) 152 -> (vector<4x4xf32>) { 153 %0 = vector.insert_strided_slice %a, %b {offsets = [2, 2], strides = [1, 1]} 154 : vector<4x4xf32> into vector<8x16xf32> 155 %1 = vector.extract_strided_slice %0 156 {offsets = [2, 2], sizes = [4, 4], strides = [1, 1]} 157 : vector<8x16xf32> to vector<4x4xf32> 158 return %1 : vector<4x4xf32> 159} 160 161// ----- 162 163// Case where the vector inserted is a subset of the vector extracted. 164// CHECK-LABEL: extract_strided_fold_insert 165// CHECK-SAME: (%[[ARG0:.*]]: vector<6x4xf32> 166// CHECK-NEXT: %[[EXT:.*]] = vector.extract_strided_slice %[[ARG0]] 167// CHECK-SAME: {offsets = [0, 0], sizes = [4, 4], strides = [1, 1]} 168// CHECK-SAME: : vector<6x4xf32> to vector<4x4xf32> 169// CHECK-NEXT: return %[[EXT]] : vector<4x4xf32> 170func.func @extract_strided_fold_insert(%a: vector<6x4xf32>, %b: vector<8x16xf32>) 171 -> (vector<4x4xf32>) { 172 %0 = vector.insert_strided_slice %a, %b {offsets = [2, 2], strides = [1, 1]} 173 : vector<6x4xf32> into vector<8x16xf32> 174 %1 = vector.extract_strided_slice %0 175 {offsets = [2, 2], sizes = [4, 4], strides = [1, 1]} 176 : vector<8x16xf32> to vector<4x4xf32> 177 return %1 : vector<4x4xf32> 178} 179 180// ----- 181 182// Negative test where the extract is not a subset of the element inserted. 183// CHECK-LABEL: extract_strided_fold_negative 184// CHECK-SAME: (%[[ARG0:.*]]: vector<4x4xf32>, %[[ARG1:.*]]: vector<8x16xf32> 185// CHECK: %[[INS:.*]] = vector.insert_strided_slice %[[ARG0]], %[[ARG1]] 186// CHECK-SAME: {offsets = [2, 2], strides = [1, 1]} 187// CHECK-SAME: : vector<4x4xf32> into vector<8x16xf32> 188// CHECK: %[[EXT:.*]] = vector.extract_strided_slice %[[INS]] 189// CHECK-SAME: {offsets = [2, 2], sizes = [6, 4], strides = [1, 1]} 190// CHECK-SAME: : vector<8x16xf32> to vector<6x4xf32> 191// CHECK-NEXT: return %[[EXT]] : vector<6x4xf32> 192func.func @extract_strided_fold_negative(%a: vector<4x4xf32>, %b: vector<8x16xf32>) 193 -> (vector<6x4xf32>) { 194 %0 = vector.insert_strided_slice %a, %b {offsets = [2, 2], strides = [1, 1]} 195 : vector<4x4xf32> into vector<8x16xf32> 196 %1 = vector.extract_strided_slice %0 197 {offsets = [2, 2], sizes = [6, 4], strides = [1, 1]} 198 : vector<8x16xf32> to vector<6x4xf32> 199 return %1 : vector<6x4xf32> 200} 201 202// ----- 203 204// Case where we need to go through 2 level of insert element. 205// CHECK-LABEL: extract_strided_fold_insert 206// CHECK-SAME: (%[[ARG0:.*]]: vector<2x8xf32>, %[[ARG1:.*]]: vector<1x4xf32>, 207// CHECK-NEXT: %[[EXT:.*]] = vector.extract_strided_slice %[[ARG1]] 208// CHECK-SAME: {offsets = [0, 0], sizes = [1, 1], strides = [1, 1]} 209// CHECK-SAME: : vector<1x4xf32> to vector<1x1xf32> 210// CHECK-NEXT: return %[[EXT]] : vector<1x1xf32> 211func.func @extract_strided_fold_insert(%a: vector<2x8xf32>, %b: vector<1x4xf32>, 212 %c : vector<1x4xf32>) -> (vector<1x1xf32>) { 213 %0 = vector.insert_strided_slice %b, %a {offsets = [0, 1], strides = [1, 1]} 214 : vector<1x4xf32> into vector<2x8xf32> 215 %1 = vector.insert_strided_slice %c, %0 {offsets = [1, 0], strides = [1, 1]} 216 : vector<1x4xf32> into vector<2x8xf32> 217 %2 = vector.extract_strided_slice %1 218 {offsets = [0, 1], sizes = [1, 1], strides = [1, 1]} 219 : vector<2x8xf32> to vector<1x1xf32> 220 return %2 : vector<1x1xf32> 221} 222 223// ----- 224 225// CHECK-LABEL: transpose_1D_identity 226// CHECK-SAME: ([[ARG:%.*]]: vector<4xf32>) 227func.func @transpose_1D_identity(%arg : vector<4xf32>) -> vector<4xf32> { 228 // CHECK-NOT: transpose 229 %0 = vector.transpose %arg, [0] : vector<4xf32> to vector<4xf32> 230 // CHECK-NEXT: return [[ARG]] 231 return %0 : vector<4xf32> 232} 233 234// ----- 235 236// CHECK-LABEL: transpose_2D_identity 237// CHECK-SAME: ([[ARG:%.*]]: vector<4x3xf32>) 238func.func @transpose_2D_identity(%arg : vector<4x3xf32>) -> vector<4x3xf32> { 239 // CHECK-NOT: transpose 240 %0 = vector.transpose %arg, [0, 1] : vector<4x3xf32> to vector<4x3xf32> 241 // CHECK-NEXT: return [[ARG]] 242 return %0 : vector<4x3xf32> 243} 244 245// ----- 246 247// CHECK-LABEL: transpose_3D_identity 248// CHECK-SAME: ([[ARG:%.*]]: vector<4x3x2xf32>) 249func.func @transpose_3D_identity(%arg : vector<4x3x2xf32>) -> vector<4x3x2xf32> { 250 // CHECK-NOT: transpose 251 %0 = vector.transpose %arg, [0, 1, 2] : vector<4x3x2xf32> to vector<4x3x2xf32> 252 // CHECK-NEXT: return [[ARG]] 253 return %0 : vector<4x3x2xf32> 254} 255 256// ----- 257 258// CHECK-LABEL: transpose_2D_sequence 259// CHECK-SAME: ([[ARG:%.*]]: vector<4x3xf32>) 260func.func @transpose_2D_sequence(%arg : vector<4x3xf32>) -> vector<4x3xf32> { 261 // CHECK-NOT: transpose 262 %0 = vector.transpose %arg, [1, 0] : vector<4x3xf32> to vector<3x4xf32> 263 %1 = vector.transpose %0, [0, 1] : vector<3x4xf32> to vector<3x4xf32> 264 %2 = vector.transpose %1, [1, 0] : vector<3x4xf32> to vector<4x3xf32> 265 %3 = vector.transpose %2, [0, 1] : vector<4x3xf32> to vector<4x3xf32> 266 // CHECK: [[ADD:%.*]] = arith.addf [[ARG]], [[ARG]] 267 %4 = arith.addf %2, %3 : vector<4x3xf32> 268 // CHECK-NEXT: return [[ADD]] 269 return %4 : vector<4x3xf32> 270} 271 272// ----- 273 274// CHECK-LABEL: transpose_3D_sequence 275// CHECK-SAME: ([[ARG:%.*]]: vector<4x3x2xf32>) 276func.func @transpose_3D_sequence(%arg : vector<4x3x2xf32>) -> vector<4x3x2xf32> { 277 // CHECK: [[T0:%.*]] = vector.transpose [[ARG]], [2, 1, 0] 278 %0 = vector.transpose %arg, [1, 2, 0] : vector<4x3x2xf32> to vector<3x2x4xf32> 279 %1 = vector.transpose %0, [1, 0, 2] : vector<3x2x4xf32> to vector<2x3x4xf32> 280 // CHECK: [[T1:%.*]] = vector.transpose %arg0, [2, 1, 0] 281 %2 = vector.transpose %1, [2, 1, 0] : vector<2x3x4xf32> to vector<4x3x2xf32> 282 %3 = vector.transpose %2, [2, 1, 0] : vector<4x3x2xf32> to vector<2x3x4xf32> 283 // CHECK: [[MUL:%.*]] = arith.mulf [[T0]], [[T1]] 284 %4 = arith.mulf %1, %3 : vector<2x3x4xf32> 285 // CHECK: [[T5:%.*]] = vector.transpose [[MUL]], [2, 1, 0] 286 %5 = vector.transpose %4, [2, 1, 0] : vector<2x3x4xf32> to vector<4x3x2xf32> 287 // CHECK-NOT: transpose 288 %6 = vector.transpose %3, [2, 1, 0] : vector<2x3x4xf32> to vector<4x3x2xf32> 289 // CHECK: [[ADD:%.*]] = arith.addf [[T5]], [[ARG]] 290 %7 = arith.addf %5, %6 : vector<4x3x2xf32> 291 // CHECK-NEXT: return [[ADD]] 292 return %7 : vector<4x3x2xf32> 293} 294 295// ----- 296 297// CHECK-LABEL: cast_transfers 298func.func @cast_transfers(%A: memref<4x8xf32>) -> (vector<4x8xf32>) { 299 %c0 = arith.constant 0 : index 300 %f0 = arith.constant 0.0 : f32 301 %0 = memref.cast %A : memref<4x8xf32> to memref<?x?xf32> 302 303 // CHECK: vector.transfer_read %{{.*}} {in_bounds = [true, true]} : memref<4x8xf32>, vector<4x8xf32> 304 %1 = vector.transfer_read %0[%c0, %c0], %f0 : memref<?x?xf32>, vector<4x8xf32> 305 306 // CHECK: vector.transfer_write %{{.*}} {in_bounds = [true, true]} : vector<4x8xf32>, memref<4x8xf32> 307 vector.transfer_write %1, %0[%c0, %c0] : vector<4x8xf32>, memref<?x?xf32> 308 return %1 : vector<4x8xf32> 309} 310 311// ----- 312 313// CHECK-LABEL: cast_transfers 314func.func @cast_transfers(%A: tensor<4x8xf32>) -> (vector<4x8xf32>) { 315 %c0 = arith.constant 0 : index 316 %f0 = arith.constant 0.0 : f32 317 %0 = tensor.cast %A : tensor<4x8xf32> to tensor<?x?xf32> 318 319 // CHECK: vector.transfer_read %{{.*}} {in_bounds = [true, true]} : tensor<4x8xf32>, vector<4x8xf32> 320 %1 = vector.transfer_read %0[%c0, %c0], %f0 : tensor<?x?xf32>, vector<4x8xf32> 321 322 return %1 : vector<4x8xf32> 323} 324 325// ----- 326 327// CHECK-LABEL: func @insert_extract_transpose_2d( 328// CHECK-SAME: %[[V:[a-zA-Z0-9]*]]: vector<2x3xf32>, 329// CHECK-SAME: %[[F0:[a-zA-Z0-9]*]]: f32, 330// CHECK-SAME: %[[F1:[a-zA-Z0-9]*]]: f32, 331// CHECK-SAME: %[[F2:[a-zA-Z0-9]*]]: f32, 332// CHECK-SAME: %[[F3:[a-zA-Z0-9]*]]: f32 333func.func @insert_extract_transpose_2d( 334 %v: vector<2x3xf32>, %f0: f32, %f1: f32, %f2: f32, %f3: f32) 335-> (f32, f32, f32) 336{ 337 %0 = vector.insert %f0, %v[0, 0] : f32 into vector<2x3xf32> 338 %1 = vector.insert %f1, %0[0, 1] : f32 into vector<2x3xf32> 339 %2 = vector.insert %f2, %1[1, 0] : f32 into vector<2x3xf32> 340 %3 = vector.insert %f3, %2[1, 1] : f32 into vector<2x3xf32> 341 %4 = vector.transpose %3, [1, 0] : vector<2x3xf32> to vector<3x2xf32> 342 %5 = vector.insert %f3, %4[1, 0] : f32 into vector<3x2xf32> 343 %6 = vector.transpose %5, [1, 0] : vector<3x2xf32> to vector<2x3xf32> 344 345 // Expected %f2 from %2 = vector.insert %f2, %1[1, 0]. 346 %r1 = vector.extract %3[1, 0] : vector<2x3xf32> 347 348 // Expected %f1 from %1 = vector.insert %f1, %0[0, 1] followed by 349 // transpose [1, 0]. 350 %r2 = vector.extract %4[1, 0] : vector<3x2xf32> 351 352 // Expected %f2 from %2 = vector.insert %f2, %1[1, 0] followed by double 353 // transpose [1, 0]. 354 %r3 = vector.extract %6[1, 0] : vector<2x3xf32> 355 356 // CHECK-NEXT: return %[[F2]], %[[F1]], %[[F2]] : f32, f32, f32 357 return %r1, %r2, %r3 : f32, f32, f32 358} 359 360// ----- 361 362// CHECK-LABEL: insert_extract_chain 363// CHECK-SAME: %[[V234:[a-zA-Z0-9]*]]: vector<2x3x4xf32> 364// CHECK-SAME: %[[V34:[a-zA-Z0-9]*]]: vector<3x4xf32> 365// CHECK-SAME: %[[V4:[a-zA-Z0-9]*]]: vector<4xf32> 366func.func @insert_extract_chain(%v234: vector<2x3x4xf32>, %v34: vector<3x4xf32>, %v4: vector<4xf32>) 367 -> (vector<4xf32>, vector<4xf32>, vector<3x4xf32>, vector<3x4xf32>) { 368 // CHECK-NEXT: %[[A34:.*]] = vector.insert 369 %A34 = vector.insert %v34, %v234[0]: vector<3x4xf32> into vector<2x3x4xf32> 370 // CHECK-NEXT: %[[B34:.*]] = vector.insert 371 %B34 = vector.insert %v34, %A34[1]: vector<3x4xf32> into vector<2x3x4xf32> 372 // CHECK-NEXT: %[[A4:.*]] = vector.insert 373 %A4 = vector.insert %v4, %B34[1, 0]: vector<4xf32> into vector<2x3x4xf32> 374 // CHECK-NEXT: %[[B4:.*]] = vector.insert 375 %B4 = vector.insert %v4, %A4[1, 1]: vector<4xf32> into vector<2x3x4xf32> 376 377 // Case 2.a. [1, 1] == insertpos ([1, 1]) 378 // Match %A4 insertionpos and fold to its source(i.e. %V4). 379 %r0 = vector.extract %B4[1, 1]: vector<2x3x4xf32> 380 381 // Case 3.a. insertpos ([1]) is a prefix of [1, 0]. 382 // Traverse %B34 to its source(i.e. %V34@[*0*]). 383 // CHECK-NEXT: %[[R1:.*]] = vector.extract %[[V34]][0] 384 %r1 = vector.extract %B34[1, 0]: vector<2x3x4xf32> 385 386 // Case 4. [1] is a prefix of insertpos ([1, 1]). 387 // Cannot traverse %B4. 388 // CHECK-NEXT: %[[R2:.*]] = vector.extract %[[B4]][1] 389 %r2 = vector.extract %B4[1]: vector<2x3x4xf32> 390 391 // Case 5. [0] is disjoint from insertpos ([1, 1]). 392 // Traverse %B4 to its dest(i.e. %A4@[0]). 393 // Traverse %A4 to its dest(i.e. %B34@[0]). 394 // Traverse %B34 to its dest(i.e. %A34@[0]). 395 // Match %A34 insertionpos and fold to its source(i.e. %V34). 396 %r3 = vector.extract %B4[0]: vector<2x3x4xf32> 397 398 // CHECK: return %[[V4]], %[[R1]], %[[R2]], %[[V34]] 399 return %r0, %r1, %r2, %r3: 400 vector<4xf32>, vector<4xf32>, vector<3x4xf32>, vector<3x4xf32> 401} 402 403// ----- 404 405// CHECK-LABEL: func @insert_extract_transpose_3d( 406// CHECK-SAME: %[[V234:[a-zA-Z0-9]*]]: vector<2x3x4xf32> 407func.func @insert_extract_transpose_3d( 408 %v234: vector<2x3x4xf32>, %v43: vector<4x3xf32>, %f0: f32) 409 -> (vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<3x4xf32>) { 410 411 %a432 = vector.transpose %v234, [2, 1, 0] : vector<2x3x4xf32> to vector<4x3x2xf32> 412 %b432 = vector.insert %f0, %a432[0, 0, 1] : f32 into vector<4x3x2xf32> 413 %c234 = vector.transpose %b432, [2, 1, 0] : vector<4x3x2xf32> to vector<2x3x4xf32> 414 // Case 1. %c234 = transpose [2,1,0] posWithSentinels [1,2,-1] -> [-1,2,1] 415 // Case 5. %b432 = insert [0,0,1] (inter([.,2,1], [.,0,1]) == 0) prop to %v432 416 // Case 1. %a432 = transpose [2,1,0] posWithSentinels [-1,2,1] -> [1,2,-1] 417 // can extract directly from %v234, the rest folds. 418 // CHECK: %[[R0:.*]] = vector.extract %[[V234]][1, 2] 419 %r0 = vector.extract %c234[1, 2] : vector<2x3x4xf32> 420 421 // CHECK-NEXT: vector.transpose 422 // CHECK-NEXT: vector.insert 423 // CHECK-NEXT: %[[F234:.*]] = vector.transpose 424 %d432 = vector.transpose %v234, [2, 1, 0] : vector<2x3x4xf32> to vector<4x3x2xf32> 425 %e432 = vector.insert %f0, %d432[0, 2, 1] : f32 into vector<4x3x2xf32> 426 %f234 = vector.transpose %e432, [2, 1, 0] : vector<4x3x2xf32> to vector<2x3x4xf32> 427 // Case 1. %c234 = transpose [2,1,0] posWithSentinels [1,2,-1] -> [-1,2,1] 428 // Case 4. %b432 = insert [0,0,1] (inter([.,2,1], [.,2,1]) != 0) 429 // Bail, cannot do better than the current. 430 // CHECK: %[[R1:.*]] = vector.extract %[[F234]] 431 %r1 = vector.extract %f234[1, 2] : vector<2x3x4xf32> 432 433 // CHECK-NEXT: vector.transpose 434 // CHECK-NEXT: vector.insert 435 // CHECK-NEXT: %[[H234:.*]] = vector.transpose 436 %g243 = vector.transpose %v234, [0, 2, 1] : vector<2x3x4xf32> to vector<2x4x3xf32> 437 %h243 = vector.insert %v43, %g243[0] : vector<4x3xf32> into vector<2x4x3xf32> 438 %i234 = vector.transpose %h243, [0, 2, 1] : vector<2x4x3xf32> to vector<2x3x4xf32> 439 // Case 1. %i234 = transpose [0,2,1] posWithSentinels [0,-1,-2] -> [0,-2,-1] 440 // Case 3.b. %b432 = insert [0] is prefix of [0,.,.] but internal transpose. 441 // Bail, cannot do better than the current. 442 // CHECK: %[[R2:.*]] = vector.extract %[[H234]][0, 1] 443 %r2 = vector.extract %i234[0, 1] : vector<2x3x4xf32> 444 445 // CHECK-NEXT: vector.transpose 446 // CHECK-NEXT: vector.insert 447 // CHECK-NEXT: %[[K234:.*]] = vector.transpose 448 %j243 = vector.transpose %v234, [0, 2, 1] : vector<2x3x4xf32> to vector<2x4x3xf32> 449 %k243 = vector.insert %v43, %j243[0] : vector<4x3xf32> into vector<2x4x3xf32> 450 %l234 = vector.transpose %k243, [0, 2, 1] : vector<2x4x3xf32> to vector<2x3x4xf32> 451 // Case 1. %i234 = transpose [0,2,1] posWithSentinels [0,-1,-2] -> [0,-2,-1] 452 // Case 2.b. %b432 = insert [0] == [0,.,.] but internal transpose. 453 // Bail, cannot do better than the current. 454 // CHECK: %[[R3:.*]] = vector.extract %[[K234]][0] 455 %r3 = vector.extract %l234[0] : vector<2x3x4xf32> 456 457 // CHECK-NEXT: return %[[R0]], %[[R1]], %[[R2]], %[[R3]] 458 return %r0, %r1, %r2, %r3: vector<4xf32>, vector<4xf32>, vector<4xf32>, vector<3x4xf32> 459} 460 461// ----- 462 463// CHECK-LABEL: fold_extracts 464// CHECK-SAME: %[[A:[a-zA-Z0-9]*]]: vector<3x4x5x6xf32> 465func.func @fold_extracts(%a : vector<3x4x5x6xf32>) -> (f32, vector<4x5x6xf32>) { 466 %b = vector.extract %a[0] : vector<3x4x5x6xf32> 467 %c = vector.extract %b[1, 2] : vector<4x5x6xf32> 468 // CHECK-NEXT: vector.extract %[[A]][0, 1, 2, 3] : vector<3x4x5x6xf32> 469 %d = vector.extract %c[3] : vector<6xf32> 470 471 // CHECK-NEXT: vector.extract %[[A]][0] : vector<3x4x5x6xf32> 472 %e = vector.extract %a[0] : vector<3x4x5x6xf32> 473 474 // CHECK-NEXT: return 475 return %d, %e : f32, vector<4x5x6xf32> 476} 477 478// ----- 479 480// CHECK-LABEL: fold_extract_transpose 481// CHECK-SAME: %[[A:[a-zA-Z0-9]*]]: vector<3x4x5x6xf32> 482// CHECK-SAME: %[[B:[a-zA-Z0-9]*]]: vector<3x6x5x6xf32> 483func.func @fold_extract_transpose( 484 %a : vector<3x4x5x6xf32>, %b : vector<3x6x5x6xf32>) -> ( 485 vector<6xf32>, vector<6xf32>, vector<6xf32>) { 486 // [3] is a proper most minor identity map in transpose. 487 // Permutation is a self inverse and we have. 488 // [0, 2, 1] ^ -1 o [0, 1, 2] = [0, 2, 1] o [0, 1, 2] 489 // = [0, 2, 1] 490 // CHECK-NEXT: vector.extract %[[A]][0, 2, 1] : vector<3x4x5x6xf32> 491 %0 = vector.transpose %a, [0, 2, 1, 3] : vector<3x4x5x6xf32> to vector<3x5x4x6xf32> 492 %1 = vector.extract %0[0, 1, 2] : vector<3x5x4x6xf32> 493 494 // [3] is a proper most minor identity map in transpose. 495 // Permutation is a not self inverse and we have. 496 // [1, 2, 0] ^ -1 o [0, 1, 2] = [2, 0, 1] o [0, 1, 2] 497 // = [2, 0, 1] 498 // CHECK-NEXT: vector.extract %[[A]][2, 0, 1] : vector<3x4x5x6xf32> 499 %2 = vector.transpose %a, [1, 2, 0, 3] : vector<3x4x5x6xf32> to vector<4x5x3x6xf32> 500 %3 = vector.extract %2[0, 1, 2] : vector<4x5x3x6xf32> 501 502 // Not a minor identity map so intra-vector level has been permuted 503 // CHECK-NEXT: vector.transpose %[[B]], [0, 2, 3, 1] 504 // CHECK-NEXT: vector.extract %{{.*}}[0, 1, 2] 505 %4 = vector.transpose %b, [0, 2, 3, 1] : vector<3x6x5x6xf32> to vector<3x5x6x6xf32> 506 %5 = vector.extract %4[0, 1, 2] : vector<3x5x6x6xf32> 507 508 return %1, %3, %5 : vector<6xf32>, vector<6xf32>, vector<6xf32> 509} 510 511// ----- 512 513// CHECK-LABEL: fold_extract_broadcast 514// CHECK-SAME: %[[A:.*]]: f32 515// CHECK: return %[[A]] : f32 516func.func @fold_extract_broadcast(%a : f32) -> f32 { 517 %b = vector.broadcast %a : f32 to vector<1x2x4xf32> 518 %r = vector.extract %b[0, 1, 2] : vector<1x2x4xf32> 519 return %r : f32 520} 521 522// ----- 523 524// CHECK-LABEL: fold_extract_broadcast_negative 525// CHECK: vector.broadcast %{{.*}} : vector<1x1xf32> to vector<1x1x4xf32> 526// CHECK: vector.extract %{{.*}}[0, 0] : vector<1x1x4xf32> 527func.func @fold_extract_broadcast_negative(%a : vector<1x1xf32>) -> vector<4xf32> { 528 %b = vector.broadcast %a : vector<1x1xf32> to vector<1x1x4xf32> 529 %r = vector.extract %b[0, 0] : vector<1x1x4xf32> 530 return %r : vector<4xf32> 531} 532 533// ----- 534 535// CHECK-LABEL: fold_extract_splat 536// CHECK-SAME: %[[A:.*]]: f32 537// CHECK: return %[[A]] : f32 538func.func @fold_extract_splat(%a : f32) -> f32 { 539 %b = vector.splat %a : vector<1x2x4xf32> 540 %r = vector.extract %b[0, 1, 2] : vector<1x2x4xf32> 541 return %r : f32 542} 543 544// ----- 545 546// CHECK-LABEL: fold_extract_broadcast_vector 547// CHECK-SAME: %[[A:.*]]: vector<4xf32> 548// CHECK: return %[[A]] : vector<4xf32> 549func.func @fold_extract_broadcast_vector(%a : vector<4xf32>) -> vector<4xf32> { 550 %b = vector.broadcast %a : vector<4xf32> to vector<1x2x4xf32> 551 %r = vector.extract %b[0, 1] : vector<1x2x4xf32> 552 return %r : vector<4xf32> 553} 554 555// ----- 556 557// CHECK-LABEL: fold_extract_broadcast 558// CHECK-SAME: %[[A:.*]]: vector<4xf32> 559// CHECK: %[[R:.*]] = vector.extract %[[A]][2] : vector<4xf32> 560// CHECK: return %[[R]] : f32 561func.func @fold_extract_broadcast(%a : vector<4xf32>) -> f32 { 562 %b = vector.broadcast %a : vector<4xf32> to vector<1x2x4xf32> 563 %r = vector.extract %b[0, 1, 2] : vector<1x2x4xf32> 564 return %r : f32 565} 566 567// ----- 568 569// CHECK-LABEL: fold_extract_broadcast 570// CHECK: %[[B:.*]] = vector.broadcast %{{.*}} : f32 to vector<4xf32> 571// CHECK: return %[[B]] : vector<4xf32> 572func.func @fold_extract_broadcast(%a : f32) -> vector<4xf32> { 573 %b = vector.broadcast %a : f32 to vector<1x2x4xf32> 574 %r = vector.extract %b[0, 1] : vector<1x2x4xf32> 575 return %r : vector<4xf32> 576} 577 578// ----- 579 580// CHECK-LABEL: fold_extract_broadcast 581// CHECK-SAME: %[[A:.*]]: vector<1xf32> 582// CHECK: %[[R:.*]] = vector.broadcast %[[A]] : vector<1xf32> to vector<8xf32> 583// CHECK: return %[[R]] : vector<8xf32> 584func.func @fold_extract_broadcast(%a : vector<1xf32>) -> vector<8xf32> { 585 %b = vector.broadcast %a : vector<1xf32> to vector<1x8xf32> 586 %r = vector.extract %b[0] : vector<1x8xf32> 587 return %r : vector<8xf32> 588} 589 590// ----- 591 592// CHECK-LABEL: func @fold_extract_shapecast 593// CHECK-SAME: (%[[A0:.*]]: vector<5x1x3x2xf32>, %[[A1:.*]]: vector<8x4x2xf32> 594// CHECK: %[[R0:.*]] = vector.extract %[[A0]][1, 0, 1, 1] : vector<5x1x3x2xf32> 595// CHECK: %[[R1:.*]] = vector.extract %[[A0]][1, 0, 2] : vector<5x1x3x2xf32> 596// CHECK: %[[R2:.*]] = vector.extract %[[A1]][7] : vector<8x4x2xf32> 597// CHECK: return %[[R0]], %[[R1]], %[[R2]], %[[A1]] : f32, vector<2xf32>, vector<4x2xf32>, vector<8x4x2xf32> 598func.func @fold_extract_shapecast(%arg0 : vector<5x1x3x2xf32>, 599 %arg1 : vector<8x4x2xf32>) 600 -> (f32, vector<2xf32>, vector<4x2xf32>, vector<8x4x2xf32>) { 601 %0 = vector.shape_cast %arg0 : vector<5x1x3x2xf32> to vector<15x2xf32> 602 %1 = vector.shape_cast %arg1 : vector<8x4x2xf32> to vector<4x2x4x2xf32> 603 %2 = vector.shape_cast %arg1 : vector<8x4x2xf32> to vector<1x8x4x2xf32> 604 %r1 = vector.extract %0[4, 1] : vector<15x2xf32> 605 %r2 = vector.extract %0[5] : vector<15x2xf32> 606 %r3 = vector.extract %1[3, 1] : vector<4x2x4x2xf32> 607 %r4 = vector.extract %2[0] : vector<1x8x4x2xf32> 608 return %r1, %r2, %r3, %r4 : f32, vector<2xf32>, vector<4x2xf32>, vector<8x4x2xf32> 609} 610 611// ----- 612 613// CHECK-LABEL: fold_extract_shapecast_negative 614// CHECK: %[[V:.*]] = vector.shape_cast %{{.*}} : vector<16xf32> to vector<2x4x2xf32> 615// CHECK: %[[R:.*]] = vector.extract %[[V]][1] : vector<2x4x2xf32> 616// CHECK: return %[[R]] : vector<4x2xf32> 617func.func @fold_extract_shapecast_negative(%arg0 : vector<16xf32>, 618 %arg1 : vector<8x4x2xf32>) -> vector<4x2xf32> { 619 %0 = vector.shape_cast %arg0 : vector<16xf32> to vector<2x4x2xf32> 620 %r = vector.extract %0[1] : vector<2x4x2xf32> 621 return %r : vector<4x2xf32> 622} 623 624// ----- 625 626// CHECK-LABEL: dont_fold_expand_collapse 627// CHECK: %[[A:.*]] = vector.shape_cast %{{.*}} : vector<1x1x64xf32> to vector<1x1x8x8xf32> 628// CHECK: %[[B:.*]] = vector.shape_cast %{{.*}} : vector<1x1x8x8xf32> to vector<8x8xf32> 629// CHECK: return %[[B]] : vector<8x8xf32> 630func.func @dont_fold_expand_collapse(%arg0: vector<1x1x64xf32>) -> vector<8x8xf32> { 631 %0 = vector.shape_cast %arg0 : vector<1x1x64xf32> to vector<1x1x8x8xf32> 632 %1 = vector.shape_cast %0 : vector<1x1x8x8xf32> to vector<8x8xf32> 633 return %1 : vector<8x8xf32> 634} 635 636// ----- 637 638// CHECK-LABEL: func @fold_broadcast_shapecast 639// CHECK-SAME: (%[[V:.+]]: vector<4xf32>) 640// CHECK: return %[[V]] 641func.func @fold_broadcast_shapecast(%arg0: vector<4xf32>) -> vector<4xf32> { 642 %0 = vector.broadcast %arg0 : vector<4xf32> to vector<1x1x4xf32> 643 %1 = vector.shape_cast %0 : vector<1x1x4xf32> to vector<4xf32> 644 return %1 : vector<4xf32> 645} 646 647// ----- 648 649// CHECK-LABEL: func @dont_fold_broadcast_shapecast_scalar 650// CHECK: vector.broadcast 651// CHECK: vector.shape_cast 652func.func @dont_fold_broadcast_shapecast_scalar(%arg0: f32) -> vector<1xf32> { 653 %0 = vector.broadcast %arg0 : f32 to vector<1x1x1xf32> 654 %1 = vector.shape_cast %0 : vector<1x1x1xf32> to vector<1xf32> 655 return %1 : vector<1xf32> 656} 657 658// ----- 659 660// CHECK-LABEL: func @dont_fold_broadcast_shapecast_diff_shape 661// CHECK: vector.broadcast 662// CHECK: vector.shape_cast 663func.func @dont_fold_broadcast_shapecast_diff_shape(%arg0: vector<4xf32>) -> vector<8xf32> { 664 %0 = vector.broadcast %arg0 : vector<4xf32> to vector<1x2x4xf32> 665 %1 = vector.shape_cast %0 : vector<1x2x4xf32> to vector<8xf32> 666 return %1 : vector<8xf32> 667} 668 669// ----- 670 671// CHECK-LABEL: fold_vector_transfers 672func.func @fold_vector_transfers(%A: memref<?x8xf32>) -> (vector<4x8xf32>, vector<4x9xf32>) { 673 %c0 = arith.constant 0 : index 674 %f0 = arith.constant 0.0 : f32 675 676 // CHECK: vector.transfer_read %{{.*}} {in_bounds = [false, true]} 677 %1 = vector.transfer_read %A[%c0, %c0], %f0 : memref<?x8xf32>, vector<4x8xf32> 678 679 // CHECK: vector.transfer_write %{{.*}} {in_bounds = [false, true]} 680 vector.transfer_write %1, %A[%c0, %c0] : vector<4x8xf32>, memref<?x8xf32> 681 682 // Both dims may be out-of-bounds, attribute is elided. 683 // CHECK: vector.transfer_read %{{.*}} 684 // CHECK-NOT: in_bounds 685 %2 = vector.transfer_read %A[%c0, %c0], %f0 : memref<?x8xf32>, vector<4x9xf32> 686 687 // Both dims may be out-of-bounds, attribute is elided. 688 // CHECK: vector.transfer_write %{{.*}} 689 // CHECK-NOT: in_bounds 690 vector.transfer_write %2, %A[%c0, %c0] : vector<4x9xf32>, memref<?x8xf32> 691 692 // CHECK: return 693 return %1, %2 : vector<4x8xf32>, vector<4x9xf32> 694} 695 696// ----- 697 698// CHECK-LABEL: bitcast_folding 699// CHECK-SAME: %[[A:.*]]: vector<4x8xf32> 700// CHECK-SAME: %[[B:.*]]: vector<2xi32> 701// CHECK: return %[[A]], %[[B]] : vector<4x8xf32>, vector<2xi32> 702func.func @bitcast_folding(%I1: vector<4x8xf32>, %I2: vector<2xi32>) -> (vector<4x8xf32>, vector<2xi32>) { 703 %0 = vector.bitcast %I1 : vector<4x8xf32> to vector<4x8xf32> 704 %1 = vector.bitcast %I2 : vector<2xi32> to vector<4xi16> 705 %2 = vector.bitcast %1 : vector<4xi16> to vector<2xi32> 706 return %0, %2 : vector<4x8xf32>, vector<2xi32> 707} 708 709// CHECK-LABEL: func @bitcast_f16_to_f32 710// bit pattern: 0x40004000 711// CHECK-DAG: %[[CST1:.+]] = arith.constant dense<2.00390625> : vector<4xf32> 712// bit pattern: 0x00000000 713// CHECK-DAG: %[[CST0:.+]] = arith.constant dense<0.000000e+00> : vector<4xf32> 714// CHECK: return %[[CST0]], %[[CST1]] 715func.func @bitcast_f16_to_f32() -> (vector<4xf32>, vector<4xf32>) { 716 %cst0 = arith.constant dense<0.0> : vector<8xf16> // bit pattern: 0x0000 717 %cst1 = arith.constant dense<2.0> : vector<8xf16> // bit pattern: 0x4000 718 %cast0 = vector.bitcast %cst0: vector<8xf16> to vector<4xf32> 719 %cast1 = vector.bitcast %cst1: vector<8xf16> to vector<4xf32> 720 return %cast0, %cast1: vector<4xf32>, vector<4xf32> 721} 722 723// ----- 724 725// CHECK-LABEL: broadcast_folding1 726// CHECK: %[[CST:.*]] = arith.constant dense<42> : vector<4xi32> 727// CHECK-NOT: vector.broadcast 728// CHECK: return %[[CST]] 729func.func @broadcast_folding1() -> vector<4xi32> { 730 %0 = arith.constant 42 : i32 731 %1 = vector.broadcast %0 : i32 to vector<4xi32> 732 return %1 : vector<4xi32> 733} 734 735// ----- 736 737// CHECK-LABEL: @broadcast_folding2 738// CHECK: %[[CST:.*]] = arith.constant dense<42> : vector<4x16xi32> 739// CHECK-NOT: vector.broadcast 740// CHECK: return %[[CST]] 741func.func @broadcast_folding2() -> vector<4x16xi32> { 742 %0 = arith.constant 42 : i32 743 %1 = vector.broadcast %0 : i32 to vector<16xi32> 744 %2 = vector.broadcast %1 : vector<16xi32> to vector<4x16xi32> 745 return %2 : vector<4x16xi32> 746} 747 748// ----- 749 750// CHECK-LABEL: @fold_consecutive_broadcasts( 751// CHECK-SAME: %[[ARG0:.*]]: i32 752// CHECK: %[[RESULT:.*]] = vector.broadcast %[[ARG0]] : i32 to vector<4x16xi32> 753// CHECK: return %[[RESULT]] 754func.func @fold_consecutive_broadcasts(%a : i32) -> vector<4x16xi32> { 755 %1 = vector.broadcast %a : i32 to vector<16xi32> 756 %2 = vector.broadcast %1 : vector<16xi32> to vector<4x16xi32> 757 return %2 : vector<4x16xi32> 758} 759 760// ----- 761 762// CHECK-LABEL: shape_cast_constant 763// CHECK-DAG: %[[CST1:.*]] = arith.constant dense<1> : vector<3x4x2xi32> 764// CHECK-DAG: %[[CST0:.*]] = arith.constant dense<2.000000e+00> : vector<20x2xf32> 765// CHECK: return %[[CST0]], %[[CST1]] : vector<20x2xf32>, vector<3x4x2xi32> 766func.func @shape_cast_constant() -> (vector<20x2xf32>, vector<3x4x2xi32>) { 767 %cst = arith.constant dense<2.000000e+00> : vector<5x4x2xf32> 768 %cst_1 = arith.constant dense<1> : vector<12x2xi32> 769 %0 = vector.shape_cast %cst : vector<5x4x2xf32> to vector<20x2xf32> 770 %1 = vector.shape_cast %cst_1 : vector<12x2xi32> to vector<3x4x2xi32> 771 return %0, %1 : vector<20x2xf32>, vector<3x4x2xi32> 772} 773 774// ----- 775 776// CHECK-LABEL: extract_strided_constant 777// CHECK-DAG: %[[CST1:.*]] = arith.constant dense<1> : vector<2x13x3xi32> 778// CHECK-DAG: %[[CST0:.*]] = arith.constant dense<2.000000e+00> : vector<12x2xf32> 779// CHECK: return %[[CST0]], %[[CST1]] : vector<12x2xf32>, vector<2x13x3xi32> 780func.func @extract_strided_constant() -> (vector<12x2xf32>, vector<2x13x3xi32>) { 781 %cst = arith.constant dense<2.000000e+00> : vector<29x7xf32> 782 %cst_1 = arith.constant dense<1> : vector<4x37x9xi32> 783 %0 = vector.extract_strided_slice %cst 784 {offsets = [2, 3], sizes = [12, 2], strides = [1, 1]} 785 : vector<29x7xf32> to vector<12x2xf32> 786 %1 = vector.extract_strided_slice %cst_1 787 {offsets = [1, 2, 5], sizes = [2, 13, 3], strides = [1, 1, 1]} 788 : vector<4x37x9xi32> to vector<2x13x3xi32> 789 return %0, %1 : vector<12x2xf32>, vector<2x13x3xi32> 790} 791 792// ----- 793 794// CHECK-LABEL: extract_strided_broadcast 795// CHECK: %[[B:.*]] = vector.broadcast %{{.*}} : vector<4xf16> to vector<2x4xf16> 796// CHECK-NEXT: return %[[B]] : vector<2x4xf16> 797func.func @extract_strided_broadcast(%arg0: vector<4xf16>) -> vector<2x4xf16> { 798 %0 = vector.broadcast %arg0 : vector<4xf16> to vector<16x4xf16> 799 %1 = vector.extract_strided_slice %0 800 {offsets = [0, 0], sizes = [2, 4], strides = [1, 1]} : 801 vector<16x4xf16> to vector<2x4xf16> 802 return %1 : vector<2x4xf16> 803} 804 805// ----- 806 807// CHECK-LABEL: extract_strided_broadcast2 808// CHECK: %[[E:.*]] = vector.extract_strided_slice %{{.*}} {offsets = [0], sizes = [2], strides = [1]} : vector<4xf16> to vector<2xf16> 809// CHECK-NEXT: %[[B:.*]] = vector.broadcast %[[E]] : vector<2xf16> to vector<2x2xf16> 810// CHECK-NEXT: return %[[B]] : vector<2x2xf16> 811func.func @extract_strided_broadcast2(%arg0: vector<4xf16>) -> vector<2x2xf16> { 812 %0 = vector.broadcast %arg0 : vector<4xf16> to vector<16x4xf16> 813 %1 = vector.extract_strided_slice %0 814 {offsets = [0, 0], sizes = [2, 2], strides = [1, 1]} : 815 vector<16x4xf16> to vector<2x2xf16> 816 return %1 : vector<2x2xf16> 817} 818 819// ----- 820 821// CHECK-LABEL: func @extract_strided_broadcast3 822// CHECK-SAME: (%[[ARG:.+]]: vector<1xf32>) 823// CHECK: %[[V:.+]] = vector.broadcast %[[ARG]] : vector<1xf32> to vector<1x4xf32> 824// CHECK: return %[[V]] 825func.func @extract_strided_broadcast3(%arg0: vector<1xf32>) -> vector<1x4xf32> { 826 %0 = vector.broadcast %arg0 : vector<1xf32> to vector<1x8xf32> 827 %1 = vector.extract_strided_slice %0 828 {offsets = [0, 4], sizes = [1, 4], strides = [1, 1]} 829 : vector<1x8xf32> to vector<1x4xf32> 830 return %1 : vector<1x4xf32> 831} 832 833// ----- 834 835// CHECK-LABEL: func @extract_strided_broadcast4 836// CHECK-SAME: (%[[ARG:.+]]: f32) 837// CHECK: %[[V:.+]] = vector.broadcast %[[ARG]] : f32 to vector<1x4xf32> 838// CHECK: return %[[V]] 839func.func @extract_strided_broadcast4(%arg0: f32) -> vector<1x4xf32> { 840 %0 = vector.broadcast %arg0 : f32 to vector<1x8xf32> 841 %1 = vector.extract_strided_slice %0 842 {offsets = [0, 4], sizes = [1, 4], strides = [1, 1]} 843 : vector<1x8xf32> to vector<1x4xf32> 844 return %1 : vector<1x4xf32> 845} 846 847// ----- 848 849// CHECK-LABEL: consecutive_shape_cast 850// CHECK: %[[C:.*]] = vector.shape_cast %{{.*}} : vector<16xf16> to vector<4x4xf16> 851// CHECK-NEXT: return %[[C]] : vector<4x4xf16> 852func.func @consecutive_shape_cast(%arg0: vector<16xf16>) -> vector<4x4xf16> { 853 %0 = vector.shape_cast %arg0 : vector<16xf16> to vector<2x8xf16> 854 %1 = vector.shape_cast %0 : vector<2x8xf16> to vector<4x4xf16> 855 return %1 : vector<4x4xf16> 856} 857 858// ----- 859 860// CHECK-LABEL: func @dead_transfer_op 861// CHECK-NOT: vector.transfer_read 862// CHECK-NOT: vector.transfer_write 863// CHECK: return 864func.func @dead_transfer_op(%arg0 : tensor<4x4xf32>, %arg1 : memref<4x4xf32>, 865 %v0 : vector<1x4xf32>) { 866 %c0 = arith.constant 0 : index 867 %cf0 = arith.constant 0.0 : f32 868 %r = vector.transfer_read %arg1[%c0, %c0], %cf0 : 869 memref<4x4xf32>, vector<1x4xf32> 870 %w = vector.transfer_write %v0, %arg0[%c0, %c0] : 871 vector<1x4xf32>, tensor<4x4xf32> 872 return 873} 874 875// ----- 876 877// CHECK-LABEL: func @dead_load 878// CHECK-NOT: vector.maskedload 879// CHECK-NOT: vector.gather 880// CHECK-NOT: vector.expandload 881// CHECK: return 882func.func @dead_load(%base: memref<?xf32>, %indices: vector<16xi32>, 883 %mask: vector<16xi1>, %passthru: vector<16xf32>) { 884 %c0 = arith.constant 0 : index 885 %0 = vector.maskedload %base[%c0], %mask, %passthru : 886 memref<?xf32>, vector<16xi1>, vector<16xf32> into vector<16xf32> 887 %1 = vector.gather %base[%c0][%indices], %mask, %passthru : 888 memref<?xf32>, vector<16xi32>, vector<16xi1>, vector<16xf32> into vector<16xf32> 889 %2 = vector.expandload %base[%c0], %mask, %passthru : 890 memref<?xf32>, vector<16xi1>, vector<16xf32> into vector<16xf32> 891 return 892} 893 894// ----- 895 896#contraction_accesses0 = [ 897 affine_map<(i, j, k) -> (i, k)>, 898 affine_map<(i, j, k) -> (k, j)>, 899 affine_map<(i, j, k) -> (i, j)> 900] 901#contraction_trait0 = { 902 indexing_maps = #contraction_accesses0, 903 iterator_types = ["parallel", "parallel", "reduction"] 904} 905 906// CHECK-LABEL: func @contractions 907// CHECK-SAME: %[[A:[0-9a-zA-Z]+]]: vector<2x3xf32> 908// CHECK-SAME: %[[B:[0-9a-zA-Z]+]]: vector<3x4xf32> 909// CHECK-SAME: %[[C:[0-9a-zA-Z]+]]: vector<2x4xf32> 910// CHECK-SAME: %[[A_I8:[0-9a-zA-Z]+]]: vector<2x3xi8> 911// CHECK-SAME: %[[B_I8:[0-9a-zA-Z]+]]: vector<3x4xi8> 912// CHECK-SAME: %[[C_I8:[0-9a-zA-Z]+]]: vector<2x4xi8> 913func.func @contractions(%a: vector<2x3xf32>, %b: vector<3x4xf32>, %c: vector<2x4xf32>, 914 %a_i8: vector<2x3xi8>, %b_i8: vector<3x4xi8>, %c_i8: vector<2x4xi8>) 915 -> (vector<2x4xf32>, vector<2x4xi8>) 916{ 917 // CHECK-NOT: arith.constant 918 %vf_0 = arith.constant dense <0.0>: vector<2x4xf32> 919 // CHECK-NOT: arith.addf 920 // CHECK: %[[D:.*]] = vector.contract {{.*}} %[[A]], %[[B]], %[[C]] 921 %0 = vector.contract #contraction_trait0 %a, %b, %vf_0: 922 vector<2x3xf32>, vector<3x4xf32> into vector<2x4xf32> 923 // CHECK-NOT: arith.addf 924 %1 = arith.addf %0, %c: vector<2x4xf32> 925 926 // CHECK-NOT: arith.constant 927 %vi8_0 = arith.constant dense <0>: vector<2x4xi8> 928 // CHECK-NOT: arith.addi 929 // CHECK: %[[D_I8:.*]] = vector.contract {{.*}} %[[A_I8]], %[[B_I8]], %[[C_I8]] 930 %i8_0 = vector.contract #contraction_trait0 %a_i8, %b_i8, %vi8_0: 931 vector<2x3xi8>, vector<3x4xi8> into vector<2x4xi8> 932 // CHECK-NOT: arith.addi 933 %i8_1 = arith.addi %i8_0, %c_i8: vector<2x4xi8> 934 935 // CHECK: return %[[D]], %[[D_I8]] 936 return %1, %i8_1: vector<2x4xf32>, vector<2x4xi8> 937} 938 939// ----- 940 941// CHECK-LABEL: func @transfer_folding_1 942// CHECK-SAME: %[[T0:[0-9a-zA-Z]+]]: tensor<2x3x4xf32> 943// CHECK-SAME: %[[T1:[0-9a-zA-Z]+]]: tensor<2x3x4xf32> 944func.func @transfer_folding_1(%t0: tensor<2x3x4xf32>, %t1: tensor<2x3x4xf32>) 945 -> (tensor<2x3x4xf32>, tensor<2x3x4xf32>, tensor<2x3x4xf32>) 946{ 947 %c0 = arith.constant 0 : index 948 %pad = arith.constant 0.0 : f32 949 %v = vector.transfer_read %t0[%c0, %c0, %c0], %pad {in_bounds = [true, true, true]} : 950 tensor<2x3x4xf32>, vector<2x3x4xf32> 951 952 %r0 = vector.transfer_write %v, %t1[%c0, %c0, %c0] {in_bounds = [true, true, true]} : 953 vector<2x3x4xf32>, tensor<2x3x4xf32> 954 955 %t2 = "test.constant"() { value = dense<6.0> : tensor<2x3x4xf32>} : () -> (tensor<2x3x4xf32>) 956 %r1 = vector.transfer_write %v, %t2[%c0, %c0, %c0] {in_bounds = [true, true, true]} : 957 vector<2x3x4xf32>, tensor<2x3x4xf32> 958 959 960 // CHECK-NEXT: some_op_that_may_have_side_effects 961 %t3 = "some_op_that_may_have_side_effects"() : () -> (tensor<2x3x4xf32>) 962 %r2 = vector.transfer_write %v, %t0[%c0, %c0, %c0] {in_bounds = [true, true, true]} : 963 vector<2x3x4xf32>, tensor<2x3x4xf32> 964 965 // CHECK-NEXT: return %[[T0]], %[[T0]], %[[T0]] 966 return %r0, %r1, %r2: tensor<2x3x4xf32>, tensor<2x3x4xf32>, tensor<2x3x4xf32> 967} 968 969// ----- 970 971// CHECK-LABEL: func @store_after_load_tensor 972// CHECK-SAME: (%[[ARG:.*]]: tensor<4x4xf32>) 973// CHECK-NOT: vector.transfer_read 974// CHECK-NOT: vector.transfer_write 975// CHECK: return %[[ARG]] : tensor<4x4xf32> 976func.func @store_after_load_tensor(%arg0 : tensor<4x4xf32>) -> tensor<4x4xf32> { 977 %c1 = arith.constant 1 : index 978 %c0 = arith.constant 0 : index 979 %cf0 = arith.constant 0.0 : f32 980 %0 = vector.transfer_read %arg0[%c1, %c0], %cf0 : 981 tensor<4x4xf32>, vector<1x4xf32> 982 %w0 = vector.transfer_write %0, %arg0[%c1, %c0] : 983 vector<1x4xf32>, tensor<4x4xf32> 984 return %w0 : tensor<4x4xf32> 985} 986 987// ----- 988 989// CHECK-LABEL: func @store_after_load_tensor_negative 990// CHECK: vector.transfer_read 991// CHECK: vector.transfer_write 992// CHECK: return 993func.func @store_after_load_tensor_negative(%arg0 : tensor<4x4xf32>) -> tensor<4x4xf32> { 994 %c1 = arith.constant 1 : index 995 %c0 = arith.constant 0 : index 996 %cf0 = arith.constant 0.0 : f32 997 %0 = vector.transfer_read %arg0[%c1, %c0], %cf0 : 998 tensor<4x4xf32>, vector<1x4xf32> 999 %w0 = vector.transfer_write %0, %arg0[%c0, %c0] : 1000 vector<1x4xf32>, tensor<4x4xf32> 1001 return %w0 : tensor<4x4xf32> 1002} 1003 1004// ----- 1005 1006// CHECK-LABEL: func @store_to_load_tensor 1007// CHECK-SAME: (%[[ARG:.*]]: tensor<4x4xf32>, %[[V0:.*]]: vector<1x4xf32>, %[[V1:.*]]: vector<1x4xf32>) 1008// CHECK-NOT: vector.transfer_write 1009// CHECK-NOT: vector.transfer_read 1010// CHECK: return %[[V0]] : vector<1x4xf32> 1011func.func @store_to_load_tensor(%arg0 : tensor<4x4xf32>, 1012 %v0 : vector<1x4xf32>, %v1 : vector<1x4xf32>) -> vector<1x4xf32> { 1013 %c1 = arith.constant 1 : index 1014 %c2 = arith.constant 2 : index 1015 %c0 = arith.constant 0 : index 1016 %cf0 = arith.constant 0.0 : f32 1017 %w0 = vector.transfer_write %v0, %arg0[%c1, %c0] {in_bounds = [true, true]} : 1018 vector<1x4xf32>, tensor<4x4xf32> 1019 %w1 = vector.transfer_write %v1, %w0[%c2, %c0] {in_bounds = [true, true]} : 1020 vector<1x4xf32>, tensor<4x4xf32> 1021 %0 = vector.transfer_read %w1[%c1, %c0], %cf0 {in_bounds = [true, true]} : 1022 tensor<4x4xf32>, vector<1x4xf32> 1023 return %0 : vector<1x4xf32> 1024} 1025 1026// ----- 1027 1028// CHECK-LABEL: func @store_to_load_negative_tensor 1029// CHECK: vector.transfer_write 1030// CHECK: vector.transfer_write 1031// CHECK: %[[V:.*]] = vector.transfer_read 1032// CHECK: return %[[V]] : vector<1x4xf32> 1033func.func @store_to_load_negative_tensor(%arg0 : tensor<4x4xf32>, 1034 %v0 : vector<1x4xf32>, %v1 : vector<1x4xf32>, %i : index) -> vector<1x4xf32> { 1035 %c1 = arith.constant 1 : index 1036 %c2 = arith.constant 2 : index 1037 %c0 = arith.constant 0 : index 1038 %cf0 = arith.constant 0.0 : f32 1039 %w0 = vector.transfer_write %v0, %arg0[%c1, %c0] {in_bounds = [true, true]} : 1040 vector<1x4xf32>, tensor<4x4xf32> 1041 %w1 = vector.transfer_write %v0, %w0[%i, %i] {in_bounds = [true, true]} : 1042 vector<1x4xf32>, tensor<4x4xf32> 1043 %0 = vector.transfer_read %w1[%c1, %c0], %cf0 {in_bounds = [true, true]} : 1044 tensor<4x4xf32>, vector<1x4xf32> 1045 return %0 : vector<1x4xf32> 1046} 1047 1048// ----- 1049 1050 1051// CHECK-LABEL: func @dead_store_tensor 1052// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index 1053// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index 1054// CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index 1055// CHECK-NOT: vector.transfer_write {{.*}}, {{.*}}[%[[C1]], %[[C0]] 1056// CHECK: vector.transfer_write {{.*}}, {{.*}}[%[[C2]], %[[C0]] 1057// CHECK: %[[VTW:.*]] = vector.transfer_write {{.*}}, {{.*}}[%[[C1]], %[[C0]] 1058// CHECK: return %[[VTW]] : tensor<4x4xf32> 1059func.func @dead_store_tensor(%arg0 : tensor<4x4xf32>, 1060 %v0 : vector<1x4xf32>, %v1 : vector<1x4xf32>, %i : index) -> tensor<4x4xf32> { 1061 %c1 = arith.constant 1 : index 1062 %c2 = arith.constant 2 : index 1063 %c0 = arith.constant 0 : index 1064 %cf0 = arith.constant 0.0 : f32 1065 %w0 = vector.transfer_write %v0, %arg0[%c1, %c0] {in_bounds = [true, true]} : 1066 vector<1x4xf32>, tensor<4x4xf32> 1067 %w1 = vector.transfer_write %v0, %w0[%c2, %c0] {in_bounds = [true, true]} : 1068 vector<1x4xf32>, tensor<4x4xf32> 1069 %w2 = vector.transfer_write %v1, %w1[%c1, %c0] {in_bounds = [true, true]} : 1070 vector<1x4xf32>, tensor<4x4xf32> 1071 return %w2 : tensor<4x4xf32> 1072} 1073 1074// ----- 1075 1076// CHECK-LABEL: func @dead_store_tensor_negative 1077// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index 1078// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index 1079// CHECK: vector.transfer_write 1080// CHECK: vector.transfer_write 1081// CHECK: vector.transfer_read 1082// CHECK: %[[VTW:.*]] = vector.transfer_write {{.*}}, {{.*}}[%[[C1]], %[[C0]]] 1083// CHECK: return %[[VTW]] : tensor<4x4xf32> 1084func.func @dead_store_tensor_negative(%arg0 : tensor<4x4xf32>, 1085 %v0 : vector<1x4xf32>, %v1 : vector<1x4xf32>, %i : index) -> tensor<4x4xf32> { 1086 %c1 = arith.constant 1 : index 1087 %c2 = arith.constant 2 : index 1088 %c0 = arith.constant 0 : index 1089 %cf0 = arith.constant 0.0 : f32 1090 %w0 = vector.transfer_write %v0, %arg0[%c1, %c0] {in_bounds = [true, true]} : 1091 vector<1x4xf32>, tensor<4x4xf32> 1092 %w1 = vector.transfer_write %v0, %w0[%c2, %c0] {in_bounds = [true, true]} : 1093 vector<1x4xf32>, tensor<4x4xf32> 1094 %0 = vector.transfer_read %w1[%i, %i], %cf0 {in_bounds = [true, true]} : 1095 tensor<4x4xf32>, vector<1x4xf32> 1096 %x = arith.addf %0, %0 : vector<1x4xf32> 1097 %w2 = vector.transfer_write %x, %w0[%c1, %c0] {in_bounds = [true, true]} : 1098 vector<1x4xf32>, tensor<4x4xf32> 1099 return %w2 : tensor<4x4xf32> 1100} 1101 1102// ----- 1103 1104// CHECK-LABEL: func @transfer_read_of_extract_slice( 1105// CHECK-SAME: %[[t:.*]]: tensor<?x?xf32>, %[[s1:.*]]: index, %[[s2:.*]]: index 1106// CHECK-DAG: %[[c4:.*]] = arith.constant 4 : index 1107// CHECK-DAG: %[[c8:.*]] = arith.constant 8 : index 1108// CHECK: %[[add:.*]] = arith.addi %[[s1]], %[[c4]] 1109// CHECK: %[[r:.*]] = vector.transfer_read %[[t]][%[[c8]], %[[add]]], %{{.*}} {in_bounds = [true, true]} : tensor<?x?xf32>, vector<5x6xf32> 1110// CHECK: return %[[r]] 1111func.func @transfer_read_of_extract_slice(%t : tensor<?x?xf32>, %s1 : index, %s2 : index) -> vector<5x6xf32> { 1112 %c3 = arith.constant 3 : index 1113 %c4 = arith.constant 4 : index 1114 %cst = arith.constant 0.0 : f32 1115 %0 = tensor.extract_slice %t[5, %s1] [10, %s2] [1, 1] : tensor<?x?xf32> to tensor<10x?xf32> 1116 %1 = vector.transfer_read %0[%c3, %c4], %cst {in_bounds = [true, true]} : tensor<10x?xf32>, vector<5x6xf32> 1117 return %1 : vector<5x6xf32> 1118} 1119 1120// ----- 1121 1122// CHECK-LABEL: func @transfer_read_of_extract_slice_rank_reducing( 1123// CHECK-SAME: %[[t:.*]]: tensor<?x?x?xf32>, %[[s1:.*]]: index, %[[s2:.*]]: index 1124// CHECK-DAG: %[[c3:.*]] = arith.constant 3 : index 1125// CHECK-DAG: %[[c5:.*]] = arith.constant 5 : index 1126// CHECK-DAG: %[[c10:.*]] = arith.constant 10 : index 1127// CHECK: %[[add:.*]] = arith.addi %[[s1]], %[[c3]] 1128// CHECK: %[[r:.*]] = vector.transfer_read %[[t]][%[[c5]], %[[add]], %[[c10]]], %{{.*}} {in_bounds = [true, true]} : tensor<?x?x?xf32>, vector<5x6xf32> 1129// CHECK: return %[[r]] 1130func.func @transfer_read_of_extract_slice_rank_reducing(%t : tensor<?x?x?xf32>, %s1 : index, %s2 : index) -> vector<5x6xf32> { 1131 %c3 = arith.constant 3 : index 1132 %c4 = arith.constant 4 : index 1133 %cst = arith.constant 0.0 : f32 1134 %0 = tensor.extract_slice %t[5, %s1, 6] [1, %s2, 12] [1, 1, 1] : tensor<?x?x?xf32> to tensor<?x12xf32> 1135 %1 = vector.transfer_read %0[%c3, %c4], %cst {in_bounds = [true, true]} : tensor<?x12xf32>, vector<5x6xf32> 1136 return %1 : vector<5x6xf32> 1137} 1138 1139// ----- 1140 1141// CHECK-LABEL: func @transfer_read_of_extract_slice_illegal_rank_reducing( 1142// CHECK: extract_slice 1143// CHECK: vector.transfer_read 1144func.func @transfer_read_of_extract_slice_illegal_rank_reducing(%t : tensor<?x?x?xf32>, %s1 : index, %s2 : index) -> vector<5x6xf32> { 1145 %c3 = arith.constant 3 : index 1146 %c4 = arith.constant 4 : index 1147 %cst = arith.constant 0.0 : f32 1148 %0 = tensor.extract_slice %t[5, %s1, 6] [%s2, 1, 12] [1, 1, 1] : tensor<?x?x?xf32> to tensor<?x12xf32> 1149 %1 = vector.transfer_read %0[%c3, %c4], %cst {in_bounds = [true, true]} : tensor<?x12xf32>, vector<5x6xf32> 1150 return %1 : vector<5x6xf32> 1151} 1152 1153// ----- 1154 1155// CHECK-LABEL: func @insert_slice_of_transfer_write( 1156// CHECK-SAME: %[[t1:.*]]: tensor<?x12xf32>, %[[v:.*]]: vector<5x6xf32>, %[[s:.*]]: index 1157// CHECK: %[[c3:.*]] = arith.constant 3 : index 1158// CHECK: %[[r:.*]] = vector.transfer_write %[[v]], %[[t1]][%[[c3]], %[[s]]] {in_bounds = [true, true]} : vector<5x6xf32>, tensor<?x12xf32> 1159// CHECK: return %[[r]] 1160func.func @insert_slice_of_transfer_write(%t1 : tensor<?x12xf32>, %v : vector<5x6xf32>, %s : index, %t2 : tensor<5x6xf32>) -> tensor<?x12xf32> { 1161 %c0 = arith.constant 0 : index 1162 %0 = vector.transfer_write %v, %t2[%c0, %c0] {in_bounds = [true, true]} : vector<5x6xf32>, tensor<5x6xf32> 1163 %1 = tensor.insert_slice %0 into %t1[3, %s] [5, 6] [1, 1] : tensor<5x6xf32> into tensor<?x12xf32> 1164 return %1 : tensor<?x12xf32> 1165} 1166 1167// ----- 1168 1169// CHECK-LABEL: func @insert_slice_of_transfer_write_illegal_rank_extending( 1170// CHECK: vector.transfer_write 1171// CHECK: insert_slice 1172func.func @insert_slice_of_transfer_write_illegal_rank_extending(%t1 : tensor<?x?x12xf32>, %v : vector<5x6xf32>, %s : index, %t2 : tensor<5x6xf32>) -> tensor<?x?x12xf32> { 1173 %c0 = arith.constant 0 : index 1174 %0 = vector.transfer_write %v, %t2[%c0, %c0] {in_bounds = [true, true]} : vector<5x6xf32>, tensor<5x6xf32> 1175 %1 = tensor.insert_slice %0 into %t1[4, 3, %s] [5, 1, 6] [1, 1, 1] : tensor<5x6xf32> into tensor<?x?x12xf32> 1176 return %1 : tensor<?x?x12xf32> 1177} 1178 1179// ----- 1180 1181// CHECK-LABEL: func @insert_slice_of_transfer_write_rank_extending( 1182// CHECK-SAME: %[[t1:.*]]: tensor<?x?x12xf32>, %[[v:.*]]: vector<5x6xf32>, %[[s:.*]]: index 1183// CHECK-DAG: %[[c3:.*]] = arith.constant 3 : index 1184// CHECK-DAG: %[[c4:.*]] = arith.constant 4 : index 1185// CHECK: %[[r:.*]] = vector.transfer_write %[[v]], %[[t1]][%[[c4]], %[[c3]], %[[s]]] {in_bounds = [true, true]} : vector<5x6xf32>, tensor<?x?x12xf32> 1186// CHECK: return %[[r]] 1187func.func @insert_slice_of_transfer_write_rank_extending(%t1 : tensor<?x?x12xf32>, %v : vector<5x6xf32>, %s : index, %t2 : tensor<5x6xf32>) -> tensor<?x?x12xf32> { 1188 %c0 = arith.constant 0 : index 1189 %0 = vector.transfer_write %v, %t2[%c0, %c0] {in_bounds = [true, true]} : vector<5x6xf32>, tensor<5x6xf32> 1190 %1 = tensor.insert_slice %0 into %t1[4, 3, %s] [1, 5, 6] [1, 1, 1] : tensor<5x6xf32> into tensor<?x?x12xf32> 1191 return %1 : tensor<?x?x12xf32> 1192} 1193 1194// ----- 1195 1196// CHECK: #[[$MAP:[0-9a-z]+]] = affine_map<(d0, d1) -> (d1, d0)> 1197 1198// CHECK-LABEL: func @swap_extract_slice_transfer_write 1199// CHECK-SAME: %[[VEC:.*]]: vector<8x4xf32> 1200// CHECK-SAME: %[[INIT_TENSOR:.*]]: tensor<4x8xf32>, 1201// CHECK-SAME: %[[ITER_ARG:.*]]: tensor<64x64xf32>, 1202// CHECK-SAME: %[[IV:.*]]: index, %[[SZ:.*]]: index) 1203func.func @swap_extract_slice_transfer_write(%arg0 : vector<8x4xf32>, 1204 %arg1 : tensor<4x8xf32>, 1205 %arg2 : tensor<64x64xf32>, 1206 %iv : index, %sz : index) -> tensor<64x64xf32> { 1207 // CHECK: %[[C0:.*]] = arith.constant 0 : index 1208 %c0 = arith.constant 0 : index 1209 1210 // CHECK: %[[T0:.*]] = tensor.extract_slice %[[ITER_ARG]] 1211 // CHECK-SAME: [%[[IV]], 16] [%[[SZ]], 8] 1212 // CHECK: %[[T1:.*]] = vector.transfer_write %[[VEC]] 1213 // CHECK-SAME: %[[T0]][%[[C0]], %[[C0]]] 1214 // CHECK-SAME: in_bounds = [true, false] 1215 // CHECK-SAME: permutation_map = #[[$MAP]] 1216 // CHECK: %[[T2:.*]] = tensor.insert_slice %[[T1]] into %[[ITER_ARG]] 1217 // CHECK-SAME: [%[[IV]], 16] [%[[SZ]], 8] 1218 %0 = vector.transfer_write %arg0, %arg1[%c0, %c0] {in_bounds = [true, true], permutation_map = affine_map<(d0, d1) -> (d1, d0)>} : vector<8x4xf32>, tensor<4x8xf32> 1219 %1 = tensor.extract_slice %0[0, 0] [%sz, 8] [1, 1] : tensor<4x8xf32> to tensor<?x8xf32> 1220 %2 = tensor.insert_slice %1 into %arg2[%iv, 16] [%sz, 8] [1, 1] : tensor<?x8xf32> into tensor<64x64xf32> 1221 1222 // CHECK: return %[[T2]] 1223 func.return %2 : tensor<64x64xf32> 1224} 1225 1226// ----- 1227 1228// CHECK-LABEL: func @do_not_swap_extract_slice_transfer_write 1229// CHECK-SAME: %[[VEC:.*]]: vector<8xf32>, 1230// CHECK-SAME: %[[VEC_SMALL:.*]]: vector<4xf32>, 1231// CHECK-SAME: %[[INIT_TENSOR:.*]]: tensor<8xf32>, 1232// CHECK-SAME: %[[ITER_ARG:.*]]: tensor<64xf32>, 1233// CHECK-SAME: %[[IV:.*]]: index, %[[SZ:.*]]: index) 1234func.func @do_not_swap_extract_slice_transfer_write(%arg0 : vector<8xf32>, 1235 %arg1 : vector<4xf32>, 1236 %arg2 : tensor<8xf32>, 1237 %arg3 : tensor<64xf32>, 1238 %iv : index, %sz : index) -> (tensor<64xf32>, tensor<64xf32>, tensor<64xf32>) { 1239 // CHECK: %[[C0:.*]] = arith.constant 0 : index 1240 %c0 = arith.constant 0 : index 1241 1242 // Don't swap if the extracted and inserted slices do not match. 1243 // CHECK: %[[T0:.*]] = vector.transfer_write %[[VEC]] 1244 // CHECK: %[[T1:.*]] = tensor.extract_slice %[[T0]] 1245 // CHECK: %[[T2:.*]] = tensor.insert_slice %[[T1]] 1246 %0 = vector.transfer_write %arg0, %arg2[%c0] {in_bounds = [true]} : vector<8xf32>, tensor<8xf32> 1247 %1 = tensor.extract_slice %0[0] [%iv] [1] : tensor<8xf32> to tensor<?xf32> 1248 %2 = tensor.insert_slice %1 into %arg3[%iv] [%sz] [1] : tensor<?xf32> into tensor<64xf32> 1249 1250 // Don't swap if the TransferWriteOp takes a small vector. 1251 // CHECK: %[[T3:.*]] = vector.transfer_write %[[VEC_SMALL]] 1252 // CHECK: %[[T4:.*]] = tensor.extract_slice %[[T3]] 1253 // CHECK: %[[T5:.*]] = tensor.insert_slice %[[T4]] 1254 %3 = vector.transfer_write %arg1, %arg2[%c0] {in_bounds = [true]} : vector<4xf32>, tensor<8xf32> 1255 %4 = tensor.extract_slice %3[0] [%sz] [1] : tensor<8xf32> to tensor<?xf32> 1256 %5 = tensor.insert_slice %4 into %arg3[%iv] [%sz] [1] : tensor<?xf32> into tensor<64xf32> 1257 1258 // Don't swap if the one of the operations is rank-reducing. 1259 // CHECK: %[[T6:.*]] = vector.transfer_write %[[VEC]] 1260 // CHECK: %[[T7:.*]] = tensor.extract_slice %[[T6]] 1261 // CHECK: %[[T8:.*]] = tensor.insert_slice %[[T7]] 1262 %6 = vector.transfer_write %arg0, %arg2[%c0] {in_bounds = [true]} : vector<8xf32>, tensor<8xf32> 1263 %7 = tensor.extract_slice %6[0] [1] [1] : tensor<8xf32> to tensor<f32> 1264 %8 = tensor.insert_slice %7 into %arg3[%iv] [1] [1] : tensor<f32> into tensor<64xf32> 1265 1266 // CHECK: return %[[T2]], %[[T5]], %[[T8]] 1267 func.return %2, %5, %8 : tensor<64xf32>, tensor<64xf32>, tensor<64xf32> 1268} 1269 1270// ----- 1271 1272// CHECK-LABEL: func @vector_multi_reduction_single_parallel( 1273// CHECK-SAME: %[[v:.*]]: vector<2xf32> 1274func.func @vector_multi_reduction_single_parallel(%arg0: vector<2xf32>) -> vector<2xf32> { 1275 %0 = vector.multi_reduction <mul>, %arg0 [] : vector<2xf32> to vector<2xf32> 1276 1277// CHECK: return %[[v]] : vector<2xf32> 1278 return %0 : vector<2xf32> 1279} 1280 1281// ----- 1282 1283// CHECK-LABEL: func @insert_strided_slice_full_range 1284// CHECK-SAME: %[[SOURCE:.+]]: vector<16x16xf16>, %{{.+}}: vector<16x16xf16> 1285func.func @insert_strided_slice_full_range(%source: vector<16x16xf16>, %dest: vector<16x16xf16>) -> vector<16x16xf16> { 1286 %0 = vector.insert_strided_slice %source, %dest {offsets = [0, 0], strides = [1, 1]} : vector<16x16xf16> into vector<16x16xf16> 1287 // CHECK: return %[[SOURCE]] 1288 return %0: vector<16x16xf16> 1289} 1290 1291// ----- 1292 1293// CHECK-LABEL: extract_strided_splat 1294// CHECK: %[[B:.*]] = vector.splat %{{.*}} : vector<2x4xf16> 1295// CHECK-NEXT: return %[[B]] : vector<2x4xf16> 1296func.func @extract_strided_splat(%arg0: f16) -> vector<2x4xf16> { 1297 %0 = vector.splat %arg0 : vector<16x4xf16> 1298 %1 = vector.extract_strided_slice %0 1299 {offsets = [1, 0], sizes = [2, 4], strides = [1, 1]} : 1300 vector<16x4xf16> to vector<2x4xf16> 1301 return %1 : vector<2x4xf16> 1302} 1303 1304// ----- 1305 1306// CHECK-LABEL: func @insert_extract_to_broadcast 1307// CHECK-SAME: (%[[ARG0:.*]]: vector<1x1x4xf32>, %[[ARG1:.*]]: vector<4xf32>) 1308// CHECK: %[[V0:.*]] = vector.extract %[[ARG0]][0, 0] : vector<1x1x4xf32> 1309// CHECK: %[[V1:.*]] = vector.broadcast %[[ARG1]] : vector<4xf32> to vector<1x1x4xf32> 1310// CHECK: return %[[V0]], %[[V1]] : vector<4xf32>, vector<1x1x4xf32> 1311func.func @insert_extract_to_broadcast(%arg0 : vector<1x1x4xf32>, 1312 %arg1 : vector<4xf32>) -> (vector<4xf32>, vector<1x1x4xf32>) { 1313 %0 = vector.extract %arg0[0, 0] : vector<1x1x4xf32> 1314 %1 = vector.insert %arg1, %arg0 [0, 0] : vector<4xf32> into vector<1x1x4xf32> 1315 return %0, %1 : vector<4xf32>, vector<1x1x4xf32> 1316} 1317 1318// ----- 1319 1320// CHECK-LABEL: extract_constant 1321// CHECK-DAG: %[[CST1:.*]] = arith.constant 1 : i32 1322// CHECK-DAG: %[[CST0:.*]] = arith.constant dense<2.000000e+00> : vector<7xf32> 1323// CHECK: return %[[CST0]], %[[CST1]] : vector<7xf32>, i32 1324func.func @extract_constant() -> (vector<7xf32>, i32) { 1325 %cst = arith.constant dense<2.000000e+00> : vector<29x7xf32> 1326 %cst_1 = arith.constant dense<1> : vector<4x37x9xi32> 1327 %0 = vector.extract %cst[2] : vector<29x7xf32> 1328 %1 = vector.extract %cst_1[1, 4, 5] : vector<4x37x9xi32> 1329 return %0, %1 : vector<7xf32>, i32 1330} 1331 1332// ----- 1333 1334// CHECK-LABEL: extract_extract_strided 1335// CHECK-SAME: %[[A:.*]]: vector<32x16x4xf16> 1336// CHECK: %[[V:.*]] = vector.extract %[[A]][9, 7] : vector<32x16x4xf16> 1337// CHECK: return %[[V]] : vector<4xf16> 1338func.func @extract_extract_strided(%arg0: vector<32x16x4xf16>) -> vector<4xf16> { 1339 %1 = vector.extract_strided_slice %arg0 1340 {offsets = [7, 3], sizes = [10, 8], strides = [1, 1]} : 1341 vector<32x16x4xf16> to vector<10x8x4xf16> 1342 %2 = vector.extract %1[2, 4] : vector<10x8x4xf16> 1343 return %2 : vector<4xf16> 1344} 1345 1346// ----- 1347 1348// CHECK-LABEL: extract_insert_strided 1349// CHECK-SAME: %[[A:.*]]: vector<6x4xf32> 1350// CHECK: %[[V:.*]] = vector.extract %[[A]][0, 2] : vector<6x4xf32> 1351// CHECK: return %[[V]] : f32 1352func.func @extract_insert_strided(%a: vector<6x4xf32>, %b: vector<8x16xf32>) 1353 -> f32 { 1354 %0 = vector.insert_strided_slice %a, %b {offsets = [2, 2], strides = [1, 1]} 1355 : vector<6x4xf32> into vector<8x16xf32> 1356 %2 = vector.extract %0[2, 4] : vector<8x16xf32> 1357 return %2 : f32 1358} 1359 1360// ----- 1361 1362// CHECK-LABEL: extract_insert_rank_reduce 1363// CHECK-SAME: %[[A:.*]]: vector<4xf32> 1364// CHECK: %[[V:.*]] = vector.extract %[[A]][2] : vector<4xf32> 1365// CHECK: return %[[V]] : f32 1366func.func @extract_insert_rank_reduce(%a: vector<4xf32>, %b: vector<8x16xf32>) 1367 -> f32 { 1368 %0 = vector.insert_strided_slice %a, %b {offsets = [2, 2], strides = [1]} 1369 : vector<4xf32> into vector<8x16xf32> 1370 %2 = vector.extract %0[2, 4] : vector<8x16xf32> 1371 return %2 : f32 1372} 1373 1374// ----- 1375 1376// CHECK-LABEL: extract_insert_negative 1377// CHECK: vector.insert_strided_slice 1378// CHECK: vector.extract 1379func.func @extract_insert_negative(%a: vector<2x15xf32>, %b: vector<12x8x16xf32>) 1380 -> vector<16xf32> { 1381 %0 = vector.insert_strided_slice %a, %b {offsets = [4, 2, 0], strides = [1, 1]} 1382 : vector<2x15xf32> into vector<12x8x16xf32> 1383 %2 = vector.extract %0[4, 2] : vector<12x8x16xf32> 1384 return %2 : vector<16xf32> 1385} 1386 1387// ----- 1388 1389// CHECK-LABEL: extract_insert_chain 1390// CHECK-SAME: (%[[A:.*]]: vector<2x16xf32>, %[[B:.*]]: vector<12x8x16xf32>, %[[C:.*]]: vector<2x16xf32>) 1391// CHECK: %[[V:.*]] = vector.extract %[[C]][0] : vector<2x16xf32> 1392// CHECK: return %[[V]] : vector<16xf32> 1393func.func @extract_insert_chain(%a: vector<2x16xf32>, %b: vector<12x8x16xf32>, %c: vector<2x16xf32>) 1394 -> vector<16xf32> { 1395 %0 = vector.insert_strided_slice %c, %b {offsets = [4, 2, 0], strides = [1, 1]} 1396 : vector<2x16xf32> into vector<12x8x16xf32> 1397 %1 = vector.insert_strided_slice %a, %0 {offsets = [0, 2, 0], strides = [1, 1]} 1398 : vector<2x16xf32> into vector<12x8x16xf32> 1399 %2 = vector.extract %1[4, 2] : vector<12x8x16xf32> 1400 return %2 : vector<16xf32> 1401} 1402 1403// ----- 1404 1405// CHECK-LABEL: extract_extract_strided2 1406// CHECK-SAME: %[[A:.*]]: vector<2x4xf32> 1407// CHECK: %[[V:.*]] = vector.extract %[[A]][1] : vector<2x4xf32> 1408// CHECK: return %[[V]] : vector<4xf32> 1409func.func @extract_extract_strided2(%A: vector<2x4xf32>) 1410 -> (vector<4xf32>) { 1411 %0 = vector.extract_strided_slice %A {offsets = [1, 0], sizes = [1, 4], strides = [1, 1]} : vector<2x4xf32> to vector<1x4xf32> 1412 %1 = vector.extract %0[0] : vector<1x4xf32> 1413 return %1 : vector<4xf32> 1414} 1415 1416// ----- 1417 1418// CHECK-LABEL: func @splat_fold 1419func.func @splat_fold() -> vector<4xf32> { 1420 %c = arith.constant 1.0 : f32 1421 %v = vector.splat %c : vector<4xf32> 1422 return %v : vector<4xf32> 1423 1424 // CHECK-NEXT: [[V:%.*]] = arith.constant dense<1.000000e+00> : vector<4xf32> 1425 // CHECK-NEXT: return [[V]] : vector<4xf32> 1426} 1427 1428// ----- 1429 1430// CHECK-LABEL: func @shuffle_1d 1431// CHECK: %[[V:.+]] = arith.constant dense<[3, 2, 5, 1]> : vector<4xi32> 1432// CHECK: return %[[V]] 1433func.func @shuffle_1d() -> vector<4xi32> { 1434 %v0 = arith.constant dense<[0, 1, 2]> : vector<3xi32> 1435 %v1 = arith.constant dense<[3, 4, 5]> : vector<3xi32> 1436 %shuffle = vector.shuffle %v0, %v1 [3, 2, 5, 1] : vector<3xi32>, vector<3xi32> 1437 return %shuffle : vector<4xi32> 1438} 1439 1440// CHECK-LABEL: func @shuffle_fold1 1441// CHECK: %arg0 : vector<4xi32> 1442func.func @shuffle_fold1(%v0 : vector<4xi32>, %v1 : vector<2xi32>) -> vector<4xi32> { 1443 %shuffle = vector.shuffle %v0, %v1 [0, 1, 2, 3] : vector<4xi32>, vector<2xi32> 1444 return %shuffle : vector<4xi32> 1445} 1446 1447// CHECK-LABEL: func @shuffle_fold2 1448// CHECK: %arg1 : vector<2xi32> 1449func.func @shuffle_fold2(%v0 : vector<4xi32>, %v1 : vector<2xi32>) -> vector<2xi32> { 1450 %shuffle = vector.shuffle %v0, %v1 [4, 5] : vector<4xi32>, vector<2xi32> 1451 return %shuffle : vector<2xi32> 1452} 1453 1454// CHECK-LABEL: func @shuffle_fold3 1455// CHECK: return %arg0 : vector<4x5x6xi32> 1456func.func @shuffle_fold3(%v0 : vector<4x5x6xi32>, %v1 : vector<2x5x6xi32>) -> vector<4x5x6xi32> { 1457 %shuffle = vector.shuffle %v0, %v1 [0, 1, 2, 3] : vector<4x5x6xi32>, vector<2x5x6xi32> 1458 return %shuffle : vector<4x5x6xi32> 1459} 1460 1461// CHECK-LABEL: func @shuffle_fold4 1462// CHECK: return %arg1 : vector<2x5x6xi32> 1463func.func @shuffle_fold4(%v0 : vector<4x5x6xi32>, %v1 : vector<2x5x6xi32>) -> vector<2x5x6xi32> { 1464 %shuffle = vector.shuffle %v0, %v1 [4, 5] : vector<4x5x6xi32>, vector<2x5x6xi32> 1465 return %shuffle : vector<2x5x6xi32> 1466} 1467 1468// CHECK-LABEL: func @shuffle_nofold1 1469// CHECK: %[[V:.+]] = vector.shuffle %arg0, %arg1 [0, 1, 2, 3, 4] : vector<4xi32>, vector<2xi32> 1470// CHECK: return %[[V]] 1471func.func @shuffle_nofold1(%v0 : vector<4xi32>, %v1 : vector<2xi32>) -> vector<5xi32> { 1472 %shuffle = vector.shuffle %v0, %v1 [0, 1, 2, 3, 4] : vector<4xi32>, vector<2xi32> 1473 return %shuffle : vector<5xi32> 1474} 1475 1476// CHECK-LABEL: func @shuffle_nofold2 1477// CHECK: %[[V:.+]] = vector.shuffle %arg0, %arg1 [0, 1, 2, 3] : vector<[4]xi32>, vector<[2]xi32> 1478// CHECK: return %[[V]] 1479func.func @shuffle_nofold2(%v0 : vector<[4]xi32>, %v1 : vector<[2]xi32>) -> vector<4xi32> { 1480 %shuffle = vector.shuffle %v0, %v1 [0, 1, 2, 3] : vector<[4]xi32>, vector<[2]xi32> 1481 return %shuffle : vector<4xi32> 1482} 1483 1484// ----- 1485 1486// CHECK-LABEL: func @transpose_scalar_broadcast1 1487// CHECK-SAME: (%[[ARG:.+]]: vector<1xf32>) 1488// CHECK: %[[V:.+]] = vector.broadcast %[[ARG]] : vector<1xf32> to vector<1x8xf32> 1489// CHECK: return %[[V]] : vector<1x8xf32> 1490func.func @transpose_scalar_broadcast1(%value: vector<1xf32>) -> vector<1x8xf32> { 1491 %bcast = vector.broadcast %value : vector<1xf32> to vector<8x1xf32> 1492 %t = vector.transpose %bcast, [1, 0] : vector<8x1xf32> to vector<1x8xf32> 1493 return %t : vector<1x8xf32> 1494} 1495 1496// ----- 1497 1498// CHECK-LABEL: func @transpose_scalar_broadcast2 1499// CHECK-SAME: (%[[ARG:.+]]: f32) 1500// CHECK: %[[V:.+]] = vector.broadcast %[[ARG]] : f32 to vector<1x8xf32> 1501// CHECK: return %[[V]] : vector<1x8xf32> 1502func.func @transpose_scalar_broadcast2(%value: f32) -> vector<1x8xf32> { 1503 %bcast = vector.broadcast %value : f32 to vector<8x1xf32> 1504 %t = vector.transpose %bcast, [1, 0] : vector<8x1xf32> to vector<1x8xf32> 1505 return %t : vector<1x8xf32> 1506} 1507 1508// ----- 1509 1510// CHECK-LABEL: func @transpose_splat_constant 1511// CHECK: %[[CST:.+]] = arith.constant dense<5.000000e+00> : vector<8x4xf32> 1512// CHECK: return %[[CST]] 1513func.func @transpose_splat_constant() -> vector<8x4xf32> { 1514 %cst = arith.constant dense<5.0> : vector<4x8xf32> 1515 %0 = vector.transpose %cst, [1, 0] : vector<4x8xf32> to vector<8x4xf32> 1516 return %0 : vector<8x4xf32> 1517} 1518 1519// CHECK-LABEL: func @transpose_splat2( 1520// CHECK-SAME: %[[VAL_0:.*]]: f32) -> vector<3x4xf32> { 1521// CHECK: %[[VAL_1:.*]] = vector.splat %[[VAL_0]] : vector<3x4xf32> 1522// CHECK: return %[[VAL_1]] : vector<3x4xf32> 1523// CHECK: } 1524func.func @transpose_splat2(%arg : f32) -> vector<3x4xf32> { 1525 %splat = vector.splat %arg : vector<4x3xf32> 1526 %0 = vector.transpose %splat, [1, 0] : vector<4x3xf32> to vector<3x4xf32> 1527 return %0 : vector<3x4xf32> 1528} 1529 1530// ----- 1531 1532// CHECK-LABEL: func @insert_element_fold 1533// CHECK: %[[V:.+]] = arith.constant dense<[0, 1, 7, 3]> : vector<4xi32> 1534// CHECK: return %[[V]] 1535func.func @insert_element_fold() -> vector<4xi32> { 1536 %v = arith.constant dense<[0, 1, 2, 3]> : vector<4xi32> 1537 %s = arith.constant 7 : i32 1538 %i = arith.constant 2 : i32 1539 %1 = vector.insertelement %s, %v[%i : i32] : vector<4xi32> 1540 return %1 : vector<4xi32> 1541} 1542 1543// ----- 1544 1545// CHECK-LABEL: func @extract_element_fold 1546// CHECK: %[[C:.+]] = arith.constant 5 : i32 1547// CHECK: return %[[C]] 1548func.func @extract_element_fold() -> i32 { 1549 %v = arith.constant dense<[1, 3, 5, 7]> : vector<4xi32> 1550 %i = arith.constant 2 : i32 1551 %1 = vector.extractelement %v[%i : i32] : vector<4xi32> 1552 return %1 : i32 1553} 1554 1555// CHECK-LABEL: func @extract_element_splat_fold 1556// CHECK-SAME: (%[[ARG:.+]]: i32) 1557// CHECK: return %[[ARG]] 1558func.func @extract_element_splat_fold(%a : i32) -> i32 { 1559 %v = vector.splat %a : vector<4xi32> 1560 %i = arith.constant 2 : i32 1561 %1 = vector.extractelement %v[%i : i32] : vector<4xi32> 1562 return %1 : i32 1563} 1564 1565// ----- 1566 1567// CHECK-LABEL: func @reduce_one_element_vector_extract 1568// CHECK-SAME: (%[[V:.+]]: vector<1xf32>) 1569// CHECK: %[[S:.+]] = vector.extract %[[V]][0] : vector<1xf32> 1570// CHECK: return %[[S]] : f32 1571func.func @reduce_one_element_vector_extract(%a : vector<1xf32>) -> f32 { 1572 %s = vector.reduction <add>, %a : vector<1xf32> into f32 1573 return %s : f32 1574} 1575 1576// ----- 1577 1578// CHECK-LABEL: func @reduce_one_element_vector_addf 1579// CHECK-SAME: (%[[V:.+]]: vector<1xf32>, %[[B:.+]]: f32) 1580// CHECK: %[[A:.+]] = vector.extract %[[V]][0] : vector<1xf32> 1581// CHECK: %[[S:.+]] = arith.addf %[[A]], %arg1 : f32 1582// CHECK: return %[[S]] 1583func.func @reduce_one_element_vector_addf(%a : vector<1xf32>, %b: f32) -> f32 { 1584 %s = vector.reduction <add>, %a, %b : vector<1xf32> into f32 1585 return %s : f32 1586} 1587 1588// ----- 1589 1590// CHECK-LABEL: func @reduce_one_element_vector_mulf 1591// CHECK-SAME: (%[[V:.+]]: vector<1xf32>, %[[B:.+]]: f32) 1592// CHECK: %[[A:.+]] = vector.extract %[[V]][0] : vector<1xf32> 1593// CHECK: %[[S:.+]] = arith.mulf %[[A]], %arg1 : f32 1594// CHECK: return %[[S]] 1595func.func @reduce_one_element_vector_mulf(%a : vector<1xf32>, %b: f32) -> f32 { 1596 %s = vector.reduction <mul>, %a, %b : vector<1xf32> into f32 1597 return %s : f32 1598} 1599 1600// ----- 1601 1602// CHECK-LABEL: func @dont_reduce_one_element_vector 1603// CHECK: vector.reduction 1604func.func @dont_reduce_one_element_vector(%a : vector<4xf32>) -> f32 { 1605 %s = vector.reduction <add>, %a : vector<4xf32> into f32 1606 return %s : f32 1607} 1608