[mlir] add support for reductions in OpenMP WsLoopOp

[ftynse]

Summary

Use a modeling similar to SCF ParallelOp to support arbitrary parallel
reductions. The two main differences are: (1) reductions are named and declared
beforehand similarly to functions using a special op that provides the neutral
element, the reduction code and optionally the atomic reduction code; (2)
reductions go through memory instead because this is closer to the OpenMP
semantics.

See https://llvm.discourse.group/t/rfc-openmp-reduction-support/3367.

[kiranchandramohan]

Thanks @ftynse for this patch. I am just making my way through this patch. Should be able to spend more time on this later today. Have a few micro-nits and questions to start.

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

27

We discussed this in the RFC, but will this dependency cause issues for any future non-llvm or out of tree lowerings?

41

Nit: Spellings.

177

Nit: Something missing at the start of this line.

mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp

50–52

Do we have to do this attachInterface for FIR or other dialect types? Or can we use OpenMP_PointerLikeTypeInterface while declaring the FIR types?

780–793

This looks trivial. Why is a custom parser and printer required here?

mlir/test/Dialect/OpenMP/invalid.mlir

108

Nit: Missing last character.

138

Nit: Missing last character.

154

Nit: will be good to have "expects" in the beginning of the expected error.

[kiranchandramohan]

A couple more questions.

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

388

Nit: spelling requies

391

Will this region finally sit outside the worksharing loop?

mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp

854

Should there be a check to ensure that the operand type is the same as the element type of the accumulator?

[ftynse]

Address review.

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

27

Absolutely not. And this dependency already exists, as you may note I'm not modifying the cmake. This only indicates to the infrastructure that this dialect might produce objects of the LLVM dialect.

391

This is a reduction _declaration_ (2.19.5.7), it has an initializer that is independent on the value that the accumulator has at the start of the workshare loop.

mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp

50–52

You can do both as long as you don't do it in the main codebase. This was the point of separable registration for interfaces -- OpenMP won't have to know about FIR. Optionally, we can also have FIR not know about OpenMP and have a completely separate registration, but it should not be the primary option.

780–793

One cannot use a keyword as anchor for an optional group in the declarative format.

854

This is already checked by TypeMatchesWith in ODS specification of the op.

[kiranchandramohan]

One more question to clarify the dependency.

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

27

The current dependency is only during translation I think. Previously we had the llvm integer type as an option for the worksharing loop indices, but that is also not there.

But if we choose the LLVM token type we now have a hard dependency isn't it?

[ftynse]

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

27

There is no token type in this version, it is unnecessary with the "variable" approach. To support the use of LLVM pointer types for variables, OpenMP must be aware of the LLVM dialect since we don't want LLVM to depend on OpenMP.

[kiranchandramohan]

mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp

935

We have this code here.

[ftynse]

Drop leftover code.

[ftynse]

mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp

935

Which is a leftover of a previous version. Thanks, this is a nice catch!

[kiranchandramohan]

Reduction is an important clause in the specification. This is also a very good reference for us to make further progress in other constructs. Thanks, @ftynse.

LGTM.

mlir/include/mlir/Dialect/OpenMP/OpenMPDialect.h

 //===- OpenMPDialect.h - MLIR Dialect for OpenMP ----------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file declares the OpenMP dialect in MLIR.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_OPENMP_OPENMPDIALECT_H_
 #define MLIR_DIALECT_OPENMP_OPENMPDIALECT_H_
 
-#include "mlir/Dialect/LLVMIR/LLVMTypes.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/SymbolTable.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 
 #include "mlir/Dialect/OpenMP/OpenMPOpsDialect.h.inc"
 #include "mlir/Dialect/OpenMP/OpenMPOpsEnums.h.inc"
+#include "mlir/Dialect/OpenMP/OpenMPTypeInterfaces.h.inc"
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/OpenMP/OpenMPOps.h.inc"
 
 #endif // MLIR_DIALECT_OPENMP_OPENMPDIALECT_H_

mlir/include/mlir/Dialect/OpenMP/OpenMPOps.td

 
 
 #ifndef OPENMP_OPS
 #define OPENMP_OPS
 
 include "mlir/IR/OpBase.td"
 include "mlir/Interfaces/SideEffectInterfaces.td"
 include "mlir/Interfaces/ControlFlowInterfaces.td"
-include "mlir/Dialect/LLVMIR/LLVMOpBase.td"
+include "mlir/IR/SymbolInterfaces.td"
 include "mlir/Dialect/OpenMP/OmpCommon.td"
+include "mlir/Dialect/LLVMIR/LLVMOpBase.td"
 
 def OpenMP_Dialect : Dialect {
   let name = "omp";
   let cppNamespace = "::mlir::omp";
+  let dependentDialects = ["::mlir::LLVM::LLVMDialect"];
 }
 
 class OpenMP_Op<string mnemonic, list<OpTrait> traits = []> :
       Op<OpenMP_Dialect, mnemonic, traits>;
 
 // Type which can be constraint accepting standard integers and indices.
 def IntLikeType : AnyTypeOf<[AnyInteger, Index]>;
 
+def OpenMP_PointerLikeTypeInterface : TypeInterface<"PointerLikeType"> {
+  let cppNamespace = "::mlir::omp";
+
+  let description = [{
+    An interface for pointer-like types suitable to contain a value that OpenMP
+    specificaiton referes to as variable.
+  }];
+
+  let methods = [
+    InterfaceMethod<
+      /*description=*/"Returns the pointee type.",
+      /*retTy=*/"::mlir::Type",
+      /*methodName=*/"getElementType"
+    >,
+  ];
+}
+
+def OpenMP_PointerLikeType : Type<
+  CPred<"$_self.isa<::mlir::omp::PointerLikeType>()">,
+  "OpenMP-compatible variable type", "::mlir::omp::PointerLikeType">;
+
 //===----------------------------------------------------------------------===//
 // 2.6 parallel Construct
 //===----------------------------------------------------------------------===//
 
 // Possible values for the default clause
 def ClauseDefaultPrivate : StrEnumAttrCase<"defprivate">;
 def ClauseDefaultFirstPrivate : StrEnumAttrCase<"deffirstprivate">;
 def ClauseDefaultShared : StrEnumAttrCase<"defshared">;
