MLIR  14.0.0git
GPUDialect.cpp
Go to the documentation of this file.
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 
14 
17 #include "mlir/IR/Attributes.h"
18 #include "mlir/IR/Builders.h"
19 #include "mlir/IR/BuiltinOps.h"
20 #include "mlir/IR/BuiltinTypes.h"
23 #include "mlir/IR/Matchers.h"
25 #include "mlir/IR/PatternMatch.h"
26 #include "mlir/IR/TypeUtilities.h"
28 #include "llvm/ADT/TypeSwitch.h"
29 
30 using namespace mlir;
31 using namespace mlir::gpu;
32 
33 #include "mlir/Dialect/GPU/GPUOpsDialect.cpp.inc"
34 
35 //===----------------------------------------------------------------------===//
36 // MMAMatrixType
37 //===----------------------------------------------------------------------===//
38 
40  StringRef operand) {
41  return Base::get(elementType.getContext(), shape, elementType, operand);
42 }
43 
46  ArrayRef<int64_t> shape, Type elementType,
47  StringRef operand) {
48  return Base::getChecked(emitError, elementType.getContext(), shape,
49  elementType, operand);
50 }
51 
52 unsigned MMAMatrixType::getNumDims() const { return getImpl()->numDims; }
53 
55  return getImpl()->getShape();
56 }
57 
58 Type MMAMatrixType::getElementType() const { return getImpl()->elementType; }
59 
60 StringRef MMAMatrixType::getOperand() const { return getImpl()->getOperand(); }
61 
63  return elementType.isF16() || elementType.isF32();
64 }
65 
68  ArrayRef<int64_t> shape, Type elementType,
69  StringRef operand) {
70  if (!operand.equals("AOp") && !operand.equals("BOp") &&
71  !operand.equals("COp"))
72  return emitError() << "operand expected to be one of AOp, BOp or COp";
73 
74  if (shape.size() != 2)
75  return emitError() << "MMAMatrixType must have exactly two dimensions";
76 
77  if (!MMAMatrixType::isValidElementType(elementType))
78  return emitError() << "MMAMatrixType elements must be F16 or F32";
79 
80  return success();
81 }
82 
83 //===----------------------------------------------------------------------===//
84 // GPUDialect
85 //===----------------------------------------------------------------------===//
86 
87 /// GPU memory space identifiers.
89  /// Generic memory space identifier.
91 
92  /// Global memory space identifier.
94 
95  /// Shared memory space identifier.
97 };
98 
99 bool GPUDialect::isKernel(Operation *op) {
100  UnitAttr isKernelAttr = op->getAttrOfType<UnitAttr>(getKernelFuncAttrName());
101  return static_cast<bool>(isKernelAttr);
102 }
103 
104 namespace {
105 /// This class defines the interface for handling inlining with gpu
106 /// operations.
107 struct GPUInlinerInterface : public DialectInlinerInterface {
109 
110  /// All gpu dialect ops can be inlined.
111  bool isLegalToInline(Operation *, Region *, bool,
112  BlockAndValueMapping &) const final {
113  return true;
114  }
115 };
116 } // namespace
117 
118 void GPUDialect::initialize() {
119  addTypes<AsyncTokenType>();
120  addTypes<MMAMatrixType>();
121  addOperations<
122 #define GET_OP_LIST
123 #include "mlir/Dialect/GPU/GPUOps.cpp.inc"
124  >();
125  addAttributes<
126 #define GET_ATTRDEF_LIST
127 #include "mlir/Dialect/GPU/GPUOpsAttributes.cpp.inc"
128  >();
129  addInterfaces<GPUInlinerInterface>();
130 }
131 
133  // Parse the main keyword for the type.
134  StringRef keyword;
135  if (parser.parseKeyword(&keyword))
136  return Type();
137  MLIRContext *context = getContext();
138 
139  // Handle 'async token' types.
140  if (keyword == "async.token")
141  return AsyncTokenType::get(context);
142 
143  if (keyword == "mma_matrix") {
144  llvm::SMLoc beginLoc = parser.getNameLoc();
145 
146  // Parse '<'.
147  if (parser.parseLess())
148  return nullptr;
149 
150  // Parse the size and elementType.
151  SmallVector<int64_t> shape;
152  Type elementType;
153  if (parser.parseDimensionList(shape, /*allowDynamic=*/false) ||
154  parser.parseType(elementType))
155  return nullptr;
156 
157  // Parse ','
158  if (parser.parseComma())
159  return nullptr;
160 
161  // Parse operand.
162  std::string operand;
163  if (failed(parser.parseOptionalString(&operand)))
164  return nullptr;
165 
166  // Parse '>'.
167  if (parser.parseGreater())
168  return nullptr;
169 
171  parser.getEncodedSourceLoc(beginLoc)),
172  shape, elementType, operand);
173  }
174 
175  parser.emitError(parser.getNameLoc(), "unknown gpu type: " + keyword);
176  return Type();
177 }
178 
179 void GPUDialect::printType(Type type, DialectAsmPrinter &os) const {
180  TypeSwitch<Type>(type)
181  .Case<AsyncTokenType>([&](Type) { os << "async.token"; })
182  .Case<MMAMatrixType>([&](MMAMatrixType fragTy) {
183  os << "mma_matrix<";
184  auto shape = fragTy.getShape();
185  for (auto dim = shape.begin(), e = shape.end() - 1; dim != e; ++dim)
186  os << *dim << 'x';
187  os << shape.back() << 'x' << fragTy.getElementType();
188  os << ", \"" << fragTy.getOperand() << "\"" << '>';
189  })
190  .Default([](Type) { llvm_unreachable("unexpected 'gpu' type kind"); });
191 }
192 
193 LogicalResult GPUDialect::verifyOperationAttribute(Operation *op,
194  NamedAttribute attr) {
195  if (!attr.getValue().isa<UnitAttr>() ||
196  attr.getName() != getContainerModuleAttrName())
197  return success();
198 
199  auto module = dyn_cast<ModuleOp>(op);
200  if (!module)
201  return op->emitError("expected '")
202  << getContainerModuleAttrName() << "' attribute to be attached to '"
203  << ModuleOp::getOperationName() << '\'';
204 
205  auto walkResult = module.walk([&module](LaunchFuncOp launchOp) -> WalkResult {
206  // Ignore launches that are nested more or less deep than functions in the
207  // module we are currently checking.
208  if (!launchOp->getParentOp() ||
209  launchOp->getParentOp()->getParentOp() != module)
210  return success();
211 
212  // Ignore launch ops with missing attributes here. The errors will be
213  // reported by the verifiers of those ops.
214  if (!launchOp->getAttrOfType<SymbolRefAttr>(
215  LaunchFuncOp::getKernelAttrName()))
216  return success();
217 
218  // Check that `launch_func` refers to a well-formed GPU kernel module.
219  StringAttr kernelModuleName = launchOp.getKernelModuleName();
220  auto kernelModule = module.lookupSymbol<GPUModuleOp>(kernelModuleName);
221  if (!kernelModule)
222  return launchOp.emitOpError()
223  << "kernel module '" << kernelModuleName.getValue()
224  << "' is undefined";
225 
226  // Check that `launch_func` refers to a well-formed kernel function.
227  Operation *kernelFunc = module.lookupSymbol(launchOp.kernelAttr());
228  if (!kernelFunc)
229  return launchOp.emitOpError("kernel function '")
230  << launchOp.kernel() << "' is undefined";
231  auto kernelConvertedFunction = dyn_cast<FunctionOpInterface>(kernelFunc);
232  if (!kernelConvertedFunction) {
233  InFlightDiagnostic diag = launchOp.emitOpError()
234  << "referenced kernel '" << launchOp.kernel()
235  << "' is not a function";
236  diag.attachNote(kernelFunc->getLoc()) << "see the kernel definition here";
237  return diag;
238  }
239 
240  if (!kernelFunc->getAttrOfType<mlir::UnitAttr>(
241  GPUDialect::getKernelFuncAttrName()))
242  return launchOp.emitOpError("kernel function is missing the '")
243  << GPUDialect::getKernelFuncAttrName() << "' attribute";
244 
245  // TODO: If the kernel isn't a GPU function (which happens during separate
246  // compilation), do not check type correspondence as it would require the
247  // verifier to be aware of the type conversion.
248  auto kernelGPUFunction = dyn_cast<gpu::GPUFuncOp>(kernelFunc);
249  if (!kernelGPUFunction)
250  return success();
251 
252  unsigned actualNumArguments = launchOp.getNumKernelOperands();
253  unsigned expectedNumArguments = kernelGPUFunction.getNumArguments();
254  if (expectedNumArguments != actualNumArguments)
255  return launchOp.emitOpError("got ")
256  << actualNumArguments << " kernel operands but expected "
257  << expectedNumArguments;
258 
259  auto functionType = kernelGPUFunction.getType();
260  for (unsigned i = 0; i < expectedNumArguments; ++i) {
261  if (launchOp.getKernelOperand(i).getType() != functionType.getInput(i)) {
262  return launchOp.emitOpError("type of function argument ")
263  << i << " does not match";
264  }
265  }
266 
267  return success();
268  });
269 
270  return walkResult.wasInterrupted() ? failure() : success();
271 }
272 
273 template <typename T>
275  auto dimension = op.dimension();
276  if (dimension != "x" && dimension != "y" && dimension != "z")
277  return op.emitError("dimension \"") << dimension << "\" is invalid";
278  return success();
279 }
280 
281 static LogicalResult verifyAllReduce(gpu::AllReduceOp allReduce) {
282  if (allReduce.body().empty() != allReduce.op().hasValue())
283  return allReduce.emitError(
284  "expected either an op attribute or a non-empty body");
285  if (!allReduce.body().empty()) {
286  if (allReduce.body().getNumArguments() != 2)
287  return allReduce.emitError("expected two region arguments");
288  for (auto argument : allReduce.body().getArguments()) {
289  if (argument.getType() != allReduce.getType())
290  return allReduce.emitError("incorrect region argument type");
291  }
292  unsigned yieldCount = 0;
293  for (Block &block : allReduce.body()) {
294  if (auto yield = dyn_cast<gpu::YieldOp>(block.getTerminator())) {
295  if (yield.getNumOperands() != 1)
296  return allReduce.emitError("expected one gpu.yield operand");
297  if (yield.getOperand(0).getType() != allReduce.getType())
298  return allReduce.emitError("incorrect gpu.yield type");
299  ++yieldCount;
300  }
301  }
302  if (yieldCount == 0)
303  return allReduce.emitError("expected gpu.yield op in region");
304  } else {
305  gpu::AllReduceOperation opName = *allReduce.op();
306  if ((opName == gpu::AllReduceOperation::AND ||
307  opName == gpu::AllReduceOperation::OR ||
308  opName == gpu::AllReduceOperation::XOR) &&
309  !allReduce.getType().isa<IntegerType>()) {
310  return allReduce.emitError()
311  << '`' << gpu::stringifyAllReduceOperation(opName) << '`'
312  << " accumulator is only compatible with Integer type";
313  }
314  }
315  return success();
316 }
317 
318 // TODO: Support optional custom attributes (without dialect prefix).
320  AllReduceOperationAttr &attr) {
321  StringRef enumStr;
322  if (!parser.parseOptionalKeyword(&enumStr)) {
323  Optional<AllReduceOperation> op = gpu::symbolizeAllReduceOperation(enumStr);
324  if (!op)
325  return parser.emitError(parser.getCurrentLocation(), "invalid op kind");
326  attr = AllReduceOperationAttr::get(parser.getContext(), *op);
327  }
328  return success();
329 }
330 
331 static void printAllReduceOperation(AsmPrinter &printer, Operation *op,
332  AllReduceOperationAttr attr) {
333  if (attr)
334  attr.print(printer);
335 }
336 
337 //===----------------------------------------------------------------------===//
338 // AsyncOpInterface
339 //===----------------------------------------------------------------------===//
340 
342  op->insertOperands(0, {token});
343  if (!op->template hasTrait<OpTrait::AttrSizedOperandSegments>())
344  return;
345  auto attrName =
347  auto sizeAttr = op->template getAttrOfType<DenseIntElementsAttr>(attrName);
348 
349  // Async dependencies is the only variadic operand.
350  if (!sizeAttr)
351  return;
352 
353  SmallVector<int32_t, 8> sizes(sizeAttr.getValues<int32_t>());
354  ++sizes.front();
355  op->setAttr(attrName, Builder(op->getContext()).getI32VectorAttr(sizes));
356 }
357 
358 //===----------------------------------------------------------------------===//
359 // LaunchOp
360 //===----------------------------------------------------------------------===//
361 
362 void LaunchOp::build(OpBuilder &builder, OperationState &result,
363  Value gridSizeX, Value gridSizeY, Value gridSizeZ,
364  Value blockSizeX, Value blockSizeY, Value blockSizeZ,
365  Value dynamicSharedMemorySize) {
366  // Add grid and block sizes as op operands, followed by the data operands.
367  result.addOperands(
368  {gridSizeX, gridSizeY, gridSizeZ, blockSizeX, blockSizeY, blockSizeZ});
369  if (dynamicSharedMemorySize)
370  result.addOperands(dynamicSharedMemorySize);
371 
372  // Create a kernel body region with kNumConfigRegionAttributes + N arguments,
373  // where the first kNumConfigRegionAttributes arguments have `index` type and
374  // the rest have the same types as the data operands.
375  Region *kernelRegion = result.addRegion();
376  Block *body = new Block();
377  for (unsigned i = 0; i < kNumConfigRegionAttributes; ++i)
378  body->addArgument(builder.getIndexType(), result.location);
379  kernelRegion->push_back(body);
380 }
381 
382 KernelDim3 LaunchOp::getBlockIds() {
383  assert(!body().empty() && "LaunchOp body must not be empty.");
384  auto args = body().getArguments();
385  return KernelDim3{args[0], args[1], args[2]};
386 }
387 
388 KernelDim3 LaunchOp::getThreadIds() {
389  assert(!body().empty() && "LaunchOp body must not be empty.");
390  auto args = body().getArguments();
391  return KernelDim3{args[3], args[4], args[5]};
392 }
393 
394 KernelDim3 LaunchOp::getGridSize() {
395  assert(!body().empty() && "LaunchOp body must not be empty.");
396  auto args = body().getArguments();
397  return KernelDim3{args[6], args[7], args[8]};
398 }
399 
400 KernelDim3 LaunchOp::getBlockSize() {
401  assert(!body().empty() && "LaunchOp body must not be empty.");
402  auto args = body().getArguments();
403  return KernelDim3{args[9], args[10], args[11]};
404 }
405 
406 KernelDim3 LaunchOp::getGridSizeOperandValues() {
407  return KernelDim3{getOperand(0), getOperand(1), getOperand(2)};
408 }
409 
410 KernelDim3 LaunchOp::getBlockSizeOperandValues() {
411  return KernelDim3{getOperand(3), getOperand(4), getOperand(5)};
412 }
413 
414 static LogicalResult verify(LaunchOp op) {
415  // Kernel launch takes kNumConfigOperands leading operands for grid/block
416  // sizes and transforms them into kNumConfigRegionAttributes region arguments
417  // for block/thread identifiers and grid/block sizes.
418  if (!op.body().empty()) {
419  if (op.body().getNumArguments() !=
420  LaunchOp::kNumConfigOperands + op.getNumOperands() -
421  (op.dynamicSharedMemorySize() ? 1 : 0))
422  return op.emitOpError("unexpected number of region arguments");
423  }
424 
425  // Block terminators without successors are expected to exit the kernel region
426  // and must be `gpu.terminator`.
427  for (Block &block : op.body()) {
428  if (block.empty())
429  continue;
430  if (block.back().getNumSuccessors() != 0)
431  continue;
432  if (!isa<gpu::TerminatorOp>(&block.back())) {
433  return block.back()
434  .emitError()
435  .append("expected '", gpu::TerminatorOp::getOperationName(),
436  "' or a terminator with successors")
437  .attachNote(op.getLoc())
438  .append("in '", LaunchOp::getOperationName(), "' body region");
439  }
440  }
441 
442  return success();
443 }
444 
445 // Pretty-print the kernel grid/block size assignment as
446 // (%iter-x, %iter-y, %iter-z) in
447 // (%size-x = %ssa-use, %size-y = %ssa-use, %size-z = %ssa-use)
448 // where %size-* and %iter-* will correspond to the body region arguments.
450  KernelDim3 operands, KernelDim3 ids) {
451  p << '(' << ids.x << ", " << ids.y << ", " << ids.z << ") in (";
452  p << size.x << " = " << operands.x << ", ";
453  p << size.y << " = " << operands.y << ", ";
454  p << size.z << " = " << operands.z << ')';
455 }
456 
457 static void printLaunchOp(OpAsmPrinter &p, LaunchOp op) {
458  // Print the launch configuration.
459  p << ' ' << op.getBlocksKeyword();
460  printSizeAssignment(p, op.getGridSize(), op.getGridSizeOperandValues(),
461  op.getBlockIds());
462  p << ' ' << op.getThreadsKeyword();
463  printSizeAssignment(p, op.getBlockSize(), op.getBlockSizeOperandValues(),
464  op.getThreadIds());
465  if (op.dynamicSharedMemorySize())
466  p << ' ' << op.getDynamicSharedMemorySizeKeyword() << ' '
467  << op.dynamicSharedMemorySize();
468 
469  p << ' ';
470  p.printRegion(op.body(), /*printEntryBlockArgs=*/false);
471  p.printOptionalAttrDict(op->getAttrs());
472 }
473 
474 // Parse the size assignment blocks for blocks and threads. These have the form
475 // (%region_arg, %region_arg, %region_arg) in
476 // (%region_arg = %operand, %region_arg = %operand, %region_arg = %operand)
477 // where %region_arg are percent-identifiers for the region arguments to be
478 // introduced further (SSA defs), and %operand are percent-identifiers for the
479 // SSA value uses.
480 static ParseResult
485  assert(indices.size() == 3 && "space for three indices expected");
487  if (parser.parseRegionArgumentList(args, /*requiredOperandCount=*/3,
489  parser.parseKeyword("in") || parser.parseLParen())
490  return failure();
491  std::move(args.begin(), args.end(), indices.begin());
492 
493  for (int i = 0; i < 3; ++i) {
494  if (i != 0 && parser.parseComma())
495  return failure();
496  if (parser.parseRegionArgument(regionSizes[i]) || parser.parseEqual() ||
497  parser.parseOperand(sizes[i]))
498  return failure();
499  }
500 
501  return parser.parseRParen();
502 }
503 
504 // Parses a Launch operation.
505 // operation ::= `gpu.launch` `blocks` `(` ssa-id-list `)` `in` ssa-reassignment
506 // `threads` `(` ssa-id-list `)` `in` ssa-reassignment
507 // region attr-dict?
508 // ssa-reassignment ::= `(` ssa-id `=` ssa-use (`,` ssa-id `=` ssa-use)* `)`
510  // Sizes of the grid and block.
512  LaunchOp::kNumConfigOperands);
514 
515  // Actual (data) operands passed to the kernel.
517 
518  // Region arguments to be created.
520  LaunchOp::kNumConfigRegionAttributes);
521  MutableArrayRef<OpAsmParser::OperandType> regionArgsRef(regionArgs);
522 
523  // Parse the size assignment segments: the first segment assigns grid sizes
524  // and defines values for block identifiers; the second segment assigns block
525  // sizes and defines values for thread identifiers. In the region argument
526  // list, identifiers precede sizes, and block-related values precede
527  // thread-related values.
528  if (parser.parseKeyword(LaunchOp::getBlocksKeyword().data()) ||
529  parseSizeAssignment(parser, sizesRef.take_front(3),
530  regionArgsRef.slice(6, 3),
531  regionArgsRef.slice(0, 3)) ||
532  parser.parseKeyword(LaunchOp::getThreadsKeyword().data()) ||
533  parseSizeAssignment(parser, sizesRef.drop_front(3),
534  regionArgsRef.slice(9, 3),
535  regionArgsRef.slice(3, 3)) ||
536  parser.resolveOperands(sizes, parser.getBuilder().getIndexType(),
537  result.operands))
538  return failure();
539 
540  OpAsmParser::OperandType dynamicSharedMemorySize;
541  if (!parser.parseOptionalKeyword(
542  LaunchOp::getDynamicSharedMemorySizeKeyword()))
543  if (parser.parseOperand(dynamicSharedMemorySize) ||
544  parser.resolveOperand(dynamicSharedMemorySize,
545  parser.getBuilder().getI32Type(),
546  result.operands))
547  return failure();
548 
549  // Introduce the body region and parse it. The region has
550  // kNumConfigRegionAttributes arguments that correspond to
551  // block/thread identifiers and grid/block sizes, all of the `index` type.
552  Type index = parser.getBuilder().getIndexType();
554  LaunchOp::kNumConfigRegionAttributes, index);
555  Region *body = result.addRegion();
556  return failure(parser.parseRegion(*body, regionArgs, dataTypes) ||
557  parser.parseOptionalAttrDict(result.attributes));
558 }
559 
560 /// Simplify the gpu.launch when the range of a thread or block ID is
561 /// trivially known to be one.
562 struct FoldLaunchArguments : public OpRewritePattern<LaunchOp> {
565  PatternRewriter &rewriter) const override {
566  // If the range implies a single value for `id`, replace `id`'s uses by
567  // zero.
568  Value zero;
569  bool simplified = false;
570  auto constPropIdUses = [&](Value id, Value size) {
571  // Check if size is trivially one.
572  if (!matchPattern(size, m_One()))
573  return;
574  if (!simplified) {
575  // Create a zero value the first time.
576  OpBuilder::InsertionGuard guard(rewriter);
577  rewriter.setInsertionPointToStart(&op.body().front());
578  zero =
579  rewriter.create<arith::ConstantIndexOp>(op.getLoc(), /*value=*/0);
580  }
581  id.replaceAllUsesWith(zero);
582  simplified = true;
583  };
584  constPropIdUses(op.getBlockIds().x, op.gridSizeX());
585  constPropIdUses(op.getBlockIds().y, op.gridSizeY());
586  constPropIdUses(op.getBlockIds().z, op.gridSizeZ());
587  constPropIdUses(op.getThreadIds().x, op.blockSizeX());
588  constPropIdUses(op.getThreadIds().y, op.blockSizeY());
589  constPropIdUses(op.getThreadIds().z, op.blockSizeZ());
590 
591  return success(simplified);
592  }
593 };
594 
595 void LaunchOp::getCanonicalizationPatterns(RewritePatternSet &rewrites,
596  MLIRContext *context) {
597  rewrites.add<FoldLaunchArguments>(context);
598 }
599 
600 //===----------------------------------------------------------------------===//
601 // LaunchFuncOp
602 //===----------------------------------------------------------------------===//
603 
604 void LaunchFuncOp::build(OpBuilder &builder, OperationState &result,
605  GPUFuncOp kernelFunc, KernelDim3 gridSize,
606  KernelDim3 blockSize, Value dynamicSharedMemorySize,
607  ValueRange kernelOperands) {
608  // Add grid and block sizes as op operands, followed by the data operands.
609  result.addOperands({gridSize.x, gridSize.y, gridSize.z, blockSize.x,
610  blockSize.y, blockSize.z});
611  if (dynamicSharedMemorySize)
612  result.addOperands(dynamicSharedMemorySize);
613  result.addOperands(kernelOperands);
614  auto kernelModule = kernelFunc->getParentOfType<GPUModuleOp>();
615  auto kernelSymbol =
616  SymbolRefAttr::get(kernelModule.getNameAttr(),
617  {SymbolRefAttr::get(kernelFunc.getNameAttr())});
618  result.addAttribute(getKernelAttrName(), kernelSymbol);
619  SmallVector<int32_t, 9> segmentSizes(9, 1);
620  segmentSizes.front() = 0; // Initially no async dependencies.
621  segmentSizes[segmentSizes.size() - 2] = dynamicSharedMemorySize ? 1 : 0;
622  segmentSizes.back() = static_cast<int32_t>(kernelOperands.size());
623  result.addAttribute(getOperandSegmentSizeAttr(),
624  builder.getI32VectorAttr(segmentSizes));
625 }
626 
627 unsigned LaunchFuncOp::getNumKernelOperands() {
628  return getNumOperands() - asyncDependencies().size() - kNumConfigOperands -
629  (dynamicSharedMemorySize() ? 1 : 0);
630 }
631 
632 StringAttr LaunchFuncOp::getKernelModuleName() {
633  return kernel().getRootReference();
634 }
635 
636 StringAttr LaunchFuncOp::getKernelName() { return kernel().getLeafReference(); }
637 
638 Value LaunchFuncOp::getKernelOperand(unsigned i) {
639  return getOperand(asyncDependencies().size() + kNumConfigOperands +
640  (dynamicSharedMemorySize() ? 1 : 0) + i);
641 }
642 
643 KernelDim3 LaunchFuncOp::getGridSizeOperandValues() {
644  auto operands = getOperands().drop_front(asyncDependencies().size());
645  return KernelDim3{operands[0], operands[1], operands[2]};
646 }
647 
648 KernelDim3 LaunchFuncOp::getBlockSizeOperandValues() {
649  auto operands = getOperands().drop_front(asyncDependencies().size());
650  return KernelDim3{operands[3], operands[4], operands[5]};
651 }
652 
653 static LogicalResult verify(LaunchFuncOp op) {
654  auto module = op->getParentOfType<ModuleOp>();
655  if (!module)
656  return op.emitOpError("expected to belong to a module");
657 
658  if (!module->getAttrOfType<UnitAttr>(
659  GPUDialect::getContainerModuleAttrName()))
660  return op.emitOpError(
661  "expected the closest surrounding module to have the '" +
662  GPUDialect::getContainerModuleAttrName() + "' attribute");
663 
664  auto kernelAttr = op->getAttrOfType<SymbolRefAttr>(op.getKernelAttrName());
665  if (!kernelAttr)
666  return op.emitOpError("symbol reference attribute '" +
667  op.getKernelAttrName() + "' must be specified");
668 
669  return success();
670 }
671 
672 static ParseResult
675  SmallVectorImpl<Type> &argTypes) {
676  if (parser.parseOptionalKeyword("args"))
677  return success();
679  SmallVector<Location> argLocations;
680  bool isVariadic = false;
682  parser, /*allowAttributes=*/false,
683  /*allowVariadic=*/false, argNames, argTypes, argAttrs, argLocations,
684  isVariadic);
685 }
686 
688  OperandRange operands, TypeRange types) {
689  if (operands.empty())
690  return;
691  printer << "args(";
692  llvm::interleaveComma(llvm::zip(operands, types), printer,
693  [&](const auto &pair) {
694  printer.printOperand(std::get<0>(pair));
695  printer << " : ";
696  printer.printType(std::get<1>(pair));
697  });
698  printer << ")";
699 }
700 
701 //===----------------------------------------------------------------------===//
702 // GPUFuncOp
703 //===----------------------------------------------------------------------===//
704 
705 /// Adds a new block argument that corresponds to buffers located in
706 /// workgroup memory.
707 BlockArgument GPUFuncOp::addWorkgroupAttribution(Type type, Location loc) {
708  auto attrName = getNumWorkgroupAttributionsAttrName();
709  auto attr = (*this)->getAttrOfType<IntegerAttr>(attrName);
710  (*this)->setAttr(attrName,
711  IntegerAttr::get(attr.getType(), attr.getValue() + 1));
712  return getBody().insertArgument(getType().getNumInputs() + attr.getInt(),
713  type, loc);
714 }
715 
716 /// Adds a new block argument that corresponds to buffers located in
717 /// private memory.
718 BlockArgument GPUFuncOp::addPrivateAttribution(Type type, Location loc) {
719  // Buffers on the private memory always come after buffers on the workgroup
720  // memory.
721  return getBody().addArgument(type, loc);
722 }
723 
724 void GPUFuncOp::build(OpBuilder &builder, OperationState &result,
725  StringRef name, FunctionType type,
726  TypeRange workgroupAttributions,
727  TypeRange privateAttributions,
728  ArrayRef<NamedAttribute> attrs) {
730  builder.getStringAttr(name));
731  result.addAttribute(getTypeAttrName(), TypeAttr::get(type));
732  result.addAttribute(getNumWorkgroupAttributionsAttrName(),
733  builder.getI64IntegerAttr(workgroupAttributions.size()));
734  result.addAttributes(attrs);
735  Region *body = result.addRegion();
736  Block *entryBlock = new Block;
737 
738  // TODO: Allow passing in proper locations here.
739  for (Type argTy : type.getInputs())
740  entryBlock->addArgument(argTy, result.location);
741  for (Type argTy : workgroupAttributions)
742  entryBlock->addArgument(argTy, result.location);
743  for (Type argTy : privateAttributions)
744  entryBlock->addArgument(argTy, result.location);
745 
746  body->getBlocks().push_back(entryBlock);
747 }
748 
749 /// Parses a GPU function memory attribution.
750 ///
751 /// memory-attribution ::= (`workgroup` `(` ssa-id-and-type-list `)`)?
752 /// (`private` `(` ssa-id-and-type-list `)`)?
753 ///
754 /// Note that this function parses only one of the two similar parts, with the
755 /// keyword provided as argument.
756 static ParseResult
757 parseAttributions(OpAsmParser &parser, StringRef keyword,
759  SmallVectorImpl<Type> &argTypes) {
760  // If we could not parse the keyword, just assume empty list and succeed.
761  if (failed(parser.parseOptionalKeyword(keyword)))
762  return success();
763 
764  if (failed(parser.parseLParen()))
765  return failure();
766 
767  // Early exit for an empty list.
768  if (succeeded(parser.parseOptionalRParen()))
769  return success();
770 
771  do {
773  Type type;
774 
775  if (parser.parseRegionArgument(arg) || parser.parseColonType(type))
776  return failure();
777 
778  args.push_back(arg);
779  argTypes.push_back(type);
780  } while (succeeded(parser.parseOptionalComma()));
781 
782  return parser.parseRParen();
783 }
784 
785 /// Parses a GPU function.
786 ///
787 /// <operation> ::= `gpu.func` symbol-ref-id `(` argument-list `)`
788 /// (`->` function-result-list)? memory-attribution `kernel`?
789 /// function-attributes? region
793  SmallVector<NamedAttrList> resultAttrs;
794  SmallVector<Type> argTypes;
795  SmallVector<Type> resultTypes;
796  SmallVector<Location> argLocations;
797  bool isVariadic;
798 
799  // Parse the function name.
800  StringAttr nameAttr;
801  if (parser.parseSymbolName(nameAttr, ::mlir::SymbolTable::getSymbolAttrName(),
802  result.attributes))
803  return failure();
804 
805  auto signatureLocation = parser.getCurrentLocation();
807  parser, /*allowVariadic=*/false, entryArgs, argTypes, argAttrs,
808  argLocations, isVariadic, resultTypes, resultAttrs)))
809  return failure();
810 
811  if (entryArgs.empty() && !argTypes.empty())
812  return parser.emitError(signatureLocation)
813  << "gpu.func requires named arguments";
814 
815  // Construct the function type. More types will be added to the region, but
816  // not to the function type.
817  Builder &builder = parser.getBuilder();
818  auto type = builder.getFunctionType(argTypes, resultTypes);
819  result.addAttribute(GPUFuncOp::getTypeAttrName(), TypeAttr::get(type));
820 
821  // Parse workgroup memory attributions.
822  if (failed(parseAttributions(parser, GPUFuncOp::getWorkgroupKeyword(),
823  entryArgs, argTypes)))
824  return failure();
825 
826  // Store the number of operands we just parsed as the number of workgroup
827  // memory attributions.
828  unsigned numWorkgroupAttrs = argTypes.size() - type.getNumInputs();
829  result.addAttribute(GPUFuncOp::getNumWorkgroupAttributionsAttrName(),
830  builder.getI64IntegerAttr(numWorkgroupAttrs));
831 
832  // Parse private memory attributions.
833  if (failed(parseAttributions(parser, GPUFuncOp::getPrivateKeyword(),
834  entryArgs, argTypes)))
835  return failure();
836 
837  // Parse the kernel attribute if present.
838  if (succeeded(parser.parseOptionalKeyword(GPUFuncOp::getKernelKeyword())))
839  result.addAttribute(GPUDialect::getKernelFuncAttrName(),
840  builder.getUnitAttr());
841 
842  // Parse attributes.
844  return failure();
845  function_interface_impl::addArgAndResultAttrs(builder, result, argAttrs,
846  resultAttrs);
847 
848  // Parse the region. If no argument names were provided, take all names
849  // (including those of attributions) from the entry block.
850  auto *body = result.addRegion();
851  return parser.parseRegion(*body, entryArgs, argTypes);
852 }
853 
854 static void printAttributions(OpAsmPrinter &p, StringRef keyword,
855  ArrayRef<BlockArgument> values) {
856  if (values.empty())
857  return;
858 
859  p << ' ' << keyword << '(';
860  llvm::interleaveComma(
861  values, p, [&p](BlockArgument v) { p << v << " : " << v.getType(); });
862  p << ')';
863 }
864 
865 /// Prints a GPU Func op.
866 static void printGPUFuncOp(OpAsmPrinter &p, GPUFuncOp op) {
867  p << ' ';
868  p.printSymbolName(op.getName());
869 
870  FunctionType type = op.getType();
872  p, op.getOperation(), type.getInputs(),
873  /*isVariadic=*/false, type.getResults());
874 
875  printAttributions(p, op.getWorkgroupKeyword(), op.getWorkgroupAttributions());
876  printAttributions(p, op.getPrivateKeyword(), op.getPrivateAttributions());
877  if (op.isKernel())
878  p << ' ' << op.getKernelKeyword();
879 
881  p, op.getOperation(), type.getNumInputs(), type.getNumResults(),
882  {op.getNumWorkgroupAttributionsAttrName(),
883  GPUDialect::getKernelFuncAttrName()});
884  p << ' ';
885  p.printRegion(op.getBody(), /*printEntryBlockArgs=*/false);
886 }
887 
888 LogicalResult GPUFuncOp::verifyType() {
889  Type type = getTypeAttr().getValue();
890  if (!type.isa<FunctionType>())
891  return emitOpError("requires '" + getTypeAttrName() +
892  "' attribute of function type");
893 
894  if (isKernel() && getType().getNumResults() != 0)
895  return emitOpError() << "expected void return type for kernel function";
896 
897  return success();
898 }
899 
901  ArrayRef<BlockArgument> attributions,
902  unsigned memorySpace) {
903  for (Value v : attributions) {
904  auto type = v.getType().dyn_cast<MemRefType>();
905  if (!type)
906  return op->emitOpError() << "expected memref type in attribution";
907 
908  if (type.getMemorySpaceAsInt() != memorySpace) {
909  return op->emitOpError()
910  << "expected memory space " << memorySpace << " in attribution";
911  }
912  }
913  return success();
914 }
915 
916 /// Verifies the body of the function.
917 LogicalResult GPUFuncOp::verifyBody() {
918  unsigned numFuncArguments = getNumArguments();
919  unsigned numWorkgroupAttributions = getNumWorkgroupAttributions();
920  unsigned numBlockArguments = front().getNumArguments();
921  if (numBlockArguments < numFuncArguments + numWorkgroupAttributions)
922  return emitOpError() << "expected at least "
923  << numFuncArguments + numWorkgroupAttributions
924  << " arguments to body region";
925 
926  ArrayRef<Type> funcArgTypes = getType().getInputs();
927  for (unsigned i = 0; i < numFuncArguments; ++i) {
928  Type blockArgType = front().getArgument(i).getType();
929  if (funcArgTypes[i] != blockArgType)
930  return emitOpError() << "expected body region argument #" << i
931  << " to be of type " << funcArgTypes[i] << ", got "
932  << blockArgType;
933  }
934 
935  if (failed(verifyAttributions(getOperation(), getWorkgroupAttributions(),
936  GPUDialect::getWorkgroupAddressSpace())) ||
937  failed(verifyAttributions(getOperation(), getPrivateAttributions(),
938  GPUDialect::getPrivateAddressSpace())))
939  return failure();
940 
941  return success();
942 }
943 
944 //===----------------------------------------------------------------------===//
945 // ReturnOp
946 //===----------------------------------------------------------------------===//
947 
948 static LogicalResult verify(gpu::ReturnOp returnOp) {
949  GPUFuncOp function = returnOp->getParentOfType<GPUFuncOp>();
950 
951  FunctionType funType = function.getType();
952 
953  if (funType.getNumResults() != returnOp.operands().size())
954  return returnOp.emitOpError()
955  .append("expected ", funType.getNumResults(), " result operands")
956  .attachNote(function.getLoc())
957  .append("return type declared here");
958 
959  for (const auto &pair : llvm::enumerate(
960  llvm::zip(function.getType().getResults(), returnOp.operands()))) {
961  Type type;
962  Value operand;
963  std::tie(type, operand) = pair.value();
964  if (type != operand.getType())
965  return returnOp.emitOpError() << "unexpected type `" << operand.getType()
966  << "' for operand #" << pair.index();
967  }
968  return success();
969 }
970 
971 //===----------------------------------------------------------------------===//
972 // GPUModuleOp
973 //===----------------------------------------------------------------------===//
974 
975 void GPUModuleOp::build(OpBuilder &builder, OperationState &result,
976  StringRef name) {
977  ensureTerminator(*result.addRegion(), builder, result.location);
978  result.attributes.push_back(builder.getNamedAttr(
980 }
981 
983  OperationState &result) {
984  StringAttr nameAttr;
985  if (parser.parseSymbolName(nameAttr, SymbolTable::getSymbolAttrName(),
986  result.attributes))
987  return failure();
988 
989  // If module attributes are present, parse them.
990  if (parser.parseOptionalAttrDictWithKeyword(result.attributes))
991  return failure();
992 
993  // Parse the module body.
994  auto *body = result.addRegion();
995  if (parser.parseRegion(*body, None, None))
996  return failure();
997 
998  // Ensure that this module has a valid terminator.
999  GPUModuleOp::ensureTerminator(*body, parser.getBuilder(), result.location);
1000  return success();
1001 }
1002 
1003 static void print(OpAsmPrinter &p, GPUModuleOp op) {
1004  p << ' ';
1005  p.printSymbolName(op.getName());
1006  p.printOptionalAttrDictWithKeyword(op->getAttrs(),
1008  p << ' ';
1009  p.printRegion(op->getRegion(0), /*printEntryBlockArgs=*/false,
1010  /*printBlockTerminators=*/false);
1011 }
1012 
1013 //===----------------------------------------------------------------------===//
1014 // GPUMemcpyOp
1015 //===----------------------------------------------------------------------===//
1016 
1017 static LogicalResult verify(MemcpyOp op) {
1018  auto srcType = op.src().getType();
1019  auto dstType = op.dst().getType();
1020 
1021  if (getElementTypeOrSelf(srcType) != getElementTypeOrSelf(dstType))
1022  return op.emitOpError("arguments have incompatible element type");
1023 
1024  if (failed(verifyCompatibleShape(srcType, dstType)))
1025  return op.emitOpError("arguments have incompatible shape");
1026 
1027  return success();
1028 }
1029 
1031  OpAsmParser &parser, Type &asyncTokenType,
1032  SmallVectorImpl<OpAsmParser::OperandType> &asyncDependencies) {
1033  auto loc = parser.getCurrentLocation();
1034  if (succeeded(parser.parseOptionalKeyword("async"))) {
1035  if (parser.getNumResults() == 0)
1036  return parser.emitError(loc, "needs to be named when marked 'async'");
1037  asyncTokenType = parser.getBuilder().getType<AsyncTokenType>();
1038  }
1039  return parser.parseOperandList(asyncDependencies,
1041 }
1042 
1044  Type asyncTokenType,
1045  OperandRange asyncDependencies) {
1046  if (asyncTokenType)
1047  printer << "async ";
1048  if (asyncDependencies.empty())
1049  return;
1050  printer << "[";
1051  llvm::interleaveComma(asyncDependencies, printer);
1052  printer << "]";
1053 }
1054 
1055 //===----------------------------------------------------------------------===//
1056 // GPU_SubgroupMmaLoadMatrixOp
1057 //===----------------------------------------------------------------------===//
1058 
1059 static LogicalResult verify(SubgroupMmaLoadMatrixOp op) {
1060  auto srcType = op.srcMemref().getType();
1061  auto resType = op.res().getType();
1062  auto resMatrixType = resType.cast<gpu::MMAMatrixType>();
1063  auto operand = resMatrixType.getOperand();
1064  auto srcMemrefType = srcType.cast<MemRefType>();
1065  auto srcMemSpace = srcMemrefType.getMemorySpaceAsInt();
1066 
1067  if (!srcMemrefType.getLayout().isIdentity())
1068  return op.emitError("expected identity layout map for source memref");
1069 
1070  if (srcMemSpace != kGenericMemorySpace && srcMemSpace != kSharedMemorySpace &&
1071  srcMemSpace != kGlobalMemorySpace)
1072  return op.emitError(
1073  "source memorySpace kGenericMemorySpace, kSharedMemorySpace or "
1074  "kGlobalMemorySpace only allowed");
1075 
1076  if (!operand.equals("AOp") && !operand.equals("BOp") &&
1077  !operand.equals("COp"))
1078  return op.emitError("only AOp, BOp and COp can be loaded");
1079 
1080  return success();
1081 }
1082 
1083 //===----------------------------------------------------------------------===//
1084 // GPU_SubgroupMmaStoreMatrixOp
1085 //===----------------------------------------------------------------------===//
1086 
1087 static LogicalResult verify(SubgroupMmaStoreMatrixOp op) {
1088  auto srcType = op.src().getType();
1089  auto dstType = op.dstMemref().getType();
1090  auto srcMatrixType = srcType.cast<gpu::MMAMatrixType>();
1091  auto dstMemrefType = dstType.cast<MemRefType>();
1092  auto dstMemSpace = dstMemrefType.getMemorySpaceAsInt();
1093  if (!dstMemrefType.getLayout().isIdentity())
1094  return op.emitError("expected identity layout map for destination memref");
1095 
1096  if (dstMemSpace != kGenericMemorySpace && dstMemSpace != kSharedMemorySpace &&
1097  dstMemSpace != kGlobalMemorySpace)
1098  return op.emitError(
1099  "destination memorySpace of kGenericMemorySpace, "
1100  "kGlobalMemorySpace or kSharedMemorySpace only allowed");
1101 
1102  if (!srcMatrixType.getOperand().equals("COp"))
1103  return op.emitError(
1104  "expected the operand matrix being stored to have 'COp' operand type");
1105 
1106  return success();
1107 }
1108 
1109 //===----------------------------------------------------------------------===//
1110 // GPU_SubgroupMmaComputeOp
1111 //===----------------------------------------------------------------------===//
1112 
1113 static LogicalResult verify(SubgroupMmaComputeOp op) {
1114  enum OperandMap { A, B, C };
1116 
1117  auto populateOpInfo = [&opTypes, &op]() {
1118  opTypes.push_back(op.opA().getType().cast<MMAMatrixType>());
1119  opTypes.push_back(op.opB().getType().cast<MMAMatrixType>());
1120  opTypes.push_back(op.opC().getType().cast<MMAMatrixType>());
1121  };
1122  populateOpInfo();
1123 
1124  if (!opTypes[A].getOperand().equals("AOp") ||
1125  !opTypes[B].getOperand().equals("BOp") ||
1126  !opTypes[C].getOperand().equals("COp"))
1127  return op.emitError("operands must be in the order AOp, BOp, COp");
1128 
1129  ArrayRef<int64_t> aShape, bShape, cShape;
1130  aShape = opTypes[A].getShape();
1131  bShape = opTypes[B].getShape();
1132  cShape = opTypes[C].getShape();
1133 
1134  if (aShape[1] != bShape[0] || aShape[0] != cShape[0] ||
1135  bShape[1] != cShape[1])
1136  return op.emitError("operand shapes do not satisfy matmul constraints");
1137 
1138  return success();
1139 }
1140 
1141 /// This is a common class used for patterns of the form
1142 /// "someop(memrefcast) -> someop". It folds the source of any memref.cast
1143 /// into the root operation directly.
1145  bool folded = false;
1146  for (OpOperand &operand : op->getOpOperands()) {
1147  auto cast = operand.get().getDefiningOp<mlir::memref::CastOp>();
1148  if (cast) {
1149  operand.set(cast.getOperand());
1150  folded = true;
1151  }
1152  }
1153  return success(folded);
1154 }
1155 
1156 LogicalResult MemcpyOp::fold(ArrayRef<Attribute> operands,
1158  return foldMemRefCast(*this);
1159 }
1160 
1161 LogicalResult MemsetOp::fold(ArrayRef<Attribute> operands,
1163  return foldMemRefCast(*this);
1164 }
1165 
1166 //===----------------------------------------------------------------------===//
1167 // GPU_AllocOp
1168 //===----------------------------------------------------------------------===//
1169 namespace {
1170 
1171 /// Folding of memref.dim(gpu.alloc(%size), %idx) -> %size similar to
1172 /// `memref::AllocOp`.
1173 struct SimplifyDimOfAllocOp : public OpRewritePattern<memref::DimOp> {
1175 
1176  LogicalResult matchAndRewrite(memref::DimOp dimOp,
1177  PatternRewriter &rewriter) const override {
1178  auto index = dimOp.index().getDefiningOp<arith::ConstantIndexOp>();
1179  if (!index)
1180  return failure();
1181 
1182  auto memrefType = dimOp.source().getType().dyn_cast<MemRefType>();
1183  if (!memrefType || !memrefType.isDynamicDim(index.value()))
1184  return failure();
1185 
1186  auto alloc = dimOp.source().getDefiningOp<AllocOp>();
1187  if (!alloc)
1188  return failure();
1189 
1190  Value substituteOp = *(alloc.dynamicSizes().begin() +
1191  memrefType.getDynamicDimIndex(index.value()));
1192  rewriter.replaceOp(dimOp, substituteOp);
1193  return success();
1194  }
1195 };
1196 
1197 } // namespace
1198 
1199 void AllocOp::getCanonicalizationPatterns(RewritePatternSet &results,
1200  MLIRContext *context) {
1201  results.add<SimplifyDimOfAllocOp>(context);
1202 }
1203 
1204 #include "mlir/Dialect/GPU/GPUOpInterfaces.cpp.inc"
1205 #include "mlir/Dialect/GPU/GPUOpsEnums.cpp.inc"
1206 
1207 #define GET_ATTRDEF_CLASSES
1208 #include "mlir/Dialect/GPU/GPUOpsAttributes.cpp.inc"
1209 
1210 #define GET_OP_CLASSES
1211 #include "mlir/Dialect/GPU/GPUOps.cpp.inc"
virtual void printSymbolName(StringRef symbolRef)
Print the given string as a symbol reference, i.e.
static bool isLegalToInline(InlinerInterface &interface, Region *src, Region *insertRegion, bool shouldCloneInlinedRegion, BlockAndValueMapping &valueMapping)
Utility to check that all of the operations within &#39;src&#39; can be inlined.
This is the representation of an operand reference.
virtual ParseResult parseOperand(OperandType &result)=0
Parse a single operand.
void printFunctionSignature(OpAsmPrinter &p, Operation *op, ArrayRef< Type > argTypes, bool isVariadic, ArrayRef< Type > resultTypes)
Prints the signature of the function-like operation op.
Include the generated interface declarations.
OpTy create(Location location, Args &&...args)
Create an operation of specific op type at the current insertion point.
Definition: Builders.h:430
This class contains a list of basic blocks and a link to the parent operation it is attached to...
Definition: Region.h:26
bool isF32() const
Definition: Types.cpp:23
ParseResult resolveOperands(ArrayRef< OperandType > operands, Type type, SmallVectorImpl< Value > &result)
Resolve a list of operands to SSA values, emitting an error on failure, or appending the results to t...
static std::string diag(llvm::Value &v)
static StringRef getSymbolAttrName()
Return the name of the attribute used for symbol names.
Definition: SymbolTable.h:55
virtual ParseResult parseLParen()=0
Parse a ( token.
A special type of RewriterBase that coordinates the application of a rewrite pattern on the current I...
Definition: PatternMatch.h:881
LogicalResult verifyCompatibleShape(ArrayRef< int64_t > shape1, ArrayRef< int64_t > shape2)
Returns success if the given two shapes are compatible.
Operation is a basic unit of execution within MLIR.
Definition: Operation.h:28
Attribute getValue() const
Return the value of the attribute.
Definition: Attributes.h:151
void addAsyncDependency(Operation *op, Value token)
Definition: GPUDialect.cpp:341
Ty getType(Args... args)
Get or construct an instance of the type &#39;ty&#39; with provided arguments.
Definition: Builders.h:83
static ParseResult parseGPUFuncOp(OpAsmParser &parser, OperationState &result)
Parses a GPU function.
Definition: GPUDialect.cpp:790
MMAMatrix represents a matrix held by a subgroup for matrix-matrix multiply accumulate operations...
Definition: GPUDialect.h:125
BlockListType & getBlocks()
Definition: Region.h:45
virtual void printType(Type type)
This class represents a diagnostic that is inflight and set to be reported.
Definition: Diagnostics.h:301
Block represents an ordered list of Operations.
Definition: Block.h:29
Generic memory space identifier.
Definition: GPUDialect.cpp:90
AttrClass getAttrOfType(StringAttr name)
Definition: Operation.h:327
virtual ParseResult parseRegionArgument(OperandType &argument)=0
Parse a region argument, this argument is resolved when calling &#39;parseRegion&#39;.
ArrayRef< int64_t > getShape() const
Get shape of the matrix.
Definition: GPUDialect.cpp:54
void addArgAndResultAttrs(Builder &builder, OperationState &result, ArrayRef< DictionaryAttr > argAttrs, ArrayRef< DictionaryAttr > resultAttrs)
Adds argument and result attributes, provided as argAttrs and resultAttrs arguments, to the list of operation attributes in result.
bool isa() const
Definition: Attributes.h:107
void push_back(Block *block)
Definition: Region.h:61
bool failed(LogicalResult result)
Utility function that returns true if the provided LogicalResult corresponds to a failure value...
Definition: LogicalResult.h:72
Shared memory space identifier.
Definition: GPUDialect.cpp:96
void printType(Type type, AsmPrinter &printer)
Prints an LLVM Dialect type.
bool succeeded(LogicalResult result)
Utility function that returns true if the provided LogicalResult corresponds to a success value...
Definition: LogicalResult.h:68
detail::constant_int_value_matcher< 1 > m_One()
Matches a constant scalar / vector splat / tensor splat integer one.
Definition: Matchers.h:243
virtual size_t getNumResults() const =0
Return the number of declared SSA results.
StringRef getTypeAttrName()
Return the name of the attribute used for function types.
static LogicalResult verifyAttributions(Operation *op, ArrayRef< BlockArgument > attributions, unsigned memorySpace)
Definition: GPUDialect.cpp:900
The OpAsmParser has methods for interacting with the asm parser: parsing things from it...
virtual Builder & getBuilder() const =0
Return a builder which provides useful access to MLIRContext, global objects like types and attribute...
NamedAttribute getNamedAttr(StringRef name, Attribute val)
Definition: Builders.cpp:81
static ParseResult parseAsyncDependencies(OpAsmParser &parser, Type &asyncTokenType, SmallVectorImpl< OpAsmParser::OperandType > &asyncDependencies)
static MMAMatrixType getChecked(function_ref< InFlightDiagnostic()> emitError, ArrayRef< int64_t > shape, Type elementType, StringRef operand)
Get MMAMatrixType at a particular location and verify construction Invariants.
Definition: GPUDialect.cpp:45
virtual ParseResult parseOptionalString(std::string *string)=0
Parse a quoted string token if present.
void insertOperands(unsigned index, ValueRange operands)
Insert the given operands into the operand list at the given &#39;index&#39;.
Definition: Operation.cpp:219
virtual ParseResult parseOptionalKeyword(StringRef keyword)=0
Parse the given keyword if present.
virtual Location getEncodedSourceLoc(llvm::SMLoc loc)=0
Re-encode the given source location as an MLIR location and return it.
virtual ParseResult parseComma()=0
Parse a , token.
virtual llvm::SMLoc getNameLoc() const =0
Return the location of the original name token.
ParseResult parseSymbolName(StringAttr &result, StringRef attrName, NamedAttrList &attrs)
Parse an -identifier and store it (without the &#39;@&#39; symbol) in a string attribute named &#39;attrName&#39;...
static ConcreteT get(MLIRContext *ctx, Args... args)
Get or create a new ConcreteT instance within the ctx.
virtual ParseResult parseOptionalComma()=0
Parse a , token if present.
SmallVector< Value, 4 > operands
static void printGPUFuncOp(OpAsmPrinter &p, GPUFuncOp op)
Prints a GPU Func op.
Definition: GPUDialect.cpp:866
ParseResult parseFunctionSignature(OpAsmParser &parser, bool allowVariadic, SmallVectorImpl< OpAsmParser::OperandType > &argNames, SmallVectorImpl< Type > &argTypes, SmallVectorImpl< NamedAttrList > &argAttrs, SmallVectorImpl< Location > &argLocations, bool &isVariadic, SmallVectorImpl< Type > &resultTypes, SmallVectorImpl< NamedAttrList > &resultAttrs)
Parses a function signature using parser.
This class defines the main interface for locations in MLIR and acts as a non-nullable wrapper around...
Definition: Location.h:48
ParseResult parseFunctionArgumentList(OpAsmParser &parser, bool allowAttributes, bool allowVariadic, SmallVectorImpl< OpAsmParser::OperandType > &argNames, SmallVectorImpl< Type > &argTypes, SmallVectorImpl< NamedAttrList > &argAttrs, SmallVectorImpl< Location > &argLocations, bool &isVariadic)
Parses function arguments using parser.
virtual ParseResult parseOperandList(SmallVectorImpl< OperandType > &result, int requiredOperandCount=-1, Delimiter delimiter=Delimiter::None)=0
Parse zero or more SSA comma-separated operand references with a specified surrounding delimiter...
static void printLaunchOp(OpAsmPrinter &p, LaunchOp op)
Definition: GPUDialect.cpp:457
void printFunctionAttributes(OpAsmPrinter &p, Operation *op, unsigned numInputs, unsigned numResults, ArrayRef< StringRef > elided={})
Prints the list of function prefixed with the "attributes" keyword.
NamedAttribute represents a combination of a name and an Attribute value.
Definition: Attributes.h:137
MLIRContext * getContext()
Return the context this operation is associated with.
Definition: Operation.h:99
MutableArrayRef< OpOperand > getOpOperands()
Definition: Operation.h:252
LogicalResult success(bool isSuccess=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:56
This class represents an efficient way to signal success or failure.
Definition: LogicalResult.h:26
static void printLaunchFuncOperands(OpAsmPrinter &printer, Operation *, OperandRange operands, TypeRange types)
Definition: GPUDialect.cpp:687
virtual ParseResult resolveOperand(const OperandType &operand, Type type, SmallVectorImpl< Value > &result)=0
Resolve an operand to an SSA value, emitting an error on failure.
LogicalResult failure(bool isFailure=true)
Utility function to generate a LogicalResult.
Definition: LogicalResult.h:62
static ParseResult parseSizeAssignment(OpAsmParser &parser, MutableArrayRef< OpAsmParser::OperandType > sizes, MutableArrayRef< OpAsmParser::OperandType > regionSizes, MutableArrayRef< OpAsmParser::OperandType > indices)
Definition: GPUDialect.cpp:481
virtual void replaceOp(Operation *op, ValueRange newValues)
This method replaces the results of the operation with the specified list of values.
virtual ParseResult parseRegion(Region &region, ArrayRef< OperandType > arguments={}, ArrayRef< Type > argTypes={}, ArrayRef< Location > argLocations={}, bool enableNameShadowing=false)=0
Parses a region.
virtual ParseResult parseGreater()=0
Parse a &#39;>&#39; token.
Type getElementTypeOrSelf(Type type)
Return the element type or return the type itself.
void addOperands(ValueRange newOperands)
virtual llvm::SMLoc getCurrentLocation()=0
Get the location of the next token and store it into the argument.
static ParseResult parseAllReduceOperation(AsmParser &parser, AllReduceOperationAttr &attr)
Definition: GPUDialect.cpp:319
Diagnostic & attachNote(Optional< Location > noteLoc=llvm::None)
Attaches a note to this diagnostic.
Definition: Diagnostics.h:335
IntegerAttr getI64IntegerAttr(int64_t value)
Definition: Builders.cpp:99
UnitAttr getUnitAttr()
Definition: Builders.cpp:85
bool isF16() const
Definition: Types.cpp:22
StringRef getOperand() const
The general form of operation this type supports is given by the equation C += A*B.
Definition: GPUDialect.cpp:60
constexpr void enumerate(std::tuple< Tys... > &tuple, CallbackT &&callback)
Definition: Matchers.h:206
DialectInlinerInterface(Dialect *dialect)
Definition: InliningUtils.h:44
StringAttr getName() const
Return the name of the attribute.
Definition: Attributes.cpp:32
static void print(OpAsmPrinter &p, GPUModuleOp op)
void addAttributes(ArrayRef< NamedAttribute > newAttributes)
Add an array of named attributes.
virtual ParseResult parseRParen()=0
Parse a ) token.
static MMAMatrixType get(ArrayRef< int64_t > shape, Type elementType, StringRef operand)
Get MMAMatrixType and verify construction Invariants.
Definition: GPUDialect.cpp:39
This is the interface that must be implemented by the dialects of operations to be inlined...
Definition: InliningUtils.h:41
static void printAttributions(OpAsmPrinter &p, StringRef keyword, ArrayRef< BlockArgument > values)
Definition: GPUDialect.cpp:854
This class provides an abstraction over the various different ranges of value types.
Definition: TypeRange.h:38
ParseResult parseKeyword(StringRef keyword, const Twine &msg="")
Parse a given keyword.
virtual ParseResult parseLess()=0
Parse a &#39;<&#39; token.
This is a pure-virtual base class that exposes the asmprinter hooks necessary to implement a custom p...
Location getLoc()
The source location the operation was defined or derived from.
Definition: Operation.h:106
This represents an operation in an abstracted form, suitable for use with the builder APIs...
static void printAllReduceOperation(AsmPrinter &printer, Operation *op, AllReduceOperationAttr attr)
Definition: GPUDialect.cpp:331
Parens surrounding zero or more operands.
Utility class for the GPU dialect to represent triples of Values accessible through ...
Definition: GPUDialect.h:36
A utility result that is used to signal how to proceed with an ongoing walk:
Definition: Visitors.h:34
This class represents an argument of a Block.
Definition: Value.h:298
LogicalResult matchAndRewrite(LaunchOp op, PatternRewriter &rewriter) const override
Definition: GPUDialect.cpp:564
static void printAsyncDependencies(OpAsmPrinter &printer, Operation *op, Type asyncTokenType, OperandRange asyncDependencies)
static ParseResult parseGPUModuleOp(OpAsmParser &parser, OperationState &result)
Definition: GPUDialect.cpp:982
static LogicalResult verifyAllReduce(gpu::AllReduceOp allReduce)
Definition: GPUDialect.cpp:281
static LogicalResult verifyIndexOp(T op)
Definition: GPUDialect.cpp:274
This base class exposes generic asm parser hooks, usable across the various derived parsers...
virtual ParseResult parseOptionalAttrDictWithKeyword(NamedAttrList &result)=0
Parse a named dictionary into &#39;result&#39; if the attributes keyword is present.
static ParseResult parseLaunchFuncOperands(OpAsmParser &parser, SmallVectorImpl< OpAsmParser::OperandType > &argNames, SmallVectorImpl< Type > &argTypes)
Definition: GPUDialect.cpp:673
Instances of the Type class are uniqued, have an immutable identifier and an optional mutable compone...
Definition: Types.h:72
static LogicalResult foldMemRefCast(Operation *op)
This is a common class used for patterns of the form "someop(memrefcast) -> someop".
virtual ParseResult parseOptionalRParen()=0
Parse a ) token if present.
MLIRContext * getContext() const
Return the MLIRContext in which this type was uniqued.
Definition: Types.cpp:19
llvm::unique_function< InFlightDiagnostic()> getDefaultDiagnosticEmitFn(MLIRContext *ctx)
Utility method to generate a callback that can be used to generate a diagnostic when checking the con...
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition: Value.h:84
void addAttribute(StringRef name, Attribute attr)
Add an attribute with the specified name.
static ParseResult parseLaunchOp(OpAsmParser &parser, OperationState &result)
Definition: GPUDialect.cpp:509
Type getElementType() const
Get elementType of a single element.
Definition: GPUDialect.cpp:58
NamedAttrList attributes
OpRewritePattern is a wrapper around RewritePattern that allows for matching and rewriting against an...
Definition: PatternMatch.h:355
void setInsertionPointToStart(Block *block)
Sets the insertion point to the start of the specified block.
Definition: Builders.h:362
virtual InFlightDiagnostic emitError(llvm::SMLoc loc, const Twine &message={})=0
Emit a diagnostic at the specified location and return failure.
virtual ParseResult parseDimensionList(SmallVectorImpl< int64_t > &dimensions, bool allowDynamic=true)=0
Parse a &#39;x&#39; separated dimension list.
virtual void printOptionalAttrDict(ArrayRef< NamedAttribute > attrs, ArrayRef< StringRef > elidedAttrs={})=0
If the specified operation has attributes, print out an attribute dictionary with their values...
RAII guard to reset the insertion point of the builder when destroyed.
Definition: Builders.h:279
void setAttr(StringAttr name, Attribute value)
If the an attribute exists with the specified name, change it to the new value.
Definition: Operation.h:347
InFlightDiagnostic emitError(Location loc)
Utility method to emit an error message using this location.
Region * addRegion()
Create a region that should be attached to the operation.
Type parseType(DialectAsmParser &parser)
Parses an LLVM dialect type.
This class is a general helper class for creating context-global objects like types, attributes, and affine expressions.
Definition: Builders.h:49
Type getType() const
Return the type of this value.
Definition: Value.h:117
RewritePatternSet & add(ConstructorArg &&arg, ConstructorArgs &&... args)
Add an instance of each of the pattern types &#39;Ts&#39; to the pattern list with the given arguments...
Definition: PatternMatch.h:930
IndexType getIndexType()
Definition: Builders.cpp:48
ImplType * getImpl() const
Utility for easy access to the storage instance.
bool matchPattern(Value value, const Pattern &pattern)
Entry point for matching a pattern over a Value.
Definition: Matchers.h:266
This is a pure-virtual base class that exposes the asmprinter hooks necessary to implement a custom p...
Specialization of arith.constant op that returns an integer of index type.
Definition: Arithmetic.h:78
virtual void printOperand(Value value)=0
Print implementations for various things an operation contains.
virtual ParseResult parseType(Type &result)=0
Parse a type.
MLIRContext is the top-level object for a collection of MLIR operations.
Definition: MLIRContext.h:55
DenseIntElementsAttr getI32VectorAttr(ArrayRef< int32_t > values)
Definition: Builders.cpp:109
virtual void printOptionalAttrDictWithKeyword(ArrayRef< NamedAttribute > attrs, ArrayRef< StringRef > elidedAttrs={})=0
If the specified operation has attributes, print out an attribute dictionary prefixed with &#39;attribute...
This class represents an operand of an operation.
Definition: Value.h:249
This class implements the operand iterators for the Operation class.
Global memory space identifier.
Definition: GPUDialect.cpp:93
static bool isValidElementType(Type elementType)
Check if a type is valid a MMAMatrixType elementType.
Definition: GPUDialect.cpp:62
unsigned getNumDims() const
Get number of dims.
Definition: GPUDialect.cpp:52
This base class exposes generic asm printer hooks, usable across the various derived printers...
virtual ParseResult parseOptionalAttrDict(NamedAttrList &result)=0
Parse a named dictionary into &#39;result&#39; if it is present.
GPUMemorySpace
GPU memory space identifiers.
Definition: GPUDialect.cpp:88
virtual ParseResult parseEqual()=0
Parse a = token.
FunctionType getFunctionType(TypeRange inputs, TypeRange results)
Definition: Builders.cpp:67
static void printSizeAssignment(OpAsmPrinter &p, KernelDim3 size, KernelDim3 operands, KernelDim3 ids)
Definition: GPUDialect.cpp:449
static ParseResult parseAttributions(OpAsmParser &parser, StringRef keyword, SmallVectorImpl< OpAsmParser::OperandType > &args, SmallVectorImpl< Type > &argTypes)
Parses a GPU function memory attribution.
Definition: GPUDialect.cpp:757
InFlightDiagnostic emitOpError(const Twine &message={})
Emit an error with the op name prefixed, like "&#39;dim&#39; op " which is convenient for verifiers...
Definition: Operation.cpp:518
bool isa() const
Definition: Types.h:234
InFlightDiagnostic emitError(const Twine &message={})
Emit an error about fatal conditions with this operation, reporting up to any diagnostic handlers tha...
Definition: Operation.cpp:231
Square brackets supporting zero or more ops, or nothing.
BlockArgument addArgument(Type type, Location loc)
Add one value to the argument list.
Definition: Block.cpp:141
MLIRContext * getContext() const
Definition: AsmPrinter.cpp:61
This class represents success/failure for operation parsing.
Definition: OpDefinition.h:36
This class helps build Operations.
Definition: Builders.h:177
Simplify the gpu.launch when the range of a thread or block ID is trivially known to be one...
Definition: GPUDialect.cpp:562
This class provides an abstraction over the different types of ranges over Values.
void push_back(NamedAttribute newAttribute)
Add an attribute with the specified name.
virtual ParseResult parseRegionArgumentList(SmallVectorImpl< OperandType > &result, int requiredOperandCount=-1, Delimiter delimiter=Delimiter::None)=0
Parse zero or more region arguments with a specified surrounding delimiter, and an optional required ...
StringAttr getStringAttr(const Twine &bytes)
Definition: Builders.cpp:201
static LogicalResult verify(function_ref< InFlightDiagnostic()> emitError, ArrayRef< int64_t > shape, Type elementType, StringRef operand)
Verify that shape and elementType are actually allowed for the MMAMatrixType.
Definition: GPUDialect.cpp:67
IntegerType getI32Type()
Definition: Builders.cpp:54
virtual ParseResult parseColonType(Type &result)=0
Parse a colon followed by a type.
The DialectAsmParser has methods for interacting with the asm parser when parsing attributes and type...
virtual void printRegion(Region &blocks, bool printEntryBlockArgs=true, bool printBlockTerminators=true, bool printEmptyBlock=false)=0
Prints a region.
static ConcreteT getChecked(const Location &loc, Args... args)
Get or create a new ConcreteT instance within the ctx, defined at the given, potentially unknown...