[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

538

conatins -> contains

556

reductioon -> reduction

567

tbale -> table

[razvanlupusoru]

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

614–615

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

588

Can you please also attach the ReductionOperator to the recipe?

[razvanlupusoru]

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

591–592

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

615

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 :)

623

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

615

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

615

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

615

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

615

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

548

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

548

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

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_i64 : i64 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_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 `init` $initRegion
+    `combiner` $combinerRegion
+  }];
+
+  let hasRegionVerifier = 1;
+}
+
+def ReductionOp : OpenACC_Op<"reduction"> {
+  let summary = "reduction op";
+
+  let description = [{
+    Indicates the value that is produced by the current reduction-participating
+    entity for a reduction requested in some ancestor.
+    The reduction is identified by the accumulator, but the value of the
+    accumulator may not be updated immediately.
+
+    Example:
+    ```
+    acc.parallel reduction(@reduction_i64 -> %a : !llvm.ptr<i64>) {
+      %1 = arith.constant 10 : i64
+      acc.reduction %1, %a : i64, !llvm.ptr<i64>
+      acc.yield
+    }
+    ```
+  }];
+
+  let arguments= (ins AnyType:$operand, OpenACC_PointerLikeType:$accumulator);
+
+  let assemblyFormat = [{
+    $operand `,` $accumulator attr-dict `:`
+        type($operand) `,` type($accumulator)
+  }];
+
+  let hasVerifier = 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
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,
-                                             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,
+  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();
 }
 
 //===----------------------------------------------------------------------===//
 // FirstprivateRecipeOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult acc::FirstprivateRecipeOp::verifyRegions() {
-  if (failed(verifyPrivateLikeRegion(*this, getInitRegion(), "init", getType(),
-                                     1, /*optional=*/false,
+  if (failed(verifyInitLikeSingleArgRegion(*this, getInitRegion(),
+                                           "privatization", "init", getType(),
                                            /*verifyYield=*/true)))
     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)))
+  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();
+
+  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();
+}
+
+LogicalResult acc::ReductionOp::verify() {
+  // TODO when actual (acc.parallel, acc.serial, acc.loop) operation support the
+  // new reduction design.
   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 init {
+} reduction {}
+
+// -----
+
+// expected-error@+1 {{expects init region with one argument of the reduction type}}
+acc.reduction.recipe @reduction_i64 : i64 init {
+^bb0(%0: i32):
+  %1 = arith.constant 0 : i64
+  acc.yield %1 : i64
+} reduction {}
+
+// -----
+
+// expected-error@+1 {{expects init region to yield a value of the reduction type}}
+acc.reduction.recipe @reduction_i64 : i64 init {
+^bb0(%0: i64):
+  %1 = arith.constant 0 : i32
+  acc.yield %1 : i32
+} reduction {}
+
+// -----
+
+// expected-error@+1 {{expects non-empty reduction region}}
+acc.reduction.recipe @reduction_i64 : i64 init {
+^bb0(%0: i64):
+  %1 = arith.constant 0 : i64
+  acc.yield %1 : i64
+} reduction {}
 
+// -----
+
+// expected-error@+1 {{expects reduction region with two arguments of the reduction type}}
+acc.reduction.recipe @reduction_i64 : i64 init {
+^bb0(%0: i64):
+  %1 = arith.constant 0 : i64
+  acc.yield %1 : i64
+} reduction {
+^bb0(%0: i32):
+  acc.yield %0 : i32
+}
+
+// -----
+
+// expected-error@+1 {{expects reduction region with two arguments of the reduction type}}
+acc.reduction.recipe @reduction_i64 : i64 init {
+^bb0(%0: i64):
+  %1 = arith.constant 0 : i64
+  acc.yield %1 : i64
+} reduction {
+^bb0(%0: i64):
+  acc.yield %0 : i64
+}
+
+// -----
+
+// expected-error@+1 {{expects reduction region to yield a value of the reduction type}}
+acc.reduction.recipe @reduction_i64 : i64 init {
+^bb0(%0: i64):
+  %1 = arith.constant 0 : i64
+  acc.yield %1 : i64
+} reduction {
+^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 init {
+^bb0(%0: i64):
+  %1 = arith.constant 0 : i64
+  acc.yield %1 : i64
+} reduction {
+^bb0(%0: i64, %1: i64):
+  %2 = arith.addi %0, %1 : i64
+  acc.yield %2 : i64
+}
+
+// CHECK-LABEL: acc.reduction.recipe @reduction_add_i64 : i64 init {
+// CHECK:       ^bb0(%{{.*}}: i64):
+// CHECK:         %[[C0:.*]] = arith.constant 0 : i64
+// CHECK:         acc.yield %[[C0]] : i64
+// CHECK:       } reduction {
+// CHECK:       ^bb0(%[[ARG0:.*]]: i64, %[[ARG1:.*]]: i64):
+// CHECK:         %[[RES:.*]] = arith.addi %[[ARG0]], %[[ARG1]] : i64
+// CHECK:         acc.yield %[[RES]] : i64
+// CHECK:       }