790 lines · plain
1// RUN: mlir-opt %s -affine-super-vectorize="virtual-vector-size=128 test-fastest-varying=0 vectorize-reductions=true" -split-input-file | FileCheck %s2 3// The inner reduction loop '%j' is vectorized.4 5func.func @vecdim_reduction(%in: memref<256x512xf32>, %out: memref<256xf32>) {6 %cst = arith.constant 0.000000e+00 : f327 affine.for %i = 0 to 256 {8 %final_red = affine.for %j = 0 to 512 iter_args(%red_iter = %cst) -> (f32) {9 %ld = affine.load %in[%i, %j] : memref<256x512xf32>10 %add = arith.addf %red_iter, %ld : f3211 affine.yield %add : f3212 }13 affine.store %final_red, %out[%i] : memref<256xf32>14 }15 return16}17 18// CHECK-LABEL: @vecdim_reduction19// CHECK: affine.for %{{.*}} = 0 to 256 {20// CHECK: %[[vzero:.*]] = arith.constant dense<0.000000e+00> : vector<128xf32>21// CHECK: %[[vred:.*]] = affine.for %{{.*}} = 0 to 512 step 128 iter_args(%[[red_iter:.*]] = %[[vzero]]) -> (vector<128xf32>) {22// CHECK: %[[ld:.*]] = vector.transfer_read %{{.*}} : memref<256x512xf32>, vector<128xf32>23// CHECK: %[[add:.*]] = arith.addf %[[red_iter]], %[[ld]] : vector<128xf32>24// CHECK: affine.yield %[[add]] : vector<128xf32>25// CHECK: }26// CHECK: %[[final_sum:.*]] = vector.reduction <add>, %[[vred:.*]] : vector<128xf32> into f3227// CHECK: affine.store %[[final_sum]], %{{.*}} : memref<256xf32>28// CHECK: }29 30// -----31 32func.func @vecdim_reduction_minf(%in: memref<256x512xf32>, %out: memref<256xf32>) {33 %cst = arith.constant 0x7F800000 : f3234 affine.for %i = 0 to 256 {35 %final_red = affine.for %j = 0 to 512 iter_args(%red_iter = %cst) -> (f32) {36 %ld = affine.load %in[%i, %j] : memref<256x512xf32>37 %min = arith.minimumf %red_iter, %ld : f3238 affine.yield %min : f3239 }40 affine.store %final_red, %out[%i] : memref<256xf32>41 }42 return43}44 45// CHECK-LABEL: @vecdim_reduction_minf46// CHECK: affine.for %{{.*}} = 0 to 256 {47// CHECK: %[[vmax:.*]] = arith.constant dense<0x7F800000> : vector<128xf32>48// CHECK: %[[vred:.*]] = affine.for %{{.*}} = 0 to 512 step 128 iter_args(%[[red_iter:.*]] = %[[vmax]]) -> (vector<128xf32>) {49// CHECK: %[[ld:.*]] = vector.transfer_read %{{.*}} : memref<256x512xf32>, vector<128xf32>50// CHECK: %[[min:.*]] = arith.minimumf %[[red_iter]], %[[ld]] : vector<128xf32>51// CHECK: affine.yield %[[min]] : vector<128xf32>52// CHECK: }53// CHECK: %[[final_min:.*]] = vector.reduction <minimumf>, %[[vred:.*]] : vector<128xf32> into f3254// CHECK: affine.store %[[final_min]], %{{.*}} : memref<256xf32>55// CHECK: }56 57// -----58 59func.func @vecdim_reduction_maxf(%in: memref<256x512xf32>, %out: memref<256xf32>) {60 %cst = arith.constant 0xFF800000 : f3261 affine.for %i = 0 to 256 {62 %final_red = affine.for %j = 0 to 512 iter_args(%red_iter = %cst) -> (f32) {63 %ld = affine.load %in[%i, %j] : memref<256x512xf32>64 %max = arith.maximumf %red_iter, %ld : f3265 affine.yield %max : f3266 }67 affine.store %final_red, %out[%i] : memref<256xf32>68 }69 return70}71 72// CHECK-LABEL: @vecdim_reduction_maxf73// CHECK: affine.for %{{.*}} = 0 to 256 {74// CHECK: %[[vmin:.*]] = arith.constant dense<0xFF800000> : vector<128xf32>75// CHECK: %[[vred:.*]] = affine.for %{{.*}} = 0 to 512 step 128 iter_args(%[[red_iter:.*]] = %[[vmin]]) -> (vector<128xf32>) {76// CHECK: %[[ld:.*]] = vector.transfer_read %{{.*}} : memref<256x512xf32>, vector<128xf32>77// CHECK: %[[max:.*]] = arith.maximumf %[[red_iter]], %[[ld]] : vector<128xf32>78// CHECK: affine.yield %[[max]] : vector<128xf32>79// CHECK: }80// CHECK: %[[final_max:.*]] = vector.reduction <maximumf>, %[[vred:.*]] : vector<128xf32> into f3281// CHECK: affine.store %[[final_max]], %{{.*}} : memref<256xf32>82// CHECK: }83 84// -----85 86func.func @vecdim_reduction_minsi(%in: memref<256x512xi32>, %out: memref<256xi32>) {87 %cst = arith.constant 2147483647 : i3288 affine.for %i = 0 to 256 {89 %final_red = affine.for %j = 0 to 512 iter_args(%red_iter = %cst) -> (i32) {90 %ld = affine.load %in[%i, %j] : memref<256x512xi32>91 %min = arith.minsi %red_iter, %ld : i3292 affine.yield %min : i3293 }94 affine.store %final_red, %out[%i] : memref<256xi32>95 }96 return97}98 99// CHECK-LABEL: @vecdim_reduction_minsi100// CHECK: affine.for %{{.*}} = 0 to 256 {101// CHECK: %[[vmax:.*]] = arith.constant dense<2147483647> : vector<128xi32>102// CHECK: %[[vred:.*]] = affine.for %{{.*}} = 0 to 512 step 128 iter_args(%[[red_iter:.*]] = %[[vmax]]) -> (vector<128xi32>) {103// CHECK: %[[ld:.*]] = vector.transfer_read %{{.*}} : memref<256x512xi32>, vector<128xi32>104// CHECK: %[[min:.*]] = arith.minsi %[[red_iter]], %[[ld]] : vector<128xi32>105// CHECK: affine.yield %[[min]] : vector<128xi32>106// CHECK: }107// CHECK: %[[final_min:.*]] = vector.reduction <minsi>, %[[vred:.*]] : vector<128xi32> into i32108// CHECK: affine.store %[[final_min]], %{{.*}} : memref<256xi32>109// CHECK: }110 111// -----112 113func.func @vecdim_reduction_maxsi(%in: memref<256x512xi32>, %out: memref<256xi32>) {114 %cst = arith.constant -2147483648 : i32115 affine.for %i = 0 to 256 {116 %final_red = affine.for %j = 0 to 512 iter_args(%red_iter = %cst) -> (i32) {117 %ld = affine.load %in[%i, %j] : memref<256x512xi32>118 %max = arith.maxsi %red_iter, %ld : i32119 affine.yield %max : i32120 }121 affine.store %final_red, %out[%i] : memref<256xi32>122 }123 return124}125 126// CHECK-LABEL: @vecdim_reduction_maxsi127// CHECK: affine.for %{{.*}} = 0 to 256 {128// CHECK: %[[vmin:.*]] = arith.constant dense<-2147483648> : vector<128xi32>129// CHECK: %[[vred:.*]] = affine.for %{{.*}} = 0 to 512 step 128 iter_args(%[[red_iter:.*]] = %[[vmin]]) -> (vector<128xi32>) {130// CHECK: %[[ld:.*]] = vector.transfer_read %{{.*}} : memref<256x512xi32>, vector<128xi32>131// CHECK: %[[max:.*]] = arith.maxsi %[[red_iter]], %[[ld]] : vector<128xi32>132// CHECK: affine.yield %[[max]] : vector<128xi32>133// CHECK: }134// CHECK: %[[final_max:.*]] = vector.reduction <maxsi>, %[[vred:.*]] : vector<128xi32> into i32135// CHECK: affine.store %[[final_max]], %{{.*}} : memref<256xi32>136// CHECK: }137 138// -----139 140func.func @vecdim_reduction_minui(%in: memref<256x512xi32>, %out: memref<256xi32>) {141 %cst = arith.constant -1 : i32142 affine.for %i = 0 to 256 {143 %final_red = affine.for %j = 0 to 512 iter_args(%red_iter = %cst) -> (i32) {144 %ld = affine.load %in[%i, %j] : memref<256x512xi32>145 %min = arith.minui %red_iter, %ld : i32146 affine.yield %min : i32147 }148 affine.store %final_red, %out[%i] : memref<256xi32>149 }150 return151}152 153// CHECK-LABEL: @vecdim_reduction_minui154// CHECK: affine.for %{{.*}} = 0 to 256 {155// CHECK: %[[vmax:.*]] = arith.constant dense<-1> : vector<128xi32>156// CHECK: %[[vred:.*]] = affine.for %{{.*}} = 0 to 512 step 128 iter_args(%[[red_iter:.*]] = %[[vmax]]) -> (vector<128xi32>) {157// CHECK: %[[ld:.*]] = vector.transfer_read %{{.*}} : memref<256x512xi32>, vector<128xi32>158// CHECK: %[[min:.*]] = arith.minui %[[red_iter]], %[[ld]] : vector<128xi32>159// CHECK: affine.yield %[[min]] : vector<128xi32>160// CHECK: }161// CHECK: %[[final_min:.*]] = vector.reduction <minui>, %[[vred:.*]] : vector<128xi32> into i32162// CHECK: affine.store %[[final_min]], %{{.*}} : memref<256xi32>163// CHECK: }164 165// -----166 167func.func @vecdim_reduction_maxui(%in: memref<256x512xi32>, %out: memref<256xi32>) {168 %cst = arith.constant 0 : i32169 affine.for %i = 0 to 256 {170 %final_red = affine.for %j = 0 to 512 iter_args(%red_iter = %cst) -> (i32) {171 %ld = affine.load %in[%i, %j] : memref<256x512xi32>172 %max = arith.maxui %red_iter, %ld : i32173 affine.yield %max : i32174 }175 affine.store %final_red, %out[%i] : memref<256xi32>176 }177 return178}179 180// CHECK-LABEL: @vecdim_reduction_maxui181// CHECK: affine.for %{{.*}} = 0 to 256 {182// CHECK: %[[vmin:.*]] = arith.constant dense<0> : vector<128xi32>183// CHECK: %[[vred:.*]] = affine.for %{{.*}} = 0 to 512 step 128 iter_args(%[[red_iter:.*]] = %[[vmin]]) -> (vector<128xi32>) {184// CHECK: %[[ld:.*]] = vector.transfer_read %{{.*}} : memref<256x512xi32>, vector<128xi32>185// CHECK: %[[max:.*]] = arith.maxui %[[red_iter]], %[[ld]] : vector<128xi32>186// CHECK: affine.yield %[[max]] : vector<128xi32>187// CHECK: }188// CHECK: %[[final_max:.*]] = vector.reduction <maxui>, %[[vred:.*]] : vector<128xi32> into i32189// CHECK: affine.store %[[final_max]], %{{.*}} : memref<256xi32>190// CHECK: }191 192// -----193 194func.func @vecdim_reduction_andi(%in: memref<256x512xi32>, %out: memref<256xi32>) {195 %cst = arith.constant -1 : i32196 affine.for %i = 0 to 256 {197 %final_red = affine.for %j = 0 to 512 iter_args(%red_iter = %cst) -> (i32) {198 %ld = affine.load %in[%i, %j] : memref<256x512xi32>199 %or = arith.andi %red_iter, %ld : i32200 affine.yield %or : i32201 }202 affine.store %final_red, %out[%i] : memref<256xi32>203 }204 return205}206 207// CHECK-LABEL: @vecdim_reduction_andi208// CHECK: affine.for %{{.*}} = 0 to 256 {209// CHECK: %[[vallone:.*]] = arith.constant dense<-1> : vector<128xi32>210// CHECK: %[[vred:.*]] = affine.for %{{.*}} = 0 to 512 step 128 iter_args(%[[red_iter:.*]] = %[[vallone]]) -> (vector<128xi32>) {211// CHECK: %[[ld:.*]] = vector.transfer_read %{{.*}} : memref<256x512xi32>, vector<128xi32>212// CHECK: %[[and:.*]] = arith.andi %[[red_iter]], %[[ld]] : vector<128xi32>213// CHECK: affine.yield %[[and]] : vector<128xi32>214// CHECK: }215// CHECK: %[[final_red:.*]] = vector.reduction <and>, %[[vred:.*]] : vector<128xi32> into i32216// CHECK: affine.store %[[final_red]], %{{.*}} : memref<256xi32>217// CHECK: }218 219// -----220 221func.func @vecdim_reduction_ori(%in: memref<256x512xi32>, %out: memref<256xi32>) {222 %cst = arith.constant 0 : i32223 affine.for %i = 0 to 256 {224 %final_red = affine.for %j = 0 to 512 iter_args(%red_iter = %cst) -> (i32) {225 %ld = affine.load %in[%i, %j] : memref<256x512xi32>226 %or = arith.ori %red_iter, %ld : i32227 affine.yield %or : i32228 }229 affine.store %final_red, %out[%i] : memref<256xi32>230 }231 return232}233 234// CHECK-LABEL: @vecdim_reduction_ori235// CHECK: affine.for %{{.*}} = 0 to 256 {236// CHECK: %[[vzero:.*]] = arith.constant dense<0> : vector<128xi32>237// CHECK: %[[vred:.*]] = affine.for %{{.*}} = 0 to 512 step 128 iter_args(%[[red_iter:.*]] = %[[vzero]]) -> (vector<128xi32>) {238// CHECK: %[[ld:.*]] = vector.transfer_read %{{.*}} : memref<256x512xi32>, vector<128xi32>239// CHECK: %[[or:.*]] = arith.ori %[[red_iter]], %[[ld]] : vector<128xi32>240// CHECK: affine.yield %[[or]] : vector<128xi32>241// CHECK: }242// CHECK: %[[final_red:.*]] = vector.reduction <or>, %[[vred:.*]] : vector<128xi32> into i32243// CHECK: affine.store %[[final_red]], %{{.*}} : memref<256xi32>244// CHECK: }245 246// -----247 248func.func @vecdim_reduction_xori(%in: memref<256x512xi32>, %out: memref<256xi32>) {249 %cst = arith.constant 0 : i32250 affine.for %i = 0 to 256 {251 %final_red = affine.for %j = 0 to 512 iter_args(%red_iter = %cst) -> (i32) {252 %ld = affine.load %in[%i, %j] : memref<256x512xi32>253 %xor = arith.xori %red_iter, %ld : i32254 affine.yield %xor : i32255 }256 affine.store %final_red, %out[%i] : memref<256xi32>257 }258 return259}260 261// CHECK-LABEL: func.func @vecdim_reduction_xori(262// CHECK-SAME: %[[input:.*]]: memref<256x512xi32>,263// CHECK-SAME: %[[output:.*]]: memref<256xi32>) {264// CHECK: %[[cst:.*]] = arith.constant 0 : i32265// CHECK: affine.for %{{.*}} = 0 to 256 {266// CHECK: %[[vzero:.*]] = arith.constant dense<0> : vector<128xi32>267// CHECK: %[[vred:.*]] = affine.for %{{.*}} = 0 to 512 step 128 iter_args(%[[red_iter:.*]] = %[[vzero]]) -> (vector<128xi32>) {268// CHECK: %[[poison:.*]] = ub.poison : i32269// CHECK: %[[ld:.*]] = vector.transfer_read %[[input]]{{\[}}%{{.*}}, %{{.*}}], %[[poison]] : memref<256x512xi32>, vector<128xi32>270// CHECK: %[[xor:.*]] = arith.xori %[[red_iter]], %[[ld]] : vector<128xi32>271// CHECK: affine.yield %[[xor]] : vector<128xi32>272// CHECK: }273// CHECK: %[[final_red:.*]] = vector.reduction <xor>, %[[vred]] : vector<128xi32> into i32274// CHECK: affine.store %[[final_red]], %[[output]]{{\[}}%{{.*}}] : memref<256xi32>275// CHECK: }276// CHECK: return277// CHECK: }278 279// -----280 281func.func @vecdim_reduction_minnumf(%in: memref<256x512xf32>, %out: memref<256xf32>) {282 %cst = arith.constant 0xFF800000 : f32283 affine.for %i = 0 to 256 {284 %final_red = affine.for %j = 0 to 512 iter_args(%red_iter = %cst) -> (f32) {285 %ld = affine.load %in[%i, %j] : memref<256x512xf32>286 %min = arith.minnumf %red_iter, %ld : f32287 affine.yield %min : f32288 }289 affine.store %final_red, %out[%i] : memref<256xf32>290 }291 return292}293 294// CHECK-LABEL: func.func @vecdim_reduction_minnumf(295// CHECK-SAME: %[[input:.*]]: memref<256x512xf32>,296// CHECK-SAME: %[[output:.*]]: memref<256xf32>) {297// CHECK: %[[cst:.*]] = arith.constant 0xFF800000 : f32298// CHECK: affine.for %{{.*}} = 0 to 256 {299// CHECK: %[[vzero:.*]] = arith.constant dense<0x7FC00000> : vector<128xf32>300// CHECK: %[[vred:.*]] = affine.for %{{.*}} = 0 to 512 step 128 iter_args(%[[red_iter:.*]] = %[[vzero]]) -> (vector<128xf32>) {301// CHECK: %[[poison:.*]] = ub.poison : f32302// CHECK: %[[ld:.*]] = vector.transfer_read %[[input]]{{\[}}%{{.*}}, %{{.*}}], %[[poison]] : memref<256x512xf32>, vector<128xf32>303// CHECK: %[[min:.*]] = arith.minnumf %[[red_iter]], %[[ld]] : vector<128xf32>304// CHECK: affine.yield %[[min]] : vector<128xf32>305// CHECK: }306// CHECK: %[[red_scalar:.*]] = vector.reduction <minnumf>, %[[vred]] : vector<128xf32> into f32307// CHECK: %[[final_red:.*]] = arith.minnumf %[[red_scalar]], %[[cst]] : f32308// CHECK: affine.store %[[final_red]], %[[output]]{{\[}}%{{.*}}] : memref<256xf32>309// CHECK: }310// CHECK: return311// CHECK: }312 313// -----314 315func.func @vecdim_reduction_maxnumf(%in: memref<256x512xf32>, %out: memref<256xf32>) {316 %cst = arith.constant 0xFF800000 : f32317 affine.for %i = 0 to 256 {318 %final_red = affine.for %j = 0 to 512 iter_args(%red_iter = %cst) -> (f32) {319 %ld = affine.load %in[%i, %j] : memref<256x512xf32>320 %max = arith.maxnumf %red_iter, %ld : f32321 affine.yield %max : f32322 }323 affine.store %final_red, %out[%i] : memref<256xf32>324 }325 return326}327 328// CHECK-LABEL: func.func @vecdim_reduction_maxnumf(329// CHECK-SAME: %[[input:.*]]: memref<256x512xf32>,330// CHECK-SAME: %[[output:.*]]: memref<256xf32>) {331// CHECK: %[[cst:.*]] = arith.constant 0xFF800000 : f32332// CHECK: affine.for %{{.*}} = 0 to 256 {333// CHECK: %[[vzero:.*]] = arith.constant dense<0xFFC00000> : vector<128xf32>334// CHECK: %[[vred:.*]] = affine.for %{{.*}} = 0 to 512 step 128 iter_args(%[[red_iter:.*]] = %[[vzero]]) -> (vector<128xf32>) {335// CHECK: %[[poison:.*]] = ub.poison : f32336// CHECK: %[[ld:.*]] = vector.transfer_read %[[input]]{{\[}}%{{.*}}, %{{.*}}], %[[poison]] : memref<256x512xf32>, vector<128xf32>337// CHECK: %[[max:.*]] = arith.maxnumf %[[red_iter]], %[[ld]] : vector<128xf32>338// CHECK: affine.yield %[[max]] : vector<128xf32>339// CHECK: }340// CHECK: %[[red_scalar:.*]] = vector.reduction <maxnumf>, %[[vred]] : vector<128xf32> into f32341// CHECK: %[[final_red:.*]] = arith.maxnumf %[[red_scalar]], %[[cst]] : f32342// CHECK: affine.store %[[final_red]], %[[output]]{{\[}}%{{.*}}] : memref<256xf32>343// CHECK: }344// CHECK: return345// CHECK: }346 347// -----348 349// The inner reduction loop '%j' is vectorized. (The order of addf's operands is350// different than in the previous test case).351 352func.func @vecdim_reduction_comm(%in: memref<256x512xf32>, %out: memref<256xf32>) {353 %cst = arith.constant 0.000000e+00 : f32354 affine.for %i = 0 to 256 {355 %final_red = affine.for %j = 0 to 512 iter_args(%red_iter = %cst) -> (f32) {356 %ld = affine.load %in[%i, %j] : memref<256x512xf32>357 %add = arith.addf %ld, %red_iter : f32358 affine.yield %add : f32359 }360 affine.store %final_red, %out[%i] : memref<256xf32>361 }362 return363}364 365// CHECK-LABEL: @vecdim_reduction_comm366// CHECK: affine.for %{{.*}} = 0 to 256 {367// CHECK: %[[vzero:.*]] = arith.constant dense<0.000000e+00> : vector<128xf32>368// CHECK: %[[vred:.*]] = affine.for %{{.*}} = 0 to 512 step 128 iter_args(%[[red_iter:.*]] = %[[vzero]]) -> (vector<128xf32>) {369// CHECK: %[[ld:.*]] = vector.transfer_read %{{.*}} : memref<256x512xf32>, vector<128xf32>370// CHECK: %[[add:.*]] = arith.addf %[[ld]], %[[red_iter]] : vector<128xf32>371// CHECK: affine.yield %[[add]] : vector<128xf32>372// CHECK: }373// CHECK: %[[final_sum:.*]] = vector.reduction <add>, %[[vred:.*]] : vector<128xf32> into f32374// CHECK: affine.store %[[final_sum]], %{{.*}} : memref<256xf32>375// CHECK: }376 377// -----378 379// The inner reduction loop '%j' is vectorized. Transforming the input before380// performing the accumulation doesn't cause any problem.381 382func.func @vecdim_reduction_expsin(%in: memref<256x512xf32>, %out: memref<256xf32>) {383 %cst = arith.constant 0.000000e+00 : f32384 affine.for %i = 0 to 256 {385 %final_red = affine.for %j = 0 to 512 iter_args(%red_iter = %cst) -> (f32) {386 %ld = affine.load %in[%i, %j] : memref<256x512xf32>387 %sin = math.sin %ld : f32388 %exp = math.exp %sin : f32389 %add = arith.addf %red_iter, %exp : f32390 affine.yield %add : f32391 }392 affine.store %final_red, %out[%i] : memref<256xf32>393 }394 return395}396 397// CHECK-LABEL: @vecdim_reduction_expsin398// CHECK: affine.for %{{.*}} = 0 to 256 {399// CHECK: %[[vzero:.*]] = arith.constant dense<0.000000e+00> : vector<128xf32>400// CHECK: %[[vred:.*]] = affine.for %{{.*}} = 0 to 512 step 128 iter_args(%[[red_iter:.*]] = %[[vzero]]) -> (vector<128xf32>) {401// CHECK: %[[ld:.*]] = vector.transfer_read %{{.*}} : memref<256x512xf32>, vector<128xf32>402// CHECK: %[[sin:.*]] = math.sin %[[ld]]403// CHECK: %[[exp:.*]] = math.exp %[[sin]]404// CHECK: %[[add:.*]] = arith.addf %[[red_iter]], %[[exp]] : vector<128xf32>405// CHECK: affine.yield %[[add]] : vector<128xf32>406// CHECK: }407// CHECK: %[[final_sum:.*]] = vector.reduction <add>, %[[vred:.*]] : vector<128xf32> into f32408// CHECK: affine.store %[[final_sum]], %{{.*}} : memref<256xf32>409// CHECK: }410 411// -----412 413// Two reductions at the same time. The inner reduction loop '%j' is vectorized.414 415func.func @two_vecdim_reductions(%in: memref<256x512xf32>, %out_sum: memref<256xf32>, %out_prod: memref<256xf32>) {416 %cst = arith.constant 1.000000e+00 : f32417 affine.for %i = 0 to 256 {418 // Note that we pass the same constant '1.0' as initial values for both419 // reductions.420 %sum, %prod = affine.for %j = 0 to 512 iter_args(%part_sum = %cst, %part_prod = %cst) -> (f32, f32) {421 %ld = affine.load %in[%i, %j] : memref<256x512xf32>422 %add = arith.addf %part_sum, %ld : f32423 %mul = arith.mulf %part_prod, %ld : f32424 affine.yield %add, %mul : f32, f32425 }426 affine.store %sum, %out_sum[%i] : memref<256xf32>427 affine.store %prod, %out_prod[%i] : memref<256xf32>428 }429 return430}431 432// CHECK-LABEL: @two_vecdim_reductions433// CHECK: %[[cst:.*]] = arith.constant 1.000000e+00 : f32434// CHECK: affine.for %{{.*}} = 0 to 256 {435// CHECK: %[[vzero:.*]] = arith.constant dense<0.000000e+00> : vector<128xf32>436// CHECK: %[[vone:.*]] = arith.constant dense<1.000000e+00> : vector<128xf32>437// CHECK: %[[vred:.*]]:2 = affine.for %{{.*}} = 0 to 512 step 128 iter_args(%[[part_sum:.*]] = %[[vzero]], %[[part_prod:.*]] = %[[vone]]) -> (vector<128xf32>, vector<128xf32>) {438// CHECK: %[[ld:.*]] = vector.transfer_read %{{.*}} : memref<256x512xf32>, vector<128xf32>439// CHECK: %[[add:.*]] = arith.addf %[[part_sum]], %[[ld]] : vector<128xf32>440// CHECK: %[[mul:.*]] = arith.mulf %[[part_prod]], %[[ld]] : vector<128xf32>441// CHECK: affine.yield %[[add]], %[[mul]] : vector<128xf32>, vector<128xf32>442// CHECK: }443// CHECK: %[[nonfinal_sum:.*]] = vector.reduction <add>, %[[vred:.*]]#0 : vector<128xf32> into f32444// Note that to compute the final sum we need to add the original initial value445// (%cst) since it is not zero.446// CHECK: %[[final_sum:.*]] = arith.addf %[[nonfinal_sum]], %[[cst]] : f32447// For the final product we don't need to do this additional step because the448// initial value equals to 1 (the neutral element for multiplication).449// CHECK: %[[final_prod:.*]] = vector.reduction <mul>, %[[vred:.*]]#1 : vector<128xf32> into f32450// CHECK: affine.store %[[final_sum]], %{{.*}} : memref<256xf32>451// CHECK: affine.store %[[final_prod]], %{{.*}} : memref<256xf32>452// CHECK: }453 454// -----455 456// The integer case.457 458func.func @two_vecdim_reductions_int(%in: memref<256x512xi64>, %out_sum: memref<256xi64>, %out_prod: memref<256xi64>) {459 %cst0 = arith.constant 0 : i64460 %cst1 = arith.constant 1 : i64461 affine.for %i = 0 to 256 {462 %sum, %prod = affine.for %j = 0 to 512 iter_args(%part_sum = %cst0, %part_prod = %cst1) -> (i64, i64) {463 %ld = affine.load %in[%i, %j] : memref<256x512xi64>464 %add = arith.addi %part_sum, %ld : i64465 %mul = arith.muli %part_prod, %ld : i64466 affine.yield %add, %mul : i64, i64467 }468 affine.store %sum, %out_sum[%i] : memref<256xi64>469 affine.store %prod, %out_prod[%i] : memref<256xi64>470 }471 return472}473 474// CHECK-LABEL: @two_vecdim_reductions475// CHECK: affine.for %{{.*}} = 0 to 256 {476// CHECK: %[[vzero:.*]] = arith.constant dense<0> : vector<128xi64>477// CHECK: %[[vone:.*]] = arith.constant dense<1> : vector<128xi64>478// CHECK: %[[vred:.*]]:2 = affine.for %{{.*}} = 0 to 512 step 128 iter_args(%[[part_sum:.*]] = %[[vzero]], %[[part_prod:.*]] = %[[vone]]) -> (vector<128xi64>, vector<128xi64>) {479// CHECK: %[[ld:.*]] = vector.transfer_read %{{.*}} : memref<256x512xi64>, vector<128xi64>480// CHECK: %[[add:.*]] = arith.addi %[[part_sum]], %[[ld]] : vector<128xi64>481// CHECK: %[[mul:.*]] = arith.muli %[[part_prod]], %[[ld]] : vector<128xi64>482// CHECK: affine.yield %[[add]], %[[mul]] : vector<128xi64>, vector<128xi64>483// CHECK: }484// CHECK: %[[final_sum:.*]] = vector.reduction <add>, %[[vred:.*]]#0 : vector<128xi64> into i64485// CHECK: %[[final_prod:.*]] = vector.reduction <mul>, %[[vred:.*]]#1 : vector<128xi64> into i64486// CHECK: affine.store %[[final_sum]], %{{.*}} : memref<256xi64>487// CHECK: affine.store %[[final_prod]], %{{.*}} : memref<256xi64>488// CHECK: }489 490// -----491 492// The outer reduction loop '%j' is vectorized.493 494func.func @vecdim_reduction_nested(%in: memref<256x512xf32>, %out: memref<1xf32>) {495 %cst = arith.constant 0.000000e+00 : f32496 %outer_red = affine.for %j = 0 to 512 iter_args(%outer_iter = %cst) -> (f32) {497 %inner_red = affine.for %i = 0 to 256 iter_args(%inner_iter = %cst) -> (f32) {498 %ld = affine.load %in[%i, %j] : memref<256x512xf32>499 %add = arith.addf %inner_iter, %ld : f32500 affine.yield %add : f32501 }502 %outer_add = arith.addf %outer_iter, %inner_red : f32503 affine.yield %outer_add : f32504 }505 affine.store %outer_red, %out[0] : memref<1xf32>506 return507}508 509// CHECK-LABEL: @vecdim_reduction_nested510// CHECK: %[[vzero:.*]] = arith.constant dense<0.000000e+00> : vector<128xf32>511// CHECK: %[[outer_red:.*]] = affine.for %{{.*}} = 0 to 512 step 128 iter_args(%[[outer_iter:.*]] = %[[vzero]]) -> (vector<128xf32>) {512// CHECK: %[[vzero:.*]] = arith.constant dense<0.000000e+00> : vector<128xf32>513// CHECK: %[[inner_red:.*]] = affine.for %{{.*}} = 0 to 256 iter_args(%[[inner_iter:.*]] = %[[vzero]]) -> (vector<128xf32>) {514// CHECK: %[[ld:.*]] = vector.transfer_read %{{.*}} : memref<256x512xf32>, vector<128xf32>515// CHECK: %[[add:.*]] = arith.addf %[[inner_iter]], %[[ld]] : vector<128xf32>516// CHECK: affine.yield %[[add]] : vector<128xf32>517// CHECK: }518// CHECK: %[[outer_add:.*]] = arith.addf %[[outer_iter]], %[[inner_red]] : vector<128xf32>519// CHECK: affine.yield %[[outer_add]] : vector<128xf32>520// CHECK: }521// CHECK: %[[final_sum:.*]] = vector.reduction <add>, %[[outer_red:.*]] : vector<128xf32> into f32522// CHECK: affine.store %[[final_sum]], %{{.*}} : memref<1xf32>523 524// -----525 526// The inner reduction loop '%j' computes partial sums as a side effect and527// is not vectorized.528 529func.func @vecdim_partial_sums_1_rejected(%in: memref<256x512xf32>, %out_sum: memref<256xf32>, %out_prod: memref<256xf32>, %out_partsum: memref<256x512xf32>) {530 %cst = arith.constant 1.000000e+00 : f32531 affine.for %i = 0 to 256 {532 %sum, %prod = affine.for %j = 0 to 512 iter_args(%part_sum = %cst, %part_prod = %cst) -> (f32, f32) {533 %ld = affine.load %in[%i, %j] : memref<256x512xf32>534 %add = arith.addf %part_sum, %ld : f32535 %mul = arith.mulf %part_prod, %ld : f32536 affine.store %add, %out_partsum[%i, %j] : memref<256x512xf32>537 affine.yield %add, %mul : f32, f32538 }539 affine.store %sum, %out_sum[%i] : memref<256xf32>540 affine.store %prod, %out_prod[%i] : memref<256xf32>541 }542 return543}544 545// CHECK-LABEL: @vecdim_partial_sums_1_rejected546// CHECK-NOT: vector547 548// -----549 550// The inner reduction loop '%j' computes partial sums as a side effect and551// is not vectorized.552 553func.func @vecdim_partial_sums_2_rejected(%in: memref<256x512xf32>, %out_sum: memref<256xf32>, %out_prod: memref<256xf32>, %out_partsum: memref<256x512xf32>) {554 %cst = arith.constant 1.000000e+00 : f32555 affine.for %i = 0 to 256 {556 %sum, %prod = affine.for %j = 0 to 512 iter_args(%part_sum = %cst, %part_prod = %cst) -> (f32, f32) {557 affine.store %part_sum, %out_partsum[%i, %j] : memref<256x512xf32>558 %ld = affine.load %in[%i, %j] : memref<256x512xf32>559 %add = arith.addf %part_sum, %ld : f32560 %mul = arith.mulf %part_prod, %ld : f32561 affine.yield %add, %mul : f32, f32562 }563 affine.store %sum, %out_sum[%i] : memref<256xf32>564 affine.store %prod, %out_prod[%i] : memref<256xf32>565 }566 return567}568 569// CHECK-LABEL: @vecdim_partial_sums_2_rejected570// CHECK-NOT: vector571 572// -----573 574// The inner reduction loop '%j' performs an unknown reduction operation and is575// not vectorized.576 577func.func @vecdim_unknown_reduction_rejected(%in: memref<256x512xf32>, %out: memref<256xf32>) {578 %cst = arith.constant 1.000000e+00 : f32579 %final_red = affine.for %j = 0 to 512 iter_args(%red_iter = %cst) -> (f32) {580 %add = arith.addf %red_iter, %red_iter : f32581 affine.yield %add : f32582 }583 affine.store %final_red, %out[0] : memref<256xf32>584 return585}586 587// CHECK-LABEL: @vecdim_unknown_reduction_rejected588// CHECK-NOT: vector589 590// -----591 592// The inner reduction loop '%j' doesn't perform any operation which is not593// recognized as a standard reduction.594 595func.func @vecdim_none_reduction_rejected(%in: memref<256x512xf32>, %out: memref<256xf32>) {596 %cst = arith.constant 1.000000e+00 : f32597 %final_red = affine.for %j = 0 to 512 iter_args(%red_iter = %cst) -> (f32) {598 affine.yield %red_iter : f32599 }600 affine.store %final_red, %out[0] : memref<256xf32>601 return602}603 604// CHECK-LABEL: @vecdim_none_reduction_rejected605// CHECK-NOT: vector606 607// -----608 609// The number of iterations is not divisable by the vector size, so a mask has610// to be applied to the last update of the accumulator.611 612func.func @vecdim_reduction_masked(%in: memref<256x512xf32>, %out: memref<256xf32>) {613 %cst = arith.constant 0.000000e+00 : f32614 affine.for %i = 0 to 256 {615 %final_red = affine.for %j = 0 to 500 iter_args(%red_iter = %cst) -> (f32) {616 %ld = affine.load %in[%i, %j] : memref<256x512xf32>617 %add = arith.addf %red_iter, %ld : f32618 affine.yield %add : f32619 }620 affine.store %final_red, %out[%i] : memref<256xf32>621 }622 return623}624 625// CHECK: #[[$map0:.*]] = affine_map<([[d0:.*]]) -> (-[[d0]] + 500)>626// CHECK-LABEL: @vecdim_reduction_masked627// CHECK: affine.for %{{.*}} = 0 to 256 {628// CHECK: %[[vzero:.*]] = arith.constant dense<0.000000e+00> : vector<128xf32>629// CHECK: %[[vred:.*]] = affine.for %[[iv:.*]] = 0 to 500 step 128 iter_args(%[[red_iter:.*]] = %[[vzero]]) -> (vector<128xf32>) {630// CHECK: %[[elems_left:.*]] = affine.apply #[[$map0]](%[[iv]])631// CHECK: %[[mask:.*]] = vector.create_mask %[[elems_left]] : vector<128xi1>632// CHECK: %[[ld:.*]] = vector.transfer_read %{{.*}} : memref<256x512xf32>, vector<128xf32>633// CHECK: %[[select:.*]] = arith.select %[[mask]], %[[ld]], %[[vzero]] : vector<128xi1>, vector<128xf32>634// CHECK: %[[add:.*]] = arith.addf %[[red_iter]], %[[select]] : vector<128xf32>635// CHECK: affine.yield %[[add]] : vector<128xf32>636// CHECK: }637// CHECK: %[[final_sum:.*]] = vector.reduction <add>, %[[vred:.*]] : vector<128xf32> into f32638// CHECK: affine.store %[[final_sum]], %{{.*}} : memref<256xf32>639// CHECK: }640 641// -----642 643// The number of iteration is not known, so a mask has to be applied.644 645func.func @vecdim_reduction_masked_unknown_ub(%in: memref<256x512xf32>, %out: memref<256xf32>, %bnd: index) {646 %cst = arith.constant 0.000000e+00 : f32647 affine.for %i = 0 to 256 {648 %final_red = affine.for %j = 0 to %bnd iter_args(%red_iter = %cst) -> (f32) {649 %ld = affine.load %in[%i, %j] : memref<256x512xf32>650 %add = arith.addf %red_iter, %ld : f32651 affine.yield %add : f32652 }653 affine.store %final_red, %out[%i] : memref<256xf32>654 }655 return656}657 658// CHECK: #[[$map1:.*]] = affine_map<([[d0:.*]]){{\[}}[[s0:.*]]{{\]}} -> (-[[d0]] + [[s0]])>659// CHECK-LABEL: @vecdim_reduction_masked_unknown_ub660// CHECK: affine.for %{{.*}} = 0 to 256 {661// CHECK: %[[vzero:.*]] = arith.constant dense<0.000000e+00> : vector<128xf32>662// CHECK: %[[vred:.*]] = affine.for %[[iv:.*]] = 0 to %[[bnd:.*]] step 128 iter_args(%[[red_iter:.*]] = %[[vzero]]) -> (vector<128xf32>) {663// CHECK: %[[elems_left:.*]] = affine.apply #[[$map1]](%[[iv]])[%[[bnd]]]664// CHECK: %[[mask:.*]] = vector.create_mask %[[elems_left]] : vector<128xi1>665// CHECK: %[[ld:.*]] = vector.transfer_read %{{.*}} : memref<256x512xf32>, vector<128xf32>666// CHECK: %[[select:.*]] = arith.select %[[mask]], %[[ld]], %[[vzero]] : vector<128xi1>, vector<128xf32>667// CHECK: %[[add:.*]] = arith.addf %[[red_iter]], %[[select]] : vector<128xf32>668// CHECK: affine.yield %[[add]] : vector<128xf32>669// CHECK: }670// CHECK: %[[final_sum:.*]] = vector.reduction <add>, %[[vred:.*]] : vector<128xf32> into f32671// CHECK: affine.store %[[final_sum]], %{{.*}} : memref<256xf32>672// CHECK: }673 674// -----675 676// The lower bound is nonzero, but the number of iterations is divisible by the677// vector size, so masking is not needed.678 679func.func @vecdim_reduction_nonzero_lb(%in: memref<256x512xf32>, %out: memref<256xf32>) {680 %cst = arith.constant 0.000000e+00 : f32681 affine.for %i = 0 to 256 {682 %final_red = affine.for %j = 127 to 511 iter_args(%red_iter = %cst) -> (f32) {683 %ld = affine.load %in[%i, %j] : memref<256x512xf32>684 %add = arith.addf %red_iter, %ld : f32685 affine.yield %add : f32686 }687 affine.store %final_red, %out[%i] : memref<256xf32>688 }689 return690}691 692// CHECK-LABEL: @vecdim_reduction_nonzero_lb693// CHECK: %{{.*}} = affine.for %{{.*}} = 127 to 511 step 128 iter_args({{.*}}) -> (vector<128xf32>) {694// CHECK-NOT: vector.create_mask695 696// -----697 698// The lower bound is unknown, so we need to create a mask.699 700func.func @vecdim_reduction_masked_unknown_lb(%in: memref<256x512xf32>, %out: memref<256xf32>, %lb: index) {701 %cst = arith.constant 0.000000e+00 : f32702 affine.for %i = 0 to 256 {703 %final_red = affine.for %j = %lb to 512 iter_args(%red_iter = %cst) -> (f32) {704 %ld = affine.load %in[%i, %j] : memref<256x512xf32>705 %add = arith.addf %red_iter, %ld : f32706 affine.yield %add : f32707 }708 affine.store %final_red, %out[%i] : memref<256xf32>709 }710 return711}712 713// CHECK: #[[$map2:.*]] = affine_map<([[d0:.*]]) -> (-[[d0]] + 512)>714// CHECK-LABEL: @vecdim_reduction_masked_unknown_lb715// CHECK: %[[vzero:.*]] = arith.constant dense<0.000000e+00> : vector<128xf32>716// CHECK: %{{.*}} = affine.for %[[iv:.*]] = %[[lb:.*]] to 512 step 128 iter_args(%[[red_iter:.*]] = {{.*}}) -> (vector<128xf32>) {717// CHECK: %[[elems_left:.*]] = affine.apply #[[$map2]](%[[iv]])718// CHECK: %[[mask:.*]] = vector.create_mask %[[elems_left]] : vector<128xi1>719// CHECK: %[[ld:.*]] = vector.transfer_read %{{.*}} : memref<256x512xf32>, vector<128xf32>720// CHECK: %[[select:.*]] = arith.select %[[mask]], %[[ld]], %[[vzero]] : vector<128xi1>, vector<128xf32>721// CHECK: %[[add:.*]] = arith.addf %[[red_iter]], %[[select]] : vector<128xf32>722// CHECK: affine.yield %[[add]] : vector<128xf32>723 724// -----725 726// The upper bound is a minimum expression.727 728func.func @vecdim_reduction_complex_ub(%in: memref<256x512xf32>, %out: memref<256xf32>, %M: index, %N: index) {729 %cst = arith.constant 0.000000e+00 : f32730 affine.for %i = 0 to 256 {731 %final_red = affine.for %j = 0 to min affine_map<(d0, d1) -> (d0, d1*2)>(%M, %N) iter_args(%red_iter = %cst) -> (f32) {732 %ld = affine.load %in[%i, %j] : memref<256x512xf32>733 %add = arith.addf %red_iter, %ld : f32734 affine.yield %add : f32735 }736 affine.store %final_red, %out[%i] : memref<256xf32>737 }738 return739}740 741// CHECK: #[[$map3:.*]] = affine_map<(d0, d1) -> (d0, d1 * 2)>742// CHECK: #[[$map3_sub:.*]] = affine_map<(d0)[s0] -> (-d0 + s0)>743// CHECK-LABEL: @vecdim_reduction_complex_ub744// CHECK: %[[vzero:.*]] = arith.constant dense<0.000000e+00> : vector<128xf32>745// CHECK: %{{.*}} = affine.for %[[iv:.*]] = 0 to min #[[$map3]](%[[M:.*]], %[[N:.*]]) step 128 iter_args(%[[red_iter:.*]] = {{.*}}) -> (vector<128xf32>) {746// CHECK: %[[ub:.*]] = affine.min #[[$map3]](%[[M]], %[[N]])747// CHECK: %[[elems_left:.*]] = affine.apply #[[$map3_sub]](%[[iv]])[%[[ub]]]748// CHECK: %[[mask:.*]] = vector.create_mask %[[elems_left]] : vector<128xi1>749// CHECK: %[[ld:.*]] = vector.transfer_read %{{.*}} : memref<256x512xf32>, vector<128xf32>750// CHECK: %[[select:.*]] = arith.select %[[mask]], %[[ld]], %[[vzero]] : vector<128xi1>, vector<128xf32>751// CHECK: %[[add:.*]] = arith.addf %[[red_iter]], %[[select]] : vector<128xf32>752// CHECK: affine.yield %[[add]] : vector<128xf32>753 754// -----755 756// The same mask is applied to both reductions.757 758func.func @vecdim_two_reductions_masked(%in: memref<256x512xf32>, %out: memref<512xf32>) {759 %cst = arith.constant 0.000000e+00 : f32760 affine.for %i = 0 to 256 {761 %final_sum, %final_expsum = affine.for %j = 0 to 500 iter_args(%sum_iter = %cst, %expsum_iter = %cst) -> (f32, f32) {762 %ld = affine.load %in[%i, %j] : memref<256x512xf32>763 %exp = math.exp %ld : f32764 %add = arith.addf %sum_iter, %ld : f32765 %eadd = arith.addf %expsum_iter, %exp : f32766 affine.yield %add, %eadd : f32, f32767 }768 affine.store %final_sum, %out[2*%i] : memref<512xf32>769 affine.store %final_expsum, %out[2*%i + 1] : memref<512xf32>770 }771 return772}773 774// CHECK: #[[$map4:.*]] = affine_map<([[d0:.*]]) -> (-[[d0]] + 500)>775// CHECK-LABEL: @vecdim_two_reductions_masked776// CHECK: affine.for %{{.*}} = 0 to 256 {777// CHECK: %[[vzero0:.*]] = arith.constant dense<0.000000e+00> : vector<128xf32>778// CHECK: %[[vzero1:.*]] = arith.constant dense<0.000000e+00> : vector<128xf32>779// CHECK: %{{.*}} = affine.for %[[iv:.*]] = 0 to 500 step 128 iter_args(%[[sum_iter:.*]] = {{.*}}, %[[esum_iter:.*]] = {{.*}}) -> (vector<128xf32>, vector<128xf32>) {780// CHECK: %[[elems_left:.*]] = affine.apply #[[$map4]](%[[iv]])781// CHECK: %[[mask:.*]] = vector.create_mask %[[elems_left]] : vector<128xi1>782// CHECK: %[[ld:.*]] = vector.transfer_read %{{.*}} : memref<256x512xf32>, vector<128xf32>783// CHECK: %[[exp:.*]] = math.exp %[[ld]] : vector<128xf32>784// CHECK: %[[select0:.*]] = arith.select %[[mask]], %[[ld]], %[[vzero0]] : vector<128xi1>, vector<128xf32>785// CHECK: %[[add:.*]] = arith.addf %[[sum_iter]], %[[select0]] : vector<128xf32>786// CHECK: %[[select1:.*]] = arith.select %[[mask]], %[[exp]], %[[vzero1]] : vector<128xi1>, vector<128xf32>787// CHECK: %[[eadd:.*]] = arith.addf %[[esum_iter]], %[[select1]] : vector<128xf32>788// CHECK: affine.yield %[[add]], %[[eadd]] : vector<128xf32>789// CHECK: }790