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;
LowerGpuOpsToROCDLOpsPass__anonfaa64c570111::LowerGpuOpsToROCDLOpsPass48 LowerGpuOpsToROCDLOpsPass(unsigned indexBitwidth) {
49 this->indexBitwidth = indexBitwidth;
50 }
51
runOnOperation__anonfaa64c570111::LowerGpuOpsToROCDLOpsPass52 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 target.addIllegalDialect<gpu::GPUDialect>();
73 target.addIllegalOp<LLVM::CosOp, LLVM::ExpOp, LLVM::FAbsOp, LLVM::FCeilOp,
74 LLVM::FFloorOp, LLVM::LogOp, LLVM::Log10Op,
75 LLVM::Log2Op, LLVM::SinOp>();
76 target.addIllegalOp<FuncOp>();
77 target.addLegalDialect<ROCDL::ROCDLDialect>();
78 // TODO: Remove once we support replacing non-root ops.
79 target.addLegalOp<gpu::YieldOp, gpu::GPUModuleOp, gpu::ModuleEndOp>();
80 if (failed(applyPartialConversion(m, target, std::move(llvmPatterns))))
81 signalPassFailure();
82 }
83 };
84
85 } // anonymous namespace
86
populateGpuToROCDLConversionPatterns(LLVMTypeConverter & converter,OwningRewritePatternList & patterns)87 void mlir::populateGpuToROCDLConversionPatterns(
88 LLVMTypeConverter &converter, OwningRewritePatternList &patterns) {
89 populateWithGenerated(converter.getDialect()->getContext(), patterns);
90 patterns.insert<
91 GPUIndexIntrinsicOpLowering<gpu::ThreadIdOp, ROCDL::ThreadIdXOp,
92 ROCDL::ThreadIdYOp, ROCDL::ThreadIdZOp>,
93 GPUIndexIntrinsicOpLowering<gpu::BlockDimOp, ROCDL::BlockDimXOp,
94 ROCDL::BlockDimYOp, ROCDL::BlockDimZOp>,
95 GPUIndexIntrinsicOpLowering<gpu::BlockIdOp, ROCDL::BlockIdXOp,
96 ROCDL::BlockIdYOp, ROCDL::BlockIdZOp>,
97 GPUIndexIntrinsicOpLowering<gpu::GridDimOp, ROCDL::GridDimXOp,
98 ROCDL::GridDimYOp, ROCDL::GridDimZOp>,
99 GPUFuncOpLowering<5>, GPUReturnOpLowering>(converter);
100 patterns.insert<OpToFuncCallLowering<AbsFOp>>(converter, "__ocml_fabs_f32",
101 "__ocml_fabs_f64");
102 patterns.insert<OpToFuncCallLowering<CeilFOp>>(converter, "__ocml_ceil_f32",
103 "__ocml_ceil_f64");
104 patterns.insert<OpToFuncCallLowering<CosOp>>(converter, "__ocml_cos_f32",
105 "__ocml_cos_f64");
106 patterns.insert<OpToFuncCallLowering<ExpOp>>(converter, "__ocml_exp_f32",
107 "__ocml_exp_f64");
108 patterns.insert<OpToFuncCallLowering<FloorFOp>>(converter, "__ocml_floor_f32",
109 "__ocml_floor_f64");
110 patterns.insert<OpToFuncCallLowering<LogOp>>(converter, "__ocml_log_f32",
111 "__ocml_log_f64");
112 patterns.insert<OpToFuncCallLowering<Log10Op>>(converter, "__ocml_log10_f32",
113 "__ocml_log10_f64");
114 patterns.insert<OpToFuncCallLowering<Log2Op>>(converter, "__ocml_log2_f32",
115 "__ocml_log2_f64");
116 patterns.insert<OpToFuncCallLowering<SinOp>>(converter, "__ocml_sin_f32",
117 "__ocml_sin_f64");
118 patterns.insert<OpToFuncCallLowering<TanhOp>>(converter, "__ocml_tanh_f32",
119 "__ocml_tanh_f64");
120 }
121
122 std::unique_ptr<OperationPass<gpu::GPUModuleOp>>
createLowerGpuOpsToROCDLOpsPass(unsigned indexBitwidth)123 mlir::createLowerGpuOpsToROCDLOpsPass(unsigned indexBitwidth) {
124 return std::make_unique<LowerGpuOpsToROCDLOpsPass>(indexBitwidth);
125 }
126