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/ControlFlowToLLVM/ControlFlowToLLVM.h" 15 #include "mlir/Conversion/GPUToROCDL/GPUToROCDLPass.h" 16 17 #include "mlir/Conversion/AMDGPUToROCDL/AMDGPUToROCDL.h" 18 #include "mlir/Conversion/ArithmeticToLLVM/ArithmeticToLLVM.h" 19 #include "mlir/Conversion/FuncToLLVM/ConvertFuncToLLVM.h" 20 #include "mlir/Conversion/LLVMCommon/ConversionTarget.h" 21 #include "mlir/Conversion/LLVMCommon/LoweringOptions.h" 22 #include "mlir/Conversion/LLVMCommon/Pattern.h" 23 #include "mlir/Conversion/LLVMCommon/TypeConverter.h" 24 #include "mlir/Conversion/MemRefToLLVM/MemRefToLLVM.h" 25 #include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h" 26 #include "mlir/Conversion/VectorToROCDL/VectorToROCDL.h" 27 #include "mlir/Dialect/Func/IR/FuncOps.h" 28 #include "mlir/Dialect/GPU/IR/GPUDialect.h" 29 #include "mlir/Dialect/GPU/Transforms/Passes.h" 30 #include "mlir/Dialect/LLVMIR/ROCDLDialect.h" 31 #include "mlir/Dialect/Math/IR/Math.h" 32 #include "mlir/Dialect/Vector/IR/VectorOps.h" 33 #include "mlir/Pass/Pass.h" 34 #include "mlir/Transforms/DialectConversion.h" 35 #include "mlir/Transforms/GreedyPatternRewriteDriver.h" 36 #include "llvm/Support/FormatVariadic.h" 37 38 #include "../GPUCommon/GPUOpsLowering.h" 39 #include "../GPUCommon/IndexIntrinsicsOpLowering.h" 40 #include "../GPUCommon/OpToFuncCallLowering.h" 41 #include "../PassDetail.h" 42 43 using namespace mlir; 44 45 namespace { 46 47 /// Import the GPU Ops to ROCDL Patterns. 48 #include "GPUToROCDL.cpp.inc" 49 50 // A pass that replaces all occurrences of GPU device operations with their 51 // corresponding ROCDL equivalent. 52 // 53 // This pass only handles device code and is not meant to be run on GPU host 54 // code. 55 struct LowerGpuOpsToROCDLOpsPass 56 : public ConvertGpuOpsToROCDLOpsBase<LowerGpuOpsToROCDLOpsPass> { 57 LowerGpuOpsToROCDLOpsPass() = default; 58 LowerGpuOpsToROCDLOpsPass(unsigned indexBitwidth, gpu::amd::Runtime runtime) { 59 this->indexBitwidth = indexBitwidth; 60 this->runtime = runtime; 61 } 62 63 void runOnOperation() override { 64 gpu::GPUModuleOp m = getOperation(); 65 66 // Request C wrapper emission. 67 for (auto func : m.getOps<func::FuncOp>()) { 68 func->setAttr(LLVM::LLVMDialect::getEmitCWrapperAttrName(), 69 UnitAttr::get(&getContext())); 70 } 71 72 /// Customize the bitwidth used for the device side index computations. 73 LowerToLLVMOptions options( 74 m.getContext(), 75 DataLayout(cast<DataLayoutOpInterface>(m.getOperation()))); 76 if (indexBitwidth != kDeriveIndexBitwidthFromDataLayout) 77 options.overrideIndexBitwidth(indexBitwidth); 78 LLVMTypeConverter converter(m.getContext(), options); 79 80 RewritePatternSet patterns(m.getContext()); 81 RewritePatternSet llvmPatterns(m.getContext()); 82 83 populateGpuRewritePatterns(patterns); 84 (void)applyPatternsAndFoldGreedily(m, std::move(patterns)); 85 86 mlir::arith::populateArithmeticToLLVMConversionPatterns(converter, 87 llvmPatterns); 88 populateAMDGPUToROCDLConversionPatterns(converter, llvmPatterns); 89 populateVectorToLLVMConversionPatterns(converter, llvmPatterns); 90 populateVectorToROCDLConversionPatterns(converter, llvmPatterns); 91 cf::populateControlFlowToLLVMConversionPatterns(converter, llvmPatterns); 92 populateFuncToLLVMConversionPatterns(converter, llvmPatterns); 93 populateMemRefToLLVMConversionPatterns(converter, llvmPatterns); 94 populateGpuToROCDLConversionPatterns(converter, llvmPatterns, runtime); 95 LLVMConversionTarget target(getContext()); 96 configureGpuToROCDLConversionLegality(target); 97 if (failed(applyPartialConversion(m, target, std::move(llvmPatterns)))) 98 signalPassFailure(); 99 } 100 }; 101 102 } // namespace 103 104 void mlir::configureGpuToROCDLConversionLegality(ConversionTarget &target) { 105 target.addIllegalOp<func::FuncOp>(); 106 target.addLegalDialect<::mlir::LLVM::LLVMDialect>(); 107 target.addLegalDialect<ROCDL::ROCDLDialect>(); 108 target.addIllegalDialect<gpu::GPUDialect>(); 109 target.addIllegalOp<LLVM::CosOp, LLVM::ExpOp, LLVM::Exp2Op, LLVM::FAbsOp, 110 LLVM::FCeilOp, LLVM::FFloorOp, LLVM::LogOp, LLVM::Log10Op, 111 LLVM::Log2Op, LLVM::PowOp, LLVM::SinOp, LLVM::SqrtOp>(); 112 113 // TODO: Remove once we support replacing non-root ops. 114 target.addLegalOp<gpu::YieldOp, gpu::GPUModuleOp, gpu::ModuleEndOp>(); 115 } 116 117 void mlir::populateGpuToROCDLConversionPatterns( 118 LLVMTypeConverter &converter, RewritePatternSet &patterns, 119 mlir::gpu::amd::Runtime runtime) { 120 using mlir::gpu::amd::Runtime; 121 122 populateWithGenerated(patterns); 123 patterns 124 .add<GPUIndexIntrinsicOpLowering<gpu::ThreadIdOp, ROCDL::ThreadIdXOp, 125 ROCDL::ThreadIdYOp, ROCDL::ThreadIdZOp>, 126 GPUIndexIntrinsicOpLowering<gpu::BlockDimOp, ROCDL::BlockDimXOp, 127 ROCDL::BlockDimYOp, ROCDL::BlockDimZOp>, 128 GPUIndexIntrinsicOpLowering<gpu::BlockIdOp, ROCDL::BlockIdXOp, 129 ROCDL::BlockIdYOp, ROCDL::BlockIdZOp>, 130 GPUIndexIntrinsicOpLowering<gpu::GridDimOp, ROCDL::GridDimXOp, 131 ROCDL::GridDimYOp, ROCDL::GridDimZOp>, 132 GPUReturnOpLowering>(converter); 133 patterns.add<GPUFuncOpLowering>( 134 converter, /*allocaAddrSpace=*/5, 135 StringAttr::get(&converter.getContext(), 136 ROCDL::ROCDLDialect::getKernelFuncAttrName())); 137 if (Runtime::HIP == runtime) { 138 patterns.add<GPUPrintfOpToHIPLowering>(converter); 139 } else if (Runtime::OpenCL == runtime) { 140 // Use address space = 4 to match the OpenCL definition of printf() 141 patterns.add<GPUPrintfOpToLLVMCallLowering>(converter, /*addressSpace=*/4); 142 } 143 144 patterns.add<OpToFuncCallLowering<math::AbsOp>>(converter, "__ocml_fabs_f32", 145 "__ocml_fabs_f64"); 146 patterns.add<OpToFuncCallLowering<math::AtanOp>>(converter, "__ocml_atan_f32", 147 "__ocml_atan_f64"); 148 patterns.add<OpToFuncCallLowering<math::Atan2Op>>( 149 converter, "__ocml_atan2_f32", "__ocml_atan2_f64"); 150 patterns.add<OpToFuncCallLowering<math::CeilOp>>(converter, "__ocml_ceil_f32", 151 "__ocml_ceil_f64"); 152 patterns.add<OpToFuncCallLowering<math::CosOp>>(converter, "__ocml_cos_f32", 153 "__ocml_cos_f64"); 154 patterns.add<OpToFuncCallLowering<math::ExpOp>>(converter, "__ocml_exp_f32", 155 "__ocml_exp_f64"); 156 patterns.add<OpToFuncCallLowering<math::Exp2Op>>(converter, "__ocml_exp2_f32", 157 "__ocml_exp2_f64"); 158 patterns.add<OpToFuncCallLowering<math::ExpM1Op>>( 159 converter, "__ocml_expm1_f32", "__ocml_expm1_f64"); 160 patterns.add<OpToFuncCallLowering<math::FloorOp>>( 161 converter, "__ocml_floor_f32", "__ocml_floor_f64"); 162 patterns.add<OpToFuncCallLowering<math::LogOp>>(converter, "__ocml_log_f32", 163 "__ocml_log_f64"); 164 patterns.add<OpToFuncCallLowering<math::Log10Op>>( 165 converter, "__ocml_log10_f32", "__ocml_log10_f64"); 166 patterns.add<OpToFuncCallLowering<math::Log1pOp>>( 167 converter, "__ocml_log1p_f32", "__ocml_log1p_f64"); 168 patterns.add<OpToFuncCallLowering<math::Log2Op>>(converter, "__ocml_log2_f32", 169 "__ocml_log2_f64"); 170 patterns.add<OpToFuncCallLowering<math::PowFOp>>(converter, "__ocml_pow_f32", 171 "__ocml_pow_f64"); 172 patterns.add<OpToFuncCallLowering<math::RsqrtOp>>( 173 converter, "__ocml_rsqrt_f32", "__ocml_rsqrt_f64"); 174 patterns.add<OpToFuncCallLowering<math::SinOp>>(converter, "__ocml_sin_f32", 175 "__ocml_sin_f64"); 176 patterns.add<OpToFuncCallLowering<math::SqrtOp>>(converter, "__ocml_sqrt_f32", 177 "__ocml_sqrt_f64"); 178 patterns.add<OpToFuncCallLowering<math::TanhOp>>(converter, "__ocml_tanh_f32", 179 "__ocml_tanh_f64"); 180 } 181 182 std::unique_ptr<OperationPass<gpu::GPUModuleOp>> 183 mlir::createLowerGpuOpsToROCDLOpsPass(unsigned indexBitwidth, 184 gpu::amd::Runtime runtime) { 185 return std::make_unique<LowerGpuOpsToROCDLOpsPass>(indexBitwidth, runtime); 186 } 187