1 //===- LowerGpuOpsToROCDLOps.cpp - MLIR GPU to ROCDL lowering passes ------===// 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 generate ROCDLIR operations for higher-level 10 // GPU operations. 11 // 12 //===----------------------------------------------------------------------===// 13 14 #include "mlir/Conversion/GPUToROCDL/GPUToROCDLPass.h" 15 16 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h" 17 #include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h" 18 #include "mlir/Dialect/GPU/GPUDialect.h" 19 #include "mlir/Dialect/GPU/Passes.h" 20 #include "mlir/Dialect/LLVMIR/ROCDLDialect.h" 21 #include "mlir/Dialect/Vector/VectorOps.h" 22 #include "mlir/Pass/Pass.h" 23 #include "mlir/Transforms/DialectConversion.h" 24 #include "llvm/Support/FormatVariadic.h" 25 26 #include "../GPUCommon/GPUOpsLowering.h" 27 #include "../GPUCommon/IndexIntrinsicsOpLowering.h" 28 #include "../GPUCommon/OpToFuncCallLowering.h" 29 #include "../PassDetail.h" 30 31 using namespace mlir; 32 33 namespace { 34 35 /// Import the GPU Ops to ROCDL Patterns. 36 #include "GPUToROCDL.cpp.inc" 37 38 // A pass that replaces all occurrences of GPU device operations with their 39 // corresponding ROCDL equivalent. 40 // 41 // This pass only handles device code and is not meant to be run on GPU host 42 // code. 43 class LowerGpuOpsToROCDLOpsPass 44 : public ConvertGpuOpsToROCDLOpsBase<LowerGpuOpsToROCDLOpsPass> { 45 public: 46 void runOnOperation() override { 47 gpu::GPUModuleOp m = getOperation(); 48 49 LLVMTypeConverter converter(m.getContext()); 50 51 OwningRewritePatternList patterns; 52 53 populateGpuRewritePatterns(m.getContext(), patterns); 54 applyPatternsAndFoldGreedily(m, patterns); 55 patterns.clear(); 56 57 populateVectorToLLVMConversionPatterns(converter, patterns); 58 populateStdToLLVMConversionPatterns(converter, patterns); 59 populateGpuToROCDLConversionPatterns(converter, patterns); 60 LLVMConversionTarget target(getContext()); 61 target.addIllegalDialect<gpu::GPUDialect>(); 62 target.addIllegalOp<LLVM::CosOp, LLVM::ExpOp, LLVM::FAbsOp, LLVM::FCeilOp, 63 LLVM::LogOp, LLVM::Log10Op, LLVM::Log2Op>(); 64 target.addIllegalOp<FuncOp>(); 65 target.addLegalDialect<ROCDL::ROCDLDialect>(); 66 // TODO(whchung): Remove once we support replacing non-root ops. 67 target.addLegalOp<gpu::YieldOp, gpu::GPUModuleOp, gpu::ModuleEndOp>(); 68 if (failed(applyPartialConversion(m, target, patterns, &converter))) 69 signalPassFailure(); 70 } 71 }; 72 73 } // anonymous namespace 74 75 void mlir::populateGpuToROCDLConversionPatterns( 76 LLVMTypeConverter &converter, OwningRewritePatternList &patterns) { 77 populateWithGenerated(converter.getDialect()->getContext(), &patterns); 78 patterns.insert< 79 GPUIndexIntrinsicOpLowering<gpu::ThreadIdOp, ROCDL::ThreadIdXOp, 80 ROCDL::ThreadIdYOp, ROCDL::ThreadIdZOp>, 81 GPUIndexIntrinsicOpLowering<gpu::BlockDimOp, ROCDL::BlockDimXOp, 82 ROCDL::BlockDimYOp, ROCDL::BlockDimZOp>, 83 GPUIndexIntrinsicOpLowering<gpu::BlockIdOp, ROCDL::BlockIdXOp, 84 ROCDL::BlockIdYOp, ROCDL::BlockIdZOp>, 85 GPUIndexIntrinsicOpLowering<gpu::GridDimOp, ROCDL::GridDimXOp, 86 ROCDL::GridDimYOp, ROCDL::GridDimZOp>, 87 GPUFuncOpLowering<5>, GPUReturnOpLowering>(converter); 88 patterns.insert<OpToFuncCallLowering<AbsFOp>>(converter, "__ocml_fabs_f32", 89 "__ocml_fabs_f64"); 90 patterns.insert<OpToFuncCallLowering<CeilFOp>>(converter, "__ocml_ceil_f32", 91 "__ocml_ceil_f64"); 92 patterns.insert<OpToFuncCallLowering<CosOp>>(converter, "__ocml_cos_f32", 93 "__ocml_cos_f64"); 94 patterns.insert<OpToFuncCallLowering<ExpOp>>(converter, "__ocml_exp_f32", 95 "__ocml_exp_f64"); 96 patterns.insert<OpToFuncCallLowering<LogOp>>(converter, "__ocml_log_f32", 97 "__ocml_log_f64"); 98 patterns.insert<OpToFuncCallLowering<Log10Op>>(converter, "__ocml_log10_f32", 99 "__ocml_log10_f64"); 100 patterns.insert<OpToFuncCallLowering<Log2Op>>(converter, "__ocml_log2_f32", 101 "__ocml_log2_f64"); 102 patterns.insert<OpToFuncCallLowering<TanhOp>>(converter, "__ocml_tanh_f32", 103 "__ocml_tanh_f64"); 104 } 105 106 std::unique_ptr<OperationPass<gpu::GPUModuleOp>> 107 mlir::createLowerGpuOpsToROCDLOpsPass() { 108 return std::make_unique<LowerGpuOpsToROCDLOpsPass>(); 109 } 110