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/Dialect/GPU/GPUDialect.h" 18 #include "mlir/Dialect/LLVMIR/ROCDLDialect.h" 19 #include "mlir/Pass/Pass.h" 20 #include "mlir/Transforms/DialectConversion.h" 21 22 #include "../GPUCommon/IndexIntrinsicsOpLowering.h" 23 #include "../GPUCommon/OpToFuncCallLowering.h" 24 #include "../PassDetail.h" 25 26 using namespace mlir; 27 28 namespace { 29 30 // A pass that replaces all occurrences of GPU device operations with their 31 // corresponding ROCDL equivalent. 32 // 33 // This pass only handles device code and is not meant to be run on GPU host 34 // code. 35 class LowerGpuOpsToROCDLOpsPass 36 : public ConvertGpuOpsToROCDLOpsBase<LowerGpuOpsToROCDLOpsPass> { 37 public: 38 void runOnOperation() override { 39 gpu::GPUModuleOp m = getOperation(); 40 41 OwningRewritePatternList patterns; 42 LLVMTypeConverter converter(m.getContext()); 43 populateStdToLLVMConversionPatterns(converter, patterns); 44 patterns.insert< 45 GPUIndexIntrinsicOpLowering<gpu::ThreadIdOp, ROCDL::ThreadIdXOp, 46 ROCDL::ThreadIdYOp, ROCDL::ThreadIdZOp>, 47 GPUIndexIntrinsicOpLowering<gpu::BlockDimOp, ROCDL::BlockDimXOp, 48 ROCDL::BlockDimYOp, ROCDL::BlockDimZOp>, 49 GPUIndexIntrinsicOpLowering<gpu::BlockIdOp, ROCDL::BlockIdXOp, 50 ROCDL::BlockIdYOp, ROCDL::BlockIdZOp>, 51 GPUIndexIntrinsicOpLowering<gpu::GridDimOp, ROCDL::GridDimXOp, 52 ROCDL::GridDimYOp, ROCDL::GridDimZOp>>( 53 converter); 54 patterns.insert<OpToFuncCallLowering<AbsFOp>>(converter, "__ocml_fabs_f32", 55 "__ocml_fabs_f64"); 56 patterns.insert<OpToFuncCallLowering<CeilFOp>>(converter, "__ocml_ceil_f32", 57 "__ocml_ceil_f64"); 58 patterns.insert<OpToFuncCallLowering<CosOp>>(converter, "__ocml_cos_f32", 59 "__ocml_cos_f64"); 60 patterns.insert<OpToFuncCallLowering<ExpOp>>(converter, "__ocml_exp_f32", 61 "__ocml_exp_f64"); 62 patterns.insert<OpToFuncCallLowering<LogOp>>(converter, "__ocml_log_f32", 63 "__ocml_log_f64"); 64 patterns.insert<OpToFuncCallLowering<Log10Op>>( 65 converter, "__ocml_log10_f32", "__ocml_log10_f64"); 66 patterns.insert<OpToFuncCallLowering<Log2Op>>(converter, "__ocml_log2_f32", 67 "__ocml_log2_f64"); 68 patterns.insert<OpToFuncCallLowering<TanhOp>>(converter, "__ocml_tanh_f32", 69 "__ocml_tanh_f64"); 70 71 ConversionTarget target(getContext()); 72 target.addLegalDialect<LLVM::LLVMDialect, ROCDL::ROCDLDialect>(); 73 target.addIllegalOp<LLVM::CosOp, LLVM::ExpOp, LLVM::FAbsOp, LLVM::FCeilOp, 74 LLVM::LogOp, LLVM::Log10Op, LLVM::Log2Op>(); 75 target.addIllegalOp<FuncOp>(); 76 if (failed(applyPartialConversion(m, target, patterns, &converter))) 77 signalPassFailure(); 78 } 79 }; 80 81 } // anonymous namespace 82 83 std::unique_ptr<OperationPass<gpu::GPUModuleOp>> 84 mlir::createLowerGpuOpsToROCDLOpsPass() { 85 return std::make_unique<LowerGpuOpsToROCDLOpsPass>(); 86 } 87