[mlir][openacc] Add reduction representation

[clementval]

Summary

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) {
}

[razvanlupusoru]

Just a few spelling errors. Some are also in the commit message.

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

540

conatins -> contains

558

reductioon -> reduction

569

tbale -> table

[razvanlupusoru]

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

616–617

Did you intend to use llvm.ptr<i64>? I would guess yes. In that case, it doesn't seem to match the signature of combiner region (unless I am misunderstanding something).

[razvanlupusoru]

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

590

Can you please also attach the ReductionOperator to the recipe?

[razvanlupusoru]

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

593–594

I feel that we should name this "combinerRegion". I worry with too many things named "reduction" (acc.reduction.recipe, acc.reduction), it makes it somewhat ambiguous.

[vzakhari]

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

617

Could you please clarify where this operation is coming from? What is the flow for generating it?

From the language we only have a reduction clause, so I am not sure why the reduction operation appears. I suppose it may appear later in the OpenACC lowering flow, for example, when a "thread-private" accumulator value needs to be added to the "global" accumulator. In other words, I am not following it :)

625

Can you please clarify the connection between the reduction and the reduction.recipe?

[clementval]

Address review comments

[clementval]

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

617

The idea is that this operation replace the actual operation that has the reduction participating variable. With the reduction operator attached to the recipe we could omit this operation and be able to track back operation that participate in the reduction during codegen and place the combiner when needed.

[vzakhari]

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

617

I see. Is this an optimization for skipping the combining, if there are no actual reduction updates in the region?

I suppose we will not be able to place the reduction operation always, for example:

#pragma acc parallel reduction(+: r)
  func(&r); // r is updated inside func

[clementval]

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

617

Yes you are correct. I would suggest to not add this operation for the moment as it is not strictly necessary.

[vzakhari]

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

617

Yes, I think we'd better add it separately, when needed. Thank you!

[clementval]

• Add reduction_operator
• Remove acc.reduction

[vzakhari]

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

550

Please update the example with the reduction operator representation. Otherwise, it looks great! Thank you!

[razvanlupusoru]

I agree with Slava about updating example. Otherwise looks great in this form. Thank you.

[clementval]

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

550

My bad, I missed it. I'll update it.

[clementval]

Update example

[clementval]

Updated. Thanks for the reviews!

[clementval]

clang-format

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

 // AccCommon requires definition of OpenACC_Dialect.
 include "mlir/Dialect/OpenACC/AccCommon.td"
 
 // Base class for OpenACC dialect ops.
 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",
     [OpenACC_ReductionOperatorAdd, OpenACC_ReductionOperatorMul,
      OpenACC_ReductionOperatorMax, OpenACC_ReductionOperatorMin,
      OpenACC_ReductionOperatorAnd, OpenACC_ReductionOperatorOr,
      OpenACC_ReductionOperatorXor, OpenACC_ReductionOperatorLogEqv,
      OpenACC_ReductionOperatorLogNeqv, OpenACC_ReductionOperatorLogAnd,
      OpenACC_ReductionOperatorLogOr
     ]> {
   let genSpecializedAttr = 0;
   let cppNamespace = "::mlir::acc";
 }
 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]>;
 
 // Simple alias to pointer-like interface to reduce verbosity.
 def OpenACC_PointerLikeType : TypeAlias<OpenACC_PointerLikeTypeInterface,
   "pointer-like type">;
 
let assemblyFormat = [{
     `copy` $copyRegion
     (`destroy` $destroyRegion^)?
   }];
 
   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
 //===----------------------------------------------------------------------===//
 
 def OpenACC_ParallelOp : OpenACC_Op<"parallel",
     [AttrSizedOperandSegments, RecursiveMemoryEffects]> {
   let summary = "parallel construct";
   let description = [{
     The "acc.parallel" operation represents a parallel construct block. It has
     one region to be executed in parallel on the current device.
 
     Example:
 
     ```mlir
     acc.parallel num_gangs(%c10) num_workers(%c10)
         private(%c : memref<10xf32>) {
       // parallel region
     }
     ```
   }];
 
   let arguments = (ins Optional<IntOrIndex>:$async,
                        UnitAttr:$asyncAttr,
                        Variadic<IntOrIndex>:$waitOperands,
                        UnitAttr:$waitAttr,
                        Optional<IntOrIndex>:$numGangs,
                        Optional<IntOrIndex>:$numWorkers,
                        Optional<IntOrIndex>:$vectorLength,
let assemblyFormat = [{
     attr-dict-with-keyword
   }];
 
   let hasVerifier = 1;
 }
 
 // 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 = [{
     `acc.yield` is a special terminator operation for block inside regions in
     acc ops (parallel and loop). It returns values to the immediately enclosing
     acc op.
   }];
 

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

 };
 
 } // namespace
 
 //===----------------------------------------------------------------------===//
 // PrivateRecipeOp
 //===----------------------------------------------------------------------===//
 
-static LogicalResult verifyPrivateLikeRegion(Operation *op, Region &region,
+
+static LogicalResult verifyInitLikeSingleArgRegion(Operation *op, Region &region,
+                                                   StringRef regionType,
                                                    StringRef regionName, Type type,
-                                             unsigned expectNbArg,
-                                             bool optionalRegion,
-                                             bool verifyYield) {
-  if (optionalRegion && region.empty())
+                                                   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";
+        return op->emitOpError() << "expects " << regionName << " region to "
+                                    "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)))
-    return failure();
-  if (failed(verifyPrivateLikeRegion(*this, getDestroyRegion(), "destroy",
-                                     getType(), 1, /*optional=*/true,
-                                     /*verifyYield=*/false)))
+  if (failed(verifyInitLikeSingleArgRegion(*this, getInitRegion(), "privatization", "init", 
+                                           getType(), /*verifyYield=*/true)))
+    return failure();
+  if (failed(verifyInitLikeSingleArgRegion(*this, getDestroyRegion(), "privatization", "destroy",
+                                           getType(), /*verifyYield=*/false, /*optional=*/true)))
     return failure();
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // FirstprivateRecipeOp
 //===----------------------------------------------------------------------===//
 
 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 (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, getCopyRegion(), "copy", getType(),
-                                     2, /*optional=*/false,
-                                     /*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();
 }
 
 //===----------------------------------------------------------------------===//
 // ParallelOp
 //===----------------------------------------------------------------------===//
 
 /// Check dataOperands for acc.parallel, acc.serial and acc.kernels.

mlir/test/Dialect/OpenACC/invalid.mlir

 ^bb0(%arg0 : i32, %arg1 : !llvm.ptr<i32>):
   llvm.store %arg0, %arg1 : !llvm.ptr<i32>
   acc.yield
 } destroy {
 ^bb0(%arg0 : f32):
   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

 // CHECK:   %[[UNDEF:.*]] = llvm.mlir.undef : !llvm.struct<(i32, i32)>
 // CHECK:   %[[UNDEF1:.*]] = llvm.insertvalue %[[C1]], %[[UNDEF]][0] : !llvm.struct<(i32, i32)> 
 // CHECK:   %[[UNDEF2:.*]] = llvm.insertvalue %[[C1]], %[[UNDEF1]][1] : !llvm.struct<(i32, i32)> 
 // CHECK:   acc.yield %[[UNDEF2]] : !llvm.struct<(i32, i32)>
 // CHECK: } destroy {
 // 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:       }