XeGPUToXeVM.cpp

Go to the documentation of this file.

//===-- 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/Dialect/XeGPU/uArch/uArchCommon.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 <string>
 
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/stride of dim rank-2 (i32)
  LeadingStride0 = 5, // Row strides of the leading (batch) dims of a >2D
  LeadingStride1 = 6, // descriptor (i32); added into offset_h by the load/store
  LeadingStride2 = 7, // lowering. Left at 0 for 2D descriptors.
};
 
// Spare payload slots above, and the resulting max lowerable descriptor rank.
static constexpr int64_t maxNdTdescLeadingDims{3};
static constexpr int64_t maxNdTdescRank{2 + maxNdTdescLeadingDims};
 
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");
}
 
/// Translates a memref memory space attribute into XeVM's numeric address
/// space, which follows the OpenCL/SPIR-V convention (0 = private, 1 =
/// global, 2 = constant, 3 = shared/local, 4 = generic). A null attribute,
/// meaning the memory space was left unspecified, maps to the default space
/// 0. Returns failure if `memSpace` is a representation this pass does not
/// know how to translate (e.g. a SPIR-V storage class or an arbitrary string
/// attribute), rather than assuming it is an `IntegerAttr` and asserting.
static FailureOr<unsigned> getNumericMemorySpace(Attribute memSpace) {
  if (!memSpace)
    return 0u;
  if (auto intAttr = llvm::dyn_cast<IntegerAttr>(memSpace))
    return static_cast<unsigned>(intAttr.getInt());
  if (auto xevmSpace = llvm::dyn_cast<xevm::AddrSpaceAttr>(memSpace))
    return static_cast<unsigned>(xevmSpace.getValue());
  if (auto gpuSpace = llvm::dyn_cast<gpu::AddressSpaceAttr>(memSpace)) {
    switch (gpuSpace.getValue()) {
    case gpu::AddressSpace::Global:
      return static_cast<unsigned>(xevm::AddrSpace::GLOBAL);
    case gpu::AddressSpace::Workgroup:
      return static_cast<unsigned>(xevm::AddrSpace::SHARED);
    case gpu::AddressSpace::Private:
      return static_cast<unsigned>(xevm::AddrSpace::PRIVATE);
    case gpu::AddressSpace::Constant:
      return static_cast<unsigned>(xevm::AddrSpace::CONSTANT);
    }
    llvm_unreachable("Unknown GPU address space");
  }
  return failure();
}
 
/// Checks if the given MemRefType refers to shared memory.
static bool isSharedMemRef(const MemRefType &memrefTy) {
  FailureOr<unsigned> addrSpace =
      getNumericMemorySpace(memrefTy.getMemorySpace());
  return succeeded(addrSpace) &&
         *addrSpace == static_cast<unsigned>(xevm::AddrSpace::SHARED);
}
 
