//===-- XeGPUToXeVM.cpp - XeGPU to XeVM dialect conversion ------*- C++ -*-===//
//
// This file is licensed under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//

#include "mlir/Conversion/XeGPUToXeVM/XeGPUToXeVM.h"
#include "mlir/Dialect/LLVMIR/LLVMTypes.h"
#include "mlir/Dialect/LLVMIR/XeVMDialect.h"

#include "mlir/Conversion/LLVMCommon/Pattern.h"
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/Index/IR/IndexDialect.h"
#include "mlir/Dialect/Index/IR/IndexOps.h"
#include "mlir/Dialect/LLVMIR/FunctionCallUtils.h"
#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/IR/SCF.h"
#include "mlir/Dialect/SCF/Transforms/Patterns.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/XeGPU/IR/XeGPU.h"
#include "mlir/Dialect/XeGPU/Utils/XeGPUUtils.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Support/LLVM.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/Support/FormatVariadic.h"

#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Types.h"

#include "llvm/ADT/TypeSwitch.h"

#include <numeric>

namespace mlir {
#define GEN_PASS_DEF_CONVERTXEGPUTOXEVMPASS
#include "mlir/Conversion/Passes.h.inc"
} // namespace mlir

using namespace mlir;

namespace {

// TODO: Below are uArch dependent values, should move away from hardcoding
static constexpr int32_t systolicDepth{8};
static constexpr int32_t executionSize{16};

// Offsets to individual fields of the 8xi32 layout nd tensor descriptor.
enum class NdTdescOffset : uint32_t {
  BasePtr = 0,    // Base pointer (i64)
  BaseShapeW = 2, // Base shape width (i32)
  BaseShapeH = 3, // Base shape height (i32)
  BasePitch = 4,  // Base pitch (i32)
};

static int32_t getNumericXeVMAddrSpace(xegpu::MemorySpace xeGpuMemspace) {
  switch (xeGpuMemspace) {
  case xegpu::MemorySpace::Global:
    return static_cast<int>(xevm::AddrSpace::GLOBAL);
  case xegpu::MemorySpace::SLM:
    return static_cast<int>(xevm::AddrSpace::SHARED);
  }
  llvm_unreachable("Unknown XeGPU memory space");
}

/// Checks if the given MemRefType refers to shared memory.
static bool isSharedMemRef(const MemRefType &memrefTy) {
  Attribute attr = memrefTy.getMemorySpace();
  if (!attr)
    return false;
  if (auto intAttr = llvm::dyn_cast<IntegerAttr>(attr))
    return intAttr.getInt() == static_cast<int>(xevm::AddrSpace::SHARED);
  if (auto xevmSpace = llvm::dyn_cast<xevm::AddrSpaceAttr>(attr))
    return xevmSpace.getValue() == xevm::AddrSpace::SHARED;
  return gpu::GPUDialect::isWorkgroupMemoryAddressSpace(attr);
}

// Get same bitwidth flat vector type of new element type.
static VectorType encodeVectorTypeTo(VectorType currentVecType,
                                     Type toElemType) {
  auto elemType = currentVecType.getElementType();
  auto currentBitWidth = elemType.getIntOrFloatBitWidth();
  auto newBitWidth = toElemType.getIntOrFloatBitWidth();
  const int size =
      currentVecType.getNumElements() * currentBitWidth / newBitWidth;
  return VectorType::get(size, toElemType);
}

static xevm::LoadCacheControl
translateLoadXeGPUCacheHint(std::optional<xegpu::CachePolicy> L1hint,
                            std::optional<xegpu::CachePolicy> L3hint) {
  auto L1hintVal = L1hint.value_or(xegpu::CachePolicy::UNCACHED);
  auto L3hintVal = L3hint.value_or(xegpu::CachePolicy::UNCACHED);
  switch (L1hintVal) {
  case xegpu::CachePolicy::CACHED:
    if (L3hintVal == xegpu::CachePolicy::CACHED)
      return xevm::LoadCacheControl::L1C_L2UC_L3C;
    else if (L3hintVal == xegpu::CachePolicy::UNCACHED)
      return xevm::LoadCacheControl::L1C_L2UC_L3UC;
    else
      llvm_unreachable("Unsupported cache control.");
  case xegpu::CachePolicy::UNCACHED:
    if (L3hintVal == xegpu::CachePolicy::CACHED)
      return xevm::LoadCacheControl::L1UC_L2UC_L3C;
    else if (L3hintVal == xegpu::CachePolicy::UNCACHED)
      return xevm::LoadCacheControl::L1UC_L2UC_L3UC;
    else
      llvm_unreachable("Unsupported cache control.");
  case xegpu::CachePolicy::STREAMING:
    if (L3hintVal == xegpu::CachePolicy::CACHED)
      return xevm::LoadCacheControl::L1S_L2UC_L3C;
    else if (L3hintVal == xegpu::CachePolicy::UNCACHED)
      return xevm::LoadCacheControl::L1S_L2UC_L3UC;
    else
      llvm_unreachable("Unsupported cache control.");
  case xegpu::CachePolicy::READ_INVALIDATE:
    return xevm::LoadCacheControl::INVALIDATE_READ;
  default:
    llvm_unreachable("Unsupported cache control.");
  }
}

static xevm::StoreCacheControl
translateStoreXeGPUCacheHint(std::optional<xegpu::CachePolicy> L1hint,
                             std::optional<xegpu::CachePolicy> L3hint) {
  auto L1hintVal = L1hint.value_or(xegpu::CachePolicy::UNCACHED);
  auto L3hintVal = L3hint.value_or(xegpu::CachePolicy::UNCACHED);
  switch (L1hintVal) {
  case xegpu::CachePolicy::UNCACHED:
    if (L3hintVal == xegpu::CachePolicy::UNCACHED)
      return xevm::StoreCacheControl::L1UC_L2UC_L3UC;
    else if (L3hintVal == xegpu::CachePolicy::WRITE_BACK)
      return xevm::StoreCacheControl::L1UC_L2UC_L3WB;
    else
      llvm_unreachable("Unsupported cache control.");
  case xegpu::CachePolicy::STREAMING:
    if (L3hintVal == xegpu::CachePolicy::UNCACHED)
      return xevm::StoreCacheControl::L1S_L2UC_L3UC;
    else if (L3hintVal == xegpu::CachePolicy::WRITE_BACK)
      return xevm::StoreCacheControl::L1S_L2UC_L3WB;
    else
      llvm_unreachable("Unsupported cache control.");
  case xegpu::CachePolicy::WRITE_BACK:
    if (L3hintVal == xegpu::CachePolicy::UNCACHED)
      return xevm::StoreCacheControl::L1WB_L2UC_L3UC;
    else if (L3hintVal == xegpu::CachePolicy::WRITE_BACK)
      return xevm::StoreCacheControl::L1WB_L2UC_L3WB;
    else
      llvm_unreachable("Unsupported cache control.");
  case xegpu::CachePolicy::WRITE_THROUGH:
    if (L3hintVal == xegpu::CachePolicy::UNCACHED)
      return xevm::StoreCacheControl::L1WT_L2UC_L3UC;
    else if (L3hintVal == xegpu::CachePolicy::WRITE_BACK)
      return xevm::StoreCacheControl::L1WT_L2UC_L3WB;
    else
      llvm_unreachable("Unsupported cache control.");
  default:
    llvm_unreachable("Unsupported cache control.");
  }
}

//
// Note:
// Block operations for tile of sub byte element types are handled by
// emulating with larger element types.
// Tensor descriptor are keep intact and only ops consuming them are
// emulated
//

class CreateNdDescToXeVMPattern
    : public OpConversionPattern<xegpu::CreateNdDescOp> {
  using OpConversionPattern::OpConversionPattern;
  LogicalResult
  matchAndRewrite(xegpu::CreateNdDescOp op,
                  xegpu::CreateNdDescOp::Adaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    SmallVector<OpFoldResult> mixedOffsets = op.getMixedOffsets();
    if (mixedOffsets.size() != 0)
      return rewriter.notifyMatchFailure(op, "Offsets not supported.");
    auto loc = op.getLoc();
    auto source = op.getSource();
    // Op is lowered to a code sequence that populates payload.
    // Payload is a 8xi32 vector. Offset to individual fields are defined in
    // NdTdescOffset enum.
    Type payloadElemTy = rewriter.getI32Type();
    VectorType payloadTy = VectorType::get(8, payloadElemTy);
    Type i64Ty = rewriter.getI64Type();
    // 4xi64 view is used for inserting the base pointer.
    VectorType payloadI64Ty = VectorType::get(4, i64Ty);
    // Initialize payload to zero.
    Value payload = arith::ConstantOp::create(
        rewriter, loc,
        DenseElementsAttr::get(payloadTy, IntegerAttr::get(payloadElemTy, 0)));

    Value baseAddr;
    Value baseShapeW;
    Value baseShapeH;

    // Source can be a memref or a pointer (ui64, ui32, i64 or i32).
    SmallVector<OpFoldResult> mixedSizes = op.getMixedSizes();
    SmallVector<OpFoldResult> mixedStrides = op.getMixedStrides();
    // Descriptor shape is expected to be 2D.
    int64_t rank = mixedSizes.size();
    auto sourceTy = source.getType();
    auto sourceMemrefTy = dyn_cast<MemRefType>(sourceTy);
    // If source is a memref, we need to extract the aligned pointer as index.
    // Pointer type is passed as i32 or i64 by type converter.
    if (sourceMemrefTy) {
      if (!sourceMemrefTy.hasRank()) {
        return rewriter.notifyMatchFailure(op, "Expected ranked Memref.");
      }
      // Access adaptor after failure check to avoid rolling back generated code
      // for materialization cast.
      baseAddr = adaptor.getSource();
    } else {
      baseAddr = adaptor.getSource();
      if (baseAddr.getType() != i64Ty) {
        // Pointer type may be i32. Cast to i64 if needed.
        baseAddr = arith::ExtUIOp::create(rewriter, loc, i64Ty, baseAddr);
      }
    }
    // 1D tensor descriptor is just the base address.
    if (rank == 1) {
      rewriter.replaceOp(op, baseAddr);
      return success();
    }
    // Utility for creating offset values from op fold result.
    auto createOffset = [&](SmallVector<OpFoldResult> &ofrVec,
                            unsigned idx) -> Value {
      Value val = getValueOrCreateConstantIntOp(rewriter, loc, ofrVec[idx]);
      val = getValueOrCreateCastToIndexLike(rewriter, loc, payloadElemTy, val);
      return val;
    };
    // Get shape values from op fold results.
    baseShapeW = createOffset(mixedSizes, 1);
    baseShapeH = createOffset(mixedSizes, 0);
    // Get pitch value from op fold results.
    Value basePitch = createOffset(mixedStrides, 0);
    // Populate payload.
    Value payLoadAsI64 =
        vector::BitCastOp::create(rewriter, loc, payloadI64Ty, payload);
    payLoadAsI64 =
        vector::InsertOp::create(rewriter, loc, baseAddr, payLoadAsI64,
                                 static_cast<int>(NdTdescOffset::BasePtr));
    payload = vector::BitCastOp::create(rewriter, loc, payloadTy, payLoadAsI64);
    payload =
        vector::InsertOp::create(rewriter, loc, baseShapeW, payload,
                                 static_cast<int>(NdTdescOffset::BaseShapeW));
    payload =
        vector::InsertOp::create(rewriter, loc, baseShapeH, payload,
                                 static_cast<int>(NdTdescOffset::BaseShapeH));
    payload =
        vector::InsertOp::create(rewriter, loc, basePitch, payload,
                                 static_cast<int>(NdTdescOffset::BasePitch));
    rewriter.replaceOp(op, payload);
    return success();
  }
};

template <
    typename OpType,
    typename = std::enable_if_t<llvm::is_one_of<
        OpType, xegpu::LoadNdOp, xegpu::StoreNdOp, xegpu::PrefetchNdOp>::value>>
class LoadStorePrefetchNdToXeVMPattern : public OpConversionPattern<OpType> {
  using OpConversionPattern<OpType>::OpConversionPattern;
  LogicalResult
  matchAndRewrite(OpType op, typename OpType::Adaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto mixedOffsets = op.getMixedOffsets();
    int64_t opOffsetsSize = mixedOffsets.size();
    auto loc = op.getLoc();
    auto ctxt = rewriter.getContext();

    auto tdesc = adaptor.getTensorDesc();
    auto tdescTy = op.getTensorDescType();
    auto tileRank = tdescTy.getRank();
    if (opOffsetsSize != tileRank)
      return rewriter.notifyMatchFailure(
          op, "Expected offset rank to match descriptor rank.");
    auto elemType = tdescTy.getElementType();
    auto elemBitSize = elemType.getIntOrFloatBitWidth();
    bool isSubByte = elemBitSize < 8;
    uint64_t wScaleFactor = 1;

    if (!isSubByte && (elemBitSize % 8 != 0))
      return rewriter.notifyMatchFailure(
          op, "Expected element type bit width to be multiple of 8.");
    auto tileW = tdescTy.getDimSize(tileRank - 1);
    // For sub byte types, only 4bits are currently supported.
    if (isSubByte) {
      if (elemBitSize != 4)
        return rewriter.notifyMatchFailure(
            op, "Only sub byte types of 4bits are supported.");
      if (tileRank != 2)
        return rewriter.notifyMatchFailure(
            op, "Sub byte types are only supported for 2D tensor descriptors.");
      auto subByteFactor = 8 / elemBitSize;
      auto tileH = tdescTy.getDimSize(0);
      // Handle special case for packed load.
      if constexpr (std::is_same_v<OpType, xegpu::LoadNdOp>) {
        if (op.getPacked().value_or(false)) {
          // packed load is implemented as packed loads of 8bit elements.
          if (tileH == systolicDepth * 4 &&
              tileW == executionSize * subByteFactor) {
            // Usage case for loading as Matrix B with pack request.
            // source is assumed to pre-packed into 8bit elements
            // Emulate with 8bit loads with pack request.
            //   scaled_tileW = executionSize
            elemType = rewriter.getIntegerType(8);
            tileW = executionSize;
            wScaleFactor = subByteFactor;
          }
        }
      }
      // If not handled by packed load case above, handle other cases.
      if (wScaleFactor == 1) {
        auto sub16BitFactor = subByteFactor * 2;
        if (tileW == executionSize * sub16BitFactor) {
          // Usage case for loading as Matrix A operand
          // Emulate with 16bit loads/stores.
          //   scaled_tileW = executionSize
          elemType = rewriter.getIntegerType(16);
          tileW = executionSize;
          wScaleFactor = sub16BitFactor;
        } else {
          return rewriter.notifyMatchFailure(
              op, "Unsupported tile shape for sub byte types.");
        }
      }
      // recompute element bit size for emulation.
      elemBitSize = elemType.getIntOrFloatBitWidth();
    }

    // Get address space from tensor descriptor memory space.
    auto ptrTypeLLVM = LLVM::LLVMPointerType::get(
        ctxt, getNumericXeVMAddrSpace(tdescTy.getMemorySpace()));
    if (tileRank == 2) {
      // Compute element byte size.
      Value elemByteSize = arith::ConstantIntOp::create(
          rewriter, loc, rewriter.getI32Type(), elemBitSize / 8);
      VectorType payloadI64Ty = VectorType::get(4, rewriter.getI64Type());
      Value payLoadAsI64 =
          vector::BitCastOp::create(rewriter, loc, payloadI64Ty, tdesc);
      Value basePtr =
          vector::ExtractOp::create(rewriter, loc, payLoadAsI64,
                                    static_cast<int>(NdTdescOffset::BasePtr));
      Value baseShapeW = vector::ExtractOp::create(
          rewriter, loc, tdesc, static_cast<int>(NdTdescOffset::BaseShapeW));
      Value baseShapeH = vector::ExtractOp::create(
          rewriter, loc, tdesc, static_cast<int>(NdTdescOffset::BaseShapeH));
      Value basePitch = vector::ExtractOp::create(
          rewriter, loc, tdesc, static_cast<int>(NdTdescOffset::BasePitch));
      // Offsets are provided by the op.
      // convert them to i32.
      Value offsetW =
          getValueOrCreateConstantIntOp(rewriter, loc, mixedOffsets[1]);
      offsetW = getValueOrCreateCastToIndexLike(rewriter, loc,
                                                rewriter.getI32Type(), offsetW);
      Value offsetH =
          getValueOrCreateConstantIntOp(rewriter, loc, mixedOffsets[0]);
      offsetH = getValueOrCreateCastToIndexLike(rewriter, loc,
                                                rewriter.getI32Type(), offsetH);
      // Convert base pointer (i64) to LLVM pointer type.
      Value basePtrLLVM =
          LLVM::IntToPtrOp::create(rewriter, loc, ptrTypeLLVM, basePtr);
      // FIXME: width or pitch is not the same as baseShapeW it should be the
      // stride of the second to last dimension in row major layout.
      // Compute width in bytes.
      Value baseShapeWInBytes =
          arith::MulIOp::create(rewriter, loc, baseShapeW, elemByteSize);
      // Compute pitch in bytes.
      Value basePitchBytes =
          arith::MulIOp::create(rewriter, loc, basePitch, elemByteSize);

      if (wScaleFactor > 1) {
        // Scale offsetW, baseShapeWInBytes for sub byte emulation.
        // Note: tileW is already scaled above.
        Value wScaleFactorValLog2 = arith::ConstantIntOp::create(
            rewriter, loc, rewriter.getI32Type(), llvm::Log2_64(wScaleFactor));
        baseShapeWInBytes = arith::ShRSIOp::create(
            rewriter, loc, baseShapeWInBytes, wScaleFactorValLog2);
        basePitchBytes = arith::ShRSIOp::create(rewriter, loc, basePitchBytes,
                                                wScaleFactorValLog2);
        offsetW =
            arith::ShRSIOp::create(rewriter, loc, offsetW, wScaleFactorValLog2);
      }
      // Get tile height from the tensor descriptor type.
      auto tileH = tdescTy.getDimSize(0);
      // Get vblocks from the tensor descriptor type.
      int32_t vblocks = tdescTy.getArrayLength();
      if constexpr (std::is_same_v<OpType, xegpu::StoreNdOp>) {
        Value src = adaptor.getValue();
        // If store value is a scalar, get value from op instead of adaptor.
        // Adaptor might have optimized away single element vector
        if (src.getType().isIntOrFloat()) {
          src = op.getValue();
        }
        VectorType srcVecTy = dyn_cast<VectorType>(src.getType());
        if (!srcVecTy)
          return rewriter.notifyMatchFailure(
              op, "Expected store value to be a vector type.");
        // Get flat vector type of integer type with matching element bit size.
        VectorType newSrcVecTy =
            encodeVectorTypeTo(srcVecTy, rewriter.getIntegerType(elemBitSize));
        if (srcVecTy != newSrcVecTy)
          src = vector::BitCastOp::create(rewriter, loc, newSrcVecTy, src);
        auto storeCacheControl =
            translateStoreXeGPUCacheHint(op.getL1Hint(), op.getL3Hint());
        xevm::BlockStore2dOp::create(
            rewriter, loc, basePtrLLVM, baseShapeWInBytes, baseShapeH,
            basePitchBytes, offsetW, offsetH, elemBitSize, tileW, tileH, src,
            xevm::StoreCacheControlAttr::get(ctxt, storeCacheControl));
        rewriter.eraseOp(op);
      } else {
        auto loadCacheControl =
            translateLoadXeGPUCacheHint(op.getL1Hint(), op.getL3Hint());
        if constexpr (std::is_same_v<OpType, xegpu::PrefetchNdOp>) {
          xevm::BlockPrefetch2dOp::create(
              rewriter, loc, basePtrLLVM, baseShapeWInBytes, baseShapeH,
              basePitchBytes, offsetW, offsetH, elemBitSize, tileW, tileH,
              vblocks, xevm::LoadCacheControlAttr::get(ctxt, loadCacheControl));
          rewriter.eraseOp(op);
        } else {
          VectorType dstVecTy = cast<VectorType>(op.getValue().getType());
          const bool vnni = op.getPacked().value_or(false);
          auto transposeValue = op.getTranspose();
          bool transpose =
              transposeValue.has_value() && transposeValue.value()[0] == 1;
          VectorType loadedTy = encodeVectorTypeTo(
              dstVecTy, vnni ? rewriter.getI32Type()
                             : rewriter.getIntegerType(elemBitSize));

          Value resultFlatVec = xevm::BlockLoad2dOp::create(
              rewriter, loc, loadedTy, basePtrLLVM, baseShapeWInBytes,
              baseShapeH, basePitchBytes, offsetW, offsetH, elemBitSize, tileW,
              tileH, vblocks, transpose, vnni,
              xevm::LoadCacheControlAttr::get(ctxt, loadCacheControl));
          resultFlatVec = vector::BitCastOp::create(
              rewriter, loc,
              encodeVectorTypeTo(loadedTy, dstVecTy.getElementType()),
              resultFlatVec);
          rewriter.replaceOp(op, resultFlatVec);
        }
      }
    } else {
      // 1D tensor descriptor.
      // `tdesc` represents base address as i64
      // Offset in number of elements, need to multiply by element byte size.
      // Compute byte offset.
      //   byteOffset = offset * elementByteSize
      Value offset =
          getValueOrCreateConstantIntOp(rewriter, loc, mixedOffsets[0]);
      offset = getValueOrCreateCastToIndexLike(rewriter, loc,
                                               rewriter.getI64Type(), offset);
      // Compute element byte size.
      Value elemByteSize = arith::ConstantIntOp::create(
          rewriter, loc, rewriter.getI64Type(), elemBitSize / 8);
      Value byteOffset =
          rewriter.createOrFold<arith::MulIOp>(loc, offset, elemByteSize);
      // Final address = basePtr + byteOffset
      Value finalAddrI64 = rewriter.createOrFold<arith::AddIOp>(
          loc, tdesc,
          getValueOrCreateCastToIndexLike(rewriter, loc, rewriter.getI64Type(),
                                          byteOffset));
      // Convert base pointer (i64) to LLVM pointer type.
      Value finalPtrLLVM =
          LLVM::IntToPtrOp::create(rewriter, loc, ptrTypeLLVM, finalAddrI64);
      if constexpr (std::is_same_v<OpType, xegpu::StoreNdOp>) {
        Value src = adaptor.getValue();
        // If store value is a scalar, get value from op instead of adaptor.
        // Adaptor might have optimized away single element vector
        if (src.getType().isIntOrFloat()) {
          src = op.getValue();
        }
        VectorType srcVecTy = dyn_cast<VectorType>(src.getType());
        if (!srcVecTy)
          return rewriter.notifyMatchFailure(
              op, "Expected store value to be a vector type.");
        // Get flat vector type of integer type with matching element bit size.
        VectorType newSrcVecTy =
            encodeVectorTypeTo(srcVecTy, rewriter.getIntegerType(elemBitSize));
        if (srcVecTy != newSrcVecTy)
          src = vector::BitCastOp::create(rewriter, loc, newSrcVecTy, src);
        auto storeCacheControl =
            translateStoreXeGPUCacheHint(op.getL1Hint(), op.getL3Hint());
        rewriter.replaceOpWithNewOp<xevm::BlockStoreOp>(
            op, finalPtrLLVM, src,
            xevm::StoreCacheControlAttr::get(ctxt, storeCacheControl));
      } else if constexpr (std::is_same_v<OpType, xegpu::LoadNdOp>) {
        auto loadCacheControl =
            translateLoadXeGPUCacheHint(op.getL1Hint(), op.getL3Hint());
        VectorType resTy = cast<VectorType>(op.getValue().getType());
        VectorType loadedTy =
            encodeVectorTypeTo(resTy, rewriter.getIntegerType(elemBitSize));
        Value load = xevm::BlockLoadOp::create(
            rewriter, loc, loadedTy, finalPtrLLVM,
            xevm::LoadCacheControlAttr::get(ctxt, loadCacheControl));
        if (loadedTy != resTy)
          load = vector::BitCastOp::create(rewriter, loc, resTy, load);
        rewriter.replaceOp(op, load);
      } else {
        return rewriter.notifyMatchFailure(
            op, "Unsupported operation: xegpu.prefetch_nd with tensor "
                "descriptor rank == 1");
      }
    }
    return success();
  }
};

// Add a builder that creates
// offset * elemByteSize + baseAddr
static Value addOffsetToBaseAddr(ConversionPatternRewriter &rewriter,
                                 Location loc, Value baseAddr, Value offset,
                                 int64_t elemByteSize) {
  Value byteSize = arith::ConstantIntOp::create(
      rewriter, loc, baseAddr.getType(), elemByteSize);
  Value byteOffset = arith::MulIOp::create(rewriter, loc, offset, byteSize);
  Value newAddr = arith::AddIOp::create(rewriter, loc, baseAddr, byteOffset);
  return newAddr;
}

template <typename OpType,
          typename = std::enable_if_t<llvm::is_one_of<
              OpType, xegpu::LoadGatherOp, xegpu::StoreScatterOp>::value>>
class LoadStoreToXeVMPattern : public OpConversionPattern<OpType> {
  using OpConversionPattern<OpType>::OpConversionPattern;
  LogicalResult
  matchAndRewrite(OpType op, typename OpType::Adaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Value offset = adaptor.getOffsets();
    if (!offset)
      return rewriter.notifyMatchFailure(op, "Expected offset to be provided.");
    auto loc = op.getLoc();
    auto ctxt = rewriter.getContext();
    auto tdescTy = op.getTensorDescType();
    Value basePtrI64;
    // Load result or Store valye Type can be vector or scalar.
    Type valOrResTy;
    if constexpr (std::is_same_v<OpType, xegpu::LoadGatherOp>)
      valOrResTy =
          this->getTypeConverter()->convertType(op.getResult().getType());
    else
      valOrResTy = adaptor.getValue().getType();
    VectorType valOrResVecTy = dyn_cast<VectorType>(valOrResTy);
    bool hasScalarVal = !valOrResVecTy;
    int64_t elemBitWidth =
        hasScalarVal ? valOrResTy.getIntOrFloatBitWidth()
                     : valOrResVecTy.getElementType().getIntOrFloatBitWidth();
    // Element type must be multiple of 8 bits.
    if (elemBitWidth % 8 != 0)
      return rewriter.notifyMatchFailure(
          op, "Expected element type bit width to be multiple of 8.");
    int64_t elemByteSize = elemBitWidth / 8;
    // Default memory space is global.
    LLVM::LLVMPointerType ptrTypeLLVM = LLVM::LLVMPointerType::get(
        ctxt, getNumericXeVMAddrSpace(xegpu::MemorySpace::Global));
    // If tensor descriptor is available, we use its memory space.
    if (tdescTy)
      ptrTypeLLVM = LLVM::LLVMPointerType::get(
          ctxt, getNumericXeVMAddrSpace(tdescTy.getMemorySpace()));
    // Base pointer can come from source (load) or dest (store).
    // If they are memrefs, we use their memory space.
    if constexpr (std::is_same_v<OpType, xegpu::LoadGatherOp>) {
      basePtrI64 = adaptor.getSource();
      if (auto memRefTy = dyn_cast<MemRefType>(op.getSource().getType())) {
        auto addrSpace = memRefTy.getMemorySpaceAsInt();
        if (addrSpace != 0)
          ptrTypeLLVM = LLVM::LLVMPointerType::get(ctxt, addrSpace);
      }
    } else {
      basePtrI64 = adaptor.getDest();
      if (auto memRefTy = dyn_cast<MemRefType>(op.getDest().getType())) {
        auto addrSpace = memRefTy.getMemorySpaceAsInt();
        if (addrSpace != 0)
          ptrTypeLLVM = LLVM::LLVMPointerType::get(ctxt, addrSpace);
      }
    }
    // Base pointer is passed as i32 or i64 by adaptor, cast to i64 if needed.
    if (basePtrI64.getType() != rewriter.getI64Type()) {
      basePtrI64 = arith::ExtUIOp::create(rewriter, loc, rewriter.getI64Type(),
                                          basePtrI64);
    }
    Value mask = adaptor.getMask();
    if (dyn_cast<VectorType>(offset.getType())) {
      // Offset needs be scalar. Single element vector is converted to scalar
      // by type converter.
      return rewriter.notifyMatchFailure(op, "Expected offset to be a scalar.");
    } else {
      // If offset is provided, we add them to the base pointer.
      // Offset is in number of elements, we need to multiply by
      // element byte size.
      basePtrI64 =
          addOffsetToBaseAddr(rewriter, loc, basePtrI64, offset, elemByteSize);
    }
    // Convert base pointer (i64) to LLVM pointer type.
    Value basePtrLLVM =
        LLVM::IntToPtrOp::create(rewriter, loc, ptrTypeLLVM, basePtrI64);

    Value maskForLane;
    VectorType maskVecTy = dyn_cast<VectorType>(mask.getType());
    if (maskVecTy) {
      // Mask needs be scalar. Single element vector is converted to scalar by
      // type converter.
      return rewriter.notifyMatchFailure(op, "Expected mask to be a scalar.");
    } else
      maskForLane = mask;
    if constexpr (std::is_same_v<OpType, xegpu::LoadGatherOp>) {
      scf::IfOp ifOp = scf::IfOp::create(rewriter, loc, {valOrResTy},
                                         maskForLane, true, true);
      // If mask is true,- then clause - load from memory and yield.
      rewriter.setInsertionPointToStart(&ifOp.getThenRegion().front());
      if (!hasScalarVal)
        valOrResTy = VectorType::get({valOrResVecTy.getNumElements()},
                                     valOrResVecTy.getElementType());
      Value loaded =
          LLVM::LoadOp::create(rewriter, loc, valOrResTy, basePtrLLVM);
      // Set cache control attribute on the load operation.
      loaded.getDefiningOp()->setAttr(
          "cache_control", xevm::LoadCacheControlAttr::get(
                               ctxt, translateLoadXeGPUCacheHint(
                                         op.getL1Hint(), op.getL3Hint())));
      scf::YieldOp::create(rewriter, loc, ValueRange{loaded});
      rewriter.setInsertionPointToStart(&ifOp.getElseRegion().front());
      // If mask is false - else clause -yield a vector of zeros.
      auto eTy = hasScalarVal ? valOrResTy : valOrResVecTy.getElementType();
      TypedAttr eVal;
      if (eTy.isFloat())
        eVal = FloatAttr::get(eTy, 0.0);
      else
        eVal = IntegerAttr::get(eTy, 0);
      if (hasScalarVal)
        loaded = arith::ConstantOp::create(rewriter, loc, eVal);
      else
        loaded = arith::ConstantOp::create(
            rewriter, loc, DenseElementsAttr::get(valOrResVecTy, eVal));
      scf::YieldOp::create(rewriter, loc, ValueRange{loaded});
      rewriter.replaceOp(op, ifOp.getResult(0));
    } else {
      // If mask is true, perform the store.
      scf::IfOp ifOp = scf::IfOp::create(rewriter, loc, maskForLane, false);
      auto body = ifOp.getBody();
      rewriter.setInsertionPointToStart(body);
      auto storeOp =
          LLVM::StoreOp::create(rewriter, loc, adaptor.getValue(), basePtrLLVM);
      // Set cache control attribute on the store operation.
      storeOp.getOperation()->setAttr(
          "cache_control", xevm::StoreCacheControlAttr::get(
                               ctxt, translateStoreXeGPUCacheHint(
                                         op.getL1Hint(), op.getL3Hint())));
      rewriter.eraseOp(op);
    }
    return success();
  }
};

class CreateMemDescOpPattern final
    : public OpConversionPattern<xegpu::CreateMemDescOp> {
public:
  using OpConversionPattern<xegpu::CreateMemDescOp>::OpConversionPattern;
  LogicalResult
  matchAndRewrite(xegpu::CreateMemDescOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {

    rewriter.replaceOp(op, adaptor.getSource());
    return success();
  }
};

template <typename OpType,
          typename = std::enable_if_t<llvm::is_one_of<
              OpType, xegpu::LoadMatrixOp, xegpu::StoreMatrixOp>::value>>
class LoadStoreMatrixToXeVMPattern : public OpConversionPattern<OpType> {
  using OpConversionPattern<OpType>::OpConversionPattern;
  LogicalResult
  matchAndRewrite(OpType op, typename OpType::Adaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {

    SmallVector<OpFoldResult> offsets = op.getMixedOffsets();
    if (offsets.empty())
      return rewriter.notifyMatchFailure(op, "Expected offset to be provided.");

    auto loc = op.getLoc();
    auto ctxt = rewriter.getContext();
    Value baseAddr32 = adaptor.getMemDesc();
    Value mdescVal = op.getMemDesc();
    // Load result or Store value Type can be vector or scalar.
    Type dataTy;
    if constexpr (std::is_same_v<OpType, xegpu::LoadMatrixOp>) {
      Type resType = op.getResult().getType();
      // Some transforms may leave unit dimension in the 2D vector, adaptors do
      // not catch it for results.
      if (auto vecType = dyn_cast<VectorType>(resType)) {
        assert(llvm::count_if(vecType.getShape(),
                              [](int64_t d) { return d != 1; }) <= 1 &&
               "Expected either 1D vector or nD with unit dimensions");
        resType = VectorType::get({vecType.getNumElements()},
                                  vecType.getElementType());
      }
      dataTy = resType;
    } else
      dataTy = adaptor.getData().getType();
    VectorType valOrResVecTy = dyn_cast<VectorType>(dataTy);
    if (!valOrResVecTy)
      valOrResVecTy = VectorType::get(1, dataTy);

    int64_t elemBitWidth =
        valOrResVecTy.getElementType().getIntOrFloatBitWidth();
    // Element type must be multiple of 8 bits.
    if (elemBitWidth % 8 != 0)
      return rewriter.notifyMatchFailure(
          op, "Expected element type bit width to be multiple of 8.");
    int64_t elemByteSize = elemBitWidth / 8;

    // Default memory space is SLM.
    LLVM::LLVMPointerType ptrTypeLLVM = LLVM::LLVMPointerType::get(
        ctxt, getNumericXeVMAddrSpace(xegpu::MemorySpace::SLM));

    auto mdescTy = cast<xegpu::MemDescType>(mdescVal.getType());

    Value linearOffset = mdescTy.getLinearOffsets(rewriter, loc, offsets);
    linearOffset = arith::IndexCastUIOp::create(
        rewriter, loc, rewriter.getI32Type(), linearOffset);
    Value basePtrI32 = addOffsetToBaseAddr(rewriter, loc, baseAddr32,
                                           linearOffset, elemByteSize);

    // convert base pointer (i32) to LLVM pointer type
    Value basePtrLLVM =
        LLVM::IntToPtrOp::create(rewriter, loc, ptrTypeLLVM, basePtrI32);

    if (op.getSubgroupBlockIoAttr()) {
      // if the attribute 'subgroup_block_io' is set to true, it lowers to
      // xevm.blockload

      Type intElemTy = rewriter.getIntegerType(elemBitWidth);
      VectorType intVecTy =
          VectorType::get(valOrResVecTy.getShape(), intElemTy);

      if constexpr (std::is_same_v<OpType, xegpu::LoadMatrixOp>) {
        Value loadOp =
            xevm::BlockLoadOp::create(rewriter, loc, intVecTy, basePtrLLVM);
        if (intVecTy != valOrResVecTy) {
          loadOp =
              vector::BitCastOp::create(rewriter, loc, valOrResVecTy, loadOp);
        }
        rewriter.replaceOp(op, loadOp);
      } else {
        Value dataToStore = adaptor.getData();
        if (valOrResVecTy != intVecTy) {
          dataToStore =
              vector::BitCastOp::create(rewriter, loc, intVecTy, dataToStore);
        }
        xevm::BlockStoreOp::create(rewriter, loc, basePtrLLVM, dataToStore,
                                   nullptr);
        rewriter.eraseOp(op);
      }
      return success();
    }

    if (valOrResVecTy.getNumElements() >= 1) {
      auto chipOpt = xegpu::getChipStr(op);
      if (!chipOpt || (*chipOpt != "pvc" && *chipOpt != "bmg")) {
        // the lowering for chunk load only works for pvc and bmg
        return rewriter.notifyMatchFailure(
            op, "The lowering is specific to pvc or bmg.");
      }
    }

    if constexpr (std::is_same_v<OpType, xegpu::LoadMatrixOp>) {
      // if the size of valOrResVecTy is 1, it lowers to a scalar load/store
      // operation. LLVM load/store does not support vector of size 1, so we
      // need to handle this case separately.
      auto scalarTy = valOrResVecTy.getElementType();
      LLVM::LoadOp loadOp;
      if (valOrResVecTy.getNumElements() == 1)
        loadOp = LLVM::LoadOp::create(rewriter, loc, scalarTy, basePtrLLVM);
      else
        loadOp =
            LLVM::LoadOp::create(rewriter, loc, valOrResVecTy, basePtrLLVM);
      rewriter.replaceOp(op, loadOp);
    } else {
      LLVM::StoreOp::create(rewriter, loc, adaptor.getData(), basePtrLLVM);
      rewriter.eraseOp(op);
    }
    return success();
  }
};

class PrefetchToXeVMPattern : public OpConversionPattern<xegpu::PrefetchOp> {
  using OpConversionPattern::OpConversionPattern;
  LogicalResult
  matchAndRewrite(xegpu::PrefetchOp op, xegpu::PrefetchOp::Adaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto loc = op.getLoc();
    auto ctxt = rewriter.getContext();
    auto tdescTy = op.getTensorDescType();
    Value basePtrI64 = adaptor.getSource();
    // Base pointer is passed as i32 or i64 by adaptor, cast to i64 if needed.
    if (basePtrI64.getType() != rewriter.getI64Type())
      basePtrI64 = arith::ExtUIOp::create(rewriter, loc, rewriter.getI64Type(),
                                          basePtrI64);
    Value offsets = adaptor.getOffsets();
    if (offsets) {
      VectorType offsetsVecTy = dyn_cast<VectorType>(offsets.getType());
      if (offsetsVecTy) {
        // Offset needs be scalar.
        return rewriter.notifyMatchFailure(op,
                                           "Expected offsets to be a scalar.");
      } else {
        int64_t elemBitWidth{0};
        int64_t elemByteSize;
        // Element byte size can come from three sources:
        if (tdescTy) {
          // If tensor descriptor is available, we use its element type to
          // determine element byte size.
          elemBitWidth = tdescTy.getElementType().getIntOrFloatBitWidth();
        } else if (auto memRefTy = dyn_cast<MemRefType>(op.getSourceType())) {
          // If memref is available, we use its element type to
          // determine element byte size.
          elemBitWidth = memRefTy.getElementType().getIntOrFloatBitWidth();
        } else {
          // Otherwise, we use the provided offset byte alignment.
          elemByteSize = *op.getOffsetAlignByte();
        }
        if (elemBitWidth != 0) {
          if (elemBitWidth % 8 != 0)
            return rewriter.notifyMatchFailure(
                op, "Expected element type bit width to be multiple of 8.");
          elemByteSize = elemBitWidth / 8;
        }
        basePtrI64 = addOffsetToBaseAddr(rewriter, loc, basePtrI64, offsets,
                                         elemByteSize);
      }
    }
    // Default memory space is global.
    LLVM::LLVMPointerType ptrTypeLLVM = LLVM::LLVMPointerType::get(
        ctxt, getNumericXeVMAddrSpace(xegpu::MemorySpace::Global));
    // If tensor descriptor is available, we use its memory space.
    if (tdescTy)
      ptrTypeLLVM = LLVM::LLVMPointerType::get(
          ctxt, getNumericXeVMAddrSpace(tdescTy.getMemorySpace()));
    // If source is a memref, we use its memory space.
    if (auto memRefTy = dyn_cast<MemRefType>(op.getSource().getType())) {
      auto addrSpace = memRefTy.getMemorySpaceAsInt();
      if (addrSpace != 0)
        ptrTypeLLVM = LLVM::LLVMPointerType::get(ctxt, addrSpace);
    }
    // Convert base pointer (i64) to LLVM pointer type.
    Value ptrLLVM =
        LLVM::IntToPtrOp::create(rewriter, loc, ptrTypeLLVM, basePtrI64);
    // Create the prefetch op with cache control attribute.
    xevm::PrefetchOp::create(
        rewriter, loc, ptrLLVM,
        xevm::LoadCacheControlAttr::get(
            ctxt, translateLoadXeGPUCacheHint(op.getL1Hint(), op.getL3Hint())));
    rewriter.eraseOp(op);
    return success();
  }
};

class FenceToXeVMPattern : public OpConversionPattern<xegpu::FenceOp> {
  using OpConversionPattern::OpConversionPattern;
  LogicalResult
  matchAndRewrite(xegpu::FenceOp op, xegpu::FenceOp::Adaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto loc = op.getLoc();
    xevm::MemScope memScope{xevm::MemScope::WORKGROUP};
    switch (op.getFenceScope()) {
    case xegpu::FenceScope::Workgroup:
      memScope = xevm::MemScope::WORKGROUP;
      break;
    case xegpu::FenceScope::GPU:
      memScope = xevm::MemScope::DEVICE;
      break;
    }
    xevm::AddrSpace addrSpace{xevm::AddrSpace::GLOBAL};
    switch (op.getMemoryKind()) {
    case xegpu::MemorySpace::Global:
      addrSpace = xevm::AddrSpace::GLOBAL;
      break;
    case xegpu::MemorySpace::SLM:
      addrSpace = xevm::AddrSpace::SHARED;
      break;
    }
    xevm::MemfenceOp::create(rewriter, loc, memScope, addrSpace);
    rewriter.eraseOp(op);
    return success();
  }
};

class DpasToXeVMPattern : public OpConversionPattern<xegpu::DpasOp> {
  using OpConversionPattern::OpConversionPattern;
  LogicalResult
  matchAndRewrite(xegpu::DpasOp op, xegpu::DpasOp::Adaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto loc = op.getLoc();
    auto ctxt = rewriter.getContext();
    auto aTy = cast<VectorType>(op.getLhs().getType());
    auto bTy = cast<VectorType>(op.getRhs().getType());
    auto resultType = cast<VectorType>(op.getResultType());

    auto encodePrecision = [&](Type type) -> xevm::ElemType {
      if (type == rewriter.getBF16Type())
        return xevm::ElemType::BF16;
      else if (type == rewriter.getF16Type())
        return xevm::ElemType::F16;
      else if (type == rewriter.getTF32Type())
        return xevm::ElemType::TF32;
      else if (type.isInteger(8)) {
        if (type.isUnsignedInteger())
          return xevm::ElemType::U8;
        return xevm::ElemType::S8;
      } else if (type == rewriter.getF32Type())
        return xevm::ElemType::F32;
      else if (type.isInteger(32))
        return xevm::ElemType::S32;
      llvm_unreachable("add more support for ElemType");
    };
    xevm::ElemType precATy = encodePrecision(aTy.getElementType());
    xevm::ElemType precBTy = encodePrecision(bTy.getElementType());
    Value c = op.getAcc();
    if (!c) {
      auto elementTy = resultType.getElementType();
      Attribute initValueAttr;
      if (isa<FloatType>(elementTy))
        initValueAttr = FloatAttr::get(elementTy, 0.0);
      else
        initValueAttr = IntegerAttr::get(elementTy, 0);
      c = arith::ConstantOp::create(
          rewriter, loc, DenseElementsAttr::get(resultType, initValueAttr));
    }

    Value aVec = op.getLhs();
    Value bVec = op.getRhs();
    auto cvecty = cast<VectorType>(c.getType());
    xevm::ElemType precCTy = encodePrecision(cvecty.getElementType());
    xevm::ElemType precDTy = encodePrecision(resultType.getElementType());
    VectorType cNty =
        VectorType::get(cvecty.getNumElements(), cvecty.getElementType());
    if (cvecty != cNty)
      c = vector::ShapeCastOp::create(rewriter, loc, cNty, c);
    Value dpasRes = xevm::MMAOp::create(
        rewriter, loc, cNty, aVec, bVec, c,
        xevm::MMAShapeAttr::get(ctxt, cvecty.getNumElements(), executionSize,
                                systolicDepth *
                                    getNumOperandsPerDword(precATy)),
        xevm::MMATypesAttr::get(ctxt, precDTy, precATy, precBTy, precCTy));
    if (cvecty != cNty)
      dpasRes = vector::ShapeCastOp::create(rewriter, loc, resultType, dpasRes);
    rewriter.replaceOp(op, dpasRes);
    return success();
  }

private:
  static unsigned getNumOperandsPerDword(xevm::ElemType pTy) {
    switch (pTy) {
    case xevm::ElemType::TF32:
      return 1;
    case xevm::ElemType::BF16:
    case xevm::ElemType::F16:
      return 2;
    case xevm::ElemType::U8:
    case xevm::ElemType::S8:
      return 4;
    default:
      llvm_unreachable("unsupported xevm::ElemType");
    }
  }
};

static std::optional<LLVM::AtomicBinOp>
matchSimpleAtomicOp(arith::AtomicRMWKind arithKind) {
  switch (arithKind) {
  case arith::AtomicRMWKind::addf:
    return LLVM::AtomicBinOp::fadd;
  case arith::AtomicRMWKind::addi:
    return LLVM::AtomicBinOp::add;
  case arith::AtomicRMWKind::assign:
    return LLVM::AtomicBinOp::xchg;
  case arith::AtomicRMWKind::maximumf:
    return LLVM::AtomicBinOp::fmax;
  case arith::AtomicRMWKind::maxs:
    return LLVM::AtomicBinOp::max;
  case arith::AtomicRMWKind::maxu:
    return LLVM::AtomicBinOp::umax;
  case arith::AtomicRMWKind::minimumf:
    return LLVM::AtomicBinOp::fmin;
  case arith::AtomicRMWKind::mins:
    return LLVM::AtomicBinOp::min;
  case arith::AtomicRMWKind::minu:
    return LLVM::AtomicBinOp::umin;
  case arith::AtomicRMWKind::ori:
    return LLVM::AtomicBinOp::_or;
  case arith::AtomicRMWKind::andi:
    return LLVM::AtomicBinOp::_and;
  default:
    return std::nullopt;
  }
}

class AtomicRMWToXeVMPattern : public OpConversionPattern<xegpu::AtomicRMWOp> {
  using OpConversionPattern::OpConversionPattern;
  LogicalResult
  matchAndRewrite(xegpu::AtomicRMWOp op, xegpu::AtomicRMWOp::Adaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto loc = op.getLoc();
    auto ctxt = rewriter.getContext();
    auto tdesc = op.getTensorDesc().getType();
    auto ptrTypeLLVM = LLVM::LLVMPointerType::get(
        ctxt, getNumericXeVMAddrSpace(tdesc.getMemorySpace()));
    Value basePtrI64 = arith::IndexCastOp::create(
        rewriter, loc, rewriter.getI64Type(), adaptor.getTensorDesc());
    Value basePtrLLVM =
        LLVM::IntToPtrOp::create(rewriter, loc, ptrTypeLLVM, basePtrI64);
    VectorType srcOrDstVecTy = cast<VectorType>(op.getValue().getType());
    VectorType srcOrDstFlatVecTy = VectorType::get(
        srcOrDstVecTy.getNumElements(), srcOrDstVecTy.getElementType());
    Value srcFlatVec = vector::ShapeCastOp::create(
        rewriter, loc, srcOrDstFlatVecTy, op.getValue());
    auto atomicKind = matchSimpleAtomicOp(op.getKind());
    assert(atomicKind.has_value());
    Value resVec = srcFlatVec;
    for (int i = 0; i < srcOrDstVecTy.getNumElements(); i++) {
      auto val = vector::ExtractOp::create(rewriter, loc, resVec, i);
      Value idx = LLVM::ConstantOp::create(rewriter, loc, rewriter.getI64Type(),
                                           rewriter.getIndexAttr(i));
      Value currPtr =
          LLVM::GEPOp::create(rewriter, loc, ptrTypeLLVM,
                              srcOrDstVecTy.getElementType(), basePtrLLVM, idx);
      Value newVal =
          LLVM::AtomicRMWOp::create(rewriter, loc, atomicKind.value(), currPtr,
                                    val, LLVM::AtomicOrdering::seq_cst);
      resVec = vector::InsertOp::create(rewriter, loc, newVal, resVec, i);
    }
    rewriter.replaceOp(op, resVec);
    return success();
  }
};

//===----------------------------------------------------------------------===//
// Pass Definition
//===----------------------------------------------------------------------===//

struct ConvertXeGPUToXeVMPass
    : public impl::ConvertXeGPUToXeVMPassBase<ConvertXeGPUToXeVMPass> {
  using Base::Base;

  void runOnOperation() override {
    LLVMTypeConverter typeConverter(&getContext());
    typeConverter.addConversion([&](VectorType type) -> Type {
      unsigned rank = type.getRank();
      auto elemType = type.getElementType();
      // If the element type is index, convert it to i64.
      if (llvm::isa<IndexType>(elemType))
        elemType = IntegerType::get(&getContext(), 64);
      // If the vector is a scalar or has a single element, return the element
      if (rank < 1 || type.getNumElements() == 1)
        return elemType;
      // Otherwise, convert the vector to a flat vector type.
      int64_t sum = llvm::product_of(type.getShape());
      return VectorType::get(sum, elemType);
    });
    typeConverter.addConversion([&](xegpu::TensorDescType type) -> Type {
      // Scattered descriptors are not supported in XeVM lowering.
      if (type.isScattered())
        return {};
      if (type.getRank() == 1)
        return IntegerType::get(&getContext(), 64);
      auto i32Type = IntegerType::get(&getContext(), 32);
      return VectorType::get(8, i32Type);
    });
    // Convert MemDescType into i32 for SLM
    typeConverter.addConversion([&](xegpu::MemDescType type) -> Type {
      return IntegerType::get(&getContext(), 32);
    });

    typeConverter.addConversion([&](MemRefType type) -> Type {
      return IntegerType::get(&getContext(), (isSharedMemRef(type) ? 32 : 64));
    });

    // LLVM type converter puts unrealized casts for the following cases:
    // add materialization casts to handle them.

    // Materialization to convert memref to i64 or i32 depending on global/SLM
    auto memrefMaterializationCast = [](OpBuilder &builder, Type type,
                                        ValueRange inputs,
                                        Location loc) -> Value {
      if (inputs.size() != 1)
        return {};
      auto input = inputs.front();
      if (auto memrefTy = dyn_cast<MemRefType>(input.getType())) {
        unsigned rank = memrefTy.getRank();
        Type indexType = builder.getIndexType();

        int64_t intOffsets;
        SmallVector<int64_t> intStrides;
        Value addr;
        Value offset;
        if (succeeded(memrefTy.getStridesAndOffset(intStrides, intOffsets)) &&
            ShapedType::isStatic(intOffsets)) {
          addr = memref::ExtractAlignedPointerAsIndexOp::create(builder, loc,
                                                                input);
          offset = arith::ConstantOp::create(builder, loc,
                                             builder.getIndexAttr(intOffsets));
        } else {

          // Result types: [base_memref, offset, stride0, stride1, ...,
          // strideN-1, size0, size1, ..., sizeN-1]
          SmallVector<Type> resultTypes{
              MemRefType::get({}, memrefTy.getElementType(),
                              MemRefLayoutAttrInterface(),
                              memrefTy.getMemorySpace()),
              indexType};
          // strides + sizes
          resultTypes.append(2 * rank, indexType);

          auto meta = memref::ExtractStridedMetadataOp::create(
              builder, loc, resultTypes, input);

          addr = memref::ExtractAlignedPointerAsIndexOp::create(
              builder, loc, meta.getBaseBuffer());
          offset = meta.getOffset();
        }

        auto addrCasted =
            arith::IndexCastUIOp::create(builder, loc, type, addr);
        auto offsetCasted =
            arith::IndexCastUIOp::create(builder, loc, type, offset);

        // Compute the final address: base address + byte offset
        auto byteSize = arith::ConstantOp::create(
            builder, loc, type,
            builder.getIntegerAttr(type,
                                   memrefTy.getElementTypeBitWidth() / 8));
        auto byteOffset =
            arith::MulIOp::create(builder, loc, offsetCasted, byteSize);
        auto addrWithOffset =
            arith::AddIOp::create(builder, loc, addrCasted, byteOffset);

        return addrWithOffset.getResult();
      }
      return {};
    };

    // Materialization to convert ui64 to i64
    auto ui64MaterializationCast = [](OpBuilder &builder, Type type,
                                      ValueRange inputs,
                                      Location loc) -> Value {
      if (inputs.size() != 1)
        return {};
      auto input = inputs.front();
      if (input.getType() == builder.getIntegerType(64, false)) {
        Value cast =
            index::CastUOp::create(builder, loc, builder.getIndexType(), input)
                .getResult();
        return arith::IndexCastUIOp::create(builder, loc, type, cast)
            .getResult();
      }
      return {};
    };

    // Materialization to convert ui32 to i32
    auto ui32MaterializationCast = [](OpBuilder &builder, Type type,
                                      ValueRange inputs,
                                      Location loc) -> Value {
      if (inputs.size() != 1)
        return {};
      auto input = inputs.front();
      if (input.getType() == builder.getIntegerType(32, false)) {
        Value cast =
            index::CastUOp::create(builder, loc, builder.getIndexType(), input)
                .getResult();
        return arith::IndexCastUIOp::create(builder, loc, type, cast)
            .getResult();
      }
      return {};
    };

    // Materialization to convert
    //   - single element 1D vector to scalar
    //   - bitcast vector of same rank
    //   - shape vector of different rank but same element type
    auto vectorMaterializationCast = [](OpBuilder &builder, Type type,
                                        ValueRange inputs,
                                        Location loc) -> Value {
      if (inputs.size() != 1)
        return {};
      auto input = inputs.front();
      if (auto vecTy = dyn_cast<VectorType>(input.getType())) {
        if (vecTy.getNumElements() == 1) {
          // If the vector has a single element, return the element type.
          Value cast =
              vector::ExtractOp::create(builder, loc, input, 0).getResult();
          if (vecTy.getElementType() == builder.getIndexType())
            cast = arith::IndexCastUIOp::create(builder, loc, type, cast)
                       .getResult();
          return cast;
        } else if (auto targetVecTy = dyn_cast<VectorType>(type)) {
          // If the target type is a vector of same rank,
          //   bitcast to the target type.
          if (targetVecTy.getRank() == vecTy.getRank())
            return vector::BitCastOp::create(builder, loc, targetVecTy, input)
                .getResult();
          else if (targetVecTy.getElementType() == vecTy.getElementType()) {
            // If the target type is a vector of different rank but same element
            // type, reshape to the target type.
            return vector::ShapeCastOp::create(builder, loc, targetVecTy, input)
                .getResult();
          }
        }
      }
      return {};
    };

    // If result type of original op is single element vector and lowered type
    // is scalar. This materialization cast creates a single element vector by
    // broadcasting the scalar value.
    auto singleElementVectorMaterializationCast =
        [](OpBuilder &builder, Type type, ValueRange inputs,
           Location loc) -> Value {
      if (inputs.size() != 1)
        return {};
      auto input = inputs.front();
      if (input.getType().isIntOrIndexOrFloat()) {
        // If the input is a scalar, and the target type is a vector of single
        // element, create a single element vector by broadcasting.
        if (auto vecTy = dyn_cast<VectorType>(type)) {
          if (vecTy.getNumElements() == 1) {
            return vector::BroadcastOp::create(builder, loc, vecTy, input)
                .getResult();
          }
        }
      }
      return {};
    };
    typeConverter.addSourceMaterialization(
        singleElementVectorMaterializationCast);
    typeConverter.addSourceMaterialization(vectorMaterializationCast);
    typeConverter.addTargetMaterialization(memrefMaterializationCast);
    typeConverter.addTargetMaterialization(ui32MaterializationCast);
    typeConverter.addTargetMaterialization(ui64MaterializationCast);
    typeConverter.addTargetMaterialization(vectorMaterializationCast);
    ConversionTarget target(getContext());
    target.addLegalDialect<xevm::XeVMDialect, LLVM::LLVMDialect,
                           vector::VectorDialect, arith::ArithDialect,
                           memref::MemRefDialect, gpu::GPUDialect,
                           index::IndexDialect>();
    target.addIllegalDialect<xegpu::XeGPUDialect>();

    RewritePatternSet patterns(&getContext());
    populateXeGPUToXeVMConversionPatterns(typeConverter, patterns);
    scf::populateSCFStructuralTypeConversionsAndLegality(typeConverter,
                                                         patterns, target);
    if (failed(applyPartialConversion(getOperation(), target,
                                      std::move(patterns))))
      signalPassFailure();
  }
};
} // namespace

//===----------------------------------------------------------------------===//
// Pattern Population
//===----------------------------------------------------------------------===//
void mlir::populateXeGPUToXeVMConversionPatterns(
    const LLVMTypeConverter &typeConverter, RewritePatternSet &patterns) {
  patterns.add<CreateNdDescToXeVMPattern,
               LoadStorePrefetchNdToXeVMPattern<xegpu::LoadNdOp>,
               LoadStorePrefetchNdToXeVMPattern<xegpu::StoreNdOp>,
               LoadStorePrefetchNdToXeVMPattern<xegpu::PrefetchNdOp>>(
      typeConverter, patterns.getContext());
  patterns.add<AtomicRMWToXeVMPattern, PrefetchToXeVMPattern,
               LoadStoreToXeVMPattern<xegpu::LoadGatherOp>,
               LoadStoreToXeVMPattern<xegpu::StoreScatterOp>>(
      typeConverter, patterns.getContext());
  patterns.add<LoadStoreMatrixToXeVMPattern<xegpu::LoadMatrixOp>,
               LoadStoreMatrixToXeVMPattern<xegpu::StoreMatrixOp>,
               CreateMemDescOpPattern>(typeConverter, patterns.getContext());
  patterns.add<FenceToXeVMPattern, DpasToXeVMPattern>(typeConverter,
                                                      patterns.getContext());
}