VectorToVectorConversions.cpp

Go to the documentation of this file.

//===-VectorToVectorConversions.cpp - Conversions within Vector -*- C++ -*-===//
//
// Copyright (C) 2023-2024 Advanced Micro Devices, Inc.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
// This file contains conversions and rewrites to the Vector dialect to make
// it compatible with the available vector instructions in AIE architectures
//===----------------------------------------------------------------------===//
 
#include "aie/Dialect/AIEVec/AIEVecUtils.h"
#include "aie/Dialect/AIEVec/Pipelines/Passes.h"
#include "aie/Dialect/AIEVec/Utils/Utils.h"
#include "mlir/Conversion/AffineToStandard/AffineToStandard.h"
#include "mlir/Dialect/Affine/Analysis/LoopAnalysis.h"
#include "mlir/Dialect/Affine/IR/AffineOps.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/Math/IR/Math.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/UB/IR/UBOps.h"
#include "mlir/Dialect/Utils/ReshapeOpsUtils.h"
#include "mlir/Dialect/Vector/Transforms/LoweringPatterns.h"
#include "mlir/IR/Matchers.h"
#include "mlir/IR/PatternMatch.h"
#include "mlir/Pass/PassManager.h"
#include "mlir/Transforms/DialectConversion.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "llvm/ADT/STLExtras.h"
#include <algorithm>
#include <utility>
 
#define DEBUG_TYPE "aievec-canonicalization"
 
using namespace mlir;
using namespace arith;
using namespace vector;
using namespace xilinx;
using namespace xilinx::aievec;
 
//============================================================================//
//================== Common AIE canonicalization analysis ====================//
//============================================================================//
 
static AIEArch decodeAIETarget(const std::string &target) {
  if (!target.empty()) {
    if (target == "aieml" || target == "aie2" || target == "aie2p")
      return AIEArch::AIE2;
    if (target != "aie")
      return AIEArch::UNKNOWN;
  }
  return AIEArch::AIE;
}
 
//============================================================================//
//============ Common AIE canonicalization conversion patterns ===============//
//============================================================================//
 
// This pattern converts a `vector.transfer_read` with an unaligned access
// into an aligned `vector.transfer_read` twice as long, followed by a
// `vector.extract_strided_slice` selecting the subvector matching the
// original `vector.transfer_read`.
struct SplitUnalignedTransferReadPattern
    : public OpConversionPattern<vector::TransferReadOp> {
  using OpConversionPattern<vector::TransferReadOp>::OpConversionPattern;
 
  SplitUnalignedTransferReadPattern(MLIRContext *context, int64_t maxVectorSize,
                                    int64_t alignment)
      : OpConversionPattern<vector::TransferReadOp>(context),
        maxVectorSize(maxVectorSize), vectorAlignment(alignment) {}
 
  LogicalResult
  matchAndRewrite(vector::TransferReadOp readOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    // Check that it's not a splat transfer read.
    if (adaptor.getPermutationMap().isConstant())
      return failure();
 
    // Check if the transfer is unaligned.
    auto vType = readOp.getVectorType();
    int64_t offset =
        getTransferReadAlignmentOffset(adaptor, vType, vectorAlignment)
            .value_or(0);
    if (offset == 0)
      return failure();
 
    // Verify that we can load a vector 2x as long as the original
    auto vLen = vType.getShape().back();
    auto longVecTy = VectorType::get(2 * vLen, vType.getElementType());
    auto longVecSize = getElementSizeInBits(vType) * 2 * vLen;
    if (longVecSize > maxVectorSize)
      return failure();
 
    // Calculate the aligned indices for the lower and higher parts.
    // TODO: Add support for cases where the offset is greater than the
    // TODO: vector length.
    auto loc = readOp.getLoc();
    Value oldInnerMostIdx = adaptor.getIndices().back();
    auto offsetCorrectionMap =
        AffineMap::get(1, 0, getAffineDimExpr(0, readOp.getContext()) - offset);
    Value newInnerMostIdx = affine::AffineApplyOp::create(
                                rewriter, readOp.getLoc(), offsetCorrectionMap,
                                SmallVector<Value, 1>({oldInnerMostIdx}))
                                .getResult();
    SmallVector<Value, 8> alignedIdx;
    alignedIdx.append(adaptor.getIndices().begin(), adaptor.getIndices().end());
    alignedIdx[alignedIdx.size() - 1] = newInnerMostIdx;
 
    // Create the aligned transfer read for a vector 2x as long that covers the
    // elements of the unaligned vector.
    auto newReadOp = vector::TransferReadOp::create(
        rewriter, loc, longVecTy, adaptor.getBase(), alignedIdx,
        adaptor.getPadding());
 
    // Create a `vector.extract_strided_slice` to extract the unaligned vector.
    rewriter.replaceOpWithNewOp<vector::ExtractStridedSliceOp>(
        readOp, newReadOp.getResult(), offset, vLen, 1);
 
    return success();
  }
 
  int64_t maxVectorSize;
  int64_t vectorAlignment;
};
 