// 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) {
  // If no hints are provided, use the default cache control.
  if (!L1hint && !L3hint)
    return xevm::LoadCacheControl::USE_DEFAULT;
  // If only one of the hints is provided, use the default for the other level.
  auto L1hintVal = L1hint.value_or(xegpu::CachePolicy::CACHED);
  auto L3hintVal = L3hint.value_or(xegpu::CachePolicy::CACHED);
  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) {
  // If no hints are provided, use the default cache control.
  if (!L1hint && !L3hint)
    return xevm::StoreCacheControl::USE_DEFAULT;
  // If only one of the hints is provided, use the default for the other level.
  auto L1hintVal = L1hint.value_or(xegpu::CachePolicy::UNCACHED);
  auto L3hintVal = L3hint.value_or(xegpu::CachePolicy::WRITE_BACK);
  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
//
 
//
// High-D (>2D) nd descriptors are lowered by viewing the source as a single
// flattened 2D plane, so the 2D-block surface covers every leading (batch)
// plane at once and a batch position becomes a row offset into it. The surface
// height is the row extent the leading dims reach:
//   base_height = size[R-2] + sum_d (size[d] - 1) * (stride[d] / stride[R-2])
//
// Leaving `base_ptr` at the true base means an out-of-range batch index lands
// past `base_height`, where the HW boundary check handles it, instead of aiming
// the surface at unmapped memory. Encoding the leading strides as row counts
// (`stride[d] / stride[R-2]`) also makes them dimensionless, so they survive
// element-type repacking (e.g. the f16 -> i32 transpose repack) without a unit
// conversion.
//
// Limitations of the flattened-plane view:
//  1. Each leading stride must be a whole number of rows. A source with gaps
//     between planes (`stride[d] % stride[R-2] != 0`) is not lowered. This can
//     only be checked when the strides are static; for dynamic strides the
//     divisibility is assumed.
//  2. `base_height` grows with the leading extent, so a source with a large
//     batch x head x sequence extent can exceed the HW 2D-block surface height.
//  3. Plane boundaries are invisible to the boundary check, which only knows
//     the surface: a tile whose rows run past `size[R-2]` spills into the
//     following rows of the surface -- the next plane, or the padding between
//     planes -- rather than being clipped there as it would be on a per-plane
//     surface. A load reads those rows instead of returning zeros, and a store
//     overwrites them, so it can corrupt the next plane. This only matters
//     when `size[R-2]` is not a multiple of the tile height.
//
 
class CreateNdDescToXeVMPattern
    : public OpConversionPattern<xegpu::CreateNdDescOp> {
  using Base::Base;
 
  // Checks the surface width and pitch against the target's 2D block memory
  // restrictions. Only these two are enforced: both are backed by the source's
  // static shape and strides, whereas the base address alignment is a property
  // of the allocation that is not visible here. The width and the pitch are
  // checked independently, each only when it is static; a value known only at
  // runtime is left alone, as guarding it inside the kernel would cost far more
  // than it is worth.
  //
  // A violation is a dead end rather than a pattern that does not apply: the op
  // has no other legalization path, so it is reported with an error that names
  // the limit it misses.
  LogicalResult checkMemoryRestrictions(xegpu::CreateNdDescOp op,
                                        int64_t rank) const {
    const auto *uArch =
        xegpu::uArch::getUArch(xegpu::getChipStr(op).value_or(""));
    // Without a known target chip there is no uArch to check against, so the
    // descriptor must be accepted as-is.
    if (!uArch)
      return success();
    const xegpu::uArch::BlockIOMemoryRestrictions &limits =
        uArch->getBlockIOMemoryRestrictions();
 
    // Recover the width and the pitch as a number of elements. A memref source
    // carries both in its type, and either one may be dynamic independently of
    // the other, so they are queried separately instead of being dropped
    // together when the source is not fully static. `getMixedSizes` /
    // `getMixedStrides` are only consulted for an integer source, where the op
    // itself holds the extents.
    std::optional<int64_t> widthElems;
    std::optional<int64_t> pitchElems;
    if (auto memrefTy = dyn_cast<MemRefType>(op.getSource().getType())) {
      if (!ShapedType::isDynamic(memrefTy.getDimSize(rank - 1)))
        widthElems = memrefTy.getDimSize(rank - 1);
      // The source is known to be strided; the caller checked that already.
      SmallVector<int64_t> staticStrides;
      int64_t staticOffset;
      if (succeeded(
              memrefTy.getStridesAndOffset(staticStrides, staticOffset)) &&
          !ShapedType::isDynamic(staticStrides[rank - 2]))
        pitchElems = staticStrides[rank - 2];
    } else {
      widthElems = getConstantIntValue(op.getMixedSizes()[rank - 1]);
      pitchElems = getConstantIntValue(op.getMixedStrides()[rank - 2]);
    }
 
    unsigned elemBitSize =
        op.getType().getElementType().getIntOrFloatBitWidth();
    // The restrictions are expressed in bytes, but a sub-byte element type can
    // leave the width or the pitch on a fraction of a byte, so the comparisons
    // are done in bits. Such a size meets no byte alignment at all and is
    // reported as a violation rather than waved through.
    auto describe = [](int64_t bits) -> std::string {
      if (bits % 8 == 0)
        return llvm::formatv("{0} bytes", bits / 8).str();
      return llvm::formatv("{0} bits", bits).str();
    };
    auto check = [&](std::optional<int64_t> elems, int64_t minBytes,
                     int64_t alignBytes,
                     llvm::StringRef what) -> LogicalResult {
      if (!elems)
        return success();
      int64_t bits = *elems * elemBitSize;
      if (bits < minBytes * 8)
        return op.emitError(llvm::formatv(
            "2D block surface {0} must be at least {1} bytes on this "
            "target, got {2}.",
            what, minBytes, describe(bits)));
      if (alignBytes > 0 && bits % (alignBytes * 8) != 0)
        return op.emitError(llvm::formatv(
            "2D block surface {0} must be a multiple of {1} bytes on this "
            "target, got {2}.",
            what, alignBytes, describe(bits)));
      return success();
    };
 
    if (failed(check(widthElems, limits.minBaseWidthBytes,
                     limits.baseWidthAlignmentBytes, "width")))
      return failure();
    // The pitch is the row stride of the innermost 2D tile.
    if (failed(check(pitchElems, limits.minBasePitchBytes,
                     limits.basePitchAlignmentBytes, "pitch")))
      return failure();
    // Consecutive rows may not overlap, so the pitch has to span the full
    // width.
    if (widthElems && pitchElems && *pitchElems < *widthElems)
      return op.emitError(
          llvm::formatv("2D block surface pitch must be at least the surface "
                        "width, got a pitch of {0} for a surface {1} wide.",
                        describe(*pitchElems * elemBitSize),
                        describe(*widthElems * elemBitSize)));
    return success();
  }
 
  LogicalResult
  matchAndRewrite(xegpu::CreateNdDescOp op,
                  xegpu::CreateNdDescOp::Adaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto loc = op.getLoc();
    auto source = op.getSource();
 
    // Check all failure conditions before generating any IR, so nothing has to
    // be rolled back.
    int64_t rank = op.getType().getRank();
    int64_t sourceRank;
    auto memrefTy = dyn_cast<MemRefType>(source.getType());
    if (memrefTy) {
      if (!memrefTy.isStrided())
        return rewriter.notifyMatchFailure(op, "Expected strided Memref.");
      sourceRank = memrefTy.getRank();
    } else if (isa<IntegerType>(source.getType())) {
      sourceRank = op.getMixedSizes().size();
    } else {
      return rewriter.notifyMatchFailure(op,
                                         "Expected ranked Memref or integer.");
    }
    if (sourceRank != rank)
      return rewriter.notifyMatchFailure(
          op, "Expected descriptor rank to match source rank; subview the "
              "source down to the descriptor rank.");
    if (rank > maxNdTdescRank)
      return rewriter.notifyMatchFailure(
          op, "Batched nd descriptor supports at most " +
                  std::to_string(maxNdTdescLeadingDims) +
                  " leading dims (rank <= " + std::to_string(maxNdTdescRank) +
                  ").");
    // Limitation 1 above; dynamic strides are assumed to divide evenly.
    if (rank > 2) {
      SmallVector<std::optional<int64_t>> constStrides(rank, std::nullopt);
      if (memrefTy) {
        SmallVector<int64_t> staticStrides;
        int64_t staticOffset;
        if (succeeded(
                memrefTy.getStridesAndOffset(staticStrides, staticOffset)))
          for (int64_t d = 0; d < rank; ++d)
            if (!ShapedType::isDynamic(staticStrides[d]))
              constStrides[d] = staticStrides[d];
      } else {
        SmallVector<OpFoldResult> mixed = op.getMixedStrides();
        for (int64_t d = 0; d < rank; ++d)
          constStrides[d] = getConstantIntValue(mixed[d]);
      }
      if (std::optional<int64_t> pitch = constStrides[rank - 2]) {
        for (int64_t d = 0; d < rank - 2; ++d) {
          std::optional<int64_t> leading = constStrides[d];
          if (leading && (*pitch == 0 || *leading % *pitch != 0))
            return rewriter.notifyMatchFailure(
                op, "Expected each leading (batch) stride to be a multiple of "
                    "the row stride; the source has gaps between planes.");
        }
      }
    }
    // Enforce the target's 2D block memory restrictions. A rank-1 descriptor is
    // just a base address and is not consumed by a 2D block op, so it is
    // exempt.
    if (rank >= 2 && failed(checkMemoryRestrictions(op, rank)))
      return failure();
 
    Type payloadElemTy = rewriter.getI32Type();
    Type i64Ty = rewriter.getI64Type();
 
    // Access the adaptor only after the failure checks, so a bail-out leaves no
    // materialization cast behind.
    Value baseAddr = adaptor.getSource();
    if (isa<IntegerType>(source.getType()) && 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();
    }
 
    SmallVector<OpFoldResult> mixedSizes;
    SmallVector<OpFoldResult> mixedStrides;
    if (memrefTy && !xegpu::hasStaticShapeAndStrides(memrefTy)) {
      auto meta =
          memref::ExtractStridedMetadataOp::create(rewriter, loc, source);
      mixedSizes = meta.getConstifiedMixedSizes();
      mixedStrides = meta.getConstifiedMixedStrides();
    } else {
      mixedSizes = op.getMixedSizes();
      mixedStrides = op.getMixedStrides();
    }
 
    // Op is lowered to a code sequence that populates payload.
    // Payload is a 8xi32 vector. Offset to individual fields are defined in
    // NdTdescOffset enum.
    VectorType payloadTy = VectorType::get(8, payloadElemTy);
    // 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)));
 
    // 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;
    };
    // The descriptor's innermost 2 dims are the 2D tile (H, W).
    Value baseShapeW = createOffset(mixedSizes, rank - 1);
    // Pitch is the stride of dim rank-2 (the row stride of the 2D tile).
    Value basePitch = createOffset(mixedStrides, rank - 2);
    // Leading (batch) strides as a number of rows (`stride[d] / pitch`).
    SmallVector<Value> leadingRowStrides;
    for (int64_t d = 0; d < rank - 2; ++d) {
      std::optional<int64_t> leading = getConstantIntValue(mixedStrides[d]);
      std::optional<int64_t> pitch =
          getConstantIntValue(mixedStrides[rank - 2]);
      if (leading && pitch && *pitch != 0)
        leadingRowStrides.push_back(arith::ConstantIntOp::create(
            rewriter, loc, payloadElemTy, *leading / *pitch));
      else
        leadingRowStrides.push_back(arith::DivUIOp::create(
            rewriter, loc, createOffset(mixedStrides, d), basePitch));
    }
    // Height of the flattened plane: every leading (batch) plane is stacked
    // into the surface, so the boundary check covers an out-of-range batch.
    //   base_height = size[R-2] + sum_d (size[d] - 1) * leadingRowStride[d]
    // For rank 2 this is just size[0].
    Value baseShapeH = createOffset(mixedSizes, rank - 2);
    if (rank > 2) {
      Value one = arith::ConstantIntOp::create(rewriter, loc, payloadElemTy, 1);
      for (int64_t d = 0; d < rank - 2; ++d) {
        Value planesBelow = rewriter.createOrFold<arith::SubIOp>(
            loc, createOffset(mixedSizes, d), one);
        Value rows = rewriter.createOrFold<arith::MulIOp>(loc, planesBelow,
                                                          leadingRowStrides[d]);
        baseShapeH =
            rewriter.createOrFold<arith::AddIOp>(loc, baseShapeH, rows);
      }
    }
    // 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));
    // The leading (batch) row strides go into the spare payload slots; the
    // load/store/prefetch lowering turns the batch offsets into a row offset
    // with them.
    for (int64_t d = 0; d < rank - 2; ++d)
      payload = vector::InsertOp::create(
          rewriter, loc, leadingRowStrides[d], payload,
          static_cast<int>(NdTdescOffset::LeadingStride0) + d);
    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.");
    if (tileRank > 2 && llvm::any_of(tdescTy.getShape().drop_back(2),
                                     [](int64_t d) { return d != 1; }))
      return rewriter.notifyMatchFailure(
          op, "Expected leading (batch) descriptor dims to be unit.");
    if (tileRank > maxNdTdescRank)
      return rewriter.notifyMatchFailure(
          op, "Expected descriptor rank <= " + std::to_string(maxNdTdescRank) +
                  ".");
    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));
 
      Value offsetW = getValueOrCreateConstantIntOp(rewriter, loc,
                                                    mixedOffsets[tileRank - 1]);
      offsetW = getValueOrCreateCastToIndexLike(rewriter, loc,
                                                rewriter.getI32Type(), offsetW);
      Value offsetH = getValueOrCreateConstantIntOp(rewriter, loc,
                                                    mixedOffsets[tileRank - 2]);
      offsetH = getValueOrCreateCastToIndexLike(rewriter, loc,
                                                rewriter.getI32Type(), offsetH);
      // Turn the leading (batch) offsets into a row offset into the flattened
      // plane, using the row-unit batch strides encoded at create time:
      //   offsetH += sum_d offset[d] * leadingRowStride[d]
      // The base pointer stays at the true base, so an out-of-range batch index
      // is caught by the HW boundary check instead of moving the surface to
      // unmapped memory.
      for (int64_t d = 0; d < tileRank - 2; ++d) {
        Value off =
            getValueOrCreateConstantIntOp(rewriter, loc, mixedOffsets[d]);
        off = getValueOrCreateCastToIndexLike(rewriter, loc,
                                              rewriter.getI32Type(), off);
        Value rowStride = vector::ExtractOp::create(
            rewriter, loc, tdesc,
            static_cast<int>(NdTdescOffset::LeadingStride0) + d);
        Value term = arith::MulIOp::create(rewriter, loc, off, rowStride);
        offsetH = arith::AddIOp::create(rewriter, loc, offsetH, term);
      }
      // 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 (second-to-last dim).
      auto tileH = tdescTy.getDimSize(tileRank - 2);
      // 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());
          bool vnni = op.getPacked().value_or(false);
          auto transposeValue = op.getTranspose();
          bool transpose =
              transposeValue.has_value() && transposeValue.value()[0] == 1;
          // Handle special case of 32x16 and 8bit element load
          // with no vnni, no transpose, no vblocks.
          // For this special case, vnni and non vnni yields the same output
          // and only the vnni variant is supported by HW.
          // Check and set vnni of the special case.
          if (elemBitSize == 8 && tileW == 16 && tileH == 32 && !vnni &&
              !transpose) {
            vnni = true;
          }
          // Handle tranpose request on small element size
          // Transpose needs to be requested on 32bit element type.
          // offsetW and tileW needs to be adjusted to account for element type
          // change.
          if (transpose && elemBitSize < 32) {
            int32_t scale = 32 / elemBitSize;
            Value scaleLog2 = arith::ConstantIntOp::create(
                rewriter, loc, rewriter.getI32Type(), llvm::Log2_64(scale));
            offsetW = arith::ShRSIOp::create(rewriter, loc, offsetW, scaleLog2);
            tileW = tileW * elemBitSize / 32;
            elemBitSize = 32;
          }
          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;
}
 
