AIEVecToLLVM.cpp

Go to the documentation of this file.

//===- AIEVecToLLVM.cpp - AIEVec to LLVM dialect conversion ---------------===//
//
// Copyright (C) 2022 Xilinx, Inc.
// Copyright (C) 2024 Advanced Micro Devices, Inc.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
 
#include "../PassDetail.h"
 
#include "aie/Conversion/AIEVecToLLVM/AIEVecToLLVM.h"
#include "aie/Dialect/AIEVec/AIE1/IR/AIEVecAIE1Ops.h"
#include "aie/Dialect/AIEVec/AIEVecUtils.h"
#include "aie/Dialect/AIEVec/IR/AIEVecMatMul.h"
#include "aie/Dialect/AIEVec/IR/AIEVecOps.h"
#include "aie/Dialect/AIEVec/IR/AIEVecShuffleModes.h"
#include "aie/Dialect/AIEVec/Utils/Utils.h"
#include "aie/Dialect/XLLVM/XLLVMDialect.h"
#include "mlir/Conversion/LLVMCommon/ConversionTarget.h"
#include "mlir/Conversion/LLVMCommon/Pattern.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/Math/IR/Math.h"
#include "mlir/Dialect/UB/IR/UBOps.h"
#include "mlir/IR/TypeUtilities.h"
#include <numeric>
#include <sstream>
 
namespace xilinx {
using namespace mlir; // For LLVM::LLVMDialect in generated getDependentDialects
#define GEN_PASS_DEF_CONVERTAIEVECTOLLVM
#include "aie/Conversion/Passes.h.inc"
} // namespace xilinx
 
using namespace mlir;
 
namespace xilinx::aievec {
 
inline static Value bitcastValueToType(OpBuilder &builder, Location loc,
                                       Value val, Type dstTy) {
  return LLVM::BitcastOp::create(builder, loc, dstTy, val).getResult();
}
 
// This function emits the instructions required to widen a 128b input vector
// into a 512b encoded as a vector<16xi32>. It first bitcasts it to a
// vector<4xi32> to respect the intrinsic signature.
inline static Value widen128bVectorValueTo512b(OpBuilder &builder, Location loc,
                                               Value val) {
  return xllvm::VectorSetI512I128IntrOp::create(
             builder, loc, VectorType::get({16}, builder.getI32Type()),
             bitcastValueToType(builder, loc, val,
                                VectorType::get({4}, builder.getI32Type())))
      .getResult();
}
 
// This function emits the instructions required to widen a 256b input vector
// into a 512b encoded as a vector<16xi32>. It first bitcasts it to a
// vector<8xi32> to respect the intrinsic signature. It will also materialize
// a constant 0, used as an insertion index.
inline static Value widen256bVectorValueTo512b(OpBuilder &builder, Location loc,
                                               Value val) {
  auto cst0 =
      LLVM::ConstantOp::create(builder, loc, builder.getI32Type(), (int32_t)0);
  return xllvm::VectorSetI512I256IntrOp::create(
             builder, loc, VectorType::get({16}, builder.getI32Type()),
             bitcastValueToType(builder, loc, val,
                                VectorType::get({8}, builder.getI32Type())),
             cst0)
      .getResult();
}
 
// This function emits the sequence of operations that forces a value into a
// specific type. This may include widening vectors to match a specific bit
// length.
static Value forceCastValueToType(OpBuilder &builder, Location loc, Value val,
                                  Type type) {
  auto valTy = val.getType();
  if (valTy == type)
    return val;
  auto srcVecTy = dyn_cast<VectorType>(valTy);
  auto dstVecTy = dyn_cast<VectorType>(type);
 
  if (srcVecTy) {
    assert(dstVecTy && "vector values cannot be forced into a non-vector type");
 
    // Flatten source vector if it's not rank-1
    auto flatSrcVecTy = getFlattenedVectorType(srcVecTy);
    if (srcVecTy != flatSrcVecTy)
      val = vector::ShapeCastOp::create(builder, loc, flatSrcVecTy, val);
 
    // Flatten destination type if it's not rank-1
    auto flatDstVecTy = getFlattenedVectorType(dstVecTy);
 
    int64_t dstVecLength =
        flatDstVecTy.getElementTypeBitWidth() * flatDstVecTy.getShape()[0];
    int64_t srcVecLength =
        flatSrcVecTy.getElementTypeBitWidth() * flatSrcVecTy.getShape()[0];
    if (srcVecLength != dstVecLength) {
      assert(srcVecLength < dstVecLength &&
             "only widening forced casts are supported");
      assert(dstVecLength == 512 &&
             (srcVecLength == 128 || srcVecLength == 256) &&
             "only 128b to 512b and 256b to 512b forced casts are supported");
      if (srcVecLength == 128)
        val = widen128bVectorValueTo512b(builder, loc, val);
      else
        val = widen256bVectorValueTo512b(builder, loc, val);
    }
 
    // Bitcast to flat destination type (bitcast only supports flat vectors)
    val = bitcastValueToType(builder, loc, val, flatDstVecTy);
 
    // Reshape back to original destination shape if needed
    if (flatDstVecTy != dstVecTy)
      val = vector::ShapeCastOp::create(builder, loc, dstVecTy, val);
 
    return val;
  }
 
  // Non-vector types can be bitcast directly
  assert(!dstVecTy && "cannot force cast scalar to vector type");
  return bitcastValueToType(builder, loc, val, type);
}
 
// This function emits the sequence of operations that forces a range of values
// to match the signature specified by the TypeRange. It can be used to convert
// the parameters of an op being converted to the types accepted by an
// intrinsic with a fixed signature that treats its inputs as "bags of bits".
static SmallVector<Value> forceCastOperandsToSignature(OpBuilder &builder,
                                                       Location loc,
                                                       ValueRange operands,
                                                       TypeRange signature) {
  return llvm::to_vector(llvm::map_range(
      llvm::zip_equal(operands, signature), [&](auto &&vt) -> Value {
        return forceCastValueToType(builder, loc, std::get<0>(vt),
                                    std::get<1>(vt));
      }));
}
 
// Utility function to get or create a noinline scalar helper function.
// This is used to create optimization barriers that prevent LLVM from
// re-vectorizing unrolled scalar operations.
//
// Parameters:
//   - module: The parent module to insert the function into
//   - rewriter: The pattern rewriter
//   - opName: Base name of the operation (e.g., "fdiv", "addf", "mulf")
//   - device: Target device ("aie2", "aie2p", etc.)
//   - argTypes: Input argument types
//   - resultType: Return type
//   - bodyBuilder: Lambda that builds the function body given (builder, loc,
//   args)
//
// Returns: The helper function (created or existing)
//
// Function naming convention: __<device>_scalar_<opName>
// Example: __aie2p_scalar_fdiv, __aie2_scalar_addf
static LLVM::LLVMFuncOp getOrCreateScalarHelperFunc(
    ModuleOp module, OpBuilder &rewriter, StringRef opName, StringRef device,
    TypeRange argTypes, Type resultType,
    const std::function<void(OpBuilder &, Location, ValueRange)> &bodyBuilder) {
 
  // Build function name: __<device>_scalar_<opName>
  std::string funcName = "__" + device.str() + "_scalar_" + opName.str();
 
  // Check if function already exists
  auto helperFunc = module.lookupSymbol<LLVM::LLVMFuncOp>(funcName);
  if (helperFunc)
    return helperFunc;
 
  // Create new function
  OpBuilder::InsertionGuard guard(rewriter);
  rewriter.setInsertionPointToStart(module.getBody());
 
  // Convert TypeRange to SmallVector<Type> for LLVM::LLVMFunctionType::get
  SmallVector<Type> argTypesVec(argTypes.begin(), argTypes.end());
 
  helperFunc = LLVM::LLVMFuncOp::create(
      rewriter, rewriter.getUnknownLoc(), funcName,
      LLVM::LLVMFunctionType::get(resultType, argTypesVec));
 
  // Mark as noinline to act as optimization barrier
  helperFunc->setAttr("passthrough", rewriter.getArrayAttr(
                                         {rewriter.getStringAttr("noinline")}));
 
  // Add function body
  auto *entryBlock = helperFunc.addEntryBlock(rewriter);
  OpBuilder::InsertionGuard bodyGuard(rewriter);
  rewriter.setInsertionPointToStart(entryBlock);
 
  // Collect function arguments
  SmallVector<Value> args;
  for (unsigned i = 0; i < argTypes.size(); ++i)
    args.push_back(entryBlock->getArgument(i));
 
  // Call the body builder with the function arguments
  bodyBuilder(rewriter, rewriter.getUnknownLoc(), args);
 
  return helperFunc;
}
 
struct BufferParams {
  uint32_t start;
  uint32_t offsets;
  uint32_t offsets_hi;
  uint32_t step;
  uint32_t square;
};
 
// sgn_x: Sign mask of matrix X. If it is one matrix X is interpreted as
// signed, else it treated as unsigned.
// sgn_y: Sign mask of matrix Y. If it is one matrix Y is interpreted as
// signed, else it treated as unsigned.
// amode/bmode/variant: config acc width, mul precision, and mul mode
// zero_acc: Zeroing of acc1. If it is one then acc1 is zeroed.
// shift16: Shift mask of acc1. If a bit is set the <<16 operation will be
// executed on acc1.
// sub_mul: Negation mask of the matrix multiplication result. If it is
// one the result of the operation will be negated.
// sub_acc1: Negation mask of acc1. If it is one acc1 will be negated.
// sub_acc2: Negation mask of acc2. If it is one acc2 will be negated.
// sub_mask: Negation mask of complex multiplications. Negates a term of a
// complex multiplication.
static inline int aiev2_vmac_compute_control(int sgn_x, int sgn_y, int amode,
                                             int bmode, int variant,
                                             int zero_acc, int shift16,
                                             int sub_mul, int sub_acc1,
                                             int sub_acc2, int sub_mask) {
  return ((unsigned)sub_mask << 16) | ((unsigned)shift16 << 10) |
         ((unsigned)sub_mul << 11) | ((unsigned)sub_acc1 << 12) |
         ((unsigned)sub_acc2 << 13) | ((unsigned)amode << 1) |
         ((unsigned)bmode << 3) | ((unsigned)variant << 5) |
         (((unsigned)sgn_x << 9) | ((unsigned)sgn_y << 8)) |
         ((unsigned)zero_acc << 0);
}
 
static std::string getVectorTypeString(VectorType type, bool abbrev = false,
                                       bool acc = false) {
  std::stringstream ss;
  auto size = getVectorLaneSize(type);
  ss << "v" << size;
  if (auto intType = dyn_cast<IntegerType>(type.getElementType())) {
    ss << (acc ? "acc" : abbrev ? "i" : "int") << intType.getWidth();
  } else if (dyn_cast<FloatType>(type.getElementType())) {
    ss << (abbrev ? "f" : "float");
  }
  return ss.str();
}
 
static std::string getMulOrFMAIntrinsicName(Operation *op) {
  std::string baseName;
  Value lhs, result;
  if (auto mulOp = dyn_cast<aievec::aie1::MulOp>(op)) {
    baseName = "mul";
    lhs = mulOp.getLhs();
    result = mulOp.getResult();
  } else if (auto fmaOp = dyn_cast<aievec::aie1::FMAOp>(op)) {
    baseName = "mac";
    lhs = fmaOp.getLhs();
    result = fmaOp.getResult();
  }
  VectorType resultType = cast<VectorType>(result.getType());
  int resultSize = getVectorLaneSize(resultType);
  std::stringstream ss;
  ss << "llvm.aie.";
  if (dyn_cast<IntegerType>(resultType.getElementType())) {
    ss << baseName;
    ss << resultSize << "."
       << getVectorTypeString(cast<VectorType>(lhs.getType()));
  } else if (dyn_cast<FloatType>(resultType.getElementType())) {
    ss << "vfp" << baseName;
  }
  return ss.str();
}
 
// Squashes the easy-to-read 16-bit square encoding into
// the 8-bit encoding the configuration register uses
static uint32_t encodeSquare(uint32_t square) {
  uint32_t out = 0;
  out |= ((square >> 0) & 0x3) << 0;
  out |= ((square >> 4) & 0x3) << 2;
  out |= ((square >> 8) & 0x3) << 4;
  out |= ((square >> 12) & 0x3) << 6;
  return out & 0xFF;
}
 
// Encode the configuration register with buffer parameters and options
// TODO: struct to handle this?
static void encodeConf(uint32_t conf[2], const BufferParams &x,
                       const BufferParams &z, bool sub) {
  conf[0] |= ((x.step & 0x3F) << 0) | ((z.step & 0x3F) << 8);
  conf[1] |= (encodeSquare(x.square) << 0) | (encodeSquare(z.square) << 8);
  conf[1] |= sub << 17;
}
 
class AddOpConversion
    : public mlir::ConvertOpToLLVMPattern<aievec::aie1::AddOp> {
public:
  using ConvertOpToLLVMPattern<aievec::aie1::AddOp>::ConvertOpToLLVMPattern;
 
  LogicalResult
  matchAndRewrite(aievec::aie1::AddOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    op.emitWarning() << "aie.add conversion is not implemented\n";
    return failure();
  }
};
 
// AIE2 version of AddElemOp conversion
class AddElemOpAIE2Conversion
    : public mlir::ConvertOpToLLVMPattern<aievec::AddElemOp> {
public:
  using ConvertOpToLLVMPattern<aievec::AddElemOp>::ConvertOpToLLVMPattern;
 
  struct DecodedAddElemOp {
    enum class Kind { FP32_FP32_FP32_16x1x1x1, UNSUPPORTED };
    Kind kind;
    int conf;
  };
 
  static DecodedAddElemOp decodeAddElemOp(OpAdaptor op) {
    auto lhs = op.getLhs();
    auto lhsVecTy = cast<VectorType>(lhs.getType());
    auto lhsScaTy = lhsVecTy.getElementType();
    unsigned lhsBitWidth = lhsScaTy.getIntOrFloatBitWidth();
 
    // Integer types
    if (llvm::isa<IntegerType>(lhsScaTy)) {
      return {DecodedAddElemOp::Kind::UNSUPPORTED, -1};
    } // Float types
    if (lhsBitWidth == 32) {
      // FP32 add_elem
      return {DecodedAddElemOp::Kind::FP32_FP32_FP32_16x1x1x1, /*conf*/ 0};
    }
 
    return {DecodedAddElemOp::Kind::UNSUPPORTED, -1};
  }
 
  LogicalResult
  matchAndRewrite(aievec::AddElemOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Location loc = op.getLoc();
    auto decodedAddElemOp = decodeAddElemOp(adaptor);
 
    if (decodedAddElemOp.kind == DecodedAddElemOp::Kind::UNSUPPORTED) {
      op.emitWarning() << "aievec.add_elem conversion is not supported.\n";
      return failure();
    }
 
    // Handle the FP32 add_elem for AIE2 - uses packed I64 representation
    if (decodedAddElemOp.kind ==
        DecodedAddElemOp::Kind::FP32_FP32_FP32_16x1x1x1) {
      auto confCst = LLVM::ConstantOp::create(
          rewriter, loc, rewriter.getI32Type(),
          rewriter.getI32IntegerAttr(decodedAddElemOp.conf));
      SmallVector<Value> operands(
          {adaptor.getLhs(), adaptor.getRhs(), confCst});
 
      auto addElemOp = xllvm::AddAccFloatAIE2IntrOp::create(
          rewriter, loc, VectorType::get({8}, rewriter.getI64Type()),
          forceCastOperandsToSignature(
              rewriter, loc, operands,
              {VectorType::get({8}, rewriter.getI64Type()),
               VectorType::get({8}, rewriter.getI64Type()),
               rewriter.getI32Type()}));
 
      // create bitcast/shape_cast for result
      auto resultVal = forceCastValueToType(rewriter, loc, addElemOp,
                                            op.getResult().getType());
      rewriter.replaceOp(op, resultVal);
      return success();
    }
 
    op.emitWarning() << "aievec.add_elem conversion is not supported.\n";
    return failure();
  }
};
 
// AIE2 version of SubElemOp conversion
class SubElemOpAIE2Conversion
    : public mlir::ConvertOpToLLVMPattern<aievec::SubElemOp> {
public:
  using ConvertOpToLLVMPattern<aievec::SubElemOp>::ConvertOpToLLVMPattern;
 
  struct DecodedSubElemOp {
    enum class Kind { FP32_FP32_FP32_16x1x1x1, UNSUPPORTED };
    Kind kind;
    int conf;
  };
 
  static DecodedSubElemOp decodeSubElemOp(OpAdaptor op) {
    auto lhs = op.getLhs();
    auto lhsVecTy = cast<VectorType>(lhs.getType());
    auto lhsScaTy = lhsVecTy.getElementType();
    unsigned lhsBitWidth = lhsScaTy.getIntOrFloatBitWidth();
 
    // Integer types
    if (llvm::isa<IntegerType>(lhsScaTy)) {
      return {DecodedSubElemOp::Kind::UNSUPPORTED, -1};
    } // Float types
    if (lhsBitWidth == 32) {
      // FP32 sub_elem
      return {DecodedSubElemOp::Kind::FP32_FP32_FP32_16x1x1x1, /*conf*/ 0};
    }
 
    return {DecodedSubElemOp::Kind::UNSUPPORTED, -1};
  }
 
  LogicalResult
  matchAndRewrite(aievec::SubElemOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Location loc = op.getLoc();
    auto decodedSubElemOp = decodeSubElemOp(adaptor);
 
    if (decodedSubElemOp.kind == DecodedSubElemOp::Kind::UNSUPPORTED) {
      op.emitWarning() << "aievec.sub_elem conversion is not supported.\n";
      return failure();
    }
 
    // Handle the FP32 sub_elem for AIE2 - uses packed I64 representation
    if (decodedSubElemOp.kind ==
        DecodedSubElemOp::Kind::FP32_FP32_FP32_16x1x1x1) {
      auto confCst = LLVM::ConstantOp::create(
          rewriter, loc, rewriter.getI32Type(),
          rewriter.getI32IntegerAttr(decodedSubElemOp.conf));
      SmallVector<Value> operands(
          {adaptor.getLhs(), adaptor.getRhs(), confCst});
 
      auto subElemOp = xllvm::SubAccFloatAIE2IntrOp::create(
          rewriter, loc, VectorType::get({8}, rewriter.getI64Type()),
          forceCastOperandsToSignature(
              rewriter, loc, operands,
              {VectorType::get({8}, rewriter.getI64Type()),
               VectorType::get({8}, rewriter.getI64Type()),
               rewriter.getI32Type()}));
 
      // create bitcast/shape_cast for result
      auto resultVal = forceCastValueToType(rewriter, loc, subElemOp,
                                            op.getResult().getType());
      rewriter.replaceOp(op, resultVal);
      return success();
    }
 
    op.emitWarning() << "aievec.sub_elem conversion is not supported.\n";
    return failure();
  }
};
 
// AIE2p version of AddElemOp and SubElemOp conversion for f32.
//
// AIE2p adds and subtracts f32 only on the 2048-bit accumulator, which holds
// 64 lanes: the ACC2048 accfloat intrinsic. An operand with fewer lanes that
// divide 64 is padded with poison, and the low lanes of the result are kept.
// The intrinsic is used even at 64 lanes, as Peano's generic fsub on
// <64 x float> gives wrong results.
template <typename SrcOpTy, typename IntrOpTy>
class AccFloatAddSubAIE2pConversion
    : public mlir::ConvertOpToLLVMPattern<SrcOpTy> {
public:
  using mlir::ConvertOpToLLVMPattern<SrcOpTy>::ConvertOpToLLVMPattern;
  using OpAdaptor = typename SrcOpTy::Adaptor;
 
  LogicalResult
  matchAndRewrite(SrcOpTy op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    constexpr int64_t accLanes = 64;
    Location loc = op.getLoc();
    auto type = cast<VectorType>(adaptor.getLhs().getType());
    int64_t lanes = type.getNumElements();
    if (!type.getElementType().isF32() || type.getRank() != 1 ||
        accLanes % lanes != 0 || lanes % 2 != 0) {
      op.emitWarning() << op->getName() << " conversion is not supported.\n";
      return failure();
    }
 
    // The padding and the extraction shuffle pairs of lanes as i64, which
    // Peano lowers to fewer register moves than f32 shuffles.
    auto accType = VectorType::get({accLanes}, rewriter.getF32Type());
    auto accI64Type = VectorType::get({accLanes / 2}, rewriter.getI64Type());
    auto i64Type = VectorType::get({lanes / 2}, rewriter.getI64Type());
    auto lowLanes = [](int64_t n, int64_t size) {
      SmallVector<int64_t> mask(size, -1);
      std::iota(mask.begin(), mask.begin() + std::min(n, size), 0);
      return mask;
    };
    Value lhs = adaptor.getLhs(), rhs = adaptor.getRhs();
    if (lanes != accLanes) {
      lhs = LLVM::BitcastOp::create(rewriter, loc, i64Type, lhs);
      rhs = LLVM::BitcastOp::create(rewriter, loc, i64Type, rhs);
      SmallVector<int64_t> padMask = lowLanes(lanes / 2, accLanes / 2);
      lhs = vector::ShuffleOp::create(rewriter, loc, lhs, lhs, padMask);
      rhs = vector::ShuffleOp::create(rewriter, loc, rhs, rhs, padMask);
      lhs = LLVM::BitcastOp::create(rewriter, loc, accType, lhs);
      rhs = LLVM::BitcastOp::create(rewriter, loc, accType, rhs);
    }
    // conf 60 selects the fp32 accumulator datapath.
    auto confCst = LLVM::ConstantOp::create(
        rewriter, loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(60));
    Value result = IntrOpTy::create(rewriter, loc, accType, lhs, rhs, confCst);
    if (lanes != accLanes) {
      result = LLVM::BitcastOp::create(rewriter, loc, accI64Type, result);
      result = vector::ShuffleOp::create(rewriter, loc, result, result,
                                         lowLanes(lanes / 2, lanes / 2));
      result = LLVM::BitcastOp::create(rewriter, loc, type, result);
    }
    rewriter.replaceOp(op, result);
    return success();
  }
};
 
using AddElemOpAIE2pConversion =
    AccFloatAddSubAIE2pConversion<aievec::AddElemOp,
                                  xllvm::AddACC2048AccFloatAIE2pIntrOp>;
using SubElemOpAIE2pConversion =
    AccFloatAddSubAIE2pConversion<aievec::SubElemOp,
                                  xllvm::SubACC2048AccFloatAIE2pIntrOp>;
 
// AIE2 version of NegOp conversion.
//
// aievec.neg operates on an accumulator. For floats the accumulator is
// v16accfloat, which is carried as vector<16xf32> in AIEVec and as <8 x i64>
// at the intrinsic boundary, so this maps directly onto the ACC512 accfloat
// negate.
class NegOpAIE2Conversion : public mlir::ConvertOpToLLVMPattern<aievec::NegOp> {
public:
  using ConvertOpToLLVMPattern<aievec::NegOp>::ConvertOpToLLVMPattern;
 
