brintos / llvm-project-archived public Read only

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