16#include "mlir/Dialect/Affine/IR/AffineOps.h"
17#include "mlir/Dialect/Arith/IR/Arith.h"
18#include "mlir/Dialect/Vector/IR/VectorOps.h"
19#include "llvm/ADT/TypeSwitch.h"
22#define DEBUG_TYPE "aievec-utils"
28static std::optional<int64_t> getLowerBoundValue(Value idx) {
29 if (
auto blkArg = dyn_cast<BlockArgument>(idx)) {
30 auto parentOp = blkArg.getOwner()->getParentOp();
31 return TypeSwitch<Operation *, std::optional<int64_t>>(parentOp)
32 .Case<affine::AffineForOp>([&blkArg](affine::AffineForOp forOp) {
33 if (forOp.getInductionVar() == blkArg &&
34 forOp.hasConstantLowerBound())
35 return std::optional<int64_t>(forOp.getConstantLowerBound());
40 return std::optional<int64_t>();
42 .Default([](
auto) {
return std::optional<int64_t>(); });
44 return TypeSwitch<Operation *, std::optional<int64_t>>(idx.getDefiningOp())
45 .Case<arith::ConstantOp>([](
auto constantOp) {
46 return std::optional<int64_t>(
47 cast<IntegerAttr>(constantOp.getValue()).getInt());
49 .Case<affine::AffineApplyOp>([](
auto applyOp) {
50 if (applyOp.getAffineMap().getNumResults() == 1) {
51 SmallVector<int64_t, 4> srcIndices;
52 for (
auto index : applyOp.getMapOperands()) {
53 std::optional<int64_t> lbv = getLowerBoundValue(index);
56 if (!lbv && !isa<BlockArgument>(index))
57 return std::optional<int64_t>();
58 srcIndices.push_back(lbv.value_or(0L));
60 return std::optional<int64_t>(
61 applyOp.getAffineMap().compose(srcIndices)[0]);
63 return std::optional<int64_t>();
65 .Default([&](
auto) {
return std::optional<int64_t>(); });
72template <
typename TransferReadLikeOp,
typename>
80 auto innerMostIndex = readOp.getIndices().back();
81 auto vectorLength = vType.getShape().back();
82 std::optional<int64_t> lbv = getLowerBoundValue(innerMostIndex);
85 int64_t vectorLengthAlignmentOffset = lbv.value() % vectorLength;
86 int64_t absoluteAlignmentOffset = alignment / vType.getElementTypeBitWidth();
87 if (vectorLengthAlignmentOffset % absoluteAlignmentOffset)
88 return vectorLengthAlignmentOffset;
92template std::optional<int64_t>
95template std::optional<int64_t>
97 VectorType vType, int64_t alignment);
100 if (vecTy.getRank() == 1)
102 auto shape = vecTy.getShape();
103 return VectorType::get(
104 {std::accumulate(shape.begin(), shape.end(), 1, std::multiplies<>())},
105 vecTy.getElementType());
mlir::VectorType getFlattenedVectorType(mlir::VectorType vecTy)
std::optional< int64_t > getTransferReadAlignmentOffset(TransferReadLikeOp readOp, mlir::VectorType vType, int64_t alignment)