  LogicalResult
  matchAndRewrite(aievec::NegOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Location loc = op.getLoc();
    auto srcVecTy = cast<VectorType>(adaptor.getSource().getType());
    auto srcScaTy = srcVecTy.getElementType();
 
    // Only the float accumulator (v16accfloat) is wired up. Integer
    // accumulators would need the ACC1024 acc32/acc64 negates instead.
    if (!isa<FloatType>(srcScaTy) || srcScaTy.getIntOrFloatBitWidth() != 32 ||
        srcVecTy.getNumElements() != 16) {
      op.emitWarning() << "aievec.neg conversion is not supported.\n";
      return failure();
    }
 
    auto v8i64Ty = VectorType::get({8}, rewriter.getI64Type());
    // conf selects the fp32 accumulator datapath. Matches the AIE API's
    // neg(v16accfloat): aiev2_compute_control(0, 0, /*amode=*/2, /*bmode=*/3,
    // 0, 0, 0, 0, 0, 0, 0) == (2 << 1) | (3 << 3) == 28.
    auto confCst = LLVM::ConstantOp::create(
        rewriter, loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(28));
    SmallVector<Value> operands({adaptor.getSource(), confCst});
 
    auto negOp = xllvm::NegACC512AccFloatAIE2IntrOp::create(
        rewriter, loc, v8i64Ty,
        forceCastOperandsToSignature(rewriter, loc, operands,
                                     {v8i64Ty, rewriter.getI32Type()}));
 
    auto resultVal =
        forceCastValueToType(rewriter, loc, negOp, op.getResult().getType());
    rewriter.replaceOp(op, resultVal);
    return success();
  }
};
 
// AIE2p version of NegOp conversion.
//
// AIE2p only exposes the negate on the 2048-bit accumulator, so a
// v16accfloat operand is widened to <64 x float>, negated, and the low 16
// lanes extracted back out. This mirrors AddElemOpAIE2pConversion.
class NegOpAIE2pConversion
    : public mlir::ConvertOpToLLVMPattern<aievec::NegOp> {
public:
  using ConvertOpToLLVMPattern<aievec::NegOp>::ConvertOpToLLVMPattern;
 
  LogicalResult
  matchAndRewrite(aievec::NegOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Location loc = op.getLoc();
    auto srcVecTy = cast<VectorType>(adaptor.getSource().getType());
    auto srcScaTy = srcVecTy.getElementType();
    unsigned laneSize = srcVecTy.getNumElements();
 
    if (!isa<FloatType>(srcScaTy) || srcScaTy.getIntOrFloatBitWidth() != 32 ||
        (laneSize != 16 && laneSize != 32)) {
      op.emitWarning() << "aievec.neg conversion is not supported.\n";
      return failure();
    }
 
    auto v64f32Ty = VectorType::get({64}, rewriter.getF32Type());
 
    // The masks below address individual lanes, but vector.shuffle indexes
    // the leading dimension, so an n-D operand has to be flattened first.
    Value srcVal = adaptor.getSource();
    auto fltSrcVecTy = getFlattenedVectorType(srcVecTy);
    if (srcVecTy != fltSrcVecTy)
      srcVal = vector::ShapeCastOp::create(rewriter, loc, fltSrcVecTy, srcVal)
                   .getResult();
 
    // Widen the accumulator to ACC2048, filling the unused lanes with poison.
    SmallVector<int64_t> expandMask;
    for (unsigned i = 0; i < laneSize; ++i)
      expandMask.push_back(i);
    for (unsigned i = laneSize; i < 64; ++i)
      expandMask.push_back(-1);
    auto srcExpanded =
        vector::ShuffleOp::create(rewriter, loc, srcVal, srcVal, expandMask);
 
    // conf selects the fp32 accumulator datapath, matching the ACC2048
    // accfloat add/sub lowerings and the AIE API's neg(v64accfloat):
    // aie2p_compute_control(0, 0, /*amode=*/2, /*bmode=*/3, /*variant=*/1,
    // 0, 0, 0, 0, 0, 0) == (2 << 1) | (3 << 3) | (1 << 5) == 60.
    auto confCst = LLVM::ConstantOp::create(
        rewriter, loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(60));
    auto negResult = xllvm::NegACC2048AccFloatAIE2pIntrOp::create(
        rewriter, loc, v64f32Ty, srcExpanded, confCst);
 
    SmallVector<int64_t> extractMask;
    for (unsigned i = 0; i < laneSize; ++i)
      extractMask.push_back(i);
    Value finalResult = vector::ShuffleOp::create(rewriter, loc, negResult,
                                                  negResult, extractMask)
                            .getResult();
 
    // Restore the operand's shape.
    if (srcVecTy != fltSrcVecTy)
      finalResult =
          vector::ShapeCastOp::create(rewriter, loc, srcVecTy, finalResult)
              .getResult();
 
    rewriter.replaceOp(op, finalResult);
    return success();
  }
};
 
class SubOpConversion
    : public mlir::ConvertOpToLLVMPattern<aievec::aie1::SubOp> {
public:
  using ConvertOpToLLVMPattern<aievec::aie1::SubOp>::ConvertOpToLLVMPattern;
 
  LogicalResult
  matchAndRewrite(aievec::aie1::SubOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    op.emitWarning() << "aie.sub conversion is not implemented\n";
    return failure();
  }
};
 
class FMAOpConversion
    : public mlir::ConvertOpToLLVMPattern<aievec::aie1::FMAOp> {
public:
  using ConvertOpToLLVMPattern<aievec::aie1::FMAOp>::ConvertOpToLLVMPattern;
 
  LogicalResult
  matchAndRewrite(aievec::aie1::FMAOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto module = op->getParentOfType<ModuleOp>();
    MLIRContext *context = rewriter.getContext();
 
    auto startType = IntegerType::get(context, 32);
    auto offsetsType = VectorType::get({2}, IntegerType::get(context, 32));
    auto confType = VectorType::get({2}, IntegerType::get(context, 32));
 
    // If the intrinsic declaration doesn't exist, create it
    std::string intrinsicName = getMulOrFMAIntrinsicName(op);
    auto func = module.lookupSymbol<LLVM::LLVMFuncOp>(
        StringAttr::get(context, intrinsicName));
 
    if (!func) {
      OpBuilder::InsertionGuard guard(rewriter);
      rewriter.setInsertionPointToStart(module.getBody());
      func = LLVM::LLVMFuncOp::create(
          rewriter, rewriter.getUnknownLoc(), intrinsicName,
          LLVM::LLVMFunctionType::get(
              op.getResult().getType(),
              {op.getLhs().getType(), op.getRhs().getType(),
               op.getAcc().getType(), startType, /* xstart */
               startType,                        /* ystart */
               startType,                        /* zstart */
               offsetsType,                      /* xoffsets */
               offsetsType,                      /* zoffsets */
               confType}));
    }
 
    // Parse the string attribute values
    BufferParams x = {};
    BufferParams z = {};
    op.getXstart().getAsInteger(0, x.start);
    op.getXoffsets().getAsInteger(0, x.offsets);
    op.getXoffsetsHi().getAsInteger(0, x.offsets_hi);
    op.getXstep().getAsInteger(0, x.step);
    op.getXsquare().getAsInteger(0, x.square);
    op.getZstart().getAsInteger(0, z.start);
    op.getZoffsets().getAsInteger(0, z.offsets);
    op.getZoffsetsHi().getAsInteger(0, z.offsets_hi);
    op.getZstep().getAsInteger(0, z.step);
    op.getZsquare().getAsInteger(0, z.square);
 
    // Encode the configuration register
    uint32_t conf[2] = {0, 0};
    encodeConf(conf, x, z, op.getFmsub());
 
    // Create the constants and replace the op
    auto xstartVal = LLVM::ConstantOp::create(
        rewriter, op->getLoc(), startType, rewriter.getI32IntegerAttr(x.start));
    auto ystartVal = LLVM::ConstantOp::create(rewriter, op->getLoc(), startType,
                                              rewriter.getI32IntegerAttr(0));
    auto zstartVal = LLVM::ConstantOp::create(
        rewriter, op->getLoc(), startType, rewriter.getI32IntegerAttr(z.start));
    auto xoffsetsVal = LLVM::ConstantOp::create(
        rewriter, op->getLoc(), offsetsType,
        rewriter.getI32VectorAttr({(int32_t)x.offsets, (int32_t)x.offsets_hi}));
    auto zoffsetsVal = LLVM::ConstantOp::create(
        rewriter, op->getLoc(), offsetsType,
        rewriter.getI32VectorAttr({(int32_t)z.offsets, (int32_t)z.offsets_hi}));
    auto confVal = LLVM::ConstantOp::create(
        rewriter, op->getLoc(), confType,
        rewriter.getI32VectorAttr({(int32_t)conf[0], (int32_t)conf[1]}));
    rewriter.replaceOpWithNewOp<LLVM::CallOp>(
        op, func,
        ValueRange{op.getLhs(), op.getRhs(), op.getAcc(), xstartVal, ystartVal,
                   zstartVal, xoffsetsVal, zoffsetsVal, confVal});
    return success();
  }
};
 
class MulOpConversion
    : public mlir::ConvertOpToLLVMPattern<aievec::aie1::MulOp> {
public:
  using ConvertOpToLLVMPattern<aievec::aie1::MulOp>::ConvertOpToLLVMPattern;
 
  LogicalResult
  matchAndRewrite(aievec::aie1::MulOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto module = op->getParentOfType<ModuleOp>();
    MLIRContext *context = rewriter.getContext();
 
    auto startType = IntegerType::get(context, 32);
    auto offsetsType = VectorType::get({2}, IntegerType::get(context, 32));
    auto confType = VectorType::get({2}, IntegerType::get(context, 32));
 
    // If the intrinsic declaration doesn't exist, create it
    std::string intrinsicName = getMulOrFMAIntrinsicName(op);
    auto func = module.lookupSymbol<LLVM::LLVMFuncOp>(
        StringAttr::get(context, intrinsicName));
 
    if (!func) {
      OpBuilder::InsertionGuard guard(rewriter);
      rewriter.setInsertionPointToStart(module.getBody());
      func = LLVM::LLVMFuncOp::create(
          rewriter, rewriter.getUnknownLoc(), intrinsicName,
          LLVM::LLVMFunctionType::get(op.getResult().getType(),
                                      {op.getLhs().getType(),
                                       op.getRhs().getType(),
                                       startType,   /* xstart */
                                       startType,   /* ystart */
                                       startType,   /* zstart */
                                       offsetsType, /* xoffsets */
                                       offsetsType, /* zoffsets */
                                       confType}));
    }
 
    // Parse the string attribute values
    BufferParams x = {};
    BufferParams z = {};
    op.getXstart().getAsInteger(0, x.start);
    op.getXoffsets().getAsInteger(0, x.offsets);
    op.getXoffsetsHi().getAsInteger(0, x.offsets_hi);
    op.getXstep().getAsInteger(0, x.step);
    op.getXsquare().getAsInteger(0, x.square);
    op.getZstart().getAsInteger(0, z.start);
    op.getZoffsets().getAsInteger(0, z.offsets);
    op.getZoffsetsHi().getAsInteger(0, z.offsets_hi);
    op.getZstep().getAsInteger(0, z.step);
    op.getZsquare().getAsInteger(0, z.square);
 
    // Encode the configuration register
    uint32_t conf[2] = {0, 0};
    encodeConf(conf, x, z, false);
 
    // Create the constants and replace the op
    auto xstartVal = LLVM::ConstantOp::create(
        rewriter, op->getLoc(), startType, rewriter.getI32IntegerAttr(x.start));
    auto ystartVal = LLVM::ConstantOp::create(rewriter, op->getLoc(), startType,
                                              rewriter.getI32IntegerAttr(0));
    auto zstartVal = LLVM::ConstantOp::create(
        rewriter, op->getLoc(), startType, rewriter.getI32IntegerAttr(z.start));
    auto xoffsetsVal = LLVM::ConstantOp::create(
        rewriter, op->getLoc(), offsetsType,
        rewriter.getI32VectorAttr({(int32_t)x.offsets, (int32_t)x.offsets_hi}));
    auto zoffsetsVal = LLVM::ConstantOp::create(
        rewriter, op->getLoc(), offsetsType,
        rewriter.getI32VectorAttr({(int32_t)z.offsets, (int32_t)z.offsets_hi}));
    auto confVal = LLVM::ConstantOp::create(
        rewriter, op->getLoc(), confType,
        rewriter.getI32VectorAttr({(int32_t)conf[0], (int32_t)conf[1]}));
    rewriter.replaceOpWithNewOp<LLVM::CallOp>(
        op, func,
        ValueRange{op.getLhs(), op.getRhs(), xstartVal, ystartVal, zstartVal,
                   xoffsetsVal, zoffsetsVal, confVal});
    return success();
  }
};
 
class MulElemOpConversion
    : public mlir::ConvertOpToLLVMPattern<aievec::MulElemOp> {
public:
  using ConvertOpToLLVMPattern<aievec::MulElemOp>::ConvertOpToLLVMPattern;
 
  MulElemOpConversion(const LLVMTypeConverter &typeConverter,
                      Aie2Fp32Emulation aie2Fp32EmulationOption)
      : ConvertOpToLLVMPattern(typeConverter),
        aie2Fp32EmulationOption(aie2Fp32EmulationOption) {}
 
  Aie2Fp32Emulation aie2Fp32EmulationOption;
 
  struct DecodedMulElemOp {
    enum class Kind {
      // DtIn0_DtIn1_DtRes_CxMxKxN
      I8_I8_I32_32x1x2x1,
      I16_I16_I32_32x1x1x1,
      I32_I32_I64_32x1x2x1,
      BF16_BF16_FP32_16x1x2x1,
      FP32_FP32_FP32_16x1x1x1,
      UNSUPPORTED
      // TODO: I16_I16_I64_16x1x2x1
    };
 
    Kind kind;
    int conf;
  };
 
  static DecodedMulElemOp decodeMulElemOp(OpAdaptor op) {
    auto lhs = op.getLhs();
    auto lhsVecTy = cast<VectorType>(lhs.getType());
    auto lhsScaTy = lhsVecTy.getElementType();
    unsigned lhsBitWidth = lhsScaTy.getIntOrFloatBitWidth();
 
    // Integer types
    if (llvm::isa<IntegerType>(lhsScaTy)) {
      if (lhsBitWidth == 8) {
        return {DecodedMulElemOp::Kind::I8_I8_I32_32x1x2x1,
                aiev2_vmac_compute_control(
                    /*sgn_x=*/1, /*sgn_y=*/1, /*amode=*/0, /*bmode=*/1,
                    /*variant=*/1, /*zero_acc=*/0, /*shift16=*/0,
                    /*sub_mul=*/0, /*sub_acc1=*/0, /*sub_acc2=*/0,
                    /*sub_mask=*/0)};
      }
      if (lhsBitWidth == 16) {
        return {DecodedMulElemOp::Kind::I16_I16_I32_32x1x1x1,
                aiev2_vmac_compute_control(
                    /*sgn_x=*/1, /*sgn_y=*/1, /*amode=*/0, /*bmode=*/3,
                    /*variant=*/1, /*zero_acc=*/0, /*shift16=*/0,
                    /*sub_mul=*/0, /*sub_acc1=*/0, /*sub_acc2=*/0,
                    /*sub_mask=*/0)};
      }
      if (lhsBitWidth == 32) {
        // emulated I32 mul_elem
        return {DecodedMulElemOp::Kind::I32_I32_I64_32x1x2x1, -1};
      }
    } else {
      // Float types
      if (lhsBitWidth == 16) {
        return {DecodedMulElemOp::Kind::BF16_BF16_FP32_16x1x2x1,
                aiev2_vmac_compute_control(
                    /*sgn_x=*/0, /*sgn_y=*/0, /*amode=*/2, /*bmode=*/3,
                    /*variant=*/1, /*zero_acc=*/0, /*shift16=*/0,
                    /*sub_mul=*/0, /*sub_acc1=*/0, /*sub_acc2=*/0,
                    /*sub_mask=*/0)};
      }
      if (lhsBitWidth == 32) {
        // emulated FP32 mul_elem
        return {DecodedMulElemOp::Kind::FP32_FP32_FP32_16x1x1x1, -1};
      }
    }
 
    return {DecodedMulElemOp::Kind::UNSUPPORTED, -1};
  }
 
  // This conversion pattern implements the below CPP emulated I32 mul_elem.
  // INTRINSIC(v16acc64)
  // mul_elem_16_2(v16int32 a0, v16int32 a1, v16int32 b0, v16int32 b1) {
  //   v32uint16 a_lo = (v32uint16)shuffle(a0, a1, 2);
  //   v32int16 a_hi = (v32int16)shuffle(a0, a1, 3);
  //   v32uint16 b_lo = (v32uint16)shuffle(b0, b1, 2);
  //   v32int16 b_hi = (v32int16)shuffle(b0, b1, 3);
  //   v16acc64 acc = ::mul_elem_16_2(a_hi, b_hi);
  //   acc = mac_elem_16_2_conf(a_hi, 1, b_lo, false, acc, 0, 1, 0, 0);
  //   acc = mac_elem_16_2_conf(a_lo, false, b_hi, 1, acc, 0, 0, 0, 0);
  //   acc = mac_elem_16_2_conf(a_lo, false, b_lo, false, acc, 0, 1, 0, 0);
  //   return acc;
  // }
  // Caller example when handling the elementwise mul of two v16int32 vectors.
  //   v16int32 v1 = LHS();
  //   v16int32 v2 = RHS();
  //   v16acc64 v3 = mul_elem_16_2(v1, broadcast_zero_s32(), v2,
  //   undef_v16int32());
  // Explantion:
  // a_lo = low_part(a0[0]--a0[15], a1[0]--a1[15])
  // a_hi = high_part(a0[0]--a0[15], a1[0]--a1[15])
  // b_lo = low_part(b0[0]--b0[15], b1[0]--b1[15])
  // b_hi = high_part(b0[0]--b0[15], b1[0]--b1[15])
  // The firt `acc` is from mul_elem_16_2(a_hi, b_hi), which performs 16 channel
  // of 1x2x1 matmul, acc[0] = a_hi[0]*b_hi[0]+a_hi[16]*b_hi[16], ... , acc[15]
  // = a_hi[15]*b_hi[15]+a_hi[31]*b_hi[31]. Then, the first MAC performs `acc`
  // left shift 16bit, and then 16 channel of 1x2x1 matmul (a_hi, b_lo)
  // accumulating to `acc`. The second MAC performs 16 channel of 1x2x1 matmul
  // (a_lo, b_hi) accumulating to `acc`. Finally, the third MAC performs 16
  // channel of 1x2x1 matmul (a_lo, b_hi) accumulating to `acc`.
  LogicalResult
  convertToEmulatedI32MulElem(aievec::MulElemOp op, OpAdaptor adaptor,
                              ConversionPatternRewriter &rewriter) const {
 
    Location loc = op.getLoc();
    auto zeroCst = LLVM::ConstantOp::create(
        rewriter, loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(0));
    auto a0 = adaptor.getLhs();
    auto a1 = xllvm::VectorBroadcast32I512IntrOp::create(
        rewriter, loc, VectorType::get({16}, rewriter.getI32Type()), zeroCst);
    auto b0 = adaptor.getRhs();
    auto b1 = xllvm::UndefV16I32IntrOp::create(
        rewriter, loc, VectorType::get({16}, rewriter.getI32Type()));
 
    // 4* Shuffle
    auto a_lo = xllvm::VectorShuffleIntrOp::create(
        rewriter, loc, VectorType::get({16}, rewriter.getI32Type()), a0, a1,
        LLVM::ConstantOp::create(rewriter, loc, rewriter.getI32Type(),
                                 rewriter.getI32IntegerAttr(2)));
    auto a_hi = xllvm::VectorShuffleIntrOp::create(
        rewriter, loc, VectorType::get({16}, rewriter.getI32Type()), a0, a1,
        LLVM::ConstantOp::create(rewriter, loc, rewriter.getI32Type(),
                                 rewriter.getI32IntegerAttr(3)));
    auto b_lo = xllvm::VectorShuffleIntrOp::create(
        rewriter, loc, VectorType::get({16}, rewriter.getI32Type()), b0, b1,
        LLVM::ConstantOp::create(rewriter, loc, rewriter.getI32Type(),
                                 rewriter.getI32IntegerAttr(2)));
    auto b_hi = xllvm::VectorShuffleIntrOp::create(
        rewriter, loc, VectorType::get({16}, rewriter.getI32Type()), b0, b1,
        LLVM::ConstantOp::create(rewriter, loc, rewriter.getI32Type(),
                                 rewriter.getI32IntegerAttr(3)));
    // MUL + 3 * MAC
    auto mulConfCst = LLVM::ConstantOp::create(
        rewriter, loc, rewriter.getI32Type(),
        rewriter.getI32IntegerAttr(aiev2_vmac_compute_control(
            /*sgn_x=*/1, /*sgn_y=*/1, /*amode=*/1, /*bmode=*/3,
            /*variant=*/2, /*zero_acc=*/0, /*shift16=*/0,
            /*sub_mul=*/0, /*sub_acc1=*/0, /*sub_acc2=*/0, /*sub_mask=*/0)));
    auto mulConfOp = xllvm::MulConfAcc64IntrOp::create(
        rewriter, loc, VectorType::get({16}, rewriter.getI64Type()),
        forceCastOperandsToSignature(
            rewriter, loc,
            /*operands=*/{a_hi, b_hi, mulConfCst},
            /*signature=*/
            {VectorType::get({64}, rewriter.getI8Type()),
             VectorType::get({16}, rewriter.getI32Type()),
             rewriter.getI32Type()}));
 
    auto createMacConfOp = [&](SmallVector<Value> operands,
                               int macConf) -> Value {
      operands.push_back(
          LLVM::ConstantOp::create(rewriter, loc, rewriter.getI32Type(),
                                   rewriter.getI32IntegerAttr(macConf)));
      return xllvm::MacConfAcc64IntrOp::create(
                 rewriter, loc, VectorType::get({16}, rewriter.getI64Type()),
                 forceCastOperandsToSignature(
                     rewriter, loc,
                     /*operands=*/operands,
                     /*signature=*/
                     {VectorType::get({64}, rewriter.getI8Type()),
                      VectorType::get({16}, rewriter.getI32Type()),
                      VectorType::get({16}, rewriter.getI64Type()),
                      rewriter.getI32Type()}))
          .getResult();
    };
    auto acc64Val = mulConfOp.getResult();
    acc64Val = createMacConfOp(
        SmallVector<Value>{a_hi, b_lo, acc64Val},
        aiev2_vmac_compute_control(
            /*sgn_x=*/1, /*sgn_y=*/0, /*amode=*/1, /*bmode=*/3,
            /*variant=*/2, /*zero_acc=*/0, /*shift16=*/1,
            /*sub_mul=*/0, /*sub_acc1=*/0, /*sub_acc2=*/0, /*sub_mask=*/0));
    acc64Val = createMacConfOp(
        SmallVector<Value>{a_lo, b_hi, acc64Val},
        aiev2_vmac_compute_control(
            /*sgn_x=*/0, /*sgn_y=*/1, /*amode=*/1, /*bmode=*/3,
            /*variant=*/2, /*zero_acc=*/0, /*shift16=*/0,
            /*sub_mul=*/0, /*sub_acc1=*/0, /*sub_acc2=*/0, /*sub_mask=*/0));
    acc64Val = createMacConfOp(
        SmallVector<Value>{a_lo, b_lo, acc64Val},
        aiev2_vmac_compute_control(
            /*sgn_x=*/0, /*sgn_y=*/0, /*amode=*/1, /*bmode=*/3,
            /*variant=*/2, /*zero_acc=*/0, /*shift16=*/1,
            /*sub_mul=*/0, /*sub_acc1=*/0, /*sub_acc2=*/0, /*sub_mask=*/0));
 
    // create bitcast/shape_cast for result
    auto resultVal =
        forceCastValueToType(rewriter, loc, acc64Val, op.getResult().getType());
    rewriter.replaceOp(op, resultVal);
    return success();
  }
 
  // This conversion pattern implements the below CPP emulated FP32 mul_elem.
  // inline v16accfloat mul_elem_16_accuracy_safe(v16float v1, v16float v2) {
  //     v32bfloat16 a = broadcast_zero_to_v32bfloat16();
  //     v32bfloat16 b = broadcast_zero_to_v32bfloat16();
  //     v32bfloat16 c = broadcast_zero_to_v32bfloat16();
  //     v32bfloat16 d = broadcast_zero_to_v32bfloat16();
  //     v32bfloat16 e = broadcast_zero_to_v32bfloat16();
  //     v32bfloat16 f = broadcast_zero_to_v32bfloat16();
  //     v32bfloat16 dummy0 = broadcast_one_to_v32bfloat16();
  //     a = insert(a,0,to_v16bfloat16((v16accfloat)v1));
  //     v16accfloat acc0 = msc_elem_16_2(a, dummy0, (v16accfloat)v1);
  //     b = insert(b,0,to_v16bfloat16(acc0));
  //     c = insert(c,0,to_v16bfloat16(msc_elem_16_2(b, dummy0, acc0)));
  //     d = insert(d,0,to_v16bfloat16((v16accfloat)v2));
  //     v16accfloat acc1 = msc_elem_16_2(d, dummy0, (v16accfloat)v2);
  //     e = insert(e,0,to_v16bfloat16(acc1));
  //     f = insert(f,0,to_v16bfloat16(msc_elem_16_2(e, dummy0, acc1)));
  //     return
  //     mac_elem_16_2(a,d,mac_elem_16_2(a,e,mac_elem_16_2(b,d,mac_elem_16_2(
  //        d,c,mac_elem_16_2(b,e,mac_elem_16_2(a,f,mac_elem_16_2(
  //           b,f,mac_elem_16_2(c,e,mul_elem_16_2(c,f)))))))));
  // }
  // Caller example when handling the elementwise mul of two v16float vectors.
  //   v16float v1 = LHS(); v16float v2 = RHS();
  //   v16accfloat v3 = mul_elem_16(v1, v2);
  // Explantion: For v32bfloat16 `a`, the first half v16bf16 contains `most
  // significant 7 bits of mantissa` from v1, and the second half v16bf16 are
  // zeros. For v16accfloat `acc0`, the MSC equals to "(original `v1` with 23
  // bits of mantissa) - (`a` with MSB 7 bits of mantissa from v1)". For
  // v32bfloat16 `b`, the first half v16bf16 contains `[7:13] bits of mantissa
  // from v1` from v1, and the second half v16bf16 are zeros. For v32bfloat16
  // `c`, the first half v16bf16 contains `[14:20] bits of mantissa from v1`
  // from v1, and the second half v16bf16 are zeros. Hence, we can represent
  // v16float in three v32bfloat16 and then perform 9 MUL/MAC in v32bfloat16 to
  // get the final elementwise multiplication result.
 
  LogicalResult
  convertToEmulatedFP32MulElem(aievec::MulElemOp op, OpAdaptor adaptor,
                               ConversionPatternRewriter &rewriter) const {
    Location loc = op.getLoc();
    auto zeroCst =
        LLVM::ConstantOp::create(rewriter, loc, rewriter.getBF16Type(),
                                 rewriter.getZeroAttr(rewriter.getBF16Type()));
    auto aZeros = xllvm::VectorBroadcast16BF512IntrOp::create(
        rewriter, loc, VectorType::get({32}, rewriter.getBF16Type()), zeroCst);
    auto bZeros = xllvm::VectorBroadcast16BF512IntrOp::create(
        rewriter, loc, VectorType::get({32}, rewriter.getBF16Type()), zeroCst);
    auto cZeros = xllvm::VectorBroadcast16BF512IntrOp::create(
        rewriter, loc, VectorType::get({32}, rewriter.getBF16Type()), zeroCst);
    auto dZeros = xllvm::VectorBroadcast16BF512IntrOp::create(
        rewriter, loc, VectorType::get({32}, rewriter.getBF16Type()), zeroCst);
    auto eZeros = xllvm::VectorBroadcast16BF512IntrOp::create(
        rewriter, loc, VectorType::get({32}, rewriter.getBF16Type()), zeroCst);
    auto fZeros = xllvm::VectorBroadcast16BF512IntrOp::create(
        rewriter, loc, VectorType::get({32}, rewriter.getBF16Type()), zeroCst);
    auto oneCst =
        LLVM::ConstantOp::create(rewriter, loc, rewriter.getBF16Type(),
                                 rewriter.getOneAttr(rewriter.getBF16Type()));
    auto dummy0 = xllvm::VectorBroadcast16BF512IntrOp::create(
        rewriter, loc, VectorType::get({32}, rewriter.getBF16Type()), oneCst);
    auto zeroCstI32 = LLVM::ConstantOp::create(
        rewriter, loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(0));
    auto mscMacMulConfCst = LLVM::ConstantOp::create(
        rewriter, loc, rewriter.getI32Type(),
        rewriter.getI32IntegerAttr(aiev2_vmac_compute_control(
            /*sgn_x=*/0, /*sgn_y=*/0, /*amode=*/2, /*bmode=*/3,
            /*variant=*/1, /*zero_acc=*/0, /*shift16=*/0,
            /*sub_mul=*/0, /*sub_acc1=*/0, /*sub_acc2=*/0, /*sub_mask=*/0)));
 
    auto extractV16FP32ToThreeV16BF16 =
        [&](Value inputV16FP32, Value aZeros, Value bZeros,
            Value cZeros) -> std::tuple<Value, Value, Value> {
      // a = insert(a,0,to_v16bfloat16((v16accfloat)v1));
      auto inputBitCasted =
          forceCastValueToType(rewriter, loc, inputV16FP32,
                               VectorType::get({8}, rewriter.getI64Type()));
      auto v1ToBF16 = xllvm::Vector16AccFloatToV16BF16AIE2IntrOp::create(
          rewriter, loc, VectorType::get({16}, rewriter.getBF16Type()),
          inputBitCasted);
      auto a = xllvm::UpdBF512BF256IntrOp::create(
          rewriter, loc, VectorType::get({32}, rewriter.getBF16Type()), aZeros,
          v1ToBF16, zeroCstI32);
 
      // v16accfloat acc0 = msc_elem_16_2(a, dummy0, (v16accfloat)v1);
      auto acc0 = xllvm::MscConfBF16IntrOp::create(
          rewriter, loc, VectorType::get({8}, rewriter.getI64Type()), a, dummy0,
          inputBitCasted, mscMacMulConfCst);
 
      // b = insert(b,0,to_v16bfloat16(acc0));
      auto acc0ToBF16 = xllvm::Vector16AccFloatToV16BF16AIE2IntrOp::create(
          rewriter, loc, VectorType::get({16}, rewriter.getBF16Type()), acc0);
      auto b = xllvm::UpdBF512BF256IntrOp::create(
          rewriter, loc, VectorType::get({32}, rewriter.getBF16Type()), bZeros,
          acc0ToBF16, zeroCstI32);
 
      // c = insert(c,0,to_v16bfloat16(msc_elem_16_2(b, dummy0, acc0)));
      auto acc0Mscb = xllvm::MscConfBF16IntrOp::create(
          rewriter, loc, VectorType::get({8}, rewriter.getI64Type()), b, dummy0,
          acc0, mscMacMulConfCst);
      auto acc0MscbToBF16 = xllvm::Vector16AccFloatToV16BF16AIE2IntrOp::create(
          rewriter, loc, VectorType::get({16}, rewriter.getBF16Type()),
          acc0Mscb);
      auto c = xllvm::UpdBF512BF256IntrOp::create(
          rewriter, loc, VectorType::get({32}, rewriter.getBF16Type()), cZeros,
          acc0MscbToBF16, zeroCstI32);
      return std::make_tuple(a.getResult(), b.getResult(), c.getResult());
    };
 
    // Get v16vfloat16 a, b, c for representing v16float v1
    auto [a, b, c] =
        extractV16FP32ToThreeV16BF16(adaptor.getLhs(), aZeros, bZeros, cZeros);
    // Get v16vfloat16 d, e, f for representing v16float v2
    auto [d, e, f] =
        extractV16FP32ToThreeV16BF16(adaptor.getRhs(), dZeros, eZeros, fZeros);
 
    // Create 1 MUL and 2/5/8 MACs depending on the Aie2Fp32EmulationOption
    auto createMacOps = [&](Value lhs, Value rhs, Value acc) -> Value {
      return xllvm::MacConfBF16IntrOp::create(
                 rewriter, loc, VectorType::get({8}, rewriter.getI64Type()),
                 lhs, rhs, acc, mscMacMulConfCst)
          .getResult();
    };
 
    Value finalMacVal;
    if (aie2Fp32EmulationOption == Aie2Fp32Emulation::AccuracyFast) {
      // Fast and Accurate option. float a*b would require 6 mac operations.
      // Input fp32 number is split in to 3 bfloat16 numbers to extract all the
      // bits of the mantissa. float a,b; both a and b are split in to 3
      // bfloat16 numbers each. Hence there would be 9 mac operations in
      // multiplication of a and b. In the 9 mac operations to emulate fp32 mul,
      // mac operations with LSBs are ignored. (3 last terms). This helps
      // improve cycle count of mul and has least impact on accuracy of result.
      // This is the default option to the aiecompiler
      auto afMul = xllvm::MulConfBF16IntrOp::create(
          rewriter, loc, VectorType::get({8}, rewriter.getI64Type()), a, f,
          mscMacMulConfCst);
      finalMacVal = createMacOps(
          a, d,
          createMacOps(
              a, e,
              createMacOps(b, d,
                           createMacOps(d, c, createMacOps(b, e, afMul)))));
    } else if (aie2Fp32EmulationOption == Aie2Fp32Emulation::AccuracyLow) {
      // Fast and least accurate option. float a*b would require 3 mac
      // operations.
      // Input fp32 number is split in to 2 bfloat16 numbers. Hence not all the
      // bits from mantissa can be used. float a,b; Both a and b are split in to
      // 2 bfloat16 numbers each. Hence there would be 4 mac operations in
      // multiplication of a and b. In the 4 mac operations to emulate fp32 mul,
      // mac operations with LSBs are ignored. (1 last term). This helps improve
      // cycle count of mul float a, b;
      auto bdMul = xllvm::MulConfBF16IntrOp::create(
          rewriter, loc, VectorType::get({8}, rewriter.getI64Type()), b, d,
          mscMacMulConfCst);
      finalMacVal = createMacOps(a, d, createMacOps(a, e, bdMul));
    } else {
      // aie2Fp32EmulationOption == Aie2Fp32Emulation::AccuracySafe
      // Most accurate option since input fp32 number is split in to 3 bfloat16
      // numbers to extract all the bits of the mantissa. float a*b would
      // require 9 mac operations due to 3 bfloat16 splits each.
      auto cfMul = xllvm::MulConfBF16IntrOp::create(
          rewriter, loc, VectorType::get({8}, rewriter.getI64Type()), c, f,
          mscMacMulConfCst);
      finalMacVal = createMacOps(
          a, d,
          createMacOps(
              a, e,
              createMacOps(
                  b, d,
                  createMacOps(
                      d, c,
                      createMacOps(
                          b, e,
                          createMacOps(
                              a, f,
                              createMacOps(b, f,
                                           createMacOps(c, e, cfMul))))))));
    }
 
    // create bitcast/shape_cast for result
    auto resultVal = forceCastValueToType(rewriter, loc, finalMacVal,
                                          op.getResult().getType());
    rewriter.replaceOp(op, resultVal);
    return success();
  }
 
  LogicalResult
  matchAndRewrite(aievec::MulElemOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Location loc = op.getLoc();
    auto decodedMulElemOp = decodeMulElemOp(adaptor);
 
    if (decodedMulElemOp.kind == DecodedMulElemOp::Kind::UNSUPPORTED) {
      op.emitWarning() << "aievec.mul_elem conversion is not supported.\n";
      return failure();
    }
 
    // Handle the emulated I32/FP32 mul_elem
    if (decodedMulElemOp.kind == DecodedMulElemOp::Kind::I32_I32_I64_32x1x2x1) {
      return convertToEmulatedI32MulElem(op, adaptor, rewriter);
    }
    if (decodedMulElemOp.kind ==
        DecodedMulElemOp::Kind::FP32_FP32_FP32_16x1x1x1) {
      return convertToEmulatedFP32MulElem(op, adaptor, rewriter);
    }
 
    // create constant for config
    auto confCst = LLVM::ConstantOp::create(
        rewriter, loc, rewriter.getI32Type(),
        rewriter.getI32IntegerAttr(decodedMulElemOp.conf));
    Value mulElemOp = nullptr;
    SmallVector<Value> operands({adaptor.getLhs(), adaptor.getRhs(), confCst});
 
    // create xllvm intrinsic
    if (decodedMulElemOp.kind == DecodedMulElemOp::Kind::I16_I16_I32_32x1x1x1 ||
        decodedMulElemOp.kind == DecodedMulElemOp::Kind::I8_I8_I32_32x1x2x1) {
      mulElemOp = xllvm::MulConfAcc32IntrOp::create(
          rewriter, loc, VectorType::get({16}, rewriter.getI64Type()),
          forceCastOperandsToSignature(
              rewriter, loc, operands,
              {VectorType::get({64}, rewriter.getI8Type()),
               VectorType::get({16}, rewriter.getI32Type()),
               rewriter.getI32Type()}));
    } else if (decodedMulElemOp.kind ==
               DecodedMulElemOp::Kind::BF16_BF16_FP32_16x1x2x1) {
      // Create zero vector using the exact pattern from working reference:
      // vbroadcast16.I512(0) -> bitcast to bf16 -> extract lower 256 bits
      auto zero32 = LLVM::ConstantOp::create(
          rewriter, loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(0));
      auto zeros_i16 = xllvm::VectorBroadcast16I512IntrOp::create(
          rewriter, loc, VectorType::get({32}, rewriter.getI16Type()), zero32);
      auto zeros_bf16 = LLVM::BitcastOp::create(
          rewriter, loc, VectorType::get({32}, rewriter.getBF16Type()),
          zeros_i16);
      auto zeroVec = xllvm::ExtBF256BF512IntrOp::create(
          rewriter, loc, VectorType::get({16}, rewriter.getBF16Type()),
          zeros_bf16, zero32);
 
      // Use set+upd pattern to match working reference
      auto idx1 = LLVM::ConstantOp::create(rewriter, loc, rewriter.getI32Type(),
                                           rewriter.getI32IntegerAttr(1));
 
      // Set lhs at lower 256 bits, then update upper 256 bits with zeros
      auto lhsSet = xllvm::VectorSetBF512BF256IntrOp::create(
          rewriter, loc, VectorType::get({32}, rewriter.getBF16Type()),
          adaptor.getLhs(), zero32);
      auto lhsConcat = xllvm::UpdBF512BF256IntrOp::create(
          rewriter, loc, VectorType::get({32}, rewriter.getBF16Type()), lhsSet,
          zeroVec, idx1);
 
      // Set rhs at lower 256 bits, then update upper 256 bits with zeros
      auto rhsSet = xllvm::VectorSetBF512BF256IntrOp::create(
          rewriter, loc, VectorType::get({32}, rewriter.getBF16Type()),
          adaptor.getRhs(), zero32);
      auto rhsConcat = xllvm::UpdBF512BF256IntrOp::create(
          rewriter, loc, VectorType::get({32}, rewriter.getBF16Type()), rhsSet,
          zeroVec, idx1);
 
      // Call bf.mul16.conf with padded vectors
      mulElemOp = xllvm::MulConfBF16IntrOp::create(
          rewriter, loc, VectorType::get({8}, rewriter.getI64Type()), lhsConcat,
          rhsConcat, confCst);
    }
 
    // create bitcast/shape_cast for result
    auto resultVal = forceCastValueToType(rewriter, loc, mulElemOp,
                                          op.getResult().getType());
    rewriter.replaceOp(op, resultVal);
    return success();
  }
};
 
// AIE2p version of MulElemOp conversion
class MulElemOpAIE2pConversion
    : public mlir::ConvertOpToLLVMPattern<aievec::MulElemOp> {
public:
  using ConvertOpToLLVMPattern<aievec::MulElemOp>::ConvertOpToLLVMPattern;
 
  struct DecodedMulElemOp {
    enum class Kind {
      BF16_BF16_FP32_16x1x1x1, // 16-lane bf16 -> 16-lane f32
      BF16_BF16_FP32_32x1x2x1, // 32-lane bf16 -> 32-lane f32
      BF16_BF16_FP32_64x1x2x1, // 64-lane bf16 -> 64-lane f32
      UNSUPPORTED
    };
    Kind kind;
    int conf;
  };
 
  static DecodedMulElemOp decodeMulElemOp(OpAdaptor op) {
    auto lhs = op.getLhs();
    auto lhsVecTy = cast<VectorType>(lhs.getType());
    auto lhsScaTy = lhsVecTy.getElementType();
    unsigned lhsBitWidth = lhsScaTy.getIntOrFloatBitWidth();
    int lhsLanes = getVectorLaneSize(lhsVecTy);
 
    // Integer types - not supported for AIE2p elementwise mul
    if (llvm::isa<IntegerType>(lhsScaTy)) {
      return {DecodedMulElemOp::Kind::UNSUPPORTED, -1};
    } // Float types
    if (lhsBitWidth == 16) {
      // BF16 mul_elem
      if (lhsLanes == 16) {
        // 16-lane         bfloat16 uses I512.I512.ACC512 intrinsic
        return {DecodedMulElemOp::Kind::BF16_BF16_FP32_16x1x1x1, /*conf*/ 60};
      }
      if (lhsLanes == 32) {
        // 32-lane         bfloat16 uses I512.I512.ACC1024 intrinsic
        return {DecodedMulElemOp::Kind::BF16_BF16_FP32_32x1x2x1, /*conf*/ 60};
      }
      if (lhsLanes == 64) {
        // 64-lane         bfloat16 uses I1024.I1024.ACC2048 intrinsic
        return {DecodedMulElemOp::Kind::BF16_BF16_FP32_64x1x2x1, /*conf*/ 60};
      }
    }
 
    return {DecodedMulElemOp::Kind::UNSUPPORTED, -1};
  }
 
  LogicalResult
  matchAndRewrite(aievec::MulElemOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Location loc = op.getLoc();
    auto decodedMulElemOp = decodeMulElemOp(adaptor);
 
    if (decodedMulElemOp.kind == DecodedMulElemOp::Kind::UNSUPPORTED) {
      op.emitWarning() << "aievec.mul_elem conversion is not supported for "
                          "AIE2p.\n";
      return failure();
    }
 
    // Create constant for config
    auto confCst = LLVM::ConstantOp::create(
        rewriter, loc, rewriter.getI32Type(),
        rewriter.getI32IntegerAttr(decodedMulElemOp.conf));
 
    Value mulElemOp = nullptr;
 
    // Handle BF16 mul_elem for AIE2p
    if (decodedMulElemOp.kind ==
        DecodedMulElemOp::Kind::BF16_BF16_FP32_16x1x1x1) {
      // 16-lane bfloat16: <16 x bfloat> x <16 x bfloat> -> <16 x float>
      // The intrinsic requires <32 x bfloat> inputs, so we need to pad
 
      // Pad LHS from 16 to 32 bfloat16 using shuffle
      SmallVector<int64_t> padMask;
      for (int i = 0; i < 16; ++i)
        padMask.push_back(i);
      for (int i = 16; i < 32; ++i)
        padMask.push_back(-1); // poison/undef
 
      auto lhsPadded = vector::ShuffleOp::create(
          rewriter, loc, adaptor.getLhs(), adaptor.getLhs(), padMask);
      auto rhsPadded = vector::ShuffleOp::create(
          rewriter, loc, adaptor.getRhs(), adaptor.getRhs(), padMask);
 
      SmallVector<Value> operands({lhsPadded, rhsPadded, confCst});
 
      // Call I512.I512.ACC512 intrinsic
      mulElemOp = xllvm::MulConfBF16I512ACC512AIE2pIntrOp::create(
          rewriter, loc, VectorType::get({16}, rewriter.getF32Type()),
          forceCastOperandsToSignature(
              rewriter, loc, operands,
              {VectorType::get({32}, rewriter.getBF16Type()),
               VectorType::get({32}, rewriter.getBF16Type()),
               rewriter.getI32Type()}));
    } else if (decodedMulElemOp.kind ==
               DecodedMulElemOp::Kind::BF16_BF16_FP32_32x1x2x1) {
      // 32-lane bfloat16: <32 x bfloat> x <32 x bfloat> -> <32 x float>
      SmallVector<Value> operands(
          {adaptor.getLhs(), adaptor.getRhs(), confCst});
      mulElemOp = xllvm::MulConfBF16I512ACC1024AIE2pIntrOp::create(
          rewriter, loc, VectorType::get({32}, rewriter.getF32Type()),
          forceCastOperandsToSignature(
              rewriter, loc, operands,
              {VectorType::get({32}, rewriter.getBF16Type()),
               VectorType::get({32}, rewriter.getBF16Type()),
               rewriter.getI32Type()}));
    } else if (decodedMulElemOp.kind ==
               DecodedMulElemOp::Kind::BF16_BF16_FP32_64x1x2x1) {
      // 64-lane bfloat16: <64 x bfloat> x <64 x bfloat> -> <64 x float>
      SmallVector<Value> operands(
          {adaptor.getLhs(), adaptor.getRhs(), confCst});
      mulElemOp = xllvm::MulConfBF16I1024ACC2048AIE2pIntrOp::create(
          rewriter, loc, VectorType::get({64}, rewriter.getF32Type()),
          forceCastOperandsToSignature(
              rewriter, loc, operands,
              {VectorType::get({64}, rewriter.getBF16Type()),
               VectorType::get({64}, rewriter.getBF16Type()),
               rewriter.getI32Type()}));
    }
 
    // create bitcast/shape_cast for result
    auto resultVal = forceCastValueToType(rewriter, loc, mulElemOp,
                                          op.getResult().getType());
    rewriter.replaceOp(op, resultVal);
    return success();
  }
};
 
// AIE2p version of FMAElemOp conversion. Uses native F32 accumulators
// and AIE2p-specific MAC intrinsics.
class FMAElemOpAIE2pConversion
    : public mlir::ConvertOpToLLVMPattern<aievec::FMAElemOp> {
public:
  using ConvertOpToLLVMPattern<aievec::FMAElemOp>::ConvertOpToLLVMPattern;
 
  LogicalResult
  matchAndRewrite(aievec::FMAElemOp fmaOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto loc = fmaOp.getLoc();
    auto lhs = adaptor.getLhs();
    auto rhs = adaptor.getRhs();
    auto acc = adaptor.getAcc();
    auto lhsTy = cast<VectorType>(lhs.getType());
    auto accTy = cast<VectorType>(acc.getType());
    auto flatLhsTy = getFlattenedVectorType(lhsTy);
    auto flatAccTy = getFlattenedVectorType(accTy);
 
    // Flatten operands, if needed
    if (lhsTy != flatLhsTy)
      lhs = vector::ShapeCastOp::create(rewriter, loc, flatLhsTy, lhs);
    if (cast<VectorType>(rhs.getType()) != flatLhsTy)
      rhs = vector::ShapeCastOp::create(rewriter, loc, flatLhsTy, rhs);
    if (accTy != flatAccTy)
      acc = vector::ShapeCastOp::create(rewriter, loc, flatAccTy, acc);
 
    if (!flatLhsTy.getElementType().isBF16()) {
      fmaOp.emitWarning()
          << "aievec.mac_elem AIE2p conversion only supports bf16 inputs.\n";
      return failure();
    }
 
    Type i32ty = rewriter.getI32Type();
    auto confCst = LLVM::ConstantOp::create(
        rewriter, loc, i32ty,
        rewriter.getI32IntegerAttr(aiev2_vmac_compute_control(
            /*sgn_x=*/0, /*sgn_y=*/0, /*amode=*/2, /*bmode=*/3,
            /*variant=*/1, /*zero_acc=*/0, /*shift16=*/0,
            /*sub_mul=*/0, /*sub_acc1=*/0, /*sub_acc2=*/0,
            /*sub_mask=*/0)));
 
    unsigned lhsLanes = flatLhsTy.getNumElements();
    Value macIntrOp = nullptr;
 
    if (lhsLanes == 16) {
      // 16-lane bf16: pad to v32bf16, use I512.I512.ACC512 intrinsic
      SmallVector<int64_t> padMask;
      for (int i = 0; i < 16; ++i)
        padMask.push_back(i);
      for (int i = 16; i < 32; ++i)
        padMask.push_back(-1); // poison
 
      auto lhsPadded =
          vector::ShuffleOp::create(rewriter, loc, lhs, lhs, padMask);
      auto rhsPadded =
          vector::ShuffleOp::create(rewriter, loc, rhs, rhs, padMask);
 
      auto v32bf16Ty = VectorType::get({32}, rewriter.getBF16Type());
      auto v16f32Ty = VectorType::get({16}, rewriter.getF32Type());
      macIntrOp = xllvm::MacConfBF16I512ACC512AIE2pIntrOp::create(
          rewriter, loc, v16f32Ty,
          forceCastOperandsToSignature(
              rewriter, loc, {lhsPadded, rhsPadded, acc, confCst},
              {v32bf16Ty, v32bf16Ty, v16f32Ty, i32ty}));
    } else if (lhsLanes == 32) {
      // 32-lane bf16: direct, use I512.I512.ACC1024 intrinsic
      auto v32bf16Ty = VectorType::get({32}, rewriter.getBF16Type());
      auto v32f32Ty = VectorType::get({32}, rewriter.getF32Type());
      macIntrOp = xllvm::MacConfBF16I512ACC1024AIE2pIntrOp::create(
          rewriter, loc, v32f32Ty,
          forceCastOperandsToSignature(
              rewriter, loc, {lhs, rhs, acc, confCst},
              {v32bf16Ty, v32bf16Ty, v32f32Ty, i32ty}));
    } else {
      fmaOp.emitWarning()
          << "aievec.mac_elem AIE2p conversion: unsupported lane count "
          << lhsLanes << ".\n";
      return failure();
    }
 
    // Recast/Reshape result
    auto resVal = forceCastValueToType(rewriter, loc, macIntrOp, flatAccTy);
    if (flatAccTy != accTy)
      resVal = vector::ShapeCastOp::create(rewriter, loc, accTy, resVal);
 
    rewriter.replaceOp(fmaOp, resVal);
    return success();
  }
};
 
// Enum to represent different AIE target architectures
enum class AIEArch {
  AIE2,
  AIE2p,
};
 
class UPSOpAIE2Conversion : public mlir::ConvertOpToLLVMPattern<aievec::UPSOp> {
public:
  using ConvertOpToLLVMPattern<aievec::UPSOp>::ConvertOpToLLVMPattern;
 
