//===- SCFToOpenMP.cpp - Structured Control Flow to OpenMP conversion -----===//
//
// 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 a pass to convert scf.parallel operations into OpenMP
// parallel loops.
//
//===----------------------------------------------------------------------===//

#include "mlir/Conversion/SCFToOpenMP/SCFToOpenMP.h"

#include "mlir/Analysis/SliceAnalysis.h"
#include "mlir/Dialect/Affine/Analysis/LoopAnalysis.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/OpenMP/OpenMPDialect.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/IR/SymbolTable.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/WalkPatternRewriteDriver.h"

namespace mlir {
#define GEN_PASS_DEF_CONVERTSCFTOOPENMPPASS
#include "mlir/Conversion/Passes.h.inc"
} // namespace mlir

using namespace mlir;

/// Matches a block containing a "simple" reduction. The expected shape of the
/// block is as follows.
///
///   ^bb(%arg0, %arg1):
///     %0 = OpTy(%arg0, %arg1)
///     scf.reduce.return %0
template <typename... OpTy>
static bool matchSimpleReduction(Block &block) {
  if (block.empty() || llvm::hasSingleElement(block) ||
      std::next(block.begin(), 2) != block.end())
    return false;

  if (block.getNumArguments() != 2)
    return false;

  SmallVector<Operation *, 4> combinerOps;
  Value reducedVal = matchReduction({block.getArguments()[1]},
                                    /*redPos=*/0, combinerOps);

  if (!reducedVal || !isa<BlockArgument>(reducedVal) || combinerOps.size() != 1)
    return false;

  return isa<OpTy...>(combinerOps[0]) &&
         isa<scf::ReduceReturnOp>(block.back()) &&
         block.front().getOperands() == block.getArguments();
}

/// Matches a block containing a select-based min/max reduction. The types of
/// select and compare operations are provided as template arguments. The
/// comparison predicates suitable for min and max are provided as function
/// arguments. If a reduction is matched, `ifMin` will be set if the reduction
/// compute the minimum and unset if it computes the maximum, otherwise it
/// remains unmodified. The expected shape of the block is as follows.
///
///   ^bb(%arg0, %arg1):
///     %0 = CompareOpTy(<one-of-predicates>, %arg0, %arg1)
///     %1 = SelectOpTy(%0, %arg0, %arg1)  // %arg0, %arg1 may be swapped here.
///     scf.reduce.return %1
template <
    typename CompareOpTy, typename SelectOpTy,
    typename Predicate = decltype(std::declval<CompareOpTy>().getPredicate())>
static bool
matchSelectReduction(Block &block, ArrayRef<Predicate> lessThanPredicates,
                     ArrayRef<Predicate> greaterThanPredicates, bool &isMin) {
  static_assert(
      llvm::is_one_of<SelectOpTy, arith::SelectOp, LLVM::SelectOp>::value,
      "only arithmetic and llvm select ops are supported");

  // Expect exactly three operations in the block.
  if (block.empty() || llvm::hasSingleElement(block) ||
      std::next(block.begin(), 2) == block.end() ||
      std::next(block.begin(), 3) != block.end())
    return false;

  // Check op kinds.
  auto compare = dyn_cast<CompareOpTy>(block.front());
  auto select = dyn_cast<SelectOpTy>(block.front().getNextNode());
  auto terminator = dyn_cast<scf::ReduceReturnOp>(block.back());
  if (!compare || !select || !terminator)
    return false;

  // Block arguments must be compared.
  if (compare->getOperands() != block.getArguments())
    return false;

  // Detect whether the comparison is less-than or greater-than, otherwise bail.
  bool isLess;
  if (llvm::is_contained(lessThanPredicates, compare.getPredicate())) {
    isLess = true;
  } else if (llvm::is_contained(greaterThanPredicates,
                                compare.getPredicate())) {
    isLess = false;
  } else {
    return false;
  }

  if (select.getCondition() != compare.getResult())
    return false;

  // Detect if the operands are swapped between cmpf and select. Match the
  // comparison type with the requested type or with the opposite of the
  // requested type if the operands are swapped. Use generic accessors because
  // std and LLVM versions of select have different operand names but identical
  // positions.
  constexpr unsigned kTrueValue = 1;
  constexpr unsigned kFalseValue = 2;
  bool sameOperands = select.getOperand(kTrueValue) == compare.getLhs() &&
                      select.getOperand(kFalseValue) == compare.getRhs();
  bool swappedOperands = select.getOperand(kTrueValue) == compare.getRhs() &&
                         select.getOperand(kFalseValue) == compare.getLhs();
  if (!sameOperands && !swappedOperands)
    return false;

  if (select.getResult() != terminator.getResult())
    return false;

  // The reduction is a min if it uses less-than predicates with same operands
  // or greather-than predicates with swapped operands. Similarly for max.
  isMin = (isLess && sameOperands) || (!isLess && swappedOperands);
  return isMin || (isLess & swappedOperands) || (!isLess && sameOperands);
}

