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