// This pattern converts a `vector.transfer_read` with a splat permutation map
// into a contiguous `vector.transfer_read` followed by a `vector.extract` to
// obtain the splat value and a `vector.broadcast` to broadcast it into a
// vector of the right size.
struct ConvertSplatTransferReadToBroadcastPattern
    : public OpConversionPattern<vector::TransferReadOp> {
  using OpConversionPattern<vector::TransferReadOp>::OpConversionPattern;
 
  ConvertSplatTransferReadToBroadcastPattern(MLIRContext *context)
      : OpConversionPattern<vector::TransferReadOp>(context) {}
 
  LogicalResult
  matchAndRewrite(vector::TransferReadOp readOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    AffineMap map = readOp.getPermutationMap();
    if (!map.isConstant())
      return failure();
 
    Value srcMemRef = adaptor.getBase();
    SmallVector<Value, 8> indices;
    Value newIdx;
    int64_t offset = 0;
    // If it's a zero-rank memory access
    if (cast<MemRefType>(srcMemRef.getType()).getRank() == 0) {
      srcMemRef = memref::ExpandShapeOp::create(
                      rewriter, readOp.getLoc(), SmallVector<int64_t, 1>({1}),
                      srcMemRef, SmallVector<ReassociationIndices, 1>({}))
                      .getResult();
      newIdx = arith::ConstantOp::create(rewriter, readOp.getLoc(),
                                         rewriter.getIndexAttr(0L));
      indices.push_back(newIdx);
    } else {
      indices.append(adaptor.getIndices().begin(), adaptor.getIndices().end());
      newIdx = indices[indices.size() - 1];
      // If the innermost index comes from an `affine.apply` op, take the base
      // as the new innermost index for the new `vector.transfer_read`, and the
      // offset as the index for the `aievec.broadcast` op.
      if (auto applyOp = newIdx.getDefiningOp<affine::AffineApplyOp>())
        if (applyOp.getAffineMap().getNumDims() == 1) {
          newIdx = applyOp.getMapOperands()[0];
          offset = applyOp.getAffineMap().compose(ArrayRef<int64_t>{0})[0];
        }
    }
    // XXX: We assume we are reading 1D vectors
    int64_t vlen = readOp.getVector().getType().getShape()[0];
    if (offset >= vlen) {
      // If the splat element is beyond the first vector, we calculate the
      // address of the vector containing the element.
      int64_t numElemsToSkip = vlen * (offset / vlen);
      offset = offset % vlen;
      auto newAddrMap = AffineMap::get(
          1, 0, getAffineDimExpr(0, readOp.getContext()) + numElemsToSkip);
      newIdx =
          affine::AffineApplyOp::create(rewriter, readOp.getLoc(), newAddrMap,
                                        SmallVector<Value, 1>({newIdx}))
              .getResult();
    }
    indices[indices.size() - 1] = newIdx;
    auto newReadOp = vector::TransferReadOp::create(
        rewriter, readOp.getLoc(), readOp.getVector().getType(), srcMemRef,
        indices, adaptor.getPadding());
    auto extractOp = vector::ExtractOp::create(rewriter, readOp.getLoc(),
                                               newReadOp.getResult(),
                                               ArrayRef<int64_t>{offset});
    rewriter.replaceOpWithNewOp<vector::BroadcastOp>(
        readOp, newReadOp.getVector().getType(), extractOp.getResult());
    return success();
  }
};
 
// This pattern swaps a UnaryOpA followed by UnaryOpB. This pattern can be used
// to improve pattern matching for mixed-type arithmetic ops, by getting sign
// extension ops closer to the single-type arithmetic operations.
template <class UnaryOpA, class UnaryOpB>
struct SwapUnaryOpsPattern : public OpRewritePattern<UnaryOpB> {
  using OpRewritePattern<UnaryOpB>::OpRewritePattern;
  // This function takes the chain of operations A->B, and returns the new type
  // between B and A after the swap.
  using InferTypeB2AFnTy = std::function<Type(UnaryOpA aOp, UnaryOpB bOp)>;
  InferTypeB2AFnTy inferTypeB2A = nullptr;
 
  SwapUnaryOpsPattern(MLIRContext *context, InferTypeB2AFnTy inferType)
      : OpRewritePattern<UnaryOpB>(context),
        inferTypeB2A(std::move(inferType)) {}
 
  LogicalResult matchAndRewrite(UnaryOpB bOp,
                                PatternRewriter &rewriter) const override {
    static_assert(
        UnaryOpA::template hasTrait<OpTrait::OneOperand>(),
        "SwapUnaryOps can only be instantiated for single-operand ops");
    static_assert(
        UnaryOpB::template hasTrait<OpTrait::OneOperand>(),
        "SwapUnaryOps can only be instantiated for single-operand ops");
    UnaryOpA aOp = bOp.getOperand().template getDefiningOp<UnaryOpA>();
    if (!aOp)
      return rewriter.notifyMatchFailure(bOp, UnaryOpB::getOperationName() +
                                                  " not preceeded by " +
                                                  UnaryOpA::getOperationName());
 
    Type newA2BTy = inferTypeB2A(aOp, bOp);
 
    auto newA =
        UnaryOpB::create(rewriter, bOp->getLoc(), SmallVector<Type>({newA2BTy}),
                         aOp->getOperands(), bOp->getAttrs());
    auto newB = UnaryOpA::create(rewriter, bOp->getLoc(),
                                 SmallVector<Type>({bOp.getResult().getType()}),
                                 newA->getResults(), aOp->getAttrs());
    rewriter.replaceOp(bOp, newB.getResult());
    return success();
  }
};
 
static SmallVector<Value> collapseInnerMostDimIndices(PatternRewriter &b,
                                                      Location loc, int numDims,
                                                      ValueRange indices,
                                                      ArrayRef<int64_t> shape,
                                                      AffineMap layout) {
  (void)layout; // Layout is verified by callers; index computation uses shape.
  auto newIdxExpr = b.getAffineDimExpr(numDims - 1);
  int64_t stride = 1;
  for (int64_t dim = numDims - 2; dim >= 0; dim--) {
    stride *= shape[shape.size() - (numDims - dim - 1)];
    newIdxExpr = newIdxExpr + b.getAffineDimExpr(dim) * stride;
  }
  auto newIndexMap = AffineMap::get(numDims, 0, newIdxExpr);
  Value newInnerMostIdxValue =
      affine::AffineApplyOp::create(b, loc, newIndexMap,
                                    indices.take_back(numDims))
          .getResult();
  SmallVector<Value> newIdxRange;
  for (auto idx : indices.drop_back(numDims))
    newIdxRange.push_back(idx);
  newIdxRange.push_back(newInnerMostIdxValue);
  return newIdxRange;
}
 
static Value collapseInnerMostShapeDims(PatternRewriter &b, Location loc,
                                        int numDims, Value val) {
  auto memRefTy = cast<MemRefType>(val.getType());
  auto shape = memRefTy.getShape();
  int64_t newInnerMostDim = std::accumulate(shape.end() - numDims, shape.end(),
                                            1, std::multiplies<>());
  SmallVector<int64_t, 4> newShape{shape.begin(), shape.end() - numDims + 1};
  newShape[shape.size() - numDims] = newInnerMostDim;
  // newShape is constructed above as exactly shape with its trailing
  // numDims dims merged into one, so this is always a valid collapse.
  auto reassocIndicesOpt = getReassociationIndicesForCollapse(shape, newShape);
  assert(reassocIndicesOpt && "newShape must be a valid collapse of shape");
  const auto &reassocIndices = *reassocIndicesOpt;
  // Let CollapseShapeOp::inferResultType compute the correct result type,
  // which preserves strided layout and dynamic offset from the source.
  auto newMemRefTy =
      memref::CollapseShapeOp::computeCollapsedType(memRefTy, reassocIndices);
  return memref::CollapseShapeOp::create(b, loc, newMemRefTy, val,
                                         reassocIndices)
      .getResult();
}
 