let description = [{
 
     `private_vars`, `firstprivate_vars`, `lastprivate_vars` and `linear_vars`
     arguments are variadic list of operands that specify the data sharing
     attributes of the list of values. The `linear_step_vars` operand
     additionally specifies the step for each associated linear operand. Note
     that the `linear_vars` and `linear_step_vars` variadic lists should contain
     the same number of elements.
 
+    Reductions can be performed in a workshare loop by specifying reduction
+    accumulator variables in `reduction_vars` and symbols referring to reduction
+    declarations in the `reductions` attribute. Each reduction is identified
+    by the accumulator it uses and accumulators must not be repeated in the same
+    reduction. The `omp.reduction` operation accepts the accumulator and a value
+    hat corresponds is considered to be produced by the current loop iteration
+    for the given reduction. If multiple values are produced for the same token,
+    i.e. there are multiple `omp.reduction`s, the last value is taken. The
+    reduction declaration specifies how to combine the values from each
+    iteration into the final value, which is available in the accumulator after
+    the loop completes.
+
     The optional `schedule_val` attribute specifies the loop schedule for this
     loop, determining how the loop is distributed across the parallel threads.
     The optional `schedule_chunk_var` associated with this determines further
     controls this distribution.
 
     The optional `collapse_val` attribute specifies the number of loops which
     are collapsed to form the worksharing loop.
 
def WsLoopOp : OpenMP_Op<"wsloop", [AttrSizedOperandSegments,
   let arguments = (ins Variadic<IntLikeType>:$lowerBound,
              Variadic<IntLikeType>:$upperBound,
              Variadic<IntLikeType>:$step,
              Variadic<AnyType>:$private_vars,
              Variadic<AnyType>:$firstprivate_vars,
              Variadic<AnyType>:$lastprivate_vars,
              Variadic<AnyType>:$linear_vars,
              Variadic<AnyType>:$linear_step_vars,
+             Variadic<OpenMP_PointerLikeType>:$reduction_vars,
+             OptionalAttr<TypedArrayAttrBase<SymbolRefAttr,
+                            "array of symbol references">>:$reductions,
              OptionalAttr<ScheduleKind>:$schedule_val,
              Optional<AnyType>:$schedule_chunk_var,
              Confined<OptionalAttr<I64Attr>, [IntMinValue<0>]>:$collapse_val,
              UnitAttr:$nowait,
              Confined<OptionalAttr<I64Attr>, [IntMinValue<0>]>:$ordered_val,
              OptionalAttr<OrderKind>:$order_val,
              UnitAttr:$inclusive);
 
   let skipDefaultBuilders = 1;
 
   let builders = [
     OpBuilder<(ins "ValueRange":$lowerBound, "ValueRange":$upperBound,
                "ValueRange":$step,
                CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes)>,
     OpBuilder<(ins "TypeRange":$resultTypes, "ValueRange":$lowerBound,
                "ValueRange":$upperBound, "ValueRange":$step,
                "ValueRange":$privateVars, "ValueRange":$firstprivateVars,
                "ValueRange":$lastprivate_vars, "ValueRange":$linear_vars,
-               "ValueRange":$linear_step_vars, "StringAttr":$schedule_val,
-               "Value":$schedule_chunk_var, "IntegerAttr":$collapse_val,
-               "UnitAttr":$nowait, "IntegerAttr":$ordered_val,
-               "StringAttr":$order_val, "UnitAttr":$inclusive, CArg<"bool",
-               "true">:$buildBody)>,
+               "ValueRange":$linear_step_vars, "ValueRange":$reduction_vars,
+               "StringAttr":$schedule_val, "Value":$schedule_chunk_var,
+               "IntegerAttr":$collapse_val, "UnitAttr":$nowait,
+               "IntegerAttr":$ordered_val, "StringAttr":$order_val,
+               "UnitAttr":$inclusive, CArg<"bool", "true">:$buildBody)>,
     OpBuilder<(ins "TypeRange":$resultTypes, "ValueRange":$operands,
                CArg<"ArrayRef<NamedAttribute>", "{}">:$attributes)>
   ];
 
   let regions = (region AnyRegion:$region);
 
   let extraClassDeclaration = [{
     /// Returns the number of loops in the workshape loop nest.
     unsigned getNumLoops() { return lowerBound().size(); }
+
+    /// Returns the number of reduction variables.
+    unsigned getNumReductionVars() { return reduction_vars().size(); }
   }];
   let parser = [{ return parseWsLoopOp(parser, result); }];
   let printer = [{ return printWsLoopOp(p, *this); }];
+  let verifier = [{ return ::verifyWsLoopOp(*this); }];
 }
 
-def YieldOp : OpenMP_Op<"yield", [NoSideEffect, ReturnLike, Terminator,
-                                  HasParent<"WsLoopOp">]> {
+def YieldOp : OpenMP_Op<"yield",
+    [NoSideEffect, ReturnLike, Terminator,
+     ParentOneOf<["WsLoopOp", "ReductionDeclareOp"]>]> {
   let summary = "loop yield and termination operation";
   let description = [{
     "omp.yield" yields SSA values from the OpenMP dialect op region and
     terminates the region. The semantics of how the values are yielded is
     defined by the parent operation.
     If "omp.yield" has any operands, the operands must match the parent
     operation's results.
   }];
def TaskwaitOp : OpenMP_Op<"taskwait"> {
   let description = [{
     The taskwait construct specifies a wait on the completion of child tasks
     of the current task.
   }];
 
   let assemblyFormat = "attr-dict";
 }
 
+//===----------------------------------------------------------------------===//
+// 2.19.5.7 declare reduction Directive
+//===----------------------------------------------------------------------===//
+
+def ReductionDeclareOp : OpenMP_Op<"reduction.declare", [Symbol]> {
+  let summary = "declares a reduction kind";
+
+  let description = [{
+    Declares an OpenMP reduction kind. This requies two mandatory and one
+    optional region.
+
+      1. The initializer region specifies how to initialize the thread-local
+         reduction value. This is usually the neutral element of the reduction.
+         For convenience, the region has an argument that contains the value
+         of the reduction accumulator at the start of the reduction. It is
+         expected to `omp.yield` the new value on all control flow paths.
+      2. The reduction region specifies how to combine two values into one, i.e.
+         the reduction operator. It accepts the two values as arguments and is
+         expected to `omp.yield` the combined value on all control flow paths.
+      3. The atomic reduction region is optional and specifies how two values
+         can be combined atomically given local accumulator variables. It is
+         expected to store the combined value in the first accumulator variable.
+
+    Note that the MLIR type system does not allow for type-polymorphic
+    reductions. Separate reduction declarations should be created for different
+    element and accumulator types.
+  }];
+
+  let arguments = (ins SymbolNameAttr:$sym_name,
+                       TypeAttr:$type);
+
+  let regions = (region AnyRegion:$initializerRegion,
+                        AnyRegion:$reductionRegion,
+                        AnyRegion:$atomicReductionRegion);
+  let verifier = "return ::verifyReductionDeclareOp(*this);";
+
+  let assemblyFormat = "$sym_name `:` $type attr-dict-with-keyword "
+                       "`init` $initializerRegion "
+                       "`combiner` $reductionRegion "
+                       "custom<AtomicReductionRegion>($atomicReductionRegion)";
+
+  let extraClassDeclaration = [{
+    PointerLikeType getAccumulatorType() {
+      if (atomicReductionRegion().empty())
+        return {};
+
+      return atomicReductionRegion().front().getArgument(0).getType();
+    }
+  }];
+}
+
+//===----------------------------------------------------------------------===//
+// 2.19.5.4 reduction clause
+//===----------------------------------------------------------------------===//
+
+def ReductionOp : OpenMP_Op<"reduction", [
+    TypesMatchWith<"value types matches accumulator element type",
+                   "accumulator", "operand",
+                 "$_self.cast<::mlir::omp::PointerLikeType>().getElementType()">
+  ]> {
+  let summary = "reduction construct";
+  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.
+  }];
+
+  let arguments= (ins AnyType:$operand, OpenMP_PointerLikeType:$accumulator);
+  let assemblyFormat =
+    "$operand `,` $accumulator attr-dict `:` type($accumulator)";
+  let verifier = "return ::verifyReductionOp(*this);";
+}
+
 #endif // OPENMP_OPS

mlir/lib/Dialect/OpenMP/IR/OpenMPDialect.cpp

 //===- OpenMPDialect.cpp - MLIR Dialect for OpenMP implementation ---------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements the OpenMP dialect and its operations.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/OpenMP/OpenMPDialect.h"
+#include "mlir/Dialect/LLVMIR/LLVMTypes.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/OperationSupport.h"
 
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/StringExtras.h"
 #include "llvm/ADT/StringRef.h"
 #include "llvm/ADT/StringSwitch.h"
 #include <cstddef>
 
 #include "mlir/Dialect/OpenMP/OpenMPOpsDialect.cpp.inc"
 #include "mlir/Dialect/OpenMP/OpenMPOpsEnums.cpp.inc"
+#include "mlir/Dialect/OpenMP/OpenMPTypeInterfaces.cpp.inc"
 
 using namespace mlir;
 using namespace mlir::omp;
 
+namespace {
+/// Model for pointer-like types that already provide a `getElementType` method.
+template <typename T>
+struct PointerLikeModel
+    : public PointerLikeType::ExternalModel<PointerLikeModel<T>, T> {
+  Type getElementType(Type pointer) const {
+    return pointer.cast<T>().getElementType();
+  }
+};
+} // end namespace
+
 void OpenMPDialect::initialize() {
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/OpenMP/OpenMPOps.cpp.inc"
       >();
+
+  LLVM::LLVMPointerType::attachInterface<
+      PointerLikeModel<LLVM::LLVMPointerType>>(*getContext());
+  MemRefType::attachInterface<PointerLikeModel<MemRefType>>(*getContext());
 }
 
 //===----------------------------------------------------------------------===//
 // ParallelOp
 //===----------------------------------------------------------------------===//
 
 void ParallelOp::build(OpBuilder &builder, OperationState &state,
                        ArrayRef<NamedAttribute> attributes) {
parseScheduleClause(OpAsmParser &parser, SmallString<8> &schedule,
   }
 
   if (parser.parseRParen())
     return failure();
 
   return success();
 }
 
+/// reduction-init ::= `reduction` `(` reduction-entry-list `)`
+/// reduction-entry-list ::= reduction-entry
+///                        | reduction-entry-list `,` reduction-entry
+/// reduction-entry ::= symbol-ref `->` ssa-id `:` type
+static ParseResult
+parseReductionVarList(OpAsmParser &parser,
+                      SmallVectorImpl<SymbolRefAttr> &symbols,
+                      SmallVectorImpl<OpAsmParser::OperandType> &operands,
+                      SmallVectorImpl<Type> &types) {
+  if (failed(parser.parseLParen()))
+    return failure();
+
+  do {
+    if (parser.parseAttribute(symbols.emplace_back()) || parser.parseArrow() ||
+        parser.parseOperand(operands.emplace_back()) ||
+        parser.parseColonType(types.emplace_back()))
+      return failure();
+  } while (succeeded(parser.parseOptionalComma()));
+  return parser.parseRParen();
+}
+
 /// Parses an OpenMP Workshare Loop operation
 ///
 /// operation ::= `omp.wsloop` loop-control clause-list
 /// loop-control ::= `(` ssa-id-list `)` `:` type `=`  loop-bounds
 /// loop-bounds := `(` ssa-id-list `)` to `(` ssa-id-list `)` steps
 /// steps := `step` `(`ssa-id-list`)`
 /// clause-list ::= clause | empty | clause-list
 /// clause ::= private | firstprivate | lastprivate | linear | schedule |
static ParseResult parseWsLoopOp(OpAsmParser &parser, OperationState &result) {
   SmallVector<Type> privateTypes;
   SmallVector<OpAsmParser::OperandType> firstprivates;
   SmallVector<Type> firstprivateTypes;
   SmallVector<OpAsmParser::OperandType> lastprivates;
   SmallVector<Type> lastprivateTypes;
   SmallVector<OpAsmParser::OperandType> linears;
   SmallVector<Type> linearTypes;
   SmallVector<OpAsmParser::OperandType> linearSteps;
+  SmallVector<SymbolRefAttr> reductionSymbols;
+  SmallVector<OpAsmParser::OperandType> reductionVars;
+  SmallVector<Type> reductionVarTypes;
   SmallString<8> schedule;
   Optional<OpAsmParser::OperandType> scheduleChunkSize;
-  std::array<int, 9> segments{numIVs, numIVs, numIVs, 0, 0, 0, 0, 0, 0};
 
   const StringRef opName = result.name.getStringRef();
   StringRef keyword;
 
   enum SegmentPos {
     lbPos = 0,
     ubPos,
     stepPos,
     privateClausePos,
     firstprivateClausePos,
     lastprivateClausePos,
     linearClausePos,
     linearStepPos,
+    reductionVarPos,
     scheduleClausePos,
   };
+  std::array<int, 10> segments{numIVs, numIVs, numIVs, 0, 0, 0, 0, 0, 0, 0};
 
   while (succeeded(parser.parseOptionalKeyword(&keyword))) {
     if (keyword == "private") {
       if (segments[privateClausePos])
         return allowedOnce(parser, "private", opName);
       if (parseOperandAndTypeList(parser, privates, privateTypes))
         return failure();
       segments[privateClausePos] = privates.size();
if (keyword == "private") {
       if (parser.parseLParen() || parser.parseKeyword(&order) ||
           parser.parseRParen())
         return failure();
       auto attr = parser.getBuilder().getStringAttr(order);
       result.addAttribute("order", attr);
     } else if (keyword == "inclusive") {
       auto attr = UnitAttr::get(parser.getBuilder().getContext());
       result.addAttribute("inclusive", attr);
+    } else if (keyword == "reduction") {
+      if (segments[reductionVarPos])
+        return allowedOnce(parser, "reduction", opName);
+      if (failed(parseReductionVarList(parser, reductionSymbols, reductionVars,
+                                       reductionVarTypes)))
+        return failure();
+      segments[reductionVarPos] = reductionVars.size();
     }
   }
 
   if (segments[privateClausePos]) {
     parser.resolveOperands(privates, privateTypes, privates[0].location,
                            result.operands);
   }
 
if (segments[linearClausePos]) {
     parser.resolveOperands(linears, linearTypes, linears[0].location,
                            result.operands);
     auto linearStepType = parser.getBuilder().getI32Type();
     SmallVector<Type> linearStepTypes(linearSteps.size(), linearStepType);
     parser.resolveOperands(linearSteps, linearStepTypes,
                            linearSteps[0].location, result.operands);
   }
 
+  if (segments[reductionVarPos]) {
+    if (failed(parser.resolveOperands(reductionVars, reductionVarTypes,
+                                      parser.getNameLoc(), result.operands))) {
+      return failure();
+    }
+    SmallVector<Attribute> reductions(reductionSymbols.begin(),
+                                      reductionSymbols.end());
+    result.addAttribute("reductions",
+                        parser.getBuilder().getArrayAttr(reductions));
+  }
+
   if (!schedule.empty()) {
     schedule[0] = llvm::toUpper(schedule[0]);
     auto attr = parser.getBuilder().getStringAttr(schedule);
     result.addAttribute("schedule_val", attr);
     if (scheduleChunkSize) {
       auto chunkSizeType = parser.getBuilder().getI32Type();
       parser.resolveOperand(*scheduleChunkSize, chunkSizeType, result.operands);
     }
   }
 
   result.addAttribute("operand_segment_sizes",
                       parser.getBuilder().getI32VectorAttr(segments));
 
   // Now parse the body.
   Region *body = result.addRegion();
   SmallVector<Type> ivTypes(numIVs, loopVarType);
-  if (parser.parseRegion(*body, ivs, ivTypes))
+  SmallVector<OpAsmParser::OperandType> blockArgs(ivs);
+  if (parser.parseRegion(*body, blockArgs, ivTypes))
     return failure();
   return success();
 }
 
 static void printWsLoopOp(OpAsmPrinter &p, WsLoopOp op) {
   auto args = op.getRegion().front().getArguments();
   p << op.getOperationName() << " (" << args << ") : " << args[0].getType()
     << " = (" << op.lowerBound() << ") to (" << op.upperBound() << ") step ("
static void printWsLoopOp(OpAsmPrinter &p, WsLoopOp op) {
 
   if (op.nowait())
     p << " nowait";
 
   if (auto ordered = op.ordered_val()) {
     p << " ordered(" << ordered << ")";
   }
 
+  if (!op.reduction_vars().empty()) {
+    p << " reduction(";
+    for (unsigned i = 0, e = op.getNumReductionVars(); i < e; ++i) {
+      if (i != 0)
+        p << ", ";
+      p << (*op.reductions())[i] << " -> " << op.reduction_vars()[i] << " : "
+        << op.reduction_vars()[i].getType();
+    }
+    p << ")";
+  }
+
   if (op.inclusive()) {
     p << " inclusive";
   }
 
   p.printRegion(op.region(), /*printEntryBlockArgs=*/false);
 }
 
 //===----------------------------------------------------------------------===//
+// ReductionOp
+//===----------------------------------------------------------------------===//
+
+static ParseResult parseAtomicReductionRegion(OpAsmParser &parser,
+                                              Region &region) {
+  if (parser.parseOptionalKeyword("atomic"))
+    return success();
+  return parser.parseRegion(region);
+}
+
+static void printAtomicReductionRegion(OpAsmPrinter &printer,
+                                       ReductionDeclareOp op, Region &region) {
+  if (region.empty())
+    return;
+  printer << "atomic ";
+  printer.printRegion(region);
+}
+
+static LogicalResult verifyReductionDeclareOp(ReductionDeclareOp op) {
+  if (op.initializerRegion().empty())
+    return op.emitOpError() << "expects non-empty initializer region";
+  Block &initializerEntryBlock = op.initializerRegion().front();
+  if (initializerEntryBlock.getNumArguments() != 1 ||
+      initializerEntryBlock.getArgument(0).getType() != op.type()) {
+    return op.emitOpError() << "expects initializer region with one argument "
+                               "of the reduction type";
+  }
+
+  for (YieldOp yieldOp : op.initializerRegion().getOps<YieldOp>()) {
+    if (yieldOp.results().size() != 1 ||
+        yieldOp.results().getTypes()[0] != op.type())
+      return op.emitOpError() << "expects initializer region to yield a value "
+                                 "of the reduction type";
+  }
+
+  if (op.reductionRegion().empty())
+    return op.emitOpError() << "expects non-empty reduction region";
+  Block &reductionEntryBlock = op.reductionRegion().front();
+  if (reductionEntryBlock.getNumArguments() != 2 ||
+      reductionEntryBlock.getArgumentTypes()[0] !=
+          reductionEntryBlock.getArgumentTypes()[1] ||
+      reductionEntryBlock.getArgumentTypes()[0] != op.type())
+    return op.emitOpError() << "expects reduction region with two arguments of "
+                               "the reduction type";
+  for (YieldOp yieldOp : op.reductionRegion().getOps<YieldOp>()) {
+    if (yieldOp.results().size() != 1 ||
+        yieldOp.results().getTypes()[0] != op.type())
+      return op.emitOpError() << "expects reduction region to yield a value "
+                                 "of the reduction type";
+  }
+
+  if (op.atomicReductionRegion().empty())
+    return success();
+
+  Block &atomicReductionEntryBlock = op.atomicReductionRegion().front();
+  if (atomicReductionEntryBlock.getNumArguments() != 2 ||
+      atomicReductionEntryBlock.getArgumentTypes()[0] !=
+          atomicReductionEntryBlock.getArgumentTypes()[1])
+    return op.emitOpError() << "expects atomic reduction region with two "
+                               "arguments of the same type";
+  auto ptrType = atomicReductionEntryBlock.getArgumentTypes()[0]
+                     .dyn_cast<PointerLikeType>();
+  if (!ptrType || ptrType.getElementType() != op.type())
+    return op.emitOpError() << "expects atomic reduction region arguments to "
+                               "be accumulators containing the reduction type";
+  return success();
+}
+
+static LogicalResult verifyReductionOp(ReductionOp op) {
+  // TODO: generalize this to an op interface when there is more than one op
+  // that supports reductions.
+  auto container = op->getParentOfType<WsLoopOp>();
+  for (unsigned i = 0, e = container.getNumReductionVars(); i < e; ++i)
+    if (container.reduction_vars()[i] == op.accumulator())
+      return success();
+
+  return op.emitOpError() << "the accumulator is not used by the parent";
+}
+
+//===----------------------------------------------------------------------===//
 // WsLoopOp
 //===----------------------------------------------------------------------===//
 
 void WsLoopOp::build(OpBuilder &builder, OperationState &state,
                      ValueRange lowerBound, ValueRange upperBound,
                      ValueRange step, ArrayRef<NamedAttribute> attributes) {
   build(builder, state, TypeRange(), lowerBound, upperBound, step,
         /*private_vars=*/ValueRange(),
         /*firstprivate_vars=*/ValueRange(), /*lastprivate_vars=*/ValueRange(),
         /*linear_vars=*/ValueRange(), /*linear_step_vars=*/ValueRange(),
-        /*schedule_val=*/nullptr, /*schedule_chunk_var=*/nullptr,
-        /*collapse_val=*/nullptr,
+        /*reduction_vars=*/ValueRange(), /*schedule_val=*/nullptr,
+        /*schedule_chunk_var=*/nullptr, /*collapse_val=*/nullptr,
         /*nowait=*/nullptr, /*ordered_val=*/nullptr, /*order_val=*/nullptr,
         /*inclusive=*/nullptr, /*buildBody=*/false);
   state.addAttributes(attributes);
 }
 
 void WsLoopOp::build(OpBuilder &, OperationState &state, TypeRange resultTypes,
                      ValueRange operands, ArrayRef<NamedAttribute> attributes) {
   state.addOperands(operands);
   state.addAttributes(attributes);
   (void)state.addRegion();
-  assert(resultTypes.size() == 0u && "mismatched number of return types");
+  assert(resultTypes.empty() && "mismatched number of return types");
   state.addTypes(resultTypes);
 }
 
 void WsLoopOp::build(OpBuilder &builder, OperationState &result,
                      TypeRange typeRange, ValueRange lowerBounds,
                      ValueRange upperBounds, ValueRange steps,
                      ValueRange privateVars, ValueRange firstprivateVars,
                      ValueRange lastprivateVars, ValueRange linearVars,
-                     ValueRange linearStepVars, StringAttr scheduleVal,
-                     Value scheduleChunkVar, IntegerAttr collapseVal,
-                     UnitAttr nowait, IntegerAttr orderedVal,
-                     StringAttr orderVal, UnitAttr inclusive, bool buildBody) {
+                     ValueRange linearStepVars, ValueRange reductionVars,
+                     StringAttr scheduleVal, Value scheduleChunkVar,
+                     IntegerAttr collapseVal, UnitAttr nowait,
+                     IntegerAttr orderedVal, StringAttr orderVal,
+                     UnitAttr inclusive, bool buildBody) {
   result.addOperands(lowerBounds);
   result.addOperands(upperBounds);
   result.addOperands(steps);
   result.addOperands(privateVars);
   result.addOperands(firstprivateVars);
   result.addOperands(linearVars);
   result.addOperands(linearStepVars);
   if (scheduleChunkVar)
result.addAttribute(
           {static_cast<int32_t>(lowerBounds.size()),
            static_cast<int32_t>(upperBounds.size()),
            static_cast<int32_t>(steps.size()),
            static_cast<int32_t>(privateVars.size()),
            static_cast<int32_t>(firstprivateVars.size()),
            static_cast<int32_t>(lastprivateVars.size()),
            static_cast<int32_t>(linearVars.size()),
            static_cast<int32_t>(linearStepVars.size()),
+           static_cast<int32_t>(reductionVars.size()),
            static_cast<int32_t>(scheduleChunkVar != nullptr ? 1 : 0)}));
 
   Region *bodyRegion = result.addRegion();
   if (buildBody) {
     OpBuilder::InsertionGuard guard(builder);
     unsigned numIVs = steps.size();
     SmallVector<Type, 8> argTypes(numIVs, steps.getType().front());
+    argTypes.append(reductionVars.size(),
+                    LLVM::LLVMTokenType::get(builder.getContext()));
     builder.createBlock(bodyRegion, {}, argTypes);
   }
 }
 
+static LogicalResult verifyWsLoopOp(WsLoopOp op) {
+  if (op.getNumReductionVars() != 0) {
+    if (!op.reductions() ||
+        op.reductions()->size() != op.getNumReductionVars()) {
+      return op.emitOpError() << "expected as many reduction symbol references "
+                                 "as reduction variables";
+    }
+  } else {
+    if (op.reductions())
+      return op.emitOpError() << "unexpected reduction symbol references";
+    return success();
+  }
+
+  DenseSet<Value> accumulators;
+  for (auto args : llvm::zip(op.reduction_vars(), *op.reductions())) {
+    Value accum = std::get<0>(args);
+    if (!accumulators.insert(accum).second) {
+      return op.emitOpError() << "accumulator variable used more than once";
+    }
+    Type varType = accum.getType().cast<PointerLikeType>();
+    auto symbolRef = std::get<1>(args).cast<SymbolRefAttr>();
+    auto decl =
+        SymbolTable::lookupNearestSymbolFrom<ReductionDeclareOp>(op, symbolRef);
+    if (!decl) {
+      return op.emitOpError() << "expected symbol reference " << symbolRef
+                              << " to point to a reduction declaration";
+    }
+
+    if (decl.getAccumulatorType() && decl.getAccumulatorType() != varType) {
+      return op.emitOpError()
+             << "expected accumulator (" << varType
+             << ") to be the same type as reduction declaration ("
+             << decl.getAccumulatorType() << ")";
+    }
+  }
+
+  return success();
+}
+
 #define GET_OP_CLASSES
 #include "mlir/Dialect/OpenMP/OpenMPOps.cpp.inc"

mlir/test/Conversion/OpenMPToLLVM/convert-to-llvmir.mlir

func @wsloop(%arg0: index, %arg1: index, %arg2: index, %arg3: index, %arg4: index, %arg5: index) {
   // CHECK: omp.parallel
   omp.parallel {
     // CHECK: omp.wsloop (%[[ARG6:.*]], %[[ARG7:.*]]) : i64 = (%[[ARG0]], %[[ARG1]]) to (%[[ARG2]], %[[ARG3]]) step (%[[ARG4]], %[[ARG5]]) {
     "omp.wsloop"(%arg0, %arg1, %arg2, %arg3, %arg4, %arg5) ( {
     ^bb0(%arg6: index, %arg7: index):  // no predecessors
       // CHECK: "test.payload"(%[[ARG6]], %[[ARG7]]) : (i64, i64) -> ()
       "test.payload"(%arg6, %arg7) : (index, index) -> ()
       omp.yield
-    }) {operand_segment_sizes = dense<[2, 2, 2, 0, 0, 0, 0, 0, 0]> : vector<9xi32>} : (index, index, index, index, index, index) -> ()
+    }) {operand_segment_sizes = dense<[2, 2, 2, 0, 0, 0, 0, 0, 0, 0]> : vector<10xi32>} : (index, index, index, index, index, index) -> ()
     omp.terminator
   }
   return
 }