// Returns true when every element of `mask` carries the same bit, so gating a
// whole contiguous block on element 0 is equivalent. Splat constants, an
// all-ones or all-zeros `vector.constant_mask`, broadcasts of a scalar, and a
// `vector.from_elements` of one repeated value qualify.
//
// TODO: this recognizes a fixed list of producers. A general uniformity query
// on the vector dialect would cover more forms and would not need updating
// every time a new producer shows up here.
static bool isUniformMask(Value mask) {
  if (!isa<VectorType>(mask.getType()))
    return true;
  DenseElementsAttr splat;
  if (matchPattern(mask, m_Constant(&splat)) && splat.isSplat())
    return true;
  if (auto constantMask = mask.getDefiningOp<vector::ConstantMaskOp>()) {
    // All-ones and all-zeros are uniform. A zero dim size is only legal when
    // every dim is zero, so those are the only two uniform cases.
    return constantMask.isAllOnesMask() ||
           llvm::all_of(constantMask.getMaskDimSizes(),
                        [](int64_t size) { return size == 0; });
  }
  if (auto broadcast = mask.getDefiningOp<vector::BroadcastOp>())
    return !isa<VectorType>(broadcast.getSource().getType());
  if (auto fromElements = mask.getDefiningOp<vector::FromElementsOp>())
    return llvm::all_equal(fromElements.getElements());
  // Flattening a rank > 1 operand to rank 1 inserts a shape cast. It keeps the
  // same elements, so it keeps uniformity.
  if (auto shapeCast = mask.getDefiningOp<vector::ShapeCastOp>())
    return isUniformMask(shapeCast.getSource());
  return false;
}
 
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();
    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));
    // 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())) {
        FailureOr<unsigned> addrSpace =
            getNumericMemorySpace(memRefTy.getMemorySpace());
        if (failed(addrSpace))
          return rewriter.notifyMatchFailure(
              op, "Unsupported memref memory space attribute.");
        if (*addrSpace != 0)
          ptrTypeLLVM = LLVM::LLVMPointerType::get(ctxt, *addrSpace);
      }
    } else {
      basePtrI64 = adaptor.getDest();
      if (auto memRefTy = dyn_cast<MemRefType>(op.getDest().getType())) {
        FailureOr<unsigned> addrSpace =
            getNumericMemorySpace(memRefTy.getMemorySpace());
        if (failed(addrSpace))
          return rewriter.notifyMatchFailure(
              op, "Unsupported memref memory space attribute.");
        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();
 
    // Coalesce a lane's multi-element access into one block access.
    //
    // Distribution gives every element its own offset and mask bit, so a lane
    // that takes D neighbouring elements arrives with `vector<D>` offsets and
    // mask. One block access needs one base offset and one mask bit, so take
    // both from element 0. A non-uniform mask cannot be reduced to one bit, so
    // it stays a vector and fails to match below.
    auto origOffsetsTy = dyn_cast<VectorType>(op.getOffsets().getType());
    if (isa<VectorType>(offset.getType()) && origOffsetsTy && valOrResVecTy &&
        origOffsetsTy.getNumElements() == valOrResVecTy.getNumElements() &&
        isUniformMask(op.getMask())) {
      offset = vector::ExtractOp::create(rewriter, loc, offset,
                                         ArrayRef<int64_t>{0});
      if (isa<VectorType>(mask.getType()))
        mask = vector::ExtractOp::create(rewriter, loc, mask,
                                         ArrayRef<int64_t>{0});
    }
 
    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()->setDiscardableAttr(
          "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()->setDiscardableAttr(
          "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, /*cache_control=*/nullptr);
        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" && *chipOpt != "cri")) {
        // the lowering for chunk load only works for pvc, bmg or cri
        return rewriter.notifyMatchFailure(
            op, "The lowering is specific to pvc, bmg or cri.");
      }
    }
 
    if constexpr (std::is_same_v<OpType, xegpu::LoadMatrixOp>) {
      // The load result type is taken from the type converter. This maps
      // element types that are not directly representable in LLVM (e.g.
      // f8E8M0FNU) to an integer storage type of the same bit width, and
      // collapses single-element vectors to a scalar, since LLVM load/store
      // does not support vectors of size 1.
      Type loadTy =
          this->getTypeConverter()->convertType(op.getResult().getType());
      auto loadOp = LLVM::LoadOp::create(rewriter, loc, loadTy, 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 Base::Base;
  LogicalResult
  matchAndRewrite(xegpu::PrefetchOp op, xegpu::PrefetchOp::Adaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto loc = op.getLoc();
    auto ctxt = rewriter.getContext();
    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 two sources:
        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 source is a memref, we use its memory space.
    if (auto memRefTy = dyn_cast<MemRefType>(op.getSource().getType())) {
      FailureOr<unsigned> addrSpace =
          getNumericMemorySpace(memRefTy.getMemorySpace());
      if (failed(addrSpace))
        return rewriter.notifyMatchFailure(
            op, "Unsupported memref memory space attribute.");
      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 Base::Base;
  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();
  }
};
 
static auto encodePrecision = [](Type type) -> xevm::ElemType {
  if (type.isBF16())
    return xevm::ElemType::BF16;
  else if (type.isF16())
    return xevm::ElemType::F16;
  else if (type.isTF32())
    return xevm::ElemType::TF32;
  else if (type.isInteger(8)) {
    if (type.isUnsignedInteger())
      return xevm::ElemType::U8;
    return xevm::ElemType::S8;
  } else if (type.isF32())
    return xevm::ElemType::F32;
  else if (type.isInteger(32))
    return xevm::ElemType::S32;
  else if (type.isF8E5M2())
    return xevm::ElemType::BF8;
  else if (type.isF8E4M3FN())
    return xevm::ElemType::F8;
  else if (mlir::isa<Float4E2M1FNType>(type))
    return xevm::ElemType::E2M1;
  llvm_unreachable("add more support for ElemType");
};
 
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:
  case xevm::ElemType::F8:
  case xevm::ElemType::BF8:
    return 4;
  case xevm::ElemType::E2M1:
    return 8;
  default:
    llvm_unreachable("unsupported xevm::ElemType");
  }
}
 
class DpasToXeVMPattern : public OpConversionPattern<xegpu::DpasOp> {
  using Base::Base;
  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());
 
    // get the correct dpasInst by getting info from chip
    auto chipStr = xegpu::getChipStr(op);
    if (!chipStr)
      return rewriter.notifyMatchFailure(op, "cannot determine target chip");
 
    const auto *uArch = mlir::xegpu::uArch::getUArch(*chipStr);
    if (!uArch)
      return rewriter.notifyMatchFailure(op, "unsupported target uArch");
 
    auto *dpasInst = const_cast<xegpu::uArch::SubgroupMatrixMultiplyAcc *>(
        llvm::dyn_cast_or_null<xegpu::uArch::SubgroupMatrixMultiplyAcc>(
            uArch->getInstruction(
                xegpu::uArch::InstructionKind::SubgroupMatrixMultiplyAcc)));
    if (!dpasInst)
      return rewriter.notifyMatchFailure(op,
                                         "DPAS not supported by target uArch");
 
    auto checkSupportedTypes = [&](VectorType vecTy,
                                   xegpu::uArch::MMAOpndKind kind) -> bool {
      auto supported = dpasInst->getSupportedTypes(*ctxt, kind);
      return llvm::find(supported, vecTy.getElementType()) != supported.end();
    };
 
    if (!checkSupportedTypes(aTy, xegpu::uArch::MMAOpndKind::MatrixA))
      return rewriter.notifyMatchFailure(
          op, "A-matrix element type not supported by target uArch");
    if (!checkSupportedTypes(bTy, xegpu::uArch::MMAOpndKind::MatrixB))
      return rewriter.notifyMatchFailure(
          op, "B-matrix element type not supported by target uArch");
    // NOTE: Supported types for MatrixC and MatrixD are identical
    if (!checkSupportedTypes(resultType, xegpu::uArch::MMAOpndKind::MatrixD))
      return rewriter.notifyMatchFailure(
          op, "result/accumulator element type not supported by target uArch");
 
    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();
  }
};
 
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 Base::Base;
  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.getI64IntegerAttr(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();
  }
};
 
