MLIR  16.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_IR_GPUDIALECT_H
15 #define MLIR_DIALECT_GPU_IR_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"
28 
29 namespace mlir {
30 namespace gpu {
31 
32 /// Utility class for the GPU dialect to represent triples of `Value`s
33 /// accessible through `.x`, `.y`, and `.z` similarly to CUDA notation.
34 struct KernelDim3 {
38 };
39 
41  : public Type::TypeBase<AsyncTokenType, Type, TypeStorage> {
42 public:
43  // Used for generic hooks in TypeBase.
44  using Base::Base;
45 };
46 
47 /// MMAMatrixType storage and uniquing. Array is uniqued based on its shape
48 /// and type.
50  MMAMatrixStorageType(unsigned numDims, const int64_t *dimShapes,
51  Type elementType, StringRef operand)
52  : dimShapes(dimShapes), numDims(numDims), elementType(elementType),
53  operand(operand) {}
54 
55  /// The hash key for uniquing.
56  using KeyTy = std::tuple<ArrayRef<int64_t>, Type, StringRef>;
57  bool operator==(const KeyTy &key) const {
58  return key == KeyTy(getShape(), elementType, operand);
59  }
60 
61  /// Construction.
63  const KeyTy &key) {
64  ArrayRef<int64_t> shape = allocator.copyInto(std::get<0>(key));
65  StringRef operand = allocator.copyInto(std::get<2>(key));
66 
67  return new (allocator.allocate<MMAMatrixStorageType>())
68  MMAMatrixStorageType(shape.size(), shape.data(), std::get<1>(key),
69  operand);
70  }
71 
73  return ArrayRef<int64_t>(dimShapes, numDims);
74  }
75 
76  StringRef getOperand() const { return operand; }
77 
78  /// Reference to the shape of the MMA matrix.
79  const int64_t *dimShapes;
80 
81  /// Number of dimensions in the MMA matrix.
82  unsigned numDims;
83 
84  /// Element type of elements held in the MMA matrix.
86 
87  /// MMA operand that this MMAMatrix holds. The general form of operation this
88  /// type supports is given by the equation C += A*B. This field specifies
89  /// which operand in the given equation is held by this type. The valid values
90  /// are "AOp", "BOp" and "COp".
91  StringRef operand;
92 };
93 
94 /// MMAMatrix represents a matrix held by a subgroup for matrix-matrix multiply
95 /// accumulate operations. MMAMatrices are taken as direct operands by these
96 /// operations and are also produced as results. These matrices are meant to
97 /// reside in the registers. A limited number of pointwise operations can be
98 /// performed on these matrices, i.e., operations which operate uniformly on
99 /// all the elements in the matrix and do not change the order of matrix
100 /// elements. The above conditions exist because the layout of matrix elements
101 /// inside the matrix is opaque i.e., the elements may be present in the
102 /// matrix in any order. The general usage of this type is shown as follows:-
103 ///
104 /// %0 = gpu.subgroup_mma_load_matrix %arg0[%c0, %c0] {leadDimension = 16 :
105 /// index} : memref<16x16xf16> -> !gpu.mma_matrix<16x16xf16, "AOp">
106 ///
107 /// The MMAMatrixType describes the shape of the matrix being loaded and the
108 /// operand being loaded too. The operand needs to be specified to aid the
109 /// lowering of this type to dialects such as NVVM where each workitem may
110 /// hold different amount of elements depending on the elementType of the
111 /// matrix. For e.g., Each workitem holds 4 vector<2xf16>s for f16 data type
112 /// and 8 f32s for f32 data type of MMAMatrix. Some other instances of usage
113 /// are:-
114 ///
115 /// %3 = gpu.subgroup_mma_compute %0, %1, %2 :
116 /// !gpu.mma_matrix<16x16xf16, "AOp">, !gpu.mma_matrix<16x16xf16, "BOp">
117 /// -> !gpu.mma_matrix<16x16xf32, "COp">
118 ///
119 ///
120 /// gpu.subgroup_mma_store_matrix %3, %arg22[%c0, %c0] {leadDimension = 16
121 /// : index}: !gpu.mma_matrix<16x16xf32, "COp">, memref<16x16xf32>
122 // TODO: consider moving this to ODS.
124  : public Type::TypeBase<MMAMatrixType, Type, MMAMatrixStorageType> {
125 public:
126  using Base::Base;
127 
128  /// Get MMAMatrixType and verify construction Invariants.
129  static MMAMatrixType get(ArrayRef<int64_t> shape, Type elementType,
130  StringRef operand);
131 
132  /// Get MMAMatrixType at a particular location and verify construction
133  /// Invariants.
135  ArrayRef<int64_t> shape, Type elementType,
136  StringRef operand);
137 
138  /// Check if a type is valid a MMAMatrixType elementType.
139  static bool isValidElementType(Type elementType);
140 
141  /// Verify that shape and elementType are actually allowed for the
142  /// MMAMatrixType.
144  ArrayRef<int64_t> shape, Type elementType,
145  StringRef operand);
146 
147  /// Get number of dims.
148  unsigned getNumDims() const;
149 
150  /// Get shape of the matrix.
151  ArrayRef<int64_t> getShape() const;
152 
153  /// Get elementType of a single element.
154  Type getElementType() const;
155 
156  /// The general form of operation this type supports is given by the equation
157  /// C += A*B. This function returns which operand in the given equation is
158  /// held by this type. String returned can be one of"AOp", "BOp" and "COp".
159  StringRef getOperand() const;
160 };
161 
162 // Adds a `gpu.async.token` to the front of the argument list.
163 void addAsyncDependency(Operation *op, Value token);
164 
165 } // namespace gpu
166 } // namespace mlir
167 
168 #include "mlir/Dialect/GPU/IR/GPUOpsEnums.h.inc"
169 
170 #include "mlir/Dialect/GPU/IR/GPUOpsDialect.h.inc"
171 
172 #include "mlir/Dialect/GPU/IR/GPUOpInterfaces.h.inc"
173 
174 #define GET_ATTRDEF_CLASSES
175 #include "mlir/Dialect/GPU/IR/GPUOpsAttributes.h.inc"
176 
177 #define GET_OP_CLASSES
178 #include "mlir/Dialect/GPU/IR/GPUOps.h.inc"
179 
180 #endif // MLIR_DIALECT_GPU_IR_GPUDIALECT_H
Include the generated interface declarations.
ArrayRef< int64_t > getShape() const
Definition: GPUDialect.h:72
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:82
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:123
This class represents a diagnostic that is inflight and set to be reported.
Definition: Diagnostics.h:307
Base storage class appearing in a Type.
Definition: TypeSupport.h:122
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:694
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:56
Type elementType
Element type of elements held in the MMA matrix.
Definition: GPUDialect.h:85
StringRef operand
MMA operand that this MMAMatrix holds.
Definition: GPUDialect.h:91
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:57
StringRef getOperand() const
Definition: GPUDialect.h:76
MMAMatrixStorageType(unsigned numDims, const int64_t *dimShapes, Type elementType, StringRef operand)
Definition: GPUDialect.h:50
Utility class for the GPU dialect to represent triples of Values accessible through ...
Definition: GPUDialect.h:34
Instances of the Type class are uniqued, have an immutable identifier and an optional mutable compone...
Definition: Types.h:74
static MMAMatrixStorageType * construct(TypeStorageAllocator &allocator, const KeyTy &key)
Construction.
Definition: GPUDialect.h:62
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:49
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:79
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.