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