[mlir][openacc] Add reduction representation

Similarly to D150622 for private clause, the reduction is currently not
modeled in a good way. This patch is inspired by the reduction representation
in the omp dialect (D105358) and make a new representation for the reduction in
the OpenACC dialect.

A new operation is introduced to model the sequences of operation needed to
initialize a local reduction value and how to combine two values during the
reduction. The operation requires two mandatory regions.

  1. The init region specifies how to initialize the local reduction
     value. The region has an argument that contains the value of the
     reduction accumulator at the start of the reduction. It is expected to
     `acc.yield` the new value.
  2. The reduction region contains a sequences of operations to combine two
     values of the reduction type into one. It has two arguments and it is
     expected to `acc.yield` the combined value.

Example:

```mlir
acc.reduction.recipe @reduction_add_i64 : i64 init reduction_operator<add> {
^bb0(%0: i64):
  // init region contains a sequence of operations to initialize the local
  // reduction value as specified in 2.5.15
  %c0 = arith.constant 0 : i64
  acc.yield %c0 : i64
} reduction {
^bb0(%0: i64, %1: i64)
  // reduction region contains a sequence of operations to combine
  // two values into one.
  %2 = arith.addi %0, %1 : i64
  acc.yield %2 : i64
}

// The reduction symbol is then used in the corresponding operation.
acc.parallel reduction(@reduction_add_i64 -> %a : i64) {
}

Reviewed By: razvanlupusoru, vzakhari

Differential Revision: https://reviews.llvm.org/D150818

mlir/include/mlir/Dialect/OpenACC/OpenACCOps.td

@@ -31,17 +31,17 @@ class OpenACC_Op<string mnemonic, list<Trait> traits = []> :
   Op<OpenACC_Dialect, mnemonic, traits>;
 
 // Reduction operation enumeration.
-def OpenACC_ReductionOperatorAdd     : I32EnumAttrCase<"redop_add", 0>;
-def OpenACC_ReductionOperatorMul     : I32EnumAttrCase<"redop_mul", 1>;
-def OpenACC_ReductionOperatorMax     : I32EnumAttrCase<"redop_max", 2>;
-def OpenACC_ReductionOperatorMin     : I32EnumAttrCase<"redop_min", 3>;
-def OpenACC_ReductionOperatorAnd     : I32EnumAttrCase<"redop_and", 4>;
-def OpenACC_ReductionOperatorOr      : I32EnumAttrCase<"redop_or", 5>;
-def OpenACC_ReductionOperatorXor     : I32EnumAttrCase<"redop_xor", 6>;
-def OpenACC_ReductionOperatorLogEqv  : I32EnumAttrCase<"redop_leqv", 7>;
-def OpenACC_ReductionOperatorLogNeqv : I32EnumAttrCase<"redop_lneqv", 8>;
-def OpenACC_ReductionOperatorLogAnd  : I32EnumAttrCase<"redop_land", 9>;
-def OpenACC_ReductionOperatorLogOr   : I32EnumAttrCase<"redop_lor", 10>;
+def OpenACC_ReductionOperatorAdd     : I32EnumAttrCase<"AccAdd", 0, "add">;
+def OpenACC_ReductionOperatorMul     : I32EnumAttrCase<"AccMul", 1, "mul">;
+def OpenACC_ReductionOperatorMax     : I32EnumAttrCase<"AccMax", 2, "max">;
+def OpenACC_ReductionOperatorMin     : I32EnumAttrCase<"AccMin", 3, "min">;
+def OpenACC_ReductionOperatorAnd     : I32EnumAttrCase<"AccIand", 4, "iand">;
+def OpenACC_ReductionOperatorOr      : I32EnumAttrCase<"AccIor", 5, "ior">;
+def OpenACC_ReductionOperatorXor     : I32EnumAttrCase<"AccXor", 6, "xor">;
+def OpenACC_ReductionOperatorLogEqv  : I32EnumAttrCase<"AccEqv", 7, "eqv">;
+def OpenACC_ReductionOperatorLogNeqv : I32EnumAttrCase<"AccNeqv", 8, "neqv">;
+def OpenACC_ReductionOperatorLogAnd  : I32EnumAttrCase<"AccLand", 9, "land">;
+def OpenACC_ReductionOperatorLogOr   : I32EnumAttrCase<"AccLor", 10, "lor">;
 
 def OpenACC_ReductionOperator : I32EnumAttr<"ReductionOperator",
     "built-in reduction operations supported by OpenACC",
@@ -57,7 +57,9 @@ def OpenACC_ReductionOperator : I32EnumAttr<"ReductionOperator",
 }
 def OpenACC_ReductionOperatorAttr : EnumAttr<OpenACC_Dialect,
                                              OpenACC_ReductionOperator,
-                                             "reduction_operator">;
+                                             "reduction_operator"> {
+  let assemblyFormat = [{ ```<` $value `>` }];
+}
 
 // Type used in operation below.
 def IntOrIndex : AnyTypeOf<[AnyInteger, Index]>;
@@ -522,6 +524,85 @@ def OpenACC_FirstprivateRecipeOp : OpenACC_Op<"firstprivate.recipe",
   let hasRegionVerifier = 1;
 }
 
+//===----------------------------------------------------------------------===//
+// 2.5.15 reduction clause
+//===----------------------------------------------------------------------===//
+
+def OpenACC_ReductionRecipeOp : OpenACC_Op<"reduction.recipe",
+    [IsolatedFromAbove, Symbol]> {
+  let summary = "reduction recipe";
+
+  let description = [{
+    Declares an OpenACC reduction recipe. The operation requires two
+    mandatory regions.
+
+      1. The initializer region specifies how to initialize the local reduction
+         value. The region has an argument that contains the value of the
+         reduction accumulator at the start of the reduction. It is expected to
+         `acc.yield` the new value.
+      2. The reduction region contains a sequences of operations to combine two
+         values of the reduction type into one. It has two arguments and it is
+         expected to `acc.yield` the combined value.
+
+    Example:
+
+    ```mlir
+    acc.reduction.recipe @reduction_add_i64 : i64 reduction_operator<add> init {
+    ^bb0(%0: i64):
+      // init region contains a sequence of operations to initialize the local
+      // reduction value as specified in 2.5.15
+      %c0 = arith.constant 0 : i64
+      acc.yield %c0 : i64
+    } combiner {
+    ^bb0(%0: i64, %1: i64)
+      // combiner region contains a sequence of operations to combine
+      // two values into one.
+      %2 = arith.addi %0, %1 : i64
+      acc.yield %2 : i64
+    }
+
+    // The reduction symbol is then used in the corresponding operation.
+    acc.parallel reduction(@reduction_add_i64 -> %a : i64) {
+    }
+    ```
+
+    The following table lists the valid operators and the initialization values
+    according to OpenACC 3.3:
+
+    |------------------------------------------------|
+    |        C/C++          |        Fortran         |
+    |-----------------------|------------------------|
+    | operator | init value | operator | init value  |
+    |     +    |      0     |     +    |      0      |
+    |     *    |      1     |     *    |      1      |
+    |    max   |    least   |    max   |    least    |
+    |    min   |   largest  |    min   |   largest   |
+    |     &    |     ~0     |   iand   | all bits on |
+    |     |    |      0     |    ior   |      0      |
+    |     ^    |      0     |   ieor   |      0      |
+    |    &&    |      1     |   .and.  |    .true.   |
+    |    ||    |      0     |    .or.  |   .false.   |
+    |          |            |   .eqv.  |    .true.   |
+    |          |            |  .neqv.  |   .false.   |
+    -------------------------------------------------|
+  }];
+
+  let arguments = (ins SymbolNameAttr:$sym_name,
+                       TypeAttr:$type,
+                       OpenACC_ReductionOperatorAttr:$reductionOperator);
+
+  let regions = (region AnyRegion:$initRegion,
+                        AnyRegion:$combinerRegion);
+
+  let assemblyFormat = [{
+    $sym_name `:` $type attr-dict-with-keyword
+    `reduction_operator` $reductionOperator
+    `init` $initRegion `combiner` $combinerRegion
+  }];
+
+  let hasRegionVerifier = 1;
+}
+
 //===----------------------------------------------------------------------===//
 // 2.5.1 parallel Construct
 //===----------------------------------------------------------------------===//
@@ -986,7 +1067,7 @@ def OpenACC_LoopOp : OpenACC_Op<"loop",
 
 // Yield operation for the acc.loop and acc.parallel operations.
 def OpenACC_YieldOp : OpenACC_Op<"yield", [ReturnLike, Terminator,
-    ParentOneOf<["ParallelOp, LoopOp, SerialOp, PrivateRecipeOp, FirstprivateRecipeOp"]>]> {
+    ParentOneOf<["FirstprivateRecipeOp, LoopOp, ParallelOp, PrivateRecipeOp, ReductionRecipeOp, SerialOp"]>]> {
   let summary = "Acc yield and termination operation";
 
   let description = [{

mlir/lib/Dialect/OpenACC/IR/OpenACC.cpp

@@ -338,48 +338,43 @@ struct RemoveConstantIfConditionWithRegion : public OpRewritePattern<OpTy> {
 // PrivateRecipeOp
 //===----------------------------------------------------------------------===//
 
-static LogicalResult verifyPrivateLikeRegion(Operation *op, Region &region,
-                                             StringRef regionName, Type type,
-                                             unsigned expectNbArg,
-                                             bool optionalRegion,
-                                             bool verifyYield) {
-  if (optionalRegion && region.empty())
+static LogicalResult verifyInitLikeSingleArgRegion(
+    Operation *op, Region &region, StringRef regionType, StringRef regionName,
+    Type type, bool verifyYield, bool optional = false) {
+  if (optional && region.empty())
     return success();
 
   if (region.empty())
     return op->emitOpError() << "expects non-empty " << regionName << " region";
   Block &firstBlock = region.front();
-  if (expectNbArg == 1 && (firstBlock.getNumArguments() != 1 ||
-                           firstBlock.getArgument(0).getType() != type))
+  if (firstBlock.getNumArguments() != 1 ||
+      firstBlock.getArgument(0).getType() != type)
     return op->emitOpError() << "expects " << regionName
                              << " region with one "
-                                "argument of the privatization type";
-  if (expectNbArg == 2 && (firstBlock.getNumArguments() != 2 ||
-                           firstBlock.getArgument(0).getType() != type))
-    return op->emitOpError() << "expects " << regionName
-                             << " region with two "
-                                "arguments of the privatization type";
+                                "argument of the "
+                             << regionType << " type";
 
   if (verifyYield) {
     for (YieldOp yieldOp : region.getOps<acc::YieldOp>()) {
       if (yieldOp.getOperands().size() != 1 ||
           yieldOp.getOperands().getTypes()[0] != type)
         return op->emitOpError() << "expects " << regionName
                                  << " region to "
-                                    "yield a value of the privatization type";
+                                    "yield a value of the "
+                                 << regionType << " type";
     }
   }
   return success();
 }
 
 LogicalResult acc::PrivateRecipeOp::verifyRegions() {
-  if (failed(verifyPrivateLikeRegion(*this, getInitRegion(), "init", getType(),
-                                     1, /*optional=*/false,
-                                     /*verifyYield=*/true)))
+  if (failed(verifyInitLikeSingleArgRegion(*this, getInitRegion(),
+                                           "privatization", "init", getType(),
+                                           /*verifyYield=*/true)))
     return failure();
-  if (failed(verifyPrivateLikeRegion(*this, getDestroyRegion(), "destroy",
-                                     getType(), 1, /*optional=*/true,
-                                     /*verifyYield=*/false)))
+  if (failed(verifyInitLikeSingleArgRegion(
+          *this, getDestroyRegion(), "privatization", "destroy", getType(),
+          /*verifyYield=*/false, /*optional=*/true)))
     return failure();
   return success();
 }
@@ -389,19 +384,55 @@ LogicalResult acc::PrivateRecipeOp::verifyRegions() {
 //===----------------------------------------------------------------------===//
 
 LogicalResult acc::FirstprivateRecipeOp::verifyRegions() {
-  if (failed(verifyPrivateLikeRegion(*this, getInitRegion(), "init", getType(),
-                                     1, /*optional=*/false,
-                                     /*verifyYield=*/true)))
+  if (failed(verifyInitLikeSingleArgRegion(*this, getInitRegion(),
+                                           "privatization", "init", getType(),
+                                           /*verifyYield=*/true)))
     return failure();
 
-  if (failed(verifyPrivateLikeRegion(*this, getCopyRegion(), "copy", getType(),
-                                     2, /*optional=*/false,
-                                     /*verifyYield=*/false)))
+  if (getCopyRegion().empty())
+    return emitOpError() << "expects non-empty copy region";
+
+  Block &firstBlock = getCopyRegion().front();
+  if (firstBlock.getNumArguments() != 2 ||
+      firstBlock.getArgument(0).getType() != getType())
+    return emitOpError() << "expects copy region with two arguments of the "
+                            "privatization type";
+
+  if (failed(verifyInitLikeSingleArgRegion(*this, getDestroyRegion(),
+                                           "privatization", "destroy",
+                                           getType(), /*verifyYield=*/false)))
     return failure();
-  if (failed(verifyPrivateLikeRegion(*this, getDestroyRegion(), "destroy",
-                                     getType(), 1, /*optional=*/true,
-                                     /*verifyYield=*/false)))
+
+  return success();
+}
+
+//===----------------------------------------------------------------------===//
+// ReductionRecipeOp
+//===----------------------------------------------------------------------===//
+
+LogicalResult acc::ReductionRecipeOp::verifyRegions() {
+  if (failed(verifyInitLikeSingleArgRegion(*this, getInitRegion(), "reduction",
+                                           "init", getType(),
+                                           /*verifyYield=*/true)))
     return failure();
+
+  if (getCombinerRegion().empty())
+    return emitOpError() << "expects non-empty combiner region";
+
+  Block &reductionBlock = getCombinerRegion().front();
+  if (reductionBlock.getNumArguments() != 2 ||
+      reductionBlock.getArgument(0).getType() != getType() ||
+      reductionBlock.getArgument(1).getType() != getType())
+    return emitOpError() << "expects combiner region with two arguments of "
+                         << "the reduction type";
+
+  for (YieldOp yieldOp : getCombinerRegion().getOps<YieldOp>()) {
+    if (yieldOp.getOperands().size() != 1 ||
+        yieldOp.getOperands().getTypes()[0] != getType())
+      return emitOpError() << "expects combiner region to yield a value "
+                              "of the reduction type";
+  }
+
   return success();
 }
 

mlir/test/Dialect/OpenACC/invalid.mlir

@@ -393,4 +393,72 @@ acc.firstprivate.recipe @privatization_i32 : i32 init {
   acc.yield
 }
 
+// -----
+
+// expected-error@+1 {{expects non-empty init region}}
+acc.reduction.recipe @reduction_i64 : i64 reduction_operator<add> init {
+} combiner {}
+
+// -----
+
+// expected-error@+1 {{expects init region with one argument of the reduction type}}
+acc.reduction.recipe @reduction_i64 : i64 reduction_operator<add> init {
+^bb0(%0: i32):
+  %1 = arith.constant 0 : i64
+  acc.yield %1 : i64
+} combiner {}
+
+// -----
+
+// expected-error@+1 {{expects init region to yield a value of the reduction type}}
+acc.reduction.recipe @reduction_i64 : i64 reduction_operator<add> init {
+^bb0(%0: i64):
+  %1 = arith.constant 0 : i32
+  acc.yield %1 : i32
+} combiner {}
+
+// -----
+
+// expected-error@+1 {{expects non-empty combiner region}}
+acc.reduction.recipe @reduction_i64 : i64 reduction_operator<add> init {
+^bb0(%0: i64):
+  %1 = arith.constant 0 : i64
+  acc.yield %1 : i64
+} combiner {}
 
+// -----
+
+// expected-error@+1 {{expects combiner region with two arguments of the reduction type}}
+acc.reduction.recipe @reduction_i64 : i64 reduction_operator<add> init {
+^bb0(%0: i64):
+  %1 = arith.constant 0 : i64
+  acc.yield %1 : i64
+} combiner {
+^bb0(%0: i32):
+  acc.yield %0 : i32
+}
+
+// -----
+
+// expected-error@+1 {{expects combiner region with two arguments of the reduction type}}
+acc.reduction.recipe @reduction_i64 : i64 reduction_operator<add> init {
+^bb0(%0: i64):
+  %1 = arith.constant 0 : i64
+  acc.yield %1 : i64
+} combiner {
+^bb0(%0: i64):
+  acc.yield %0 : i64
+}
+
+// -----
+
+// expected-error@+1 {{expects combiner region to yield a value of the reduction type}}
+acc.reduction.recipe @reduction_i64 : i64 reduction_operator<add> init {
+^bb0(%0: i64):
+  %1 = arith.constant 0 : i64
+  acc.yield %1 : i64
+} combiner {
+^bb0(%0: i64, %1: i64):
+  %2 = arith.constant 0 : i32
+  acc.yield %2 : i32
+}

mlir/test/Dialect/OpenACC/ops.mlir

@@ -1335,3 +1335,25 @@ acc.private.recipe @privatization_struct_i32_i64 : !llvm.struct<(i32, i32)> init
 // CHECK: ^bb0(%[[ARG0:.*]]: !llvm.struct<(i32, i32)>):
 // CHECK:   func.call @destroy_struct(%[[ARG0]]) : (!llvm.struct<(i32, i32)>) -> ()
 // CHECK:   acc.terminator
+
+// -----
+
+acc.reduction.recipe @reduction_add_i64 : i64 reduction_operator<add> init {
+^bb0(%0: i64):
+  %1 = arith.constant 0 : i64
+  acc.yield %1 : i64
+} combiner {
+^bb0(%0: i64, %1: i64):
+  %2 = arith.addi %0, %1 : i64
+  acc.yield %2 : i64
+}
+
+// CHECK-LABEL: acc.reduction.recipe @reduction_add_i64 : i64 reduction_operator <add> init {
+// CHECK:       ^bb0(%{{.*}}: i64):
+// CHECK:         %[[C0:.*]] = arith.constant 0 : i64
+// CHECK:         acc.yield %[[C0]] : i64
+// CHECK:       } combiner {
+// CHECK:       ^bb0(%[[ARG0:.*]]: i64, %[[ARG1:.*]]: i64):
+// CHECK:         %[[RES:.*]] = arith.addi %[[ARG0]], %[[ARG1]] : i64
+// CHECK:         acc.yield %[[RES]] : i64
+// CHECK:       }