Utils.cpp

Go to the documentation of this file.

//===- Utils.cpp - Utilities to support AIE vectorization -----------------===//
//
// Copyright (C) 2023 Advanced Micro Devices, Inc.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
//
// This file implements utilities for the AIEVec dialect
//
//===----------------------------------------------------------------------===//
 
#include "aie/Dialect/AIEVec/Utils/Utils.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "llvm/ADT/TypeSwitch.h"
#include <numeric>
 
#define DEBUG_TYPE "aievec-utils"
 
using namespace mlir;
 
namespace xilinx::aievec {
 
static std::optional<int64_t> getLowerBoundValue(Value idx) {
  if (auto blkArg = dyn_cast<BlockArgument>(idx)) {
    auto *parentOp = blkArg.getOwner()->getParentOp();
    return TypeSwitch<Operation *, std::optional<int64_t>>(parentOp)
        .Case<affine::AffineForOp>([&blkArg](affine::AffineForOp forOp) {
          if (forOp.getInductionVar() == blkArg &&
              forOp.hasConstantLowerBound())
            return std::optional<int64_t>(forOp.getConstantLowerBound());
          // If it's an iteration argument or the lower bound is an
          // affine expression.
          // TODO: Compute the value of the lower bound affine expression
          // TODO: if it's constant.
          return std::optional<int64_t>();
        })
        .Default([](auto) { return std::optional<int64_t>(); });
  }
  return TypeSwitch<Operation *, std::optional<int64_t>>(idx.getDefiningOp())
      .Case<arith::ConstantOp>([](auto constantOp) {
        return std::optional<int64_t>(
            cast<IntegerAttr>(constantOp.getValue()).getInt());
      })
      .Case<affine::AffineApplyOp>([](auto applyOp) {
        if (applyOp.getAffineMap().getNumResults() == 1) {
          auto affineMap = applyOp.getAffineMap();
 
          // Check if the map has symbols - if so, try symbol-aware evaluation
          if (affineMap.getNumSymbols() > 0) {
            auto operands = applyOp.getMapOperands();
            unsigned numDims = affineMap.getNumDims();
 
            SmallVector<int64_t, 4> dimValues;
            SmallVector<int64_t, 4> symbolValues;
 
            // Collect dimension values
            for (unsigned i = 0; i < numDims; i++) {
              std::optional<int64_t> lbv = getLowerBoundValue(operands[i]);
              if (!lbv && !isa<BlockArgument>(operands[i]))
                return std::optional<int64_t>();
              dimValues.push_back(lbv.value_or(0L));
            }
 
            // Collect symbol values
            for (unsigned i = numDims; i < operands.size(); i++) {
              std::optional<int64_t> lbv = getLowerBoundValue(operands[i]);
              if (!lbv && !isa<BlockArgument>(operands[i]))
                return std::optional<int64_t>();
              symbolValues.push_back(lbv.value_or(0L));
            }
 
            // Manually replace and constant fold
            auto expr = affineMap.getResult(0);
            auto ctx = affineMap.getContext();
 
            // Convert int64_t values to AffineConstantExpr
            SmallVector<AffineExpr, 4> dimExprs;
            SmallVector<AffineExpr, 4> symbolExprs;
            for (auto val : dimValues)
              dimExprs.push_back(getAffineConstantExpr(val, ctx));
            for (auto val : symbolValues)
              symbolExprs.push_back(getAffineConstantExpr(val, ctx));
 
            expr = expr.replaceDimsAndSymbols(dimExprs, symbolExprs);
            if (auto constExpr = dyn_cast<AffineConstantExpr>(expr)) {
              return std::optional<int64_t>(constExpr.getValue());
            }
 
            // If expression couldn't be constant-folded, return nullopt
            return std::optional<int64_t>();
          }
 
          auto operands = applyOp.getMapOperands();
          SmallVector<int64_t, 4> srcIndices;
          for (auto index : operands) {
            std::optional<int64_t> lbv = getLowerBoundValue(index);
            // XXX: We assume block arguments to either have well-defined
            // XXX: compile-time values, or to be aligned.
            if (!lbv && !isa<BlockArgument>(index))
              return std::optional<int64_t>();
            srcIndices.push_back(lbv.value_or(0L));
          }
          return std::optional<int64_t>(affineMap.compose(srcIndices)[0]);
        }
        return std::optional<int64_t>();
      })
      .Default([&](auto) { return std::optional<int64_t>(); });
}
 
// Return the offset of a given transfer read operation with regards to the
// specified vector type. If the read is aligned to the specified alignment
// parameter (in bits), then the offset is 0. Otherwise, the offset is the
// number of elements past the immediately preceding aligned vector length.
template <typename TransferReadLikeOp, typename>
std::optional<int64_t> getTransferReadAlignmentOffset(TransferReadLikeOp readOp,
                                                      VectorType vType,
                                                      int64_t alignment) {
  // TODO: Add support for cases where the index is not comming from an
  // TODO: `affine.apply` op or when the affine map has more than one
  // TODO: dimension. We also need to address the case where the index is an
  // TODO: induction variable.
  auto innerMostIndex = readOp.getIndices().back();
  auto vectorLength = vType.getShape().back();
  std::optional<int64_t> lbv = getLowerBoundValue(innerMostIndex);
  if (!lbv)
    return std::nullopt;
  int64_t vectorLengthAlignmentOffset = lbv.value() % vectorLength;
  int64_t absoluteAlignmentOffset = alignment / vType.getElementTypeBitWidth();
  if (vectorLengthAlignmentOffset % absoluteAlignmentOffset)
    return vectorLengthAlignmentOffset;
  return 0;
}
 
template std::optional<int64_t>
getTransferReadAlignmentOffset(vector::TransferReadOp readOp, VectorType vType,
                               int64_t alignment);
template std::optional<int64_t>
getTransferReadAlignmentOffset(vector::TransferReadOp::Adaptor readOp,
                               VectorType vType, int64_t alignment);
 
VectorType getFlattenedVectorType(VectorType vecTy) {
  if (vecTy.getRank() == 1)
    return vecTy;
  auto shape = vecTy.getShape();
  return VectorType::get(
      {std::accumulate(shape.begin(), shape.end(), 1, std::multiplies<>())},
      vecTy.getElementType());
}
 
} // namespace xilinx::aievec