371 lines · plain
1// RUN: mlir-opt %s -convert-vector-to-llvm -split-input-file | FileCheck '-D$IDX_TYPE=i32' %s2// RUN: mlir-opt %s --convert-vector-to-llvm='force-32bit-vector-indices=0' | FileCheck '-D$IDX_TYPE=i64' %s3 4func.func @transfer_read_write_1d(%A : memref<?xf32>, %base: index) -> vector<17xf32> {5 %f7 = arith.constant 7.0: f326 %f = vector.transfer_read %A[%base], %f77 {permutation_map = affine_map<(d0) -> (d0)>} :8 memref<?xf32>, vector<17xf32>9 vector.transfer_write %f, %A[%base]10 {permutation_map = affine_map<(d0) -> (d0)>} :11 vector<17xf32>, memref<?xf32>12 return %f: vector<17xf32>13}14// CHECK-LABEL: func @transfer_read_write_1d15// CHECK-SAME: %[[MEM:.*]]: memref<?xf32>,16// CHECK-SAME: %[[BASE:.*]]: index) -> vector<17xf32>17// 1. Create pass-through vector.18// CHECK-DAG: %[[PASS_THROUGH:.*]] = arith.constant dense<7.000000e+00> : vector<17xf32>19//20// 2. Create a vector with linear indices [ 0 .. vector_length - 1 ].21// CHECK-DAG: %[[linearIndex:.*]] = arith.constant dense22// CHECK-SAME: <[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]> : vector<17x[[$IDX_TYPE]]>23//24// 3. Let dim be the memref dimension, compute the in-bound index (dim - offset)25// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index26// CHECK: %[[DIM:.*]] = memref.dim %[[MEM]], %[[C0]] : memref<?xf32>27// CHECK: %[[BOUND:.*]] = arith.subi %[[DIM]], %[[BASE]] : index28//29// 4. Create bound vector to compute in-bound mask:30// [ 0 .. vector_length - 1 ] < [ dim - offset .. dim - offset ]31// CHECK: %[[btrunc:.*]] = arith.index_cast %[[BOUND]] :32// CMP32-SAME: index to i3233// CMP64-SAME: index to i6434// CHECK: %[[boundVecInsert:.*]] = llvm.insertelement %[[btrunc]]35// CHECK: %[[boundVect:.*]] = llvm.shufflevector %[[boundVecInsert]]36// CHECK: %[[mask:.*]] = arith.cmpi sgt, %[[boundVect]], %[[linearIndex]] : vector<17x[[$IDX_TYPE]]>37// CMP64-SAME: : vector<17xi64>38//39// 5. Bitcast to vector form.40// CHECK: %[[gep:.*]] = llvm.getelementptr %{{.*}} :41// CHECK-SAME: (!llvm.ptr, i64) -> !llvm.ptr, f3242//43// 6. Rewrite as a masked read.44// CHECK: %[[loaded:.*]] = llvm.intr.masked.load %[[gep]], %[[mask]],45// CHECK-SAME: %[[PASS_THROUGH]] {alignment = 4 : i32} :46// CHECK-SAME: -> vector<17xf32>47//48// 1. Let dim be the memref dimension, compute the in-bound index (dim - offset)49// CHECK: %[[DIM_b:.*]] = memref.dim %[[MEM]], %[[C0]] : memref<?xf32>50// CHECK: %[[BOUND_b:.*]] = arith.subi %[[DIM_b]], %[[BASE]] : index51//52// 2. Create bound vector to compute in-bound mask:53// [ 0 .. vector_length - 1 ] < [ dim - offset .. dim - offset ]54// CHECK: %[[btrunc_b:.*]] = arith.index_cast %[[BOUND_b]]55// CMP32-SAME: index to i3256// CHECK: %[[boundVecInsert_b:.*]] = llvm.insertelement %[[btrunc_b]]57// CHECK: %[[boundVect_b:.*]] = llvm.shufflevector %[[boundVecInsert_b]]58// CHECK: %[[mask_b:.*]] = arith.cmpi sgt, %[[boundVect_b]],59// CHECK-SAME: %[[linearIndex]] : vector<17x[[$IDX_TYPE]]>60//61// 3. Bitcast to vector form.62// CHECK: %[[gep_b:.*]] = llvm.getelementptr {{.*}} :63// CHECK-SAME: (!llvm.ptr, i64) -> !llvm.ptr, f3264//65// 4. Rewrite as a masked write.66// CHECK: llvm.intr.masked.store %[[loaded]], %[[gep_b]], %[[mask_b]]67// CHECK-SAME: {alignment = 4 : i32} :68// CHECK-SAME: vector<17xf32>, vector<17xi1> into !llvm.ptr69 70func.func @transfer_read_write_1d_scalable(%A : memref<?xf32>, %base: index) -> vector<[17]xf32> {71 %f7 = arith.constant 7.0: f3272 %f = vector.transfer_read %A[%base], %f773 {permutation_map = affine_map<(d0) -> (d0)>} :74 memref<?xf32>, vector<[17]xf32>75 vector.transfer_write %f, %A[%base]76 {permutation_map = affine_map<(d0) -> (d0)>} :77 vector<[17]xf32>, memref<?xf32>78 return %f: vector<[17]xf32>79}80// CHECK-LABEL: func @transfer_read_write_1d_scalable81// CHECK-SAME: %[[MEM:.*]]: memref<?xf32>,82// CHECK-SAME: %[[BASE:.*]]: index) -> vector<[17]xf32>83// 1. Create pass-through vector.84// CHECK-DAG: %[[PASS_THROUGH:.*]] = arith.constant dense<7.000000e+00> : vector<[17]xf32>85//86// 2. Let dim be the memref dimension, compute the in-bound index (dim - offset)87// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : index88// CHECK: %[[DIM:.*]] = memref.dim %[[MEM]], %[[C0]] : memref<?xf32>89// CHECK: %[[BOUND:.*]] = arith.subi %[[DIM]], %[[BASE]] : index90//91// 3. Create a vector with linear indices [ 0 .. vector_length - 1 ].92// CHECK: %[[linearIndex:.*]] = llvm.intr.stepvector : vector<[17]x[[$IDX_TYPE]]>93//94// 4. Create bound vector to compute in-bound mask:95// [ 0 .. vector_length - 1 ] < [ dim - offset .. dim - offset ]96// CHECK: %[[btrunc:.*]] = arith.index_cast %[[BOUND]] : index to [[$IDX_TYPE]]97// CHECK: %[[boundVecInsert:.*]] = llvm.insertelement %[[btrunc]]98// CHECK: %[[boundVect:.*]] = llvm.shufflevector %[[boundVecInsert]]99// CHECK: %[[mask:.*]] = arith.cmpi slt, %[[linearIndex]], %[[boundVect]]100// CHECK-SAME: : vector<[17]x[[$IDX_TYPE]]>101//102// 5. Bitcast to vector form.103// CHECK: %[[gep:.*]] = llvm.getelementptr %{{.*}} :104// CHECK-SAME: (!llvm.ptr, i64) -> !llvm.ptr, f32105//106// 6. Rewrite as a masked read.107// CHECK: %[[loaded:.*]] = llvm.intr.masked.load %[[gep]], %[[mask]],108// CHECK-SAME: %[[PASS_THROUGH]] {alignment = 4 : i32} :109// CHECK-SAME: -> vector<[17]xf32>110//111// 1. Let dim be the memref dimension, compute the in-bound index (dim - offset)112// CHECK: %[[DIM_b:.*]] = memref.dim %[[MEM]], %[[C0]] : memref<?xf32>113// CHECK: %[[BOUND_b:.*]] = arith.subi %[[DIM_b]], %[[BASE]] : index114//115// 2. Create a vector with linear indices [ 0 .. vector_length - 1 ].116// CHECK: %[[linearIndex_b:.*]] = llvm.intr.stepvector : vector<[17]x[[$IDX_TYPE]]>117//118// 3. Create bound vector to compute in-bound mask:119// [ 0 .. vector_length - 1 ] < [ dim - offset .. dim - offset ]120// CHECK: %[[btrunc_b:.*]] = arith.index_cast %[[BOUND_b]] : index to [[$IDX_TYPE]]121// CHECK: %[[boundVecInsert_b:.*]] = llvm.insertelement %[[btrunc_b]]122// CHECK: %[[boundVect_b:.*]] = llvm.shufflevector %[[boundVecInsert_b]]123// CHECK: %[[mask_b:.*]] = arith.cmpi slt, %[[linearIndex_b]],124// CHECK-SAME: %[[boundVect_b]] : vector<[17]x[[$IDX_TYPE]]>125//126// 4. Bitcast to vector form.127// CHECK: %[[gep_b:.*]] = llvm.getelementptr {{.*}} :128// CHECK-SAME: (!llvm.ptr, i64) -> !llvm.ptr, f32129//130// 5. Rewrite as a masked write.131// CHECK: llvm.intr.masked.store %[[loaded]], %[[gep_b]], %[[mask_b]]132// CHECK-SAME: {alignment = 4 : i32} :133// CHECK-SAME: vector<[17]xf32>, vector<[17]xi1> into !llvm.ptr134 135// -----136 137func.func @transfer_read_write_index_1d(%A : memref<?xindex>, %base: index) -> vector<17xindex> {138 %f7 = arith.constant 7: index139 %f = vector.transfer_read %A[%base], %f7140 {permutation_map = affine_map<(d0) -> (d0)>} :141 memref<?xindex>, vector<17xindex>142 vector.transfer_write %f, %A[%base]143 {permutation_map = affine_map<(d0) -> (d0)>} :144 vector<17xindex>, memref<?xindex>145 return %f: vector<17xindex>146}147// CHECK-LABEL: func @transfer_read_write_index_1d148// CHECK-SAME: %[[BASE:[a-zA-Z0-9]*]]: index) -> vector<17xindex>149// CHECK: %[[SPLAT:.*]] = arith.constant dense<7> : vector<17xindex>150// CHECK: %{{.*}} = builtin.unrealized_conversion_cast %[[SPLAT]] : vector<17xindex> to vector<17xi64>151 152// CHECK: %[[loaded:.*]] = llvm.intr.masked.load %{{.*}}, %{{.*}}, %{{.*}} {alignment = 8 : i32} :153// CHECK-SAME: (!llvm.ptr, vector<17xi1>, vector<17xi64>) -> vector<17xi64>154 155// CHECK: llvm.intr.masked.store %[[loaded]], %{{.*}}, %{{.*}} {alignment = 8 : i32} :156// CHECK-SAME: vector<17xi64>, vector<17xi1> into !llvm.ptr157 158func.func @transfer_read_write_index_1d_scalable(%A : memref<?xindex>, %base: index) -> vector<[17]xindex> {159 %f7 = arith.constant 7: index160 %f = vector.transfer_read %A[%base], %f7161 {permutation_map = affine_map<(d0) -> (d0)>} :162 memref<?xindex>, vector<[17]xindex>163 vector.transfer_write %f, %A[%base]164 {permutation_map = affine_map<(d0) -> (d0)>} :165 vector<[17]xindex>, memref<?xindex>166 return %f: vector<[17]xindex>167}168// CHECK-LABEL: func @transfer_read_write_index_1d169// CHECK-SAME: %[[BASE:[a-zA-Z0-9]*]]: index) -> vector<[17]xindex>170// CHECK: %[[SPLAT:.*]] = arith.constant dense<7> : vector<[17]xindex>171// CHECK: %{{.*}} = builtin.unrealized_conversion_cast %[[SPLAT]] : vector<[17]xindex> to vector<[17]xi64>172 173// CHECK: %[[loaded:.*]] = llvm.intr.masked.load %{{.*}}, %{{.*}}, %{{.*}} {alignment = 8 : i32} :174// CHECK-SAME: (!llvm.ptr, vector<[17]xi1>, vector<[17]xi64>) -> vector<[17]xi64>175 176// CHECK: llvm.intr.masked.store %[[loaded]], %{{.*}}, %{{.*}} {alignment = 8 : i32} :177// CHECK-SAME: vector<[17]xi64>, vector<[17]xi1> into !llvm.ptr178 179// -----180 181func.func @transfer_read_2d_to_1d(%A : memref<?x?xf32>, %base0: index, %base1: index) -> vector<17xf32> {182 %f7 = arith.constant 7.0: f32183 %f = vector.transfer_read %A[%base0, %base1], %f7184 {permutation_map = affine_map<(d0, d1) -> (d1)>} :185 memref<?x?xf32>, vector<17xf32>186 return %f: vector<17xf32>187}188// CHECK-LABEL: func @transfer_read_2d_to_1d189// CHECK-SAME: %[[BASE_0:[a-zA-Z0-9]*]]: index, %[[BASE_1:[a-zA-Z0-9]*]]: index) -> vector<17xf32>190//191// Create a vector with linear indices [ 0 .. vector_length - 1 ].192// CHECK-DAG: %[[linearIndex:.*]] = arith.constant dense<[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16]> :193// CHECK-SAME: vector<17x[[$IDX_TYPE]]>194//195// CHECK-DAG: %[[c1:.*]] = arith.constant 1 : index196// CHECK: %[[DIM:.*]] = memref.dim %{{.*}}, %[[c1]] : memref<?x?xf32>197//198// Compute the in-bound index (dim - offset)199// CHECK: %[[BOUND:.*]] = arith.subi %[[DIM]], %[[BASE_1]] : index200//201// Create bound vector to compute in-bound mask:202// [ 0 .. vector_length - 1 ] < [ dim - offset .. dim - offset ]203// CHECK: %[[btrunc:.*]] = arith.index_cast %[[BOUND]] : index to [[$IDX_TYPE]]204// CHECK: %[[boundVecInsert:.*]] = llvm.insertelement %[[btrunc]]205// CHECK: %[[boundVect:.*]] = llvm.shufflevector %[[boundVecInsert]]206// CHECK: %[[mask:.*]] = arith.cmpi sgt, %[[boundVect]], %[[linearIndex]]207 208func.func @transfer_read_2d_to_1d_scalable(%A : memref<?x?xf32>, %base0: index, %base1: index) -> vector<[17]xf32> {209 %f7 = arith.constant 7.0: f32210 %f = vector.transfer_read %A[%base0, %base1], %f7211 {permutation_map = affine_map<(d0, d1) -> (d1)>} :212 memref<?x?xf32>, vector<[17]xf32>213 return %f: vector<[17]xf32>214}215// CHECK-LABEL: func @transfer_read_2d_to_1d216// CHECK-SAME: %[[BASE_0:[a-zA-Z0-9]*]]: index, %[[BASE_1:[a-zA-Z0-9]*]]: index) -> vector<[17]xf32>217// CHECK: %[[c1:.*]] = arith.constant 1 : index218// CHECK: %[[DIM:.*]] = memref.dim %{{.*}}, %[[c1]] : memref<?x?xf32>219//220// Compute the in-bound index (dim - offset)221// CHECK: %[[BOUND:.*]] = arith.subi %[[DIM]], %[[BASE_1]] : index222//223// Create a vector with linear indices [ 0 .. vector_length - 1 ].224// CHECK: %[[linearIndex:.*]] = llvm.intr.stepvector : vector<[17]x[[$IDX_TYPE]]>225//226// Create bound vector to compute in-bound mask:227// [ 0 .. vector_length - 1 ] < [ dim - offset .. dim - offset ]228// CHECK: %[[btrunc:.*]] = arith.index_cast %[[BOUND]] : index to [[$IDX_TYPE]]229// CHECK: %[[boundVecInsert:.*]] = llvm.insertelement %[[btrunc]]230// CHECK: %[[boundVect:.*]] = llvm.shufflevector %[[boundVecInsert]]231// CHECK: %[[mask:.*]] = arith.cmpi slt, %[[linearIndex]], %[[boundVect]]232 233// -----234 235func.func @transfer_read_write_1d_non_zero_addrspace(%A : memref<?xf32, 3>, %base: index) -> vector<17xf32> {236 %f7 = arith.constant 7.0: f32237 %f = vector.transfer_read %A[%base], %f7238 {permutation_map = affine_map<(d0) -> (d0)>} :239 memref<?xf32, 3>, vector<17xf32>240 vector.transfer_write %f, %A[%base]241 {permutation_map = affine_map<(d0) -> (d0)>} :242 vector<17xf32>, memref<?xf32, 3>243 return %f: vector<17xf32>244}245// CHECK-LABEL: func @transfer_read_write_1d_non_zero_addrspace246// CHECK-SAME: %[[BASE:[a-zA-Z0-9]*]]: index) -> vector<17xf32>247//248// CHECK: %[[c0:.*]] = arith.constant 0 : index249//250// 1. Check address space for GEP is correct.251// CHECK: %[[gep:.*]] = llvm.getelementptr {{.*}} :252// CHECK-SAME: (!llvm.ptr<3>, i64) -> !llvm.ptr<3>, f32253//254// 2. Check address space of the memref is correct.255// CHECK: %[[DIM:.*]] = memref.dim %{{.*}}, %[[c0]] : memref<?xf32, 3>256//257// 3. Check address space for GEP is correct.258// CHECK: %[[gep_b:.*]] = llvm.getelementptr {{.*}} :259// CHECK-SAME: (!llvm.ptr<3>, i64) -> !llvm.ptr<3>, f32260 261func.func @transfer_read_write_1d_non_zero_addrspace_scalable(%A : memref<?xf32, 3>, %base: index) -> vector<[17]xf32> {262 %f7 = arith.constant 7.0: f32263 %f = vector.transfer_read %A[%base], %f7264 {permutation_map = affine_map<(d0) -> (d0)>} :265 memref<?xf32, 3>, vector<[17]xf32>266 vector.transfer_write %f, %A[%base]267 {permutation_map = affine_map<(d0) -> (d0)>} :268 vector<[17]xf32>, memref<?xf32, 3>269 return %f: vector<[17]xf32>270}271// CHECK-LABEL: func @transfer_read_write_1d_non_zero_addrspace_scalable272// CHECK-SAME: %[[BASE:[a-zA-Z0-9]*]]: index) -> vector<[17]xf32>273//274// CHECK: %[[c0:.*]] = arith.constant 0 : index275//276// 1. Check address space for GEP is correct.277// CHECK: %[[gep:.*]] = llvm.getelementptr {{.*}} :278// CHECK-SAME: (!llvm.ptr<3>, i64) -> !llvm.ptr<3>, f32279//280// 2. Check address space of the memref is correct.281// CHECK: %[[DIM:.*]] = memref.dim %{{.*}}, %[[c0]] : memref<?xf32, 3>282//283// 3. Check address space for GEP is correct.284// CHECK: %[[gep_b:.*]] = llvm.getelementptr {{.*}} :285// CHECK-SAME: (!llvm.ptr<3>, i64) -> !llvm.ptr<3>, f32286 287// -----288 289func.func @transfer_read_1d_inbounds(%A : memref<?xf32>, %base: index) -> vector<17xf32> {290 %f7 = arith.constant 7.0: f32291 %f = vector.transfer_read %A[%base], %f7 {in_bounds = [true]} :292 memref<?xf32>, vector<17xf32>293 return %f: vector<17xf32>294}295// CHECK-LABEL: func @transfer_read_1d_inbounds296// CHECK-SAME: %[[BASE:[a-zA-Z0-9]*]]: index) -> vector<17xf32>297//298// 1. Bitcast to vector form.299// CHECK: %[[gep:.*]] = llvm.getelementptr {{.*}} :300// CHECK-SAME: (!llvm.ptr, i64) -> !llvm.ptr, f32301//302// 2. Rewrite as a load.303// CHECK: %[[loaded:.*]] = llvm.load %[[gep]] {alignment = 4 : i64} : !llvm.ptr -> vector<17xf32>304 305func.func @transfer_read_1d_inbounds_scalable(%A : memref<?xf32>, %base: index) -> vector<[17]xf32> {306 %f7 = arith.constant 7.0: f32307 %f = vector.transfer_read %A[%base], %f7 {in_bounds = [true]} :308 memref<?xf32>, vector<[17]xf32>309 return %f: vector<[17]xf32>310}311// CHECK-LABEL: func @transfer_read_1d_inbounds_scalable312// CHECK-SAME: %[[BASE:[a-zA-Z0-9]*]]: index) -> vector<[17]xf32>313//314// 1. Bitcast to vector form.315// CHECK: %[[gep:.*]] = llvm.getelementptr {{.*}} :316// CHECK-SAME: (!llvm.ptr, i64) -> !llvm.ptr, f32317//318// 2. Rewrite as a load.319// CHECK: %[[loaded:.*]] = llvm.load %[[gep]] {alignment = 4 : i64} : !llvm.ptr -> vector<[17]xf32>320 321// -----322 323// CHECK-LABEL: func @transfer_read_write_1d_mask324// CHECK: %[[mask1:.*]] = arith.constant dense<[false, false, true, false, true]>325// CHECK: %[[cmpi:.*]] = arith.cmpi sgt326// CHECK: %[[mask2:.*]] = arith.andi %[[cmpi]], %[[mask1]]327// CHECK: %[[r:.*]] = llvm.intr.masked.load %{{.*}}, %[[mask2]]328// CHECK: %[[cmpi_1:.*]] = arith.cmpi sgt329// CHECK: %[[mask3:.*]] = arith.andi %[[cmpi_1]], %[[mask1]]330// CHECK: llvm.intr.masked.store %[[r]], %{{.*}}, %[[mask3]]331// CHECK: return %[[r]]332func.func @transfer_read_write_1d_mask(%A : memref<?xf32>, %base : index) -> vector<5xf32> {333 %m = arith.constant dense<[0, 0, 1, 0, 1]> : vector<5xi1>334 %f7 = arith.constant 7.0: f32335 %f = vector.transfer_read %A[%base], %f7, %m : memref<?xf32>, vector<5xf32>336 vector.transfer_write %f, %A[%base], %m : vector<5xf32>, memref<?xf32>337 return %f: vector<5xf32>338}339 340// CHECK-LABEL: func @transfer_read_write_1d_mask_scalable341// CHECK-SAME: %[[mask:[a-zA-Z0-9]*]]: vector<[5]xi1>342// CHECK: %[[cmpi:.*]] = arith.cmpi slt343// CHECK: %[[mask1:.*]] = arith.andi %[[cmpi]], %[[mask]]344// CHECK: %[[r:.*]] = llvm.intr.masked.load %{{.*}}, %[[mask1]]345// CHECK: %[[cmpi_1:.*]] = arith.cmpi slt346// CHECK: %[[mask2:.*]] = arith.andi %[[cmpi_1]], %[[mask]]347// CHECK: llvm.intr.masked.store %[[r]], %{{.*}}, %[[mask2]]348// CHECK: return %[[r]]349func.func @transfer_read_write_1d_mask_scalable(%A : memref<?xf32>, %base : index, %m : vector<[5]xi1>) -> vector<[5]xf32> {350 %f7 = arith.constant 7.0: f32351 %f = vector.transfer_read %A[%base], %f7, %m : memref<?xf32>, vector<[5]xf32>352 vector.transfer_write %f, %A[%base], %m : vector<[5]xf32>, memref<?xf32>353 return %f: vector<[5]xf32>354}355 356// -----357 358// Can't lower xfer_read/xfer_write on tensors, but this shouldn't crash359 360// CHECK-LABEL: func @transfer_read_write_tensor361// CHECK: vector.transfer_read362// CHECK: vector.transfer_write363func.func @transfer_read_write_tensor(%A: tensor<?xf32>, %base : index) -> vector<4xf32> {364 %f7 = arith.constant 7.0: f32365 %c0 = arith.constant 0: index366 %f = vector.transfer_read %A[%base], %f7 : tensor<?xf32>, vector<4xf32>367 %w = vector.transfer_write %f, %A[%c0] : vector<4xf32>, tensor<?xf32>368 "test.some_use"(%w) : (tensor<?xf32>) -> ()369 return %f : vector<4xf32>370}371