mlir/test/Dialect/OpenMP/invalid.mlir

 
 func @proc_bind_once() {
   // expected-error@+1 {{at most one proc_bind clause can appear on the omp.parallel operation}}
   omp.parallel proc_bind(close) proc_bind(spread) {
   }
 
   return
 }
+
+// -----
+
+// expected-error @below {{op expects initializer region with one argument of the reduction type}}
+omp.reduction.declare @add_f32 : f64
+init {
+^bb0(%arg: f32):
+  %0 = constant 0.0 : f32
+  omp.yield (%0 : f32)
+}
+combiner {
+^bb1(%arg0: f32, %arg1: f32):
+  %1 = addf %arg0, %arg1 : f32
+  omp.yield (%1 : f32)
+}
+
+// -----
+
+// expected-error @below {{expects initializer region to yield a value of the reduction typ}}
+omp.reduction.declare @add_f32 : f32
+init {
+^bb0(%arg: f32):
+  %0 = constant 0.0 : f64
+  omp.yield (%0 : f64)
+}
+combiner {
+^bb1(%arg0: f32, %arg1: f32):
+  %1 = addf %arg0, %arg1 : f32
+  omp.yield (%1 : f32)
+}
+
+// -----
+
+// expected-error @below {{expects reduction region with two arguments of the reduction type}}
+omp.reduction.declare @add_f32 : f32
+init {
+^bb0(%arg: f32):
+  %0 = constant 0.0 : f32
+  omp.yield (%0 : f32)
+}
+combiner {
+^bb1(%arg0: f64, %arg1: f64):
+  %1 = addf %arg0, %arg1 : f64
+  omp.yield (%1 : f64)
+}
+
+// -----
+
+// expected-error @below {{expects reduction region to yield a value of the reduction typ}}
+omp.reduction.declare @add_f32 : f32
+init {
+^bb0(%arg: f32):
+  %0 = constant 0.0 : f32
+  omp.yield (%0 : f32)
+}
+combiner {
+^bb1(%arg0: f32, %arg1: f32):
+  %1 = addf %arg0, %arg1 : f32
+  %2 = fpext %1 : f32 to f64
+  omp.yield (%2 : f64)
+}
+
+// -----
+
+// expected-error @below {{atomic reduction region with two arguments of the same type}}
+omp.reduction.declare @add_f32 : f32
+init {
+^bb0(%arg: f32):
+  %0 = constant 0.0 : f32
+  omp.yield (%0 : f32)
+}
+combiner {
+^bb1(%arg0: f32, %arg1: f32):
+  %1 = addf %arg0, %arg1 : f32
+  omp.yield (%1 : f32)
+}
+atomic {
+^bb2(%arg0: memref<f32>, %arg1: memref<f64>):
+  omp.yield
+}
+
+// -----
+
+// expected-error @below {{expects atomic reduction region arguments to be accumulators containing the reduction type}}
+omp.reduction.declare @add_f32 : f32
+init {
+^bb0(%arg: f32):
+  %0 = constant 0.0 : f32
+  omp.yield (%0 : f32)
+}
+combiner {
+^bb1(%arg0: f32, %arg1: f32):
+  %1 = addf %arg0, %arg1 : f32
+  omp.yield (%1 : f32)
+}
+atomic {
+^bb2(%arg0: memref<f64>, %arg1: memref<f64>):
+  omp.yield
+}
+
+// -----
+
+omp.reduction.declare @add_f32 : f32
+init {
+^bb0(%arg: f32):
+  %0 = constant 0.0 : f32
+  omp.yield (%0 : f32)
+}
+combiner {
+^bb1(%arg0: f32, %arg1: f32):
+  %1 = addf %arg0, %arg1 : f32
+  omp.yield (%1 : f32)
+}
+
+func @foo(%lb : index, %ub : index, %step : index) {
+  %c1 = constant 1 : i32
+  %0 = llvm.alloca %c1 x i32 : (i32) -> !llvm.ptr<f32>
+  %1 = llvm.alloca %c1 x i32 : (i32) -> !llvm.ptr<f32>
+
+  omp.wsloop (%iv) : index = (%lb) to (%ub) step (%step)
+  reduction(@add_f32 -> %0 : !llvm.ptr<f32>) {
+    %2 = constant 2.0 : f32
+    // expected-error @below {{accumulator is not used by the parent}}
+    omp.reduction %2, %1 : !llvm.ptr<f32>
+    omp.yield
+  }
+  return
+}
+
+// -----
+
+func @foo(%lb : index, %ub : index, %step : index) {
+  %c1 = constant 1 : i32
+  %0 = llvm.alloca %c1 x i32 : (i32) -> !llvm.ptr<f32>
+  %1 = llvm.alloca %c1 x i32 : (i32) -> !llvm.ptr<f32>
+
+  // expected-error @below {{expected symbol reference @foo to point to a reduction declaration}}
+  omp.wsloop (%iv) : index = (%lb) to (%ub) step (%step)
+  reduction(@foo -> %0 : !llvm.ptr<f32>) {
+    %2 = constant 2.0 : f32
+    omp.reduction %2, %1 : !llvm.ptr<f32>
+    omp.yield
+  }
+  return
+}
+
+// -----
+
+omp.reduction.declare @add_f32 : f32
+init {
+^bb0(%arg: f32):
+  %0 = constant 0.0 : f32
+  omp.yield (%0 : f32)
+}
+combiner {
+^bb1(%arg0: f32, %arg1: f32):
+  %1 = addf %arg0, %arg1 : f32
+  omp.yield (%1 : f32)
+}
+
+func @foo(%lb : index, %ub : index, %step : index) {
+  %c1 = constant 1 : i32
+  %0 = llvm.alloca %c1 x i32 : (i32) -> !llvm.ptr<f32>
+
+  // expected-error @below {{accumulator variable used more than once}}
+  omp.wsloop (%iv) : index = (%lb) to (%ub) step (%step)
+  reduction(@add_f32 -> %0 : !llvm.ptr<f32>, @add_f32 -> %0 : !llvm.ptr<f32>) {
+    %2 = constant 2.0 : f32
+    omp.reduction %2, %0 : !llvm.ptr<f32>
+    omp.yield
+  }
+  return
+}
+
+// -----
+
+omp.reduction.declare @add_f32 : f32
+init {
+^bb0(%arg: f32):
+  %0 = constant 0.0 : f32
+  omp.yield (%0 : f32)
+}
+combiner {
+^bb1(%arg0: f32, %arg1: f32):
+  %1 = addf %arg0, %arg1 : f32
+  omp.yield (%1 : f32)
+}
+atomic {
+^bb2(%arg2: !llvm.ptr<f32>, %arg3: !llvm.ptr<f32>):
+  %2 = llvm.load %arg3 : !llvm.ptr<f32>
+  llvm.atomicrmw fadd %arg2, %2 monotonic : f32
+  omp.yield
+}
+
+func @foo(%lb : index, %ub : index, %step : index, %mem : memref<1xf32>) {
+  %c1 = constant 1 : i32
+
+  // expected-error @below {{expected accumulator ('memref<1xf32>') to be the same type as reduction declaration ('!llvm.ptr<f32>')}}
+  omp.wsloop (%iv) : index = (%lb) to (%ub) step (%step)
+  reduction(@add_f32 -> %mem : memref<1xf32>) {
+    %2 = constant 2.0 : f32
+    omp.reduction %2, %mem : memref<1xf32>
+    omp.yield
+  }
+  return
+}

