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