// RUN: mlir-opt %s -affine-super-vectorize="virtual-vector-size=128 test-fastest-varying=0 vectorize-reductions=true" -split-input-file | FileCheck %s

// The inner reduction loop '%j' is vectorized.

func @vecdim_reduction(%in: memref<256x512xf32>, %out: memref<256xf32>) {
 %cst = arith.constant 0.000000e+00 : f32
 affine.for %i = 0 to 256 {
   %final_red = affine.for %j = 0 to 512 iter_args(%red_iter = %cst) -> (f32) {
     %ld = affine.load %in[%i, %j] : memref<256x512xf32>
     %add = arith.addf %red_iter, %ld : f32
     affine.yield %add : f32
   }
   affine.store %final_red, %out[%i] : memref<256xf32>
 }
 return
}

// CHECK-LABEL: @vecdim_reduction
// CHECK:       affine.for %{{.*}} = 0 to 256 {
// CHECK:         %[[vzero:.*]] = arith.constant dense<0.000000e+00> : vector<128xf32>
// CHECK:         %[[vred:.*]] = affine.for %{{.*}} = 0 to 512 step 128 iter_args(%[[red_iter:.*]] = %[[vzero]]) -> (vector<128xf32>) {
// CHECK:           %[[ld:.*]] = vector.transfer_read %{{.*}} : memref<256x512xf32>, vector<128xf32>
// CHECK:           %[[add:.*]] = arith.addf %[[red_iter]], %[[ld]] : vector<128xf32>
// CHECK:           affine.yield %[[add]] : vector<128xf32>
// CHECK:         }
// CHECK:         %[[final_sum:.*]] = vector.reduction "add", %[[vred:.*]] : vector<128xf32> into f32
// CHECK:         affine.store %[[final_sum]], %{{.*}} : memref<256xf32>
// CHECK:       }

// -----

func @vecdim_reduction_minf(%in: memref<256x512xf32>, %out: memref<256xf32>) {
 %cst = arith.constant 0x7F800000 : f32
 affine.for %i = 0 to 256 {
   %final_red = affine.for %j = 0 to 512 iter_args(%red_iter = %cst) -> (f32) {
     %ld = affine.load %in[%i, %j] : memref<256x512xf32>
     %min = arith.minf %red_iter, %ld : f32
     affine.yield %min : f32
   }
   affine.store %final_red, %out[%i] : memref<256xf32>
 }
 return
}

// CHECK-LABEL: @vecdim_reduction_minf
// CHECK:       affine.for %{{.*}} = 0 to 256 {
// CHECK:         %[[vmax:.*]] = arith.constant dense<0x7F800000> : vector<128xf32>
// CHECK:         %[[vred:.*]] = affine.for %{{.*}} = 0 to 512 step 128 iter_args(%[[red_iter:.*]] = %[[vmax]]) -> (vector<128xf32>) {
// CHECK:           %[[ld:.*]] = vector.transfer_read %{{.*}} : memref<256x512xf32>, vector<128xf32>
// CHECK:           %[[min:.*]] = arith.minf %[[red_iter]], %[[ld]] : vector<128xf32>
// CHECK:           affine.yield %[[min]] : vector<128xf32>
// CHECK:         }
// CHECK:         %[[final_min:.*]] = vector.reduction "minf", %[[vred:.*]] : vector<128xf32> into f32
// CHECK:         affine.store %[[final_min]], %{{.*}} : memref<256xf32>
// CHECK:       }

// -----

func @vecdim_reduction_maxf(%in: memref<256x512xf32>, %out: memref<256xf32>) {
 %cst = arith.constant 0xFF800000 : f32
 affine.for %i = 0 to 256 {
   %final_red = affine.for %j = 0 to 512 iter_args(%red_iter = %cst) -> (f32) {
     %ld = affine.load %in[%i, %j] : memref<256x512xf32>
     %max = arith.maxf %red_iter, %ld : f32
     affine.yield %max : f32
   }
   affine.store %final_red, %out[%i] : memref<256xf32>
 }
 return
}

// CHECK-LABEL: @vecdim_reduction_maxf
// CHECK:       affine.for %{{.*}} = 0 to 256 {
// CHECK:         %[[vmin:.*]] = arith.constant dense<0xFF800000> : vector<128xf32>
// CHECK:         %[[vred:.*]] = affine.for %{{.*}} = 0 to 512 step 128 iter_args(%[[red_iter:.*]] = %[[vmin]]) -> (vector<128xf32>) {
// CHECK:           %[[ld:.*]] = vector.transfer_read %{{.*}} : memref<256x512xf32>, vector<128xf32>
// CHECK:           %[[max:.*]] = arith.maxf %[[red_iter]], %[[ld]] : vector<128xf32>
// CHECK:           affine.yield %[[max]] : vector<128xf32>
// CHECK:         }
// CHECK:         %[[final_max:.*]] = vector.reduction "maxf", %[[vred:.*]] : vector<128xf32> into f32
// CHECK:         affine.store %[[final_max]], %{{.*}} : memref<256xf32>
// CHECK:       }