/// Check if a memref has contiguous row-major strides (each stride equals the
/// product of trailing dimensions). Dynamic strides are accepted when the
/// corresponding dimension size is 1. Memrefs with dynamic offsets are fine.
static bool hasContiguousRowMajorStrides(MemRefType memRefTy) {
  SmallVector<int64_t> strides;
  int64_t offset;
  if (failed(memRefTy.getStridesAndOffset(strides, offset)))
    return false;
  auto shape = memRefTy.getShape();
  int64_t expected = 1;
  for (int i = shape.size() - 1; i >= 0; --i) {
    if (strides[i] != ShapedType::kDynamic && strides[i] != expected)
      return false;
    if (shape[i] != ShapedType::kDynamic)
      expected *= shape[i];
    else
      expected = ShapedType::kDynamic;
  }
  return true;
}
 
// This pattern flatten multidimensional `vector.transfer_read` operations
// replacing them with a `memref.collapse_shape`, a 1D `vector.transfer_read`,
// and a `vector.shape_cast`.
struct FlattenMultDimTransferReadPattern
    : public OpConversionPattern<vector::TransferReadOp> {
  using OpConversionPattern<vector::TransferReadOp>::OpConversionPattern;
 
  LogicalResult
  matchAndRewrite(vector::TransferReadOp readOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    // We can only deal with unmasked transfer ops with an identity permutation
    // map.
    if (!adaptor.getPermutationMap().isMinorIdentity() || adaptor.getMask())
      return failure();
    VectorType vectorTy = readOp.getVector().getType();
    if (vectorTy.getRank() < 2)
      return failure();
    // Work only on bufferized reads
    MemRefType memRefTy = dyn_cast<MemRefType>(adaptor.getBase().getType());
    if (!memRefTy)
      return failure();
    auto memRefShape = memRefTy.getShape();
    auto vecShape = vectorTy.getShape();
 
    auto newVectorTy =
        VectorType::get({std::accumulate(vecShape.begin(), vecShape.end(), 1,
                                         std::multiplies<>())},
                        vectorTy.getElementType());
    if (!hasContiguousRowMajorStrides(memRefTy))
      return failure();
 
    AffineMap layout = memRefTy.getLayout().getAffineMap();
    auto newIndices =
        collapseInnerMostDimIndices(rewriter, readOp.getLoc(), vecShape.size(),
                                    adaptor.getIndices(), memRefShape, layout);
    auto newSource = collapseInnerMostShapeDims(
        rewriter, readOp.getLoc(), vecShape.size(), adaptor.getBase());
    auto newVector = vector::TransferReadOp::create(
        rewriter, readOp.getLoc(), newVectorTy, newSource, newIndices,
        /*padding*/
        arith::getZeroConstant(rewriter, readOp.getLoc(),
                               newVectorTy.getElementType()));
 
    auto inBoundsArrayAttrOpt = adaptor.getInBounds();
    if (inBoundsArrayAttrOpt) {
      SmallVector<bool> inBounds =
          llvm::to_vector(inBoundsArrayAttrOpt.getAsValueRange<BoolAttr>());
      SmallVector<bool> newInBounds({false});
      newInBounds[0] = llvm::all_of(inBounds, [](bool v) { return v; });
      newVector.getProperties().setInBounds(
          rewriter.getBoolArrayAttr(newInBounds));
    }
 
    rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(readOp, vectorTy,
                                                     newVector);
 
    return success();
  }
};
 
// This pattern flatten multidimensional `vector.transfer_write` operations
// replacing them with a `memref.collapse_shape`, a `vector.shape_cast`, and a
// 1D `vector.transfer_write`,
struct FlattenMultDimTransferWritePattern
    : public OpConversionPattern<vector::TransferWriteOp> {
  using OpConversionPattern<vector::TransferWriteOp>::OpConversionPattern;
 
  LogicalResult
  matchAndRewrite(vector::TransferWriteOp writeOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    // We can only deal with unmasked transfer ops with an identity permutation
    // map.
    if (!adaptor.getPermutationMap().isMinorIdentity() || adaptor.getMask())
      return failure();
    VectorType vectorTy = cast<VectorType>(adaptor.getValueToStore().getType());
    if (vectorTy.getRank() < 2)
      return failure();
    // Work only on bufferized reads
    MemRefType memRefTy = dyn_cast<MemRefType>(adaptor.getBase().getType());
    if (!memRefTy)
      return failure();
    auto memRefShape = memRefTy.getShape();
    auto vecShape = vectorTy.getShape();
 
    if (!hasContiguousRowMajorStrides(memRefTy))
      return failure();
 
    auto newVectorTy =
        VectorType::get({std::accumulate(vecShape.begin(), vecShape.end(), 1,
                                         std::multiplies<>())},
                        vectorTy.getElementType());
    AffineMap layout = memRefTy.getLayout().getAffineMap();
    auto newVector =
        vector::ShapeCastOp::create(rewriter, writeOp.getLoc(), newVectorTy,
                                    adaptor.getValueToStore())
            .getResult();
    auto newIndices =
        collapseInnerMostDimIndices(rewriter, writeOp.getLoc(), vecShape.size(),
                                    adaptor.getIndices(), memRefShape, layout);
    auto newSource = collapseInnerMostShapeDims(
        rewriter, writeOp.getLoc(), vecShape.size(), adaptor.getBase());
 
    auto newOp = rewriter.replaceOpWithNewOp<vector::TransferWriteOp>(
        writeOp, newVector, newSource, newIndices);
 
    auto inBoundsArrayAttrOpt = adaptor.getInBounds();
    if (inBoundsArrayAttrOpt) {
      SmallVector<bool> inBounds =
          llvm::to_vector(inBoundsArrayAttrOpt.getAsValueRange<BoolAttr>());
      SmallVector<bool> newInBounds({false});
      newInBounds[0] = llvm::all_of(inBounds, [](bool v) { return v; });
      newOp.getProperties().setInBounds(rewriter.getBoolArrayAttr(newInBounds));
    }
 
    return success();
  }
};
 
/// Utility function to check if all provided indices are constant zero values.
/// @return success() if all indices are constant zeros, failure() otherwise
static LogicalResult isAllZeroOffsetAccess(mlir::OperandRange indices) {
  if (!llvm::all_of(indices, [](Value val) {
        IntegerAttr attr;
        if (!matchPattern(val, m_Constant(&attr)))
          return false;
        return attr.getInt() == 0;
      })) {
    return failure();
  }
  return success();
}
 