  LogicalResult
  matchAndRewrite(aievec::UPSOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Location loc = op.getLoc();
 
    Value result = op.getResult();
    VectorType resultType = cast<VectorType>(result.getType());
    VectorType flatResTy = getFlattenedVectorType(resultType);
    Type resultScaTy = resultType.getElementType();
    unsigned resultBitWidth = resultScaTy.getIntOrFloatBitWidth();
    int resultLanes = getVectorLaneSize(resultType);
    int resultVectorSize = resultBitWidth * resultLanes;
 
    Value opSrcVal = adaptor.getSource();
    auto srcVecTy = cast<VectorType>(opSrcVal.getType());
    auto fltSrcVecTy = getFlattenedVectorType(srcVecTy);
    if (srcVecTy != fltSrcVecTy)
      opSrcVal = vector::ShapeCastOp::create(rewriter, op.getLoc(), fltSrcVecTy,
                                             opSrcVal)
                     .getResult();
 
    // create xllvm intrinsic
    // Integer types
    Value upsIntrOp = nullptr;
    if (llvm::isa<IntegerType>(resultScaTy)) {
      // create constant for sign
      auto signCst = LLVM::ConstantOp::create(
          rewriter, loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(1));
      auto shiftCst =
          LLVM::ConstantOp::create(rewriter, loc, rewriter.getI32Type(),
                                   rewriter.getI32IntegerAttr(op.getShift()));
 
      SmallVector<Value> operands({opSrcVal, shiftCst, signCst});
      if (resultVectorSize == 512) {
        if (resultBitWidth == 32) {
          // v16int16 -> v16acc32
          upsIntrOp = xllvm::Acc32V16I256UpsAIE2IntrOp::create(
              rewriter, loc, VectorType::get({8}, rewriter.getI64Type()),
              forceCastOperandsToSignature(
                  rewriter, loc, operands,
                  {VectorType::get({16}, rewriter.getI16Type()),
                   rewriter.getI32Type(), rewriter.getI32Type()}));
        } else if (resultBitWidth == 64) {
          // v8int32 -> v8acc64
          upsIntrOp = xllvm::Acc64V8I256UpsAIE2IntrOp::create(
              rewriter, loc, VectorType::get({8}, rewriter.getI64Type()),
              forceCastOperandsToSignature(
                  rewriter, loc, operands,
                  {VectorType::get({8}, rewriter.getI32Type()),
                   rewriter.getI32Type(), rewriter.getI32Type()}));
        }
      } else if (resultVectorSize == 1024) {
        Value src = opSrcVal;
        VectorType srcType = cast<VectorType>(src.getType());
        Type srcScaType = srcType.getElementType();
        unsigned srcBitWidth = srcScaType.getIntOrFloatBitWidth();
 
        if (resultBitWidth == 32 && srcBitWidth == 16) {
          // v32int16 -> v32acc32
          upsIntrOp = xllvm::Acc32V32I512UpsAIE2IntrOp::create(
              rewriter, loc, VectorType::get({16}, rewriter.getI64Type()),
              forceCastOperandsToSignature(
                  rewriter, loc, operands,
                  {VectorType::get({32}, rewriter.getI16Type()),
                   rewriter.getI32Type(), rewriter.getI32Type()}));
        } else if (resultBitWidth == 64 && srcBitWidth == 32) {
          // v16int32 -> v16acc64
          upsIntrOp = xllvm::Acc64V16I512UpsAIE2IntrOp::create(
              rewriter, loc, VectorType::get({16}, rewriter.getI64Type()),
              forceCastOperandsToSignature(
                  rewriter, loc, operands,
                  {VectorType::get({16}, rewriter.getI32Type()),
                   rewriter.getI32Type(), rewriter.getI32Type()}));
        } else if (resultBitWidth == 64 && srcBitWidth == 16) {
          // v16int16 -> v16acc64
          upsIntrOp = xllvm::Acc64V16I256UpsAIE2IntrOp::create(
              rewriter, loc, VectorType::get({16}, rewriter.getI64Type()),
              forceCastOperandsToSignature(
                  rewriter, loc, operands,
                  {VectorType::get({16}, rewriter.getI16Type()),
                   rewriter.getI32Type(), rewriter.getI32Type()}));
        } else if (resultBitWidth == 32 && srcBitWidth == 8) {
          // v32int8 -> v32acc32
          upsIntrOp = xllvm::Acc32V32I256UpsAIE2IntrOp::create(
              rewriter, loc, VectorType::get({16}, rewriter.getI64Type()),
              forceCastOperandsToSignature(
                  rewriter, loc, operands,
                  {VectorType::get({32}, rewriter.getI8Type()),
                   rewriter.getI32Type(), rewriter.getI32Type()}));
        }
      }
    } else {
      // Float types
      // AIE2p uses native F32 types, AIE2 uses packed I64 types
      if (resultVectorSize == 512) {
        // v16bfloat16 -> v16accfloat
        upsIntrOp = xllvm::Vector16BF16ToV16AccFloatAIE2IntrOp::create(
            rewriter, loc, VectorType::get({8}, rewriter.getI64Type()),
            forceCastOperandsToSignature(
                rewriter, loc, {opSrcVal},
                {VectorType::get({16}, rewriter.getBF16Type())}));
      } else if (resultVectorSize == 1024) {
        // v32bfloat16 -> v32accfloat
        // The CPP example of the implementation is below:
        //   INTRINSIC(v32accfloat) ups_to_v32accfloat(v32bfloat16 a) {
        //     v16accfloat x0 = ups_to_v16accfloat(extract_v16bfloat16(a, 0));
        //     v16accfloat x1 = ups_to_v16accfloat(extract_v16bfloat16(a, 1));
        //     return concat(x0, x1);
        //   }
        auto indexZeroCst =
            LLVM::ConstantOp::create(rewriter, loc, rewriter.getI32Type(),
                                     rewriter.getI32IntegerAttr(0));
        auto indexOneCst =
            LLVM::ConstantOp::create(rewriter, loc, rewriter.getI32Type(),
                                     rewriter.getI32IntegerAttr(1));
        auto extractUps = [&](Value source, Value index) -> Value {
          auto extOp = xllvm::ExtI256I512IntrOp::create(
              rewriter, loc, VectorType::get({8}, rewriter.getI32Type()),
              forceCastOperandsToSignature(
                  rewriter, loc, {source, index},
                  {VectorType::get({16}, rewriter.getI32Type()),
                   rewriter.getI32Type()}));
          return xllvm::Vector16BF16ToV16AccFloatAIE2IntrOp::create(
              rewriter, loc, VectorType::get({8}, rewriter.getI64Type()),
              forceCastOperandsToSignature(
                  rewriter, loc, {extOp},
                  {VectorType::get({16}, rewriter.getBF16Type())}));
        };
        auto resLo = extractUps(opSrcVal, indexZeroCst);
        auto resHi = extractUps(opSrcVal, indexOneCst);
        // Concat the two 512-bit vector to a 1024-bit vector.
        // Note that given sources a0 and a1, the result is [a1; a0].
        upsIntrOp = xllvm::ConcatI1024I512IntrOp::create(
            rewriter, loc, VectorType::get({32}, rewriter.getI32Type()),
            forceCastOperandsToSignature(
                rewriter, loc, {resLo, resHi},
                {VectorType::get({16}, rewriter.getI32Type()),
                 VectorType::get({16}, rewriter.getI32Type())}));
      }
    }
 
    if (!upsIntrOp) {
      op.emitWarning() << "aievec.ups is not supported.\n";
      return failure();
    }
 
    // create bitcast for result if needed
    if (flatResTy != upsIntrOp.getType())
      upsIntrOp = LLVM::BitcastOp::create(rewriter, loc, flatResTy, upsIntrOp);
 
    if (flatResTy != resultType)
      upsIntrOp =
          vector::ShapeCastOp::create(rewriter, loc, resultType, upsIntrOp);
 
    rewriter.replaceOp(op, upsIntrOp);
 
    return success();
  }
};
 
// TODO: Split the op at AIEVec dialect level
class UPSOpAIE2pConversion
    : public mlir::ConvertOpToLLVMPattern<aievec::UPSOp> {
public:
  using ConvertOpToLLVMPattern<aievec::UPSOp>::ConvertOpToLLVMPattern;
 
  LogicalResult
  matchAndRewrite(aievec::UPSOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Location loc = op.getLoc();
 
    Value result = op.getResult();
    VectorType resultType = cast<VectorType>(result.getType());
    VectorType flatResTy = getFlattenedVectorType(resultType);
    Type resultScaTy = resultType.getElementType();
    unsigned resultBitWidth = resultScaTy.getIntOrFloatBitWidth();
    int resultLanes = getVectorLaneSize(resultType);
    int resultVectorSize = resultBitWidth * resultLanes;
 
    Value opSrcVal = adaptor.getSource();
    auto srcVecTy = cast<VectorType>(opSrcVal.getType());
    auto fltSrcVecTy = getFlattenedVectorType(srcVecTy);
    if (srcVecTy != fltSrcVecTy)
      opSrcVal = vector::ShapeCastOp::create(rewriter, op.getLoc(), fltSrcVecTy,
                                             opSrcVal)
                     .getResult();
 
    // create xllvm intrinsic
    // Integer types
    Value upsIntrOp = nullptr;
    if (llvm::isa<IntegerType>(resultScaTy)) {
      // create constant for sign
      auto signCst = LLVM::ConstantOp::create(
          rewriter, loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(1));
      auto shiftCst =
          LLVM::ConstantOp::create(rewriter, loc, rewriter.getI32Type(),
                                   rewriter.getI32IntegerAttr(op.getShift()));
 
      SmallVector<Value> operands({opSrcVal, shiftCst, signCst});
      if (resultVectorSize == 512) {
        if (resultBitWidth == 32) {
          // v16int16 -> v16acc32
          upsIntrOp = xllvm::Acc32V16I256UpsAIE2pIntrOp::create(
              rewriter, loc, VectorType::get({16}, rewriter.getI32Type()),
              forceCastOperandsToSignature(
                  rewriter, loc, operands,
                  {VectorType::get({16}, rewriter.getI16Type()),
                   rewriter.getI32Type(), rewriter.getI32Type()}));
        } else if (resultBitWidth == 64) {
          // v8int32 -> v8acc64
          upsIntrOp = xllvm::Acc64V8I256UpsAIE2pIntrOp::create(
              rewriter, loc, VectorType::get({8}, rewriter.getI64Type()),
              forceCastOperandsToSignature(
                  rewriter, loc, operands,
                  {VectorType::get({8}, rewriter.getI32Type()),
                   rewriter.getI32Type(), rewriter.getI32Type()}));
        }
      } else if (resultVectorSize == 1024) {
        Value src = opSrcVal;
        VectorType srcType = cast<VectorType>(src.getType());
        Type srcScaType = srcType.getElementType();
        unsigned srcBitWidth = srcScaType.getIntOrFloatBitWidth();
        int srcLanes = getVectorLaneSize(srcType);
        int srcVectorSize = srcBitWidth * srcLanes;
 
        if (resultBitWidth == 32 && srcBitWidth == 16 && srcVectorSize == 512) {
          // v32int16 -> v32acc32
          upsIntrOp = xllvm::Acc32V32I512UpsAIE2pIntrOp::create(
              rewriter, loc, VectorType::get({32}, rewriter.getI32Type()),
              forceCastOperandsToSignature(
                  rewriter, loc, operands,
                  {VectorType::get({32}, rewriter.getI16Type()),
                   rewriter.getI32Type(), rewriter.getI32Type()}));
        } else if (resultBitWidth == 64 && srcBitWidth == 32 &&
                   srcVectorSize == 512) {
          // v16int32 -> v16acc64
          upsIntrOp = xllvm::Acc64V16I512UpsAIE2pIntrOp::create(
              rewriter, loc, VectorType::get({16}, rewriter.getI64Type()),
              forceCastOperandsToSignature(
                  rewriter, loc, operands,
                  {VectorType::get({16}, rewriter.getI32Type()),
                   rewriter.getI32Type(), rewriter.getI32Type()}));
        } else if (resultBitWidth == 64 && srcBitWidth == 16 &&
                   srcVectorSize == 256) {
          // v16int16 -> v16acc64
          upsIntrOp = xllvm::Acc64V16I256UpsAIE2pIntrOp::create(
              rewriter, loc, VectorType::get({16}, rewriter.getI64Type()),
              forceCastOperandsToSignature(
                  rewriter, loc, operands,
                  {VectorType::get({16}, rewriter.getI16Type()),
                   rewriter.getI32Type(), rewriter.getI32Type()}));
        } else if (resultBitWidth == 32 && srcBitWidth == 8 &&
                   srcVectorSize == 256) {
          // v32int8 -> v32acc32
          upsIntrOp = xllvm::Acc32V32I256UpsAIE2pIntrOp::create(
              rewriter, loc, VectorType::get({32}, rewriter.getI32Type()),
              forceCastOperandsToSignature(
                  rewriter, loc, operands,
                  {VectorType::get({32}, rewriter.getI8Type()),
                   rewriter.getI32Type(), rewriter.getI32Type()}));
        }
      } else if (resultVectorSize == 2048) {
        Value src = opSrcVal;
        VectorType srcType = cast<VectorType>(src.getType());
        Type srcScaType = srcType.getElementType();
        unsigned srcBitWidth = srcScaType.getIntOrFloatBitWidth();
        int srcLanes = getVectorLaneSize(srcType);
        int srcVectorSize = srcBitWidth * srcLanes;
 
        if (resultBitWidth == 32 && srcBitWidth == 8 && srcVectorSize == 512) {
          // v64int8 -> v64acc32
          upsIntrOp = xllvm::Acc32V64I512UpsAIE2pIntrOp::create(
              rewriter, loc, VectorType::get({64}, rewriter.getI32Type()),
              forceCastOperandsToSignature(
                  rewriter, loc, operands,
                  {VectorType::get({64}, rewriter.getI8Type()),
                   rewriter.getI32Type(), rewriter.getI32Type()}));
        } else if (resultBitWidth == 64 && srcBitWidth == 16 &&
                   srcVectorSize == 512) {
          // v32int16 -> v32acc64
          upsIntrOp = xllvm::Acc64V32I512UpsAIE2pIntrOp::create(
              rewriter, loc, VectorType::get({32}, rewriter.getI64Type()),
              forceCastOperandsToSignature(
                  rewriter, loc, operands,
                  {VectorType::get({32}, rewriter.getI16Type()),
                   rewriter.getI32Type(), rewriter.getI32Type()}));
        } else if (resultBitWidth == 32 && srcBitWidth == 16 &&
                   srcVectorSize == 1024) {
          // v64int16 -> v64acc32
          // Extract 2 chunks of v32int16 and convert each to v32acc32
          auto index0Cst =
              LLVM::ConstantOp::create(rewriter, loc, rewriter.getI32Type(),
                                       rewriter.getI32IntegerAttr(0));
          auto index1Cst =
              LLVM::ConstantOp::create(rewriter, loc, rewriter.getI32Type(),
                                       rewriter.getI32IntegerAttr(1));
 
          auto extractUps2048 = [&](Value source, Value index, Value shiftCst,
                                    Value signCst) -> Value {
            // Use vector::ShuffleOp to extract 512-bit from 1024-bit
            // Cast source to v32xi32 for shuffling
            auto v32i32Source = forceCastValueToType(
                rewriter, loc, source,
                VectorType::get({32}, rewriter.getI32Type()));
 
            // Determine shuffle mask based on index
            // index 0: elements [0-15]
            // index 1: elements [16-31]
            SmallVector<int64_t> shuffleMask;
            if (auto constIndex = index.getDefiningOp<LLVM::ConstantOp>()) {
              auto indexAttr = cast<IntegerAttr>(constIndex.getValue());
              int64_t idxVal = indexAttr.getInt();
              int startIdx = idxVal * 16;
              for (int i = 0; i < 16; ++i) {
                shuffleMask.push_back(startIdx + i);
              }
            } else {
              // Default to index 0 if not constant
              for (int i = 0; i < 16; ++i) {
                shuffleMask.push_back(i);
              }
            }
 
            auto extOp = vector::ShuffleOp::create(rewriter, loc, v32i32Source,
                                                   v32i32Source, shuffleMask);
 
            return xllvm::Acc32V32I512UpsAIE2pIntrOp::create(
                rewriter, loc, VectorType::get({32}, rewriter.getI32Type()),
                forceCastOperandsToSignature(
                    rewriter, loc, {extOp, shiftCst, signCst},
                    {VectorType::get({32}, rewriter.getI16Type()),
                     rewriter.getI32Type(), rewriter.getI32Type()}));
          };
 
          auto res0 = extractUps2048(opSrcVal, index0Cst, shiftCst, signCst);
          auto res1 = extractUps2048(opSrcVal, index1Cst, shiftCst, signCst);
 
          // Concat two 1024-bit vectors to a 2048-bit vector using
          // vector::ShuffleOp
          SmallVector<int64_t> concatMask;
          for (int i = 0; i < 64; ++i) {
            concatMask.push_back(i);
          }
          upsIntrOp =
              vector::ShuffleOp::create(rewriter, loc, res0, res1, concatMask);
        }
      }
    } else {
      // Float types
      // AIE2p uses native F32 types, AIE2 uses packed I64 types
      if (resultVectorSize == 512) {
        // v16bfloat16 -> v16accfloat
        upsIntrOp = xllvm::Vector16BF16ToV16AccFloatAIE2pIntrOp::create(
            rewriter, loc, VectorType::get({16}, rewriter.getF32Type()),
            forceCastOperandsToSignature(
                rewriter, loc, {opSrcVal},
                {VectorType::get({16}, rewriter.getBF16Type())}));
      } else if (resultVectorSize == 1024) {
        // v32bfloat16 -> v32accfloat
        upsIntrOp = xllvm::Vector32BF16ToV32AccFloatAIE2pIntrOp::create(
            rewriter, loc, VectorType::get({32}, rewriter.getF32Type()),
            forceCastOperandsToSignature(
                rewriter, loc, {opSrcVal},
                {VectorType::get({32}, rewriter.getBF16Type())}));
      } else if (resultVectorSize == 2048) {
        // v64bfloat16 -> v64accfloat
        // Extract 2 chunks of v32bfloat16 and convert each to v32accfloat
        auto index0Cst =
            LLVM::ConstantOp::create(rewriter, loc, rewriter.getI32Type(),
                                     rewriter.getI32IntegerAttr(0));
        auto index1Cst =
            LLVM::ConstantOp::create(rewriter, loc, rewriter.getI32Type(),
                                     rewriter.getI32IntegerAttr(1));
 
        auto extractUps2048 = [&](Value source, Value index) -> Value {
          // Use vector::ShuffleOp to extract 512-bit from 1024-bit
          // Cast source to v32xi32 for shuffling
          auto v32i32Source = forceCastValueToType(
              rewriter, loc, source,
              VectorType::get({32}, rewriter.getI32Type()));
 
          // Determine shuffle mask based on index
          // index 0: elements [0-15]
          // index 1: elements [16-31]
          SmallVector<int64_t> shuffleMask;
          if (auto constIndex = index.getDefiningOp<LLVM::ConstantOp>()) {
            auto indexAttr = cast<IntegerAttr>(constIndex.getValue());
            int64_t idxVal = indexAttr.getInt();
            int startIdx = idxVal * 16;
            for (int i = 0; i < 16; ++i) {
              shuffleMask.push_back(startIdx + i);
            }
          } else {
            // Default to index 0 if not constant
            for (int i = 0; i < 16; ++i) {
              shuffleMask.push_back(i);
            }
          }
 
          auto extOp = vector::ShuffleOp::create(rewriter, loc, v32i32Source,
                                                 v32i32Source, shuffleMask);
 
          return xllvm::Vector32BF16ToV32AccFloatAIE2pIntrOp::create(
              rewriter, loc, VectorType::get({32}, rewriter.getF32Type()),
              forceCastOperandsToSignature(
                  rewriter, loc, {extOp},
                  {VectorType::get({32}, rewriter.getBF16Type())}));
        };
 
        auto res0 = extractUps2048(opSrcVal, index0Cst);
        auto res1 = extractUps2048(opSrcVal, index1Cst);
 
        // Concat two 1024-bit vectors to a 2048-bit vector using
        // vector::ShuffleOp
        auto v32i32Res0 = forceCastValueToType(
            rewriter, loc, res0, VectorType::get({32}, rewriter.getI32Type()));
        auto v32i32Res1 = forceCastValueToType(
            rewriter, loc, res1, VectorType::get({32}, rewriter.getI32Type()));
 
        SmallVector<int64_t> concatMask;
        for (int i = 0; i < 64; ++i) {
          concatMask.push_back(i);
        }
        upsIntrOp = vector::ShuffleOp::create(rewriter, loc, v32i32Res0,
                                              v32i32Res1, concatMask);
      }
    }
 
    if (!upsIntrOp) {
      op.emitWarning() << "aievec.ups is not supported.\n";
      return failure();
    }
 
    // create bitcast for result if needed
    if (flatResTy != upsIntrOp.getType())
      upsIntrOp = LLVM::BitcastOp::create(rewriter, loc, flatResTy, upsIntrOp);
 
    if (flatResTy != resultType)
      upsIntrOp =
          vector::ShapeCastOp::create(rewriter, loc, resultType, upsIntrOp);
 
    rewriter.replaceOp(op, upsIntrOp);
 
    return success();
  }
};
 
class SRSOpAIE2Conversion : public mlir::ConvertOpToLLVMPattern<aievec::SRSOp> {
public:
  using ConvertOpToLLVMPattern<aievec::SRSOp>::ConvertOpToLLVMPattern;
 