// -----

func @vecdim_reduction_minsi(%in: memref<256x512xi32>, %out: memref<256xi32>) {
 %cst = arith.constant 2147483647 : i32
 affine.for %i = 0 to 256 {
   %final_red = affine.for %j = 0 to 512 iter_args(%red_iter = %cst) -> (i32) {
     %ld = affine.load %in[%i, %j] : memref<256x512xi32>
     %min = arith.minsi %red_iter, %ld : i32
     affine.yield %min : i32
   }
   affine.store %final_red, %out[%i] : memref<256xi32>
 }
 return
}

// CHECK-LABEL: @vecdim_reduction_minsi
// CHECK:       affine.for %{{.*}} = 0 to 256 {
// CHECK:         %[[vmax:.*]] = arith.constant dense<2147483647> : vector<128xi32>
// CHECK:         %[[vred:.*]] = affine.for %{{.*}} = 0 to 512 step 128 iter_args(%[[red_iter:.*]] = %[[vmax]]) -> (vector<128xi32>) {
// CHECK:           %[[ld:.*]] = vector.transfer_read %{{.*}} : memref<256x512xi32>, vector<128xi32>
// CHECK:           %[[min:.*]] = arith.minsi %[[red_iter]], %[[ld]] : vector<128xi32>
// CHECK:           affine.yield %[[min]] : vector<128xi32>
// CHECK:         }
// CHECK:         %[[final_min:.*]] = vector.reduction "minsi", %[[vred:.*]] : vector<128xi32> into i32
// CHECK:         affine.store %[[final_min]], %{{.*}} : memref<256xi32>
// CHECK:       }

// -----

func @vecdim_reduction_maxsi(%in: memref<256x512xi32>, %out: memref<256xi32>) {
 %cst = arith.constant -2147483648 : i32
 affine.for %i = 0 to 256 {
   %final_red = affine.for %j = 0 to 512 iter_args(%red_iter = %cst) -> (i32) {
     %ld = affine.load %in[%i, %j] : memref<256x512xi32>
     %max = arith.maxsi %red_iter, %ld : i32
     affine.yield %max : i32
   }
   affine.store %final_red, %out[%i] : memref<256xi32>
 }
 return
}

// CHECK-LABEL: @vecdim_reduction_maxsi
// CHECK:       affine.for %{{.*}} = 0 to 256 {
// CHECK:         %[[vmin:.*]] = arith.constant dense<-2147483648> : vector<128xi32>
// CHECK:         %[[vred:.*]] = affine.for %{{.*}} = 0 to 512 step 128 iter_args(%[[red_iter:.*]] = %[[vmin]]) -> (vector<128xi32>) {
// CHECK:           %[[ld:.*]] = vector.transfer_read %{{.*}} : memref<256x512xi32>, vector<128xi32>
// CHECK:           %[[max:.*]] = arith.maxsi %[[red_iter]], %[[ld]] : vector<128xi32>
// CHECK:           affine.yield %[[max]] : vector<128xi32>
// CHECK:         }
// CHECK:         %[[final_max:.*]] = vector.reduction "maxsi", %[[vred:.*]] : vector<128xi32> into i32
// CHECK:         affine.store %[[final_max]], %{{.*}} : memref<256xi32>
// CHECK:       }

// -----

func @vecdim_reduction_minui(%in: memref<256x512xi32>, %out: memref<256xi32>) {
 %cst = arith.constant -1 : i32
 affine.for %i = 0 to 256 {
   %final_red = affine.for %j = 0 to 512 iter_args(%red_iter = %cst) -> (i32) {
     %ld = affine.load %in[%i, %j] : memref<256x512xi32>
     %min = arith.minui %red_iter, %ld : i32
     affine.yield %min : i32
   }
   affine.store %final_red, %out[%i] : memref<256xi32>
 }
 return
}

// CHECK-LABEL: @vecdim_reduction_minui
// CHECK:       affine.for %{{.*}} = 0 to 256 {
// CHECK:         %[[vmax:.*]] = arith.constant dense<-1> : vector<128xi32>
// CHECK:         %[[vred:.*]] = affine.for %{{.*}} = 0 to 512 step 128 iter_args(%[[red_iter:.*]] = %[[vmax]]) -> (vector<128xi32>) {
// CHECK:           %[[ld:.*]] = vector.transfer_read %{{.*}} : memref<256x512xi32>, vector<128xi32>
// CHECK:           %[[min:.*]] = arith.minui %[[red_iter]], %[[ld]] : vector<128xi32>
// CHECK:           affine.yield %[[min]] : vector<128xi32>
// CHECK:         }
// CHECK:         %[[final_min:.*]] = vector.reduction "minui", %[[vred:.*]] : vector<128xi32> into i32
// CHECK:         affine.store %[[final_min]], %{{.*}} : memref<256xi32>
// CHECK:       }