/// Utility function to convert OpFoldResult offsets from a SubView operation
/// into a vector of Values.
static SmallVector<Value> opFoldResultsToValues(PatternRewriter &rewriter,
                                                Location loc,
                                                memref::SubViewOp subViewOp) {
  OpBuilder::InsertionGuard g(rewriter);
  rewriter.setInsertionPoint(subViewOp);
  SmallVector<Value> newIndices;
  for (OpFoldResult offset : subViewOp.getMixedOffsets()) {
    Value indexVal;
    if (auto attr = dyn_cast<Attribute>(offset)) {
      indexVal = arith::ConstantIndexOp::create(
          rewriter, loc, cast<IntegerAttr>(attr).getInt());
    } else {
      indexVal = cast<Value>(offset);
    }
    newIndices.push_back(indexVal);
  }
  return newIndices;
}
 
/// Pattern to canonicalize trivial vector.transfer_read operations on subviews.
///
/// This pattern eliminates unnecessary memref.subview operations when the
/// transfer_read accesses the subview with all-zero indices. It transforms:
///
/// INPUT:
///   %subview = memref.subview %memref [offset0, offset1, ...]
///   %result = vector.transfer_read %subview[0, 0, ...]
///
/// OUTPUT:
///   %result = vector.transfer_read %memref[offset0, offset1, ...]
///
/// The pattern only matches when:
/// - The base of transfer_read is defined by a memref.subview operation
/// - All indices in the transfer_read are constant zeros
struct CanonicalizeTrivialReadAccessSubviewOpPattern
    : public OpRewritePattern<vector::TransferReadOp> {
  using OpRewritePattern<vector::TransferReadOp>::OpRewritePattern;
 
  LogicalResult matchAndRewrite(vector::TransferReadOp readOp,
                                PatternRewriter &rewriter) const override {
    // Check if the base memref comes from a subview operation
    auto subViewOp = dyn_cast_if_present<memref::SubViewOp>(
        readOp.getBase().getDefiningOp());
    if (!subViewOp)
      return failure();
 
    // Verify that all access indices are zero
    if (failed(isAllZeroOffsetAccess(readOp.getIndices())))
      return failure();
 
    // Convert subview offsets to explicit index values
    SmallVector<Value> newIndices =
        opFoldResultsToValues(rewriter, readOp.getLoc(), subViewOp);
 
    // Replace with direct access to the original memref using subview offsets
    rewriter.replaceOpWithNewOp<vector::TransferReadOp>(
        readOp, readOp.getType(), subViewOp.getSource(), newIndices,
        readOp.getPadding(), readOp.getInBoundsValues());
    return success();
  }
};
 
/// Pattern to canonicalize trivial vector.transfer_write operations on
/// subviews.
///
/// This pattern eliminates unnecessary memref.subview operations when the
/// transfer_write accesses the subview with all-zero indices. It transforms:
///
/// INPUT:
///   %subview = memref.subview %memref [offset0, offset1, ...]
///   vector.transfer_write %value, %subview[0, 0, ...]
///
/// OUTPUT:
///   vector.transfer_write %value, %memref[offset0, offset1, ...]
///
/// The pattern only matches when:
/// - The base of transfer_write is defined by a memref.subview operation
/// - All indices in the transfer_write are constant zeros
struct CanonicalizeTrivialWriteAccessSubviewOpPattern
    : public OpRewritePattern<vector::TransferWriteOp> {
  using OpRewritePattern<vector::TransferWriteOp>::OpRewritePattern;
 
  LogicalResult matchAndRewrite(vector::TransferWriteOp writeOp,
                                PatternRewriter &rewriter) const override {
    // Check if the base memref comes from a subview operation
    auto subViewOp = dyn_cast_if_present<memref::SubViewOp>(
        writeOp.getBase().getDefiningOp());
    if (!subViewOp)
      return failure();
 
    // Verify that all access indices are zero
    if (failed(isAllZeroOffsetAccess(writeOp.getIndices())))
      return failure();
 
    // Convert subview offsets to explicit index values
    SmallVector<Value> newIndices =
        opFoldResultsToValues(rewriter, writeOp.getLoc(), subViewOp);
 
    // Create new transfer_write with direct access to original memref
    vector::TransferWriteOp::create(rewriter, writeOp.getLoc(),
                                    writeOp.getVector(), subViewOp.getSource(),
                                    newIndices, writeOp.getInBoundsValues());
 
    // Remove the original transfer_write operation
    rewriter.eraseOp(writeOp);
    return success();
  }
};
 
// Carry an n-D vector through an scf.for loop as a rank-1 vector. The LLVM
// lowering turns an n-D vector into an array of rank-1 rows, so a loop that
// carries one splits it into rows and rebuilds it on every iteration instead
// of keeping it in registers. AIE2P lowers n-D elementwise ops by flattening
// them to rank 1, so the shape casts added here at the loop boundary fold
// against the ones that flattening adds. Each match rewrites one iter_arg.
struct FlattenVectorIterArgsPattern : OpRewritePattern<scf::ForOp> {
  using OpRewritePattern::OpRewritePattern;
 
  LogicalResult matchAndRewrite(scf::ForOp forOp,
                                PatternRewriter &rewriter) const override {
    for (OpOperand &init : forOp.getInitArgsMutable()) {
      auto vecType = dyn_cast<VectorType>(init.get().getType());
      if (!vecType || vecType.getRank() < 2 || vecType.isScalable())
        continue;
      Value flat = vector::ShapeCastOp::create(rewriter, forOp.getLoc(),
                                               getFlattenedVectorType(vecType),
                                               init.get());
      auto shapeCast = [](OpBuilder &builder, Location loc, Type type,
                          Value source) -> Value {
        return vector::ShapeCastOp::create(builder, loc, type, source);
      };
      rewriter.replaceOp(forOp, scf::replaceAndCastForOpIterArg(
                                    rewriter, forOp, init, flat, shapeCast));
      return success();
    }
    return failure();
  }
};
 
//============================================================================//
//============ AIE2 canonicalization conversion patterns ===============//
//============================================================================//
 
