| //===- LowerMatrixIntrinsics.cpp - Lower matrix intrinsics -----*- C++ -*-===// |
| // |
| // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions. |
| // See https://llvm.org/LICENSE.txt for license information. |
| // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception |
| // |
| //===----------------------------------------------------------------------===// |
| // |
| // Lower matrix intrinsics to vector operations. |
| // |
| // TODO: |
| // * Implement multiply & add fusion |
| // * Add remark, summarizing the available matrix optimization opportunities. |
| // |
| //===----------------------------------------------------------------------===// |
| |
| #include "llvm/Transforms/Scalar/LowerMatrixIntrinsics.h" |
| #include "llvm/ADT/GraphTraits.h" |
| #include "llvm/ADT/PostOrderIterator.h" |
| #include "llvm/ADT/SmallVector.h" |
| #include "llvm/Analysis/TargetTransformInfo.h" |
| #include "llvm/Analysis/VectorUtils.h" |
| #include "llvm/IR/CFG.h" |
| #include "llvm/IR/DataLayout.h" |
| #include "llvm/IR/Function.h" |
| #include "llvm/IR/IRBuilder.h" |
| #include "llvm/IR/Instructions.h" |
| #include "llvm/IR/IntrinsicInst.h" |
| #include "llvm/IR/PatternMatch.h" |
| #include "llvm/InitializePasses.h" |
| #include "llvm/Pass.h" |
| #include "llvm/Support/Debug.h" |
| #include "llvm/Transforms/Scalar.h" |
| |
| using namespace llvm; |
| using namespace PatternMatch; |
| |
| #define DEBUG_TYPE "lower-matrix-intrinsics" |
| |
| static cl::opt<bool> EnableShapePropagation("matrix-propagate-shape", |
| cl::init(true)); |
| |
| static cl::opt<bool> AllowContractEnabled( |
| "matrix-allow-contract", cl::init(false), cl::Hidden, |
| cl::desc("Allow the use of FMAs if available and profitable. This may " |
| "result in different results, due to less rounding error.")); |
| |
| namespace { |
| |
| // Given an element poitner \p BasePtr to the start of a (sub) matrix, compute |
| // the start address of column \p Col with type (\p EltType x \p NumRows) |
| // assuming \p Stride elements between start two consecutive columns. |
| // \p Stride must be >= \p NumRows. |
| // |
| // Consider a 4x4 matrix like below |
| // |
| // 0 1 2 3 |
| // 0 v_0_0 v_0_1 v_0_2 v_0_3 |
| // 1 v_1_0 v_1_1 v_1_2 v_1_3 |
| // 2 v_2_0 v_2_1 v_2_2 v_2_3 |
| // 3 v_3_0 v_3_1 v_3_2 v_3_3 |
| |
| // To compute the column addresses for a 2x3 sub-matrix at row 1 and column 1, |
| // we need a pointer to the first element of the submatrix as base pointer. |
| // Then we can use computeColumnAddr to compute the addresses for the columns |
| // of the sub-matrix. |
| // |
| // Column 0: computeColumnAddr(Base, 0 (column), 4 (stride), 2 (num rows), ..) |
| // -> just returns Base |
| // Column 1: computeColumnAddr(Base, 1 (column), 4 (stride), 2 (num rows), ..) |
| // -> returns Base + (1 * 4) |
| // Column 2: computeColumnAddr(Base, 2 (column), 4 (stride), 2 (num rows), ..) |
| // -> returns Base + (2 * 4) |
| // |
| // The graphic below illustrates the number of elements in a column (marked |
| // with |) and the number of skipped elements (marked with }). |
| // |
| // v_0_0 v_0_1 {v_0_2 {v_0_3 |
| // Base Col 1 Col 2 |
| // | | | |
| // v_1_0 |v_1_1 |v_1_2 |v_1_3 |
| // v_2_0 |v_2_1 |v_2_2 |v_2_3 |
| // v_3_0 {v_3_1 {v_3_2 v_3_3 |
| // |
| Value *computeColumnAddr(Value *BasePtr, Value *Col, Value *Stride, |
| unsigned NumRows, Type *EltType, |
| IRBuilder<> &Builder) { |
| |
| assert((!isa<ConstantInt>(Stride) || |
| cast<ConstantInt>(Stride)->getZExtValue() >= NumRows) && |
| "Stride must be >= the number of rows."); |
| unsigned AS = cast<PointerType>(BasePtr->getType())->getAddressSpace(); |
| |
| // Compute the start of the column with index Col as Col * Stride. |
| Value *ColumnStart = Builder.CreateMul(Col, Stride, "col.start"); |
| |
| // Get pointer to the start of the selected column. Skip GEP creation, |
| // if we select column 0. |
| if (isa<ConstantInt>(ColumnStart) && cast<ConstantInt>(ColumnStart)->isZero()) |
| ColumnStart = BasePtr; |
| else |
| ColumnStart = Builder.CreateGEP(EltType, BasePtr, ColumnStart, "col.gep"); |
| |
| // Cast elementwise column start pointer to a pointer to a column |
| // (EltType x NumRows)*. |
| Type *ColumnType = VectorType::get(EltType, NumRows); |
| Type *ColumnPtrType = PointerType::get(ColumnType, AS); |
| return Builder.CreatePointerCast(ColumnStart, ColumnPtrType, "col.cast"); |
| } |
| |
| /// LowerMatrixIntrinsics contains the methods used to lower matrix intrinsics. |
| /// |
| /// Currently, the lowering for each matrix intrinsic is done as follows: |
| /// 1. Propagate the shape information from intrinsics to connected |
| /// instructions. |
| /// 2. Lower instructions with shape information. |
| /// 2.1. Get column vectors for each argument. If we already lowered the |
| /// definition of an argument, use the produced column vectors directly. |
| /// If not, split the operand vector containing an embedded matrix into |
| /// a set of column vectors, |
| /// 2.2. Lower the instruction in terms of columnwise operations, which yields |
| /// a set of column vectors containing result matrix. Note that we lower |
| /// all instructions that have shape information. Besides the intrinsics, |
| /// this includes stores for example. |
| /// 2.3. Update uses of the lowered instruction. If we have shape information |
| /// for a user, there is nothing to do, as we will look up the result |
| /// column matrix when lowering the user. For other uses, we embed the |
| /// result matrix in a flat vector and update the use. |
| /// 2.4. Cache the result column matrix for the instruction we lowered |
| /// 3. After we lowered all instructions in a function, remove the now |
| /// obsolete instructions. |
| /// |
| class LowerMatrixIntrinsics { |
| Function &Func; |
| const DataLayout &DL; |
| const TargetTransformInfo &TTI; |
| |
| /// Wrapper class representing a matrix as a set of column vectors. |
| /// All column vectors must have the same vector type. |
| class ColumnMatrixTy { |
| SmallVector<Value *, 16> Columns; |
| |
| public: |
| ColumnMatrixTy() : Columns() {} |
| ColumnMatrixTy(ArrayRef<Value *> Cols) |
| : Columns(Cols.begin(), Cols.end()) {} |
| |
| Value *getColumn(unsigned i) const { return Columns[i]; } |
| |
| void setColumn(unsigned i, Value *V) { Columns[i] = V; } |
| |
| size_t getNumColumns() const { return Columns.size(); } |
| size_t getNumRows() const { |
| assert(Columns.size() > 0 && "Cannot call getNumRows without columns"); |
| return cast<VectorType>(Columns[0]->getType())->getNumElements(); |
| } |
| |
| const SmallVectorImpl<Value *> &getColumnVectors() const { return Columns; } |
| |
| SmallVectorImpl<Value *> &getColumnVectors() { return Columns; } |
| |
| void addColumn(Value *V) { Columns.push_back(V); } |
| |
| iterator_range<SmallVector<Value *, 8>::iterator> columns() { |
| return make_range(Columns.begin(), Columns.end()); |
| } |
| |
| /// Embed the columns of the matrix into a flat vector by concatenating |
| /// them. |
| Value *embedInVector(IRBuilder<> &Builder) const { |
| return Columns.size() == 1 ? Columns[0] |
| : concatenateVectors(Builder, Columns); |
| } |
| }; |
| |
| struct ShapeInfo { |
| unsigned NumRows; |
| unsigned NumColumns; |
| |
| ShapeInfo(unsigned NumRows = 0, unsigned NumColumns = 0) |
| : NumRows(NumRows), NumColumns(NumColumns) {} |
| |
| ShapeInfo(Value *NumRows, Value *NumColumns) |
| : NumRows(cast<ConstantInt>(NumRows)->getZExtValue()), |
| NumColumns(cast<ConstantInt>(NumColumns)->getZExtValue()) {} |
| |
| bool operator==(const ShapeInfo &other) { |
| return NumRows == other.NumRows && NumColumns == other.NumColumns; |
| } |
| bool operator!=(const ShapeInfo &other) { return !(*this == other); } |
| |
| /// Returns true if shape-information is defined, meaning both dimensions |
| /// are != 0. |
| operator bool() const { |
| assert(NumRows == 0 || NumColumns != 0); |
| return NumRows != 0; |
| } |
| }; |
| |
| /// Maps instructions to their shape information. The shape information |
| /// describes the shape to be used while lowering. This matches the shape of |
| /// the result value of the instruction, with the only exceptions being store |
| /// instructions and the matrix_columnwise_store intrinsics. For those, the |
| /// shape information indicates that those instructions should be lowered |
| /// using shape information as well. |
| DenseMap<Value *, ShapeInfo> ShapeMap; |
| |
| /// List of instructions to remove. While lowering, we are not replacing all |
| /// users of a lowered instruction, if shape information is available and |
| /// those need to be removed after we finished lowering. |
| SmallVector<Instruction *, 16> ToRemove; |
| |
| /// Map from instructions to their produced column matrix. |
| DenseMap<Value *, ColumnMatrixTy> Inst2ColumnMatrix; |
| |
| public: |
| LowerMatrixIntrinsics(Function &F, TargetTransformInfo &TTI) |
| : Func(F), DL(F.getParent()->getDataLayout()), TTI(TTI) {} |
| |
| /// Return the set of column vectors that a matrix value is lowered to. |
| /// |
| /// If we lowered \p MatrixVal, just return the cache result column matrix. |
| /// Otherwie split the flat vector \p MatrixVal containing a matrix with |
| /// shape \p SI into column vectors. |
| ColumnMatrixTy getMatrix(Value *MatrixVal, const ShapeInfo &SI, |
| IRBuilder<> Builder) { |
| VectorType *VType = dyn_cast<VectorType>(MatrixVal->getType()); |
| assert(VType && "MatrixVal must be a vector type"); |
| assert(VType->getNumElements() == SI.NumRows * SI.NumColumns && |
| "The vector size must match the number of matrix elements"); |
| |
| // Check if we lowered MatrixVal using shape information. In that case, |
| // return the existing column matrix, if it matches the requested shape |
| // information. If there is a mis-match, embed the result in a flat |
| // vector and split it later. |
| auto Found = Inst2ColumnMatrix.find(MatrixVal); |
| if (Found != Inst2ColumnMatrix.end()) { |
| ColumnMatrixTy &M = Found->second; |
| // Return the found matrix, if its shape matches the requested shape |
| // information |
| if (SI.NumRows == M.getNumRows() && SI.NumColumns == M.getNumColumns()) |
| return M; |
| |
| MatrixVal = M.embedInVector(Builder); |
| } |
| |
| // Otherwise split MatrixVal. |
| SmallVector<Value *, 16> SplitVecs; |
| Value *Undef = UndefValue::get(VType); |
| for (unsigned MaskStart = 0; MaskStart < VType->getNumElements(); |
| MaskStart += SI.NumRows) { |
| Constant *Mask = createSequentialMask(Builder, MaskStart, SI.NumRows, 0); |
| Value *V = Builder.CreateShuffleVector(MatrixVal, Undef, Mask, "split"); |
| SplitVecs.push_back(V); |
| } |
| |
| return {SplitVecs}; |
| } |
| |
| /// If \p V already has a known shape return false. Otherwise set the shape |
| /// for instructions that support it. |
| bool setShapeInfo(Value *V, ShapeInfo Shape) { |
| assert(Shape && "Shape not set"); |
| if (isa<UndefValue>(V) || !supportsShapeInfo(V)) |
| return false; |
| |
| auto SIter = ShapeMap.find(V); |
| if (SIter != ShapeMap.end()) { |
| LLVM_DEBUG(dbgs() << " not overriding existing shape: " |
| << SIter->second.NumRows << " " |
| << SIter->second.NumColumns << " for " << *V << "\n"); |
| return false; |
| } |
| |
| ShapeMap.insert({V, Shape}); |
| LLVM_DEBUG(dbgs() << " " << Shape.NumRows << " x " << Shape.NumColumns |
| << " for " << *V << "\n"); |
| return true; |
| } |
| |
| bool isUniformShape(Value *V) { |
| Instruction *I = dyn_cast<Instruction>(V); |
| if (!I) |
| return true; |
| |
| switch (I->getOpcode()) { |
| case Instruction::FAdd: |
| case Instruction::FSub: |
| case Instruction::FMul: // Scalar multiply. |
| case Instruction::Add: |
| case Instruction::Mul: |
| case Instruction::Sub: |
| return true; |
| default: |
| return false; |
| } |
| } |
| |
| /// Returns true if shape information can be used for \p V. The supported |
| /// instructions must match the instructions that can be lowered by this pass. |
| bool supportsShapeInfo(Value *V) { |
| Instruction *Inst = dyn_cast<Instruction>(V); |
| if (!Inst) |
| return false; |
| |
| IntrinsicInst *II = dyn_cast<IntrinsicInst>(Inst); |
| if (II) |
| switch (II->getIntrinsicID()) { |
| case Intrinsic::matrix_multiply: |
| case Intrinsic::matrix_transpose: |
| case Intrinsic::matrix_columnwise_load: |
| case Intrinsic::matrix_columnwise_store: |
| return true; |
| default: |
| return false; |
| } |
| return isUniformShape(V) || isa<StoreInst>(V) || isa<LoadInst>(V); |
| } |
| |
| /// Propagate the shape information of instructions to their users. |
| /// The work list contains instructions for which we can compute the shape, |
| /// either based on the information provided by matrix intrinsics or known |
| /// shapes of operands. |
| SmallVector<Instruction *, 32> |
| propagateShapeForward(SmallVectorImpl<Instruction *> &WorkList) { |
| SmallVector<Instruction *, 32> NewWorkList; |
| // Pop an element for which we guaranteed to have at least one of the |
| // operand shapes. Add the shape for this and then add users to the work |
| // list. |
| LLVM_DEBUG(dbgs() << "Forward-propagate shapes:\n"); |
| while (!WorkList.empty()) { |
| Instruction *Inst = WorkList.back(); |
| WorkList.pop_back(); |
| |
| // New entry, set the value and insert operands |
| bool Propagate = false; |
| |
| Value *MatrixA; |
| Value *MatrixB; |
| Value *M; |
| Value *N; |
| Value *K; |
| if (match(Inst, m_Intrinsic<Intrinsic::matrix_multiply>( |
| m_Value(MatrixA), m_Value(MatrixB), m_Value(M), |
| m_Value(N), m_Value(K)))) { |
| Propagate = setShapeInfo(Inst, {M, K}); |
| } else if (match(Inst, m_Intrinsic<Intrinsic::matrix_transpose>( |
| m_Value(MatrixA), m_Value(M), m_Value(N)))) { |
| // Flip dimensions. |
| Propagate = setShapeInfo(Inst, {N, M}); |
| } else if (match(Inst, m_Intrinsic<Intrinsic::matrix_columnwise_store>( |
| m_Value(MatrixA), m_Value(), m_Value(), |
| m_Value(M), m_Value(N)))) { |
| Propagate = setShapeInfo(Inst, {N, M}); |
| } else if (match(Inst, |
| m_Intrinsic<Intrinsic::matrix_columnwise_load>( |
| m_Value(), m_Value(), m_Value(M), m_Value(N)))) { |
| Propagate = setShapeInfo(Inst, {M, N}); |
| } else if (match(Inst, m_Store(m_Value(MatrixA), m_Value()))) { |
| auto OpShape = ShapeMap.find(MatrixA); |
| if (OpShape != ShapeMap.end()) |
| setShapeInfo(Inst, OpShape->second); |
| continue; |
| } else if (isUniformShape(Inst)) { |
| // Find the first operand that has a known shape and use that. |
| for (auto &Op : Inst->operands()) { |
| auto OpShape = ShapeMap.find(Op.get()); |
| if (OpShape != ShapeMap.end()) { |
| Propagate |= setShapeInfo(Inst, OpShape->second); |
| break; |
| } |
| } |
| } |
| |
| if (Propagate) { |
| NewWorkList.push_back(Inst); |
| for (auto *User : Inst->users()) |
| if (ShapeMap.count(User) == 0) |
| WorkList.push_back(cast<Instruction>(User)); |
| } |
| } |
| |
| return NewWorkList; |
| } |
| |
| /// Propagate the shape to operands of instructions with shape information. |
| /// \p Worklist contains the instruction for which we already know the shape. |
| SmallVector<Instruction *, 32> |
| propagateShapeBackward(SmallVectorImpl<Instruction *> &WorkList) { |
| SmallVector<Instruction *, 32> NewWorkList; |
| |
| auto pushInstruction = [](Value *V, |
| SmallVectorImpl<Instruction *> &WorkList) { |
| Instruction *I = dyn_cast<Instruction>(V); |
| if (I) |
| WorkList.push_back(I); |
| }; |
| // Pop an element with known shape. Traverse the operands, if their shape |
| // derives from the result shape and is unknown, add it and add them to the |
| // worklist. |
| LLVM_DEBUG(dbgs() << "Backward-propagate shapes:\n"); |
| while (!WorkList.empty()) { |
| Value *V = WorkList.back(); |
| WorkList.pop_back(); |
| |
| size_t BeforeProcessingV = WorkList.size(); |
| if (!isa<Instruction>(V)) |
| continue; |
| |
| Value *MatrixA; |
| Value *MatrixB; |
| Value *M; |
| Value *N; |
| Value *K; |
| if (match(V, m_Intrinsic<Intrinsic::matrix_multiply>( |
| m_Value(MatrixA), m_Value(MatrixB), m_Value(M), |
| m_Value(N), m_Value(K)))) { |
| if (setShapeInfo(MatrixA, {M, N})) |
| pushInstruction(MatrixA, WorkList); |
| |
| if (setShapeInfo(MatrixB, {N, K})) |
| pushInstruction(MatrixB, WorkList); |
| |
| } else if (match(V, m_Intrinsic<Intrinsic::matrix_transpose>( |
| m_Value(MatrixA), m_Value(M), m_Value(N)))) { |
| // Flip dimensions. |
| if (setShapeInfo(MatrixA, {M, N})) |
| pushInstruction(MatrixA, WorkList); |
| } else if (match(V, m_Intrinsic<Intrinsic::matrix_columnwise_store>( |
| m_Value(MatrixA), m_Value(), m_Value(), |
| m_Value(M), m_Value(N)))) { |
| if (setShapeInfo(MatrixA, {M, N})) { |
| pushInstruction(MatrixA, WorkList); |
| } |
| } else if (isa<LoadInst>(V) || |
| match(V, m_Intrinsic<Intrinsic::matrix_columnwise_load>())) { |
| // Nothing to do, no matrix input. |
| } else if (isa<StoreInst>(V)) { |
| // Nothing to do. We forward-propagated to this so we would just |
| // backward propagate to an instruction with an already known shape. |
| } else if (isUniformShape(V)) { |
| // Propagate to all operands. |
| ShapeInfo Shape = ShapeMap[V]; |
| for (Use &U : cast<Instruction>(V)->operands()) { |
| if (setShapeInfo(U.get(), Shape)) |
| pushInstruction(U.get(), WorkList); |
| } |
| } |
| // After we discovered new shape info for new instructions in the |
| // worklist, we use their users as seeds for the next round of forward |
| // propagation. |
| for (size_t I = BeforeProcessingV; I != WorkList.size(); I++) |
| for (User *U : WorkList[I]->users()) |
| if (isa<Instruction>(U) && V != U) |
| NewWorkList.push_back(cast<Instruction>(U)); |
| } |
| return NewWorkList; |
| } |
| |
| bool Visit() { |
| if (EnableShapePropagation) { |
| SmallVector<Instruction *, 32> WorkList; |
| |
| // Initially only the shape of matrix intrinsics is known. |
| // Initialize the work list with ops carrying shape information. |
| for (BasicBlock &BB : Func) |
| for (Instruction &Inst : BB) { |
| IntrinsicInst *II = dyn_cast<IntrinsicInst>(&Inst); |
| if (!II) |
| continue; |
| |
| switch (II->getIntrinsicID()) { |
| case Intrinsic::matrix_multiply: |
| case Intrinsic::matrix_transpose: |
| case Intrinsic::matrix_columnwise_load: |
| case Intrinsic::matrix_columnwise_store: |
| WorkList.push_back(&Inst); |
| break; |
| default: |
| break; |
| } |
| } |
| // Propagate shapes until nothing changes any longer. |
| while (!WorkList.empty()) { |
| WorkList = propagateShapeForward(WorkList); |
| WorkList = propagateShapeBackward(WorkList); |
| } |
| } |
| |
| ReversePostOrderTraversal<Function *> RPOT(&Func); |
| bool Changed = false; |
| for (auto *BB : RPOT) { |
| for (Instruction &Inst : make_early_inc_range(*BB)) { |
| IRBuilder<> Builder(&Inst); |
| |
| if (CallInst *CInst = dyn_cast<CallInst>(&Inst)) |
| Changed |= VisitCallInst(CInst); |
| |
| Value *Op1; |
| Value *Op2; |
| if (auto *BinOp = dyn_cast<BinaryOperator>(&Inst)) |
| Changed |= VisitBinaryOperator(BinOp); |
| if (match(&Inst, m_Load(m_Value(Op1)))) |
| Changed |= VisitLoad(&Inst, Op1, Builder); |
| else if (match(&Inst, m_Store(m_Value(Op1), m_Value(Op2)))) |
| Changed |= VisitStore(&Inst, Op1, Op2, Builder); |
| } |
| } |
| |
| for (Instruction *Inst : reverse(ToRemove)) |
| Inst->eraseFromParent(); |
| |
| return Changed; |
| } |
| |
| LoadInst *createColumnLoad(Value *ColumnPtr, Type *EltType, |
| IRBuilder<> Builder) { |
| unsigned Align = DL.getABITypeAlignment(EltType); |
| return Builder.CreateAlignedLoad(ColumnPtr, Align, "col.load"); |
| } |
| |
| StoreInst *createColumnStore(Value *ColumnValue, Value *ColumnPtr, |
| Type *EltType, IRBuilder<> Builder) { |
| unsigned Align = DL.getABITypeAlignment(EltType); |
| return Builder.CreateAlignedStore(ColumnValue, ColumnPtr, Align); |
| } |
| |
| |
| /// Turns \p BasePtr into an elementwise pointer to \p EltType. |
| Value *createElementPtr(Value *BasePtr, Type *EltType, IRBuilder<> &Builder) { |
| unsigned AS = cast<PointerType>(BasePtr->getType())->getAddressSpace(); |
| Type *EltPtrType = PointerType::get(EltType, AS); |
| return Builder.CreatePointerCast(BasePtr, EltPtrType); |
| } |
| |
| /// Replace intrinsic calls |
| bool VisitCallInst(CallInst *Inst) { |
| if (!Inst->getCalledFunction() || !Inst->getCalledFunction()->isIntrinsic()) |
| return false; |
| |
| switch (Inst->getCalledFunction()->getIntrinsicID()) { |
| case Intrinsic::matrix_multiply: |
| LowerMultiply(Inst); |
| break; |
| case Intrinsic::matrix_transpose: |
| LowerTranspose(Inst); |
| break; |
| case Intrinsic::matrix_columnwise_load: |
| LowerColumnwiseLoad(Inst); |
| break; |
| case Intrinsic::matrix_columnwise_store: |
| LowerColumnwiseStore(Inst); |
| break; |
| default: |
| return false; |
| } |
| return true; |
| } |
| |
| void LowerLoad(Instruction *Inst, Value *Ptr, Value *Stride, |
| ShapeInfo Shape) { |
| IRBuilder<> Builder(Inst); |
| auto VType = cast<VectorType>(Inst->getType()); |
| Value *EltPtr = createElementPtr(Ptr, VType->getElementType(), Builder); |
| ColumnMatrixTy Result; |
| // Distance between start of one column and the start of the next |
| for (unsigned C = 0, E = Shape.NumColumns; C < E; ++C) { |
| Value *GEP = |
| computeColumnAddr(EltPtr, Builder.getInt32(C), Stride, Shape.NumRows, |
| VType->getElementType(), Builder); |
| Value *Column = createColumnLoad(GEP, VType->getElementType(), Builder); |
| Result.addColumn(Column); |
| } |
| |
| finalizeLowering(Inst, Result, Builder); |
| } |
| |
| /// Lowers llvm.matrix.columnwise.load. |
| /// |
| /// The intrinsic loads a matrix from memory using a stride between columns. |
| void LowerColumnwiseLoad(CallInst *Inst) { |
| Value *Ptr = Inst->getArgOperand(0); |
| Value *Stride = Inst->getArgOperand(1); |
| LowerLoad(Inst, Ptr, Stride, |
| {Inst->getArgOperand(2), Inst->getArgOperand(3)}); |
| } |
| |
| void LowerStore(Instruction *Inst, Value *Matrix, Value *Ptr, Value *Stride, |
| ShapeInfo Shape) { |
| IRBuilder<> Builder(Inst); |
| auto VType = cast<VectorType>(Matrix->getType()); |
| Value *EltPtr = createElementPtr(Ptr, VType->getElementType(), Builder); |
| auto LM = getMatrix(Matrix, Shape, Builder); |
| for (auto C : enumerate(LM.columns())) { |
| Value *GEP = |
| computeColumnAddr(EltPtr, Builder.getInt32(C.index()), Stride, |
| Shape.NumRows, VType->getElementType(), Builder); |
| createColumnStore(C.value(), GEP, VType->getElementType(), Builder); |
| } |
| |
| ToRemove.push_back(Inst); |
| } |
| |
| /// Lowers llvm.matrix.columnwise.store. |
| /// |
| /// The intrinsic store a matrix back memory using a stride between columns. |
| void LowerColumnwiseStore(CallInst *Inst) { |
| Value *Matrix = Inst->getArgOperand(0); |
| Value *Ptr = Inst->getArgOperand(1); |
| Value *Stride = Inst->getArgOperand(2); |
| LowerStore(Inst, Matrix, Ptr, Stride, |
| {Inst->getArgOperand(3), Inst->getArgOperand(4)}); |
| } |
| |
| /// Extract a column vector of \p NumElts starting at index (\p I, \p J) from |
| /// the matrix \p LM represented as a vector of column vectors. |
| Value *extractVector(const ColumnMatrixTy &LM, unsigned I, unsigned J, |
| unsigned NumElts, IRBuilder<> Builder) { |
| Value *Col = LM.getColumn(J); |
| Value *Undef = UndefValue::get(Col->getType()); |
| Constant *Mask = createSequentialMask(Builder, I, NumElts, 0); |
| return Builder.CreateShuffleVector(Col, Undef, Mask, "block"); |
| } |
| |
| // Set elements I..I+NumElts-1 to Block |
| Value *insertVector(Value *Col, unsigned I, Value *Block, |
| IRBuilder<> Builder) { |
| |
| // First, bring Block to the same size as Col |
| unsigned BlockNumElts = |
| cast<VectorType>(Block->getType())->getNumElements(); |
| unsigned NumElts = cast<VectorType>(Col->getType())->getNumElements(); |
| assert(NumElts >= BlockNumElts && "Too few elements for current block"); |
| |
| Value *ExtendMask = |
| createSequentialMask(Builder, 0, BlockNumElts, NumElts - BlockNumElts); |
| Value *Undef = UndefValue::get(Block->getType()); |
| Block = Builder.CreateShuffleVector(Block, Undef, ExtendMask); |
| |
| // If Col is 7 long and I is 2 and BlockNumElts is 2 the mask is: 0, 1, 7, |
| // 8, 4, 5, 6 |
| SmallVector<Constant *, 16> Mask; |
| unsigned i; |
| for (i = 0; i < I; i++) |
| Mask.push_back(Builder.getInt32(i)); |
| |
| unsigned VecNumElts = cast<VectorType>(Col->getType())->getNumElements(); |
| for (; i < I + BlockNumElts; i++) |
| Mask.push_back(Builder.getInt32(i - I + VecNumElts)); |
| |
| for (; i < VecNumElts; i++) |
| Mask.push_back(Builder.getInt32(i)); |
| |
| Value *MaskVal = ConstantVector::get(Mask); |
| |
| return Builder.CreateShuffleVector(Col, Block, MaskVal); |
| } |
| |
| Value *createMulAdd(Value *Sum, Value *A, Value *B, bool UseFPOp, |
| IRBuilder<> &Builder, bool AllowContraction) { |
| |
| if (!Sum) |
| return UseFPOp ? Builder.CreateFMul(A, B) : Builder.CreateMul(A, B); |
| |
| if (UseFPOp) { |
| if (AllowContraction) { |
| // Use fmuladd for floating point operations and let the backend decide |
| // if that's profitable. |
| Value *FMulAdd = Intrinsic::getDeclaration( |
| Func.getParent(), Intrinsic::fmuladd, A->getType()); |
| return Builder.CreateCall(FMulAdd, {A, B, Sum}); |
| } |
| Value *Mul = Builder.CreateFMul(A, B); |
| return Builder.CreateFAdd(Sum, Mul); |
| } |
| |
| Value *Mul = Builder.CreateMul(A, B); |
| return Builder.CreateAdd(Sum, Mul); |
| } |
| |
| /// Cache \p Matrix as result of \p Inst and update the uses of \p Inst. For |
| /// users with shape information, there's nothing to do: the will use the |
| /// cached value when they are lowered. For other users, \p Matrix is |
| /// flattened and the uses are updated to use it. Also marks \p Inst for |
| /// deletion. |
| void finalizeLowering(Instruction *Inst, ColumnMatrixTy Matrix, |
| IRBuilder<> &Builder) { |
| Inst2ColumnMatrix.insert(std::make_pair(Inst, Matrix)); |
| |
| ToRemove.push_back(Inst); |
| Value *Flattened = nullptr; |
| for (auto I = Inst->use_begin(), E = Inst->use_end(); I != E;) { |
| Use &U = *I++; |
| if (ShapeMap.find(U.getUser()) == ShapeMap.end()) { |
| if (!Flattened) |
| Flattened = Matrix.embedInVector(Builder); |
| U.set(Flattened); |
| } |
| } |
| } |
| |
| /// Lowers llvm.matrix.multiply. |
| void LowerMultiply(CallInst *MatMul) { |
| IRBuilder<> Builder(MatMul); |
| auto *EltType = cast<VectorType>(MatMul->getType())->getElementType(); |
| ShapeInfo LShape(MatMul->getArgOperand(2), MatMul->getArgOperand(3)); |
| ShapeInfo RShape(MatMul->getArgOperand(3), MatMul->getArgOperand(4)); |
| |
| const ColumnMatrixTy &Lhs = |
| getMatrix(MatMul->getArgOperand(0), LShape, Builder); |
| const ColumnMatrixTy &Rhs = |
| getMatrix(MatMul->getArgOperand(1), RShape, Builder); |
| |
| const unsigned R = LShape.NumRows; |
| const unsigned M = LShape.NumColumns; |
| const unsigned C = RShape.NumColumns; |
| assert(M == RShape.NumRows); |
| |
| // Initialize the output |
| ColumnMatrixTy Result; |
| for (unsigned J = 0; J < C; ++J) |
| Result.addColumn(UndefValue::get(VectorType::get(EltType, R))); |
| |
| const unsigned VF = std::max(TTI.getRegisterBitWidth(true) / |
| EltType->getPrimitiveSizeInBits(), |
| uint64_t(1)); |
| |
| bool AllowContract = AllowContractEnabled || (isa<FPMathOperator>(MatMul) && |
| MatMul->hasAllowContract()); |
| // Multiply columns from the first operand with scalars from the second |
| // operand. Then move along the K axes and accumulate the columns. With |
| // this the adds can be vectorized without reassociation. |
| for (unsigned J = 0; J < C; ++J) { |
| unsigned BlockSize = VF; |
| for (unsigned I = 0; I < R; I += BlockSize) { |
| // Gradually lower the vectorization factor to cover the remainder. |
| while (I + BlockSize > R) |
| BlockSize /= 2; |
| |
| Value *Sum = nullptr; |
| for (unsigned K = 0; K < M; ++K) { |
| Value *L = extractVector(Lhs, I, K, BlockSize, Builder); |
| Value *RH = Builder.CreateExtractElement(Rhs.getColumn(J), K); |
| Value *Splat = Builder.CreateVectorSplat(BlockSize, RH, "splat"); |
| Sum = createMulAdd(Sum, L, Splat, EltType->isFloatingPointTy(), |
| Builder, AllowContract); |
| } |
| Result.setColumn(J, insertVector(Result.getColumn(J), I, Sum, Builder)); |
| } |
| } |
| finalizeLowering(MatMul, Result, Builder); |
| } |
| |
| /// Lowers llvm.matrix.transpose. |
| void LowerTranspose(CallInst *Inst) { |
| ColumnMatrixTy Result; |
| IRBuilder<> Builder(Inst); |
| Value *InputVal = Inst->getArgOperand(0); |
| VectorType *VectorTy = cast<VectorType>(InputVal->getType()); |
| ShapeInfo ArgShape(Inst->getArgOperand(1), Inst->getArgOperand(2)); |
| ColumnMatrixTy InputMatrix = getMatrix(InputVal, ArgShape, Builder); |
| |
| for (unsigned Row = 0; Row < ArgShape.NumRows; ++Row) { |
| // Build a single column vector for this row. First initialize it. |
| Value *ResultColumn = UndefValue::get( |
| VectorType::get(VectorTy->getElementType(), ArgShape.NumColumns)); |
| |
| // Go through the elements of this row and insert it into the resulting |
| // column vector. |
| for (auto C : enumerate(InputMatrix.columns())) { |
| Value *Elt = Builder.CreateExtractElement(C.value(), Row); |
| // We insert at index Column since that is the row index after the |
| // transpose. |
| ResultColumn = |
| Builder.CreateInsertElement(ResultColumn, Elt, C.index()); |
| } |
| Result.addColumn(ResultColumn); |
| } |
| |
| finalizeLowering(Inst, Result, Builder); |
| } |
| |
| /// Lower load instructions, if shape information is available. |
| bool VisitLoad(Instruction *Inst, Value *Ptr, IRBuilder<> &Builder) { |
| auto I = ShapeMap.find(Inst); |
| if (I == ShapeMap.end()) |
| return false; |
| |
| LowerLoad(Inst, Ptr, Builder.getInt32(I->second.NumRows), I->second); |
| return true; |
| } |
| |
| bool VisitStore(Instruction *Inst, Value *StoredVal, Value *Ptr, |
| IRBuilder<> &Builder) { |
| auto I = ShapeMap.find(StoredVal); |
| if (I == ShapeMap.end()) |
| return false; |
| |
| LowerStore(Inst, StoredVal, Ptr, Builder.getInt32(I->second.NumRows), I->second); |
| return true; |
| } |
| |
| /// Lower binary operators, if shape information is available. |
| bool VisitBinaryOperator(BinaryOperator *Inst) { |
| auto I = ShapeMap.find(Inst); |
| if (I == ShapeMap.end()) |
| return false; |
| |
| Value *Lhs = Inst->getOperand(0); |
| Value *Rhs = Inst->getOperand(1); |
| |
| IRBuilder<> Builder(Inst); |
| ShapeInfo &Shape = I->second; |
| |
| ColumnMatrixTy LoweredLhs = getMatrix(Lhs, Shape, Builder); |
| ColumnMatrixTy LoweredRhs = getMatrix(Rhs, Shape, Builder); |
| |
| // Add each column and store the result back into the opmapping |
| ColumnMatrixTy Result; |
| auto BuildColumnOp = [&Builder, Inst](Value *LHS, Value *RHS) { |
| switch (Inst->getOpcode()) { |
| case Instruction::Add: |
| return Builder.CreateAdd(LHS, RHS); |
| case Instruction::Mul: |
| return Builder.CreateMul(LHS, RHS); |
| case Instruction::Sub: |
| return Builder.CreateSub(LHS, RHS); |
| case Instruction::FAdd: |
| return Builder.CreateFAdd(LHS, RHS); |
| case Instruction::FMul: |
| return Builder.CreateFMul(LHS, RHS); |
| case Instruction::FSub: |
| return Builder.CreateFSub(LHS, RHS); |
| default: |
| llvm_unreachable("Unsupported binary operator for matrix"); |
| } |
| }; |
| for (unsigned C = 0; C < Shape.NumColumns; ++C) |
| Result.addColumn( |
| BuildColumnOp(LoweredLhs.getColumn(C), LoweredRhs.getColumn(C))); |
| |
| finalizeLowering(Inst, Result, Builder); |
| return true; |
| } |
| }; |
| } // namespace |
| |
| PreservedAnalyses LowerMatrixIntrinsicsPass::run(Function &F, |
| FunctionAnalysisManager &AM) { |
| auto &TTI = AM.getResult<TargetIRAnalysis>(F); |
| LowerMatrixIntrinsics LMT(F, TTI); |
| if (LMT.Visit()) { |
| PreservedAnalyses PA; |
| PA.preserveSet<CFGAnalyses>(); |
| return PA; |
| } |
| return PreservedAnalyses::all(); |
| } |
| |
| namespace { |
| |
| class LowerMatrixIntrinsicsLegacyPass : public FunctionPass { |
| public: |
| static char ID; |
| |
| LowerMatrixIntrinsicsLegacyPass() : FunctionPass(ID) { |
| initializeLowerMatrixIntrinsicsLegacyPassPass( |
| *PassRegistry::getPassRegistry()); |
| } |
| |
| bool runOnFunction(Function &F) override { |
| auto *TTI = &getAnalysis<TargetTransformInfoWrapperPass>().getTTI(F); |
| LowerMatrixIntrinsics LMT(F, *TTI); |
| bool C = LMT.Visit(); |
| return C; |
| } |
| |
| void getAnalysisUsage(AnalysisUsage &AU) const override { |
| AU.addRequired<TargetTransformInfoWrapperPass>(); |
| AU.setPreservesCFG(); |
| } |
| }; |
| } // namespace |
| |
| static const char pass_name[] = "Lower the matrix intrinsics"; |
| char LowerMatrixIntrinsicsLegacyPass::ID = 0; |
| INITIALIZE_PASS_BEGIN(LowerMatrixIntrinsicsLegacyPass, DEBUG_TYPE, pass_name, |
| false, false) |
| INITIALIZE_PASS_END(LowerMatrixIntrinsicsLegacyPass, DEBUG_TYPE, pass_name, |
| false, false) |
| |
| Pass *llvm::createLowerMatrixIntrinsicsPass() { |
| return new LowerMatrixIntrinsicsLegacyPass(); |
| } |