  LogicalResult
  matchAndRewrite(aievec::SRSOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Location loc = op.getLoc();
 
    Value result = op.getResult();
    VectorType resultType = cast<VectorType>(result.getType());
    Type resultScaTy = resultType.getElementType();
    unsigned resultBitWidth = resultScaTy.getIntOrFloatBitWidth();
    int resultLanes = getVectorLaneSize(resultType);
    int resultVectorSize = resultBitWidth * resultLanes;
 
    // Integer types
    Value srsIntrOp = nullptr;
    if (llvm::isa<IntegerType>(resultScaTy)) {
      // create constant for sign from the op's sign attribute
      auto signCst =
          LLVM::ConstantOp::create(rewriter, loc, rewriter.getI32Type(),
                                   rewriter.getI32IntegerAttr(op.getSign()));
 
      // create xllvm intrinsic
      SmallVector<Value> operands(
          {adaptor.getSource(), adaptor.getShift(), signCst});
      if (resultVectorSize == 512) {
        if (resultBitWidth == 16) {
          srsIntrOp = xllvm::I512V32Acc32SrsAIE2IntrOp::create(
              rewriter, loc, VectorType::get({32}, rewriter.getI16Type()),
              forceCastOperandsToSignature(
                  rewriter, loc, operands,
                  {VectorType::get({16}, rewriter.getI64Type()),
                   rewriter.getI32Type(), rewriter.getI32Type()}));
        } else if (resultBitWidth == 32) {
          srsIntrOp = xllvm::I512V16Acc64SrsAIE2IntrOp::create(
              rewriter, loc, VectorType::get({16}, rewriter.getI32Type()),
              forceCastOperandsToSignature(
                  rewriter, loc, operands,
                  {VectorType::get({16}, rewriter.getI64Type()),
                   rewriter.getI32Type(), rewriter.getI32Type()}));
        }
      } else if (resultVectorSize == 256) {
        Value src = adaptor.getSource();
        VectorType srcType = cast<VectorType>(src.getType());
        Type srcScaType = srcType.getElementType();
        unsigned srcBitWidth = srcScaType.getIntOrFloatBitWidth();
 
        if (resultBitWidth == 16 && srcBitWidth == 32) {
          srsIntrOp = xllvm::I256V16Acc32SrsAIE2IntrOp::create(
              rewriter, loc, VectorType::get({16}, rewriter.getI16Type()),
              forceCastOperandsToSignature(
                  rewriter, loc, operands,
                  {VectorType::get({8}, rewriter.getI64Type()),
                   rewriter.getI32Type(), rewriter.getI32Type()}));
        } else if (resultBitWidth == 8 && srcBitWidth == 32) {
          srsIntrOp = xllvm::I256V32Acc32SrsAIE2IntrOp::create(
              rewriter, loc, VectorType::get({32}, rewriter.getI8Type()),
              forceCastOperandsToSignature(
                  rewriter, loc, operands,
                  {VectorType::get({16}, rewriter.getI64Type()),
                   rewriter.getI32Type(), rewriter.getI32Type()}));
        } else if (resultBitWidth == 16 && srcBitWidth == 64) {
          srsIntrOp = xllvm::I256V16Acc64SrsAIE2IntrOp::create(
              rewriter, loc, VectorType::get({16}, rewriter.getI16Type()),
              forceCastOperandsToSignature(
                  rewriter, loc, operands,
                  {VectorType::get({16}, rewriter.getI64Type()),
                   rewriter.getI32Type(), rewriter.getI32Type()}));
        } else if (resultBitWidth == 32 && srcBitWidth == 64) {
          srsIntrOp = xllvm::I256V8Acc64SrsAIE2IntrOp::create(
              rewriter, loc, VectorType::get({8}, rewriter.getI32Type()),
              forceCastOperandsToSignature(
                  rewriter, loc, operands,
                  {VectorType::get({8}, rewriter.getI64Type()),
                   rewriter.getI32Type(), rewriter.getI32Type()}));
        }
      }
    } else {
      // Float types
      if (resultVectorSize == 256) {
        srsIntrOp = xllvm::Vector16AccFloatToV16BF16AIE2IntrOp::create(
            rewriter, loc, VectorType::get({16}, rewriter.getBF16Type()),
            forceCastOperandsToSignature(
                rewriter, loc, {adaptor.getSource()},
                {VectorType::get({8}, rewriter.getI64Type())}));
      } else if (resultVectorSize == 512) {
        // v32accfloat -> v32bfloat16
        // The CPP example of the implementation is below:
        //   v32bfloat16 to_v32bfloat16(v32accfloat acc) {
        //     v16bfloat16 x0 = to_v16bfloat16(extract_v16accfloat(acc, 0));
        //     v16bfloat16 x1 = to_v16bfloat16(extract_v16accfloat(acc, 1));
        //     return concat(x0, x1);
        //   }
        auto indexZeroCst =
            LLVM::ConstantOp::create(rewriter, loc, rewriter.getI32Type(),
                                     rewriter.getI32IntegerAttr(0));
        auto indexOneCst =
            LLVM::ConstantOp::create(rewriter, loc, rewriter.getI32Type(),
                                     rewriter.getI32IntegerAttr(1));
        auto extractSrs = [&](Value source, Value index) -> Value {
          auto extOp = xllvm::ExtI512I1024IntrOp::create(
              rewriter, loc, VectorType::get({16}, rewriter.getI32Type()),
              forceCastOperandsToSignature(
                  rewriter, loc, {source, index},
                  {VectorType::get({32}, rewriter.getI32Type()),
                   rewriter.getI32Type()}));
          return xllvm::Vector16AccFloatToV16BF16AIE2IntrOp::create(
              rewriter, loc, VectorType::get({16}, rewriter.getBF16Type()),
              forceCastOperandsToSignature(
                  rewriter, loc, {extOp},
                  {VectorType::get({8}, rewriter.getI64Type())}));
        };
        auto resLo = extractSrs(adaptor.getSource(), indexZeroCst);
        auto resHi = extractSrs(adaptor.getSource(), indexOneCst);
        // Concat the two 256-bit vector to a 512-bit vector.
        // Note that given sources a0 and a1, the result is [a1; a0].
        srsIntrOp = xllvm::ConcatI512I256IntrOp::create(
            rewriter, loc, VectorType::get({16}, rewriter.getI32Type()),
            forceCastOperandsToSignature(
                rewriter, loc, {resLo, resHi},
                {VectorType::get({8}, rewriter.getI32Type()),
                 VectorType::get({8}, rewriter.getI32Type())}));
      }
    }
 
    if (!srsIntrOp) {
      op.emitWarning() << "aievec.srs is not supported.\n";
      return failure();
    }
 
    // create bitcast/shape_cast for result if needed
    auto resultVal = forceCastValueToType(rewriter, loc, srsIntrOp,
                                          op.getResult().getType());
    rewriter.replaceOp(op, resultVal);
 
    return success();
  }
};
 
// TODO: Split the op at AIEVec dialect level
class SRSOpAIE2pConversion
    : public mlir::ConvertOpToLLVMPattern<aievec::SRSOp> {
public:
  using ConvertOpToLLVMPattern<aievec::SRSOp>::ConvertOpToLLVMPattern;
 
  LogicalResult
  matchAndRewrite(aievec::SRSOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Location loc = op.getLoc();
 
    Value result = op.getResult();
    VectorType resultType = cast<VectorType>(result.getType());
    Type resultScaTy = resultType.getElementType();
    unsigned resultBitWidth = resultScaTy.getIntOrFloatBitWidth();
    int resultLanes = getVectorLaneSize(resultType);
    int resultVectorSize = resultBitWidth * resultLanes;
 
    // Integer types
    Value srsIntrOp = nullptr;
    if (llvm::isa<IntegerType>(resultScaTy)) {
      // create constant for sign from the op's sign attribute
      auto signCst =
          LLVM::ConstantOp::create(rewriter, loc, rewriter.getI32Type(),
                                   rewriter.getI32IntegerAttr(op.getSign()));
 
      // create xllvm intrinsic
      SmallVector<Value> operands(
          {adaptor.getSource(), adaptor.getShift(), signCst});
      if (resultVectorSize == 512) {
        Value src = adaptor.getSource();
        VectorType srcType = cast<VectorType>(src.getType());
        Type srcScaType = srcType.getElementType();
        unsigned srcBitWidth = srcScaType.getIntOrFloatBitWidth();
 
        if (resultBitWidth == 16 && srcBitWidth == 32) {
          // v32acc32 -> v32int16
          srsIntrOp = xllvm::I512V32Acc32SrsAIE2pIntrOp::create(
              rewriter, loc, VectorType::get({32}, rewriter.getI16Type()),
              forceCastOperandsToSignature(
                  rewriter, loc, operands,
                  {VectorType::get({32}, rewriter.getI32Type()),
                   rewriter.getI32Type(), rewriter.getI32Type()}));
        } else if (resultBitWidth == 16 && srcBitWidth == 64) {
          // v32acc64 -> v32int16
          srsIntrOp = xllvm::I512V32Acc64SrsAIE2pIntrOp::create(
              rewriter, loc, VectorType::get({32}, rewriter.getI16Type()),
              forceCastOperandsToSignature(
                  rewriter, loc, operands,
                  {VectorType::get({32}, rewriter.getI64Type()),
                   rewriter.getI32Type(), rewriter.getI32Type()}));
        } else if (resultBitWidth == 32 && srcBitWidth == 64) {
          // v16acc64 -> v16int32
          srsIntrOp = xllvm::I512V16Acc64SrsAIE2pIntrOp::create(
              rewriter, loc, VectorType::get({16}, rewriter.getI32Type()),
              forceCastOperandsToSignature(
                  rewriter, loc, operands,
                  {VectorType::get({16}, rewriter.getI64Type()),
                   rewriter.getI32Type(), rewriter.getI32Type()}));
        } else if (resultBitWidth == 8 && srcBitWidth == 32) {
          // v64acc32 -> v64int8
          srsIntrOp = xllvm::I512V64Acc32SrsAIE2pIntrOp::create(
              rewriter, loc, VectorType::get({64}, rewriter.getI8Type()),
              forceCastOperandsToSignature(
                  rewriter, loc, operands,
                  {VectorType::get({64}, rewriter.getI32Type()),
                   rewriter.getI32Type(), rewriter.getI32Type()}));
        }
      } else if (resultVectorSize == 256) {
        Value src = adaptor.getSource();
        VectorType srcType = cast<VectorType>(src.getType());
        Type srcScaType = srcType.getElementType();
        unsigned srcBitWidth = srcScaType.getIntOrFloatBitWidth();
 
        if (resultBitWidth == 16 && srcBitWidth == 32) {
          // v16acc32 -> v16int16
          srsIntrOp = xllvm::I256V16Acc32SrsAIE2pIntrOp::create(
              rewriter, loc, VectorType::get({16}, rewriter.getI16Type()),
              forceCastOperandsToSignature(
                  rewriter, loc, operands,
                  {VectorType::get({16}, rewriter.getI32Type()),
                   rewriter.getI32Type(), rewriter.getI32Type()}));
        } else if (resultBitWidth == 8 && srcBitWidth == 32) {
          // v32acc32 -> v32int8
          srsIntrOp = xllvm::I256V32Acc32SrsAIE2pIntrOp::create(
              rewriter, loc, VectorType::get({32}, rewriter.getI8Type()),
              forceCastOperandsToSignature(
                  rewriter, loc, operands,
                  {VectorType::get({32}, rewriter.getI32Type()),
                   rewriter.getI32Type(), rewriter.getI32Type()}));
        } else if (resultBitWidth == 16 && srcBitWidth == 64) {
          // v16acc64 -> v16int16
          srsIntrOp = xllvm::I256V16Acc64SrsAIE2pIntrOp::create(
              rewriter, loc, VectorType::get({16}, rewriter.getI16Type()),
              forceCastOperandsToSignature(
                  rewriter, loc, operands,
                  {VectorType::get({16}, rewriter.getI64Type()),
                   rewriter.getI32Type(), rewriter.getI32Type()}));
        } else if (resultBitWidth == 32 && srcBitWidth == 64) {
          // v8acc64 -> v8int32
          srsIntrOp = xllvm::I256V8Acc64SrsAIE2pIntrOp::create(
              rewriter, loc, VectorType::get({8}, rewriter.getI32Type()),
              forceCastOperandsToSignature(
                  rewriter, loc, operands,
                  {VectorType::get({8}, rewriter.getI64Type()),
                   rewriter.getI32Type(), rewriter.getI32Type()}));
        }
      } else if (resultVectorSize == 1024) {
        Value src = adaptor.getSource();
        VectorType srcType = cast<VectorType>(src.getType());
        Type srcScaType = srcType.getElementType();
        unsigned srcBitWidth = srcScaType.getIntOrFloatBitWidth();
 
        if (resultBitWidth == 16 && srcBitWidth == 32) {
          // v64acc32 -> v64int16
          // Extract 2 chunks of v32acc32 and convert each to v32int16
          auto index0Cst =
              LLVM::ConstantOp::create(rewriter, loc, rewriter.getI32Type(),
                                       rewriter.getI32IntegerAttr(0));
          auto index1Cst =
              LLVM::ConstantOp::create(rewriter, loc, rewriter.getI32Type(),
                                       rewriter.getI32IntegerAttr(1));
 
          auto extractSrs1024 = [&](Value source, Value index, Value shiftCst,
                                    Value signCst) -> Value {
            // Use vector::ShuffleOp to extract 1024-bit from 2048-bit
            // Cast source to v64xi32 for shuffling
            auto v64i32Source = forceCastValueToType(
                rewriter, loc, source,
                VectorType::get({64}, rewriter.getI32Type()));
 
            // Determine shuffle mask based on index
            // index 0: elements [0-31]
            // index 1: elements [32-63]
            SmallVector<int64_t> shuffleMask;
            if (auto constIndex = index.getDefiningOp<LLVM::ConstantOp>()) {
              auto indexAttr = cast<IntegerAttr>(constIndex.getValue());
              int64_t idxVal = indexAttr.getInt();
              int startIdx = idxVal * 32;
              for (int i = 0; i < 32; ++i) {
                shuffleMask.push_back(startIdx + i);
              }
            } else {
              // Default to index 0 if not constant
              for (int i = 0; i < 32; ++i) {
                shuffleMask.push_back(i);
              }
            }
 
            auto extOp = vector::ShuffleOp::create(rewriter, loc, v64i32Source,
                                                   v64i32Source, shuffleMask);
 
            return xllvm::I512V32Acc32SrsAIE2pIntrOp::create(
                rewriter, loc, VectorType::get({32}, rewriter.getI16Type()),
                forceCastOperandsToSignature(
                    rewriter, loc, {extOp, shiftCst, signCst},
                    {VectorType::get({32}, rewriter.getI32Type()),
                     rewriter.getI32Type(), rewriter.getI32Type()}));
          };
 
          auto res0 =
              extractSrs1024(src, index0Cst, adaptor.getShift(), signCst);
          auto res1 =
              extractSrs1024(src, index1Cst, adaptor.getShift(), signCst);
 
          // Concat two 512-bit vectors to a 1024-bit vector using
          // vector::ShuffleOp
          auto v16i32Res0 = forceCastValueToType(
              rewriter, loc, res0,
              VectorType::get({16}, rewriter.getI32Type()));
          auto v16i32Res1 = forceCastValueToType(
              rewriter, loc, res1,
              VectorType::get({16}, rewriter.getI32Type()));
 
          SmallVector<int64_t> concatMask;
          for (int i = 0; i < 32; ++i) {
            concatMask.push_back(i);
          }
          srsIntrOp = vector::ShuffleOp::create(rewriter, loc, v16i32Res0,
                                                v16i32Res1, concatMask);
        }
      }
    } else {
      // Float types
      // AIE2p uses native F32 types, AIE2 uses packed I64 types
      if (resultVectorSize == 256) {
        // v16accfloat -> v16bfloat16
        srsIntrOp = xllvm::Vector16AccFloatToV16BF16AIE2pIntrOp::create(
            rewriter, loc, VectorType::get({16}, rewriter.getBF16Type()),
            forceCastOperandsToSignature(
                rewriter, loc, {adaptor.getSource()},
                {VectorType::get({16}, rewriter.getF32Type())}));
      } else if (resultVectorSize == 512) {
        // v32accfloat -> v32bfloat16
        srsIntrOp = xllvm::Vector32AccFloatToV32BF16AIE2pIntrOp::create(
            rewriter, loc, VectorType::get({32}, rewriter.getBF16Type()),
            forceCastOperandsToSignature(
                rewriter, loc, {adaptor.getSource()},
                {VectorType::get({32}, rewriter.getF32Type())}));
      } else if (resultVectorSize == 1024) {
        // v64accfloat -> v64bfloat16
        // Extract 2 chunks of v32accfloat and convert each to v32bfloat16
        auto index0Cst =
            LLVM::ConstantOp::create(rewriter, loc, rewriter.getI32Type(),
                                     rewriter.getI32IntegerAttr(0));
        auto index1Cst =
            LLVM::ConstantOp::create(rewriter, loc, rewriter.getI32Type(),
                                     rewriter.getI32IntegerAttr(1));
 
        auto extractSrs1024 = [&](Value source, Value index) -> Value {
          // Use vector::ShuffleOp to extract 1024-bit from 2048-bit
          // Cast source to v64xi32 for shuffling
          auto v64i32Source = forceCastValueToType(
              rewriter, loc, source,
              VectorType::get({64}, rewriter.getI32Type()));
 
          // Determine shuffle mask based on index
          // index 0: elements [0-31]
          // index 1: elements [32-63]
          SmallVector<int64_t> shuffleMask;
          if (auto constIndex = index.getDefiningOp<LLVM::ConstantOp>()) {
            auto indexAttr = cast<IntegerAttr>(constIndex.getValue());
            int64_t idxVal = indexAttr.getInt();
            int startIdx = idxVal * 32;
            for (int i = 0; i < 32; ++i) {
              shuffleMask.push_back(startIdx + i);
            }
          } else {
            // Default to index 0 if not constant
            for (int i = 0; i < 32; ++i) {
              shuffleMask.push_back(i);
            }
          }
 
          auto extOp = vector::ShuffleOp::create(rewriter, loc, v64i32Source,
                                                 v64i32Source, shuffleMask);
 
          return xllvm::Vector32AccFloatToV32BF16AIE2pIntrOp::create(
              rewriter, loc, VectorType::get({32}, rewriter.getBF16Type()),
              forceCastOperandsToSignature(
                  rewriter, loc, {extOp},
                  {VectorType::get({32}, rewriter.getF32Type())}));
        };
 
        auto res0 = extractSrs1024(adaptor.getSource(), index0Cst);
        auto res1 = extractSrs1024(adaptor.getSource(), index1Cst);
 
        // Concat two 512-bit vectors to a 1024-bit vector using
        // vector::ShuffleOp
        auto v16i32Res0 = forceCastValueToType(
            rewriter, loc, res0, VectorType::get({16}, rewriter.getI32Type()));
        auto v16i32Res1 = forceCastValueToType(
            rewriter, loc, res1, VectorType::get({16}, rewriter.getI32Type()));
 
        SmallVector<int64_t> concatMask;
        for (int i = 0; i < 32; ++i) {
          concatMask.push_back(i);
        }
        srsIntrOp = vector::ShuffleOp::create(rewriter, loc, v16i32Res0,
                                              v16i32Res1, concatMask);
      }
    }
 
    if (!srsIntrOp) {
      op.emitWarning() << "aievec.srs is not supported.\n";
      return failure();
    }
 
    // create bitcast/shape_cast for result if needed
    auto resultVal = forceCastValueToType(rewriter, loc, srsIntrOp,
                                          op.getResult().getType());
    rewriter.replaceOp(op, resultVal);
 
    return success();
  }
};
 
class UPDOpConversion : public mlir::ConvertOpToLLVMPattern<aievec::UPDOp> {
public:
  using ConvertOpToLLVMPattern<aievec::UPDOp>::ConvertOpToLLVMPattern;
 
  static std::string getIntrinsicName(aievec::UPDOp op, int loadSize) {
    auto resultType = cast<VectorType>(op.getResult().getType());
    std::stringstream ss;
    ss << "llvm.aie.upd.";
    ss << (loadSize == 128 ? 'v' : loadSize == 256 ? 'w' : 'x') << ".";
    ss << getVectorTypeString(resultType) << ".";
    // The index affects which intrinsic to call
    ss << (op.getIndex() == 0 ? "lo" : "hi");
    return ss.str();
  }
 
  LogicalResult
  matchAndRewrite(aievec::UPDOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto module = op->getParentOfType<ModuleOp>();
    MLIRContext *context = rewriter.getContext();
 
    // A bit more complicated: load the vector, then update result vector
    // AIE1 is capable of 128-bit on one bank and 256-bit loads on even-odd
    // banks Identify size of update
    int vecSizeInBits =
        getVectorSizeInBits(cast<VectorType>(op.getResult().getType()));
 
    auto ptr = this->getStridedElementPtr(
        rewriter, op->getLoc(), cast<MemRefType>(op.getSource().getType()),
        adaptor.getSource(), adaptor.getIndices());
 
    // TODO: handle the offset field
 
    if (vecSizeInBits <= 256) {
      // Total <=256-bit updates are much simpler:
      // we can do a direct load into the vector register
      // look at the indices to calculate the address
      auto vectorPtrType = LLVM::LLVMPointerType::get(
          getContext(),
          cast<MemRefType>(op.getSource().getType()).getMemorySpaceAsInt());
      auto castedPtr =
          LLVM::BitcastOp::create(rewriter, op->getLoc(), vectorPtrType, ptr);
      auto vecType = cast<VectorType>(op.getResult().getType());
      rewriter.replaceOpWithNewOp<LLVM::LoadOp>(op, vecType, castedPtr, 1);
    } else {
      // Total >256-bit updates will require upd ops to fill the whole vector
      // each UDP op represents one of these 256-bit loads and updates
 
      // Determine the load size
      // TODO: no examples of 1024-bit output vectors: doesn't feel right
      // to attempt a 512-bit load to do an update like this
      int loadSize = vecSizeInBits == 256   ? 128
                     : vecSizeInBits == 512 ? 256
                                            : 512;
 
      // Create a vectorType for the load proper
      // Load half of the final result vector
      auto resultType = cast<VectorType>(op.getResult().getType());
      int lanes = getVectorLaneSize(resultType);
      auto loadType =
          VectorType::get({(int64_t)lanes / 2}, resultType.getElementType());
 
      // Load the vector
      auto vectorPtrType = LLVM::LLVMPointerType::get(
          getContext(),
          cast<MemRefType>(op.getSource().getType()).getMemorySpaceAsInt());
      auto castedPtr =
          LLVM::BitcastOp::create(rewriter, op->getLoc(), vectorPtrType, ptr);
      auto loadValue =
          LLVM::LoadOp::create(rewriter, op->getLoc(), loadType, castedPtr, 1);
 
      // Get set up for the intrinsic
      std::string intrinsicName = getIntrinsicName(op, loadSize);
 
      // If the intrinsic declaration doesn't exist, create it
      auto func = module.lookupSymbol<LLVM::LLVMFuncOp>(
          StringAttr::get(context, intrinsicName));
 
      if (!func) {
        OpBuilder::InsertionGuard guard(rewriter);
        rewriter.setInsertionPointToStart(module.getBody());
        func = LLVM::LLVMFuncOp::create(
            rewriter, rewriter.getUnknownLoc(), intrinsicName,
            LLVM::LLVMFunctionType::get(resultType, {resultType, loadType}));
      }
 
      // Determine what the destination is
      Value destValue;
      if (adaptor.getVector()) {
        // This UPD is using an existing destination vector
        destValue = adaptor.getVector();
      } else {
        // If this UPD is not working off of an existing destination vector,
        // create an undefined vector as the destination
 
        // TODO: determine if the undef intrinsic is needed or if an LLVM
        // undef suffices destValue =
        // LLVM::UndefOp::create(rewriter, op->getLoc(), resultType);
 
        std::stringstream ss;
        ss << "llvm.aie." << getVectorTypeString(resultType) << ".undef";
        std::string intrinsicName = ss.str();
 
        auto func = module.lookupSymbol<LLVM::LLVMFuncOp>(
            StringAttr::get(rewriter.getContext(), intrinsicName));
 
        if (!func) {
          OpBuilder::InsertionGuard guard(rewriter);
          rewriter.setInsertionPointToStart(module.getBody());
          func = LLVM::LLVMFuncOp::create(
              rewriter, rewriter.getUnknownLoc(), intrinsicName,
              LLVM::LLVMFunctionType::get(resultType, {}));
        }
        destValue =
            LLVM::CallOp::create(rewriter, op->getLoc(), func, ValueRange{})
                ->getOpResult(0);
      }
 
      // Create our call
      rewriter.replaceOpWithNewOp<LLVM::CallOp>(
          op, func, ValueRange{destValue, loadValue});
    }
 
    return success();
  }
};
 
class ConcatOpConversion
    : public mlir::ConvertOpToLLVMPattern<aievec::ConcatOp> {
public:
  using ConvertOpToLLVMPattern<aievec::ConcatOp>::ConvertOpToLLVMPattern;
 
  LogicalResult
  matchAndRewrite(aievec::ConcatOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Location loc = op.getLoc();
 
    SmallVector<Value> sources = adaptor.getSources();
    Value src = sources.front();
    VectorType srcType = cast<VectorType>(src.getType());
    Type srcScalarType = srcType.getElementType();
    unsigned srcBitWidth = srcScalarType.getIntOrFloatBitWidth();
    int srcLanes = getVectorLaneSize(srcType);
    int srcVectorSize = srcBitWidth * srcLanes;
 
    Value result = op.getResult();
    VectorType resultType = cast<VectorType>(result.getType());
    Type resultScaTy = resultType.getElementType();
    unsigned resultBitWidth = resultScaTy.getIntOrFloatBitWidth();
    int resultLanes = getVectorLaneSize(resultType);
    int resultVectorSize = resultBitWidth * resultLanes;
 
    if (sources.size() != 2 && sources.size() != 4) {
      op.emitWarning() << "aievec.concat with " << sources.size()
                       << " operands is not supported.\n";
      return failure();
    }
 
    // create xllvm intrinsic
    Value concatOp = nullptr;
    if (srcVectorSize == 256 && resultVectorSize == 512) {
      concatOp = xllvm::ConcatI512I256IntrOp::create(
          rewriter, loc, VectorType::get({16}, rewriter.getI32Type()),
          forceCastOperandsToSignature(
              rewriter, loc, adaptor.getSources(),
              {VectorType::get({8}, rewriter.getI32Type()),
               VectorType::get({8}, rewriter.getI32Type())}));
    } else if (srcVectorSize == 256 && resultVectorSize == 1024) {
      concatOp = xllvm::ConcatI1024I256IntrOp::create(
          rewriter, loc, VectorType::get({32}, rewriter.getI32Type()),
          forceCastOperandsToSignature(
              rewriter, loc, adaptor.getSources(),
              {VectorType::get({8}, rewriter.getI32Type()),
               VectorType::get({8}, rewriter.getI32Type()),
               VectorType::get({8}, rewriter.getI32Type()),
               VectorType::get({8}, rewriter.getI32Type())}));
    } else if (srcVectorSize == 512 && resultVectorSize == 1024) {
      concatOp = xllvm::ConcatI1024I512IntrOp::create(
          rewriter, loc, VectorType::get({32}, rewriter.getI32Type()),
          forceCastOperandsToSignature(
              rewriter, loc, adaptor.getSources(),
              {VectorType::get({16}, rewriter.getI32Type()),
               VectorType::get({16}, rewriter.getI32Type())}));
    } else {
      op.emitWarning() << "aievec.concat with " << srcVectorSize
                       << "-bit operands, and " << resultVectorSize
                       << "-bit result is not supported.\n";
      return failure();
    }
 
    // create bitcast/shape_cast for result
    auto resultVal =
        forceCastValueToType(rewriter, loc, concatOp, op.getResult().getType());
    rewriter.replaceOp(op, resultVal);
 
    return success();
  }
};
 
class ExtOpConversion : public mlir::ConvertOpToLLVMPattern<aievec::ExtOp> {
public:
  using ConvertOpToLLVMPattern<aievec::ExtOp>::ConvertOpToLLVMPattern;
 