struct ConvertLeadingUnitDimInsertToReshapePattern
    : public OpRewritePattern<vector::InsertOp> {
 
  using OpRewritePattern<vector::InsertOp>::OpRewritePattern;
 
  LogicalResult matchAndRewrite(vector::InsertOp insOp,
                                PatternRewriter &rewriter) const override {
    auto insSrcTy = dyn_cast<VectorType>(insOp.getValueToStoreType());
    if (!insSrcTy)
      return failure();
 
    auto srcShape = insSrcTy.getShape();
    auto dstShape = insOp.getDestVectorType().getShape();
 
    unsigned long numLeadUnitDimDst = 0;
    while (numLeadUnitDimDst < dstShape.size() &&
           dstShape[numLeadUnitDimDst] == 1)
      numLeadUnitDimDst++;
 
    if (!numLeadUnitDimDst)
      return failure();
 
    unsigned long numLeadUnitDimSrc = 0;
    while (numLeadUnitDimSrc < srcShape.size() &&
           srcShape[numLeadUnitDimSrc] == 1)
      numLeadUnitDimSrc++;
 
    SmallVector<int64_t> nonLeadUnitDimDstShape(
        dstShape.begin() + numLeadUnitDimDst, dstShape.end());
    SmallVector<int64_t> nonLeadUnitDimSrcShape(
        srcShape.begin() + numLeadUnitDimSrc, srcShape.end());
 
    if (nonLeadUnitDimSrcShape != nonLeadUnitDimDstShape)
      return failure();
 
    rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(
        insOp, insOp.getDestVectorType(), insOp.getValueToStore());
    return success();
  }
};
 
//============================================================================//
//================ Common AIE canonicalization configuration =================//
//============================================================================//
static void
configureCommonAIECanonicalizeLegalizations(ConversionTarget &target) {
  target.addLegalDialect<arith::ArithDialect, affine::AffineDialect,
                         memref::MemRefDialect, vector::VectorDialect,
                         ub::UBDialect>();
}
 
static void
populateCommonAIECanonicalizeConversionPatterns(RewritePatternSet &patterns) {
  patterns.add<ConvertSplatTransferReadToBroadcastPattern>(
      patterns.getContext());
}
 
//============================================================================//
//============== AIEv1-specific canonicalization configuration ===============//
//============================================================================//
 
static void configureAIEv1CanonicalizeLegalizations(ConversionTarget &target) {
  target.addDynamicallyLegalOp<vector::TransferReadOp>(
      [](vector::TransferReadOp op) {
        return !op.getPermutationMap().isConstant() &&
               getTransferReadAlignmentOffset(op, op.getVectorType(), 128)
                       .value_or(0) == 0;
      });
}
 
static void
populateAIEv1CanonicalizeConversionPatterns(RewritePatternSet &patterns) {
  patterns.add<SplitUnalignedTransferReadPattern>(patterns.getContext(), 512,
                                                  128);
}
 
//============================================================================//
//============== AIE2-specific canonicalization configuration ===============//
//============================================================================//
 
static void configureAIE2CanonicalizeLegalizations(ConversionTarget &target) {
  target.addDynamicallyLegalOp<vector::TransferReadOp>(
      [](vector::TransferReadOp op) {
        return !op.getPermutationMap().isConstant() &&
               getTransferReadAlignmentOffset(op, op.getVectorType(), 256)
                       .value_or(0) == 0 &&
               op.getVector().getType().getRank() < 2;
      });
  target.addDynamicallyLegalOp<vector::TransferWriteOp>(
      [](vector::TransferWriteOp op) {
        return cast<VectorType>(op.getVector().getType()).getRank() < 2;
      });
}
 
static void
populateAIE2CanonicalizeConversionPatterns(RewritePatternSet &patterns) {
  patterns.add<SplitUnalignedTransferReadPattern>(patterns.getContext(), 1024,
                                                  256);
  patterns.add<FlattenMultDimTransferReadPattern,
               FlattenMultDimTransferWritePattern>(patterns.getContext());
}
 
//============================================================================//
//=================== Common AIE Canonicalization Passes =====================//
//============================================================================//
 
//===----------------------------------------------------------------------===//
// BF16 Emulation: Emulate f32 vector arithmetic using bf16 operations.
//===----------------------------------------------------------------------===//
 
// Smart truncation helper: if the value was produced by arith.extf from bf16,
// reuse the bf16 source directly to avoid redundant extf->truncf chains.
static Value smartTruncF32ToBF16(PatternRewriter &rewriter, Location loc,
                                 Value val, Type bf16Type) {
  if (auto extfOp = val.getDefiningOp<arith::ExtFOp>()) {
    if (extfOp.getIn().getType() == bf16Type)
      return extfOp.getIn();
  }
  return arith::TruncFOp::create(rewriter, loc, bf16Type, val);
}
 
// `v` rounded to the nearest bf16, still in its own semantics.
static APFloat roundedToBF16(APFloat v) {
  bool losesInfo = false;
  const llvm::fltSemantics &sem = v.getSemantics();
  v.convert(APFloat::BFloat(), APFloat::rmNearestTiesToEven, &losesInfo);
  v.convert(sem, APFloat::rmNearestTiesToEven, &losesInfo);
  return v;
}
 
// True if `val` is an f32 vector whose elements are all bf16 values.
static bool holdsBF16(Value val) {
  if (auto extfOp = val.getDefiningOp<arith::ExtFOp>())
    return getElementTypeOrSelf(extfOp.getIn().getType()).isBF16();
  DenseFPElementsAttr dense;
  return matchPattern(val, m_Constant(&dense)) &&
         llvm::all_of(dense.getValues<APFloat>(), [](const APFloat &v) {
           return roundedToBF16(v).bitwiseIsEqual(v);
         });
}
 
// `val` rounded to bf16, still as f32. Constants are rounded here.
static Value roundToBF16(PatternRewriter &rewriter, Location loc, Value val) {
  if (holdsBF16(val))
    return val;
  auto vecType = cast<VectorType>(val.getType());
  DenseFPElementsAttr dense;
  if (matchPattern(val, m_Constant(&dense))) {
    SmallVector<APFloat> values;
    for (const APFloat &v : dense.getValues<APFloat>())
      values.push_back(roundedToBF16(v));
    return arith::ConstantOp::create(rewriter, loc,
                                     DenseElementsAttr::get(vecType, values));
  }
  auto bf16VecType =
      VectorType::get(vecType.getShape(), rewriter.getBF16Type());
  Value bf16 = smartTruncF32ToBF16(rewriter, loc, val, bf16VecType);
  return arith::ExtFOp::create(rewriter, loc, vecType, bf16,
                               /*fastmath=*/nullptr);
}
 
static bool isF32Vector(Type type) {
  auto vecType = dyn_cast<VectorType>(type);
  return vecType && vecType.getElementType().isF32();
}
 
