Bug Summary

File:llvm/lib/Transforms/Scalar/LowerMatrixIntrinsics.cpp
Warning:line 1900, column 38
Called C++ object pointer is null

Annotated Source Code

Press '?' to see keyboard shortcuts

clang -cc1 -cc1 -triple x86_64-pc-linux-gnu -analyze -disable-free -disable-llvm-verifier -discard-value-names -main-file-name LowerMatrixIntrinsics.cpp -analyzer-store=region -analyzer-opt-analyze-nested-blocks -analyzer-checker=core -analyzer-checker=apiModeling -analyzer-checker=unix -analyzer-checker=deadcode -analyzer-checker=cplusplus -analyzer-checker=security.insecureAPI.UncheckedReturn -analyzer-checker=security.insecureAPI.getpw -analyzer-checker=security.insecureAPI.gets -analyzer-checker=security.insecureAPI.mktemp -analyzer-checker=security.insecureAPI.mkstemp -analyzer-checker=security.insecureAPI.vfork -analyzer-checker=nullability.NullPassedToNonnull -analyzer-checker=nullability.NullReturnedFromNonnull -analyzer-output plist -w -setup-static-analyzer -analyzer-config-compatibility-mode=true -mrelocation-model pic -pic-level 2 -mframe-pointer=none -fmath-errno -fno-rounding-math -mconstructor-aliases -munwind-tables -target-cpu x86-64 -tune-cpu generic -debugger-tuning=gdb -ffunction-sections -fdata-sections -fcoverage-compilation-dir=/build/llvm-toolchain-snapshot-14~++20210903100615+fd66b44ec19e/build-llvm/lib/Transforms/Scalar -resource-dir /usr/lib/llvm-14/lib/clang/14.0.0 -D _GNU_SOURCE -D __STDC_CONSTANT_MACROS -D __STDC_FORMAT_MACROS -D __STDC_LIMIT_MACROS -I /build/llvm-toolchain-snapshot-14~++20210903100615+fd66b44ec19e/build-llvm/lib/Transforms/Scalar -I /build/llvm-toolchain-snapshot-14~++20210903100615+fd66b44ec19e/llvm/lib/Transforms/Scalar -I /build/llvm-toolchain-snapshot-14~++20210903100615+fd66b44ec19e/build-llvm/include -I /build/llvm-toolchain-snapshot-14~++20210903100615+fd66b44ec19e/llvm/include -D NDEBUG -internal-isystem /usr/lib/gcc/x86_64-linux-gnu/10/../../../../include/c++/10 -internal-isystem /usr/lib/gcc/x86_64-linux-gnu/10/../../../../include/x86_64-linux-gnu/c++/10 -internal-isystem /usr/lib/gcc/x86_64-linux-gnu/10/../../../../include/c++/10/backward -internal-isystem /usr/lib/llvm-14/lib/clang/14.0.0/include -internal-isystem /usr/local/include -internal-isystem /usr/lib/gcc/x86_64-linux-gnu/10/../../../../x86_64-linux-gnu/include -internal-externc-isystem /usr/include/x86_64-linux-gnu -internal-externc-isystem /include -internal-externc-isystem /usr/include -O2 -Wno-unused-parameter -Wwrite-strings -Wno-missing-field-initializers -Wno-long-long -Wno-maybe-uninitialized -Wno-class-memaccess -Wno-redundant-move -Wno-pessimizing-move -Wno-noexcept-type -Wno-comment -std=c++14 -fdeprecated-macro -fdebug-compilation-dir=/build/llvm-toolchain-snapshot-14~++20210903100615+fd66b44ec19e/build-llvm/lib/Transforms/Scalar -fdebug-prefix-map=/build/llvm-toolchain-snapshot-14~++20210903100615+fd66b44ec19e=. -ferror-limit 19 -fvisibility-inlines-hidden -stack-protector 2 -fgnuc-version=4.2.1 -vectorize-loops -vectorize-slp -analyzer-output=html -analyzer-config stable-report-filename=true -faddrsig -D__GCC_HAVE_DWARF2_CFI_ASM=1 -o /tmp/scan-build-2021-09-04-040900-46481-1 -x c++ /build/llvm-toolchain-snapshot-14~++20210903100615+fd66b44ec19e/llvm/lib/Transforms/Scalar/LowerMatrixIntrinsics.cpp
1//===- LowerMatrixIntrinsics.cpp - Lower matrix intrinsics -----*- C++ -*-===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7//===----------------------------------------------------------------------===//
8//
9// Lower matrix intrinsics to vector operations.
10//
11// TODO:
12// * Improve fusion:
13// * Support more cases, e.g. multiply-add, multiply-sub, operands/results
14// transposed.
15// * Improve cost-modeling, e.g. choose different number of rows/columns
16// columns for tiles, consider cost of copies on alias.
17//
18//===----------------------------------------------------------------------===//
19
20#include "llvm/Transforms/Scalar/LowerMatrixIntrinsics.h"
21#include "llvm/ADT/GraphTraits.h"
22#include "llvm/ADT/PostOrderIterator.h"
23#include "llvm/ADT/SmallVector.h"
24#include "llvm/Analysis/AliasAnalysis.h"
25#include "llvm/Analysis/DomTreeUpdater.h"
26#include "llvm/Analysis/OptimizationRemarkEmitter.h"
27#include "llvm/Analysis/TargetTransformInfo.h"
28#include "llvm/Analysis/ValueTracking.h"
29#include "llvm/Analysis/VectorUtils.h"
30#include "llvm/IR/CFG.h"
31#include "llvm/IR/DataLayout.h"
32#include "llvm/IR/DebugInfoMetadata.h"
33#include "llvm/IR/Function.h"
34#include "llvm/IR/IRBuilder.h"
35#include "llvm/IR/Instructions.h"
36#include "llvm/IR/IntrinsicInst.h"
37#include "llvm/IR/MatrixBuilder.h"
38#include "llvm/IR/PatternMatch.h"
39#include "llvm/InitializePasses.h"
40#include "llvm/Pass.h"
41#include "llvm/Support/Alignment.h"
42#include "llvm/Support/CommandLine.h"
43#include "llvm/Support/Debug.h"
44#include "llvm/Transforms/Scalar.h"
45#include "llvm/Transforms/Utils/BasicBlockUtils.h"
46#include "llvm/Transforms/Utils/LoopUtils.h"
47#include "llvm/Transforms/Utils/MatrixUtils.h"
48
49using namespace llvm;
50using namespace PatternMatch;
51
52#define DEBUG_TYPE"lower-matrix-intrinsics" "lower-matrix-intrinsics"
53
54static cl::opt<bool>
55 FuseMatrix("fuse-matrix", cl::init(true), cl::Hidden,
56 cl::desc("Enable/disable fusing matrix instructions."));
57// TODO: Allow and use non-square tiles.
58static cl::opt<unsigned> TileSize(
59 "fuse-matrix-tile-size", cl::init(4), cl::Hidden,
60 cl::desc(
61 "Tile size for matrix instruction fusion using square-shaped tiles."));
62static cl::opt<bool> TileUseLoops("fuse-matrix-use-loops", cl::init(false),
63 cl::Hidden,
64 cl::desc("Generate loop nest for tiling."));
65static cl::opt<bool> ForceFusion(
66 "force-fuse-matrix", cl::init(false), cl::Hidden,
67 cl::desc("Force matrix instruction fusion even if not profitable."));
68static cl::opt<bool> AllowContractEnabled(
69 "matrix-allow-contract", cl::init(false), cl::Hidden,
70 cl::desc("Allow the use of FMAs if available and profitable. This may "
71 "result in different results, due to less rounding error."));
72
73enum class MatrixLayoutTy { ColumnMajor, RowMajor };
74
75static cl::opt<MatrixLayoutTy> MatrixLayout(
76 "matrix-default-layout", cl::init(MatrixLayoutTy::ColumnMajor),
77 cl::desc("Sets the default matrix layout"),
78 cl::values(clEnumValN(MatrixLayoutTy::ColumnMajor, "column-major",llvm::cl::OptionEnumValue { "column-major", int(MatrixLayoutTy
::ColumnMajor), "Use column-major layout" }
79 "Use column-major layout")llvm::cl::OptionEnumValue { "column-major", int(MatrixLayoutTy
::ColumnMajor), "Use column-major layout" }
,
80 clEnumValN(MatrixLayoutTy::RowMajor, "row-major",llvm::cl::OptionEnumValue { "row-major", int(MatrixLayoutTy::
RowMajor), "Use row-major layout" }
81 "Use row-major layout")llvm::cl::OptionEnumValue { "row-major", int(MatrixLayoutTy::
RowMajor), "Use row-major layout" }
));
82
83/// Helper function to either return Scope, if it is a subprogram or the
84/// attached subprogram for a local scope.
85static DISubprogram *getSubprogram(DIScope *Scope) {
86 if (auto *Subprogram = dyn_cast<DISubprogram>(Scope))
87 return Subprogram;
88 return cast<DILocalScope>(Scope)->getSubprogram();
89}
90
91namespace {
92
93// Given an element pointer \p BasePtr to the start of a (sub) matrix, compute
94// the start address of vector \p VecIdx with type (\p EltType x \p NumElements)
95// assuming \p Stride elements between start two consecutive vectors.
96// \p Stride must be >= \p NumElements.
97// For column-major matrixes, the function computes the address of a column
98// vectors and \p NumElements must be set to the number of elements in a column
99// (= number of rows of the matrix). For row-major matrixes, the function
100// computes the address of a row vector and \p NumElements must be set to the
101// number of elements in a column (= number of columns of the matrix).
102//
103// Consider a 4x4 matrix in column-mjaor layout like below
104//
105// 0 1 2 3
106// 0 v_0_0 v_0_1 v_0_2 v_0_3
107// 1 v_1_0 v_1_1 v_1_2 v_1_3
108// 2 v_2_0 v_2_1 v_2_2 v_2_3
109// 3 v_3_0 v_3_1 v_3_2 v_3_3
110
111// To compute the column addresses for a 2x3 sub-matrix at row 1 and column 1,
112// we need a pointer to the first element of the submatrix as base pointer.
113// Then we can use computeVectorAddr to compute the addresses for the columns
114// of the sub-matrix.
115//
116// Column 0: computeVectorAddr(Base, 0 (column), 4 (stride), 2 (num rows), ..)
117// -> just returns Base
118// Column 1: computeVectorAddr(Base, 1 (column), 4 (stride), 2 (num rows), ..)
119// -> returns Base + (1 * 4)
120// Column 2: computeVectorAddr(Base, 2 (column), 4 (stride), 2 (num rows), ..)
121// -> returns Base + (2 * 4)
122//
123// The graphic below illustrates the number of elements in a column (marked
124// with |) and the number of skipped elements (marked with }).
125//
126// v_0_0 v_0_1 {v_0_2 {v_0_3
127// Base Col 1 Col 2
128// | | |
129// v_1_0 |v_1_1 |v_1_2 |v_1_3
130// v_2_0 |v_2_1 |v_2_2 |v_2_3
131// v_3_0 {v_3_1 {v_3_2 v_3_3
132//
133Value *computeVectorAddr(Value *BasePtr, Value *VecIdx, Value *Stride,
134 unsigned NumElements, Type *EltType,
135 IRBuilder<> &Builder) {
136
137 assert((!isa<ConstantInt>(Stride) ||(static_cast<void> (0))
138 cast<ConstantInt>(Stride)->getZExtValue() >= NumElements) &&(static_cast<void> (0))
139 "Stride must be >= the number of elements in the result vector.")(static_cast<void> (0));
140 unsigned AS = cast<PointerType>(BasePtr->getType())->getAddressSpace();
141
142 // Compute the start of the vector with index VecIdx as VecIdx * Stride.
143 Value *VecStart = Builder.CreateMul(VecIdx, Stride, "vec.start");
144
145 // Get pointer to the start of the selected vector. Skip GEP creation,
146 // if we select vector 0.
147 if (isa<ConstantInt>(VecStart) && cast<ConstantInt>(VecStart)->isZero())
148 VecStart = BasePtr;
149 else
150 VecStart = Builder.CreateGEP(EltType, BasePtr, VecStart, "vec.gep");
151
152 // Cast elementwise vector start pointer to a pointer to a vector
153 // (EltType x NumElements)*.
154 auto *VecType = FixedVectorType::get(EltType, NumElements);
155 Type *VecPtrType = PointerType::get(VecType, AS);
156 return Builder.CreatePointerCast(VecStart, VecPtrType, "vec.cast");
157}
158
159/// LowerMatrixIntrinsics contains the methods used to lower matrix intrinsics.
160///
161/// Currently, the lowering for each matrix intrinsic is done as follows:
162/// 1. Propagate the shape information from intrinsics to connected
163/// instructions.
164/// 2. Lower instructions with shape information (assuming column-major layout).
165/// The lowering works similarly using row-major layout.
166/// 2.1. Get column vectors for each argument. If we already lowered the
167/// definition of an argument, use the produced column vectors directly.
168/// If not, split the operand vector containing an embedded matrix into
169/// a set of column vectors,
170/// 2.2. Lower the instruction in terms of column major operations, which
171/// yields a set of column vectors containing result matrix. Note that we
172/// lower all instructions that have shape information. Besides the
173/// intrinsics, this includes stores for example.
174/// 2.3. Update uses of the lowered instruction. If we have shape information
175/// for a user, there is nothing to do, as we will look up the result
176/// column matrix when lowering the user. For other uses, we embed the
177/// result matrix in a flat vector and update the use.
178/// 2.4. Cache the result column matrix for the instruction we lowered
179/// 3. After we lowered all instructions in a function, remove the now
180/// obsolete instructions.
181///
182class LowerMatrixIntrinsics {
183 Function &Func;
184 const DataLayout &DL;
185 const TargetTransformInfo &TTI;
186 AliasAnalysis *AA;
187 DominatorTree *DT;
188 LoopInfo *LI;
189 OptimizationRemarkEmitter *ORE;
190
191 /// Contains estimates of the number of operations (loads, stores, compute) required to lower a matrix operation.
192 struct OpInfoTy {
193 /// Number of stores emitted to generate this matrix.
194 unsigned NumStores = 0;
195 /// Number of loads emitted to generate this matrix.
196 unsigned NumLoads = 0;
197 /// Number of compute operations emitted to generate this matrix.
198 unsigned NumComputeOps = 0;
199 /// Most of the time transposes can be fused with matrix multiplies or can
200 /// be folded away via algebraic simplifications. This is the number of
201 /// transposes that we failed to make "free" via such optimizations.
202 unsigned NumExposedTransposes = 0;
203
204 OpInfoTy &operator+=(const OpInfoTy &RHS) {
205 NumStores += RHS.NumStores;
206 NumLoads += RHS.NumLoads;
207 NumComputeOps += RHS.NumComputeOps;
208 NumExposedTransposes += RHS.NumExposedTransposes;
209 return *this;
210 }
211 };
212
213 /// Wrapper class representing a matrix as a set of vectors, either in row or
214 /// column major layout. All vectors must have the same vector type.
215 class MatrixTy {
216 SmallVector<Value *, 16> Vectors;
217
218 OpInfoTy OpInfo;
219
220 bool IsColumnMajor = true;
221
222 public:
223 MatrixTy()
224 : Vectors(),
225 IsColumnMajor(MatrixLayout == MatrixLayoutTy::ColumnMajor) {}
226 MatrixTy(ArrayRef<Value *> Vectors)
227 : Vectors(Vectors.begin(), Vectors.end()),
228 IsColumnMajor(MatrixLayout == MatrixLayoutTy::ColumnMajor) {}
229 MatrixTy(unsigned NumRows, unsigned NumColumns, Type *EltTy)
230 : IsColumnMajor(MatrixLayout == MatrixLayoutTy::ColumnMajor) {
231
232 unsigned D = isColumnMajor() ? NumColumns : NumRows;
233 for (unsigned J = 0; J < D; ++J)
234 addVector(UndefValue::get(FixedVectorType::get(
235 EltTy, isColumnMajor() ? NumRows : NumColumns)));
236 }
237
238 Value *getVector(unsigned i) const { return Vectors[i]; }
239 Value *getColumn(unsigned i) const {
240 assert(isColumnMajor() && "only supported for column-major matrixes")(static_cast<void> (0));
241 return Vectors[i];
242 }
243 Value *getRow(unsigned i) const {
244 assert(!isColumnMajor() && "only supported for row-major matrixes")(static_cast<void> (0));
245 return Vectors[i];
246 }
247
248 void setVector(unsigned i, Value *V) { Vectors[i] = V; }
249
250 Type *getElementType() const { return getVectorTy()->getElementType(); }
251
252 unsigned getNumVectors() const {
253 if (isColumnMajor())
254 return getNumColumns();
255 return getNumRows();
256 }
257
258 unsigned getNumColumns() const {
259 if (isColumnMajor())
260 return Vectors.size();
261 else {
262 assert(Vectors.size() > 0 && "Cannot call getNumRows without columns")(static_cast<void> (0));
263 return cast<FixedVectorType>(Vectors[0]->getType())->getNumElements();
264 }
265 }
266 unsigned getNumRows() const {
267 if (isColumnMajor()) {
268 assert(Vectors.size() > 0 && "Cannot call getNumRows without columns")(static_cast<void> (0));
269 return cast<FixedVectorType>(Vectors[0]->getType())->getNumElements();
270 } else
271 return Vectors.size();
272 }
273
274 void addVector(Value *V) { Vectors.push_back(V); }
275 VectorType *getColumnTy() {
276 assert(isColumnMajor() && "only supported for column-major matrixes")(static_cast<void> (0));
277 return getVectorTy();
278 }
279
280 VectorType *getVectorTy() const {
281 return cast<VectorType>(Vectors[0]->getType());
282 }
283
284 iterator_range<SmallVector<Value *, 8>::iterator> columns() {
285 assert(isColumnMajor() &&(static_cast<void> (0))
286 "columns() only supported for column-major matrixes")(static_cast<void> (0));
287 return make_range(Vectors.begin(), Vectors.end());
288 }
289
290 iterator_range<SmallVector<Value *, 8>::iterator> vectors() {
291 return make_range(Vectors.begin(), Vectors.end());
292 }
293
294 /// Embed the vectors of the matrix into a flat vector by concatenating
295 /// them.
296 Value *embedInVector(IRBuilder<> &Builder) const {
297 return Vectors.size() == 1 ? Vectors[0]
298 : concatenateVectors(Builder, Vectors);
299 }
300
301 MatrixTy &addNumLoads(unsigned N) {
302 OpInfo.NumLoads += N;
303 return *this;
304 }
305
306 void setNumLoads(unsigned N) { OpInfo.NumLoads = N; }
307
308 MatrixTy &addNumStores(unsigned N) {
309 OpInfo.NumStores += N;
310 return *this;
311 }
312
313 MatrixTy &addNumExposedTransposes(unsigned N) {
314 OpInfo.NumExposedTransposes += N;
315 return *this;
316 }
317
318 MatrixTy &addNumComputeOps(unsigned N) {
319 OpInfo.NumComputeOps += N;
320 return *this;
321 }
322
323 unsigned getNumStores() const { return OpInfo.NumStores; }
324 unsigned getNumLoads() const { return OpInfo.NumLoads; }
325 unsigned getNumComputeOps() const { return OpInfo.NumComputeOps; }
326
327 const OpInfoTy &getOpInfo() const { return OpInfo; }
328
329 bool isColumnMajor() const { return IsColumnMajor; }
330
331 unsigned getStride() const {
332 if (isColumnMajor())
333 return getNumRows();
334 return getNumColumns();
335 }
336
337 /// Extract a vector of \p NumElts starting at index (\p I, \p J). If the
338 /// matrix is column-major, the result vector is extracted from a column
339 /// vector, otherwise from a row vector.
340 Value *extractVector(unsigned I, unsigned J, unsigned NumElts,
341 IRBuilder<> &Builder) const {
342 Value *Vec = isColumnMajor() ? getColumn(J) : getRow(I);
343 return Builder.CreateShuffleVector(
344 Vec, createSequentialMask(isColumnMajor() ? I : J, NumElts, 0),
345 "block");
346 }
347 };
348
349 struct ShapeInfo {
350 unsigned NumRows;
351 unsigned NumColumns;
352
353 bool IsColumnMajor;
354
355 ShapeInfo(unsigned NumRows = 0, unsigned NumColumns = 0)
356 : NumRows(NumRows), NumColumns(NumColumns),
357 IsColumnMajor(MatrixLayout == MatrixLayoutTy::ColumnMajor) {}
358
359 ShapeInfo(Value *NumRows, Value *NumColumns)
360 : ShapeInfo(cast<ConstantInt>(NumRows)->getZExtValue(),
361 cast<ConstantInt>(NumColumns)->getZExtValue()) {}
362
363 bool operator==(const ShapeInfo &other) {
364 return NumRows == other.NumRows && NumColumns == other.NumColumns;
365 }
366 bool operator!=(const ShapeInfo &other) { return !(*this == other); }
367
368 /// Returns true if shape-information is defined, meaning both dimensions
369 /// are != 0.
370 operator bool() const {
371 assert(NumRows == 0 || NumColumns != 0)(static_cast<void> (0));
372 return NumRows != 0;
373 }
374
375 unsigned getStride() const {
376 if (IsColumnMajor)
377 return NumRows;
378 return NumColumns;
379 }
380
381 unsigned getNumVectors() const {
382 if (IsColumnMajor)
383 return NumColumns;
384 return NumRows;
385 }
386 };
387
388 /// Maps instructions to their shape information. The shape information
389 /// describes the shape to be used while lowering. This matches the shape of
390 /// the result value of the instruction, with the only exceptions being store
391 /// instructions and the matrix_column_major_store intrinsics. For those, the
392 /// shape information indicates that those instructions should be lowered
393 /// using shape information as well. A ValueMap is used so that when
394 /// sub-passes like optimizeTransposes performs RAUW the map stays
395 /// up-to-date.
396 ValueMap<Value *, ShapeInfo> ShapeMap;
397
398 /// List of instructions to remove. While lowering, we are not replacing all
399 /// users of a lowered instruction, if shape information is available and
400 /// those need to be removed after we finished lowering.
401 SmallVector<Instruction *, 16> ToRemove;
402
403 /// Map from instructions to their produced column matrix.
404 MapVector<Value *, MatrixTy> Inst2ColumnMatrix;
405
406private:
407 static FastMathFlags getFastMathFlags(Instruction *Inst) {
408 FastMathFlags FMF;
409
410 if (isa<FPMathOperator>(*Inst))
411 FMF = Inst->getFastMathFlags();
412
413 FMF.setAllowContract(AllowContractEnabled || FMF.allowContract());
414
415 return FMF;
416 }
417
418public:
419 LowerMatrixIntrinsics(Function &F, TargetTransformInfo &TTI,
420 AliasAnalysis *AA, DominatorTree *DT, LoopInfo *LI,
421 OptimizationRemarkEmitter *ORE)
422 : Func(F), DL(F.getParent()->getDataLayout()), TTI(TTI), AA(AA), DT(DT),
423 LI(LI), ORE(ORE) {}
424
425 unsigned getNumOps(Type *VT) {
426 assert(isa<VectorType>(VT) && "Expected vector type")(static_cast<void> (0));
427 return getNumOps(VT->getScalarType(),
428 cast<FixedVectorType>(VT)->getNumElements());
429 }
430
431 /// Is this the minimal version executed in the backend pipelines.
432 bool isMinimal() const {
433 return !DT;
434 }
435
436 /// Return the estimated number of vector ops required for an operation on
437 /// \p VT * N.
438 unsigned getNumOps(Type *ST, unsigned N) {
439 return std::ceil((ST->getPrimitiveSizeInBits() * N).getFixedSize() /
440 double(TTI.getRegisterBitWidth(
441 TargetTransformInfo::RGK_FixedWidthVector)
442 .getFixedSize()));
443 }
444
445 /// Return the set of vectors that a matrix value is lowered to.
446 ///
447 /// If we lowered \p MatrixVal, just return the cache result matrix. Otherwise
448 /// split the flat vector \p MatrixVal containing a matrix with shape \p SI
449 /// into vectors.
450 MatrixTy getMatrix(Value *MatrixVal, const ShapeInfo &SI,
451 IRBuilder<> &Builder) {
452 VectorType *VType = dyn_cast<VectorType>(MatrixVal->getType());
453 assert(VType && "MatrixVal must be a vector type")(static_cast<void> (0));
454 assert(cast<FixedVectorType>(VType)->getNumElements() ==(static_cast<void> (0))
455 SI.NumRows * SI.NumColumns &&(static_cast<void> (0))
456 "The vector size must match the number of matrix elements")(static_cast<void> (0));
457
458 // Check if we lowered MatrixVal using shape information. In that case,
459 // return the existing matrix, if it matches the requested shape
460 // information. If there is a mis-match, embed the result in a flat
461 // vector and split it later.
462 auto Found = Inst2ColumnMatrix.find(MatrixVal);
463 if (Found != Inst2ColumnMatrix.end()) {
464 MatrixTy &M = Found->second;
465 // Return the found matrix, if its shape matches the requested shape
466 // information
467 if (SI.NumRows == M.getNumRows() && SI.NumColumns == M.getNumColumns())
468 return M;
469
470 MatrixVal = M.embedInVector(Builder);
471 }
472
473 // Otherwise split MatrixVal.
474 SmallVector<Value *, 16> SplitVecs;
475 for (unsigned MaskStart = 0;
476 MaskStart < cast<FixedVectorType>(VType)->getNumElements();
477 MaskStart += SI.getStride()) {
478 Value *V = Builder.CreateShuffleVector(
479 MatrixVal, createSequentialMask(MaskStart, SI.getStride(), 0),
480 "split");
481 SplitVecs.push_back(V);
482 }
483
484 return {SplitVecs};
485 }
486
487 /// If \p V already has a known shape return false. Otherwise set the shape
488 /// for instructions that support it.
489 bool setShapeInfo(Value *V, ShapeInfo Shape) {
490 assert(Shape && "Shape not set")(static_cast<void> (0));
491 if (isa<UndefValue>(V) || !supportsShapeInfo(V))
492 return false;
493
494 auto SIter = ShapeMap.find(V);
495 if (SIter != ShapeMap.end()) {
496 LLVM_DEBUG(dbgs() << " not overriding existing shape: "do { } while (false)
497 << SIter->second.NumRows << " "do { } while (false)
498 << SIter->second.NumColumns << " for " << *V << "\n")do { } while (false);
499 return false;
500 }
501
502 ShapeMap.insert({V, Shape});
503 LLVM_DEBUG(dbgs() << " " << Shape.NumRows << " x " << Shape.NumColumnsdo { } while (false)
504 << " for " << *V << "\n")do { } while (false);
505 return true;
506 }
507
508 bool isUniformShape(Value *V) {
509 Instruction *I = dyn_cast<Instruction>(V);
510 if (!I)
511 return true;
512
513 switch (I->getOpcode()) {
514 case Instruction::FAdd:
515 case Instruction::FSub:
516 case Instruction::FMul: // Scalar multiply.
517 case Instruction::FNeg:
518 case Instruction::Add:
519 case Instruction::Mul:
520 case Instruction::Sub:
521 return true;
522 default:
523 return false;
524 }
525 }
526
527 /// Returns true if shape information can be used for \p V. The supported
528 /// instructions must match the instructions that can be lowered by this pass.
529 bool supportsShapeInfo(Value *V) {
530 Instruction *Inst = dyn_cast<Instruction>(V);
531 if (!Inst)
532 return false;
533
534 IntrinsicInst *II = dyn_cast<IntrinsicInst>(Inst);
535 if (II)
536 switch (II->getIntrinsicID()) {
537 case Intrinsic::matrix_multiply:
538 case Intrinsic::matrix_transpose:
539 case Intrinsic::matrix_column_major_load:
540 case Intrinsic::matrix_column_major_store:
541 return true;
542 default:
543 return false;
544 }
545 return isUniformShape(V) || isa<StoreInst>(V) || isa<LoadInst>(V);
546 }
547
548 /// Propagate the shape information of instructions to their users.
549 /// The work list contains instructions for which we can compute the shape,
550 /// either based on the information provided by matrix intrinsics or known
551 /// shapes of operands.
552 SmallVector<Instruction *, 32>
553 propagateShapeForward(SmallVectorImpl<Instruction *> &WorkList) {
554 SmallVector<Instruction *, 32> NewWorkList;
555 // Pop an element for which we guaranteed to have at least one of the
556 // operand shapes. Add the shape for this and then add users to the work
557 // list.
558 LLVM_DEBUG(dbgs() << "Forward-propagate shapes:\n")do { } while (false);
559 while (!WorkList.empty()) {
560 Instruction *Inst = WorkList.pop_back_val();
561
562 // New entry, set the value and insert operands
563 bool Propagate = false;
564
565 Value *MatrixA;
566 Value *MatrixB;
567 Value *M;
568 Value *N;
569 Value *K;
570 if (match(Inst, m_Intrinsic<Intrinsic::matrix_multiply>(
571 m_Value(MatrixA), m_Value(MatrixB), m_Value(M),
572 m_Value(N), m_Value(K)))) {
573 Propagate = setShapeInfo(Inst, {M, K});
574 } else if (match(Inst, m_Intrinsic<Intrinsic::matrix_transpose>(
575 m_Value(MatrixA), m_Value(M), m_Value(N)))) {
576 // Flip dimensions.
577 Propagate = setShapeInfo(Inst, {N, M});
578 } else if (match(Inst, m_Intrinsic<Intrinsic::matrix_column_major_store>(
579 m_Value(MatrixA), m_Value(), m_Value(),
580 m_Value(), m_Value(M), m_Value(N)))) {
581 Propagate = setShapeInfo(Inst, {N, M});
582 } else if (match(Inst, m_Intrinsic<Intrinsic::matrix_column_major_load>(
583 m_Value(), m_Value(), m_Value(), m_Value(M),
584 m_Value(N)))) {
585 Propagate = setShapeInfo(Inst, {M, N});
586 } else if (match(Inst, m_Store(m_Value(MatrixA), m_Value()))) {
587 auto OpShape = ShapeMap.find(MatrixA);
588 if (OpShape != ShapeMap.end())
589 setShapeInfo(Inst, OpShape->second);
590 continue;
591 } else if (isUniformShape(Inst)) {
592 // Find the first operand that has a known shape and use that.
593 for (auto &Op : Inst->operands()) {
594 auto OpShape = ShapeMap.find(Op.get());
595 if (OpShape != ShapeMap.end()) {
596 Propagate |= setShapeInfo(Inst, OpShape->second);
597 break;
598 }
599 }
600 }
601
602 if (Propagate) {
603 NewWorkList.push_back(Inst);
604 for (auto *User : Inst->users())
605 if (ShapeMap.count(User) == 0)
606 WorkList.push_back(cast<Instruction>(User));
607 }
608 }
609
610 return NewWorkList;
611 }
612
613 /// Propagate the shape to operands of instructions with shape information.
614 /// \p Worklist contains the instruction for which we already know the shape.
615 SmallVector<Instruction *, 32>
616 propagateShapeBackward(SmallVectorImpl<Instruction *> &WorkList) {
617 SmallVector<Instruction *, 32> NewWorkList;
618
619 auto pushInstruction = [](Value *V,
620 SmallVectorImpl<Instruction *> &WorkList) {
621 Instruction *I = dyn_cast<Instruction>(V);
622 if (I)
623 WorkList.push_back(I);
624 };
625 // Pop an element with known shape. Traverse the operands, if their shape
626 // derives from the result shape and is unknown, add it and add them to the
627 // worklist.
628 LLVM_DEBUG(dbgs() << "Backward-propagate shapes:\n")do { } while (false);
629 while (!WorkList.empty()) {
630 Value *V = WorkList.pop_back_val();
631
632 size_t BeforeProcessingV = WorkList.size();
633 if (!isa<Instruction>(V))
634 continue;
635
636 Value *MatrixA;
637 Value *MatrixB;
638 Value *M;
639 Value *N;
640 Value *K;
641 if (match(V, m_Intrinsic<Intrinsic::matrix_multiply>(
642 m_Value(MatrixA), m_Value(MatrixB), m_Value(M),
643 m_Value(N), m_Value(K)))) {
644 if (setShapeInfo(MatrixA, {M, N}))
645 pushInstruction(MatrixA, WorkList);
646
647 if (setShapeInfo(MatrixB, {N, K}))
648 pushInstruction(MatrixB, WorkList);
649
650 } else if (match(V, m_Intrinsic<Intrinsic::matrix_transpose>(
651 m_Value(MatrixA), m_Value(M), m_Value(N)))) {
652 // Flip dimensions.
653 if (setShapeInfo(MatrixA, {M, N}))
654 pushInstruction(MatrixA, WorkList);
655 } else if (match(V, m_Intrinsic<Intrinsic::matrix_column_major_store>(
656 m_Value(MatrixA), m_Value(), m_Value(), m_Value(),
657 m_Value(M), m_Value(N)))) {
658 if (setShapeInfo(MatrixA, {M, N})) {
659 pushInstruction(MatrixA, WorkList);
660 }
661 } else if (isa<LoadInst>(V) ||
662 match(V, m_Intrinsic<Intrinsic::matrix_column_major_load>())) {
663 // Nothing to do, no matrix input.
664 } else if (isa<StoreInst>(V)) {
665 // Nothing to do. We forward-propagated to this so we would just
666 // backward propagate to an instruction with an already known shape.
667 } else if (isUniformShape(V)) {
668 // Propagate to all operands.
669 ShapeInfo Shape = ShapeMap[V];
670 for (Use &U : cast<Instruction>(V)->operands()) {
671 if (setShapeInfo(U.get(), Shape))
672 pushInstruction(U.get(), WorkList);
673 }
674 }
675 // After we discovered new shape info for new instructions in the
676 // worklist, we use their users as seeds for the next round of forward
677 // propagation.
678 for (size_t I = BeforeProcessingV; I != WorkList.size(); I++)
679 for (User *U : WorkList[I]->users())
680 if (isa<Instruction>(U) && V != U)
681 NewWorkList.push_back(cast<Instruction>(U));
682 }
683 return NewWorkList;
684 }
685
686 /// Try moving transposes in order to fold them away or into multiplies.
687 void optimizeTransposes() {
688 auto ReplaceAllUsesWith = [this](Instruction &Old, Value *New) {
689 // We need to remove Old from the ShapeMap otherwise RAUW will replace it
690 // with New. We should only add New it it supportsShapeInfo so we insert
691 // it conditionally instead.
692 auto S = ShapeMap.find(&Old);
693 if (S != ShapeMap.end()) {
694 ShapeMap.erase(S);
695 if (supportsShapeInfo(New))
696 ShapeMap.insert({New, S->second});
697 }
698 Old.replaceAllUsesWith(New);
699 };
700
701 // First sink all transposes inside matmuls, hoping that we end up with NN,
702 // NT or TN variants.
703 for (BasicBlock &BB : reverse(Func)) {
704 for (auto II = BB.rbegin(); II != BB.rend();) {
705 Instruction &I = *II;
706 // We may remove II. By default continue on the next/prev instruction.
707 ++II;
708 // If we were to erase II, move again.
709 auto EraseFromParent = [&II](Value *V) {
710 auto *Inst = cast<Instruction>(V);
711 if (Inst->use_empty()) {
712 if (Inst == &*II) {
713 ++II;
714 }
715 Inst->eraseFromParent();
716 }
717 };
718
719 // If we're creating a new instruction, continue from there.
720 Instruction *NewInst = nullptr;
721
722 IRBuilder<> IB(&I);
723 MatrixBuilder<IRBuilder<>> Builder(IB);
724
725 Value *TA, *TAMA, *TAMB;
726 ConstantInt *R, *K, *C;
727 if (match(&I, m_Intrinsic<Intrinsic::matrix_transpose>(m_Value(TA)))) {
728
729 // Transpose of a transpose is a nop
730 Value *TATA;
731 if (match(TA,
732 m_Intrinsic<Intrinsic::matrix_transpose>(m_Value(TATA)))) {
733 ReplaceAllUsesWith(I, TATA);
734 EraseFromParent(&I);
735 EraseFromParent(TA);
736 }
737
738 // (A * B)^t -> B^t * A^t
739 // RxK KxC CxK KxR
740 else if (match(TA, m_Intrinsic<Intrinsic::matrix_multiply>(
741 m_Value(TAMA), m_Value(TAMB), m_ConstantInt(R),
742 m_ConstantInt(K), m_ConstantInt(C)))) {
743 Value *T0 = Builder.CreateMatrixTranspose(TAMB, K->getZExtValue(),
744 C->getZExtValue(),
745 TAMB->getName() + "_t");
746 // We are being run after shape prop, add shape for newly created
747 // instructions so that we lower them later.
748 setShapeInfo(T0, {C, K});
749 Value *T1 = Builder.CreateMatrixTranspose(TAMA, R->getZExtValue(),
750 K->getZExtValue(),
751 TAMA->getName() + "_t");
752 setShapeInfo(T1, {K, R});
753 NewInst = Builder.CreateMatrixMultiply(T0, T1, C->getZExtValue(),
754 K->getZExtValue(),
755 R->getZExtValue(), "mmul");
756 ReplaceAllUsesWith(I, NewInst);
757 EraseFromParent(&I);
758 EraseFromParent(TA);
759 }
760 }
761
762 // If we replaced I with a new instruction, continue from there.
763 if (NewInst)
764 II = std::next(BasicBlock::reverse_iterator(NewInst));
765 }
766 }
767
768 // If we have a TT matmul, lift the transpose. We may be able to fold into
769 // consuming multiply.
770 for (BasicBlock &BB : Func) {
771 for (BasicBlock::iterator II = BB.begin(); II != BB.end();) {
772 Instruction *I = &*II;
773 // We may remove I.
774 ++II;
775 Value *A, *B, *AT, *BT;
776 ConstantInt *R, *K, *C;
777 // A^t * B ^t -> (B * A)^t
778 if (match(&*I, m_Intrinsic<Intrinsic::matrix_multiply>(
779 m_Value(A), m_Value(B), m_ConstantInt(R),
780 m_ConstantInt(K), m_ConstantInt(C))) &&
781 match(A, m_Intrinsic<Intrinsic::matrix_transpose>(m_Value(AT))) &&
782 match(B, m_Intrinsic<Intrinsic::matrix_transpose>(m_Value((BT))))) {
783 IRBuilder<> IB(&*I);
784 MatrixBuilder<IRBuilder<>> Builder(IB);
785 Value *M = Builder.CreateMatrixMultiply(
786 BT, AT, C->getZExtValue(), K->getZExtValue(), R->getZExtValue());
787 setShapeInfo(M, {C, R});
788 Instruction *NewInst = Builder.CreateMatrixTranspose(
789 M, C->getZExtValue(), R->getZExtValue());
790 ReplaceAllUsesWith(*I, NewInst);
791 if (I->use_empty())
792 I->eraseFromParent();
793 if (A->use_empty())
794 cast<Instruction>(A)->eraseFromParent();
795 if (A != B && B->use_empty())
796 cast<Instruction>(B)->eraseFromParent();
797 }
798 }
799 }
800 }
801
802 bool Visit() {
803 SmallVector<Instruction *, 32> WorkList;
804
805 // Initially only the shape of matrix intrinsics is known.
806 // Initialize the work list with ops carrying shape information.
807 for (BasicBlock &BB : Func)
808 for (Instruction &Inst : BB) {
809 IntrinsicInst *II = dyn_cast<IntrinsicInst>(&Inst);
810 if (!II)
811 continue;
812
813 switch (II->getIntrinsicID()) {
814 case Intrinsic::matrix_multiply:
815 case Intrinsic::matrix_transpose:
816 case Intrinsic::matrix_column_major_load:
817 case Intrinsic::matrix_column_major_store:
818 WorkList.push_back(&Inst);
819 break;
820 default:
821 break;
822 }
823 }
824
825 // Avoid unnecessary work if there are no matrix intrinsics in the function.
826 if (WorkList.empty())
827 return false;
828
829 // Propagate shapes until nothing changes any longer.
830 while (!WorkList.empty()) {
831 WorkList = propagateShapeForward(WorkList);
832 WorkList = propagateShapeBackward(WorkList);
833 }
834
835 if (!isMinimal()) {
836 optimizeTransposes();
837 LLVM_DEBUG({do { } while (false)
838 dbgs() << "Dump after matrix transpose optimization:\n";do { } while (false)
839 Func.dump();do { } while (false)
840 })do { } while (false);
841 }
842
843 bool Changed = false;
844 SmallVector<CallInst *, 16> MaybeFusableInsts;
845 SmallVector<Instruction *, 16> MatrixInsts;
846
847 // First, collect all instructions with shape information and candidates for
848 // fusion (currently only matrix multiplies).
849 ReversePostOrderTraversal<Function *> RPOT(&Func);
850 for (auto *BB : RPOT)
851 for (Instruction &I : *BB) {
852 if (ShapeMap.find(&I) == ShapeMap.end())
853 continue;
854 if (match(&I, m_Intrinsic<Intrinsic::matrix_multiply>()))
855 MaybeFusableInsts.push_back(cast<CallInst>(&I));
856 MatrixInsts.push_back(&I);
857 }
858
859 // Second, try to fuse candidates.
860 SmallPtrSet<Instruction *, 16> FusedInsts;
861 for (CallInst *CI : MaybeFusableInsts)
862 LowerMatrixMultiplyFused(CI, FusedInsts);
863 Changed = !FusedInsts.empty();
864
865 // Third, lower remaining instructions with shape information.
866 for (Instruction *Inst : MatrixInsts) {
867 if (FusedInsts.count(Inst))
868 continue;
869
870 IRBuilder<> Builder(Inst);
871
872 if (CallInst *CInst = dyn_cast<CallInst>(Inst))
873 Changed |= VisitCallInst(CInst);
874
875 Value *Op1;
876 Value *Op2;
877 if (auto *BinOp = dyn_cast<BinaryOperator>(Inst))
878 Changed |= VisitBinaryOperator(BinOp);
879 if (auto *UnOp = dyn_cast<UnaryOperator>(Inst))
880 Changed |= VisitUnaryOperator(UnOp);
881 if (match(Inst, m_Load(m_Value(Op1))))
882 Changed |= VisitLoad(cast<LoadInst>(Inst), Op1, Builder);
883 else if (match(Inst, m_Store(m_Value(Op1), m_Value(Op2))))
884 Changed |= VisitStore(cast<StoreInst>(Inst), Op1, Op2, Builder);
885 }
886
887 if (ORE) {
888 RemarkGenerator RemarkGen(Inst2ColumnMatrix, *ORE, Func);
889 RemarkGen.emitRemarks();
890 }
891
892 // Delete the instructions backwards, as it has a reduced likelihood of
893 // having to update as many def-use and use-def chains.
894 //
895 // Because we add to ToRemove during fusion we can't guarantee that defs
896 // are before uses. Change uses to undef temporarily as these should get
897 // removed as well.
898 //
899 // For verification, we keep track of where we changed uses to undefs in
900 // UndefedInsts and then check that we in fact remove them.
901 SmallSet<Instruction *, 16> UndefedInsts;
902 for (auto *Inst : reverse(ToRemove)) {
903 for (auto I = Inst->use_begin(), E = Inst->use_end(); I != E;) {
904 Use &U = *I++;
905 if (auto *Undefed = dyn_cast<Instruction>(U.getUser()))
906 UndefedInsts.insert(Undefed);
907 U.set(UndefValue::get(Inst->getType()));
908 }
909 Inst->eraseFromParent();
910 UndefedInsts.erase(Inst);
911 }
912 if (!UndefedInsts.empty()) {
913 // If we didn't remove all undefed instructions, it's a hard error.
914 dbgs() << "Undefed but present instructions:\n";
915 for (auto *I : UndefedInsts)
916 dbgs() << *I << "\n";
917 llvm_unreachable("Undefed but instruction not removed")__builtin_unreachable();
918 }
919
920 return Changed;
921 }
922
923 /// Turns \p BasePtr into an elementwise pointer to \p EltType.
924 Value *createElementPtr(Value *BasePtr, Type *EltType, IRBuilder<> &Builder) {
925 unsigned AS = cast<PointerType>(BasePtr->getType())->getAddressSpace();
926 Type *EltPtrType = PointerType::get(EltType, AS);
927 return Builder.CreatePointerCast(BasePtr, EltPtrType);
928 }
929
930 /// Replace intrinsic calls
931 bool VisitCallInst(CallInst *Inst) {
932 if (!Inst->getCalledFunction() || !Inst->getCalledFunction()->isIntrinsic())
933 return false;
934
935 switch (Inst->getCalledFunction()->getIntrinsicID()) {
936 case Intrinsic::matrix_multiply:
937 LowerMultiply(Inst);
938 break;
939 case Intrinsic::matrix_transpose:
940 LowerTranspose(Inst);
941 break;
942 case Intrinsic::matrix_column_major_load:
943 LowerColumnMajorLoad(Inst);
944 break;
945 case Intrinsic::matrix_column_major_store:
946 LowerColumnMajorStore(Inst);
947 break;
948 default:
949 return false;
950 }
951 return true;
952 }
953
954 /// Compute the alignment for a column/row \p Idx with \p Stride between them.
955 /// The address at \p Idx == 0 has alignment \p A. If \p Stride is a
956 /// ConstantInt, reduce the initial alignment based on the byte offset. For
957 /// non-ConstantInt strides, return the common alignment of the initial
958 /// alignment and the element size in bytes.
959 Align getAlignForIndex(unsigned Idx, Value *Stride, Type *ElementTy,
960 MaybeAlign A) const {
961 Align InitialAlign = DL.getValueOrABITypeAlignment(A, ElementTy);
962 if (Idx == 0)
963 return InitialAlign;
964
965 TypeSize ElementSizeInBits = DL.getTypeSizeInBits(ElementTy);
966 if (auto *ConstStride = dyn_cast<ConstantInt>(Stride)) {
967 uint64_t StrideInBytes =
968 ConstStride->getZExtValue() * ElementSizeInBits / 8;
969 return commonAlignment(InitialAlign, Idx * StrideInBytes);
970 }
971 return commonAlignment(InitialAlign, ElementSizeInBits / 8);
972 }
973
974 /// Load a matrix with \p Shape starting at \p Ptr and using \p Stride between
975 /// vectors.
976 MatrixTy loadMatrix(Type *Ty, Value *Ptr, MaybeAlign MAlign, Value *Stride,
977 bool IsVolatile, ShapeInfo Shape, IRBuilder<> &Builder) {
978 auto *VType = cast<VectorType>(Ty);
979 Type *EltTy = VType->getElementType();
980 Type *VecTy = FixedVectorType::get(EltTy, Shape.getStride());
981 Value *EltPtr = createElementPtr(Ptr, EltTy, Builder);
982 MatrixTy Result;
983 for (unsigned I = 0, E = Shape.getNumVectors(); I < E; ++I) {
984 Value *GEP = computeVectorAddr(
985 EltPtr, Builder.getIntN(Stride->getType()->getScalarSizeInBits(), I),
986 Stride, Shape.getStride(), EltTy, Builder);
987 Value *Vector = Builder.CreateAlignedLoad(
988 VecTy, GEP, getAlignForIndex(I, Stride, EltTy, MAlign),
989 IsVolatile, "col.load");
990
991 Result.addVector(Vector);
992 }
993 return Result.addNumLoads(getNumOps(Result.getVectorTy()) *
994 Result.getNumVectors());
995 }
996
997 /// Loads a sub-matrix with shape \p ResultShape from a \p R x \p C matrix,
998 /// starting at \p MatrixPtr[I][J].
999 MatrixTy loadMatrix(Value *MatrixPtr, MaybeAlign Align, bool IsVolatile,
1000 ShapeInfo MatrixShape, Value *I, Value *J,
1001 ShapeInfo ResultShape, Type *EltTy,
1002 IRBuilder<> &Builder) {
1003
1004 Value *Offset = Builder.CreateAdd(
1005 Builder.CreateMul(J, Builder.getInt64(MatrixShape.getStride())), I);
1006
1007 unsigned AS = cast<PointerType>(MatrixPtr->getType())->getAddressSpace();
1008 Value *EltPtr =
1009 Builder.CreatePointerCast(MatrixPtr, PointerType::get(EltTy, AS));
1010 Value *TileStart = Builder.CreateGEP(EltTy, EltPtr, Offset);
1011 auto *TileTy = FixedVectorType::get(EltTy, ResultShape.NumRows *
1012 ResultShape.NumColumns);
1013 Type *TilePtrTy = PointerType::get(TileTy, AS);
1014 Value *TilePtr =
1015 Builder.CreatePointerCast(TileStart, TilePtrTy, "col.cast");
1016
1017 return loadMatrix(TileTy, TilePtr, Align,
1018 Builder.getInt64(MatrixShape.getStride()), IsVolatile,
1019 ResultShape, Builder);
1020 }
1021
1022 /// Lower a load instruction with shape information.
1023 void LowerLoad(Instruction *Inst, Value *Ptr, MaybeAlign Align, Value *Stride,
1024 bool IsVolatile, ShapeInfo Shape) {
1025 IRBuilder<> Builder(Inst);
1026 finalizeLowering(Inst,
1027 loadMatrix(Inst->getType(), Ptr, Align, Stride, IsVolatile,
1028 Shape, Builder),
1029 Builder);
1030 }
1031
1032 /// Lowers llvm.matrix.column.major.load.
1033 ///
1034 /// The intrinsic loads a matrix from memory using a stride between columns.
1035 void LowerColumnMajorLoad(CallInst *Inst) {
1036 assert(MatrixLayout == MatrixLayoutTy::ColumnMajor &&(static_cast<void> (0))
1037 "Intrinsic only supports column-major layout!")(static_cast<void> (0));
1038 Value *Ptr = Inst->getArgOperand(0);
1039 Value *Stride = Inst->getArgOperand(1);
1040 LowerLoad(Inst, Ptr, Inst->getParamAlign(0), Stride,
1041 cast<ConstantInt>(Inst->getArgOperand(2))->isOne(),
1042 {Inst->getArgOperand(3), Inst->getArgOperand(4)});
1043 }
1044
1045 /// Stores a sub-matrix \p StoreVal into the \p R x \p C matrix starting at \p
1046 /// MatrixPtr[I][J].
1047 void storeMatrix(const MatrixTy &StoreVal, Value *MatrixPtr,
1048 MaybeAlign MAlign, bool IsVolatile, ShapeInfo MatrixShape,
1049 Value *I, Value *J, Type *EltTy, IRBuilder<> &Builder) {
1050 Value *Offset = Builder.CreateAdd(
1051 Builder.CreateMul(J, Builder.getInt64(MatrixShape.getStride())), I);
1052
1053 unsigned AS = cast<PointerType>(MatrixPtr->getType())->getAddressSpace();
1054 Value *EltPtr =
1055 Builder.CreatePointerCast(MatrixPtr, PointerType::get(EltTy, AS));
1056 Value *TileStart = Builder.CreateGEP(EltTy, EltPtr, Offset);
1057 auto *TileTy = FixedVectorType::get(EltTy, StoreVal.getNumRows() *
1058 StoreVal.getNumColumns());
1059 Type *TilePtrTy = PointerType::get(TileTy, AS);
1060 Value *TilePtr =
1061 Builder.CreatePointerCast(TileStart, TilePtrTy, "col.cast");
1062
1063 storeMatrix(TileTy, StoreVal, TilePtr, MAlign,
1064 Builder.getInt64(MatrixShape.getStride()), IsVolatile, Builder);
1065 }
1066
1067 /// Store matrix \p StoreVal starting at \p Ptr and using \p Stride between
1068 /// vectors.
1069 MatrixTy storeMatrix(Type *Ty, MatrixTy StoreVal, Value *Ptr,
1070 MaybeAlign MAlign, Value *Stride, bool IsVolatile,
1071 IRBuilder<> &Builder) {
1072 auto VType = cast<VectorType>(Ty);
1073 Value *EltPtr = createElementPtr(Ptr, VType->getElementType(), Builder);
1074 for (auto Vec : enumerate(StoreVal.vectors())) {
1075 Value *GEP = computeVectorAddr(
1076 EltPtr,
1077 Builder.getIntN(Stride->getType()->getScalarSizeInBits(),
1078 Vec.index()),
1079 Stride, StoreVal.getStride(), VType->getElementType(), Builder);
1080 Builder.CreateAlignedStore(Vec.value(), GEP,
1081 getAlignForIndex(Vec.index(), Stride,
1082 VType->getElementType(),
1083 MAlign),
1084 IsVolatile);
1085 }
1086 return MatrixTy().addNumStores(getNumOps(StoreVal.getVectorTy()) *
1087 StoreVal.getNumVectors());
1088 }
1089
1090 /// Lower a store instruction with shape information.
1091 void LowerStore(Instruction *Inst, Value *Matrix, Value *Ptr, MaybeAlign A,
1092 Value *Stride, bool IsVolatile, ShapeInfo Shape) {
1093 IRBuilder<> Builder(Inst);
1094 auto StoreVal = getMatrix(Matrix, Shape, Builder);
1095 finalizeLowering(Inst,
1096 storeMatrix(Matrix->getType(), StoreVal, Ptr, A, Stride,
1097 IsVolatile, Builder),
1098 Builder);
1099 }
1100
1101 /// Lowers llvm.matrix.column.major.store.
1102 ///
1103 /// The intrinsic store a matrix back memory using a stride between columns.
1104 void LowerColumnMajorStore(CallInst *Inst) {
1105 assert(MatrixLayout == MatrixLayoutTy::ColumnMajor &&(static_cast<void> (0))
1106 "Intrinsic only supports column-major layout!")(static_cast<void> (0));
1107 Value *Matrix = Inst->getArgOperand(0);
1108 Value *Ptr = Inst->getArgOperand(1);
1109 Value *Stride = Inst->getArgOperand(2);
1110 LowerStore(Inst, Matrix, Ptr, Inst->getParamAlign(1), Stride,
1111 cast<ConstantInt>(Inst->getArgOperand(3))->isOne(),
1112 {Inst->getArgOperand(4), Inst->getArgOperand(5)});
1113 }
1114
1115 // Set elements I..I+NumElts-1 to Block
1116 Value *insertVector(Value *Col, unsigned I, Value *Block,
1117 IRBuilder<> &Builder) {
1118
1119 // First, bring Block to the same size as Col
1120 unsigned BlockNumElts =
1121 cast<FixedVectorType>(Block->getType())->getNumElements();
1122 unsigned NumElts = cast<FixedVectorType>(Col->getType())->getNumElements();
1123 assert(NumElts >= BlockNumElts && "Too few elements for current block")(static_cast<void> (0));
1124
1125 Block = Builder.CreateShuffleVector(
1126 Block, createSequentialMask(0, BlockNumElts, NumElts - BlockNumElts));
1127
1128 // If Col is 7 long and I is 2 and BlockNumElts is 2 the mask is: 0, 1, 7,
1129 // 8, 4, 5, 6
1130 SmallVector<int, 16> Mask;
1131 unsigned i;
1132 for (i = 0; i < I; i++)
1133 Mask.push_back(i);
1134
1135 unsigned VecNumElts =
1136 cast<FixedVectorType>(Col->getType())->getNumElements();
1137 for (; i < I + BlockNumElts; i++)
1138 Mask.push_back(i - I + VecNumElts);
1139
1140 for (; i < VecNumElts; i++)
1141 Mask.push_back(i);
1142
1143 return Builder.CreateShuffleVector(Col, Block, Mask);
1144 }
1145
1146 Value *createMulAdd(Value *Sum, Value *A, Value *B, bool UseFPOp,
1147 IRBuilder<> &Builder, bool AllowContraction,
1148 unsigned &NumComputeOps) {
1149 NumComputeOps += getNumOps(A->getType());
1150 if (!Sum)
1151 return UseFPOp ? Builder.CreateFMul(A, B) : Builder.CreateMul(A, B);
1152
1153 if (UseFPOp) {
1154 if (AllowContraction) {
1155 // Use fmuladd for floating point operations and let the backend decide
1156 // if that's profitable.
1157 Function *FMulAdd = Intrinsic::getDeclaration(
1158 Func.getParent(), Intrinsic::fmuladd, A->getType());
1159 return Builder.CreateCall(FMulAdd, {A, B, Sum});
1160 }
1161 NumComputeOps += getNumOps(A->getType());
1162 Value *Mul = Builder.CreateFMul(A, B);
1163 return Builder.CreateFAdd(Sum, Mul);
1164 }
1165
1166 NumComputeOps += getNumOps(A->getType());
1167 Value *Mul = Builder.CreateMul(A, B);
1168 return Builder.CreateAdd(Sum, Mul);
1169 }
1170
1171 /// Cache \p Matrix as result of \p Inst and update the uses of \p Inst. For
1172 /// users with shape information, there's nothing to do: they will use the
1173 /// cached value when they are lowered. For other users, \p Matrix is
1174 /// flattened and the uses are updated to use it. Also marks \p Inst for
1175 /// deletion.
1176 void finalizeLowering(Instruction *Inst, MatrixTy Matrix,
1177 IRBuilder<> &Builder) {
1178 auto inserted = Inst2ColumnMatrix.insert(std::make_pair(Inst, Matrix));
1179 (void)inserted;
1180 assert(inserted.second && "multiple matrix lowering mapping")(static_cast<void> (0));
1181
1182 ToRemove.push_back(Inst);
1183 Value *Flattened = nullptr;
1184 for (Use &U : llvm::make_early_inc_range(Inst->uses())) {
1185 if (ShapeMap.find(U.getUser()) == ShapeMap.end()) {
1186 if (!Flattened)
1187 Flattened = Matrix.embedInVector(Builder);
1188 U.set(Flattened);
1189 }
1190 }
1191 }
1192
1193 /// Compute \p Result += \p A * \p B for input matrices with left-associating
1194 /// addition.
1195 ///
1196 /// We can fold a transpose into the operand that is used to extract scalars.
1197 /// This is the first operands with row-major and the second with
1198 /// column-major. If \p IsScalarMatrixTransposed we assume the appropriate
1199 /// operand is transposed.
1200 void emitMatrixMultiply(MatrixTy &Result, const MatrixTy &A,
1201 const MatrixTy &B, IRBuilder<> &Builder, bool IsTiled,
1202 bool IsScalarMatrixTransposed, FastMathFlags FMF) {
1203 const unsigned VF = std::max<unsigned>(
1204 TTI.getRegisterBitWidth(TargetTransformInfo::RGK_FixedWidthVector)
1205 .getFixedSize() /
1206 Result.getElementType()->getPrimitiveSizeInBits().getFixedSize(),
1207 1U);
1208 unsigned R = Result.getNumRows();
1209 unsigned C = Result.getNumColumns();
1210 unsigned M = A.getNumColumns();
1211
1212 bool IsFP = Result.getElementType()->isFloatingPointTy();
1213 assert(A.isColumnMajor() == B.isColumnMajor() &&(static_cast<void> (0))
1214 Result.isColumnMajor() == A.isColumnMajor() &&(static_cast<void> (0))
1215 "operands must agree on matrix layout")(static_cast<void> (0));
1216 unsigned NumComputeOps = 0;
1217
1218 Builder.setFastMathFlags(FMF);
1219
1220 if (A.isColumnMajor()) {
1221 // Multiply columns from the first operand with scalars from the second
1222 // operand. Then move along the K axes and accumulate the columns. With
1223 // this the adds can be vectorized without reassociation.
1224 for (unsigned J = 0; J < C; ++J) {
1225 unsigned BlockSize = VF;
1226 // If Result is zero, we don't need to accumulate in the K==0 iteration.
1227 bool isSumZero = isa<ConstantAggregateZero>(Result.getColumn(J));
1228
1229 for (unsigned I = 0; I < R; I += BlockSize) {
1230 // Gradually lower the vectorization factor to cover the remainder.
1231 while (I + BlockSize > R)
1232 BlockSize /= 2;
1233
1234 Value *Sum = IsTiled ? Result.extractVector(I, J, BlockSize, Builder)
1235 : nullptr;
1236 for (unsigned K = 0; K < M; ++K) {
1237 Value *L = A.extractVector(I, K, BlockSize, Builder);
1238 Value *RH = Builder.CreateExtractElement(
1239 B.getColumn(IsScalarMatrixTransposed ? K : J),
1240 IsScalarMatrixTransposed ? J : K);
1241 Value *Splat = Builder.CreateVectorSplat(BlockSize, RH, "splat");
1242 Sum =
1243 createMulAdd(isSumZero && K == 0 ? nullptr : Sum, L, Splat,
1244 IsFP, Builder, FMF.allowContract(), NumComputeOps);
1245 }
1246 Result.setVector(J,
1247 insertVector(Result.getVector(J), I, Sum, Builder));
1248 }
1249 }
1250 } else {
1251 // Multiply rows from the second operand with scalars from the first
1252 // operand. Then move along the K axes and accumulate the rows. With this
1253 // the adds can be vectorized without reassociation.
1254 for (unsigned I = 0; I < R; ++I) {
1255 unsigned BlockSize = VF;
1256 bool isSumZero = isa<ConstantAggregateZero>(Result.getRow(I));
1257 for (unsigned J = 0; J < C; J += BlockSize) {
1258 // Gradually lower the vectorization factor to cover the remainder.
1259 while (J + BlockSize > C)
1260 BlockSize /= 2;
1261
1262 Value *Sum = nullptr;
1263 for (unsigned K = 0; K < M; ++K) {
1264 Value *R = B.extractVector(K, J, BlockSize, Builder);
1265 Value *LH = Builder.CreateExtractElement(
1266 A.getVector(IsScalarMatrixTransposed ? K : I),
1267 IsScalarMatrixTransposed ? I : K);
1268 Value *Splat = Builder.CreateVectorSplat(BlockSize, LH, "splat");
1269 Sum =
1270 createMulAdd(isSumZero && K == 0 ? nullptr : Sum, Splat, R,
1271 IsFP, Builder, FMF.allowContract(), NumComputeOps);
1272 }
1273 Result.setVector(I,
1274 insertVector(Result.getVector(I), J, Sum, Builder));
1275 }
1276 }
1277 }
1278 Result.addNumComputeOps(NumComputeOps);
1279 }
1280
1281 /// Ensure that the memory in \p Load does not alias \p Store by potentially
1282 /// copying it to a new location. This new or otherwise the original location
1283 /// is returned.
1284 Value *getNonAliasingPointer(LoadInst *Load, StoreInst *Store,
1285 CallInst *MatMul) {
1286 MemoryLocation StoreLoc = MemoryLocation::get(Store);
1287 MemoryLocation LoadLoc = MemoryLocation::get(Load);
1288
1289 // If we can statically determine noalias we're good.
1290 if (AA->isNoAlias(LoadLoc, StoreLoc))
1291 return Load->getPointerOperand();
1292
1293 // Create code to check if the memory locations of the Load and Store
1294 // overlap and if they do, copy Load's operand to a new buffer.
1295
1296 // First, create new blocks for 2n part of the check and the copy.
1297 BasicBlock *Check0 = MatMul->getParent();
1298 // FIXME: Use lazy DTU and update SplitBlock to accept a DTU instead of a
1299 // DT. Manually collect dominator tree updates, to avoid unnecessary work,
1300 // as we adjust Check0 and Check1's branches.
1301 SmallVector<DominatorTree::UpdateType, 4> DTUpdates;
1302 for (BasicBlock *Succ : successors(Check0))
1303 DTUpdates.push_back({DT->Delete, Check0, Succ});
1304
1305 BasicBlock *Check1 =
1306 SplitBlock(MatMul->getParent(), MatMul, (DomTreeUpdater *)nullptr, LI,
1307 nullptr, "alias_cont");
1308 BasicBlock *Copy =
1309 SplitBlock(MatMul->getParent(), MatMul, (DomTreeUpdater *)nullptr, LI,
1310 nullptr, "copy");
1311 BasicBlock *Fusion =
1312 SplitBlock(MatMul->getParent(), MatMul, (DomTreeUpdater *)nullptr, LI,
1313 nullptr, "no_alias");
1314
1315 // Check if the loaded memory location begins before the end of the store
1316 // location. If the condition holds, they might overlap, otherwise they are
1317 // guaranteed to not overlap.
1318 IRBuilder<> Builder(MatMul);
1319 Check0->getTerminator()->eraseFromParent();
1320 Builder.SetInsertPoint(Check0);
1321 Type *IntPtrTy = Builder.getIntPtrTy(Load->getModule()->getDataLayout());
1322 Value *StoreBegin = Builder.CreatePtrToInt(
1323 const_cast<Value *>(StoreLoc.Ptr), IntPtrTy, "store.begin");
1324 Value *StoreEnd = Builder.CreateAdd(
1325 StoreBegin, ConstantInt::get(IntPtrTy, StoreLoc.Size.getValue()),
1326 "store.end", true, true);
1327 Value *LoadBegin = Builder.CreatePtrToInt(const_cast<Value *>(LoadLoc.Ptr),
1328 IntPtrTy, "load.begin");
1329 Builder.CreateCondBr(Builder.CreateICmpULT(LoadBegin, StoreEnd), Check1,
1330 Fusion);
1331
1332 // Check if the store begins before the end of the load location. If the
1333 // condition holds, they alias, otherwise they are guaranteed to not
1334 // overlap.
1335 Check1->getTerminator()->eraseFromParent();
1336 Builder.SetInsertPoint(Check1, Check1->begin());
1337 Value *LoadEnd = Builder.CreateAdd(
1338 LoadBegin, ConstantInt::get(IntPtrTy, LoadLoc.Size.getValue()),
1339 "load.end", true, true);
1340 Builder.CreateCondBr(Builder.CreateICmpULT(StoreBegin, LoadEnd), Copy,
1341 Fusion);
1342
1343 // Copy load operand to new alloca.
1344 Builder.SetInsertPoint(Copy, Copy->begin());
1345 AllocaInst *NewLd =
1346 Builder.CreateAlloca(Load->getType(), Load->getPointerAddressSpace());
1347 Builder.CreateMemCpy(NewLd, NewLd->getAlign(),
1348 Load->getPointerOperand(), Load->getAlign(),
1349 LoadLoc.Size.getValue());
1350 Builder.SetInsertPoint(Fusion, Fusion->begin());
1351 PHINode *PHI = Builder.CreatePHI(Load->getPointerOperandType(), 3);
1352 PHI->addIncoming(Load->getPointerOperand(), Check0);
1353 PHI->addIncoming(Load->getPointerOperand(), Check1);
1354 PHI->addIncoming(NewLd, Copy);
1355
1356 // Adjust DT.
1357 DTUpdates.push_back({DT->Insert, Check0, Check1});
1358 DTUpdates.push_back({DT->Insert, Check0, Fusion});
1359 DTUpdates.push_back({DT->Insert, Check1, Copy});
1360 DTUpdates.push_back({DT->Insert, Check1, Fusion});
1361 DT->applyUpdates(DTUpdates);
1362 return PHI;
1363 }
1364
1365 bool isFusionProfitable(CallInst *MatMul) {
1366 if (ForceFusion)
1367 return true;
1368
1369 ShapeInfo LShape(MatMul->getArgOperand(2), MatMul->getArgOperand(3));
1370 ShapeInfo RShape(MatMul->getArgOperand(3), MatMul->getArgOperand(4));
1371
1372 const unsigned R = LShape.NumRows;
1373 const unsigned C = RShape.NumColumns;
1374 const unsigned M = LShape.NumColumns;
1375 auto *EltType = cast<VectorType>(MatMul->getType())->getElementType();
1376
1377 const unsigned VF = std::max<unsigned>(
1378 TTI.getRegisterBitWidth(TargetTransformInfo::RGK_FixedWidthVector)
1379 .getFixedSize() /
1380 EltType->getPrimitiveSizeInBits().getFixedSize(),
1381 1U);
1382
1383 // Cost model for tiling
1384 //
1385 // For tiling to be beneficial, we need reuse either along the R or
1386 // the C axis. We vectorize along the R axis so that means at least
1387 // 3 elements.
1388 // TODO: Also consider cost of copying if operands alias.
1389 if (R <= VF && C == 1)
1390 return false;
1391 // Then we need enough elements to exceed the number of vector
1392 // registers we have. Note that this is an oversimplification since
1393 // fusing also takes some extra loads which may exceed the number of
1394 // reloads necessary.
1395 unsigned Op0Regs = (R + VF - 1) / VF * M;
1396 unsigned Op1Regs = (M + VF - 1) / VF * C;
1397 return Op0Regs + Op1Regs > TTI.getNumberOfRegisters(true);
1398 }
1399
1400 MatrixTy getZeroMatrix(Type *EltType, unsigned R, unsigned C) {
1401 MatrixTy Res;
1402 auto *ColumType = FixedVectorType::get(EltType, R);
1403 for (unsigned I = 0; I < C; ++I)
1404 Res.addVector(ConstantAggregateZero::get(ColumType));
1405 return Res;
1406 }
1407
1408 void createTiledLoops(CallInst *MatMul, Value *LPtr, ShapeInfo LShape,
1409 Value *RPtr, ShapeInfo RShape, StoreInst *Store) {
1410 auto *EltType = cast<VectorType>(MatMul->getType())->getElementType();
1411
1412 // Create the main tiling loop nest.
1413 TileInfo TI(LShape.NumRows, RShape.NumColumns, LShape.NumColumns, TileSize);
1414 DomTreeUpdater DTU(DT, DomTreeUpdater::UpdateStrategy::Lazy);
1415 Instruction *InsertI = cast<Instruction>(MatMul);
1416 BasicBlock *Start = InsertI->getParent();
1417 BasicBlock *End =
1418 SplitBlock(InsertI->getParent(), InsertI, DT, LI, nullptr, "continue");
1419 IRBuilder<> Builder(MatMul);
1420 BasicBlock *InnerBody = TI.CreateTiledLoops(Start, End, Builder, DTU, *LI);
1421
1422 Type *TileVecTy =
1423 FixedVectorType::get(MatMul->getType()->getScalarType(), TileSize);
1424 MatrixTy TileResult;
1425 // Insert in the inner loop header.
1426 Builder.SetInsertPoint(TI.InnerLoopHeader->getTerminator());
1427 // Create PHI nodes for the result columns to accumulate across iterations.
1428 SmallVector<PHINode *, 4> ColumnPhis;
1429 for (unsigned I = 0; I < TileSize; I++) {
1430 auto *Phi = Builder.CreatePHI(TileVecTy, 2, "result.vec." + Twine(I));
1431 Phi->addIncoming(ConstantAggregateZero::get(TileVecTy),
1432 TI.RowLoopHeader->getSingleSuccessor());
1433 TileResult.addVector(Phi);
1434 ColumnPhis.push_back(Phi);
1435 }
1436
1437 // Insert in the inner loop body, which computes
1438 // Res += Load(CurrentRow, K) * Load(K, CurrentColumn)
1439 Builder.SetInsertPoint(InnerBody->getTerminator());
1440 // Load tiles of the operands.
1441 MatrixTy A = loadMatrix(LPtr, {}, false, LShape, TI.CurrentRow, TI.CurrentK,
1442 {TileSize, TileSize}, EltType, Builder);
1443 MatrixTy B = loadMatrix(RPtr, {}, false, RShape, TI.CurrentK, TI.CurrentCol,
1444 {TileSize, TileSize}, EltType, Builder);
1445 emitMatrixMultiply(TileResult, A, B, Builder, true, false,
1446 getFastMathFlags(MatMul));
1447 // Store result after the inner loop is done.
1448 Builder.SetInsertPoint(TI.RowLoopLatch->getTerminator());
1449 storeMatrix(TileResult, Store->getPointerOperand(), Store->getAlign(),
1450 Store->isVolatile(), {LShape.NumRows, RShape.NumColumns},
1451 TI.CurrentRow, TI.CurrentCol, EltType, Builder);
1452
1453 for (unsigned I = 0; I < TileResult.getNumVectors(); I++)
1454 ColumnPhis[I]->addIncoming(TileResult.getVector(I), TI.InnerLoopLatch);
1455
1456 // Force unrolling of a few iterations of the inner loop, to make sure there
1457 // is enough work per iteration.
1458 // FIXME: The unroller should make this decision directly instead, but
1459 // currently the cost-model is not up to the task.
1460 unsigned InnerLoopUnrollCount = std::min(10u, LShape.NumColumns / TileSize);
1461 addStringMetadataToLoop(LI->getLoopFor(TI.InnerLoopHeader),
1462 "llvm.loop.unroll.count", InnerLoopUnrollCount);
1463 }
1464
1465 void emitSIMDTiling(CallInst *MatMul, LoadInst *LoadOp0, LoadInst *LoadOp1,
1466 StoreInst *Store,
1467 SmallPtrSetImpl<Instruction *> &FusedInsts) {
1468 assert(MatrixLayout == MatrixLayoutTy::ColumnMajor &&(static_cast<void> (0))
1469 "Tiling only supported for column-major matrixes at the moment!")(static_cast<void> (0));
1470 if (!isFusionProfitable(MatMul))
1471 return;
1472
1473 ShapeInfo LShape(MatMul->getArgOperand(2), MatMul->getArgOperand(3));
1474 ShapeInfo RShape(MatMul->getArgOperand(3), MatMul->getArgOperand(4));
1475
1476 const unsigned R = LShape.NumRows;
1477 const unsigned C = RShape.NumColumns;
1478 const unsigned M = LShape.NumColumns;
1479 auto *EltType = cast<VectorType>(MatMul->getType())->getElementType();
1480
1481 Value *APtr = getNonAliasingPointer(LoadOp0, Store, MatMul);
1482 Value *BPtr = getNonAliasingPointer(LoadOp1, Store, MatMul);
1483 Value *CPtr = Store->getPointerOperand();
1484
1485 if (TileUseLoops && (R % TileSize == 0 && C % TileSize == 0))
1486 createTiledLoops(MatMul, APtr, LShape, BPtr, RShape, Store);
1487 else {
1488 IRBuilder<> Builder(Store);
1489 for (unsigned J = 0; J < C; J += TileSize)
1490 for (unsigned I = 0; I < R; I += TileSize) {
1491 const unsigned TileR = std::min(R - I, unsigned(TileSize));
1492 const unsigned TileC = std::min(C - J, unsigned(TileSize));
1493 MatrixTy Res = getZeroMatrix(EltType, TileR, TileC);
1494
1495 for (unsigned K = 0; K < M; K += TileSize) {
1496 const unsigned TileM = std::min(M - K, unsigned(TileSize));
1497 MatrixTy A =
1498 loadMatrix(APtr, LoadOp0->getAlign(), LoadOp0->isVolatile(),
1499 LShape, Builder.getInt64(I), Builder.getInt64(K),
1500 {TileR, TileM}, EltType, Builder);
1501 MatrixTy B =
1502 loadMatrix(BPtr, LoadOp1->getAlign(), LoadOp1->isVolatile(),
1503 RShape, Builder.getInt64(K), Builder.getInt64(J),
1504 {TileM, TileC}, EltType, Builder);
1505 emitMatrixMultiply(Res, A, B, Builder, true, false,
1506 getFastMathFlags(MatMul));
1507 }
1508 storeMatrix(Res, CPtr, Store->getAlign(), Store->isVolatile(), {R, M},
1509 Builder.getInt64(I), Builder.getInt64(J), EltType,
1510 Builder);
1511 }
1512 }
1513
1514 // Mark eliminated instructions as fused and remove them.
1515 FusedInsts.insert(Store);
1516 FusedInsts.insert(MatMul);
1517 Store->eraseFromParent();
1518 MatMul->eraseFromParent();
1519 if (LoadOp0->hasNUses(0)) {
1520 FusedInsts.insert(LoadOp0);
1521 LoadOp0->eraseFromParent();
1522 }
1523 if (LoadOp1 != LoadOp0 && LoadOp1->hasNUses(0)) {
1524 FusedInsts.insert(LoadOp1);
1525 LoadOp1->eraseFromParent();
1526 }
1527 }
1528
1529 /// Try to lower matrix multiply chains by fusing operations.
1530 ///
1531 /// Call finalizeLowering on lowered instructions. Instructions that are
1532 /// completely eliminated by fusion are added to \p FusedInsts.
1533 void LowerMatrixMultiplyFused(CallInst *MatMul,
1534 SmallPtrSetImpl<Instruction *> &FusedInsts) {
1535 if (!FuseMatrix || !DT)
1536 return;
1537
1538 assert(AA && LI && "Analyses should be available")(static_cast<void> (0));
1539
1540 Value *A = MatMul->getArgOperand(0);
1541 Value *B = MatMul->getArgOperand(1);
1542
1543 // We can fold the transpose into the operand that is used to fetch scalars.
1544 Value *T;
1545 if (MatrixLayout == MatrixLayoutTy::ColumnMajor
1546 ? match(B, m_Intrinsic<Intrinsic::matrix_transpose>(m_Value(T)))
1547 : match(A, m_Intrinsic<Intrinsic::matrix_transpose>(m_Value(T)))) {
1548 IRBuilder<> Builder(MatMul);
1549 auto *EltType = cast<VectorType>(MatMul->getType())->getElementType();
1550 ShapeInfo LShape(MatMul->getArgOperand(2), MatMul->getArgOperand(3));
1551 ShapeInfo RShape(MatMul->getArgOperand(3), MatMul->getArgOperand(4));
1552 const unsigned R = LShape.NumRows;
1553 const unsigned M = LShape.NumColumns;
1554 const unsigned C = RShape.NumColumns;
1555
1556 MatrixTy MA;
1557 MatrixTy MB;
1558
1559 Value *Transpose;
1560 if (MatrixLayout == MatrixLayoutTy::ColumnMajor) {
1561 MA = getMatrix(A, ShapeInfo(R, M), Builder);
1562 MB = getMatrix(T, ShapeInfo(C, M), Builder);
1563 Transpose = B;
1564 } else {
1565 MA = getMatrix(T, ShapeInfo(R, M), Builder);
1566 MB = getMatrix(B, ShapeInfo(C, M), Builder);
1567 Transpose = A;
1568 }
1569
1570 // Initialize the output
1571 MatrixTy Result(R, C, EltType);
1572
1573 emitMatrixMultiply(Result, MA, MB, Builder, false, true,
1574 getFastMathFlags(MatMul));
1575
1576 FusedInsts.insert(MatMul);
1577 if (Transpose->hasOneUse()) {
1578 FusedInsts.insert(cast<Instruction>(Transpose));
1579 ToRemove.push_back(cast<Instruction>(Transpose));
1580 // TODO: add a fake entry for the folded instruction so that this is
1581 // included in the expression in the remark.
1582 Inst2ColumnMatrix[Transpose] = MatrixTy(M, C, EltType);
1583 }
1584 finalizeLowering(MatMul, Result, Builder);
1585 return;
1586 }
1587
1588 if (!MatMul->hasOneUse() || MatrixLayout != MatrixLayoutTy::ColumnMajor)
1589 return;
1590
1591 // Lower {ld, ld} -> matmul -> st chains. No need to call finalizeLowering
1592 // since the single store user will be lowered as part of this.
1593 auto *LoadOp0 = dyn_cast<LoadInst>(A);
1594 auto *LoadOp1 = dyn_cast<LoadInst>(B);
1595 auto *Store = dyn_cast<StoreInst>(*MatMul->user_begin());
1596 if (LoadOp0 && LoadOp1 && Store) {
1597 // The store address must dominate the MatMul instruction, otherwise
1598 // we create invalid IR.
1599 SetVector<Value *> WorkList;
1600 WorkList.insert(Store->getOperand(1));
1601 SmallVector<Instruction *> ToHoist;
1602 for (unsigned I = 0; I != WorkList.size(); ++I) {
1603 Value *Current = WorkList[I];
1604 auto *CurrI = dyn_cast<Instruction>(Current);
1605 if (!CurrI)
1606 continue;
1607 if (isa<PHINode>(CurrI))
1608 return;
1609 if (DT->dominates(CurrI, MatMul))
1610 continue;
1611 if (CurrI->mayHaveSideEffects() || CurrI->mayReadFromMemory())
1612 return;
1613 ToHoist.push_back(CurrI);
1614 WorkList.insert(CurrI->op_begin(), CurrI->op_end());
1615 }
1616
1617 sort(ToHoist, [this](Instruction *A, Instruction *B) {
1618 return DT->dominates(A, B);
1619 });
1620 for (Instruction *I : ToHoist)
1621 I->moveBefore(MatMul);
1622
1623 emitSIMDTiling(MatMul, LoadOp0, LoadOp1, Store, FusedInsts);
1624 return;
1625 }
1626 }
1627
1628 /// Lowers llvm.matrix.multiply.
1629 void LowerMultiply(CallInst *MatMul) {
1630 IRBuilder<> Builder(MatMul);
1631 auto *EltType = cast<VectorType>(MatMul->getType())->getElementType();
1632 ShapeInfo LShape(MatMul->getArgOperand(2), MatMul->getArgOperand(3));
1633 ShapeInfo RShape(MatMul->getArgOperand(3), MatMul->getArgOperand(4));
1634
1635 const MatrixTy &Lhs = getMatrix(MatMul->getArgOperand(0), LShape, Builder);
1636 const MatrixTy &Rhs = getMatrix(MatMul->getArgOperand(1), RShape, Builder);
1637 assert(Lhs.getElementType() == Rhs.getElementType() &&(static_cast<void> (0))
1638 "Matrix multiply argument element types do not match.")(static_cast<void> (0));
1639
1640 const unsigned R = LShape.NumRows;
1641 const unsigned C = RShape.NumColumns;
1642 assert(LShape.NumColumns == RShape.NumRows)(static_cast<void> (0));
1643
1644 // Initialize the output
1645 MatrixTy Result(R, C, EltType);
1646 assert(Lhs.getElementType() == Result.getElementType() &&(static_cast<void> (0))
1647 "Matrix multiply result element type does not match arguments.")(static_cast<void> (0));
1648
1649 emitMatrixMultiply(Result, Lhs, Rhs, Builder, false, false,
1650 getFastMathFlags(MatMul));
1651 finalizeLowering(MatMul, Result, Builder);
1652 }
1653
1654 /// Lowers llvm.matrix.transpose.
1655 void LowerTranspose(CallInst *Inst) {
1656 MatrixTy Result;
1657 IRBuilder<> Builder(Inst);
1658 Value *InputVal = Inst->getArgOperand(0);
1659 VectorType *VectorTy = cast<VectorType>(InputVal->getType());
1660 ShapeInfo ArgShape(Inst->getArgOperand(1), Inst->getArgOperand(2));
1661 MatrixTy InputMatrix = getMatrix(InputVal, ArgShape, Builder);
1662
1663 const unsigned NewNumVecs =
1664 InputMatrix.isColumnMajor() ? ArgShape.NumRows : ArgShape.NumColumns;
1665 const unsigned NewNumElts =
1666 InputMatrix.isColumnMajor() ? ArgShape.NumColumns : ArgShape.NumRows;
1667
1668 for (unsigned I = 0; I < NewNumVecs; ++I) {
1669 // Build a single result vector. First initialize it.
1670 Value *ResultVector = UndefValue::get(
1671 FixedVectorType::get(VectorTy->getElementType(), NewNumElts));
1672 // Go through the old elements and insert it into the resulting vector.
1673 for (auto J : enumerate(InputMatrix.vectors())) {
1674 Value *Elt = Builder.CreateExtractElement(J.value(), I);
1675 // Row and column indices are transposed.
1676 ResultVector =
1677 Builder.CreateInsertElement(ResultVector, Elt, J.index());
1678 }
1679 Result.addVector(ResultVector);
1680 }
1681
1682 // TODO: Improve estimate of operations needed for transposes. Currently we
1683 // just count the insertelement/extractelement instructions, but do not
1684 // account for later simplifications/combines.
1685 finalizeLowering(
1686 Inst,
1687 Result.addNumComputeOps(2 * ArgShape.NumRows * ArgShape.NumColumns)
1688 .addNumExposedTransposes(1),
1689 Builder);
1690 }
1691
1692 /// Lower load instructions, if shape information is available.
1693 bool VisitLoad(LoadInst *Inst, Value *Ptr, IRBuilder<> &Builder) {
1694 auto I = ShapeMap.find(Inst);
1695 if (I == ShapeMap.end())
1696 return false;
1697
1698 LowerLoad(Inst, Ptr, Inst->getAlign(),
1699 Builder.getInt64(I->second.getStride()), Inst->isVolatile(),
1700 I->second);
1701 return true;
1702 }
1703
1704 bool VisitStore(StoreInst *Inst, Value *StoredVal, Value *Ptr,
1705 IRBuilder<> &Builder) {
1706 auto I = ShapeMap.find(StoredVal);
1707 if (I == ShapeMap.end())
1708 return false;
1709
1710 LowerStore(Inst, StoredVal, Ptr, Inst->getAlign(),
1711 Builder.getInt64(I->second.getStride()), Inst->isVolatile(),
1712 I->second);
1713 return true;
1714 }
1715
1716 /// Lower binary operators, if shape information is available.
1717 bool VisitBinaryOperator(BinaryOperator *Inst) {
1718 auto I = ShapeMap.find(Inst);
1719 if (I == ShapeMap.end())
1720 return false;
1721
1722 Value *Lhs = Inst->getOperand(0);
1723 Value *Rhs = Inst->getOperand(1);
1724
1725 IRBuilder<> Builder(Inst);
1726 ShapeInfo &Shape = I->second;
1727
1728 MatrixTy Result;
1729 MatrixTy A = getMatrix(Lhs, Shape, Builder);
1730 MatrixTy B = getMatrix(Rhs, Shape, Builder);
1731 assert(A.isColumnMajor() == B.isColumnMajor() &&(static_cast<void> (0))
1732 Result.isColumnMajor() == A.isColumnMajor() &&(static_cast<void> (0))
1733 "operands must agree on matrix layout")(static_cast<void> (0));
1734
1735 Builder.setFastMathFlags(getFastMathFlags(Inst));
1736
1737 // Helper to perform binary op on vectors.
1738 auto BuildVectorOp = [&Builder, Inst](Value *LHS, Value *RHS) {
1739 switch (Inst->getOpcode()) {
1740 case Instruction::Add:
1741 return Builder.CreateAdd(LHS, RHS);
1742 case Instruction::Mul:
1743 return Builder.CreateMul(LHS, RHS);
1744 case Instruction::Sub:
1745 return Builder.CreateSub(LHS, RHS);
1746 case Instruction::FAdd:
1747 return Builder.CreateFAdd(LHS, RHS);
1748 case Instruction::FMul:
1749 return Builder.CreateFMul(LHS, RHS);
1750 case Instruction::FSub:
1751 return Builder.CreateFSub(LHS, RHS);
1752 default:
1753 llvm_unreachable("Unsupported binary operator for matrix")__builtin_unreachable();
1754 }
1755 };
1756
1757 for (unsigned I = 0; I < Shape.getNumVectors(); ++I)
1758 Result.addVector(BuildVectorOp(A.getVector(I), B.getVector(I)));
1759
1760 finalizeLowering(Inst,
1761 Result.addNumComputeOps(getNumOps(Result.getVectorTy()) *
1762 Result.getNumVectors()),
1763 Builder);
1764 return true;
1765 }
1766
1767 /// Lower unary operators, if shape information is available.
1768 bool VisitUnaryOperator(UnaryOperator *Inst) {
1769 auto I = ShapeMap.find(Inst);
1770 if (I == ShapeMap.end())
1771 return false;
1772
1773 Value *Op = Inst->getOperand(0);
1774
1775 IRBuilder<> Builder(Inst);
1776 ShapeInfo &Shape = I->second;
1777
1778 MatrixTy Result;
1779 MatrixTy M = getMatrix(Op, Shape, Builder);
1780
1781 Builder.setFastMathFlags(getFastMathFlags(Inst));
1782
1783 // Helper to perform unary op on vectors.
1784 auto BuildVectorOp = [&Builder, Inst](Value *Op) {
1785 switch (Inst->getOpcode()) {
1786 case Instruction::FNeg:
1787 return Builder.CreateFNeg(Op);
1788 default:
1789 llvm_unreachable("Unsupported unary operator for matrix")__builtin_unreachable();
1790 }
1791 };
1792
1793 for (unsigned I = 0; I < Shape.getNumVectors(); ++I)
1794 Result.addVector(BuildVectorOp(M.getVector(I)));
1795
1796 finalizeLowering(Inst,
1797 Result.addNumComputeOps(getNumOps(Result.getVectorTy()) *
1798 Result.getNumVectors()),
1799 Builder);
1800 return true;
1801 }
1802
1803 /// Helper to linearize a matrix expression tree into a string. Currently
1804 /// matrix expressions are linarized by starting at an expression leaf and
1805 /// linearizing bottom up.
1806 struct ExprLinearizer {
1807 unsigned LengthToBreak = 100;
1808 std::string Str;
1809 raw_string_ostream Stream;
1810 unsigned LineLength = 0;
1811 const DataLayout &DL;
1812
1813 /// Mapping from instructions to matrixes. It is used to identify
1814 /// matrix instructions.
1815 const MapVector<Value *, MatrixTy> &Inst2Matrix;
1816
1817 /// Mapping from values to the leaves of all expressions that the value is
1818 /// part of.
1819 const DenseMap<Value *, SmallPtrSet<Value *, 2>> &Shared;
1820
1821 /// Set of matrix expressions in the scope of a given DISubprogram.
1822 const SmallSetVector<Value *, 32> &ExprsInSubprogram;
1823
1824 /// Leaf node of the expression to linearize.
1825 Value *Leaf;
1826
1827 /// Used to keep track of sub-expressions that get reused while linearizing
1828 /// the expression. Re-used sub-expressions are marked as (reused).
1829 SmallPtrSet<Value *, 8> ReusedExprs;
1830
1831 ExprLinearizer(const DataLayout &DL,
1832 const MapVector<Value *, MatrixTy> &Inst2Matrix,
1833 const DenseMap<Value *, SmallPtrSet<Value *, 2>> &Shared,
1834 const SmallSetVector<Value *, 32> &ExprsInSubprogram,
1835 Value *Leaf)
1836 : Str(), Stream(Str), DL(DL), Inst2Matrix(Inst2Matrix), Shared(Shared),
1837 ExprsInSubprogram(ExprsInSubprogram), Leaf(Leaf) {}
1838
1839 void indent(unsigned N) {
1840 LineLength += N;
1841 for (unsigned i = 0; i < N; i++)
1842 Stream << " ";
1843 }
1844
1845 void lineBreak() {
1846 Stream << "\n";
1847 LineLength = 0;
1848 }
1849
1850 void maybeIndent(unsigned Indent) {
1851 if (LineLength >= LengthToBreak)
1852 lineBreak();
1853
1854 if (LineLength == 0)
1855 indent(Indent);
1856 }
1857
1858 void write(StringRef S) {
1859 LineLength += S.size();
1860 Stream << S;
1861 }
1862
1863 Value *getUnderlyingObjectThroughLoads(Value *V) {
1864 if (Value *Ptr = getPointerOperand(V))
1865 return getUnderlyingObjectThroughLoads(Ptr);
1866 else if (V->getType()->isPointerTy())
1867 return getUnderlyingObject(V);
1868 return V;
1869 }
1870
1871 /// Returns true if \p V is a matrix value in the given subprogram.
1872 bool isMatrix(Value *V) const { return ExprsInSubprogram.count(V); }
1873
1874 /// If \p V is a matrix value, print its shape as as NumRows x NumColumns to
1875 /// \p SS.
1876 void prettyPrintMatrixType(Value *V, raw_string_ostream &SS) {
1877 auto M = Inst2Matrix.find(V);
1878 if (M == Inst2Matrix.end())
1879 SS << "unknown";
1880 else {
1881 SS << M->second.getNumRows();
1882 SS << "x";
1883 SS << M->second.getNumColumns();
1884 }
1885 }
1886
1887 /// Write the called function name. Handles calls to llvm.matrix.*
1888 /// specially: we write the name, followed by the dimensions of the input
1889 /// matrixes, followed by the scalar type name.
1890 void writeFnName(CallInst *CI) {
1891 if (!CI->getCalledFunction())
9
Taking false branch
1892 write("<no called fn>");
1893 else {
1894 StringRef Name = CI->getCalledFunction()->getName();
1895 if (!Name.startswith("llvm.matrix")) {
10
Assuming the condition is false
11
Taking false branch
1896 write(Name);
1897 return;
1898 }
1899 IntrinsicInst *II = dyn_cast<IntrinsicInst>(CI);
12
Assuming 'CI' is not a 'IntrinsicInst'
13
'II' initialized to a null pointer value
1900 write(Intrinsic::getBaseName(II->getIntrinsicID())
14
Called C++ object pointer is null
1901 .drop_front(StringRef("llvm.matrix.").size()));
1902 write(".");
1903 std::string Tmp;
1904 raw_string_ostream SS(Tmp);
1905
1906 switch (II->getIntrinsicID()) {
1907 case Intrinsic::matrix_multiply:
1908 prettyPrintMatrixType(II->getOperand(0), SS);
1909 SS << ".";
1910 prettyPrintMatrixType(II->getOperand(1), SS);
1911 SS << "." << *II->getType()->getScalarType();
1912 break;
1913 case Intrinsic::matrix_transpose:
1914 prettyPrintMatrixType(II->getOperand(0), SS);
1915 SS << "." << *II->getType()->getScalarType();
1916 break;
1917 case Intrinsic::matrix_column_major_load:
1918 prettyPrintMatrixType(II, SS);
1919 SS << "." << *II->getType()->getScalarType();
1920 break;
1921 case Intrinsic::matrix_column_major_store:
1922 prettyPrintMatrixType(II->getOperand(0), SS);
1923 SS << "." << *II->getOperand(0)->getType()->getScalarType();
1924 break;
1925 default:
1926 llvm_unreachable("Unhandled case")__builtin_unreachable();
1927 }
1928 SS.flush();
1929 write(Tmp);
1930 }
1931 }
1932
1933 unsigned getNumShapeArgs(CallInst *CI) const {
1934 if (IntrinsicInst *II = dyn_cast<IntrinsicInst>(CI)) {
1935 switch (II->getIntrinsicID()) {
1936 case Intrinsic::matrix_multiply:
1937 return 3;
1938 case Intrinsic::matrix_transpose:
1939 return 2;
1940 case Intrinsic::matrix_column_major_load:
1941 case Intrinsic::matrix_column_major_store:
1942 return 3;
1943 default:
1944 return 0;
1945 }
1946 }
1947 return 0;
1948 }
1949
1950 /// Special printing for values: for pointers, we print if they refer to an
1951 /// (function) external address or a stack address, for other values we
1952 /// either print the constant or "scalar"/"matrix" for other values.
1953 void write(Value *V) {
1954 V = getUnderlyingObjectThroughLoads(V);
1955 if (V->getType()->isPointerTy()) {
1956 if (isa<AllocaInst>(V)) {
1957 Stream << "stack addr";
1958 LineLength += StringRef("stack addr").size();
1959 } else {
1960 Stream << "addr";
1961 LineLength += StringRef("addr").size();
1962 }
1963 if (!V->getName().empty()) {
1964 Stream << " %" << V->getName() << "";
1965 LineLength += V->getName().size() + 2;
1966 }
1967 return;
1968 }
1969
1970 std::string Tmp;
1971 raw_string_ostream TmpStream(Tmp);
1972
1973 if (auto *CI = dyn_cast<ConstantInt>(V))
1974 TmpStream << CI->getValue();
1975 else if (isa<Constant>(V))
1976 TmpStream << "constant";
1977 else {
1978 if (isMatrix(V))
1979 TmpStream << "matrix";
1980 else
1981 TmpStream << "scalar";
1982 }
1983 TmpStream.flush();
1984 Tmp = std::string(StringRef(Tmp).trim());
1985 LineLength += Tmp.size();
1986 Stream << Tmp;
1987 }
1988
1989 /// Linearize expression \p Expr starting at an indentation of \p Indent.
1990 /// Expressions that are re-used multiple times are prefixed with (reused)
1991 /// at the re-used root instruction.
1992 void linearizeExpr(Value *Expr, unsigned Indent, bool ParentReused,
1993 bool ParentShared) {
1994 auto *I = cast<Instruction>(Expr);
2
'Expr' is a 'Instruction'
1995 maybeIndent(Indent);
1996 SmallVector<Value *, 8> Ops;
1997
1998 // Is Expr shared with other expression leaves?
1999 bool ExprShared = false;
2000
2001 // Deal with shared subtrees. Mark them as shared, if required.
2002 if (!ParentShared
2.1
'ParentShared' is false
) {
3
Taking true branch
2003 auto SI = Shared.find(Expr);
2004 assert(SI != Shared.end() && SI->second.count(Leaf))(static_cast<void> (0));
2005
2006 for (Value *S : SI->second) {
2007 if (S == Leaf)
2008 continue;
2009 DebugLoc DL = cast<Instruction>(S)->getDebugLoc();
2010 write("shared with remark at line " + std::to_string(DL.getLine()) +
2011 " column " + std::to_string(DL.getCol()) + " (");
2012 }
2013 ExprShared = SI->second.size() > 1;
4
Assuming the condition is false
2014 }
2015
2016 bool Reused = !ReusedExprs.insert(Expr).second;
5
Assuming field 'second' is true
2017 if (Reused
5.1
'Reused' is false
&& !ParentReused)
2018 write("(reused) ");
2019
2020 if (auto *CI
6.1
'CI' is non-null
= dyn_cast<CallInst>(I)) {
6
Assuming 'I' is a 'CallInst'
7
Taking true branch
2021 writeFnName(CI);
8
Calling 'ExprLinearizer::writeFnName'
2022
2023 Ops.append(CI->arg_begin(), CI->arg_end() - getNumShapeArgs(CI));
2024 } else if (isa<BitCastInst>(Expr)) {
2025 // Special case bitcasts, which are used to materialize matrixes from
2026 // non-matrix ops.
2027 write("matrix");
2028 return;
2029 } else {
2030 Ops.append(I->value_op_begin(), I->value_op_end());
2031 write(std::string(I->getOpcodeName()));
2032 }
2033
2034 write(std::string("("));
2035
2036 unsigned NumOpsToBreak = 1;
2037 if (match(Expr, m_Intrinsic<Intrinsic::matrix_column_major_load>()))
2038 NumOpsToBreak = 2;
2039
2040 for (Value *Op : Ops) {
2041 if (Ops.size() > NumOpsToBreak)
2042 lineBreak();
2043
2044 maybeIndent(Indent + 1);
2045 if (isMatrix(Op))
2046 linearizeExpr(Op, Indent + 1, Reused, ExprShared);
2047 else
2048 write(Op);
2049 if (Op != Ops.back())
2050 write(", ");
2051 }
2052
2053 write(")");
2054 }
2055
2056 const std::string &getResult() {
2057 Stream.flush();
2058 return Str;
2059 }
2060 };
2061
2062 /// Generate remarks for matrix operations in a function. To generate remarks
2063 /// for matrix expressions, the following approach is used:
2064 /// 1. Use the inlined-at debug information to group matrix operations to the
2065 /// DISubprograms they are contained in.
2066 /// 2. Collect leaves of matrix expressions (done in
2067 /// RemarkGenerator::getExpressionLeaves) for each subprogram - expression
2068 // mapping. Leaves are lowered matrix instructions without other matrix
2069 // users (like stores) in the current subprogram.
2070 /// 3. For each leaf, create a remark containing a linearizied version of the
2071 /// matrix expression. The expression is linearized by a recursive
2072 /// bottom-up traversal of the matrix operands, starting at a leaf. Note
2073 /// that multiple leaves can share sub-expressions. Shared subexpressions
2074 /// are explicitly marked as shared().
2075 struct RemarkGenerator {
2076 const MapVector<Value *, MatrixTy> &Inst2Matrix;
2077 OptimizationRemarkEmitter &ORE;
2078 Function &Func;
2079 const DataLayout &DL;
2080
2081 RemarkGenerator(const MapVector<Value *, MatrixTy> &Inst2Matrix,
2082 OptimizationRemarkEmitter &ORE, Function &Func)
2083 : Inst2Matrix(Inst2Matrix), ORE(ORE), Func(Func),
2084 DL(Func.getParent()->getDataLayout()) {}
2085
2086 /// Return all leaves of the expressions in \p ExprsInSubprogram. Those are
2087 /// instructions in Inst2Matrix returning void or without any users in
2088 /// \p ExprsInSubprogram. Currently that should only include stores.
2089 SmallVector<Value *, 4>
2090 getExpressionLeaves(const SmallSetVector<Value *, 32> &ExprsInSubprogram) {
2091 SmallVector<Value *, 4> Leaves;
2092 for (auto *Expr : ExprsInSubprogram)
2093 if (Expr->getType()->isVoidTy() ||
2094 !any_of(Expr->users(), [&ExprsInSubprogram](User *U) {
2095 return ExprsInSubprogram.count(U);
2096 }))
2097 Leaves.push_back(Expr);
2098 return Leaves;
2099 }
2100
2101 /// Recursively traverse expression \p V starting at \p Leaf and add \p Leaf
2102 /// to all visited expressions in \p Shared. Limit the matrix operations to
2103 /// the ones in \p ExprsInSubprogram.
2104 void collectSharedInfo(Value *Leaf, Value *V,
2105 const SmallSetVector<Value *, 32> &ExprsInSubprogram,
2106 DenseMap<Value *, SmallPtrSet<Value *, 2>> &Shared) {
2107
2108 if (!ExprsInSubprogram.count(V))
2109 return;
2110
2111 auto I = Shared.insert({V, {}});
2112 I.first->second.insert(Leaf);
2113
2114 for (Value *Op : cast<Instruction>(V)->operand_values())
2115 collectSharedInfo(Leaf, Op, ExprsInSubprogram, Shared);
2116 }
2117
2118 /// Calculate the number of exclusive and shared op counts for expression
2119 /// starting at \p V. Expressions used multiple times are counted once.
2120 /// Limit the matrix operations to the ones in \p ExprsInSubprogram.
2121 std::pair<OpInfoTy, OpInfoTy>
2122 sumOpInfos(Value *Root, SmallPtrSetImpl<Value *> &ReusedExprs,
2123 const SmallSetVector<Value *, 32> &ExprsInSubprogram,
2124 DenseMap<Value *, SmallPtrSet<Value *, 2>> &Shared) const {
2125 if (!ExprsInSubprogram.count(Root))
2126 return {};
2127
2128 // Already counted this expression. Stop.
2129 if (!ReusedExprs.insert(Root).second)
2130 return {};
2131
2132 OpInfoTy SharedCount;
2133 OpInfoTy Count;
2134
2135 auto I = Shared.find(Root);
2136 auto CM = Inst2Matrix.find(Root);
2137 if (I->second.size() == 1)
2138 Count = CM->second.getOpInfo();
2139 else
2140 SharedCount = CM->second.getOpInfo();
2141
2142 for (Value *Op : cast<Instruction>(Root)->operand_values()) {
2143 auto C = sumOpInfos(Op, ReusedExprs, ExprsInSubprogram, Shared);
2144 Count += C.first;
2145 SharedCount += C.second;
2146 }
2147 return {Count, SharedCount};
2148 }
2149
2150 void emitRemarks() {
2151 if (!ORE.allowExtraAnalysis(DEBUG_TYPE"lower-matrix-intrinsics"))
2152 return;
2153
2154 // Map matrix operations to their containting subprograms, by traversing
2155 // the inlinedAt chain. If the function does not have a DISubprogram, we
2156 // only map them to the containing function.
2157 MapVector<DISubprogram *, SmallVector<Value *, 8>> Subprog2Exprs;
2158 for (auto &KV : Inst2Matrix) {
2159 if (Func.getSubprogram()) {
2160 auto *I = cast<Instruction>(KV.first);
2161 DILocation *Context = I->getDebugLoc();
2162 while (Context) {
2163 auto I =
2164 Subprog2Exprs.insert({getSubprogram(Context->getScope()), {}});
2165 I.first->second.push_back(KV.first);
2166 Context = DebugLoc(Context).getInlinedAt();
2167 }
2168 } else {
2169 auto I = Subprog2Exprs.insert({nullptr, {}});
2170 I.first->second.push_back(KV.first);
2171 }
2172 }
2173 for (auto &KV : Subprog2Exprs) {
2174 SmallSetVector<Value *, 32> ExprsInSubprogram(KV.second.begin(),
2175 KV.second.end());
2176 auto Leaves = getExpressionLeaves(ExprsInSubprogram);
2177
2178 DenseMap<Value *, SmallPtrSet<Value *, 2>> Shared;
2179 for (Value *Leaf : Leaves)
2180 collectSharedInfo(Leaf, Leaf, ExprsInSubprogram, Shared);
2181
2182 // Generate remarks for each leaf.
2183 for (auto *L : Leaves) {
2184
2185 DebugLoc Loc = cast<Instruction>(L)->getDebugLoc();
2186 DILocation *Context = cast<Instruction>(L)->getDebugLoc();
2187 while (Context) {
2188 if (getSubprogram(Context->getScope()) == KV.first) {
2189 Loc = Context;
2190 break;
2191 }
2192 Context = DebugLoc(Context).getInlinedAt();
2193 }
2194
2195 SmallPtrSet<Value *, 8> ReusedExprs;
2196 OpInfoTy Counts, SharedCounts;
2197 std::tie(Counts, SharedCounts) =
2198 sumOpInfos(L, ReusedExprs, ExprsInSubprogram, Shared);
2199
2200 OptimizationRemark Rem(DEBUG_TYPE"lower-matrix-intrinsics", "matrix-lowered", Loc,
2201 cast<Instruction>(L)->getParent());
2202
2203 Rem << "Lowered with ";
2204 Rem << ore::NV("NumStores", Counts.NumStores) << " stores, "
2205 << ore::NV("NumLoads", Counts.NumLoads) << " loads, "
2206 << ore::NV("NumComputeOps", Counts.NumComputeOps)
2207 << " compute ops, "
2208 << ore::NV("NumExposedTransposes", Counts.NumExposedTransposes)
2209 << " exposed transposes";
2210
2211 if (SharedCounts.NumStores > 0 || SharedCounts.NumLoads > 0 ||
2212 SharedCounts.NumComputeOps > 0) {
2213 Rem << ",\nadditionally "
2214 << ore::NV("NumStores", SharedCounts.NumStores) << " stores, "
2215 << ore::NV("NumLoads", SharedCounts.NumLoads) << " loads, "
2216 << ore::NV("NumFPOps", SharedCounts.NumComputeOps)
2217 << " compute ops"
2218 << " are shared with other expressions";
2219 }
2220
2221 Rem << ("\n" + linearize(L, Shared, ExprsInSubprogram, DL));
2222 ORE.emit(Rem);
2223 }
2224 }
2225 }
2226
2227 std::string
2228 linearize(Value *L,
2229 const DenseMap<Value *, SmallPtrSet<Value *, 2>> &Shared,
2230 const SmallSetVector<Value *, 32> &ExprsInSubprogram,
2231 const DataLayout &DL) {
2232 ExprLinearizer Lin(DL, Inst2Matrix, Shared, ExprsInSubprogram, L);
2233 Lin.linearizeExpr(L, 0, false, false);
1
Calling 'ExprLinearizer::linearizeExpr'
2234 return Lin.getResult();
2235 }
2236 };
2237};
2238} // namespace
2239
2240PreservedAnalyses LowerMatrixIntrinsicsPass::run(Function &F,
2241 FunctionAnalysisManager &AM) {
2242 auto &TTI = AM.getResult<TargetIRAnalysis>(F);
2243 OptimizationRemarkEmitter *ORE = nullptr;
2244 AAResults *AA = nullptr;
2245 DominatorTree *DT = nullptr;
2246 LoopInfo *LI = nullptr;
2247
2248 if (!Minimal) {
2249 ORE = &AM.getResult<OptimizationRemarkEmitterAnalysis>(F);
2250 AA = &AM.getResult<AAManager>(F);
2251 DT = &AM.getResult<DominatorTreeAnalysis>(F);
2252 LI = &AM.getResult<LoopAnalysis>(F);
2253 }
2254
2255 LowerMatrixIntrinsics LMT(F, TTI, AA, DT, LI, ORE);
2256 if (LMT.Visit()) {
2257 PreservedAnalyses PA;
2258 if (!Minimal) {
2259 PA.preserve<LoopAnalysis>();
2260 PA.preserve<DominatorTreeAnalysis>();
2261 }
2262 return PA;
2263 }
2264 return PreservedAnalyses::all();
2265}
2266
2267namespace {
2268
2269class LowerMatrixIntrinsicsLegacyPass : public FunctionPass {
2270public:
2271 static char ID;
2272
2273 LowerMatrixIntrinsicsLegacyPass() : FunctionPass(ID) {
2274 initializeLowerMatrixIntrinsicsLegacyPassPass(
2275 *PassRegistry::getPassRegistry());
2276 }
2277
2278 bool runOnFunction(Function &F) override {
2279 auto &TTI = getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
2280 auto &ORE = getAnalysis<OptimizationRemarkEmitterWrapperPass>().getORE();
2281 auto &AA = getAnalysis<AAResultsWrapperPass>().getAAResults();
2282 auto &DT = getAnalysis<DominatorTreeWrapperPass>().getDomTree();
2283 auto &LI = getAnalysis<LoopInfoWrapperPass>().getLoopInfo();
2284 LowerMatrixIntrinsics LMT(F, TTI, &AA, &DT, &LI, &ORE);
2285 bool C = LMT.Visit();
2286 return C;
2287 }
2288
2289 void getAnalysisUsage(AnalysisUsage &AU) const override {
2290 AU.addRequired<TargetTransformInfoWrapperPass>();
2291 AU.addRequired<OptimizationRemarkEmitterWrapperPass>();
2292 AU.addRequired<AAResultsWrapperPass>();
2293 AU.addRequired<DominatorTreeWrapperPass>();
2294 AU.addPreserved<DominatorTreeWrapperPass>();
2295 AU.addRequired<LoopInfoWrapperPass>();
2296 AU.addPreserved<LoopInfoWrapperPass>();
2297 }
2298};
2299} // namespace
2300
2301static const char pass_name[] = "Lower the matrix intrinsics";
2302char LowerMatrixIntrinsicsLegacyPass::ID = 0;
2303INITIALIZE_PASS_BEGIN(LowerMatrixIntrinsicsLegacyPass, DEBUG_TYPE, pass_name,static void *initializeLowerMatrixIntrinsicsLegacyPassPassOnce
(PassRegistry &Registry) {
2304 false, false)static void *initializeLowerMatrixIntrinsicsLegacyPassPassOnce
(PassRegistry &Registry) {
2305INITIALIZE_PASS_DEPENDENCY(OptimizationRemarkEmitterWrapperPass)initializeOptimizationRemarkEmitterWrapperPassPass(Registry);
2306INITIALIZE_PASS_DEPENDENCY(AAResultsWrapperPass)initializeAAResultsWrapperPassPass(Registry);
2307INITIALIZE_PASS_DEPENDENCY(DominatorTreeWrapperPass)initializeDominatorTreeWrapperPassPass(Registry);
2308INITIALIZE_PASS_DEPENDENCY(LoopInfoWrapperPass)initializeLoopInfoWrapperPassPass(Registry);
2309INITIALIZE_PASS_END(LowerMatrixIntrinsicsLegacyPass, DEBUG_TYPE, pass_name,PassInfo *PI = new PassInfo( pass_name, "lower-matrix-intrinsics"
, &LowerMatrixIntrinsicsLegacyPass::ID, PassInfo::NormalCtor_t
(callDefaultCtor<LowerMatrixIntrinsicsLegacyPass>), false
, false); Registry.registerPass(*PI, true); return PI; } static
llvm::once_flag InitializeLowerMatrixIntrinsicsLegacyPassPassFlag
; void llvm::initializeLowerMatrixIntrinsicsLegacyPassPass(PassRegistry
&Registry) { llvm::call_once(InitializeLowerMatrixIntrinsicsLegacyPassPassFlag
, initializeLowerMatrixIntrinsicsLegacyPassPassOnce, std::ref
(Registry)); }
2310 false, false)PassInfo *PI = new PassInfo( pass_name, "lower-matrix-intrinsics"
, &LowerMatrixIntrinsicsLegacyPass::ID, PassInfo::NormalCtor_t
(callDefaultCtor<LowerMatrixIntrinsicsLegacyPass>), false
, false); Registry.registerPass(*PI, true); return PI; } static
llvm::once_flag InitializeLowerMatrixIntrinsicsLegacyPassPassFlag
; void llvm::initializeLowerMatrixIntrinsicsLegacyPassPass(PassRegistry
&Registry) { llvm::call_once(InitializeLowerMatrixIntrinsicsLegacyPassPassFlag
, initializeLowerMatrixIntrinsicsLegacyPassPassOnce, std::ref
(Registry)); }
2311
2312Pass *llvm::createLowerMatrixIntrinsicsPass() {
2313 return new LowerMatrixIntrinsicsLegacyPass();
2314}
2315
2316namespace {
2317
2318/// A lightweight version of the matrix lowering pass that only requires TTI.
2319/// Advanced features that require DT, AA or ORE like tiling are disabled. This
2320/// is used to lower matrix intrinsics if the main lowering pass is not run, for
2321/// example with -O0.
2322class LowerMatrixIntrinsicsMinimalLegacyPass : public FunctionPass {
2323public:
2324 static char ID;
2325
2326 LowerMatrixIntrinsicsMinimalLegacyPass() : FunctionPass(ID) {
2327 initializeLowerMatrixIntrinsicsMinimalLegacyPassPass(
2328 *PassRegistry::getPassRegistry());
2329 }
2330
2331 bool runOnFunction(Function &F) override {
2332 auto &TTI = getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F);
2333 LowerMatrixIntrinsics LMT(F, TTI, nullptr, nullptr, nullptr, nullptr);
2334 bool C = LMT.Visit();
2335 return C;
2336 }
2337
2338 void getAnalysisUsage(AnalysisUsage &AU) const override {
2339 AU.addRequired<TargetTransformInfoWrapperPass>();
2340 AU.setPreservesCFG();
2341 }
2342};
2343} // namespace
2344
2345static const char pass_name_minimal[] = "Lower the matrix intrinsics (minimal)";
2346char LowerMatrixIntrinsicsMinimalLegacyPass::ID = 0;
2347INITIALIZE_PASS_BEGIN(LowerMatrixIntrinsicsMinimalLegacyPass,static void *initializeLowerMatrixIntrinsicsMinimalLegacyPassPassOnce
(PassRegistry &Registry) {
2348 "lower-matrix-intrinsics-minimal", pass_name_minimal,static void *initializeLowerMatrixIntrinsicsMinimalLegacyPassPassOnce
(PassRegistry &Registry) {
2349 false, false)static void *initializeLowerMatrixIntrinsicsMinimalLegacyPassPassOnce
(PassRegistry &Registry) {
2350INITIALIZE_PASS_END(LowerMatrixIntrinsicsMinimalLegacyPass,PassInfo *PI = new PassInfo( pass_name_minimal, "lower-matrix-intrinsics-minimal"
, &LowerMatrixIntrinsicsMinimalLegacyPass::ID, PassInfo::
NormalCtor_t(callDefaultCtor<LowerMatrixIntrinsicsMinimalLegacyPass
>), false, false); Registry.registerPass(*PI, true); return
PI; } static llvm::once_flag InitializeLowerMatrixIntrinsicsMinimalLegacyPassPassFlag
; void llvm::initializeLowerMatrixIntrinsicsMinimalLegacyPassPass
(PassRegistry &Registry) { llvm::call_once(InitializeLowerMatrixIntrinsicsMinimalLegacyPassPassFlag
, initializeLowerMatrixIntrinsicsMinimalLegacyPassPassOnce, std
::ref(Registry)); }
2351 "lower-matrix-intrinsics-minimal", pass_name_minimal, false,PassInfo *PI = new PassInfo( pass_name_minimal, "lower-matrix-intrinsics-minimal"
, &LowerMatrixIntrinsicsMinimalLegacyPass::ID, PassInfo::
NormalCtor_t(callDefaultCtor<LowerMatrixIntrinsicsMinimalLegacyPass
>), false, false); Registry.registerPass(*PI, true); return
PI; } static llvm::once_flag InitializeLowerMatrixIntrinsicsMinimalLegacyPassPassFlag
; void llvm::initializeLowerMatrixIntrinsicsMinimalLegacyPassPass
(PassRegistry &Registry) { llvm::call_once(InitializeLowerMatrixIntrinsicsMinimalLegacyPassPassFlag
, initializeLowerMatrixIntrinsicsMinimalLegacyPassPassOnce, std
::ref(Registry)); }
2352 false)PassInfo *PI = new PassInfo( pass_name_minimal, "lower-matrix-intrinsics-minimal"
, &LowerMatrixIntrinsicsMinimalLegacyPass::ID, PassInfo::
NormalCtor_t(callDefaultCtor<LowerMatrixIntrinsicsMinimalLegacyPass
>), false, false); Registry.registerPass(*PI, true); return
PI; } static llvm::once_flag InitializeLowerMatrixIntrinsicsMinimalLegacyPassPassFlag
; void llvm::initializeLowerMatrixIntrinsicsMinimalLegacyPassPass
(PassRegistry &Registry) { llvm::call_once(InitializeLowerMatrixIntrinsicsMinimalLegacyPassPassFlag
, initializeLowerMatrixIntrinsicsMinimalLegacyPassPassOnce, std
::ref(Registry)); }
2353
2354Pass *llvm::createLowerMatrixIntrinsicsMinimalPass() {
2355 return new LowerMatrixIntrinsicsMinimalLegacyPass();
2356}