mlir/lib/Dialect/LLVMIR/IR/ROCDLDialect.cpp

   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/GPU/IR/CompilationInterfaces.h"
  20 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
  21 #include "mlir/IR/Builders.h"
  22 #include "mlir/IR/BuiltinTypes.h"
  23 #include "mlir/IR/DialectImplementation.h"
  24 #include "mlir/IR/MLIRContext.h"
  25 #include "mlir/IR/Operation.h"
  26 #include "llvm/ADT/TypeSwitch.h"
  27 #include "llvm/AsmParser/Parser.h"
  28 #include "llvm/IR/Attributes.h"
  29 #include "llvm/IR/Function.h"
  30 #include "llvm/IR/Type.h"
  31 #include "llvm/Support/SourceMgr.h"
  32
  33 using namespace mlir;
  34 using namespace ROCDL;
  35
  36 #include "mlir/Dialect/LLVMIR/ROCDLOpsDialect.cpp.inc"
  37
  38 //===----------------------------------------------------------------------===//
  39 // Parsing for ROCDL ops
  40 //===----------------------------------------------------------------------===//
  41
  42 // <operation> ::=
  43 //     `llvm.amdgcn.buffer.load.* %rsrc, %vindex, %offset, %glc, %slc :
  44 //     result_type`
  45 ParseResult MubufLoadOp::parse(OpAsmParser &parser, OperationState &result) {
  46   SmallVector<OpAsmParser::UnresolvedOperand, 8> ops;
  47   Type type;
  48   if (parser.parseOperandList(ops, 5) || parser.parseColonType(type) ||
  49       parser.addTypeToList(type, result.types))
  50     return failure();
  51
  52   MLIRContext *context = parser.getContext();
  53   auto int32Ty = IntegerType::get(context, 32);
  54   auto int1Ty = IntegerType::get(context, 1);
  55   auto i32x4Ty = LLVM::getFixedVectorType(int32Ty, 4);
  56   return parser.resolveOperands(ops,
  57                                 {i32x4Ty, int32Ty, int32Ty, int1Ty, int1Ty},
  58                                 parser.getNameLoc(), result.operands);
  59 }
  60
  61 void MubufLoadOp::print(OpAsmPrinter &p) {
  62   p << " " << getOperands() << " : " << (*this)->getResultTypes();
  63 }
  64
  65 // <operation> ::=
  66 //     `llvm.amdgcn.buffer.store.* %vdata, %rsrc, %vindex, %offset, %glc, %slc :
  67 //     result_type`
  68 ParseResult MubufStoreOp::parse(OpAsmParser &parser, OperationState &result) {
  69   SmallVector<OpAsmParser::UnresolvedOperand, 8> ops;
  70   Type type;
  71   if (parser.parseOperandList(ops, 6) || parser.parseColonType(type))
  72     return failure();
  73
  74   MLIRContext *context = parser.getContext();
  75   auto int32Ty = IntegerType::get(context, 32);
  76   auto int1Ty = IntegerType::get(context, 1);
  77   auto i32x4Ty = LLVM::getFixedVectorType(int32Ty, 4);
  78
  79   if (parser.resolveOperands(ops,
  80                              {type, i32x4Ty, int32Ty, int32Ty, int1Ty, int1Ty},
  81                              parser.getNameLoc(), result.operands))
  82     return failure();
  83   return success();
  84 }
  85
  86 void MubufStoreOp::print(OpAsmPrinter &p) {
  87   p << " " << getOperands() << " : " << getVdata().getType();
  88 }
  89
  90 // <operation> ::=
  91 //     `llvm.amdgcn.raw.buffer.load.* %rsrc, %offset, %soffset, %aux
  92 //     : result_type`
  93 ParseResult RawBufferLoadOp::parse(OpAsmParser &parser,
  94                                    OperationState &result) {
  95   SmallVector<OpAsmParser::UnresolvedOperand, 4> ops;
  96   Type type;
  97   if (parser.parseOperandList(ops, 4) || parser.parseColonType(type) ||
  98       parser.addTypeToList(type, result.types))
  99     return failure();
 100
 101   auto bldr = parser.getBuilder();
 102   auto int32Ty = bldr.getI32Type();
 103   auto i32x4Ty = VectorType::get({4}, int32Ty);
 104   return parser.resolveOperands(ops, {i32x4Ty, int32Ty, int32Ty, int32Ty},
 105                                 parser.getNameLoc(), result.operands);
 106 }
 107
 108 void RawBufferLoadOp::print(OpAsmPrinter &p) {
 109   p << " " << getOperands() << " : " << getRes().getType();
 110 }
 111
 112 // <operation> ::=
 113 //     `llvm.amdgcn.raw.buffer.store.* %vdata, %rsrc,  %offset,
 114 //     %soffset, %aux : result_type`
 115 ParseResult RawBufferStoreOp::parse(OpAsmParser &parser,
 116                                     OperationState &result) {
 117   SmallVector<OpAsmParser::UnresolvedOperand, 5> ops;
 118   Type type;
 119   if (parser.parseOperandList(ops, 5) || parser.parseColonType(type))
 120     return failure();
 121
 122   auto bldr = parser.getBuilder();
 123   auto int32Ty = bldr.getI32Type();
 124   auto i32x4Ty = VectorType::get({4}, int32Ty);
 125
 126   if (parser.resolveOperands(ops, {type, i32x4Ty, int32Ty, int32Ty, int32Ty},
 127                              parser.getNameLoc(), result.operands))
 128     return failure();
 129   return success();
 130 }
 131
 132 void RawBufferStoreOp::print(OpAsmPrinter &p) {
 133   p << " " << getOperands() << " : " << getVdata().getType();
 134 }
 135
 136 // <operation> ::=
 137 //     `llvm.amdgcn.raw.buffer.atomic.fadd.* %vdata, %rsrc,  %offset,
 138 //     %soffset, %aux : result_type`
 139 ParseResult RawBufferAtomicFAddOp::parse(OpAsmParser &parser,
 140                                          OperationState &result) {
 141   SmallVector<OpAsmParser::UnresolvedOperand, 5> ops;
 142   Type type;
 143   if (parser.parseOperandList(ops, 5) || parser.parseColonType(type))
 144     return failure();
 145
 146   auto bldr = parser.getBuilder();
 147   auto int32Ty = bldr.getI32Type();
 148   auto i32x4Ty = VectorType::get({4}, int32Ty);
 149
 150   if (parser.resolveOperands(ops, {type, i32x4Ty, int32Ty, int32Ty, int32Ty},
 151                              parser.getNameLoc(), result.operands))
 152     return failure();
 153   return success();
 154 }
 155
 156 void RawBufferAtomicFAddOp::print(mlir::OpAsmPrinter &p) {
 157   p << " " << getOperands() << " : " << getVdata().getType();
 158 }
 159
 160 // <operation> ::=
 161 //     `llvm.amdgcn.raw.buffer.atomic.fmax.* %vdata, %rsrc,  %offset,
 162 //     %soffset, %aux : result_type`
 163 ParseResult RawBufferAtomicFMaxOp::parse(OpAsmParser &parser,
 164                                          OperationState &result) {
 165   SmallVector<OpAsmParser::UnresolvedOperand, 5> ops;
 166   Type type;
 167   if (parser.parseOperandList(ops, 5) || parser.parseColonType(type))
 168     return failure();
 169
 170   auto bldr = parser.getBuilder();
 171   auto int32Ty = bldr.getI32Type();
 172   auto i32x4Ty = VectorType::get({4}, int32Ty);
 173
 174   if (parser.resolveOperands(ops, {type, i32x4Ty, int32Ty, int32Ty, int32Ty},
 175                              parser.getNameLoc(), result.operands))
 176     return failure();
 177   return success();
 178 }
 179
 180 void RawBufferAtomicFMaxOp::print(mlir::OpAsmPrinter &p) {
 181   p << " " << getOperands() << " : " << getVdata().getType();
 182 }
 183
 184 // <operation> ::=
 185 //     `llvm.amdgcn.raw.buffer.atomic.smax.* %vdata, %rsrc,  %offset,
 186 //     %soffset, %aux : result_type`
 187 ParseResult RawBufferAtomicSMaxOp::parse(OpAsmParser &parser,
 188                                          OperationState &result) {
 189   SmallVector<OpAsmParser::UnresolvedOperand, 5> ops;
 190   Type type;
 191   if (parser.parseOperandList(ops, 5) || parser.parseColonType(type))
 192     return failure();
 193
 194   auto bldr = parser.getBuilder();
 195   auto int32Ty = bldr.getI32Type();
 196   auto i32x4Ty = VectorType::get({4}, int32Ty);
 197
 198   if (parser.resolveOperands(ops, {type, i32x4Ty, int32Ty, int32Ty, int32Ty},
 199                              parser.getNameLoc(), result.operands))
 200     return failure();
 201   return success();
 202 }
 203
 204 void RawBufferAtomicSMaxOp::print(mlir::OpAsmPrinter &p) {
 205   p << " " << getOperands() << " : " << getVdata().getType();
 206 }
 207
 208 // <operation> ::=
 209 //     `llvm.amdgcn.raw.buffer.atomic.umin.* %vdata, %rsrc,  %offset,
 210 //     %soffset, %aux : result_type`
 211 ParseResult RawBufferAtomicUMinOp::parse(OpAsmParser &parser,
 212                                          OperationState &result) {
 213   SmallVector<OpAsmParser::UnresolvedOperand, 5> ops;
 214   Type type;
 215   if (parser.parseOperandList(ops, 5) || parser.parseColonType(type))
 216     return failure();
 217
 218   auto bldr = parser.getBuilder();
 219   auto int32Ty = bldr.getI32Type();
 220   auto i32x4Ty = VectorType::get({4}, int32Ty);
 221
 222   if (parser.resolveOperands(ops, {type, i32x4Ty, int32Ty, int32Ty, int32Ty},
 223                              parser.getNameLoc(), result.operands))
 224     return failure();
 225   return success();
 226 }
 227
 228 void RawBufferAtomicUMinOp::print(mlir::OpAsmPrinter &p) {
 229   p << " " << getOperands() << " : " << getVdata().getType();
 230 }
 231
 232 //===----------------------------------------------------------------------===//
 233 // ROCDLDialect initialization, type parsing, and registration.
 234 //===----------------------------------------------------------------------===//
 235
 236 // TODO: This should be the llvm.rocdl dialect once this is supported.
 237 void ROCDLDialect::initialize() {
 238   addOperations<
 239 #define GET_OP_LIST
 240 #include "mlir/Dialect/LLVMIR/ROCDLOps.cpp.inc"
 241       >();
 242
 243   addAttributes<
 244 #define GET_ATTRDEF_LIST
 245 #include "mlir/Dialect/LLVMIR/ROCDLOpsAttributes.cpp.inc"
 246       >();
 247
 248   // Support unknown operations because not all ROCDL operations are registered.
 249   allowUnknownOperations();
 250   declarePromisedInterface<ROCDLTargetAttr, gpu::TargetAttrInterface>();
 251 }
 252
 253 LogicalResult ROCDLDialect::verifyOperationAttribute(Operation *op,
 254                                                      NamedAttribute attr) {
 255   // Kernel function attribute should be attached to functions.
 256   if (attr.getName() == ROCDLDialect::getKernelFuncAttrName()) {
 257     if (!isa<LLVM::LLVMFuncOp>(op)) {
 258       return op->emitError() << "'" << ROCDLDialect::getKernelFuncAttrName()
 259                              << "' attribute attached to unexpected op";
 260     }
 261   }
 262   return success();
 263 }
 264
 265 //===----------------------------------------------------------------------===//
 266 // ROCDL target attribute.
 267 //===----------------------------------------------------------------------===//
 268 LogicalResult
 269 ROCDLTargetAttr::verify(function_ref<InFlightDiagnostic()> emitError,
 270                         int optLevel, StringRef triple, StringRef chip,
 271                         StringRef features, StringRef abiVersion,
 272                         DictionaryAttr flags, ArrayAttr files) {
 273   if (optLevel < 0 || optLevel > 3) {
 274     emitError() << "The optimization level must be a number between 0 and 3.";
 275     return failure();
 276   }
 277   if (triple.empty()) {
 278     emitError() << "The target triple cannot be empty.";
 279     return failure();
 280   }
 281   if (chip.empty()) {
 282     emitError() << "The target chip cannot be empty.";
 283     return failure();
 284   }
 285   if (abiVersion != "400" && abiVersion != "500") {
 286     emitError() << "Invalid ABI version, it must be either `400` or `500`.";
 287     return failure();
 288   }
 289   if (files && !llvm::all_of(files, [](::mlir::Attribute attr) {
 290         return attr && mlir::isa<StringAttr>(attr);
 291       })) {
 292     emitError() << "All the elements in the `link` array must be strings.";
 293     return failure();
 294   }
 295   return success();
 296 }
 297
 298 #define GET_OP_CLASSES
 299 #include "mlir/Dialect/LLVMIR/ROCDLOps.cpp.inc"
 300
 301 #define GET_ATTRDEF_CLASSES
 302 #include "mlir/Dialect/LLVMIR/ROCDLOpsAttributes.cpp.inc"