687 lines · plain
1// RUN: mlir-opt -split-input-file \2// RUN: -transform-preload-library='transform-library-paths=%p/../td/vectorize-with-patterns.mlir' \3// RUN: -transform-interpreter=entry-point=vectorize_with_patterns %s | FileCheck %s4 5//===----------------------------------------------------------------------===//6// Contiguous load7//===----------------------------------------------------------------------===//8 9#map = affine_map<(d0, d1, d2) -> (d0, d1, d2)>10func.func @vectorize_nd_tensor_extract_transfer_read_basic(11 %arg0: tensor<3x3x3xf32>,12 %arg1: tensor<1x1x3xf32>) -> tensor<1x1x3xf32> {13 14 %res = linalg.generic {15 indexing_maps = [#map],16 iterator_types = ["parallel", "parallel", "parallel"]17 } outs(%arg1 : tensor<1x1x3xf32>) {18 ^bb0(%out: f32):19 %1 = linalg.index 0 : index20 %2 = linalg.index 1 : index21 %3 = linalg.index 2 : index22 %4 = tensor.extract %arg0[%1, %2, %3] : tensor<3x3x3xf32>23 linalg.yield %4 : f3224 } -> tensor<1x1x3xf32>25 26 return %res : tensor<1x1x3xf32>27}28 29// CHECK-LABEL: func.func @vectorize_nd_tensor_extract_transfer_read_basic30// CHECK-SAME: %[[ARG0:.*]]: tensor<3x3x3xf32>31// CHECK-SAME: %[[ARG1:.*]]: tensor<1x1x3xf32>32 33// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index34// CHECK-DAG: %[[PV:.+]] = ub.poison : f3235 36// CHECK: %[[READ:.*]] = vector.transfer_read %[[ARG0]][%[[C0]], %[[C0]], %[[C0]]], %[[PV]] {in_bounds = [true, true, true]} : tensor<3x3x3xf32>, vector<1x1x3xf32>37// CHECK: vector.transfer_write %[[READ]], %[[ARG1]][%[[C0]], %[[C0]], %[[C0]]] {in_bounds = [true, true, true]} : vector<1x1x3xf32>, tensor<1x1x3xf32>38 39 // -----40 41func.func @vectorize_nd_tensor_extract_transfer_read_complex(%6: tensor<45x80x16xf32>, %arg0: index, %arg2: index, %arg1: index, %arg4: index, %extracted_slice : tensor<1x4xf32>) -> tensor<1x4xf32> {42 %c79 = arith.constant 79 : index43 %25 = linalg.generic {44 indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>],45 iterator_types = ["parallel", "parallel"]46 } outs(%extracted_slice : tensor<1x4xf32>) {47 ^bb0(%out: f32):48 %26 = linalg.index 0 : index49 %27 = arith.addi %arg0, %26 : index50 %28 = arith.addi %27, %arg2 : index51 %29 = linalg.index 1 : index52 %30 = arith.addi %arg1, %29 : index53 %31 = arith.addi %30, %arg4 : index54 %extracted = tensor.extract %6[%28, %c79, %31] : tensor<45x80x16xf32>55 linalg.yield %extracted : f3256 } -> tensor<1x4xf32>57 return %25 : tensor<1x4xf32>58}59 60 61// CHECK-LABEL: func.func @vectorize_nd_tensor_extract_transfer_read_complex(62// CHECK-SAME: %[[ARG0:.*]]: tensor<45x80x16xf32>,63// CHECK-SAME: %[[ARG1:.*]]: index, %[[ARG2:.*]]: index, %[[ARG3:.*]]: index, %[[ARG4:.*]]: index,64// CHECK-SAME: %[[ARG5:.*]]: tensor<1x4xf32>) -> tensor<1x4xf32> {65 66// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index67// CHECK-DAG: %[[PV:.*]] = ub.poison : f3268// CHECK-DAG: %[[C79:.*]] = arith.constant 79 : index69// CHECK: %[[ADD1:.*]] = arith.addi %[[ARG1]], %[[ARG2]] : index70// CHECK: %[[ADD2:.*]] = arith.addi %[[ARG3]], %[[ARG4]] : index71 72// CHECK: %[[READ:.*]] = vector.transfer_read %[[ARG0]]{{\[}}%[[ADD1]], %[[C79]], %[[ADD2]]], %[[PV]] {in_bounds = [true, true]} : tensor<45x80x16xf32>, vector<1x4xf32>73// CHECK: %[[WRITE:.*]] = vector.transfer_write %[[READ]], %[[ARG5]]{{\[}}%[[C0]], %[[C0]]] {in_bounds = [true, true]} : vector<1x4xf32>, tensor<1x4xf32>74// CHECK: return %[[WRITE]] : tensor<1x4xf32>75// CHECK: }76 77// -----78 79// The vectorizer converts `affine.apply` so that the subsequent Ops can be vectorised based on the converted ops. Contiguous load.80func.func @vectorize_nd_tensor_extract_with_affine_apply_contiguous(%6: tensor<80x16xf32>, %arg0: index, %extracted_slice : tensor<1x4xf32>) -> tensor<1x4xf32> {81 %c79 = arith.constant 79 : index82 %1 = linalg.generic {83 indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>],84 iterator_types = ["parallel", "parallel"]85 } outs(%extracted_slice : tensor<1x4xf32>) {86 ^bb0(%out: f32):87 %2 = linalg.index 1 : index88 %3 = affine.apply affine_map<(d0, d1) -> (d0 + d1)>(%2, %arg0)89 %extracted = tensor.extract %6[%c79, %3] : tensor<80x16xf32>90 linalg.yield %extracted : f3291 } -> tensor<1x4xf32>92 return %1 : tensor<1x4xf32>93}94 95// CHECK-LABEL: func.func @vectorize_nd_tensor_extract_with_affine_apply_contiguous(96// CHECK-SAME: %[[ARG0:.*]]: tensor<80x16xf32>,97// CHECK-SAME: %[[ARG1:.*]]: index,98// CHECK-SAME: %[[ARG2:.*]]: tensor<1x4xf32>) -> tensor<1x4xf32> {99 100// CHECK-DAG: %[[PV:.*]] = ub.poison : f32101// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index102// CHECK-DAG: %[[C79:.*]] = arith.constant 79 : index103 104// CHECK: %[[READ:.*]] = vector.transfer_read %[[ARG0]]{{\[}}%[[C79]], %[[ARG1]]], %[[PV]] {in_bounds = [true, true]} : tensor<80x16xf32>, vector<1x4xf32>105// CHECK: %[[WRITE:.*]] = vector.transfer_write %[[READ]], %[[ARG2]]{{\[}}%[[C0]], %[[C0]]] {in_bounds = [true, true]} : vector<1x4xf32>, tensor<1x4xf32>106// CHECK: return %[[WRITE]] : tensor<1x4xf32>107// CHECK: }108 109// -----110 111func.func @vectorize_nd_tensor_extract_with_tensor_extract(%input_1: tensor<1x20xi32>, %input_2: tensor<257x24xf32>, %arg0 : index, %arg1 : index, %arg2 : index, %arg3 : index) -> tensor<1x1x4xf32> {112 %c0 = arith.constant 0 : index113 %c256 = arith.constant 256 : index114 %output = tensor.empty() : tensor<1x1x4xf32>115 %1 = linalg.generic {indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>], iterator_types = ["parallel", "parallel", "parallel"]} outs(%output : tensor<1x1x4xf32>) {116 ^bb0(%out: f32):117 %13 = linalg.index 0 : index118 %14 = affine.apply affine_map<(d0, d1, d2) -> (d0 + d1 + d2)>(%arg0, %13, %arg2)119 %15 = linalg.index 2 : index120 %16 = linalg.index 1 : index121 %17 = affine.apply affine_map<(d0, d1, d2, d3) -> (d0 + d1 * 24 + d2 + d3)>(%arg1, %16, %15, %arg3)122 %extracted_0 = tensor.extract %input_1[%c0, %14] : tensor<1x20xi32>123 %18 = arith.index_cast %extracted_0 : i32 to index124 %19 = arith.maxsi %18, %c0 : index125 %20 = arith.minsi %19, %c256 : index126 %extracted_1 = tensor.extract %input_2[%20, %17] : tensor<257x24xf32>127 linalg.yield %extracted_1 : f32128 } -> tensor<1x1x4xf32>129 return %1 : tensor<1x1x4xf32>130}131 132// CHECK-LABEL: func.func @vectorize_nd_tensor_extract_with_tensor_extract(133// CHECK-SAME: %[[INPUT_1:.*]]: tensor<1x20xi32>,134// CHECK-SAME: %[[INPUT_2:.*]]: tensor<257x24xf32>,135// CHECK-SAME: %[[INPUT_3:.*]]: index, %[[INPUT_4:.*]]: index, %[[INPUT_5:.*]]: index,136// CHECK: %[[EXTRACTED_0_IDX_0:.*]] = arith.constant 0 : index137// CHECK: %[[SCALAR:.*]] = arith.addi %[[INPUT_3]], %[[INPUT_5]] : index138// First `vector.transfer_read` from the generic Op - loop invariant scalar load.139// CHECK: vector.transfer_read %[[INPUT_1]][%[[EXTRACTED_0_IDX_0]], %[[SCALAR]]]140// CHECK-SAME: tensor<1x20xi32>, vector<i32>141// The following `tensor.extract` from the generic Op s a contiguous load (all Ops used142// for address calculation also satisfy the required conditions).143// CHECK: vector.transfer_read %[[INPUT_2]][%{{.*}}, %{{.*}}, %{{.*}} {in_bounds = [true, true]} : tensor<257x24xf32>, vector<1x4xf32>144 145// Make sure that non-linear arithmetic operations (e.g. arith.maxsi) are allowed when calculating indices for load operations. Contiguous load.146func.func @vectorize_nd_tensor_extract_with_maxsi_contiguous(%arg0: tensor<80x16xf32>, %extracted_slice : tensor<1x4xf32>) -> tensor<1x4xf32> {147 %c16 = arith.constant 16 : index148 %1 = linalg.generic {149 indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>],150 iterator_types = ["parallel", "parallel"]151 } outs(%extracted_slice : tensor<1x4xf32>) {152 ^bb0(%out: f32):153 %2 = linalg.index 0 : index154 %3 = linalg.index 1 : index155 %4 = arith.maxsi %2, %c16 : index156 %extracted = tensor.extract %arg0[%4, %3] : tensor<80x16xf32>157 linalg.yield %extracted : f32158 } -> tensor<1x4xf32>159 return %1 : tensor<1x4xf32>160}161 162// CHECK-LABEL: func.func @vectorize_nd_tensor_extract_with_maxsi_contiguous(163// CHECK-SAME: %[[VAL_0:.*]]: tensor<80x16xf32>,164// CHECK-SAME: %[[VAL_1:.*]]: tensor<1x4xf32>) -> tensor<1x4xf32> {165// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index166// CHECK-DAG: %[[C16:.*]] = arith.constant 16 : index167// CHECK-DAG: %[[PV:.*]] = ub.poison : f32168 169// CHECK: %[[VAL_8:.*]] = vector.transfer_read %[[VAL_0]]{{\[}}%[[C16]], %[[C0]]], %[[PV]] {in_bounds = [true, true]} : tensor<80x16xf32>, vector<1x4xf32>170// CHECK: %[[VAL_9:.*]] = vector.transfer_write %[[VAL_8]], %[[VAL_1]]{{\[}}%[[C0]], %[[C0]]] {in_bounds = [true, true]} : vector<1x4xf32>, tensor<1x4xf32>171// CHECK: return %[[VAL_9]] : tensor<1x4xf32>172// CHECK: }173 174// -----175 176//===----------------------------------------------------------------------===//177// Gather load178//===----------------------------------------------------------------------===//179 180#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d2)>181#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>182func.func @vectorize_1d_tensor_extract(%arg0: tensor<3xf32>, %arg1: tensor<4x3xi32>, %arg2: tensor<4x7x3x2xf32>) -> tensor<4x7x3x2xf32> {183 %1 = linalg.generic {184 indexing_maps = [#map0, #map1],185 iterator_types = ["parallel", "parallel", "parallel", "parallel"]186 } ins(%arg1 : tensor<4x3xi32>) outs(%arg2 : tensor<4x7x3x2xf32>) {187 ^bb0(%arg3: i32, %arg4: f32):188 %2 = arith.index_cast %arg3 : i32 to index189 %3 = tensor.extract %arg0[%2] : tensor<3xf32>190 linalg.yield %3 : f32191 } -> tensor<4x7x3x2xf32>192 return %1 : tensor<4x7x3x2xf32>193}194// CHECK-LABEL: func.func @vectorize_1d_tensor_extract195// CHECK-SAME: %[[ARG0:.*]]: tensor<3xf32>196// CHECK-SAME: %[[ARG1:.*]]: tensor<4x3xi32>197// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index198// CHECK-DAG: %[[MASK:.*]] = arith.constant dense<true> : vector<4x7x3x2xi1>199// CHECK-DAG: %[[PASSTHRU:.*]] = arith.constant dense<0.000000e+00> : vector<4x7x3x2xf32>200// CHECK: %[[V0:.*]] = vector.transfer_read %[[ARG1]]201// CHECK: %[[CAST:.*]] = arith.index_cast %[[V0]]202// CHECK: %[[BROADCAST:.*]] = vector.broadcast %[[CAST]]203// CHECK: %[[INDICES:.*]] = vector.transpose %[[BROADCAST]]204// CHECK: %[[GATHER:.*]] = vector.gather %[[ARG0]][%[[C0]]] [%[[INDICES]]], %[[MASK]], %[[PASSTHRU]]205// CHECK: vector.transfer_write %[[GATHER]]206 207// -----208 209#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d2)>210#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>211#map2 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>212func.func @vectorize_nd_tensor_extract_index_from_tensor(%arg0: tensor<3x3xf32>, %arg1: tensor<4x3xi32>, %arg2: tensor<4x3xi32>, %arg3: tensor<4x7x2xf32>, %arg4: tensor<4x7x3x2xf32>) -> tensor<4x7x3x2xf32> {213 %2 = linalg.generic {214 indexing_maps = [#map0, #map0, #map1, #map2],215 iterator_types = ["parallel", "parallel", "parallel", "parallel"]216 } ins(%arg1, %arg2, %arg3 : tensor<4x3xi32>, tensor<4x3xi32>, tensor<4x7x2xf32>) outs(%arg4 : tensor<4x7x3x2xf32>) {217 ^bb0(%arg5: i32, %arg6: i32, %arg7: f32, %arg8: f32):218 %3 = arith.index_cast %arg5 : i32 to index219 %4 = arith.index_cast %arg6 : i32 to index220 %7 = tensor.extract %arg0[%3, %4] : tensor<3x3xf32>221 linalg.yield %7 : f32222 } -> tensor<4x7x3x2xf32>223 return %2 : tensor<4x7x3x2xf32>224}225// CHECK-LABEL: func.func @vectorize_nd_tensor_extract_index_from_tensor226// CHECK-SAME: %[[ARG0:.*]]: tensor<3x3xf32>227// CHECK-SAME: %[[ARG1:arg1]]: tensor<4x3xi32>228// CHECK-SAME: %[[ARG2:arg2]]: tensor<4x3xi32>229// CHECK-SAME: %[[ARG3:.*]]: tensor<4x7x2xf32>230// CHECK-SAME: %[[ARG4:.*]]: tensor<4x7x3x2xf32>231// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index232// CHECK-DAG: %[[PV:.*]] = ub.poison : i32233// CHECK-DAG: %[[CST:.*]] = arith.constant dense<3> : vector<4x3xindex>234// CHECK-DAG: %[[CST_1:.*]] = arith.constant dense<true> : vector<4x7x3x2xi1>235// CHECK-DAG: %[[PASSTHRU:.*]] = arith.constant dense<0.000000e+00> : vector<4x7x3x2xf32>236// CHECK: %[[V0:.*]] = vector.transfer_read %[[ARG1]][%[[C0]], %[[C0]]], %[[PV]] {in_bounds = [true, true]} : tensor<4x3xi32>, vector<4x3xi32>237// CHECK: %[[V1:.*]] = vector.transfer_read %[[ARG2]][%[[C0]], %[[C0]]], %[[PV]] {in_bounds = [true, true]} : tensor<4x3xi32>, vector<4x3xi32>238// CHECK: %[[CAST:.*]] = arith.index_cast %[[V0]] : vector<4x3xi32> to vector<4x3xindex>239// CHECK: %[[CAST_1:.*]] = arith.index_cast %[[V1]] : vector<4x3xi32> to vector<4x3xindex>240// CHECK: %[[MULI:.*]] = arith.muli %[[CAST]], %[[CST]] : vector<4x3xindex>241// CHECK: %[[ADDI:.*]] = arith.addi %[[CAST_1]], %[[MULI]] : vector<4x3xindex>242// CHECK: %[[B:.*]] = vector.broadcast %[[ADDI]] : vector<4x3xindex> to vector<7x2x4x3xindex>243// CHECK: %[[T:.*]] = vector.transpose %[[B]], [2, 0, 3, 1] : vector<7x2x4x3xindex> to vector<4x7x3x2xindex>244// CHECK: %[[GATHER:.*]] = vector.gather %[[ARG0]][%[[C0]], %[[C0]]] [%[[T]]], %[[CST_1]], %[[PASSTHRU]] : tensor<3x3xf32>, vector<4x7x3x2xindex>, vector<4x7x3x2xi1>, vector<4x7x3x2xf32> into vector<4x7x3x2xf32>245// CHECK: vector.transfer_write %[[GATHER]], %[[ARG4]][%[[C0]], %[[C0]], %[[C0]], %[[C0]]] {in_bounds = [true, true, true, true]} : vector<4x7x3x2xf32>, tensor<4x7x3x2xf32>246 247// -----248 249#map = affine_map<(d0, d1) -> (d0, d1)>250#map1 = affine_map<(d0, d1, d2) -> (d0 + d1 + d2)>251func.func @vectorize_nd_tensor_extract_load_1d_column_vector_using_gather_load(%arg0: tensor<8x128x768xf32>, %arg1 : index) -> tensor<8x1xf32> {252 %c0 = arith.constant 0 : index253 %0 = tensor.empty() : tensor<8x1xf32>254 %1 = linalg.generic {255 indexing_maps = [#map],256 iterator_types = ["parallel", "parallel"]257 } outs(%0 : tensor<8x1xf32>) {258 ^bb0(%arg5: f32):259 %2 = linalg.index 0 : index260 %3 = linalg.index 1 : index261 %4 = affine.apply #map1(%arg1, %3, %arg1)262 %extracted = tensor.extract %arg0[%2, %c0, %4] : tensor<8x128x768xf32>263 linalg.yield %extracted : f32264 } -> tensor<8x1xf32>265 return %1 : tensor<8x1xf32>266}267 268// CHECK-LABEL: func.func @vectorize_nd_tensor_extract_load_1d_column_vector_using_gather_load269// CHECK-SAME: %[[ARG0:.*]]: tensor<8x128x768xf32>270// CHECK-SAME: %[[ARG1:.*]]: index271// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index272// CHECK-DAG: %[[PASSTHRU:.*]] = arith.constant dense<0.000000e+00> : vector<8x1xf32>273// CHECK-DAG: %[[CST_0:.*]] = arith.constant dense<true> : vector<8x1xi1>274// CHECK-DAG: %[[CST_1:.*]] = arith.constant dense<[0, 98304, 196608, 294912, 393216, 491520, 589824, 688128]> : vector<8xindex>275// CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<8x1xf32>276// CHECK: %[[ADDI_ARG1:.*]] = arith.addi %[[ARG1]], %[[ARG1]] : index277// CHECK: %[[B1:.*]] = vector.broadcast %[[CST_1]] : vector<8xindex> to vector<1x8xindex>278// CHECK: %[[T:.*]] = vector.transpose %[[B1]], [1, 0] : vector<1x8xindex> to vector<8x1xindex>279// CHECK: %[[B2:.*]] = vector.broadcast %[[ADDI_ARG1]] : index to vector<8x1xindex>280// CHECK: %[[ADDI:.*]] = arith.addi %[[B2]], %[[T]] : vector<8x1xindex>281// CHECK: %[[GATHER:.*]] = vector.gather %[[ARG0]][%[[C0]], %[[C0]], %[[C0]]] [%[[ADDI]]], %[[CST_0]], %[[PASSTHRU]] : tensor<8x128x768xf32>, vector<8x1xindex>, vector<8x1xi1>, vector<8x1xf32> into vector<8x1xf32>282// CHECK: vector.transfer_write %[[GATHER]], %[[EMPTY]][%[[C0]], %[[C0]]] {in_bounds = [true, true]} : vector<8x1xf32>, tensor<8x1xf32>283 284// -----285 286// Reading a 1D column vector (hence a candidate for a contiguous load), but given287// %1, it's a gather load.288 289#map = affine_map<(d0, d1) -> (d0, d1)>290func.func @index_from_output_column_vector_gather_load(%src: tensor<8x128xf32>) -> tensor<8x1xf32> {291 %c0 = arith.constant 0 : index292 %0 = tensor.empty() : tensor<8x1xf32>293 %res = linalg.generic {294 indexing_maps = [#map],295 iterator_types = ["parallel", "parallel"]296 } outs(%0 : tensor<8x1xf32>) {297 ^bb0(%arg1: f32):298 %1 = linalg.index 0 : index299 %extracted = tensor.extract %src[%1, %c0] : tensor<8x128xf32>300 linalg.yield %extracted : f32301 } -> tensor<8x1xf32>302 return %res : tensor<8x1xf32>303}304 305// CHECK-LABEL: func.func @index_from_output_column_vector_gather_load(306// CHECK-SAME: %[[SRC:.*]]: tensor<8x128xf32>) -> tensor<8x1xf32> {307// CHECK: %[[IDX_VEC:.*]] = arith.constant dense<[0, 128, 256, 384, 512, 640, 768, 896]> : vector<8xindex>308// CHECK: %[[C0:.*]] = arith.constant 0 : index309// CHECK: %[[PASS_THRU:.*]] = arith.constant dense<0.000000e+00> : vector<8x1xf32>310// CHECK: %[[MASK:.*]] = arith.constant dense<true> : vector<8x1xi1>311// CHECK: %[[OUT:.*]] = tensor.empty() : tensor<8x1xf32>312// CHECK: %[[B:.*]] = vector.broadcast %[[IDX_VEC]] : vector<8xindex> to vector<1x8xindex>313// CHECK: %[[TR:.*]] = vector.transpose %[[B]], [1, 0] : vector<1x8xindex> to vector<8x1xindex>314// CHECK: %[[GATHER:.*]] = vector.gather %[[SRC]]{{\[}}%[[C0]], %[[C0]]] {{\[}}%[[TR]]], %[[MASK]], %[[PASS_THRU]] : tensor<8x128xf32>, vector<8x1xindex>, vector<8x1xi1>, vector<8x1xf32> into vector<8x1xf32>315// CHECK: %[[RES:.*]] = vector.transfer_write %[[GATHER]], %[[OUT]]{{\[}}%[[C0]], %[[C0]]] {in_bounds = [true, true]} : vector<8x1xf32>, tensor<8x1xf32>316// CHECK: return %[[RES]] : tensor<8x1xf32>317 318// -----319 320// Same as above, but the access indices have been swapped and hence this _is_321// a contiguous load. Currently not supported and lowered as vector.gather322// instead.323// TODO: Make sure that this is lowered as a contiguous load.324 325#map = affine_map<(d0, d1) -> (d0, d1)>326func.func @index_from_output_column_vector_contiguous_load(%src: tensor<8x128xf32>) -> tensor<8x1xf32> {327 %c0 = arith.constant 0 : index328 %0 = tensor.empty() : tensor<8x1xf32>329 %res = linalg.generic {330 indexing_maps = [#map],331 iterator_types = ["parallel", "parallel"]332 } outs(%0 : tensor<8x1xf32>) {333 ^bb0(%arg1: f32):334 %1 = linalg.index 0 : index335 %extracted = tensor.extract %src[%c0, %1] : tensor<8x128xf32>336 linalg.yield %extracted : f32337 } -> tensor<8x1xf32>338 return %res : tensor<8x1xf32>339}340 341// CHECK-LABEL: func.func @index_from_output_column_vector_contiguous_load(342// CHECK-SAME: %[[SRC:.*]]: tensor<8x128xf32>) -> tensor<8x1xf32> {343// CHECK: %[[C0:.*]] = arith.constant 0 : index344// CHECK: %[[PASS_THRU:.*]] = arith.constant dense<0.000000e+00> : vector<8x1xf32>345// CHECK: %[[MASK:.*]] = arith.constant dense<true> : vector<8x1xi1>346// CHECK: %[[IDX_VEC:.*]] = arith.constant dense<[0, 1, 2, 3, 4, 5, 6, 7]> : vector<8xindex>347// CHECK: %[[OUT:.*]] = tensor.empty() : tensor<8x1xf32>348// CHECK: %[[B:.*]] = vector.broadcast %[[IDX_VEC]] : vector<8xindex> to vector<1x8xindex>349// CHECK: %[[TR:.*]] = vector.transpose %[[B]], [1, 0] : vector<1x8xindex> to vector<8x1xindex>350// CHECK: %[[GATHER:.*]] = vector.gather %[[SRC]]{{\[}}%[[C0]], %[[C0]]] {{\[}}%[[TR]]], %[[MASK]], %[[PASS_THRU]] : tensor<8x128xf32>, vector<8x1xindex>, vector<8x1xi1>, vector<8x1xf32> into vector<8x1xf32>351// CHECK: %[[RES:.*]] = vector.transfer_write %[[GATHER]], %[[OUT]]{{\[}}%[[C0]], %[[C0]]] {in_bounds = [true, true]} : vector<8x1xf32>, tensor<8x1xf32>352// CHECK: return %[[RES]] : tensor<8x1xf32>353 354// -----355 356#map = affine_map<(d0) -> (d0)>357func.func @vectorize_nd_tensor_extract_contiguous_and_gather(%arg0: tensor<6xf32>, %arg1: tensor<5xi32>) -> tensor<5xf32> {358 %c5 = arith.constant 5 : index359 %c0 = arith.constant 0 : index360 %0 = tensor.empty() : tensor<5xf32>361 %1 = linalg.generic {indexing_maps = [#map], iterator_types = ["parallel"]} outs(%0 : tensor<5xf32>) {362 ^bb0(%out: f32):363 %2 = linalg.index 0 : index364 %extracted = tensor.extract %arg1[%2] : tensor<5xi32>365 %3 = arith.index_cast %extracted : i32 to index366 %4 = arith.maxsi %3, %c0 : index367 %5 = arith.minsi %4, %c5 : index368 %extracted_0 = tensor.extract %arg0[%5] : tensor<6xf32>369 linalg.yield %extracted_0 : f32370 } -> tensor<5xf32>371 return %1 : tensor<5xf32>372}373 374// CHECK-LABEL: func.func @vectorize_nd_tensor_extract_contiguous_and_gather(375// CHECK-SAME: %[[VAL_0:.*]]: tensor<6xf32>376// CHECK-SAME: %[[VAL_1:.*]]: tensor<5xi32>377// CHECK-DAG: %[[VAL_2:.*]] = arith.constant 0 : index378// CHECK-DAG: %[[VAL_3:.*]] = ub.poison : i32379// CHECK-DAG: %[[VAL_4:.*]] = arith.constant dense<0> : vector<5xindex>380// CHECK-DAG: %[[VAL_5:.*]] = arith.constant dense<5> : vector<5xindex>381// CHECK-DAG: %[[VAL_6:.*]] = arith.constant dense<true> : vector<5xi1>382// CHECK-DAG: %[[VAL_7:.*]] = arith.constant dense<0.000000e+00> : vector<5xf32>383// CHECK: %[[VAL_8:.*]] = tensor.empty() : tensor<5xf32>384// CHECK: %[[VAL_9:.*]] = vector.transfer_read %[[VAL_1]]{{\[}}%{{.*}}], %[[VAL_3]] {in_bounds = [true]} : tensor<5xi32>, vector<5xi32>385// CHECK: %[[VAL_10:.*]] = arith.index_cast %[[VAL_9]] : vector<5xi32> to vector<5xindex>386// CHECK: %[[VAL_11:.*]] = arith.maxsi %[[VAL_10]], %[[VAL_4]] : vector<5xindex>387// CHECK: %[[VAL_12:.*]] = arith.minsi %[[VAL_11]], %[[VAL_5]] : vector<5xindex>388// CHECK: %[[VAL_13:.*]] = vector.gather %[[VAL_0]]{{\[}}%[[VAL_2]]] {{\[}}%[[VAL_12]]], %[[VAL_6]], %[[VAL_7]] : tensor<6xf32>, vector<5xindex>, vector<5xi1>, vector<5xf32> into vector<5xf32>389// CHECK: %[[VAL_14:.*]] = vector.transfer_write %[[VAL_13]], %[[VAL_8]]{{\[}}%[[VAL_2]]] {in_bounds = [true]} : vector<5xf32>, tensor<5xf32>390// CHECK: return %[[VAL_14]] : tensor<5xf32>391 392// The vectorizer converts `affine.apply` so that the subsequent Ops can be vectorised based on the converted ops. Gather load.393func.func @vectorize_nd_tensor_extract_with_affine_apply_gather(%6: tensor<80x16xf32>, %arg0: index, %extracted_slice : tensor<1x4xf32>) -> tensor<1x4xf32> {394 %c16 = arith.constant 16 : index395 %1 = linalg.generic {396 indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>],397 iterator_types = ["parallel", "parallel"]398 } outs(%extracted_slice : tensor<1x4xf32>) {399 ^bb0(%out: f32):400 %2 = linalg.index 1 : index401 %3 = affine.apply affine_map<(d0, d1) -> (d0 + d1)>(%2, %arg0)402 %extracted = tensor.extract %6[%3, %c16] : tensor<80x16xf32>403 linalg.yield %extracted : f32404 } -> tensor<1x4xf32>405 return %1 : tensor<1x4xf32>406}407 408// CHECK-LABEL: func.func @vectorize_nd_tensor_extract_with_affine_apply_gather(409// CHECK-SAME: %[[VAL_0:.*]]: tensor<80x16xf32>,410// CHECK-SAME: %[[VAL_1:.*]]: index,411// CHECK-SAME: %[[VAL_2:.*]]: tensor<1x4xf32>) -> tensor<1x4xf32> {412// CHECK-DAG: %[[VAL_3:.*]] = arith.constant dense<[0, 1, 2, 3]> : vector<4xindex>413// CHECK-DAG: %[[VAL_4:.*]] = arith.constant dense<true> : vector<1x4xi1>414// CHECK-DAG: %[[VAL_5:.*]] = arith.constant dense<0.000000e+00> : vector<1x4xf32>415// CHECK-DAG: %[[VAL_6:.*]] = arith.constant 0 : index416// CHECK-DAG: %[[VAL_7:.*]] = arith.constant dense<16> : vector<4xindex>417// CHECK: %[[VAL_8:.*]] = vector.broadcast %[[VAL_1]] : index to vector<4xindex>418// CHECK: %[[VAL_9:.*]] = arith.addi %[[VAL_8]], %[[VAL_3]] : vector<4xindex>419// CHECK: %[[VAL_10:.*]] = arith.muli %[[VAL_9]], %[[VAL_7]] : vector<4xindex>420// CHECK: %[[VAL_11:.*]] = arith.addi %[[VAL_10]], %[[VAL_7]] : vector<4xindex>421// CHECK: %[[VAL_12:.*]] = vector.broadcast %[[VAL_11]] : vector<4xindex> to vector<1x4xindex>422// CHECK: %[[VAL_13:.*]] = vector.gather %[[VAL_0]]{{\[}}%[[VAL_6]], %[[VAL_6]]] {{\[}}%[[VAL_12]]], %[[VAL_4]], %[[VAL_5]] : tensor<80x16xf32>, vector<1x4xindex>, vector<1x4xi1>, vector<1x4xf32> into vector<1x4xf32>423// CHECK: %[[VAL_14:.*]] = vector.transfer_write %[[VAL_13]], %[[VAL_2]]{{\[}}%[[VAL_6]], %[[VAL_6]]] {in_bounds = [true, true]} : vector<1x4xf32>, tensor<1x4xf32>424// CHECK: return %[[VAL_14]] : tensor<1x4xf32>425// CHECK: }426 427// Make sure that non-linear arithmetic operations (e.g. arith.maxsi) are allowed when calculating indices for load operations. Gather load.428func.func @vectorize_nd_tensor_extract_with_maxsi_gather(%arg0: tensor<80x16xf32>, %extracted_slice : tensor<1x4xf32>) -> tensor<1x4xf32> {429 %c79 = arith.constant 79 : index430 %1 = linalg.generic {431 indexing_maps = [affine_map<(d0, d1) -> (d0, d1)>],432 iterator_types = ["parallel", "parallel"]433 } outs(%extracted_slice : tensor<1x4xf32>) {434 ^bb0(%out: f32):435 %2 = linalg.index 1 : index436 %3 = arith.maxsi %2, %c79 : index437 %extracted = tensor.extract %arg0[%3, %2] : tensor<80x16xf32>438 linalg.yield %extracted : f32439 } -> tensor<1x4xf32>440 return %1 : tensor<1x4xf32>441}442 443// CHECK-LABEL: func.func @vectorize_nd_tensor_extract_with_maxsi_gather(444// CHECK-SAME: %[[VAL_0:.*]]: tensor<80x16xf32>,445// CHECK-SAME: %[[VAL_1:.*]]: tensor<1x4xf32>) -> tensor<1x4xf32> {446// CHECK-DAG: %[[VAL_2:.*]] = arith.constant dense<[1264, 1265, 1266, 1267]> : vector<4xindex>447// CHECK-DAG: %[[VAL_4:.*]] = arith.constant dense<true> : vector<1x4xi1>448// CHECK-DAG: %[[VAL_5:.*]] = arith.constant dense<0.000000e+00> : vector<1x4xf32>449// CHECK-DAG: %[[VAL_6:.*]] = arith.constant 0 : index450// CHECK: %[[VAL_7:.*]] = vector.broadcast %[[VAL_2]] : vector<4xindex> to vector<1x4xindex>451// CHECK: %[[VAL_9:.*]] = vector.gather %[[VAL_0]]{{\[}}%[[VAL_6]], %[[VAL_6]]] {{\[}}%[[VAL_7]]], %[[VAL_4]], %[[VAL_5]] : tensor<80x16xf32>, vector<1x4xindex>, vector<1x4xi1>, vector<1x4xf32> into vector<1x4xf32>452// CHECK: %[[VAL_10:.*]] = vector.transfer_write %[[VAL_9]], %[[VAL_1]]{{\[}}%[[VAL_6]], %[[VAL_6]]] {in_bounds = [true, true]} : vector<1x4xf32>, tensor<1x4xf32>453// CHECK: return %[[VAL_10]] : tensor<1x4xf32>454// CHECK: }455 456// -----457 458// The vectorizer assumes it's a gather load whenever using a block argument to calculate an index.459#map = affine_map<(d0) -> (d0)>460func.func @vectorize_nd_tensor_extract_block_arg(%arg0: tensor<5x6xf32>, %arg1: tensor<5xindex>) -> tensor<5xf32> {461 %0 = tensor.empty() : tensor<5xf32>462 %1 = linalg.generic {indexing_maps = [#map, #map], iterator_types = ["parallel"]} ins(%arg1: tensor<5xindex>) outs(%0 : tensor<5xf32>) {463 ^bb0(%in: index, %out: f32):464 %2 = linalg.index 0 : index465 %extracted_0 = tensor.extract %arg0[%in, %2] : tensor<5x6xf32>466 linalg.yield %extracted_0 : f32467 } -> tensor<5xf32>468 return %1 : tensor<5xf32>469}470 471// CHECK-LABEL: func.func @vectorize_nd_tensor_extract_block_arg(472// CHECK-SAME: %[[VAL_0:.*]]: tensor<5x6xf32>,473// CHECK-SAME: %[[VAL_1:.*]]: tensor<5xindex>) -> tensor<5xf32> {474// CHECK-DAG: %[[PAD:.*]] = ub.poison : index475// CHECK-DAG: %[[VAL_2:.*]] = arith.constant 0 : index476// CHECK-DAG: %[[VAL_3:.*]] = arith.constant dense<[0, 1, 2, 3, 4]> : vector<5xindex>477// CHECK-DAG: %[[VAL_4:.*]] = arith.constant dense<true> : vector<5xi1>478// CHECK-DAG: %[[VAL_5:.*]] = arith.constant dense<0.000000e+00> : vector<5xf32>479// CHECK-DAG: %[[VAL_6:.*]] = arith.constant dense<6> : vector<5xindex>480// CHECK: %[[VAL_7:.*]] = tensor.empty() : tensor<5xf32>481// CHECK: %[[VAL_8:.*]] = vector.transfer_read %[[VAL_1]]{{\[}}%[[VAL_2]]], %[[PAD]] {in_bounds = [true]} : tensor<5xindex>, vector<5xindex>482// CHECK: %[[VAL_9:.*]] = arith.muli %[[VAL_8]], %[[VAL_6]] : vector<5xindex>483// CHECK: %[[VAL_10:.*]] = arith.addi %[[VAL_9]], %[[VAL_3]] : vector<5xindex>484// CHECK: %[[VAL_11:.*]] = vector.gather %[[VAL_0]]{{\[}}%[[VAL_2]], %[[VAL_2]]] {{\[}}%[[VAL_10]]], %[[VAL_4]], %[[VAL_5]] : tensor<5x6xf32>, vector<5xindex>, vector<5xi1>, vector<5xf32> into vector<5xf32>485// CHECK: %[[VAL_12:.*]] = vector.transfer_write %[[VAL_11]], %[[VAL_7]]{{\[}}%[[VAL_2]]] {in_bounds = [true]} : vector<5xf32>, tensor<5xf32>486// CHECK: return %[[VAL_12]] : tensor<5xf32>487// CHECK: }488 489// -----490 491#map = affine_map<(d0, d1, d2) -> (d0, d1, d2)>492#map1 = affine_map<(d0, d1, d2) -> (d0 + d1 + d2)>493func.func @vectorize_reverse_like_tensor_extract(%arg0: tensor<1x2x3xf32>, %arg1: tensor<1x1x3xf32>, %arg2: index) -> tensor<1x1x3xf32> {494 %c1 = arith.constant 1 : index495 %c0 = arith.constant 0 : index496 %c2 = arith.constant 2 : index497 %0 = linalg.generic {indexing_maps = [#map], iterator_types = ["parallel", "parallel", "parallel"]} outs(%arg1 : tensor<1x1x3xf32>) {498 ^bb0(%out: f32):499 %1 = linalg.index 1 : index500 %2 = linalg.index 0 : index501 %3 = affine.apply #map1(%1, %2, %arg2)502 %4 = linalg.index 2 : index503 %5 = arith.subi %c2, %4 : index504 %extracted = tensor.extract %arg0[%c0, %3, %5] : tensor<1x2x3xf32>505 linalg.yield %extracted : f32506 } -> tensor<1x1x3xf32>507 return %0 : tensor<1x1x3xf32>508}509// CHECK-LABEL: func.func @vectorize_reverse_like_tensor_extract510// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]*]]511// CHECK-SAME: %[[ARG1:[0-9a-zA-Z]*]]512// CHECK-SAME: %[[ARG2:[0-9a-zA-Z]*]]513// CHECK-DAG: %[[C3:.+]] = arith.constant 3 : index514// CHECK-DAG: %[[C0:.+]] = arith.constant 0 : index515// CHECK-DAG: %[[MASK:.*]] = arith.constant dense<true> : vector<1x1x3xi1>516// CHECK-DAG: %[[PASSTHRU:.*]] = arith.constant dense<0.000000e+00> : vector<1x1x3xf32>517// CHECK-DAG: %[[INIT_IDX:.+]] = arith.constant dense<[2, 1, 0]> : vector<3xindex>518// CHECK: %[[T0:.+]] = arith.muli %[[ARG2]], %[[C3]] : index519// CHECK: %[[T1:.+]] = vector.broadcast %[[T0]] : index to vector<1x1x3xindex>520// CHECK: %[[T2:.+]] = vector.broadcast %[[INIT_IDX]]521// CHECK: %[[T3:.+]] = arith.addi %[[T2]], %[[T1]]522// CHECK: %[[GATHER:.*]] = vector.gather %[[ARG0]][%[[C0]], %[[C0]], %[[C0]]] [%[[T3]]], %[[MASK]], %[[PASSTHRU]]523// CHECK: vector.transfer_write %[[GATHER]]524 525//===----------------------------------------------------------------------===//526// Scalar load + broadcast527//===----------------------------------------------------------------------===//528 529// -----530 531#map = affine_map<(d0, d1, d2) -> (d0, d1, d2)>532func.func @vectorize_nd_tensor_extract_scalar_broadcast(%src: tensor<3x3xf32>, %init: tensor<1x1x3xf32>) -> tensor<1x1x3xf32> {533 %c0 = arith.constant 1 : index534 %c1 = arith.constant 2 : index535 536 %res = linalg.generic {537 indexing_maps = [#map],538 iterator_types = ["parallel", "parallel", "parallel"]539 } outs(%init : tensor<1x1x3xf32>) {540 ^bb0(%arg4: f32):541 %1 = tensor.extract %src[%c0, %c1] : tensor<3x3xf32>542 linalg.yield %1 : f32543 } -> tensor<1x1x3xf32>544 545 return %res : tensor<1x1x3xf32>546}547 548// CHECK-LABEL: func.func @vectorize_nd_tensor_extract_scalar_broadcast(549// CHECK-SAME: %[[SRC:.*]]: tensor<3x3xf32>,550// CHECK-SAME: %[[INIT:.*]]: tensor<1x1x3xf32>) -> tensor<1x1x3xf32> {551// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index552// CHECK-DAG: %[[C1:.*]] = arith.constant 1 : index553// CHECK-DAG: %[[C2:.*]] = arith.constant 2 : index554// CHECK-DAG: %[[PAD:.*]] = ub.poison : f32555// CHECK: %[[READ:.*]] = vector.transfer_read %[[SRC]][%[[C1]], %[[C2]]], %[[PAD]] : tensor<3x3xf32>, vector<f32>556// CHECK: %[[READ_BCAST:.*]] = vector.broadcast %[[READ]] : vector<f32> to vector<1x1x3xf32>557// CHECK: vector.transfer_write %[[READ_BCAST]], %[[INIT]][%[[C0]], %[[C0]], %[[C0]]] {in_bounds = [true, true, true]} : vector<1x1x3xf32>, tensor<1x1x3xf32>558 559// -----560 561#map = affine_map<() -> ()>562func.func @extract_scalar_from_0d_into_0d(%src: tensor<f32>, %init: tensor<f32>) -> tensor<f32> {563 %res = linalg.generic {564 indexing_maps = [#map],565 iterator_types = []566 } outs(%init : tensor<f32>) {567 ^bb0(%in: f32):568 %1 = tensor.extract %src[] : tensor<f32>569 linalg.yield %1 : f32570 } -> tensor<f32>571 572 return %res : tensor<f32>573}574 575// CHECK-LABEL: func.func @extract_scalar_from_0d_into_0d(576// CHECK-SAME: %[[SRC:.*]]: tensor<f32>,577// CHECK-SAME: %[[INIT:.*]]: tensor<f32>) -> tensor<f32> {578// CHECK: %[[PAD:.*]] = ub.poison : f32579// CHECK: %[[READ:.*]] = vector.transfer_read %[[SRC]][], %[[PAD]] : tensor<f32>, vector<f32>580// CHECK: vector.transfer_write %[[READ]], %[[INIT]][] : vector<f32>, tensor<f32>581 582// -----583 584#map = affine_map<(n) -> (n)>585func.func @extract_scalar_from_0d_into_1d(%src: tensor<f32>, %init: tensor<1xf32>) -> tensor<1xf32> {586 %res = linalg.generic {587 indexing_maps = [#map],588 iterator_types = ["parallel"]589 } outs(%init : tensor<1xf32>) {590 ^bb0(%in: f32):591 %1 = tensor.extract %src[] : tensor<f32>592 linalg.yield %1 : f32593 } -> tensor<1xf32>594 595 return %res : tensor<1xf32>596}597// CHECK-LABEL: func.func @extract_scalar_from_0d_into_1d(598// CHECK-SAME: %[[SRC:.*]]: tensor<f32>,599// CHECK-SAME: %[[INIT:.*]]: tensor<1xf32>) -> tensor<1xf32> {600// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index601// CHECK-DAG: %[[PAD:.*]] = ub.poison : f32602// CHECK: %[[READ:.*]] = vector.transfer_read %[[SRC]][], %[[PAD]] : tensor<f32>, vector<f32>603// CHECK: %[[READ_BCAST:.*]] = vector.broadcast %[[READ]] : vector<f32> to vector<1xf32>604// CHECK: vector.transfer_write %[[READ_BCAST]], %[[INIT]][%[[C0]]] {in_bounds = [true]} : vector<1xf32>, tensor<1xf32>605 606// -----607 608#map1 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>609func.func @vectorize_0d_tensor_extract(%src: tensor<f32>, %init: tensor<1x1x3xf32>) -> tensor<1x1x3xf32> {610 %res = linalg.generic {611 indexing_maps = [#map1],612 iterator_types = ["parallel", "parallel", "parallel"]613 } outs(%init : tensor<1x1x3xf32>) {614 ^bb0(%arg4: f32):615 %1 = tensor.extract %src[] : tensor<f32>616 linalg.yield %1 : f32617 } -> tensor<1x1x3xf32>618 return %res : tensor<1x1x3xf32>619}620 621// CHECK-LABEL: func.func @vectorize_0d_tensor_extract(622// CHECK-SAME: %[[SRC:.*]]: tensor<f32>623// CHECK: %[[READ:.*]] = vector.transfer_read %[[SRC]][], %{{.+}} : tensor<f32>624// CHECK: vector.broadcast %[[READ]] : vector<f32> to vector<1x1x3xf32>625 626// -----627 628func.func @scalar_read_with_broadcast_from_column_tensor(%init: tensor<1x1x4xi32>) -> tensor<1x1x4xi32> {629 %c4 = arith.constant 4 : index630 %c0 = arith.constant 0 : index631 %src = arith.constant dense<[[0], [1], [2], [3], [4], [5], [6], [7], [8], [9], [10], [11], [12], [13], [14]]> : tensor<15x1xi32>632 633 %res = linalg.generic {634 indexing_maps = [affine_map<(d0, d1, d2) -> (d0, d1, d2)>],635 iterator_types = ["parallel", "parallel", "parallel"]636 } outs(%init : tensor<1x1x4xi32>) {637 638 ^bb0(%out: i32):639 %idx = linalg.index 0 : index640 %extracted = tensor.extract %src[%idx, %c0] : tensor<15x1xi32>641 linalg.yield %extracted : i32642 } -> tensor<1x1x4xi32>643 644 return %res : tensor<1x1x4xi32>645}646 647// CHECK-LABEL: func.func @scalar_read_with_broadcast_from_column_tensor648// CHECK-SAME: %[[INIT:.*]]: tensor<1x1x4xi32>) -> tensor<1x1x4xi32> {649// CHECK-DAG: %[[PAD:.*]] = ub.poison : i32650// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index651// CHECK-DAG: %[[SRC:.*]] = arith.constant dense<{{\[\[}}0], [1], [2], [3], [4], [5], [6], [7], [8], [9], [10], [11], [12], [13], [14]]> : tensor<15x1xi32>652// CHECK: %[[READ:.*]] = vector.transfer_read %[[SRC]]{{\[}}%[[C0]], %[[C0]]], %[[PAD]] : tensor<15x1xi32>, vector<i32>653// CHECK: %[[READ_BCAST:.*]] = vector.broadcast %[[READ]] : vector<i32> to vector<1x1x4xi32>654// CHECK: %[[RES:.*]] = vector.transfer_write %[[READ_BCAST]], %[[INIT]][%[[C0]], %[[C0]], %[[C0]]] {in_bounds = [true, true, true]} : vector<1x1x4xi32>, tensor<1x1x4xi32>655 656// -----657 658// TODO: Currently this fails to vectorise when the indices are non-constant.659 660#map = affine_map<(d0, d1, d2) -> (d0, d1, d2)>661func.func @vectorize_nd_tensor_extract_transfer_read_basic_column(662 %src: tensor<3x3x3xf32>,663 %init: tensor<3x1x1xf32>) -> tensor<3x1x1xf32> {664 665 %c0 = arith.constant 0 : index666 667 %res = linalg.generic {668 indexing_maps = [#map],669 iterator_types = ["parallel", "parallel", "parallel"]670 } outs(%init : tensor<3x1x1xf32>) {671 ^bb0(%out: f32):672 %1 = tensor.extract %src[%c0, %c0, %c0] : tensor<3x3x3xf32>673 linalg.yield %1 : f32674 } -> tensor<3x1x1xf32>675 676 return %res : tensor<3x1x1xf32>677}678 679// CHECK-LABEL: func.func @vectorize_nd_tensor_extract_transfer_read_basic_column(680// CHECK-SAME: %[[SRC:.*]]: tensor<3x3x3xf32>,681// CHECK-SAME: %[[INIT:.*]]: tensor<3x1x1xf32>)682// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index683// CHECK-DAG: %[[PV:.*]] = ub.poison : f32684// CHECK: %[[READ:.*]] = vector.transfer_read %[[SRC]][%[[C0]], %[[C0]], %[[C0]]], %[[PV]] : tensor<3x3x3xf32>, vector<f32>685// CHECK: %[[READ_BCAST:.*]] = vector.broadcast %[[READ]] : vector<f32> to vector<3x1x1xf32>686// CHECK: vector.transfer_write %[[READ_BCAST]], %[[INIT]]{{\[}}%[[C0]], %[[C0]], %[[C0]]] {in_bounds = [true, true, true]} : vector<3x1x1xf32>, tensor<3x1x1xf32>687