class DpasMxToXeVMPattern : public OpConversionPattern<xegpu::DpasMxOp> {
  using Base::Base;
  LogicalResult
  matchAndRewrite(xegpu::DpasMxOp op, xegpu::DpasMxOp::Adaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto loc = op.getLoc();
    auto ctxt = rewriter.getContext();
    auto aTy = op.getA().getType();
    auto bTy = op.getB().getType();
    auto resVecTy =
        cast<VectorType>(getTypeConverter()->convertType(op.getType()));
 
    auto chipStr = xegpu::getChipStr(op);
    if (!chipStr)
      return rewriter.notifyMatchFailure(op, "cannot determine target chip");
 
    const auto *uArch = xegpu::uArch::getUArch(*chipStr);
    if (!uArch)
      return rewriter.notifyMatchFailure(op, "unsupported target uArch");
 
    // TODO: Add supported shape check
 
    xevm::ElemType precATy = encodePrecision(aTy.getElementType());
    xevm::ElemType precBTy = encodePrecision(bTy.getElementType());
    Value c = adaptor.getAcc();
    if (!c) {
      auto elementTy = resVecTy.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(resVecTy, initValueAttr));
    }
 
    Value aVec = adaptor.getA();
    Value bVec = adaptor.getB();
    auto aVecTy = cast<VectorType>(aVec.getType());
    auto bVecTy = cast<VectorType>(bVec.getType());
    if (aVecTy.getElementTypeBitWidth() == 4)
      aVec = vector::BitCastOp::create(
          rewriter, loc,
          VectorType::get(aVecTy.getNumElements() / 2, rewriter.getI8Type()),
          aVec);
    if (bVecTy.getElementTypeBitWidth() == 4)
      bVec = vector::BitCastOp::create(
          rewriter, loc,
          VectorType::get(bVecTy.getNumElements() / 2, rewriter.getI8Type()),
          bVec);
    auto cVecTy = cast<VectorType>(c.getType());
    xevm::ElemType precCTy = encodePrecision(cVecTy.getElementType());
    xevm::ElemType precDTy = encodePrecision(resVecTy.getElementType());
    Value scaleA = adaptor.getScaleA();
    Value scaleB = adaptor.getScaleB();
    Value dpasMxRes = xevm::MMAMxOp::create(
        rewriter, loc, resVecTy, aVec, bVec, scaleA, scaleB, c,
        xevm::MMAShapeAttr::get(ctxt, cVecTy.getNumElements(), executionSize,
                                systolicDepth *
                                    getNumOperandsPerDword(precATy)),
        xevm::MMATypesAttr::get(ctxt, precDTy, precATy, precBTy, precCTy));
    rewriter.replaceOp(op, dpasMxRes);
    return success();
  }
};
 