// -----

func @vecdim_reduction_maxui(%in: memref<256x512xi32>, %out: memref<256xi32>) {
 %cst = arith.constant 0 : i32
 affine.for %i = 0 to 256 {
   %final_red = affine.for %j = 0 to 512 iter_args(%red_iter = %cst) -> (i32) {
     %ld = affine.load %in[%i, %j] : memref<256x512xi32>
     %max = arith.maxui %red_iter, %ld : i32
     affine.yield %max : i32
   }
   affine.store %final_red, %out[%i] : memref<256xi32>
 }
 return
}

// CHECK-LABEL: @vecdim_reduction_maxui
// CHECK:       affine.for %{{.*}} = 0 to 256 {
// CHECK:         %[[vmin:.*]] = arith.constant dense<0> : vector<128xi32>
// CHECK:         %[[vred:.*]] = affine.for %{{.*}} = 0 to 512 step 128 iter_args(%[[red_iter:.*]] = %[[vmin]]) -> (vector<128xi32>) {
// CHECK:           %[[ld:.*]] = vector.transfer_read %{{.*}} : memref<256x512xi32>, vector<128xi32>
// CHECK:           %[[max:.*]] = arith.maxui %[[red_iter]], %[[ld]] : vector<128xi32>
// CHECK:           affine.yield %[[max]] : vector<128xi32>
// CHECK:         }
// CHECK:         %[[final_max:.*]] = vector.reduction "maxui", %[[vred:.*]] : vector<128xi32> into i32
// CHECK:         affine.store %[[final_max]], %{{.*}} : memref<256xi32>
// CHECK:       }

// -----

// The inner reduction loop '%j' is vectorized. (The order of addf's operands is
// different than in the previous test case).

func @vecdim_reduction_comm(%in: memref<256x512xf32>, %out: memref<256xf32>) {
 %cst = arith.constant 0.000000e+00 : f32
 affine.for %i = 0 to 256 {
   %final_red = affine.for %j = 0 to 512 iter_args(%red_iter = %cst) -> (f32) {
     %ld = affine.load %in[%i, %j] : memref<256x512xf32>
     %add = arith.addf %ld, %red_iter : f32
     affine.yield %add : f32
   }
   affine.store %final_red, %out[%i] : memref<256xf32>
 }
 return
}

// CHECK-LABEL: @vecdim_reduction_comm
// CHECK:       affine.for %{{.*}} = 0 to 256 {
// CHECK:         %[[vzero:.*]] = arith.constant dense<0.000000e+00> : vector<128xf32>
// CHECK:         %[[vred:.*]] = affine.for %{{.*}} = 0 to 512 step 128 iter_args(%[[red_iter:.*]] = %[[vzero]]) -> (vector<128xf32>) {
// CHECK:           %[[ld:.*]] = vector.transfer_read %{{.*}} : memref<256x512xf32>, vector<128xf32>
// CHECK:           %[[add:.*]] = arith.addf %[[ld]], %[[red_iter]] : vector<128xf32>
// CHECK:           affine.yield %[[add]] : vector<128xf32>
// CHECK:         }
// CHECK:         %[[final_sum:.*]] = vector.reduction "add", %[[vred:.*]] : vector<128xf32> into f32
// CHECK:         affine.store %[[final_sum]], %{{.*}} : memref<256xf32>
// CHECK:       }

// -----

// The inner reduction loop '%j' is vectorized. Transforming the input before
// performing the accumulation doesn't cause any problem.

func @vecdim_reduction_expsin(%in: memref<256x512xf32>, %out: memref<256xf32>) {
 %cst = arith.constant 0.000000e+00 : f32
 affine.for %i = 0 to 256 {
   %final_red = affine.for %j = 0 to 512 iter_args(%red_iter = %cst) -> (f32) {
     %ld = affine.load %in[%i, %j] : memref<256x512xf32>
     %sin = math.sin %ld : f32
     %exp = math.exp %sin : f32
     %add = arith.addf %red_iter, %exp : f32
     affine.yield %add : f32
   }
   affine.store %final_red, %out[%i] : memref<256xf32>
 }
 return
}

