1 //===- LowerGpuOpsToROCDLOps.cpp - MLIR GPU to ROCDL lowering passes ------===//
2 //
3 // Copyright 2019 The MLIR Authors.
4 //
5 // Licensed under the Apache License, Version 2.0 (the "License");
6 // you may not use this file except in compliance with the License.
7 // You may obtain a copy of the License at
8 //
9 //   http://www.apache.org/licenses/LICENSE-2.0
10 //
11 // Unless required by applicable law or agreed to in writing, software
12 // distributed under the License is distributed on an "AS IS" BASIS,
13 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14 // See the License for the specific language governing permissions and
15 // limitations under the License.
16 // =============================================================================
17 //
18 // This file implements a pass to generate ROCDLIR operations for higher-level
19 // GPU operations.
20 //
21 //===----------------------------------------------------------------------===//
22 
23 #include "mlir/Conversion/GPUToROCDL/GPUToROCDLPass.h"
24 
25 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h"
26 #include "mlir/Dialect/GPU/GPUDialect.h"
27 #include "mlir/Dialect/LLVMIR/ROCDLDialect.h"
28 #include "mlir/Pass/Pass.h"
29 #include "mlir/Transforms/DialectConversion.h"
30 
31 #include "../GPUCommon/IndexIntrinsicsOpLowering.h"
32 
33 using namespace mlir;
34 
35 namespace {
36 
37 // A pass that replaces all occurrences of GPU device operations with their
38 // corresponding ROCDL equivalent.
39 //
40 // This pass only handles device code and is not meant to be run on GPU host
41 // code.
42 class LowerGpuOpsToROCDLOpsPass : public ModulePass<LowerGpuOpsToROCDLOpsPass> {
43 public:
44   void runOnModule() override {
45     ModuleOp m = getModule();
46     if (!m.getAttrOfType<UnitAttr>(gpu::GPUDialect::getKernelModuleAttrName()))
47       return;
48 
49     OwningRewritePatternList patterns;
50     LLVMTypeConverter converter(m.getContext());
51     populateStdToLLVMConversionPatterns(converter, patterns);
52     patterns.insert<
53         GPUIndexIntrinsicOpLowering<gpu::ThreadId, ROCDL::ThreadIdXOp,
54                                     ROCDL::ThreadIdYOp, ROCDL::ThreadIdZOp>,
55         GPUIndexIntrinsicOpLowering<gpu::BlockDim, ROCDL::BlockDimXOp,
56                                     ROCDL::BlockDimYOp, ROCDL::BlockDimZOp>,
57         GPUIndexIntrinsicOpLowering<gpu::BlockId, ROCDL::BlockIdXOp,
58                                     ROCDL::BlockIdYOp, ROCDL::BlockIdZOp>,
59         GPUIndexIntrinsicOpLowering<gpu::GridDim, ROCDL::GridDimXOp,
60                                     ROCDL::GridDimYOp, ROCDL::GridDimZOp>>(
61         converter);
62 
63     ConversionTarget target(getContext());
64     target.addLegalDialect<LLVM::LLVMDialect, ROCDL::ROCDLDialect>();
65     target.addDynamicallyLegalOp<FuncOp>(
66         [&](FuncOp op) { return converter.isSignatureLegal(op.getType()); });
67     if (failed(applyPartialConversion(m, target, patterns, &converter)))
68       signalPassFailure();
69   }
70 };
71 
72 } // anonymous namespace
73 
74 std::unique_ptr<OpPassBase<ModuleOp>> mlir::createLowerGpuOpsToROCDLOpsPass() {
75   return std::make_unique<LowerGpuOpsToROCDLOpsPass>();
76 }
77 
78 static PassRegistration<LowerGpuOpsToROCDLOpsPass>
79     pass("lower-gpu-ops-to-rocdl-ops",
80          "Generate ROCDL operations for gpu operations");
81