//===- VectorToXeGPU.cpp - Convert vector to XeGPU dialect ------*- C++ -*-===//
//
// Part of the LLVM Project, 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
//
//===----------------------------------------------------------------------===//
//
// This file implements lowering of vector operations to XeGPU dialect ops.
//
//===----------------------------------------------------------------------===//

#include "mlir/Conversion/VectorToXeGPU/VectorToXeGPU.h"

#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/XeGPU/IR/XeGPU.h"
#include "mlir/Dialect/XeGPU/Utils/XeGPUUtils.h"
#include "mlir/Pass/Pass.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "llvm/ADT/TypeSwitch.h"

#include <algorithm>
#include <optional>

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

using namespace mlir;

namespace {

// Return true if value represents a zero constant.
static bool isZeroConstant(Value val) {
  auto constant = val.getDefiningOp<arith::ConstantOp>();
  if (!constant)
    return false;

  return TypeSwitch<Attribute, bool>(constant.getValue())
      .Case<FloatAttr>(
          [](auto floatAttr) { return floatAttr.getValue().isZero(); })
      .Case<IntegerAttr>(
          [](auto intAttr) { return intAttr.getValue().isZero(); })
      .Default([](auto) { return false; });
}

static LogicalResult storeLoadPreconditions(PatternRewriter &rewriter,
                                            Operation *op, VectorType vecTy) {
  // Validate only vector as the basic vector store and load ops guarantee
  // XeGPU-compatible memref source.
  unsigned vecRank = vecTy.getRank();
  if (!(vecRank == 1 || vecRank == 2))
    return rewriter.notifyMatchFailure(op, "Expects 1D or 2D vector");

  return success();
}

static LogicalResult transferPreconditions(PatternRewriter &rewriter,
                                           VectorTransferOpInterface xferOp) {
  if (xferOp.getMask())
    return rewriter.notifyMatchFailure(xferOp,
                                       "Masked transfer is not supported");

  auto srcTy = dyn_cast<MemRefType>(xferOp.getShapedType());
  if (!srcTy)
    return rewriter.notifyMatchFailure(xferOp, "Expects memref source");

  // Validate further transfer op semantics.
  SmallVector<int64_t> strides;
  int64_t offset;
  if (failed(srcTy.getStridesAndOffset(strides, offset)) || strides.back() != 1)
    return rewriter.notifyMatchFailure(
        xferOp, "Buffer must be contiguous in the innermost dimension");

  VectorType vecTy = xferOp.getVectorType();
  unsigned vecRank = vecTy.getRank();
  if (xferOp.hasOutOfBoundsDim() && vecRank < 2)
    return rewriter.notifyMatchFailure(
        xferOp, "Boundary check is available only for block instructions.");

  AffineMap map = xferOp.getPermutationMap();
  if (!map.isProjectedPermutation(/*allowZeroInResults=*/false))
    return rewriter.notifyMatchFailure(xferOp, "Unsupported permutation map");
  unsigned numInputDims = map.getNumInputs();
  for (AffineExpr expr : map.getResults().take_back(vecRank)) {
    auto dim = dyn_cast<AffineDimExpr>(expr);
    if (dim.getPosition() < (numInputDims - vecRank))
      return rewriter.notifyMatchFailure(
          xferOp, "Only the innermost dimensions can be accessed");
  }

  return success();
}

static xegpu::CreateNdDescOp
createNdDescriptor(PatternRewriter &rewriter, Location loc,
                   xegpu::TensorDescType descType, TypedValue<MemRefType> src,
                   Operation::operand_range offsets) {
  MemRefType srcTy = src.getType();
  auto [strides, offset] = srcTy.getStridesAndOffset();

  xegpu::CreateNdDescOp ndDesc;
  if (srcTy.hasStaticShape()) {
    ndDesc = xegpu::CreateNdDescOp::create(rewriter, loc, descType, src,
                                           getAsOpFoldResult(offsets));
  } else {
    // In case of any dynamic shapes, source's shape and strides have to be
    // explicitly provided.
    SmallVector<Value> sourceDims;
    unsigned srcRank = srcTy.getRank();
    for (unsigned i = 0; i < srcRank; ++i)
      sourceDims.push_back(memref::DimOp::create(rewriter, loc, src, i));

    SmallVector<int64_t> constOffsets;
    SmallVector<Value> dynOffsets;
    for (Value offset : offsets) {
      std::optional<int64_t> staticVal = getConstantIntValue(offset);
      if (!staticVal)
        dynOffsets.push_back(offset);
      constOffsets.push_back(staticVal.value_or(ShapedType::kDynamic));
    }

    SmallVector<Value> dynShapes;
    for (auto [idx, shape] : llvm::enumerate(srcTy.getShape())) {
      if (shape == ShapedType::kDynamic)
        dynShapes.push_back(sourceDims[idx]);
    }

    // Compute strides in reverse order.
    SmallVector<Value> dynStrides;
    Value accStride = arith::ConstantIndexOp::create(rewriter, loc, 1);
    // Last stride is guaranteed to be static and unit.
    for (int i = static_cast<int>(strides.size()) - 2; i >= 0; --i) {
      accStride =
          arith::MulIOp::create(rewriter, loc, accStride, sourceDims[i + 1]);
      if (strides[i] == ShapedType::kDynamic)
        dynStrides.push_back(accStride);
    }
    std::reverse(dynStrides.begin(), dynStrides.end());

    ndDesc = xegpu::CreateNdDescOp::create(
        rewriter, loc, descType, src, dynOffsets, dynShapes, dynStrides,
        DenseI64ArrayAttr::get(rewriter.getContext(), constOffsets),
        DenseI64ArrayAttr::get(rewriter.getContext(), srcTy.getShape()),
        DenseI64ArrayAttr::get(rewriter.getContext(), strides));
  }

  return ndDesc;
}

// Adjusts the strides of a memref according to a given permutation map for
// vector operations.
//
// This function updates the innermost strides in the `strides` array to
// reflect the permutation specified by `permMap`. The permutation is computed
// using the inverse and broadcasting-aware version of the permutation map,
// and is applied to the relevant strides. This ensures that memory accesses
// are consistent with the logical permutation of vector elements.
//
// Example:
//   Suppose we have a memref of rank 4 with strides `[s0, s1, s2, s3]`.
//   If the permutation map swaps the last two dimensions (e.g., [0, 1] -> [1,
//   0]), then after calling this function, the last two strides will be
//   swapped:
//     Original strides: [s0, s1, s2, s3]
//     After permutation: [s0, s1, s3, s2]
//
static void adjustStridesForPermutation(AffineMap permMap,
                                        SmallVectorImpl<Value> &strides) {

  AffineMap invMap = inverseAndBroadcastProjectedPermutation(permMap);
  SmallVector<unsigned> perms;
  invMap.isPermutationOfMinorIdentityWithBroadcasting(perms);
  SmallVector<int64_t> perms64(perms.begin(), perms.end());
  strides = applyPermutation(strides, perms64);
}

// Computes memory strides and a memref offset for vector transfer operations,
// handling both static and dynamic memrefs while applying permutation
// transformations for XeGPU lowering.
template <
    typename OpType,
    typename = std::enable_if_t<llvm::is_one_of<
        std::decay_t<OpType>, vector::TransferReadOp, vector::TransferWriteOp,
        vector::GatherOp, vector::ScatterOp>::value>>
static std::pair<SmallVector<Value>, Value>
computeMemrefMeta(OpType xferOp, PatternRewriter &rewriter) {
  SmallVector<Value> strides;
  Value baseMemref = xferOp.getBase();
  MemRefType memrefType = dyn_cast<MemRefType>(baseMemref.getType());

  Location loc = xferOp.getLoc();
  Value offsetVal = nullptr;
  if (memrefType.hasStaticShape()) {
    int64_t offset;
    SmallVector<int64_t> intStrides;
    if (failed(memrefType.getStridesAndOffset(intStrides, offset)))
      return {{}, offsetVal};
    bool hasDynamicStrides = llvm::any_of(intStrides, [](int64_t strideVal) {
      return ShapedType::isDynamic(strideVal);
    });

    if (!hasDynamicStrides)
      for (int64_t s : intStrides)
        strides.push_back(arith::ConstantIndexOp::create(rewriter, loc, s));

    if (!ShapedType::isDynamic(offset))
      offsetVal = arith::ConstantIndexOp::create(rewriter, loc, offset);
  }

  if (strides.empty() || !offsetVal) {
    // For dynamic shape memref, use memref.extract_strided_metadata to get
    // stride values
    unsigned rank = memrefType.getRank();
    Type indexType = rewriter.getIndexType();

    // Result types: [base_memref, offset, stride0, stride1, ..., strideN-1,
    // size0, size1, ..., sizeN-1]
    SmallVector<Type> resultTypes;
    resultTypes.push_back(MemRefType::get(
        {}, memrefType.getElementType())); // base memref (unranked)
    resultTypes.push_back(indexType);      // offset

    for (unsigned i = 0; i < rank; ++i)
      resultTypes.push_back(indexType); // strides

    for (unsigned i = 0; i < rank; ++i)
      resultTypes.push_back(indexType); // sizes

    auto meta = memref::ExtractStridedMetadataOp::create(
        rewriter, loc, resultTypes, baseMemref);

    if (strides.empty())
      strides.append(meta.getStrides().begin(), meta.getStrides().end());

    if (!offsetVal)
      offsetVal = meta.getOffset();
  }

  if constexpr (llvm::is_one_of<std::decay_t<OpType>, vector::TransferReadOp,
                                vector::TransferWriteOp>::value) {
    AffineMap permMap = xferOp.getPermutationMap();
    // Adjust strides according to the permutation map (e.g., for transpose)
    adjustStridesForPermutation(permMap, strides);
  }

  return {strides, offsetVal};
}

// This function compute the vectors of localOffsets for scattered load/stores.
// It is used in the lowering of vector.transfer_read/write to
// load_gather/store_scatter Example:
//   %0 = vector.transfer_read %expand_shape[%block_id_y, %c0, %c0, %c0, %c0],
//               %cst {in_bounds = [true, true, true, true]}>} :
//               memref<8x4x2x6x32xbf16>, vector<4x2x6x32xbf16>
//
//   %6 = vector.step: vector<4xindex>
//   %7 = vector.step: vector<2xindex>
//   %8 = vector.step: vector<6xindex>
//   %9 = vector.step: vector<32xindex>
//   %10 = arith.mul %6, 384
//   %11 = arith.mul %7, 192
//   %12 = arith.mul %8, 32
//   %13 = arith.mul %9, 1
//   %14 = vector.shape_cast %10: vector<4xindex> -> vector<4x1x1x1xbf16>
//   %15 = vector.shape_cast %11: vector<2xindex> -> vector<1x2x1x1xbf16>
//   %16 = vector.shape_cast %12: vector<6xindex> -> vector<1x1x6x1xbf16>
//   %17 = vector.shape_cast %13: vector<32xindex> -> vector<1x1x1x32xbf16>
//   %18 = vector.broadcast %14: vector<4x1x1x1xbf16> -> vector<4x2x6x32xindex>
//   %19 = vector.broadcast %15: vector<1x2x1x1xbf16> -> vector<4x2x6x32xindex>
//   %20 = vector.broadcast %16: vector<1x1x6x1xbf16> -> vector<4x2x6x32xindex>
//   %21 = vector.broadcast %17: vector<1x1x1x32xbf16> -> vector<4x2x6x32xindex>
//   %22 = arith.add %18, %19
//   %23 = arith.add %20, %21
//   %local_offsets = arith.add %22, %23
//   %orig_offset = %block_id_y * 4x2x6x32 // consider using affine map
//   %offsets =  memref_offset + orig_offset + local_offsets
static Value computeOffsets(VectorTransferOpInterface xferOp,
                            PatternRewriter &rewriter, ArrayRef<Value> strides,
                            Value baseOffset) {
  Location loc = xferOp.getLoc();
  VectorType vectorType = xferOp.getVectorType();
  SmallVector<Value> indices(xferOp.getIndices().begin(),
                             xferOp.getIndices().end());
  ArrayRef<int64_t> vectorShape = vectorType.getShape();

  // Create vector.step operations for each dimension
  SmallVector<Value> stepVectors;
  llvm::map_to_vector(vectorShape, [&](int64_t dim) {
    auto stepType = VectorType::get({dim}, rewriter.getIndexType());
    auto stepOp = vector::StepOp::create(rewriter, loc, stepType);
    stepVectors.push_back(stepOp);
    return stepOp;
  });

  // Multiply step vectors by corresponding strides
  size_t memrefRank = strides.size();
  size_t vectorRank = vectorShape.size();
  SmallVector<Value> strideMultiplied;
  for (size_t i = 0; i < vectorRank; ++i) {
    size_t memrefDim = memrefRank - vectorRank + i;
    Value strideValue = strides[memrefDim];
    auto mulType = dyn_cast<VectorType>(stepVectors[i].getType());
    auto bcastOp =
        vector::BroadcastOp::create(rewriter, loc, mulType, strideValue);
    auto mulOp = arith::MulIOp::create(rewriter, loc, stepVectors[i], bcastOp);
    strideMultiplied.push_back(mulOp);
  }

  // Shape cast each multiplied vector to add singleton dimensions
  SmallVector<Value> shapeCasted;
  for (size_t i = 0; i < vectorRank; ++i) {
    SmallVector<int64_t> newShape(vectorRank, 1);
    newShape[i] = vectorShape[i];
    auto newType = VectorType::get(newShape, rewriter.getIndexType());
    auto castOp = vector::ShapeCastOp::create(rewriter, loc, newType,
                                              strideMultiplied[i]);
    shapeCasted.push_back(castOp);
  }

  // Broadcast each shape-casted vector to full vector shape
  SmallVector<Value> broadcasted;
  auto fullIndexVectorType =
      VectorType::get(vectorShape, rewriter.getIndexType());
  for (Value shapeCastVal : shapeCasted) {
    auto broadcastOp = vector::BroadcastOp::create(
        rewriter, loc, fullIndexVectorType, shapeCastVal);
    broadcasted.push_back(broadcastOp);
  }

  // Add all broadcasted vectors together to compute local offsets
  Value localOffsets = broadcasted[0];
  for (size_t i = 1; i < broadcasted.size(); ++i)
    localOffsets =
        arith::AddIOp::create(rewriter, loc, localOffsets, broadcasted[i]);

  // Compute base offset from transfer read indices
  for (size_t i = 0; i < indices.size(); ++i) {
    Value strideVal = strides[i];
    Value offsetContrib =
        arith::MulIOp::create(rewriter, loc, indices[i], strideVal);
    baseOffset =
        arith::AddIOp::create(rewriter, loc, baseOffset, offsetContrib);
  }
  // Broadcast base offset to match vector shape
  Value bcastBase = vector::BroadcastOp::create(
      rewriter, loc, fullIndexVectorType, baseOffset);
  localOffsets = arith::AddIOp::create(rewriter, loc, bcastBase, localOffsets);
  return localOffsets;
}

// Compute the element-wise offsets for vector.gather or vector.scatter ops.
//
// This function linearizes the base offsets of the gather/scatter operation
// and combines them with the per-element indices to produce a final vector of
// memory offsets.
template <
    typename OpType,
    typename = std::enable_if_t<llvm::is_one_of<
        std::decay_t<OpType>, vector::GatherOp, vector::ScatterOp>::value>>
static Value computeOffsets(PatternRewriter &rewriter, OpType gatScatOp,
                            ArrayRef<Value> strides, Value baseOffset) {
  Location loc = gatScatOp.getLoc();
  SmallVector<Value> offsets = gatScatOp.getOffsets();
  for (size_t i = 0; i < offsets.size(); ++i) {
    Value offsetContrib =
        arith::MulIOp::create(rewriter, loc, offsets[i], strides[i]);
    baseOffset =
        arith::AddIOp::create(rewriter, loc, baseOffset, offsetContrib);
  }
  Value indices = gatScatOp.getIndices();
  VectorType vecType = cast<VectorType>(indices.getType());

  Value strideVector =
      vector::BroadcastOp::create(rewriter, loc, vecType, strides.back())
          .getResult();
  Value stridedIndices =
      arith::MulIOp::create(rewriter, loc, strideVector, indices).getResult();

  Value baseVector =
      vector::BroadcastOp::create(
          rewriter, loc,
          VectorType::get(vecType.getShape(), rewriter.getIndexType()),
          baseOffset)
          .getResult();
  return arith::AddIOp::create(rewriter, loc, baseVector, stridedIndices)
      .getResult();
}

template <
    typename OpType,
    typename = std::enable_if_t<llvm::is_one_of<
        std::decay_t<OpType>, vector::TransferReadOp, vector::TransferWriteOp,
        vector::GatherOp, vector::ScatterOp>::value>>
// Convert memref to i64 base pointer
static Value memrefToIndexPtr(OpType xferOp, PatternRewriter &rewriter) {
  Location loc = xferOp.getLoc();
  auto indexPtr = memref::ExtractAlignedPointerAsIndexOp::create(
                      rewriter, loc, xferOp.getBase())
                      .getResult();
  return arith::IndexCastOp::create(rewriter, loc, rewriter.getI64Type(),
                                    indexPtr)
      .getResult();
}

static LogicalResult lowerToScatteredLoadOp(vector::TransferReadOp readOp,
                                            PatternRewriter &rewriter) {

  Location loc = readOp.getLoc();
  VectorType vectorType = readOp.getVectorType();
  ArrayRef<int64_t> vectorShape = vectorType.getShape();
  auto memrefType = dyn_cast<MemRefType>(readOp.getShapedType());
  if (!memrefType)
    return rewriter.notifyMatchFailure(readOp, "Expected memref source");

  auto meta = computeMemrefMeta(readOp, rewriter);
  if (meta.first.empty())
    return rewriter.notifyMatchFailure(readOp, "Failed to compute strides");

  Value localOffsets =
      computeOffsets(readOp, rewriter, meta.first, meta.second);

  Value flatMemref = memrefToIndexPtr(readOp, rewriter);

  Value mask = vector::ConstantMaskOp::create(
      rewriter, loc, VectorType::get(vectorShape, rewriter.getI1Type()),
      vectorShape);
  auto gatherOp = xegpu::LoadGatherOp::create(
      rewriter, loc, vectorType, flatMemref, localOffsets, mask,
      /*chunk_size=*/IntegerAttr{},
      /*l1_hint=*/xegpu::CachePolicyAttr{},
      /*l2_hint=*/xegpu::CachePolicyAttr{},
      /*l3_hint=*/xegpu::CachePolicyAttr{});

  rewriter.replaceOp(readOp, gatherOp.getResult());
  return success();
}

static LogicalResult lowerToScatteredStoreOp(vector::TransferWriteOp writeOp,
                                             PatternRewriter &rewriter) {

  Location loc = writeOp.getLoc();
  VectorType vectorType = writeOp.getVectorType();
  ArrayRef<int64_t> vectorShape = vectorType.getShape();

  auto memrefType = dyn_cast<MemRefType>(writeOp.getShapedType());
  if (!memrefType)
    return rewriter.notifyMatchFailure(writeOp, "Expected memref source");

  auto meta = computeMemrefMeta(writeOp, rewriter);
  if (meta.first.empty())
    return rewriter.notifyMatchFailure(writeOp, "Failed to compute strides");

  Value localOffsets =
      computeOffsets(writeOp, rewriter, meta.first, meta.second);

  Value flatMemref = memrefToIndexPtr(writeOp, rewriter);

  Value mask = vector::ConstantMaskOp::create(
      rewriter, loc, VectorType::get(vectorShape, rewriter.getI1Type()),
      vectorShape);
  xegpu::StoreScatterOp::create(rewriter, loc, writeOp.getVector(), flatMemref,
                                localOffsets, mask,
                                /*chunk_size=*/IntegerAttr{},
                                /*l1_hint=*/xegpu::CachePolicyAttr{},
                                /*l2_hint=*/xegpu::CachePolicyAttr{},
                                /*l3_hint=*/xegpu::CachePolicyAttr{});
  rewriter.eraseOp(writeOp);
  return success();
}

struct TransferReadLowering : public OpRewritePattern<vector::TransferReadOp> {
  using OpRewritePattern<vector::TransferReadOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(vector::TransferReadOp readOp,
                                PatternRewriter &rewriter) const override {
    Location loc = readOp.getLoc();

    if (failed(transferPreconditions(rewriter, readOp)))
      return failure();

    // TODO:This check needs to be replaced with proper uArch capability check
    auto chip = xegpu::getChipStr(readOp);
    if (chip != "pvc" && chip != "bmg") {
      // lower to scattered load Op if the target HW doesn't have 2d block load
      // support
      // TODO: add support for OutOfBound access
      if (readOp.hasOutOfBoundsDim())
        return failure();
      return lowerToScatteredLoadOp(readOp, rewriter);
    }

    // Perform common data transfer checks.
    VectorType vecTy = readOp.getVectorType();
    if (failed(storeLoadPreconditions(rewriter, readOp, vecTy)))
      return failure();

    bool isOutOfBounds = readOp.hasOutOfBoundsDim();
    if (isOutOfBounds && !isZeroConstant(readOp.getPadding()))
      return rewriter.notifyMatchFailure(
          readOp, "Unsupported non-zero padded out-of-bounds read");

    AffineMap readMap = readOp.getPermutationMap();
    bool isTransposeLoad = !readMap.isMinorIdentity();

    Type elementType = vecTy.getElementType();
    unsigned minTransposeBitWidth = 32;
    if (isTransposeLoad &&
        elementType.getIntOrFloatBitWidth() < minTransposeBitWidth)
      return rewriter.notifyMatchFailure(
          readOp, "Unsupported data type for transposition");

    // If load is transposed, get the base shape for the tensor descriptor.
    SmallVector<int64_t> descShape(vecTy.getShape());
    if (isTransposeLoad)
      std::reverse(descShape.begin(), descShape.end());
    auto descType = xegpu::TensorDescType::get(
        descShape, elementType, /*array_length=*/1,
        /*boundary_check=*/isOutOfBounds, xegpu::MemorySpace::Global);

    xegpu::CreateNdDescOp ndDesc =
        createNdDescriptor(rewriter, loc, descType,
                           dyn_cast<TypedValue<MemRefType>>(readOp.getBase()),
                           readOp.getIndices());

    DenseI64ArrayAttr transposeAttr =
        !isTransposeLoad ? nullptr
                         : DenseI64ArrayAttr::get(rewriter.getContext(),
                                                  ArrayRef<int64_t>{1, 0});
    // By default, no specific caching policy is assigned.
    xegpu::CachePolicyAttr hint = nullptr;
    auto loadOp = xegpu::LoadNdOp::create(rewriter, loc, vecTy, ndDesc,
                                          /*packed=*/nullptr, transposeAttr,
                                          /*l1_hint=*/hint,
                                          /*l2_hint=*/hint, /*l3_hint=*/hint);
    rewriter.replaceOp(readOp, loadOp);

    return success();
  }
};

struct TransferWriteLowering
    : public OpRewritePattern<vector::TransferWriteOp> {
  using OpRewritePattern<vector::TransferWriteOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(vector::TransferWriteOp writeOp,
                                PatternRewriter &rewriter) const override {
    Location loc = writeOp.getLoc();

    if (failed(transferPreconditions(rewriter, writeOp)))
      return failure();

    // TODO:This check needs to be replaced with proper uArch capability check
    auto chip = xegpu::getChipStr(writeOp);
    if (chip != "pvc" && chip != "bmg") {
      // lower to scattered store Op if the target HW doesn't have 2d block
      // store support
      // TODO: add support for OutOfBound access
      if (writeOp.hasOutOfBoundsDim())
        return failure();
      return lowerToScatteredStoreOp(writeOp, rewriter);
    }

    // Perform common data transfer checks.
    VectorType vecTy = writeOp.getVectorType();
    if (failed(storeLoadPreconditions(rewriter, writeOp, vecTy)))
      return failure();

    AffineMap map = writeOp.getPermutationMap();
    if (!map.isMinorIdentity())
      return rewriter.notifyMatchFailure(writeOp, "Expects identity map");

    auto descType = xegpu::TensorDescType::get(
        vecTy.getShape(), vecTy.getElementType(),
        /*array_length=*/1, /*boundary_check=*/writeOp.hasOutOfBoundsDim(),
        xegpu::MemorySpace::Global);
    xegpu::CreateNdDescOp ndDesc =
        createNdDescriptor(rewriter, loc, descType,
                           dyn_cast<TypedValue<MemRefType>>(writeOp.getBase()),
                           writeOp.getIndices());

    // By default, no specific caching policy is assigned.
    xegpu::CachePolicyAttr hint = nullptr;
    auto storeOp =
        xegpu::StoreNdOp::create(rewriter, loc, writeOp.getVector(), ndDesc,
                                 /*l1_hint=*/hint,
                                 /*l2_hint=*/hint, /*l3_hint=*/hint);
    rewriter.replaceOp(writeOp, storeOp);

    return success();
  }
};

struct GatherLowering : public OpRewritePattern<vector::GatherOp> {
  using OpRewritePattern<vector::GatherOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(vector::GatherOp gatherOp,
                                PatternRewriter &rewriter) const override {
    auto srcTy = dyn_cast<MemRefType>(gatherOp.getBase().getType());
    if (!srcTy)
      return rewriter.notifyMatchFailure(gatherOp, "Expects memref source");

    Location loc = gatherOp.getLoc();
    VectorType vectorType = gatherOp.getVectorType();

    auto meta = computeMemrefMeta(gatherOp, rewriter);
    if (meta.first.empty())
      return rewriter.notifyMatchFailure(gatherOp, "Failed to compute strides");

    Value localOffsets =
        computeOffsets(rewriter, gatherOp, meta.first, meta.second);
    Value flatMemref = memrefToIndexPtr(gatherOp, rewriter);

    auto xeGatherOp = xegpu::LoadGatherOp::create(
        rewriter, loc, vectorType, flatMemref, localOffsets, gatherOp.getMask(),
        /*chunk_size=*/IntegerAttr{},
        /*l1_hint=*/xegpu::CachePolicyAttr{},
        /*l2_hint=*/xegpu::CachePolicyAttr{},
        /*l3_hint=*/xegpu::CachePolicyAttr{});

    auto selectOp =
        arith::SelectOp::create(rewriter, loc, gatherOp.getMask(),
                                xeGatherOp.getResult(), gatherOp.getPassThru());
    rewriter.replaceOp(gatherOp, selectOp.getResult());
    return success();
  }
};

struct ScatterLowering : public OpRewritePattern<vector::ScatterOp> {
  using OpRewritePattern<vector::ScatterOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(vector::ScatterOp scatterOp,
                                PatternRewriter &rewriter) const override {
    auto srcTy = dyn_cast<MemRefType>(scatterOp.getBase().getType());
    if (!srcTy)
      return rewriter.notifyMatchFailure(scatterOp, "Expects memref source");

    Location loc = scatterOp.getLoc();
    auto meta = computeMemrefMeta(scatterOp, rewriter);
    if (meta.first.empty())
      return rewriter.notifyMatchFailure(scatterOp,
                                         "Failed to compute strides");

    Value localOffsets =
        computeOffsets(rewriter, scatterOp, meta.first, meta.second);
    Value flatMemref = memrefToIndexPtr(scatterOp, rewriter);

    xegpu::StoreScatterOp::create(rewriter, loc, scatterOp.getValueToStore(),
                                  flatMemref, localOffsets, scatterOp.getMask(),
                                  /*chunk_size=*/IntegerAttr{},
                                  /*l1_hint=*/xegpu::CachePolicyAttr{},
                                  /*l2_hint=*/xegpu::CachePolicyAttr{},
                                  /*l3_hint=*/xegpu::CachePolicyAttr{});
    rewriter.eraseOp(scatterOp);
    return success();
  }
};

struct LoadLowering : public OpRewritePattern<vector::LoadOp> {
  using OpRewritePattern<vector::LoadOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(vector::LoadOp loadOp,
                                PatternRewriter &rewriter) const override {
    Location loc = loadOp.getLoc();

    VectorType vecTy = loadOp.getResult().getType();
    if (failed(storeLoadPreconditions(rewriter, loadOp, vecTy)))
      return failure();

    // Boundary check is available only for block instructions.
    bool boundaryCheck = vecTy.getRank() > 1;

    auto descType = xegpu::TensorDescType::get(
        vecTy.getShape(), vecTy.getElementType(), /*array_length=*/1,
        boundaryCheck, xegpu::MemorySpace::Global);
    xegpu::CreateNdDescOp ndDesc = createNdDescriptor(
        rewriter, loc, descType, loadOp.getBase(), loadOp.getIndices());

    // By default, no specific caching policy is assigned.
    xegpu::CachePolicyAttr hint = nullptr;
    auto loadNdOp = xegpu::LoadNdOp::create(
        rewriter, loc, vecTy, ndDesc, /*packed=*/nullptr, /*transpose=*/nullptr,
        /*l1_hint=*/hint,
        /*l2_hint=*/hint, /*l3_hint=*/hint);
    rewriter.replaceOp(loadOp, loadNdOp);

    return success();
  }
};

struct StoreLowering : public OpRewritePattern<vector::StoreOp> {
  using OpRewritePattern<vector::StoreOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(vector::StoreOp storeOp,
                                PatternRewriter &rewriter) const override {
    Location loc = storeOp.getLoc();

    TypedValue<VectorType> vector = storeOp.getValueToStore();
    VectorType vecTy = vector.getType();
    if (failed(storeLoadPreconditions(rewriter, storeOp, vecTy)))
      return failure();

    // Boundary check is available only for block instructions.
    bool boundaryCheck = vecTy.getRank() > 1;

    auto descType = xegpu::TensorDescType::get(
        vecTy.getShape(), vecTy.getElementType(),
        /*array_length=*/1, boundaryCheck, xegpu::MemorySpace::Global);
    xegpu::CreateNdDescOp ndDesc = createNdDescriptor(
        rewriter, loc, descType, storeOp.getBase(), storeOp.getIndices());

    // By default, no specific caching policy is assigned.
    xegpu::CachePolicyAttr hint = nullptr;
    auto storeNdOp =
        xegpu::StoreNdOp::create(rewriter, loc, vector, ndDesc,
                                 /*l1_hint=*/hint,
                                 /*l2_hint=*/hint, /*l3_hint=*/hint);
    rewriter.replaceOp(storeOp, storeNdOp);

    return success();
  }
};

struct ContractionLowering : public OpRewritePattern<vector::ContractionOp> {
  using OpRewritePattern<vector::ContractionOp>::OpRewritePattern;