// CHECK-LABEL: @vecdim_reduction_expsin
// CHECK:       affine.for %{{.*}} = 0 to 256 {
// CHECK:         %[[vzero:.*]] = arith.constant dense<0.000000e+00> : vector<128xf32>
// CHECK:         %[[vred:.*]] = affine.for %{{.*}} = 0 to 512 step 128 iter_args(%[[red_iter:.*]] = %[[vzero]]) -> (vector<128xf32>) {
// CHECK:           %[[ld:.*]] = vector.transfer_read %{{.*}} : memref<256x512xf32>, vector<128xf32>
// CHECK:           %[[sin:.*]] = math.sin %[[ld]]
// CHECK:           %[[exp:.*]] = math.exp %[[sin]]
// CHECK:           %[[add:.*]] = arith.addf %[[red_iter]], %[[exp]] : vector<128xf32>
// CHECK:           affine.yield %[[add]] : vector<128xf32>
// CHECK:         }
// CHECK:         %[[final_sum:.*]] = vector.reduction "add", %[[vred:.*]] : vector<128xf32> into f32
// CHECK:         affine.store %[[final_sum]], %{{.*}} : memref<256xf32>
// CHECK:       }

// -----

// Two reductions at the same time. The inner reduction loop '%j' is vectorized.

func @two_vecdim_reductions(%in: memref<256x512xf32>, %out_sum: memref<256xf32>, %out_prod: memref<256xf32>) {
 %cst = arith.constant 1.000000e+00 : f32
 affine.for %i = 0 to 256 {
   // Note that we pass the same constant '1.0' as initial values for both
   // reductions.
   %sum, %prod = affine.for %j = 0 to 512 iter_args(%part_sum = %cst, %part_prod = %cst) -> (f32, f32) {
     %ld = affine.load %in[%i, %j] : memref<256x512xf32>
     %add = arith.addf %part_sum, %ld : f32
     %mul = arith.mulf %part_prod, %ld : f32
     affine.yield %add, %mul : f32, f32
   }
   affine.store %sum, %out_sum[%i] : memref<256xf32>
   affine.store %prod, %out_prod[%i] : memref<256xf32>
 }
 return
}

// CHECK-LABEL: @two_vecdim_reductions
// CHECK:       %[[cst:.*]] = arith.constant 1.000000e+00 : f32
// CHECK:       affine.for %{{.*}} = 0 to 256 {
// CHECK:         %[[vzero:.*]] = arith.constant dense<0.000000e+00> : vector<128xf32>
// CHECK:         %[[vone:.*]] = arith.constant dense<1.000000e+00> : vector<128xf32>
// CHECK:         %[[vred:.*]]:2 = affine.for %{{.*}} = 0 to 512 step 128 iter_args(%[[part_sum:.*]] = %[[vzero]], %[[part_prod:.*]] = %[[vone]]) -> (vector<128xf32>, vector<128xf32>) {
// CHECK:           %[[ld:.*]] = vector.transfer_read %{{.*}} : memref<256x512xf32>, vector<128xf32>
// CHECK:           %[[add:.*]] = arith.addf %[[part_sum]], %[[ld]] : vector<128xf32>
// CHECK:           %[[mul:.*]] = arith.mulf %[[part_prod]], %[[ld]] : vector<128xf32>
// CHECK:           affine.yield %[[add]], %[[mul]] : vector<128xf32>, vector<128xf32>
// CHECK:         }
// CHECK:         %[[nonfinal_sum:.*]] = vector.reduction "add", %[[vred:.*]]#0 : vector<128xf32> into f32
// Note that to compute the final sum we need to add the original initial value
// (%cst) since it is not zero.
// CHECK:         %[[final_sum:.*]] = arith.addf %[[nonfinal_sum]], %[[cst]] : f32
// For the final product we don't need to do this additional step because the
// initial value equals to 1 (the neutral element for multiplication).
// CHECK:         %[[final_prod:.*]] = vector.reduction "mul", %[[vred:.*]]#1 : vector<128xf32> into f32
// CHECK:         affine.store %[[final_sum]], %{{.*}} : memref<256xf32>
// CHECK:         affine.store %[[final_prod]], %{{.*}} : memref<256xf32>
// CHECK:       }

// -----

// The integer case.

func @two_vecdim_reductions_int(%in: memref<256x512xi64>, %out_sum: memref<256xi64>, %out_prod: memref<256xi64>) {
 %cst0 = arith.constant 0 : i64
 %cst1 = arith.constant 1 : i64
 affine.for %i = 0 to 256 {
   %sum, %prod = affine.for %j = 0 to 512 iter_args(%part_sum = %cst0, %part_prod = %cst1) -> (i64, i64) {
     %ld = affine.load %in[%i, %j] : memref<256x512xi64>
     %add = arith.addi %part_sum, %ld : i64
     %mul = arith.muli %part_prod, %ld : i64
     affine.yield %add, %mul : i64, i64
   }
   affine.store %sum, %out_sum[%i] : memref<256xi64>
   affine.store %prod, %out_prod[%i] : memref<256xi64>
 }
 return
}