  LogicalResult
  matchAndRewrite(aievec::ExtOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Location loc = op.getLoc();
 
    Value src = adaptor.getSource();
    VectorType srcType = cast<VectorType>(src.getType());
    Type srcScalarType = srcType.getElementType();
    unsigned srcBitWidth = srcScalarType.getIntOrFloatBitWidth();
    int srcLanes = getVectorLaneSize(srcType);
    int srcVectorSize = srcBitWidth * srcLanes;
 
    Value result = op.getResult();
    VectorType resultType = cast<VectorType>(result.getType());
    Type resultScaTy = resultType.getElementType();
    unsigned resultBitWidth = resultScaTy.getIntOrFloatBitWidth();
    int resultLanes = getVectorLaneSize(resultType);
    int resultVectorSize = resultBitWidth * resultLanes;
 
    // create constant for index
    auto indexCst =
        LLVM::ConstantOp::create(rewriter, loc, rewriter.getI32Type(),
                                 rewriter.getI32IntegerAttr(op.getIndex()));
 
    // create xllvm intrinsic
    SmallVector<Value> operands({adaptor.getSource(), indexCst});
    Value extOp = nullptr;
    // Integer types
    if (resultVectorSize == 256 && srcVectorSize == 512) {
      extOp = xllvm::ExtI256I512IntrOp::create(
          rewriter, loc, VectorType::get({8}, rewriter.getI32Type()),
          forceCastOperandsToSignature(
              rewriter, loc, operands,
              {VectorType::get({16}, rewriter.getI32Type()),
               rewriter.getI32Type()}));
    } else if (resultVectorSize == 512 && srcVectorSize == 1024) {
      extOp = xllvm::ExtI512I1024IntrOp::create(
          rewriter, loc, VectorType::get({16}, rewriter.getI32Type()),
          forceCastOperandsToSignature(
              rewriter, loc, operands,
              {VectorType::get({32}, rewriter.getI32Type()),
               rewriter.getI32Type()}));
    } else if (resultVectorSize == 256 && srcVectorSize == 1024) {
      extOp = xllvm::ExtI256I1024IntrOp::create(
          rewriter, loc, VectorType::get({8}, rewriter.getI32Type()),
          forceCastOperandsToSignature(
              rewriter, loc, operands,
              {VectorType::get({32}, rewriter.getI32Type()),
               rewriter.getI32Type()}));
    } else if (resultVectorSize == 128 && srcVectorSize == 512) {
      auto shiftOp = adaptor.getSource();
      if (op.getIndex() > 0) {
        auto undefOp = xllvm::UndefV16I32IntrOp::create(
            rewriter, loc, VectorType::get({16}, rewriter.getI32Type()));
        auto stepCst =
            LLVM::ConstantOp::create(rewriter, loc, rewriter.getI32Type(),
                                     rewriter.getI32IntegerAttr(0));
        auto shiftCst = LLVM::ConstantOp::create(
            rewriter, loc, rewriter.getI32Type(),
            rewriter.getI32IntegerAttr(op.getIndex() * 16));
        SmallVector<Value> shiftOperands{adaptor.getSource(), undefOp, stepCst,
                                         shiftCst};
        // Right shift the source vector in index * 16 bytes (i.e. in index *
        // 128 bits). The integer index is expected to be 0 to 3.
        shiftOp = xllvm::VectorShiftI512I512IntrOp::create(
            rewriter, loc, VectorType::get({16}, rewriter.getI32Type()),
            forceCastOperandsToSignature(
                rewriter, loc, shiftOperands,
                {VectorType::get({16}, rewriter.getI32Type()),
                 VectorType::get({16}, rewriter.getI32Type()),
                 rewriter.getI32Type(), rewriter.getI32Type()}));
      }
      // The underlying intrinsic takes a source vector and extract the lowest
      // 128-bit. i.e. it always extracts the input vector with index = 0.
      extOp = xllvm::ExtI128I512IntrOp::create(
          rewriter, loc, VectorType::get({4}, rewriter.getI32Type()),
          forceCastOperandsToSignature(
              rewriter, loc, /*operands=*/{shiftOp},
              {VectorType::get({16}, rewriter.getI32Type())}));
    } else {
      op.emitWarning() << "aievec.ext with " << srcVectorSize
                       << "-bit source, and " << resultVectorSize
                       << "-bit result is not supported.\n";
      return failure();
    }
 
    // create bitcast/shape_cast for result
    auto resultVal =
        forceCastValueToType(rewriter, loc, extOp, op.getResult().getType());
    rewriter.replaceOp(op, resultVal);
 
    return success();
  }
};
 
// AIE2p version of ExtOp conversion using vector.shuffle
class ExtOpAIE2pConversion
    : public mlir::ConvertOpToLLVMPattern<aievec::ExtOp> {
public:
  using ConvertOpToLLVMPattern<aievec::ExtOp>::ConvertOpToLLVMPattern;
 
  LogicalResult
  matchAndRewrite(aievec::ExtOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Location loc = op.getLoc();
 
    Value src = adaptor.getSource();
    VectorType srcType = cast<VectorType>(src.getType());
    VectorType resultType = cast<VectorType>(op.getResult().getType());
 
    int srcLanes = getVectorLaneSize(srcType);
    int resultLanes = getVectorLaneSize(resultType);
 
    // The result is one of the equal parts the source divides into.
    if (srcLanes % resultLanes != 0) {
      op.emitWarning() << "aievec.ext of a part that does not divide the "
                          "source is not supported for AIE2p.\n";
      return failure();
    }
 
    // Part `index`: [index * resultLanes, (index + 1) * resultLanes)
    SmallVector<int64_t> shuffleMask;
    int startIdx = op.getIndex() * resultLanes;
    for (int i = 0; i < resultLanes; ++i) {
      shuffleMask.push_back(startIdx + i);
    }
 
    auto extracted =
        vector::ShuffleOp::create(rewriter, loc, src, src, shuffleMask);
 
    rewriter.replaceOp(op, extracted);
    return success();
  }
};
 
class SelectOpConversion
    : public mlir::ConvertOpToLLVMPattern<aievec::aie1::SelectOp> {
public:
  using ConvertOpToLLVMPattern<aievec::aie1::SelectOp>::ConvertOpToLLVMPattern;
 
  static std::string getIntrinsicName(aievec::aie1::SelectOp op) {
    auto xbuffType = cast<VectorType>(op.getXbuff().getType());
    std::stringstream ss;
    ss << "llvm.aie.prim." << getVectorTypeString(xbuffType) << ".select";
    return ss.str();
  }
 
  LogicalResult
  matchAndRewrite(aievec::aie1::SelectOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto module = op->getParentOfType<ModuleOp>();
    MLIRContext *context = rewriter.getContext();
 
    auto selectType = IntegerType::get(context, 32);
    auto startType = IntegerType::get(context, 32);
    auto offsetsType = VectorType::get({2}, IntegerType::get(context, 32));
    auto confType = VectorType::get({2}, IntegerType::get(context, 32));
 
    // If the intrinsic declaration doesn't exist, create it
    std::string intrinsicName = getIntrinsicName(op);
    auto func = module.lookupSymbol<LLVM::LLVMFuncOp>(
        StringAttr::get(context, intrinsicName));
 
    if (!func) {
      OpBuilder::InsertionGuard guard(rewriter);
      rewriter.setInsertionPointToStart(module.getBody());
      func = LLVM::LLVMFuncOp::create(
          rewriter, rewriter.getUnknownLoc(), intrinsicName,
          LLVM::LLVMFunctionType::get(op.getResult().getType(),
                                      {op.getXbuff().getType(), selectType,
                                       startType,   /* xstart */
                                       startType,   /* ystart */
                                       offsetsType, /* xoffsets */
                                       offsetsType, /* yoffsets */
                                       confType}));
    }
 
    // Parse the string attribute values
    uint32_t select = 0;
    BufferParams x = {};
    BufferParams y = {};
    BufferParams z = {};
 
    op.getSelect().getAsInteger(0, select);
    op.getXstart().getAsInteger(0, x.start);
    op.getXoffsets().getAsInteger(0, x.offsets);
    op.getXoffsetsHi().getAsInteger(0, x.offsets_hi);
    op.getXsquare().getAsInteger(0, x.square);
    op.getYstart().getAsInteger(0, y.start);
    op.getYoffsets().getAsInteger(0, y.offsets);
    op.getYoffsetsHi().getAsInteger(0, y.offsets_hi);
    op.getYsquare().getAsInteger(0, y.square);
 
    // Encode the configuration register
    uint32_t conf[2] = {0, 0};
    encodeConf(conf, x, z, false);
    conf[1] |= encodeSquare(y.square) << 21;
 
    // Create the constants and replace the op
    auto selectVal = LLVM::ConstantOp::create(
        rewriter, op->getLoc(), selectType, rewriter.getI32IntegerAttr(select));
    auto xstartVal = LLVM::ConstantOp::create(
        rewriter, op->getLoc(), startType, rewriter.getI32IntegerAttr(x.start));
    auto ystartVal = LLVM::ConstantOp::create(
        rewriter, op->getLoc(), startType, rewriter.getI32IntegerAttr(y.start));
    auto xoffsetsVal = LLVM::ConstantOp::create(
        rewriter, op->getLoc(), offsetsType,
        rewriter.getI32VectorAttr({(int32_t)x.offsets, (int32_t)x.offsets_hi}));
    auto yoffsetsVal = LLVM::ConstantOp::create(
        rewriter, op->getLoc(), offsetsType,
        rewriter.getI32VectorAttr({(int32_t)y.offsets, (int32_t)y.offsets_hi}));
    auto confVal = LLVM::ConstantOp::create(
        rewriter, op->getLoc(), confType,
        rewriter.getI32VectorAttr({(int32_t)conf[0], (int32_t)conf[1]}));
    rewriter.replaceOpWithNewOp<LLVM::CallOp>(
        op, func,
        ValueRange{op.getXbuff(), selectVal, xstartVal, ystartVal, xoffsetsVal,
                   yoffsetsVal, confVal});
    return success();
  }
};
 
class PackOpConversion : public mlir::ConvertOpToLLVMPattern<aievec::PackOp> {
public:
  using ConvertOpToLLVMPattern<aievec::PackOp>::ConvertOpToLLVMPattern;
 
  static std::string getIntrinsicName(aievec::PackOp op) {
    auto sourceType = cast<VectorType>(op.getSource().getType());
    std::stringstream ss;
    ss << "llvm.aie.pack." << getVectorTypeString(sourceType);
    return ss.str();
  }
 
  LogicalResult
  matchAndRewrite(aievec::PackOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto module = op->getParentOfType<ModuleOp>();
    MLIRContext *context = rewriter.getContext();
 
    // If the intrinsic declaration doesn't exist, create it
    std::string intrinsicName = getIntrinsicName(op);
    auto func = module.lookupSymbol<LLVM::LLVMFuncOp>(
        StringAttr::get(context, intrinsicName));
 
    if (!func) {
      OpBuilder::InsertionGuard guard(rewriter);
      rewriter.setInsertionPointToStart(module.getBody());
      func = LLVM::LLVMFuncOp::create(
          rewriter, rewriter.getUnknownLoc(), intrinsicName,
          LLVM::LLVMFunctionType::get(op.getResult().getType(),
                                      {op.getSource().getType()}));
    }
 
    rewriter.replaceOpWithNewOp<LLVM::CallOp>(op, func,
                                              ValueRange{op.getSource()});
    return success();
  }
};
 
class UnpackOpConversion
    : public mlir::ConvertOpToLLVMPattern<aievec::UnpackOp> {
public:
  using ConvertOpToLLVMPattern<aievec::UnpackOp>::ConvertOpToLLVMPattern;
 
  LogicalResult
  matchAndRewrite(aievec::UnpackOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    op.emitWarning() << "aie.unpack conversion is not implemented\n";
    return failure();
  }
};
 
// AIE2P unpack int4-packed: source is vector<N/2 x i8> carrying N packed
// nibbles, result is vector<N x i8> (one byte per nibble). Maps directly
// to llvm.aie2p.unpack.I512.I8.I4 (N=64) or .I1024.I8.I4 (N=128).
// AWQ packed weights are unsigned; the signed offset is folded into the
// per-group zero point so the intrinsic is invoked with sign=0.
class UnpackI4OpAIE2pConversion
    : public mlir::ConvertOpToLLVMPattern<aievec::UnpackOp> {
public:
  using ConvertOpToLLVMPattern<aievec::UnpackOp>::ConvertOpToLLVMPattern;
 
  LogicalResult
  matchAndRewrite(aievec::UnpackOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Location loc = op.getLoc();
    auto srcTy = cast<VectorType>(op.getSource().getType());
    auto dstTy = cast<VectorType>(op.getResult().getType());
    if (!srcTy.getElementType().isInteger(8) ||
        !dstTy.getElementType().isInteger(8))
      return failure();
 
    unsigned srcLanes = getVectorLaneSize(srcTy);
    unsigned dstLanes = getVectorLaneSize(dstTy);
    if (dstLanes != 2 * srcLanes)
      return failure();
    if (dstLanes != 64 && dstLanes != 128)
      return failure();
 
    auto signCst = LLVM::ConstantOp::create(
        rewriter, loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(0));
 
    Value unpacked;
    if (dstLanes == 64) {
      unpacked = xllvm::UnpackI512I8I4AIE2pIntrOp::create(
                     rewriter, loc, VectorType::get({64}, rewriter.getI8Type()),
                     adaptor.getSource(), signCst)
                     ->getResult(0);
    } else {
      unpacked =
          xllvm::UnpackI1024I8I4AIE2pIntrOp::create(
              rewriter, loc, VectorType::get({128}, rewriter.getI8Type()),
              adaptor.getSource(), signCst)
              ->getResult(0);
    }
    rewriter.replaceOp(op, unpacked);
    return success();
  }
};
 
class BroadcastOpConversion
    : public mlir::ConvertOpToLLVMPattern<aievec::BroadcastOp> {
public:
  using ConvertOpToLLVMPattern<aievec::BroadcastOp>::ConvertOpToLLVMPattern;
 
  LogicalResult
  matchAndRewrite(aievec::BroadcastOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    op.emitWarning() << "aie.broadcast conversion is not implemented\n";
    return failure();
  }
};
 
// Helper to pad a vector from N lanes to dstLanes using vector.shuffle
// with poison (-1) for the upper portion.
static Value padVectorWithPoison(ConversionPatternRewriter &rewriter,
                                 Location loc, Value vec, int srcLanes,
                                 int dstLanes) {
  SmallVector<int64_t> padMask;
  for (int i = 0; i < srcLanes; ++i)
    padMask.push_back(i);
  for (int i = srcLanes; i < dstLanes; ++i)
    padMask.push_back(-1);
  return vector::ShuffleOp::create(rewriter, loc, vec, vec, padMask);
}
 
// Helper to extract the first N lanes from a wider vector using
// vector.shuffle.
static Value extractLowerLanes(ConversionPatternRewriter &rewriter,
                               Location loc, Value vec, int lanes) {
  SmallVector<int64_t> extractMask;
  for (int i = 0; i < lanes; ++i)
    extractMask.push_back(i);
  return vector::ShuffleOp::create(rewriter, loc, vec, vec, extractMask);
}
 
class MaxOpConversion : public mlir::ConvertOpToLLVMPattern<aievec::MaxOp> {
public:
  using ConvertOpToLLVMPattern<aievec::MaxOp>::ConvertOpToLLVMPattern;
 
  LogicalResult
  matchAndRewrite(aievec::MaxOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Location loc = op.getLoc();
 
    VectorType resultType = cast<VectorType>(op.getResult().getType());
    Type resultScaTy = resultType.getElementType();
    unsigned resultBitWidth = resultScaTy.getIntOrFloatBitWidth();
    int resultLanes = getVectorLaneSize(resultType);
    int resultVectorSize = resultBitWidth * resultLanes;
 
    // aievec.max op has the AllTypesMatch constraint on lhs/rhs/res
    if (resultVectorSize != 512 && resultVectorSize != 256) {
      op.emitWarning() << "aievec.max conversion with " << resultVectorSize
                       << "-bit result is not supported.\n";
      return failure();
    }
 
    // create xllvm intrinsic
    Value maxOp = nullptr;
    if (llvm::isa<IntegerType>(resultScaTy)) {
      // create constant for third operand `cmp`
      // Note: `cmp` is implicitly treated as `sign` to the vmax intrinsic
      auto cmpCst = LLVM::ConstantOp::create(
          rewriter, loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(1));
      SmallVector<Value> operands{adaptor.getLhs(), adaptor.getRhs(), cmpCst};
      if (resultBitWidth == 8) {
        maxOp = xllvm::VectorMaxLt8IntrOp::create(
            rewriter, loc,
            mlir::LLVM::LLVMStructType::getLiteral(
                rewriter.getContext(),
                {VectorType::get({64}, rewriter.getI8Type()),
                 VectorType::get({2}, rewriter.getI32Type())}),
            forceCastOperandsToSignature(
                rewriter, loc, operands,
                {VectorType::get({64}, rewriter.getI8Type()),
                 VectorType::get({64}, rewriter.getI8Type()),
                 rewriter.getI32Type()}));
      } else if (resultBitWidth == 16) {
        maxOp = xllvm::VectorMaxLt16IntrOp::create(
            rewriter, loc,
            mlir::LLVM::LLVMStructType::getLiteral(
                rewriter.getContext(),
                {VectorType::get({32}, rewriter.getI16Type()),
                 rewriter.getI32Type()}),
            forceCastOperandsToSignature(
                rewriter, loc, operands,
                {VectorType::get({32}, rewriter.getI16Type()),
                 VectorType::get({32}, rewriter.getI16Type()),
                 rewriter.getI32Type()}));
      } else if (resultBitWidth == 32) {
        maxOp = xllvm::VectorMaxLt32IntrOp::create(
            rewriter, loc,
            mlir::LLVM::LLVMStructType::getLiteral(
                rewriter.getContext(),
                {VectorType::get({16}, rewriter.getI32Type()),
                 rewriter.getI32Type()}),
            forceCastOperandsToSignature(
                rewriter, loc, operands,
                {VectorType::get({16}, rewriter.getI32Type()),
                 VectorType::get({16}, rewriter.getI32Type()),
                 rewriter.getI32Type()}));
      }
    } else {
      if (resultBitWidth == 16) {
        auto v32bf16Ty = VectorType::get({32}, rewriter.getBF16Type());
        Value lhs = adaptor.getLhs(), rhs = adaptor.getRhs();
 
        // Pad 16-lane to 32-lane if needed
        if (resultLanes == 16) {
          lhs = padVectorWithPoison(rewriter, loc, lhs, 16, 32);
          rhs = padVectorWithPoison(rewriter, loc, rhs, 16, 32);
        }
 
        maxOp = xllvm::VectorMaxLtBf16IntrOp::create(
            rewriter, loc,
            mlir::LLVM::LLVMStructType::getLiteral(
                rewriter.getContext(), {v32bf16Ty, rewriter.getI32Type()}),
            forceCastOperandsToSignature(rewriter, loc, {lhs, rhs},
                                         {v32bf16Ty, v32bf16Ty}));
      }
    }
 
    if (!maxOp) {
      op.emitWarning() << "aievec.max conversion fails due to unsupported "
                          "element data type.\n";
      return failure();
    }
 
    // Extract the vector result from the struct
    Value resultVec = LLVM::ExtractValueOp::create(rewriter, loc, maxOp,
                                                   /*position=*/0);
    // Truncate back to 16 lanes if padded
    if (resultLanes == 16 && !llvm::isa<IntegerType>(resultScaTy))
      resultVec = extractLowerLanes(rewriter, loc, resultVec, 16);
 
    rewriter.replaceOp(op, resultVec);
 
    return success();
  }
};
 
class MinOpConversion : public mlir::ConvertOpToLLVMPattern<aievec::MinOp> {
public:
  using ConvertOpToLLVMPattern<aievec::MinOp>::ConvertOpToLLVMPattern;
 
  LogicalResult
  matchAndRewrite(aievec::MinOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Location loc = op.getLoc();
 
    VectorType resultType = cast<VectorType>(op.getResult().getType());
    Type resultScaTy = resultType.getElementType();
    unsigned resultBitWidth = resultScaTy.getIntOrFloatBitWidth();
    int resultLanes = getVectorLaneSize(resultType);
    int resultVectorSize = resultBitWidth * resultLanes;
 
    // aievec.min op has the AllTypesMatch constraint on lhs/rhs/res
    if (resultVectorSize != 512 && resultVectorSize != 256) {
      op.emitWarning() << "aievec.min conversion with " << resultVectorSize
                       << "-bit result is not supported.\n";
      return failure();
    }
 
    // create xllvm intrinsic
    Value minOp = nullptr;
    if (llvm::isa<IntegerType>(resultScaTy)) {
      // create constant for third operand `cmp`
      // Note: `cmp` is implicitly treated as `sign` to the vmin intrinsic
      auto cmpCst = LLVM::ConstantOp::create(
          rewriter, loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(1));
      SmallVector<Value> operands{adaptor.getLhs(), adaptor.getRhs(), cmpCst};
      if (resultBitWidth == 8) {
        minOp = xllvm::VectorMinGe8IntrOp::create(
            rewriter, loc,
            mlir::LLVM::LLVMStructType::getLiteral(
                rewriter.getContext(),
                {VectorType::get({64}, rewriter.getI8Type()),
                 VectorType::get({2}, rewriter.getI32Type())}),
            forceCastOperandsToSignature(
                rewriter, loc, operands,
                {VectorType::get({64}, rewriter.getI8Type()),
                 VectorType::get({64}, rewriter.getI8Type()),
                 rewriter.getI32Type()}));
      } else if (resultBitWidth == 16) {
        minOp = xllvm::VectorMinGe16IntrOp::create(
            rewriter, loc,
            mlir::LLVM::LLVMStructType::getLiteral(
                rewriter.getContext(),
                {VectorType::get({32}, rewriter.getI16Type()),
                 rewriter.getI32Type()}),
            forceCastOperandsToSignature(
                rewriter, loc, operands,
                {VectorType::get({32}, rewriter.getI16Type()),
                 VectorType::get({32}, rewriter.getI16Type()),
                 rewriter.getI32Type()}));
      } else if (resultBitWidth == 32) {
        minOp = xllvm::VectorMinGe32IntrOp::create(
            rewriter, loc,
            mlir::LLVM::LLVMStructType::getLiteral(
                rewriter.getContext(),
                {VectorType::get({16}, rewriter.getI32Type()),
                 rewriter.getI32Type()}),
            forceCastOperandsToSignature(
                rewriter, loc, operands,
                {VectorType::get({16}, rewriter.getI32Type()),
                 VectorType::get({16}, rewriter.getI32Type()),
                 rewriter.getI32Type()}));
      }
    } else {
      if (resultBitWidth == 16) {
        auto v32bf16Ty = VectorType::get({32}, rewriter.getBF16Type());
        Value lhs = adaptor.getLhs(), rhs = adaptor.getRhs();
 
        // Pad 16-lane to 32-lane if needed
        if (resultLanes == 16) {
          lhs = padVectorWithPoison(rewriter, loc, lhs, 16, 32);
          rhs = padVectorWithPoison(rewriter, loc, rhs, 16, 32);
        }
 
        minOp = xllvm::VectorMinGeBf16IntrOp::create(
            rewriter, loc,
            mlir::LLVM::LLVMStructType::getLiteral(
                rewriter.getContext(), {v32bf16Ty, rewriter.getI32Type()}),
            forceCastOperandsToSignature(rewriter, loc, {lhs, rhs},
                                         {v32bf16Ty, v32bf16Ty}));
      }
    }
 
    if (!minOp) {
      op.emitWarning() << "aievec.min conversion fails due to unsupported "
                          "element data type.\n";
      return failure();
    }
 
    // Extract the vector result from the struct
    Value resultVec = LLVM::ExtractValueOp::create(rewriter, loc, minOp,
                                                   /*position=*/0);
    // Truncate back to 16 lanes if padded
    if (resultLanes == 16 && !llvm::isa<IntegerType>(resultScaTy))
      resultVec = extractLowerLanes(rewriter, loc, resultVec, 16);
 
    rewriter.replaceOp(op, resultVec);
 
    return success();
  }
};
 
// AIE2p version of MaxOp conversion
class MaxOpAIE2pConversion
    : public mlir::ConvertOpToLLVMPattern<aievec::MaxOp> {
public:
  using ConvertOpToLLVMPattern<aievec::MaxOp>::ConvertOpToLLVMPattern;
 
  LogicalResult
  matchAndRewrite(aievec::MaxOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Location loc = op.getLoc();
 
    VectorType resultType = cast<VectorType>(op.getResult().getType());
    Type resultScaTy = resultType.getElementType();
    unsigned resultBitWidth = resultScaTy.getIntOrFloatBitWidth();
    int resultLanes = getVectorLaneSize(resultType);
    int resultVectorSize = resultBitWidth * resultLanes;
 
    // aievec.max op has the AllTypesMatch constraint on lhs/rhs/res
    if (resultVectorSize != 512 && resultVectorSize != 256) {
      op.emitWarning() << "aievec.max conversion with " << resultVectorSize
                       << "-bit result is not supported.\n";
      return failure();
    }
 
    // create xllvm intrinsic
    Value maxOp = nullptr;
    if (llvm::isa<IntegerType>(resultScaTy)) {
      // create constant for third operand `cmp`
      // Note: `cmp` is implicitly treated as `sign` to the vmax intrinsic
      auto cmpCst = LLVM::ConstantOp::create(
          rewriter, loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(1));
      SmallVector<Value> operands{adaptor.getLhs(), adaptor.getRhs(), cmpCst};
      if (resultBitWidth == 8) {
        maxOp = xllvm::VectorMaxLt8AIE2pIntrOp::create(
            rewriter, loc,
            mlir::LLVM::LLVMStructType::getLiteral(
                rewriter.getContext(),
                {VectorType::get({64}, rewriter.getI8Type()),
                 VectorType::get({2}, rewriter.getI32Type())}),
            forceCastOperandsToSignature(
                rewriter, loc, operands,
                {VectorType::get({64}, rewriter.getI8Type()),
                 VectorType::get({64}, rewriter.getI8Type()),
                 rewriter.getI32Type()}));
      } else if (resultBitWidth == 16) {
        maxOp = xllvm::VectorMaxLt16AIE2pIntrOp::create(
            rewriter, loc,
            mlir::LLVM::LLVMStructType::getLiteral(
                rewriter.getContext(),
                {VectorType::get({32}, rewriter.getI16Type()),
                 rewriter.getI32Type()}),
            forceCastOperandsToSignature(
                rewriter, loc, operands,
                {VectorType::get({32}, rewriter.getI16Type()),
                 VectorType::get({32}, rewriter.getI16Type()),
                 rewriter.getI32Type()}));
      } else if (resultBitWidth == 32) {
        maxOp = xllvm::VectorMaxLt32AIE2pIntrOp::create(
            rewriter, loc,
            mlir::LLVM::LLVMStructType::getLiteral(
                rewriter.getContext(),
                {VectorType::get({16}, rewriter.getI32Type()),
                 rewriter.getI32Type()}),
            forceCastOperandsToSignature(
                rewriter, loc, operands,
                {VectorType::get({16}, rewriter.getI32Type()),
                 VectorType::get({16}, rewriter.getI32Type()),
                 rewriter.getI32Type()}));
      }
    } else {
      if (resultBitWidth == 16) {
        auto v32bf16Ty = VectorType::get({32}, rewriter.getBF16Type());
        Value lhs = adaptor.getLhs(), rhs = adaptor.getRhs();
 
        if (resultLanes == 16) {
          lhs = padVectorWithPoison(rewriter, loc, lhs, 16, 32);
          rhs = padVectorWithPoison(rewriter, loc, rhs, 16, 32);
        }
 
        maxOp = xllvm::VectorMaxLtBf16AIE2pIntrOp::create(
            rewriter, loc,
            mlir::LLVM::LLVMStructType::getLiteral(
                rewriter.getContext(), {v32bf16Ty, rewriter.getI32Type()}),
            forceCastOperandsToSignature(rewriter, loc, {lhs, rhs},
                                         {v32bf16Ty, v32bf16Ty}));
      }
    }
 
    if (!maxOp) {
      op.emitWarning() << "aievec.max conversion fails due to unsupported "
                          "element data type.\n";
      return failure();
    }
 
    Value resultVec = LLVM::ExtractValueOp::create(rewriter, loc, maxOp,
                                                   /*position=*/0);
    if (resultLanes == 16 && !llvm::isa<IntegerType>(resultScaTy))
      resultVec = extractLowerLanes(rewriter, loc, resultVec, 16);
 
    rewriter.replaceOp(op, resultVec);
 
    return success();
  }
};
 
// AIE2p version of MinOp conversion
class MinOpAIE2pConversion
    : public mlir::ConvertOpToLLVMPattern<aievec::MinOp> {
public:
  using ConvertOpToLLVMPattern<aievec::MinOp>::ConvertOpToLLVMPattern;
 
  LogicalResult
  matchAndRewrite(aievec::MinOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Location loc = op.getLoc();
 
    VectorType resultType = cast<VectorType>(op.getResult().getType());
    Type resultScaTy = resultType.getElementType();
    unsigned resultBitWidth = resultScaTy.getIntOrFloatBitWidth();
    int resultLanes = getVectorLaneSize(resultType);
    int resultVectorSize = resultBitWidth * resultLanes;
 
    // aievec.min op has the AllTypesMatch constraint on lhs/rhs/res
    if (resultVectorSize != 512 && resultVectorSize != 256) {
      op.emitWarning() << "aievec.min conversion with " << resultVectorSize
                       << "-bit result is not supported.\n";
      return failure();
    }
 
    // create xllvm intrinsic
    Value minOp = nullptr;
    if (llvm::isa<IntegerType>(resultScaTy)) {
      // create constant for third operand `cmp`
      // Note: `cmp` is implicitly treated as `sign` to the vmin intrinsic
      auto cmpCst = LLVM::ConstantOp::create(
          rewriter, loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(1));
      SmallVector<Value> operands{adaptor.getLhs(), adaptor.getRhs(), cmpCst};
      if (resultBitWidth == 8) {
        minOp = xllvm::VectorMinGe8AIE2pIntrOp::create(
            rewriter, loc,
            mlir::LLVM::LLVMStructType::getLiteral(
                rewriter.getContext(),
                {VectorType::get({64}, rewriter.getI8Type()),
                 VectorType::get({2}, rewriter.getI32Type())}),
            forceCastOperandsToSignature(
                rewriter, loc, operands,
                {VectorType::get({64}, rewriter.getI8Type()),
                 VectorType::get({64}, rewriter.getI8Type()),
                 rewriter.getI32Type()}));
      } else if (resultBitWidth == 16) {
        minOp = xllvm::VectorMinGe16AIE2pIntrOp::create(
            rewriter, loc,
            mlir::LLVM::LLVMStructType::getLiteral(
                rewriter.getContext(),
                {VectorType::get({32}, rewriter.getI16Type()),
                 rewriter.getI32Type()}),
            forceCastOperandsToSignature(
                rewriter, loc, operands,
                {VectorType::get({32}, rewriter.getI16Type()),
                 VectorType::get({32}, rewriter.getI16Type()),
                 rewriter.getI32Type()}));
      } else if (resultBitWidth == 32) {
        minOp = xllvm::VectorMinGe32AIE2pIntrOp::create(
            rewriter, loc,
            mlir::LLVM::LLVMStructType::getLiteral(
                rewriter.getContext(),
                {VectorType::get({16}, rewriter.getI32Type()),
                 rewriter.getI32Type()}),
            forceCastOperandsToSignature(
                rewriter, loc, operands,
                {VectorType::get({16}, rewriter.getI32Type()),
                 VectorType::get({16}, rewriter.getI32Type()),
                 rewriter.getI32Type()}));
      }
    } else {
      if (resultBitWidth == 16) {
        auto v32bf16Ty = VectorType::get({32}, rewriter.getBF16Type());
        Value lhs = adaptor.getLhs(), rhs = adaptor.getRhs();
 
        if (resultLanes == 16) {
          lhs = padVectorWithPoison(rewriter, loc, lhs, 16, 32);
          rhs = padVectorWithPoison(rewriter, loc, rhs, 16, 32);
        }
 
        minOp = xllvm::VectorMinGeBf16AIE2pIntrOp::create(
            rewriter, loc,
            mlir::LLVM::LLVMStructType::getLiteral(
                rewriter.getContext(), {v32bf16Ty, rewriter.getI32Type()}),
            forceCastOperandsToSignature(rewriter, loc, {lhs, rhs},
                                         {v32bf16Ty, v32bf16Ty}));
      }
    }
 
    if (!minOp) {
      op.emitWarning() << "aievec.min conversion fails due to unsupported "
                          "element data type.\n";
      return failure();
    }
 
    Value resultVec = LLVM::ExtractValueOp::create(rewriter, loc, minOp,
                                                   /*position=*/0);
    if (resultLanes == 16 && !llvm::isa<IntegerType>(resultScaTy))
      resultVec = extractLowerLanes(rewriter, loc, resultVec, 16);
 
    rewriter.replaceOp(op, resultVec);
 
    return success();
  }
};
 
// The i32 bitmask of the lanes where `lhs rel rhs` holds, with `rel` one of
// lt, ge, gt, le, eq and ne. `LtOp` and `GeOp` are vmax.lt and vmin.ge
// intrinsics, which return the lanes where a < b and a >= b in field 1 of their
// struct result. gt and le swap the operands; eq and ne combine the compares in
// both directions. `extraOperands` follow the two vectors in each call.
template <typename LtOp, typename GeOp>
static Value buildCmpBitmask(OpBuilder &builder, Location loc, StringRef rel,
                             Value lhs, Value rhs, ValueRange extraOperands) {
  auto structTy = LLVM::LLVMStructType::getLiteral(
      builder.getContext(), {lhs.getType(), builder.getI32Type()});
  auto compare = [&](bool lessThan, Value a, Value b) -> Value {
    SmallVector<Value> operands{a, b};
    llvm::append_range(operands, extraOperands);
    Value result = lessThan
                       ? Value(LtOp::create(builder, loc, structTy, operands))
                       : Value(GeOp::create(builder, loc, structTy, operands));
    return LLVM::ExtractValueOp::create(builder, loc, result, 1);
  };
  if (rel == "lt")
    return compare(/*lessThan=*/true, lhs, rhs);
  if (rel == "ge")
    return compare(/*lessThan=*/false, lhs, rhs);
  if (rel == "gt")
    return compare(/*lessThan=*/true, rhs, lhs);
  if (rel == "le")
    return compare(/*lessThan=*/false, rhs, lhs);
  assert((rel == "eq" || rel == "ne") && "unexpected compare predicate");
  // Separate statements fix the order in which the two compares are emitted.
  bool isEq = rel == "eq";
  Value forward = compare(/*lessThan=*/!isEq, lhs, rhs);
  Value backward = compare(/*lessThan=*/!isEq, rhs, lhs);
  if (isEq)
    return LLVM::AndOp::create(builder, loc, forward, backward);
  return LLVM::OrOp::create(builder, loc, forward, backward);
}
 
// ----- CmpOp conversion for AIE2 and AIE2P -----
// Implements aievec.cmp using vmax.lt / vmin.ge intrinsics. The integer
// intrinsics take a last operand that selects a signed (1) or unsigned (0)
// compare, which follows the predicate.
template <typename MaxLtBf16IntrOp, typename MinGeBf16IntrOp,
          typename MaxLt32IntrOp, typename MinGe32IntrOp,
          typename MaxLt16IntrOp, typename MinGe16IntrOp>
class CmpOpConversionBase : public mlir::ConvertOpToLLVMPattern<aievec::CmpOp> {
public:
  using ConvertOpToLLVMPattern<aievec::CmpOp>::ConvertOpToLLVMPattern;
 