  LogicalResult matchAndRewrite(vector::ContractionOp contractOp,
                                PatternRewriter &rewriter) const override {
    Location loc = contractOp.getLoc();

    if (contractOp.getKind() != vector::CombiningKind::ADD)
      return rewriter.notifyMatchFailure(contractOp,
                                         "Expects add combining kind");

    TypedValue<Type> acc = contractOp.getAcc();
    VectorType accType = dyn_cast<VectorType>(acc.getType());
    if (!accType || accType.getRank() != 2)
      return rewriter.notifyMatchFailure(contractOp, "Expects acc 2D vector");

    // Accept only plain 2D data layout.
    // VNNI packing is applied to DPAS as a separate lowering step.
    TypedValue<VectorType> lhs = contractOp.getLhs();
    TypedValue<VectorType> rhs = contractOp.getRhs();
    if (lhs.getType().getRank() != 2 || rhs.getType().getRank() != 2)
      return rewriter.notifyMatchFailure(contractOp,
                                         "Expects lhs and rhs 2D vectors");

    if (!isRowMajorMatmul(contractOp.getIndexingMapsAttr()))
      return rewriter.notifyMatchFailure(contractOp, "Invalid indexing maps");

    auto dpasOp = xegpu::DpasOp::create(rewriter, loc,
                                        TypeRange{contractOp.getResultType()},
                                        ValueRange{lhs, rhs, acc});
    rewriter.replaceOp(contractOp, dpasOp);

    return success();
  }
};

struct ConvertVectorToXeGPUPass
    : public impl::ConvertVectorToXeGPUBase<ConvertVectorToXeGPUPass> {
  void runOnOperation() override {
    RewritePatternSet patterns(&getContext());
    populateVectorToXeGPUConversionPatterns(patterns);
    if (failed(applyPatternsGreedily(getOperation(), std::move(patterns))))
      return signalPassFailure();
  }
};

} // namespace

void mlir::populateVectorToXeGPUConversionPatterns(
    RewritePatternSet &patterns) {
  patterns
      .add<TransferReadLowering, TransferWriteLowering, LoadLowering,
           ScatterLowering, GatherLowering, StoreLowering, ContractionLowering>(
          patterns.getContext());
}