1 //===- ROCDLDialect.cpp - ROCDL IR Ops and Dialect registration -----------===//
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 defines the types and operation details for the ROCDL IR dialect in
10 // MLIR, and the LLVM IR dialect.  It also registers the dialect.
11 //
12 // The ROCDL dialect only contains GPU specific additions on top of the general
13 // LLVM dialect.
14 //
15 //===----------------------------------------------------------------------===//
16 
17 #include "mlir/Dialect/LLVMIR/ROCDLDialect.h"
18 
19 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
20 #include "mlir/IR/Builders.h"
21 #include "mlir/IR/MLIRContext.h"
22 #include "mlir/IR/Operation.h"
23 #include "mlir/IR/StandardTypes.h"
24 #include "llvm/AsmParser/Parser.h"
25 #include "llvm/IR/Attributes.h"
26 #include "llvm/IR/Function.h"
27 #include "llvm/IR/Type.h"
28 #include "llvm/Support/SourceMgr.h"
29 
30 using namespace mlir;
31 using namespace ROCDL;
32 
33 //===----------------------------------------------------------------------===//
34 // Parsing for ROCDL ops
35 //===----------------------------------------------------------------------===//
36 
37 static LLVM::LLVMDialect *getLlvmDialect(OpAsmParser &parser) {
38   return parser.getBuilder()
39       .getContext()
40       ->getRegisteredDialect<LLVM::LLVMDialect>();
41 }
42 
43 // <operation> ::=
44 //     `llvm.amdgcn.buffer.load.* %rsrc, %vindex, %offset, %glc, %slc :
45 //     result_type`
46 static ParseResult parseROCDLMubufLoadOp(OpAsmParser &parser,
47                                          OperationState &result) {
48   SmallVector<OpAsmParser::OperandType, 8> ops;
49   Type type;
50   if (parser.parseOperandList(ops, 5) || parser.parseColonType(type) ||
51       parser.addTypeToList(type, result.types))
52     return failure();
53 
54   auto int32Ty = LLVM::LLVMType::getInt32Ty(getLlvmDialect(parser));
55   auto int1Ty = LLVM::LLVMType::getInt1Ty(getLlvmDialect(parser));
56   auto i32x4Ty = LLVM::LLVMType::getVectorTy(int32Ty, 4);
57   return parser.resolveOperands(ops,
58                                 {i32x4Ty, int32Ty, int32Ty, int1Ty, int1Ty},
59                                 parser.getNameLoc(), result.operands);
60 }
61 
62 // <operation> ::=
63 //     `llvm.amdgcn.buffer.store.* %vdata, %rsrc, %vindex, %offset, %glc, %slc :
64 //     result_type`
65 static ParseResult parseROCDLMubufStoreOp(OpAsmParser &parser,
66                                           OperationState &result) {
67   SmallVector<OpAsmParser::OperandType, 8> ops;
68   Type type;
69   if (parser.parseOperandList(ops, 6) || parser.parseColonType(type))
70     return failure();
71 
72   auto int32Ty = LLVM::LLVMType::getInt32Ty(getLlvmDialect(parser));
73   auto int1Ty = LLVM::LLVMType::getInt1Ty(getLlvmDialect(parser));
74   auto i32x4Ty = LLVM::LLVMType::getVectorTy(int32Ty, 4);
75 
76   if (parser.resolveOperands(ops,
77                              {type, i32x4Ty, int32Ty, int32Ty, int1Ty, int1Ty},
78                              parser.getNameLoc(), result.operands))
79     return failure();
80   return success();
81 }
82 
83 //===----------------------------------------------------------------------===//
84 // ROCDLDialect initialization, type parsing, and registration.
85 //===----------------------------------------------------------------------===//
86 
87 // TODO(herhut): This should be the llvm.rocdl dialect once this is supported.
88 ROCDLDialect::ROCDLDialect(MLIRContext *context) : Dialect("rocdl", context) {
89   addOperations<
90 #define GET_OP_LIST
91 #include "mlir/Dialect/LLVMIR/ROCDLOps.cpp.inc"
92       >();
93 
94   // Support unknown operations because not all ROCDL operations are registered.
95   allowUnknownOperations();
96 }
97 
98 namespace mlir {
99 namespace ROCDL {
100 #define GET_OP_CLASSES
101 #include "mlir/Dialect/LLVMIR/ROCDLOps.cpp.inc"
102 } // namespace ROCDL
103 } // namespace mlir
104 
105