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