forked from OSchip/llvm-project
1836 lines
77 KiB
C++
1836 lines
77 KiB
C++
//===- MatmulOptimizer.cpp -----------------------------------------------===//
|
||
//
|
||
// 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
|
||
//
|
||
//===----------------------------------------------------------------------===//
|
||
|
||
#include "polly/MatmulOptimizer.h"
|
||
#include "polly/DependenceInfo.h"
|
||
#include "polly/Options.h"
|
||
#include "polly/ScheduleTreeTransform.h"
|
||
#include "polly/ScopInfo.h"
|
||
#include "polly/ScopPass.h"
|
||
#include "polly/Simplify.h"
|
||
#include "polly/Support/GICHelper.h"
|
||
#include "polly/Support/ISLTools.h"
|
||
#include "llvm/ADT/ArrayRef.h"
|
||
#include "llvm/ADT/DenseSet.h"
|
||
#include "llvm/ADT/Sequence.h"
|
||
#include "llvm/ADT/SetOperations.h"
|
||
#include "llvm/ADT/SmallVector.h"
|
||
#include "llvm/ADT/StringRef.h"
|
||
#include "llvm/ADT/iterator_range.h"
|
||
#include "llvm/Analysis/TargetTransformInfo.h"
|
||
#include "llvm/IR/DataLayout.h"
|
||
#include "llvm/IR/Function.h"
|
||
#include "llvm/IR/Module.h"
|
||
#include "llvm/Support/CommandLine.h"
|
||
#include "llvm/Support/Debug.h"
|
||
#include "llvm/Support/TypeSize.h"
|
||
#include "llvm/Support/raw_ostream.h"
|
||
#include "isl/ctx.h"
|
||
#include "isl/schedule_node.h"
|
||
#include "isl/schedule_type.h"
|
||
#include "isl/union_map.h"
|
||
#include "isl/union_set.h"
|
||
#include <algorithm>
|
||
#include <cassert>
|
||
#include <cmath>
|
||
#include <cstdint>
|
||
#include <string>
|
||
#include <vector>
|
||
|
||
#define DEBUG_TYPE "polly-opt-isl"
|
||
|
||
using namespace llvm;
|
||
using namespace polly;
|
||
|
||
namespace llvm {
|
||
class Value;
|
||
}
|
||
|
||
static cl::opt<int> LatencyVectorFma(
|
||
"polly-target-latency-vector-fma",
|
||
cl::desc("The minimal number of cycles between issuing two "
|
||
"dependent consecutive vector fused multiply-add "
|
||
"instructions."),
|
||
cl::Hidden, cl::init(8), cl::cat(PollyCategory));
|
||
|
||
static cl::opt<int> ThroughputVectorFma(
|
||
"polly-target-throughput-vector-fma",
|
||
cl::desc("A throughput of the processor floating-point arithmetic units "
|
||
"expressed in the number of vector fused multiply-add "
|
||
"instructions per clock cycle."),
|
||
cl::Hidden, cl::init(1), cl::cat(PollyCategory));
|
||
|
||
static cl::opt<int> FirstCacheLevelSize(
|
||
"polly-target-1st-cache-level-size",
|
||
cl::desc("The size of the first cache level specified in bytes."),
|
||
cl::Hidden, cl::init(-1), cl::cat(PollyCategory));
|
||
|
||
static cl::opt<int> FirstCacheLevelDefaultSize(
|
||
"polly-target-1st-cache-level-default-size",
|
||
cl::desc("The default size of the first cache level specified in bytes"
|
||
" (if not enough were provided by the TargetTransformInfo)."),
|
||
cl::Hidden, cl::init(32768), cl::cat(PollyCategory));
|
||
|
||
static cl::opt<int> SecondCacheLevelSize(
|
||
"polly-target-2nd-cache-level-size",
|
||
cl::desc("The size of the second level specified in bytes."), cl::Hidden,
|
||
cl::init(-1), cl::cat(PollyCategory));
|
||
|
||
static cl::opt<int> SecondCacheLevelDefaultSize(
|
||
"polly-target-2nd-cache-level-default-size",
|
||
cl::desc("The default size of the second cache level specified in bytes"
|
||
" (if not enough were provided by the TargetTransformInfo)."),
|
||
cl::Hidden, cl::init(262144), cl::cat(PollyCategory));
|
||
|
||
// This option, along with --polly-target-2nd-cache-level-associativity,
|
||
// --polly-target-1st-cache-level-size, and --polly-target-2st-cache-level-size
|
||
// represent the parameters of the target cache, which do not have typical
|
||
// values that can be used by default. However, to apply the pattern matching
|
||
// optimizations, we use the values of the parameters of Intel Core i7-3820
|
||
// SandyBridge in case the parameters are not specified or not provided by the
|
||
// TargetTransformInfo.
|
||
static cl::opt<int> FirstCacheLevelAssociativity(
|
||
"polly-target-1st-cache-level-associativity",
|
||
cl::desc("The associativity of the first cache level."), cl::Hidden,
|
||
cl::init(-1), cl::cat(PollyCategory));
|
||
|
||
static cl::opt<int> FirstCacheLevelDefaultAssociativity(
|
||
"polly-target-1st-cache-level-default-associativity",
|
||
cl::desc("The default associativity of the first cache level"
|
||
" (if not enough were provided by the TargetTransformInfo)."),
|
||
cl::Hidden, cl::init(8), cl::cat(PollyCategory));
|
||
|
||
static cl::opt<int> SecondCacheLevelAssociativity(
|
||
"polly-target-2nd-cache-level-associativity",
|
||
cl::desc("The associativity of the second cache level."), cl::Hidden,
|
||
cl::init(-1), cl::cat(PollyCategory));
|
||
|
||
static cl::opt<int> SecondCacheLevelDefaultAssociativity(
|
||
"polly-target-2nd-cache-level-default-associativity",
|
||
cl::desc("The default associativity of the second cache level"
|
||
" (if not enough were provided by the TargetTransformInfo)."),
|
||
cl::Hidden, cl::init(8), cl::cat(PollyCategory));
|
||
|
||
static cl::opt<int> VectorRegisterBitwidth(
|
||
"polly-target-vector-register-bitwidth",
|
||
cl::desc("The size in bits of a vector register (if not set, this "
|
||
"information is taken from LLVM's target information."),
|
||
cl::Hidden, cl::init(-1), cl::cat(PollyCategory));
|
||
|
||
static cl::opt<int> PollyPatternMatchingNcQuotient(
|
||
"polly-pattern-matching-nc-quotient",
|
||
cl::desc("Quotient that is obtained by dividing Nc, the parameter of the"
|
||
"macro-kernel, by Nr, the parameter of the micro-kernel"),
|
||
cl::Hidden, cl::init(256), cl::cat(PollyCategory));
|
||
|
||
static cl::opt<bool>
|
||
PMBasedTCOpts("polly-tc-opt",
|
||
cl::desc("Perform optimizations of tensor contractions based "
|
||
"on pattern matching"),
|
||
cl::init(false), cl::ZeroOrMore, cl::cat(PollyCategory));
|
||
|
||
static cl::opt<bool>
|
||
PMBasedMMMOpts("polly-matmul-opt",
|
||
cl::desc("Perform optimizations of matrix multiplications "
|
||
"based on pattern matching"),
|
||
cl::init(true), cl::ZeroOrMore, cl::cat(PollyCategory));
|
||
|
||
static cl::opt<int> OptComputeOut(
|
||
"polly-tc-dependences-computeout",
|
||
cl::desc("Bound the dependence analysis by a maximal amount of "
|
||
"computational steps (0 means no bound)"),
|
||
cl::Hidden, cl::init(500000), cl::ZeroOrMore, cl::cat(PollyCategory));
|
||
|
||
namespace {
|
||
/// Parameters of the micro kernel.
|
||
///
|
||
/// Parameters, which determine sizes of rank-1 (i.e., outer product) update
|
||
/// used in the optimized matrix multiplication.
|
||
struct MicroKernelParamsTy {
|
||
int Mr;
|
||
int Nr;
|
||
};
|
||
|
||
/// Parameters of the macro kernel.
|
||
///
|
||
/// Parameters, which determine sizes of blocks of partitioned matrices
|
||
/// used in the optimized matrix multiplication.
|
||
struct MacroKernelParamsTy {
|
||
int Mc;
|
||
int Nc;
|
||
int Kc;
|
||
};
|
||
|
||
/// Parameters of the matrix multiplication operands.
|
||
///
|
||
/// Parameters, which describe access relations that represent operands of the
|
||
/// matrix multiplication.
|
||
struct MatMulInfoTy {
|
||
MemoryAccess *A = nullptr;
|
||
MemoryAccess *B = nullptr;
|
||
MemoryAccess *ReadFromC = nullptr;
|
||
MemoryAccess *WriteToC = nullptr;
|
||
int i = -1;
|
||
int j = -1;
|
||
int k = -1;
|
||
};
|
||
|
||
/// Parameters of the tensor contraction operands.
|
||
///
|
||
/// A general d-dimensional tensor T ∈ R ^ Nu0 x ... x Nud−1 can be defined
|
||
/// as the set of scalar elements indexed by the set of indices u0 ... ud,
|
||
///
|
||
/// T ≡ {Anu0...nud−1 ∈ R | (u0,...,ud−1) ∈ Nu0 x ... x Nud−1}.
|
||
///
|
||
/// Let A, B, and C be dA, dB, and dC-dimensional tensors, respectively.
|
||
/// Let the free and the contracted indices of the tensor A be grouped into
|
||
/// two bundles I = i0...ir−1 and P = p0...pt−1, respectively. Similarly,
|
||
/// the free and the contracted indices of B are grouped into bundles
|
||
/// J = j0..js−1 and P and the free indices of C are grouped into
|
||
/// bundles I and J.
|
||
///
|
||
/// Tensor contraction (TC) of tensors A, B into tensor C can be represented as
|
||
/// C(shuffle(I,J))=∑α·A(shuffle(I,P))·B(shuffle(P,J))+β·C(shuffle(I,J)),
|
||
/// where ∑ is a summation over all contracted indices of P,
|
||
/// α, β ∈ R, Npi is the length of the tensor dimension that corresponds
|
||
/// to the index pi, A(shuffle(I, P)), B(shuffle(P, J)), C(shuffle(I, J)) are
|
||
/// accesses to tensors A, B, C, respectively,
|
||
/// shuffle(I, J), shuffle(I, P), and shuffle(P, J) are permutations of
|
||
/// the enclosed indices.
|
||
///
|
||
/// Multiplication of C(shuffle(I,J)) by β can be moved into a different SCoP
|
||
/// statement by loop distribution, which is done by the isl scheduler.
|
||
// If β is not equal to one, the optimization of TC of Polly requires
|
||
/// such a transformation.
|
||
///
|
||
/// TCInfoTy contains parameters, which describe access relations that represent
|
||
/// operands of the tensor contraction.
|
||
struct TCInfoTy {
|
||
/// @{
|
||
/// Memory accesses that represent reading from tensors, which are operands of
|
||
/// the tensor contraction.
|
||
MemoryAccess *A = nullptr;
|
||
MemoryAccess *B = nullptr;
|
||
/// @}
|
||
|
||
/// @{
|
||
/// Memory accesses that represent reading from and writing into the tensor,
|
||
/// which contains the result of the tensor contraction.
|
||
MemoryAccess *ReadFromC = nullptr;
|
||
MemoryAccess *WriteToC = nullptr;
|
||
/// @}
|
||
|
||
/// @{
|
||
/// Input dimensions of the schedule space, which represent free
|
||
/// indices of tensors.
|
||
SmallDenseSet<int> I;
|
||
SmallDenseSet<int> J;
|
||
/// @}
|
||
|
||
/// Input dimension of the schedule space, which represents contracted
|
||
/// indices of tensors.
|
||
SmallDenseSet<int> P;
|
||
|
||
/// @{
|
||
/// Sizes of tensor dimensions for corresponding input dimensions of
|
||
/// the schedule space. The size of the tensor dimension can be larger than
|
||
/// the size of the corresponding input dimension of the schedule space.
|
||
/// This does not correspond to a tensor contraction. However, such a pattern
|
||
/// will be optimized by the transformation.
|
||
SmallVector<int> DimensionSizes;
|
||
SmallVector<int> ADimensions;
|
||
SmallVector<int> BDimensions;
|
||
SmallVector<int> CDimensions;
|
||
/// @}
|
||
|
||
/// @{
|
||
/// Permutations of indices of I, J, and P, which describe operands of
|
||
/// the tensor contraction and its result.
|
||
SmallVector<int> OrderedI;
|
||
SmallVector<int> OrderedJ;
|
||
SmallVector<int> OrderedP;
|
||
/// @}
|
||
};
|
||
|
||
/// Create an isl::union_set, which describes the option of the form
|
||
/// [isolate[] -> unroll[x]].
|
||
///
|
||
/// @param Ctx An isl::ctx, which is used to create the isl::union_set.
|
||
static isl::union_set getUnrollIsolatedSetOptions(isl::ctx Ctx) {
|
||
isl::space Space = isl::space(Ctx, 0, 0, 1);
|
||
isl::map UnrollIsolatedSetOption = isl::map::universe(Space);
|
||
isl::id DimInId = isl::id::alloc(Ctx, "isolate", nullptr);
|
||
isl::id DimOutId = isl::id::alloc(Ctx, "unroll", nullptr);
|
||
UnrollIsolatedSetOption =
|
||
UnrollIsolatedSetOption.set_tuple_id(isl::dim::in, DimInId);
|
||
UnrollIsolatedSetOption =
|
||
UnrollIsolatedSetOption.set_tuple_id(isl::dim::out, DimOutId);
|
||
return UnrollIsolatedSetOption.wrap();
|
||
}
|
||
|
||
/// Permute the two dimensions of the isl map.
|
||
///
|
||
/// Permute @p DstPos and @p SrcPos dimensions of the isl map @p Map that
|
||
/// have type @p DimType.
|
||
///
|
||
/// @param Map The isl map to be modified.
|
||
/// @param DimType The type of the dimensions.
|
||
/// @param DstPos The first dimension.
|
||
/// @param SrcPos The second dimension.
|
||
/// @return The modified map.
|
||
static isl::map permuteDimensions(isl::map Map, isl::dim DimType,
|
||
unsigned DstPos, unsigned SrcPos) {
|
||
assert(DstPos < unsignedFromIslSize(Map.dim(DimType)) &&
|
||
SrcPos < unsignedFromIslSize(Map.dim(DimType)));
|
||
if (DstPos == SrcPos)
|
||
return Map;
|
||
isl::id DimId;
|
||
if (Map.has_tuple_id(DimType))
|
||
DimId = Map.get_tuple_id(DimType);
|
||
auto FreeDim = DimType == isl::dim::in ? isl::dim::out : isl::dim::in;
|
||
isl::id FreeDimId;
|
||
if (Map.has_tuple_id(FreeDim))
|
||
FreeDimId = Map.get_tuple_id(FreeDim);
|
||
auto MaxDim = std::max(DstPos, SrcPos);
|
||
auto MinDim = std::min(DstPos, SrcPos);
|
||
Map = Map.move_dims(FreeDim, 0, DimType, MaxDim, 1);
|
||
Map = Map.move_dims(FreeDim, 0, DimType, MinDim, 1);
|
||
Map = Map.move_dims(DimType, MinDim, FreeDim, 1, 1);
|
||
Map = Map.move_dims(DimType, MaxDim, FreeDim, 0, 1);
|
||
if (!DimId.is_null())
|
||
Map = Map.set_tuple_id(DimType, DimId);
|
||
if (!FreeDimId.is_null())
|
||
Map = Map.set_tuple_id(FreeDim, FreeDimId);
|
||
return Map;
|
||
}
|
||
|
||
/// Check the form of the access relation.
|
||
///
|
||
/// Check that the access relation @p AccMap has the form M[i][j], where i
|
||
/// is a @p FirstPos and j is a @p SecondPos.
|
||
///
|
||
/// @param AccMap The access relation to be checked.
|
||
/// @param FirstPos The index of the input dimension that is mapped to
|
||
/// the first output dimension.
|
||
/// @param SecondPos The index of the input dimension that is mapped to the
|
||
/// second output dimension.
|
||
/// @return True in case @p AccMap has the expected form and false,
|
||
/// otherwise.
|
||
static bool isMatMulOperandAcc(isl::set Domain, isl::map AccMap, int &FirstPos,
|
||
int &SecondPos) {
|
||
isl::space Space = AccMap.get_space();
|
||
isl::map Universe = isl::map::universe(Space);
|
||
|
||
if (unsignedFromIslSize(Space.dim(isl::dim::out)) != 2)
|
||
return false;
|
||
|
||
// MatMul has the form:
|
||
// for (i = 0; i < N; i++)
|
||
// for (j = 0; j < M; j++)
|
||
// for (k = 0; k < P; k++)
|
||
// C[i, j] += A[i, k] * B[k, j]
|
||
//
|
||
// Permutation of three outer loops: 3! = 6 possibilities.
|
||
int FirstDims[] = {0, 0, 1, 1, 2, 2};
|
||
int SecondDims[] = {1, 2, 2, 0, 0, 1};
|
||
for (int i = 0; i < 6; i += 1) {
|
||
auto PossibleMatMul =
|
||
Universe.equate(isl::dim::in, FirstDims[i], isl::dim::out, 0)
|
||
.equate(isl::dim::in, SecondDims[i], isl::dim::out, 1);
|
||
|
||
AccMap = AccMap.intersect_domain(Domain);
|
||
PossibleMatMul = PossibleMatMul.intersect_domain(Domain);
|
||
|
||
// If AccMap spans entire domain (Non-partial write),
|
||
// compute FirstPos and SecondPos.
|
||
// If AccMap != PossibleMatMul here (the two maps have been gisted at
|
||
// this point), it means that the writes are not complete, or in other
|
||
// words, it is a Partial write and Partial writes must be rejected.
|
||
if (AccMap.is_equal(PossibleMatMul)) {
|
||
if (FirstPos != -1 && FirstPos != FirstDims[i])
|
||
continue;
|
||
FirstPos = FirstDims[i];
|
||
if (SecondPos != -1 && SecondPos != SecondDims[i])
|
||
continue;
|
||
SecondPos = SecondDims[i];
|
||
return true;
|
||
}
|
||
}
|
||
|
||
return false;
|
||
}
|
||
|
||
/// Does the memory access represent a non-scalar operand of the matrix
|
||
/// multiplication.
|
||
///
|
||
/// Check that the memory access @p MemAccess is the read access to a non-scalar
|
||
/// operand of the matrix multiplication or its result.
|
||
///
|
||
/// @param MemAccess The memory access to be checked.
|
||
/// @param MMI Parameters of the matrix multiplication operands.
|
||
/// @return True in case the memory access represents the read access
|
||
/// to a non-scalar operand of the matrix multiplication and
|
||
/// false, otherwise.
|
||
static bool isMatMulNonScalarReadAccess(MemoryAccess *MemAccess,
|
||
MatMulInfoTy &MMI) {
|
||
if (!MemAccess->isLatestArrayKind() || !MemAccess->isRead())
|
||
return false;
|
||
auto AccMap = MemAccess->getLatestAccessRelation();
|
||
isl::set StmtDomain = MemAccess->getStatement()->getDomain();
|
||
if (isMatMulOperandAcc(StmtDomain, AccMap, MMI.i, MMI.j) && !MMI.ReadFromC) {
|
||
MMI.ReadFromC = MemAccess;
|
||
return true;
|
||
}
|
||
if (isMatMulOperandAcc(StmtDomain, AccMap, MMI.i, MMI.k) && !MMI.A) {
|
||
MMI.A = MemAccess;
|
||
return true;
|
||
}
|
||
if (isMatMulOperandAcc(StmtDomain, AccMap, MMI.k, MMI.j) && !MMI.B) {
|
||
MMI.B = MemAccess;
|
||
return true;
|
||
}
|
||
return false;
|
||
}
|
||
|
||
/// Check accesses to operands of the matrix multiplication.
|
||
///
|
||
/// Check that accesses of the SCoP statement, which corresponds to
|
||
/// the partial schedule @p PartialSchedule, are scalar in terms of loops
|
||
/// containing the matrix multiplication, in case they do not represent
|
||
/// accesses to the non-scalar operands of the matrix multiplication or
|
||
/// its result.
|
||
///
|
||
/// @param PartialSchedule The partial schedule of the SCoP statement.
|
||
/// @param MMI Parameters of the matrix multiplication operands.
|
||
/// @return True in case the corresponding SCoP statement
|
||
/// represents matrix multiplication and false,
|
||
/// otherwise.
|
||
static bool containsOnlyMatrMultAcc(isl::map PartialSchedule,
|
||
MatMulInfoTy &MMI) {
|
||
auto InputDimId = PartialSchedule.get_tuple_id(isl::dim::in);
|
||
auto *Stmt = static_cast<ScopStmt *>(InputDimId.get_user());
|
||
unsigned OutDimNum = unsignedFromIslSize(PartialSchedule.range_tuple_dim());
|
||
assert(OutDimNum > 2 && "In case of the matrix multiplication the loop nest "
|
||
"and, consequently, the corresponding scheduling "
|
||
"functions have at least three dimensions.");
|
||
auto MapI =
|
||
permuteDimensions(PartialSchedule, isl::dim::out, MMI.i, OutDimNum - 1);
|
||
auto MapJ =
|
||
permuteDimensions(PartialSchedule, isl::dim::out, MMI.j, OutDimNum - 1);
|
||
auto MapK =
|
||
permuteDimensions(PartialSchedule, isl::dim::out, MMI.k, OutDimNum - 1);
|
||
|
||
auto Accesses = getAccessesInOrder(*Stmt);
|
||
for (auto *MemA = Accesses.begin(); MemA != Accesses.end() - 1; MemA++) {
|
||
auto *MemAccessPtr = *MemA;
|
||
if (MemAccessPtr->isLatestArrayKind() && MemAccessPtr != MMI.WriteToC &&
|
||
!isMatMulNonScalarReadAccess(MemAccessPtr, MMI) &&
|
||
!(MemAccessPtr->isStrideZero(MapI) &&
|
||
MemAccessPtr->isStrideZero(MapJ) && MemAccessPtr->isStrideZero(MapK)))
|
||
return false;
|
||
}
|
||
return true;
|
||
}
|
||
|
||
/// Check for dependencies corresponding to the matrix multiplication.
|
||
///
|
||
/// Check that there is only true dependence of the form
|
||
/// S(..., k, ...) -> S(..., k + 1, …), where S is the SCoP statement
|
||
/// represented by @p Schedule and k is @p Pos. Such a dependence corresponds
|
||
/// to the dependency produced by the matrix multiplication.
|
||
///
|
||
/// @param Schedule The schedule of the SCoP statement.
|
||
/// @param D The SCoP dependencies.
|
||
/// @param Pos The parameter to describe an acceptable true dependence.
|
||
/// In case it has a negative value, try to determine its
|
||
/// acceptable value.
|
||
/// @return True in case dependencies correspond to the matrix multiplication
|
||
/// and false, otherwise.
|
||
static bool containsOnlyMatMulDep(isl::map Schedule, const Dependences *D,
|
||
int &Pos) {
|
||
isl::union_map Dep = D->getDependences(Dependences::TYPE_RAW);
|
||
isl::union_map Red = D->getDependences(Dependences::TYPE_RED);
|
||
if (!Red.is_null())
|
||
Dep = Dep.unite(Red);
|
||
auto DomainSpace = Schedule.get_space().domain();
|
||
auto Space = DomainSpace.map_from_domain_and_range(DomainSpace);
|
||
auto Deltas = Dep.extract_map(Space).deltas();
|
||
int DeltasDimNum = unsignedFromIslSize(Deltas.dim(isl::dim::set));
|
||
for (int i = 0; i < DeltasDimNum; i++) {
|
||
auto Val = Deltas.plain_get_val_if_fixed(isl::dim::set, i);
|
||
Pos = Pos < 0 && Val.is_one() ? i : Pos;
|
||
if (Val.is_nan() || !(Val.is_zero() || (i == Pos && Val.is_one())))
|
||
return false;
|
||
}
|
||
if (DeltasDimNum == 0 || Pos < 0)
|
||
return false;
|
||
return true;
|
||
}
|
||
|
||
/// Check if the SCoP statement could probably be optimized with analytical
|
||
/// modeling.
|
||
///
|
||
/// containsMatrMult tries to determine whether the following conditions
|
||
/// are true:
|
||
/// 1. The last memory access modeling an array, MA1, represents writing to
|
||
/// memory and has the form S(..., i1, ..., i2, ...) -> M(i1, i2) or
|
||
/// S(..., i2, ..., i1, ...) -> M(i1, i2), where S is the SCoP statement
|
||
/// under consideration.
|
||
/// 2. There is only one loop-carried true dependency, and it has the
|
||
/// form S(..., i3, ...) -> S(..., i3 + 1, ...), and there are no
|
||
/// loop-carried or anti dependencies.
|
||
/// 3. SCoP contains three access relations, MA2, MA3, and MA4 that represent
|
||
/// reading from memory and have the form S(..., i3, ...) -> M(i1, i3),
|
||
/// S(..., i3, ...) -> M(i3, i2), S(...) -> M(i1, i2), respectively,
|
||
/// and all memory accesses of the SCoP that are different from MA1, MA2,
|
||
/// MA3, and MA4 have stride 0, if the innermost loop is exchanged with any
|
||
/// of loops i1, i2 and i3.
|
||
///
|
||
/// @param PartialSchedule The PartialSchedule that contains a SCoP statement
|
||
/// to check.
|
||
/// @D The SCoP dependencies.
|
||
/// @MMI Parameters of the matrix multiplication operands.
|
||
static bool containsMatrMult(isl::map PartialSchedule, const Dependences *D,
|
||
MatMulInfoTy &MMI) {
|
||
auto InputDimsId = PartialSchedule.get_tuple_id(isl::dim::in);
|
||
auto *Stmt = static_cast<ScopStmt *>(InputDimsId.get_user());
|
||
if (Stmt->size() <= 1)
|
||
return false;
|
||
|
||
auto Accesses = getAccessesInOrder(*Stmt);
|
||
for (auto *MemA = Accesses.end() - 1; MemA != Accesses.begin(); MemA--) {
|
||
auto *MemAccessPtr = *MemA;
|
||
if (!MemAccessPtr->isLatestArrayKind())
|
||
continue;
|
||
if (!MemAccessPtr->isWrite())
|
||
return false;
|
||
auto AccMap = MemAccessPtr->getLatestAccessRelation();
|
||
if (!isMatMulOperandAcc(Stmt->getDomain(), AccMap, MMI.i, MMI.j))
|
||
return false;
|
||
MMI.WriteToC = MemAccessPtr;
|
||
break;
|
||
}
|
||
|
||
if (!containsOnlyMatMulDep(PartialSchedule, D, MMI.k))
|
||
return false;
|
||
|
||
if (!MMI.WriteToC || !containsOnlyMatrMultAcc(PartialSchedule, MMI))
|
||
return false;
|
||
|
||
if (!MMI.A || !MMI.B || !MMI.ReadFromC)
|
||
return false;
|
||
return true;
|
||
}
|
||
|
||
/// Permute two dimensions of the band node.
|
||
///
|
||
/// Permute FirstDim and SecondDim dimensions of the Node.
|
||
///
|
||
/// @param Node The band node to be modified.
|
||
/// @param FirstDim The first dimension to be permuted.
|
||
/// @param SecondDim The second dimension to be permuted.
|
||
static isl::schedule_node permuteBandNodeDimensions(isl::schedule_node Node,
|
||
unsigned FirstDim,
|
||
unsigned SecondDim) {
|
||
assert(isl_schedule_node_get_type(Node.get()) == isl_schedule_node_band &&
|
||
(unsigned)isl_schedule_node_band_n_member(Node.get()) >
|
||
std::max(FirstDim, SecondDim));
|
||
auto PartialSchedule =
|
||
isl::manage(isl_schedule_node_band_get_partial_schedule(Node.get()));
|
||
auto PartialScheduleFirstDim = PartialSchedule.at(FirstDim);
|
||
auto PartialScheduleSecondDim = PartialSchedule.at(SecondDim);
|
||
PartialSchedule =
|
||
PartialSchedule.set_union_pw_aff(SecondDim, PartialScheduleFirstDim);
|
||
PartialSchedule =
|
||
PartialSchedule.set_union_pw_aff(FirstDim, PartialScheduleSecondDim);
|
||
Node = isl::manage(isl_schedule_node_delete(Node.release()));
|
||
return Node.insert_partial_schedule(PartialSchedule);
|
||
}
|
||
|
||
static isl::schedule_node
|
||
createMicroKernel(isl::schedule_node Node,
|
||
MicroKernelParamsTy MicroKernelParams) {
|
||
Node = applyRegisterTiling(Node, {MicroKernelParams.Mr, MicroKernelParams.Nr},
|
||
1);
|
||
Node = Node.parent().parent();
|
||
return permuteBandNodeDimensions(Node, 0, 1).child(0).child(0);
|
||
}
|
||
|
||
/// Create the BLIS macro-kernel.
|
||
///
|
||
/// We create the BLIS macro-kernel by applying a combination of tiling
|
||
/// of dimensions of the band node and interchanging of two innermost
|
||
/// modified dimensions. The values of of MacroKernelParams's fields are used
|
||
/// as tile sizes.
|
||
///
|
||
/// @param Node The schedule node to be modified.
|
||
/// @param MacroKernelParams Parameters of the macro kernel
|
||
/// to be used as tile sizes.
|
||
static isl::schedule_node
|
||
createMacroKernel(isl::schedule_node Node,
|
||
MacroKernelParamsTy MacroKernelParams) {
|
||
assert(isl_schedule_node_get_type(Node.get()) == isl_schedule_node_band);
|
||
if (MacroKernelParams.Mc == 1 && MacroKernelParams.Nc == 1 &&
|
||
MacroKernelParams.Kc == 1)
|
||
return Node;
|
||
int DimOutNum = isl_schedule_node_band_n_member(Node.get());
|
||
std::vector<int> TileSizes(DimOutNum, 1);
|
||
TileSizes[DimOutNum - 3] = MacroKernelParams.Mc;
|
||
TileSizes[DimOutNum - 2] = MacroKernelParams.Nc;
|
||
TileSizes[DimOutNum - 1] = MacroKernelParams.Kc;
|
||
Node = tileNode(Node, "1st level tiling", TileSizes, 1);
|
||
Node = Node.parent().parent();
|
||
Node = permuteBandNodeDimensions(Node, DimOutNum - 2, DimOutNum - 1);
|
||
Node = permuteBandNodeDimensions(Node, DimOutNum - 3, DimOutNum - 1);
|
||
|
||
return Node.child(0).child(0);
|
||
}
|
||
|
||
/// Get the size of the widest type of the matrix multiplication operands
|
||
/// in bytes, including alignment padding.
|
||
///
|
||
/// @param MMI Parameters of the matrix multiplication operands.
|
||
/// @return The size of the widest type of the matrix multiplication operands
|
||
/// in bytes, including alignment padding.
|
||
static uint64_t getMatMulAlignTypeSize(MatMulInfoTy MMI) {
|
||
auto *S = MMI.A->getStatement()->getParent();
|
||
auto &DL = S->getFunction().getParent()->getDataLayout();
|
||
auto ElementSizeA = DL.getTypeAllocSize(MMI.A->getElementType());
|
||
auto ElementSizeB = DL.getTypeAllocSize(MMI.B->getElementType());
|
||
auto ElementSizeC = DL.getTypeAllocSize(MMI.WriteToC->getElementType());
|
||
return std::max({ElementSizeA, ElementSizeB, ElementSizeC});
|
||
}
|
||
|
||
/// Get the size of the widest type of the matrix multiplication operands
|
||
/// in bits.
|
||
///
|
||
/// @param MMI Parameters of the matrix multiplication operands.
|
||
/// @return The size of the widest type of the matrix multiplication operands
|
||
/// in bits.
|
||
static uint64_t getMatMulTypeSize(MatMulInfoTy MMI) {
|
||
auto *S = MMI.A->getStatement()->getParent();
|
||
auto &DL = S->getFunction().getParent()->getDataLayout();
|
||
auto ElementSizeA = DL.getTypeSizeInBits(MMI.A->getElementType());
|
||
auto ElementSizeB = DL.getTypeSizeInBits(MMI.B->getElementType());
|
||
auto ElementSizeC = DL.getTypeSizeInBits(MMI.WriteToC->getElementType());
|
||
return std::max({ElementSizeA, ElementSizeB, ElementSizeC});
|
||
}
|
||
|
||
/// Get parameters of the BLIS micro kernel.
|
||
///
|
||
/// We choose the Mr and Nr parameters of the micro kernel to be large enough
|
||
/// such that no stalls caused by the combination of latencies and dependencies
|
||
/// are introduced during the updates of the resulting matrix of the matrix
|
||
/// multiplication. However, they should also be as small as possible to
|
||
/// release more registers for entries of multiplied matrices.
|
||
///
|
||
/// @param TTI Target Transform Info.
|
||
/// @param MMI Parameters of the matrix multiplication operands.
|
||
/// @return The structure of type MicroKernelParamsTy.
|
||
/// @see MicroKernelParamsTy
|
||
static MicroKernelParamsTy getMicroKernelParams(const TargetTransformInfo *TTI,
|
||
MatMulInfoTy MMI) {
|
||
assert(TTI && "The target transform info should be provided.");
|
||
|
||
// Nvec - Number of double-precision floating-point numbers that can be hold
|
||
// by a vector register. Use 2 by default.
|
||
long RegisterBitwidth = VectorRegisterBitwidth;
|
||
|
||
if (RegisterBitwidth == -1)
|
||
RegisterBitwidth =
|
||
TTI->getRegisterBitWidth(TargetTransformInfo::RGK_FixedWidthVector);
|
||
auto ElementSize = getMatMulTypeSize(MMI);
|
||
assert(ElementSize > 0 && "The element size of the matrix multiplication "
|
||
"operands should be greater than zero.");
|
||
auto Nvec = RegisterBitwidth / ElementSize;
|
||
if (Nvec == 0)
|
||
Nvec = 2;
|
||
int Nr = ceil(sqrt((double)(Nvec * LatencyVectorFma * ThroughputVectorFma)) /
|
||
Nvec) *
|
||
Nvec;
|
||
int Mr = ceil((double)(Nvec * LatencyVectorFma * ThroughputVectorFma / Nr));
|
||
return {Mr, Nr};
|
||
}
|
||
|
||
/// Determine parameters of the target cache.
|
||
///
|
||
/// @param TTI Target Transform Info.
|
||
static void getTargetCacheParameters(const llvm::TargetTransformInfo *TTI) {
|
||
auto L1DCache = llvm::TargetTransformInfo::CacheLevel::L1D;
|
||
auto L2DCache = llvm::TargetTransformInfo::CacheLevel::L2D;
|
||
if (FirstCacheLevelSize == -1) {
|
||
if (TTI->getCacheSize(L1DCache))
|
||
FirstCacheLevelSize = TTI->getCacheSize(L1DCache).value();
|
||
else
|
||
FirstCacheLevelSize = static_cast<int>(FirstCacheLevelDefaultSize);
|
||
}
|
||
if (SecondCacheLevelSize == -1) {
|
||
if (TTI->getCacheSize(L2DCache))
|
||
SecondCacheLevelSize = TTI->getCacheSize(L2DCache).value();
|
||
else
|
||
SecondCacheLevelSize = static_cast<int>(SecondCacheLevelDefaultSize);
|
||
}
|
||
if (FirstCacheLevelAssociativity == -1) {
|
||
if (TTI->getCacheAssociativity(L1DCache))
|
||
FirstCacheLevelAssociativity =
|
||
TTI->getCacheAssociativity(L1DCache).value();
|
||
else
|
||
FirstCacheLevelAssociativity =
|
||
static_cast<int>(FirstCacheLevelDefaultAssociativity);
|
||
}
|
||
if (SecondCacheLevelAssociativity == -1) {
|
||
if (TTI->getCacheAssociativity(L2DCache))
|
||
SecondCacheLevelAssociativity =
|
||
TTI->getCacheAssociativity(L2DCache).value();
|
||
else
|
||
SecondCacheLevelAssociativity =
|
||
static_cast<int>(SecondCacheLevelDefaultAssociativity);
|
||
}
|
||
}
|
||
|
||
/// Get parameters of the BLIS macro kernel.
|
||
///
|
||
/// During the computation of matrix multiplication, blocks of partitioned
|
||
/// matrices are mapped to different layers of the memory hierarchy.
|
||
/// To optimize data reuse, blocks should be ideally kept in cache between
|
||
/// iterations. Since parameters of the macro kernel determine sizes of these
|
||
/// blocks, there are upper and lower bounds on these parameters.
|
||
///
|
||
/// @param TTI Target Transform Info.
|
||
/// @param MicroKernelParams Parameters of the micro-kernel
|
||
/// to be taken into account.
|
||
/// @param MMI Parameters of the matrix multiplication operands.
|
||
/// @return The structure of type MacroKernelParamsTy.
|
||
/// @see MacroKernelParamsTy
|
||
/// @see MicroKernelParamsTy
|
||
static MacroKernelParamsTy
|
||
getMacroKernelParams(const llvm::TargetTransformInfo *TTI,
|
||
const MicroKernelParamsTy &MicroKernelParams,
|
||
MatMulInfoTy MMI) {
|
||
getTargetCacheParameters(TTI);
|
||
// According to www.cs.utexas.edu/users/flame/pubs/TOMS-BLIS-Analytical.pdf,
|
||
// it requires information about the first two levels of a cache to determine
|
||
// all the parameters of a macro-kernel. It also checks that an associativity
|
||
// degree of a cache level is greater than two. Otherwise, another algorithm
|
||
// for determination of the parameters should be used.
|
||
if (!(MicroKernelParams.Mr > 0 && MicroKernelParams.Nr > 0 &&
|
||
FirstCacheLevelSize > 0 && SecondCacheLevelSize > 0 &&
|
||
FirstCacheLevelAssociativity > 2 && SecondCacheLevelAssociativity > 2))
|
||
return {1, 1, 1};
|
||
// The quotient should be greater than zero.
|
||
if (PollyPatternMatchingNcQuotient <= 0)
|
||
return {1, 1, 1};
|
||
int Car = floor(
|
||
(FirstCacheLevelAssociativity - 1) /
|
||
(1 + static_cast<double>(MicroKernelParams.Nr) / MicroKernelParams.Mr));
|
||
|
||
// Car can be computed to be zero since it is floor to int.
|
||
// On Mac OS, division by 0 does not raise a signal. This causes negative
|
||
// tile sizes to be computed. Prevent division by Cac==0 by early returning
|
||
// if this happens.
|
||
if (Car == 0)
|
||
return {1, 1, 1};
|
||
|
||
auto ElementSize = getMatMulAlignTypeSize(MMI);
|
||
assert(ElementSize > 0 && "The element size of the matrix multiplication "
|
||
"operands should be greater than zero.");
|
||
int Kc = (Car * FirstCacheLevelSize) /
|
||
(MicroKernelParams.Mr * FirstCacheLevelAssociativity * ElementSize);
|
||
double Cac =
|
||
static_cast<double>(Kc * ElementSize * SecondCacheLevelAssociativity) /
|
||
SecondCacheLevelSize;
|
||
int Mc = floor((SecondCacheLevelAssociativity - 2) / Cac);
|
||
int Nc = PollyPatternMatchingNcQuotient * MicroKernelParams.Nr;
|
||
|
||
assert(Mc > 0 && Nc > 0 && Kc > 0 &&
|
||
"Matrix block sizes should be greater than zero");
|
||
return {Mc, Nc, Kc};
|
||
}
|
||
|
||
/// Create an access relation that is specific to
|
||
/// the matrix multiplication pattern.
|
||
///
|
||
/// Create an access relation of the following form:
|
||
/// [O0, O1, O2, O3, O4, O5, O6, O7, O8] -> [OI, O5, OJ]
|
||
/// where I is @p FirstDim, J is @p SecondDim.
|
||
///
|
||
/// It can be used, for example, to create relations that helps to consequently
|
||
/// access elements of operands of a matrix multiplication after creation of
|
||
/// the BLIS micro and macro kernels.
|
||
///
|
||
/// @see ScheduleTreeOptimizer::createMicroKernel
|
||
/// @see ScheduleTreeOptimizer::createMacroKernel
|
||
///
|
||
/// Subsequently, the described access relation is applied to the range of
|
||
/// @p MapOldIndVar, that is used to map original induction variables to
|
||
/// the ones, which are produced by schedule transformations. It helps to
|
||
/// define relations using a new space and, at the same time, keep them
|
||
/// in the original one.
|
||
///
|
||
/// @param MapOldIndVar The relation, which maps original induction variables
|
||
/// to the ones, which are produced by schedule
|
||
/// transformations.
|
||
/// @param FirstDim, SecondDim The input dimensions that are used to define
|
||
/// the specified access relation.
|
||
/// @return The specified access relation.
|
||
static isl::map getMatMulAccRel(isl::map MapOldIndVar, unsigned FirstDim,
|
||
unsigned SecondDim) {
|
||
auto AccessRelSpace = isl::space(MapOldIndVar.ctx(), 0, 9, 3);
|
||
auto AccessRel = isl::map::universe(AccessRelSpace);
|
||
AccessRel = AccessRel.equate(isl::dim::in, FirstDim, isl::dim::out, 0);
|
||
AccessRel = AccessRel.equate(isl::dim::in, 5, isl::dim::out, 1);
|
||
AccessRel = AccessRel.equate(isl::dim::in, SecondDim, isl::dim::out, 2);
|
||
return MapOldIndVar.apply_range(AccessRel);
|
||
}
|
||
|
||
static isl::schedule_node createExtensionNode(isl::schedule_node Node,
|
||
isl::map ExtensionMap) {
|
||
auto Extension = isl::union_map(ExtensionMap);
|
||
auto NewNode = isl::schedule_node::from_extension(Extension);
|
||
return Node.graft_before(NewNode);
|
||
}
|
||
|
||
static isl::schedule_node optimizePackedB(isl::schedule_node Node,
|
||
ScopStmt *Stmt, isl::map MapOldIndVar,
|
||
MicroKernelParamsTy MicroParams,
|
||
MacroKernelParamsTy MacroParams,
|
||
MatMulInfoTy &MMI) {
|
||
Scop *S = Stmt->getParent();
|
||
isl::set Domain = Stmt->getDomain();
|
||
|
||
// Create packed array.
|
||
unsigned FirstDimSize = MacroParams.Nc / MicroParams.Nr;
|
||
unsigned SecondDimSize = MacroParams.Kc;
|
||
unsigned ThirdDimSize = MicroParams.Nr;
|
||
ScopArrayInfo *PackedB =
|
||
S->createScopArrayInfo(MMI.B->getElementType(), "Packed_B",
|
||
{FirstDimSize, SecondDimSize, ThirdDimSize});
|
||
|
||
// Compute the access relation for copying from B to PackedB.
|
||
isl::map AccRelB = MMI.B->getLatestAccessRelation();
|
||
isl::map AccRelPackedB = getMatMulAccRel(MapOldIndVar, 3, 7);
|
||
AccRelPackedB =
|
||
AccRelPackedB.set_tuple_id(isl::dim::out, PackedB->getBasePtrId());
|
||
|
||
// Create the copy statement and redirect access.
|
||
ScopStmt *CopyStmt = S->addScopStmt(AccRelB, AccRelPackedB, Domain);
|
||
MMI.B->setNewAccessRelation(AccRelPackedB);
|
||
|
||
unsigned Dim = unsignedFromIslSize(MapOldIndVar.range_tuple_dim());
|
||
assert(Dim >= 2);
|
||
// Insert into the schedule tree.
|
||
isl::map ExtMap = MapOldIndVar.project_out(isl::dim::out, 2, Dim - 2);
|
||
ExtMap = ExtMap.reverse();
|
||
ExtMap = ExtMap.fix_si(isl::dim::out, MMI.i, 0);
|
||
ExtMap = ExtMap.intersect_range(Domain);
|
||
ExtMap = ExtMap.set_tuple_id(isl::dim::out, CopyStmt->getDomainId());
|
||
return createExtensionNode(Node, ExtMap);
|
||
}
|
||
|
||
static isl::schedule_node optimizePackedA(isl::schedule_node Node, ScopStmt *,
|
||
isl::map MapOldIndVar,
|
||
MicroKernelParamsTy MicroParams,
|
||
MacroKernelParamsTy MacroParams,
|
||
MatMulInfoTy &MMI) {
|
||
isl::id InputDimsId = MapOldIndVar.get_tuple_id(isl::dim::in);
|
||
ScopStmt *Stmt = static_cast<ScopStmt *>(InputDimsId.get_user());
|
||
isl::set Domain = Stmt->getDomain();
|
||
isl::id DomainId = Domain.get_tuple_id();
|
||
|
||
// Create the packed array.
|
||
unsigned FirstDimSize = MacroParams.Mc / MicroParams.Mr;
|
||
unsigned SecondDimSize = MacroParams.Kc;
|
||
unsigned ThirdDimSize = MicroParams.Mr;
|
||
ScopArrayInfo *PackedA = Stmt->getParent()->createScopArrayInfo(
|
||
MMI.A->getElementType(), "Packed_A",
|
||
{FirstDimSize, SecondDimSize, ThirdDimSize});
|
||
|
||
// Compute the access relation for copying from A to PackedA.
|
||
isl::map AccRelA = MMI.A->getLatestAccessRelation();
|
||
isl::map AccRelPackedA = getMatMulAccRel(MapOldIndVar, 4, 6);
|
||
AccRelPackedA =
|
||
AccRelPackedA.set_tuple_id(isl::dim::out, PackedA->getBasePtrId());
|
||
// { MemrefA[] -> PackedA[] }
|
||
isl::map PackedATranslator = AccRelPackedA.apply_domain(AccRelA);
|
||
|
||
// Compute the domain for the copy statement.
|
||
// Construct the copy statement domain out of the 3 outermost scatter
|
||
// dimensions (to match the 3 band nodes surrounding the extension node) and
|
||
// the array elements to copy (one statement instance per array element).
|
||
// { Scatter[] }
|
||
isl::set ScatterDomain = MapOldIndVar.intersect_domain(Domain).range();
|
||
// { Scatter[] -> OutermostScatter[] }
|
||
isl::map OuterDomainMap =
|
||
makeIdentityMap(ScatterDomain, true).project_out(isl::dim::out, 3, 6);
|
||
// { Scatter[] -> MemrefA[] }
|
||
isl::map CopyFrom = MapOldIndVar.reverse().apply_range(AccRelA);
|
||
// { Scatter[] -> CopyStmt[] }
|
||
isl::map DomainTranslator = OuterDomainMap.range_product(CopyFrom);
|
||
// { CopyStmt[] }
|
||
isl::set CopyDomain = DomainTranslator.range();
|
||
|
||
// Translate the access relations to the new domain.
|
||
// { CopyStmt[] -> MemrefA[] }
|
||
CopyFrom = CopyFrom.apply_domain(DomainTranslator);
|
||
// { CopyStmt[] -> PackedA[] }
|
||
isl::map CopyTo = CopyFrom.apply_range(PackedATranslator);
|
||
|
||
// Create the copy statement and redirect access.
|
||
ScopStmt *CopyStmt =
|
||
Stmt->getParent()->addScopStmt(CopyFrom, CopyTo, CopyDomain);
|
||
MMI.A->setNewAccessRelation(AccRelPackedA);
|
||
|
||
// Insert into the schedule tree.
|
||
// { Scatter[] -> CopyStmt[] }
|
||
isl::map ExtScatterCopy = makeIdentityMap(CopyStmt->getDomain(), true);
|
||
ExtScatterCopy = ExtScatterCopy.project_out(isl::dim::in, 3, 2);
|
||
return createExtensionNode(Node, ExtScatterCopy);
|
||
}
|
||
|
||
/// Apply the packing transformation.
|
||
///
|
||
/// The packing transformation can be described as a data-layout
|
||
/// transformation that requires to introduce a new array, copy data
|
||
/// to the array, and change memory access locations to reference the array.
|
||
/// It can be used to ensure that elements of the new array are read in-stride
|
||
/// access, aligned to cache lines boundaries, and preloaded into certain cache
|
||
/// levels.
|
||
///
|
||
/// As an example let us consider the packing of the array A that would help
|
||
/// to read its elements with in-stride access. An access to the array A
|
||
/// is represented by an access relation that has the form
|
||
/// S[i, j, k] -> A[i, k]. The scheduling function of the SCoP statement S has
|
||
/// the form S[i,j, k] -> [floor((j mod Nc) / Nr), floor((i mod Mc) / Mr),
|
||
/// k mod Kc, j mod Nr, i mod Mr].
|
||
///
|
||
/// To ensure that elements of the array A are read in-stride access, we add
|
||
/// a new array Packed_A[Mc/Mr][Kc][Mr] to the SCoP, using
|
||
/// Scop::createScopArrayInfo, change the access relation
|
||
/// S[i, j, k] -> A[i, k] to
|
||
/// S[i, j, k] -> Packed_A[floor((i mod Mc) / Mr), k mod Kc, i mod Mr], using
|
||
/// MemoryAccess::setNewAccessRelation, and copy the data to the array, using
|
||
/// the copy statement created by Scop::addScopStmt.
|
||
///
|
||
/// @param Node The schedule node to be optimized.
|
||
/// @param MapOldIndVar The relation, which maps original induction variables
|
||
/// to the ones, which are produced by schedule
|
||
/// transformations.
|
||
/// @param MicroParams, MacroParams Parameters of the BLIS kernel
|
||
/// to be taken into account.
|
||
/// @param MMI Parameters of the matrix multiplication operands.
|
||
/// @return The optimized schedule node.
|
||
static isl::schedule_node
|
||
optimizeDataLayoutMatrMulPattern(isl::schedule_node Node, isl::map MapOldIndVar,
|
||
MicroKernelParamsTy MicroParams,
|
||
MacroKernelParamsTy MacroParams,
|
||
MatMulInfoTy &MMI) {
|
||
isl::id InputDimsId = MapOldIndVar.get_tuple_id(isl::dim::in);
|
||
ScopStmt *Stmt = static_cast<ScopStmt *>(InputDimsId.get_user());
|
||
|
||
Node = Node.parent().parent().parent().parent().parent().parent();
|
||
Node = isl::manage(isl_schedule_node_band_split(Node.release(), 2));
|
||
|
||
Node = Node.child(0);
|
||
Node =
|
||
optimizePackedB(Node, Stmt, MapOldIndVar, MicroParams, MacroParams, MMI);
|
||
|
||
Node = Node.child(0);
|
||
Node =
|
||
optimizePackedA(Node, Stmt, MapOldIndVar, MicroParams, MacroParams, MMI);
|
||
|
||
return Node.child(0).child(0).child(0).child(0).child(0);
|
||
}
|
||
|
||
/// Get a relation mapping induction variables produced by schedule
|
||
/// transformations to the original ones.
|
||
///
|
||
/// @param Node The schedule node produced as the result of creation
|
||
/// of the BLIS kernels.
|
||
/// @param MicroKernelParams, MacroKernelParams Parameters of the BLIS kernel
|
||
/// to be taken into account.
|
||
/// @return The relation mapping original induction variables to the ones
|
||
/// produced by schedule transformation.
|
||
/// @see ScheduleTreeOptimizer::createMicroKernel
|
||
/// @see ScheduleTreeOptimizer::createMacroKernel
|
||
/// @see getMacroKernelParams
|
||
static isl::map
|
||
getInductionVariablesSubstitution(isl::schedule_node Node,
|
||
MicroKernelParamsTy MicroKernelParams,
|
||
MacroKernelParamsTy MacroKernelParams) {
|
||
auto Child = Node.child(0);
|
||
auto UnMapOldIndVar = Child.get_prefix_schedule_union_map();
|
||
auto MapOldIndVar = isl::map::from_union_map(UnMapOldIndVar);
|
||
unsigned Dim = unsignedFromIslSize(MapOldIndVar.range_tuple_dim());
|
||
if (Dim > 9u)
|
||
return MapOldIndVar.project_out(isl::dim::out, 0, Dim - 9);
|
||
return MapOldIndVar;
|
||
}
|
||
|
||
/// Isolate a set of partial tile prefixes and unroll the isolated part.
|
||
///
|
||
/// The set should ensure that it contains only partial tile prefixes that have
|
||
/// exactly Mr x Nr iterations of the two innermost loops produced by
|
||
/// the optimization of the matrix multiplication. Mr and Nr are parameters of
|
||
/// the micro-kernel.
|
||
///
|
||
/// In case of parametric bounds, this helps to auto-vectorize the unrolled
|
||
/// innermost loops, using the SLP vectorizer.
|
||
///
|
||
/// @param Node The schedule node to be modified.
|
||
/// @param MicroKernelParams Parameters of the micro-kernel
|
||
/// to be taken into account.
|
||
/// @return The modified isl_schedule_node.
|
||
static isl::schedule_node
|
||
isolateAndUnrollMatMulInnerLoops(isl::schedule_node Node,
|
||
MicroKernelParamsTy MicroKernelParams) {
|
||
isl::schedule_node Child = Node.child(0);
|
||
isl::union_map UnMapOldIndVar = Child.get_prefix_schedule_relation();
|
||
isl::set Prefix = isl::map::from_union_map(UnMapOldIndVar).range();
|
||
unsigned Dims = unsignedFromIslSize(Prefix.tuple_dim());
|
||
assert(Dims >= 1);
|
||
Prefix = Prefix.project_out(isl::dim::set, Dims - 1, 1);
|
||
Prefix = getPartialTilePrefixes(Prefix, MicroKernelParams.Nr);
|
||
Prefix = getPartialTilePrefixes(Prefix, MicroKernelParams.Mr);
|
||
|
||
isl::union_set IsolateOption =
|
||
getIsolateOptions(Prefix.add_dims(isl::dim::set, 3), 3);
|
||
isl::ctx Ctx = Node.ctx();
|
||
auto Options = IsolateOption.unite(getDimOptions(Ctx, "unroll"));
|
||
Options = Options.unite(getUnrollIsolatedSetOptions(Ctx));
|
||
Node = Node.as<isl::schedule_node_band>().set_ast_build_options(Options);
|
||
Node = Node.parent().parent().parent();
|
||
IsolateOption = getIsolateOptions(Prefix, 3);
|
||
Options = IsolateOption.unite(getDimOptions(Ctx, "separate"));
|
||
Node = Node.as<isl::schedule_node_band>().set_ast_build_options(Options);
|
||
Node = Node.child(0).child(0).child(0);
|
||
return Node;
|
||
}
|
||
|
||
/// Insert "Loop Vectorizer Disabled" mark node.
|
||
///
|
||
/// @param Node The child of the mark node to be inserted.
|
||
/// @return The modified isl_schedule_node.
|
||
static isl::schedule_node markLoopVectorizerDisabled(isl::schedule_node Node) {
|
||
auto Id = isl::id::alloc(Node.ctx(), "Loop Vectorizer Disabled", nullptr);
|
||
return Node.insert_mark(Id).child(0);
|
||
}
|
||
|
||
/// Restore the initial ordering of dimensions of the band node
|
||
///
|
||
/// In case the band node represents all the dimensions of the iteration
|
||
/// domain, recreate the band node to restore the initial ordering of the
|
||
/// dimensions.
|
||
///
|
||
/// @param Node The band node to be modified.
|
||
/// @return The modified schedule node.
|
||
static isl::schedule_node
|
||
getBandNodeWithOriginDimOrder(isl::schedule_node Node) {
|
||
assert(isl_schedule_node_get_type(Node.get()) == isl_schedule_node_band);
|
||
if (isl_schedule_node_get_type(Node.child(0).get()) != isl_schedule_node_leaf)
|
||
return Node;
|
||
auto Domain = Node.get_universe_domain();
|
||
assert(isl_union_set_n_set(Domain.get()) == 1);
|
||
if (Node.get_schedule_depth().release() != 0 ||
|
||
(unsignedFromIslSize(isl::set(Domain).tuple_dim()) !=
|
||
unsignedFromIslSize(Node.as<isl::schedule_node_band>().n_member())))
|
||
return Node;
|
||
Node = isl::manage(isl_schedule_node_delete(Node.copy()));
|
||
auto PartialSchedulePwAff = Domain.identity_union_pw_multi_aff();
|
||
auto PartialScheduleMultiPwAff =
|
||
isl::multi_union_pw_aff(PartialSchedulePwAff);
|
||
PartialScheduleMultiPwAff =
|
||
PartialScheduleMultiPwAff.reset_tuple_id(isl::dim::set);
|
||
return Node.insert_partial_schedule(PartialScheduleMultiPwAff);
|
||
}
|
||
|
||
static isl::schedule_node optimizeMatMulPattern(isl::schedule_node Node,
|
||
const TargetTransformInfo *TTI,
|
||
MatMulInfoTy &MMI) {
|
||
assert(TTI && "The target transform info should be provided.");
|
||
int DimOutNum = isl_schedule_node_band_n_member(Node.get());
|
||
assert(DimOutNum > 2 && "In case of the matrix multiplication the loop nest "
|
||
"and, consequently, the corresponding scheduling "
|
||
"functions have at least three dimensions.");
|
||
Node = getBandNodeWithOriginDimOrder(Node);
|
||
Node = permuteBandNodeDimensions(Node, MMI.i, DimOutNum - 3);
|
||
int NewJ = MMI.j == DimOutNum - 3 ? MMI.i : MMI.j;
|
||
int NewK = MMI.k == DimOutNum - 3 ? MMI.i : MMI.k;
|
||
Node = permuteBandNodeDimensions(Node, NewJ, DimOutNum - 2);
|
||
NewK = NewK == DimOutNum - 2 ? NewJ : NewK;
|
||
Node = permuteBandNodeDimensions(Node, NewK, DimOutNum - 1);
|
||
auto MicroKernelParams = getMicroKernelParams(TTI, MMI);
|
||
auto MacroKernelParams = getMacroKernelParams(TTI, MicroKernelParams, MMI);
|
||
Node = createMacroKernel(Node, MacroKernelParams);
|
||
Node = createMicroKernel(Node, MicroKernelParams);
|
||
if (MacroKernelParams.Mc == 1 || MacroKernelParams.Nc == 1 ||
|
||
MacroKernelParams.Kc == 1)
|
||
return Node;
|
||
auto MapOldIndVar = getInductionVariablesSubstitution(Node, MicroKernelParams,
|
||
MacroKernelParams);
|
||
if (MapOldIndVar.is_null())
|
||
return Node;
|
||
Node = markLoopVectorizerDisabled(Node.parent()).child(0);
|
||
Node = isolateAndUnrollMatMulInnerLoops(Node, MicroKernelParams);
|
||
return optimizeDataLayoutMatrMulPattern(Node, MapOldIndVar, MicroKernelParams,
|
||
MacroKernelParams, MMI);
|
||
}
|
||
|
||
/// Check if this node contains a partial schedule that could
|
||
/// probably be optimized with analytical modeling.
|
||
///
|
||
/// isMatrMultPattern tries to determine whether the following conditions
|
||
/// are true:
|
||
/// 1. the partial schedule contains only one statement.
|
||
/// 2. there are exactly three input dimensions.
|
||
/// 3. all memory accesses of the statement will have stride 0 or 1, if we
|
||
/// interchange loops (switch the variable used in the inner loop to
|
||
/// the outer loop).
|
||
/// 4. all memory accesses of the statement except from the last one, are
|
||
/// read memory access and the last one is write memory access.
|
||
/// 5. all subscripts of the last memory access of the statement don't
|
||
/// contain the variable used in the inner loop.
|
||
/// If this is the case, we could try to use an approach that is similar to
|
||
/// the one used to get close-to-peak performance of matrix multiplications.
|
||
///
|
||
/// @param Node The node to check.
|
||
/// @param D The SCoP dependencies.
|
||
/// @param MMI Parameters of the matrix multiplication operands.
|
||
static bool isMatrMultPattern(isl::schedule_node Node, const Dependences *D,
|
||
MatMulInfoTy &MMI) {
|
||
auto PartialSchedule = isl::manage(
|
||
isl_schedule_node_band_get_partial_schedule_union_map(Node.get()));
|
||
if (isl_schedule_node_band_n_member(Node.get()) < 3 ||
|
||
Node.get_schedule_depth().release() != 0 ||
|
||
isl_union_map_n_map(PartialSchedule.get()) != 1)
|
||
return false;
|
||
auto NewPartialSchedule = isl::map::from_union_map(PartialSchedule);
|
||
if (containsMatrMult(NewPartialSchedule, D, MMI))
|
||
return true;
|
||
return false;
|
||
}
|
||
|
||
/// Get the dimension size.
|
||
///
|
||
/// Return the size of the dimension @p Pos, which is obtained from @p SAI.
|
||
/// Return -1 in the case of the first dimension of a multi-dimensional array,
|
||
/// since the ScopArrayInfo class does not carry size information.
|
||
///
|
||
/// @param SAI The information about the array.
|
||
/// @param Pos The position of the dimension.
|
||
/// @return The size of the dimension.
|
||
static int getDimSize(const ScopArrayInfo *SAI, unsigned Pos) {
|
||
if (Pos == 0)
|
||
return -1;
|
||
const llvm::SCEV *SCEVDimSize = SAI->getDimensionSize(Pos);
|
||
assert(SCEVDimSize);
|
||
auto *ConstantDimSize = dyn_cast<const SCEVConstant>(SCEVDimSize);
|
||
assert(ConstantDimSize);
|
||
auto *IntDimSize = dyn_cast<ConstantInt>(ConstantDimSize->getValue());
|
||
assert(IntDimSize);
|
||
return IntDimSize->getSExtValue();
|
||
}
|
||
|
||
/// Check whether the access relation has the specified form.
|
||
///
|
||
/// Check that the access relation @p AccMap has the form T[I0, …, In], where
|
||
/// indexes I0, …, In are specified by @p Dimensions.
|
||
///
|
||
/// @param Domain The domain of the access relation.
|
||
/// @param AccMap The access relation to be checked.
|
||
/// @param Dimensions The permutation of the subset of the input dimensions.
|
||
/// @return True if @p AccMap has the expected form and false,
|
||
/// otherwise.
|
||
static bool isCorrectAccessMap(isl::set Domain, isl::map AccMap,
|
||
ArrayRef<int> Dimensions) {
|
||
isl::space Space = AccMap.get_space();
|
||
if (unsignedFromIslSize(Space.dim(isl::dim::out)) != Dimensions.size())
|
||
return false;
|
||
|
||
// Create an access relation of the following form:
|
||
// [I0, …, Im] -> [Il, …, In], where indexes
|
||
// Il, …, In are specified by @p Dimensions.
|
||
isl::map PossibleTensor = isl::map::universe(Space);
|
||
unsigned DimInSize = unsignedFromIslSize(Space.dim(isl::dim::in));
|
||
for (unsigned i = 0; i < Dimensions.size(); i++) {
|
||
const int InPos = Dimensions[i];
|
||
if ((InPos >= static_cast<int>(DimInSize)) || (InPos < 0))
|
||
return false;
|
||
PossibleTensor =
|
||
PossibleTensor.equate(isl::dim::in, InPos, isl::dim::out, i);
|
||
}
|
||
|
||
AccMap = AccMap.intersect_domain(Domain);
|
||
PossibleTensor = PossibleTensor.intersect_domain(Domain);
|
||
|
||
// If AccMap != PossibleTensor here (the two maps have been gisted at
|
||
// this point), it means that the writes are not complete, or in other
|
||
// words, it is a Partial write and Partial writes must be rejected.
|
||
return AccMap.is_equal(PossibleTensor);
|
||
}
|
||
|
||
/// Check whether the access represents the tensor contraction operand.
|
||
///
|
||
/// Check that the access relation @p AccMap has the form T[i1, …, in].
|
||
/// Obtained indexes i1, …, in, their sizes and their permutation are stored
|
||
/// into @p IndexSet, @p DimensionSizes, and @p Dimensions, respectively.
|
||
///
|
||
/// @param Domain The domain of the access relation.
|
||
/// @param AccMap The access relation to be checked.
|
||
/// @param IndexSet The subset of the input dimensions.
|
||
/// @param DimensionSizes Sizes of the input dimensions of @p Dimensions.
|
||
/// @param Dimensions The permutation of the subset of the input dimensions.
|
||
/// @return True if @p AccMap has the expected form and false,
|
||
/// otherwise.
|
||
static bool isTCOperandAcc(isl::set Domain, isl::map AccMap,
|
||
SmallDenseSet<int> &IndexSet,
|
||
SmallVectorImpl<int> &DimensionSizes,
|
||
SmallVectorImpl<int> &Dimensions) {
|
||
isl::id Id = AccMap.get_tuple_id(isl::dim::out);
|
||
const ScopArrayInfo *SAI = ScopArrayInfo::getFromId(Id);
|
||
assert(SAI && "AccMap should represent memory access");
|
||
|
||
// Fix values of output dimensions with respect to their positions.
|
||
// In the case of the tensor contraction, values of output dimensions are
|
||
// fixed and form a permutation of a subset of values of input dimensions.
|
||
//
|
||
// For example, in the case of Stmt[i][j][k] -> A[k][i], which represents
|
||
// the operand of the tensor contraction, we get the following map by fixing
|
||
// the output dimensions Stmt[1][j][0] -> A[0][1].
|
||
//
|
||
// We store the permutation of the subset of the input dimensions {2, 0} into
|
||
// @p Dimensions.
|
||
//
|
||
// The obtained permutation and the isCorrectAccessMap function are used to
|
||
// check whether the access relation @p AccMap represents the tensor
|
||
// contraction operand. For example, in the case of
|
||
// Stmt[i][j][k] -> A[i-1][j+1], we get Stmt[1][0][k] -> A[0][1] and,
|
||
// consequently, {1, 0}, which is rejected by isCorrectAccessMap,
|
||
// since it corresponds to Stmt[i][j][k] -> A[j][i].
|
||
isl::map CheckMap = isl::manage(AccMap.copy());
|
||
unsigned OutDimNum = unsignedFromIslSize(CheckMap.dim(isl::dim::out));
|
||
for (unsigned i = 0; i < OutDimNum; i++)
|
||
CheckMap = CheckMap.fix_si(isl::dim::out, i, i);
|
||
|
||
// Try to obtain the permutation and sizes of corresponding input dimensions.
|
||
Dimensions.assign(OutDimNum, -1);
|
||
for (unsigned i : rangeIslSize(0, CheckMap.dim(isl::dim::in))) {
|
||
isl::val Val = getConstant(CheckMap, isl::dim::in, i);
|
||
if (!Val.is_int())
|
||
continue;
|
||
int OutPos = -1;
|
||
llvm::APInt ValAPInt = APIntFromVal(Val);
|
||
if (ValAPInt.isSignedIntN(32))
|
||
OutPos = ValAPInt.getSExtValue();
|
||
if ((OutPos < 0) || (OutPos >= static_cast<int>(OutDimNum)) ||
|
||
IndexSet.count(i))
|
||
return false;
|
||
IndexSet.insert(i);
|
||
Dimensions[OutPos] = i;
|
||
if (DimensionSizes[i] <= 0)
|
||
DimensionSizes[i] = getDimSize(SAI, OutPos);
|
||
}
|
||
|
||
return isCorrectAccessMap(Domain, AccMap, Dimensions);
|
||
}
|
||
|
||
/// Find the intersection of two sets.
|
||
///
|
||
/// Find the intersection of the set @p A and the set @p B.
|
||
///
|
||
/// @param A, B Sets to intersect.
|
||
/// @return The set intersection.
|
||
static SmallDenseSet<int> intersect(const SmallDenseSet<int> &A,
|
||
const SmallDenseSet<int> &B) {
|
||
SmallDenseSet<int> Intersection = A;
|
||
set_intersect(Intersection, B);
|
||
return Intersection;
|
||
}
|
||
|
||
/// Check whether the set is a superset.
|
||
///
|
||
/// Check that the set @p A is a superset of @p B.
|
||
///
|
||
/// @param A, B Sets to be checked.
|
||
/// @return True if the set A is a superset of B.
|
||
static bool isSuperset(const SmallDenseSet<int> &A,
|
||
const SmallDenseSet<int> &B) {
|
||
return intersect(A, B).size() == B.size();
|
||
}
|
||
|
||
/// Find the union of two sets.
|
||
///
|
||
/// Find the union of the set @p A and the set @p B.
|
||
///
|
||
/// @param A, B Sets to unite.
|
||
/// @return The set union.
|
||
static SmallDenseSet<int> unite(const SmallDenseSet<int> &A,
|
||
const SmallDenseSet<int> &B) {
|
||
SmallDenseSet<int> Union = A;
|
||
set_union(Union, B);
|
||
return Union;
|
||
}
|
||
|
||
/// Determine the access that writes to the tensor, which contains
|
||
/// the result of the tensor contraction.
|
||
///
|
||
/// @param Domain The domain of the statement.
|
||
/// @param Stmt The statement, which writes to memory.
|
||
/// @param TCI The information about the tensor contraction.
|
||
/// @param IandJIndexSet The set, which contains free indexes of tensors.
|
||
/// @return The determined MemoryAccess, or nullptr if there is no necessary
|
||
/// access within the SCoP.
|
||
static MemoryAccess *getWriteAccess(isl::set Domain, ScopStmt *Stmt,
|
||
TCInfoTy &TCI,
|
||
SmallDenseSet<int> &IandJIndexSet) {
|
||
TCI.WriteToC = nullptr;
|
||
SmallVector<MemoryAccess *, 32> Accesses = getAccessesInOrder(*Stmt);
|
||
for (MemoryAccess *MemA : reverse(Accesses)) {
|
||
// A TC-like does not contain write scalar memory accesses
|
||
if (!MemA->isLatestArrayKind())
|
||
return nullptr;
|
||
// The last memory access should be a write memory access.
|
||
if (!MemA->isWrite())
|
||
return nullptr;
|
||
|
||
isl::map AccMap = MemA->getLatestAccessRelation();
|
||
if (!isTCOperandAcc(Domain, AccMap, IandJIndexSet, TCI.DimensionSizes,
|
||
TCI.CDimensions))
|
||
return nullptr;
|
||
|
||
return MemA;
|
||
}
|
||
return nullptr;
|
||
}
|
||
|
||
/// Determine an access, which reads elements of an operand of the tensor
|
||
/// contraction
|
||
///
|
||
/// @param MemAccessPtr The access, which reads elements of the tensor.
|
||
/// @param IndexSet The set, which contains indexes of the tensors.
|
||
/// @param IandJIndexSet The set, which contains free indexes of tensors.
|
||
/// @param Dimensions The permutation of the subset of the input dimensions.
|
||
/// @param TCI The information about the tensor contraction.
|
||
/// @return True if the memory access @p MemAccessPtr corresponds
|
||
/// to the tensor contraction.
|
||
static bool setReadAccess(MemoryAccess *MemAccessPtr,
|
||
const SmallDenseSet<int> &IndexSet,
|
||
const SmallDenseSet<int> &IandJIndexSet,
|
||
ArrayRef<int> Dimensions, TCInfoTy &TCI) {
|
||
if (!TCI.A) {
|
||
// Probably IndexSet is a union of I and P sets.
|
||
if (!isSuperset(IndexSet, TCI.P))
|
||
return false;
|
||
|
||
// Obtain the set I.
|
||
TCI.I = set_difference(IndexSet, TCI.P);
|
||
if (!isSuperset(IandJIndexSet, TCI.I))
|
||
return false;
|
||
|
||
// Obtain the set J.
|
||
TCI.J = set_difference(IandJIndexSet, TCI.I);
|
||
|
||
// Set the first operand of the tensor contraction.
|
||
TCI.A = MemAccessPtr;
|
||
llvm::replace(TCI.ADimensions, TCI.ADimensions.begin(),
|
||
TCI.ADimensions.end(), Dimensions.begin(), Dimensions.end());
|
||
return true;
|
||
}
|
||
|
||
if (!TCI.B) {
|
||
// IndexSet should be a union of J and P sets.
|
||
if (unite(TCI.P, TCI.J) != IndexSet)
|
||
return false;
|
||
|
||
// Set the second operand of the tensor contraction.
|
||
TCI.B = MemAccessPtr;
|
||
llvm::replace(TCI.BDimensions, TCI.BDimensions.begin(),
|
||
TCI.BDimensions.end(), Dimensions.begin(), Dimensions.end());
|
||
return true;
|
||
}
|
||
|
||
return false;
|
||
}
|
||
|
||
/// Check that all memory accesses of the statement, except from the last
|
||
/// one, are read memory accesses, which read elements of operands of the tensor
|
||
/// contraction and its result.
|
||
///
|
||
/// @param Domain The domain of the statement.
|
||
/// @param Stmt The statement, which writes to memory.
|
||
/// @param TCI The information about the tensor contraction.
|
||
/// @param IandJIndexSet The set, which contains free indexes of tensors.
|
||
/// @return True if all read memory accesses of the statement @p Stmt correspond
|
||
/// to the tensor contraction.
|
||
static bool setReadAccesses(isl::set Domain, ScopStmt *Stmt, TCInfoTy &TCI,
|
||
SmallDenseSet<int> &IandJIndexSet) {
|
||
TCI.A = nullptr;
|
||
TCI.B = nullptr;
|
||
TCI.ReadFromC = nullptr;
|
||
SmallVector<MemoryAccess *, 32> Accesses = getAccessesInOrder(*Stmt);
|
||
for (auto *MemA = Accesses.begin(); *MemA != TCI.WriteToC; MemA++) {
|
||
MemoryAccess *MemAccessPtr = *MemA;
|
||
|
||
// All memory accesses, except from the last one, should be read memory
|
||
// accesses.
|
||
if (MemAccessPtr->isWrite())
|
||
return false;
|
||
|
||
isl::map AccMap = MemAccessPtr->getLatestAccessRelation();
|
||
|
||
if (!MemAccessPtr->isLatestArrayKind()) {
|
||
// Check whether the scalar read memory access is not partial.
|
||
if (!Domain.is_subset(AccMap.domain()))
|
||
return false;
|
||
continue;
|
||
return false;
|
||
}
|
||
|
||
// There is only one memory access, which reads elements of the result of
|
||
// the tensor contraction.
|
||
if (AccMap.is_equal(TCI.WriteToC->getLatestAccessRelation())) {
|
||
if (TCI.ReadFromC)
|
||
return false;
|
||
TCI.ReadFromC = MemAccessPtr;
|
||
continue;
|
||
}
|
||
|
||
SmallVector<int> Dimensions;
|
||
SmallDenseSet<int> IndexSet;
|
||
if (!isTCOperandAcc(Domain, AccMap, IndexSet, TCI.DimensionSizes,
|
||
Dimensions))
|
||
return false;
|
||
|
||
if (!setReadAccess(MemAccessPtr, IndexSet, IandJIndexSet, Dimensions, TCI))
|
||
return false;
|
||
}
|
||
|
||
// Check that there are read memory accesses, which read elements of operands
|
||
// of the tensor contraction and its result.
|
||
return TCI.ReadFromC && TCI.A && TCI.B;
|
||
}
|
||
|
||
/// Check accesses to operands of the tensor contraction.
|
||
///
|
||
/// Check that accesses of the SCoP statement, which corresponds to
|
||
/// the partial schedule @p PartialSchedule, represent accesses
|
||
/// to the non-scalar operands of the tensor contraction.
|
||
///
|
||
/// @param Domain The domain of the SCoP statement.
|
||
/// @param PartialSchedule The partial schedule of the SCoP statement.
|
||
/// @param TCI Parameters of the tensor contraction operands.
|
||
/// @return True if the corresponding SCoP statement
|
||
/// represents tensor contraction and false,
|
||
/// otherwise.
|
||
static bool containsOnlyTCAcc(isl::set Domain, isl::map PartialSchedule,
|
||
TCInfoTy &TCI) {
|
||
isl::id InputDimsId = PartialSchedule.get_tuple_id(isl::dim::in);
|
||
ScopStmt *Stmt = static_cast<ScopStmt *>(InputDimsId.get_user());
|
||
|
||
// In region statements, the order of memory accesses execution is not
|
||
// predictable at compile-time.
|
||
if ((Stmt->size() <= 1) || Stmt->isRegionStmt())
|
||
return false;
|
||
|
||
unsigned DimNum = unsignedFromIslSize(PartialSchedule.dim(isl::dim::in));
|
||
TCI.DimensionSizes.resize(DimNum);
|
||
SmallDenseSet<int> IandJIndexSet;
|
||
|
||
TCI.WriteToC = getWriteAccess(Domain, Stmt, TCI, IandJIndexSet);
|
||
if (!TCI.WriteToC)
|
||
return false;
|
||
|
||
if (intersect(IandJIndexSet, TCI.P).size() != 0)
|
||
return false;
|
||
|
||
if (!setReadAccesses(Domain, Stmt, TCI, IandJIndexSet))
|
||
return false;
|
||
|
||
return true;
|
||
}
|
||
|
||
/// Check that dependency corresponds to the tensor contraction carried over
|
||
/// loop dimension @p Dim.
|
||
///
|
||
/// Check that the dependency has the form
|
||
/// S(..., ki, max(k(i + 1)), ..., max(kn), ...) ->
|
||
/// S(..., ki + 1, min(k(i + 1)), ..., min(kn), ...), where S is the SCoP
|
||
/// statement. For this purpose, we analyze the set @p DepDelta, which
|
||
/// represents the differences between image elements and domain elements of
|
||
/// the corresponding map.
|
||
///
|
||
/// @param DepDelta The set contains the differences between image elements
|
||
/// and corresponding domain elements of the map, which
|
||
/// represents the dependency.
|
||
/// @param Dim The position of the index ki.
|
||
/// @param BoundDeltas In the case of indexes of ki, the difference between
|
||
/// image elements and corresponding domain elements
|
||
/// corresponds to the difference between lexicographic
|
||
/// minimum and lexicographic maximum of the corresponding
|
||
/// dimension of the domain of the statement.
|
||
/// @param IndexSet Obtained indexes ki, which describe the dependency.
|
||
/// @return True if dependencies correspond to the tensor contraction
|
||
/// and false, otherwise.
|
||
static bool isReductionCarriedOverDim(isl::set DepDelta, unsigned Dim,
|
||
isl::pw_multi_aff BoundDeltas,
|
||
const SmallDenseSet<int> &IndexSet) {
|
||
isl::space Space = DepDelta.get_space();
|
||
isl::set Superset = isl::set::universe(Space);
|
||
for (unsigned i = 0; i < Dim; i += 1)
|
||
Superset = Superset.fix_si(isl::dim::set, i, 0);
|
||
Superset = Superset.fix_si(isl::dim::set, Dim, 1);
|
||
|
||
// Check that the difference between the image element and the domain element
|
||
// is equal to one in the case of the index ki. Image elements and
|
||
// corresponding domain elements should be equal in the case of positions,
|
||
// which are lower than the specified position.
|
||
if (!DepDelta.is_subset(Superset))
|
||
return false;
|
||
|
||
// Compute a set, which is used to analyze how values of
|
||
// the domain are related to the map that describes the dependency.
|
||
isl_pw_multi_aff *DepDeltaPW = isl_pw_multi_aff_from_set(DepDelta.copy());
|
||
BoundDeltas = BoundDeltas.add(isl::manage(DepDeltaPW));
|
||
isl_set *ComplementRawSet = isl_set_from_pw_multi_aff(BoundDeltas.release());
|
||
isl::set Complement = isl::manage(ComplementRawSet);
|
||
|
||
for (unsigned i : rangeIslSize(Dim + 1, DepDelta.dim(isl::dim::set))) {
|
||
if (!IndexSet.count(i)) {
|
||
// Check the difference between the image element and the domain element
|
||
// in the case of indexes, which do not describe the dependency.
|
||
if (DepDelta.plain_get_val_if_fixed(isl::dim::set, i).is_zero())
|
||
continue;
|
||
return false;
|
||
}
|
||
|
||
// In the case of other indexes, which describe the dependency,
|
||
// the difference between the image element and the domain element
|
||
// should be equal to the difference between lexicographic minimum and
|
||
// lexicographic maximum of the domain of the statement.
|
||
if (!Complement.plain_get_val_if_fixed(isl::dim::set, i).is_zero())
|
||
return false;
|
||
}
|
||
|
||
return true;
|
||
}
|
||
|
||
/// Check whether dependencies are over the complete domain.
|
||
///
|
||
/// In the case of the tensor contraction RAW, WAW, WAR dependencies
|
||
/// have the form
|
||
/// S(..., ki, max(k(i + 1)), ..., max(kn), ...) ->
|
||
/// S(..., ki + 1, min(k(i + 1)), ..., min(kn), ...), where S is the SCoP
|
||
/// statement. Consequently, the domain of the dependencies
|
||
/// can be described as
|
||
/// Domain / Domain ∩ S(…, max(kn),…) ∩ S(…, max(k(i + 1)),…),
|
||
/// where Domain is the domain of the statement S.
|
||
///
|
||
/// For example, in the case of the following tensor contraction,
|
||
/// corresponding domains will have the following form.
|
||
///
|
||
/// An example of the tensor contraction:
|
||
/// for (i = 0; i < 1024; i++)
|
||
/// for (j = 0; j < 1024; j++)
|
||
/// for (l = 0; l < 64; ++l)
|
||
/// for (w = 0; w < 64; ++w)
|
||
/// C[i][j] += A[i][l][w] * B[w][j][l];
|
||
///
|
||
/// The domain of the statement:
|
||
/// { S[i0, i1, i2, i3] : i0 >= 0 and i0 <= 1023 and
|
||
/// i1 >= 0 and i1 <= 1023 and
|
||
/// i2 >= 0 and i2 <= 63 and
|
||
/// i3 >= 0 and i3 <= 63 }
|
||
///
|
||
/// The domain of the dependencies:
|
||
/// { S[i0, i1, i2, i3] : (i0 >= 0 and i0 <= 1023 and
|
||
/// i1 >= 0 and i1 <= 1023 and
|
||
/// i2 >= 0 and i2 <= 63 and
|
||
/// i3 >= 0 and i3 <= 62) or
|
||
/// (i3 = 63 and i0 >= 0 and i0 <= 1023 and
|
||
/// i1 >= 0 and i1 <= 1023 and
|
||
/// i2 >= 0 and i2 <= 62) }
|
||
///
|
||
/// @param Domain The domain of the statement.
|
||
/// @param DepsForStmt RAW and RED dependencies for the statement.
|
||
/// @param UpperBound The lexicographic maximum of the elements in
|
||
/// the @p Domain.
|
||
/// @param IndexSet Obtained indexes ki, which describe the dependencies.
|
||
/// @return True if dependencies are over the complete domain
|
||
/// and false, otherwise.
|
||
static bool areDepsOverCompleteDomain(isl::set Domain, isl::map DepsForStmt,
|
||
isl::pw_multi_aff UpperBound,
|
||
SmallDenseSet<int> &IndexSet) {
|
||
isl_set *UpperBoundRawSet = isl_set_from_pw_multi_aff(UpperBound.copy());
|
||
isl::set UpperBoundSet = isl::manage(UpperBoundRawSet);
|
||
|
||
isl::set DomainRed = isl::manage(Domain.copy());
|
||
for (const auto It : IndexSet) {
|
||
isl::val FixedVal = UpperBoundSet.plain_get_val_if_fixed(isl::dim::set, It);
|
||
if (FixedVal.is_nan())
|
||
return false;
|
||
DomainRed = isl::manage(
|
||
isl_set_fix_val(DomainRed.copy(), isl_dim_set, It, FixedVal.release()));
|
||
}
|
||
return DepsForStmt.domain().intersect(Domain).is_equal(
|
||
Domain.subtract(DomainRed));
|
||
}
|
||
|
||
/// Check that dependencies correspond to the tensor contraction.
|
||
///
|
||
/// Check that there are only true dependencies of the form
|
||
/// S(..., ki, max(k(i + 1)), ..., max(kn), ...) ->
|
||
/// S(..., ki + 1, min(k(i + 1)), ..., min(kn), ...), where S is the SCoP
|
||
/// statement represented by @p Schedule. Such dependencies are produced by
|
||
/// the tensor contraction. Obtained indexes ki are stored into @p IndexSet.
|
||
///
|
||
/// The form of anti and output dependencies is specified implicitly by
|
||
/// the form the SCoP statement, which is checked by subsequent analysis.
|
||
///
|
||
/// @param Schedule The schedule of the SCoP statement.
|
||
/// @param D The SCoP dependencies.
|
||
/// @param Domain The domain of the statement.
|
||
/// @param IndexSet Obtained indexes ki, which describe the dependencies.
|
||
/// @return True if dependencies correspond to the tensor contraction
|
||
/// and false, otherwise.
|
||
static bool containsOnlyTcDeps(isl::map Schedule, const Dependences *D,
|
||
SmallDenseSet<int> &IndexSet, isl::set Domain) {
|
||
IslMaxOperationsGuard MaxOpGuard(Schedule.ctx().get(), OptComputeOut);
|
||
|
||
isl::union_map Dep =
|
||
D->getDependences(Dependences::TYPE_RAW | Dependences::TYPE_RED);
|
||
|
||
isl::space DomainSpace = Schedule.get_space().domain();
|
||
isl::space Space = DomainSpace.map_from_domain_and_range(DomainSpace);
|
||
isl::map DepsForStmt = Dep.extract_map(Space);
|
||
isl::set DepDeltas = DepsForStmt.deltas();
|
||
isl::size DeltasDimNum = DepDeltas.dim(isl::dim::set);
|
||
isl::pw_multi_aff LowerBound = Domain.lexmin_pw_multi_aff();
|
||
isl::pw_multi_aff UpperBound = Domain.lexmax_pw_multi_aff();
|
||
isl::pw_multi_aff BoundDeltas = UpperBound.sub(LowerBound);
|
||
|
||
for (int i : reverse(rangeIslSize(0, DeltasDimNum))) {
|
||
// In the case of the tensor contraction, the difference between image
|
||
// elements and domain elements lies on a hyperplane where a dimension
|
||
// has the fixed value one.
|
||
isl::set Intersection = DepDeltas.fix_si(isl::dim::set, i, 1);
|
||
if (Intersection.is_empty())
|
||
continue;
|
||
|
||
if (!isReductionCarriedOverDim(Intersection, i, BoundDeltas, IndexSet))
|
||
return false;
|
||
|
||
IndexSet.insert(i);
|
||
DepDeltas = DepDeltas.subtract(Intersection);
|
||
}
|
||
|
||
// In the case of the tensor contraction, all dependencies should have
|
||
// the previously described form.
|
||
if ((unsignedFromIslSize(DeltasDimNum) == 0) || !DepDeltas.is_empty())
|
||
return false;
|
||
|
||
return areDepsOverCompleteDomain(Domain, DepsForStmt, UpperBound, IndexSet);
|
||
}
|
||
|
||
/// Check if the SCoP statement could probably be optimized with analytical
|
||
/// modeling.
|
||
///
|
||
/// containsTCInfoTy tries to determine whether the following conditions
|
||
/// are true:
|
||
///
|
||
/// 1. The last memory access modeling an array, MA1, represents writing to
|
||
/// memory and has the form S(..., I, ..., J, ...) -> M(shuffle(I, J)),
|
||
/// where S is the SCoP statement under consideration and shuffle(I, J)
|
||
/// is a permutation of indexes of sets I and J.
|
||
/// 2. There are only true dependencies of the form
|
||
/// S(..., ki, max(k(i + 1)), ..., max(kn), ...) ->
|
||
/// S(..., ki + 1, min(k(i + 1)), ..., min(kn), ...), where S is the SCoP
|
||
/// statement represented by @p Schedule and ki are indexes of the set P.
|
||
/// 3. SCoP contains an arbitrary number of reads from constants and only three
|
||
/// access relations, MA2, MA3, and MA4 that represent reading from memory
|
||
/// and have the form
|
||
/// S(..., I, ..., P, ...) -> M(shuffle(I, P)),
|
||
/// S(..., P, ..., J, ...) -> M(shuffle(J, P)),
|
||
/// S(...) -> M(shuffle(I, J)), respectively.
|
||
///
|
||
/// @param PartialSchedule The PartialSchedule that contains a SCoP statement
|
||
/// to check.
|
||
/// @param D The SCoP dependencies.
|
||
/// @param TCI Parameters of the tensor contraction operands.
|
||
/// @param Domain The domain of the statement.
|
||
/// @return True if dependencies and memory accesses correspond to the tensor
|
||
/// contraction and false, otherwise.
|
||
static bool containsTCInfoTy(isl::map PartialSchedule, const Dependences *D,
|
||
TCInfoTy &TCI, isl::set Domain) {
|
||
if (!containsOnlyTcDeps(PartialSchedule, D, TCI.P, Domain))
|
||
return false;
|
||
|
||
// TODO: handle cases of scalar multiplication if needed.
|
||
if (TCI.P.size() == 0)
|
||
return false;
|
||
|
||
if (!containsOnlyTCAcc(Domain, PartialSchedule, TCI))
|
||
return false;
|
||
|
||
// TODO: handle cases of GEMV if needed.
|
||
if ((TCI.I.size() == 0) || (TCI.J.size() == 0))
|
||
return false;
|
||
|
||
return true;
|
||
}
|
||
|
||
/// Check if this node contains a partial schedule that could
|
||
/// probably be optimized with analytical modeling.
|
||
///
|
||
/// isTCPattern is used to determine whether the SCoP represents a TC-like
|
||
/// kernel [1], which is a perfectly nested set of loops, with a data usage
|
||
/// pattern that is similar to that produced by the tensor contraction.
|
||
///
|
||
/// A TC-like kernel can be defined as follows:
|
||
///
|
||
/// 1. It satisfies the requirements of the polyhedral model.
|
||
/// 2. Without loss of generality, it contains three nonempty bundles of
|
||
/// one-dimensional for-loops with induction variables that are grouped into
|
||
/// bundles I = i0...i(r-1), J = j0..j(s-1), and P = p0...p(t-1), and they
|
||
/// are incremented by one.
|
||
/// 3. The innermost loop body can be represented as a statement of the form
|
||
/// C(shuffle(I, J)) = E(A(shuffle(I, P)), B(shuffle(P, J)),
|
||
/// C(shuffle(I, J))), where A(shuffle(I, P)), B(shuffle(P, J)),
|
||
/// C(shuffle(I, J)) are accesses to tensors A, B, C, respectively,
|
||
/// shuffle(I, J), shuffle(I, P), and shuffle(P, J) are permutations of the
|
||
/// enclosed indices, and E is an expression that contains reads from
|
||
/// the tensors A, B, C, and an arbitrary number of reads from constants
|
||
/// with respect to bundles I, J, and P.
|
||
///
|
||
/// TC can be considered as a particular case of a TC-like kernel.
|
||
///
|
||
/// The order of loops with indexes from P should be preserved. Otherwise,
|
||
/// isTCPattern should check if a commutative operation is used.
|
||
///
|
||
/// isTCPattern performs the following steps to check whether the SCoP
|
||
/// corresponds to a definition of a TC-like kernel:
|
||
///
|
||
/// 1. Checks that the node is the innermost band node.
|
||
/// 2. Checks that the partial schedule contains only one statement.
|
||
/// 3. Check that all ancestors of the node contain all band nodes for
|
||
/// the statement and only mark nodes interleave such band nodes. This
|
||
/// corresponds to a straightforward implementation of TC.
|
||
/// 4. Analyses the dependencies to determine contraction dimensions.
|
||
/// 5. Check that the last memory access modeling an array, represents writing
|
||
/// to the result of the TC-like kernel.
|
||
/// 6. Check that SCoP contains only three access relations that represent
|
||
/// reading of the operands of the TC-like kernel and an arbitrary number of
|
||
/// reads from constants.
|
||
///
|
||
/// [1] - Gareev R., Grosser T., Kruse M. High-Performance Generalized Tensor
|
||
/// Operations: A Compiler-Oriented Approach // ACM Transactions
|
||
/// Architecture and Code Optimization (TACO). 2018.
|
||
/// Vol. 15, no. 3. P. 34:1–34:27. DOI: 10.1145/3235029.
|
||
///
|
||
/// If this is the case, we could logically represent tensors as matrices and
|
||
/// apply algorithms, which are used to get close-to-peak performance of
|
||
/// matrix multiplications in manually tuned BLAS libraries (e.g., BLIS).
|
||
///
|
||
/// @param Node The node to check.
|
||
/// @param D The SCoP dependencies.
|
||
/// @param TCI Parameters of the tensor contraction operands.
|
||
static bool isTCPattern(isl::schedule_node Node, const Dependences *D,
|
||
TCInfoTy &TCI) {
|
||
Node = Node.child(0);
|
||
isl::union_map PartialSchedule = Node.get_prefix_schedule_union_map();
|
||
isl::union_set Domain = Node.domain();
|
||
Node = Node.parent();
|
||
|
||
// The partial schedule should contain only one statement.
|
||
// TODO: This constraint should not be intrinsic to the algorithm.
|
||
if (isl_union_set_n_set(Domain.get()) != 1)
|
||
return false;
|
||
|
||
isl_schedule_node_type NodeType = isl_schedule_node_get_type(Node.get());
|
||
|
||
// Check that all ancestors of the node contain all band nodes for
|
||
// the statement, which represents the TC-like kernel, and only mark nodes
|
||
// interleave such band nodes. This corresponds to a straightforward
|
||
// implementation of TC with/without DeLICM applied.
|
||
//
|
||
// For example, this covers the matrix multiplication pattern after a full
|
||
// run of -polly-optree and -polly-delicm, where the write access is not
|
||
// through the original memory access, but trough a PHI node that was
|
||
// delicmed. Subsequently, such band nodes will be replaced by a single band
|
||
// node.
|
||
//
|
||
// The corresponding schedule can be the following, where Stmt_for_body8
|
||
// contains the matrix multiplication:
|
||
//
|
||
// domain: "{ Stmt_for_body8[i0, i1, i2] : 0 <= i0 <= 1599 and
|
||
// 0 <= i1 <= 1799 and
|
||
// 0 <= i2 <= 2199;
|
||
// Stmt_for_body3[i0, i1] : 0 <= i0 <= 1599 and
|
||
// 0 <= i1 <= 1799;
|
||
// Stmt_for_body3_last[i0, i1] : 0 <= i0 <= 1599 and
|
||
// 0 <= i1 <= 1799 }"
|
||
// child:
|
||
// sequence:
|
||
// - filter: "{ Stmt_for_body3[i0, i1] }"
|
||
// child:
|
||
// schedule: "[{ Stmt_for_body3[i0, i1] -> [(i0)] },
|
||
// { Stmt_for_body3[i0, i1] -> [(i1)] }]"
|
||
// permutable: 1
|
||
// coincident: [ 1, 1 ]
|
||
// - filter: "{ Stmt_for_body3_last[i0, i1] }"
|
||
// child:
|
||
// schedule: "[{ Stmt_for_body3_last[i0, i1] -> [(i0)] },
|
||
// { Stmt_for_body3_last[i0, i1] -> [(i1)] }]"
|
||
// permutable: 1
|
||
// coincident: [ 1, 1 ]
|
||
// - filter: "{ Stmt_for_body8[i0, i1, i2] }"
|
||
// child:
|
||
// schedule: "[{ Stmt_for_body8[i0, i1, i2] -> [(i0)] },
|
||
// { Stmt_for_body8[i0, i1, i2] -> [(i1)] },
|
||
// { Stmt_for_body8[i0, i1, i2] -> [(i2)] }]"
|
||
// permutable: 1
|
||
// coincident: [ 1, 1, 0 ]
|
||
//
|
||
while (NodeType != isl_schedule_node_domain) {
|
||
if (NodeType == isl_schedule_node_filter) {
|
||
if (!Node.parent().isa<isl::schedule_node_sequence>() ||
|
||
!Node.parent().parent().isa<isl::schedule_node_domain>())
|
||
return false;
|
||
break;
|
||
}
|
||
|
||
if ((NodeType != isl_schedule_node_band) &&
|
||
(NodeType != isl_schedule_node_mark))
|
||
return false;
|
||
|
||
Node = Node.parent();
|
||
NodeType = isl_schedule_node_get_type(Node.get());
|
||
}
|
||
|
||
isl::map PartialScheduleMap = isl::map::from_union_map(PartialSchedule);
|
||
if (containsTCInfoTy(PartialScheduleMap, D, TCI, isl::set(Domain)))
|
||
return true;
|
||
|
||
return false;
|
||
}
|
||
|
||
} // namespace
|
||
|
||
isl::schedule_node
|
||
polly::tryOptimizeMatMulPattern(isl::schedule_node Node,
|
||
const llvm::TargetTransformInfo *TTI,
|
||
const Dependences *D) {
|
||
TCInfoTy TCI;
|
||
if (PMBasedTCOpts && isTCPattern(Node, D, TCI))
|
||
LLVM_DEBUG(dbgs() << "The tensor contraction pattern was detected\n");
|
||
MatMulInfoTy MMI;
|
||
if (PMBasedMMMOpts && isMatrMultPattern(Node, D, MMI)) {
|
||
LLVM_DEBUG(dbgs() << "The matrix multiplication pattern was detected\n");
|
||
return optimizeMatMulPattern(Node, TTI, MMI);
|
||
}
|
||
return {};
|
||
}
|