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