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