mlir/test/Dialect/OpenMP/ops.mlir

 
 // CHECK-LABEL: omp_wsloop
 func @omp_wsloop(%lb : index, %ub : index, %step : index, %data_var : memref<i32>, %linear_var : i32, %chunk_var : i32) -> () {
 
   // CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) private(%{{.*}} : memref<i32>, %{{.*}} : memref<i32>) collapse(2) ordered(1)
   "omp.wsloop" (%lb, %ub, %step, %data_var, %data_var) ({
     ^bb0(%iv: index):
       omp.yield
-  }) {operand_segment_sizes = dense<[1,1,1,2,0,0,0,0,0]> : vector<9xi32>, collapse_val = 2, ordered_val = 1} :
+  }) {operand_segment_sizes = dense<[1,1,1,2,0,0,0,0,0,0]> : vector<10xi32>, collapse_val = 2, ordered_val = 1} :
     (index, index, index, memref<i32>, memref<i32>) -> ()
 
   // CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) linear(%{{.*}} = %{{.*}} : memref<i32>) schedule(static)
   "omp.wsloop" (%lb, %ub, %step, %data_var, %linear_var) ({
     ^bb0(%iv: index):
       omp.yield
-  }) {operand_segment_sizes = dense<[1,1,1,0,0,0,1,1,0]> : vector<9xi32>, schedule_val = "Static"} :
+  }) {operand_segment_sizes = dense<[1,1,1,0,0,0,1,1,0,0]> : vector<10xi32>, schedule_val = "Static"} :
     (index, index, index, memref<i32>, i32) -> ()
 
   // CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) linear(%{{.*}} = %{{.*}} : memref<i32>, %{{.*}} = %{{.*}} : memref<i32>) schedule(static)
   "omp.wsloop" (%lb, %ub, %step, %data_var, %data_var, %linear_var, %linear_var) ({
     ^bb0(%iv: index):
       omp.yield
-  }) {operand_segment_sizes = dense<[1,1,1,0,0,0,2,2,0]> : vector<9xi32>, schedule_val = "Static"} :
+  }) {operand_segment_sizes = dense<[1,1,1,0,0,0,2,2,0,0]> : vector<10xi32>, schedule_val = "Static"} :
     (index, index, index, memref<i32>, memref<i32>, i32, i32) -> ()
 
   // CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) private(%{{.*}} : memref<i32>) firstprivate(%{{.*}} : memref<i32>) lastprivate(%{{.*}} : memref<i32>) linear(%{{.*}} = %{{.*}} : memref<i32>) schedule(dynamic = %{{.*}}) collapse(3) ordered(2)
   "omp.wsloop" (%lb, %ub, %step, %data_var, %data_var, %data_var, %data_var, %linear_var, %chunk_var) ({
     ^bb0(%iv: index):
       omp.yield
-  }) {operand_segment_sizes = dense<[1,1,1,1,1,1,1,1,1]> : vector<9xi32>, schedule_val = "Dynamic", collapse_val = 3, ordered_val = 2} :
+  }) {operand_segment_sizes = dense<[1,1,1,1,1,1,1,1,0,1]> : vector<10xi32>, schedule_val = "Dynamic", collapse_val = 3, ordered_val = 2} :
     (index, index, index, memref<i32>, memref<i32>, memref<i32>, memref<i32>, i32, i32) -> ()
 
   // CHECK: omp.wsloop (%{{.*}}) : index = (%{{.*}}) to (%{{.*}}) step (%{{.*}}) private(%{{.*}} : memref<i32>) schedule(auto) nowait
   "omp.wsloop" (%lb, %ub, %step, %data_var) ({
     ^bb0(%iv: index):
       omp.yield
-  }) {operand_segment_sizes = dense<[1,1,1,1,0,0,0,0,0]> : vector<9xi32>, nowait, schedule_val = "Auto"} :
+  }) {operand_segment_sizes = dense<[1,1,1,1,0,0,0,0,0,0]> : vector<10xi32>, nowait, schedule_val = "Auto"} :
     (index, index, index, memref<i32>) -> ()
 
   return
 }
 
 // CHECK-LABEL: omp_wsloop_pretty
 func @omp_wsloop_pretty(%lb : index, %ub : index, %step : index,
                  %data_var : memref<i32>, %linear_var : i32, %chunk_var : i32) -> () {
"omp.target"(%if_cond, %device, %num_threads) ({
        omp.terminator
     }) {operand_segment_sizes = dense<[1,1,1]>: vector<3xi32>, nowait } : ( i1, si32, si32  ) -> ()
 
     // CHECK: omp.barrier
     omp.barrier
 
     return
 }
