[mlir][openacc] Use new reduction design in acc.parallel

[clementval]

Summary

After D150818 the reduction clause is represented
with a acc.reduction.recipe operation and an operand.
This patch updates the acc.parallel op for the new design.

[jeanPerier]

Looks great!

[razvanlupusoru]

Other than the name change, this looks great to me.

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

638–640

Same feedback as on another change - I would prefer this list being called reductionRecipes.

[clementval]

Rename reductions to reductionRecipes

[razvanlupusoru]

This is great. Thank you for the renaming.

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

let arguments = (ins Optional<IntOrIndex>:$async,
                        Variadic<IntOrIndex>:$waitOperands,
                        UnitAttr:$waitAttr,
                        Optional<IntOrIndex>:$numGangs,
                        Optional<IntOrIndex>:$numWorkers,
                        Optional<IntOrIndex>:$vectorLength,
                        Optional<I1>:$ifCond,
                        Optional<I1>:$selfCond,
                        UnitAttr:$selfAttr,
-                       OptionalAttr<OpenACC_ReductionOperatorAttr>:$reductionOp,
                        Variadic<AnyType>:$reductionOperands,
+                       OptionalAttr<SymbolRefArrayAttr>:$reductionRecipes,
                        Variadic<OpenACC_PointerLikeTypeInterface>:$gangPrivateOperands,
                        OptionalAttr<SymbolRefArrayAttr>:$privatizations,
                        Variadic<AnyType>:$gangFirstPrivateOperands,
                        Variadic<OpenACC_PointerLikeTypeInterface>:$dataClauseOperands,
                        OptionalAttr<DefaultValueAttr>:$defaultAttr);
 
   let regions = (region AnyRegion:$region);
 
   let extraClassDeclaration = [{
     /// The number of data operands.
     unsigned getNumDataOperands();
 
     /// The i-th data operand passed.
     Value getDataOperand(unsigned i);
   }];
 
   let assemblyFormat = [{
     oilist(
         `dataOperands` `(` $dataClauseOperands `:` type($dataClauseOperands) `)`
       | `async` `(` $async `:` type($async) `)`
       | `firstprivate` `(` $gangFirstPrivateOperands `:`
             type($gangFirstPrivateOperands) `)`
       | `num_gangs` `(` $numGangs `:` type($numGangs) `)`
       | `num_workers` `(` $numWorkers `:` type($numWorkers) `)`
-      | `private` `(` custom<PrivatizationList>(
+      | `private` `(` custom<SymOperandList>(
             $gangPrivateOperands, type($gangPrivateOperands), $privatizations)
         `)`
       | `vector_length` `(` $vectorLength `:` type($vectorLength) `)`
       | `wait` `(` $waitOperands `:` type($waitOperands) `)`
       | `self` `(` $selfCond `)`
       | `if` `(` $ifCond `)`
-      | `reduction` `(` $reductionOperands `:` type($reductionOperands) `)`
+      | `reduction` `(` custom<SymOperandList>(
+            $reductionOperands, type($reductionOperands), $reductionRecipes)
+        `)`
     )
     $region attr-dict-with-keyword
   }];
 
   let hasVerifier = 1;
 }
 
 //===----------------------------------------------------------------------===//
