MLIR-AIE
Utils.cpp
Go to the documentation of this file.
1//===- Utils.cpp - Utilities to support AIE vectorization -----------------===//
2//
3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
4// See https://llvm.org/LICENSE.txt for license information.
5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
6//
7// (c) Copyright 2023, Advanced Micro Devices, Inc.
8//
9//===----------------------------------------------------------------------===//
10//
11// This file implements utilities for the AIEVec dialect
12//
13//===----------------------------------------------------------------------===//
14
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"
20#include <numeric>
21
22#define DEBUG_TYPE "aievec-utils"
23
24using namespace mlir;
25
26namespace xilinx::aievec {
27
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());
36 // If it's an iteration argument or the lower bound is an
37 // affine expression.
38 // TODO: Compute the value of the lower bound affine expression
39 // TODO: if it's constant.
40 return std::optional<int64_t>();
41 })
42 .Default([](auto) { return std::optional<int64_t>(); });
43 }
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());
48 })
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);
54 // XXX: We assume block arguments to either have well-defined
55 // XXX: compile-time values, or to be aligned.
56 if (!lbv && !isa<BlockArgument>(index))
57 return std::optional<int64_t>();
58 srcIndices.push_back(lbv.value_or(0L));
59 }
60 return std::optional<int64_t>(
61 applyOp.getAffineMap().compose(srcIndices)[0]);
62 }
63 return std::optional<int64_t>();
64 })
65 .Default([&](auto) { return std::optional<int64_t>(); });
66}
67
68// Return the offset of a given transfer read operation with regards to the
69// specified vector type. If the read is aligned to the specified alignment
70// parameter (in bits), then the offset is 0. Otherwise, the offset is the
71// number of elements past the immediately preceding aligned vector length.
72template <typename TransferReadLikeOp, typename>
73std::optional<int64_t> getTransferReadAlignmentOffset(TransferReadLikeOp readOp,
74 VectorType vType,
75 int64_t alignment) {
76 // TODO: Add support for cases where the index is not comming from an
77 // TODO: `affine.apply` op or when the affine map has more than one
78 // TODO: dimension. We also need to address the case where the index is an
79 // TODO: induction variable.
80 auto innerMostIndex = readOp.getIndices().back();
81 auto vectorLength = vType.getShape().back();
82 std::optional<int64_t> lbv = getLowerBoundValue(innerMostIndex);
83 if (!lbv)
84 return std::nullopt;
85 int64_t vectorLengthAlignmentOffset = lbv.value() % vectorLength;
86 int64_t absoluteAlignmentOffset = alignment / vType.getElementTypeBitWidth();
87 if (vectorLengthAlignmentOffset % absoluteAlignmentOffset)
88 return vectorLengthAlignmentOffset;
89 return 0;
90}
91
92template std::optional<int64_t>
93getTransferReadAlignmentOffset(vector::TransferReadOp readOp, VectorType vType,
94 int64_t alignment);
95template std::optional<int64_t>
96getTransferReadAlignmentOffset(vector::TransferReadOp::Adaptor readOp,
97 VectorType vType, int64_t alignment);
98
99VectorType getFlattenedVectorType(VectorType vecTy) {
100 if (vecTy.getRank() == 1)
101 return vecTy;
102 auto shape = vecTy.getShape();
103 return VectorType::get(
104 {std::accumulate(shape.begin(), shape.end(), 1, std::multiplies<>())},
105 vecTy.getElementType());
106}
107
108} // namespace xilinx::aievec
mlir::VectorType getFlattenedVectorType(mlir::VectorType vecTy)
std::optional< int64_t > getTransferReadAlignmentOffset(TransferReadLikeOp readOp, mlir::VectorType vType, int64_t alignment)