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/Conversion/VectorToROCDL/VectorToROCDL.h" 19 #include "mlir/Dialect/GPU/GPUDialect.h" 20 #include "mlir/Dialect/GPU/Passes.h" 21 #include "mlir/Dialect/LLVMIR/ROCDLDialect.h" 22 #include "mlir/Dialect/Vector/VectorOps.h" 23 #include "mlir/Pass/Pass.h" 24 #include "mlir/Transforms/DialectConversion.h" 25 #include "mlir/Transforms/GreedyPatternRewriteDriver.h" 26 #include "llvm/Support/FormatVariadic.h" 27 28 #include "../GPUCommon/GPUOpsLowering.h" 29 #include "../GPUCommon/IndexIntrinsicsOpLowering.h" 30 #include "../GPUCommon/OpToFuncCallLowering.h" 31 #include "../PassDetail.h" 32 33 using namespace mlir; 34 35 namespace { 36 37 /// Import the GPU Ops to ROCDL Patterns. 38 #include "GPUToROCDL.cpp.inc" 39 40 // A pass that replaces all occurrences of GPU device operations with their 41 // corresponding ROCDL equivalent. 42 // 43 // This pass only handles device code and is not meant to be run on GPU host 44 // code. 45 struct LowerGpuOpsToROCDLOpsPass 46 : public ConvertGpuOpsToROCDLOpsBase<LowerGpuOpsToROCDLOpsPass> { 47 LowerGpuOpsToROCDLOpsPass() = default; 48 LowerGpuOpsToROCDLOpsPass(unsigned indexBitwidth) { 49 this->indexBitwidth = indexBitwidth; 50 } 51 52 void runOnOperation() override { 53 gpu::GPUModuleOp m = getOperation(); 54 55 /// Customize the bitwidth used for the device side index computations. 56 LowerToLLVMOptions options = {/*useBarePtrCallConv =*/false, 57 /*emitCWrappers =*/true, 58 /*indexBitwidth =*/indexBitwidth, 59 /*useAlignedAlloc =*/false}; 60 LLVMTypeConverter converter(m.getContext(), options); 61 62 OwningRewritePatternList patterns, llvmPatterns; 63 64 populateGpuRewritePatterns(m.getContext(), patterns); 65 applyPatternsAndFoldGreedily(m, std::move(patterns)); 66 67 populateVectorToLLVMConversionPatterns(converter, llvmPatterns); 68 populateVectorToROCDLConversionPatterns(converter, llvmPatterns); 69 populateStdToLLVMConversionPatterns(converter, llvmPatterns); 70 populateGpuToROCDLConversionPatterns(converter, llvmPatterns); 71 LLVMConversionTarget target(getContext()); 72 configureGpuToROCDLConversionLegality(target); 73 if (failed(applyPartialConversion(m, target, std::move(llvmPatterns)))) 74 signalPassFailure(); 75 } 76 }; 77 78 } // anonymous namespace 79 80 void mlir::configureGpuToROCDLConversionLegality(ConversionTarget &target) { 81 target.addIllegalOp<FuncOp>(); 82 target.addLegalDialect<::mlir::LLVM::LLVMDialect>(); 83 target.addLegalDialect<ROCDL::ROCDLDialect>(); 84 target.addIllegalDialect<gpu::GPUDialect>(); 85 target.addIllegalOp<LLVM::CosOp, LLVM::ExpOp, LLVM::FAbsOp, LLVM::FCeilOp, 86 LLVM::FFloorOp, LLVM::LogOp, LLVM::Log10Op, LLVM::Log2Op, 87 LLVM::PowOp, LLVM::SinOp, LLVM::SqrtOp>(); 88 89 // TODO: Remove once we support replacing non-root ops. 90 target.addLegalOp<gpu::YieldOp, gpu::GPUModuleOp, gpu::ModuleEndOp>(); 91 } 92 93 void mlir::populateGpuToROCDLConversionPatterns( 94 LLVMTypeConverter &converter, OwningRewritePatternList &patterns) { 95 populateWithGenerated(converter.getDialect()->getContext(), patterns); 96 patterns.insert< 97 GPUIndexIntrinsicOpLowering<gpu::ThreadIdOp, ROCDL::ThreadIdXOp, 98 ROCDL::ThreadIdYOp, ROCDL::ThreadIdZOp>, 99 GPUIndexIntrinsicOpLowering<gpu::BlockDimOp, ROCDL::BlockDimXOp, 100 ROCDL::BlockDimYOp, ROCDL::BlockDimZOp>, 101 GPUIndexIntrinsicOpLowering<gpu::BlockIdOp, ROCDL::BlockIdXOp, 102 ROCDL::BlockIdYOp, ROCDL::BlockIdZOp>, 103 GPUIndexIntrinsicOpLowering<gpu::GridDimOp, ROCDL::GridDimXOp, 104 ROCDL::GridDimYOp, ROCDL::GridDimZOp>, 105 GPUFuncOpLowering<5>, GPUReturnOpLowering>(converter); 106 patterns.insert<OpToFuncCallLowering<AbsFOp>>(converter, "__ocml_fabs_f32", 107 "__ocml_fabs_f64"); 108 patterns.insert<OpToFuncCallLowering<AtanOp>>(converter, "__ocml_atan_f32", 109 "__ocml_atan_f64"); 110 patterns.insert<OpToFuncCallLowering<Atan2Op>>(converter, "__ocml_atan2_f32", 111 "__ocml_atan2_f64"); 112 patterns.insert<OpToFuncCallLowering<CeilFOp>>(converter, "__ocml_ceil_f32", 113 "__ocml_ceil_f64"); 114 patterns.insert<OpToFuncCallLowering<CosOp>>(converter, "__ocml_cos_f32", 115 "__ocml_cos_f64"); 116 patterns.insert<OpToFuncCallLowering<ExpOp>>(converter, "__ocml_exp_f32", 117 "__ocml_exp_f64"); 118 patterns.insert<OpToFuncCallLowering<FloorFOp>>(converter, "__ocml_floor_f32", 119 "__ocml_floor_f64"); 120 patterns.insert<OpToFuncCallLowering<LogOp>>(converter, "__ocml_log_f32", 121 "__ocml_log_f64"); 122 patterns.insert<OpToFuncCallLowering<Log10Op>>(converter, "__ocml_log10_f32", 123 "__ocml_log10_f64"); 124 patterns.insert<OpToFuncCallLowering<Log2Op>>(converter, "__ocml_log2_f32", 125 "__ocml_log2_f64"); 126 patterns.insert<OpToFuncCallLowering<PowFOp>>(converter, "__ocml_pow_f32", 127 "__ocml_pow_f64"); 128 patterns.insert<OpToFuncCallLowering<RsqrtOp>>(converter, "__ocml_rsqrt_f32", 129 "__ocml_rsqrt_f64"); 130 patterns.insert<OpToFuncCallLowering<SinOp>>(converter, "__ocml_sin_f32", 131 "__ocml_sin_f64"); 132 patterns.insert<OpToFuncCallLowering<SqrtOp>>(converter, "__ocml_sqrt_f32", 133 "__ocml_sqrt_f64"); 134 patterns.insert<OpToFuncCallLowering<TanhOp>>(converter, "__ocml_tanh_f32", 135 "__ocml_tanh_f64"); 136 } 137 138 std::unique_ptr<OperationPass<gpu::GPUModuleOp>> 139 mlir::createLowerGpuOpsToROCDLOpsPass(unsigned indexBitwidth) { 140 return std::make_unique<LowerGpuOpsToROCDLOpsPass>(indexBitwidth); 141 } 142