//===----------------------------------------------------------------------===//
// arith.extf / arith.truncf to xevm.extf / xevm.truncf
//===----------------------------------------------------------------------===//
//
// Micro-scaling (MX) GEMM lowering breaks arith.scaling_extf/scaling_truncf
// into plain arith.extf/arith.truncf whose narrow side uses one of the MX float
// formats (f8E5M2, f8E4M3FN or f4E2M1FN). These narrow floats have no native
// LLVM support, so the conversions are mapped onto the dedicated xevm.extf /
// xevm.truncf ops which lower to hardware builtins. The f8E8M0FNU scale type is
// intentionally not handled here: it is expanded into integer arithmetic by
// arith-expand before this pass runs.
 
// xevm.extf / xevm.truncf only convert between the MX narrow floats and
// f16/bf16, and the underlying builtins operate on exactly 16 f16/bf16 values.
static constexpr int64_t kXeVMExtfTruncfNumElems = 16;
 
// Maps a narrow MX float element type to the matching xevm.extf source enum.
static std::optional<xevm::ExtfSrcElemTypes> getExtfNarrowType(Type etype) {
  if (isa<Float8E5M2Type>(etype))
    return xevm::ExtfSrcElemTypes::BF8;
  if (isa<Float8E4M3FNType>(etype))
    return xevm::ExtfSrcElemTypes::F8;
  if (isa<Float4E2M1FNType>(etype))
    return xevm::ExtfSrcElemTypes::E2M1;
  return std::nullopt;
}
 