/// Returns the float semantics for the given float type.
static const llvm::fltSemantics &fltSemanticsForType(FloatType type) {
  if (type.isF16())
    return llvm::APFloat::IEEEhalf();
  if (type.isF32())
    return llvm::APFloat::IEEEsingle();
  if (type.isF64())
    return llvm::APFloat::IEEEdouble();
  if (type.isF128())
    return llvm::APFloat::IEEEquad();
  if (type.isBF16())
    return llvm::APFloat::BFloat();
  if (type.isF80())
    return llvm::APFloat::x87DoubleExtended();
  llvm_unreachable("unknown float type");
}

/// Helper to create a splat attribute for vector types, or return the scalar
/// attribute for scalar types.
static Attribute getSplatOrScalarAttr(Type type, Attribute val) {
  if (auto vecType = dyn_cast<VectorType>(type))
    return DenseElementsAttr::get(vecType, val);
  return val;
}

/// Returns an attribute with the minimum (if `min` is set) or the maximum value
/// (otherwise) for the given float type.
static Attribute minMaxValueForFloat(Type type, bool min) {
  Type elType = getElementTypeOrSelf(type);
  auto fltType = cast<FloatType>(elType);
  auto val = llvm::APFloat::getLargest(fltSemanticsForType(fltType), min);

  return getSplatOrScalarAttr(type, FloatAttr::get(elType, val));
}

/// Returns an attribute with the signed integer minimum (if `min` is set) or
/// the maximum value (otherwise) for the given integer type, regardless of its
/// signedness semantics (only the width is considered).
static Attribute minMaxValueForSignedInt(Type type, bool min) {
  Type elType = getElementTypeOrSelf(type);
  auto intType = cast<IntegerType>(elType);
  unsigned bitwidth = intType.getWidth();
  auto val = min ? llvm::APInt::getSignedMinValue(bitwidth)
                 : llvm::APInt::getSignedMaxValue(bitwidth);

  return getSplatOrScalarAttr(type, IntegerAttr::get(elType, val));
}

/// Returns an attribute with the unsigned integer minimum (if `min` is set) or
/// the maximum value (otherwise) for the given integer type, regardless of its
/// signedness semantics (only the width is considered).
static Attribute minMaxValueForUnsignedInt(Type type, bool min) {
  Type elType = getElementTypeOrSelf(type);
  auto intType = cast<IntegerType>(elType);
  unsigned bitwidth = intType.getWidth();
  auto val =
      min ? llvm::APInt::getZero(bitwidth) : llvm::APInt::getAllOnes(bitwidth);

  return getSplatOrScalarAttr(type, IntegerAttr::get(elType, val));
}

