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