forked from OSchip/llvm-project
895 lines
33 KiB
C++
895 lines
33 KiB
C++
//===- 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();
|
|
}
|