578 lines · plain
1// RUN: mlir-opt %s -split-input-file -test-vector-linearize -verify-diagnostics | FileCheck %s2 3// CHECK-LABEL: test_linearize4// CHECK-SAME: (%[[ORIG_ARG:.*]]: vector<2x2xf32>)5func.func @test_linearize(%arg0: vector<2x2xf32>) -> vector<2x2xf32> {6 7 // CHECK: %[[ARG:.*]] = vector.shape_cast %[[ORIG_ARG]] : vector<2x2xf32> to vector<4xf32>8 // CHECK: %[[CST:.*]] = arith.constant dense<[1.000000e+00, 2.000000e+00, 3.000000e+00, 4.000000e+00]> : vector<4xf32>9 // CHECK: %[[RES:.*]] = vector.shape_cast %[[CST]] : vector<4xf32> to vector<2x2xf32>10 %0 = arith.constant dense<[[1.0, 2.0], [3.0, 4.0]]> : vector<2x2xf32>11 12 // CHECK: %{{.*}} = math.sin %[[ARG]] : vector<4xf32>13 %1 = math.sin %arg0 : vector<2x2xf32>14 15 // CHECK: %{{.*}} = arith.addf %[[ARG]], %[[CST]] : vector<4xf32>16 %2 = arith.addf %arg0, %0 : vector<2x2xf32>17 18 // CHECK: return %[[RES]] : vector<2x2xf32>19 return %0 : vector<2x2xf32>20}21 22// -----23 24// CHECK-LABEL: test_linearize_poison25func.func @test_linearize_poison() -> vector<2x2xf32> {26 27 // CHECK: %[[POISON:.*]] = ub.poison : vector<4xf32>28 // CHECK: %[[RES:.*]] = vector.shape_cast %[[POISON]] : vector<4xf32> to vector<2x2xf32>29 %0 = ub.poison : vector<2x2xf32>30 31 // CHECK: return %[[RES]] : vector<2x2xf32>32 return %0 : vector<2x2xf32>33}34 35// -----36 37// CHECK-LABEL: test_partial_linearize38// CHECK-SAME: (%[[ORIG_ARG:.*]]: vector<2x2xf32>, %[[ORIG_ARG2:.*]]: vector<4x4xf32>)39func.func @test_partial_linearize(%arg0: vector<2x2xf32>, %arg1: vector<4x4xf32>) -> vector<2x2xf32> {40 41 // CHECK-DAG: %[[ARG:.*]] = vector.shape_cast %[[ORIG_ARG]] : vector<2x2xf32> to vector<4xf32>42 // CHECK-DAG: %[[ARG2:.*]] = vector.shape_cast %[[ORIG_ARG2]] : vector<4x4xf32> to vector<16xf32>43 // CHECK: %[[CST:.*]] = arith.constant dense<{{.*}}> : vector<4xf32>44 // CHECK: %[[RES:.*]] = vector.shape_cast %[[CST]] : vector<4xf32> to vector<2x2xf32>45 %0 = arith.constant dense<[[1.0, 2.0], [3.0, 4.0]]> : vector<2x2xf32>46 47 // CHECK: %[[C2:.*]] = arith.constant dense<{{.*}}> : vector<16xf32>48 %5 = arith.constant dense<[[1.0, 2.0, 3.0, 4.0], [1.0, 2.0,3.0, 4.0], [1.0, 2.0, 3.0, 4.0], [1.0, 2.0, 5.0, 6.0]]> : vector<4x4xf32>49 50 // Arith and math ops are handled in generic way, check some of them51 // CHECK: %[[SIN:.*]] = math.sin %[[ARG]] : vector<4xf32>52 %1 = math.sin %arg0 : vector<2x2xf32>53 54 // CHECK: %[[SIN1:.*]] = math.sin %[[ARG2]] : vector<16xf32>55 %6 = math.sin %arg1 : vector<4x4xf32>56 57 // CHECK: %{{.*}} = arith.addf %[[ARG]], %[[CST]] : vector<4xf32>58 %2 = arith.addf %arg0, %0 : vector<2x2xf32>59 60 // CHECK: %[[ADD2:.*]] = arith.addf %[[ARG2]], %[[C2]] : vector<16xf32>61 %7 = arith.addf %arg1, %5 : vector<4x4xf32>62 63 // CHECK: return %[[RES]] : vector<2x2xf32>64 return %0 : vector<2x2xf32>65}66 67// -----68 69// vectorizable operation (arith.mulf) with tensor result types.70 71// CHECK-LABEL: test_tensor_no_linearize72func.func @test_tensor_no_linearize(%arg0: tensor<2x2xf32>, %arg1: tensor<2x2xf32>) -> (tensor<2x2xf32>, tensor<2x2xf32>) {73 74 // CHECK: %[[MULF:.*]] = arith.mulf %arg0, %arg1 : tensor<2x2xf32>75 %0 = arith.mulf %arg0, %arg1 : tensor<2x2xf32>76 77 return %0, %arg0 : tensor<2x2xf32>, tensor<2x2xf32>78}79 80// -----81 82// CHECK-LABEL: func.func @test_scalable_linearize(83// CHECK-SAME: %[[ARG_0:.*]]: vector<2x[2]xf32>) -> vector<2x[2]xf32> {84func.func @test_scalable_linearize(%arg0: vector<2x[2]xf32>) -> vector<2x[2]xf32> {85 86 // CHECK: %[[SC:.*]] = vector.shape_cast %[[ARG_0]] : vector<2x[2]xf32> to vector<[4]xf32>87 // CHECK: %[[CST:.*]] = arith.constant dense<3.000000e+00> : vector<[4]xf32>88 %0 = arith.constant dense<[[3., 3.], [3., 3.]]> : vector<2x[2]xf32>89 90 // CHECK: %[[SIN:.*]] = math.sin %[[SC]] : vector<[4]xf32>91 %1 = math.sin %arg0 : vector<2x[2]xf32>92 93 // CHECK: %[[ADDF:.*]] = arith.addf %[[SIN]], %[[CST]] : vector<[4]xf32>94 %2 = arith.addf %0, %1 : vector<2x[2]xf32>95 96 // CHECK: %[[RES:.*]] = vector.shape_cast %[[ADDF]] : vector<[4]xf32> to vector<2x[2]xf32>97 // CHECK: return %[[RES]] : vector<2x[2]xf32>98 return %2 : vector<2x[2]xf32>99}100 101// -----102 103// CHECK-LABEL: func.func @test_scalable_no_linearize(104// CHECK-SAME: %[[VAL_0:.*]]: vector<[2]x[2]xf32>) -> vector<[2]x[2]xf32> {105func.func @test_scalable_no_linearize(%arg0: vector<[2]x[2]xf32>) -> vector<[2]x[2]xf32> {106 107 // CHECK: %[[CST:.*]] = arith.constant dense<2.000000e+00> : vector<[2]x[2]xf32>108 %0 = arith.constant dense<[[2., 2.], [2., 2.]]> : vector<[2]x[2]xf32>109 110 // CHECK: %[[SIN:.*]] = math.sin %[[VAL_0]] : vector<[2]x[2]xf32>111 %1 = math.sin %arg0 : vector<[2]x[2]xf32>112 113 // CHECK: %[[RES:.*]] = arith.addf %[[CST]], %[[SIN]] : vector<[2]x[2]xf32>114 %2 = arith.addf %0, %1 : vector<[2]x[2]xf32>115 116 // CHECK: return %[[RES]] : vector<[2]x[2]xf32>117 return %2 : vector<[2]x[2]xf32>118}119 120// -----121 122// CHECK-LABEL: func.func @test_0d_vector123func.func @test_0d_vector() -> vector<f32> {124 125 // CHECK: %[[CST:.+]] = arith.constant dense<0.000000e+00> : vector<f32>126 %0 = arith.constant dense<0.0> : vector<f32>127 128 // CHECK: return %[[CST]]129 return %0 : vector<f32>130}131 132// -----133 134// CHECK-LABEL: test_extract_strided_slice_2D135// CHECK-SAME: (%[[ORIG_ARG:.*]]: vector<4x8xf32>) -> vector<2x2xf32> {136func.func @test_extract_strided_slice_2D(%arg0 : vector<4x8xf32>) -> vector<2x2xf32> {137 138 // CHECK: %[[ARG:.*]] = vector.shape_cast %[[ORIG_ARG]] : vector<4x8xf32> to vector<32xf32>139 // CHECK: %[[SHUFFLE:.*]] = vector.shuffle %[[ARG]], %[[ARG]]140 // CHECK-SAME: [4, 5, 12, 13] : vector<32xf32>, vector<32xf32>141 // CHECK: %[[RES:.*]] = vector.shape_cast %[[SHUFFLE]] : vector<4xf32> to vector<2x2xf32>142 // CHECK: return %[[RES]] : vector<2x2xf32143 %0 = vector.extract_strided_slice %arg0 { sizes = [2, 2], strides = [1, 1], offsets = [0, 4]}144 : vector<4x8xf32> to vector<2x2xf32>145 return %0 : vector<2x2xf32>146}147 148// -----149 150// CHECK-LABEL: func.func @test_extract_strided_slice_2D_scalable(151// CHECK-SAME: %[[VAL_0:.*]]: vector<4x[8]xf32>) -> vector<2x[8]xf32> {152func.func @test_extract_strided_slice_2D_scalable(%arg0: vector<4x[8]xf32>) -> vector<2x[8]xf32> {153 154 // CHECK-NOT: vector.shuffle155 // CHECK-NOT: vector.shape_cast156 // CHECK: %[[RES:.*]] = vector.extract_strided_slice %[[VAL_0]]157 %0 = vector.extract_strided_slice %arg0 { sizes = [2, 8], strides = [1, 1], offsets = [1, 0] } : vector<4x[8]xf32> to vector<2x[8]xf32>158 159 // CHECK: return %[[RES]] : vector<2x[8]xf32>160 return %0 : vector<2x[8]xf32>161}162 163// -----164 165// CHECK-LABEL: test_extract_strided_slice_3D166// CHECK-SAME: (%[[ORIG_ARG:.*]]: vector<2x8x2xf32>) -> vector<1x4x2xf32> {167func.func @test_extract_strided_slice_3D(%arg0 : vector<2x8x2xf32>) -> vector<1x4x2xf32> {168 169 // CHECK: %[[ARG:.*]] = vector.shape_cast %[[ORIG_ARG]] : vector<2x8x2xf32> to vector<32xf32>170 // CHECK: %[[SHUFFLE:.*]] = vector.shuffle %[[ARG]], %[[ARG]]171 // CHECK-SAME: [20, 21, 22, 23, 24, 25, 26, 27] : vector<32xf32>, vector<32xf32>172 // CHECK: %[[RES:.*]] = vector.shape_cast %[[SHUFFLE]] : vector<8xf32> to vector<1x4x2xf32>173 // CHECK: return %[[RES]] : vector<1x4x2xf32>174 %0 = vector.extract_strided_slice %arg0 { offsets = [1, 2], strides = [1, 1], sizes = [1, 4] }175 : vector<2x8x2xf32> to vector<1x4x2xf32>176 return %0 : vector<1x4x2xf32>177}178 179// -----180 181// Test of insert_strided_slice -> shuffle.182// This is a contiguous insertion of 4 elements at offset 6 into a vector of 12 elements.183// CHECK-LABEL: insert_strided_slice_2D_into_4D184func.func @insert_strided_slice_2D_into_4D(%arg0 : vector<2x2xi8>, %arg1 : vector<2x1x3x2xi8>) -> vector<2x1x3x2xi8> {185 186// CHECK-DAG: %[[ARG0:.*]] = vector.shape_cast {{.*}} to vector<4xi8>187// CHECK-DAG: %[[ARG1:.*]] = vector.shape_cast {{.*}} to vector<12xi8>188// CHECK: vector.shuffle %[[ARG1]], %[[ARG0]]189// CHECK-SAME: [0, 1, 2, 3, 4, 5, 12, 13, 14, 15, 10, 11] : vector<12xi8>, vector<4xi8>190 %0 = vector.insert_strided_slice %arg0, %arg1 {offsets = [1, 0, 0, 0], strides = [1, 1]} : vector<2x2xi8> into vector<2x1x3x2xi8>191 192// CHECK: %[[RES:.*]] = vector.shape_cast {{.*}} to vector<2x1x3x2xi8>193// CHECK: return %[[RES]] : vector<2x1x3x2xi8>194 return %0 : vector<2x1x3x2xi8>195}196 197// -----198 199// Test of insert_strided_slice -> shuffle.200// [[[0, 1], [2, 3], [4, 5]], [[6, 7], [8, 9], [10, 11]], [[12, 13], [14, 15]], [[16, 17]]]201// ^ ^202// | |203// where the 2 elements are inserted into the 3x3x2 vector204// CHECK-LABEL: insert_strided_slice_3D205func.func @insert_strided_slice_3D(%arg0 : vector<1x2x1xi8>, %arg1 : vector<3x3x2xi8>) -> vector<3x3x2xi8> {206 207// CHECK-DAG: %[[ARG0:.*]] = vector.shape_cast {{.*}} to vector<2xi8>208// CHECK-DAG: %[[ARG1:.*]] = vector.shape_cast {{.*}} to vector<18xi8>209// CHECK: vector.shuffle %[[ARG1]], %[[ARG0]]210// CHECK-SAME: [0, 1, 2, 3, 4, 5, 6, 7, 8, 18, 10, 19, 12, 13, 14, 15, 16, 17] : vector<18xi8>, vector<2xi8>211 %0 = vector.insert_strided_slice %arg0, %arg1 {offsets = [1, 1, 1], sizes = [1, 2, 1], strides = [1, 1, 1]} : vector<1x2x1xi8> into vector<3x3x2xi8>212 213// CHECK: %[[RES:.*]] = vector.shape_cast {{.*}} to vector<3x3x2xi8>214// CHECK: return %[[RES]] : vector<3x3x2xi8>215 return %0 : vector<3x3x2xi8>216}217 218// -----219 220// CHECK-LABEL: insert_strided_slice_2D_higher_offsets221func.func @insert_strided_slice_2D_higher_offsets(%arg0 : vector<2x1xi8>, %arg1 : vector<2x2xi8>, %arg2 : vector<5x2xi8>) -> vector<5x2xi8> {222 223 // CHECK: [0, 1, 2, 3, 10, 11, 12, 13, 8, 9]224 // ^^^ ^^^ ^^^ ^^^225 // insertion indices226 %0 = vector.insert_strided_slice %arg1, %arg2 {offsets = [2, 0], sizes = [2, 2], strides = [1, 1]} : vector<2x2xi8> into vector<5x2xi8>227 228 // CHECK: [0, 1, 2, 3, 10, 5, 11, 7, 8, 9]229 // ^^^ ^^^230 %1 = vector.insert_strided_slice %arg0, %0 {offsets = [2, 0], sizes = [2, 1], strides = [1, 1]} : vector<2x1xi8> into vector<5x2xi8>231 232 // CHECK: [0, 1, 2, 3, 4, 5, 6, 10, 8, 11]233 // ^^^ ^^^234 %2 = vector.insert_strided_slice %arg0, %1 {offsets = [3, 1], sizes = [2, 1], strides = [1, 1]} : vector<2x1xi8> into vector<5x2xi8>235 236 return %2 : vector<5x2xi8>237}238 239// -----240 241// CHECK-LABEL: negative_insert_strided_slice_scalable242// CHECK-NOT: vector.shuffle243// CHECK: return244func.func @negative_insert_strided_slice_scalable(%arg0 : vector<1x[2]xi8>, %arg1 : vector<2x[2]xi8>) -> vector<2x[2]xi8> {245 %0 = vector.insert_strided_slice %arg0, %arg1 {offsets = [0, 0], strides = [1,1]} : vector<1x[2]xi8> into vector<2x[2]xi8>246 return %0 : vector<2x[2]xi8>247}248 249// -----250 251// CHECK-LABEL: test_vector_shuffle252// CHECK-SAME: (%[[ORIG_ARG0:.*]]: vector<4x2xf32>, %[[ORIG_ARG1:.*]]: vector<4x2xf32>) -> vector<8x2xf32> {253func.func @test_vector_shuffle(%arg0: vector<4x2xf32>, %arg1: vector<4x2xf32>) -> vector<8x2xf32> {254 255 // CHECK-DAG: %[[ARG0:.*]] = vector.shape_cast %[[ORIG_ARG0]] : vector<4x2xf32> to vector<8xf32>256 // CHECK-DAG: %[[ARG1:.*]] = vector.shape_cast %[[ORIG_ARG1]] : vector<4x2xf32> to vector<8xf32>257 // CHECK: %[[SHUFFLE:.*]] = vector.shuffle %[[ARG0]], %[[ARG1]]258 // CHECK-SAME: [0, 1, 8, 9, 2, 3, 10, 11, 4, 5, 12, 13, 6, 7, 14, 15] : vector<8xf32>, vector<8xf32>259 // CHECK: %[[RES:.*]] = vector.shape_cast %[[SHUFFLE]] : vector<16xf32> to vector<8x2xf32>260 // CHECK: return %[[RES]] : vector<8x2xf32>261 %0 = vector.shuffle %arg0, %arg1 [0, 4, 1, 5, 2, 6, 3, 7] : vector<4x2xf32>, vector<4x2xf32>262 return %0 : vector<8x2xf32>263}264 265// -----266 267// CHECK-LABEL: test_vector_extract_scalar268// CHECK-SAME: (%[[ARG:.*]]: vector<2x4xi32>) -> i32 {269func.func @test_vector_extract_scalar(%arg0 : vector<2x4xi32>) -> i32 {270 271 // CHECK: %[[SRC_1D:.*]] = vector.shape_cast %[[ARG]] : vector<2x4xi32> to vector<8xi32>272 // CHECK: %[[EXTRACT_1D:.*]] = vector.extract %[[SRC_1D]][6] : i32 from vector<8xi32>273 // CHECK: return %[[EXTRACT_1D]] : i32274 %0 = vector.extract %arg0[1, 2] : i32 from vector<2x4xi32>275 return %0 : i32276}277 278// -----279 280// CHECK-LABEL: test_vector_extract281// CHECK-SAME: (%[[ORIG_ARG:.*]]: vector<2x8x2xf32>) -> vector<8x2xf32> {282func.func @test_vector_extract(%arg0: vector<2x8x2xf32>) -> vector<8x2xf32> {283 284 // CHECK: %[[ARG:.*]] = vector.shape_cast %[[ORIG_ARG]] : vector<2x8x2xf32> to vector<32xf32>285 // CHECK: %[[SHUFFLE:.*]] = vector.shuffle %[[ARG]], %[[ARG]]286 // CHECK-SAME: [16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31] : vector<32xf32>, vector<32xf32>287 // CHECK: %[[RES:.*]] = vector.shape_cast %[[SHUFFLE]] : vector<16xf32> to vector<8x2xf32>288 // CHECK: return %[[RES]] : vector<8x2xf32>289 %0 = vector.extract %arg0[1]: vector<8x2xf32> from vector<2x8x2xf32>290 return %0 : vector<8x2xf32>291}292 293// -----294 295// CHECK-LABEL: func.func @test_vector_extract_scalable(296// CHECK-SAME: %[[VAL_0:.*]]: vector<2x8x[2]xf32>) -> vector<8x[2]xf32> {297func.func @test_vector_extract_scalable(%arg0: vector<2x8x[2]xf32>) -> vector<8x[2]xf32> {298 299 // CHECK-NOT: vector.shuffle300 // CHECK-NOT: vector.shape_cast301 // CHECK: %[[RES:.*]] = vector.extract %[[VAL_0]][1] : vector<8x[2]xf32> from vector<2x8x[2]xf32>302 %0 = vector.extract %arg0[1]: vector<8x[2]xf32> from vector<2x8x[2]xf32>303 304 // CHECK: return %[[RES]] : vector<8x[2]xf32>305 return %0 : vector<8x[2]xf32>306}307 308// -----309 310// CHECK-LABEL: test_vector_insert_scalar311// CHECK-SAME: (%[[DEST:.*]]: vector<2x4xf32>, %[[SRC:.*]]: f32) -> vector<2x4xf32> {312func.func @test_vector_insert_scalar(%arg0: vector<2x4xf32>, %arg1: f32) -> vector<2x4xf32> {313 314 // CHECK: %[[DEST_1D:.*]] = vector.shape_cast %[[DEST]] : vector<2x4xf32> to vector<8xf32>315 // CHECK: %[[INSERT_1D:.*]] = vector.insert %[[SRC]], %[[DEST_1D]] [6] : f32 into vector<8xf32>316 // CHECK: %[[RES:.*]] = vector.shape_cast %[[INSERT_1D]] : vector<8xf32> to vector<2x4xf32>317 // CHECK: return %[[RES]] : vector<2x4xf32>318 %0 = vector.insert %arg1, %arg0[1, 2]: f32 into vector<2x4xf32>319 return %0 : vector<2x4xf32>320}321 322// -----323 324// CHECK-LABEL: test_vector_insert325// CHECK-SAME: (%[[DEST:.*]]: vector<2x8x4xf32>, %[[SRC:.*]]: vector<8x4xf32>) -> vector<2x8x4xf32> {326func.func @test_vector_insert(%arg0: vector<2x8x4xf32>, %arg1: vector<8x4xf32>) -> vector<2x8x4xf32> {327 328 // CHECK-DAG: %[[ARG_SRC:.*]] = vector.shape_cast %[[SRC]] : vector<8x4xf32> to vector<32xf32>329 // CHECK-DAG: %[[ARG_DEST:.*]] = vector.shape_cast %[[DEST]] : vector<2x8x4xf32> to vector<64xf32>330 // CHECK: %[[SHUFFLE:.*]] = vector.shuffle %[[ARG_DEST]], %[[ARG_SRC]]331 // CHECK-SAME: [64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87,332 // CHECK-SAME: 88, 89, 90, 91, 92, 93, 94, 95, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48,333 // CHECK-SAME: 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63] : vector<64xf32>, vector<32xf32>334 // CHECK: %[[RES:.*]] = vector.shape_cast %[[SHUFFLE]] : vector<64xf32> to vector<2x8x4xf32>335 // CHECK: return %[[RES]] : vector<2x8x4xf32>336 %0 = vector.insert %arg1, %arg0[0]: vector<8x4xf32> into vector<2x8x4xf32>337 return %0 : vector<2x8x4xf32>338}339 340// -----341 342// CHECK-LABEL: func.func @test_vector_insert_scalable(343// CHECK-SAME: %[[VAL_0:.*]]: vector<2x8x[4]xf32>, %[[VAL_1:.*]]: vector<8x[4]xf32>) -> vector<2x8x[4]xf32> {344func.func @test_vector_insert_scalable(%arg0: vector<2x8x[4]xf32>, %arg1: vector<8x[4]xf32>) -> vector<2x8x[4]xf32> {345 346 // CHECK-NOT: vector.shuffle347 // CHECK-NOT: vector.shape_cast348 // CHECK: %[[RES:.*]] = vector.insert %[[VAL_1]], %[[VAL_0]] [0] : vector<8x[4]xf32> into vector<2x8x[4]xf32>349 350 %0 = vector.insert %arg1, %arg0[0]: vector<8x[4]xf32> into vector<2x8x[4]xf32>351 // CHECK: return %[[RES]] : vector<2x8x[4]xf32>352 return %0 : vector<2x8x[4]xf32>353}354 355// -----356 357// CHECK-LABEL: test_vector_bitcast358// CHECK-SAME: %[[ARG_0:.*]]: vector<4x4xf32>359func.func @test_vector_bitcast(%arg0: vector<4x4xf32>) -> vector<4x8xf16> {360 361 // CHECK: %[[DOWNCAST:.*]] = vector.shape_cast %[[ARG_0]] : vector<4x4xf32> to vector<16xf32>362 // CHECK: %[[BITCAST:.*]] = vector.bitcast %[[DOWNCAST]] : vector<16xf32> to vector<32xf16>363 // CHECK: %[[UPCAST:.*]] = vector.shape_cast %[[BITCAST]] : vector<32xf16> to vector<4x8xf16>364 %1 = vector.bitcast %arg0 : vector<4x4xf32> to vector<4x8xf16>365 return %1 : vector<4x8xf16>366}367 368// -----369 370// CHECK-LABEL: test_vector_bitcast371// CHECK-SAME: %[[ARG_0:.*]]: vector<4x2xf32>372func.func @test_vector_bitcast(%arg0: vector<4x2xf32>) -> vector<4x4xf16> {373 374 // CHECK: %[[DOWNCAST:.*]] = vector.shape_cast %[[ARG_0]] : vector<4x2xf32> to vector<8xf32>375 // CHECK: %[[BITCAST:.*]] = vector.bitcast %[[DOWNCAST]] : vector<8xf32> to vector<16xf16>376 // CHECK: %[[UPCAST:.*]] = vector.shape_cast %[[BITCAST]] : vector<16xf16> to vector<4x4xf16>377 %1 = vector.bitcast %arg0 : vector<4x2xf32> to vector<4x4xf16>378 return %1 : vector<4x4xf16>379}380 381// -----382 383// CHECK-LABEL: test_vector_bitcast384// CHECK-SAME: %[[ARG_0:.*]]: vector<4x[2]xf32>385func.func @test_vector_bitcast(%arg0: vector<4x[2]xf32>) -> vector<4x[4]xf16> {386 387 // CHECK: %[[DOWNCAST:.*]] = vector.shape_cast %[[ARG_0]] : vector<4x[2]xf32> to vector<[8]xf32>388 // CHECK: %[[BITCAST:.*]] = vector.bitcast %[[DOWNCAST]] : vector<[8]xf32> to vector<[16]xf16>389 // CHECK: %[[UPCAST:.*]] = vector.shape_cast %[[BITCAST]] : vector<[16]xf16> to vector<4x[4]xf16>390 %1 = vector.bitcast %arg0 : vector<4x[2]xf32> to vector<4x[4]xf16>391 return %1 : vector<4x[4]xf16>392}393 394// -----395 396// CHECK-LABEL: test_vector_bitcast397// CHECK-SAME: %[[ARG_0:.*]]: vector<[4]x2xf32>398func.func @test_vector_bitcast(%arg0: vector<[4]x2xf32>) -> vector<[4]x4xf16> {399 400 // CHECK: %[[DOWNCAST:.*]] = vector.shape_cast %[[ARG_0]] : vector<[4]x2xf32> to vector<[8]xf32>401 // CHECK: %[[BITCAST:.*]] = vector.bitcast %[[DOWNCAST]] : vector<[8]xf32> to vector<[16]xf16>402 // CHECK: %[[UPCAST:.*]] = vector.shape_cast %[[BITCAST]] : vector<[16]xf16> to vector<[4]x4xf16>403 %1 = vector.bitcast %arg0 : vector<[4]x2xf32> to vector<[4]x4xf16>404 return %1 : vector<[4]x4xf16>405}406 407// -----408 409// CHECK-LABEL: test_linearize_across_for410func.func @test_linearize_across_for(%arg0 : vector<4xi8>) -> vector<4xi8> {411 %0 = vector.shape_cast %arg0 : vector<4xi8> to vector<2x2xi8>412 %c0 = arith.constant 0 : index413 %c1 = arith.constant 1 : index414 %c4 = arith.constant 4 : index415 416 // CHECK: scf.for {{.*}} -> (vector<4xi8>)417 %1 = scf.for %i = %c0 to %c4 step %c1 iter_args(%arg1 = %0) -> (vector<2x2xi8>) {418 419 // CHECK: arith.addi {{.*}} : vector<4xi8>420 %2 = arith.addi %arg1, %0 : vector<2x2xi8>421 422 // CHECK: scf.yield {{.*}} : vector<4xi8>423 scf.yield %2 : vector<2x2xi8>424 }425 %3 = vector.shape_cast %1 : vector<2x2xi8> to vector<4xi8>426 return %3 : vector<4xi8>427}428 429// -----430 431// CHECK-LABEL: linearize_vector_broadcast_scalar_source432// CHECK-SAME: (%[[ARG:.*]]: i32) -> vector<4x2xi32>433func.func @linearize_vector_broadcast_scalar_source(%arg0: i32) -> vector<4x2xi32> {434 435 // CHECK: %[[BROADCAST:.*]] = vector.broadcast %[[ARG]] : i32 to vector<8xi32>436 // CHECK: %[[CAST:.*]] = vector.shape_cast %[[BROADCAST]] : vector<8xi32> to vector<4x2xi32>437 // CHECK: return %[[CAST]] : vector<4x2xi32>438 %0 = vector.broadcast %arg0 : i32 to vector<4x2xi32>439 return %0 : vector<4x2xi32>440}441 442// -----443 444// CHECK-LABEL: linearize_vector_broadcast_rank_two_source445// CHECK-SAME: (%[[ARG:.*]]: vector<1x1xi32>) -> vector<4x2xi32>446func.func @linearize_vector_broadcast_rank_two_source(%arg0: vector<1x1xi32>) -> vector<4x2xi32> {447 448 // CHECK: %[[CAST0:.*]] = vector.shape_cast %[[ARG]] : vector<1x1xi32> to vector<1xi32>449 // CHECK: %[[BROADCAST:.*]] = vector.broadcast %[[CAST0]] : vector<1xi32> to vector<8xi32>450 // CHECK: %[[CAST1:.*]] = vector.shape_cast %[[BROADCAST]] : vector<8xi32> to vector<4x2xi32>451 // CHECK: return %[[CAST1]] : vector<4x2xi32>452 %0 = vector.broadcast %arg0 : vector<1x1xi32> to vector<4x2xi32>453 return %0 : vector<4x2xi32>454}455 456// -----457 458// CHECK-LABEL: linearize_scalable_vector_broadcast459// CHECK-SAME: (%[[ARG:.*]]: i32) -> vector<4x[2]xi32>460func.func @linearize_scalable_vector_broadcast(%arg0: i32) -> vector<4x[2]xi32> {461 462 // CHECK: %[[BROADCAST:.*]] = vector.broadcast %[[ARG]] : i32 to vector<[8]xi32>463 // CHECK: %[[CAST:.*]] = vector.shape_cast %[[BROADCAST]] : vector<[8]xi32> to vector<4x[2]xi32>464 // CHECK: return %[[CAST]] : vector<4x[2]xi32>465 %0 = vector.broadcast %arg0 : i32 to vector<4x[2]xi32>466 return %0 : vector<4x[2]xi32>467 468}469 470// -----471 472// CHECK-LABEL: linearize_create_mask473// CHECK-SAME: (%[[ARG0:.*]]: index, %[[ARG1:.*]]: index) -> vector<1x16xi1>474func.func @linearize_create_mask(%arg0 : index, %arg1 : index) -> vector<1x16xi1> {475 476 // CHECK: %[[C0:.*]] = arith.constant 0 : index477 // CHECK: %[[CMP:.*]] = arith.cmpi sgt, %[[ARG0]], %[[C0]] : index478 // CHECK: %[[INDEXCAST:.*]] = arith.index_cast %[[CMP]] : i1 to index479 // CHECK: %[[MULI:.*]] = arith.andi %[[INDEXCAST]], %[[ARG1]] : index480 // CHECK: %[[MASK_1D:.*]] = vector.create_mask %[[MULI]] : vector<16xi1>481 // CHECK: %[[CAST:.*]] = vector.shape_cast %[[MASK_1D]] : vector<16xi1> to vector<1x16xi1>482 // CHECK: return %[[CAST]] : vector<1x16xi1>483 %0 = vector.create_mask %arg0, %arg1 : vector<1x16xi1>484 return %0 : vector<1x16xi1>485}486 487// -----488// CHECK-LABEL: linearize_scalable_create_mask489func.func @linearize_scalable_create_mask(%arg0 : index, %arg1 : index) -> vector<1x[16]xi1> {490 491 // CHECK: %[[MASK_1D:.*]] = vector.create_mask {{%.*}} : vector<[16]xi1>492 %0 = vector.create_mask %arg0, %arg1 : vector<1x[16]xi1>493 return %0 : vector<1x[16]xi1>494}495 496// CHECK-LABEL: linearize_vector_load497// CHECK-SAME: (%[[ARG0:.*]]: memref<2x8xf32>) -> vector<1x4xf32>498func.func @linearize_vector_load(%arg0: memref<2x8xf32>) -> vector<1x4xf32> {499 // CHECK: %[[CST0:.*]] = arith.constant 0 : index500 // CHECK: %[[LOAD:.*]] = vector.load %[[ARG0]][%[[CST0]], %[[CST0]]] : memref<2x8xf32>, vector<4xf32>501 // CHECK: %[[CAST:.*]] = vector.shape_cast %[[LOAD]] : vector<4xf32> to vector<1x4xf32>502 // CHECK: return %[[CAST]] : vector<1x4xf32>503 %c0 = arith.constant 0 : index504 %0 = vector.load %arg0[%c0, %c0] : memref<2x8xf32>, vector<1x4xf32>505 return %0 : vector<1x4xf32>506}507 508// CHECK-LABEL: linearize_vector_store509// CHECK-SAME: (%[[ARG0:.*]]: memref<2x8xf32>, %[[ARG1:.*]]: vector<1x4xf32>)510func.func @linearize_vector_store(%arg0: memref<2x8xf32>, %arg1: vector<1x4xf32>) {511 // CHECK: %[[CAST:.*]] = vector.shape_cast %arg1 : vector<1x4xf32> to vector<4xf32>512 // CHECK: %[[CST0:.*]] = arith.constant 0 : index513 // CHECK: vector.store %[[CAST]], %[[ARG0]][%[[CST0]], %[[CST0]]] : memref<2x8xf32>, vector<4xf32>514 %c0 = arith.constant 0 : index515 vector.store %arg1, %arg0[%c0, %c0] : memref<2x8xf32>, vector<1x4xf32>516 return517}518 519// CHECK-LABEL: linearize_vector_load_scalable520// CHECK-SAME: (%[[ARG0:.*]]: memref<2x8xf32>) -> vector<1x[4]xf32>521func.func @linearize_vector_load_scalable(%arg0: memref<2x8xf32>) -> vector<1x[4]xf32> {522 // CHECK: %[[CST0:.*]] = arith.constant 0 : index523 // CHECK: %[[LOAD:.*]] = vector.load %[[ARG0]][%[[CST0]], %[[CST0]]] : memref<2x8xf32>, vector<[4]xf32>524 // CHECK: %[[CAST:.*]] = vector.shape_cast %[[LOAD]] : vector<[4]xf32> to vector<1x[4]xf32>525 // CHECK: return %[[CAST]] : vector<1x[4]xf32>526 %c0 = arith.constant 0 : index527 %0 = vector.load %arg0[%c0, %c0] : memref<2x8xf32>, vector<1x[4]xf32>528 return %0 : vector<1x[4]xf32>529}530 531// CHECK-LABEL: linearize_vector_store_scalable532// CHECK-SAME: (%[[ARG0:.*]]: memref<2x8xf32>, %[[ARG1:.*]]: vector<1x[4]xf32>)533func.func @linearize_vector_store_scalable(%arg0: memref<2x8xf32>, %arg1: vector<1x[4]xf32>) {534 // CHECK: %[[CAST:.*]] = vector.shape_cast %arg1 : vector<1x[4]xf32> to vector<[4]xf32>535 // CHECK: %[[CST0:.*]] = arith.constant 0 : index536 // CHECK: vector.store %[[CAST]], %[[ARG0]][%[[CST0]], %[[CST0]]] : memref<2x8xf32>, vector<[4]xf32>537 %c0 = arith.constant 0 : index538 vector.store %arg1, %arg0[%c0, %c0] : memref<2x8xf32>, vector<1x[4]xf32>539 return540}541 542// -----543 544// Test pattern LinearizeVectorFromElements.545 546// CHECK-LABEL: test_vector_from_elements547// CHECK-SAME: %[[ARG_0:.*]]: f32, %[[ARG_1:.*]]: f32, %[[ARG_2:.*]]: f32, %[[ARG_3:.*]]: f32548func.func @test_vector_from_elements(%arg0: f32, %arg1: f32, %arg2: f32, %arg3: f32) -> vector<2x2xf32> {549 // CHECK: %[[FROM_ELEMENTS:.*]] = vector.from_elements %[[ARG_0]], %[[ARG_1]], %[[ARG_2]], %[[ARG_3]] : vector<4xf32>550 // CHECK: %[[CAST:.*]] = vector.shape_cast %[[FROM_ELEMENTS]] : vector<4xf32> to vector<2x2xf32>551 // CHECK: return %[[CAST]] : vector<2x2xf32>552 %1 = vector.from_elements %arg0, %arg1, %arg2, %arg3 : vector<2x2xf32>553 return %1 : vector<2x2xf32>554}555 556// -----557 558// CHECK-LABEL: func.func @to_elements_1d(559// CHECK-SAME: %[[ARG0:.+]]: vector<2xf32>560// CHECK: %[[RES:.+]]:2 = vector.to_elements %[[ARG0]] : vector<2xf32>561// CHECK: return %[[RES]]#0, %[[RES]]#1562func.func @to_elements_1d(%arg0: vector<2xf32>) -> (f32, f32) {563 %0:2 = vector.to_elements %arg0 : vector<2xf32>564 return %0#0, %0#1 : f32, f32565}566 567// -----568 569// CHECK-LABEL: func.func @to_elements_2d(570// CHECK-SAME: %[[ARG0:.+]]: vector<2x2xf32>571// CHECK: %[[CAST:.+]] = vector.shape_cast %[[ARG0]]572// CHECK: %[[RES:.+]]:4 = vector.to_elements %[[CAST]] : vector<4xf32>573// CHECK: return %[[RES]]#0, %[[RES]]#1, %[[RES]]#2, %[[RES]]#3574func.func @to_elements_2d(%arg0: vector<2x2xf32>) -> (f32, f32, f32, f32) {575 %0:4 = vector.to_elements %arg0 : vector<2x2xf32>576 return %0#0, %0#1, %0#2, %0#3 : f32, f32, f32, f32577}578