| //===- 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", 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>{}, | |
| /* 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 |