+
+// CHECK: omp.reduction.declare
+// CHECK-LABEL: @add_f32
+// CHECK: : f32
+// CHECK: init
+// CHECK: ^{{.+}}(%{{.+}}: f32):
+// CHECK:   omp.yield
+// CHECK: combiner
+// CHECK: ^{{.+}}(%{{.+}}: f32, %{{.+}}: f32):
+// CHECK:   omp.yield
+// CHECK: atomic
+// CHECK: ^{{.+}}(%{{.+}}: !llvm.ptr<f32>, %{{.+}}: !llvm.ptr<f32>):
+// CHECK:  omp.yield
+omp.reduction.declare @add_f32 : f32
+init {
+^bb0(%arg: f32):
+  %0 = constant 0.0 : f32
+  omp.yield (%0 : f32)
+}
+combiner {
+^bb1(%arg0: f32, %arg1: f32):
+  %1 = addf %arg0, %arg1 : f32
+  omp.yield (%1 : f32)
+}
+atomic {
+^bb2(%arg2: !llvm.ptr<f32>, %arg3: !llvm.ptr<f32>):
+  %2 = llvm.load %arg3 : !llvm.ptr<f32>
+  llvm.atomicrmw fadd %arg2, %2 monotonic : f32
+  omp.yield
+}
+
+func @reduction(%lb : index, %ub : index, %step : index) {
+  %c1 = constant 1 : i32
+  %0 = llvm.alloca %c1 x i32 : (i32) -> !llvm.ptr<f32>
+  // CHECK: reduction(@add_f32 -> %{{.+}} : !llvm.ptr<f32>)
+  omp.wsloop (%iv) : index = (%lb) to (%ub) step (%step)
+  reduction(@add_f32 -> %0 : !llvm.ptr<f32>) {
+    %1 = constant 2.0 : f32
+    // CHECK: omp.reduction %{{.+}}, %{{.+}}
+    omp.reduction %1, %0 : !llvm.ptr<f32>
+    omp.yield
+  }
+  return
+}
+
+// CHECK: omp.reduction.declare
+// CHECK-LABEL: @add2_f32
+omp.reduction.declare @add2_f32 : f32
+// CHECK: init
+init {
+^bb0(%arg: f32):
+  %0 = constant 0.0 : f32
+  omp.yield (%0 : f32)
+}
+// CHECK: combiner
+combiner {
+^bb1(%arg0: f32, %arg1: f32):
+  %1 = addf %arg0, %arg1 : f32
+  omp.yield (%1 : f32)
+}
+// CHECK-NOT: atomic
+
+func @reduction2(%lb : index, %ub : index, %step : index) {
+  %0 = memref.alloca() : memref<1xf32>
+  // CHECK: reduction
+  omp.wsloop (%iv) : index = (%lb) to (%ub) step (%step)
+  reduction(@add2_f32 -> %0 : memref<1xf32>) {
+    %1 = constant 2.0 : f32
+    // CHECK: omp.reduction
+    omp.reduction %1, %0 : memref<1xf32>
+    omp.yield
+  }
+  return
+}
+