def OpenACC_SerialOp : OpenACC_Op<"serial",
   }];
 
   let assemblyFormat = [{
     oilist(
         `dataOperands` `(` $dataClauseOperands `:` type($dataClauseOperands) `)`
       | `async` `(` $async `:` type($async) `)`
       | `firstprivate` `(` $gangFirstPrivateOperands `:`
             type($gangFirstPrivateOperands) `)`
-      | `private` `(` custom<PrivatizationList>(
+      | `private` `(` custom<SymOperandList>(
             $gangPrivateOperands, type($gangPrivateOperands), $privatizations)
         `)`
       | `wait` `(` $waitOperands `:` type($waitOperands) `)`
       | `self` `(` $selfCond `)`
       | `if` `(` $ifCond `)`
       | `reduction` `(` $reductionOperands `:` type($reductionOperands) `)`
     )
     $region attr-dict-with-keyword
def OpenACC_LoopOp : OpenACC_Op<"loop",
   }];
 
   let hasCustomAssemblyFormat = 1;
   let assemblyFormat = [{
     oilist(
         `gang` `` custom<GangClause>($gangNum, type($gangNum), $gangStatic, type($gangStatic), $hasGang)
       | `worker` `` custom<WorkerClause>($workerNum, type($workerNum), $hasWorker)
       | `vector` `` custom<VectorClause>($vectorLength, type($vectorLength), $hasVector)
-      | `private` `(` custom<PrivatizationList>(
+      | `private` `(` custom<SymOperandList>(
             $privateOperands, type($privateOperands), $privatizations)
         `)`
       | `tile` `(` $tileOperands `:` type($tileOperands) `)`
       | `reduction` `(` $reductionOperands `:` type($reductionOperands) `)`
     )
     $region
     ( `(` type($results)^ `)` )?
     attr-dict-with-keyword

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

LogicalResult acc::ReductionRecipeOp::verifyRegions() {
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // Custom parser and printer verifier for private clause
 //===----------------------------------------------------------------------===//
 
-static ParseResult parsePrivatizationList(
+static ParseResult parseSymOperandList(
     mlir::OpAsmParser &parser,
     llvm::SmallVectorImpl<mlir::OpAsmParser::UnresolvedOperand> &operands,
-    llvm::SmallVectorImpl<Type> &types, mlir::ArrayAttr &privatizationSymbols) {
-  llvm::SmallVector<SymbolRefAttr> privatizationVec;
+    llvm::SmallVectorImpl<Type> &types, mlir::ArrayAttr &symbols) {
+  llvm::SmallVector<SymbolRefAttr> attributes;
   if (failed(parser.parseCommaSeparatedList([&]() {
-        if (parser.parseAttribute(privatizationVec.emplace_back()) ||
+        if (parser.parseAttribute(attributes.emplace_back()) ||
             parser.parseArrow() ||
             parser.parseOperand(operands.emplace_back()) ||
             parser.parseColonType(types.emplace_back()))
           return failure();
         return success();
       })))
     return failure();
-  llvm::SmallVector<mlir::Attribute> privatizations(privatizationVec.begin(),
-                                                    privatizationVec.end());
-  privatizationSymbols = ArrayAttr::get(parser.getContext(), privatizations);
+  llvm::SmallVector<mlir::Attribute> arrayAttr(attributes.begin(),
+                                               attributes.end());
+  symbols = ArrayAttr::get(parser.getContext(), arrayAttr);
   return success();
 }
 
-static void
-printPrivatizationList(mlir::OpAsmPrinter &p, mlir::Operation *op,
-                       mlir::OperandRange privateOperands,
-                       mlir::TypeRange privateTypes,
-                       std::optional<mlir::ArrayAttr> privatizations) {
-  for (unsigned i = 0, e = privatizations->size(); i < e; ++i) {
+static void printSymOperandList(mlir::OpAsmPrinter &p, mlir::Operation *op,
+                                mlir::OperandRange operands,
+                                mlir::TypeRange types,
+                                std::optional<mlir::ArrayAttr> attributes) {
+  for (unsigned i = 0, e = attributes->size(); i < e; ++i) {
     if (i != 0)
       p << ", ";
-    p << (*privatizations)[i] << " -> " << privateOperands[i] << " : "
-      << privateOperands[i].getType();
+    p << (*attributes)[i] << " -> " << operands[i] << " : "
+      << operands[i].getType();
   }
 }
 
 //===----------------------------------------------------------------------===//
 // ParallelOp
 //===----------------------------------------------------------------------===//
 
 /// Check dataOperands for acc.parallel, acc.serial and acc.kernels.
 template <typename Op>
 static LogicalResult checkDataOperands(Op op,
                                        const mlir::ValueRange &operands) {
   for (mlir::Value operand : operands)
     if (!mlir::isa<acc::AttachOp, acc::CopyinOp, acc::CopyoutOp, acc::CreateOp,
                    acc::DeleteOp, acc::DetachOp, acc::DevicePtrOp,
                    acc::GetDevicePtrOp, acc::NoCreateOp, acc::PresentOp>(
             operand.getDefiningOp()))
       return op.emitError(
           "expect data entry/exit operation or acc.getdeviceptr "
           "as defining op");
   return success();
 }
 
+template <typename Op>
 static LogicalResult
-checkPrivatizationList(Operation *op,
-                       std::optional<mlir::ArrayAttr> privatizations,
-                       mlir::OperandRange privateOperands) {
-  if (!privateOperands.empty()) {
-    if (!privatizations || privatizations->size() != privateOperands.size())
-      return op->emitOpError() << "expected as many privatizations symbol "
-                                  "reference as private operands";
+checkSymOperandList(Operation *op, std::optional<mlir::ArrayAttr> attributes,
+                    mlir::OperandRange operands, llvm::StringRef operandName,
+                    llvm::StringRef symbolName) {
+  if (!operands.empty()) {
+    if (!attributes || attributes->size() != operands.size())
+      return op->emitOpError()
+             << "expected as many " << symbolName << " symbol reference as "
+             << operandName << " operands";
   } else {
-    if (privatizations)
-      return op->emitOpError() << "unexpected privatizations symbol reference";
+    if (attributes)
+      return op->emitOpError()
+             << "unexpected " << symbolName << " symbol reference";
     return success();
   }
 
-  llvm::DenseSet<Value> privates;
-  for (auto args : llvm::zip(privateOperands, *privatizations)) {
-    mlir::Value privateOperand = std::get<0>(args);
+  llvm::DenseSet<Value> set;
+  for (auto args : llvm::zip(operands, *attributes)) {
+    mlir::Value operand = std::get<0>(args);
 
-    if (!privates.insert(privateOperand).second)
-      return op->emitOpError() << "private operand appears more than once";
+    if (!set.insert(operand).second)
+      return op->emitOpError()
+             << operandName << " operand appears more than once";
 
-    mlir::Type varType = privateOperand.getType();
+    mlir::Type varType = operand.getType();
     auto symbolRef = std::get<1>(args).cast<SymbolRefAttr>();
-    auto decl =
-        SymbolTable::lookupNearestSymbolFrom<PrivateRecipeOp>(op, symbolRef);
+    auto decl = SymbolTable::lookupNearestSymbolFrom<Op>(op, symbolRef);
     if (!decl)
-      return op->emitOpError() << "expected symbol reference " << symbolRef
-                               << " to point to a private declaration";
+      return op->emitOpError()
+             << "expected symbol reference " << symbolRef << " to point to a "
+             << operandName << " declaration";
 
     if (decl.getType() && decl.getType() != varType)
       return op->emitOpError()
-             << "expected private (" << varType
-             << ") to be the same type as private declaration ("
-             << decl.getType() << ")";
+             << "expected private (" << varType << ") to be the same type as "
+             << operandName << " declaration (" << decl.getType() << ")";
   }
 
   return success();
 }
 
 unsigned ParallelOp::getNumDataOperands() {
   return getReductionOperands().size() + getGangPrivateOperands().size() +
          getGangFirstPrivateOperands().size() + getDataClauseOperands().size();
 }
 
 Value ParallelOp::getDataOperand(unsigned i) {
   unsigned numOptional = getAsync() ? 1 : 0;
   numOptional += getNumGangs() ? 1 : 0;
   numOptional += getNumWorkers() ? 1 : 0;
   numOptional += getVectorLength() ? 1 : 0;
   numOptional += getIfCond() ? 1 : 0;
   numOptional += getSelfCond() ? 1 : 0;
   return getOperand(getWaitOperands().size() + numOptional + i);
 }
 
 LogicalResult acc::ParallelOp::verify() {
-  if (failed(checkPrivatizationList(*this, getPrivatizations(),
-                                    getGangPrivateOperands())))
+  if (failed(checkSymOperandList<mlir::acc::PrivateRecipeOp>(
+          *this, getPrivatizations(), getGangPrivateOperands(), "private",
+          "privatizations")))
+    return failure();
+  if (failed(checkSymOperandList<mlir::acc::ReductionRecipeOp>(
+          *this, getReductionRecipes(), getReductionOperands(), "reduction",
+          "reductions")))
     return failure();
   return checkDataOperands<acc::ParallelOp>(*this, getDataClauseOperands());
 }
 
 //===----------------------------------------------------------------------===//
 // SerialOp
 //===----------------------------------------------------------------------===//
 
 unsigned SerialOp::getNumDataOperands() {
   return getReductionOperands().size() + getGangPrivateOperands().size() +
          getGangFirstPrivateOperands().size() + getDataClauseOperands().size();
 }
 
 Value SerialOp::getDataOperand(unsigned i) {
   unsigned numOptional = getAsync() ? 1 : 0;
   numOptional += getIfCond() ? 1 : 0;
   numOptional += getSelfCond() ? 1 : 0;
   return getOperand(getWaitOperands().size() + numOptional + i);
 }
 
 LogicalResult acc::SerialOp::verify() {
+
   return checkDataOperands<acc::SerialOp>(*this, getDataClauseOperands());
 }
 
 //===----------------------------------------------------------------------===//
 // KernelsOp
 //===----------------------------------------------------------------------===//
 
 unsigned KernelsOp::getNumDataOperands() {
return emitError() << "only one of \"" << acc::LoopOp::getAutoAttrStrName()
                        << getSeqAttrName()
                        << " can be present at the same time";
   }
 
   // Gang, worker and vector are incompatible with seq.
   if (getSeq() && (getHasGang() || getHasWorker() || getHasVector()))
     return emitError("gang, worker or vector cannot appear with the seq attr");
 
-  if (failed(checkPrivatizationList(*this, getPrivatizations(),
-                                    getPrivateOperands())))
+  if (failed(checkSymOperandList<mlir::acc::PrivateRecipeOp>(
+          *this, getPrivatizations(), getPrivateOperands(), "private",
+          "privatizations")))
     return failure();
 
   // Check non-empty body().
   if (getRegion().empty())
     return emitError("expected non-empty body.");
 
   return success();
 }

mlir/test/Dialect/OpenACC/ops.mlir

 // 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:       }
+
+func.func @acc_reduc_test(%a : i64) -> () {
+  acc.parallel reduction(@reduction_add_i64 -> %a : i64) {
+  }
+  return
+}
+
+// CHECK-LABEL: func.func @acc_reduc_test(
+// CHECK-SAME:    %[[ARG0:.*]]: i64)
+// CHECK:         acc.parallel reduction(@reduction_add_i64 -> %[[ARG0]] : i64)