/// Creates an OpenMP reduction declaration and inserts it into the provided
/// symbol table. The declaration has a constant initializer with the neutral
/// value `initValue`, and the `reductionIndex`-th reduction combiner carried
/// over from `reduce`.
static omp::DeclareReductionOp
createDecl(PatternRewriter &builder, SymbolTable &symbolTable,
           scf::ReduceOp reduce, int64_t reductionIndex, Attribute initValue) {
  OpBuilder::InsertionGuard guard(builder);
  Type type = reduce.getOperands()[reductionIndex].getType();
  auto decl = omp::DeclareReductionOp::create(builder, reduce.getLoc(),
                                              "__scf_reduction",
                                              /*sym_visibility=*/nullptr, type,
                                              /*byref_element_type=*/{});
  symbolTable.insert(decl);

  builder.createBlock(&decl.getInitializerRegion(),
                      decl.getInitializerRegion().end(), {type},
                      {reduce.getOperands()[reductionIndex].getLoc()});
  builder.setInsertionPointToEnd(&decl.getInitializerRegion().back());
  Value init =
      LLVM::ConstantOp::create(builder, reduce.getLoc(), type, initValue);
  omp::YieldOp::create(builder, reduce.getLoc(), init);

  Operation *terminator =
      &reduce.getReductions()[reductionIndex].front().back();
  assert(isa<scf::ReduceReturnOp>(terminator) &&
         "expected reduce op to be terminated by reduce return");
  builder.setInsertionPoint(terminator);
  builder.replaceOpWithNewOp<omp::YieldOp>(terminator,
                                           terminator->getOperands());
  builder.inlineRegionBefore(reduce.getReductions()[reductionIndex],
                             decl.getReductionRegion(),
                             decl.getReductionRegion().end());
  return decl;
}

/// Adds an atomic reduction combiner to the given OpenMP reduction declaration
/// using llvm.atomicrmw of the given kind.
static omp::DeclareReductionOp addAtomicRMW(OpBuilder &builder,
                                            LLVM::AtomicBinOp atomicKind,
                                            omp::DeclareReductionOp decl,
                                            scf::ReduceOp reduce,
                                            int64_t reductionIndex) {
  OpBuilder::InsertionGuard guard(builder);
  auto ptrType = LLVM::LLVMPointerType::get(builder.getContext());
  Location reduceOperandLoc = reduce.getOperands()[reductionIndex].getLoc();
  builder.createBlock(&decl.getAtomicReductionRegion(),
                      decl.getAtomicReductionRegion().end(), {ptrType, ptrType},
                      {reduceOperandLoc, reduceOperandLoc});
  Block *atomicBlock = &decl.getAtomicReductionRegion().back();
  builder.setInsertionPointToEnd(atomicBlock);
  Value loaded = LLVM::LoadOp::create(builder, reduce.getLoc(), decl.getType(),
                                      atomicBlock->getArgument(1));
  LLVM::AtomicRMWOp::create(builder, reduce.getLoc(), atomicKind,
                            atomicBlock->getArgument(0), loaded,
                            LLVM::AtomicOrdering::monotonic);
  omp::YieldOp::create(builder, reduce.getLoc(), ArrayRef<Value>());
  return decl;
}

/// Returns true if the type is supported by llvm.atomicrmw.
/// LLVM IR currently does not support atomic operations on vector types.
/// See LLVM Language Reference Manual on 'atomicrmw'.
static bool supportsAtomic(Type type) { return !isa<VectorType>(type); }

