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 Ops: m_Value(), Ops: m_Value(), Ops: m_Value(V&: M), Ops: m_Value(V&: N), Ops: m_Value(V&: K))))
331 return ShapeInfo(M, K);
332 if (match(V: I, P: m_Intrinsic<Intrinsic::matrix_transpose>(Ops: m_Value(), Ops: m_Value(V&: M),
333 Ops: 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 Ops: m_Value(), Ops: m_Value(), Ops: m_Value(), Ops: m_Value(), Ops: m_Value(V&: M),
339 Ops: m_Value(V&: N))))
340 return ShapeInfo(N, M);
341 if (match(V: I, P: m_Intrinsic<Intrinsic::matrix_column_major_load>(
342 Ops: m_Value(), Ops: m_Value(), Ops: m_Value(), Ops: m_Value(V&: M), Ops: 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 Ops: m_Value(V&: MatrixA), Ops: m_Value(V&: MatrixB), Ops: m_Value(V&: M),
814 Ops: m_Value(V&: N), Ops: 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 Ops: m_Value(V&: MatrixA), Ops: m_Value(V&: M), Ops: 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 Ops: m_Value(V&: MatrixA), Ops: m_Value(), Ops: m_Value(), Ops: m_Value(),
828 Ops: m_Value(V&: M), Ops: 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 Ops: m_Value(V&: TA), Ops: m_ConstantInt(CI&: R), Ops: 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 Ops: m_Value(V&: TATA), Ops: m_Specific(V: C), Ops: 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 Ops: m_Value(V&: TAMA), Ops: m_Value(V&: TAMB), Ops: m_ConstantInt(CI&: R),
952 Ops: m_ConstantInt(CI&: K), Ops: 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 Ops: m_Value(V&: A), Ops: m_Value(V&: B), Ops: m_ConstantInt(CI&: R),
1034 Ops: m_ConstantInt(CI&: K), Ops: m_ConstantInt(CI&: C))) &&
1035 match(V: A, P: m_Intrinsic<Intrinsic::matrix_transpose>(Ops: m_Value(V&: AT))) &&
1036 match(V: B, P: m_Intrinsic<Intrinsic::matrix_transpose>(Ops: 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 Ops: m_Value(V&: AT), Ops: m_ConstantInt(CI&: R), Ops: m_ConstantInt(CI&: C))) &&
1054 match(V: B, P: m_Intrinsic<Intrinsic::matrix_transpose>(
1055 Ops: m_Value(V&: BT), Ops: m_ConstantInt(), Ops: 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 Ops: m_Value(), Ops: 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 Ops: m_Value(V&: Arg)))) {
1724 auto *MatLoad = cast<IntrinsicInst>(Val: Op);
1725 bool IsVolatile = cast<ConstantInt>(Val: MatLoad->getArgOperand(i: 2))->isOne();
1726 // Preserve the volatile flag and alignment of the original load.
1727 Align Alignment = getAlignForIndex(
1728 Idx: 0, Stride: MatLoad->getArgOperand(i: 1),
1729 ElementTy: cast<FixedVectorType>(Val: Op->getType())->getElementType(),
1730 A: MatLoad->getParamAlign(ArgNo: 0));
1731 auto *NewLoad = Builder.CreateAlignedLoad(Ty: Op->getType(), Ptr: Arg, Align: Alignment,
1732 isVolatile: IsVolatile);
1733 Op->replaceAllUsesWith(V: NewLoad);
1734 eraseFromParentAndRemoveFromShapeMap(Inst: cast<Instruction>(Val: Op));
1735 return;
1736 } else if (match(V: Op, P: m_Intrinsic<Intrinsic::matrix_transpose>(
1737 Ops: m_Value(V&: Arg)))) {
1738 ToRemove.push_back(Elt: cast<Instruction>(Val: Op));
1739 Op->replaceAllUsesWith(V: Arg);
1740 return;
1741 }
1742 };
1743
1744 for (auto *V : ToFlatten)
1745 FlattenArg(V);
1746
1747 LHS = MatMul->getArgOperand(i: 0);
1748
1749 // Insert mul/fmul and llvm.vector.reduce.fadd
1750 Value *Mul =
1751 IsIntVec ? Builder.CreateMul(LHS, RHS) : Builder.CreateFMul(L: LHS, R: RHS);
1752
1753 Value *Result;
1754 if (IsIntVec)
1755 Result = Builder.CreateAddReduce(Src: Mul);
1756 else {
1757 Result = Builder.CreateFAddReduce(
1758 Acc: ConstantFP::get(
1759 Ty: cast<FixedVectorType>(Val: LHS->getType())->getElementType(), V: 0.0),
1760 Src: Mul);
1761 cast<Instruction>(Val: Result)->setFastMathFlags(FMF);
1762 }
1763
1764 // pack scalar back into a matrix and then replace matmul inst
1765 Result = Builder.CreateInsertElement(Vec: PoisonValue::get(T: MatMul->getType()),
1766 NewElt: Result, Idx: uint64_t(0));
1767 MatMul->replaceAllUsesWith(V: Result);
1768 FusedInsts.insert(Ptr: MatMul);
1769 ToRemove.push_back(Elt: MatMul);
1770 }
1771
1772 /// Given \p Remainder iterations of the the matmul inner loop,
1773 /// potentially lower \p Blocksize that is used for the underlying
1774 /// vector.
1775 unsigned capBlockSize(unsigned BlockSize, unsigned Remainder, Type *EltType) {
1776 if (BlockSize <= Remainder)
1777 return BlockSize;
1778
1779 // If the remainder is also a legal type just use it.
1780 auto *VecTy = FixedVectorType::get(ElementType: EltType, NumElts: Remainder);
1781 if (TTI.isTypeLegal(Ty: VecTy))
1782 return Remainder;
1783
1784 // Similarly, if the vector is small enough that we don't want
1785 // to split further.
1786 if (VecTy->getPrimitiveSizeInBits() <= SplitMatmulRemainderOverThreshold)
1787 return Remainder;
1788
1789 // Gradually lower the vectorization factor to cover the
1790 // remainder.
1791 do {
1792 BlockSize /= 2;
1793 } while (BlockSize > Remainder);
1794 return BlockSize;
1795 }
1796
1797 /// Compute \p Result += \p A * \p B for input matrices with left-associating
1798 /// addition.
1799 ///
1800 /// We can fold a transpose into the operand that is used to extract scalars.
1801 /// This is the first operands with row-major and the second with
1802 /// column-major. If \p IsScalarMatrixTransposed we assume the appropriate
1803 /// operand is transposed.
1804 void emitMatrixMultiply(MatrixTy &Result, const MatrixTy &A,
1805 const MatrixTy &B, IRBuilder<> &Builder, bool IsTiled,
1806 bool IsScalarMatrixTransposed, FastMathFlags FMF) {
1807 const unsigned VF = std::max<unsigned>(
1808 a: TTI.getRegisterBitWidth(K: TargetTransformInfo::RGK_FixedWidthVector)
1809 .getFixedValue() /
1810 Result.getElementType()->getPrimitiveSizeInBits().getFixedValue(),
1811 b: 1U);
1812 unsigned R = Result.getNumRows();
1813 unsigned C = Result.getNumColumns();
1814 unsigned M = A.getNumColumns();
1815
1816 bool IsFP = Result.getElementType()->isFloatingPointTy();
1817 assert(A.isColumnMajor() == B.isColumnMajor() &&
1818 Result.isColumnMajor() == A.isColumnMajor() &&
1819 "operands must agree on matrix layout");
1820 unsigned NumComputeOps = 0;
1821
1822 Builder.setFastMathFlags(FMF);
1823
1824 if (A.isColumnMajor()) {
1825 // Multiply columns from the first operand with scalars from the second
1826 // operand. Then move along the K axes and accumulate the columns. With
1827 // this the adds can be vectorized without reassociation.
1828 for (unsigned J = 0; J < C; ++J) {
1829 unsigned BlockSize = VF;
1830 // If Result is zero, we don't need to accumulate in the K==0 iteration.
1831 bool isSumZero = isa<ConstantAggregateZero>(Val: Result.getColumn(i: J));
1832
1833 for (unsigned I = 0; I < R; I += BlockSize) {
1834 // Lower block size to make sure we stay within bounds.
1835 BlockSize = capBlockSize(BlockSize, Remainder: R - I, EltType: Result.getElementType());
1836 Value *Sum = IsTiled ? Result.extractVector(I, J, NumElts: BlockSize, Builder)
1837 : nullptr;
1838 for (unsigned K = 0; K < M; ++K) {
1839 Value *L = A.extractVector(I, J: K, NumElts: BlockSize, Builder);
1840 Value *RH = Builder.CreateExtractElement(
1841 Vec: B.getColumn(i: IsScalarMatrixTransposed ? K : J),
1842 Idx: IsScalarMatrixTransposed ? J : K);
1843 Value *Splat = Builder.CreateVectorSplat(NumElts: BlockSize, V: RH, Name: "splat");
1844 Sum =
1845 createMulAdd(Sum: isSumZero && K == 0 ? nullptr : Sum, A: L, B: Splat,
1846 UseFPOp: IsFP, Builder, AllowContraction: FMF.allowContract(), NumComputeOps);
1847 }
1848 Result.setVector(i: J,
1849 V: insertVector(Col: Result.getVector(i: J), I, Block: Sum, Builder));
1850 }
1851 }
1852 } else {
1853 // Multiply rows from the second operand with scalars from the first
1854 // operand. Then move along the K axes and accumulate the rows. With this
1855 // the adds can be vectorized without reassociation.
1856 for (unsigned I = 0; I < R; ++I) {
1857 unsigned BlockSize = VF;
1858 bool isSumZero = isa<ConstantAggregateZero>(Val: Result.getRow(i: I));
1859 for (unsigned J = 0; J < C; J += BlockSize) {
1860 // Lower the vectorization factor to cover the remainder.
1861 BlockSize = capBlockSize(BlockSize, Remainder: C - J, EltType: Result.getElementType());
1862
1863 Value *Sum = nullptr;
1864 for (unsigned K = 0; K < M; ++K) {
1865 Value *R = B.extractVector(I: K, J, NumElts: BlockSize, Builder);
1866 Value *LH = Builder.CreateExtractElement(
1867 Vec: A.getVector(i: IsScalarMatrixTransposed ? K : I),
1868 Idx: IsScalarMatrixTransposed ? I : K);
1869 Value *Splat = Builder.CreateVectorSplat(NumElts: BlockSize, V: LH, Name: "splat");
1870 Sum =
1871 createMulAdd(Sum: isSumZero && K == 0 ? nullptr : Sum, A: Splat, B: R,
1872 UseFPOp: IsFP, Builder, AllowContraction: FMF.allowContract(), NumComputeOps);
1873 }
1874 Result.setVector(i: I,
1875 V: insertVector(Col: Result.getVector(i: I), I: J, Block: Sum, Builder));
1876 }
1877 }
1878 }
1879 Result.addNumComputeOps(N: NumComputeOps);
1880 }
1881
1882 /// Ensure that the memory in \p Load does not alias \p Store by potentially
1883 /// copying it to a new location. This new or otherwise the original location
1884 /// is returned.
1885 std::pair<Value *, AllocaInst *>
1886 getNonAliasingPointer(LoadInst *Load, StoreInst *Store, CallInst *MatMul) {
1887 MemoryLocation StoreLoc = MemoryLocation::get(SI: Store);
1888 MemoryLocation LoadLoc = MemoryLocation::get(LI: Load);
1889
1890 // If we can statically determine noalias we're good.
1891 if (AA->isNoAlias(LocA: LoadLoc, LocB: StoreLoc))
1892 return {Load->getPointerOperand(), nullptr};
1893
1894 // If the pointers are in different address spaces, we cannot compare them
1895 // at runtime. Conservatively copy the load operand to a new buffer.
1896 IRBuilder<> AllocaBuilder(&Func.getEntryBlock().front());
1897 if (Load->getPointerAddressSpace() != Store->getPointerAddressSpace()) {
1898 auto *VT = cast<FixedVectorType>(Val: Load->getType());
1899 auto *ArrayTy =
1900 ArrayType::get(ElementType: VT->getElementType(), NumElements: VT->getNumElements());
1901 AllocaInst *Alloca =
1902 AllocaBuilder.CreateAlloca(Ty: ArrayTy, AddrSpace: Load->getPointerAddressSpace());
1903 IRBuilder<> Builder(MatMul);
1904 Builder.CreateLifetimeStart(Ptr: Alloca);
1905 Builder.CreateMemCpy(Dst: Alloca, DstAlign: Alloca->getAlign(),
1906 Src: Load->getPointerOperand(), SrcAlign: Load->getAlign(),
1907 Size: LoadLoc.Size.getValue());
1908 return {Alloca, Alloca};
1909 }
1910
1911 // Create code to check if the memory locations of the Load and Store
1912 // overlap and if they do, copy Load's operand to a new buffer.
1913
1914 // First, create new blocks for 2n part of the check and the copy.
1915 BasicBlock *Check0 = MatMul->getParent();
1916 // FIXME: Use lazy DTU and update SplitBlock to accept a DTU instead of a
1917 // DT. Manually collect dominator tree updates, to avoid unnecessary work,
1918 // as we adjust Check0 and Check1's branches.
1919 SmallVector<DominatorTree::UpdateType, 4> DTUpdates;
1920 for (BasicBlock *Succ : successors(BB: Check0))
1921 DTUpdates.push_back(Elt: {DT->Delete, Check0, Succ});
1922
1923 BasicBlock *Check1 =
1924 SplitBlock(Old: MatMul->getParent(), SplitPt: MatMul, DTU: (DomTreeUpdater *)nullptr, LI,
1925 MSSAU: nullptr, BBName: "alias_cont");
1926 BasicBlock *Copy =
1927 SplitBlock(Old: MatMul->getParent(), SplitPt: MatMul, DTU: (DomTreeUpdater *)nullptr, LI,
1928 MSSAU: nullptr, BBName: "copy");
1929 BasicBlock *Fusion =
1930 SplitBlock(Old: MatMul->getParent(), SplitPt: MatMul, DTU: (DomTreeUpdater *)nullptr, LI,
1931 MSSAU: nullptr, BBName: "no_alias");
1932
1933 // Check if the loaded memory location begins before the end of the store
1934 // location. If the condition holds, they might overlap, otherwise they are
1935 // guaranteed to not overlap.
1936 IRBuilder<> Builder(MatMul);
1937 Check0->getTerminator()->eraseFromParent();
1938 Builder.SetInsertPoint(Check0);
1939 Type *AddrTy = DL.getAddressType(PtrTy: Store->getPointerOperand()->getType());
1940 Value *StoreBegin = Store->getPointerOperand();
1941 Value *StoreEnd = Builder.CreatePtrAdd(
1942 Ptr: StoreBegin, Offset: ConstantInt::get(Ty: AddrTy, V: StoreLoc.Size.getValue()),
1943 Name: "store.end",
1944 NW: GEPNoWrapFlags::inBounds() | GEPNoWrapFlags::noUnsignedWrap());
1945 Value *LoadBegin = Load->getPointerOperand();
1946 CondBrInst *BR1 = Builder.CreateCondBr(
1947 Cond: Builder.CreateICmpULT(LHS: LoadBegin, RHS: StoreEnd), True: Check1, False: Fusion);
1948 setExplicitlyUnknownBranchWeightsIfProfiled(I&: *BR1, DEBUG_TYPE);
1949
1950 // Check if the store begins before the end of the load location. If the
1951 // condition holds, they alias, otherwise they are guaranteed to not
1952 // overlap.
1953 Check1->getTerminator()->eraseFromParent();
1954 Builder.SetInsertPoint(TheBB: Check1, IP: Check1->begin());
1955
1956 auto *VT = cast<FixedVectorType>(Val: Load->getType());
1957 // Use an array type for the alloca, to avoid potentially huge alignment
1958 // requirements for large vector types.
1959 auto *ArrayTy = ArrayType::get(ElementType: VT->getElementType(), NumElements: VT->getNumElements());
1960 AllocaInst *Alloca =
1961 AllocaBuilder.CreateAlloca(Ty: ArrayTy, AddrSpace: Load->getPointerAddressSpace());
1962 Builder.CreateLifetimeStart(Ptr: Alloca);
1963
1964 Value *LoadEnd = Builder.CreatePtrAdd(
1965 Ptr: LoadBegin, Offset: ConstantInt::get(Ty: AddrTy, V: LoadLoc.Size.getValue()),
1966 Name: "load.end",
1967 NW: GEPNoWrapFlags::inBounds() | GEPNoWrapFlags::noUnsignedWrap());
1968 CondBrInst *BR2 = Builder.CreateCondBr(
1969 Cond: Builder.CreateICmpULT(LHS: StoreBegin, RHS: LoadEnd), True: Copy, False: Fusion);
1970 setExplicitlyUnknownBranchWeightsIfProfiled(I&: *BR2, DEBUG_TYPE);
1971
1972 // Copy load operand to new alloca.
1973 Builder.SetInsertPoint(TheBB: Copy, IP: Copy->begin());
1974 Builder.CreateMemCpy(Dst: Alloca, DstAlign: Alloca->getAlign(), Src: Load->getPointerOperand(),
1975 SrcAlign: Load->getAlign(), Size: LoadLoc.Size.getValue());
1976 Builder.SetInsertPoint(TheBB: Fusion, IP: Fusion->begin());
1977 PHINode *PHI = Builder.CreatePHI(Ty: Load->getPointerOperandType(), NumReservedValues: 3);
1978 PHI->addIncoming(V: Load->getPointerOperand(), BB: Check0);
1979 PHI->addIncoming(V: Load->getPointerOperand(), BB: Check1);
1980 PHI->addIncoming(V: Alloca, BB: Copy);
1981
1982 // Adjust DT.
1983 DTUpdates.push_back(Elt: {DT->Insert, Check0, Check1});
1984 DTUpdates.push_back(Elt: {DT->Insert, Check0, Fusion});
1985 DTUpdates.push_back(Elt: {DT->Insert, Check1, Copy});
1986 DTUpdates.push_back(Elt: {DT->Insert, Check1, Fusion});
1987 DT->applyUpdates(Updates: DTUpdates);
1988 return {PHI, Alloca};
1989 }
1990
1991 bool isFusionProfitable(CallInst *MatMul) {
1992 if (ForceFusion)
1993 return true;
1994
1995 ShapeInfo LShape(MatMul->getArgOperand(i: 2), MatMul->getArgOperand(i: 3));
1996 ShapeInfo RShape(MatMul->getArgOperand(i: 3), MatMul->getArgOperand(i: 4));
1997
1998 const unsigned R = LShape.NumRows;
1999 const unsigned C = RShape.NumColumns;
2000 const unsigned M = LShape.NumColumns;
2001 auto *EltType = cast<FixedVectorType>(Val: MatMul->getType())->getElementType();
2002
2003 const unsigned VF = std::max<unsigned>(
2004 a: TTI.getRegisterBitWidth(K: TargetTransformInfo::RGK_FixedWidthVector)
2005 .getFixedValue() /
2006 EltType->getPrimitiveSizeInBits().getFixedValue(),
2007 b: 1U);
2008
2009 // Cost model for tiling
2010 //
2011 // For tiling to be beneficial, we need reuse either along the R or
2012 // the C axis. We vectorize along the R axis so that means at least
2013 // 3 elements.
2014 // TODO: Also consider cost of copying if operands alias.
2015 if (R <= VF && C == 1)
2016 return false;
2017 // Then we need enough elements to exceed the number of vector
2018 // registers we have. Note that this is an oversimplification since
2019 // fusing also takes some extra loads which may exceed the number of
2020 // reloads necessary.
2021 unsigned Op0Regs = (R + VF - 1) / VF * M;
2022 unsigned Op1Regs = (M + VF - 1) / VF * C;
2023 return Op0Regs + Op1Regs >
2024 TTI.getNumberOfRegisters(ClassID: TTI.getRegisterClassForType(Vector: true));
2025 }
2026
2027 MatrixTy getZeroMatrix(Type *EltType, unsigned R, unsigned C) {
2028 MatrixTy Res;
2029 auto *ColumType = FixedVectorType::get(ElementType: EltType, NumElts: R);
2030 for (unsigned I = 0; I < C; ++I)
2031 Res.addVector(V: ConstantAggregateZero::get(Ty: ColumType));
2032 return Res;
2033 }
2034
2035 void createTiledLoops(CallInst *MatMul, Value *LPtr, ShapeInfo LShape,
2036 Value *RPtr, ShapeInfo RShape, StoreInst *Store) {
2037 auto *EltType = cast<FixedVectorType>(Val: MatMul->getType())->getElementType();
2038
2039 // Create the main tiling loop nest.
2040 TileInfo TI(LShape.NumRows, RShape.NumColumns, LShape.NumColumns, TileSize);
2041 DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Lazy);
2042 Instruction *InsertI = cast<Instruction>(Val: MatMul);
2043 BasicBlock *Start = InsertI->getParent();
2044 BasicBlock *End =
2045 SplitBlock(Old: InsertI->getParent(), SplitPt: InsertI, DT, LI, MSSAU: nullptr, BBName: "continue");
2046 IRBuilder<> Builder(MatMul);
2047 BasicBlock *InnerBody = TI.CreateTiledLoops(Start, End, B&: Builder, DTU, LI&: *LI);
2048
2049 Type *TileVecTy =
2050 FixedVectorType::get(ElementType: MatMul->getType()->getScalarType(), NumElts: TileSize);
2051 MatrixTy TileResult;
2052 // Insert in the inner loop header.
2053 Builder.SetInsertPoint(TI.KLoop.Header->getTerminator());
2054 // Create PHI nodes for the result columns to accumulate across iterations.
2055 SmallVector<PHINode *, 4> ColumnPhis;
2056 for (unsigned I = 0; I < TileSize; I++) {
2057 auto *Phi = Builder.CreatePHI(Ty: TileVecTy, NumReservedValues: 2, Name: "result.vec." + Twine(I));
2058 Phi->addIncoming(V: ConstantAggregateZero::get(Ty: TileVecTy),
2059 BB: TI.RowLoop.Header->getSingleSuccessor());
2060 TileResult.addVector(V: Phi);
2061 ColumnPhis.push_back(Elt: Phi);
2062 }
2063
2064 // Insert in the inner loop body, which computes
2065 // Res += Load(CurrentRow, K) * Load(K, CurrentColumn)
2066 Builder.SetInsertPoint(InnerBody->getTerminator());
2067 // Load tiles of the operands.
2068 MatrixTy A =
2069 loadMatrix(MatrixPtr: LPtr, Align: {}, IsVolatile: false, MatrixShape: LShape, I: TI.RowLoop.Index, J: TI.KLoop.Index,
2070 ResultShape: {TileSize, TileSize}, EltTy: EltType, Builder);
2071 MatrixTy B =
2072 loadMatrix(MatrixPtr: RPtr, Align: {}, IsVolatile: false, MatrixShape: RShape, I: TI.KLoop.Index, J: TI.ColumnLoop.Index,
2073 ResultShape: {TileSize, TileSize}, EltTy: EltType, Builder);
2074 emitMatrixMultiply(Result&: TileResult, A, B, Builder, IsTiled: true, IsScalarMatrixTransposed: false,
2075 FMF: getFastMathFlags(Inst: MatMul));
2076 // Store result after the inner loop is done.
2077 Builder.SetInsertPoint(TI.RowLoop.Latch->getTerminator());
2078 storeMatrix(StoreVal: TileResult, MatrixPtr: Store->getPointerOperand(), MAlign: Store->getAlign(),
2079 IsVolatile: Store->isVolatile(), MatrixShape: {LShape.NumRows, RShape.NumColumns},
2080 I: TI.RowLoop.Index, J: TI.ColumnLoop.Index, EltTy: EltType, Builder);
2081
2082 for (unsigned I = 0; I < TileResult.getNumVectors(); I++)
2083 ColumnPhis[I]->addIncoming(V: TileResult.getVector(i: I), BB: TI.KLoop.Latch);
2084
2085 // Force unrolling of a few iterations of the inner loop, to make sure there
2086 // is enough work per iteration.
2087 // FIXME: The unroller should make this decision directly instead, but
2088 // currently the cost-model is not up to the task.
2089 unsigned InnerLoopUnrollCount = std::min(a: 10u, b: LShape.NumColumns / TileSize);
2090 addStringMetadataToLoop(TheLoop: LI->getLoopFor(BB: TI.KLoop.Header),
2091 MDString: "llvm.loop.unroll.count", V: InnerLoopUnrollCount);
2092 }
2093
2094 void emitSIMDTiling(CallInst *MatMul, LoadInst *LoadOp0, LoadInst *LoadOp1,
2095 StoreInst *Store,
2096 SmallPtrSetImpl<Instruction *> &FusedInsts) {
2097 assert(MatrixLayout == MatrixLayoutTy::ColumnMajor &&
2098 "Tiling only supported for column-major matrixes at the moment!");
2099 if (!isFusionProfitable(MatMul))
2100 return;
2101
2102 ShapeInfo LShape(MatMul->getArgOperand(i: 2), MatMul->getArgOperand(i: 3));
2103 ShapeInfo RShape(MatMul->getArgOperand(i: 3), MatMul->getArgOperand(i: 4));
2104
2105 const unsigned R = LShape.NumRows;
2106 const unsigned C = RShape.NumColumns;
2107 const unsigned M = LShape.NumColumns;
2108 auto *EltType = cast<FixedVectorType>(Val: MatMul->getType())->getElementType();
2109
2110 auto [APtr, AAlloca] = getNonAliasingPointer(Load: LoadOp0, Store, MatMul);
2111 auto [BPtr, BAlloca] = getNonAliasingPointer(Load: LoadOp1, Store, MatMul);
2112 Value *CPtr = Store->getPointerOperand();
2113
2114 // Use loop-based tiling when the number of expected operations exceeds
2115 // threshold.
2116 unsigned NumOps = getNumNativeVectorOps(EltType, R, M, C);
2117 bool UseLoops =
2118 (NumOps > TileLoopsThreshold) && R % TileSize == 0 && C % TileSize == 0;
2119 if (UseLoops)
2120 createTiledLoops(MatMul, LPtr: APtr, LShape, RPtr: BPtr, RShape, Store);
2121 else {
2122 IRBuilder<> Builder(Store);
2123 for (unsigned J = 0; J < C; J += TileSize)
2124 for (unsigned I = 0; I < R; I += TileSize) {
2125 const unsigned TileR = std::min(a: R - I, b: unsigned(TileSize));
2126 const unsigned TileC = std::min(a: C - J, b: unsigned(TileSize));
2127 MatrixTy Res = getZeroMatrix(EltType, R: TileR, C: TileC);
2128
2129 for (unsigned K = 0; K < M; K += TileSize) {
2130 const unsigned TileM = std::min(a: M - K, b: unsigned(TileSize));
2131 MatrixTy A =
2132 loadMatrix(MatrixPtr: APtr, Align: LoadOp0->getAlign(), IsVolatile: LoadOp0->isVolatile(),
2133 MatrixShape: LShape, I: getIndex(Ptr: APtr, V: I), J: getIndex(Ptr: APtr, V: K),
2134 ResultShape: {TileR, TileM}, EltTy: EltType, Builder);
2135 MatrixTy B =
2136 loadMatrix(MatrixPtr: BPtr, Align: LoadOp1->getAlign(), IsVolatile: LoadOp1->isVolatile(),
2137 MatrixShape: RShape, I: getIndex(Ptr: BPtr, V: K), J: getIndex(Ptr: BPtr, V: J),
2138 ResultShape: {TileM, TileC}, EltTy: EltType, Builder);
2139 emitMatrixMultiply(Result&: Res, A, B, Builder, IsTiled: true, IsScalarMatrixTransposed: false,
2140 FMF: getFastMathFlags(Inst: MatMul));
2141 }
2142 storeMatrix(StoreVal: Res, MatrixPtr: CPtr, MAlign: Store->getAlign(), IsVolatile: Store->isVolatile(), MatrixShape: {R, M},
2143 I: getIndex(Ptr: CPtr, V: I), J: getIndex(Ptr: CPtr, V: J), EltTy: EltType, Builder);
2144 }
2145 }
2146
2147 // End the lifetime of the allocas used for alias-safe copies.
2148 {
2149 IRBuilder<> Builder(Store);
2150 if (AAlloca)
2151 Builder.CreateLifetimeEnd(Ptr: AAlloca);
2152 if (BAlloca)
2153 Builder.CreateLifetimeEnd(Ptr: BAlloca);
2154 }
2155
2156 // Mark eliminated instructions as fused and remove them.
2157 FusedInsts.insert(Ptr: Store);
2158 FusedInsts.insert(Ptr: MatMul);
2159 eraseFromParentAndRemoveFromShapeMap(Inst: Store);
2160 eraseFromParentAndRemoveFromShapeMap(Inst: MatMul);
2161 if (LoadOp0->use_empty()) {
2162 FusedInsts.insert(Ptr: LoadOp0);
2163 eraseFromParentAndRemoveFromShapeMap(Inst: LoadOp0);
2164 }
2165 if (LoadOp1 != LoadOp0 && LoadOp1->use_empty()) {
2166 FusedInsts.insert(Ptr: LoadOp1);
2167 eraseFromParentAndRemoveFromShapeMap(Inst: LoadOp1);
2168 }
2169 }
2170
2171 /// Try to lower matrix multiply chains by fusing operations.
2172 ///
2173 /// Call finalizeLowering on lowered instructions. Instructions that are
2174 /// completely eliminated by fusion are added to \p FusedInsts.
2175 void
2176 LowerMatrixMultiplyFused(CallInst *MatMul,
2177 SmallPtrSetImpl<Instruction *> &FusedInsts,
2178 SmallVector<IntrinsicInst *, 16> &LifetimeEnds) {
2179 if (!FuseMatrix || !DT || TileSize == 0)
2180 return;
2181
2182 assert(AA && LI && "Analyses should be available");
2183
2184 Value *A = MatMul->getArgOperand(i: 0);
2185 Value *B = MatMul->getArgOperand(i: 1);
2186
2187 // We can fold the transpose into the operand that is used to fetch scalars.
2188 Value *T;
2189 if (MatrixLayout == MatrixLayoutTy::ColumnMajor
2190 ? match(V: B, P: m_Intrinsic<Intrinsic::matrix_transpose>(Ops: m_Value(V&: T)))
2191 : match(V: A, P: m_Intrinsic<Intrinsic::matrix_transpose>(Ops: m_Value(V&: T)))) {
2192 IRBuilder<> Builder(MatMul);
2193 auto *EltType =
2194 cast<FixedVectorType>(Val: MatMul->getType())->getElementType();
2195 ShapeInfo LShape(MatMul->getArgOperand(i: 2), MatMul->getArgOperand(i: 3));
2196 ShapeInfo RShape(MatMul->getArgOperand(i: 3), MatMul->getArgOperand(i: 4));
2197 const unsigned R = LShape.NumRows;
2198 const unsigned M = LShape.NumColumns;
2199 const unsigned C = RShape.NumColumns;
2200
2201 MatrixTy MA;
2202 MatrixTy MB;
2203
2204 Value *Transpose;
2205 if (MatrixLayout == MatrixLayoutTy::ColumnMajor) {
2206 MA = getMatrix(MatrixVal: A, SI: ShapeInfo(R, M), Builder);
2207 MB = getMatrix(MatrixVal: T, SI: ShapeInfo(C, M), Builder);
2208 Transpose = B;
2209 } else {
2210 MA = getMatrix(MatrixVal: T, SI: ShapeInfo(R, M), Builder);
2211 MB = getMatrix(MatrixVal: B, SI: ShapeInfo(C, M), Builder);
2212 Transpose = A;
2213 }
2214
2215 // Initialize the output
2216 MatrixTy Result(R, C, EltType);
2217
2218 emitMatrixMultiply(Result, A: MA, B: MB, Builder, IsTiled: false, IsScalarMatrixTransposed: true,
2219 FMF: getFastMathFlags(Inst: MatMul));
2220
2221 FusedInsts.insert(Ptr: MatMul);
2222 if (Transpose->hasOneUse()) {
2223 FusedInsts.insert(Ptr: cast<Instruction>(Val: Transpose));
2224 ToRemove.push_back(Elt: cast<Instruction>(Val: Transpose));
2225 // TODO: add a fake entry for the folded instruction so that this is
2226 // included in the expression in the remark.
2227 Inst2ColumnMatrix[Transpose] = MatrixTy(M, C, EltType);
2228 }
2229 finalizeLowering(Inst: MatMul, Matrix: Result, Builder);
2230 return;
2231 }
2232
2233 if (!MatMul->hasOneUse() || MatrixLayout != MatrixLayoutTy::ColumnMajor)
2234 return;
2235
2236 // Lower {ld, ld} -> matmul -> st chains. No need to call finalizeLowering
2237 // since the single store user will be lowered as part of this.
2238 auto *LoadOp0 = dyn_cast<LoadInst>(Val: A);
2239 auto *LoadOp1 = dyn_cast<LoadInst>(Val: B);
2240 auto *Store = dyn_cast<StoreInst>(Val: *MatMul->user_begin());
2241 if (LoadOp0 && LoadOp1 && Store) {
2242 // The store address must dominate the MatMul instruction, otherwise
2243 // we create invalid IR.
2244 SetVector<Value *> WorkList;
2245 WorkList.insert(X: Store->getOperand(i_nocapture: 1));
2246 SmallVector<Instruction *> ToHoist;
2247 for (unsigned I = 0; I != WorkList.size(); ++I) {
2248 Value *Current = WorkList[I];
2249 auto *CurrI = dyn_cast<Instruction>(Val: Current);
2250 if (!CurrI)
2251 continue;
2252 if (isa<PHINode>(Val: CurrI))
2253 return;
2254 if (DT->dominates(Def: CurrI, User: MatMul))
2255 continue;
2256 if (CurrI->mayHaveSideEffects() || CurrI->mayReadFromMemory())
2257 return;
2258 ToHoist.push_back(Elt: CurrI);
2259 WorkList.insert_range(R: CurrI->operands());
2260 }
2261
2262 sort(C&: ToHoist, Comp: [this](Instruction *A, Instruction *B) {
2263 return DT->dominates(Def: A, User: B);
2264 });
2265 for (Instruction *I : ToHoist)
2266 I->moveBefore(InsertPos: MatMul->getIterator());
2267
2268 // Deal with lifetime.end calls that might be between Load0/Load1 and the
2269 // store. To avoid introducing loads to dead objects (i.e. after the
2270 // lifetime has been termined by @llvm.lifetime.end), either sink them
2271 // after the store if in the same block, or remove the lifetime.end marker
2272 // otherwise. This might pessimize further optimizations, by extending the
2273 // lifetime of the object until the function returns, but should be
2274 // conservatively correct.
2275 MemoryLocation Load0Loc = MemoryLocation::get(LI: LoadOp0);
2276 MemoryLocation Load1Loc = MemoryLocation::get(LI: LoadOp1);
2277 BasicBlock *StoreParent = Store->getParent();
2278 bool FusableOpsInSameBlock = LoadOp0->getParent() == StoreParent &&
2279 LoadOp1->getParent() == StoreParent;
2280 for (unsigned Idx = 0; Idx != LifetimeEnds.size();) {
2281 IntrinsicInst *End = LifetimeEnds[Idx];
2282 llvm::scope_exit Inc([&Idx]() { Idx++; });
2283 // If the lifetime.end is guaranteed to be before the loads or after the
2284 // store, it won't interfere with fusion.
2285 if (DT->dominates(Def: End, User: LoadOp0) && DT->dominates(Def: End, User: LoadOp1))
2286 continue;
2287 if (DT->dominates(Def: Store, User: End))
2288 continue;
2289 // If all fusable ops are in the same block and the lifetime.end is in a
2290 // different block, it won't interfere with fusion.
2291 if (FusableOpsInSameBlock && End->getParent() != StoreParent)
2292 continue;
2293
2294 // If the loads don't alias the lifetime.end, it won't interfere with
2295 // fusion.
2296 MemoryLocation EndLoc = MemoryLocation::getForArgument(Call: End, ArgIdx: 0, TLI: nullptr);
2297 if (!EndLoc.Ptr)
2298 continue;
2299 if (AA->isNoAlias(LocA: Load0Loc, LocB: EndLoc) && AA->isNoAlias(LocA: Load1Loc, LocB: EndLoc))
2300 continue;
2301
2302 // If both lifetime.end and the store are in the same block, extend the
2303 // lifetime until after the store, so the new lifetime covers the loads
2304 // we introduce later.
2305 if (End->getParent() == StoreParent) {
2306 End->moveAfter(MovePos: Store);
2307 continue;
2308 }
2309
2310 // Otherwise remove the conflicting lifetime.end marker.
2311 ToRemove.push_back(Elt: End);
2312 std::swap(a&: LifetimeEnds[Idx], b&: LifetimeEnds.back());
2313 LifetimeEnds.pop_back();
2314 Inc.release();
2315 }
2316
2317 emitSIMDTiling(MatMul, LoadOp0, LoadOp1, Store, FusedInsts);
2318 return;
2319 }
2320 }
2321
2322 /// Lowers llvm.matrix.multiply.
2323 MatrixTy LowerMultiply(CallInst *MatMul, IRBuilder<> &Builder) {
2324 auto *EltType = cast<FixedVectorType>(Val: MatMul->getType())->getElementType();
2325 ShapeInfo LShape(MatMul->getArgOperand(i: 2), MatMul->getArgOperand(i: 3));
2326 ShapeInfo RShape(MatMul->getArgOperand(i: 3), MatMul->getArgOperand(i: 4));
2327
2328 const MatrixTy &Lhs = getMatrix(MatrixVal: MatMul->getArgOperand(i: 0), SI: LShape, Builder);
2329 const MatrixTy &Rhs = getMatrix(MatrixVal: MatMul->getArgOperand(i: 1), SI: RShape, Builder);
2330 assert(Lhs.getElementType() == Rhs.getElementType() &&
2331 "Matrix multiply argument element types do not match.");
2332
2333 const unsigned R = LShape.NumRows;
2334 const unsigned C = RShape.NumColumns;
2335 assert(LShape.NumColumns == RShape.NumRows);
2336
2337 // Initialize the output
2338 MatrixTy Result(R, C, EltType);
2339 assert(Lhs.getElementType() == Result.getElementType() &&
2340 "Matrix multiply result element type does not match arguments.");
2341
2342 emitMatrixMultiply(Result, A: Lhs, B: Rhs, Builder, IsTiled: false, IsScalarMatrixTransposed: false,
2343 FMF: getFastMathFlags(Inst: MatMul));
2344 return Result;
2345 }
2346
2347 /// Lowers llvm.matrix.transpose.
2348 MatrixTy LowerTranspose(CallInst *Inst, IRBuilder<> &Builder) {
2349 MatrixTy Result;
2350 Value *InputVal = Inst->getArgOperand(i: 0);
2351 FixedVectorType *VectorTy = cast<FixedVectorType>(Val: InputVal->getType());
2352 ShapeInfo ArgShape(Inst->getArgOperand(i: 1), Inst->getArgOperand(i: 2));
2353 MatrixTy InputMatrix = getMatrix(MatrixVal: InputVal, SI: ArgShape, Builder);
2354
2355 const unsigned NewNumVecs =
2356 InputMatrix.isColumnMajor() ? ArgShape.NumRows : ArgShape.NumColumns;
2357 const unsigned NewNumElts =
2358 InputMatrix.isColumnMajor() ? ArgShape.NumColumns : ArgShape.NumRows;
2359
2360 for (unsigned I = 0; I < NewNumVecs; ++I) {
2361 // Build a single result vector. First initialize it.
2362 Value *ResultVector = PoisonValue::get(
2363 T: FixedVectorType::get(ElementType: VectorTy->getElementType(), NumElts: NewNumElts));
2364 // Go through the old elements and insert it into the resulting vector.
2365 for (auto J : enumerate(First: InputMatrix.vectors())) {
2366 Value *Elt = Builder.CreateExtractElement(Vec: J.value(), Idx: I);
2367 // Row and column indices are transposed.
2368 ResultVector =
2369 Builder.CreateInsertElement(Vec: ResultVector, NewElt: Elt, Idx: J.index());
2370 }
2371 Result.addVector(V: ResultVector);
2372 }
2373
2374 // TODO: Improve estimate of operations needed for transposes. Currently we
2375 // just count the insertelement/extractelement instructions, but do not
2376 // account for later simplifications/combines.
2377 return Result.addNumComputeOps(N: 2 * ArgShape.NumRows * ArgShape.NumColumns)
2378 .addNumExposedTransposes(N: 1);
2379 }
2380
2381 /// Lower load instructions.
2382 MatrixTy VisitLoad(LoadInst *Inst, const ShapeInfo &SI, Value *Ptr,
2383 IRBuilder<> &Builder) {
2384 return LowerLoad(Inst, Ptr, Align: Inst->getAlign(), Stride: getIndex(Ptr, V: SI.getStride()),
2385 IsVolatile: Inst->isVolatile(), Shape: SI, Builder);
2386 }
2387
2388 MatrixTy VisitStore(StoreInst *Inst, const ShapeInfo &SI, Value *StoredVal,
2389 Value *Ptr, IRBuilder<> &Builder) {
2390 return LowerStore(Inst, Matrix: StoredVal, Ptr, A: Inst->getAlign(),
2391 Stride: getIndex(Ptr, V: SI.getStride()), IsVolatile: Inst->isVolatile(), Shape: SI,
2392 Builder);
2393 }
2394
2395 MatrixTy VisitPHI(PHINode *Inst, const ShapeInfo &SI, IRBuilder<> &Builder) {
2396 auto BlockIP = Inst->getParent()->getFirstInsertionPt();
2397 Builder.SetInsertPoint(BlockIP);
2398 MatrixTy PhiM = getMatrix(MatrixVal: Inst, SI, Builder);
2399
2400 for (auto [IncomingV, IncomingB] :
2401 llvm::zip_equal(t: Inst->incoming_values(), u: Inst->blocks())) {
2402 // getMatrix() may insert some instructions to help with reshaping. The
2403 // safest place for those is at the top of the block after the rest of the
2404 // PHI's. Even better, if we can put it in the incoming block.
2405 Builder.SetInsertPoint(BlockIP);
2406 if (auto *IncomingInst = dyn_cast<Instruction>(Val&: IncomingV))
2407 if (auto MaybeIP = IncomingInst->getInsertionPointAfterDef())
2408 Builder.SetInsertPoint(*MaybeIP);
2409
2410 MatrixTy OpM = getMatrix(MatrixVal: IncomingV, SI, Builder);
2411
2412 for (unsigned VI = 0, VE = PhiM.getNumVectors(); VI != VE; ++VI) {
2413 PHINode *NewPHI = cast<PHINode>(Val: PhiM.getVector(i: VI));
2414 NewPHI->addIncoming(V: OpM.getVector(i: VI), BB: IncomingB);
2415 }
2416 }
2417
2418 // finalizeLowering() may also insert instructions in some cases. The safe
2419 // place for those is at the end of the initial block of PHIs.
2420 Builder.SetInsertPoint(BlockIP);
2421 return PhiM;
2422 }
2423
2424 /// Lower binary operators.
2425 MatrixTy VisitBinaryOperator(BinaryOperator *Inst, const ShapeInfo &SI,
2426 IRBuilder<> &Builder) {
2427 Value *Lhs = Inst->getOperand(i_nocapture: 0);
2428 Value *Rhs = Inst->getOperand(i_nocapture: 1);
2429
2430 MatrixTy Result;
2431 MatrixTy A = getMatrix(MatrixVal: Lhs, SI, Builder);
2432 MatrixTy B = getMatrix(MatrixVal: Rhs, SI, Builder);
2433 assert(A.isColumnMajor() == B.isColumnMajor() &&
2434 Result.isColumnMajor() == A.isColumnMajor() &&
2435 "operands must agree on matrix layout");
2436
2437 Builder.setFastMathFlags(getFastMathFlags(Inst));
2438
2439 for (auto [AV, BV] : llvm::zip_equal(t: A.vectors(), u: B.vectors()))
2440 Result.addVector(V: Builder.CreateBinOp(Opc: Inst->getOpcode(), LHS: AV, RHS: BV));
2441
2442 return Result.addNumComputeOps(N: getNumOps(VT: Result.getVectorTy()) *
2443 Result.getNumVectors());
2444 }
2445
2446 /// Lower unary operators.
2447 MatrixTy VisitUnaryOperator(UnaryOperator *Inst, const ShapeInfo &SI,
2448 IRBuilder<> &Builder) {
2449 Value *Op = Inst->getOperand(i_nocapture: 0);
2450
2451 MatrixTy Result;
2452 MatrixTy M = getMatrix(MatrixVal: Op, SI, Builder);
2453
2454 Builder.setFastMathFlags(getFastMathFlags(Inst));
2455
2456 // Helper to perform unary op on vectors.
2457 auto BuildVectorOp = [&Builder, Inst](Value *Op) {
2458 switch (Inst->getOpcode()) {
2459 case Instruction::FNeg:
2460 return Builder.CreateFNeg(V: Op);
2461 default:
2462 llvm_unreachable("Unsupported unary operator for matrix");
2463 }
2464 };
2465
2466 for (auto *Vector : M.vectors())
2467 Result.addVector(V: BuildVectorOp(Vector));
2468
2469 return Result.addNumComputeOps(N: getNumOps(VT: Result.getVectorTy()) *
2470 Result.getNumVectors());
2471 }
2472
2473 /// Lower cast instructions.
2474 MatrixTy VisitCastInstruction(CastInst *Inst, const ShapeInfo &Shape,
2475 IRBuilder<> &Builder) {
2476 Value *Op = Inst->getOperand(i_nocapture: 0);
2477
2478 MatrixTy Result;
2479 MatrixTy M = getMatrix(MatrixVal: Op, SI: Shape, Builder);
2480
2481 Builder.setFastMathFlags(getFastMathFlags(Inst));
2482
2483 auto *OrigVTy = cast<VectorType>(Val: Inst->getType());
2484 auto *NewVTy = VectorType::get(ElementType: OrigVTy->getElementType(),
2485 EC: ElementCount::getFixed(MinVal: M.getStride()));
2486
2487 for (auto *Vector : M.vectors())
2488 Result.addVector(V: Builder.CreateCast(Op: Inst->getOpcode(), V: Vector, DestTy: NewVTy));
2489
2490 return Result.addNumComputeOps(N: getNumOps(VT: Result.getVectorTy()) *
2491 Result.getNumVectors());
2492 }
2493
2494 /// Lower selects.
2495 MatrixTy VisitSelectInst(SelectInst *Inst, const ShapeInfo &Shape,
2496 IRBuilder<> &Builder) {
2497 Value *Cond = Inst->getOperand(i_nocapture: 0);
2498 Value *OpA = Inst->getOperand(i_nocapture: 1);
2499 Value *OpB = Inst->getOperand(i_nocapture: 2);
2500
2501 MatrixTy Result;
2502 MatrixTy A = getMatrix(MatrixVal: OpA, SI: Shape, Builder);
2503 MatrixTy B = getMatrix(MatrixVal: OpB, SI: Shape, Builder);
2504
2505 SmallVector<Value*> CondV;
2506 Instruction *MDFrom = nullptr;
2507 if (isa<FixedVectorType>(Val: Cond->getType())) {
2508 MatrixTy C = getMatrix(MatrixVal: Cond, SI: Shape, Builder);
2509 llvm::copy(Range: C.vectors(), Out: std::back_inserter(x&: CondV));
2510 } else {
2511 CondV.resize(N: A.getNumVectors());
2512 llvm::fill(Range&: CondV, Value&: Cond);
2513 if (!ProfcheckDisableMetadataFixes)
2514 MDFrom = Inst;
2515 }
2516
2517 for (auto [CV, AV, BV] : llvm::zip_equal(t&: CondV, u: A.vectors(), args: B.vectors())) {
2518 assert(!(isa<VectorType>(CV->getType()) && static_cast<bool>(MDFrom)) &&
2519 "If we have a vector conditional, we should be propagating "
2520 "profile information.");
2521 Result.addVector(V: Builder.CreateSelect(C: CV, True: AV, False: BV, Name: "", MDFrom));
2522 }
2523
2524 return Result.addNumComputeOps(N: getNumOps(VT: Result.getVectorTy()) *
2525 Result.getNumVectors());
2526 }
2527
2528 /// Helper to linearize a matrix expression tree into a string. Currently
2529 /// matrix expressions are linarized by starting at an expression leaf and
2530 /// linearizing bottom up.
2531 struct ExprLinearizer {
2532 unsigned LengthToBreak = 100;
2533 std::string Str;
2534 raw_string_ostream Stream;
2535 unsigned LineLength = 0;
2536 const DataLayout &DL;
2537
2538 /// Mapping from instructions to matrixes. It is used to identify
2539 /// matrix instructions.
2540 const MapVector<Value *, MatrixTy> &Inst2Matrix;
2541
2542 /// Mapping from values to the leaves of all expressions that the value is
2543 /// part of.
2544 const DenseMap<Value *, SmallPtrSet<Value *, 2>> &Shared;
2545
2546 /// Set of matrix expressions in the scope of a given DISubprogram.
2547 const SmallSetVector<Value *, 32> &ExprsInSubprogram;
2548
2549 /// Leaf node of the expression to linearize.
2550 Value *Leaf;
2551
2552 /// Used to keep track of sub-expressions that get reused while linearizing
2553 /// the expression. Re-used sub-expressions are marked as (reused).
2554 SmallPtrSet<Value *, 8> ReusedExprs;
2555
2556 ExprLinearizer(const DataLayout &DL,
2557 const MapVector<Value *, MatrixTy> &Inst2Matrix,
2558 const DenseMap<Value *, SmallPtrSet<Value *, 2>> &Shared,
2559 const SmallSetVector<Value *, 32> &ExprsInSubprogram,
2560 Value *Leaf)
2561 : Stream(Str), DL(DL), Inst2Matrix(Inst2Matrix), Shared(Shared),
2562 ExprsInSubprogram(ExprsInSubprogram), Leaf(Leaf) {}
2563
2564 void indent(unsigned N) {
2565 LineLength += N;
2566 for (unsigned i = 0; i < N; i++)
2567 Stream << " ";
2568 }
2569
2570 void lineBreak() {
2571 Stream << "\n";
2572 LineLength = 0;
2573 }
2574
2575 void maybeIndent(unsigned Indent) {
2576 if (LineLength >= LengthToBreak)
2577 lineBreak();
2578
2579 if (LineLength == 0)
2580 indent(N: Indent);
2581 }
2582
2583 void write(StringRef S) {
2584 LineLength += S.size();
2585 Stream << S;
2586 }
2587
2588 Value *getUnderlyingObjectThroughLoads(Value *V) {
2589 if (Value *Ptr = getPointerOperand(V))
2590 return getUnderlyingObjectThroughLoads(V: Ptr);
2591 else if (V->getType()->isPointerTy())
2592 return getUnderlyingObject(V);
2593 return V;
2594 }
2595
2596 /// Returns true if \p V is a matrix value in the given subprogram.
2597 bool isMatrix(Value *V) const { return ExprsInSubprogram.count(key: V); }
2598
2599 /// If \p V is a matrix value, print its shape as NumRows x NumColumns to
2600 /// \p SS.
2601 void prettyPrintMatrixType(Value *V, raw_string_ostream &SS) {
2602 auto M = Inst2Matrix.find(Key: V);
2603 if (M == Inst2Matrix.end())
2604 SS << "unknown";
2605 else {
2606 SS << M->second.getNumRows();
2607 SS << "x";
2608 SS << M->second.getNumColumns();
2609 }
2610 }
2611
2612 /// Write the called function name. Handles calls to llvm.matrix.*
2613 /// specially: we write the name, followed by the dimensions of the input
2614 /// matrixes, followed by the scalar type name.
2615 void writeFnName(CallInst *CI) {
2616 if (!CI->getCalledFunction())
2617 write(S: "<no called fn>");
2618 else {
2619 StringRef Name = CI->getCalledFunction()->getName();
2620 if (!Name.starts_with(Prefix: "llvm.matrix")) {
2621 write(S: Name);
2622 return;
2623 }
2624 auto *II = cast<IntrinsicInst>(Val: CI);
2625 write(S: Intrinsic::getBaseName(id: II->getIntrinsicID())
2626 .drop_front(N: StringRef("llvm.matrix.").size()));
2627 write(S: ".");
2628 std::string Tmp;
2629 raw_string_ostream SS(Tmp);
2630
2631 switch (II->getIntrinsicID()) {
2632 case Intrinsic::matrix_multiply:
2633 prettyPrintMatrixType(V: II->getOperand(i_nocapture: 0), SS);
2634 SS << ".";
2635 prettyPrintMatrixType(V: II->getOperand(i_nocapture: 1), SS);
2636 SS << "." << *II->getType()->getScalarType();
2637 break;
2638 case Intrinsic::matrix_transpose:
2639 prettyPrintMatrixType(V: II->getOperand(i_nocapture: 0), SS);
2640 SS << "." << *II->getType()->getScalarType();
2641 break;
2642 case Intrinsic::matrix_column_major_load:
2643 prettyPrintMatrixType(V: II, SS);
2644 SS << "." << *II->getType()->getScalarType();
2645 break;
2646 case Intrinsic::matrix_column_major_store:
2647 prettyPrintMatrixType(V: II->getOperand(i_nocapture: 0), SS);
2648 SS << "." << *II->getOperand(i_nocapture: 0)->getType()->getScalarType();
2649 break;
2650 default:
2651 llvm_unreachable("Unhandled case");
2652 }
2653 write(S: Tmp);
2654 }
2655 }
2656
2657 unsigned getNumShapeArgs(CallInst *CI) const {
2658 if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(Val: CI)) {
2659 switch (II->getIntrinsicID()) {
2660 case Intrinsic::matrix_multiply:
2661 return 3;
2662 case Intrinsic::matrix_transpose:
2663 return 2;
2664 case Intrinsic::matrix_column_major_load:
2665 case Intrinsic::matrix_column_major_store:
2666 return 3;
2667 default:
2668 return 0;
2669 }
2670 }
2671 return 0;
2672 }
2673
2674 /// Special printing for values: for pointers, we print if they refer to an
2675 /// (function) external address or a stack address, for other values we
2676 /// either print the constant or "scalar"/"matrix" for other values.
2677 void write(Value *V) {
2678 V = getUnderlyingObjectThroughLoads(V);
2679 if (V->getType()->isPointerTy()) {
2680 if (isa<AllocaInst>(Val: V)) {
2681 Stream << "stack addr";
2682 LineLength += StringRef("stack addr").size();
2683 } else {
2684 Stream << "addr";
2685 LineLength += StringRef("addr").size();
2686 }
2687 if (!V->getName().empty()) {
2688 Stream << " %" << V->getName() << "";
2689 LineLength += V->getName().size() + 2;
2690 }
2691 return;
2692 }
2693
2694 std::string Tmp;
2695 raw_string_ostream TmpStream(Tmp);
2696
2697 if (auto *CI = dyn_cast<ConstantInt>(Val: V))
2698 TmpStream << CI->getValue();
2699 else if (isa<Constant>(Val: V))
2700 TmpStream << "constant";
2701 else {
2702 if (isMatrix(V))
2703 TmpStream << "matrix";
2704 else
2705 TmpStream << "scalar";
2706 }
2707 Tmp = std::string(StringRef(Tmp).trim());
2708 LineLength += Tmp.size();
2709 Stream << Tmp;
2710 }
2711
2712 /// Linearize expression \p Expr starting at an indentation of \p Indent.
2713 /// Expressions that are re-used multiple times are prefixed with (reused)
2714 /// at the re-used root instruction.
2715 void linearizeExpr(Value *Expr, unsigned Indent, bool ParentReused,
2716 bool ParentShared) {
2717 auto *I = cast<Instruction>(Val: Expr);
2718 maybeIndent(Indent);
2719 SmallVector<Value *, 8> Ops;
2720
2721 // Is Expr shared with other expression leaves?
2722 bool ExprShared = false;
2723
2724 // Deal with shared subtrees. Mark them as shared, if required.
2725 if (!ParentShared) {
2726 auto SI = Shared.find(Val: Expr);
2727 assert(SI != Shared.end() && SI->second.count(Leaf));
2728
2729 for (Value *S : SI->second) {
2730 if (S == Leaf)
2731 continue;
2732 DebugLoc DL = cast<Instruction>(Val: S)->getDebugLoc();
2733 write(S: "shared with remark at line " + std::to_string(val: DL.getLine()) +
2734 " column " + std::to_string(val: DL.getCol()) + " (");
2735 }
2736 ExprShared = SI->second.size() > 1;
2737 }
2738
2739 bool Reused = !ReusedExprs.insert(Ptr: Expr).second;
2740 if (Reused && !ParentReused)
2741 write(S: "(reused) ");
2742
2743 if (auto *CI = dyn_cast<CallInst>(Val: I)) {
2744 writeFnName(CI);
2745
2746 Ops.append(in_start: CI->arg_begin(), in_end: CI->arg_end() - getNumShapeArgs(CI));
2747 } else if (isa<BitCastInst>(Val: Expr)) {
2748 // Special case bitcasts, which are used to materialize matrixes from
2749 // non-matrix ops.
2750 write(S: "matrix");
2751 return;
2752 } else {
2753 Ops.append(in_start: I->value_op_begin(), in_end: I->value_op_end());
2754 write(S: I->getOpcodeName());
2755 }
2756
2757 write(S: "(");
2758
2759 unsigned NumOpsToBreak = 1;
2760 if (match(V: Expr, P: m_Intrinsic<Intrinsic::matrix_column_major_load>()))
2761 NumOpsToBreak = 2;
2762
2763 for (Value *Op : Ops) {
2764 if (Ops.size() > NumOpsToBreak)
2765 lineBreak();
2766
2767 maybeIndent(Indent: Indent + 1);
2768 if (isMatrix(V: Op))
2769 linearizeExpr(Expr: Op, Indent: Indent + 1, ParentReused: Reused, ParentShared: ExprShared);
2770 else
2771 write(V: Op);
2772 if (Op != Ops.back())
2773 write(S: ", ");
2774 }
2775
2776 write(S: ")");
2777 }
2778
2779 const std::string &getResult() {
2780 return Str;
2781 }
2782 };
2783
2784 /// Generate remarks for matrix operations in a function. To generate remarks
2785 /// for matrix expressions, the following approach is used:
2786 /// 1. Use the inlined-at debug information to group matrix operations to the
2787 /// DISubprograms they are contained in.
2788 /// 2. Collect leaves of matrix expressions (done in
2789 /// RemarkGenerator::getExpressionLeaves) for each subprogram - expression
2790 // mapping. Leaves are lowered matrix instructions without other matrix
2791 // users (like stores) in the current subprogram.
2792 /// 3. For each leaf, create a remark containing a linearizied version of the
2793 /// matrix expression. The expression is linearized by a recursive
2794 /// bottom-up traversal of the matrix operands, starting at a leaf. Note
2795 /// that multiple leaves can share sub-expressions. Shared subexpressions
2796 /// are explicitly marked as shared().
2797 struct RemarkGenerator {
2798 const MapVector<Value *, MatrixTy> &Inst2Matrix;
2799 OptimizationRemarkEmitter &ORE;
2800 Function &Func;
2801 const DataLayout &DL;
2802
2803 RemarkGenerator(const MapVector<Value *, MatrixTy> &Inst2Matrix,
2804 OptimizationRemarkEmitter &ORE, Function &Func)
2805 : Inst2Matrix(Inst2Matrix), ORE(ORE), Func(Func),
2806 DL(Func.getDataLayout()) {}
2807
2808 /// Return all leaves of the expressions in \p ExprsInSubprogram. Those are
2809 /// instructions in Inst2Matrix returning void or without any users in
2810 /// \p ExprsInSubprogram. Currently that should only include stores.
2811 SmallVector<Value *, 4>
2812 getExpressionLeaves(const SmallSetVector<Value *, 32> &ExprsInSubprogram) {
2813 SmallVector<Value *, 4> Leaves;
2814 for (auto *Expr : ExprsInSubprogram)
2815 if (Expr->getType()->isVoidTy() ||
2816 !any_of(Range: Expr->users(), P: [&ExprsInSubprogram](User *U) {
2817 return ExprsInSubprogram.count(key: U);
2818 }))
2819 Leaves.push_back(Elt: Expr);
2820 return Leaves;
2821 }
2822
2823 /// Recursively traverse expression \p V starting at \p Leaf and add \p Leaf
2824 /// to all visited expressions in \p Shared. Limit the matrix operations to
2825 /// the ones in \p ExprsInSubprogram.
2826 void collectSharedInfo(Value *Leaf, Value *V,
2827 const SmallSetVector<Value *, 32> &ExprsInSubprogram,
2828 DenseMap<Value *, SmallPtrSet<Value *, 2>> &Shared) {
2829
2830 if (!ExprsInSubprogram.count(key: V))
2831 return;
2832
2833 Shared[V].insert(Ptr: Leaf);
2834
2835 for (Value *Op : cast<Instruction>(Val: V)->operand_values())
2836 collectSharedInfo(Leaf, V: Op, ExprsInSubprogram, Shared);
2837 }
2838
2839 /// Calculate the number of exclusive and shared op counts for expression
2840 /// starting at \p V. Expressions used multiple times are counted once.
2841 /// Limit the matrix operations to the ones in \p ExprsInSubprogram.
2842 std::pair<OpInfoTy, OpInfoTy>
2843 sumOpInfos(Value *Root, SmallPtrSetImpl<Value *> &ReusedExprs,
2844 const SmallSetVector<Value *, 32> &ExprsInSubprogram,
2845 DenseMap<Value *, SmallPtrSet<Value *, 2>> &Shared) const {
2846 if (!ExprsInSubprogram.count(key: Root))
2847 return {};
2848
2849 // Already counted this expression. Stop.
2850 if (!ReusedExprs.insert(Ptr: Root).second)
2851 return {};
2852
2853 OpInfoTy SharedCount;
2854 OpInfoTy Count;
2855
2856 auto I = Shared.find(Val: Root);
2857 auto CM = Inst2Matrix.find(Key: Root);
2858 if (I->second.size() == 1)
2859 Count = CM->second.getOpInfo();
2860 else
2861 SharedCount = CM->second.getOpInfo();
2862
2863 for (Value *Op : cast<Instruction>(Val: Root)->operand_values()) {
2864 auto C = sumOpInfos(Root: Op, ReusedExprs, ExprsInSubprogram, Shared);
2865 Count += C.first;
2866 SharedCount += C.second;
2867 }
2868 return {Count, SharedCount};
2869 }
2870
2871 void emitRemarks() {
2872 if (!ORE.allowExtraAnalysis(DEBUG_TYPE))
2873 return;
2874
2875 // Map matrix operations to their containting subprograms, by traversing
2876 // the inlinedAt chain. If the function does not have a DISubprogram, we
2877 // only map them to the containing function.
2878 MapVector<DISubprogram *, SmallVector<Value *, 8>> Subprog2Exprs;
2879 for (const auto &KV : Inst2Matrix) {
2880 if (Func.getSubprogram()) {
2881 auto *I = cast<Instruction>(Val: KV.first);
2882 DILocation *Context = I->getDebugLoc();
2883 while (Context) {
2884 Subprog2Exprs[getSubprogram(Scope: Context->getScope())].push_back(
2885 Elt: KV.first);
2886 Context = DebugLoc(Context).getInlinedAt();
2887 }
2888 } else {
2889 Subprog2Exprs[nullptr].push_back(Elt: KV.first);
2890 }
2891 }
2892 for (auto &KV : Subprog2Exprs) {
2893 SmallSetVector<Value *, 32> ExprsInSubprogram(KV.second.begin(),
2894 KV.second.end());
2895 auto Leaves = getExpressionLeaves(ExprsInSubprogram);
2896
2897 DenseMap<Value *, SmallPtrSet<Value *, 2>> Shared;
2898 for (Value *Leaf : Leaves)
2899 collectSharedInfo(Leaf, V: Leaf, ExprsInSubprogram, Shared);
2900
2901 // Generate remarks for each leaf.
2902 for (auto *L : Leaves) {
2903
2904 DebugLoc Loc = cast<Instruction>(Val: L)->getDebugLoc();
2905 DILocation *Context = cast<Instruction>(Val: L)->getDebugLoc();
2906 while (Context) {
2907 if (getSubprogram(Scope: Context->getScope()) == KV.first) {
2908 Loc = Context;
2909 break;
2910 }
2911 Context = DebugLoc(Context).getInlinedAt();
2912 }
2913
2914 SmallPtrSet<Value *, 8> ReusedExprs;
2915 OpInfoTy Counts, SharedCounts;
2916 std::tie(args&: Counts, args&: SharedCounts) =
2917 sumOpInfos(Root: L, ReusedExprs, ExprsInSubprogram, Shared);
2918
2919 OptimizationRemark Rem(DEBUG_TYPE, "matrix-lowered", Loc,
2920 cast<Instruction>(Val: L)->getParent());
2921
2922 Rem << "Lowered with ";
2923 Rem << ore::NV("NumStores", Counts.NumStores) << " stores, "
2924 << ore::NV("NumLoads", Counts.NumLoads) << " loads, "
2925 << ore::NV("NumComputeOps", Counts.NumComputeOps)
2926 << " compute ops, "
2927 << ore::NV("NumExposedTransposes", Counts.NumExposedTransposes)
2928 << " exposed transposes";
2929
2930 if (SharedCounts.NumStores > 0 || SharedCounts.NumLoads > 0 ||
2931 SharedCounts.NumComputeOps > 0) {
2932 Rem << ",\nadditionally "
2933 << ore::NV("NumStores", SharedCounts.NumStores) << " stores, "
2934 << ore::NV("NumLoads", SharedCounts.NumLoads) << " loads, "
2935 << ore::NV("NumFPOps", SharedCounts.NumComputeOps)
2936 << " compute ops"
2937 << " are shared with other expressions";
2938 }
2939
2940 Rem << ("\n" + linearize(L, Shared, ExprsInSubprogram, DL));
2941 ORE.emit(OptDiag&: Rem);
2942 }
2943 }
2944 }
2945
2946 std::string
2947 linearize(Value *L,
2948 const DenseMap<Value *, SmallPtrSet<Value *, 2>> &Shared,
2949 const SmallSetVector<Value *, 32> &ExprsInSubprogram,
2950 const DataLayout &DL) {
2951 ExprLinearizer Lin(DL, Inst2Matrix, Shared, ExprsInSubprogram, L);
2952 Lin.linearizeExpr(Expr: L, Indent: 0, ParentReused: false, ParentShared: false);
2953 return Lin.getResult();
2954 }
2955 };
2956};
2957} // namespace
2958
2959PreservedAnalyses LowerMatrixIntrinsicsPass::run(Function &F,
2960 FunctionAnalysisManager &AM) {
2961 auto &TTI = AM.getResult<TargetIRAnalysis>(IR&: F);
2962
2963 LowerMatrixIntrinsics LMT(F, TTI, Minimal ? nullptr : &AM);
2964 if (LMT.Visit()) {
2965 PreservedAnalyses PA;
2966 if (!Minimal) {
2967 PA.preserve<LoopAnalysis>();
2968 PA.preserve<DominatorTreeAnalysis>();
2969 }
2970 return PA;
2971 }
2972 return PreservedAnalyses::all();
2973}
2974
2975void LowerMatrixIntrinsicsPass::printPipeline(
2976 raw_ostream &OS, function_ref<StringRef(StringRef)> MapClassName2PassName) {
2977 static_cast<PassInfoMixin<LowerMatrixIntrinsicsPass> *>(this)->printPipeline(
2978 OS, MapClassName2PassName);
2979 OS << '<';
2980 if (Minimal)
2981 OS << "minimal";
2982 OS << '>';
2983}
2984