  LogicalResult
  matchAndRewrite(aievec::CmpOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Location loc = op.getLoc();
    auto vecTy = cast<VectorType>(op.getLhs().getType());
    auto elTy = vecTy.getElementType();
    unsigned elWidth = elTy.getIntOrFloatBitWidth();
    unsigned lanes = getVectorLaneSize(vecTy);
    StringRef pred = op.getPred();
 
    bool isUnsigned = llvm::is_contained({"ult", "uge", "ugt", "ule"}, pred);
    if (!isUnsigned &&
        !llvm::is_contained({"slt", "sge", "sgt", "sle", "eq", "ne"}, pred))
      return failure();
    // The relation without its signedness prefix: slt -> lt, uge -> ge.
    StringRef rel = pred.size() == 3 ? pred.drop_front() : pred;
 
    Value lhs = adaptor.getLhs(), rhs = adaptor.getRhs();
    Value bitmask;
    if (elWidth == 16 && isa<FloatType>(elTy) && (lanes == 16 || lanes == 32)) {
      // bf16 vectors (16 or 32 lanes); 16 lanes are padded to 32
      auto v32bf16Ty = VectorType::get({32}, rewriter.getBF16Type());
      if (lanes == 16) {
        lhs = padVectorWithPoison(rewriter, loc, lhs, 16, 32);
        rhs = padVectorWithPoison(rewriter, loc, rhs, 16, 32);
      }
      SmallVector<Value> operands = forceCastOperandsToSignature(
          rewriter, loc, {lhs, rhs}, {v32bf16Ty, v32bf16Ty});
      bitmask = buildCmpBitmask<MaxLtBf16IntrOp, MinGeBf16IntrOp>(
          rewriter, loc, rel, operands[0], operands[1], {});
      // Mask off upper bits for 16-lane inputs
      if (lanes == 16) {
        auto mask =
            LLVM::ConstantOp::create(rewriter, loc, rewriter.getI32Type(),
                                     rewriter.getI32IntegerAttr(0xFFFF));
        bitmask = LLVM::AndOp::create(rewriter, loc, bitmask, mask);
      }
    } else if (isa<IntegerType>(elTy) && elWidth * lanes == 512 &&
               (elWidth == 32 || elWidth == 16)) {
      // v16i32 and v32i16 vectors. eq and ne do not depend on the signedness
      // and use the signed form.
      auto intVecTy =
          VectorType::get({lanes}, rewriter.getIntegerType(elWidth));
      Value sign = LLVM::ConstantOp::create(
          rewriter, loc, rewriter.getI32Type(),
          rewriter.getI32IntegerAttr(isUnsigned ? 0 : 1));
      SmallVector<Value> operands = forceCastOperandsToSignature(
          rewriter, loc, {lhs, rhs}, {intVecTy, intVecTy});
      bitmask = elWidth == 32
                    ? buildCmpBitmask<MaxLt32IntrOp, MinGe32IntrOp>(
                          rewriter, loc, rel, operands[0], operands[1], sign)
                    : buildCmpBitmask<MaxLt16IntrOp, MinGe16IntrOp>(
                          rewriter, loc, rel, operands[0], operands[1], sign);
    } else {
      return failure();
    }
 
    // Cast i32 bitmask to the op's unsigned integer result type
    rewriter.replaceOpWithNewOp<UnrealizedConversionCastOp>(
        op, op.getResult().getType(), bitmask);
    return success();
  }
};
 
using CmpOpConversion =
    CmpOpConversionBase<xllvm::VectorMaxLtBf16IntrOp,
                        xllvm::VectorMinGeBf16IntrOp,
                        xllvm::VectorMaxLt32IntrOp, xllvm::VectorMinGe32IntrOp,
                        xllvm::VectorMaxLt16IntrOp, xllvm::VectorMinGe16IntrOp>;
 
using CmpOpAIE2pConversion = CmpOpConversionBase<
    xllvm::VectorMaxLtBf16AIE2pIntrOp, xllvm::VectorMinGeBf16AIE2pIntrOp,
    xllvm::VectorMaxLt32AIE2pIntrOp, xllvm::VectorMinGe32AIE2pIntrOp,
    xllvm::VectorMaxLt16AIE2pIntrOp, xllvm::VectorMinGe16AIE2pIntrOp>;
 
// ----- SelOp conversion for AIE2/AIE2p -----
// Implements aievec.sel using vsel16/vsel32 intrinsics.
// For bf16 vectors, bitcasts to i16, calls vsel16, bitcasts back.
template <typename Sel16IntrOp, typename Sel32IntrOp>
class SelOpConversionBase : public mlir::ConvertOpToLLVMPattern<aievec::SelOp> {
public:
  using ConvertOpToLLVMPattern<aievec::SelOp>::ConvertOpToLLVMPattern;
 
  LogicalResult
  matchAndRewrite(aievec::SelOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Location loc = op.getLoc();
    auto resultType = cast<VectorType>(op.getResult().getType());
    auto elTy = resultType.getElementType();
    unsigned elWidth = elTy.getIntOrFloatBitWidth();
    unsigned lanes = getVectorLaneSize(resultType);
    auto i32Ty = rewriter.getI32Type();
 
    // Cast the sel bitmask from unsigned integer to i32
    Value selMask = adaptor.getSel();
    if (selMask.getType() != i32Ty)
      selMask =
          UnrealizedConversionCastOp::create(rewriter, loc, i32Ty, selMask)
              .getResult(0);
 
    // Handle bf16 vectors (16 or 32 lanes) via bitcast to i16
    if (elWidth == 16 && isa<FloatType>(elTy)) {
      auto v32i16Ty = VectorType::get({32}, rewriter.getI16Type());
      auto v32bf16Ty = VectorType::get({32}, rewriter.getBF16Type());
 
      Value lhs = adaptor.getLhs(), rhs = adaptor.getRhs();
      bool needExtract = false;
 
      if (lanes == 16) {
        lhs = padVectorWithPoison(rewriter, loc, lhs, 16, 32);
        rhs = padVectorWithPoison(rewriter, loc, rhs, 16, 32);
        needExtract = true;
      } else if (lanes != 32) {
        return failure();
      }
 
      // Bitcast bf16 -> i16
      auto lhsCast = forceCastValueToType(rewriter, loc, lhs, v32bf16Ty);
      auto lhsI16 = LLVM::BitcastOp::create(rewriter, loc, v32i16Ty, lhsCast);
      auto rhsCast = forceCastValueToType(rewriter, loc, rhs, v32bf16Ty);
      auto rhsI16 = LLVM::BitcastOp::create(rewriter, loc, v32i16Ty, rhsCast);
 
      auto selResult = Sel16IntrOp::create(
          rewriter, loc, v32i16Ty,
          forceCastOperandsToSignature(rewriter, loc, {lhsI16, rhsI16, selMask},
                                       {v32i16Ty, v32i16Ty, i32Ty}));
 
      // Bitcast i16 -> bf16
      Value result =
          LLVM::BitcastOp::create(rewriter, loc, v32bf16Ty, selResult);
 
      if (needExtract)
        result = extractLowerLanes(rewriter, loc, result, 16);
 
      rewriter.replaceOp(op, result);
      return success();
    }
 
    // Handle i32 vectors (16 lanes)
    if (elWidth == 32 && isa<IntegerType>(elTy) && lanes == 16) {
      auto v16i32Ty = VectorType::get({16}, rewriter.getI32Type());
      auto selResult = Sel32IntrOp::create(
          rewriter, loc, v16i32Ty,
          forceCastOperandsToSignature(
              rewriter, loc, {adaptor.getLhs(), adaptor.getRhs(), selMask},
              {v16i32Ty, v16i32Ty, i32Ty}));
      rewriter.replaceOp(op, selResult->getResult(0));
      return success();
    }
 
    // Handle i16 vectors (32 lanes)
    if (elWidth == 16 && isa<IntegerType>(elTy) && lanes == 32) {
      auto v32i16Ty = VectorType::get({32}, rewriter.getI16Type());
      auto selResult = Sel16IntrOp::create(
          rewriter, loc, v32i16Ty,
          forceCastOperandsToSignature(
              rewriter, loc, {adaptor.getLhs(), adaptor.getRhs(), selMask},
              {v32i16Ty, v32i16Ty, i32Ty}));
      rewriter.replaceOp(op, selResult->getResult(0));
      return success();
    }
 
    return failure();
  }
};
 
using SelOpConversion =
    SelOpConversionBase<xllvm::VectorSel16IntrOp, xllvm::VectorSel32IntrOp>;
 
using SelOpAIE2pConversion = SelOpConversionBase<xllvm::VectorSel16AIE2pIntrOp,
                                                 xllvm::VectorSel32AIE2pIntrOp>;
 
class BroadcastScalarOpConversion
    : public mlir::ConvertOpToLLVMPattern<aievec::BroadcastScalarOp> {
public:
  using ConvertOpToLLVMPattern<
      aievec::BroadcastScalarOp>::ConvertOpToLLVMPattern;
 
  LogicalResult
  matchAndRewrite(aievec::BroadcastScalarOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Location loc = op.getLoc();
 
    Value result = op.getResult();
    VectorType resultType = cast<VectorType>(result.getType());
    Type resultScaTy = resultType.getElementType();
    unsigned resultBitWidth = resultScaTy.getIntOrFloatBitWidth();
    int resultLanes = getVectorLaneSize(resultType);
    int resultVectorSize = resultBitWidth * resultLanes;
 
    if (resultVectorSize != 512) {
      op.emitWarning()
          << "aievec.broadcast_scalar conversion with result vector size "
          << resultVectorSize << " is not implemented.\n";
      return failure();
    }
 
    // Integer types
    if (llvm::isa<IntegerType>(resultScaTy)) {
      Value src = adaptor.getSource();
      Type srcType = src.getType();
      unsigned srcBitWidth = srcType.getIntOrFloatBitWidth();
 
      if (srcBitWidth < 32) {
        src = LLVM::SExtOp::create(rewriter, loc, rewriter.getI32Type(), src);
      }
 
      if (resultBitWidth == 8) {
        rewriter.replaceOpWithNewOp<xllvm::VectorBroadcast8I512IntrOp>(
            op, VectorType::get({64}, rewriter.getI8Type()), src);
      } else if (resultBitWidth == 16) {
        rewriter.replaceOpWithNewOp<xllvm::VectorBroadcast16I512IntrOp>(
            op, VectorType::get({32}, rewriter.getI16Type()), src);
      } else if (resultBitWidth == 32) {
        rewriter.replaceOpWithNewOp<xllvm::VectorBroadcast32I512IntrOp>(
            op, VectorType::get({16}, rewriter.getI32Type()), src);
      } else {
        op.emitWarning()
            << "aievec.broadcast_scalar conversion with result bitwidth "
            << resultBitWidth << " is not implemented.\n";
        return failure();
      }
    } else {
      // Float types
      if (resultBitWidth == 16) {
        rewriter.replaceOpWithNewOp<xllvm::VectorBroadcast16BF512IntrOp>(
            op, VectorType::get({32}, rewriter.getBF16Type()),
            adaptor.getSource());
      } else if (resultBitWidth == 32) {
        // Use llvm.aie2.vbroadcast32.I512 with bitcasts (float -> i32 -> float)
        // Following the pattern: %0 = bitcast float %b to i32
        //                        %1 = tail call <16 x i32>
        //                        @llvm.aie2.vbroadcast32.I512(i32 %0) %2 =
        //                        bitcast <16 x i32> %1 to <16 x float>
        auto srcAsI32 = bitcastValueToType(rewriter, loc, adaptor.getSource(),
                                           rewriter.getI32Type());
        auto broadcastI32 = xllvm::VectorBroadcast32I512IntrOp::create(
            rewriter, loc, VectorType::get({16}, rewriter.getI32Type()),
            srcAsI32);
        auto resultF32 =
            bitcastValueToType(rewriter, loc, broadcastI32,
                               VectorType::get({16}, rewriter.getF32Type()));
        rewriter.replaceOp(op, resultF32);
      } else {
        op.emitWarning()
            << "aievec.broadcast_scalar conversion with result bitwidth "
            << resultBitWidth << " is not implemented.\n";
        return failure();
      }
    }
 
    return success();
  }
};
 
// AIE2p version of BroadcastScalarOp conversion using insertelement +
// shufflevector
class BroadcastScalarOpAIE2pConversion
    : public mlir::ConvertOpToLLVMPattern<aievec::BroadcastScalarOp> {
public:
  using ConvertOpToLLVMPattern<
      aievec::BroadcastScalarOp>::ConvertOpToLLVMPattern;
 
  LogicalResult
  matchAndRewrite(aievec::BroadcastScalarOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Location loc = op.getLoc();
 
    Value result = op.getResult();
    VectorType resultType = cast<VectorType>(result.getType());
    Type resultScaTy = resultType.getElementType();
    unsigned resultBitWidth = resultScaTy.getIntOrFloatBitWidth();
    int resultLanes = getVectorLaneSize(resultType);
    int resultVectorSize = resultBitWidth * resultLanes;
 
    // Support both 256-bit and 512-bit vectors for AIE2p
    if (resultVectorSize != 256 && resultVectorSize != 512) {
      op.emitWarning()
          << "aievec.broadcast_scalar conversion with result vector size "
          << resultVectorSize << " is not implemented for AIE2p.\n";
      return failure();
    }
 
    Value src = adaptor.getSource();
    Type srcType = src.getType();
 
    // For integer types, extend or truncate to match result element type
    if (llvm::isa<IntegerType>(resultScaTy)) {
      unsigned srcBitWidth = srcType.getIntOrFloatBitWidth();
      if (srcBitWidth < resultBitWidth) {
        src = LLVM::SExtOp::create(rewriter, loc, resultScaTy, src);
      } else if (srcBitWidth > resultBitWidth) {
        src = LLVM::TruncOp::create(rewriter, loc, resultScaTy, src);
      }
    }
 
    // Create poison vector of the result type
    auto poisonVec = LLVM::PoisonOp::create(rewriter, loc, resultType);
 
    // Insert scalar at position 0
    auto idx0 = LLVM::ConstantOp::create(rewriter, loc, rewriter.getI64Type(),
                                         rewriter.getI64IntegerAttr(0));
    auto insertedVec = LLVM::InsertElementOp::create(rewriter, loc, resultType,
                                                     poisonVec, src, idx0);
 
    // Create shufflevector mask with all zeros (broadcast position 0 to all
    // lanes)
    SmallVector<int64_t> broadcastMask(resultLanes, 0);
    auto broadcastVec = vector::ShuffleOp::create(rewriter, loc, insertedVec,
                                                  insertedVec, broadcastMask);
 
    rewriter.replaceOp(op, broadcastVec);
    return success();
  }
};
 
class ShiftOpConversion : public mlir::ConvertOpToLLVMPattern<aievec::ShiftOp> {
public:
  using ConvertOpToLLVMPattern<aievec::ShiftOp>::ConvertOpToLLVMPattern;
 