/// Pattern to drop `bf16 -> f32 -> bf16` round trips.
///
/// Arith used to fold these but no longer does: widening a signaling NaN quiets
/// it, so the original bit pattern cannot be recovered. This pass has always
/// made that trade anyway -- `smartTruncF32ToBF16` collapses the same pair on
/// sight -- so fold here to keep the pipeline's pre-existing behavior rather
/// than leaving the round trips for the backend.
struct FoldBF16RoundTripPattern : public OpRewritePattern<arith::TruncFOp> {
  using OpRewritePattern::OpRewritePattern;
 
  LogicalResult matchAndRewrite(arith::TruncFOp op,
                                PatternRewriter &rewriter) const override {
    auto resultType = dyn_cast<VectorType>(op.getType());
    auto intermediateType = dyn_cast<VectorType>(op.getIn().getType());
    if (!resultType || !intermediateType ||
        !resultType.getElementType().isBF16() ||
        !intermediateType.getElementType().isF32())
      return failure();
 
    auto extfOp = op.getIn().getDefiningOp<arith::ExtFOp>();
    if (!extfOp || extfOp.getIn().getType() != resultType)
      return failure();
 
    rewriter.replaceOp(op, extfOp.getIn());
    return success();
  }
};
 
/// Pattern to emulate f32 binary vector arithmetic ops in bf16.
/// For an op like: %r = arith.addf %a, %b : vector<16xf32>
/// Produces:
///   %a_bf16 = arith.truncf %a : vector<16xf32> to vector<16xbf16>
///   %b_bf16 = arith.truncf %b : vector<16xf32> to vector<16xbf16>
///   %r_bf16 = arith.addf %a_bf16, %b_bf16 : vector<16xbf16>
///   %r = arith.extf %r_bf16 : vector<16xbf16> to vector<16xf32>
template <typename OpTy>
struct EmulateBinaryF32InBF16Pattern : public OpRewritePattern<OpTy> {
  using OpRewritePattern<OpTy>::OpRewritePattern;
 
  LogicalResult matchAndRewrite(OpTy op,
                                PatternRewriter &rewriter) const override {
    auto resultType = dyn_cast<VectorType>(op.getType());
    if (!resultType || !resultType.getElementType().isF32())
      return failure();
 
    Location loc = op.getLoc();
    auto bf16VecType =
        VectorType::get(resultType.getShape(), rewriter.getBF16Type());
 
    Value lhsBF16 =
        smartTruncF32ToBF16(rewriter, loc, op.getLhs(), bf16VecType);
    Value rhsBF16 =
        smartTruncF32ToBF16(rewriter, loc, op.getRhs(), bf16VecType);
 
    Value newResult =
        OpTy::create(rewriter, loc, bf16VecType, lhsBF16, rhsBF16);
    auto extOp = arith::ExtFOp::create(rewriter, loc, resultType, newResult,
                                       /*fastmath=*/nullptr);
    rewriter.replaceOp(op, extOp);
    return success();
  }
};
 
/// Pattern to emulate f32 comparison ops in bf16.
/// Result type stays vector<Nxi1>, only operands are truncated.
struct EmulateCmpFF32InBF16Pattern : public OpRewritePattern<arith::CmpFOp> {
  using OpRewritePattern::OpRewritePattern;
 
  LogicalResult matchAndRewrite(arith::CmpFOp op,
                                PatternRewriter &rewriter) const override {
    auto lhsType = dyn_cast<VectorType>(op.getLhs().getType());
    if (!lhsType || !lhsType.getElementType().isF32())
      return failure();
 
    Location loc = op.getLoc();
    auto bf16VecType =
        VectorType::get(lhsType.getShape(), rewriter.getBF16Type());
 
    Value lhsBF16 =
        smartTruncF32ToBF16(rewriter, loc, op.getLhs(), bf16VecType);
    Value rhsBF16 =
        smartTruncF32ToBF16(rewriter, loc, op.getRhs(), bf16VecType);
 
    rewriter.replaceOpWithNewOp<arith::CmpFOp>(op, op.getPredicate(), lhsBF16,
                                               rhsBF16);
    return success();
  }
};
 
/// Pattern to emulate f32 select ops in bf16.
/// Condition stays vector<Nxi1>, true/false values are truncated.
struct EmulateSelectF32InBF16Pattern
    : public OpRewritePattern<arith::SelectOp> {
  using OpRewritePattern::OpRewritePattern;
 
  LogicalResult matchAndRewrite(arith::SelectOp op,
                                PatternRewriter &rewriter) const override {
    auto resultType = dyn_cast<VectorType>(op.getType());
    if (!resultType || !resultType.getElementType().isF32())
      return failure();
 
    Location loc = op.getLoc();
    auto bf16VecType =
        VectorType::get(resultType.getShape(), rewriter.getBF16Type());
 
    Value trueValBF16 =
        smartTruncF32ToBF16(rewriter, loc, op.getTrueValue(), bf16VecType);
    Value falseValBF16 =
        smartTruncF32ToBF16(rewriter, loc, op.getFalseValue(), bf16VecType);
 
    Value newResult = arith::SelectOp::create(rewriter, loc, op.getCondition(),
                                              trueValBF16, falseValBF16);
    auto extOp = arith::ExtFOp::create(rewriter, loc, resultType, newResult,
                                       /*fastmath=*/nullptr);
    rewriter.replaceOp(op, extOp);
    return success();
  }
};
 
/// Pattern to emulate an f32 multiply with bf16 inputs. The product of two
/// bf16 values fits in f32 except at the ends of the exponent range, so only
/// the inputs are rounded:
///   %r = arith.mulf %a, %b : vector<16xf32>
/// becomes
///   %r = arith.mulf (extf (truncf %a)), (extf (truncf %b)) : vector<16xf32>
struct EmulateMulFInBF16Pattern : public OpRewritePattern<arith::MulFOp> {
  using OpRewritePattern::OpRewritePattern;
 
  LogicalResult matchAndRewrite(arith::MulFOp op,
                                PatternRewriter &rewriter) const override {
    if (!isF32Vector(op.getType()) ||
        (holdsBF16(op.getLhs()) && holdsBF16(op.getRhs())))
      return failure();
    Location loc = op.getLoc();
    Value lhs = roundToBF16(rewriter, loc, op.getLhs());
    Value rhs = roundToBF16(rewriter, loc, op.getRhs());
    rewriter.replaceOpWithNewOp<arith::MulFOp>(op, lhs, rhs,
                                               op.getFastmathAttr());
    return success();
  }
};
 
