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