// Maps a narrow MX float element type to the matching xevm.truncf dest enum.
static std::optional<xevm::TruncfDstElemTypes> getTruncfNarrowType(Type etype) {
  if (isa<Float8E5M2Type>(etype))
    return xevm::TruncfDstElemTypes::BF8;
  if (isa<Float8E4M3FNType>(etype))
    return xevm::TruncfDstElemTypes::F8;
  if (isa<Float4E2M1FNType>(etype))
    return xevm::TruncfDstElemTypes::E2M1;
  return std::nullopt;
}
 
// Returns true if `op` is an arith.extf that can be lowered to xevm.extf, i.e.
// a rank-1 widening from an MX narrow float to a 16-element f16/bf16 vector.
static bool isXeVMExtf(arith::ExtFOp op) {
  auto srcTy = dyn_cast<VectorType>(op.getIn().getType());
  auto dstTy = dyn_cast<VectorType>(op.getType());
  if (!srcTy || !dstTy || srcTy.getRank() != 1 || dstTy.getRank() != 1)
    return false;
  if (dstTy.getNumElements() != kXeVMExtfTruncfNumElems)
    return false;
  Type dstETy = dstTy.getElementType();
  if (!dstETy.isF16() && !dstETy.isBF16())
    return false;
  return getExtfNarrowType(srcTy.getElementType()).has_value();
}
 
// Returns true if `op` is an arith.truncf that can be lowered to xevm.truncf,
// i.e. a rank-1 truncation from an f16/bf16 vector to an MX narrow float. The
// source has to hold a whole number of the fixed-size groups xevm.truncf
// converts at a time; wider vectors are converted in several steps.
static bool isXeVMTruncf(arith::TruncFOp op) {
  auto srcTy = dyn_cast<VectorType>(op.getIn().getType());
  auto dstTy = dyn_cast<VectorType>(op.getType());
  if (!srcTy || !dstTy || srcTy.getRank() != 1 || dstTy.getRank() != 1)
    return false;
  int64_t numElems = srcTy.getNumElements();
  if (numElems == 0 || numElems % kXeVMExtfTruncfNumElems != 0)
    return false;
  Type srcETy = srcTy.getElementType();
  if (!srcETy.isF16() && !srcETy.isBF16())
    return false;
  return getTruncfNarrowType(dstTy.getElementType()).has_value();
}
 
class ExtfToXeVMPattern : public OpConversionPattern<arith::ExtFOp> {
  using Base::Base;
  LogicalResult
  matchAndRewrite(arith::ExtFOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    if (!isXeVMExtf(op))
      return rewriter.notifyMatchFailure(op, "not a xevm.extf compatible extf");
    Location loc = op.getLoc();
    MLIRContext *ctx = op.getContext();
    auto srcVecTy = cast<VectorType>(op.getIn().getType());
    auto dstVecTy = cast<VectorType>(op.getType());
    xevm::ExtfSrcElemTypes srcEnum =
        *getExtfNarrowType(srcVecTy.getElementType());
    xevm::ExtfDstElemTypes dstEnum = dstVecTy.getElementType().isF16()
                                         ? xevm::ExtfDstElemTypes::F16
                                         : xevm::ExtfDstElemTypes::BF16;
    // The narrow float operand has already been type-converted to an integer
    // vector of the same bit width (i4 for fp4, i8 for fp8). xevm.extf takes
    // the values packed into an i8 vector, so re-pack fp4 (i4) operands.
    Value src = adaptor.getIn();
    auto convSrcTy = cast<VectorType>(src.getType());
    if (convSrcTy.getElementTypeBitWidth() == 4)
      src = vector::BitCastOp::create(
          rewriter, loc,
          VectorType::get(convSrcTy.getNumElements() / 2, rewriter.getI8Type()),
          src);
    Type resTy = getTypeConverter()->convertType(dstVecTy);
    Value res = xevm::ExtfOp::create(
        rewriter, loc, resTy, src, xevm::ExtfSrcElemTypeAttr::get(ctx, srcEnum),
        xevm::ExtfDstElemTypeAttr::get(ctx, dstEnum));
    rewriter.replaceOp(op, res);
    return success();
  }
};
 
class TruncfToXeVMPattern : public OpConversionPattern<arith::TruncFOp> {
  using Base::Base;
  LogicalResult
  matchAndRewrite(arith::TruncFOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    if (!isXeVMTruncf(op))
      return rewriter.notifyMatchFailure(op,
                                         "not a xevm.truncf compatible truncf");
    Location loc = op.getLoc();
    MLIRContext *ctx = op.getContext();
    auto srcVecTy = cast<VectorType>(op.getIn().getType());
    auto dstVecTy = cast<VectorType>(op.getType());
    xevm::TruncfSrcElemTypes srcEnum = srcVecTy.getElementType().isF16()
                                           ? xevm::TruncfSrcElemTypes::F16
                                           : xevm::TruncfSrcElemTypes::BF16;
    xevm::TruncfDstElemTypes dstEnum =
        *getTruncfNarrowType(dstVecTy.getElementType());
    auto srcEnumAttr = xevm::TruncfSrcElemTypeAttr::get(ctx, srcEnum);
    auto dstEnumAttr = xevm::TruncfDstElemTypeAttr::get(ctx, dstEnum);
 
    // xevm.truncf lowers to instructions that convert a fixed number of
    // elements at a time, so a wider source is converted one group at a time
    // and the packed results are concatenated. Each group produces the narrow
    // floats packed into an i8 vector.
    int64_t numGroups = srcVecTy.getNumElements() / kXeVMExtfTruncfNumElems;
    int64_t groupBytes =
        kXeVMExtfTruncfNumElems * dstVecTy.getElementTypeBitWidth() / 8;
    Type groupTy = VectorType::get(groupBytes, rewriter.getI8Type());
 
    Value src = adaptor.getIn();
    Value packed;
    if (numGroups == 1) {
      packed = xevm::TruncfOp::create(rewriter, loc, groupTy, src, srcEnumAttr,
                                      dstEnumAttr);
    } else {
      auto packedTy =
          VectorType::get(groupBytes * numGroups, rewriter.getI8Type());
      packed = arith::ConstantOp::create(rewriter, loc, packedTy,
                                         rewriter.getZeroAttr(packedTy));
      for (int64_t group = 0; group < numGroups; group++) {
        Value slice = vector::ExtractStridedSliceOp::create(
            rewriter, loc, src, group * kXeVMExtfTruncfNumElems,
            kXeVMExtfTruncfNumElems, /*strides=*/1);
        Value converted = xevm::TruncfOp::create(rewriter, loc, groupTy, slice,
                                                 srcEnumAttr, dstEnumAttr);
        packed = vector::InsertStridedSliceOp::create(
            rewriter, loc, converted, packed, group * groupBytes,
            /*strides=*/1);
      }
    }
    // Re-shape to the type-converted result type (i4 vector for fp4).
    Type resTy = getTypeConverter()->convertType(dstVecTy);
    if (packed.getType() != resTy)
      packed = vector::BitCastOp::create(rewriter, loc, resTy, packed);
    rewriter.replaceOp(op, packed);
    return success();
  }
};
 