/// Pattern to fold an f32 multiply into the add that is its only user, as a
/// vector.fma with bf16 inputs. Only when both ops allow contraction: the
/// fused result can differ where the product underflows or overflows.
struct FuseMulAddFInBF16Pattern : public OpRewritePattern<arith::AddFOp> {
  using OpRewritePattern::OpRewritePattern;
 
  LogicalResult matchAndRewrite(arith::AddFOp op,
                                PatternRewriter &rewriter) const override {
    auto contracts = [](arith::FastMathFlagsAttr flags) {
      return flags && bitEnumContainsAll(flags.getValue(),
                                         arith::FastMathFlags::contract);
    };
    if (!isF32Vector(op.getType()) || !contracts(op.getFastmathAttr()))
      return failure();
    for (auto [mulOperand, acc] : {std::pair(op.getLhs(), op.getRhs()),
                                   std::pair(op.getRhs(), op.getLhs())}) {
      auto mulOp = mulOperand.getDefiningOp<arith::MulFOp>();
      if (!mulOp || !mulOp->hasOneUse() || !contracts(mulOp.getFastmathAttr()))
        continue;
      Location loc = op.getLoc();
      Value lhs = roundToBF16(rewriter, loc, mulOp.getLhs());
      Value rhs = roundToBF16(rewriter, loc, mulOp.getRhs());
      rewriter.replaceOpWithNewOp<vector::FMAOp>(op, lhs, rhs, acc);
      rewriter.eraseOp(mulOp);
      return success();
    }
    return failure();
  }
};
 
/// Pattern to emulate f32 vector.fma with bf16 inputs. The accumulator stays
/// f32.
struct EmulateFMAF32InBF16Pattern : public OpRewritePattern<vector::FMAOp> {
  using OpRewritePattern::OpRewritePattern;
 
  LogicalResult matchAndRewrite(vector::FMAOp op,
                                PatternRewriter &rewriter) const override {
    if (!isF32Vector(op.getType()) ||
        (holdsBF16(op.getLhs()) && holdsBF16(op.getRhs())))
      return failure();
    Location loc = op.getLoc();
    Value lhs = roundToBF16(rewriter, loc, op.getLhs());
    Value rhs = roundToBF16(rewriter, loc, op.getRhs());
    rewriter.replaceOpWithNewOp<vector::FMAOp>(op, lhs, rhs, op.getAcc());
    return success();
  }
};
 
/// Pattern to emulate f32 unary vector ops in bf16.
template <typename OpTy>
struct EmulateUnaryF32InBF16Pattern : public OpRewritePattern<OpTy> {
  using OpRewritePattern<OpTy>::OpRewritePattern;
 
  LogicalResult matchAndRewrite(OpTy op,
                                PatternRewriter &rewriter) const override {
    auto resultType = dyn_cast<VectorType>(op.getType());
    if (!resultType || !resultType.getElementType().isF32())
      return failure();
 
    Location loc = op.getLoc();
    auto bf16VecType =
        VectorType::get(resultType.getShape(), rewriter.getBF16Type());
 
    Value inputBF16 =
        smartTruncF32ToBF16(rewriter, loc, op->getOperand(0), bf16VecType);
 
    Value newResult = OpTy::create(rewriter, loc, bf16VecType, inputBF16);
    auto extOp = arith::ExtFOp::create(rewriter, loc, resultType, newResult,
                                       /*fastmath=*/nullptr);
    rewriter.replaceOp(op, extOp);
    return success();
  }
};
 
struct BF16EmulationPass
    : public PassWrapper<BF16EmulationPass, OperationPass<>> {
 
  void runOnOperation() override {
    auto *op = getOperation();
    MLIRContext *context = &getContext();
    RewritePatternSet patterns(context);
 
    // Multiplies take bf16 inputs; sums stay f32, which the accumulator
    // adds natively.
    patterns.add<FuseMulAddFInBF16Pattern>(context, /*benefit=*/2);
    patterns.add<EmulateMulFInBF16Pattern, EmulateFMAF32InBF16Pattern>(context);
    patterns.add<EmulateBinaryF32InBF16Pattern<arith::MaximumFOp>,
                 EmulateBinaryF32InBF16Pattern<arith::MinimumFOp>>(context);
 
    // Note: arith.divf is NOT demoted because bf16 vector divf is unsupported
    // on all AIE targets (Peano does not legalize G_FDIV on <16 x s16>).
 
    // Special-case ops (excluding ReductionOp — its scalar accumulator
    // and result lower to fp_to_bf16/bf16_to_fp which older Peano versions
    // cannot select on AIE2P; reductions are intentionally left in f32
    // to avoid these scalar bf16 conversions)
    patterns.add<EmulateCmpFF32InBF16Pattern, EmulateSelectF32InBF16Pattern>(
        context);
 
    // Unary ops
    patterns.add<EmulateUnaryF32InBF16Pattern<arith::NegFOp>,
                 EmulateUnaryF32InBF16Pattern<math::TanhOp>>(context);
 
    // Clean up the round trips the demotion itself introduces.
    patterns.add<FoldBF16RoundTripPattern>(context);
 
    (void)applyPatternsGreedily(op, std::move(patterns));
  }
};
 
std::unique_ptr<::mlir::Pass> xilinx::aievec::createBF16EmulationPass() {
  return std::make_unique<BF16EmulationPass>();
}
 
// A broadcast of a vector that only adds leading dims, as the vector repeated
// by a 1-D shuffle (mask [i % n]). The AIE2P shuffle lowering then sees the
// repetition, alone or together with a shuffle that reads it, as one gather.
struct BroadcastLeadingDimsToRepeatPattern
    : public OpRewritePattern<vector::BroadcastOp> {
  using OpRewritePattern::OpRewritePattern;
 
  LogicalResult matchAndRewrite(vector::BroadcastOp op,
                                PatternRewriter &rewriter) const override {
    auto srcType = dyn_cast<VectorType>(op.getSourceType());
    VectorType resType = op.getResultVectorType();
    if (!srcType || srcType.getRank() != 1 || resType.getRank() < 2 ||
        resType.getShape().back() != srcType.getNumElements() ||
        !srcType.getElementType().isIntOrFloat() || resType.isScalable())
      return failure();
    int64_t n = srcType.getNumElements();
    SmallVector<int64_t> mask;
    for (int64_t i = 0; i < resType.getNumElements(); ++i)
      mask.push_back(i % n);
    Value src = op.getSource();
    Value flat =
        vector::ShuffleOp::create(rewriter, op.getLoc(), src, src, mask);
    rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(op, resType, flat);
    return success();
  }
};
 
