RPCS3/llvm-mirror
1
0
Fork 0
mirror of https://github.com/RPCS3/llvm-mirror.git synced 2024-11-23 03:02:36 +01:00
Code Issues Projects Releases Wiki Activity
d02332a857
llvm-mirror/lib/Transforms/Scalar/LowerMatrixIntrinsics.cpp
Florian Hahn 4d6151bec8 [Matrix] Update shape propagation to iterate until done. 
This patch updates the shape propagation to iterate until no new shape
information is discovered.

As initial seed for the forward propagation, we use the matrix intrinsic
instructions. Both propagateShapeForward and propagateShapeBackward
return new work lists, with the instructions to be used for the next
iteration. When propagating forward, we record all instructions we added
new shape information for. When propagating backward, we record all
users of instructions we added new shape information for.

Reviewers: anemet, Gerolf, reames, hfinkel, andrew.w.kaylor

Reviewed By: anemet

Differential Revision: https://reviews.llvm.org/D70901
2020-01-09 10:52:52 +00:00

895 lines
33 KiB
C++
Raw Blame History

//===- 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();
}
Reference in New Issue View Git Blame Copy Permalink