mlir/test/Target/LLVMIR/openmp-llvm.mlir

^bb0(%arg1: i64):
       // tested there. Just check that the right functions are called.
       // CHECK: call i32 @__kmpc_global_thread_num
       // CHECK: call void @__kmpc_for_static_init_{{.*}}(%struct.ident_t* @[[$wsloop_loc_struct]],
       %3 = llvm.mlir.constant(2.000000e+00 : f32) : f32
       %4 = llvm.getelementptr %arg0[%arg1] : (!llvm.ptr<f32>, i64) -> !llvm.ptr<f32>
       llvm.store %3, %4 : !llvm.ptr<f32>
       omp.yield
       // CHECK: call void @__kmpc_for_static_fini(%struct.ident_t* @[[$wsloop_loc_struct]],
-    }) {operand_segment_sizes = dense<[1, 1, 1, 0, 0, 0, 0, 0, 0]> : vector<9xi32>} : (i64, i64, i64) -> ()
+    }) {operand_segment_sizes = dense<[1, 1, 1, 0, 0, 0, 0, 0, 0, 0]> : vector<10xi32>} : (i64, i64, i64) -> ()
     omp.terminator
   }
   llvm.return
 }
 
 // CHECK-LABEL: @wsloop_inclusive_1
 llvm.func @wsloop_inclusive_1(%arg0: !llvm.ptr<f32>) {
   %0 = llvm.mlir.constant(42 : index) : i64
   %1 = llvm.mlir.constant(10 : index) : i64
   %2 = llvm.mlir.constant(1 : index) : i64
   // CHECK: store i64 31, i64* %{{.*}}upperbound
   "omp.wsloop"(%1, %0, %2) ( {
   ^bb0(%arg1: i64):
     %3 = llvm.mlir.constant(2.000000e+00 : f32) : f32
     %4 = llvm.getelementptr %arg0[%arg1] : (!llvm.ptr<f32>, i64) -> !llvm.ptr<f32>
     llvm.store %3, %4 : !llvm.ptr<f32>
     omp.yield
-  }) {operand_segment_sizes = dense<[1, 1, 1, 0, 0, 0, 0, 0, 0]> : vector<9xi32>} : (i64, i64, i64) -> ()
+  }) {operand_segment_sizes = dense<[1, 1, 1, 0, 0, 0, 0, 0, 0, 0]> : vector<10xi32>} : (i64, i64, i64) -> ()
   llvm.return
 }
 
 // CHECK-LABEL: @wsloop_inclusive_2
 llvm.func @wsloop_inclusive_2(%arg0: !llvm.ptr<f32>) {
   %0 = llvm.mlir.constant(42 : index) : i64
   %1 = llvm.mlir.constant(10 : index) : i64
   %2 = llvm.mlir.constant(1 : index) : i64
   // CHECK: store i64 32, i64* %{{.*}}upperbound
   "omp.wsloop"(%1, %0, %2) ( {
   ^bb0(%arg1: i64):
     %3 = llvm.mlir.constant(2.000000e+00 : f32) : f32
     %4 = llvm.getelementptr %arg0[%arg1] : (!llvm.ptr<f32>, i64) -> !llvm.ptr<f32>
     llvm.store %3, %4 : !llvm.ptr<f32>
     omp.yield
-  }) {inclusive, operand_segment_sizes = dense<[1, 1, 1, 0, 0, 0, 0, 0, 0]> : vector<9xi32>} : (i64, i64, i64) -> ()
+  }) {inclusive, operand_segment_sizes = dense<[1, 1, 1, 0, 0, 0, 0, 0, 0, 0]> : vector<10xi32>} : (i64, i64, i64) -> ()
   llvm.return
 }
 
 llvm.func @body(i64)
 
 llvm.func @test_omp_wsloop_dynamic(%lb : i64, %ub : i64, %step : i64) -> () {
  omp.wsloop (%iv) : i64 = (%lb) to (%ub) step (%step) schedule(dynamic) {
   // CHECK: call void @__kmpc_dispatch_init_8u