struct VectorBroadcastLoweringPass
    : public PassWrapper<VectorBroadcastLoweringPass, OperationPass<>> {
  VectorBroadcastLoweringPass(bool aie2p = false) : aie2p(aie2p) {}
 
  void runOnOperation() override {
    auto *op = getOperation();
    MLIRContext *context = &getContext();
    RewritePatternSet patterns(context);
    populateVectorBroadcastLoweringPatterns(patterns);
    patterns.add<ConvertLeadingUnitDimInsertToReshapePattern>(
        patterns.getContext());
    // Ahead of the generic lowering, which builds the result a row at a time.
    if (aie2p)
      patterns.add<BroadcastLeadingDimsToRepeatPattern>(patterns.getContext(),
                                                        /*benefit=*/2);
 
    (void)applyPatternsGreedily(op, std::move(patterns));
  }
 
private:
  bool aie2p;
};
 
static std::unique_ptr<::mlir::Pass>
createVectorBroadcastLoweringPass(bool aie2p = false) {
  return std::make_unique<VectorBroadcastLoweringPass>(aie2p);
}
 
// This pass converts standard vector ops into a subset of `Vector` ops more
// amenable to being converted to `AIEVec`. So far, this process consists of
// two steps:
//    1) Replace splat transfer reads with contiguous transfer reads followed
//       by `extract` + `splat` operations.
//    2) Split unaligned transfer reads into a wider aligned transfer read
//       followed by a `vector.extract_strided_slice` operation.
struct CanonicalizeVectorForAIEVecPass
    : public PassWrapper<CanonicalizeVectorForAIEVecPass, OperationPass<>> {
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(CanonicalizeVectorForAIEVecPass)
 
  CanonicalizeVectorForAIEVecPass() = default;
  CanonicalizeVectorForAIEVecPass(const CanonicalizeVectorForAIEVecPass &pass)
      : PassWrapper(pass) {}
 
  CanonicalizeVectorForAIEVecPass(
      const CanonicalizeVectorForAIEVecOptions &options)
      : CanonicalizeVectorForAIEVecPass() {
    aieTarget = options.aieTarget;
  }
 
  // In case we want to register this pass as a standalone pass for test
  // purposes.
  StringRef getArgument() const final {
    return "test-canonicalize-vector-for-aievec";
  }
 
  StringRef getDescription() const final {
    return "Canonicalize vector operations for AIEVec conversion";
  }
 
  void getDependentDialects(DialectRegistry &registry) const override {
    registry
        .insert<arith::ArithDialect, memref::MemRefDialect,
                vector::VectorDialect, affine::AffineDialect, ub::UBDialect>();
  }
 
  Option<std::string> aieTarget{
      *this, "aie-target",
      llvm::cl::desc(
          "Select AIE version: \"aie\", \"aie2\", or \"aie2p\". This will "
          "determine the vector size and available operations."),
      llvm::cl::init("aie")};
 
  void runOnOperation() override {
    auto *op = getOperation();
    MLIRContext *context = &getContext();
    RewritePatternSet patterns(context);
    ConversionTarget target(*context);
 
    AIEArch aieVersion = decodeAIETarget(aieTarget);
    if (aieVersion == AIEArch::UNKNOWN) {
      op->emitError() << "unknown AIE target '" << aieTarget << "'";
      signalPassFailure();
      return;
    }
 
    populateCommonAIECanonicalizeConversionPatterns(patterns);
    configureCommonAIECanonicalizeLegalizations(target);
    if (aieVersion == AIEArch::AIE) {
      populateAIEv1CanonicalizeConversionPatterns(patterns);
      configureAIEv1CanonicalizeLegalizations(target);
    } else {
      populateAIE2CanonicalizeConversionPatterns(patterns);
      configureAIE2CanonicalizeLegalizations(target);
    }
 
    {
      RewritePatternSet patterns(context);
      patterns.add<CanonicalizeTrivialReadAccessSubviewOpPattern,
                   CanonicalizeTrivialWriteAccessSubviewOpPattern>(context);
      if (aieTarget == "aie2p")
        patterns.add<FlattenVectorIterArgsPattern>(context);
      (void)applyPatternsGreedily(op, std::move(patterns));
    }
 
    if (failed(applyPartialConversion(op, target, std::move(patterns)))) {
      signalPassFailure();
    }
  }
};
 
static std::unique_ptr<::mlir::Pass> createCanonicalizeVectorForAIEVecPass(
    const CanonicalizeVectorForAIEVecOptions &options) {
  return std::make_unique<CanonicalizeVectorForAIEVecPass>(options);
}
 
struct ReorderOperationsPass
    : public PassWrapper<ReorderOperationsPass, OperationPass<>> {
 
  void runOnOperation() override {
    auto *op = getOperation();
    MLIRContext *context = &getContext();
    RewritePatternSet patterns(context);
 
    patterns.add<SwapUnaryOpsPattern<arith::ExtSIOp, vector::BroadcastOp>>(
        patterns.getContext(),
        [](arith::ExtSIOp extOp, vector::BroadcastOp bcastOp) -> Type {
          Type extInElemTy = extOp.getIn().getType();
          auto extInVecTy = dyn_cast<VectorType>(extInElemTy);
          if (extInVecTy)
            extInElemTy = extInVecTy.getElementType();
          return VectorType::get(bcastOp.getResultVectorType().getShape(),
                                 extInElemTy);
        });
 
    (void)applyPatternsGreedily(op, std::move(patterns));
  }
};
 
static std::unique_ptr<::mlir::Pass> createReorderOperationsPass() {
  return std::make_unique<ReorderOperationsPass>();
}
 
//============================================================================//
//=============== Main Vector2Vector Pipeline Configuration ==================//
//============================================================================//
 
void xilinx::aievec::buildCanonicalizeVectorForAIEVec(
    OpPassManager &pm, const CanonicalizeVectorForAIEVecOptions &options) {
  // Add `Vector` code canonicalization passes
  // TODO: Add passes to unroll vector with unsupported types
  // TODO: Add passes to split vectors that won't fit in registers
 
  // If bf16-emulation is enabled, demote f32 vector arithmetic to bf16 first.
  if (options.enableBF16Emulation)
    pm.addPass(createBF16EmulationPass());
 
  pm.addPass(createReorderOperationsPass());
  pm.addPass(createCopyRemovalPass());
  pm.addPass(createVectorBroadcastLoweringPass(options.aieTarget == "aie2p"));
  pm.addPass(createCanonicalizeVectorForAIEVecPass(options));
}