/// Creates an OpenMP reduction declaration that corresponds to the given SCF
/// reduction and returns it. Recognizes common reductions in order to identify
/// the neutral value, necessary for the OpenMP declaration. If the reduction
/// cannot be recognized, returns null.
static omp::DeclareReductionOp declareReduction(PatternRewriter &builder,
                                                scf::ReduceOp reduce,
                                                int64_t reductionIndex) {
  Operation *container = SymbolTable::getNearestSymbolTable(reduce);
  SymbolTable symbolTable(container);

  // Insert reduction declarations in the symbol-table ancestor before the
  // ancestor of the current insertion point.
  Operation *insertionPoint = reduce;
  while (insertionPoint->getParentOp() != container)
    insertionPoint = insertionPoint->getParentOp();
  OpBuilder::InsertionGuard guard(builder);
  builder.setInsertionPoint(insertionPoint);

  assert(llvm::hasSingleElement(reduce.getReductions()[reductionIndex]) &&
         "expected reduction region to have a single element");

  // Match simple binary reductions that can be expressed with atomicrmw.
  Type type = reduce.getOperands()[reductionIndex].getType();
  Block &reduction = reduce.getReductions()[reductionIndex].front();

  // Handle scalar element type extraction for vector bitwidth safety.
  Type elType = getElementTypeOrSelf(type);

  // Arithmetic Reductions
  if (matchSimpleReduction<arith::AddFOp, LLVM::FAddOp>(reduction)) {
    omp::DeclareReductionOp decl = createDecl(
        builder, symbolTable, reduce, reductionIndex,
        getSplatOrScalarAttr(type, builder.getFloatAttr(elType, 0.0)));
    return supportsAtomic(type) ? addAtomicRMW(builder, LLVM::AtomicBinOp::fadd,
                                               decl, reduce, reductionIndex)
                                : decl;
  }
  if (matchSimpleReduction<arith::AddIOp, LLVM::AddOp>(reduction)) {
    omp::DeclareReductionOp decl = createDecl(
        builder, symbolTable, reduce, reductionIndex,
        getSplatOrScalarAttr(type, builder.getIntegerAttr(elType, 0)));
    return supportsAtomic(type) ? addAtomicRMW(builder, LLVM::AtomicBinOp::add,
                                               decl, reduce, reductionIndex)
                                : decl;
  }
  if (matchSimpleReduction<arith::OrIOp, LLVM::OrOp>(reduction)) {
    omp::DeclareReductionOp decl = createDecl(
        builder, symbolTable, reduce, reductionIndex,
        getSplatOrScalarAttr(type, builder.getIntegerAttr(elType, 0)));
    return supportsAtomic(type) ? addAtomicRMW(builder, LLVM::AtomicBinOp::_or,
                                               decl, reduce, reductionIndex)
                                : decl;
  }
  if (matchSimpleReduction<arith::XOrIOp, LLVM::XOrOp>(reduction)) {
    omp::DeclareReductionOp decl = createDecl(
        builder, symbolTable, reduce, reductionIndex,
        getSplatOrScalarAttr(type, builder.getIntegerAttr(elType, 0)));
    return supportsAtomic(type) ? addAtomicRMW(builder, LLVM::AtomicBinOp::_xor,
                                               decl, reduce, reductionIndex)
                                : decl;
  }
  if (matchSimpleReduction<arith::AndIOp, LLVM::AndOp>(reduction)) {
    APInt allOnes = llvm::APInt::getAllOnes(elType.getIntOrFloatBitWidth());
    omp::DeclareReductionOp decl = createDecl(
        builder, symbolTable, reduce, reductionIndex,
        getSplatOrScalarAttr(type, builder.getIntegerAttr(elType, allOnes)));
    return supportsAtomic(type) ? addAtomicRMW(builder, LLVM::AtomicBinOp::_and,
                                               decl, reduce, reductionIndex)
                                : decl;
  }

  // Match simple binary reductions that cannot be expressed with atomicrmw.
  // TODO: add atomic region using cmpxchg (which needs atomic load to be
  // available as an op).
  if (matchSimpleReduction<arith::MulFOp, LLVM::FMulOp>(reduction)) {
    return createDecl(
        builder, symbolTable, reduce, reductionIndex,
        getSplatOrScalarAttr(type, builder.getFloatAttr(elType, 1.0)));
  }

  if (matchSimpleReduction<arith::MulIOp, LLVM::MulOp>(reduction)) {
    return createDecl(
        builder, symbolTable, reduce, reductionIndex,
        getSplatOrScalarAttr(type, builder.getIntegerAttr(elType, 1)));
  }

  // Match select-based min/max reductions.
  bool isMin;
  // Floating Point Min/Max
  if (matchSelectReduction<arith::CmpFOp, arith::SelectOp,
                           arith::CmpFPredicate>(
          reduction, {arith::CmpFPredicate::OLT, arith::CmpFPredicate::OLE},
          {arith::CmpFPredicate::OGT, arith::CmpFPredicate::OGE}, isMin) ||
      matchSelectReduction<arith::CmpFOp, arith::SelectOp,
                           arith::CmpFPredicate>(
          reduction, {arith::CmpFPredicate::OGT, arith::CmpFPredicate::OGE},
          {arith::CmpFPredicate::OLT, arith::CmpFPredicate::OLE}, isMin)) {
    return createDecl(builder, symbolTable, reduce, reductionIndex,
                      minMaxValueForFloat(type, !isMin));
  }

  // Integer Min/Max
  if (matchSelectReduction<arith::CmpIOp, arith::SelectOp,
                           arith::CmpIPredicate>(
          reduction, {arith::CmpIPredicate::slt, arith::CmpIPredicate::sle},
          {arith::CmpIPredicate::sgt, arith::CmpIPredicate::sge}, isMin) ||
      matchSelectReduction<arith::CmpIOp, arith::SelectOp,
                           arith::CmpIPredicate>(
          reduction, {arith::CmpIPredicate::sgt, arith::CmpIPredicate::sge},
          {arith::CmpIPredicate::slt, arith::CmpIPredicate::sle}, isMin)) {
    omp::DeclareReductionOp decl =
        createDecl(builder, symbolTable, reduce, reductionIndex,
                   minMaxValueForSignedInt(type, !isMin));
    return supportsAtomic(type) ? addAtomicRMW(builder,
                                               isMin ? LLVM::AtomicBinOp::min
                                                     : LLVM::AtomicBinOp::max,
                                               decl, reduce, reductionIndex)
                                : decl;
  }

  // Unsigned Integer Min/Max
  if (matchSelectReduction<arith::CmpIOp, arith::SelectOp,
                           arith::CmpIPredicate>(
          reduction, {arith::CmpIPredicate::ult, arith::CmpIPredicate::ule},
          {arith::CmpIPredicate::ugt, arith::CmpIPredicate::uge}, isMin) ||
      matchSelectReduction<arith::CmpIOp, arith::SelectOp,
                           arith::CmpIPredicate>(
          reduction, {arith::CmpIPredicate::ugt, arith::CmpIPredicate::uge},
          {arith::CmpIPredicate::ult, arith::CmpIPredicate::ule}, isMin)) {
    omp::DeclareReductionOp decl =
        createDecl(builder, symbolTable, reduce, reductionIndex,
                   minMaxValueForUnsignedInt(type, !isMin));
    return supportsAtomic(type) ? addAtomicRMW(builder,
                                               isMin ? LLVM::AtomicBinOp::umin
                                                     : LLVM::AtomicBinOp::umax,
                                               decl, reduce, reductionIndex)
                                : decl;
  }

  return nullptr;
}