// CHECK-LABEL: @two_vecdim_reductions
// CHECK:       affine.for %{{.*}} = 0 to 256 {
// CHECK:         %[[vzero:.*]] = arith.constant dense<0> : vector<128xi64>
// CHECK:         %[[vone:.*]] = arith.constant dense<1> : vector<128xi64>
// CHECK:         %[[vred:.*]]:2 = affine.for %{{.*}} = 0 to 512 step 128 iter_args(%[[part_sum:.*]] = %[[vzero]], %[[part_prod:.*]] = %[[vone]]) -> (vector<128xi64>, vector<128xi64>) {
// CHECK:           %[[ld:.*]] = vector.transfer_read %{{.*}} : memref<256x512xi64>, vector<128xi64>
// CHECK:           %[[add:.*]] = arith.addi %[[part_sum]], %[[ld]] : vector<128xi64>
// CHECK:           %[[mul:.*]] = arith.muli %[[part_prod]], %[[ld]] : vector<128xi64>
// CHECK:           affine.yield %[[add]], %[[mul]] : vector<128xi64>, vector<128xi64>
// CHECK:         }
// CHECK:         %[[final_sum:.*]] = vector.reduction "add", %[[vred:.*]]#0 : vector<128xi64> into i64
// CHECK:         %[[final_prod:.*]] = vector.reduction "mul", %[[vred:.*]]#1 : vector<128xi64> into i64
// CHECK:         affine.store %[[final_sum]], %{{.*}} : memref<256xi64>
// CHECK:         affine.store %[[final_prod]], %{{.*}} : memref<256xi64>
// CHECK:       }

// -----

// The outer reduction loop '%j' is vectorized.

func @vecdim_reduction_nested(%in: memref<256x512xf32>, %out: memref<1xf32>) {
 %cst = arith.constant 0.000000e+00 : f32
 %outer_red = affine.for %j = 0 to 512 iter_args(%outer_iter = %cst) -> (f32) {
   %inner_red = affine.for %i = 0 to 256 iter_args(%inner_iter = %cst) -> (f32) {
     %ld = affine.load %in[%i, %j] : memref<256x512xf32>
     %add = arith.addf %inner_iter, %ld : f32
     affine.yield %add : f32
   }
   %outer_add = arith.addf %outer_iter, %inner_red : f32
   affine.yield %outer_add : f32
 }
 affine.store %outer_red, %out[0] : memref<1xf32>
 return
}

// CHECK-LABEL: @vecdim_reduction_nested
// CHECK:       %[[vzero:.*]] = arith.constant dense<0.000000e+00> : vector<128xf32>
// CHECK:       %[[outer_red:.*]] = affine.for %{{.*}} = 0 to 512 step 128 iter_args(%[[outer_iter:.*]] = %[[vzero]]) -> (vector<128xf32>) {
// CHECK:         %[[vzero:.*]] = arith.constant dense<0.000000e+00> : vector<128xf32>
// CHECK:         %[[inner_red:.*]] = affine.for %{{.*}} = 0 to 256 iter_args(%[[inner_iter:.*]] = %[[vzero]]) -> (vector<128xf32>) {
// CHECK:           %[[ld:.*]] = vector.transfer_read %{{.*}} : memref<256x512xf32>, vector<128xf32>
// CHECK:           %[[add:.*]] = arith.addf %[[inner_iter]], %[[ld]] : vector<128xf32>
// CHECK:           affine.yield %[[add]] : vector<128xf32>
// CHECK:         }
// CHECK:         %[[outer_add:.*]] = arith.addf %[[outer_iter]], %[[inner_red]] : vector<128xf32>
// CHECK:         affine.yield %[[outer_add]] : vector<128xf32>
// CHECK:       }
// CHECK:       %[[final_sum:.*]] = vector.reduction "add", %[[outer_red:.*]] : vector<128xf32> into f32
// CHECK:       affine.store %[[final_sum]], %{{.*}} : memref<1xf32>

// -----

// The inner reduction loop '%j' computes partial sums as a side effect and
// is not vectorized.

func @vecdim_partial_sums_1_rejected(%in: memref<256x512xf32>, %out_sum: memref<256xf32>, %out_prod: memref<256xf32>, %out_partsum: memref<256x512xf32>) {
 %cst = arith.constant 1.000000e+00 : f32
 affine.for %i = 0 to 256 {
   %sum, %prod = affine.for %j = 0 to 512 iter_args(%part_sum = %cst, %part_prod = %cst) -> (f32, f32) {
     %ld = affine.load %in[%i, %j] : memref<256x512xf32>
     %add = arith.addf %part_sum, %ld : f32
     %mul = arith.mulf %part_prod, %ld : f32
     affine.store %add, %out_partsum[%i, %j] : memref<256x512xf32>
     affine.yield %add, %mul : f32, f32
   }
   affine.store %sum, %out_sum[%i] : memref<256xf32>
   affine.store %prod, %out_prod[%i] : memref<256xf32>
 }
 return
}

// CHECK-LABEL: @vecdim_partial_sums_1_rejected
// CHECK-NOT:   vector

// -----

