[mlir][Vector] Support 0-D vectors in ReductionOp

[guraypp]

Summary

This commit adds support for 0-D vectors in ReductionOp.

[dcaballe]

Thanks, Guray! LGTM!

mlir/include/mlir/Dialect/Vector/IR/VectorOps.td

270–271

nit: split line?

[guraypp]

Splitted the large line in two

[nicolasvasilache]

LGTM conditioned on trimming unnecessary tests.

mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir

1120

Can we add just 1 test ? The behavior you want to test here is already covered and replicating the same test for different elemental types is unnecessary.

mlir/test/Dialect/Vector/ops.mlir

549

same here re duplication plz

[guraypp]

trimmed the extra tests

mlir/include/mlir/Dialect/Vector/IR/VectorOps.td

def Vector_ContractionOp :
   let hasVerifier = 1;
 }
 
 def Vector_ReductionOp :
   Vector_Op<"reduction", [NoSideEffect,
      PredOpTrait<"source operand and result have same element type",
                  TCresVTEtIsSameAsOpBase<0, 0>>,
      DeclareOpInterfaceMethods<VectorUnrollOpInterface,
                                ["getShapeForUnroll"]>]>,
-    Arguments<(ins Vector_CombiningKindAttr:$kind, AnyVector:$vector,
-                   Optional<AnyType>:$acc)>,
+    Arguments<(ins Vector_CombiningKindAttr:$kind,
+               AnyVectorOfAnyRank:$vector,
     Results<(outs AnyType:$dest)> {
   let summary = "reduction operation";
   let description = [{
     Reduces an 1-D vector "horizontally" into a scalar using the given
     operation (add/mul/min/max for int/fp and and/or/xor for int only).
     Reductions also allow an optional fused accumulator.
 
     Note that these operations are restricted to 1-D vectors to remain

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

 
 void vector::ReductionOp::build(OpBuilder &builder, OperationState &result,
                                 CombiningKind kind, Value vector, Value acc) {
   build(builder, result, vector.getType().cast<VectorType>().getElementType(),
         kind, vector, acc);
 }
 
 LogicalResult ReductionOp::verify() {
-  // Verify for 1-D vector.
+  // Verify for 0-D and 1-D vector.
   int64_t rank = getVectorType().getRank();
-  if (rank != 1)
+  if (rank > 1)
     return emitOpError("unsupported reduction rank: ") << rank;
 
   // Verify supported reduction kind.
   Type eltType = getDest().getType();
   if (!isSupportedCombiningKind(getKind(), eltType))
     return emitOpError("unsupported reduction type '")
            << eltType << "' for kind '" << stringifyCombiningKind(getKind())
            << "'";

mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir

func.func @vector_fma(%a: vector<8xf32>, %b: vector<2x4xf32>, %c: vector<1x1x1xf32>) -> (vector<8xf32>, vector<2x4xf32>, vector<1x1x1xf32>) {
   //  CHECK-SAME:   (vector<1xf32>, vector<1xf32>, vector<1xf32>) -> vector<1xf32>
   %2 = vector.fma %c, %c, %c : vector<1x1x1xf32>
 
   return %0, %1, %2: vector<8xf32>, vector<2x4xf32>, vector<1x1x1xf32>
 }
 
 // -----
 
+func.func @reduce_0d_f32(%arg0: vector<f32>) -> f32 {
+  %0 = vector.reduction <add>, %arg0 : vector<f32> into f32
+  return %0 : f32
+}
+// CHECK-LABEL: @reduce_0d_f32(
+// CHECK-SAME: %[[A:.*]]: vector<f32>)
+//      CHECK: %[[CA:.*]] = builtin.unrealized_conversion_cast %[[A]] : vector<f32> to vector<1xf32>
+//      CHECK: %[[C:.*]] = llvm.mlir.constant(0.000000e+00 : f32) : f32
+//      CHECK: %[[V:.*]] = "llvm.intr.vector.reduce.fadd"(%[[C]], %[[CA]])
+// CHECK-SAME: {reassoc = false} : (f32, vector<1xf32>) -> f32
+//      CHECK: return %[[V]] : f32
+
+// -----
+
 func.func @reduce_f16(%arg0: vector<16xf16>) -> f16 {
   %0 = vector.reduction <add>, %arg0 : vector<16xf16> into f16
   return %0 : f16
 }
 // CHECK-LABEL: @reduce_f16(
 // CHECK-SAME: %[[A:.*]]: vector<16xf16>)
 //      CHECK: %[[C:.*]] = llvm.mlir.constant(0.000000e+00 : f16) : f16
 //      CHECK: %[[V:.*]] = "llvm.intr.vector.reduce.fadd"(%[[C]], %[[A]])

mlir/test/Dialect/Vector/ops.mlir

func.func @reduce_int(%arg0: vector<16xi32>) -> i32 {
   // CHECK:    vector.reduction <or>, %{{.*}} : vector<16xi32> into i32
   vector.reduction <or>, %arg0 : vector<16xi32> into i32
   // CHECK:    %[[X:.*]] = vector.reduction <xor>, %{{.*}} : vector<16xi32> into i32
   %0 = vector.reduction <xor>, %arg0 : vector<16xi32> into i32
   // CHECK:    return %[[X]] : i32
   return %0 : i32
 }
 
+// CHECK-LABEL: @reduce_int
+func.func @reduce_int_0d(%arg0: vector<i32>) -> i32 {
+  // CHECK:    vector.reduction <add>, %{{.*}} : vector<i32> into i32
+  vector.reduction <add>, %arg0 : vector<i32> into i32
+  // CHECK:    vector.reduction <mul>, %{{.*}} : vector<i32> into i32
+  vector.reduction <mul>, %arg0 : vector<i32> into i32
+  // CHECK:    vector.reduction <minui>, %{{.*}} : vector<i32> into i32
+  vector.reduction <minui>, %arg0 : vector<i32> into i32
+  // CHECK:    vector.reduction <minsi>, %{{.*}} : vector<i32> into i32
+  vector.reduction <minsi>, %arg0 : vector<i32> into i32
+  // CHECK:    vector.reduction <maxui>, %{{.*}} : vector<i32> into i32
+  vector.reduction <maxui>, %arg0 : vector<i32> into i32
+  // CHECK:    vector.reduction <maxsi>, %{{.*}} : vector<i32> into i32
+  vector.reduction <maxsi>, %arg0 : vector<i32> into i32
+  // CHECK:    vector.reduction <and>, %{{.*}} : vector<i32> into i32
+  vector.reduction <and>, %arg0 : vector<i32> into i32
+  // CHECK:    vector.reduction <or>, %{{.*}} : vector<i32> into i32
+  vector.reduction <or>, %arg0 : vector<i32> into i32
+  // CHECK:    %[[X:.*]] = vector.reduction <xor>, %{{.*}} : vector<i32> into i32
+  %0 = vector.reduction <xor>, %arg0 : vector<i32> into i32
+  // CHECK:    return %[[X]] : i32
+  return %0 : i32
+}
+
 // CHECK-LABEL: @transpose_fp
 func.func @transpose_fp(%arg0: vector<3x7xf32>) -> vector<7x3xf32> {
   // CHECK: %[[X:.*]] = vector.transpose %{{.*}}, [1, 0] : vector<3x7xf32> to vector<7x3xf32>
   %0 = vector.transpose %arg0, [1, 0] : vector<3x7xf32> to vector<7x3xf32>
   // CHECK: return %[[X]] : vector<7x3xf32>
   return %0 : vector<7x3xf32>
 }
 

mlir/test/Integration/Dialect/Vector/CPU/test-0-d-vectors.mlir

func.func @create_mask_0d(%zero : index, %one : index) {
 
   %one_mask = vector.create_mask %one : vector<i1>
   // CHECK: ( 1 )
   vector.print %one_mask : vector<i1>
 
   return
 }
 
+func.func @reduce_add(%arg0: vector<f32>) -> f32 {
+  %0 = vector.reduction <add>, %arg0 : vector<f32> into f32    
+  return %0 : f32
+}
+
 func.func @entry() {
   %0 = arith.constant 42.0 : f32
   %1 = arith.constant dense<0.0> : vector<f32>
   %2 = call  @insert_element_0d(%0, %1) : (f32, vector<f32>) -> (vector<f32>)
   call  @extract_element_0d(%2) : (vector<f32>) -> ()
 
   %3 = arith.constant dense<42.0> : vector<f32>
   call  @print_vector_0d(%3) : (vector<f32>) -> ()
func.func @entry() {
   %smaller = arith.constant dense<42> : vector<i32>
   %bigger = arith.constant dense<4242> : vector<i32>
   call  @arith_cmpi_0d(%smaller, %bigger) : (vector<i32>, vector<i32>) -> ()
 
   %zero_idx = arith.constant 0 : index
   %one_idx = arith.constant 1 : index
   call  @create_mask_0d(%zero_idx, %one_idx) : (index, index) -> ()
 
+  %red_array = arith.constant dense<5.0> : vector<f32>
+  %red_res = call  @reduce_add(%red_array) : (vector<f32>) -> (f32)  
+  vector.print %red_res : f32
+  // CHECK: 5
+
   return
 }