1//===- CreateAsyncGroups.cpp - Create async device copies -----------------===//
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#include "mlir/Dialect/NVGPU/Transforms/Transforms.h"
10
11#include "mlir/Dialect/Arith/IR/Arith.h"
12#include "mlir/Dialect/NVGPU/IR/NVGPUDialect.h"
13#include "mlir/Dialect/NVGPU/Transforms/Utils.h"
14#include "mlir/Dialect/Vector/IR/VectorOps.h"
15#include "mlir/IR/BuiltinAttributes.h"
16#include "mlir/IR/BuiltinTypes.h"
17
18using namespace mlir;
19
20/// Return "true" if the given vector transfer op is contiguous and suitable
21/// for replacement with an async copy.
22template <typename OpTy>
23static bool isContiguousXferOp(OpTy op) {
24 return op.getPermutationMap().isMinorIdentity() && op.isDimInBounds(0) &&
25 op.hasPureBufferSemantics() &&
26 cast<MemRefType>(nvgpu::getMemrefOperand(op).getType())
27 .isLastDimUnitStride();
28}
29
30/// Return "true" if the given op is a contiguous and suitable
31/// vector.transfer_write or vector.store op.
32static bool isContiguousStore(Operation *write) {
33 if (auto transferWrite = dyn_cast<vector::TransferWriteOp>(Val: write))
34 return isContiguousXferOp(op: transferWrite) && !transferWrite.getMask();
35 // vector.store are always contiguous.
36 return isa<vector::StoreOp>(Val: write);
37}
38
39/// Return "true" if the given op is a contiguous and suitable
40/// vector.transfer_read or vector.load op.
41static bool isContiguousRead(Operation *read) {
42 if (auto transferRead = dyn_cast<vector::TransferReadOp>(Val: read))
43 return isContiguousXferOp(op: transferRead);
44 // vector.load are always contiguous.
45 return isa<vector::LoadOp>(Val: read);
46}
47
48namespace {
49/// A vector.create_mask op and extract position.
50struct TransferMask {
51 vector::CreateMaskOp createMaskOp;
52 SmallVector<int64_t> extractPosition;
53};
54} // namespace
55
56/// If the given vector load op has a mask that is defined by
57/// vector.create_mask, return that op.
58static FailureOr<TransferMask> getMaskOp(Operation *loadOp) {
59 auto transferRead = dyn_cast<vector::TransferReadOp>(Val: loadOp);
60 if (!transferRead || !transferRead.getMask())
61 return TransferMask{.createMaskOp: {}, .extractPosition: {}};
62 assert(transferRead.getMask().getType().getRank() == 1 &&
63 "expected 1-D mask");
64
65 // Case 1: Mask is the result of a vector.create_mask.
66 if (auto maskOp =
67 transferRead.getMask().getDefiningOp<vector::CreateMaskOp>())
68 return TransferMask{.createMaskOp: maskOp, .extractPosition: {}};
69
70 // Case 2: Mask is the result of a vector.extract(vector.create_mask).
71 if (auto extractOp =
72 transferRead.getMask().getDefiningOp<vector::ExtractOp>())
73 if (auto maskOp =
74 extractOp.getVector().getDefiningOp<vector::CreateMaskOp>())
75 return TransferMask{.createMaskOp: maskOp,
76 .extractPosition: SmallVector<int64_t>(extractOp.getStaticPosition())};
77
78 // All other cases: not supported.
79 return failure();
80}
81
82/// Build an SSA value that represents the number of read elements.
83static Value buildNumReadElements(OpBuilder &b, Location loc,
84 Operation *readOp) {
85 FailureOr<TransferMask> transferMask = getMaskOp(loadOp: readOp);
86 assert(succeeded(transferMask) && "invalid transfer mask");
87
88 // No mask => no num_read_elements.
89 if (!transferMask->createMaskOp)
90 return Value();
91
92 // No extract: return size of "ones" segment in the mask.
93 if (transferMask->extractPosition.empty()) {
94 assert(transferMask->createMaskOp.getNumOperands() == 1 &&
95 "expected single operand");
96 return transferMask->createMaskOp.getOperand(i: 0);
97 }
98
99 // vector.extract(vector.create_mask).
100 // If extract_pos < num_ones, take number of elements from the least
101 // significant dimension. (Do this for all dimensions and bit-AND the
102 // conditions.)
103 assert(transferMask->createMaskOp.getVectorType().getRank() -
104 transferMask->extractPosition.size() ==
105 1 &&
106 "expected N-D -> (N-1)-D extract");
107 Value cond;
108 // Note: There is one more `sz` than `pos`. The loop end with the last `pos`.
109 for (auto [pos, sz] : llvm::zip(t&: transferMask->extractPosition,
110 u: transferMask->createMaskOp->getOperands())) {
111 Value cmp =
112 b.create<arith::CmpIOp>(location: loc, args: arith::CmpIPredicate::slt,
113 args: b.create<arith::ConstantIndexOp>(location: loc, args&: pos), args&: sz);
114 if (!cond) {
115 cond = cmp;
116 continue;
117 }
118 cond = b.create<arith::AndIOp>(location: loc, args&: cmp, args&: cond);
119 }
120 return b.create<arith::SelectOp>(
121 location: loc, args&: cond, args: transferMask->createMaskOp->getOperands().back(),
122 args: b.create<arith::ConstantIndexOp>(location: loc, args: 0));
123}
124
125/// Return "true" if the conversion to async copy is supported by "async copy".
126static bool resultsInSupportedAsyncCopy(MemRefType memrefType,
127 VectorType vecType) {
128 assert(vecType.getRank() == 1 && "expected 1-D vector");
129 constexpr int64_t kSupportedCpAsyncAlignmentsInBytes[3] = {4, 8, 16};
130
131 // Condition 1: the copy size must be supported.
132 bool supportedCopySize = false;
133 int64_t numElements = vecType.getNumElements();
134 Type elementType = vecType.getElementType();
135 for (int64_t alignmentInBytes : kSupportedCpAsyncAlignmentsInBytes) {
136 if (alignmentInBytes * 8 ==
137 numElements * elementType.getIntOrFloatBitWidth()) {
138 supportedCopySize = true;
139 break;
140 }
141 }
142 if (!supportedCopySize)
143 return false;
144
145 // TODO: Condition 2: the alignments must be supported. For cp.async the
146 // NVIDIA doc (section 6.4.1) says: "The address must be naturally aligned to
147 // a multiple of the access size. If an address is not properly aligned, the
148 // resulting behavior is undefined.".
149 return true;
150}
151
152void nvgpu::createAsyncGroups(RewriterBase &rewriter, Operation *op,
153 bool bypassL1) {
154 llvm::SmallSetVector<Operation *, 16> copyToSharedMem;
155
156 // Look for all the copy that can be converted to async copy ops.
157 op->walk(callback: [&](Operation *writeOp) {
158 // Look for contiguous 1D vector store into shared memory.
159 if (!isContiguousStore(write: writeOp))
160 return;
161 Value vectorVal = nvgpu::getValueStored(op: writeOp);
162 if (cast<VectorType>(Val: vectorVal.getType()).getRank() != 1)
163 return;
164 Value storeBase = nvgpu::getMemrefOperand(op: writeOp);
165 if (!nvgpu::NVGPUDialect::hasSharedMemoryAddressSpace(
166 type: cast<MemRefType>(Val: storeBase.getType())))
167 return;
168
169 // The stored vector must originate from a contiguous 1D vector load.
170 Operation *readOp = vectorVal.getDefiningOp();
171 if (readOp == nullptr || !isContiguousRead(read: readOp))
172 return;
173 Value loadBase = nvgpu::getMemrefOperand(op: readOp);
174 // Should be reading from global memory (not shared memory).
175 if (nvgpu::NVGPUDialect::hasSharedMemoryAddressSpace(
176 type: cast<MemRefType>(Val: loadBase.getType())))
177 return;
178
179 // Look for compatible mask and padding.
180 if (auto transferRead = dyn_cast<vector::TransferReadOp>(Val: readOp)) {
181 if (Value mask = transferRead.getMask()) {
182 if (getConstantIntValue(ofr: transferRead.getPadding()) ==
183 static_cast<int64_t>(0))
184 return;
185 if (failed(Result: getMaskOp(loadOp: readOp)))
186 return;
187 }
188 }
189
190 // Check whether both accesses are supported before we emit: this is
191 // necessary to ensure the correctness of DeviceAsyncCopyOp.
192 VectorType vecType = cast<VectorType>(Val: vectorVal.getType());
193
194 if (!resultsInSupportedAsyncCopy(memrefType: cast<MemRefType>(Val: loadBase.getType()),
195 vecType) ||
196 !resultsInSupportedAsyncCopy(memrefType: cast<MemRefType>(Val: storeBase.getType()),
197 vecType))
198 return;
199
200 copyToSharedMem.insert(X: writeOp);
201 return;
202 });
203
204 while (!copyToSharedMem.empty()) {
205 // Start a group with the first write.
206 SmallVector<Operation *> group;
207 Operation *writeOp = *copyToSharedMem.begin();
208 copyToSharedMem.remove(X: writeOp);
209 group.push_back(Elt: writeOp);
210 Operation *nextNode = writeOp;
211
212 // Look in the next nodes for more copies to add to the same group.
213 while ((nextNode = nextNode->getNextNode())) {
214 // Ignore ops without side effects.
215 auto memInterface = dyn_cast<MemoryEffectOpInterface>(Val: nextNode);
216 if (memInterface && memInterface.hasNoEffect() &&
217 !nextNode->hasTrait<OpTrait::HasRecursiveMemoryEffects>())
218 continue;
219 // Ignore read from a different address space.
220 if (isa<vector::TransferReadOp, vector::LoadOp>(Val: nextNode)) {
221 Operation *readOp = nextNode;
222 Value memrefOperand = nvgpu::getMemrefOperand(op: readOp);
223 if (!nvgpu::NVGPUDialect::hasSharedMemoryAddressSpace(
224 type: cast<MemRefType>(Val: memrefOperand.getType()))) {
225 continue;
226 }
227 }
228 if (copyToSharedMem.count(key: nextNode)) {
229 // Found another copy, add it to the group.
230 copyToSharedMem.remove(X: nextNode);
231 group.push_back(Elt: nextNode);
232 continue;
233 }
234 // If the op is something else stop the accumulating op in the group.
235 break;
236 }
237
238 // Emit the group.
239 SmallVector<Value> tokens;
240 for (Operation *writeOp : group) {
241 rewriter.setInsertionPoint(writeOp);
242 Value vectorVal = nvgpu::getValueStored(op: writeOp);
243 auto vectorType = cast<VectorType>(Val: vectorVal.getType());
244 int64_t numElements = vectorType.getNumElements();
245 Operation *readOp = vectorVal.getDefiningOp();
246 Value storeBase = nvgpu::getMemrefOperand(op: writeOp);
247 Value loadBase = nvgpu::getMemrefOperand(op: readOp);
248 Value numReadElements =
249 buildNumReadElements(b&: rewriter, loc: writeOp->getLoc(), readOp);
250 auto dstMemref = cast<MemRefType>(Val: storeBase.getType());
251 int64_t sizeInBytes =
252 (dstMemref.getElementTypeBitWidth() * numElements) / 8;
253 // bypass_l1 only possible with 16 byte transfer.
254 Value token = rewriter.create<nvgpu::DeviceAsyncCopyOp>(
255 location: writeOp->getLoc(), args: nvgpu::DeviceAsyncTokenType::get(ctx: op->getContext()),
256 /*dst=*/args&: storeBase, /*dstIndices=*/args: nvgpu::getIndices(op: writeOp),
257 /*src=*/args&: loadBase,
258 /*srcIndices=*/args: nvgpu::getIndices(op: readOp),
259 /*dstElements=*/args: rewriter.getIndexAttr(value: numElements),
260 /*srcElements=*/args&: numReadElements,
261 /*bypassL1=*/args: bypassL1 && sizeInBytes == 16 ? rewriter.getUnitAttr()
262 : UnitAttr());
263 tokens.push_back(Elt: token);
264 }
265
266 // Create the group and wait for it right after.
267 Value groupToken = rewriter.create<nvgpu::DeviceAsyncCreateGroupOp>(
268 location: op->getLoc(), args: nvgpu::DeviceAsyncTokenType::get(ctx: op->getContext()),
269 args&: tokens);
270 rewriter.create<nvgpu::DeviceAsyncWaitOp>(location: op->getLoc(), args&: groupToken,
271 args: nullptr);
272 // Clean up old stores.
273 for (Operation *writeOp : group)
274 rewriter.eraseOp(op: writeOp);
275 }
276}
277

source code of mlir/lib/Dialect/NVGPU/Transforms/CreateAsyncGroups.cpp