// Lowers `xegpu.lane_shuffle` to `xevm.bitcast_shuffle`.
//
// `xevm.bitcast_shuffle` concatenates the components of its operand across the
// subgroup, the first component of every lane first, and then hands chunks the
// size of a result component back out to the lanes in order. Numbering the
// elements of a `vector<NxT>` fragment held by lane `i` of a subgroup of size
// `S` by their logical position, that concatenation is exactly the `pack` mode
// input numbering `j * S + i`. Taking the result as a single `N * width(T)` bit
// scalar then hands lane `i` the logical positions `i * N .. i * N + N - 1`,
// which is the `pack` mode output numbering.
//
// So `pack` is a vector-to-scalar `xevm.bitcast_shuffle` followed by a bitcast
// back to the fragment type, and `unpack`, being its inverse, is a bitcast to
// the scalar followed by a scalar-to-vector `xevm.bitcast_shuffle`.
//
// `xevm.bitcast_shuffle` only accepts the integer types `i8`, `i16`, `i32` and
// `i64`, since it is bit-preserving and so does not depend on how the bits are
// interpreted. A fragment of a floating point type is therefore bitcast to a
// same-width integer vector on the way in and back on the way out.
class LaneShuffleToXeVMPattern
    : public OpConversionPattern<xegpu::LaneShuffleOp> {
  using Base::Base;
  LogicalResult
  matchAndRewrite(xegpu::LaneShuffleOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto vecTy = dyn_cast<VectorType>(adaptor.getSource().getType());
    if (!vecTy)
      return rewriter.notifyMatchFailure(op, "Expected a vector fragment.");
    // The shuffle redistributes whole bytes between the lanes, so sub-byte
    // element types, fp4 in particular, cannot be shuffled. Widths without a
    // matching integer type the op accepts are rejected for the same reason.
    unsigned elemBits = vecTy.getElementTypeBitWidth();
    if (elemBits != 8 && elemBits != 16 && elemBits != 32 && elemBits != 64)
      return rewriter.notifyMatchFailure(
          op, "Expected an element type of 8, 16, 32 or 64 bits.");
    int64_t fragmentBits = vecTy.getNumElements() * elemBits;
    if (fragmentBits > 64 || !llvm::isPowerOf2_64(fragmentBits))
      return rewriter.notifyMatchFailure(
          op, "Expected a fragment of 8, 16, 32 or 64 bits.");
 
    Location loc = op.getLoc();
    Type packedTy = rewriter.getIntegerType(fragmentBits);
    // The integer vector type the shuffle actually operates on. Equal to the
    // fragment type when that is already an integer vector.
    VectorType shuffleTy =
        VectorType::get(vecTy.getShape(), rewriter.getIntegerType(elemBits));
 
    Value res;
    if (op.getMode() == xegpu::LaneShuffleMode::Pack) {
      Value src = adaptor.getSource();
      if (shuffleTy != vecTy)
        src = LLVM::BitcastOp::create(rewriter, loc, shuffleTy, src);
      res = xevm::BitcastShuffleOp::create(rewriter, loc, packedTy, src);
      res = LLVM::BitcastOp::create(rewriter, loc, vecTy, res);
    } else {
      Value packed =
          LLVM::BitcastOp::create(rewriter, loc, packedTy, adaptor.getSource());
      res = xevm::BitcastShuffleOp::create(rewriter, loc, shuffleTy, packed);
      if (shuffleTy != vecTy)
        res = LLVM::BitcastOp::create(rewriter, loc, vecTy, res);
    }
    rewriter.replaceOp(op, res);
    return success();
  }
};
 
//===----------------------------------------------------------------------===//
// Pass Definition
//===----------------------------------------------------------------------===//
 
