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