// The inner reduction loop '%j' computes partial sums as a side effect and
// is not vectorized.

func @vecdim_partial_sums_2_rejected(%in: memref<256x512xf32>, %out_sum: memref<256xf32>, %out_prod: memref<256xf32>, %out_partsum: memref<256x512xf32>) {
 %cst = arith.constant 1.000000e+00 : f32
 affine.for %i = 0 to 256 {
   %sum, %prod = affine.for %j = 0 to 512 iter_args(%part_sum = %cst, %part_prod = %cst) -> (f32, f32) {
     affine.store %part_sum, %out_partsum[%i, %j] : memref<256x512xf32>
     %ld = affine.load %in[%i, %j] : memref<256x512xf32>
     %add = arith.addf %part_sum, %ld : f32
     %mul = arith.mulf %part_prod, %ld : f32
     affine.yield %add, %mul : f32, f32
   }
   affine.store %sum, %out_sum[%i] : memref<256xf32>
   affine.store %prod, %out_prod[%i] : memref<256xf32>
 }
 return
}

// CHECK-LABEL: @vecdim_partial_sums_2_rejected
// CHECK-NOT:   vector

// -----

// The inner reduction loop '%j' performs an unknown reduction operation and is
// not vectorized.

func @vecdim_unknown_reduction_rejected(%in: memref<256x512xf32>, %out: memref<256xf32>) {
 %cst = arith.constant 1.000000e+00 : f32
 %final_red = affine.for %j = 0 to 512 iter_args(%red_iter = %cst) -> (f32) {
   %add = arith.addf %red_iter, %red_iter : f32
   affine.yield %add : f32
 }
 affine.store %final_red, %out[0] : memref<256xf32>
 return
}

// CHECK-LABEL: @vecdim_unknown_reduction_rejected
// CHECK-NOT:   vector

// -----

// The inner reduction loop '%j' doesn't perform any operation which is not
// recognized as a standard reduction.

func @vecdim_none_reduction_rejected(%in: memref<256x512xf32>, %out: memref<256xf32>) {
 %cst = arith.constant 1.000000e+00 : f32
 %final_red = affine.for %j = 0 to 512 iter_args(%red_iter = %cst) -> (f32) {
   affine.yield %red_iter : f32
 }
 affine.store %final_red, %out[0] : memref<256xf32>
 return
}

// CHECK-LABEL: @vecdim_none_reduction_rejected
// CHECK-NOT:   vector

// -----

// The number of iterations is not divisable by the vector size, so a mask has
// to be applied to the last update of the accumulator.

func @vecdim_reduction_masked(%in: memref<256x512xf32>, %out: memref<256xf32>) {
 %cst = arith.constant 0.000000e+00 : f32
 affine.for %i = 0 to 256 {
   %final_red = affine.for %j = 0 to 500 iter_args(%red_iter = %cst) -> (f32) {
     %ld = affine.load %in[%i, %j] : memref<256x512xf32>
     %add = arith.addf %red_iter, %ld : f32
     affine.yield %add : f32
   }
   affine.store %final_red, %out[%i] : memref<256xf32>
 }
 return
}

// CHECK:       #[[$map0:.*]] = affine_map<([[d0:.*]]) -> (-[[d0]] + 500)>
// CHECK-LABEL: @vecdim_reduction_masked
// CHECK:       affine.for %{{.*}} = 0 to 256 {
// CHECK:         %[[vzero:.*]] = arith.constant dense<0.000000e+00> : vector<128xf32>
// CHECK:         %[[vred:.*]] = affine.for %[[iv:.*]] = 0 to 500 step 128 iter_args(%[[red_iter:.*]] = %[[vzero]]) -> (vector<128xf32>) {
// CHECK:           %[[elems_left:.*]] = affine.apply #[[$map0]](%[[iv]])
// CHECK:           %[[mask:.*]] = vector.create_mask %[[elems_left]] : vector<128xi1>
// CHECK:           %[[ld:.*]] = vector.transfer_read %{{.*}} : memref<256x512xf32>, vector<128xf32>
// CHECK:           %[[add:.*]] = arith.addf %[[red_iter]], %[[ld]] : vector<128xf32>
// CHECK:           %[[new_acc:.*]] = select %[[mask]], %[[add]], %[[red_iter]] : vector<128xi1>, vector<128xf32>
// CHECK:           affine.yield %[[new_acc]] : vector<128xf32>
// CHECK:         }
// CHECK:         %[[final_sum:.*]] = vector.reduction "add", %[[vred:.*]] : vector<128xf32> into f32
// CHECK:         affine.store %[[final_sum]], %{{.*}} : memref<256xf32>
// CHECK:       }

// -----

// The number of iteration is not known, so a mask has to be applied.