struct ConvertXeGPUToXeVMPass
    : public impl::ConvertXeGPUToXeVMPassBase<ConvertXeGPUToXeVMPass> {
  using Base::Base;
 
  void runOnOperation() override {
    MLIRContext *context = &getContext();
 
    // XeVM type converter is based on LLVM type converter with the
    // following customizations.
    // First, type conversion rules are added for xegpu custom types,
    // TensorDescType and MemDescType.
    // Second, MemRefType is lowered to single integer type
    // Third, VectorType of single element or 0D is converted to vector
    // element type. Otherwise, vector type is flatten to 1D.
    LowerToLLVMOptions options(context);
    options.overrideIndexBitwidth(this->use64bitIndex ? 64 : 32);
    LLVMTypeConverter typeConverter(context, options);
 
    Type xevmIndexType = typeConverter.convertType(IndexType::get(context));
    Type i32Type = IntegerType::get(context, 32);
    typeConverter.addConversion([&](VectorType type) -> Type {
      auto elemType = typeConverter.convertType(type.getElementType());
      // If the vector rank is 0 or has a single element, return the element
      unsigned rank = type.getRank();
      if (rank == 0 || 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 {
      if (type.getRank() == 1)
        return xevmIndexType;
      return VectorType::get(8, i32Type);
    });
    // SLM access related type conversions.
    // TODO: LLVM DLTI provides clean way of representing different pointer size
    // based on address space. Currently pointer size of SLM access is hard
    // coded to 32bit. Update to use DLTI when switching overall XeGPU lowering
    // to use DLTI instead of use64bitIndex option used above.
 
    // Convert MemDescType into i32 for SLM
    typeConverter.addConversion(
        [&](xegpu::MemDescType type) -> Type { return i32Type; });
 
    typeConverter.addConversion([&](MemRefType type) -> Type {
      return isSharedMemRef(type) ? i32Type : xevmIndexType;
    });
 
    // 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
    // Applies only to target materialization.
    // Note: int type to memref materialization is not required as xegpu ops
    // currently do not produce memrefs as result.
    auto memrefToIntMaterializationCast = [](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
    // Applies only to target materialization.
    // Note: i64 to ui64 materialization is not required as xegpu ops
    // currently do not produce ui64 as result.
    auto ui64ToI64MaterializationCast = [](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
    // Applies only to target materialization.
    // Note: i32 to ui32 materialization is not required as xegpu ops
    // currently do not produce ui32 as result.
    auto ui32ToI32MaterializationCast = [](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 between vector types
    //   - Add shape cast for different shapes
    //   - Add bitcast for different element types
    // Applies to both source and target materialization.
    auto vectorToVectorMaterializationCast = [](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 (auto targetVecTy = dyn_cast<VectorType>(type)) {
          Value cast = input;
          // If the target type has a different shape, add a shape cast
          // If the target type has a different element type, add a bitcast
          if (targetVecTy.getShape() != vecTy.getShape()) {
            cast = vector::ShapeCastOp::create(
                       builder, loc,
                       VectorType::get(targetVecTy.getShape(),
                                       vecTy.getElementType()),
                       cast)
                       .getResult();
          }
          if (targetVecTy.getElementType() != vecTy.getElementType()) {
            cast = vector::BitCastOp::create(builder, loc, targetVecTy, cast)
                       .getResult();
          }
          return cast;
        }
      }
      return {};
    };
 
    // Materialization to convert
    //   - single element vector to single element of vector element type
    // Applies only to target materialization.
    auto vectorToSingleElementMaterializationCast =
        [](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())) {
        // Source needs to be single element vector
        auto rank = vecTy.getRank();
        if (rank != 0 && vecTy.getNumElements() != 1)
          return {};
        auto inElemTy = vecTy.getElementType();
        // extract scalar
        Value cast = input;
        if (rank == 0) {
          cast = vector::ExtractOp::create(builder, loc, cast, {}).getResult();
        } else {
          cast = vector::ExtractOp::create(builder, loc, cast,
                                           SmallVector<int64_t>(rank, 0))
                     .getResult();
        }
        // Extracted element type may need conversion
        // Two cases
        // 1. Index type to integer type
        // 2. Other element type mismatch
        if (inElemTy.isIndex()) {
          cast = arith::IndexCastUIOp::create(builder, loc, type, cast)
                     .getResult();
        } else if (inElemTy != type) {
          cast = arith::BitcastOp::create(builder, loc, type, cast).getResult();
        }
        return cast;
      }
      return {};
    };
 
    // Materialization to convert
    //   - single element of vector element type to single element vector
    // If result type of original op is single element vector and lowered type
    // is scalar. This materialization cast creates a single element vector by
    // First convert element type if needed and then broadcast to single
    // element vector.
    // Applies only to source materialization.
    auto singleElementToVectorMaterializationCast =
        [](OpBuilder &builder, Type type, ValueRange inputs,
           Location loc) -> Value {
      if (inputs.size() != 1)
        return {};
      auto input = inputs.front();
      auto inTy = input.getType();
      if (!inTy.isIntOrFloat())
        return {};
      // If the target type is a vector of rank 0 or single element vector
      // of element type matching input type, broadcast input to target type.
      if (auto vecTy = dyn_cast<VectorType>(type)) {
        if (vecTy.getRank() != 0 && vecTy.getNumElements() != 1)
          return {};
        auto outElemTy = vecTy.getElementType();
        Value cast = input;
        if (outElemTy.isIndex()) {
          cast = arith::IndexCastUIOp::create(builder, loc,
                                              builder.getIndexType(), cast)
                     .getResult();
        } else if (inTy != outElemTy) {
          cast = arith::BitcastOp::create(builder, loc, outElemTy, cast)
                     .getResult();
        }
        return vector::BroadcastOp::create(builder, loc, vecTy, cast)
            .getResult();
      }
      return {};
    };
    typeConverter.addSourceMaterialization(
        singleElementToVectorMaterializationCast);
    typeConverter.addSourceMaterialization(vectorToVectorMaterializationCast);
    typeConverter.addTargetMaterialization(memrefToIntMaterializationCast);
    typeConverter.addTargetMaterialization(ui32ToI32MaterializationCast);
    typeConverter.addTargetMaterialization(ui64ToI64MaterializationCast);
    typeConverter.addTargetMaterialization(
        vectorToSingleElementMaterializationCast);
    typeConverter.addTargetMaterialization(vectorToVectorMaterializationCast);
    ConversionTarget target(*context);
    target.addLegalDialect<xevm::XeVMDialect, LLVM::LLVMDialect,
                           vector::VectorDialect, arith::ArithDialect,
                           memref::MemRefDialect, gpu::GPUDialect,
                           index::IndexDialect>();
    target.addIllegalDialect<xegpu::XeGPUDialect>();
    // arith.extf/arith.truncf between MX narrow floats and f16/bf16 are routed
    // to xevm.extf/xevm.truncf; all other arith float casts stay legal.
    target.addDynamicallyLegalOp<arith::ExtFOp>(
        [](arith::ExtFOp op) { return !isXeVMExtf(op); });
    target.addDynamicallyLegalOp<arith::TruncFOp>(
        [](arith::TruncFOp op) { return !isXeVMTruncf(op); });
 
    RewritePatternSet patterns(context);
    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());
  patterns.add<DpasMxToXeVMPattern>(typeConverter, patterns.getContext());
  patterns.add<ExtfToXeVMPattern, TruncfToXeVMPattern>(typeConverter,
                                                       patterns.getContext());
  patterns.add<LaneShuffleToXeVMPattern>(typeConverter, patterns.getContext());
}