MLIR  15.0.0git
GPUDialect.h
Go to the documentation of this file.
1 //===- GPUDialect.h - MLIR Dialect for GPU Kernels --------------*- C++ -*-===//
2 //
3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4 // See https://llvm.org/LICENSE.txt for license information.
5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6 //
7 //===----------------------------------------------------------------------===//
8 //
9 // This file defines the GPU kernel-related operations and puts them in the
10 // corresponding dialect.
11 //
12 //===----------------------------------------------------------------------===//
13 
14 #ifndef MLIR_DIALECT_GPU_GPUDIALECT_H
15 #define MLIR_DIALECT_GPU_GPUDIALECT_H
16 
18 #include "mlir/IR/Builders.h"
19 #include "mlir/IR/BuiltinTypes.h"
20 #include "mlir/IR/Dialect.h"
22 #include "mlir/IR/OpDefinition.h"
24 #include "mlir/IR/SymbolTable.h"
27 
28 namespace mlir {
29 namespace gpu {
30 
31 /// Utility class for the GPU dialect to represent triples of `Value`s
32 /// accessible through `.x`, `.y`, and `.z` similarly to CUDA notation.
33 struct KernelDim3 {
37 };
38 
40  : public Type::TypeBase<AsyncTokenType, Type, TypeStorage> {
41 public:
42  // Used for generic hooks in TypeBase.
43  using Base::Base;
44 };
45 
46 /// MMAMatrixType storage and uniquing. Array is uniqued based on its shape
47 /// and type.
49  MMAMatrixStorageType(unsigned numDims, const int64_t *dimShapes,
50  Type elementType, StringRef operand)
51  : dimShapes(dimShapes), numDims(numDims), elementType(elementType),
52  operand(operand) {}
53 
54  /// The hash key for uniquing.
55  using KeyTy = std::tuple<ArrayRef<int64_t>, Type, StringRef>;
56  bool operator==(const KeyTy &key) const {
57  return key == KeyTy(getShape(), elementType, operand);
58  }
59 
60  /// Construction.
62  const KeyTy &key) {
63  ArrayRef<int64_t> shape = allocator.copyInto(std::get<0>(key));
64  StringRef operand = allocator.copyInto(std::get<2>(key));
65 
66  return new (allocator.allocate<MMAMatrixStorageType>())
67  MMAMatrixStorageType(shape.size(), shape.data(), std::get<1>(key),
68  operand);
69  }
70 
72  return ArrayRef<int64_t>(dimShapes, numDims);
73  }
74 
75  StringRef getOperand() const { return operand; }
76 
77  /// Reference to the shape of the MMA matrix.
78  const int64_t *dimShapes;
79 
80  /// Number of dimensions in the MMA matrix.
81  unsigned numDims;
82 
83  /// Element type of elements held in the MMA matrix.
85 
86  /// MMA operand that this MMAMatrix holds. The general form of operation this
87  /// type supports is given by the equation C += A*B. This field specifies
88  /// which operand in the given equation is held by this type. The valid values
89  /// are "AOp", "BOp" and "COp".
90  StringRef operand;
91 };
92 
93 /// MMAMatrix represents a matrix held by a subgroup for matrix-matrix multiply
94 /// accumulate operations. MMAMatrices are taken as direct operands by these
95 /// operations and are also produced as results. These matrices are meant to
96 /// reside in the registers. A limited number of pointwise operations can be
97 /// performed on these matrices, i.e., operations which operate uniformly on
98 /// all the elements in the matrix and do not change the order of matrix
99 /// elements. The above conditions exist because the layout of matrix elements
100 /// inside the matrix is opaque i.e., the elements may be present in the
101 /// matrix in any order. The general usage of this type is shown as follows:-
102 ///
103 /// %0 = gpu.subgroup_mma_load_matrix %arg0[%c0, %c0] {leadDimension = 16 :
104 /// index} : memref<16x16xf16> -> !gpu.mma_matrix<16x16xf16, "AOp">
105 ///
106 /// The MMAMatrixType describes the shape of the matrix being loaded and the
107 /// operand being loaded too. The operand needs to be specified to aid the
108 /// lowering of this type to dialects such as NVVM where each workitem may
109 /// hold different amount of elements depending on the elementType of the
110 /// matrix. For e.g., Each workitem holds 4 vector<2xf16>s for f16 data type
111 /// and 8 f32s for f32 data type of MMAMatrix. Some other instances of usage
112 /// are:-
113 ///
114 /// %3 = gpu.subgroup_mma_compute %0, %1, %2 :
115 /// !gpu.mma_matrix<16x16xf16, "AOp">, !gpu.mma_matrix<16x16xf16, "BOp">
116 /// -> !gpu.mma_matrix<16x16xf32, "COp">
117 ///
118 ///
119 /// gpu.subgroup_mma_store_matrix %3, %arg22[%c0, %c0] {leadDimension = 16
120 /// : index}: !gpu.mma_matrix<16x16xf32, "COp">, memref<16x16xf32>
121 // TODO: consider moving this to ODS.
123  : public Type::TypeBase<MMAMatrixType, Type, MMAMatrixStorageType> {
124 public:
125  using Base::Base;
126 
127  /// Get MMAMatrixType and verify construction Invariants.
128  static MMAMatrixType get(ArrayRef<int64_t> shape, Type elementType,
129  StringRef operand);
130 
131  /// Get MMAMatrixType at a particular location and verify construction
132  /// Invariants.
134  ArrayRef<int64_t> shape, Type elementType,
135  StringRef operand);
136 
137  /// Check if a type is valid a MMAMatrixType elementType.
138  static bool isValidElementType(Type elementType);
139 
140  /// Verify that shape and elementType are actually allowed for the
141  /// MMAMatrixType.
143  ArrayRef<int64_t> shape, Type elementType,
144  StringRef operand);
145 
146  /// Get number of dims.
147  unsigned getNumDims() const;
148 
149  /// Get shape of the matrix.
150  ArrayRef<int64_t> getShape() const;
151 
152  /// Get elementType of a single element.
153  Type getElementType() const;
154 
155  /// The general form of operation this type supports is given by the equation
156  /// C += A*B. This function returns which operand in the given equation is
157  /// held by this type. String returned can be one of"AOp", "BOp" and "COp".
158  StringRef getOperand() const;
159 };
160 
161 // Adds a `gpu.async.token` to the front of the argument list.
162 void addAsyncDependency(Operation *op, Value token);
163 
164 } // namespace gpu
165 } // namespace mlir
166 
167 #include "mlir/Dialect/GPU/GPUOpsEnums.h.inc"
168 
169 #include "mlir/Dialect/GPU/GPUOpsDialect.h.inc"
170 
171 #include "mlir/Dialect/GPU/GPUOpInterfaces.h.inc"
172 
173 #define GET_ATTRDEF_CLASSES
174 #include "mlir/Dialect/GPU/GPUOpsAttributes.h.inc"
175 
176 #define GET_OP_CLASSES
177 #include "mlir/Dialect/GPU/GPUOps.h.inc"
178 
179 #endif // MLIR_DIALECT_GPU_GPUDIALECT_H
Include the generated interface declarations.
ArrayRef< int64_t > getShape() const
Definition: GPUDialect.h:71
Operation is a basic unit of execution within MLIR.
Definition: Operation.h:28
unsigned numDims
Number of dimensions in the MMA matrix.
Definition: GPUDialect.h:81
void addAsyncDependency(Operation *op, Value token)
Definition: GPUDialect.cpp:371
MMAMatrix represents a matrix held by a subgroup for matrix-matrix multiply accumulate operations...
Definition: GPUDialect.h:122
This class represents a diagnostic that is inflight and set to be reported.
Definition: Diagnostics.h:311
Base storage class appearing in a Type.
Definition: TypeSupport.h:121
This is a utility allocator used to allocate memory for instances of derived types.
static Type getElementType(Type type, ArrayRef< int32_t > indices, function_ref< InFlightDiagnostic(StringRef)> emitErrorFn)
Walks the given type hierarchy with the given indices, potentially down to component granularity...
Definition: SPIRVOps.cpp:687
static ArrayRef< int64_t > getShape(Type type)
Returns the shape of the given type.
Definition: Traits.cpp:117
std::tuple< ArrayRef< int64_t >, Type, StringRef > KeyTy
The hash key for uniquing.
Definition: GPUDialect.h:55
Type elementType
Element type of elements held in the MMA matrix.
Definition: GPUDialect.h:84
StringRef operand
MMA operand that this MMAMatrix holds.
Definition: GPUDialect.h:90
T * allocate()
Allocate an instance of the provided type.
This class represents an efficient way to signal success or failure.
Definition: LogicalResult.h:26
bool operator==(const KeyTy &key) const
Definition: GPUDialect.h:56
StringRef getOperand() const
Definition: GPUDialect.h:75
MMAMatrixStorageType(unsigned numDims, const int64_t *dimShapes, Type elementType, StringRef operand)
Definition: GPUDialect.h:49
Utility class for the GPU dialect to represent triples of Values accessible through ...
Definition: GPUDialect.h:33
Instances of the Type class are uniqued, have an immutable identifier and an optional mutable compone...
Definition: Types.h:72
static MMAMatrixStorageType * construct(TypeStorageAllocator &allocator, const KeyTy &key)
Construction.
Definition: GPUDialect.h:61
This class represents an instance of an SSA value in the MLIR system, representing a computable value...
Definition: Value.h:85
MMAMatrixType storage and uniquing.
Definition: GPUDialect.h:48
InFlightDiagnostic emitError(Location loc)
Utility method to emit an error message using this location.
Utility class for implementing users of storage classes uniqued by a StorageUniquer.
const int64_t * dimShapes
Reference to the shape of the MMA matrix.
Definition: GPUDialect.h:78
LogicalResult verify(Operation *op, bool verifyRecursively=true)
Perform (potentially expensive) checks of invariants, used to detect compiler bugs, on this operation and any nested operations.
Definition: Verifier.cpp:372
ArrayRef< T > copyInto(ArrayRef< T > elements)
Copy the specified array of elements into memory managed by our bump pointer allocator.