func @vecdim_reduction_masked_unknown_ub(%in: memref<256x512xf32>, %out: memref<256xf32>, %bnd: index) {
 %cst = arith.constant 0.000000e+00 : f32
 affine.for %i = 0 to 256 {
   %final_red = affine.for %j = 0 to %bnd iter_args(%red_iter = %cst) -> (f32) {
     %ld = affine.load %in[%i, %j] : memref<256x512xf32>
     %add = arith.addf %red_iter, %ld : f32
     affine.yield %add : f32
   }
   affine.store %final_red, %out[%i] : memref<256xf32>
 }
 return
}

// CHECK:       #[[$map1:.*]] = affine_map<([[d0:.*]]){{\[}}[[s0:.*]]{{\]}} -> (-[[d0]] + [[s0]])>
// CHECK-LABEL: @vecdim_reduction_masked_unknown_ub
// CHECK:       affine.for %{{.*}} = 0 to 256 {
// CHECK:         %[[vzero:.*]] = arith.constant dense<0.000000e+00> : vector<128xf32>
// CHECK:         %[[vred:.*]] = affine.for %[[iv:.*]] = 0 to %[[bnd:.*]] step 128 iter_args(%[[red_iter:.*]] = %[[vzero]]) -> (vector<128xf32>) {
// CHECK:           %[[elems_left:.*]] = affine.apply #[[$map1]](%[[iv]])[%[[bnd]]]
// CHECK:           %[[mask:.*]] = vector.create_mask %[[elems_left]] : vector<128xi1>
// CHECK:           %[[ld:.*]] = vector.transfer_read %{{.*}} : memref<256x512xf32>, vector<128xf32>
// CHECK:           %[[add:.*]] = arith.addf %[[red_iter]], %[[ld]] : vector<128xf32>
// CHECK:           %[[new_acc:.*]] = select %[[mask]], %[[add]], %[[red_iter]] : vector<128xi1>, vector<128xf32>
// CHECK:           affine.yield %[[new_acc]] : vector<128xf32>
// CHECK:         }
// CHECK:         %[[final_sum:.*]] = vector.reduction "add", %[[vred:.*]] : vector<128xf32> into f32
// CHECK:         affine.store %[[final_sum]], %{{.*}} : memref<256xf32>
// CHECK:       }

// -----

// The lower bound is nonzero, but the number of iterations is divisible by the
// vector size, so masking is not needed.

func @vecdim_reduction_nonzero_lb(%in: memref<256x512xf32>, %out: memref<256xf32>) {
 %cst = arith.constant 0.000000e+00 : f32
 affine.for %i = 0 to 256 {
   %final_red = affine.for %j = 127 to 511 iter_args(%red_iter = %cst) -> (f32) {
     %ld = affine.load %in[%i, %j] : memref<256x512xf32>
     %add = arith.addf %red_iter, %ld : f32
     affine.yield %add : f32
   }
   affine.store %final_red, %out[%i] : memref<256xf32>
 }
 return
}

// CHECK-LABEL: @vecdim_reduction_nonzero_lb
// CHECK:         %{{.*}} = affine.for %{{.*}} = 127 to 511 step 128 iter_args({{.*}}) -> (vector<128xf32>) {
// CHECK-NOT:     vector.create_mask

// -----

// The lower bound is unknown, so we need to create a mask.

func @vecdim_reduction_masked_unknown_lb(%in: memref<256x512xf32>, %out: memref<256xf32>, %lb: index) {
 %cst = arith.constant 0.000000e+00 : f32
 affine.for %i = 0 to 256 {
   %final_red = affine.for %j = %lb to 512 iter_args(%red_iter = %cst) -> (f32) {
     %ld = affine.load %in[%i, %j] : memref<256x512xf32>
     %add = arith.addf %red_iter, %ld : f32
     affine.yield %add : f32
   }
   affine.store %final_red, %out[%i] : memref<256xf32>
 }
 return
}

// CHECK:       #[[$map2:.*]] = affine_map<([[d0:.*]]) -> (-[[d0]] + 512)>
// CHECK-LABEL: @vecdim_reduction_masked_unknown_lb
// CHECK:         %{{.*}} = affine.for %[[iv:.*]] = %[[lb:.*]] to 512 step 128 iter_args(%[[red_iter:.*]] = {{.*}}) -> (vector<128xf32>) {
// CHECK:           %[[elems_left:.*]] = affine.apply #[[$map2]](%[[iv]])
// CHECK:           %[[mask:.*]] = vector.create_mask %[[elems_left]] : vector<128xi1>
// CHECK:           %[[ld:.*]] = vector.transfer_read %{{.*}} : memref<256x512xf32>, vector<128xf32>
// CHECK:           %[[add:.*]] = arith.addf %[[red_iter]], %[[ld]] : vector<128xf32>
// CHECK:           %[[new_acc:.*]] = select %[[mask]], %[[add]], %[[red_iter]] : vector<128xi1>, vector<128xf32>
// CHECK:           affine.yield %[[new_acc]] : vector<128xf32>

