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