1 //===- SCFToOpenMP.cpp - Structured Control Flow to OpenMP conversion -----===// 2 // 3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. 4 // See https://llvm.org/LICENSE.txt for license information. 5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception 6 // 7 //===----------------------------------------------------------------------===// 8 // 9 // This file implements a pass to convert scf.parallel operations into OpenMP 10 // parallel loops. 11 // 12 //===----------------------------------------------------------------------===// 13 14 #include "mlir/Conversion/SCFToOpenMP/SCFToOpenMP.h" 15 #include "../PassDetail.h" 16 #include "mlir/Analysis/SliceAnalysis.h" 17 #include "mlir/Dialect/Affine/Analysis/LoopAnalysis.h" 18 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h" 19 #include "mlir/Dialect/Func/IR/FuncOps.h" 20 #include "mlir/Dialect/LLVMIR/LLVMDialect.h" 21 #include "mlir/Dialect/MemRef/IR/MemRef.h" 22 #include "mlir/Dialect/OpenMP/OpenMPDialect.h" 23 #include "mlir/Dialect/SCF/SCF.h" 24 #include "mlir/IR/ImplicitLocOpBuilder.h" 25 #include "mlir/IR/SymbolTable.h" 26 #include "mlir/Transforms/DialectConversion.h" 27 28 using namespace mlir; 29 30 /// Matches a block containing a "simple" reduction. The expected shape of the 31 /// block is as follows. 32 /// 33 /// ^bb(%arg0, %arg1): 34 /// %0 = OpTy(%arg0, %arg1) 35 /// scf.reduce.return %0 36 template <typename... OpTy> 37 static bool matchSimpleReduction(Block &block) { 38 if (block.empty() || llvm::hasSingleElement(block) || 39 std::next(block.begin(), 2) != block.end()) 40 return false; 41 42 if (block.getNumArguments() != 2) 43 return false; 44 45 SmallVector<Operation *, 4> combinerOps; 46 Value reducedVal = matchReduction({block.getArguments()[1]}, 47 /*redPos=*/0, combinerOps); 48 49 if (!reducedVal || !reducedVal.isa<BlockArgument>() || 50 combinerOps.size() != 1) 51 return false; 52 53 return isa<OpTy...>(combinerOps[0]) && 54 isa<scf::ReduceReturnOp>(block.back()) && 55 block.front().getOperands() == block.getArguments(); 56 } 57 58 /// Matches a block containing a select-based min/max reduction. The types of 59 /// select and compare operations are provided as template arguments. The 60 /// comparison predicates suitable for min and max are provided as function 61 /// arguments. If a reduction is matched, `ifMin` will be set if the reduction 62 /// compute the minimum and unset if it computes the maximum, otherwise it 63 /// remains unmodified. The expected shape of the block is as follows. 64 /// 65 /// ^bb(%arg0, %arg1): 66 /// %0 = CompareOpTy(<one-of-predicates>, %arg0, %arg1) 67 /// %1 = SelectOpTy(%0, %arg0, %arg1) // %arg0, %arg1 may be swapped here. 68 /// scf.reduce.return %1 69 template < 70 typename CompareOpTy, typename SelectOpTy, 71 typename Predicate = decltype(std::declval<CompareOpTy>().getPredicate())> 72 static bool 73 matchSelectReduction(Block &block, ArrayRef<Predicate> lessThanPredicates, 74 ArrayRef<Predicate> greaterThanPredicates, bool &isMin) { 75 static_assert( 76 llvm::is_one_of<SelectOpTy, arith::SelectOp, LLVM::SelectOp>::value, 77 "only arithmetic and llvm select ops are supported"); 78 79 // Expect exactly three operations in the block. 80 if (block.empty() || llvm::hasSingleElement(block) || 81 std::next(block.begin(), 2) == block.end() || 82 std::next(block.begin(), 3) != block.end()) 83 return false; 84 85 // Check op kinds. 86 auto compare = dyn_cast<CompareOpTy>(block.front()); 87 auto select = dyn_cast<SelectOpTy>(block.front().getNextNode()); 88 auto terminator = dyn_cast<scf::ReduceReturnOp>(block.back()); 89 if (!compare || !select || !terminator) 90 return false; 91 92 // Block arguments must be compared. 93 if (compare->getOperands() != block.getArguments()) 94 return false; 95 96 // Detect whether the comparison is less-than or greater-than, otherwise bail. 97 bool isLess; 98 if (llvm::find(lessThanPredicates, compare.getPredicate()) != 99 lessThanPredicates.end()) { 100 isLess = true; 101 } else if (llvm::find(greaterThanPredicates, compare.getPredicate()) != 102 greaterThanPredicates.end()) { 103 isLess = false; 104 } else { 105 return false; 106 } 107 108 if (select.getCondition() != compare.getResult()) 109 return false; 110 111 // Detect if the operands are swapped between cmpf and select. Match the 112 // comparison type with the requested type or with the opposite of the 113 // requested type if the operands are swapped. Use generic accessors because 114 // std and LLVM versions of select have different operand names but identical 115 // positions. 116 constexpr unsigned kTrueValue = 1; 117 constexpr unsigned kFalseValue = 2; 118 bool sameOperands = select.getOperand(kTrueValue) == compare.getLhs() && 119 select.getOperand(kFalseValue) == compare.getRhs(); 120 bool swappedOperands = select.getOperand(kTrueValue) == compare.getRhs() && 121 select.getOperand(kFalseValue) == compare.getLhs(); 122 if (!sameOperands && !swappedOperands) 123 return false; 124 125 if (select.getResult() != terminator.getResult()) 126 return false; 127 128 // The reduction is a min if it uses less-than predicates with same operands 129 // or greather-than predicates with swapped operands. Similarly for max. 130 isMin = (isLess && sameOperands) || (!isLess && swappedOperands); 131 return isMin || (isLess & swappedOperands) || (!isLess && sameOperands); 132 } 133 134 /// Returns the float semantics for the given float type. 135 static const llvm::fltSemantics &fltSemanticsForType(FloatType type) { 136 if (type.isF16()) 137 return llvm::APFloat::IEEEhalf(); 138 if (type.isF32()) 139 return llvm::APFloat::IEEEsingle(); 140 if (type.isF64()) 141 return llvm::APFloat::IEEEdouble(); 142 if (type.isF128()) 143 return llvm::APFloat::IEEEquad(); 144 if (type.isBF16()) 145 return llvm::APFloat::BFloat(); 146 if (type.isF80()) 147 return llvm::APFloat::x87DoubleExtended(); 148 llvm_unreachable("unknown float type"); 149 } 150 151 /// Returns an attribute with the minimum (if `min` is set) or the maximum value 152 /// (otherwise) for the given float type. 153 static Attribute minMaxValueForFloat(Type type, bool min) { 154 auto fltType = type.cast<FloatType>(); 155 return FloatAttr::get( 156 type, llvm::APFloat::getLargest(fltSemanticsForType(fltType), min)); 157 } 158 159 /// Returns an attribute with the signed integer minimum (if `min` is set) or 160 /// the maximum value (otherwise) for the given integer type, regardless of its 161 /// signedness semantics (only the width is considered). 162 static Attribute minMaxValueForSignedInt(Type type, bool min) { 163 auto intType = type.cast<IntegerType>(); 164 unsigned bitwidth = intType.getWidth(); 165 return IntegerAttr::get(type, min ? llvm::APInt::getSignedMinValue(bitwidth) 166 : llvm::APInt::getSignedMaxValue(bitwidth)); 167 } 168 169 /// Returns an attribute with the unsigned integer minimum (if `min` is set) or 170 /// the maximum value (otherwise) for the given integer type, regardless of its 171 /// signedness semantics (only the width is considered). 172 static Attribute minMaxValueForUnsignedInt(Type type, bool min) { 173 auto intType = type.cast<IntegerType>(); 174 unsigned bitwidth = intType.getWidth(); 175 return IntegerAttr::get(type, min ? llvm::APInt::getNullValue(bitwidth) 176 : llvm::APInt::getAllOnesValue(bitwidth)); 177 } 178 179 /// Creates an OpenMP reduction declaration and inserts it into the provided 180 /// symbol table. The declaration has a constant initializer with the neutral 181 /// value `initValue`, and the reduction combiner carried over from `reduce`. 182 static omp::ReductionDeclareOp createDecl(PatternRewriter &builder, 183 SymbolTable &symbolTable, 184 scf::ReduceOp reduce, 185 Attribute initValue) { 186 OpBuilder::InsertionGuard guard(builder); 187 auto decl = builder.create<omp::ReductionDeclareOp>( 188 reduce.getLoc(), "__scf_reduction", reduce.getOperand().getType()); 189 symbolTable.insert(decl); 190 191 Type type = reduce.getOperand().getType(); 192 builder.createBlock(&decl.initializerRegion(), decl.initializerRegion().end(), 193 {type}, {reduce.getOperand().getLoc()}); 194 builder.setInsertionPointToEnd(&decl.initializerRegion().back()); 195 Value init = 196 builder.create<LLVM::ConstantOp>(reduce.getLoc(), type, initValue); 197 builder.create<omp::YieldOp>(reduce.getLoc(), init); 198 199 Operation *terminator = &reduce.getRegion().front().back(); 200 assert(isa<scf::ReduceReturnOp>(terminator) && 201 "expected reduce op to be terminated by redure return"); 202 builder.setInsertionPoint(terminator); 203 builder.replaceOpWithNewOp<omp::YieldOp>(terminator, 204 terminator->getOperands()); 205 builder.inlineRegionBefore(reduce.getRegion(), decl.reductionRegion(), 206 decl.reductionRegion().end()); 207 return decl; 208 } 209 210 /// Adds an atomic reduction combiner to the given OpenMP reduction declaration 211 /// using llvm.atomicrmw of the given kind. 212 static omp::ReductionDeclareOp addAtomicRMW(OpBuilder &builder, 213 LLVM::AtomicBinOp atomicKind, 214 omp::ReductionDeclareOp decl, 215 scf::ReduceOp reduce) { 216 OpBuilder::InsertionGuard guard(builder); 217 Type type = reduce.getOperand().getType(); 218 Type ptrType = LLVM::LLVMPointerType::get(type); 219 Location reduceOperandLoc = reduce.getOperand().getLoc(); 220 builder.createBlock(&decl.atomicReductionRegion(), 221 decl.atomicReductionRegion().end(), {ptrType, ptrType}, 222 {reduceOperandLoc, reduceOperandLoc}); 223 Block *atomicBlock = &decl.atomicReductionRegion().back(); 224 builder.setInsertionPointToEnd(atomicBlock); 225 Value loaded = builder.create<LLVM::LoadOp>(reduce.getLoc(), 226 atomicBlock->getArgument(1)); 227 builder.create<LLVM::AtomicRMWOp>(reduce.getLoc(), type, atomicKind, 228 atomicBlock->getArgument(0), loaded, 229 LLVM::AtomicOrdering::monotonic); 230 builder.create<omp::YieldOp>(reduce.getLoc(), ArrayRef<Value>()); 231 return decl; 232 } 233 234 /// Creates an OpenMP reduction declaration that corresponds to the given SCF 235 /// reduction and returns it. Recognizes common reductions in order to identify 236 /// the neutral value, necessary for the OpenMP declaration. If the reduction 237 /// cannot be recognized, returns null. 238 static omp::ReductionDeclareOp declareReduction(PatternRewriter &builder, 239 scf::ReduceOp reduce) { 240 Operation *container = SymbolTable::getNearestSymbolTable(reduce); 241 SymbolTable symbolTable(container); 242 243 // Insert reduction declarations in the symbol-table ancestor before the 244 // ancestor of the current insertion point. 245 Operation *insertionPoint = reduce; 246 while (insertionPoint->getParentOp() != container) 247 insertionPoint = insertionPoint->getParentOp(); 248 OpBuilder::InsertionGuard guard(builder); 249 builder.setInsertionPoint(insertionPoint); 250 251 assert(llvm::hasSingleElement(reduce.getRegion()) && 252 "expected reduction region to have a single element"); 253 254 // Match simple binary reductions that can be expressed with atomicrmw. 255 Type type = reduce.getOperand().getType(); 256 Block &reduction = reduce.getRegion().front(); 257 if (matchSimpleReduction<arith::AddFOp, LLVM::FAddOp>(reduction)) { 258 omp::ReductionDeclareOp decl = createDecl(builder, symbolTable, reduce, 259 builder.getFloatAttr(type, 0.0)); 260 return addAtomicRMW(builder, LLVM::AtomicBinOp::fadd, decl, reduce); 261 } 262 if (matchSimpleReduction<arith::AddIOp, LLVM::AddOp>(reduction)) { 263 omp::ReductionDeclareOp decl = createDecl(builder, symbolTable, reduce, 264 builder.getIntegerAttr(type, 0)); 265 return addAtomicRMW(builder, LLVM::AtomicBinOp::add, decl, reduce); 266 } 267 if (matchSimpleReduction<arith::OrIOp, LLVM::OrOp>(reduction)) { 268 omp::ReductionDeclareOp decl = createDecl(builder, symbolTable, reduce, 269 builder.getIntegerAttr(type, 0)); 270 return addAtomicRMW(builder, LLVM::AtomicBinOp::_or, decl, reduce); 271 } 272 if (matchSimpleReduction<arith::XOrIOp, LLVM::XOrOp>(reduction)) { 273 omp::ReductionDeclareOp decl = createDecl(builder, symbolTable, reduce, 274 builder.getIntegerAttr(type, 0)); 275 return addAtomicRMW(builder, LLVM::AtomicBinOp::_xor, decl, reduce); 276 } 277 if (matchSimpleReduction<arith::AndIOp, LLVM::AndOp>(reduction)) { 278 omp::ReductionDeclareOp decl = createDecl( 279 builder, symbolTable, reduce, 280 builder.getIntegerAttr( 281 type, llvm::APInt::getAllOnesValue(type.getIntOrFloatBitWidth()))); 282 return addAtomicRMW(builder, LLVM::AtomicBinOp::_and, decl, reduce); 283 } 284 285 // Match simple binary reductions that cannot be expressed with atomicrmw. 286 // TODO: add atomic region using cmpxchg (which needs atomic load to be 287 // available as an op). 288 if (matchSimpleReduction<arith::MulFOp, LLVM::FMulOp>(reduction)) { 289 return createDecl(builder, symbolTable, reduce, 290 builder.getFloatAttr(type, 1.0)); 291 } 292 293 // Match select-based min/max reductions. 294 bool isMin; 295 if (matchSelectReduction<arith::CmpFOp, arith::SelectOp>( 296 reduction, {arith::CmpFPredicate::OLT, arith::CmpFPredicate::OLE}, 297 {arith::CmpFPredicate::OGT, arith::CmpFPredicate::OGE}, isMin) || 298 matchSelectReduction<LLVM::FCmpOp, LLVM::SelectOp>( 299 reduction, {LLVM::FCmpPredicate::olt, LLVM::FCmpPredicate::ole}, 300 {LLVM::FCmpPredicate::ogt, LLVM::FCmpPredicate::oge}, isMin)) { 301 return createDecl(builder, symbolTable, reduce, 302 minMaxValueForFloat(type, !isMin)); 303 } 304 if (matchSelectReduction<arith::CmpIOp, arith::SelectOp>( 305 reduction, {arith::CmpIPredicate::slt, arith::CmpIPredicate::sle}, 306 {arith::CmpIPredicate::sgt, arith::CmpIPredicate::sge}, isMin) || 307 matchSelectReduction<LLVM::ICmpOp, LLVM::SelectOp>( 308 reduction, {LLVM::ICmpPredicate::slt, LLVM::ICmpPredicate::sle}, 309 {LLVM::ICmpPredicate::sgt, LLVM::ICmpPredicate::sge}, isMin)) { 310 omp::ReductionDeclareOp decl = createDecl( 311 builder, symbolTable, reduce, minMaxValueForSignedInt(type, !isMin)); 312 return addAtomicRMW(builder, 313 isMin ? LLVM::AtomicBinOp::min : LLVM::AtomicBinOp::max, 314 decl, reduce); 315 } 316 if (matchSelectReduction<arith::CmpIOp, arith::SelectOp>( 317 reduction, {arith::CmpIPredicate::ult, arith::CmpIPredicate::ule}, 318 {arith::CmpIPredicate::ugt, arith::CmpIPredicate::uge}, isMin) || 319 matchSelectReduction<LLVM::ICmpOp, LLVM::SelectOp>( 320 reduction, {LLVM::ICmpPredicate::ugt, LLVM::ICmpPredicate::ule}, 321 {LLVM::ICmpPredicate::ugt, LLVM::ICmpPredicate::uge}, isMin)) { 322 omp::ReductionDeclareOp decl = createDecl( 323 builder, symbolTable, reduce, minMaxValueForUnsignedInt(type, !isMin)); 324 return addAtomicRMW( 325 builder, isMin ? LLVM::AtomicBinOp::umin : LLVM::AtomicBinOp::umax, 326 decl, reduce); 327 } 328 329 return nullptr; 330 } 331 332 namespace { 333 334 struct ParallelOpLowering : public OpRewritePattern<scf::ParallelOp> { 335 using OpRewritePattern<scf::ParallelOp>::OpRewritePattern; 336 337 LogicalResult matchAndRewrite(scf::ParallelOp parallelOp, 338 PatternRewriter &rewriter) const override { 339 // Replace SCF yield with OpenMP yield. 340 { 341 OpBuilder::InsertionGuard guard(rewriter); 342 rewriter.setInsertionPointToEnd(parallelOp.getBody()); 343 assert(llvm::hasSingleElement(parallelOp.getRegion()) && 344 "expected scf.parallel to have one block"); 345 rewriter.replaceOpWithNewOp<omp::YieldOp>( 346 parallelOp.getBody()->getTerminator(), ValueRange()); 347 } 348 349 // Declare reductions. 350 // TODO: consider checking it here is already a compatible reduction 351 // declaration and use it instead of redeclaring. 352 SmallVector<Attribute> reductionDeclSymbols; 353 for (auto reduce : parallelOp.getOps<scf::ReduceOp>()) { 354 omp::ReductionDeclareOp decl = declareReduction(rewriter, reduce); 355 if (!decl) 356 return failure(); 357 reductionDeclSymbols.push_back( 358 SymbolRefAttr::get(rewriter.getContext(), decl.sym_name())); 359 } 360 361 // Allocate reduction variables. Make sure the we don't overflow the stack 362 // with local `alloca`s by saving and restoring the stack pointer. 363 Location loc = parallelOp.getLoc(); 364 Value one = rewriter.create<LLVM::ConstantOp>( 365 loc, rewriter.getIntegerType(64), rewriter.getI64IntegerAttr(1)); 366 SmallVector<Value> reductionVariables; 367 reductionVariables.reserve(parallelOp.getNumReductions()); 368 for (Value init : parallelOp.getInitVals()) { 369 assert((LLVM::isCompatibleType(init.getType()) || 370 init.getType().isa<LLVM::PointerElementTypeInterface>()) && 371 "cannot create a reduction variable if the type is not an LLVM " 372 "pointer element"); 373 Value storage = rewriter.create<LLVM::AllocaOp>( 374 loc, LLVM::LLVMPointerType::get(init.getType()), one, 0); 375 rewriter.create<LLVM::StoreOp>(loc, init, storage); 376 reductionVariables.push_back(storage); 377 } 378 379 // Replace the reduction operations contained in this loop. Must be done 380 // here rather than in a separate pattern to have access to the list of 381 // reduction variables. 382 for (auto pair : 383 llvm::zip(parallelOp.getOps<scf::ReduceOp>(), reductionVariables)) { 384 OpBuilder::InsertionGuard guard(rewriter); 385 scf::ReduceOp reduceOp = std::get<0>(pair); 386 rewriter.setInsertionPoint(reduceOp); 387 rewriter.replaceOpWithNewOp<omp::ReductionOp>( 388 reduceOp, reduceOp.getOperand(), std::get<1>(pair)); 389 } 390 391 // Create the parallel wrapper. 392 auto ompParallel = rewriter.create<omp::ParallelOp>(loc); 393 { 394 395 OpBuilder::InsertionGuard guard(rewriter); 396 rewriter.createBlock(&ompParallel.region()); 397 398 { 399 auto scope = rewriter.create<memref::AllocaScopeOp>(parallelOp.getLoc(), 400 TypeRange()); 401 rewriter.create<omp::TerminatorOp>(loc); 402 OpBuilder::InsertionGuard allocaGuard(rewriter); 403 rewriter.createBlock(&scope.getBodyRegion()); 404 rewriter.setInsertionPointToStart(&scope.getBodyRegion().front()); 405 406 // Replace the loop. 407 auto loop = rewriter.create<omp::WsLoopOp>( 408 parallelOp.getLoc(), parallelOp.getLowerBound(), 409 parallelOp.getUpperBound(), parallelOp.getStep()); 410 rewriter.create<memref::AllocaScopeReturnOp>(loc); 411 412 rewriter.inlineRegionBefore(parallelOp.getRegion(), loop.region(), 413 loop.region().begin()); 414 if (!reductionVariables.empty()) { 415 loop.reductionsAttr( 416 ArrayAttr::get(rewriter.getContext(), reductionDeclSymbols)); 417 loop.reduction_varsMutable().append(reductionVariables); 418 } 419 } 420 } 421 422 // Load loop results. 423 SmallVector<Value> results; 424 results.reserve(reductionVariables.size()); 425 for (Value variable : reductionVariables) { 426 Value res = rewriter.create<LLVM::LoadOp>(loc, variable); 427 results.push_back(res); 428 } 429 rewriter.replaceOp(parallelOp, results); 430 431 return success(); 432 } 433 }; 434 435 /// Applies the conversion patterns in the given function. 436 static LogicalResult applyPatterns(ModuleOp module) { 437 ConversionTarget target(*module.getContext()); 438 target.addIllegalOp<scf::ReduceOp, scf::ReduceReturnOp, scf::ParallelOp>(); 439 target.addLegalDialect<omp::OpenMPDialect, LLVM::LLVMDialect, 440 memref::MemRefDialect>(); 441 442 RewritePatternSet patterns(module.getContext()); 443 patterns.add<ParallelOpLowering>(module.getContext()); 444 FrozenRewritePatternSet frozen(std::move(patterns)); 445 return applyPartialConversion(module, target, frozen); 446 } 447 448 /// A pass converting SCF operations to OpenMP operations. 449 struct SCFToOpenMPPass : public ConvertSCFToOpenMPBase<SCFToOpenMPPass> { 450 /// Pass entry point. 451 void runOnOperation() override { 452 if (failed(applyPatterns(getOperation()))) 453 signalPassFailure(); 454 } 455 }; 456 457 } // namespace 458 459 std::unique_ptr<OperationPass<ModuleOp>> mlir::createConvertSCFToOpenMPPass() { 460 return std::make_unique<SCFToOpenMPPass>(); 461 } 462