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/Dialect/GPU/GPUDialect.h"
18 #include "mlir/Dialect/LLVMIR/ROCDLDialect.h"
19 #include "mlir/Pass/Pass.h"
20 #include "mlir/Transforms/DialectConversion.h"
21 
22 #include "../GPUCommon/IndexIntrinsicsOpLowering.h"
23 #include "../GPUCommon/OpToFuncCallLowering.h"
24 
25 using namespace mlir;
26 
27 namespace {
28 
29 // A pass that replaces all occurrences of GPU device operations with their
30 // corresponding ROCDL equivalent.
31 //
32 // This pass only handles device code and is not meant to be run on GPU host
33 // code.
34 class LowerGpuOpsToROCDLOpsPass
35     : public OperationPass<LowerGpuOpsToROCDLOpsPass, gpu::GPUModuleOp> {
36 public:
37   void runOnOperation() override {
38     gpu::GPUModuleOp m = getOperation();
39 
40     OwningRewritePatternList patterns;
41     LLVMTypeConverter converter(m.getContext());
42     populateStdToLLVMConversionPatterns(converter, patterns);
43     patterns.insert<
44         GPUIndexIntrinsicOpLowering<gpu::ThreadIdOp, ROCDL::ThreadIdXOp,
45                                     ROCDL::ThreadIdYOp, ROCDL::ThreadIdZOp>,
46         GPUIndexIntrinsicOpLowering<gpu::BlockDimOp, ROCDL::BlockDimXOp,
47                                     ROCDL::BlockDimYOp, ROCDL::BlockDimZOp>,
48         GPUIndexIntrinsicOpLowering<gpu::BlockIdOp, ROCDL::BlockIdXOp,
49                                     ROCDL::BlockIdYOp, ROCDL::BlockIdZOp>,
50         GPUIndexIntrinsicOpLowering<gpu::GridDimOp, ROCDL::GridDimXOp,
51                                     ROCDL::GridDimYOp, ROCDL::GridDimZOp>>(
52         converter);
53     patterns.insert<OpToFuncCallLowering<AbsFOp>>(converter, "__ocml_fabs_f32",
54                                                   "__ocml_fabs_f64");
55     patterns.insert<OpToFuncCallLowering<CeilFOp>>(converter, "__ocml_ceil_f32",
56                                                    "__ocml_ceil_f64");
57     patterns.insert<OpToFuncCallLowering<CosOp>>(converter, "__ocml_cos_f32",
58                                                  "__ocml_cos_f64");
59     patterns.insert<OpToFuncCallLowering<ExpOp>>(converter, "__ocml_exp_f32",
60                                                  "__ocml_exp_f64");
61     patterns.insert<OpToFuncCallLowering<TanhOp>>(converter, "__ocml_tanh_f32",
62                                                   "__ocml_tanh_f64");
63 
64     ConversionTarget target(getContext());
65     target.addLegalDialect<LLVM::LLVMDialect, ROCDL::ROCDLDialect>();
66     target.addIllegalOp<LLVM::FAbsOp, LLVM::FCeilOp, LLVM::CosOp,
67                         LLVM::ExpOp>();
68     target.addDynamicallyLegalOp<LLVM::CallOp>(
69         gpu::filterIllegalLLVMIntrinsics({"tanh", "tanhf"}, m.getContext()));
70     target.addIllegalOp<FuncOp>();
71     if (failed(applyPartialConversion(m, target, patterns, &converter)))
72       signalPassFailure();
73   }
74 };
75 
76 } // anonymous namespace
77 
78 std::unique_ptr<OpPassBase<gpu::GPUModuleOp>>
79 mlir::createLowerGpuOpsToROCDLOpsPass() {
80   return std::make_unique<LowerGpuOpsToROCDLOpsPass>();
81 }
82 
83 static PassRegistration<LowerGpuOpsToROCDLOpsPass>
84     pass("convert-gpu-to-rocdl",
85          "Generate ROCDL operations for gpu operations");
86