  LogicalResult
  matchAndRewrite(aievec::ShiftOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Location loc = op.getLoc();
 
    Value result = op.getResult();
    VectorType resultType = cast<VectorType>(result.getType());
    Type resultScaTy = resultType.getElementType();
    unsigned resultBitWidth = resultScaTy.getIntOrFloatBitWidth();
    int resultLanes = getVectorLaneSize(resultType);
    int resultVectorSize = resultBitWidth * resultLanes;
 
    if (resultVectorSize != 512) {
      op.emitWarning() << "aievec.shift conversion with result vector size "
                       << resultVectorSize << " is not implemented.\n";
      return failure();
    }
 
    // assume step is always zero
    auto stepCst = LLVM::ConstantOp::create(
        rewriter, loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(0));
 
    // create xllvm intrinsic
    Value shiftOp = nullptr;
    SmallVector<Value> operands(
        {adaptor.getLhs(), adaptor.getRhs(), stepCst, adaptor.getShift()});
    if (llvm::isa<IntegerType>(resultScaTy)) {
      // Integer types
      shiftOp = xllvm::VectorShiftI512I512IntrOp::create(
          rewriter, loc, VectorType::get({16}, rewriter.getI32Type()),
          forceCastOperandsToSignature(
              rewriter, loc, operands,
              {VectorType::get({16}, rewriter.getI32Type()),
               VectorType::get({16}, rewriter.getI32Type()),
               rewriter.getI32Type(), rewriter.getI32Type()}));
    } else {
      // Float types
      shiftOp = xllvm::VectorShiftBF512BF512IntrOp::create(
          rewriter, loc, VectorType::get({32}, rewriter.getBF16Type()),
          forceCastOperandsToSignature(
              rewriter, loc, operands,
              {VectorType::get({32}, rewriter.getBF16Type()),
               VectorType::get({32}, rewriter.getBF16Type()),
               rewriter.getI32Type(), rewriter.getI32Type()}));
    }
 
    // create bitcast/shape_cast for result
    auto resultVal =
        forceCastValueToType(rewriter, loc, shiftOp, op.getResult().getType());
    rewriter.replaceOp(op, resultVal);
 
    return success();
  }
};
 
// AIE2p version of ShiftOp conversion
class ShiftOpAIE2pConversion
    : public mlir::ConvertOpToLLVMPattern<aievec::ShiftOp> {
public:
  using ConvertOpToLLVMPattern<aievec::ShiftOp>::ConvertOpToLLVMPattern;
 
  LogicalResult
  matchAndRewrite(aievec::ShiftOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Location loc = op.getLoc();
 
    Value result = op.getResult();
    VectorType resultType = cast<VectorType>(result.getType());
    Type resultScaTy = resultType.getElementType();
    unsigned resultBitWidth = resultScaTy.getIntOrFloatBitWidth();
    int resultLanes = getVectorLaneSize(resultType);
    int resultVectorSize = resultBitWidth * resultLanes;
 
    if (resultVectorSize != 512) {
      op.emitWarning() << "aievec.shift conversion with result vector size "
                       << resultVectorSize << " is not implemented.\n";
      return failure();
    }
 
    // assume step is always zero
    auto stepCst = LLVM::ConstantOp::create(
        rewriter, loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(0));
 
    // create xllvm intrinsic
    Value shiftOp = nullptr;
    SmallVector<Value> operands(
        {adaptor.getLhs(), adaptor.getRhs(), stepCst, adaptor.getShift()});
    if (llvm::isa<IntegerType>(resultScaTy)) {
      // Integer types - use AIE2p intrinsic
      shiftOp = xllvm::VectorShiftI512I512AIE2pIntrOp::create(
          rewriter, loc, VectorType::get({16}, rewriter.getI32Type()),
          forceCastOperandsToSignature(
              rewriter, loc, operands,
              {VectorType::get({16}, rewriter.getI32Type()),
               VectorType::get({16}, rewriter.getI32Type()),
               rewriter.getI32Type(), rewriter.getI32Type()}));
    } else {
      // Float types - use AIE2p intrinsic
      shiftOp = xllvm::VectorShiftBF512BF512AIE2pIntrOp::create(
          rewriter, loc, VectorType::get({32}, rewriter.getBF16Type()),
          forceCastOperandsToSignature(
              rewriter, loc, operands,
              {VectorType::get({32}, rewriter.getBF16Type()),
               VectorType::get({32}, rewriter.getBF16Type()),
               rewriter.getI32Type(), rewriter.getI32Type()}));
    }
 
    // create bitcast/shape_cast for result
    auto resultVal =
        forceCastValueToType(rewriter, loc, shiftOp, op.getResult().getType());
    rewriter.replaceOp(op, resultVal);
 
    return success();
  }
};
 
class ExtractElemOpConversion
    : public mlir::ConvertOpToLLVMPattern<aievec::ExtElemOp> {
public:
  using ConvertOpToLLVMPattern<aievec::ExtElemOp>::ConvertOpToLLVMPattern;
 
  LogicalResult
  matchAndRewrite(aievec::ExtElemOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Location loc = op.getLoc();
 
    Type resultType = op.getResult().getType();
    unsigned resultBitWidth = resultType.getIntOrFloatBitWidth();
 
    Value src = adaptor.getSource();
    VectorType srcType = cast<VectorType>(src.getType());
    Type srcScalarType = srcType.getElementType();
    unsigned srcBitWidth = srcScalarType.getIntOrFloatBitWidth();
    int srcLanes = getVectorLaneSize(srcType);
    int srcVectorSize = srcBitWidth * srcLanes;
 
    if (srcVectorSize != 512) {
      op.emitWarning() << "aievec.ext_elem conversion with source vector size "
                       << srcVectorSize << " is not supported.\n";
      return failure();
    }
 
    // create constant for sign
    auto signCst = LLVM::ConstantOp::create(
        rewriter, loc, rewriter.getI32Type(), rewriter.getI32IntegerAttr(1));
 
    // create xllvm intrinsic
    Value extElemOp = nullptr;
    SmallVector<Value> operands(
        {adaptor.getSource(), adaptor.getIndex(), signCst});
    if (resultBitWidth == 8) {
      extElemOp = xllvm::VectorExtractElem8I512IntrOp::create(
          rewriter, loc, rewriter.getI32Type(),
          forceCastOperandsToSignature(
              rewriter, loc, operands,
              {VectorType::get({64}, rewriter.getI8Type()),
               rewriter.getI32Type(), rewriter.getI32Type()}));
    } else if (resultBitWidth == 16) {
      extElemOp = xllvm::VectorExtractElem16I512IntrOp::create(
          rewriter, loc, rewriter.getI32Type(),
          forceCastOperandsToSignature(
              rewriter, loc, operands,
              {VectorType::get({32}, rewriter.getI16Type()),
               rewriter.getI32Type(), rewriter.getI32Type()}));
    } else if (resultBitWidth == 32) {
      extElemOp = xllvm::VectorExtractElem32I512IntrOp::create(
          rewriter, loc, rewriter.getI32Type(),
          forceCastOperandsToSignature(
              rewriter, loc, operands,
              {VectorType::get({16}, rewriter.getI32Type()),
               rewriter.getI32Type(), rewriter.getI32Type()}));
    } else {
      op.emitWarning() << "aievec.ext_elem conversion with result bit width "
                       << resultBitWidth << " is not implemented.\n";
      return failure();
    }
 
    // create truncation op (and bitcast op)
    if (llvm::isa<IntegerType>(resultType)) {
      if (resultBitWidth < 32) {
        // Two-step truncation to avoid direct i32→i8 which the AIE2
        // backend cannot legalize after SLP vectorization.
        if (resultBitWidth < 16) {
          auto i16Ty = rewriter.getI16Type();
          auto trunc16 = LLVM::TruncOp::create(rewriter, loc, i16Ty, extElemOp);
          rewriter.replaceOpWithNewOp<LLVM::TruncOp>(op, resultType,
                                                     trunc16.getResult());
        } else {
          rewriter.replaceOpWithNewOp<LLVM::TruncOp>(op, resultType, extElemOp);
        }
      } else {
        rewriter.replaceOp(op, extElemOp);
      }
    } else {
      // Float types
      if (resultBitWidth == 16) {
        extElemOp = LLVM::TruncOp::create(rewriter, loc, rewriter.getI16Type(),
                                          extElemOp);
      }
      rewriter.replaceOpWithNewOp<LLVM::BitcastOp>(op, resultType, extElemOp);
    }
 
    return success();
  }
};
 
class FMAElemOpConversion
    : public mlir::ConvertOpToLLVMPattern<aievec::FMAElemOp> {
public:
  using ConvertOpToLLVMPattern<aievec::FMAElemOp>::ConvertOpToLLVMPattern;
 
  LogicalResult
  matchAndRewrite(aievec::FMAElemOp fmaOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto loc = fmaOp.getLoc();
    auto lhs = adaptor.getLhs();
    auto rhs = adaptor.getRhs();
    auto acc = adaptor.getAcc();
    auto lhsTy = cast<VectorType>(lhs.getType());
    auto rhsTy = cast<VectorType>(rhs.getType());
    auto accTy = cast<VectorType>(acc.getType());
    auto flatLhsTy = getFlattenedVectorType(lhsTy);
    auto flatRhsTy = getFlattenedVectorType(rhsTy);
    auto flatAccTy = getFlattenedVectorType(accTy);
 
    // Flatten operands, if needed
    if (lhsTy != flatLhsTy)
      lhs = vector::ShapeCastOp::create(rewriter, loc, flatLhsTy, lhs);
    if (rhsTy != flatRhsTy)
      rhs = vector::ShapeCastOp::create(rewriter, loc, flatRhsTy, rhs);
    if (accTy != flatAccTy)
      acc = vector::ShapeCastOp::create(rewriter, loc, flatAccTy, acc);
 
    // Build vmac configuration constant
    Type i32ty = rewriter.getI32Type();
    auto confCst = LLVM::ConstantOp::create(
        rewriter, loc, i32ty,
        rewriter.getI32IntegerAttr(aiev2_vmac_compute_control(
            /*sgn_x=*/0, /*sgn_y=*/0, /*amode=*/2, /*bmode=*/3,
            /*variant=*/1, /*zero_acc=*/0, /*shift16=*/0,
            /*sub_mul=*/0, /*sub_acc1=*/0, /*sub_acc2=*/0,
            /*sub_mask=*/0)));
 
    // Pad 16-lane bf16 operands to 32-lane using set+upd intrinsics.
    // forceCastOperandsToSignature only does bitwise reinterpretation, which
    // leaves garbage in the upper lanes. The MAC intrinsic requires properly
    // zero-padded v32bf16 inputs.
    auto v32bf16Ty = VectorType::get({32}, rewriter.getBF16Type());
    if (flatLhsTy.getElementType().isBF16() &&
        flatLhsTy.getNumElements() < 32) {
      auto zero32 = LLVM::ConstantOp::create(rewriter, loc, i32ty,
                                             rewriter.getI32IntegerAttr(0));
      auto zeros_i16 = xllvm::VectorBroadcast16I512IntrOp::create(
          rewriter, loc, VectorType::get({32}, rewriter.getI16Type()), zero32);
      auto zeros_bf16 =
          LLVM::BitcastOp::create(rewriter, loc, v32bf16Ty, zeros_i16);
      auto zeroVec = xllvm::ExtBF256BF512IntrOp::create(
          rewriter, loc, VectorType::get({16}, rewriter.getBF16Type()),
          zeros_bf16, zero32);
 
      auto idx1 = LLVM::ConstantOp::create(rewriter, loc, i32ty,
                                           rewriter.getI32IntegerAttr(1));
 
      auto lhsSet = xllvm::VectorSetBF512BF256IntrOp::create(
          rewriter, loc, v32bf16Ty, lhs, zero32);
      lhs = xllvm::UpdBF512BF256IntrOp::create(rewriter, loc, v32bf16Ty, lhsSet,
                                               zeroVec, idx1);
 
      auto rhsSet = xllvm::VectorSetBF512BF256IntrOp::create(
          rewriter, loc, v32bf16Ty, rhs, zero32);
      rhs = xllvm::UpdBF512BF256IntrOp::create(rewriter, loc, v32bf16Ty, rhsSet,
                                               zeroVec, idx1);
    }
 
    // Insert vmac intrinsic
    auto v8i64Ty = VectorType::get({8}, rewriter.getI64Type());
    auto macIntrOp = xllvm::MacConfBF16IntrOp::create(
        rewriter, loc, v8i64Ty,
        forceCastOperandsToSignature(rewriter, loc, {lhs, rhs, acc, confCst},
                                     {v32bf16Ty, v32bf16Ty, v8i64Ty, i32ty}));
 
    // Recast/Reshape result
    auto resVal =
        forceCastValueToType(rewriter, loc, macIntrOp.getResult(), flatAccTy);
    if (flatAccTy != accTy)
      resVal = vector::ShapeCastOp::create(rewriter, loc, accTy, resVal);
 
    rewriter.replaceOp(fmaOp, resVal);
    return success();
  }
};
 
// Look through vector.shape_cast ops to find the defining op. The
// VectorToAIEVec pass inserts shape_casts (via reshapeLeadingUnitDims) between
// the extension ops and the matmul, which would otherwise hide the signedness.
static Value lookThroughShapeCasts(Value v) {
  while (auto castOp = v.getDefiningOp<vector::ShapeCastOp>())
    v = castOp.getSource();
  return v;
}
 
// If `v` is defined by a signedness-carrying unrealized_conversion_cast --
// which is how VectorToAIEVec attaches si*/ui* to an otherwise signless
// operand -- return the signedness it carries and advance `v` past it, so the
// intrinsic consumes the signless input and the cast is left dead for
// canonicalization rather than leaking to LLVM translation.
//
// A cast whose result element type is not explicitly signed or unsigned is
// left in place: it is some other type conversion, and dropping it would
// silently change the operand.
static std::optional<bool> peelSignednessCast(Value &v) {
  auto castOp = v.getDefiningOp<UnrealizedConversionCastOp>();
  if (!castOp)
    return std::nullopt;
  auto vecTy = dyn_cast<VectorType>(v.getType());
  auto intTy =
      vecTy ? dyn_cast<IntegerType>(vecTy.getElementType()) : IntegerType();
  if (!intTy || intTy.isSignless())
    return std::nullopt;
  v = castOp.getInputs()[0];
  return intTy.isSigned();
}
 
// The result of resolving one integer matmul operand.
struct ResolvedMatMulOperand {
  // The value the MAC intrinsic should consume, with any widening op and any
  // signedness-carrying cast peeled off.
  Value narrowed;
  // Whether that value is to be treated as signed, i.e. whether the MAC
  // configuration word's signX (lhs) / signY (rhs) bit should be set.
  bool isSigned;
};
 
// Recover the signedness of an integer matmul operand. MLIR integers are
// signless, so it has to come from context; in priority order:
//
//   1. an arith.extsi / arith.extui feeding the operand -- peel it and take
//      its signedness;
//   2. an unrealized_conversion_cast to a signed/unsigned element type, which
//      is how VectorToAIEVec carries signedness across the peel it performs.
//      Peel it too, so the intrinsic consumes the signless value and the cast
//      is left dead for canonicalization rather than leaking to LLVM;
//   3. an explicitly signed or unsigned element type;
//   4. otherwise the element type is signless -- treat it as signed.
//
// Case 4 matches the conventional reading of a signless integer, and is the
// more useful default: the shapes that reach it are the ones whose operand
// type is already a legal AIE2 narrow type (so nothing was extended), which in
// practice means i16/i32 operands carrying signed data.
//
// Both aievec.matmul and aievec.matmul_aie2p resolve signedness through this
// function. They previously each had their own copy and had drifted apart.
static ResolvedMatMulOperand resolveMatMulOperand(Value v) {
  std::optional<bool> castSigned = peelSignednessCast(v);
 
  Value orig = lookThroughShapeCasts(v);
  if (auto extSIOp = orig.getDefiningOp<arith::ExtSIOp>())
    return {lookThroughShapeCasts(extSIOp.getIn()), true};
  if (auto extUIOp = orig.getDefiningOp<arith::ExtUIOp>())
    return {lookThroughShapeCasts(extUIOp.getIn()), false};
  if (castSigned)
    return {v, *castSigned};
 
  auto vecTy = dyn_cast<VectorType>(v.getType());
  auto intTy =
      vecTy ? dyn_cast<IntegerType>(vecTy.getElementType()) : IntegerType();
  return {v, !(intTy && intTy.isUnsigned())};
}
 
// Transposes the rows x cols tile held row-major in the flat vector `v` with
// the shuffle modes getTileTransposeModes picks for it. A 256-bit tile is
// padded to 512 bits for the shuffle, and a 1024-bit one is shuffled from its
// two 512-bit halves into the two halves of the result.
static Value transposeTile(OpBuilder &b, Location loc, Value v, int64_t rows,
                           int64_t cols, bool aie2p) {
  auto vecTy = cast<VectorType>(v.getType());
  std::optional<SmallVector<aievec::ShuffleMode, 2>> modes =
      aievec::getTileTransposeModes(aie2p, vecTy.getElementType(), rows, cols);
  assert(modes && "the matmul verifier admits only transposable tiles");
  if (modes->empty())
    return v;
  Type i32ty = b.getI32Type();
  auto v16i32Ty = VectorType::get({16}, i32ty);
  SmallVector<Value, 2> modeCsts;
  auto createModeConstants = [&]() {
    for (aievec::ShuffleMode mode : *modes)
      modeCsts.push_back(LLVM::ConstantOp::create(
          b, loc, i32ty, b.getI32IntegerAttr(static_cast<int32_t>(mode))));
  };
  auto shuffle = [&](Value x, Value y, Value modeCst) -> Value {
    if (aie2p)
      return xllvm::VectorShuffleAIE2pIntrOp::create(b, loc, v16i32Ty, x, y,
                                                     modeCst);
    return xllvm::VectorShuffleIntrOp::create(b, loc, v16i32Ty, x, y, modeCst);
  };
  int64_t numElems = vecTy.getNumElements();
  if (modes->size() == 2) {
    auto halves = forceCastValueToType(b, loc, v, VectorType::get({32}, i32ty));
    auto lo = vector::ShuffleOp::create(
        b, loc, halves, halves, llvm::to_vector(llvm::seq<int64_t>(0, 16)));
    auto hi = vector::ShuffleOp::create(
        b, loc, halves, halves, llvm::to_vector(llvm::seq<int64_t>(16, 32)));
    createModeConstants();
    Value resLo = shuffle(lo, hi, modeCsts[0]);
    Value resHi = shuffle(lo, hi, modeCsts[1]);
    auto res = vector::ShuffleOp::create(
        b, loc, resLo, resHi, llvm::to_vector(llvm::seq<int64_t>(0, 32)));
    return forceCastValueToType(b, loc, res, vecTy);
  }
  bool half = numElems * vecTy.getElementTypeBitWidth() == 256;
  Value x = v;
  if (half) {
    SmallVector<int64_t> padMask =
        llvm::to_vector(llvm::seq<int64_t>(0, numElems));
    padMask.append(numElems, vector::ShuffleOp::kPoisonIndex);
    x = vector::ShuffleOp::create(b, loc, v, v, padMask);
  }
  Value xi = forceCastValueToType(b, loc, x, v16i32Ty);
  Value y =
      aie2p ? xi
            : xllvm::UndefV16I32IntrOp::create(b, loc, v16i32Ty).getResult();
  createModeConstants();
  Value res =
      forceCastValueToType(b, loc, shuffle(xi, y, modeCsts[0]), x.getType());
  if (half)
    res = vector::ShuffleOp::create(
        b, loc, res, res, llvm::to_vector(llvm::seq<int64_t>(0, numElems)));
  return res;
}
 
// Brings the flattened operand `flat`, stored as `storedTy` (transposed when
// `stored` is set), into the layout the recipe reads it in.
static Value materializeOperandLayout(OpBuilder &b, Location loc, Value flat,
                                      VectorType storedTy, bool stored,
                                      bool native, bool aie2p) {
  if (stored == native)
    return flat;
  return transposeTile(b, loc, flat, storedTy.getDimSize(0),
                       storedTy.getDimSize(1), aie2p);
}
 
class MatMulOpConversion
    : public mlir::ConvertOpToLLVMPattern<aievec::MatMulOp> {
  using ConvertOpToLLVMPattern<aievec::MatMulOp>::ConvertOpToLLVMPattern;
 
  enum class Accumulator { I32, I64, BF16 };
 
  // The accumulator kind and `mac.conf` modes of each AIE2 recipe.
  static std::tuple<Accumulator, int, int, int>
  getAccAndModes(aievec::MatMulRecipeKind kind) {
    using K = aievec::MatMulRecipeKind;
    switch (kind) {
    case K::AIE2_I8xI4_4x16x8:
      return {Accumulator::I32, 0, 0, 0};
    case K::AIE2_I8xI8_4x8x8:
      return {Accumulator::I32, 0, 1, 0};
    case K::AIE2_I16xI8_4x4x8:
      return {Accumulator::I32, 0, 2, 0};
    case K::AIE2_I16xI16_4x2x8:
      return {Accumulator::I32, 0, 3, 0};
    case K::AIE2_I16xI8_2x8x8_Acc64:
      return {Accumulator::I64, 1, 2, 0};
    case K::AIE2_I16xI8_4x8x4_Acc64:
      return {Accumulator::I64, 1, 2, 1};
    case K::AIE2_I16xI16_2x4x8_Acc64:
      return {Accumulator::I64, 1, 3, 0};
    case K::AIE2_I16xI16_4x4x4_Acc64:
      return {Accumulator::I64, 1, 3, 1};
    case K::AIE2_I32xI16_4x2x4_Acc64:
      return {Accumulator::I64, 1, 0, 0};
    case K::AIE2_BF16_4x8x4:
      return {Accumulator::BF16, 2, 3, 0};
    default:
      llvm_unreachable("not an AIE2 recipe");
    }
  }
 
  LogicalResult
  matchAndRewrite(aievec::MatMulOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto lhsTy = cast<VectorType>(op.getLhs().getType());
    auto rhsTy = cast<VectorType>(op.getRhs().getType());
    auto accTy = cast<VectorType>(op.getAcc().getType());
    std::optional<aievec::MatMulRecipe> recipe =
        aievec::checkMatMulTile(/*aie2p=*/false, lhsTy, rhsTy, accTy,
                                op.getLhsTransposed(), op.getRhsTransposed());
    if (!recipe)
      return rewriter.notifyMatchFailure(op, "not an AIE2 matmul tile");
    auto [accKind, amode, bmode, variant] = getAccAndModes(recipe->kind);
 
    // Recover operand signedness, and narrow each operand to the value the MAC
    // should consume. An operand is only narrowed to a value of its own
    // element width, as the recipe was chosen for that width.
    Value lhs = adaptor.getLhs();
    Value rhs = adaptor.getRhs();
    int signX = 0, signY = 0;
    if (accKind != Accumulator::BF16) {
      auto narrow = [](Value v, VectorType ty, int &sign) {
        ResolvedMatMulOperand resolved = resolveMatMulOperand(v);
        sign = resolved.isSigned ? 1 : 0;
        if (getElementTypeOrSelf(resolved.narrowed).getIntOrFloatBitWidth() ==
            ty.getElementTypeBitWidth())
          return resolved.narrowed;
        return v;
      };
      lhs = narrow(lhs, lhsTy, signX);
      rhs = narrow(rhs, rhsTy, signY);
    }
    int conf = aiev2_vmac_compute_control(
        /*sgn_x=*/signX, /*sgn_y=*/signY, amode, bmode, variant,
        /*zero_acc=*/0, /*shift16=*/0, /*sub_mul=*/0, /*sub_acc1=*/0,
        /*sub_acc2=*/0, /*sub_mask=*/0);
 
    Location loc = op.getLoc();
    // Flatten the inputs, and transpose those the multiply reads the other
    // way.
    auto flatten = [&](Value v) -> Value {
      auto flatTy = getFlattenedVectorType(cast<VectorType>(v.getType()));
      if (v.getType() == flatTy)
        return v;
      return vector::ShapeCastOp::create(rewriter, loc, flatTy, v);
    };
    lhs = flatten(lhs);
    rhs = flatten(rhs);
    auto accFlattenedVecTy = getFlattenedVectorType(accTy);
    Value acc = vector::ShapeCastOp::create(rewriter, loc, accFlattenedVecTy,
                                            adaptor.getAcc());
    lhs = materializeOperandLayout(rewriter, loc, lhs, lhsTy,
                                   op.getLhsTransposed(), recipe->lhsKM,
                                   /*aie2p=*/false);
    rhs = materializeOperandLayout(rewriter, loc, rhs, rhsTy,
                                   op.getRhsTransposed(), recipe->rhsNK,
                                   /*aie2p=*/false);
 
    Type i32ty = rewriter.getI32Type();
    auto confCst = LLVM::ConstantOp::create(rewriter, loc, i32ty,
                                            rewriter.getI32IntegerAttr(conf));
    SmallVector<Value> operands({lhs, rhs, acc, confCst});
    Value matMulResVal;
    if (accKind == Accumulator::BF16)
      matMulResVal =
          xllvm::MacConfBF16IntrOp::create(
              rewriter, loc, VectorType::get({8}, rewriter.getI64Type()),
              forceCastOperandsToSignature(
                  rewriter, loc, operands,
                  {VectorType::get({32}, rewriter.getBF16Type()),
                   VectorType::get({32}, rewriter.getBF16Type()),
                   VectorType::get({8}, rewriter.getI64Type()), i32ty}))
              .getResult();
    else {
      SmallVector<Type> intrFuncSig(
          {VectorType::get({64}, rewriter.getI8Type()),
           VectorType::get({16}, i32ty),
           VectorType::get({16}, rewriter.getI64Type()), i32ty});
      VectorType v16xi64ty = VectorType::get({16}, rewriter.getI64Type());
      if (accKind == Accumulator::I32)
        matMulResVal = xllvm::MacConfAcc32IntrOp::create(
                           rewriter, loc, v16xi64ty,
                           forceCastOperandsToSignature(rewriter, loc, operands,
                                                        intrFuncSig))
                           .getResult();
      else
        matMulResVal = xllvm::MacConfAcc64IntrOp::create(
                           rewriter, loc, v16xi64ty,
                           forceCastOperandsToSignature(rewriter, loc, operands,
                                                        intrFuncSig))
                           .getResult();
    }
 
    auto castFromAcc =
        bitcastValueToType(rewriter, loc, matMulResVal, accFlattenedVecTy);
 
    rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(op, op.getType(),
                                                     castFromAcc);
 
    return success();
  }
};
 
// Widens a v64bf16 to v64accfloat by multiplying it by 1.0, which is exact
// and runs on the multiplier instead of the conversion unit.
static Value widenBF16ToAccFloat(OpBuilder &rewriter, Location loc, Type i32ty,
                                 Value v64bf16) {
  auto v64bf16Ty = VectorType::get({64}, rewriter.getBF16Type());
  auto ones = LLVM::ConstantOp::create(
      rewriter, loc, v64bf16Ty,
      DenseElementsAttr::get(
          v64bf16Ty, rewriter.getFloatAttr(rewriter.getBF16Type(), 1.0)));
  // bf16 x bf16 elementwise multiply into an fp32 accumulator.
  auto conf = LLVM::ConstantOp::create(rewriter, loc, i32ty,
                                       rewriter.getI32IntegerAttr(60));
  return xllvm::MulConfBF16I1024ACC2048AIE2pIntrOp::create(
      rewriter, loc, VectorType::get({64}, rewriter.getF32Type()), v64bf16,
      ones, conf);
}
 
// Helper function to perform BFP16-based 8×8 matmul via mac_8x8_8x8T_conf
// LHS: v64accfloat, RHS: v64accfloat (transposed 8×8), ACC: v64i32
// Returns: v64i32 result
static Value performBFP16_8x8MatMul(OpBuilder &rewriter, Location loc,
                                    Type i32ty, Value lhs64f32,
                                    Value rhs64f32Transposed, Value acc64i32,
                                    Value confCst) {
  auto v64i32Ty = VectorType::get({64}, rewriter.getI32Type());
 
  // Convert both to BFP16 format
  auto bfpStructTy = mlir::LLVM::LLVMStructType::getLiteral(
      rewriter.getContext(), {VectorType::get({64}, rewriter.getI8Type()),
                              VectorType::get({8}, rewriter.getI8Type())});
 
  auto lhsBFP = xllvm::Vector64AccFloatToV64BFP16EBS8AIE2pIntrOp::create(
      rewriter, loc, bfpStructTy, lhs64f32);
  auto rhsBFP = xllvm::Vector64AccFloatToV64BFP16EBS8AIE2pIntrOp::create(
      rewriter, loc, bfpStructTy, rhs64f32Transposed);
 
  // Extract mantissa and exponent
  auto lhsData = LLVM::ExtractValueOp::create(rewriter, loc, lhsBFP, 0);
  auto lhsExp = LLVM::ExtractValueOp::create(rewriter, loc, lhsBFP, 1);
  auto rhsData = LLVM::ExtractValueOp::create(rewriter, loc, rhsBFP, 0);
  auto rhsExp = LLVM::ExtractValueOp::create(rewriter, loc, rhsBFP, 1);
 
  // Perform BFP16 matmul
  return xllvm::MacConfBFP576ACC2048AIE2pIntrOp::create(
      rewriter, loc, v64i32Ty, lhsData, lhsExp, rhsData, rhsExp, acc64i32,
      confCst);
}
 
// Helper function to perform 8×8×4 BF16 matmul following mac_8x8_8x4_bf16
// LHS: 64 bfloat16 (8×8 matrix stored [k, m]), RHS: 32 bfloat16 (8×4 matrix)
// ACC: 32 float (8×4 result)
static Value perform8x8x4MatMul(OpBuilder &rewriter, Location loc, Type i32ty,
                                Value lhsKM64bf16, Value rhs32bf16,
                                Value acc32f32) {
  auto v32bf16Ty = VectorType::get({32}, rewriter.getBF16Type());
  auto v32f32Ty = VectorType::get({32}, rewriter.getF32Type());
 
  // Rows 0-3 and 4-7 of the [k, m] lhs: one column of A per 8 values.
  auto xaBF16 =
      vector::ShuffleOp::create(rewriter, loc, lhsKM64bf16, lhsKM64bf16,
                                llvm::to_vector(llvm::seq<int64_t>(0, 32)));
  auto xbBF16 =
      vector::ShuffleOp::create(rewriter, loc, lhsKM64bf16, lhsKM64bf16,
                                llvm::to_vector(llvm::seq<int64_t>(32, 64)));
 
  // Helper to extract and broadcast 8 elements to 32, then shuffle with T16_4x8
  auto extractBroadcastShuffle = [&](Value src, int idx) -> Value {
    SmallVector<int64_t> extractMask;
    int startIdx = idx * 8;
    for (int i = 0; i < 8; ++i)
      extractMask.push_back(startIdx + i);
    // Broadcast by repeating 4 times to get 32 elements
    for (int rep = 0; rep < 3; ++rep) {
      for (int i = 0; i < 8; ++i)
        extractMask.push_back(startIdx + i);
    }
    auto broadcasted =
        vector::ShuffleOp::create(rewriter, loc, src, src, extractMask);
 
    // Apply T16_4x8 shuffle pattern (mode 29)
    auto broadI32 = forceCastValueToType(rewriter, loc, broadcasted,
                                         VectorType::get({16}, i32ty));
    auto shuffleMode4x8 = LLVM::ConstantOp::create(
        rewriter, loc, i32ty, rewriter.getI32IntegerAttr(29));
    auto shuffled = xllvm::VectorShuffleAIE2pIntrOp::create(
        rewriter, loc, VectorType::get({16}, i32ty), broadI32, broadI32,
        shuffleMode4x8);
 
    return forceCastValueToType(rewriter, loc, shuffled, v32bf16Ty);
  };
 
  // Prepare 8 row vectors from xa and xb
  SmallVector<Value> rowVectors;
  for (int i = 0; i < 4; ++i)
    rowVectors.push_back(extractBroadcastShuffle(xaBF16, i));
  for (int i = 0; i < 4; ++i)
    rowVectors.push_back(extractBroadcastShuffle(xbBF16, i));
 
  // Helper to extract and broadcast 4 elements to 32 (for RHS columns)
  auto extractBroadcast4 = [&](Value src, int idx) -> Value {
    SmallVector<int64_t> mask;
    int startIdx = idx * 4;
    // Repeat the 4 elements 8 times to get 32 elements
    for (int rep = 0; rep < 8; ++rep) {
      for (int i = 0; i < 4; ++i)
        mask.push_back(startIdx + i);
    }
    return vector::ShuffleOp::create(rewriter, loc, src, src, mask);
  };
 
  // Prepare 8 column vectors from RHS
  SmallVector<Value> colVectors;
  for (int i = 0; i < 8; ++i)
    colVectors.push_back(extractBroadcast4(rhs32bf16, i));
 
  // Perform 8 MAC operations with conf=60 (no zero_acc)
  auto conf60 = LLVM::ConstantOp::create(rewriter, loc, i32ty,
                                         rewriter.getI32IntegerAttr(60));
 
  Value acc = acc32f32;
  for (int i = 0; i < 8; ++i) {
    acc = xllvm::MacConfBF16I512ACC1024AIE2pIntrOp::create(
        rewriter, loc, v32f32Ty, rowVectors[i], colVectors[i], acc, conf60);
  }
 
  return acc;
}
 
class MatMulOpAIE2pConversion
    : public mlir::ConvertOpToLLVMPattern<aievec::MatMulOp_AIE2P> {
  using ConvertOpToLLVMPattern<aievec::MatMulOp_AIE2P>::ConvertOpToLLVMPattern;
 
  LogicalResult
  matchAndRewrite(aievec::MatMulOp_AIE2P op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    using K = aievec::MatMulRecipeKind;
    auto lhsTy = cast<VectorType>(op.getLhs().getType());
    auto rhsTy = cast<VectorType>(op.getRhs().getType());
    auto accTy = cast<VectorType>(op.getAcc().getType());
    std::optional<aievec::MatMulRecipe> recipe =
        aievec::checkMatMulTile(/*aie2p=*/true, lhsTy, rhsTy, accTy,
                                op.getLhsTransposed(), op.getRhsTransposed());
    if (!recipe)
      return rewriter.notifyMatchFailure(op, "not an AIE2P matmul tile");
 
    // Signedness follows the same rule as aievec.matmul; see
    // resolveMatMulOperand. Only the signedness is taken from it here; the
    // operand values are only advanced past a signedness-carrying cast, not
    // past a widening op, as the recipe was chosen for the operand types.
    Value lhsOperand = adaptor.getLhs();
    Value rhsOperand = adaptor.getRhs();
    bool lhsIsSigned = resolveMatMulOperand(lhsOperand).isSigned;
    bool rhsIsSigned = resolveMatMulOperand(rhsOperand).isSigned;
    peelSignednessCast(lhsOperand);
    peelSignednessCast(rhsOperand);
 
    Location loc = op.getLoc();
    // Flatten the inputs. An operand the multiply reads the other way is
    // transposed where the recipe first uses it.
    Value lhsFlat = vector::ShapeCastOp::create(
        rewriter, loc,
        getFlattenedVectorType(cast<VectorType>(lhsOperand.getType())),
        lhsOperand);
    Value rhsFlat = vector::ShapeCastOp::create(
        rewriter, loc,
        getFlattenedVectorType(cast<VectorType>(rhsOperand.getType())),
        rhsOperand);
    auto accFlattenedVecTy = getFlattenedVectorType(accTy);
    Value acc = vector::ShapeCastOp::create(rewriter, loc, accFlattenedVecTy,
                                            adaptor.getAcc());
    auto getLhs = [&]() {
      return materializeOperandLayout(rewriter, loc, lhsFlat, lhsTy,
                                      op.getLhsTransposed(), recipe->lhsKM,
                                      /*aie2p=*/true);
    };
    auto getRhs = [&]() {
      return materializeOperandLayout(rewriter, loc, rhsFlat, rhsTy,
                                      op.getRhsTransposed(), recipe->rhsNK,
                                      /*aie2p=*/true);
    };
    Type i32ty = rewriter.getI32Type();
 
    Value matMulResVal;
    switch (recipe->kind) {
    case K::AIE2P_I8_8x8x8: {
      // <8x8xi8> x <8x8xi8> + <8x8xi32>
      // Signature: <32 x i64> @llvm.aie2p.I512.I512.ACC2048.mac.conf(
      //              <16 x i32>, <32 x i16>, <32 x i64>, i32)
      // Base conf for amode=0/bmode=1 is bmode<<3 = 8; signedness lives in
      // signX (bit 9) / signY (bit 8).
      int conf =
          8 | ((lhsIsSigned ? 1 : 0) << 9) | ((rhsIsSigned ? 1 : 0) << 8);
      auto confCst = LLVM::ConstantOp::create(rewriter, loc, i32ty,
                                              rewriter.getI32IntegerAttr(conf));
      matMulResVal =
          xllvm::MacConfI512ACC2048AIE2pIntrOp::create(
              rewriter, loc, VectorType::get({32}, rewriter.getI64Type()),
              forceCastOperandsToSignature(
                  rewriter, loc, {getLhs(), getRhs(), acc, confCst},
                  {VectorType::get({16}, rewriter.getI32Type()),
                   VectorType::get({32}, rewriter.getI16Type()),
                   VectorType::get({32}, rewriter.getI64Type()), i32ty}))
              .getResult();
      break;
    }
    case K::AIE2P_I16_8x2x8: {
      // <8x2xi16> x <2x8xi16> + <8x8xi32>
      // Input vectors are 16 lanes each, need to pad to 64 lanes for intrinsic
      // Signature: <32 x i64> @llvm.aie2p.I1024.I1024.ACC2048.mac.conf(
      //              <32 x i32>, <64 x i16>, <32 x i64>, i32)
      SmallVector<int64_t> padMask = llvm::to_vector(llvm::seq<int64_t>(0, 16));
      padMask.append(48, vector::ShuffleOp::kPoisonIndex);
      Value lhs = getLhs(), rhs = getRhs();
      auto lhsPadded =
          vector::ShuffleOp::create(rewriter, loc, lhs, lhs, padMask);
      auto rhsPadded =
          vector::ShuffleOp::create(rewriter, loc, rhs, rhs, padMask);
      auto confCst = LLVM::ConstantOp::create(rewriter, loc, i32ty,
                                              rewriter.getI32IntegerAttr(24));
      matMulResVal =
          xllvm::MacConfI1024ACC2048AIE2pIntrOp::create(
              rewriter, loc, VectorType::get({32}, rewriter.getI64Type()),
              forceCastOperandsToSignature(
                  rewriter, loc, {lhsPadded, rhsPadded, acc, confCst},
                  {VectorType::get({32}, rewriter.getI32Type()),
                   VectorType::get({64}, rewriter.getI16Type()),
                   VectorType::get({32}, rewriter.getI64Type()), i32ty}))
              .getResult();
      break;
    }
    case K::AIE2P_BFP16_8x8x8: {
      // <8x8xbf16> x <8x8xbf16> + <8x8xf32>
      // This implements the 8×8×8 BF16 matmul using BFP16 format
      // Following the aie_api reference implementation that converts to BFP16
      auto v32f32Ty = VectorType::get({32}, rewriter.getF32Type());
 
      // Step 1: Convert LHS v64bfloat16 to v64accfloat (in two v32 chunks)
      Value lhs = getLhs();
      auto lhs32bf16_lo = vector::ShuffleOp::create(
          rewriter, loc, lhs, lhs, llvm::to_vector(llvm::seq<int64_t>(0, 32)));
      auto lhs32bf16_hi = vector::ShuffleOp::create(
          rewriter, loc, lhs, lhs, llvm::to_vector(llvm::seq<int64_t>(32, 64)));
      auto lhs32f32_lo = xllvm::Vector32BF16ToV32AccFloatAIE2pIntrOp::create(
          rewriter, loc, v32f32Ty, lhs32bf16_lo);
      auto lhs32f32_hi = xllvm::Vector32BF16ToV32AccFloatAIE2pIntrOp::create(
          rewriter, loc, v32f32Ty, lhs32bf16_hi);
 
      // Concat to v64accfloat
      auto lhs64f32 =
          vector::ShuffleOp::create(rewriter, loc, lhs32f32_lo, lhs32f32_hi,
                                    llvm::to_vector(llvm::seq<int64_t>(0, 64)));
 
      // Step 2: Widen the [n, k] RHS to accfloat
      auto rhsNK = widenBF16ToAccFloat(rewriter, loc, i32ty, getRhs());
 
      // Step 3: BFP16 8×8 matmul
      auto conf780 = LLVM::ConstantOp::create(rewriter, loc, i32ty,
                                              rewriter.getI32IntegerAttr(780));
      matMulResVal = performBFP16_8x8MatMul(
          rewriter, loc, i32ty, lhs64f32, rhsNK,
          forceCastValueToType(rewriter, loc, acc,
                               VectorType::get({64}, rewriter.getI32Type())),
          conf780);
      break;
    }
    case K::AIE2P_BFP16_4x8x8: {
      // <4x8xbf16> x <8x8xbf16> + <4x8xf32>
      // LHS: 32 lanes, RHS: 64 lanes, ACC: 32 lanes
      // Similar to 8×8×8 but only use first 32 lanes of LHS
      auto v32f32Ty = VectorType::get({32}, rewriter.getF32Type());
      SmallVector<int64_t> padMask = llvm::to_vector(llvm::seq<int64_t>(0, 32));
      padMask.append(32, vector::ShuffleOp::kPoisonIndex);
 
      // Step 1: Convert LHS v32bfloat16 to v32accfloat, then pad to v64accfloat
      auto lhs32f32 = xllvm::Vector32BF16ToV32AccFloatAIE2pIntrOp::create(
          rewriter, loc, v32f32Ty, getLhs());
      auto lhs64f32 =
          vector::ShuffleOp::create(rewriter, loc, lhs32f32, lhs32f32, padMask);
 
      // Step 2: Widen the [n, k] RHS to accfloat
      auto rhsNK = widenBF16ToAccFloat(rewriter, loc, i32ty, getRhs());
 
      // Step 3: Pad ACC from 32 to 64 i32
      auto acc32i32 = forceCastValueToType(
          rewriter, loc, acc, VectorType::get({32}, rewriter.getI32Type()));
      auto acc64i32 =
          vector::ShuffleOp::create(rewriter, loc, acc32i32, acc32i32, padMask);
 
      // Step 4: BFP16 8×8 matmul
      auto confCst = LLVM::ConstantOp::create(rewriter, loc, i32ty,
                                              rewriter.getI32IntegerAttr(780));
      auto result64i32 = performBFP16_8x8MatMul(rewriter, loc, i32ty, lhs64f32,
                                                rhsNK, acc64i32, confCst);
 
      // Step 5: Extract first 32 elements
      matMulResVal =
          vector::ShuffleOp::create(rewriter, loc, result64i32, result64i32,
                                    llvm::to_vector(llvm::seq<int64_t>(0, 32)));
      break;
    }
    case K::AIE2P_BF16Outer_8x1x8: {
      // <8x1xbf16> x <1x8xbf16> + <8x8xf32>
      // Outer product: grow_replicate both to 64, transpose LHS, use
      // mac_elem_64_conf
      auto v64f32Ty = VectorType::get({64}, rewriter.getF32Type());
 
      // Replicate LHS and RHS from 8 to 64 elements (8 times each)
      SmallVector<int64_t> replicateMask;
      for (int rep = 0; rep < 8; ++rep)
        for (int i = 0; i < 8; ++i)
          replicateMask.push_back(i);
      Value lhs = lhsFlat, rhs = rhsFlat;
      auto lhs64bf16 =
          vector::ShuffleOp::create(rewriter, loc, lhs, lhs, replicateMask);
 
      // Transpose LHS as 8×8 matrix
      SmallVector<int64_t> transposeMask;
      for (int c = 0; c < 8; ++c)
        for (int r = 0; r < 8; ++r)
          transposeMask.push_back(r * 8 + c);
      auto lhs64bf16Transposed = vector::ShuffleOp::create(
          rewriter, loc, lhs64bf16, lhs64bf16, transposeMask);
      auto rhs64bf16 =
          vector::ShuffleOp::create(rewriter, loc, rhs, rhs, replicateMask);
 
      // mac_elem_64_conf: 64 bf16 lanes, so the I1024 operand form.
      auto confCst = LLVM::ConstantOp::create(rewriter, loc, i32ty,
                                              rewriter.getI32IntegerAttr(60));
      matMulResVal = xllvm::MacConfBF16I1024ACC2048AIE2pIntrOp::create(
          rewriter, loc, v64f32Ty, lhs64bf16Transposed, rhs64bf16, acc,
          confCst);
      break;
    }
    case K::AIE2P_BF16Elem_4x8x4: {
      // 4×8×4 matmul: <4x8xbf16> x <8x4xbf16> + <4x4xf32>
      // Following the reference pattern: a.grow<64>(), b,
      // acc.grow<32>().extract<16>(0). Pad the [m, k] LHS to 8 rows, transpose
      // it to [k, m] and pad ACC 16→32 for the 8×8×4 multiply, then extract
      // 32→16.
      SmallVector<int64_t> lhsPadMask =
          llvm::to_vector(llvm::seq<int64_t>(0, 32));
      lhsPadMask.append(32, vector::ShuffleOp::kPoisonIndex);
      Value lhs = getLhs();
      Value lhsPadded =
          vector::ShuffleOp::create(rewriter, loc, lhs, lhs, lhsPadMask);
      lhsPadded = transposeTile(rewriter, loc, lhsPadded, 8, 8,
                                /*aie2p=*/true);
 
      SmallVector<int64_t> accPadMask =
          llvm::to_vector(llvm::seq<int64_t>(0, 16));
      accPadMask.append(16, vector::ShuffleOp::kPoisonIndex);
      auto accPadded =
          vector::ShuffleOp::create(rewriter, loc, acc, acc, accPadMask);
 
      Value acc32 = perform8x8x4MatMul(rewriter, loc, i32ty, lhsPadded,
                                       getRhs(), accPadded);
 
      matMulResVal =
          vector::ShuffleOp::create(rewriter, loc, acc32, acc32,
                                    llvm::to_vector(llvm::seq<int64_t>(0, 16)));
      break;
    }
    case K::AIE2P_BF16Elem_8x8x4: {
      // 8×8×4 matmul: <8x8xbf16> [k, m] x <8x4xbf16> + <8x4xf32>
      Value lhs = getLhs();
      Value rhs = getRhs();
      matMulResVal = perform8x8x4MatMul(rewriter, loc, i32ty, lhs, rhs, acc);
      break;
    }
    default:
      return rewriter.notifyMatchFailure(op, "not an AIE2P recipe");
    }
 
    // Cast from flattened result back to original accumulator shape
    auto castFromAcc =
        forceCastValueToType(rewriter, loc, matMulResVal, accFlattenedVecTy);
    // Reshape back to original shape
    rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(op, op.getType(),
                                                     castFromAcc);
    return success();
  }
};
 
// This pattern folds aievec.cast op. For AIE2, the accumulators are in 32/64
// bits, and the vectors are in 4/8/16/32 bits. Hence, we don't have to
// explicitly express the casting between accumulators and vectors at the LLVM
// dialect level. The backend LLVM compiler will decide the correct accumulator
// or vector registers given the ops and intrinsics.
class FoldAIECastOps : public mlir::ConvertOpToLLVMPattern<aievec::CastOp> {
  using ConvertOpToLLVMPattern<aievec::CastOp>::ConvertOpToLLVMPattern;
 
  // Helper to check if a value is a constant zero
  static bool isConstantZero(Value val) {
    DenseElementsAttr denseAttr;
 
    // Check for both arith.constant and llvm.mlir.constant
    if (auto arithConstOp = val.getDefiningOp<arith::ConstantOp>()) {
      denseAttr = dyn_cast<DenseElementsAttr>(arithConstOp.getValue());
    } else if (auto llvmConstOp = val.getDefiningOp<LLVM::ConstantOp>()) {
      denseAttr = dyn_cast<DenseElementsAttr>(llvmConstOp.getValue());
    }
 
    if (!denseAttr || !denseAttr.isSplat())
      return false;
 
    auto splatAttr = denseAttr.getSplatValue<Attribute>();
    if (auto floatAttr = dyn_cast<FloatAttr>(splatAttr))
      return floatAttr.getValue().isZero();
    if (auto intAttr = dyn_cast<IntegerAttr>(splatAttr))
      return intAttr.getValue().isZero();
 
    return false;
  }
 
