1 //===- SCFToGPUPass.cpp - Convert a loop nest to a GPU kernel -----------===// 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 #include "mlir/Conversion/SCFToGPU/SCFToGPUPass.h" 10 #include "../PassDetail.h" 11 #include "mlir/Conversion/SCFToGPU/SCFToGPU.h" 12 #include "mlir/Dialect/Affine/IR/AffineOps.h" 13 #include "mlir/Dialect/GPU/GPUDialect.h" 14 #include "mlir/Dialect/SCF/SCF.h" 15 #include "mlir/Dialect/StandardOps/IR/Ops.h" 16 #include "mlir/Transforms/DialectConversion.h" 17 18 #include "llvm/ADT/ArrayRef.h" 19 #include "llvm/Support/CommandLine.h" 20 21 using namespace mlir; 22 using namespace mlir::scf; 23 24 namespace { 25 // A pass that traverses top-level loops in the function and converts them to 26 // GPU launch operations. Nested launches are not allowed, so this does not 27 // walk the function recursively to avoid considering nested loops. 28 struct ForLoopMapper : public ConvertAffineForToGPUBase<ForLoopMapper> { 29 ForLoopMapper() = default; 30 ForLoopMapper(unsigned numBlockDims, unsigned numThreadDims) { 31 this->numBlockDims = numBlockDims; 32 this->numThreadDims = numThreadDims; 33 } 34 35 void runOnFunction() override { 36 for (Operation &op : llvm::make_early_inc_range(getFunction().getOps())) { 37 if (auto forOp = dyn_cast<AffineForOp>(&op)) { 38 if (failed(convertAffineLoopNestToGPULaunch(forOp, numBlockDims, 39 numThreadDims))) 40 signalPassFailure(); 41 } 42 } 43 } 44 }; 45 46 struct ParallelLoopToGpuPass 47 : public ConvertParallelLoopToGpuBase<ParallelLoopToGpuPass> { 48 void runOnOperation() override { 49 OwningRewritePatternList patterns; 50 populateParallelLoopToGPUPatterns(patterns, &getContext()); 51 ConversionTarget target(getContext()); 52 target.addLegalDialect<StandardOpsDialect>(); 53 target.addLegalDialect<AffineDialect>(); 54 target.addLegalDialect<gpu::GPUDialect>(); 55 target.addLegalDialect<scf::SCFDialect>(); 56 target.addIllegalOp<scf::ParallelOp>(); 57 if (failed(applyPartialConversion(getOperation(), target, patterns))) 58 signalPassFailure(); 59 } 60 }; 61 62 } // namespace 63 64 std::unique_ptr<OperationPass<FuncOp>> 65 mlir::createAffineForToGPUPass(unsigned numBlockDims, unsigned numThreadDims) { 66 return std::make_unique<ForLoopMapper>(numBlockDims, numThreadDims); 67 } 68 std::unique_ptr<OperationPass<FuncOp>> mlir::createAffineForToGPUPass() { 69 return std::make_unique<ForLoopMapper>(); 70 } 71 72 std::unique_ptr<Pass> mlir::createParallelLoopToGpuPass() { 73 return std::make_unique<ParallelLoopToGpuPass>(); 74 } 75