namespace {

struct ParallelOpLowering : public OpRewritePattern<scf::ParallelOp> {
  static constexpr unsigned kUseOpenMPDefaultNumThreads = 0;
  unsigned numThreads;

  ParallelOpLowering(MLIRContext *context,
                     unsigned numThreads = kUseOpenMPDefaultNumThreads)
      : OpRewritePattern<scf::ParallelOp>(context), numThreads(numThreads) {}

  LogicalResult matchAndRewrite(scf::ParallelOp parallelOp,
                                PatternRewriter &rewriter) const override {
    // Bail out early if any reduction init value has a type that is not
    // compatible with LLVM (e.g. index), since we cannot allocate a reduction
    // variable for such types.
    for (Value init : parallelOp.getInitVals()) {
      if (!LLVM::isCompatibleType(init.getType()) &&
          !isa<LLVM::PointerElementTypeInterface>(init.getType()))
        return rewriter.notifyMatchFailure(
            parallelOp, "reduction init type is not an LLVM-compatible type");
    }

    // Declare reductions.
    // TODO: consider checking it here is already a compatible reduction
    // declaration and use it instead of redeclaring.
    SmallVector<Attribute> reductionSyms;
    SmallVector<omp::DeclareReductionOp> ompReductionDecls;
    auto reduce = cast<scf::ReduceOp>(parallelOp.getBody()->getTerminator());
    for (int64_t i = 0, e = parallelOp.getNumReductions(); i < e; ++i) {
      omp::DeclareReductionOp decl = declareReduction(rewriter, reduce, i);
      ompReductionDecls.push_back(decl);
      if (!decl)
        return failure();
      reductionSyms.push_back(
          SymbolRefAttr::get(rewriter.getContext(), decl.getSymName()));
    }

    // Allocate reduction variables. Make sure the we don't overflow the stack
    // with local `alloca`s by saving and restoring the stack pointer.
    Location loc = parallelOp.getLoc();
    Value one =
        LLVM::ConstantOp::create(rewriter, loc, rewriter.getIntegerType(64),
                                 rewriter.getI64IntegerAttr(1));
    SmallVector<Value> reductionVariables;
    reductionVariables.reserve(parallelOp.getNumReductions());
    auto ptrType = LLVM::LLVMPointerType::get(parallelOp.getContext());
    for (Value init : parallelOp.getInitVals()) {
      Value storage = LLVM::AllocaOp::create(rewriter, loc, ptrType,
                                             init.getType(), one, 0);
      LLVM::StoreOp::create(rewriter, loc, init, storage);
      reductionVariables.push_back(storage);
    }

    // Replace the reduction operations contained in this loop. Must be done
    // here rather than in a separate pattern to have access to the list of
    // reduction variables.
    for (auto [x, y, rD] : llvm::zip_equal(
             reductionVariables, reduce.getOperands(), ompReductionDecls)) {
      OpBuilder::InsertionGuard guard(rewriter);
      rewriter.setInsertionPoint(reduce);
      Region &redRegion = rD.getReductionRegion();
      // The SCF dialect by definition contains only structured operations
      // and hence the SCF reduction region will contain a single block.
      // The ompReductionDecls region is a copy of the SCF reduction region
      // and hence has the same property.
      assert(redRegion.hasOneBlock() &&
             "expect reduction region to have one block");
      Value pvtRedVar = parallelOp.getRegion().addArgument(x.getType(), loc);
      Value pvtRedVal = LLVM::LoadOp::create(rewriter, reduce.getLoc(),
                                             rD.getType(), pvtRedVar);
      // Make a copy of the reduction combiner region in the body
      mlir::OpBuilder builder(rewriter.getContext());
      builder.setInsertionPoint(reduce);
      mlir::IRMapping mapper;
      assert(redRegion.getNumArguments() == 2 &&
             "expect reduction region to have two arguments");
      mapper.map(redRegion.getArgument(0), pvtRedVal);
      mapper.map(redRegion.getArgument(1), y);
      for (auto &op : redRegion.getOps()) {
        Operation *cloneOp = builder.clone(op, mapper);
        if (auto yieldOp = dyn_cast<omp::YieldOp>(*cloneOp)) {
          assert(yieldOp && yieldOp.getResults().size() == 1 &&
                 "expect YieldOp in reduction region to return one result");
          Value redVal = yieldOp.getResults()[0];
          LLVM::StoreOp::create(rewriter, loc, redVal, pvtRedVar);
          rewriter.eraseOp(yieldOp);
          break;
        }
      }
    }
    rewriter.eraseOp(reduce);

    SmallVector<Value> numThreadsVars;
    if (numThreads > 0) {
      Value numThreadsVar = LLVM::ConstantOp::create(
          rewriter, loc, rewriter.getI32IntegerAttr(numThreads));
      numThreadsVars.push_back(numThreadsVar);
    }
    // Create the parallel wrapper.
    auto ompParallel = omp::ParallelOp::create(
        rewriter, loc,
        /* allocate_vars = */ llvm::SmallVector<Value>{},
        /* allocator_vars = */ llvm::SmallVector<Value>{},
        /* allocate_alignments = */ nullptr,
        /* allocate_private_indices = */ nullptr,
        /* if_expr = */ Value{},
        /* num_threads_vars = */ numThreadsVars,
        /* private_vars = */ ValueRange(),
        /* private_syms = */ nullptr,
        /* private_needs_barrier = */ nullptr,
        /* proc_bind_kind = */ omp::ClauseProcBindKindAttr{},
        /* reduction_mod = */ nullptr,
        /* reduction_vars = */ llvm::SmallVector<Value>{},
        /* reduction_byref = */ DenseBoolArrayAttr{},
        /* reduction_syms = */ ArrayAttr{});
    {

      OpBuilder::InsertionGuard guard(rewriter);
      rewriter.createBlock(&ompParallel.getRegion());

      // Replace the loop.
      {
        OpBuilder::InsertionGuard allocaGuard(rewriter);
        // Create worksharing loop wrapper.
        auto wsloopOp = omp::WsloopOp::create(rewriter, parallelOp.getLoc());
        if (!reductionVariables.empty()) {
          wsloopOp.setReductionSymsAttr(
              ArrayAttr::get(rewriter.getContext(), reductionSyms));
          wsloopOp.getReductionVarsMutable().append(reductionVariables);
          llvm::SmallVector<bool> reductionByRef;
          // false because these reductions always reduce scalars and so do
          // not need to pass by reference
          reductionByRef.resize(reductionVariables.size(), false);
          wsloopOp.setReductionByref(
              DenseBoolArrayAttr::get(rewriter.getContext(), reductionByRef));
        }
        omp::TerminatorOp::create(rewriter, loc); // omp.parallel terminator.

        // The wrapper's entry block arguments will define the reduction
        // variables.
        llvm::SmallVector<mlir::Type> reductionTypes;
        reductionTypes.reserve(reductionVariables.size());
        llvm::transform(reductionVariables, std::back_inserter(reductionTypes),
                        [](mlir::Value v) { return v.getType(); });
        rewriter.createBlock(
            &wsloopOp.getRegion(), {}, reductionTypes,
            llvm::SmallVector<mlir::Location>(reductionVariables.size(),
                                              parallelOp.getLoc()));

        // Create loop nest and populate region with contents of scf.parallel.
        auto loopOp = omp::LoopNestOp::create(
            rewriter, parallelOp.getLoc(), parallelOp.getLowerBound().size(),
            parallelOp.getLowerBound(), parallelOp.getUpperBound(),
            parallelOp.getStep(), /*loop_inclusive=*/false,
            /*tile_sizes=*/nullptr);

        rewriter.inlineRegionBefore(parallelOp.getRegion(), loopOp.getRegion(),
                                    loopOp.getRegion().begin());

        // Remove reduction-related block arguments from omp.loop_nest and
        // redirect uses to the corresponding omp.wsloop block argument.
        mlir::Block &loopOpEntryBlock = loopOp.getRegion().front();
        unsigned numLoops = parallelOp.getNumLoops();
        rewriter.replaceAllUsesWith(
            loopOpEntryBlock.getArguments().drop_front(numLoops),
            wsloopOp.getRegion().getArguments());
        loopOpEntryBlock.eraseArguments(
            numLoops, loopOpEntryBlock.getNumArguments() - numLoops);

        Block *ops =
            rewriter.splitBlock(&loopOpEntryBlock, loopOpEntryBlock.begin());
        rewriter.setInsertionPointToStart(&loopOpEntryBlock);

        auto scope = memref::AllocaScopeOp::create(
            rewriter, parallelOp.getLoc(), TypeRange());
        omp::YieldOp::create(rewriter, loc, ValueRange());
        Block *scopeBlock = rewriter.createBlock(&scope.getBodyRegion());
        rewriter.mergeBlocks(ops, scopeBlock);
        rewriter.setInsertionPointToEnd(&*scope.getBodyRegion().begin());
        memref::AllocaScopeReturnOp::create(rewriter, loc, ValueRange());
      }
    }

    // Load loop results.
    SmallVector<Value> results;
    results.reserve(reductionVariables.size());
    for (auto [variable, type] :
         llvm::zip(reductionVariables, parallelOp.getResultTypes())) {
      Value res = LLVM::LoadOp::create(rewriter, loc, type, variable);
      results.push_back(res);
    }
    rewriter.replaceOp(parallelOp, results);

    return success();
  }
};

/// Applies the conversion patterns in the given function.
static LogicalResult applyPatterns(ModuleOp module, unsigned numThreads) {
  RewritePatternSet patterns(module.getContext());
  patterns.add<ParallelOpLowering>(module.getContext(), numThreads);
  FrozenRewritePatternSet frozen(std::move(patterns));
  walkAndApplyPatterns(module, frozen);
  auto status = module.walk([](Operation *op) {
    if (isa<scf::ReduceOp, scf::ReduceReturnOp, scf::ParallelOp>(op)) {
      op->emitError("unconverted operation found");
      return WalkResult::interrupt();
    }
    return WalkResult::advance();
  });
  return failure(status.wasInterrupted());
}

/// A pass converting SCF operations to OpenMP operations.
struct SCFToOpenMPPass
    : public impl::ConvertSCFToOpenMPPassBase<SCFToOpenMPPass> {

  using Base::Base;

  /// Pass entry point.
  void runOnOperation() override {
    if (failed(applyPatterns(getOperation(), numThreads)))
      signalPassFailure();
  }
};

} // namespace