1 //===- GPUDialect.cpp - MLIR Dialect for GPU Kernels implementation -------===//
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 the GPU kernel-related dialect and its operations.
10 //
11 //===----------------------------------------------------------------------===//
12 
13 #include "mlir/Dialect/GPU/GPUDialect.h"
14 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
15 #include "mlir/Dialect/StandardOps/IR/Ops.h"
16 #include "mlir/IR/Builders.h"
17 #include "mlir/IR/Function.h"
18 #include "mlir/IR/FunctionImplementation.h"
19 #include "mlir/IR/Module.h"
20 #include "mlir/IR/OpImplementation.h"
21 #include "mlir/IR/PatternMatch.h"
22 #include "mlir/IR/StandardTypes.h"
23 
24 using namespace mlir;
25 using namespace mlir::gpu;
26 
27 //===----------------------------------------------------------------------===//
28 // GPUDialect
29 //===----------------------------------------------------------------------===//
30 
31 StringRef GPUDialect::getDialectName() { return "gpu"; }
32 
33 bool GPUDialect::isKernel(Operation *op) {
34   UnitAttr isKernelAttr = op->getAttrOfType<UnitAttr>(getKernelFuncAttrName());
35   return static_cast<bool>(isKernelAttr);
36 }
37 
38 GPUDialect::GPUDialect(MLIRContext *context)
39     : Dialect(getDialectName(), context) {
40   addOperations<
41 #define GET_OP_LIST
42 #include "mlir/Dialect/GPU/GPUOps.cpp.inc"
43       >();
44 }
45 
46 LogicalResult GPUDialect::verifyOperationAttribute(Operation *op,
47                                                    NamedAttribute attr) {
48   if (!attr.second.isa<UnitAttr>() ||
49       !attr.first.is(getContainerModuleAttrName()))
50     return success();
51 
52   auto module = dyn_cast<ModuleOp>(op);
53   if (!module)
54     return op->emitError("expected '")
55            << getContainerModuleAttrName() << "' attribute to be attached to '"
56            << ModuleOp::getOperationName() << '\'';
57 
58   auto walkResult = module.walk([&module](LaunchFuncOp launchOp) -> WalkResult {
59     // Ignore launches that are nested more or less deep than functions in the
60     // module we are currently checking.
61     if (!launchOp.getParentOp() ||
62         launchOp.getParentOp()->getParentOp() != module)
63       return success();
64 
65     // Ignore launch ops with missing attributes here. The errors will be
66     // reported by the verifiers of those ops.
67     if (!launchOp.getAttrOfType<StringAttr>(
68             LaunchFuncOp::getKernelAttrName()) ||
69         !launchOp.getAttrOfType<SymbolRefAttr>(
70             LaunchFuncOp::getKernelModuleAttrName()))
71       return success();
72 
73     // Check that `launch_func` refers to a well-formed GPU kernel module.
74     StringRef kernelModuleName = launchOp.getKernelModuleName();
75     auto kernelModule = module.lookupSymbol<GPUModuleOp>(kernelModuleName);
76     if (!kernelModule)
77       return launchOp.emitOpError()
78              << "kernel module '" << kernelModuleName << "' is undefined";
79 
80     // Check that `launch_func` refers to a well-formed kernel function.
81     StringRef kernelName = launchOp.kernel();
82     Operation *kernelFunc = kernelModule.lookupSymbol(kernelName);
83     auto kernelGPUFunction = dyn_cast_or_null<gpu::GPUFuncOp>(kernelFunc);
84     auto kernelLLVMFunction = dyn_cast_or_null<LLVM::LLVMFuncOp>(kernelFunc);
85     if (!kernelGPUFunction && !kernelLLVMFunction)
86       return launchOp.emitOpError("kernel function '")
87              << kernelName << "' is undefined";
88     if (!kernelFunc->getAttrOfType<mlir::UnitAttr>(
89             GPUDialect::getKernelFuncAttrName()))
90       return launchOp.emitOpError("kernel function is missing the '")
91              << GPUDialect::getKernelFuncAttrName() << "' attribute";
92 
93     // TODO(ntv,zinenko,herhut): if the kernel function has been converted to
94     // the LLVM dialect but the caller hasn't (which happens during the
95     // separate compilation), do not check type correspondance as it would
96     // require the verifier to be aware of the LLVM type conversion.
97     if (kernelLLVMFunction)
98       return success();
99 
100     unsigned actualNumArguments = launchOp.getNumKernelOperands();
101     unsigned expectedNumArguments = kernelGPUFunction.getNumArguments();
102     if (expectedNumArguments != actualNumArguments)
103       return launchOp.emitOpError("got ")
104              << actualNumArguments << " kernel operands but expected "
105              << expectedNumArguments;
106 
107     auto functionType = kernelGPUFunction.getType();
108     for (unsigned i = 0; i < expectedNumArguments; ++i) {
109       if (launchOp.getKernelOperand(i).getType() != functionType.getInput(i)) {
110         return launchOp.emitOpError("type of function argument ")
111                << i << " does not match";
112       }
113     }
114 
115     return success();
116   });
117 
118   return walkResult.wasInterrupted() ? failure() : success();
119 }
120 
121 template <typename T> static LogicalResult verifyIndexOp(T op) {
122   auto dimension = op.dimension();
123   if (dimension != "x" && dimension != "y" && dimension != "z")
124     return op.emitError("dimension \"") << dimension << "\" is invalid";
125   return success();
126 }
127 
128 static LogicalResult verifyAllReduce(gpu::AllReduceOp allReduce) {
129   if (allReduce.body().empty() != allReduce.op().hasValue())
130     return allReduce.emitError(
131         "expected either an op attribute or a non-empty body");
132   if (!allReduce.body().empty()) {
133     if (allReduce.body().front().getNumArguments() != 2)
134       return allReduce.emitError("expected two region arguments");
135     for (auto argument : allReduce.body().front().getArguments()) {
136       if (argument.getType() != allReduce.getType())
137         return allReduce.emitError("incorrect region argument type");
138     }
139     unsigned yieldCount = 0;
140     for (Block &block : allReduce.body()) {
141       if (auto yield = dyn_cast<gpu::YieldOp>(block.getTerminator())) {
142         if (yield.getNumOperands() != 1)
143           return allReduce.emitError("expected one gpu.yield operand");
144         if (yield.getOperand(0).getType() != allReduce.getType())
145           return allReduce.emitError("incorrect gpu.yield type");
146         ++yieldCount;
147       }
148     }
149     if (yieldCount == 0)
150       return allReduce.emitError("expected gpu.yield op in region");
151   } else {
152     StringRef opName = *allReduce.op();
153     if ((opName == "and" || opName == "or" || opName == "xor") &&
154         !allReduce.getType().isa<IntegerType>()) {
155       return allReduce.emitError()
156              << '`' << opName << '`'
157              << " accumulator is only compatible with Integer type";
158     }
159   }
160   return success();
161 }
162 
163 static LogicalResult verifyShuffleOp(gpu::ShuffleOp shuffleOp) {
164   auto type = shuffleOp.value().getType();
165   if (shuffleOp.result().getType() != type) {
166     return shuffleOp.emitOpError()
167            << "requires the same type for value operand and result";
168   }
169   if (!type.isSignlessIntOrFloat() || type.getIntOrFloatBitWidth() != 32) {
170     return shuffleOp.emitOpError()
171            << "requires value operand type to be f32 or i32";
172   }
173   return success();
174 }
175 
176 static void printShuffleOp(OpAsmPrinter &p, ShuffleOp op) {
177   p << ShuffleOp::getOperationName() << ' ' << op.getOperands() << ' '
178     << op.mode() << " : " << op.value().getType();
179 }
180 
181 static ParseResult parseShuffleOp(OpAsmParser &parser, OperationState &state) {
182   SmallVector<OpAsmParser::OperandType, 3> operandInfo;
183   if (parser.parseOperandList(operandInfo, 3))
184     return failure();
185 
186   StringRef mode;
187   if (parser.parseKeyword(&mode))
188     return failure();
189   state.addAttribute("mode", parser.getBuilder().getStringAttr(mode));
190 
191   Type valueType;
192   Type int32Type = parser.getBuilder().getIntegerType(32);
193   Type int1Type = parser.getBuilder().getI1Type();
194   if (parser.parseColonType(valueType) ||
195       parser.resolveOperands(operandInfo, {valueType, int32Type, int32Type},
196                              parser.getCurrentLocation(), state.operands) ||
197       parser.addTypesToList({valueType, int1Type}, state.types))
198     return failure();
199   return success();
200 }
201 
202 //===----------------------------------------------------------------------===//
203 // LaunchOp
204 //===----------------------------------------------------------------------===//
205 
206 void LaunchOp::build(Builder *builder, OperationState &result, Value gridSizeX,
207                      Value gridSizeY, Value gridSizeZ, Value blockSizeX,
208                      Value blockSizeY, Value blockSizeZ) {
209   // Add grid and block sizes as op operands, followed by the data operands.
210   result.addOperands(
211       {gridSizeX, gridSizeY, gridSizeZ, blockSizeX, blockSizeY, blockSizeZ});
212 
213   // Create a kernel body region with kNumConfigRegionAttributes + N arguments,
214   // where the first kNumConfigRegionAttributes arguments have `index` type and
215   // the rest have the same types as the data operands.
216   Region *kernelRegion = result.addRegion();
217   Block *body = new Block();
218   body->addArguments(
219       std::vector<Type>(kNumConfigRegionAttributes, builder->getIndexType()));
220   kernelRegion->push_back(body);
221 }
222 
223 KernelDim3 LaunchOp::getBlockIds() {
224   assert(!body().getBlocks().empty() && "FuncOp body must not be empty.");
225   auto args = body().getBlocks().front().getArguments();
226   return KernelDim3{args[0], args[1], args[2]};
227 }
228 
229 KernelDim3 LaunchOp::getThreadIds() {
230   assert(!body().getBlocks().empty() && "FuncOp body must not be empty.");
231   auto args = body().getBlocks().front().getArguments();
232   return KernelDim3{args[3], args[4], args[5]};
233 }
234 
235 KernelDim3 LaunchOp::getGridSize() {
236   assert(!body().getBlocks().empty() && "FuncOp body must not be empty.");
237   auto args = body().getBlocks().front().getArguments();
238   return KernelDim3{args[6], args[7], args[8]};
239 }
240 
241 KernelDim3 LaunchOp::getBlockSize() {
242   assert(!body().getBlocks().empty() && "FuncOp body must not be empty.");
243   auto args = body().getBlocks().front().getArguments();
244   return KernelDim3{args[9], args[10], args[11]};
245 }
246 
247 KernelDim3 LaunchOp::getGridSizeOperandValues() {
248   return KernelDim3{getOperand(0), getOperand(1), getOperand(2)};
249 }
250 
251 KernelDim3 LaunchOp::getBlockSizeOperandValues() {
252   return KernelDim3{getOperand(3), getOperand(4), getOperand(5)};
253 }
254 
255 static LogicalResult verify(LaunchOp op) {
256   // Kernel launch takes kNumConfigOperands leading operands for grid/block
257   // sizes and transforms them into kNumConfigRegionAttributes region arguments
258   // for block/thread identifiers and grid/block sizes.
259   if (!op.body().empty()) {
260     Block &entryBlock = op.body().front();
261     if (entryBlock.getNumArguments() !=
262         LaunchOp::kNumConfigOperands + op.getNumOperands())
263       return op.emitOpError("unexpected number of region arguments");
264   }
265 
266   // Block terminators without successors are expected to exit the kernel region
267   // and must be `gpu.terminator`.
268   for (Block &block : op.body()) {
269     if (block.empty())
270       continue;
271     if (block.back().getNumSuccessors() != 0)
272       continue;
273     if (!isa<gpu::TerminatorOp>(&block.back())) {
274       return block.back()
275           .emitError()
276           .append("expected '", gpu::TerminatorOp::getOperationName(),
277                   "' or a terminator with successors")
278           .attachNote(op.getLoc())
279           .append("in '", LaunchOp::getOperationName(), "' body region");
280     }
281   }
282 
283   return success();
284 }
285 
286 // Pretty-print the kernel grid/block size assignment as
287 //   (%iter-x, %iter-y, %iter-z) in
288 //   (%size-x = %ssa-use, %size-y = %ssa-use, %size-z = %ssa-use)
289 // where %size-* and %iter-* will correspond to the body region arguments.
290 static void printSizeAssignment(OpAsmPrinter &p, KernelDim3 size,
291                                 ValueRange operands, KernelDim3 ids) {
292   p << '(' << ids.x << ", " << ids.y << ", " << ids.z << ") in (";
293   p << size.x << " = " << operands[0] << ", ";
294   p << size.y << " = " << operands[1] << ", ";
295   p << size.z << " = " << operands[2] << ')';
296 }
297 
298 static void printLaunchOp(OpAsmPrinter &p, LaunchOp op) {
299   ValueRange operands = op.getOperands();
300 
301   // Print the launch configuration.
302   p << LaunchOp::getOperationName() << ' ' << op.getBlocksKeyword();
303   printSizeAssignment(p, op.getGridSize(), operands.take_front(3),
304                       op.getBlockIds());
305   p << ' ' << op.getThreadsKeyword();
306   printSizeAssignment(p, op.getBlockSize(), operands.slice(3, 3),
307                       op.getThreadIds());
308 
309   p.printRegion(op.body(), /*printEntryBlockArgs=*/false);
310   p.printOptionalAttrDict(op.getAttrs());
311 }
312 
313 // Parse the size assignment blocks for blocks and threads.  These have the form
314 //   (%region_arg, %region_arg, %region_arg) in
315 //   (%region_arg = %operand, %region_arg = %operand, %region_arg = %operand)
316 // where %region_arg are percent-identifiers for the region arguments to be
317 // introduced further (SSA defs), and %operand are percent-identifiers for the
318 // SSA value uses.
319 static ParseResult
320 parseSizeAssignment(OpAsmParser &parser,
321                     MutableArrayRef<OpAsmParser::OperandType> sizes,
322                     MutableArrayRef<OpAsmParser::OperandType> regionSizes,
323                     MutableArrayRef<OpAsmParser::OperandType> indices) {
324   assert(indices.size() == 3 && "space for three indices expected");
325   SmallVector<OpAsmParser::OperandType, 3> args;
326   if (parser.parseRegionArgumentList(args, /*requiredOperandCount=*/3,
327                                      OpAsmParser::Delimiter::Paren) ||
328       parser.parseKeyword("in") || parser.parseLParen())
329     return failure();
330   std::move(args.begin(), args.end(), indices.begin());
331 
332   for (int i = 0; i < 3; ++i) {
333     if (i != 0 && parser.parseComma())
334       return failure();
335     if (parser.parseRegionArgument(regionSizes[i]) || parser.parseEqual() ||
336         parser.parseOperand(sizes[i]))
337       return failure();
338   }
339 
340   return parser.parseRParen();
341 }
342 
343 // Parses a Launch operation.
344 // operation ::= `gpu.launch` `blocks` `(` ssa-id-list `)` `in` ssa-reassignment
345 //                           `threads` `(` ssa-id-list `)` `in` ssa-reassignment
346 //                            region attr-dict?
347 // ssa-reassignment ::= `(` ssa-id `=` ssa-use (`,` ssa-id `=` ssa-use)* `)`
348 static ParseResult parseLaunchOp(OpAsmParser &parser, OperationState &result) {
349   // Sizes of the grid and block.
350   SmallVector<OpAsmParser::OperandType, LaunchOp::kNumConfigOperands> sizes(
351       LaunchOp::kNumConfigOperands);
352   MutableArrayRef<OpAsmParser::OperandType> sizesRef(sizes);
353 
354   // Actual (data) operands passed to the kernel.
355   SmallVector<OpAsmParser::OperandType, 4> dataOperands;
356 
357   // Region arguments to be created.
358   SmallVector<OpAsmParser::OperandType, 16> regionArgs(
359       LaunchOp::kNumConfigRegionAttributes);
360   MutableArrayRef<OpAsmParser::OperandType> regionArgsRef(regionArgs);
361 
362   // Parse the size assignment segments: the first segment assigns grid sizes
363   // and defines values for block identifiers; the second segment assigns block
364   // sizes and defines values for thread identifiers.  In the region argument
365   // list, identifiers precede sizes, and block-related values precede
366   // thread-related values.
367   if (parser.parseKeyword(LaunchOp::getBlocksKeyword().data()) ||
368       parseSizeAssignment(parser, sizesRef.take_front(3),
369                           regionArgsRef.slice(6, 3),
370                           regionArgsRef.slice(0, 3)) ||
371       parser.parseKeyword(LaunchOp::getThreadsKeyword().data()) ||
372       parseSizeAssignment(parser, sizesRef.drop_front(3),
373                           regionArgsRef.slice(9, 3),
374                           regionArgsRef.slice(3, 3)) ||
375       parser.resolveOperands(sizes, parser.getBuilder().getIndexType(),
376                              result.operands))
377     return failure();
378 
379   // Introduce the body region and parse it. The region has
380   // kNumConfigRegionAttributes arguments that correspond to
381   // block/thread identifiers and grid/block sizes, all of the `index` type.
382   Type index = parser.getBuilder().getIndexType();
383   SmallVector<Type, LaunchOp::kNumConfigRegionAttributes> dataTypes(
384       LaunchOp::kNumConfigRegionAttributes, index);
385   Region *body = result.addRegion();
386   return failure(parser.parseRegion(*body, regionArgs, dataTypes) ||
387                  parser.parseOptionalAttrDict(result.attributes));
388 }
389 
390 //===----------------------------------------------------------------------===//
391 // LaunchFuncOp
392 //===----------------------------------------------------------------------===//
393 
394 void LaunchFuncOp::build(Builder *builder, OperationState &result,
395                          GPUFuncOp kernelFunc, Value gridSizeX, Value gridSizeY,
396                          Value gridSizeZ, Value blockSizeX, Value blockSizeY,
397                          Value blockSizeZ, ValueRange kernelOperands) {
398   // Add grid and block sizes as op operands, followed by the data operands.
399   result.addOperands(
400       {gridSizeX, gridSizeY, gridSizeZ, blockSizeX, blockSizeY, blockSizeZ});
401   result.addOperands(kernelOperands);
402   result.addAttribute(getKernelAttrName(),
403                       builder->getStringAttr(kernelFunc.getName()));
404   auto kernelModule = kernelFunc.getParentOfType<GPUModuleOp>();
405   result.addAttribute(getKernelModuleAttrName(),
406                       builder->getSymbolRefAttr(kernelModule.getName()));
407 }
408 
409 void LaunchFuncOp::build(Builder *builder, OperationState &result,
410                          GPUFuncOp kernelFunc, KernelDim3 gridSize,
411                          KernelDim3 blockSize, ValueRange kernelOperands) {
412   build(builder, result, kernelFunc, gridSize.x, gridSize.y, gridSize.z,
413         blockSize.x, blockSize.y, blockSize.z, kernelOperands);
414 }
415 
416 StringRef LaunchFuncOp::kernel() {
417   return getAttrOfType<StringAttr>(getKernelAttrName()).getValue();
418 }
419 
420 unsigned LaunchFuncOp::getNumKernelOperands() {
421   return getNumOperands() - kNumConfigOperands;
422 }
423 
424 StringRef LaunchFuncOp::getKernelModuleName() {
425   return getAttrOfType<SymbolRefAttr>(getKernelModuleAttrName())
426       .getRootReference();
427 }
428 
429 Value LaunchFuncOp::getKernelOperand(unsigned i) {
430   return getOperation()->getOperand(i + kNumConfigOperands);
431 }
432 
433 KernelDim3 LaunchFuncOp::getGridSizeOperandValues() {
434   return KernelDim3{getOperand(0), getOperand(1), getOperand(2)};
435 }
436 
437 KernelDim3 LaunchFuncOp::getBlockSizeOperandValues() {
438   return KernelDim3{getOperand(3), getOperand(4), getOperand(5)};
439 }
440 
441 static LogicalResult verify(LaunchFuncOp op) {
442   auto module = op.getParentOfType<ModuleOp>();
443   if (!module)
444     return op.emitOpError("expected to belong to a module");
445 
446   if (!module.getAttrOfType<UnitAttr>(GPUDialect::getContainerModuleAttrName()))
447     return op.emitOpError(
448         "expected the closest surrounding module to have the '" +
449         GPUDialect::getContainerModuleAttrName() + "' attribute");
450 
451   auto kernelAttr = op.getAttrOfType<StringAttr>(op.getKernelAttrName());
452   if (!kernelAttr)
453     return op.emitOpError("string attribute '" + op.getKernelAttrName() +
454                           "' must be specified");
455 
456   auto kernelModuleAttr =
457       op.getAttrOfType<SymbolRefAttr>(op.getKernelModuleAttrName());
458   if (!kernelModuleAttr)
459     return op.emitOpError("symbol reference attribute '" +
460                           op.getKernelModuleAttrName() + "' must be specified");
461 
462   return success();
463 }
464 
465 //===----------------------------------------------------------------------===//
466 // GPUFuncOp
467 //===----------------------------------------------------------------------===//
468 
469 /// Adds a workgroup attribution to "op" of the MemRef type with the given shape
470 /// and element type.
471 Value GPUFuncOp::addWorkgroupAttribution(ArrayRef<int64_t> shape,
472                                          Type elementType) {
473   unsigned pos = getNumFuncArguments() + getNumWorkgroupAttributions();
474   Block &bodyBlock = body().front();
475   Value attribution = bodyBlock.insertArgument(
476       std::next(bodyBlock.args_begin(), pos),
477       MemRefType::get(shape, elementType, /*affineMapComposition=*/{},
478                       GPUDialect::getWorkgroupAddressSpace()));
479   auto numWorkgroupBuffersAttr =
480       getAttrOfType<IntegerAttr>(getNumWorkgroupAttributionsAttrName());
481   setAttr(getNumWorkgroupAttributionsAttrName(),
482           IntegerAttr::get(numWorkgroupBuffersAttr.getType(),
483                            numWorkgroupBuffersAttr.getValue() + 1));
484   return attribution;
485 }
486 
487 void GPUFuncOp::build(Builder *builder, OperationState &result, StringRef name,
488                       FunctionType type, ArrayRef<Type> workgroupAttributions,
489                       ArrayRef<Type> privateAttributions,
490                       ArrayRef<NamedAttribute> attrs) {
491   result.addAttribute(SymbolTable::getSymbolAttrName(),
492                       builder->getStringAttr(name));
493   result.addAttribute(getTypeAttrName(), TypeAttr::get(type));
494   result.addAttribute(getNumWorkgroupAttributionsAttrName(),
495                       builder->getI64IntegerAttr(workgroupAttributions.size()));
496   result.addAttributes(attrs);
497   Region *body = result.addRegion();
498   Block *entryBlock = new Block;
499   entryBlock->addArguments(type.getInputs());
500   entryBlock->addArguments(workgroupAttributions);
501   entryBlock->addArguments(privateAttributions);
502 
503   body->getBlocks().push_back(entryBlock);
504 }
505 
506 /// Parses a GPU function memory attribution.
507 ///
508 /// memory-attribution ::= (`workgroup` `(` ssa-id-and-type-list `)`)?
509 ///                        (`private` `(` ssa-id-and-type-list `)`)?
510 ///
511 /// Note that this function parses only one of the two similar parts, with the
512 /// keyword provided as argument.
513 static ParseResult
514 parseAttributions(OpAsmParser &parser, StringRef keyword,
515                   SmallVectorImpl<OpAsmParser::OperandType> &args,
516                   SmallVectorImpl<Type> &argTypes) {
517   // If we could not parse the keyword, just assume empty list and succeed.
518   if (failed(parser.parseOptionalKeyword(keyword)))
519     return success();
520 
521   if (failed(parser.parseLParen()))
522     return failure();
523 
524   // Early exit for an empty list.
525   if (succeeded(parser.parseOptionalRParen()))
526     return success();
527 
528   do {
529     OpAsmParser::OperandType arg;
530     Type type;
531 
532     if (parser.parseRegionArgument(arg) || parser.parseColonType(type))
533       return failure();
534 
535     args.push_back(arg);
536     argTypes.push_back(type);
537   } while (succeeded(parser.parseOptionalComma()));
538 
539   return parser.parseRParen();
540 }
541 
542 /// Parses a GPU function.
543 ///
544 /// <operation> ::= `gpu.func` symbol-ref-id `(` argument-list `)`
545 ///                 (`->` function-result-list)? memory-attribution `kernel`?
546 ///                 function-attributes? region
547 static ParseResult parseGPUFuncOp(OpAsmParser &parser, OperationState &result) {
548   SmallVector<OpAsmParser::OperandType, 8> entryArgs;
549   SmallVector<SmallVector<NamedAttribute, 2>, 1> argAttrs;
550   SmallVector<SmallVector<NamedAttribute, 2>, 1> resultAttrs;
551   SmallVector<Type, 8> argTypes;
552   SmallVector<Type, 4> resultTypes;
553   bool isVariadic;
554 
555   // Parse the function name.
556   StringAttr nameAttr;
557   if (parser.parseSymbolName(nameAttr, ::mlir::SymbolTable::getSymbolAttrName(),
558                              result.attributes))
559     return failure();
560 
561   auto signatureLocation = parser.getCurrentLocation();
562   if (failed(impl::parseFunctionSignature(
563           parser, /*allowVariadic=*/false, entryArgs, argTypes, argAttrs,
564           isVariadic, resultTypes, resultAttrs)))
565     return failure();
566 
567   if (entryArgs.empty() && !argTypes.empty())
568     return parser.emitError(signatureLocation)
569            << "gpu.func requires named arguments";
570 
571   // Construct the function type. More types will be added to the region, but
572   // not to the function type.
573   Builder &builder = parser.getBuilder();
574   auto type = builder.getFunctionType(argTypes, resultTypes);
575   result.addAttribute(GPUFuncOp::getTypeAttrName(), TypeAttr::get(type));
576 
577   // Parse workgroup memory attributions.
578   if (failed(parseAttributions(parser, GPUFuncOp::getWorkgroupKeyword(),
579                                entryArgs, argTypes)))
580     return failure();
581 
582   // Store the number of operands we just parsed as the number of workgroup
583   // memory attributions.
584   unsigned numWorkgroupAttrs = argTypes.size() - type.getNumInputs();
585   result.addAttribute(GPUFuncOp::getNumWorkgroupAttributionsAttrName(),
586                       builder.getI64IntegerAttr(numWorkgroupAttrs));
587 
588   // Parse private memory attributions.
589   if (failed(parseAttributions(parser, GPUFuncOp::getPrivateKeyword(),
590                                entryArgs, argTypes)))
591     return failure();
592 
593   // Parse the kernel attribute if present.
594   if (succeeded(parser.parseOptionalKeyword(GPUFuncOp::getKernelKeyword())))
595     result.addAttribute(GPUDialect::getKernelFuncAttrName(),
596                         builder.getUnitAttr());
597 
598   // Parse attributes.
599   if (failed(parser.parseOptionalAttrDictWithKeyword(result.attributes)))
600     return failure();
601   mlir::impl::addArgAndResultAttrs(builder, result, argAttrs, resultAttrs);
602 
603   // Parse the region. If no argument names were provided, take all names
604   // (including those of attributions) from the entry block.
605   auto *body = result.addRegion();
606   return parser.parseRegion(*body, entryArgs, argTypes);
607 }
608 
609 static void printAttributions(OpAsmPrinter &p, StringRef keyword,
610                               ArrayRef<BlockArgument> values) {
611   if (values.empty())
612     return;
613 
614   p << ' ' << keyword << '(';
615   interleaveComma(values, p,
616                   [&p](BlockArgument v) { p << v << " : " << v.getType(); });
617   p << ')';
618 }
619 
620 /// Prints a GPU Func op.
621 static void printGPUFuncOp(OpAsmPrinter &p, GPUFuncOp op) {
622   p << GPUFuncOp::getOperationName() << ' ';
623   p.printSymbolName(op.getName());
624 
625   FunctionType type = op.getType();
626   impl::printFunctionSignature(p, op.getOperation(), type.getInputs(),
627                                /*isVariadic=*/false, type.getResults());
628 
629   printAttributions(p, op.getWorkgroupKeyword(), op.getWorkgroupAttributions());
630   printAttributions(p, op.getPrivateKeyword(), op.getPrivateAttributions());
631   if (op.isKernel())
632     p << ' ' << op.getKernelKeyword();
633 
634   impl::printFunctionAttributes(p, op.getOperation(), type.getNumInputs(),
635                                 type.getNumResults(),
636                                 {op.getNumWorkgroupAttributionsAttrName(),
637                                  GPUDialect::getKernelFuncAttrName()});
638   p.printRegion(op.getBody(), /*printEntryBlockArgs=*/false);
639 }
640 
641 void GPUFuncOp::setType(FunctionType newType) {
642   auto oldType = getType();
643   assert(newType.getNumResults() == oldType.getNumResults() &&
644          "unimplemented: changes to the number of results");
645 
646   SmallVector<char, 16> nameBuf;
647   for (int i = newType.getNumInputs(), e = oldType.getNumInputs(); i < e; i++)
648     removeAttr(getArgAttrName(i, nameBuf));
649 
650   setAttr(getTypeAttrName(), TypeAttr::get(newType));
651 }
652 
653 /// Hook for FunctionLike verifier.
654 LogicalResult GPUFuncOp::verifyType() {
655   Type type = getTypeAttr().getValue();
656   if (!type.isa<FunctionType>())
657     return emitOpError("requires '" + getTypeAttrName() +
658                        "' attribute of function type");
659 
660   if (isKernel() && getType().getNumResults() != 0)
661     return emitOpError() << "expected void return type for kernel function";
662 
663   return success();
664 }
665 
666 static LogicalResult verifyAttributions(Operation *op,
667                                         ArrayRef<BlockArgument> attributions,
668                                         unsigned memorySpace) {
669   for (Value v : attributions) {
670     auto type = v.getType().dyn_cast<MemRefType>();
671     if (!type)
672       return op->emitOpError() << "expected memref type in attribution";
673 
674     if (type.getMemorySpace() != memorySpace) {
675       return op->emitOpError()
676              << "expected memory space " << memorySpace << " in attribution";
677     }
678   }
679   return success();
680 }
681 
682 /// Verifies the body of the function.
683 LogicalResult GPUFuncOp::verifyBody() {
684   unsigned numFuncArguments = getNumArguments();
685   unsigned numWorkgroupAttributions = getNumWorkgroupAttributions();
686   unsigned numBlockArguments = front().getNumArguments();
687   if (numBlockArguments < numFuncArguments + numWorkgroupAttributions)
688     return emitOpError() << "expected at least "
689                          << numFuncArguments + numWorkgroupAttributions
690                          << " arguments to body region";
691 
692   ArrayRef<Type> funcArgTypes = getType().getInputs();
693   for (unsigned i = 0; i < numFuncArguments; ++i) {
694     Type blockArgType = front().getArgument(i).getType();
695     if (funcArgTypes[i] != blockArgType)
696       return emitOpError() << "expected body region argument #" << i
697                            << " to be of type " << funcArgTypes[i] << ", got "
698                            << blockArgType;
699   }
700 
701   if (failed(verifyAttributions(getOperation(), getWorkgroupAttributions(),
702                                 GPUDialect::getWorkgroupAddressSpace())) ||
703       failed(verifyAttributions(getOperation(), getPrivateAttributions(),
704                                 GPUDialect::getPrivateAddressSpace())))
705     return failure();
706 
707   return success();
708 }
709 
710 //===----------------------------------------------------------------------===//
711 // ReturnOp
712 //===----------------------------------------------------------------------===//
713 
714 static ParseResult parseReturnOp(OpAsmParser &parser, OperationState &result) {
715   llvm::SmallVector<OpAsmParser::OperandType, 4> operands;
716   llvm::SmallVector<Type, 4> types;
717   if (parser.parseOperandList(operands) ||
718       parser.parseOptionalColonTypeList(types) ||
719       parser.resolveOperands(operands, types, parser.getCurrentLocation(),
720                              result.operands))
721     return failure();
722 
723   return success();
724 }
725 
726 static LogicalResult verify(gpu::ReturnOp returnOp) {
727   GPUFuncOp function = returnOp.getParentOfType<GPUFuncOp>();
728 
729   FunctionType funType = function.getType();
730 
731   if (funType.getNumResults() != returnOp.operands().size())
732     return returnOp.emitOpError()
733         .append("expected ", funType.getNumResults(), " result operands")
734         .attachNote(function.getLoc())
735         .append("return type declared here");
736 
737   for (auto pair : llvm::enumerate(
738            llvm::zip(function.getType().getResults(), returnOp.operands()))) {
739     Type type;
740     Value operand;
741     std::tie(type, operand) = pair.value();
742     if (type != operand.getType())
743       return returnOp.emitOpError() << "unexpected type `" << operand.getType()
744                                     << "' for operand #" << pair.index();
745   }
746   return success();
747 }
748 
749 //===----------------------------------------------------------------------===//
750 // GPUModuleOp
751 //===----------------------------------------------------------------------===//
752 
753 void GPUModuleOp::build(Builder *builder, OperationState &result,
754                         StringRef name) {
755   ensureTerminator(*result.addRegion(), *builder, result.location);
756   result.attributes.push_back(builder->getNamedAttr(
757       ::mlir::SymbolTable::getSymbolAttrName(), builder->getStringAttr(name)));
758 }
759 
760 static ParseResult parseGPUModuleOp(OpAsmParser &parser,
761                                     OperationState &result) {
762   StringAttr nameAttr;
763   if (parser.parseSymbolName(nameAttr, SymbolTable::getSymbolAttrName(),
764                              result.attributes))
765     return failure();
766 
767   // If module attributes are present, parse them.
768   if (parser.parseOptionalAttrDictWithKeyword(result.attributes))
769     return failure();
770 
771   // Parse the module body.
772   auto *body = result.addRegion();
773   if (parser.parseRegion(*body, None, None))
774     return failure();
775 
776   // Ensure that this module has a valid terminator.
777   GPUModuleOp::ensureTerminator(*body, parser.getBuilder(), result.location);
778   return success();
779 }
780 
781 static void print(OpAsmPrinter &p, GPUModuleOp op) {
782   p << op.getOperationName() << ' ';
783   p.printSymbolName(op.getName());
784   p.printOptionalAttrDictWithKeyword(op.getAttrs(),
785                                      {SymbolTable::getSymbolAttrName()});
786   p.printRegion(op.getOperation()->getRegion(0), /*printEntryBlockArgs=*/false,
787                 /*printBlockTerminators=*/false);
788 }
789 
790 // Namespace avoids ambiguous ReturnOpOperandAdaptor.
791 namespace mlir {
792 namespace gpu {
793 #define GET_OP_CLASSES
794 #include "mlir/Dialect/GPU/GPUOps.cpp.inc"
795 } // namespace gpu
796 } // namespace mlir
797