// -----

// The upper bound is a minimum expression.

func @vecdim_reduction_complex_ub(%in: memref<256x512xf32>, %out: memref<256xf32>, %M: index, %N: index) {
 %cst = arith.constant 0.000000e+00 : f32
 affine.for %i = 0 to 256 {
   %final_red = affine.for %j = 0 to min affine_map<(d0, d1) -> (d0, d1*2)>(%M, %N) iter_args(%red_iter = %cst) -> (f32) {
     %ld = affine.load %in[%i, %j] : memref<256x512xf32>
     %add = arith.addf %red_iter, %ld : f32
     affine.yield %add : f32
   }
   affine.store %final_red, %out[%i] : memref<256xf32>
 }
 return
}

// CHECK:       #[[$map3:.*]] = affine_map<([[d0:.*]], [[d1:.*]]) -> ([[d0]], [[d1]] * 2)>
// CHECK:       #[[$map3_sub:.*]] = affine_map<([[d0:.*]], [[d1:.*]]) -> ([[d0]] - [[d1]])>
// CHECK-LABEL: @vecdim_reduction_complex_ub
// CHECK:         %{{.*}} = affine.for %[[iv:.*]] = 0 to min #[[$map3]](%[[M:.*]], %[[N:.*]]) step 128 iter_args(%[[red_iter:.*]] = {{.*}}) -> (vector<128xf32>) {
// CHECK:           %[[ub:.*]] = affine.min #[[$map3]](%[[M]], %[[N]])
// CHECK:           %[[elems_left:.*]] = affine.apply #[[$map3_sub]](%[[ub]], %[[iv]])
// CHECK:           %[[mask:.*]] = vector.create_mask %[[elems_left]] : vector<128xi1>
// CHECK:           %[[ld:.*]] = vector.transfer_read %{{.*}} : memref<256x512xf32>, vector<128xf32>
// CHECK:           %[[add:.*]] = arith.addf %[[red_iter]], %[[ld]] : vector<128xf32>
// CHECK:           %[[new_acc:.*]] = select %[[mask]], %[[add]], %[[red_iter]] : vector<128xi1>, vector<128xf32>
// CHECK:           affine.yield %[[new_acc]] : vector<128xf32>

// -----

// The same mask is applied to both reductions.

func @vecdim_two_reductions_masked(%in: memref<256x512xf32>, %out: memref<512xf32>) {
 %cst = arith.constant 0.000000e+00 : f32
 affine.for %i = 0 to 256 {
   %final_sum, %final_expsum = affine.for %j = 0 to 500 iter_args(%sum_iter = %cst, %expsum_iter = %cst) -> (f32, f32) {
     %ld = affine.load %in[%i, %j] : memref<256x512xf32>
     %exp = math.exp %ld : f32
     %add = arith.addf %sum_iter, %ld : f32
     %eadd = arith.addf %expsum_iter, %exp : f32
     affine.yield %add, %eadd : f32, f32
   }
   affine.store %final_sum, %out[2*%i] : memref<512xf32>
   affine.store %final_expsum, %out[2*%i + 1] : memref<512xf32>
 }
 return
}

// CHECK:       #[[$map4:.*]] = affine_map<([[d0:.*]]) -> (-[[d0]] + 500)>
// CHECK-LABEL: @vecdim_two_reductions_masked
// CHECK:       affine.for %{{.*}} = 0 to 256 {
// CHECK:         %{{.*}} = affine.for %[[iv:.*]] = 0 to 500 step 128 iter_args(%[[sum_iter:.*]] = {{.*}}, %[[esum_iter:.*]] = {{.*}}) -> (vector<128xf32>, vector<128xf32>) {
// CHECK:           %[[elems_left:.*]] = affine.apply #[[$map4]](%[[iv]])
// CHECK:           %[[mask:.*]] = vector.create_mask %[[elems_left]] : vector<128xi1>
// CHECK:           %[[ld:.*]] = vector.transfer_read %{{.*}} : memref<256x512xf32>, vector<128xf32>
// CHECK:           %[[exp:.*]] = math.exp %[[ld]] : vector<128xf32>
// CHECK:           %[[add:.*]] = arith.addf %[[sum_iter]], %[[ld]] : vector<128xf32>
// CHECK:           %[[eadd:.*]] = arith.addf %[[esum_iter]], %[[exp]] : vector<128xf32>
// CHECK:           %[[new_acc:.*]] = select %[[mask]], %[[add]], %[[sum_iter]] : vector<128xi1>, vector<128xf32>
// CHECK:           %[[new_eacc:.*]] = select %[[mask]], %[[eadd]], %[[esum_iter]] : vector<128xi1>, vector<128xf32>
// CHECK:           affine.yield %[[new_acc]], %[[new_eacc]] : vector<128xf32>
// CHECK:         }