1//===- LowerMatrixIntrinsics.cpp - Lower matrix intrinsics -----*- 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// Lower matrix intrinsics to vector operations.
10//
11// TODO:
12// * Improve fusion:
13// * Support more cases, e.g. multiply-add, multiply-sub, operands/results
14// transposed.
15// * Improve cost-modeling, e.g. choose different number of rows/columns
16// columns for tiles, consider cost of copies on alias.
17//
18//===----------------------------------------------------------------------===//
19
20#include "llvm/Transforms/Scalar/LowerMatrixIntrinsics.h"
21#include "llvm/ADT/PostOrderIterator.h"
22#include "llvm/ADT/STLExtras.h"
23#include "llvm/ADT/ScopeExit.h"
24#include "llvm/ADT/SmallVector.h"
25#include "llvm/ADT/Statistic.h"
26#include "llvm/Analysis/AliasAnalysis.h"
27#include "llvm/Analysis/DomTreeUpdater.h"
28#include "llvm/Analysis/LoopInfo.h"
29#include "llvm/Analysis/OptimizationRemarkEmitter.h"
30#include "llvm/Analysis/TargetTransformInfo.h"
31#include "llvm/Analysis/ValueTracking.h"
32#include "llvm/Analysis/VectorUtils.h"
33#include "llvm/IR/CFG.h"
34#include "llvm/IR/DataLayout.h"
35#include "llvm/IR/DebugInfoMetadata.h"
36#include "llvm/IR/DerivedTypes.h"
37#include "llvm/IR/Function.h"
38#include "llvm/IR/IRBuilder.h"
39#include "llvm/IR/InstrTypes.h"
40#include "llvm/IR/Instructions.h"
41#include "llvm/IR/IntrinsicInst.h"
42#include "llvm/IR/MatrixBuilder.h"
43#include "llvm/IR/PatternMatch.h"
44#include "llvm/IR/ProfDataUtils.h"
45#include "llvm/Support/Alignment.h"
46#include "llvm/Support/CommandLine.h"
47#include "llvm/Support/Compiler.h"
48#include "llvm/Support/Debug.h"
49#include "llvm/Transforms/Utils/BasicBlockUtils.h"
50#include "llvm/Transforms/Utils/LoopUtils.h"
51#include "llvm/Transforms/Utils/MatrixUtils.h"
52
53#include <cmath>
54
55using namespace llvm;
56using namespace PatternMatch;
57
58#define DEBUG_TYPE "lower-matrix-intrinsics"
59
60STATISTIC(FlattenedMatrices, "Number of matrix flattenings");
61STATISTIC(ReshapedMatrices, "Number of matrix reshapes");
62STATISTIC(SplitMatrices, "Number of matrix splits");
63
64static cl::opt<bool>
65 FuseMatrix("fuse-matrix", cl::init(Val: true), cl::Hidden,
66 cl::desc("Enable/disable fusing matrix instructions."));
67// TODO: Allow and use non-square tiles.
68static cl::opt<unsigned> TileSize(
69 "fuse-matrix-tile-size", cl::init(Val: 4), cl::Hidden,
70 cl::desc(
71 "Tile size for matrix instruction fusion using square-shaped tiles."));
72static cl::opt<unsigned>
73 TileLoopsThreshold("fuse-matrix-loops-threshold", cl::init(Val: 200), cl::Hidden,
74 cl::desc("Generate loop nests for tiling when expected "
75 "number of operations exceeds threshold."));
76static cl::opt<bool> ForceFusion(
77 "force-fuse-matrix", cl::init(Val: false), cl::Hidden,
78 cl::desc("Force matrix instruction fusion even if not profitable."));
79static cl::opt<bool> AllowContractEnabled(
80 "matrix-allow-contract", cl::init(Val: false), cl::Hidden,
81 cl::desc("Allow the use of FMAs if available and profitable. This may "
82 "result in different results, due to less rounding error."));
83
84static cl::opt<bool>
85 VerifyShapeInfo("verify-matrix-shapes", cl::Hidden,
86 cl::desc("Enable/disable matrix shape verification."),
87 cl::init(Val: false));
88
89enum class MatrixLayoutTy { ColumnMajor, RowMajor };
90
91static cl::opt<MatrixLayoutTy> MatrixLayout(
92 "matrix-default-layout", cl::init(Val: MatrixLayoutTy::ColumnMajor),
93 cl::desc("Sets the default matrix layout"),
94 cl::values(clEnumValN(MatrixLayoutTy::ColumnMajor, "column-major",
95 "Use column-major layout"),
96 clEnumValN(MatrixLayoutTy::RowMajor, "row-major",
97 "Use row-major layout")));
98
99static cl::opt<bool> PrintAfterTransposeOpt("matrix-print-after-transpose-opt",
100 cl::init(Val: false));
101
102static cl::opt<unsigned> SplitMatmulRemainderOverThreshold(
103 "matrix-split-matmul-remainder-over-threshold", cl::Hidden,
104 cl::desc("Illegal remainder vectors over this size in bits should be split "
105 "in the inner loop of matmul"),
106 cl::init(Val: 0));
107
108namespace llvm {
109extern cl::opt<bool> ProfcheckDisableMetadataFixes;
110} // end namespace llvm
111
112/// Helper function to either return Scope, if it is a subprogram or the
113/// attached subprogram for a local scope.
114static DISubprogram *getSubprogram(DIScope *Scope) {
115 if (auto *Subprogram = dyn_cast<DISubprogram>(Val: Scope))
116 return Subprogram;
117 return cast<DILocalScope>(Val: Scope)->getSubprogram();
118}
119
120/// Return true if V is a splat of a value (which is used when multiplying a
121/// matrix with a scalar).
122static bool isSplat(Value *V) {
123 if (auto *SV = dyn_cast<ShuffleVectorInst>(Val: V))
124 return SV->isZeroEltSplat();
125 return false;
126}
127
128/// Match any mul operation (fp or integer).
129template <typename LTy, typename RTy>
130static auto m_AnyMul(const LTy &L, const RTy &R) {
131 return m_CombineOr(m_Mul(L, R), m_FMul(L, R));
132}
133
134/// Match any add operation (fp or integer).
135template <typename LTy, typename RTy>
136static auto m_AnyAdd(const LTy &L, const RTy &R) {
137 return m_CombineOr(m_Add(L, R), m_FAdd(L, R));
138}
139
140// Given an element pointer \p BasePtr to the start of a (sub) matrix, compute
141// the start address of vector \p VecIdx with type (\p EltType x \p NumElements)
142// assuming \p Stride elements between start two consecutive vectors.
143// \p Stride must be >= \p NumElements.
144// For column-major matrixes, the function computes the address of a column
145// vectors and \p NumElements must be set to the number of elements in a column
146// (= number of rows of the matrix). For row-major matrixes, the function
147// computes the address of a row vector and \p NumElements must be set to the
148// number of elements in a column (= number of columns of the matrix).
149//
150// Consider a 4x4 matrix in column-mjaor layout like below
151//
152// 0 1 2 3
153// 0 v_0_0 v_0_1 v_0_2 v_0_3
154// 1 v_1_0 v_1_1 v_1_2 v_1_3
155// 2 v_2_0 v_2_1 v_2_2 v_2_3
156// 3 v_3_0 v_3_1 v_3_2 v_3_3
157
158// To compute the column addresses for a 2x3 sub-matrix at row 1 and column 1,
159// we need a pointer to the first element of the submatrix as base pointer.
160// Then we can use computeVectorAddr to compute the addresses for the columns
161// of the sub-matrix.
162//
163// Column 0: computeVectorAddr(Base, 0 (column), 4 (stride), 2 (num rows), ..)
164// -> just returns Base
165// Column 1: computeVectorAddr(Base, 1 (column), 4 (stride), 2 (num rows), ..)
166// -> returns Base + (1 * 4)
167// Column 2: computeVectorAddr(Base, 2 (column), 4 (stride), 2 (num rows), ..)
168// -> returns Base + (2 * 4)
169//
170// The graphic below illustrates the number of elements in a column (marked
171// with |) and the number of skipped elements (marked with }).
172//
173// v_0_0 v_0_1 {v_0_2 {v_0_3
174// Base Col 1 Col 2
175// | | |
176// v_1_0 |v_1_1 |v_1_2 |v_1_3
177// v_2_0 |v_2_1 |v_2_2 |v_2_3
178// v_3_0 {v_3_1 {v_3_2 v_3_3
179//
180static Value *computeVectorAddr(Value *BasePtr, Value *VecIdx, Value *Stride,
181 unsigned NumElements, Type *EltType,
182 IRBuilder<> &Builder) {
183
184 assert((!isa<ConstantInt>(Stride) ||
185 cast<ConstantInt>(Stride)->getZExtValue() >= NumElements) &&
186 "Stride must be >= the number of elements in the result vector.");
187
188 // Compute the start of the vector with index VecIdx as VecIdx * Stride.
189 Value *VecStart = Builder.CreateMul(LHS: VecIdx, RHS: Stride, Name: "vec.start");
190
191 // Get pointer to the start of the selected vector. Skip GEP creation,
192 // if we select vector 0.
193 if (isa<ConstantInt>(Val: VecStart) && cast<ConstantInt>(Val: VecStart)->isZero())
194 VecStart = BasePtr;
195 else
196 VecStart = Builder.CreateInBoundsGEP(Ty: EltType, Ptr: BasePtr, IdxList: VecStart, Name: "vec.gep");
197
198 return VecStart;
199}
200
201namespace {
202struct ShapeInfo {
203 unsigned NumRows;
204 unsigned NumColumns;
205
206 bool IsColumnMajor;
207
208 ShapeInfo(unsigned NumRows = 0, unsigned NumColumns = 0)
209 : NumRows(NumRows), NumColumns(NumColumns),
210 IsColumnMajor(MatrixLayout == MatrixLayoutTy::ColumnMajor) {}
211
212 ShapeInfo(Value *NumRows, Value *NumColumns)
213 : ShapeInfo(cast<ConstantInt>(Val: NumRows)->getZExtValue(),
214 cast<ConstantInt>(Val: NumColumns)->getZExtValue()) {}
215
216 bool operator==(const ShapeInfo &other) {
217 return NumRows == other.NumRows && NumColumns == other.NumColumns;
218 }
219 bool operator!=(const ShapeInfo &other) { return !(*this == other); }
220
221 /// Returns true if shape-information is defined, meaning both dimensions
222 /// are != 0.
223 operator bool() const {
224 assert(NumRows == 0 || NumColumns != 0);
225 return NumRows != 0;
226 }
227
228 unsigned getStride() const {
229 if (IsColumnMajor)
230 return NumRows;
231 return NumColumns;
232 }
233
234 unsigned getNumVectors() const {
235 if (IsColumnMajor)
236 return NumColumns;
237 return NumRows;
238 }
239
240 /// Returns the transposed shape.
241 ShapeInfo t() const { return ShapeInfo(NumColumns, NumRows); }
242
243 friend raw_ostream &operator<<(raw_ostream &OS, ShapeInfo SI);
244
245 LLVM_DUMP_METHOD void dump() const { dbgs() << *this << '\n'; }
246};
247
248raw_ostream &operator<<(raw_ostream &OS, ShapeInfo SI) {
249 return OS << SI.NumRows << 'x' << SI.NumColumns;
250}
251
252} // namespace
253
254static bool isShapePreserving(Value *V) {
255 Instruction *I = dyn_cast<Instruction>(Val: V);
256 if (!I)
257 return true;
258
259 if (isa<SelectInst>(Val: I))
260 return true;
261
262 if (I->isBinaryOp())
263 return true;
264
265 if (auto *Cast = dyn_cast<CastInst>(Val: V)) {
266 switch (Cast->getOpcode()) {
267 case llvm::Instruction::Trunc:
268 case llvm::Instruction::ZExt:
269 case llvm::Instruction::SExt:
270 case llvm::Instruction::FPToUI:
271 case llvm::Instruction::FPToSI:
272 case llvm::Instruction::UIToFP:
273 case llvm::Instruction::SIToFP:
274 case llvm::Instruction::FPTrunc:
275 case llvm::Instruction::FPExt:
276 return true;
277 case llvm::Instruction::AddrSpaceCast:
278 case CastInst::PtrToAddr:
279 case CastInst::PtrToInt:
280 case CastInst::IntToPtr:
281 return false;
282 case CastInst::BitCast: {
283 if (auto *SrcVTy = dyn_cast<FixedVectorType>(Val: Cast->getSrcTy()))
284 if (auto *DestVTy = dyn_cast<FixedVectorType>(Val: Cast->getDestTy()))
285 return SrcVTy->getNumElements() == DestVTy->getNumElements();
286 return false;
287 }
288 case llvm::Instruction::CastOpsEnd:
289 llvm_unreachable("not an actual cast op");
290 }
291 llvm_unreachable("unhandled cast opcode");
292 }
293
294 if (auto *II = dyn_cast<IntrinsicInst>(Val: V))
295 switch (II->getIntrinsicID()) {
296 case Intrinsic::abs:
297 case Intrinsic::fabs:
298 return true;
299 default:
300 return false;
301 }
302
303 switch (I->getOpcode()) {
304 case Instruction::PHI:
305 case Instruction::FNeg:
306 return true;
307 default:
308 return false;
309 }
310}
311
312/// Return an iterator over the operands of \p I that should share shape
313/// information with \p I.
314static iterator_range<Use *> getShapedOperandsForInst(Instruction *I) {
315 assert(isShapePreserving(I) &&
316 "Can't retrieve shaped operands for an instruction that does not "
317 "preserve shape information");
318 auto Ops = I->operands();
319 return isa<SelectInst>(Val: I) ? drop_begin(RangeOrContainer&: Ops) : Ops;
320}
321
322/// Return the ShapeInfo for the result of \p I, it it can be determined.
323static std::optional<ShapeInfo>
324computeShapeInfoForInst(Instruction *I,
325 const DenseMap<Value *, ShapeInfo> &ShapeMap) {
326 Value *M;
327 Value *N;
328 Value *K;
329 if (match(V: I, P: m_Intrinsic<Intrinsic::matrix_multiply>(
330 Op0: m_Value(), Op1: m_Value(), Op2: m_Value(V&: M), Op3: m_Value(V&: N), Op4: m_Value(V&: K))))
331 return ShapeInfo(M, K);
332 if (match(V: I, P: m_Intrinsic<Intrinsic::matrix_transpose>(Op0: m_Value(), Op1: m_Value(V&: M),
333 Op2: m_Value(V&: N)))) {
334 // Flip dimensions.
335 return ShapeInfo(N, M);
336 }
337 if (match(V: I, P: m_Intrinsic<Intrinsic::matrix_column_major_store>(
338 Op0: m_Value(), Op1: m_Value(), Op2: m_Value(), Op3: m_Value(), Op4: m_Value(V&: M),
339 Op5: m_Value(V&: N))))
340 return ShapeInfo(N, M);
341 if (match(V: I, P: m_Intrinsic<Intrinsic::matrix_column_major_load>(
342 Op0: m_Value(), Op1: m_Value(), Op2: m_Value(), Op3: m_Value(V&: M), Op4: m_Value(V&: N))))
343 return ShapeInfo(M, N);
344 Value *MatrixA;
345 if (match(V: I, P: m_Store(ValueOp: m_Value(V&: MatrixA), PointerOp: m_Value()))) {
346 auto OpShape = ShapeMap.find(Val: MatrixA);
347 if (OpShape != ShapeMap.end())
348 return OpShape->second;
349 }
350
351 if (isShapePreserving(V: I)) {
352 auto ShapedOps = getShapedOperandsForInst(I);
353 // Find the first operand that has a known shape and use that.
354 for (auto &Op : ShapedOps) {
355 auto OpShape = ShapeMap.find(Val: Op.get());
356 if (OpShape != ShapeMap.end())
357 return OpShape->second;
358 }
359 }
360 return std::nullopt;
361}
362
363namespace {
364
365/// LowerMatrixIntrinsics contains the methods used to lower matrix intrinsics.
366///
367/// Currently, the lowering for each matrix intrinsic is done as follows:
368/// 1. Propagate the shape information from intrinsics to connected
369/// instructions.
370/// 2. Lower instructions with shape information (assuming column-major layout).
371/// The lowering works similarly using row-major layout.
372/// 2.1. Get column vectors for each argument. If we already lowered the
373/// definition of an argument, use the produced column vectors directly.
374/// If not, split the operand vector containing an embedded matrix into
375/// a set of column vectors,
376/// 2.2. Lower the instruction in terms of column major operations, which
377/// yields a set of column vectors containing result matrix. Note that we
378/// lower all instructions that have shape information. Besides the
379/// intrinsics, this includes stores for example.
380/// 2.3. Update uses of the lowered instruction. If we have shape information
381/// for a user, there is nothing to do, as we will look up the result
382/// column matrix when lowering the user. For other uses, we embed the
383/// result matrix in a flat vector and update the use.
384/// 2.4. Cache the result column matrix for the instruction we lowered
385/// 3. After we lowered all instructions in a function, remove the now
386/// obsolete instructions.
387///
388class LowerMatrixIntrinsics {
389 Function &Func;
390 const DataLayout &DL;
391 const TargetTransformInfo &TTI;
392 FunctionAnalysisManager *AM;
393 AliasAnalysis *AA = nullptr;
394 DominatorTree *DT = nullptr;
395 LoopInfo *LI = nullptr;
396 OptimizationRemarkEmitter *ORE = nullptr;
397
398 /// Contains estimates of the number of operations (loads, stores, compute)
399 /// required to lower a matrix operation.
400 struct OpInfoTy {
401 /// Number of stores emitted to generate this matrix.
402 unsigned NumStores = 0;
403 /// Number of loads emitted to generate this matrix.
404 unsigned NumLoads = 0;
405 /// Number of compute operations emitted to generate this matrix.
406 unsigned NumComputeOps = 0;
407 /// Most of the time transposes can be fused with matrix multiplies or can
408 /// be folded away via algebraic simplifications. This is the number of
409 /// transposes that we failed to make "free" via such optimizations.
410 unsigned NumExposedTransposes = 0;
411
412 OpInfoTy &operator+=(const OpInfoTy &RHS) {
413 NumStores += RHS.NumStores;
414 NumLoads += RHS.NumLoads;
415 NumComputeOps += RHS.NumComputeOps;
416 NumExposedTransposes += RHS.NumExposedTransposes;
417 return *this;
418 }
419 };
420
421 /// Wrapper class representing a matrix as a set of vectors, either in row or
422 /// column major layout. All vectors must have the same vector type.
423 class MatrixTy {
424 SmallVector<Value *, 16> Vectors;
425
426 OpInfoTy OpInfo;
427
428 bool IsColumnMajor = true;
429
430 public:
431 MatrixTy() : IsColumnMajor(MatrixLayout == MatrixLayoutTy::ColumnMajor) {}
432 MatrixTy(ArrayRef<Value *> Vectors)
433 : Vectors(Vectors),
434 IsColumnMajor(MatrixLayout == MatrixLayoutTy::ColumnMajor) {}
435 MatrixTy(unsigned NumRows, unsigned NumColumns, Type *EltTy)
436 : IsColumnMajor(MatrixLayout == MatrixLayoutTy::ColumnMajor) {
437
438 unsigned D = isColumnMajor() ? NumColumns : NumRows;
439 for (unsigned J = 0; J < D; ++J)
440 addVector(V: PoisonValue::get(T: FixedVectorType::get(
441 ElementType: EltTy, NumElts: isColumnMajor() ? NumRows : NumColumns)));
442 }
443
444 Value *getVector(unsigned i) const { return Vectors[i]; }
445 Value *getColumn(unsigned i) const {
446 assert(isColumnMajor() && "only supported for column-major matrixes");
447 return Vectors[i];
448 }
449 Value *getRow(unsigned i) const {
450 assert(!isColumnMajor() && "only supported for row-major matrixes");
451 return Vectors[i];
452 }
453
454 void setVector(unsigned i, Value *V) { Vectors[i] = V; }
455
456 Type *getElementType() const { return getVectorTy()->getElementType(); }
457
458 unsigned getNumVectors() const {
459 if (isColumnMajor())
460 return getNumColumns();
461 return getNumRows();
462 }
463
464 unsigned getNumColumns() const {
465 if (isColumnMajor())
466 return Vectors.size();
467 else {
468 assert(Vectors.size() > 0 && "Cannot call getNumRows without columns");
469 return getVectorTy()->getNumElements();
470 }
471 }
472 unsigned getNumRows() const {
473 if (isColumnMajor()) {
474 assert(Vectors.size() > 0 && "Cannot call getNumRows without columns");
475 return getVectorTy()->getNumElements();
476 } else
477 return Vectors.size();
478 }
479
480 void addVector(Value *V) { Vectors.push_back(Elt: V); }
481 FixedVectorType *getColumnTy() {
482 assert(isColumnMajor() && "only supported for column-major matrixes");
483 return getVectorTy();
484 }
485
486 FixedVectorType *getVectorTy() const {
487 return cast<FixedVectorType>(Val: Vectors[0]->getType());
488 }
489
490 iterator_range<SmallVector<Value *, 8>::iterator> columns() {
491 assert(isColumnMajor() &&
492 "columns() only supported for column-major matrixes");
493 return make_range(x: Vectors.begin(), y: Vectors.end());
494 }
495
496 iterator_range<SmallVector<Value *, 8>::iterator> vectors() {
497 return make_range(x: Vectors.begin(), y: Vectors.end());
498 }
499
500 /// Embed the vectors of the matrix into a flat vector by concatenating
501 /// them.
502 Value *embedInVector(IRBuilder<> &Builder) const {
503 return Vectors.size() == 1 ? Vectors[0]
504 : concatenateVectors(Builder, Vecs: Vectors);
505 }
506
507 MatrixTy &addNumLoads(unsigned N) {
508 OpInfo.NumLoads += N;
509 return *this;
510 }
511
512 void setNumLoads(unsigned N) { OpInfo.NumLoads = N; }
513
514 MatrixTy &addNumStores(unsigned N) {
515 OpInfo.NumStores += N;
516 return *this;
517 }
518
519 MatrixTy &addNumExposedTransposes(unsigned N) {
520 OpInfo.NumExposedTransposes += N;
521 return *this;
522 }
523
524 MatrixTy &addNumComputeOps(unsigned N) {
525 OpInfo.NumComputeOps += N;
526 return *this;
527 }
528
529 unsigned getNumStores() const { return OpInfo.NumStores; }
530 unsigned getNumLoads() const { return OpInfo.NumLoads; }
531 unsigned getNumComputeOps() const { return OpInfo.NumComputeOps; }
532
533 const OpInfoTy &getOpInfo() const { return OpInfo; }
534
535 bool isColumnMajor() const { return IsColumnMajor; }
536
537 unsigned getStride() const {
538 if (isColumnMajor())
539 return getNumRows();
540 return getNumColumns();
541 }
542
543 ShapeInfo shape() const { return {getNumRows(), getNumColumns()}; }
544
545 /// Extract a vector of \p NumElts starting at index (\p I, \p J). If the
546 /// matrix is column-major, the result vector is extracted from a column
547 /// vector, otherwise from a row vector.
548 Value *extractVector(unsigned I, unsigned J, unsigned NumElts,
549 IRBuilder<> &Builder) const {
550 Value *Vec = isColumnMajor() ? getColumn(i: J) : getRow(i: I);
551 assert(cast<FixedVectorType>(Vec->getType())->getNumElements() >=
552 NumElts &&
553 "Extracted vector will contain poison values");
554 return Builder.CreateShuffleVector(
555 V: Vec, Mask: createSequentialMask(Start: isColumnMajor() ? I : J, NumInts: NumElts, NumUndefs: 0),
556 Name: "block");
557 }
558 };
559
560 /// Maps instructions to their shape information. The shape information
561 /// describes the shape to be used while lowering. This matches the shape of
562 /// the result value of the instruction, with the only exceptions being store
563 /// instructions and the matrix_column_major_store intrinsics. For those, the
564 /// shape information indicates that those instructions should be lowered
565 /// using shape information as well. Note that extra care is needed when
566 /// erasing or RAUW'ing a value that is present in ShapeMap. If the
567 /// replacement is also a matrix operation, use
568 /// updateShapeAndReplaceAllUsesWith to make sure the replacement is added to
569 /// ShapeMap. We don't use ValueMap, as there are also cases where we do not
570 /// want to add shape information for a replacement instruction. When directly
571 /// erasing a value with an entry in ShapeMap, use
572 /// eraseFromParentAndRemoveFromShapeMap to make sure ShapeMap is also updated
573 /// accordingly.
574 DenseMap<Value *, ShapeInfo> ShapeMap;
575
576 /// List of instructions to remove. While lowering, we are not replacing all
577 /// users of a lowered instruction, if shape information is available and
578 /// those need to be removed after we finished lowering.
579 SmallVector<Instruction *, 16> ToRemove;
580
581 /// Map from instructions to their produced column matrix.
582 MapVector<Value *, MatrixTy> Inst2ColumnMatrix;
583
584private:
585 static FastMathFlags getFastMathFlags(Instruction *Inst) {
586 FastMathFlags FMF;
587
588 if (isa<FPMathOperator>(Val: *Inst))
589 FMF = Inst->getFastMathFlags();
590
591 FMF.setAllowContract(AllowContractEnabled || FMF.allowContract());
592
593 return FMF;
594 }
595
596public:
597 LowerMatrixIntrinsics(Function &F, TargetTransformInfo &TTI,
598 FunctionAnalysisManager *AM)
599 : Func(F), DL(F.getDataLayout()), TTI(TTI), AM(AM) {}
600
601 unsigned getNumOps(Type *VT) {
602 assert(isa<FixedVectorType>(VT) && "Expected vector type");
603 return getNumOps(ST: VT->getScalarType(),
604 N: cast<FixedVectorType>(Val: VT)->getNumElements());
605 }
606
607 /// Is this the minimal version executed in the backend pipelines.
608 bool isMinimal() const {
609 return !DT;
610 }
611
612 /// Return the estimated number of vector ops required for an operation on
613 /// \p VT * N.
614 unsigned getNumOps(Type *ST, unsigned N) {
615 return std::ceil(x: (ST->getPrimitiveSizeInBits() * N).getFixedValue() /
616 double(TTI.getRegisterBitWidth(
617 K: TargetTransformInfo::RGK_FixedWidthVector)
618 .getFixedValue()));
619 }
620
621 /// Estimate the number of native vector operations for a multiply of matrices
622 /// with dimensions \p R x \p M and \p M x \p C. Native ops are computed as
623 /// ceil(ElementCount * ElementBits / RegisterBits).
624 ///
625 /// Native vector ops per operation type (VF = native vector elements):
626 /// FMAs: C * ceil(R/VF) * M (one FMA per VF output elements)
627 /// A loads: ceil(R/VF) * M (A has M columns, ceil(R/VF) native loads each)
628 /// B loads: ceil(M/VF) * C (B has C columns, ceil(M/VF) native loads each)
629 /// Stores: C * ceil(R/VF) (one store per VF output elements)
630 unsigned getNumNativeVectorOps(Type *EltType, unsigned R, unsigned M,
631 unsigned C) {
632 unsigned NumFMAs = C * getNumOps(ST: EltType, N: R) * M;
633 unsigned NumALoads = getNumOps(ST: EltType, N: R) * M;
634 unsigned NumBLoads = getNumOps(ST: EltType, N: M) * C;
635 unsigned NumStores = getNumOps(ST: EltType, N: R) * C;
636 return NumFMAs + NumALoads + NumBLoads + NumStores;
637 }
638
639 /// Return the set of vectors that a matrix value is lowered to.
640 ///
641 /// If we lowered \p MatrixVal, just return the cache result matrix. Otherwise
642 /// split the flat vector \p MatrixVal containing a matrix with shape \p SI
643 /// into vectors.
644 MatrixTy getMatrix(Value *MatrixVal, const ShapeInfo &SI,
645 IRBuilder<> &Builder) {
646 FixedVectorType *VType = cast<FixedVectorType>(Val: MatrixVal->getType());
647 assert(VType->getNumElements() == SI.NumRows * SI.NumColumns &&
648 "The vector size must match the number of matrix elements");
649
650 // Check if we lowered MatrixVal using shape information. In that case,
651 // return the existing matrix, if it matches the requested shape
652 // information. If there is a mis-match, embed the result in a flat
653 // vector and split it later.
654 auto Found = Inst2ColumnMatrix.find(Key: MatrixVal);
655 if (Found != Inst2ColumnMatrix.end()) {
656 MatrixTy &M = Found->second;
657 // Return the found matrix, if its shape matches the requested shape
658 // information
659 if (SI.NumRows == M.getNumRows() && SI.NumColumns == M.getNumColumns())
660 return M;
661
662 MatrixVal = M.embedInVector(Builder);
663 }
664
665 // Otherwise split MatrixVal.
666 SmallVector<Value *, 16> SplitVecs;
667 for (unsigned MaskStart = 0; MaskStart < VType->getNumElements();
668 MaskStart += SI.getStride()) {
669 Value *V = Builder.CreateShuffleVector(
670 V: MatrixVal, Mask: createSequentialMask(Start: MaskStart, NumInts: SI.getStride(), NumUndefs: 0),
671 Name: "split");
672 SplitVecs.push_back(Elt: V);
673 }
674
675 if (Instruction *Inst = dyn_cast<Instruction>(Val: MatrixVal)) {
676 if (Found != Inst2ColumnMatrix.end()) {
677 // FIXME: re: "at least": SplitVecs.size() doesn't count the shuffles
678 // that embedInVector created.
679 LLVM_DEBUG(dbgs() << "matrix reshape from " << Found->second.shape()
680 << " to " << SI << " using at least "
681 << SplitVecs.size() << " shuffles on behalf of:\n"
682 << *Inst << '\n');
683 ReshapedMatrices++;
684 } else if (!ShapeMap.contains(Val: MatrixVal)) {
685 LLVM_DEBUG(
686 dbgs()
687 << "splitting a " << SI << " matrix with " << SplitVecs.size()
688 << " shuffles beacuse we do not have a shape-aware lowering for "
689 "its def:\n"
690 << *Inst << '\n');
691 (void)Inst;
692 SplitMatrices++;
693 } else {
694 // The ShapeMap has it, so it's a case where we're being lowered
695 // before the def, and we expect that InstCombine will clean things up
696 // afterward.
697 }
698 }
699
700 return {SplitVecs};
701 }
702
703 /// If \p V already has a known shape return false. Otherwise set the shape
704 /// for instructions that support it.
705 bool setShapeInfo(Value *V, ShapeInfo Shape) {
706 assert(Shape && "Shape not set");
707 if (isa<UndefValue>(Val: V) || !supportsShapeInfo(V))
708 return false;
709
710 auto SIter = ShapeMap.find(Val: V);
711 if (SIter != ShapeMap.end()) {
712 if (VerifyShapeInfo && (SIter->second.NumRows != Shape.NumRows ||
713 SIter->second.NumColumns != Shape.NumColumns)) {
714 errs() << "Conflicting shapes (" << SIter->second.NumRows << "x"
715 << SIter->second.NumColumns << " vs " << Shape.NumRows << "x"
716 << Shape.NumColumns << ") for " << *V << "\n";
717 report_fatal_error(
718 reason: "Matrix shape verification failed, compilation aborted!");
719 }
720
721 LLVM_DEBUG(dbgs() << " not overriding existing shape: "
722 << SIter->second.NumRows << " "
723 << SIter->second.NumColumns << " for " << *V << "\n");
724 return false;
725 }
726
727 ShapeMap.insert(KV: {V, Shape});
728 LLVM_DEBUG(dbgs() << " " << Shape.NumRows << " x " << Shape.NumColumns
729 << " for " << *V << "\n");
730 return true;
731 }
732
733 /// Returns true if shape information can be used for \p V. The supported
734 /// instructions must match the instructions that can be lowered by this pass.
735 bool supportsShapeInfo(Value *V) {
736 Instruction *Inst = dyn_cast<Instruction>(Val: V);
737 if (!Inst)
738 return false;
739
740 IntrinsicInst *II = dyn_cast<IntrinsicInst>(Val: Inst);
741 if (II)
742 switch (II->getIntrinsicID()) {
743 case Intrinsic::matrix_multiply:
744 case Intrinsic::matrix_transpose:
745 case Intrinsic::matrix_column_major_load:
746 case Intrinsic::matrix_column_major_store:
747 return true;
748 default:
749 break;
750 }
751 return isShapePreserving(V) || isa<StoreInst>(Val: V) || isa<LoadInst>(Val: V);
752 }
753
754 /// Propagate the shape information of instructions to their users.
755 /// The work list contains instructions for which we can compute the shape,
756 /// either based on the information provided by matrix intrinsics or known
757 /// shapes of operands.
758 SmallVector<Instruction *, 32>
759 propagateShapeForward(SmallVectorImpl<Instruction *> &WorkList) {
760 SmallVector<Instruction *, 32> NewWorkList;
761 // Pop an element for which we guaranteed to have at least one of the
762 // operand shapes. Add the shape for this and then add users to the work
763 // list.
764 LLVM_DEBUG(dbgs() << "Forward-propagate shapes:\n");
765 while (!WorkList.empty()) {
766 Instruction *Inst = WorkList.pop_back_val();
767
768 // New entry, set the value and insert operands
769 bool Propagate = false;
770 if (auto SI = computeShapeInfoForInst(I: Inst, ShapeMap))
771 Propagate = setShapeInfo(V: Inst, Shape: *SI);
772
773 if (Propagate) {
774 NewWorkList.push_back(Elt: Inst);
775 for (auto *User : Inst->users())
776 if (ShapeMap.count(Val: User) == 0)
777 WorkList.push_back(Elt: cast<Instruction>(Val: User));
778 }
779 }
780
781 return NewWorkList;
782 }
783
784 /// Propagate the shape to operands of instructions with shape information.
785 /// \p Worklist contains the instruction for which we already know the shape.
786 SmallVector<Instruction *, 32>
787 propagateShapeBackward(SmallVectorImpl<Instruction *> &WorkList) {
788 SmallVector<Instruction *, 32> NewWorkList;
789
790 auto pushInstruction = [](Value *V,
791 SmallVectorImpl<Instruction *> &WorkList) {
792 Instruction *I = dyn_cast<Instruction>(Val: V);
793 if (I)
794 WorkList.push_back(Elt: I);
795 };
796 // Pop an element with known shape. Traverse the operands, if their shape
797 // derives from the result shape and is unknown, add it and add them to the
798 // worklist.
799 LLVM_DEBUG(dbgs() << "Backward-propagate shapes:\n");
800 while (!WorkList.empty()) {
801 Value *V = WorkList.pop_back_val();
802
803 size_t BeforeProcessingV = WorkList.size();
804 if (!isa<Instruction>(Val: V))
805 continue;
806
807 Value *MatrixA;
808 Value *MatrixB;
809 Value *M;
810 Value *N;
811 Value *K;
812 if (match(V, P: m_Intrinsic<Intrinsic::matrix_multiply>(
813 Op0: m_Value(V&: MatrixA), Op1: m_Value(V&: MatrixB), Op2: m_Value(V&: M),
814 Op3: m_Value(V&: N), Op4: m_Value(V&: K)))) {
815 if (setShapeInfo(V: MatrixA, Shape: {M, N}))
816 pushInstruction(MatrixA, WorkList);
817
818 if (setShapeInfo(V: MatrixB, Shape: {N, K}))
819 pushInstruction(MatrixB, WorkList);
820
821 } else if (match(V, P: m_Intrinsic<Intrinsic::matrix_transpose>(
822 Op0: m_Value(V&: MatrixA), Op1: m_Value(V&: M), Op2: m_Value(V&: N)))) {
823 // Flip dimensions.
824 if (setShapeInfo(V: MatrixA, Shape: {M, N}))
825 pushInstruction(MatrixA, WorkList);
826 } else if (match(V, P: m_Intrinsic<Intrinsic::matrix_column_major_store>(
827 Op0: m_Value(V&: MatrixA), Op1: m_Value(), Op2: m_Value(), Op3: m_Value(),
828 Op4: m_Value(V&: M), Op5: m_Value(V&: N)))) {
829 if (setShapeInfo(V: MatrixA, Shape: {M, N})) {
830 pushInstruction(MatrixA, WorkList);
831 }
832 } else if (isa<LoadInst>(Val: V) ||
833 match(V, P: m_Intrinsic<Intrinsic::matrix_column_major_load>())) {
834 // Nothing to do, no matrix input.
835 } else if (isa<StoreInst>(Val: V)) {
836 // Nothing to do. We forward-propagated to this so we would just
837 // backward propagate to an instruction with an already known shape.
838 } else if (isShapePreserving(V)) {
839 auto ShapedOps = getShapedOperandsForInst(I: cast<Instruction>(Val: V));
840 // Propagate to all operands.
841 ShapeInfo Shape = ShapeMap[V];
842 for (Use &U : ShapedOps) {
843 if (setShapeInfo(V: U.get(), Shape))
844 pushInstruction(U.get(), WorkList);
845 }
846 }
847 // After we discovered new shape info for new instructions in the
848 // worklist, we use their users as seeds for the next round of forward
849 // propagation.
850 for (size_t I = BeforeProcessingV; I != WorkList.size(); I++)
851 for (User *U : WorkList[I]->users())
852 if (isa<Instruction>(Val: U) && V != U)
853 NewWorkList.push_back(Elt: cast<Instruction>(Val: U));
854 }
855 return NewWorkList;
856 }
857
858 /// (Op0 op Op1)^T -> Op0^T op Op1^T
859 /// Transpose \p Op0 and \p Op1 of shape \p Shape0 and \p Shape1, then use
860 /// them on both sides of \p Operation.
861 Instruction *distributeTransposes(
862 Value *Op0, ShapeInfo Shape0, Value *Op1, ShapeInfo Shape1,
863 MatrixBuilder &Builder,
864 function_ref<Instruction *(Value *, ShapeInfo, Value *, ShapeInfo)>
865 Operation) {
866 Value *T0 = Builder.CreateMatrixTranspose(
867 Matrix: Op0, Rows: Shape0.NumRows, Columns: Shape0.NumColumns, Name: Op0->getName() + "_t");
868 // We are being run after shape prop, add shape for newly created
869 // instructions so that we lower them later.
870 setShapeInfo(V: T0, Shape: Shape0.t());
871 Value *T1 = Builder.CreateMatrixTranspose(
872 Matrix: Op1, Rows: Shape1.NumRows, Columns: Shape1.NumColumns, Name: Op1->getName() + "_t");
873 setShapeInfo(V: T1, Shape: Shape1.t());
874 return Operation(T0, Shape0.t(), T1, Shape1.t());
875 }
876
877 /// Erase \p Inst from both ShapeMap (if an entry exists) and erase \p Inst
878 /// itself.
879 void eraseFromParentAndRemoveFromShapeMap(Instruction *Inst) {
880 ShapeMap.erase(Val: Inst);
881 Inst->eraseFromParent();
882 }
883
884 /// Erase \p V from \p BB and move \II forward to avoid invalidating
885 /// iterators.
886 void eraseFromParentAndMove(Value *V, BasicBlock::reverse_iterator &II,
887 BasicBlock &BB) {
888 auto *Inst = cast<Instruction>(Val: V);
889 // Still used, don't erase.
890 if (!Inst->use_empty())
891 return;
892 if (II != BB.rend() && Inst == &*II)
893 ++II;
894 eraseFromParentAndRemoveFromShapeMap(Inst);
895 }
896
897 /// Add a new entry to ShapeMap for \p New with \p Old's shape info, erase the
898 /// entry for \p Old and replace all uses of \p Old with \p New.
899 void updateShapeAndReplaceAllUsesWith(Instruction &Old, Value *New) {
900 // We need to remove Old from the ShapeMap otherwise RAUW will replace it
901 // with New. We should only add New it it supportsShapeInfo so we insert
902 // it conditionally instead.
903 auto S = ShapeMap.find(Val: &Old);
904 if (S != ShapeMap.end()) {
905 ShapeInfo Shape = S->second;
906 ShapeMap.erase(I: S);
907 if (supportsShapeInfo(V: New))
908 ShapeMap.insert(KV: {New, Shape});
909 }
910 Old.replaceAllUsesWith(V: New);
911 }
912
913 /// Sink a top-level transpose inside matmuls and adds.
914 /// This creates and erases instructions as needed, and returns the newly
915 /// created instruction while updating the iterator to avoid invalidation. If
916 /// this returns nullptr, no new instruction was created.
917 Instruction *sinkTranspose(Instruction &I, BasicBlock::reverse_iterator &II,
918 bool &Changed) {
919 BasicBlock &BB = *I.getParent();
920 IRBuilder<> IB(&I);
921 MatrixBuilder Builder(IB);
922
923 Value *TA, *TAMA, *TAMB;
924 ConstantInt *R, *K, *C;
925 if (!match(V: &I, P: m_Intrinsic<Intrinsic::matrix_transpose>(
926 Op0: m_Value(V&: TA), Op1: m_ConstantInt(CI&: R), Op2: m_ConstantInt(CI&: C))))
927 return nullptr;
928
929 // Transpose of a transpose is a nop when the shapes match.
930 Value *TATA;
931 if (match(V: TA, P: m_Intrinsic<Intrinsic::matrix_transpose>(
932 Op0: m_Value(V&: TATA), Op1: m_Specific(V: C), Op2: m_Specific(V: R)))) {
933 updateShapeAndReplaceAllUsesWith(Old&: I, New: TATA);
934 eraseFromParentAndMove(V: &I, II, BB);
935 eraseFromParentAndMove(V: TA, II, BB);
936 Changed = true;
937 return nullptr;
938 }
939
940 // k^T -> k
941 if (isSplat(V: TA)) {
942 updateShapeAndReplaceAllUsesWith(Old&: I, New: TA);
943 eraseFromParentAndMove(V: &I, II, BB);
944 Changed = true;
945 return nullptr;
946 }
947
948 // (A * B)^t -> B^t * A^t
949 // RxK KxC CxK KxR
950 if (match(V: TA, P: m_Intrinsic<Intrinsic::matrix_multiply>(
951 Op0: m_Value(V&: TAMA), Op1: m_Value(V&: TAMB), Op2: m_ConstantInt(CI&: R),
952 Op3: m_ConstantInt(CI&: K), Op4: m_ConstantInt(CI&: C)))) {
953 auto NewInst = distributeTransposes(
954 Op0: TAMB, Shape0: {K, C}, Op1: TAMA, Shape1: {R, K}, Builder,
955 Operation: [&](Value *T0, ShapeInfo Shape0, Value *T1, ShapeInfo Shape1) {
956 return Builder.CreateMatrixMultiply(LHS: T0, RHS: T1, LHSRows: Shape0.NumRows,
957 LHSColumns: Shape0.NumColumns,
958 RHSColumns: Shape1.NumColumns, Name: "mmul");
959 });
960 updateShapeAndReplaceAllUsesWith(Old&: I, New: NewInst);
961 eraseFromParentAndMove(V: &I, II, BB);
962 eraseFromParentAndMove(V: TA, II, BB);
963 Changed = true;
964 return NewInst;
965 }
966
967 // Same as above, but with a mul, which occurs when multiplied
968 // with a scalar.
969 // (A * k)^t -> A^t * k
970 // R x C RxC
971 if (match(V: TA, P: m_AnyMul(L: m_Value(V&: TAMA), R: m_Value(V&: TAMB))) &&
972 (isSplat(V: TAMA) || isSplat(V: TAMB))) {
973 IRBuilder<> LocalBuilder(&I);
974 // We know that the transposed operand is of shape RxC.
975 // An when multiplied with a scalar, the shape is preserved.
976 auto NewInst = distributeTransposes(
977 Op0: TAMA, Shape0: {R, C}, Op1: TAMB, Shape1: {R, C}, Builder,
978 Operation: [&](Value *T0, ShapeInfo Shape0, Value *T1, ShapeInfo Shape1) {
979 bool IsFP = I.getType()->isFPOrFPVectorTy();
980 auto *Mul = IsFP ? LocalBuilder.CreateFMul(L: T0, R: T1, Name: "mmul")
981 : LocalBuilder.CreateMul(LHS: T0, RHS: T1, Name: "mmul");
982 auto *Result = cast<Instruction>(Val: Mul);
983 setShapeInfo(V: Result, Shape: Shape0);
984 return Result;
985 });
986 updateShapeAndReplaceAllUsesWith(Old&: I, New: NewInst);
987 eraseFromParentAndMove(V: &I, II, BB);
988 eraseFromParentAndMove(V: TA, II, BB);
989 Changed = true;
990 return NewInst;
991 }
992
993 // (A + B)^t -> A^t + B^t
994 // RxC RxC CxR CxR
995 if (match(V: TA, P: m_AnyAdd(L: m_Value(V&: TAMA), R: m_Value(V&: TAMB)))) {
996 IRBuilder<> LocalBuilder(&I);
997 auto NewInst = distributeTransposes(
998 Op0: TAMA, Shape0: {R, C}, Op1: TAMB, Shape1: {R, C}, Builder,
999 Operation: [&](Value *T0, ShapeInfo Shape0, Value *T1, ShapeInfo Shape1) {
1000 bool IsFP = I.getType()->isFPOrFPVectorTy();
1001 auto *Add = IsFP ? LocalBuilder.CreateFAdd(L: T0, R: T1, Name: "madd")
1002 : LocalBuilder.CreateAdd(LHS: T0, RHS: T1, Name: "madd");
1003
1004 auto *Result = cast<Instruction>(Val: Add);
1005 setShapeInfo(V: Result, Shape: Shape0);
1006 return Result;
1007 });
1008 updateShapeAndReplaceAllUsesWith(Old&: I, New: NewInst);
1009 eraseFromParentAndMove(V: &I, II, BB);
1010 eraseFromParentAndMove(V: TA, II, BB);
1011 Changed = true;
1012 return NewInst;
1013 }
1014
1015 return nullptr;
1016 }
1017
1018 bool liftTranspose(Instruction &I) {
1019 // Erase dead Instructions after lifting transposes from binops.
1020 auto CleanupBinOp = [this](Instruction &T, Value *A, Value *B) {
1021 if (T.use_empty())
1022 eraseFromParentAndRemoveFromShapeMap(Inst: &T);
1023 if (A->use_empty())
1024 eraseFromParentAndRemoveFromShapeMap(Inst: cast<Instruction>(Val: A));
1025 if (A != B && B->use_empty())
1026 eraseFromParentAndRemoveFromShapeMap(Inst: cast<Instruction>(Val: B));
1027 };
1028
1029 Value *A, *B, *AT, *BT;
1030 ConstantInt *R, *K, *C;
1031 // A^t * B ^t -> (B * A)^t
1032 if (match(V: &I, P: m_Intrinsic<Intrinsic::matrix_multiply>(
1033 Op0: m_Value(V&: A), Op1: m_Value(V&: B), Op2: m_ConstantInt(CI&: R),
1034 Op3: m_ConstantInt(CI&: K), Op4: m_ConstantInt(CI&: C))) &&
1035 match(V: A, P: m_Intrinsic<Intrinsic::matrix_transpose>(Op0: m_Value(V&: AT))) &&
1036 match(V: B, P: m_Intrinsic<Intrinsic::matrix_transpose>(Op0: m_Value(V&: (BT))))) {
1037 IRBuilder<> IB(&I);
1038 MatrixBuilder Builder(IB);
1039 Value *M = Builder.CreateMatrixMultiply(
1040 LHS: BT, RHS: AT, LHSRows: C->getZExtValue(), LHSColumns: K->getZExtValue(), RHSColumns: R->getZExtValue());
1041 setShapeInfo(V: M, Shape: {C, R});
1042 Instruction *NewInst = Builder.CreateMatrixTranspose(Matrix: M, Rows: C->getZExtValue(),
1043 Columns: R->getZExtValue());
1044 updateShapeAndReplaceAllUsesWith(Old&: I, New: NewInst);
1045 CleanupBinOp(I, A, B);
1046 return true;
1047 }
1048 // A^t + B ^t -> (A + B)^t. Pick rows and columns from first transpose. If
1049 // the shape of the second transpose is different, there's a shape conflict
1050 // which gets resolved by picking the shape of the first operand.
1051 else if (match(V: &I, P: m_FAdd(L: m_Value(V&: A), R: m_Value(V&: B))) &&
1052 match(V: A, P: m_Intrinsic<Intrinsic::matrix_transpose>(
1053 Op0: m_Value(V&: AT), Op1: m_ConstantInt(CI&: R), Op2: m_ConstantInt(CI&: C))) &&
1054 match(V: B, P: m_Intrinsic<Intrinsic::matrix_transpose>(
1055 Op0: m_Value(V&: BT), Op1: m_ConstantInt(), Op2: m_ConstantInt()))) {
1056 IRBuilder<> Builder(&I);
1057 auto *Add = Builder.CreateFAdd(L: AT, R: BT, Name: "mfadd");
1058 MatrixBuilder MBuilder(Builder);
1059 Instruction *NewInst = MBuilder.CreateMatrixTranspose(
1060 Matrix: Add, Rows: R->getZExtValue(), Columns: C->getZExtValue(), Name: "mfadd_t");
1061 updateShapeAndReplaceAllUsesWith(Old&: I, New: NewInst);
1062 assert(computeShapeInfoForInst(NewInst, ShapeMap) ==
1063 computeShapeInfoForInst(&I, ShapeMap) &&
1064 "Shape of new instruction doesn't match original shape.");
1065 CleanupBinOp(I, A, B);
1066 if (auto *AddI = dyn_cast<Instruction>(Val: Add)) {
1067 setShapeInfo(V: AddI, Shape: {R, C});
1068 assert(
1069 computeShapeInfoForInst(AddI, ShapeMap).value_or(ShapeMap[AddI]) ==
1070 ShapeMap[AddI] &&
1071 "Shape of updated addition doesn't match cached shape.");
1072 }
1073 return true;
1074 }
1075 return false;
1076 }
1077
1078 /// Try moving transposes in order to fold them away or into multiplies.
1079 bool optimizeTransposes() {
1080 bool Changed = false;
1081 // First sink all transposes inside matmuls and adds, hoping that we end up
1082 // with NN, NT or TN variants.
1083 for (BasicBlock &BB : reverse(C&: Func)) {
1084 for (auto II = BB.rbegin(); II != BB.rend();) {
1085 Instruction &I = *II;
1086 // We may remove II. By default continue on the next/prev instruction.
1087 ++II;
1088 if (Instruction *NewInst = sinkTranspose(I, II, Changed))
1089 II = std::next(x: BasicBlock::reverse_iterator(NewInst));
1090 }
1091 }
1092
1093 // If we have a TT matmul or a TT add, lift the transpose. We may be able
1094 // to fold into consuming multiply or add.
1095 for (BasicBlock &BB : Func) {
1096 for (Instruction &I : llvm::make_early_inc_range(Range&: BB)) {
1097 Changed |= liftTranspose(I);
1098 }
1099 }
1100 return Changed;
1101 }
1102
1103 bool Visit() {
1104 SmallVector<Instruction *, 32> WorkList;
1105
1106 // Initially only the shape of matrix intrinsics is known.
1107 // Initialize the work list with ops carrying shape information.
1108 for (BasicBlock &BB : Func)
1109 for (Instruction &Inst : BB) {
1110 IntrinsicInst *II = dyn_cast<IntrinsicInst>(Val: &Inst);
1111 if (!II)
1112 continue;
1113
1114 switch (II->getIntrinsicID()) {
1115 case Intrinsic::matrix_multiply:
1116 case Intrinsic::matrix_transpose:
1117 case Intrinsic::matrix_column_major_load:
1118 case Intrinsic::matrix_column_major_store:
1119 WorkList.push_back(Elt: &Inst);
1120 break;
1121 default:
1122 break;
1123 }
1124 }
1125
1126 // Avoid unnecessary work if there are no matrix intrinsics in the function.
1127 if (WorkList.empty())
1128 return false;
1129
1130 if (AM) {
1131 ORE = &AM->getResult<OptimizationRemarkEmitterAnalysis>(IR&: Func);
1132 AA = &AM->getResult<AAManager>(IR&: Func);
1133 DT = &AM->getResult<DominatorTreeAnalysis>(IR&: Func);
1134 LI = &AM->getResult<LoopAnalysis>(IR&: Func);
1135 }
1136
1137 // Propagate shapes until nothing changes any longer.
1138 while (!WorkList.empty()) {
1139 WorkList = propagateShapeForward(WorkList);
1140 WorkList = propagateShapeBackward(WorkList);
1141 }
1142
1143 bool Changed = false;
1144 if (!isMinimal()) {
1145 Changed |= optimizeTransposes();
1146 if (PrintAfterTransposeOpt) {
1147 dbgs() << "Dump after matrix transpose optimization:\n";
1148 Func.print(OS&: dbgs());
1149 }
1150 }
1151
1152 SmallVector<CallInst *, 16> MaybeFusableInsts;
1153 SmallVector<Instruction *, 16> MatrixInsts;
1154 SmallVector<IntrinsicInst *, 16> LifetimeEnds;
1155
1156 // First, collect all instructions with shape information and candidates for
1157 // fusion (currently only matrix multiplies).
1158 ReversePostOrderTraversal<Function *> RPOT(&Func);
1159 for (auto *BB : RPOT)
1160 for (Instruction &I : *BB) {
1161 if (match(V: &I, P: m_Intrinsic<Intrinsic::lifetime_end>()))
1162 LifetimeEnds.push_back(Elt: cast<IntrinsicInst>(Val: &I));
1163 if (!ShapeMap.contains(Val: &I))
1164 continue;
1165 if (match(V: &I, P: m_Intrinsic<Intrinsic::matrix_multiply>()))
1166 MaybeFusableInsts.push_back(Elt: cast<CallInst>(Val: &I));
1167 MatrixInsts.push_back(Elt: &I);
1168 }
1169
1170 // Second, try to lower any dot products
1171 SmallPtrSet<Instruction *, 16> FusedInsts;
1172 for (CallInst *CI : MaybeFusableInsts)
1173 lowerDotProduct(MatMul: CI, FusedInsts, FMF: getFastMathFlags(Inst: CI));
1174
1175 // Third, try to fuse candidates.
1176 for (CallInst *CI : MaybeFusableInsts)
1177 if (!FusedInsts.contains(Ptr: CI))
1178 LowerMatrixMultiplyFused(MatMul: CI, FusedInsts, LifetimeEnds);
1179
1180 Changed |= !FusedInsts.empty();
1181
1182 // Fourth, pre-process all the PHINode's. The incoming values will be
1183 // assigned later in VisitPHI.
1184 for (Instruction *Inst : MatrixInsts) {
1185 if (FusedInsts.count(Ptr: Inst))
1186 continue;
1187
1188 auto *PHI = dyn_cast<PHINode>(Val: Inst);
1189 if (!PHI)
1190 continue;
1191
1192 const ShapeInfo &SI = ShapeMap.at(Val: Inst);
1193 auto *EltTy = cast<FixedVectorType>(Val: PHI->getType())->getElementType();
1194 MatrixTy PhiM(SI.NumRows, SI.NumColumns, EltTy);
1195
1196 IRBuilder<> Builder(Inst);
1197 for (unsigned VI = 0, VE = PhiM.getNumVectors(); VI != VE; ++VI)
1198 PhiM.setVector(i: VI, V: Builder.CreatePHI(Ty: PhiM.getVectorTy(),
1199 NumReservedValues: PHI->getNumIncomingValues(),
1200 Name: PHI->getName()));
1201 assert(!Inst2ColumnMatrix.contains(PHI) && "map already contains phi?");
1202 Inst2ColumnMatrix[PHI] = PhiM;
1203 }
1204
1205 // Fifth, lower remaining instructions with shape information.
1206 for (Instruction *Inst : MatrixInsts) {
1207 if (FusedInsts.count(Ptr: Inst))
1208 continue;
1209
1210 const ShapeInfo &SI = ShapeMap.at(Val: Inst);
1211
1212 Value *Op1;
1213 Value *Op2;
1214 MatrixTy Result;
1215 IRBuilder<> Builder(Inst);
1216 if (auto *BinOp = dyn_cast<BinaryOperator>(Val: Inst))
1217 Result = VisitBinaryOperator(Inst: BinOp, SI, Builder);
1218 else if (auto *Cast = dyn_cast<CastInst>(Val: Inst))
1219 Result = VisitCastInstruction(Inst: Cast, Shape: SI, Builder);
1220 else if (auto *UnOp = dyn_cast<UnaryOperator>(Val: Inst))
1221 Result = VisitUnaryOperator(Inst: UnOp, SI, Builder);
1222 else if (auto *Intr = dyn_cast<IntrinsicInst>(Val: Inst))
1223 Result = VisitIntrinsicInst(Inst: Intr, SI, Builder);
1224 else if (auto *Select = dyn_cast<SelectInst>(Val: Inst))
1225 Result = VisitSelectInst(Inst: Select, Shape: SI, Builder);
1226 else if (match(V: Inst, P: m_Load(Op: m_Value(V&: Op1))))
1227 Result = VisitLoad(Inst: cast<LoadInst>(Val: Inst), SI, Ptr: Op1, Builder);
1228 else if (match(V: Inst, P: m_Store(ValueOp: m_Value(V&: Op1), PointerOp: m_Value(V&: Op2))))
1229 Result = VisitStore(Inst: cast<StoreInst>(Val: Inst), SI, StoredVal: Op1, Ptr: Op2, Builder);
1230 else if (auto *PHI = dyn_cast<PHINode>(Val: Inst))
1231 Result = VisitPHI(Inst: PHI, SI, Builder);
1232 else
1233 continue;
1234
1235 finalizeLowering(Inst, Matrix: Result, Builder);
1236 Changed = true;
1237 }
1238
1239 if (ORE) {
1240 RemarkGenerator RemarkGen(Inst2ColumnMatrix, *ORE, Func);
1241 RemarkGen.emitRemarks();
1242 }
1243
1244 // Delete the instructions backwards, as it has a reduced likelihood of
1245 // having to update as many def-use and use-def chains.
1246 //
1247 // Because we add to ToRemove during fusion we can't guarantee that defs
1248 // are before uses. Change uses to poison temporarily as these should get
1249 // removed as well.
1250 //
1251 // For verification, we keep track of where we changed uses to poison in
1252 // PoisonedInsts and then check that we in fact remove them.
1253 SmallPtrSet<Instruction *, 16> PoisonedInsts;
1254 for (auto *Inst : reverse(C&: ToRemove)) {
1255 for (Use &U : llvm::make_early_inc_range(Range: Inst->uses())) {
1256 if (auto *Poisoned = dyn_cast<Instruction>(Val: U.getUser()))
1257 PoisonedInsts.insert(Ptr: Poisoned);
1258 U.set(PoisonValue::get(T: Inst->getType()));
1259 }
1260 Inst->eraseFromParent();
1261 PoisonedInsts.erase(Ptr: Inst);
1262 }
1263 if (!PoisonedInsts.empty()) {
1264 // If we didn't remove all poisoned instructions, it's a hard error.
1265 dbgs() << "Poisoned but present instructions:\n";
1266 for (auto *I : PoisonedInsts)
1267 dbgs() << *I << "\n";
1268 llvm_unreachable("Poisoned but instruction not removed");
1269 }
1270
1271 return Changed;
1272 }
1273
1274 /// Replace intrinsic calls.
1275 MatrixTy VisitIntrinsicInst(IntrinsicInst *Inst, const ShapeInfo &SI,
1276 IRBuilder<> &Builder) {
1277 assert(Inst->getCalledFunction() &&
1278 Inst->getCalledFunction()->isIntrinsic());
1279
1280 switch (Inst->getCalledFunction()->getIntrinsicID()) {
1281 case Intrinsic::matrix_multiply:
1282 return LowerMultiply(MatMul: Inst, Builder);
1283 case Intrinsic::matrix_transpose:
1284 return LowerTranspose(Inst, Builder);
1285 case Intrinsic::matrix_column_major_load:
1286 return LowerColumnMajorLoad(Inst, Builder);
1287 case Intrinsic::matrix_column_major_store:
1288 return LowerColumnMajorStore(Inst, Builder);
1289 case Intrinsic::abs:
1290 case Intrinsic::fabs: {
1291 MatrixTy Result;
1292 MatrixTy M = getMatrix(MatrixVal: Inst->getOperand(i_nocapture: 0), SI, Builder);
1293 Builder.setFastMathFlags(getFastMathFlags(Inst));
1294
1295 for (auto *Vector : M.vectors()) {
1296 switch (Inst->getIntrinsicID()) {
1297 case Intrinsic::abs:
1298 Result.addVector(V: Builder.CreateBinaryIntrinsic(ID: Intrinsic::abs, LHS: Vector,
1299 RHS: Inst->getOperand(i_nocapture: 1)));
1300 continue;
1301 case Intrinsic::fabs:
1302 Result.addVector(
1303 V: Builder.CreateUnaryIntrinsic(ID: Inst->getIntrinsicID(), Op: Vector));
1304 continue;
1305 default:
1306 llvm_unreachable("unexpected intrinsic");
1307 }
1308 }
1309
1310 return Result.addNumComputeOps(N: getNumOps(VT: Result.getVectorTy()) *
1311 Result.getNumVectors());
1312 }
1313 default:
1314 break;
1315 }
1316 llvm_unreachable(
1317 "only intrinsics supporting shape info should be seen here");
1318 }
1319
1320 /// Compute the alignment for a column/row \p Idx with \p Stride between them.
1321 /// The address at \p Idx == 0 has alignment \p A. If \p Stride is a
1322 /// ConstantInt, reduce the initial alignment based on the byte offset. For
1323 /// non-ConstantInt strides, return the common alignment of the initial
1324 /// alignment and the element size in bytes.
1325 Align getAlignForIndex(unsigned Idx, Value *Stride, Type *ElementTy,
1326 MaybeAlign A) const {
1327 Align InitialAlign = DL.getValueOrABITypeAlignment(Alignment: A, Ty: ElementTy);
1328 if (Idx == 0)
1329 return InitialAlign;
1330
1331 TypeSize ElementSizeInBits = DL.getTypeSizeInBits(Ty: ElementTy);
1332 if (auto *ConstStride = dyn_cast<ConstantInt>(Val: Stride)) {
1333 uint64_t StrideInBytes =
1334 ConstStride->getZExtValue() * ElementSizeInBits / 8;
1335 return commonAlignment(A: InitialAlign, Offset: Idx * StrideInBytes);
1336 }
1337 return commonAlignment(A: InitialAlign, Offset: ElementSizeInBits / 8);
1338 }
1339
1340 IntegerType *getIndexType(Value *Ptr) const {
1341 return cast<IntegerType>(Val: DL.getIndexType(PtrTy: Ptr->getType()));
1342 }
1343
1344 Value *getIndex(Value *Ptr, uint64_t V) const {
1345 return ConstantInt::get(Ty: getIndexType(Ptr), V);
1346 }
1347
1348 Value *castToIndexType(Value *Ptr, Value *V, IRBuilder<> &Builder) const {
1349 assert(isa<IntegerType>(V->getType()) &&
1350 "Attempted to cast non-integral type to integer index");
1351 // In case the data layout's index type differs in width from the type of
1352 // the value we're given, truncate or zero extend to the appropriate width.
1353 // We zero extend here as indices are unsigned.
1354 return Builder.CreateZExtOrTrunc(V, DestTy: getIndexType(Ptr),
1355 Name: V->getName() + ".cast");
1356 }
1357
1358 /// Load a matrix with \p Shape starting at \p Ptr and using \p Stride between
1359 /// vectors.
1360 MatrixTy loadMatrix(Type *Ty, Value *Ptr, MaybeAlign MAlign, Value *Stride,
1361 bool IsVolatile, ShapeInfo Shape, IRBuilder<> &Builder) {
1362 auto *VType = cast<FixedVectorType>(Val: Ty);
1363 Type *EltTy = VType->getElementType();
1364 Type *VecTy = FixedVectorType::get(ElementType: EltTy, NumElts: Shape.getStride());
1365 Value *EltPtr = Ptr;
1366 MatrixTy Result;
1367 Stride = castToIndexType(Ptr, V: Stride, Builder);
1368 for (unsigned I = 0, E = Shape.getNumVectors(); I < E; ++I) {
1369 Value *GEP = computeVectorAddr(
1370 BasePtr: EltPtr, VecIdx: Builder.getIntN(N: Stride->getType()->getScalarSizeInBits(), C: I),
1371 Stride, NumElements: Shape.getStride(), EltType: EltTy, Builder);
1372 Value *Vector = Builder.CreateAlignedLoad(
1373 Ty: VecTy, Ptr: GEP, Align: getAlignForIndex(Idx: I, Stride, ElementTy: EltTy, A: MAlign),
1374 isVolatile: IsVolatile, Name: "col.load");
1375
1376 Result.addVector(V: Vector);
1377 }
1378 return Result.addNumLoads(N: getNumOps(VT: Result.getVectorTy()) *
1379 Result.getNumVectors());
1380 }
1381
1382 /// Loads a sub-matrix with shape \p ResultShape from a \p R x \p C matrix,
1383 /// starting at \p MatrixPtr[I][J].
1384 MatrixTy loadMatrix(Value *MatrixPtr, MaybeAlign Align, bool IsVolatile,
1385 ShapeInfo MatrixShape, Value *I, Value *J,
1386 ShapeInfo ResultShape, Type *EltTy,
1387 IRBuilder<> &Builder) {
1388 Value *Offset = Builder.CreateAdd(
1389 LHS: Builder.CreateMul(LHS: J, RHS: getIndex(Ptr: MatrixPtr, V: MatrixShape.getStride())), RHS: I);
1390
1391 Value *TileStart = Builder.CreateInBoundsGEP(Ty: EltTy, Ptr: MatrixPtr, IdxList: Offset);
1392 auto *TileTy = FixedVectorType::get(ElementType: EltTy, NumElts: ResultShape.NumRows *
1393 ResultShape.NumColumns);
1394
1395 return loadMatrix(Ty: TileTy, Ptr: TileStart, MAlign: Align,
1396 Stride: getIndex(Ptr: MatrixPtr, V: MatrixShape.getStride()), IsVolatile,
1397 Shape: ResultShape, Builder);
1398 }
1399
1400 /// Lower a load instruction with shape information.
1401 MatrixTy LowerLoad(Instruction *Inst, Value *Ptr, MaybeAlign Align,
1402 Value *Stride, bool IsVolatile, ShapeInfo Shape,
1403 IRBuilder<> &Builder) {
1404 return loadMatrix(Ty: Inst->getType(), Ptr, MAlign: Align, Stride, IsVolatile, Shape,
1405 Builder);
1406 }
1407
1408 /// Lowers llvm.matrix.column.major.load.
1409 ///
1410 /// The intrinsic loads a matrix from memory using a stride between columns.
1411 MatrixTy LowerColumnMajorLoad(CallInst *Inst, IRBuilder<> &Builder) {
1412 assert(MatrixLayout == MatrixLayoutTy::ColumnMajor &&
1413 "Intrinsic only supports column-major layout!");
1414 Value *Ptr = Inst->getArgOperand(i: 0);
1415 Value *Stride = Inst->getArgOperand(i: 1);
1416 return LowerLoad(Inst, Ptr, Align: Inst->getParamAlign(ArgNo: 0), Stride,
1417 IsVolatile: cast<ConstantInt>(Val: Inst->getArgOperand(i: 2))->isOne(),
1418 Shape: {Inst->getArgOperand(i: 3), Inst->getArgOperand(i: 4)}, Builder);
1419 }
1420
1421 /// Stores a sub-matrix \p StoreVal into the \p R x \p C matrix starting at \p
1422 /// MatrixPtr[I][J].
1423 void storeMatrix(const MatrixTy &StoreVal, Value *MatrixPtr,
1424 MaybeAlign MAlign, bool IsVolatile, ShapeInfo MatrixShape,
1425 Value *I, Value *J, Type *EltTy, IRBuilder<> &Builder) {
1426 Value *Offset = Builder.CreateAdd(
1427 LHS: Builder.CreateMul(LHS: J, RHS: getIndex(Ptr: MatrixPtr, V: MatrixShape.getStride())), RHS: I);
1428
1429 Value *TileStart = Builder.CreateInBoundsGEP(Ty: EltTy, Ptr: MatrixPtr, IdxList: Offset);
1430 auto *TileTy = FixedVectorType::get(ElementType: EltTy, NumElts: StoreVal.getNumRows() *
1431 StoreVal.getNumColumns());
1432
1433 storeMatrix(Ty: TileTy, StoreVal, Ptr: TileStart, MAlign,
1434 Stride: getIndex(Ptr: MatrixPtr, V: MatrixShape.getStride()), IsVolatile,
1435 Builder);
1436 }
1437
1438 /// Store matrix \p StoreVal starting at \p Ptr and using \p Stride between
1439 /// vectors.
1440 MatrixTy storeMatrix(Type *Ty, MatrixTy StoreVal, Value *Ptr,
1441 MaybeAlign MAlign, Value *Stride, bool IsVolatile,
1442 IRBuilder<> &Builder) {
1443 auto *VType = cast<FixedVectorType>(Val: Ty);
1444 Value *EltPtr = Ptr;
1445 Stride = castToIndexType(Ptr, V: Stride, Builder);
1446 for (auto Vec : enumerate(First: StoreVal.vectors())) {
1447 Value *GEP = computeVectorAddr(
1448 BasePtr: EltPtr,
1449 VecIdx: Builder.getIntN(N: Stride->getType()->getScalarSizeInBits(),
1450 C: Vec.index()),
1451 Stride, NumElements: StoreVal.getStride(), EltType: VType->getElementType(), Builder);
1452 Builder.CreateAlignedStore(Val: Vec.value(), Ptr: GEP,
1453 Align: getAlignForIndex(Idx: Vec.index(), Stride,
1454 ElementTy: VType->getElementType(),
1455 A: MAlign),
1456 isVolatile: IsVolatile);
1457 }
1458 return MatrixTy().addNumStores(N: getNumOps(VT: StoreVal.getVectorTy()) *
1459 StoreVal.getNumVectors());
1460 }
1461
1462 /// Lower a store instruction with shape information.
1463 MatrixTy LowerStore(Instruction *Inst, Value *Matrix, Value *Ptr,
1464 MaybeAlign A, Value *Stride, bool IsVolatile,
1465 ShapeInfo Shape, IRBuilder<> &Builder) {
1466 auto StoreVal = getMatrix(MatrixVal: Matrix, SI: Shape, Builder);
1467 return storeMatrix(Ty: Matrix->getType(), StoreVal, Ptr, MAlign: A, Stride, IsVolatile,
1468 Builder);
1469 }
1470
1471 /// Lowers llvm.matrix.column.major.store.
1472 ///
1473 /// The intrinsic store a matrix back memory using a stride between columns.
1474 MatrixTy LowerColumnMajorStore(CallInst *Inst, IRBuilder<> &Builder) {
1475 assert(MatrixLayout == MatrixLayoutTy::ColumnMajor &&
1476 "Intrinsic only supports column-major layout!");
1477 Value *Matrix = Inst->getArgOperand(i: 0);
1478 Value *Ptr = Inst->getArgOperand(i: 1);
1479 Value *Stride = Inst->getArgOperand(i: 2);
1480 return LowerStore(Inst, Matrix, Ptr, A: Inst->getParamAlign(ArgNo: 1), Stride,
1481 IsVolatile: cast<ConstantInt>(Val: Inst->getArgOperand(i: 3))->isOne(),
1482 Shape: {Inst->getArgOperand(i: 4), Inst->getArgOperand(i: 5)},
1483 Builder);
1484 }
1485
1486 // Set elements I..I+NumElts-1 to Block
1487 Value *insertVector(Value *Col, unsigned I, Value *Block,
1488 IRBuilder<> &Builder) {
1489
1490 // First, bring Block to the same size as Col
1491 unsigned BlockNumElts =
1492 cast<FixedVectorType>(Val: Block->getType())->getNumElements();
1493 unsigned NumElts = cast<FixedVectorType>(Val: Col->getType())->getNumElements();
1494 assert(NumElts >= BlockNumElts && "Too few elements for current block");
1495
1496 Block = Builder.CreateShuffleVector(
1497 V: Block, Mask: createSequentialMask(Start: 0, NumInts: BlockNumElts, NumUndefs: NumElts - BlockNumElts));
1498
1499 // If Col is 7 long and I is 2 and BlockNumElts is 2 the mask is: 0, 1, 7,
1500 // 8, 4, 5, 6
1501 SmallVector<int, 16> Mask;
1502 unsigned i;
1503 for (i = 0; i < I; i++)
1504 Mask.push_back(Elt: i);
1505
1506 unsigned VecNumElts =
1507 cast<FixedVectorType>(Val: Col->getType())->getNumElements();
1508 for (; i < I + BlockNumElts; i++)
1509 Mask.push_back(Elt: i - I + VecNumElts);
1510
1511 for (; i < VecNumElts; i++)
1512 Mask.push_back(Elt: i);
1513
1514 return Builder.CreateShuffleVector(V1: Col, V2: Block, Mask);
1515 }
1516
1517 Value *createMulAdd(Value *Sum, Value *A, Value *B, bool UseFPOp,
1518 IRBuilder<> &Builder, bool AllowContraction,
1519 unsigned &NumComputeOps) {
1520 NumComputeOps += getNumOps(VT: A->getType());
1521 if (!Sum)
1522 return UseFPOp ? Builder.CreateFMul(L: A, R: B) : Builder.CreateMul(LHS: A, RHS: B);
1523
1524 if (UseFPOp) {
1525 if (AllowContraction) {
1526 // Use fmuladd for floating point operations and let the backend decide
1527 // if that's profitable.
1528 return Builder.CreateIntrinsic(ID: Intrinsic::fmuladd, OverloadTypes: A->getType(),
1529 Args: {A, B, Sum});
1530 }
1531 NumComputeOps += getNumOps(VT: A->getType());
1532 Value *Mul = Builder.CreateFMul(L: A, R: B);
1533 return Builder.CreateFAdd(L: Sum, R: Mul);
1534 }
1535
1536 NumComputeOps += getNumOps(VT: A->getType());
1537 Value *Mul = Builder.CreateMul(LHS: A, RHS: B);
1538 return Builder.CreateAdd(LHS: Sum, RHS: Mul);
1539 }
1540
1541 /// Cache \p Matrix as result of \p Inst and update the uses of \p Inst. For
1542 /// users with shape information, there's nothing to do: they will use the
1543 /// cached value when they are lowered. For other users, \p Matrix is
1544 /// flattened and the uses are updated to use it. Also marks \p Inst for
1545 /// deletion.
1546 void finalizeLowering(Instruction *Inst, MatrixTy Matrix,
1547 IRBuilder<> &Builder) {
1548 auto inserted = Inst2ColumnMatrix.insert(KV: std::make_pair(x&: Inst, y&: Matrix));
1549 (void)inserted;
1550 assert((inserted.second || isa<PHINode>(Inst)) &&
1551 "multiple matrix lowering mapping");
1552
1553 ToRemove.push_back(Elt: Inst);
1554 Value *Flattened = nullptr;
1555 for (Use &U : llvm::make_early_inc_range(Range: Inst->uses())) {
1556 if (ShapeMap.contains(Val: U.getUser()))
1557 continue;
1558
1559 if (!Flattened) {
1560 Flattened = Matrix.embedInVector(Builder);
1561 LLVM_DEBUG(
1562 if (Instruction *User = dyn_cast<Instruction>(U.getUser())) dbgs()
1563 << "flattening a " << Matrix.shape() << " matrix:\n"
1564 << *Inst
1565 << "\nbecause we do not have a shape-aware lowering for its "
1566 "user:\n"
1567 << *User << '\n';);
1568 FlattenedMatrices++;
1569 }
1570 U.set(Flattened);
1571 }
1572 }
1573
1574 /// Special case for MatMul lowering. Prevents scalar loads of row-major
1575 /// vectors Lowers to vector reduction add instead of sequential add if
1576 /// reassocation is enabled.
1577 void lowerDotProduct(CallInst *MatMul,
1578 SmallPtrSet<Instruction *, 16> &FusedInsts,
1579 FastMathFlags FMF) {
1580 if (FusedInsts.contains(Ptr: MatMul) ||
1581 MatrixLayout != MatrixLayoutTy::ColumnMajor)
1582 return;
1583 ShapeInfo LShape(MatMul->getArgOperand(i: 2), MatMul->getArgOperand(i: 3));
1584 ShapeInfo RShape(MatMul->getArgOperand(i: 3), MatMul->getArgOperand(i: 4));
1585
1586 if (LShape.NumRows != 1 || RShape.NumColumns != 1) // not a dot product
1587 return;
1588
1589 Value *LHS = MatMul->getArgOperand(i: 0);
1590 Value *RHS = MatMul->getArgOperand(i: 1);
1591
1592 Type *ElementType = cast<FixedVectorType>(Val: LHS->getType())->getElementType();
1593 bool IsIntVec = ElementType->isIntegerTy();
1594
1595 // Floating point reductions require reassocation.
1596 if (!IsIntVec && !FMF.allowReassoc())
1597 return;
1598
1599 auto CanBeFlattened = [](Value *Op) {
1600 if (match(V: Op, P: m_BinOp()))
1601 return true;
1602 return match(
1603 V: Op, P: m_OneUse(SubPattern: m_CombineOr(
1604 Ps: m_Load(Op: m_Value()),
1605 Ps: m_CombineOr(Ps: m_Intrinsic<Intrinsic::matrix_transpose>(),
1606 Ps: m_Intrinsic<Intrinsic::matrix_column_major_load>(
1607 Op0: m_Value(), Op1: m_One())))));
1608 };
1609 // Returns the cost benefit of using \p Op with the dot product lowering. If
1610 // the returned cost is < 0, the argument is cheaper to use in the
1611 // dot-product lowering.
1612 auto GetCostForArg = [this, &CanBeFlattened](Value *Op, unsigned N) {
1613 if (!ShapeMap.contains(Val: Op))
1614 return InstructionCost::getInvalid();
1615
1616 if (!isa<Instruction>(Val: Op))
1617 return InstructionCost(0);
1618
1619 FixedVectorType *VecTy = cast<FixedVectorType>(Val: Op->getType());
1620 Type *EltTy = VecTy->getElementType();
1621
1622 if (!CanBeFlattened(Op)) {
1623 InstructionCost EmbedCost(0);
1624 // Roughly estimate the cost for embedding the columns into a vector.
1625 for (unsigned I = 1; I < N; ++I)
1626 EmbedCost += TTI.getShuffleCost(
1627 Kind: TTI::SK_Splice, DstTy: FixedVectorType::get(ElementType: EltTy, NumElts: 1),
1628 SrcTy: FixedVectorType::get(ElementType: EltTy, NumElts: 1), Mask: {}, CostKind: TTI::TCK_RecipThroughput);
1629 return EmbedCost;
1630 }
1631
1632 if (match(V: Op, P: m_BinOp()) && ShapeMap.contains(Val: Op)) {
1633 InstructionCost OriginalCost =
1634 TTI.getArithmeticInstrCost(Opcode: cast<Instruction>(Val: Op)->getOpcode(),
1635 Ty: EltTy) *
1636 N;
1637 InstructionCost NewCost = TTI.getArithmeticInstrCost(
1638 Opcode: cast<Instruction>(Val: Op)->getOpcode(), Ty: VecTy);
1639 return NewCost - OriginalCost;
1640 }
1641
1642 if (match(V: Op, P: m_Intrinsic<Intrinsic::matrix_transpose>())) {
1643 // The transpose can be skipped for the dot product lowering, roughly
1644 // estimate the savings as the cost of embedding the columns in a
1645 // vector.
1646 InstructionCost EmbedCost(0);
1647 for (unsigned I = 1; I < N; ++I)
1648 EmbedCost -= TTI.getShuffleCost(
1649 Kind: TTI::SK_Splice, DstTy: FixedVectorType::get(ElementType: EltTy, NumElts: 1),
1650 SrcTy: FixedVectorType::get(ElementType: EltTy, NumElts: 1), Mask: {}, CostKind: TTI::TCK_RecipThroughput);
1651 return EmbedCost;
1652 }
1653
1654 // Costs for loads.
1655 if (N == 1)
1656 return InstructionCost(0);
1657
1658 return TTI.getMemoryOpCost(Opcode: Instruction::Load, Src: VecTy, Alignment: Align(1), AddressSpace: 0) -
1659 N * TTI.getMemoryOpCost(Opcode: Instruction::Load, Src: EltTy, Alignment: Align(1), AddressSpace: 0);
1660 };
1661
1662 // Iterate over LHS and operations feeding LHS and check if it is profitable
1663 // to flatten the visited ops. For each op, we compute the difference
1664 // between the flattened and matrix versions.
1665 SmallPtrSet<Value *, 4> Seen;
1666 SmallVector<Value *> WorkList;
1667 SmallVector<Value *> ToFlatten;
1668 WorkList.push_back(Elt: LHS);
1669 InstructionCost LHSCost(0);
1670 while (!WorkList.empty()) {
1671 Value *Op = WorkList.pop_back_val();
1672 if (!Seen.insert(Ptr: Op).second)
1673 continue;
1674
1675 InstructionCost OpCost = GetCostForArg(Op, LShape.NumColumns);
1676 if (OpCost + LHSCost >= LHSCost)
1677 continue;
1678
1679 LHSCost += OpCost;
1680 ToFlatten.push_back(Elt: Op);
1681 if (auto *I = dyn_cast<Instruction>(Val: Op))
1682 WorkList.append(in_start: I->op_begin(), in_end: I->op_end());
1683 }
1684
1685 // We compare the costs of a vector.reduce.add to sequential add.
1686 int AddOpCode = IsIntVec ? Instruction::Add : Instruction::FAdd;
1687 int MulOpCode = IsIntVec ? Instruction::Mul : Instruction::FMul;
1688 InstructionCost ReductionCost =
1689 TTI.getArithmeticReductionCost(
1690 Opcode: AddOpCode, Ty: cast<FixedVectorType>(Val: LHS->getType()),
1691 FMF: IsIntVec ? std::nullopt : std::optional(FMF)) +
1692 TTI.getArithmeticInstrCost(Opcode: MulOpCode, Ty: LHS->getType());
1693 InstructionCost SequentialAddCost =
1694 TTI.getArithmeticInstrCost(Opcode: AddOpCode, Ty: ElementType) *
1695 (LShape.NumColumns - 1) +
1696 TTI.getArithmeticInstrCost(Opcode: MulOpCode, Ty: ElementType) *
1697 (LShape.NumColumns);
1698 if ((LHSCost + ReductionCost - SequentialAddCost) > InstructionCost(0))
1699 return;
1700
1701 FusedInsts.insert(Ptr: MatMul);
1702 IRBuilder<> Builder(MatMul);
1703 auto FlattenArg = [&Builder, &FusedInsts, &CanBeFlattened,
1704 this](Value *Op) {
1705 // Matmul must be the only user of loads because we don't use LowerLoad
1706 // for row vectors (LowerLoad results in scalar loads and shufflevectors
1707 // instead of single vector load).
1708 if (!CanBeFlattened(Op))
1709 return;
1710
1711 if (match(V: Op, P: m_BinOp())) {
1712 auto It = ShapeMap.find(Val: Op);
1713 if (It != ShapeMap.end()) {
1714 It->second = It->second.t();
1715 return;
1716 }
1717 }
1718
1719 FusedInsts.insert(Ptr: cast<Instruction>(Val: Op));
1720 // If vector uses the builtin load, lower to a LoadInst
1721 Value *Arg;
1722 if (match(V: Op, P: m_Intrinsic<Intrinsic::matrix_column_major_load>(
1723 Op0: m_Value(V&: Arg)))) {
1724 auto *NewLoad = Builder.CreateLoad(Ty: Op->getType(), Ptr: Arg);
1725 Op->replaceAllUsesWith(V: NewLoad);
1726 eraseFromParentAndRemoveFromShapeMap(Inst: cast<Instruction>(Val: Op));
1727 return;
1728 } else if (match(V: Op, P: m_Intrinsic<Intrinsic::matrix_transpose>(
1729 Op0: m_Value(V&: Arg)))) {
1730 ToRemove.push_back(Elt: cast<Instruction>(Val: Op));
1731 Op->replaceAllUsesWith(V: Arg);
1732 return;
1733 }
1734 };
1735
1736 for (auto *V : ToFlatten)
1737 FlattenArg(V);
1738
1739 LHS = MatMul->getArgOperand(i: 0);
1740
1741 // Insert mul/fmul and llvm.vector.reduce.fadd
1742 Value *Mul =
1743 IsIntVec ? Builder.CreateMul(LHS, RHS) : Builder.CreateFMul(L: LHS, R: RHS);
1744
1745 Value *Result;
1746 if (IsIntVec)
1747 Result = Builder.CreateAddReduce(Src: Mul);
1748 else {
1749 Result = Builder.CreateFAddReduce(
1750 Acc: ConstantFP::get(
1751 Ty: cast<FixedVectorType>(Val: LHS->getType())->getElementType(), V: 0.0),
1752 Src: Mul);
1753 cast<Instruction>(Val: Result)->setFastMathFlags(FMF);
1754 }
1755
1756 // pack scalar back into a matrix and then replace matmul inst
1757 Result = Builder.CreateInsertElement(Vec: PoisonValue::get(T: MatMul->getType()),
1758 NewElt: Result, Idx: uint64_t(0));
1759 MatMul->replaceAllUsesWith(V: Result);
1760 FusedInsts.insert(Ptr: MatMul);
1761 ToRemove.push_back(Elt: MatMul);
1762 }
1763
1764 /// Given \p Remainder iterations of the the matmul inner loop,
1765 /// potentially lower \p Blocksize that is used for the underlying
1766 /// vector.
1767 unsigned capBlockSize(unsigned BlockSize, unsigned Remainder, Type *EltType) {
1768 if (BlockSize <= Remainder)
1769 return BlockSize;
1770
1771 // If the remainder is also a legal type just use it.
1772 auto *VecTy = FixedVectorType::get(ElementType: EltType, NumElts: Remainder);
1773 if (TTI.isTypeLegal(Ty: VecTy))
1774 return Remainder;
1775
1776 // Similarly, if the vector is small enough that we don't want
1777 // to split further.
1778 if (VecTy->getPrimitiveSizeInBits() <= SplitMatmulRemainderOverThreshold)
1779 return Remainder;
1780
1781 // Gradually lower the vectorization factor to cover the
1782 // remainder.
1783 do {
1784 BlockSize /= 2;
1785 } while (BlockSize > Remainder);
1786 return BlockSize;
1787 }
1788
1789 /// Compute \p Result += \p A * \p B for input matrices with left-associating
1790 /// addition.
1791 ///
1792 /// We can fold a transpose into the operand that is used to extract scalars.
1793 /// This is the first operands with row-major and the second with
1794 /// column-major. If \p IsScalarMatrixTransposed we assume the appropriate
1795 /// operand is transposed.
1796 void emitMatrixMultiply(MatrixTy &Result, const MatrixTy &A,
1797 const MatrixTy &B, IRBuilder<> &Builder, bool IsTiled,
1798 bool IsScalarMatrixTransposed, FastMathFlags FMF) {
1799 const unsigned VF = std::max<unsigned>(
1800 a: TTI.getRegisterBitWidth(K: TargetTransformInfo::RGK_FixedWidthVector)
1801 .getFixedValue() /
1802 Result.getElementType()->getPrimitiveSizeInBits().getFixedValue(),
1803 b: 1U);
1804 unsigned R = Result.getNumRows();
1805 unsigned C = Result.getNumColumns();
1806 unsigned M = A.getNumColumns();
1807
1808 bool IsFP = Result.getElementType()->isFloatingPointTy();
1809 assert(A.isColumnMajor() == B.isColumnMajor() &&
1810 Result.isColumnMajor() == A.isColumnMajor() &&
1811 "operands must agree on matrix layout");
1812 unsigned NumComputeOps = 0;
1813
1814 Builder.setFastMathFlags(FMF);
1815
1816 if (A.isColumnMajor()) {
1817 // Multiply columns from the first operand with scalars from the second
1818 // operand. Then move along the K axes and accumulate the columns. With
1819 // this the adds can be vectorized without reassociation.
1820 for (unsigned J = 0; J < C; ++J) {
1821 unsigned BlockSize = VF;
1822 // If Result is zero, we don't need to accumulate in the K==0 iteration.
1823 bool isSumZero = isa<ConstantAggregateZero>(Val: Result.getColumn(i: J));
1824
1825 for (unsigned I = 0; I < R; I += BlockSize) {
1826 // Lower block size to make sure we stay within bounds.
1827 BlockSize = capBlockSize(BlockSize, Remainder: R - I, EltType: Result.getElementType());
1828 Value *Sum = IsTiled ? Result.extractVector(I, J, NumElts: BlockSize, Builder)
1829 : nullptr;
1830 for (unsigned K = 0; K < M; ++K) {
1831 Value *L = A.extractVector(I, J: K, NumElts: BlockSize, Builder);
1832 Value *RH = Builder.CreateExtractElement(
1833 Vec: B.getColumn(i: IsScalarMatrixTransposed ? K : J),
1834 Idx: IsScalarMatrixTransposed ? J : K);
1835 Value *Splat = Builder.CreateVectorSplat(NumElts: BlockSize, V: RH, Name: "splat");
1836 Sum =
1837 createMulAdd(Sum: isSumZero && K == 0 ? nullptr : Sum, A: L, B: Splat,
1838 UseFPOp: IsFP, Builder, AllowContraction: FMF.allowContract(), NumComputeOps);
1839 }
1840 Result.setVector(i: J,
1841 V: insertVector(Col: Result.getVector(i: J), I, Block: Sum, Builder));
1842 }
1843 }
1844 } else {
1845 // Multiply rows from the second operand with scalars from the first
1846 // operand. Then move along the K axes and accumulate the rows. With this
1847 // the adds can be vectorized without reassociation.
1848 for (unsigned I = 0; I < R; ++I) {
1849 unsigned BlockSize = VF;
1850 bool isSumZero = isa<ConstantAggregateZero>(Val: Result.getRow(i: I));
1851 for (unsigned J = 0; J < C; J += BlockSize) {
1852 // Lower the vectorization factor to cover the remainder.
1853 BlockSize = capBlockSize(BlockSize, Remainder: C - J, EltType: Result.getElementType());
1854
1855 Value *Sum = nullptr;
1856 for (unsigned K = 0; K < M; ++K) {
1857 Value *R = B.extractVector(I: K, J, NumElts: BlockSize, Builder);
1858 Value *LH = Builder.CreateExtractElement(
1859 Vec: A.getVector(i: IsScalarMatrixTransposed ? K : I),
1860 Idx: IsScalarMatrixTransposed ? I : K);
1861 Value *Splat = Builder.CreateVectorSplat(NumElts: BlockSize, V: LH, Name: "splat");
1862 Sum =
1863 createMulAdd(Sum: isSumZero && K == 0 ? nullptr : Sum, A: Splat, B: R,
1864 UseFPOp: IsFP, Builder, AllowContraction: FMF.allowContract(), NumComputeOps);
1865 }
1866 Result.setVector(i: I,
1867 V: insertVector(Col: Result.getVector(i: I), I: J, Block: Sum, Builder));
1868 }
1869 }
1870 }
1871 Result.addNumComputeOps(N: NumComputeOps);
1872 }
1873
1874 /// Ensure that the memory in \p Load does not alias \p Store by potentially
1875 /// copying it to a new location. This new or otherwise the original location
1876 /// is returned.
1877 std::pair<Value *, AllocaInst *>
1878 getNonAliasingPointer(LoadInst *Load, StoreInst *Store, CallInst *MatMul) {
1879 MemoryLocation StoreLoc = MemoryLocation::get(SI: Store);
1880 MemoryLocation LoadLoc = MemoryLocation::get(LI: Load);
1881
1882 // If we can statically determine noalias we're good.
1883 if (AA->isNoAlias(LocA: LoadLoc, LocB: StoreLoc))
1884 return {Load->getPointerOperand(), nullptr};
1885
1886 // If the pointers are in different address spaces, we cannot compare them
1887 // at runtime. Conservatively copy the load operand to a new buffer.
1888 IRBuilder<> AllocaBuilder(&Func.getEntryBlock().front());
1889 if (Load->getPointerAddressSpace() != Store->getPointerAddressSpace()) {
1890 auto *VT = cast<FixedVectorType>(Val: Load->getType());
1891 auto *ArrayTy =
1892 ArrayType::get(ElementType: VT->getElementType(), NumElements: VT->getNumElements());
1893 AllocaInst *Alloca =
1894 AllocaBuilder.CreateAlloca(Ty: ArrayTy, AddrSpace: Load->getPointerAddressSpace());
1895 IRBuilder<> Builder(MatMul);
1896 Builder.CreateLifetimeStart(Ptr: Alloca);
1897 Builder.CreateMemCpy(Dst: Alloca, DstAlign: Alloca->getAlign(),
1898 Src: Load->getPointerOperand(), SrcAlign: Load->getAlign(),
1899 Size: LoadLoc.Size.getValue());
1900 return {Alloca, Alloca};
1901 }
1902
1903 // Create code to check if the memory locations of the Load and Store
1904 // overlap and if they do, copy Load's operand to a new buffer.
1905
1906 // First, create new blocks for 2n part of the check and the copy.
1907 BasicBlock *Check0 = MatMul->getParent();
1908 // FIXME: Use lazy DTU and update SplitBlock to accept a DTU instead of a
1909 // DT. Manually collect dominator tree updates, to avoid unnecessary work,
1910 // as we adjust Check0 and Check1's branches.
1911 SmallVector<DominatorTree::UpdateType, 4> DTUpdates;
1912 for (BasicBlock *Succ : successors(BB: Check0))
1913 DTUpdates.push_back(Elt: {DT->Delete, Check0, Succ});
1914
1915 BasicBlock *Check1 =
1916 SplitBlock(Old: MatMul->getParent(), SplitPt: MatMul, DTU: (DomTreeUpdater *)nullptr, LI,
1917 MSSAU: nullptr, BBName: "alias_cont");
1918 BasicBlock *Copy =
1919 SplitBlock(Old: MatMul->getParent(), SplitPt: MatMul, DTU: (DomTreeUpdater *)nullptr, LI,
1920 MSSAU: nullptr, BBName: "copy");
1921 BasicBlock *Fusion =
1922 SplitBlock(Old: MatMul->getParent(), SplitPt: MatMul, DTU: (DomTreeUpdater *)nullptr, LI,
1923 MSSAU: nullptr, BBName: "no_alias");
1924
1925 // Check if the loaded memory location begins before the end of the store
1926 // location. If the condition holds, they might overlap, otherwise they are
1927 // guaranteed to not overlap.
1928 IRBuilder<> Builder(MatMul);
1929 Check0->getTerminator()->eraseFromParent();
1930 Builder.SetInsertPoint(Check0);
1931 Type *AddrTy = DL.getAddressType(PtrTy: Store->getPointerOperand()->getType());
1932 Value *StoreBegin = Store->getPointerOperand();
1933 Value *StoreEnd = Builder.CreatePtrAdd(
1934 Ptr: StoreBegin, Offset: ConstantInt::get(Ty: AddrTy, V: StoreLoc.Size.getValue()),
1935 Name: "store.end",
1936 NW: GEPNoWrapFlags::inBounds() | GEPNoWrapFlags::noUnsignedWrap());
1937 Value *LoadBegin = Load->getPointerOperand();
1938 CondBrInst *BR1 = Builder.CreateCondBr(
1939 Cond: Builder.CreateICmpULT(LHS: LoadBegin, RHS: StoreEnd), True: Check1, False: Fusion);
1940 setExplicitlyUnknownBranchWeightsIfProfiled(I&: *BR1, DEBUG_TYPE);
1941
1942 // Check if the store begins before the end of the load location. If the
1943 // condition holds, they alias, otherwise they are guaranteed to not
1944 // overlap.
1945 Check1->getTerminator()->eraseFromParent();
1946 Builder.SetInsertPoint(TheBB: Check1, IP: Check1->begin());
1947
1948 auto *VT = cast<FixedVectorType>(Val: Load->getType());
1949 // Use an array type for the alloca, to avoid potentially huge alignment
1950 // requirements for large vector types.
1951 auto *ArrayTy = ArrayType::get(ElementType: VT->getElementType(), NumElements: VT->getNumElements());
1952 AllocaInst *Alloca =
1953 AllocaBuilder.CreateAlloca(Ty: ArrayTy, AddrSpace: Load->getPointerAddressSpace());
1954 Builder.CreateLifetimeStart(Ptr: Alloca);
1955
1956 Value *LoadEnd = Builder.CreatePtrAdd(
1957 Ptr: LoadBegin, Offset: ConstantInt::get(Ty: AddrTy, V: LoadLoc.Size.getValue()),
1958 Name: "load.end",
1959 NW: GEPNoWrapFlags::inBounds() | GEPNoWrapFlags::noUnsignedWrap());
1960 CondBrInst *BR2 = Builder.CreateCondBr(
1961 Cond: Builder.CreateICmpULT(LHS: StoreBegin, RHS: LoadEnd), True: Copy, False: Fusion);
1962 setExplicitlyUnknownBranchWeightsIfProfiled(I&: *BR2, DEBUG_TYPE);
1963
1964 // Copy load operand to new alloca.
1965 Builder.SetInsertPoint(TheBB: Copy, IP: Copy->begin());
1966 Builder.CreateMemCpy(Dst: Alloca, DstAlign: Alloca->getAlign(), Src: Load->getPointerOperand(),
1967 SrcAlign: Load->getAlign(), Size: LoadLoc.Size.getValue());
1968 Builder.SetInsertPoint(TheBB: Fusion, IP: Fusion->begin());
1969 PHINode *PHI = Builder.CreatePHI(Ty: Load->getPointerOperandType(), NumReservedValues: 3);
1970 PHI->addIncoming(V: Load->getPointerOperand(), BB: Check0);
1971 PHI->addIncoming(V: Load->getPointerOperand(), BB: Check1);
1972 PHI->addIncoming(V: Alloca, BB: Copy);
1973
1974 // Adjust DT.
1975 DTUpdates.push_back(Elt: {DT->Insert, Check0, Check1});
1976 DTUpdates.push_back(Elt: {DT->Insert, Check0, Fusion});
1977 DTUpdates.push_back(Elt: {DT->Insert, Check1, Copy});
1978 DTUpdates.push_back(Elt: {DT->Insert, Check1, Fusion});
1979 DT->applyUpdates(Updates: DTUpdates);
1980 return {PHI, Alloca};
1981 }
1982
1983 bool isFusionProfitable(CallInst *MatMul) {
1984 if (ForceFusion)
1985 return true;
1986
1987 ShapeInfo LShape(MatMul->getArgOperand(i: 2), MatMul->getArgOperand(i: 3));
1988 ShapeInfo RShape(MatMul->getArgOperand(i: 3), MatMul->getArgOperand(i: 4));
1989
1990 const unsigned R = LShape.NumRows;
1991 const unsigned C = RShape.NumColumns;
1992 const unsigned M = LShape.NumColumns;
1993 auto *EltType = cast<FixedVectorType>(Val: MatMul->getType())->getElementType();
1994
1995 const unsigned VF = std::max<unsigned>(
1996 a: TTI.getRegisterBitWidth(K: TargetTransformInfo::RGK_FixedWidthVector)
1997 .getFixedValue() /
1998 EltType->getPrimitiveSizeInBits().getFixedValue(),
1999 b: 1U);
2000
2001 // Cost model for tiling
2002 //
2003 // For tiling to be beneficial, we need reuse either along the R or
2004 // the C axis. We vectorize along the R axis so that means at least
2005 // 3 elements.
2006 // TODO: Also consider cost of copying if operands alias.
2007 if (R <= VF && C == 1)
2008 return false;
2009 // Then we need enough elements to exceed the number of vector
2010 // registers we have. Note that this is an oversimplification since
2011 // fusing also takes some extra loads which may exceed the number of
2012 // reloads necessary.
2013 unsigned Op0Regs = (R + VF - 1) / VF * M;
2014 unsigned Op1Regs = (M + VF - 1) / VF * C;
2015 return Op0Regs + Op1Regs >
2016 TTI.getNumberOfRegisters(ClassID: TTI.getRegisterClassForType(Vector: true));
2017 }
2018
2019 MatrixTy getZeroMatrix(Type *EltType, unsigned R, unsigned C) {
2020 MatrixTy Res;
2021 auto *ColumType = FixedVectorType::get(ElementType: EltType, NumElts: R);
2022 for (unsigned I = 0; I < C; ++I)
2023 Res.addVector(V: ConstantAggregateZero::get(Ty: ColumType));
2024 return Res;
2025 }
2026
2027 void createTiledLoops(CallInst *MatMul, Value *LPtr, ShapeInfo LShape,
2028 Value *RPtr, ShapeInfo RShape, StoreInst *Store) {
2029 auto *EltType = cast<FixedVectorType>(Val: MatMul->getType())->getElementType();
2030
2031 // Create the main tiling loop nest.
2032 TileInfo TI(LShape.NumRows, RShape.NumColumns, LShape.NumColumns, TileSize);
2033 DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Lazy);
2034 Instruction *InsertI = cast<Instruction>(Val: MatMul);
2035 BasicBlock *Start = InsertI->getParent();
2036 BasicBlock *End =
2037 SplitBlock(Old: InsertI->getParent(), SplitPt: InsertI, DT, LI, MSSAU: nullptr, BBName: "continue");
2038 IRBuilder<> Builder(MatMul);
2039 BasicBlock *InnerBody = TI.CreateTiledLoops(Start, End, B&: Builder, DTU, LI&: *LI);
2040
2041 Type *TileVecTy =
2042 FixedVectorType::get(ElementType: MatMul->getType()->getScalarType(), NumElts: TileSize);
2043 MatrixTy TileResult;
2044 // Insert in the inner loop header.
2045 Builder.SetInsertPoint(TI.KLoop.Header->getTerminator());
2046 // Create PHI nodes for the result columns to accumulate across iterations.
2047 SmallVector<PHINode *, 4> ColumnPhis;
2048 for (unsigned I = 0; I < TileSize; I++) {
2049 auto *Phi = Builder.CreatePHI(Ty: TileVecTy, NumReservedValues: 2, Name: "result.vec." + Twine(I));
2050 Phi->addIncoming(V: ConstantAggregateZero::get(Ty: TileVecTy),
2051 BB: TI.RowLoop.Header->getSingleSuccessor());
2052 TileResult.addVector(V: Phi);
2053 ColumnPhis.push_back(Elt: Phi);
2054 }
2055
2056 // Insert in the inner loop body, which computes
2057 // Res += Load(CurrentRow, K) * Load(K, CurrentColumn)
2058 Builder.SetInsertPoint(InnerBody->getTerminator());
2059 // Load tiles of the operands.
2060 MatrixTy A =
2061 loadMatrix(MatrixPtr: LPtr, Align: {}, IsVolatile: false, MatrixShape: LShape, I: TI.RowLoop.Index, J: TI.KLoop.Index,
2062 ResultShape: {TileSize, TileSize}, EltTy: EltType, Builder);
2063 MatrixTy B =
2064 loadMatrix(MatrixPtr: RPtr, Align: {}, IsVolatile: false, MatrixShape: RShape, I: TI.KLoop.Index, J: TI.ColumnLoop.Index,
2065 ResultShape: {TileSize, TileSize}, EltTy: EltType, Builder);
2066 emitMatrixMultiply(Result&: TileResult, A, B, Builder, IsTiled: true, IsScalarMatrixTransposed: false,
2067 FMF: getFastMathFlags(Inst: MatMul));
2068 // Store result after the inner loop is done.
2069 Builder.SetInsertPoint(TI.RowLoop.Latch->getTerminator());
2070 storeMatrix(StoreVal: TileResult, MatrixPtr: Store->getPointerOperand(), MAlign: Store->getAlign(),
2071 IsVolatile: Store->isVolatile(), MatrixShape: {LShape.NumRows, RShape.NumColumns},
2072 I: TI.RowLoop.Index, J: TI.ColumnLoop.Index, EltTy: EltType, Builder);
2073
2074 for (unsigned I = 0; I < TileResult.getNumVectors(); I++)
2075 ColumnPhis[I]->addIncoming(V: TileResult.getVector(i: I), BB: TI.KLoop.Latch);
2076
2077 // Force unrolling of a few iterations of the inner loop, to make sure there
2078 // is enough work per iteration.
2079 // FIXME: The unroller should make this decision directly instead, but
2080 // currently the cost-model is not up to the task.
2081 unsigned InnerLoopUnrollCount = std::min(a: 10u, b: LShape.NumColumns / TileSize);
2082 addStringMetadataToLoop(TheLoop: LI->getLoopFor(BB: TI.KLoop.Header),
2083 MDString: "llvm.loop.unroll.count", V: InnerLoopUnrollCount);
2084 }
2085
2086 void emitSIMDTiling(CallInst *MatMul, LoadInst *LoadOp0, LoadInst *LoadOp1,
2087 StoreInst *Store,
2088 SmallPtrSetImpl<Instruction *> &FusedInsts) {
2089 assert(MatrixLayout == MatrixLayoutTy::ColumnMajor &&
2090 "Tiling only supported for column-major matrixes at the moment!");
2091 if (!isFusionProfitable(MatMul))
2092 return;
2093
2094 ShapeInfo LShape(MatMul->getArgOperand(i: 2), MatMul->getArgOperand(i: 3));
2095 ShapeInfo RShape(MatMul->getArgOperand(i: 3), MatMul->getArgOperand(i: 4));
2096
2097 const unsigned R = LShape.NumRows;
2098 const unsigned C = RShape.NumColumns;
2099 const unsigned M = LShape.NumColumns;
2100 auto *EltType = cast<FixedVectorType>(Val: MatMul->getType())->getElementType();
2101
2102 auto [APtr, AAlloca] = getNonAliasingPointer(Load: LoadOp0, Store, MatMul);
2103 auto [BPtr, BAlloca] = getNonAliasingPointer(Load: LoadOp1, Store, MatMul);
2104 Value *CPtr = Store->getPointerOperand();
2105
2106 // Use loop-based tiling when the number of expected operations exceeds
2107 // threshold.
2108 unsigned NumOps = getNumNativeVectorOps(EltType, R, M, C);
2109 bool UseLoops =
2110 (NumOps > TileLoopsThreshold) && R % TileSize == 0 && C % TileSize == 0;
2111 if (UseLoops)
2112 createTiledLoops(MatMul, LPtr: APtr, LShape, RPtr: BPtr, RShape, Store);
2113 else {
2114 IRBuilder<> Builder(Store);
2115 for (unsigned J = 0; J < C; J += TileSize)
2116 for (unsigned I = 0; I < R; I += TileSize) {
2117 const unsigned TileR = std::min(a: R - I, b: unsigned(TileSize));
2118 const unsigned TileC = std::min(a: C - J, b: unsigned(TileSize));
2119 MatrixTy Res = getZeroMatrix(EltType, R: TileR, C: TileC);
2120
2121 for (unsigned K = 0; K < M; K += TileSize) {
2122 const unsigned TileM = std::min(a: M - K, b: unsigned(TileSize));
2123 MatrixTy A =
2124 loadMatrix(MatrixPtr: APtr, Align: LoadOp0->getAlign(), IsVolatile: LoadOp0->isVolatile(),
2125 MatrixShape: LShape, I: getIndex(Ptr: APtr, V: I), J: getIndex(Ptr: APtr, V: K),
2126 ResultShape: {TileR, TileM}, EltTy: EltType, Builder);
2127 MatrixTy B =
2128 loadMatrix(MatrixPtr: BPtr, Align: LoadOp1->getAlign(), IsVolatile: LoadOp1->isVolatile(),
2129 MatrixShape: RShape, I: getIndex(Ptr: BPtr, V: K), J: getIndex(Ptr: BPtr, V: J),
2130 ResultShape: {TileM, TileC}, EltTy: EltType, Builder);
2131 emitMatrixMultiply(Result&: Res, A, B, Builder, IsTiled: true, IsScalarMatrixTransposed: false,
2132 FMF: getFastMathFlags(Inst: MatMul));
2133 }
2134 storeMatrix(StoreVal: Res, MatrixPtr: CPtr, MAlign: Store->getAlign(), IsVolatile: Store->isVolatile(), MatrixShape: {R, M},
2135 I: getIndex(Ptr: CPtr, V: I), J: getIndex(Ptr: CPtr, V: J), EltTy: EltType, Builder);
2136 }
2137 }
2138
2139 // End the lifetime of the allocas used for alias-safe copies.
2140 {
2141 IRBuilder<> Builder(Store);
2142 if (AAlloca)
2143 Builder.CreateLifetimeEnd(Ptr: AAlloca);
2144 if (BAlloca)
2145 Builder.CreateLifetimeEnd(Ptr: BAlloca);
2146 }
2147
2148 // Mark eliminated instructions as fused and remove them.
2149 FusedInsts.insert(Ptr: Store);
2150 FusedInsts.insert(Ptr: MatMul);
2151 eraseFromParentAndRemoveFromShapeMap(Inst: Store);
2152 eraseFromParentAndRemoveFromShapeMap(Inst: MatMul);
2153 if (LoadOp0->use_empty()) {
2154 FusedInsts.insert(Ptr: LoadOp0);
2155 eraseFromParentAndRemoveFromShapeMap(Inst: LoadOp0);
2156 }
2157 if (LoadOp1 != LoadOp0 && LoadOp1->use_empty()) {
2158 FusedInsts.insert(Ptr: LoadOp1);
2159 eraseFromParentAndRemoveFromShapeMap(Inst: LoadOp1);
2160 }
2161 }
2162
2163 /// Try to lower matrix multiply chains by fusing operations.
2164 ///
2165 /// Call finalizeLowering on lowered instructions. Instructions that are
2166 /// completely eliminated by fusion are added to \p FusedInsts.
2167 void
2168 LowerMatrixMultiplyFused(CallInst *MatMul,
2169 SmallPtrSetImpl<Instruction *> &FusedInsts,
2170 SmallVector<IntrinsicInst *, 16> &LifetimeEnds) {
2171 if (!FuseMatrix || !DT || TileSize == 0)
2172 return;
2173
2174 assert(AA && LI && "Analyses should be available");
2175
2176 Value *A = MatMul->getArgOperand(i: 0);
2177 Value *B = MatMul->getArgOperand(i: 1);
2178
2179 // We can fold the transpose into the operand that is used to fetch scalars.
2180 Value *T;
2181 if (MatrixLayout == MatrixLayoutTy::ColumnMajor
2182 ? match(V: B, P: m_Intrinsic<Intrinsic::matrix_transpose>(Op0: m_Value(V&: T)))
2183 : match(V: A, P: m_Intrinsic<Intrinsic::matrix_transpose>(Op0: m_Value(V&: T)))) {
2184 IRBuilder<> Builder(MatMul);
2185 auto *EltType =
2186 cast<FixedVectorType>(Val: MatMul->getType())->getElementType();
2187 ShapeInfo LShape(MatMul->getArgOperand(i: 2), MatMul->getArgOperand(i: 3));
2188 ShapeInfo RShape(MatMul->getArgOperand(i: 3), MatMul->getArgOperand(i: 4));
2189 const unsigned R = LShape.NumRows;
2190 const unsigned M = LShape.NumColumns;
2191 const unsigned C = RShape.NumColumns;
2192
2193 MatrixTy MA;
2194 MatrixTy MB;
2195
2196 Value *Transpose;
2197 if (MatrixLayout == MatrixLayoutTy::ColumnMajor) {
2198 MA = getMatrix(MatrixVal: A, SI: ShapeInfo(R, M), Builder);
2199 MB = getMatrix(MatrixVal: T, SI: ShapeInfo(C, M), Builder);
2200 Transpose = B;
2201 } else {
2202 MA = getMatrix(MatrixVal: T, SI: ShapeInfo(R, M), Builder);
2203 MB = getMatrix(MatrixVal: B, SI: ShapeInfo(C, M), Builder);
2204 Transpose = A;
2205 }
2206
2207 // Initialize the output
2208 MatrixTy Result(R, C, EltType);
2209
2210 emitMatrixMultiply(Result, A: MA, B: MB, Builder, IsTiled: false, IsScalarMatrixTransposed: true,
2211 FMF: getFastMathFlags(Inst: MatMul));
2212
2213 FusedInsts.insert(Ptr: MatMul);
2214 if (Transpose->hasOneUse()) {
2215 FusedInsts.insert(Ptr: cast<Instruction>(Val: Transpose));
2216 ToRemove.push_back(Elt: cast<Instruction>(Val: Transpose));
2217 // TODO: add a fake entry for the folded instruction so that this is
2218 // included in the expression in the remark.
2219 Inst2ColumnMatrix[Transpose] = MatrixTy(M, C, EltType);
2220 }
2221 finalizeLowering(Inst: MatMul, Matrix: Result, Builder);
2222 return;
2223 }
2224
2225 if (!MatMul->hasOneUse() || MatrixLayout != MatrixLayoutTy::ColumnMajor)
2226 return;
2227
2228 // Lower {ld, ld} -> matmul -> st chains. No need to call finalizeLowering
2229 // since the single store user will be lowered as part of this.
2230 auto *LoadOp0 = dyn_cast<LoadInst>(Val: A);
2231 auto *LoadOp1 = dyn_cast<LoadInst>(Val: B);
2232 auto *Store = dyn_cast<StoreInst>(Val: *MatMul->user_begin());
2233 if (LoadOp0 && LoadOp1 && Store) {
2234 // The store address must dominate the MatMul instruction, otherwise
2235 // we create invalid IR.
2236 SetVector<Value *> WorkList;
2237 WorkList.insert(X: Store->getOperand(i_nocapture: 1));
2238 SmallVector<Instruction *> ToHoist;
2239 for (unsigned I = 0; I != WorkList.size(); ++I) {
2240 Value *Current = WorkList[I];
2241 auto *CurrI = dyn_cast<Instruction>(Val: Current);
2242 if (!CurrI)
2243 continue;
2244 if (isa<PHINode>(Val: CurrI))
2245 return;
2246 if (DT->dominates(Def: CurrI, User: MatMul))
2247 continue;
2248 if (CurrI->mayHaveSideEffects() || CurrI->mayReadFromMemory())
2249 return;
2250 ToHoist.push_back(Elt: CurrI);
2251 WorkList.insert_range(R: CurrI->operands());
2252 }
2253
2254 sort(C&: ToHoist, Comp: [this](Instruction *A, Instruction *B) {
2255 return DT->dominates(Def: A, User: B);
2256 });
2257 for (Instruction *I : ToHoist)
2258 I->moveBefore(InsertPos: MatMul->getIterator());
2259
2260 // Deal with lifetime.end calls that might be between Load0/Load1 and the
2261 // store. To avoid introducing loads to dead objects (i.e. after the
2262 // lifetime has been termined by @llvm.lifetime.end), either sink them
2263 // after the store if in the same block, or remove the lifetime.end marker
2264 // otherwise. This might pessimize further optimizations, by extending the
2265 // lifetime of the object until the function returns, but should be
2266 // conservatively correct.
2267 MemoryLocation Load0Loc = MemoryLocation::get(LI: LoadOp0);
2268 MemoryLocation Load1Loc = MemoryLocation::get(LI: LoadOp1);
2269 BasicBlock *StoreParent = Store->getParent();
2270 bool FusableOpsInSameBlock = LoadOp0->getParent() == StoreParent &&
2271 LoadOp1->getParent() == StoreParent;
2272 for (unsigned Idx = 0; Idx != LifetimeEnds.size();) {
2273 IntrinsicInst *End = LifetimeEnds[Idx];
2274 llvm::scope_exit Inc([&Idx]() { Idx++; });
2275 // If the lifetime.end is guaranteed to be before the loads or after the
2276 // store, it won't interfere with fusion.
2277 if (DT->dominates(Def: End, User: LoadOp0) && DT->dominates(Def: End, User: LoadOp1))
2278 continue;
2279 if (DT->dominates(Def: Store, User: End))
2280 continue;
2281 // If all fusable ops are in the same block and the lifetime.end is in a
2282 // different block, it won't interfere with fusion.
2283 if (FusableOpsInSameBlock && End->getParent() != StoreParent)
2284 continue;
2285
2286 // If the loads don't alias the lifetime.end, it won't interfere with
2287 // fusion.
2288 MemoryLocation EndLoc = MemoryLocation::getForArgument(Call: End, ArgIdx: 0, TLI: nullptr);
2289 if (!EndLoc.Ptr)
2290 continue;
2291 if (AA->isNoAlias(LocA: Load0Loc, LocB: EndLoc) && AA->isNoAlias(LocA: Load1Loc, LocB: EndLoc))
2292 continue;
2293
2294 // If both lifetime.end and the store are in the same block, extend the
2295 // lifetime until after the store, so the new lifetime covers the loads
2296 // we introduce later.
2297 if (End->getParent() == StoreParent) {
2298 End->moveAfter(MovePos: Store);
2299 continue;
2300 }
2301
2302 // Otherwise remove the conflicting lifetime.end marker.
2303 ToRemove.push_back(Elt: End);
2304 std::swap(a&: LifetimeEnds[Idx], b&: LifetimeEnds.back());
2305 LifetimeEnds.pop_back();
2306 Inc.release();
2307 }
2308
2309 emitSIMDTiling(MatMul, LoadOp0, LoadOp1, Store, FusedInsts);
2310 return;
2311 }
2312 }
2313
2314 /// Lowers llvm.matrix.multiply.
2315 MatrixTy LowerMultiply(CallInst *MatMul, IRBuilder<> &Builder) {
2316 auto *EltType = cast<FixedVectorType>(Val: MatMul->getType())->getElementType();
2317 ShapeInfo LShape(MatMul->getArgOperand(i: 2), MatMul->getArgOperand(i: 3));
2318 ShapeInfo RShape(MatMul->getArgOperand(i: 3), MatMul->getArgOperand(i: 4));
2319
2320 const MatrixTy &Lhs = getMatrix(MatrixVal: MatMul->getArgOperand(i: 0), SI: LShape, Builder);
2321 const MatrixTy &Rhs = getMatrix(MatrixVal: MatMul->getArgOperand(i: 1), SI: RShape, Builder);
2322 assert(Lhs.getElementType() == Rhs.getElementType() &&
2323 "Matrix multiply argument element types do not match.");
2324
2325 const unsigned R = LShape.NumRows;
2326 const unsigned C = RShape.NumColumns;
2327 assert(LShape.NumColumns == RShape.NumRows);
2328
2329 // Initialize the output
2330 MatrixTy Result(R, C, EltType);
2331 assert(Lhs.getElementType() == Result.getElementType() &&
2332 "Matrix multiply result element type does not match arguments.");
2333
2334 emitMatrixMultiply(Result, A: Lhs, B: Rhs, Builder, IsTiled: false, IsScalarMatrixTransposed: false,
2335 FMF: getFastMathFlags(Inst: MatMul));
2336 return Result;
2337 }
2338
2339 /// Lowers llvm.matrix.transpose.
2340 MatrixTy LowerTranspose(CallInst *Inst, IRBuilder<> &Builder) {
2341 MatrixTy Result;
2342 Value *InputVal = Inst->getArgOperand(i: 0);
2343 FixedVectorType *VectorTy = cast<FixedVectorType>(Val: InputVal->getType());
2344 ShapeInfo ArgShape(Inst->getArgOperand(i: 1), Inst->getArgOperand(i: 2));
2345 MatrixTy InputMatrix = getMatrix(MatrixVal: InputVal, SI: ArgShape, Builder);
2346
2347 const unsigned NewNumVecs =
2348 InputMatrix.isColumnMajor() ? ArgShape.NumRows : ArgShape.NumColumns;
2349 const unsigned NewNumElts =
2350 InputMatrix.isColumnMajor() ? ArgShape.NumColumns : ArgShape.NumRows;
2351
2352 for (unsigned I = 0; I < NewNumVecs; ++I) {
2353 // Build a single result vector. First initialize it.
2354 Value *ResultVector = PoisonValue::get(
2355 T: FixedVectorType::get(ElementType: VectorTy->getElementType(), NumElts: NewNumElts));
2356 // Go through the old elements and insert it into the resulting vector.
2357 for (auto J : enumerate(First: InputMatrix.vectors())) {
2358 Value *Elt = Builder.CreateExtractElement(Vec: J.value(), Idx: I);
2359 // Row and column indices are transposed.
2360 ResultVector =
2361 Builder.CreateInsertElement(Vec: ResultVector, NewElt: Elt, Idx: J.index());
2362 }
2363 Result.addVector(V: ResultVector);
2364 }
2365
2366 // TODO: Improve estimate of operations needed for transposes. Currently we
2367 // just count the insertelement/extractelement instructions, but do not
2368 // account for later simplifications/combines.
2369 return Result.addNumComputeOps(N: 2 * ArgShape.NumRows * ArgShape.NumColumns)
2370 .addNumExposedTransposes(N: 1);
2371 }
2372
2373 /// Lower load instructions.
2374 MatrixTy VisitLoad(LoadInst *Inst, const ShapeInfo &SI, Value *Ptr,
2375 IRBuilder<> &Builder) {
2376 return LowerLoad(Inst, Ptr, Align: Inst->getAlign(), Stride: getIndex(Ptr, V: SI.getStride()),
2377 IsVolatile: Inst->isVolatile(), Shape: SI, Builder);
2378 }
2379
2380 MatrixTy VisitStore(StoreInst *Inst, const ShapeInfo &SI, Value *StoredVal,
2381 Value *Ptr, IRBuilder<> &Builder) {
2382 return LowerStore(Inst, Matrix: StoredVal, Ptr, A: Inst->getAlign(),
2383 Stride: getIndex(Ptr, V: SI.getStride()), IsVolatile: Inst->isVolatile(), Shape: SI,
2384 Builder);
2385 }
2386
2387 MatrixTy VisitPHI(PHINode *Inst, const ShapeInfo &SI, IRBuilder<> &Builder) {
2388 auto BlockIP = Inst->getParent()->getFirstInsertionPt();
2389 Builder.SetInsertPoint(BlockIP);
2390 MatrixTy PhiM = getMatrix(MatrixVal: Inst, SI, Builder);
2391
2392 for (auto [IncomingV, IncomingB] :
2393 llvm::zip_equal(t: Inst->incoming_values(), u: Inst->blocks())) {
2394 // getMatrix() may insert some instructions to help with reshaping. The
2395 // safest place for those is at the top of the block after the rest of the
2396 // PHI's. Even better, if we can put it in the incoming block.
2397 Builder.SetInsertPoint(BlockIP);
2398 if (auto *IncomingInst = dyn_cast<Instruction>(Val&: IncomingV))
2399 if (auto MaybeIP = IncomingInst->getInsertionPointAfterDef())
2400 Builder.SetInsertPoint(*MaybeIP);
2401
2402 MatrixTy OpM = getMatrix(MatrixVal: IncomingV, SI, Builder);
2403
2404 for (unsigned VI = 0, VE = PhiM.getNumVectors(); VI != VE; ++VI) {
2405 PHINode *NewPHI = cast<PHINode>(Val: PhiM.getVector(i: VI));
2406 NewPHI->addIncoming(V: OpM.getVector(i: VI), BB: IncomingB);
2407 }
2408 }
2409
2410 // finalizeLowering() may also insert instructions in some cases. The safe
2411 // place for those is at the end of the initial block of PHIs.
2412 Builder.SetInsertPoint(BlockIP);
2413 return PhiM;
2414 }
2415
2416 /// Lower binary operators.
2417 MatrixTy VisitBinaryOperator(BinaryOperator *Inst, const ShapeInfo &SI,
2418 IRBuilder<> &Builder) {
2419 Value *Lhs = Inst->getOperand(i_nocapture: 0);
2420 Value *Rhs = Inst->getOperand(i_nocapture: 1);
2421
2422 MatrixTy Result;
2423 MatrixTy A = getMatrix(MatrixVal: Lhs, SI, Builder);
2424 MatrixTy B = getMatrix(MatrixVal: Rhs, SI, Builder);
2425 assert(A.isColumnMajor() == B.isColumnMajor() &&
2426 Result.isColumnMajor() == A.isColumnMajor() &&
2427 "operands must agree on matrix layout");
2428
2429 Builder.setFastMathFlags(getFastMathFlags(Inst));
2430
2431 for (auto [AV, BV] : llvm::zip_equal(t: A.vectors(), u: B.vectors()))
2432 Result.addVector(V: Builder.CreateBinOp(Opc: Inst->getOpcode(), LHS: AV, RHS: BV));
2433
2434 return Result.addNumComputeOps(N: getNumOps(VT: Result.getVectorTy()) *
2435 Result.getNumVectors());
2436 }
2437
2438 /// Lower unary operators.
2439 MatrixTy VisitUnaryOperator(UnaryOperator *Inst, const ShapeInfo &SI,
2440 IRBuilder<> &Builder) {
2441 Value *Op = Inst->getOperand(i_nocapture: 0);
2442
2443 MatrixTy Result;
2444 MatrixTy M = getMatrix(MatrixVal: Op, SI, Builder);
2445
2446 Builder.setFastMathFlags(getFastMathFlags(Inst));
2447
2448 // Helper to perform unary op on vectors.
2449 auto BuildVectorOp = [&Builder, Inst](Value *Op) {
2450 switch (Inst->getOpcode()) {
2451 case Instruction::FNeg:
2452 return Builder.CreateFNeg(V: Op);
2453 default:
2454 llvm_unreachable("Unsupported unary operator for matrix");
2455 }
2456 };
2457
2458 for (auto *Vector : M.vectors())
2459 Result.addVector(V: BuildVectorOp(Vector));
2460
2461 return Result.addNumComputeOps(N: getNumOps(VT: Result.getVectorTy()) *
2462 Result.getNumVectors());
2463 }
2464
2465 /// Lower cast instructions.
2466 MatrixTy VisitCastInstruction(CastInst *Inst, const ShapeInfo &Shape,
2467 IRBuilder<> &Builder) {
2468 Value *Op = Inst->getOperand(i_nocapture: 0);
2469
2470 MatrixTy Result;
2471 MatrixTy M = getMatrix(MatrixVal: Op, SI: Shape, Builder);
2472
2473 Builder.setFastMathFlags(getFastMathFlags(Inst));
2474
2475 auto *OrigVTy = cast<VectorType>(Val: Inst->getType());
2476 auto *NewVTy = VectorType::get(ElementType: OrigVTy->getElementType(),
2477 EC: ElementCount::getFixed(MinVal: M.getStride()));
2478
2479 for (auto *Vector : M.vectors())
2480 Result.addVector(V: Builder.CreateCast(Op: Inst->getOpcode(), V: Vector, DestTy: NewVTy));
2481
2482 return Result.addNumComputeOps(N: getNumOps(VT: Result.getVectorTy()) *
2483 Result.getNumVectors());
2484 }
2485
2486 /// Lower selects.
2487 MatrixTy VisitSelectInst(SelectInst *Inst, const ShapeInfo &Shape,
2488 IRBuilder<> &Builder) {
2489 Value *Cond = Inst->getOperand(i_nocapture: 0);
2490 Value *OpA = Inst->getOperand(i_nocapture: 1);
2491 Value *OpB = Inst->getOperand(i_nocapture: 2);
2492
2493 MatrixTy Result;
2494 MatrixTy A = getMatrix(MatrixVal: OpA, SI: Shape, Builder);
2495 MatrixTy B = getMatrix(MatrixVal: OpB, SI: Shape, Builder);
2496
2497 SmallVector<Value*> CondV;
2498 Instruction *MDFrom = nullptr;
2499 if (isa<FixedVectorType>(Val: Cond->getType())) {
2500 MatrixTy C = getMatrix(MatrixVal: Cond, SI: Shape, Builder);
2501 llvm::copy(Range: C.vectors(), Out: std::back_inserter(x&: CondV));
2502 } else {
2503 CondV.resize(N: A.getNumVectors());
2504 llvm::fill(Range&: CondV, Value&: Cond);
2505 if (!ProfcheckDisableMetadataFixes)
2506 MDFrom = Inst;
2507 }
2508
2509 for (auto [CV, AV, BV] : llvm::zip_equal(t&: CondV, u: A.vectors(), args: B.vectors())) {
2510 assert(!(isa<VectorType>(CV->getType()) && static_cast<bool>(MDFrom)) &&
2511 "If we have a vector conditional, we should be propagating "
2512 "profile information.");
2513 Result.addVector(V: Builder.CreateSelect(C: CV, True: AV, False: BV, Name: "", MDFrom));
2514 }
2515
2516 return Result.addNumComputeOps(N: getNumOps(VT: Result.getVectorTy()) *
2517 Result.getNumVectors());
2518 }
2519
2520 /// Helper to linearize a matrix expression tree into a string. Currently
2521 /// matrix expressions are linarized by starting at an expression leaf and
2522 /// linearizing bottom up.
2523 struct ExprLinearizer {
2524 unsigned LengthToBreak = 100;
2525 std::string Str;
2526 raw_string_ostream Stream;
2527 unsigned LineLength = 0;
2528 const DataLayout &DL;
2529
2530 /// Mapping from instructions to matrixes. It is used to identify
2531 /// matrix instructions.
2532 const MapVector<Value *, MatrixTy> &Inst2Matrix;
2533
2534 /// Mapping from values to the leaves of all expressions that the value is
2535 /// part of.
2536 const DenseMap<Value *, SmallPtrSet<Value *, 2>> &Shared;
2537
2538 /// Set of matrix expressions in the scope of a given DISubprogram.
2539 const SmallSetVector<Value *, 32> &ExprsInSubprogram;
2540
2541 /// Leaf node of the expression to linearize.
2542 Value *Leaf;
2543
2544 /// Used to keep track of sub-expressions that get reused while linearizing
2545 /// the expression. Re-used sub-expressions are marked as (reused).
2546 SmallPtrSet<Value *, 8> ReusedExprs;
2547
2548 ExprLinearizer(const DataLayout &DL,
2549 const MapVector<Value *, MatrixTy> &Inst2Matrix,
2550 const DenseMap<Value *, SmallPtrSet<Value *, 2>> &Shared,
2551 const SmallSetVector<Value *, 32> &ExprsInSubprogram,
2552 Value *Leaf)
2553 : Stream(Str), DL(DL), Inst2Matrix(Inst2Matrix), Shared(Shared),
2554 ExprsInSubprogram(ExprsInSubprogram), Leaf(Leaf) {}
2555
2556 void indent(unsigned N) {
2557 LineLength += N;
2558 for (unsigned i = 0; i < N; i++)
2559 Stream << " ";
2560 }
2561
2562 void lineBreak() {
2563 Stream << "\n";
2564 LineLength = 0;
2565 }
2566
2567 void maybeIndent(unsigned Indent) {
2568 if (LineLength >= LengthToBreak)
2569 lineBreak();
2570
2571 if (LineLength == 0)
2572 indent(N: Indent);
2573 }
2574
2575 void write(StringRef S) {
2576 LineLength += S.size();
2577 Stream << S;
2578 }
2579
2580 Value *getUnderlyingObjectThroughLoads(Value *V) {
2581 if (Value *Ptr = getPointerOperand(V))
2582 return getUnderlyingObjectThroughLoads(V: Ptr);
2583 else if (V->getType()->isPointerTy())
2584 return getUnderlyingObject(V);
2585 return V;
2586 }
2587
2588 /// Returns true if \p V is a matrix value in the given subprogram.
2589 bool isMatrix(Value *V) const { return ExprsInSubprogram.count(key: V); }
2590
2591 /// If \p V is a matrix value, print its shape as NumRows x NumColumns to
2592 /// \p SS.
2593 void prettyPrintMatrixType(Value *V, raw_string_ostream &SS) {
2594 auto M = Inst2Matrix.find(Key: V);
2595 if (M == Inst2Matrix.end())
2596 SS << "unknown";
2597 else {
2598 SS << M->second.getNumRows();
2599 SS << "x";
2600 SS << M->second.getNumColumns();
2601 }
2602 }
2603
2604 /// Write the called function name. Handles calls to llvm.matrix.*
2605 /// specially: we write the name, followed by the dimensions of the input
2606 /// matrixes, followed by the scalar type name.
2607 void writeFnName(CallInst *CI) {
2608 if (!CI->getCalledFunction())
2609 write(S: "<no called fn>");
2610 else {
2611 StringRef Name = CI->getCalledFunction()->getName();
2612 if (!Name.starts_with(Prefix: "llvm.matrix")) {
2613 write(S: Name);
2614 return;
2615 }
2616 auto *II = cast<IntrinsicInst>(Val: CI);
2617 write(S: Intrinsic::getBaseName(id: II->getIntrinsicID())
2618 .drop_front(N: StringRef("llvm.matrix.").size()));
2619 write(S: ".");
2620 std::string Tmp;
2621 raw_string_ostream SS(Tmp);
2622
2623 switch (II->getIntrinsicID()) {
2624 case Intrinsic::matrix_multiply:
2625 prettyPrintMatrixType(V: II->getOperand(i_nocapture: 0), SS);
2626 SS << ".";
2627 prettyPrintMatrixType(V: II->getOperand(i_nocapture: 1), SS);
2628 SS << "." << *II->getType()->getScalarType();
2629 break;
2630 case Intrinsic::matrix_transpose:
2631 prettyPrintMatrixType(V: II->getOperand(i_nocapture: 0), SS);
2632 SS << "." << *II->getType()->getScalarType();
2633 break;
2634 case Intrinsic::matrix_column_major_load:
2635 prettyPrintMatrixType(V: II, SS);
2636 SS << "." << *II->getType()->getScalarType();
2637 break;
2638 case Intrinsic::matrix_column_major_store:
2639 prettyPrintMatrixType(V: II->getOperand(i_nocapture: 0), SS);
2640 SS << "." << *II->getOperand(i_nocapture: 0)->getType()->getScalarType();
2641 break;
2642 default:
2643 llvm_unreachable("Unhandled case");
2644 }
2645 write(S: Tmp);
2646 }
2647 }
2648
2649 unsigned getNumShapeArgs(CallInst *CI) const {
2650 if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(Val: CI)) {
2651 switch (II->getIntrinsicID()) {
2652 case Intrinsic::matrix_multiply:
2653 return 3;
2654 case Intrinsic::matrix_transpose:
2655 return 2;
2656 case Intrinsic::matrix_column_major_load:
2657 case Intrinsic::matrix_column_major_store:
2658 return 3;
2659 default:
2660 return 0;
2661 }
2662 }
2663 return 0;
2664 }
2665
2666 /// Special printing for values: for pointers, we print if they refer to an
2667 /// (function) external address or a stack address, for other values we
2668 /// either print the constant or "scalar"/"matrix" for other values.
2669 void write(Value *V) {
2670 V = getUnderlyingObjectThroughLoads(V);
2671 if (V->getType()->isPointerTy()) {
2672 if (isa<AllocaInst>(Val: V)) {
2673 Stream << "stack addr";
2674 LineLength += StringRef("stack addr").size();
2675 } else {
2676 Stream << "addr";
2677 LineLength += StringRef("addr").size();
2678 }
2679 if (!V->getName().empty()) {
2680 Stream << " %" << V->getName() << "";
2681 LineLength += V->getName().size() + 2;
2682 }
2683 return;
2684 }
2685
2686 std::string Tmp;
2687 raw_string_ostream TmpStream(Tmp);
2688
2689 if (auto *CI = dyn_cast<ConstantInt>(Val: V))
2690 TmpStream << CI->getValue();
2691 else if (isa<Constant>(Val: V))
2692 TmpStream << "constant";
2693 else {
2694 if (isMatrix(V))
2695 TmpStream << "matrix";
2696 else
2697 TmpStream << "scalar";
2698 }
2699 Tmp = std::string(StringRef(Tmp).trim());
2700 LineLength += Tmp.size();
2701 Stream << Tmp;
2702 }
2703
2704 /// Linearize expression \p Expr starting at an indentation of \p Indent.
2705 /// Expressions that are re-used multiple times are prefixed with (reused)
2706 /// at the re-used root instruction.
2707 void linearizeExpr(Value *Expr, unsigned Indent, bool ParentReused,
2708 bool ParentShared) {
2709 auto *I = cast<Instruction>(Val: Expr);
2710 maybeIndent(Indent);
2711 SmallVector<Value *, 8> Ops;
2712
2713 // Is Expr shared with other expression leaves?
2714 bool ExprShared = false;
2715
2716 // Deal with shared subtrees. Mark them as shared, if required.
2717 if (!ParentShared) {
2718 auto SI = Shared.find(Val: Expr);
2719 assert(SI != Shared.end() && SI->second.count(Leaf));
2720
2721 for (Value *S : SI->second) {
2722 if (S == Leaf)
2723 continue;
2724 DebugLoc DL = cast<Instruction>(Val: S)->getDebugLoc();
2725 write(S: "shared with remark at line " + std::to_string(val: DL.getLine()) +
2726 " column " + std::to_string(val: DL.getCol()) + " (");
2727 }
2728 ExprShared = SI->second.size() > 1;
2729 }
2730
2731 bool Reused = !ReusedExprs.insert(Ptr: Expr).second;
2732 if (Reused && !ParentReused)
2733 write(S: "(reused) ");
2734
2735 if (auto *CI = dyn_cast<CallInst>(Val: I)) {
2736 writeFnName(CI);
2737
2738 Ops.append(in_start: CI->arg_begin(), in_end: CI->arg_end() - getNumShapeArgs(CI));
2739 } else if (isa<BitCastInst>(Val: Expr)) {
2740 // Special case bitcasts, which are used to materialize matrixes from
2741 // non-matrix ops.
2742 write(S: "matrix");
2743 return;
2744 } else {
2745 Ops.append(in_start: I->value_op_begin(), in_end: I->value_op_end());
2746 write(S: I->getOpcodeName());
2747 }
2748
2749 write(S: "(");
2750
2751 unsigned NumOpsToBreak = 1;
2752 if (match(V: Expr, P: m_Intrinsic<Intrinsic::matrix_column_major_load>()))
2753 NumOpsToBreak = 2;
2754
2755 for (Value *Op : Ops) {
2756 if (Ops.size() > NumOpsToBreak)
2757 lineBreak();
2758
2759 maybeIndent(Indent: Indent + 1);
2760 if (isMatrix(V: Op))
2761 linearizeExpr(Expr: Op, Indent: Indent + 1, ParentReused: Reused, ParentShared: ExprShared);
2762 else
2763 write(V: Op);
2764 if (Op != Ops.back())
2765 write(S: ", ");
2766 }
2767
2768 write(S: ")");
2769 }
2770
2771 const std::string &getResult() {
2772 return Str;
2773 }
2774 };
2775
2776 /// Generate remarks for matrix operations in a function. To generate remarks
2777 /// for matrix expressions, the following approach is used:
2778 /// 1. Use the inlined-at debug information to group matrix operations to the
2779 /// DISubprograms they are contained in.
2780 /// 2. Collect leaves of matrix expressions (done in
2781 /// RemarkGenerator::getExpressionLeaves) for each subprogram - expression
2782 // mapping. Leaves are lowered matrix instructions without other matrix
2783 // users (like stores) in the current subprogram.
2784 /// 3. For each leaf, create a remark containing a linearizied version of the
2785 /// matrix expression. The expression is linearized by a recursive
2786 /// bottom-up traversal of the matrix operands, starting at a leaf. Note
2787 /// that multiple leaves can share sub-expressions. Shared subexpressions
2788 /// are explicitly marked as shared().
2789 struct RemarkGenerator {
2790 const MapVector<Value *, MatrixTy> &Inst2Matrix;
2791 OptimizationRemarkEmitter &ORE;
2792 Function &Func;
2793 const DataLayout &DL;
2794
2795 RemarkGenerator(const MapVector<Value *, MatrixTy> &Inst2Matrix,
2796 OptimizationRemarkEmitter &ORE, Function &Func)
2797 : Inst2Matrix(Inst2Matrix), ORE(ORE), Func(Func),
2798 DL(Func.getDataLayout()) {}
2799
2800 /// Return all leaves of the expressions in \p ExprsInSubprogram. Those are
2801 /// instructions in Inst2Matrix returning void or without any users in
2802 /// \p ExprsInSubprogram. Currently that should only include stores.
2803 SmallVector<Value *, 4>
2804 getExpressionLeaves(const SmallSetVector<Value *, 32> &ExprsInSubprogram) {
2805 SmallVector<Value *, 4> Leaves;
2806 for (auto *Expr : ExprsInSubprogram)
2807 if (Expr->getType()->isVoidTy() ||
2808 !any_of(Range: Expr->users(), P: [&ExprsInSubprogram](User *U) {
2809 return ExprsInSubprogram.count(key: U);
2810 }))
2811 Leaves.push_back(Elt: Expr);
2812 return Leaves;
2813 }
2814
2815 /// Recursively traverse expression \p V starting at \p Leaf and add \p Leaf
2816 /// to all visited expressions in \p Shared. Limit the matrix operations to
2817 /// the ones in \p ExprsInSubprogram.
2818 void collectSharedInfo(Value *Leaf, Value *V,
2819 const SmallSetVector<Value *, 32> &ExprsInSubprogram,
2820 DenseMap<Value *, SmallPtrSet<Value *, 2>> &Shared) {
2821
2822 if (!ExprsInSubprogram.count(key: V))
2823 return;
2824
2825 Shared[V].insert(Ptr: Leaf);
2826
2827 for (Value *Op : cast<Instruction>(Val: V)->operand_values())
2828 collectSharedInfo(Leaf, V: Op, ExprsInSubprogram, Shared);
2829 }
2830
2831 /// Calculate the number of exclusive and shared op counts for expression
2832 /// starting at \p V. Expressions used multiple times are counted once.
2833 /// Limit the matrix operations to the ones in \p ExprsInSubprogram.
2834 std::pair<OpInfoTy, OpInfoTy>
2835 sumOpInfos(Value *Root, SmallPtrSetImpl<Value *> &ReusedExprs,
2836 const SmallSetVector<Value *, 32> &ExprsInSubprogram,
2837 DenseMap<Value *, SmallPtrSet<Value *, 2>> &Shared) const {
2838 if (!ExprsInSubprogram.count(key: Root))
2839 return {};
2840
2841 // Already counted this expression. Stop.
2842 if (!ReusedExprs.insert(Ptr: Root).second)
2843 return {};
2844
2845 OpInfoTy SharedCount;
2846 OpInfoTy Count;
2847
2848 auto I = Shared.find(Val: Root);
2849 auto CM = Inst2Matrix.find(Key: Root);
2850 if (I->second.size() == 1)
2851 Count = CM->second.getOpInfo();
2852 else
2853 SharedCount = CM->second.getOpInfo();
2854
2855 for (Value *Op : cast<Instruction>(Val: Root)->operand_values()) {
2856 auto C = sumOpInfos(Root: Op, ReusedExprs, ExprsInSubprogram, Shared);
2857 Count += C.first;
2858 SharedCount += C.second;
2859 }
2860 return {Count, SharedCount};
2861 }
2862
2863 void emitRemarks() {
2864 if (!ORE.allowExtraAnalysis(DEBUG_TYPE))
2865 return;
2866
2867 // Map matrix operations to their containting subprograms, by traversing
2868 // the inlinedAt chain. If the function does not have a DISubprogram, we
2869 // only map them to the containing function.
2870 MapVector<DISubprogram *, SmallVector<Value *, 8>> Subprog2Exprs;
2871 for (const auto &KV : Inst2Matrix) {
2872 if (Func.getSubprogram()) {
2873 auto *I = cast<Instruction>(Val: KV.first);
2874 DILocation *Context = I->getDebugLoc();
2875 while (Context) {
2876 Subprog2Exprs[getSubprogram(Scope: Context->getScope())].push_back(
2877 Elt: KV.first);
2878 Context = DebugLoc(Context).getInlinedAt();
2879 }
2880 } else {
2881 Subprog2Exprs[nullptr].push_back(Elt: KV.first);
2882 }
2883 }
2884 for (auto &KV : Subprog2Exprs) {
2885 SmallSetVector<Value *, 32> ExprsInSubprogram(KV.second.begin(),
2886 KV.second.end());
2887 auto Leaves = getExpressionLeaves(ExprsInSubprogram);
2888
2889 DenseMap<Value *, SmallPtrSet<Value *, 2>> Shared;
2890 for (Value *Leaf : Leaves)
2891 collectSharedInfo(Leaf, V: Leaf, ExprsInSubprogram, Shared);
2892
2893 // Generate remarks for each leaf.
2894 for (auto *L : Leaves) {
2895
2896 DebugLoc Loc = cast<Instruction>(Val: L)->getDebugLoc();
2897 DILocation *Context = cast<Instruction>(Val: L)->getDebugLoc();
2898 while (Context) {
2899 if (getSubprogram(Scope: Context->getScope()) == KV.first) {
2900 Loc = Context;
2901 break;
2902 }
2903 Context = DebugLoc(Context).getInlinedAt();
2904 }
2905
2906 SmallPtrSet<Value *, 8> ReusedExprs;
2907 OpInfoTy Counts, SharedCounts;
2908 std::tie(args&: Counts, args&: SharedCounts) =
2909 sumOpInfos(Root: L, ReusedExprs, ExprsInSubprogram, Shared);
2910
2911 OptimizationRemark Rem(DEBUG_TYPE, "matrix-lowered", Loc,
2912 cast<Instruction>(Val: L)->getParent());
2913
2914 Rem << "Lowered with ";
2915 Rem << ore::NV("NumStores", Counts.NumStores) << " stores, "
2916 << ore::NV("NumLoads", Counts.NumLoads) << " loads, "
2917 << ore::NV("NumComputeOps", Counts.NumComputeOps)
2918 << " compute ops, "
2919 << ore::NV("NumExposedTransposes", Counts.NumExposedTransposes)
2920 << " exposed transposes";
2921
2922 if (SharedCounts.NumStores > 0 || SharedCounts.NumLoads > 0 ||
2923 SharedCounts.NumComputeOps > 0) {
2924 Rem << ",\nadditionally "
2925 << ore::NV("NumStores", SharedCounts.NumStores) << " stores, "
2926 << ore::NV("NumLoads", SharedCounts.NumLoads) << " loads, "
2927 << ore::NV("NumFPOps", SharedCounts.NumComputeOps)
2928 << " compute ops"
2929 << " are shared with other expressions";
2930 }
2931
2932 Rem << ("\n" + linearize(L, Shared, ExprsInSubprogram, DL));
2933 ORE.emit(OptDiag&: Rem);
2934 }
2935 }
2936 }
2937
2938 std::string
2939 linearize(Value *L,
2940 const DenseMap<Value *, SmallPtrSet<Value *, 2>> &Shared,
2941 const SmallSetVector<Value *, 32> &ExprsInSubprogram,
2942 const DataLayout &DL) {
2943 ExprLinearizer Lin(DL, Inst2Matrix, Shared, ExprsInSubprogram, L);
2944 Lin.linearizeExpr(Expr: L, Indent: 0, ParentReused: false, ParentShared: false);
2945 return Lin.getResult();
2946 }
2947 };
2948};
2949} // namespace
2950
2951PreservedAnalyses LowerMatrixIntrinsicsPass::run(Function &F,
2952 FunctionAnalysisManager &AM) {
2953 auto &TTI = AM.getResult<TargetIRAnalysis>(IR&: F);
2954
2955 LowerMatrixIntrinsics LMT(F, TTI, Minimal ? nullptr : &AM);
2956 if (LMT.Visit()) {
2957 PreservedAnalyses PA;
2958 if (!Minimal) {
2959 PA.preserve<LoopAnalysis>();
2960 PA.preserve<DominatorTreeAnalysis>();
2961 }
2962 return PA;
2963 }
2964 return PreservedAnalyses::all();
2965}
2966
2967void LowerMatrixIntrinsicsPass::printPipeline(
2968 raw_ostream &OS, function_ref<StringRef(StringRef)> MapClassName2PassName) {
2969 static_cast<PassInfoMixin<LowerMatrixIntrinsicsPass> *>(this)->printPipeline(
2970 OS, MapClassName2PassName);
2971 OS << '<';
2972 if (Minimal)
2973 OS << "minimal";
2974 OS << '>';
2975}
2976