  LogicalResult
  matchAndRewrite(aievec::CastOp castOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    // Special handling for isResAcc=true with zero constant source
    // The backend cannot handle zeroinitializer for accumulator types,
    // so we must use the vbroadcast.zero.acc1024 intrinsic instead
    if (!castOp.getIsResAcc() || !isConstantZero(adaptor.getSource())) {
      // Default behavior: fold the cast
      rewriter.replaceOp(castOp, adaptor.getSource());
      return success();
    }
 
    Location loc = castOp.getLoc();
    auto srcVecType = cast<VectorType>(castOp.getSource().getType());
    Type srcElemType = srcVecType.getElementType();
    int lanes = getVectorLaneSize(srcVecType);
 
    // For f32 vectors (accfloat), use vbroadcast.zero.acc1024
    if (srcElemType.isF32() && lanes == 16) {
      // Call vbroadcast.zero.acc1024 to get vector<16xi64>
      auto zeroAcc1024 = xllvm::VectorBroadcastZeroAcc1024IntrOp::create(
          rewriter, loc, VectorType::get({16}, rewriter.getI64Type()));
 
      // Extract lower 8 elements to get vector<8xi64> (512-bit accumulator)
      SmallVector<int64_t> extractMask = {0, 1, 2, 3, 4, 5, 6, 7};
      auto zeroAcc512 = vector::ShuffleOp::create(rewriter, loc, zeroAcc1024,
                                                  zeroAcc1024, extractMask);
 
      // Bitcast back to vector<16xf32> to match the cast result type
      auto result = LLVM::BitcastOp::create(
          rewriter, loc, VectorType::get({16}, rewriter.getF32Type()),
          zeroAcc512);
 
      rewriter.replaceOp(castOp, result);
      return success();
    }
 
    // Fallback: fold the cast (should not reach here for supported cases)
    rewriter.replaceOp(castOp, adaptor.getSource());
    return success();
  }
};
 
// AIE2p version of FoldAIECastOps.
// On AIE2P the int/float accumulators share the same physical register, so
// aievec.cast between e.g. v32i32 (acc32) and v32f32 (accfloat) is a no-op
// at hardware level. LLVM IR still needs a typed value, so emit an
// llvm.bitcast when the source/destination LLVM types differ but have the
// same bitwidth. Replace with the source value directly when types already
// match (the cast becomes a pure type alias).
class FoldAIECastOpsAIE2p
    : public mlir::ConvertOpToLLVMPattern<aievec::CastOp> {
  using ConvertOpToLLVMPattern<aievec::CastOp>::ConvertOpToLLVMPattern;
 
  LogicalResult
  matchAndRewrite(aievec::CastOp castOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Value src = adaptor.getSource();
    Type srcTy = src.getType();
    Type dstTy = getTypeConverter()->convertType(castOp.getResult().getType());
    if (srcTy == dstTy) {
      rewriter.replaceOp(castOp, src);
      return success();
    }
    rewriter.replaceOpWithNewOp<LLVM::BitcastOp>(castOp, dstTy, src);
    return success();
  }
};
 
class ShuffleOpConversion
    : public mlir::ConvertOpToLLVMPattern<aievec::ShuffleOp> {
  using ConvertOpToLLVMPattern<aievec::ShuffleOp>::ConvertOpToLLVMPattern;
 
  LogicalResult
  matchAndRewrite(aievec::ShuffleOp shuffleOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    if (isAIE2POnlyShuffleMode(shuffleOp.getMode()))
      return shuffleOp.emitOpError()
             << "shuffle mode '" << stringifyShuffleMode(shuffleOp.getMode())
             << "' is not available on AIE2";
    auto loc = shuffleOp.getLoc();
    auto lhs = adaptor.getLhs();
    auto rhs = adaptor.getRhs();
    auto i32ty = rewriter.getI32Type();
    auto v16xi32ty = VectorType::get({16}, i32ty);
    if (!rhs)
      rhs = xllvm::UndefV16I32IntrOp::create(rewriter, loc, v16xi32ty);
 
    auto modeAttrVal =
        LLVM::ConstantOp::create(rewriter, loc, i32ty,
                                 static_cast<int32_t>(shuffleOp.getMode()))
            .getResult();
    auto vShuffleVal = xllvm::VectorShuffleIntrOp::create(
                           rewriter, loc, v16xi32ty,
                           forceCastOperandsToSignature(
                               rewriter, loc,
                               /*operands=*/{lhs, rhs, modeAttrVal},
                               /*signature=*/{v16xi32ty, v16xi32ty, i32ty}))
                           .getResult();
 
    vShuffleVal = forceCastValueToType(rewriter, loc, vShuffleVal,
                                       shuffleOp.getResult().getType());
 
    rewriter.replaceOp(shuffleOp, vShuffleVal);
 
    return success();
  }
};
 
// aievec.shuffle on AIE2P: AIE2P keeps AIE2's mode numbers and adds its own;
// the intrinsic is AIE2P's.
class ShuffleOpAIE2pConversion
    : public mlir::ConvertOpToLLVMPattern<aievec::ShuffleOp> {
  using ConvertOpToLLVMPattern<aievec::ShuffleOp>::ConvertOpToLLVMPattern;
 
  LogicalResult
  matchAndRewrite(aievec::ShuffleOp shuffleOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto loc = shuffleOp.getLoc();
    auto i32ty = rewriter.getI32Type();
    auto v16xi32ty = VectorType::get({16}, i32ty);
    Value lhs = adaptor.getLhs();
    Value rhs = adaptor.getRhs();
    if (!rhs)
      rhs = lhs;
    auto mode = LLVM::ConstantOp::create(
        rewriter, loc, i32ty, static_cast<int32_t>(shuffleOp.getMode()));
    Value v = xllvm::VectorShuffleAIE2pIntrOp::create(
                  rewriter, loc, v16xi32ty,
                  forceCastOperandsToSignature(rewriter, loc, {lhs, rhs, mode},
                                               {v16xi32ty, v16xi32ty, i32ty}))
                  .getResult();
    rewriter.replaceOp(shuffleOp,
                       forceCastValueToType(rewriter, loc, v,
                                            shuffleOp.getResult().getType()));
    return success();
  }
};
 
// Convert aievec.inv to xllvm.intr.aie2p.inv intrinsic for AIE2P
// Scalar f32: direct conversion to xllvm.intr.aie2p.inv
// Vector f32: unroll into scalar xllvm.intr.aie2p.inv operations
class InvOpAIE2pConversion
    : public mlir::ConvertOpToLLVMPattern<aievec::InvOp> {
public:
  using ConvertOpToLLVMPattern<aievec::InvOp>::ConvertOpToLLVMPattern;
 
  LogicalResult
  matchAndRewrite(aievec::InvOp invOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto loc = invOp.getLoc();
    auto operandType = adaptor.getSource().getType();
 
    // Handle scalar f32 inverse
    if (operandType.isF32()) {
      auto invResult = xllvm::InvAIE2pIntrOp::create(
          rewriter, loc, rewriter.getF32Type(), adaptor.getSource());
      rewriter.replaceOp(invOp, invResult);
      return success();
    }
 
    // Handle vector<N x f32> inverse
    auto vecType = dyn_cast<VectorType>(operandType);
    if (!vecType || !vecType.getElementType().isF32())
      return failure();
 
    // Unroll vector inverse into scalar operations
    int numElements = getVectorLaneSize(vecType);
    Value result = LLVM::PoisonOp::create(rewriter, loc, vecType);
 
    for (int i = 0; i < numElements; ++i) {
      // Extract element i
      auto indexCst = LLVM::ConstantOp::create(
          rewriter, loc, rewriter.getI64Type(), rewriter.getI64IntegerAttr(i));
      auto extractedElem = LLVM::ExtractElementOp::create(
          rewriter, loc, adaptor.getSource(), indexCst);
 
      // Call xllvm.intr.aie2p.inv on the scalar
      auto invResult = xllvm::InvAIE2pIntrOp::create(
          rewriter, loc, rewriter.getF32Type(), extractedElem);
 
      // Insert result back into vector
      result = LLVM::InsertElementOp::create(rewriter, loc, vecType, result,
                                             invResult, indexCst);
    }
 
    rewriter.replaceOp(invOp, result);
    return success();
  }
};
 
// Apply `IntrOpTy`, an AIE2P intrinsic from v16 accfloat to v16 bf16 such as
// exp2 or tanh, to a 16- or 32-lane bf16 or f32 vector. A 32-lane vector is
// processed as two 16-lane halves. bf16 lanes are widened to accfloat before
// the intrinsic, and the bf16 results are widened again when `resultElemType`
// is f32.
template <typename IntrOpTy>
static Value applyAccFloatUnaryIntrinsic(ConversionPatternRewriter &rewriter,
                                         Location loc, Value source,
                                         Type resultElemType) {
  auto srcType = cast<VectorType>(source.getType());
  unsigned laneSize = getVectorLaneSize(srcType);
  auto v16bf16Ty = VectorType::get({16}, rewriter.getBF16Type());
  auto v16f32Ty = VectorType::get({16}, rewriter.getF32Type());
  auto toF32 = [&](Value v16bf16) -> Value {
    return xllvm::Vector16BF16ToV16AccFloatAIE2pIntrOp::create(
        rewriter, loc, v16f32Ty, v16bf16);
  };
 
  SmallVector<Value, 2> parts;
  if (laneSize == 16) {
    parts.push_back(source);
  } else {
    for (int64_t half = 0; half < 2; ++half) {
      SmallVector<int64_t> mask =
          llvm::to_vector(llvm::seq<int64_t>(16 * half, 16 * (half + 1)));
      parts.push_back(
          vector::ShuffleOp::create(rewriter, loc, source, source, mask));
    }
  }
  if (srcType.getElementType().isBF16())
    for (Value &part : parts)
      part = toF32(part);
  for (Value &part : parts)
    part = IntrOpTy::create(rewriter, loc, v16bf16Ty, part);
  if (resultElemType.isF32())
    for (Value &part : parts)
      part = toF32(part);
 
  if (parts.size() == 1)
    return parts.front();
  SmallVector<int64_t> combineMask = llvm::to_vector(llvm::seq<int64_t>(0, 32));
  return vector::ShuffleOp::create(rewriter, loc, parts[0], parts[1],
                                   combineMask);
}
 
// Convert aievec.exp to xllvm.exp2 intrinsic for AIE2P
// Uses the identity: exp(x) = exp2(x * log2(e))
// Supports both lane-16 and lane-32 bf16 vectors
class ExpOpAIE2pConversion
    : public mlir::ConvertOpToLLVMPattern<aievec::ExpOp> {
public:
  using ConvertOpToLLVMPattern<aievec::ExpOp>::ConvertOpToLLVMPattern;
 
  LogicalResult
  matchAndRewrite(aievec::ExpOp expOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto loc = expOp.getLoc();
    auto srcType = cast<VectorType>(adaptor.getSource().getType());
    auto srcElemType = srcType.getElementType();
    unsigned laneSize = getVectorLaneSize(srcType);
 
    // Support v16bfloat16 and v32bfloat16
    if ((laneSize != 16 && laneSize != 32) || !srcElemType.isBF16())
      return expOp.emitWarning()
             << "aievec.exp conversion only supports v16bfloat16 and "
                "v32bfloat16.\n";
 
    // Step 1: Create bf16 constant for log2(e) ≈ 1.442695
    auto log2eBF16Const = LLVM::ConstantOp::create(
        rewriter, loc, rewriter.getBF16Type(),
        rewriter.getFloatAttr(rewriter.getBF16Type(), 1.442695));
 
    // Broadcast log2(e) to match input lane size
    SmallVector<int64_t> broadcastMask;
    for (unsigned i = 0; i < laneSize; ++i)
      broadcastMask.push_back(0);
 
    auto v1bf16 = LLVM::UndefOp::create(
        rewriter, loc, VectorType::get({1}, rewriter.getBF16Type()));
    auto v1bf16Inserted = LLVM::InsertElementOp::create(
        rewriter, loc, v1bf16, log2eBF16Const,
        LLVM::ConstantOp::create(rewriter, loc, rewriter.getI32Type(), 0));
 
    auto log2eVec = vector::ShuffleOp::create(rewriter, loc, v1bf16Inserted,
                                              v1bf16Inserted, broadcastMask);
 
    // Step 2: Multiply input by log2(e) in bf16 domain using MulElemOp
    // For lane-16: uses I512.I512.ACC512
    // For lane-32: uses I512.I512.ACC1024
    auto resultF32Ty =
        VectorType::get({(int64_t)laneSize}, rewriter.getF32Type());
    auto mulResult = aievec::MulElemOp::create(rewriter, loc, resultF32Ty,
                                               adaptor.getSource(), log2eVec);
 
    // Step 3: exp2 of the product
    Value exp2Result = applyAccFloatUnaryIntrinsic<xllvm::Exp2AIE2pIntrOp>(
        rewriter, loc, mulResult, rewriter.getBF16Type());
 
    rewriter.replaceOp(expOp, exp2Result);
 
    return success();
  }
};
 
// Convert aievec.tanh or aievec.exp2 to the matching xllvm intrinsic for
// AIE2P. Supports lane-16 and lane-32 bf16 and f32 vectors.
template <typename SrcOpTy, typename IntrOpTy>
class AccFloatUnaryOpAIE2pConversion
    : public mlir::ConvertOpToLLVMPattern<SrcOpTy> {
public:
  using ConvertOpToLLVMPattern<SrcOpTy>::ConvertOpToLLVMPattern;
  using OpAdaptor = typename SrcOpTy::Adaptor;
 
  LogicalResult
  matchAndRewrite(SrcOpTy op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto srcType = cast<VectorType>(adaptor.getSource().getType());
    Type srcElemType = srcType.getElementType();
    unsigned laneSize = getVectorLaneSize(srcType);
 
    if ((laneSize != 16 && laneSize != 32) ||
        (!srcElemType.isBF16() && !srcElemType.isF32()))
      return op.emitWarning()
             << op->getName()
             << " conversion only supports v16 and v32 of bfloat16 or float.\n";
 
    rewriter.replaceOp(
        op, applyAccFloatUnaryIntrinsic<IntrOpTy>(
                rewriter, op.getLoc(), adaptor.getSource(), srcElemType));
    return success();
  }
};
 
// Convert math.rsqrt (scalar f32 or vector f32) to xllvm.intr.aie2p.invsqrt
// Scalar f32: direct conversion to xllvm.intr.aie2p.invsqrt
// Vector f32: unroll into scalar xllvm.intr.aie2p.invsqrt operations
class RsqrtOpAIE2pConversion
    : public mlir::ConvertOpToLLVMPattern<math::RsqrtOp> {
public:
  using ConvertOpToLLVMPattern<math::RsqrtOp>::ConvertOpToLLVMPattern;
 
  LogicalResult
  matchAndRewrite(math::RsqrtOp rsqrtOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto loc = rsqrtOp.getLoc();
    auto operandType = adaptor.getOperand().getType();
 
    // Handle scalar f32 rsqrt
    if (operandType.isF32()) {
      auto rsqrtResult = xllvm::InvsqrtAIE2pIntrOp::create(
          rewriter, loc, rewriter.getF32Type(), adaptor.getOperand());
      rewriter.replaceOp(rsqrtOp, rsqrtResult);
      return success();
    }
 
    // Handle vector<N x f32> rsqrt
    auto vecType = dyn_cast<VectorType>(operandType);
    if (!vecType || !vecType.getElementType().isF32())
      return failure();
 
    // Unroll vector rsqrt into scalar operations
    int numElements = getVectorLaneSize(vecType);
    Value result = LLVM::PoisonOp::create(rewriter, loc, vecType);
 
    for (int i = 0; i < numElements; ++i) {
      // Extract element i
      auto indexCst = LLVM::ConstantOp::create(
          rewriter, loc, rewriter.getI64Type(), rewriter.getI64IntegerAttr(i));
      auto extractedElem = LLVM::ExtractElementOp::create(
          rewriter, loc, adaptor.getOperand(), indexCst);
 
      // Call xllvm.intr.aie2p.invsqrt on the scalar
      auto rsqrtResult = xllvm::InvsqrtAIE2pIntrOp::create(
          rewriter, loc, rewriter.getF32Type(), extractedElem);
 
      // Insert result back into vector
      result = LLVM::InsertElementOp::create(rewriter, loc, vecType, result,
                                             rsqrtResult, indexCst);
    }
 
    rewriter.replaceOp(rsqrtOp, result);
    return success();
  }
};
 
// A one-element vector built from a scalar is lowered to llvm.bitcast. As an
// insertelement, LLVM folds it with the extract of a lane into a shufflevector
// with a one-element result, which Peano's AIE combines do not handle.
static bool isOneLaneVectorOfScalar(Type resultType, Type sourceType) {
  auto vecType = dyn_cast<VectorType>(resultType);
  return vecType && !vecType.isScalable() && vecType.getRank() == 1 &&
         vecType.getNumElements() == 1 && !isa<VectorType>(sourceType);
}
 
static bool isOneLaneBroadcastOfScalar(vector::BroadcastOp op) {
  return isOneLaneVectorOfScalar(op.getType(), op.getSourceType());
}
 
static bool isOneLaneInsertOfScalar(vector::InsertOp op) {
  return isOneLaneVectorOfScalar(op.getType(),
                                 op.getValueToStore().getType()) &&
         !op.hasDynamicPosition();
}
 
class OneLaneBroadcastOpConversion
    : public mlir::ConvertOpToLLVMPattern<vector::BroadcastOp> {
public:
  using ConvertOpToLLVMPattern<vector::BroadcastOp>::ConvertOpToLLVMPattern;
 
  LogicalResult
  matchAndRewrite(vector::BroadcastOp bcastOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    if (!isOneLaneBroadcastOfScalar(bcastOp))
      return failure();
    rewriter.replaceOpWithNewOp<LLVM::BitcastOp>(bcastOp, bcastOp.getType(),
                                                 adaptor.getSource());
    return success();
  }
};
 
class OneLaneInsertOpConversion
    : public mlir::ConvertOpToLLVMPattern<vector::InsertOp> {
public:
  using ConvertOpToLLVMPattern<vector::InsertOp>::ConvertOpToLLVMPattern;
 
  LogicalResult
  matchAndRewrite(vector::InsertOp insertOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    if (!isOneLaneInsertOfScalar(insertOp))
      return failure();
    rewriter.replaceOpWithNewOp<LLVM::BitcastOp>(insertOp, insertOp.getType(),
                                                 adaptor.getValueToStore());
    return success();
  }
};
 
// Convert arith.divf for vector<N x f32> to unrolled scalar divisions
// Uses a noinline helper function call as a barrier to prevent LLVM
// re-vectorization. Scalar f32 divisions are handled by downstream passes.
class FdivOpConversion : public mlir::ConvertOpToLLVMPattern<arith::DivFOp> {
public:
  using ConvertOpToLLVMPattern<arith::DivFOp>::ConvertOpToLLVMPattern;
 
  FdivOpConversion(const LLVMTypeConverter &typeConverter, StringRef device)
      : ConvertOpToLLVMPattern(typeConverter), deviceName(device.str()) {}
 
  std::string deviceName;
 
  LogicalResult
  matchAndRewrite(arith::DivFOp divOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto loc = divOp.getLoc();
    auto lhsType = adaptor.getLhs().getType();
 
    // Only handle vector<N x f32> fdiv
    // Scalar f32 fdiv is handled by downstream passes
    auto vecType = dyn_cast<VectorType>(lhsType);
    if (!vecType || !vecType.getElementType().isF32())
      return failure();
 
    auto rhsType = adaptor.getRhs().getType();
    auto rhsVecType = dyn_cast<VectorType>(rhsType);
    if (!rhsVecType || rhsVecType != vecType)
      return failure();
 
    // For AIE2P, implement a/b as a * inv(b) using the hardware reciprocal
    // intrinsic. For AIE2, use scalar fdiv helper function.
    auto module = divOp->getParentOfType<ModuleOp>();
    auto f32Ty = rewriter.getF32Type();
 
    // Select device-specific body for the scalar helper function.
    // AIE2P: use inv(b) * a (hardware reciprocal intrinsic is reliable).
    // AIE2: use scalar fdiv directly.
    std::function<void(OpBuilder &, Location, ValueRange)> bodyBuilder;
    if (deviceName == "aie2p") {
      bodyBuilder = [](OpBuilder &builder, Location loc, ValueRange args) {
        auto invResult = xllvm::InvAIE2pIntrOp::create(
            builder, loc, builder.getF32Type(), args[1]);
        auto mulResult =
            arith::MulFOp::create(builder, loc, args[0], invResult);
        LLVM::ReturnOp::create(builder, loc, ValueRange{mulResult});
      };
    } else {
      bodyBuilder = [](OpBuilder &builder, Location loc, ValueRange args) {
        auto divResult = arith::DivFOp::create(builder, loc, args[0], args[1]);
        LLVM::ReturnOp::create(builder, loc, ValueRange{divResult});
      };
    }
 
    // Get or create the noinline scalar helper (acts as barrier to prevent
    // LLVM from re-vectorizing the scalar ops).
    auto helperFunc =
        getOrCreateScalarHelperFunc(module, rewriter, "fdiv", deviceName,
                                    /*argTypes=*/{f32Ty, f32Ty},
                                    /*resultType=*/f32Ty, bodyBuilder);
 
    // Unroll vector fdiv into scalar helper function calls.
    int numElements = getVectorLaneSize(vecType);
    Value result = LLVM::PoisonOp::create(rewriter, loc, vecType);
 
    for (int i = 0; i < numElements; ++i) {
      auto indexCst = LLVM::ConstantOp::create(
          rewriter, loc, rewriter.getI64Type(), rewriter.getI64IntegerAttr(i));
      auto lhsElem = LLVM::ExtractElementOp::create(rewriter, loc,
                                                    adaptor.getLhs(), indexCst);
      auto rhsElem = LLVM::ExtractElementOp::create(rewriter, loc,
                                                    adaptor.getRhs(), indexCst);
 
      auto divResult = LLVM::CallOp::create(rewriter, loc, helperFunc,
                                            ValueRange{lhsElem, rhsElem})
                           ->getResult(0);
 
      result = LLVM::InsertElementOp::create(rewriter, loc, vecType, result,
                                             divResult, indexCst);
    }
 
    rewriter.replaceOp(divOp, result);
    return success();
  }
};
 
void populateAIEVecToLLVMCommonConversionPatterns(
    mlir::LLVMTypeConverter &converter, mlir::RewritePatternSet &patterns) {
  // clang-format off
  // Patterns that work for all backends (AIE1, AIE2, AIE2p)
  patterns.add<AddOpConversion,
               SubOpConversion,
               FMAOpConversion,
               MulOpConversion,
               UPDOpConversion,
               SelectOpConversion,
               PackOpConversion,
               UnpackOpConversion,
               BroadcastOpConversion,
               BroadcastScalarOpConversion,
               MatMulOpConversion,
               ShuffleOpConversion>(converter);
  // clang-format on
}
 
static void populateAIEVecToLLVMAIE2ConversionPatterns(
    mlir::LLVMTypeConverter &converter, mlir::RewritePatternSet &patterns,
    Aie2Fp32Emulation aie2Fp32EmulationOption) {
  // Patterns specific to AIE2 backend
  patterns.add<AddElemOpAIE2Conversion, SubElemOpAIE2Conversion>(converter);
  patterns.add<NegOpAIE2Conversion>(converter);
  patterns.add<MulElemOpConversion>(converter, aie2Fp32EmulationOption);
  patterns.add<UPSOpAIE2Conversion, SRSOpAIE2Conversion>(converter);
  patterns.add<ShiftOpConversion>(converter);
  patterns.add<MaxOpConversion, MinOpConversion>(converter);
  patterns.add<CmpOpConversion, SelOpConversion>(converter);
  patterns.add<ExtOpConversion>(converter);
  patterns.add<ExtractElemOpConversion>(converter);
  patterns.add<ConcatOpConversion>(converter);
  patterns.add<FMAElemOpConversion>(converter);
  patterns.add<FoldAIECastOps>(converter);
  patterns.add<FdivOpConversion>(converter, "aie2");
  patterns.add<OneLaneBroadcastOpConversion, OneLaneInsertOpConversion>(
      converter);
}
 
// AIE2p version of ExtractElemOp conversion using LLVM extractelement
class ExtractElemOpAIE2pConversion
    : public mlir::ConvertOpToLLVMPattern<aievec::ExtElemOp> {
public:
  using ConvertOpToLLVMPattern<aievec::ExtElemOp>::ConvertOpToLLVMPattern;
 
  LogicalResult
  matchAndRewrite(aievec::ExtElemOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Location loc = op.getLoc();
 
    // AIE2p doesn't have dedicated vextract intrinsics, so use LLVM
    // extractelement
    Value extracted = LLVM::ExtractElementOp::create(
        rewriter, loc, adaptor.getSource(), adaptor.getIndex());
 
    rewriter.replaceOp(op, extracted);
    return success();
  }
};
 
// AIE2p version of ConcatOp conversion using vector.shuffle
class ConcatOpAIE2pConversion
    : public mlir::ConvertOpToLLVMPattern<aievec::ConcatOp> {
public:
  using ConvertOpToLLVMPattern<aievec::ConcatOp>::ConvertOpToLLVMPattern;
 
  // AIE2p has no concat intrinsic: the sources are joined two at a time with
  // vector.shuffle. When a round has an odd number, the last one is first
  // widened with poison to the width of the others.
  LogicalResult
  matchAndRewrite(aievec::ConcatOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Location loc = op.getLoc();
    SmallVector<Value> parts(adaptor.getSources());
    if (parts.empty()) {
      op.emitWarning() << "aievec.concat with no sources is not supported.\n";
      return failure();
    }
    auto lanesOf = [](Value v) {
      return cast<VectorType>(v.getType()).getNumElements();
    };
    while (parts.size() > 1) {
      SmallVector<Value> next;
      for (size_t i = 0; i + 1 < parts.size(); i += 2) {
        Value lhs = parts[i], rhs = parts[i + 1];
        int64_t lhsLanes = lanesOf(lhs), rhsLanes = lanesOf(rhs);
        if (rhsLanes < lhsLanes) {
          SmallVector<int64_t> padMask(lhsLanes, -1);
          std::iota(padMask.begin(), padMask.begin() + rhsLanes, 0);
          rhs = vector::ShuffleOp::create(rewriter, loc, rhs, rhs, padMask);
        }
        SmallVector<int64_t> mask(lhsLanes + rhsLanes);
        std::iota(mask.begin(), mask.end(), 0);
        next.push_back(
            vector::ShuffleOp::create(rewriter, loc, lhs, rhs, mask));
      }
      if (parts.size() % 2)
        next.push_back(parts.back());
      parts = std::move(next);
    }
    rewriter.replaceOp(op, parts.front());
    return success();
  }
};
 
void populateAIEVecToLLVMAIE2pConversionPatterns(
    mlir::LLVMTypeConverter &converter, mlir::RewritePatternSet &patterns) {
  // Takes precedence over the common pattern, which emits the AIE2 intrinsic.
  patterns.add<ShuffleOpAIE2pConversion>(converter, /*benefit=*/2);
  // Patterns specific to AIE2p backend
  patterns.add<AddElemOpAIE2pConversion, SubElemOpAIE2pConversion>(converter);
  patterns.add<NegOpAIE2pConversion>(converter);
  patterns.add<MulElemOpAIE2pConversion>(converter);
  patterns.add<FMAElemOpAIE2pConversion>(converter);
  patterns.add<UPSOpAIE2pConversion, SRSOpAIE2pConversion>(converter);
  patterns.add<MatMulOpAIE2pConversion>(converter);
  patterns.add<ShiftOpAIE2pConversion>(converter);
  patterns.add<MaxOpAIE2pConversion, MinOpAIE2pConversion>(converter);
  patterns.add<CmpOpAIE2pConversion, SelOpAIE2pConversion>(converter);
  patterns.add<ExtOpAIE2pConversion>(converter);
  patterns.add<ExtractElemOpAIE2pConversion>(converter);
  patterns.add<ConcatOpAIE2pConversion>(converter);
  // Higher benefit so it matches before the generic stub UnpackOpConversion
  // returns failure for i4->i8 inputs.
  patterns.add<UnpackI4OpAIE2pConversion>(converter, /*benefit=*/2);
  patterns.add<ExpOpAIE2pConversion>(converter);
  patterns.add<
      AccFloatUnaryOpAIE2pConversion<aievec::TanhOp, xllvm::TanhAIE2pIntrOp>,
      AccFloatUnaryOpAIE2pConversion<aievec::Exp2Op, xllvm::Exp2AIE2pIntrOp>>(
      converter);
  patterns.add<InvOpAIE2pConversion>(converter);
  patterns.add<BroadcastScalarOpAIE2pConversion>(converter);
  patterns.add<RsqrtOpAIE2pConversion>(converter);
  patterns.add<FdivOpConversion>(converter, "aie2p");
  patterns.add<FoldAIECastOpsAIE2p>(converter);
  patterns.add<OneLaneBroadcastOpConversion, OneLaneInsertOpConversion>(
      converter);
}
 
void populateAIEVecToLLVMConversionPatterns(
    mlir::LLVMTypeConverter &converter, mlir::RewritePatternSet &patterns,
    Aie2Fp32Emulation aie2Fp32EmulationOption, StringRef aieTarget) {
  populateAIEVecToLLVMCommonConversionPatterns(converter, patterns);
  if (aieTarget == "aie2p")
    populateAIEVecToLLVMAIE2pConversionPatterns(converter, patterns);
  else
    populateAIEVecToLLVMAIE2ConversionPatterns(converter, patterns,
                                               aie2Fp32EmulationOption);
}
 
// Configure legalization rules shared by AIE2 and AIE2p
static void configureAIEVecToLLVMLegalizations(LLVMConversionTarget &target) {
  // Vector f32 divf is illegal (needs unrolling to scalar divf)
  // Scalar f32 divf is legal (handled by downstream passes)
  target.addDynamicallyLegalOp<arith::DivFOp>([](arith::DivFOp divOp) {
    auto resultType = divOp.getType();
    if (auto vecType = dyn_cast<VectorType>(resultType)) {
      // Vector f32 divf is illegal and needs conversion
      return !vecType.getElementType().isF32();
    }
    // Scalar divf is legal
    return true;
  });
 
  // One-element vectors built from a scalar are illegal (lowered to bitcast)
  target.addDynamicallyLegalOp<vector::BroadcastOp>(
      [](vector::BroadcastOp op) { return !isOneLaneBroadcastOfScalar(op); });
  target.addDynamicallyLegalOp<vector::InsertOp>(
      [](vector::InsertOp op) { return !isOneLaneInsertOfScalar(op); });
}
 
struct ConvertAIEVecToLLVMPass
    : xilinx::impl::ConvertAIEVecToLLVMBase<ConvertAIEVecToLLVMPass> {
  ConvertAIEVecToLLVMPass() = default;
  ConvertAIEVecToLLVMPass(const xilinx::ConvertAIEVecToLLVMOptions &options) {
    aieTarget = options.aieTarget;
    aie2Fp32Emulation = options.aie2Fp32Emulation;
  }
 
  void runOnOperation() override {
    RewritePatternSet patterns(&getContext());
    LLVMTypeConverter converter(&getContext());
 
    // Don't convert vector types, we want to handle multi-dimensional
    // vector on our own.
    converter.addConversion(
        [&](VectorType type) -> std::optional<Type> { return type; });
 
    populateAIEVecToLLVMConversionPatterns(converter, patterns,
                                           aie2Fp32Emulation, aieTarget);
 
    LLVMConversionTarget target(getContext());
    target.addIllegalDialect<xilinx::aievec::AIEVecDialect,
                             xilinx::aievec::aie1::AIEVecAIE1Dialect>();
    target.addLegalDialect<arith::ArithDialect, vector::VectorDialect,
                           xilinx::xllvm::XLLVMDialect, ub::UBDialect>();
 
    // Configure legalizations for AIE2/AIE2p
    configureAIEVecToLLVMLegalizations(target);
 
    if (failed(applyPartialConversion(getOperation(), target,
                                      std::move(patterns))))
      signalPassFailure();
  }
};
 
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
createConvertAIEVecToLLVMPass() {
  return std::make_unique<ConvertAIEVecToLLVMPass>();
}
 
std::unique_ptr<mlir::OperationPass<mlir::ModuleOp>>
createConvertAIEVecToLLVMPass(
    const xilinx::ConvertAIEVecToLLVMOptions &options) {
  return std::make_unique<ConvertAIEVecToLLVMPass>(options);
}
 
} // namespace xilinx::aievec