VectorToXeGPU.cpp

Go to the documentation of this file.

//===- 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/Conversion/VectorToGPU/VectorToGPU.h"
 
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/UB/IR/UBOps.h"
#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/Vector/Utils/VectorUtils.h"
#include "mlir/Dialect/XeGPU/IR/XeGPU.h"
#include "mlir/Dialect/XeGPU/Utils/XeGPUUtils.h"
#include "mlir/Dialect/XeGPU/uArch/uArchBase.h"
#include "mlir/Dialect/XeGPU/uArch/uArchCommon.h"
#include "mlir/IR/Matchers.h"
#include "mlir/Interfaces/SideEffectInterfaces.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 floatAttr) { return floatAttr.getValue().isZero(); })
      .Case([](IntegerAttr intAttr) { return intAttr.getValue().isZero(); })
      .Default(false);
}
 
// Return true if the transfer padding value is compatible with the implicit
// padding of an nd block load. LoadNdOp fills out-of-bounds elements with zero,
// so a zero constant matches its semantics exactly. A poison padding means the
// out-of-bounds elements are "don't care", so any implicit padding (including
// zero) is also acceptable.
static bool isZeroOrPoisonPadding(Value val) {
  return isZeroConstant(val) || val.getDefiningOp<ub::PoisonOp>();
}
 
// Return true if the permutation map keeps every dimension in place except the
// innermost two, which are swapped, e.g.:
//   (d0, d1) -> (d1, d0)
//   (d0, d1, d2) -> (d2, d1)
//   (d0, d1, d2, d3) -> (d0, d1, d3, d2)
// This is the only non-identity permutation an nd block load can realize (by
// loading the untransposed block and applying a trailing vector.transpose).
static bool isInnermostTwoDimsTransposed(AffineMap map) {
  unsigned numResults = map.getNumResults();
  if (numResults < 2)
    return false;
  MLIRContext *ctx = map.getContext();
  unsigned numInputs = map.getNumInputs();
  // All but the innermost two results must match the minor-identity map.
  for (unsigned i = 0; i + 2 < numResults; ++i)
    if (map.getResult(i) != getAffineDimExpr(numInputs - numResults + i, ctx))
      return false;
  // The innermost two results must be the last two input dims, swapped.
  return map.getResult(numResults - 2) ==
             getAffineDimExpr(numInputs - 1, ctx) &&
         map.getResult(numResults - 1) == getAffineDimExpr(numInputs - 2, ctx);
}
 
// Return true if `uArch` can transfer a `shape`-shaped tile of `elemTy`
// elements with the subgroup 2D block instruction `instKind`.
//
// A 2D block instruction accesses the two innermost dimensions of the tile;
// leading dimensions are unrolled into a sequence of 2D accesses. The
// innermost 2D tile does not have to match a hardware block exactly, because
// the later XeGPU passes (layout propagation and blocking) split it into
// hardware-sized blocks - but that split only exists when each of the two
// innermost extents is a multiple of a supported block extent. Transfers
// without such a split cannot be expressed with block instructions at all and
// must use a different lowering.
static bool isSupportedBlockShape(const xegpu::uArch::uArch *uArch,
                                  xegpu::uArch::InstructionKind instKind,
                                  ArrayRef<int64_t> shape, Type elemTy,
                                  bool hasTranspose = false) {
  if (shape.size() < 2)
    return false;
  if (!uArch || !uArch->isSupportedInstruction(instKind))
    return false;
 
  const auto *blockInst = dyn_cast<xegpu::uArch::BlockIOInstructionInterface>(
      uArch->getInstruction(instKind));
  if (!blockInst)
    return false;
 
  int width = static_cast<int>(shape.back());
  int height = static_cast<int>(shape[shape.size() - 2]);
 
  // A tile is accessible when both of its extents are multiples of a supported
  // block extent, so that the later XeGPU passes can split it into hardware
  // blocks. A missing entry means this variant does not support the element
  // type.
  auto fitsBlockShapes = [&](bool hasTransform) {
    std::optional<xegpu::uArch::BlockIOInstructionInterface::BlockShapes>
        blockShapes = blockInst->getBlockWidthHeightCount(elemTy, hasTransform,
                                                          hasTranspose);
    if (!blockShapes)
      return false;
    auto [widths, heights, counts] = *blockShapes;
    return xegpu::getLargestDivisor(width, widths) != -1 &&
           xegpu::getLargestDivisor(height, heights) != -1;
  };
 
  // Whether a tile is eventually loaded with the transformed (VNNI) variant is
  // only decided later, when layout propagation knows whether it feeds the B
  // operand of a matrix operation, so accept a tile that either variant can
  // access.
  return fitsBlockShapes(/*hasTransform=*/false) ||
         fitsBlockShapes(/*hasTransform=*/true);
}
 
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");
  if (isa<VectorType>(srcTy.getElementType()))
    return rewriter.notifyMatchFailure(xferOp,
                                       "Expects a memref of scalar elements");
 
  // Validate further transfer op semantics.
  SmallVector<int64_t> strides;
  int64_t offset;
  if (failed(srcTy.getStridesAndOffset(strides, offset)))
    return rewriter.notifyMatchFailure(xferOp,
                                       "The memref strides cannot be inferred");
  if (strides.empty())
    return rewriter.notifyMatchFailure(xferOp, "0D memref is not supported");
  if (strides.back() != 1)
    return rewriter.notifyMatchFailure(
        xferOp, "Buffer must be contiguous in the innermost dimension");
 
  VectorType vecTy = xferOp.getVectorType();
  unsigned vecRank = vecTy.getRank();
  if (vecRank == 0)
    return rewriter.notifyMatchFailure(xferOp, "0D vectors are not supported");
 
  AffineMap map = xferOp.getPermutationMap();
  if (!map.isProjectedPermutation(/*allowZeroInResults=*/true))
    return rewriter.notifyMatchFailure(xferOp, "Unsupported permutation map");
  // The transfer is built from a run of innermost source dimensions, as wide as
  // the vector, with one dimension of slack for a 1D transfer: dropping the
  // broadcast results of a map like (d1, 0) leaves a rank-1 transfer that still
  // reads d1, and the scattered path reads along any dimension anyway.
  unsigned numInputDims = map.getNumInputs();
  unsigned numAccessibleDims = std::max(vecRank, 2u);
  unsigned firstAccessibleDim =
      numInputDims > numAccessibleDims ? numInputDims - numAccessibleDims : 0;
  for (AffineExpr expr : map.getResults().take_back(vecRank)) {
    // A broadcast result names no source dimension, so it constrains nothing.
    auto dim = dyn_cast<AffineDimExpr>(expr);
    if (!dim)
      continue;
    if (dim.getPosition() < firstAccessibleDim)
      return rewriter.notifyMatchFailure(
          xferOp, "Only the innermost dimensions can be accessed");
  }
 
  return success();
}
 
// Returns the memref strides in vector-dimension order: the stride at vector
// dimension `v` is the stride of the memref dimension the permutation map has
// `v` read, so that stepping one element along `v` steps that far in memory.
//
// Example:
//   For a memref of rank 3 with strides `[s0, s1, s2]` read into a rank-2
//   vector through `(d0, d1, d2) -> (d2, d1)`, the result is `[s2, s1]`.
//
static SmallVector<Value> getStridesInVectorOrder(AffineMap permMap,
                                                  ArrayRef<Value> strides) {
  return llvm::map_to_vector(permMap.getResults(), [&](AffineExpr expr) {
    return strides[cast<AffineDimExpr>(expr).getPosition()];
  });
}
 
// 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();
  }
 
  // Strides are returned in original memref order; permutation is applied in
  // computeOffsets only where offsets are indexed in vector order.
  return {strides, offsetVal};
}
 
// Adds the transfer's indices to `baseOffset`, the memref's own element offset:
// the element at `indices` sits at `baseOffset + sum(indices[d] * strides[d])`.
static Value computeBaseOffset(VectorTransferOpInterface xferOp,
                               PatternRewriter &rewriter,
                               ArrayRef<Value> strides, Value baseOffset) {
  Location loc = xferOp.getLoc();
  for (auto [index, stride] : llvm::zip_equal(xferOp.getIndices(), strides)) {
    Value contrib = arith::MulIOp::create(rewriter, loc, index, stride);
    baseOffset = arith::AddIOp::create(rewriter, loc, baseOffset, contrib);
  }
  return baseOffset;
}
 
// 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();
  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;
  });
 
  SmallVector<Value> vectorStrides =
      getStridesInVectorOrder(xferOp.getPermutationMap(), strides);
 
  // Multiply step vectors by corresponding strides
  size_t vectorRank = vectorShape.size();
  SmallVector<Value> strideMultiplied;
  for (size_t i = 0; i < vectorRank; ++i) {
    Value strideValue = vectorStrides[i];
    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]);
 
  // Broadcast base offset to match vector shape
  baseOffset = computeBaseOffset(xferOp, rewriter, strides, baseOffset);
  Value bcastBase = vector::BroadcastOp::create(
      rewriter, loc, fullIndexVectorType, baseOffset);
  localOffsets = arith::AddIOp::create(rewriter, loc, bcastBase, localOffsets);
  return localOffsets;
}
 
// Returns the size of memref dimension `dim` as a value.
static Value getMemrefDimSize(VectorTransferOpInterface xferOp, unsigned dim,
                              PatternRewriter &rewriter) {
  Location loc = xferOp.getLoc();
  auto memrefTy = cast<MemRefType>(xferOp.getShapedType());
  if (memrefTy.isDynamicDim(dim))
    return memref::DimOp::create(rewriter, loc, xferOp.getBase(), dim)
        .getResult();
  return arith::ConstantIndexOp::create(rewriter, loc, memrefTy.getDimSize(dim))
      .getResult();
}
 
// Builds the mask a scattered access needs to stay within the source bounds.
//
// The element at vector position `i` along vector dimension `v` accesses memref
// dimension `d` (the one `v` maps to) at `indices[d] + i`, so it is in bounds
// iff `i < dim(d) - indices[d]`. Dimensions the transfer declares in-bounds are
// skipped; non-vector dimensions are always in bounds by the op's contract. The
// per-dimension masks are spread over the full vector shape and combined, the
// same way `computeOffsets` spreads the per-dimension offsets.
//
// Example, for a `vector<8xf32>` read of a `memref<?xf32>` at `%off`:
//   %dim = memref.dim %src, %c0
//   %limit = arith.subi %dim, %off
//   %step = vector.step : vector<8xindex>
//   %bcast = vector.broadcast %limit : index to vector<8xindex>
//   %mask = arith.cmpi slt, %step, %bcast : vector<8xindex>
static Value computeInBoundsMask(VectorTransferOpInterface xferOp,
                                 PatternRewriter &rewriter) {
  Location loc = xferOp.getLoc();
  ArrayRef<int64_t> vectorShape = xferOp.getVectorType().getShape();
  auto maskType = VectorType::get(vectorShape, rewriter.getI1Type());
  AffineMap map = xferOp.getPermutationMap();
  OperandRange indices = xferOp.getIndices();
 
  Value mask;
  for (unsigned v = 0, e = vectorShape.size(); v < e; ++v) {
    if (xferOp.isDimInBounds(v))
      continue;
    unsigned d = cast<AffineDimExpr>(map.getResult(v)).getPosition();
    Value bound = getMemrefDimSize(xferOp, d, rewriter);
    // A negative limit masks the dimension off entirely, which is what a
    // starting index past the end of the dimension should do.
    Value limit = arith::SubIOp::create(rewriter, loc, bound, indices[d]);
    auto stepType = VectorType::get({vectorShape[v]}, rewriter.getIndexType());
    Value step = vector::StepOp::create(rewriter, loc, stepType);
    Value limitVec =
        vector::BroadcastOp::create(rewriter, loc, stepType, limit);
    Value dimMask = arith::CmpIOp::create(
        rewriter, loc, arith::CmpIPredicate::slt, step, limitVec);
    if (vectorShape.size() > 1) {
      SmallVector<int64_t> expandedShape(vectorShape.size(), 1);
      expandedShape[v] = vectorShape[v];
      dimMask = vector::ShapeCastOp::create(
          rewriter, loc, VectorType::get(expandedShape, rewriter.getI1Type()),
          dimMask);
      dimMask = vector::BroadcastOp::create(rewriter, loc, maskType, dimMask);
    }
    mask = mask
               ? arith::AndIOp::create(rewriter, loc, mask, dimMask).getResult()
               : dimMask;
  }
  if (mask)
    return mask;
  return vector::ConstantMaskOp::create(rewriter, loc, maskType, vectorShape);
}
 
// Builds that mask as a scalar `i1`, for a transfer of a single element. The
// element sits at `indices`, so `i < dim(d) - indices[d]` at `i == 0` reduces
// to `indices[d] < dim(d)`, which needs neither the limit subtraction nor the
// step vector. The dimensions checked are the same.
//
// Example, for a `vector<1xf32>` read of a `memref<?xf32>` at `%off`:
//   %dim = memref.dim %src, %c0
//   %mask = arith.cmpi slt, %off, %dim : index
static Value computeUnitInBoundsMask(VectorTransferOpInterface xferOp,
                                     PatternRewriter &rewriter) {
  Location loc = xferOp.getLoc();
  AffineMap map = xferOp.getPermutationMap();
  OperandRange indices = xferOp.getIndices();
 
  Value mask;
  for (unsigned v = 0, e = xferOp.getVectorType().getRank(); v < e; ++v) {
    if (xferOp.isDimInBounds(v))
      continue;
    unsigned d = cast<AffineDimExpr>(map.getResult(v)).getPosition();
    Value bound = getMemrefDimSize(xferOp, d, rewriter);
    Value dimMask = arith::CmpIOp::create(
        rewriter, loc, arith::CmpIPredicate::slt, indices[d], bound);
    mask = mask
               ? arith::AndIOp::create(rewriter, loc, mask, dimMask).getResult()
               : dimMask;
  }
  if (mask)
    return mask;
  return arith::ConstantOp::create(rewriter, loc, rewriter.getBoolAttr(true));
}
 
// 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();
}
 
// Collapses shapes of a nD memref to the target rank while applying offsets for
// the collapsed dimensions. Returns the new memref value and the remaining
// offsets for the last targetRank dimensions. For example:
//   input: %memref = memref<2x4x8x32xf32>, offsets=[%i0, %i1, %i2, %i3],
//   output: %memref[%i0, %i1, 0, 0] -> memref<8x32xf32>, offsets: [%i2, %i3]
static std::pair<Value, SmallVector<OpFoldResult>>
convertMemrefAndOffsetsToTargetRank(PatternRewriter &rewriter, Location loc,
                                    Value memref,
                                    SmallVector<OpFoldResult> offsets,
                                    int64_t targetRank) {
  auto memrefType = cast<MemRefType>(memref.getType());
  unsigned rank = memrefType.getRank();
 
  if (rank <= targetRank)
    return {memref, offsets};
 
  int64_t numCombinedDims = rank - targetRank;
  SmallVector<OpFoldResult> subviewOffsets;
  SmallVector<OpFoldResult> subviewSizes;
  SmallVector<OpFoldResult> subviewStrides;
 
  // For the combined dimensions: use the provided offsets, size=1, stride=1
  for (unsigned i = 0; i < numCombinedDims; ++i) {
    subviewOffsets.push_back(offsets[i]);
    subviewSizes.push_back(rewriter.getI64IntegerAttr(1));
    subviewStrides.push_back(rewriter.getI64IntegerAttr(1));
  }
 
  // For the last targetRank dimensions: offset=0, use full size, stride=1
  SmallVector<int64_t> resultShape;
  auto originalShape = memrefType.getShape();
  auto meta = memref::ExtractStridedMetadataOp::create(rewriter, loc, memref);
  for (unsigned i = numCombinedDims; i < rank; ++i) {
    subviewOffsets.push_back(rewriter.getI64IntegerAttr(0));
    if (ShapedType::isDynamic(originalShape[i])) {
      subviewSizes.push_back(meta.getSizes()[i]);
      resultShape.push_back(ShapedType::kDynamic);
    } else {
      subviewSizes.push_back(rewriter.getI64IntegerAttr(originalShape[i]));
      resultShape.push_back(originalShape[i]);
    }
    subviewStrides.push_back(rewriter.getI64IntegerAttr(1));
  }
 
  auto resultType = memref::SubViewOp::inferRankReducedResultType(
      resultShape, memrefType, subviewOffsets, subviewSizes, subviewStrides);
  auto subviewOp =
      memref::SubViewOp::create(rewriter, loc, resultType, memref,
                                subviewOffsets, subviewSizes, subviewStrides);
 
  // Return the remaining offsets for the last targetRank dimensions
  SmallVector<OpFoldResult> newOffsets(offsets.begin() + numCombinedDims,
                                       offsets.end());
  return {subviewOp.getResult(), newOffsets};
}
 
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();
}
 
// Returns true if every use of `vec` extracts a scalar element from it. An
// unused value does not qualify.
static bool isUsedAsScalar(Value vec) {
  if (vec.use_empty())
    return false;
  return llvm::all_of(vec.getUsers(), [](Operation *user) {
    auto extractOp = dyn_cast<vector::ExtractOp>(user);
    return extractOp && !isa<VectorType>(extractOp.getResult().getType());
  });
}
 
// Lowers a transfer of a single element to a scalar `xegpu.load`.
//
// The transfer touches one location, which its own indices name, so the load
// takes a scalar offset and a scalar mask - no descriptor, no lane vectors -
// and the result is broadcast back to the transfer's unit-size vector type.
//
//   %off = arith.addi %base, %contrib : index
//   %inb = arith.cmpi slt, %idx, %dim : index
//   %val = xegpu.load %src[%off], %inb : i64, index, i1 -> f32
//   %vec = vector.broadcast %val : f32 to vector<1xf32>
static LogicalResult lowerToScalarLoadOp(vector::TransferReadOp readOp,
                                         PatternRewriter &rewriter) {
  Location loc = readOp.getLoc();
  VectorType vectorType = readOp.getVectorType();
  if (!isa<MemRefType>(readOp.getShapedType()))
    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 offset = computeBaseOffset(readOp, rewriter, meta.first, meta.second);
  Value flatMemref = memrefToIndexPtr(readOp, rewriter);
 
  Value mask = computeUnitInBoundsMask(readOp, rewriter);
  auto loadOp = xegpu::LoadGatherOp::create(
      rewriter, loc, vectorType.getElementType(), flatMemref, offset, mask,
      /*l1_hint=*/xegpu::CachePolicyAttr{},
      /*l2_hint=*/xegpu::CachePolicyAttr{},
      /*l3_hint=*/xegpu::CachePolicyAttr{},
      /*layout=*/nullptr, /*contiguity=*/nullptr);
 
  // A masked-off xegpu.load is unspecified, so the padding has to be applied
  // explicitly, as in lowerToScatteredLoadOp. A poison padding is "don't care".
  Value scalar = loadOp.getResult();
  if (readOp.hasOutOfBoundsDim() &&
      !readOp.getPadding().getDefiningOp<ub::PoisonOp>())
    scalar = arith::SelectOp::create(rewriter, loc, mask, scalar,
                                     readOp.getPadding());
 
  rewriter.replaceOpWithNewOp<vector::BroadcastOp>(readOp, vectorType, scalar);
  return success();
}
 
static LogicalResult lowerToScatteredLoadOp(vector::TransferReadOp readOp,
                                            PatternRewriter &rewriter) {
 
  Location loc = readOp.getLoc();
  VectorType vectorType = readOp.getVectorType();
  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 = computeInBoundsMask(readOp, rewriter);
  auto gatherOp = xegpu::LoadGatherOp::create(
      rewriter, loc, vectorType, flatMemref, localOffsets, mask,
      /*l1_hint=*/xegpu::CachePolicyAttr{},
      /*l2_hint=*/xegpu::CachePolicyAttr{},
      /*l3_hint=*/xegpu::CachePolicyAttr{},
      /*layout=*/nullptr, /*contiguity=*/nullptr);
 
  // The masked-off lanes of an xegpu.load are unspecified, so the transfer's
  // padding has to be applied explicitly. A poison padding leaves them "don't
  // care", so the select can be skipped.
  Value result = gatherOp.getResult();
  if (readOp.hasOutOfBoundsDim() &&
      !readOp.getPadding().getDefiningOp<ub::PoisonOp>()) {
    Value padding = vector::BroadcastOp::create(rewriter, loc, vectorType,
                                                readOp.getPadding());
    result = arith::SelectOp::create(rewriter, loc, mask, result, padding);
  }
 
  rewriter.replaceOp(readOp, result);
  return success();
}
 
// An nd block load fetches a 2D block, which a 1D transfer does not ask for -
// but a leading unit dimension gives it one:
//
//   %v = vector.transfer_read %src[%i0, %i1] : memref<32x64xf32>, vector<8xf32>
// becomes
//   %b = vector.transfer_read %src[%i0, %i1]
//     : memref<32x64xf32>, vector<1x8xf32>
//   %v = vector.shape_cast %b : vector<1x8xf32> to vector<8xf32>
//
// Only a transfer running along the innermost source dimension qualifies. That
// dimension is contiguous, which transferPreconditions has already checked, so
// the <1xN> block it asks for is one the hardware can fetch. A transfer along
// any other dimension would need the unit dimension on the other side, and a
// <Nx1> block is not a shape the layouts downstream can distribute - the
// scattered path is the one that fits it.
static LogicalResult lowerToUnitDimRead(vector::TransferReadOp readOp,
                                        PatternRewriter &rewriter) {
  VectorType vecTy = readOp.getVectorType();
  AffineMap map = readOp.getPermutationMap();
  int64_t numInputs = map.getNumInputs();
  int64_t pos = cast<AffineDimExpr>(map.getResult(0)).getPosition();
  if (numInputs < 2 || pos != numInputs - 1)
    return rewriter.notifyMatchFailure(
        readOp, "1D transfer does not run along the innermost dimension");
 
  Value block = vector::TransferReadOp::create(
      rewriter, readOp.getLoc(),
      VectorType::get({1, vecTy.getDimSize(0)}, vecTy.getElementType()),
      readOp.getBase(), readOp.getIndices(),
      AffineMapAttr::get(
          AffineMap::getMinorIdentityMap(numInputs, 2, map.getContext())),
      readOp.getPadding(), /*mask=*/Value(),
      rewriter.getBoolArrayAttr({true, readOp.isDimInBounds(0)}));
 
  rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(readOp, vecTy, block);
  return success();
}
 
// Lowers a transfer whose permutation map broadcasts some of the vector
// dimensions into a transfer of just the data in memory - the map with its
// broadcast results dropped - plus the vector ops that expand it back:
//
//   %r = vector.transfer_read %src[%i0, %i1, %i2], %pad
//     {permutation_map = affine_map<(d0, d1, d2) -> (d2, 0)>,
//      in_bounds = [false, false]} : memref<256x32x512xf16>, vector<64x64xf16>
// becomes
//   %v = vector.transfer_read %src[%i0, %i1, %i2], %pad
//     {permutation_map = affine_map<(d0, d1, d2) -> (d2)>,
//      in_bounds = [false]} : memref<256x32x512xf16>, vector<64xf16>
//   %c = vector.shape_cast %v : vector<64xf16> to vector<1x64xf16>
//   %t = vector.transpose %c, [1, 0] : vector<1x64xf16> to vector<64x1xf16>
//   %r = vector.broadcast %t : vector<64x1xf16> to vector<64x64xf16>
//
// A broadcast result names no source dimension, so only the dimensions the map
// does name hold data; dropping the rest leaves the transfer that reads it, at
// the extents and bounds checks those dimensions had. The dropped results come
// back as unit dimensions, which the transpose moves to the positions those
// results held, and vector.broadcast stretches to their final extents.
static LogicalResult lowerBroadcastRead(vector::TransferReadOp readOp,
                                        PatternRewriter &rewriter) {
  Location loc = readOp.getLoc();
  VectorType vecTy = readOp.getVectorType();
  ArrayRef<int64_t> vecShape = vecTy.getShape();
  AffineMap map = readOp.getPermutationMap();
  int64_t vecRank = vecTy.getRank();
 
  SmallVector<int64_t> dataPos, bcastPos;
  for (int64_t i = 0; i < vecRank; ++i) {
    if (isa<AffineDimExpr>(map.getResult(i)))
      dataPos.push_back(i);
    else
      bcastPos.push_back(i);
  }
 
  SmallVector<AffineExpr> readExprs;
  SmallVector<int64_t> readShape;
  SmallVector<bool> readInBounds;
  for (int64_t i : dataPos) {
    readExprs.push_back(map.getResult(i));
    readShape.push_back(vecShape[i]);
    readInBounds.push_back(readOp.isDimInBounds(i));
  }
  // An all-broadcast map leaves no data dimension to read. Read one element of
  // the innermost source dimension instead - a 0D transfer is not supported,
  // and one element is all the data such a transfer holds.
  if (readExprs.empty()) {
    readExprs.push_back(rewriter.getAffineDimExpr(map.getNumInputs() - 1));
    readShape.push_back(1);
    readInBounds.push_back(true);
  }
 
  auto readTy = VectorType::get(readShape, vecTy.getElementType());
  Value read = vector::TransferReadOp::create(
      rewriter, loc, readTy, readOp.getBase(), readOp.getIndices(),
      AffineMapAttr::get(
          AffineMap::get(map.getNumInputs(), 0, readExprs, map.getContext())),
      readOp.getPadding(), /*mask=*/Value(),
      rewriter.getBoolArrayAttr(readInBounds));
 
  // Grow the transfer back to the requested rank, one unit dimension per
  // dropped result. The shape_cast adds them innermost, where they leave the
  // dimension the data came in on innermost too, and the transpose moves them
  // to the positions their results held. Only unit dimensions move, so the
  // transpose permutes dimensions without shuffling elements between them.
  SmallVector<int64_t> order(bcastPos);
  llvm::append_range(order, dataPos);
  SmallVector<int64_t> castShape(bcastPos.size(), 1);
  for (int64_t i : dataPos)
    castShape.push_back(vecShape[i]);
  auto castTy = VectorType::get(castShape, vecTy.getElementType());
  Value cast = castTy == readTy
                   ? read
                   : vector::ShapeCastOp::create(rewriter, loc, castTy, read);
 
  if (!llvm::is_sorted(order)) {
    SmallVector<int64_t> perm(vecRank);
    for (int64_t i = 0; i < vecRank; ++i)
      perm[order[i]] = i;
    cast = vector::TransposeOp::create(rewriter, loc, cast, perm);
  }
 
  if (cast.getType() == vecTy)
    rewriter.replaceOp(readOp, cast);
  else
    rewriter.replaceOpWithNewOp<vector::BroadcastOp>(readOp, vecTy, cast);
  return success();
}
 
static LogicalResult lowerToScatteredStoreOp(vector::TransferWriteOp writeOp,
                                             PatternRewriter &rewriter) {
 
  Location loc = writeOp.getLoc();
  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);
 
  // Out-of-bounds elements are simply not stored, so the mask is all this
  // needs - no counterpart to the read's padding.
  Value mask = computeInBoundsMask(writeOp, rewriter);
  xegpu::StoreScatterOp::create(rewriter, loc, writeOp.getVector(), flatMemref,
                                localOffsets, mask,
                                /*l1_hint=*/xegpu::CachePolicyAttr{},
                                /*l2_hint=*/xegpu::CachePolicyAttr{},
                                /*l3_hint=*/xegpu::CachePolicyAttr{},
                                /*layout=*/nullptr, /*contiguity=*/nullptr);
  rewriter.eraseOp(writeOp);
  return success();
}
 
// Rewrites a transfer_read whose permutation map broadcasts some vector
// dimensions into a transfer of just the data in memory plus the vector ops
// that expand it back (see lowerBroadcastRead), so that the conversion
// patterns only see maps without broadcast results.
struct LegalizeBroadcastTransferRead
    : public OpRewritePattern<vector::TransferReadOp> {
  using Base::Base;
 
  LogicalResult matchAndRewrite(vector::TransferReadOp readOp,
                                PatternRewriter &rewriter) const override {
    if (readOp.getPermutationMap().getBroadcastDims().empty())
      return rewriter.notifyMatchFailure(readOp, "No broadcast dimension");
    if (failed(transferPreconditions(rewriter, readOp)))
      return failure();
    return lowerBroadcastRead(readOp, rewriter);
  }
};
 
// Returns true if the target's 2D block load can fetch the `tileShape` tile of
// a non-SLM transfer_read: a scalar int/float element type, out-of-bounds
// padding that matches load_nd's implicit zero padding, and a tile shape the
// block load can access.
static bool is2DBlockLoadable(vector::TransferReadOp readOp,
                              ArrayRef<int64_t> tileShape,
                              bool hasTranspose = false) {
  auto readMemTy = cast<MemRefType>(readOp.getShapedType());
  if (!readMemTy.getElementType().isIntOrFloat())
    return false;
  if (readOp.hasOutOfBoundsDim() && !isZeroOrPoisonPadding(readOp.getPadding()))
    return false;
  const xegpu::uArch::uArch *uArch =
      xegpu::uArch::getUArch(xegpu::getChipStr(readOp).value_or(""));
  return isSupportedBlockShape(
      uArch, xegpu::uArch::InstructionKind::Subgroup2DBlockLoad, tileShape,
      readOp.getVectorType().getElementType(), hasTranspose);
}
 
// Gives a 1D transfer_read the leading unit dimension an nd block load needs
// (see lowerToUnitDimRead). Only worth it where the target's 2D block load can
// access the resulting <1xN> tile; otherwise TransferReadLowering takes the 1D
// vector as is on the scattered path.
struct PromoteTransferReadToUnitDim
    : public OpRewritePattern<vector::TransferReadOp> {
  using Base::Base;
 
  LogicalResult matchAndRewrite(vector::TransferReadOp readOp,
                                PatternRewriter &rewriter) const override {
    VectorType vecTy = readOp.getVectorType();
    if (vecTy.getRank() != 1)
      return rewriter.notifyMatchFailure(readOp, "Not a 1D transfer");
    // A broadcast map is left to LegalizeBroadcastTransferRead.
    if (!readOp.getPermutationMap().getBroadcastDims().empty())
      return rewriter.notifyMatchFailure(readOp, "Broadcast permutation map");
    if (failed(transferPreconditions(rewriter, readOp)))
      return failure();
 
    // SLM reads lower to load_matrix, which takes the 1D vector as is.
    if (xegpu::XeGPUDialect::isSharedMemory(
            cast<MemRefType>(readOp.getShapedType())))
      return rewriter.notifyMatchFailure(readOp, "SLM transfer");
    // A single element only extracted to a scalar becomes a scalar load.
    if (vecTy.getNumElements() == 1 && isUsedAsScalar(readOp.getResult()))
      return rewriter.notifyMatchFailure(readOp, "Used as a scalar");
    if (!is2DBlockLoadable(readOp, {1, vecTy.getDimSize(0)}))
      return rewriter.notifyMatchFailure(
          readOp, "<1xN> tile is not accessible by a 2D block load");
 
    return lowerToUnitDimRead(readOp, rewriter);
  }
};
 
struct TransferReadLowering : public OpRewritePattern<vector::TransferReadOp> {
  using Base::Base;
 
  LogicalResult matchAndRewrite(vector::TransferReadOp readOp,
                                PatternRewriter &rewriter) const override {
    Location loc = readOp.getLoc();
 
    if (failed(transferPreconditions(rewriter, readOp)))
      return failure();
 
    // Broadcast maps are rewritten away by the transfer_read legalization that
    // runs ahead of this pattern.
    if (!readOp.getPermutationMap().getBroadcastDims().empty())
      return rewriter.notifyMatchFailure(
          readOp, "Broadcast permutation map must be legalized first");
 
    auto readMemTy = cast<MemRefType>(readOp.getShapedType());
    VectorType loadedVecTy = readOp.getVectorType();
    bool isOutOfBounds = readOp.hasOutOfBoundsDim();
    // Check if the memref has address space 3 (shared local memory)
    bool isSharedMemory = xegpu::XeGPUDialect::isSharedMemory(readMemTy);
    // Handle the SLM case.
    if (isSharedMemory) {
      // load_matrix supports 1D and 2D loads from SLM.
      if (loadedVecTy.getRank() != 1 && loadedVecTy.getRank() != 2)
        return rewriter.notifyMatchFailure(
            readOp, "Only 1D and 2D vector loads are supported for SLM");
      AffineMap readMap = readOp.getPermutationMap();
      if (!readMap.isMinorIdentity())
        return rewriter.notifyMatchFailure(
            readOp,
            "Non identity transposition is not supported for SLM loads.");
      // Out of bounds case is not supported for SLM loads.
      if (isOutOfBounds)
        return rewriter.notifyMatchFailure(
            readOp, "Out-of-bounds access is not supported for SLM loads");
 
      // Create mem_desc for SLM
      auto memDescType =
          xegpu::MemDescType::get(rewriter.getContext(), readMemTy.getShape(),
                                  readMemTy.getElementType(),
                                  /*mem_layout=*/nullptr);
      auto createMemDescOp = xegpu::CreateMemDescOp::create(
          rewriter, loc, memDescType, readOp.getBase());
      // Convert indices to OpFoldResult for LoadMatrixOp
      SmallVector<OpFoldResult> indices =
          getAsOpFoldResult(readOp.getIndices());
      auto loadMatrixOp = xegpu::LoadMatrixOp::create(
          rewriter, loc, loadedVecTy, createMemDescOp.getResult(), indices,
          /*layout=*/nullptr);
 
      rewriter.replaceOp(readOp, loadMatrixOp.getResult());
      return success();
    }
 
    // A transfer of a single element that is only ever extracted to a scalar
    // becomes a scalar load, at any rank: the vector is a wrapper its consumers
    // undo, so neither an nd descriptor nor a lane-vector gather buys anything.
    // A unit-size vector genuinely used as a vector keeps the paths below.
    if (loadedVecTy.getNumElements() == 1 && isUsedAsScalar(readOp.getResult()))
      return lowerToScalarLoadOp(readOp, rewriter);
 
    // An nd block load can realize a minor-identity map directly, or an
    // innermost-two-dims transpose via a trailing vector.transpose. Any other
    // permutation (e.g. a mid-vector transpose of a high-dim load) is left to
    // the scattered path, which permutes strides explicitly.
    AffineMap readMap = readOp.getPermutationMap();
    bool isTransposeLoad = isInnermostTwoDimsTransposed(readMap);
 
    // A block load transfers the tile as it is laid out in memory, so a
    // transposing read describes a tile whose innermost two dims are swapped.
    Type elementType = loadedVecTy.getElementType();
    SmallVector<int64_t> descShape(loadedVecTy.getShape());
    if (isTransposeLoad) {
      size_t rank = descShape.size();
      assert(rank >= 2 && "Transpose requires at least 2 dimensions");
      std::swap(descShape[rank - 1], descShape[rank - 2]);
    }
    // Swapping a unit dimension moves no data - e.g. a <1x64> tile read as
    // <64x1> - so the tile is loaded as is and the trailing vector.transpose is
    // only a reshape, with no need for the transposing block load variant.
    bool needsTransposedBlockLoad = isTransposeLoad &&
                                    descShape[descShape.size() - 1] != 1 &&
                                    descShape[descShape.size() - 2] != 1;
 
    // Prefer an nd block load. It requires a vector of rank >= 2 backed by a
    // scalar-element memref, a map the block load can realize, and a tile shape
    // the target's 2D block load can access. 1D vectors use the scattered
    // xegpu.load path instead, which has a richer interface (e.g. layout
    // capabilities). Out-of-bounds reads are allowed as long as the padding
    // matches load_nd's implicit zero padding.
    bool canLowerToLoadNd =
        loadedVecTy.getRank() > 1 &&
        (readMap.isMinorIdentity() || isTransposeLoad) &&
        is2DBlockLoadable(readOp, descShape,
                          /*hasTranspose=*/needsTransposedBlockLoad);
 
    if (canLowerToLoadNd) {
      // The load produces the memory-ordered tile; the transpose below restores
      // the shape the transfer_read asked for.
      if (isTransposeLoad)
        loadedVecTy = VectorType::get(descShape, elementType);
      auto descType = xegpu::TensorDescType::get(
          descShape, elementType, /*array_length=*/1,
          /*boundary_check=*/isOutOfBounds, xegpu::MemorySpace::Global);
      auto [src, indices] = convertMemrefAndOffsetsToTargetRank(
          rewriter, loc, readOp.getBase(),
          getAsOpFoldResult(readOp.getIndices()), loadedVecTy.getRank());
      // By default, no specific caching policy is assigned.
      xegpu::CachePolicyAttr hint = nullptr;
      xegpu::CreateNdDescOp ndDesc = xegpu::CreateNdDescOp::create(
          rewriter, loc, descType, dyn_cast<TypedValue<MemRefType>>(src));
 
      Operation *loadedOp =
          xegpu::LoadNdOp::create(rewriter, loc, loadedVecTy, ndDesc, indices,
                                  /*packed=*/nullptr, /*transpose=*/nullptr,
                                  /*l1_hint=*/hint,
                                  /*l2_hint=*/hint, /*l3_hint=*/hint,
                                  /*layout=*/nullptr);
      if (isTransposeLoad) {
        // Undo the innermost-two-dims swap with a trailing vector.transpose:
        // keep the leading dimensions in place and interchange only the last
        // two.
        int64_t rank = loadedVecTy.getRank();
        SmallVector<int64_t> perm(llvm::to_vector(llvm::seq<int64_t>(0, rank)));
        std::swap(perm[rank - 1], perm[rank - 2]);
        loadedOp = vector::TransposeOp::create(rewriter, loc,
                                               loadedOp->getResult(0), perm);
      }
      rewriter.replaceOp(readOp, loadedOp);
      return success();
    }
 
    // Fall back to a scattered load. It supports arbitrary permutations and any
    // rank, and masks off the out-of-bounds elements.
    return lowerToScatteredLoadOp(readOp, rewriter);
  }
};
 
struct TransferWriteLowering
    : public OpRewritePattern<vector::TransferWriteOp> {
  using Base::Base;
 
  LogicalResult matchAndRewrite(vector::TransferWriteOp writeOp,
                                PatternRewriter &rewriter) const override {
    Location loc = writeOp.getLoc();
 
    if (failed(transferPreconditions(rewriter, writeOp)))
      return failure();
    // Perform common data transfer checks.
    VectorType vecTy = writeOp.getVectorType();
    auto writeMemTy = cast<MemRefType>(writeOp.getShapedType());
    // Check if the memref has address space 3 (shared local memory)
    bool isSharedMemory = xegpu::XeGPUDialect::isSharedMemory(writeMemTy);
 
    // For shared local memory (address space 3), use create_mem_desc +
    // store_matrix
    if (isSharedMemory) {
      // store_matrix supports 1D and 2D stores to SLM.
      if (vecTy.getRank() != 1 && vecTy.getRank() != 2)
        return rewriter.notifyMatchFailure(
            writeOp, "Only 1D and 2D vector stores are supported for SLM");
      // store_matrix writes the payload along the innermost dimensions of the
      // mem_desc, so it cannot express a transposed or strided store, e.g. a
      // 1D vector written along a non-innermost dimension.
      if (!writeOp.getPermutationMap().isMinorIdentity())
        return rewriter.notifyMatchFailure(
            writeOp,
            "Non identity transposition is not supported for SLM stores.");
      // Out of bounds case is not supported for SLM stores.
      if (writeOp.hasOutOfBoundsDim())
        return rewriter.notifyMatchFailure(
            writeOp, "Out-of-bounds access is not supported for SLM stores");
      // Create mem_desc for SLM
      auto memDescType =
          xegpu::MemDescType::get(rewriter.getContext(), writeMemTy.getShape(),
                                  writeMemTy.getElementType(),
                                  /*mem_layout=*/nullptr);
 
      auto createMemDescOp = xegpu::CreateMemDescOp::create(
          rewriter, loc, memDescType, writeOp.getBase());
 
      // Convert indices to OpFoldResult for StoreMatrixOp
      SmallVector<OpFoldResult> indices =
          getAsOpFoldResult(writeOp.getIndices());
 
      xegpu::StoreMatrixOp::create(rewriter, loc, writeOp.getVector(),
                                   createMemDescOp.getResult(), indices,
                                   /*layout=*/nullptr);
 
      rewriter.eraseOp(writeOp);
      return success();
    }
 
    const xegpu::uArch::uArch *uArch =
        xegpu::uArch::getUArch(xegpu::getChipStr(writeOp).value_or(""));
 
    // Prefer an nd block store. It requires a vector of rank >= 2 backed by a
    // scalar-element memref, a minor-identity map (block stores have no
    // transpose support), and a tile shape the target's 2D block store can
    // access. 1D vectors use the scattered xegpu.store path instead, which has
    // a richer interface. Out-of-bounds writes are handled by the descriptor's
    // boundary check.
    AffineMap map = writeOp.getPermutationMap();
    bool canLowerToStoreNd =
        vecTy.getRank() > 1 && map.isMinorIdentity() &&
        writeMemTy.getElementType().isIntOrFloat() &&
        isSupportedBlockShape(
            uArch, xegpu::uArch::InstructionKind::Subgroup2DBlockStore,
            vecTy.getShape(), vecTy.getElementType());
 
    if (canLowerToStoreNd) {
      auto [src, indices] = convertMemrefAndOffsetsToTargetRank(
          rewriter, loc, writeOp.getBase(),
          getAsOpFoldResult(writeOp.getIndices()), vecTy.getRank());
 
      auto descType = xegpu::TensorDescType::get(
          vecTy.getShape(), vecTy.getElementType(),
          /*array_length=*/1, /*boundary_check=*/writeOp.hasOutOfBoundsDim(),
          xegpu::MemorySpace::Global);
      // By default, no specific caching policy is assigned.
      xegpu::CachePolicyAttr hint = nullptr;
      xegpu::CreateNdDescOp ndDesc = xegpu::CreateNdDescOp::create(
          rewriter, loc, descType, dyn_cast<TypedValue<MemRefType>>(src));
 
      auto storeOp = xegpu::StoreNdOp::create(
          rewriter, loc, writeOp.getVector(), ndDesc, indices,
          /*l1_hint=*/hint,
          /*l2_hint=*/hint, /*l3_hint=*/hint,
          /*layout=*/nullptr);
      rewriter.replaceOp(writeOp, storeOp);
      return success();
    }
 
    // Fall back to a scattered store. It supports arbitrary permutations and
    // any rank, and masks off the out-of-bounds elements.
    return lowerToScatteredStoreOp(writeOp, rewriter);
  }
};
 
struct GatherLowering : public OpRewritePattern<vector::GatherOp> {
  using Base::Base;
 
  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");
 
    if (gatherOp.getVectorType().isScalable())
      return rewriter.notifyMatchFailure(gatherOp,
                                         "Scalable vectors are not supported");
 
    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(),
        /*l1_hint=*/xegpu::CachePolicyAttr{},
        /*l2_hint=*/xegpu::CachePolicyAttr{},
        /*l3_hint=*/xegpu::CachePolicyAttr{},
        /*layout=*/nullptr, /*contiguity=*/nullptr);
 
    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 Base::Base;
 
  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");
 
    if (scatterOp.getVectorType().isScalable())
      return rewriter.notifyMatchFailure(scatterOp,
                                         "Scalable vectors are not supported");
 
    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(),
                                  /*l1_hint=*/xegpu::CachePolicyAttr{},
                                  /*l2_hint=*/xegpu::CachePolicyAttr{},
                                  /*l3_hint=*/xegpu::CachePolicyAttr{},
                                  /*layout=*/nullptr,
                                  /*contiguity=*/nullptr);
    rewriter.eraseOp(scatterOp);
    return success();
  }
};
 
struct LoadLowering : public OpRewritePattern<vector::LoadOp> {
  using Base::Base;
 
  LogicalResult matchAndRewrite(vector::LoadOp loadOp,
                                PatternRewriter &rewriter) const override {
    Location loc = loadOp.getLoc();
 
    VectorType vecTy = loadOp.getResult().getType();
    MemRefType memTy = loadOp.getBase().getType();
    // The plain vector.load lowering only supports 1D/2D block loads.
    if (vecTy.getRank() != 1 && vecTy.getRank() != 2)
      return rewriter.notifyMatchFailure(loadOp, "Expects 1D or 2D vector");
    if (!memTy.getElementType().isIntOrFloat())
      return rewriter.notifyMatchFailure(
          loadOp, "Unsupported memref element type: expected integer or float");
 
    // Boundary check is available only for block instructions.
    bool boundaryCheck = vecTy.getRank() > 1;
    // By default, no specific caching policy is assigned.
    xegpu::CachePolicyAttr hint = nullptr;
 
    auto [src, indices] = convertMemrefAndOffsetsToTargetRank(
        rewriter, loc, loadOp.getBase(), getAsOpFoldResult(loadOp.getIndices()),
        vecTy.getRank());
 
    auto descType = xegpu::TensorDescType::get(
        vecTy.getShape(), vecTy.getElementType(), /*array_length=*/1,
        boundaryCheck, xegpu::MemorySpace::Global);
 
    xegpu::CreateNdDescOp ndDesc = xegpu::CreateNdDescOp::create(
        rewriter, loc, descType, dyn_cast<TypedValue<MemRefType>>(src));
    auto loadNdOp =
        xegpu::LoadNdOp::create(rewriter, loc, vecTy, ndDesc, indices,
                                /*packed=*/nullptr, /*transpose=*/nullptr,
                                /*l1_hint=*/hint,
                                /*l2_hint=*/hint, /*l3_hint=*/hint,
                                /*layout=*/nullptr);
    rewriter.replaceOp(loadOp, loadNdOp);
 
    return success();
  }
};
 
struct StoreLowering : public OpRewritePattern<vector::StoreOp> {
  using Base::Base;
 
  LogicalResult matchAndRewrite(vector::StoreOp storeOp,
                                PatternRewriter &rewriter) const override {
    Location loc = storeOp.getLoc();
 
    TypedValue<VectorType> vector = storeOp.getValueToStore();
    VectorType vecTy = vector.getType();
    MemRefType memTy = storeOp.getBase().getType();
    // The plain vector.store lowering only supports 1D/2D block stores.
    if (vecTy.getRank() != 1 && vecTy.getRank() != 2)
      return rewriter.notifyMatchFailure(storeOp, "Expects 1D or 2D vector");
    if (!memTy.getElementType().isIntOrFloat())
      return rewriter.notifyMatchFailure(
          storeOp,
          "Unsupported memref element type: expected integer or float");
 
    // Boundary check is available only for block instructions.
    bool boundaryCheck = vecTy.getRank() > 1;
 
    auto [src, indices] = convertMemrefAndOffsetsToTargetRank(
        rewriter, loc, storeOp.getBase(),
        getAsOpFoldResult(storeOp.getIndices()), vecTy.getRank());
 
    auto descType = xegpu::TensorDescType::get(
        vecTy.getShape(), vecTy.getElementType(),
        /*array_length=*/1, boundaryCheck, xegpu::MemorySpace::Global);
 
    // By default, no specific caching policy is assigned.
    xegpu::CachePolicyAttr hint = nullptr;
    xegpu::CreateNdDescOp ndDesc = xegpu::CreateNdDescOp::create(
        rewriter, loc, descType, dyn_cast<TypedValue<MemRefType>>(src));
 
    auto storeNdOp =
        xegpu::StoreNdOp::create(rewriter, loc, vector, ndDesc, indices,
                                 /*l1_hint=*/hint,
                                 /*l2_hint=*/hint, /*l3_hint=*/hint,
                                 /*layout=*/nullptr);
 
    rewriter.replaceOp(storeOp, storeNdOp);
 
    return success();
  }
};
 
// If `indexingMaps` describe a (batched) row-major matmul
//   lhs[b..., m, k], rhs[b..., k, n], acc[b..., m, n]
// return the number of leading batch dims (0 for a plain 2D matmul);
// otherwise return std::nullopt.
static std::optional<int64_t>
getRowMajorMatmulBatchRank(ArrayAttr indexingMaps) {
  if (indexingMaps.size() != 3)
    return std::nullopt;
 
  AffineMap mapA = cast<AffineMapAttr>(indexingMaps[0]).getValue();
  AffineMap mapB = cast<AffineMapAttr>(indexingMaps[1]).getValue();
  AffineMap mapC = cast<AffineMapAttr>(indexingMaps[2]).getValue();
 
  // The result map exposes the batch dims followed by the core (m, n) dims.
  if (mapC.getNumResults() < 2)
    return std::nullopt;
  int64_t batchRank = mapC.getNumResults() - 2;
 
  // A single `k` reduction gives batchRank + 3 iteration dims; each operand
  // map exposes batchRank + 2 dims (batch dims + 2 core dims).
  unsigned numDims = static_cast<unsigned>(batchRank) + 3;
  unsigned numOperandResults = static_cast<unsigned>(batchRank) + 2;
  if (mapA.getNumInputs() != numDims || mapB.getNumInputs() != numDims ||
      mapC.getNumInputs() != numDims)
    return std::nullopt;
  if (mapA.getNumResults() != numOperandResults ||
      mapB.getNumResults() != numOperandResults)
    return std::nullopt;
 
  // Reconstruct the canonical maps from the batch/m/n dims of the result and
  // the k dim of lhs, then compare against the actual maps.
  MLIRContext *context = indexingMaps.getContext();
  ArrayRef<AffineExpr> batchDims = mapC.getResults().take_front(batchRank);
  AffineExpr m = mapC.getResult(batchRank);
  AffineExpr n = mapC.getResult(batchRank + 1);
  AffineExpr k = mapA.getResult(batchRank + 1);
 
  SmallVector<AffineExpr> aDims = llvm::to_vector(batchDims);
  aDims.push_back(m);
  aDims.push_back(k);
  SmallVector<AffineExpr> bDims = llvm::to_vector(batchDims);
  bDims.push_back(k);
  bDims.push_back(n);
  SmallVector<AffineExpr> cDims = llvm::to_vector(batchDims);
  cDims.push_back(m);
  cDims.push_back(n);
 
  auto expected = ArrayAttr::get(
      context,
      {AffineMapAttr::get(AffineMap::get(numDims, 0, aDims, context)),
       AffineMapAttr::get(AffineMap::get(numDims, 0, bDims, context)),
       AffineMapAttr::get(AffineMap::get(numDims, 0, cDims, context))});
  if (indexingMaps != expected)
    return std::nullopt;
  return batchRank;
}
 
struct ContractionLowering : public vector::AddContractionOpRewritePattern {
  using Base::Base;
 
  LogicalResult
  matchAndRewriteAddContractOp(vector::ContractionOp contractOp,
                               PatternRewriter &rewriter) const override {
    Location loc = contractOp.getLoc();
 
    TypedValue<VectorType> lhs = contractOp.getLhs();
    TypedValue<VectorType> rhs = contractOp.getRhs();
    TypedValue<Type> acc = contractOp.getAcc();
    VectorType accType = dyn_cast<VectorType>(acc.getType());
    if (!accType)
      return rewriter.notifyMatchFailure(contractOp, "Expects vector acc");
 
    std::optional<int64_t> batchRank =
        getRowMajorMatmulBatchRank(contractOp.getIndexingMapsAttr());
    if (!batchRank)
      return rewriter.notifyMatchFailure(
          contractOp,
          "Expects a (batched) row-major matmul: leading dims must "
          "be batch dims shared by lhs, rhs, and acc; innermost two "
          "dims must be (M, K), (K, N), and (M, N)");
 
    // xegpu.dpas operands are limited to 2 batch + 2 core dims.
    if (*batchRank > 2)
      return rewriter.notifyMatchFailure(contractOp,
                                         "Expects operands of rank 4 or less");
 
    auto dpasOp = xegpu::DpasOp::create(
        rewriter, loc, contractOp.getResultType(), lhs, rhs, acc,
        /*layout_a=*/nullptr, /*layout_b=*/nullptr, /*layout_cd=*/nullptr);
    rewriter.replaceOp(contractOp, dpasOp);
 
    return success();
  }
};
 
// Returns the `vector.shape_cast` that flattened a value of type `ndType` into
// `flat`, if that is how `flat` was produced.
static vector::ShapeCastOp getFlattenCast(Value flat, VectorType ndType) {
  auto shapeCast = flat.getDefiningOp<vector::ShapeCastOp>();
  if (shapeCast && shapeCast.getSourceVectorType() == ndType)
    return shapeCast;
  return nullptr;
}
 
static DenseElementsAttr getDenseConstant(Value flat) {
  DenseElementsAttr elements;
  if (matchPattern(flat, m_Constant(&elements)))
    return elements;
  return nullptr;
}
 
static vector::BroadcastOp getSplatBroadcast(Value flat) {
  auto broadcast = flat.getDefiningOp<vector::BroadcastOp>();
  if (broadcast && !isa<VectorType>(broadcast.getSourceType()))
    return broadcast;
  return nullptr;
}
 
// Returns true if the cast `unflatten` creates will fold away, i.e. if `flat`
// is a flattening cast, a constant or a splat.
static bool canUnflatten(Value flat, VectorType ndType) {
  return getFlattenCast(flat, ndType) || getDenseConstant(flat) ||
         getSplatBroadcast(flat);
}
 
// Reshape `flat` to `ndType`; `canUnflatten` must hold so the cast folds away.
static Value unflatten(PatternRewriter &rewriter, Value flat,
                       VectorType ndType) {
  assert(canUnflatten(flat, ndType) && "expected the cast to fold away");
  return vector::ShapeCastOp::create(rewriter, flat.getLoc(), ndType, flat);
}
 
// Restore the N-D form of a flattened `vector.gather` / `vector.scatter`.
//
// XeGPU layouts are expressed in terms of the N-D shape of the accessed data,
// so a flattened gather/scatter forces layout propagation to reason through the
// surrounding `vector.shape_cast` ops. That adds complexity and tends to yield
// layouts that lower to unoptimized code.
//
// Before:
//   %cst = arith.constant dense<0.0> : vector<8192xbf16>
//   %fi = vector.shape_cast %idx : vector<128x64xindex> to vector<8192xindex>
//   %fm = vector.shape_cast %mask : vector<128x64xi1> to vector<8192xi1>
//   %fr = vector.gather %src[%c0] [%fi], %fm, %cst : memref<?xbf16>,
//       vector<8192xindex>, vector<8192xi1>, vector<8192xbf16>
//       into vector<8192xbf16>
//   %res = vector.shape_cast %fr : vector<8192xbf16> to vector<128x64xbf16>
//
// After:
//   %cst = arith.constant dense<0.0> : vector<128x64xbf16>
//   %res = vector.gather %src[%c0] [%idx], %mask, %cst : memref<?xbf16>,
//       vector<128x64xindex>, vector<128x64xi1>, vector<128x64xbf16>
//       into vector<128x64xbf16>
template <typename OpTy>
struct UnflattenGatherScatter : public OpRewritePattern<OpTy> {
  using OpRewritePattern<OpTy>::OpRewritePattern;
 
  LogicalResult matchAndRewrite(OpTy op,
                                PatternRewriter &rewriter) const override {
    constexpr bool isGather = std::is_same_v<OpTy, vector::GatherOp>;
 
    if (!isa<MemRefType>(op.getBase().getType()))
      return rewriter.notifyMatchFailure(op, "expects a memref source");
 
    if (op.getIndexVectorType().getRank() != 1)
      return rewriter.notifyMatchFailure(op, "index vector is not 1-D");
 
    // The N-D shape comes from the index operand's producer: this only undoes
    // a flattening that already happened, it never invents a shape.
    auto indexCast =
        op.getIndices().template getDefiningOp<vector::ShapeCastOp>();
    if (!indexCast || indexCast.getSourceVectorType().getRank() < 2)
      return rewriter.notifyMatchFailure(
          op, "index vector is not a shape_cast of an N-D vector");
    VectorType ndIndexType = indexCast.getSourceVectorType();
    VectorType ndMaskType =
        ndIndexType.cloneWith(std::nullopt, rewriter.getI1Type());
    VectorType ndType = ndIndexType.cloneWith(
        std::nullopt, op.getVectorType().getElementType());
 
    // Check everything before creating any IR: a partially applied rewrite
    // would leave dead ops behind.
    if (!canUnflatten(op.getMask(), ndMaskType))
      return rewriter.notifyMatchFailure(op, "cannot un-flatten the mask");
 
    if constexpr (isGather) {
      if (!canUnflatten(op.getPassThru(), ndType))
        return rewriter.notifyMatchFailure(op,
                                           "cannot un-flatten the pass-thru");
 
      Value mask = unflatten(rewriter, op.getMask(), ndMaskType);
      Value passThru = unflatten(rewriter, op.getPassThru(), ndType);
      auto ndGather = vector::GatherOp::create(
          rewriter, op.getLoc(), ndType, op.getBase(), op.getOffsets(),
          indexCast.getSource(), mask, passThru, op.getAlignmentAttr());
      ndGather->setDiscardableAttrs(op->getDiscardableAttrDictionary());
      rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(op, op.getVectorType(),
                                                       ndGather);
    } else {
      if (!canUnflatten(op.getValueToStore(), ndType))
        return rewriter.notifyMatchFailure(
            op, "cannot un-flatten the stored value");
 
      Value mask = unflatten(rewriter, op.getMask(), ndMaskType);
      Value valueToStore = unflatten(rewriter, op.getValueToStore(), ndType);
      // Only operand types change, and this keeps the optional tensor result
      // untouched.
      rewriter.modifyOpInPlace(op, [&] {
        op.getIndicesMutable().assign(indexCast.getSource());
        op.getMaskMutable().assign(mask);
        op.getValueToStoreMutable().assign(valueToStore);
      });
    }
    return success();
  }
};
 
// Rewrite a broadcast that is immediately transposed into a shape_cast that
// inserts the new dims where they finally belong, followed by a broadcast that
// stretches them in place.
//
// This lets layout propagation give %v the slice dim of the row reduction
// that typically produces it (dims = [1] below, not [0]):
//   #row = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>
//   #col = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1],
//                        order = [0, 1]>
//
// Before:
//   // %v : #xegpu.slice<#col, dims = [0]>
//   %b = vector.broadcast %v : vector<16xf32> to vector<64x16xf32>  // #col
//   %t = vector.transpose %b, [1, 0]
//          : vector<64x16xf32> to vector<16x64xf32>                // #row
//
// After:
//   // %v : #xegpu.slice<#row, dims = [1]>
//   %c = vector.shape_cast %v : vector<16xf32> to vector<16x1xf32>  // #row
//   %t = vector.broadcast %c
//          : vector<16x1xf32> to vector<16x64xf32>                 // #row
struct TransposedBroadcastToStretch
    : public OpRewritePattern<vector::TransposeOp> {
  using OpRewritePattern<vector::TransposeOp>::OpRewritePattern;
 
  LogicalResult matchAndRewrite(vector::TransposeOp op,
                                PatternRewriter &rewriter) const override {
    auto bcast = op.getVector().getDefiningOp<vector::BroadcastOp>();
    if (!bcast)
      return rewriter.notifyMatchFailure(op, "not transposing a broadcast");
 
    // Scalar sources are left to FoldTransposeBroadcast.
    auto srcTy = dyn_cast<VectorType>(bcast.getSourceType());
    if (!srcTy || srcTy.getRank() < 1)
      return rewriter.notifyMatchFailure(op,
                                         "broadcast source is not a vector");
 
    VectorType bcastTy = bcast.getResultVectorType();
    if (srcTy.isScalable() || bcastTy.isScalable())
      return rewriter.notifyMatchFailure(op, "scalable vectors");
 
    int64_t srcRank = srcTy.getRank();
    int64_t numNewDims = bcastTy.getRank() - srcRank;
    if (numNewDims < 1)
      return rewriter.notifyMatchFailure(op, "broadcast adds no dims");
 
    // Only a pure prepend; no source dim may be stretched.
    if (srcTy.getShape() != bcastTy.getShape().take_back(srcRank))
      return rewriter.notifyMatchFailure(op,
                                         "broadcast stretches a source dim");
 
    // Source dims must move to the front and new dims to the back, both in
    // order.
    ArrayRef<int64_t> perm = op.getPermutation();
    for (int64_t i = 0; i < srcRank; ++i)
      if (perm[i] != numNewDims + i)
        return rewriter.notifyMatchFailure(
            op, "source dims do not move to the front in order");
    for (int64_t i = 0; i < numNewDims; ++i)
      if (perm[srcRank + i] != i)
        return rewriter.notifyMatchFailure(
            op, "broadcast dims do not stay trailing in order");
 
    // Source shape followed by one unit dim per new dim.
    SmallVector<int64_t> castShape(srcTy.getShape());
    castShape.append(numNewDims, 1);
    auto castTy = VectorType::get(castShape, srcTy.getElementType());
 
    DictionaryAttr attrs = op->getDiscardableAttrDictionary();
    auto cast = vector::ShapeCastOp::create(rewriter, op.getLoc(), castTy,
                                            bcast.getSource());
    auto stretch = rewriter.replaceOpWithNewOp<vector::BroadcastOp>(
        op, op.getResultVectorType(), cast);
    stretch->setDiscardableAttrs(attrs);
    return success();
  }
};
 
// Bring the vector ops into a form the conversion patterns lower directly (see
// populateVectorToXeGPULegalizationPatterns).
static LogicalResult legalizeForVectorToXeGPU(Operation *root) {
  RewritePatternSet patterns(root->getContext());
  populateVectorToXeGPULegalizationPatterns(patterns);
  return applyPatternsGreedily(root, std::move(patterns));
}
 
// Returns `memrefTy` with its memory space replaced by `newMemSpace`.
static MemRefType withMemorySpace(MemRefType memrefTy, Attribute newMemSpace) {
  return MemRefType::get(memrefTy.getShape(), memrefTy.getElementType(),
                         memrefTy.getLayout(), newMemSpace);
}
 
// Rewrite every `memref.alloca` not already in shared local memory (SLM) to
// be in SLM (address space 3), and propagate the new memory space through
// memref-producing aliasing users (e.g. memref.cast, memref.subview,
// memref.expand_shape, ...). Consumers that take a memref operand but
// produce a non-memref result (e.g. vector.transfer_read, vector.load) are
// left untouched: their operand type simply reflects the new memory space.
//
// This makes `xegpu.load_matrix`/`xegpu.store_matrix` lowering work end-to-end
// for IR coming from bufferization, which by default assigns memory space 0/1
// to allocations.
static void promoteAllocasToSLM(Operation *root) {
  MLIRContext *ctx = root->getContext();
  Attribute slmAttr = IntegerAttr::get(IntegerType::get(ctx, 64), 3);
 
  // A user is treated as a memref-producing alias (e.g. memref.cast,
  // memref.subview, memref.expand_shape, ...) if it is side-effect free and
  // produces at least one memref result. This excludes ops like memref.copy
  // that have memory effects.
  auto isMemrefResultOp = [](Operation *op) {
    if (!isMemoryEffectFree(op))
      return false;
    return llvm::any_of(op->getResultTypes(),
                        [](Type t) { return isa<MemRefType>(t); });
  };
 
  // Update `v`'s type to have SLM memory space, then walk forward through
  // memref-producing users and update their result types accordingly.
  std::function<void(Value)> propagate = [&](Value v) {
    auto memrefTy = dyn_cast<MemRefType>(v.getType());
    if (!memrefTy || xegpu::XeGPUDialect::isSharedMemory(memrefTy))
      return;
    v.setType(withMemorySpace(memrefTy, slmAttr));
    for (Operation *user : v.getUsers()) {
      if (!isMemrefResultOp(user))
        continue;
      for (Value result : user->getResults())
        propagate(result);
    }
  };
 
  SmallVector<memref::AllocaOp> allocas;
  root->walk([&](memref::AllocaOp op) {
    auto memrefTy = dyn_cast<MemRefType>(op.getResult().getType());
    if (!memrefTy || xegpu::XeGPUDialect::isSharedMemory(memrefTy))
      return;
    allocas.push_back(op);
  });
 
  for (memref::AllocaOp alloca : allocas) {
    OpBuilder builder(alloca);
    auto memrefTy = cast<MemRefType>(alloca.getResult().getType());
    auto newTy = withMemorySpace(memrefTy, slmAttr);
    auto newOp = memref::AllocaOp::create(
        builder, alloca.getLoc(), newTy, alloca.getDynamicSizes(),
        alloca.getSymbolOperands(), alloca.getAlignmentAttr());
    alloca.getResult().replaceAllUsesWith(newOp.getResult());
    alloca.erase();
    // Propagate the new memory space through memref-producing consumers.
    for (Operation *user : newOp.getResult().getUsers()) {
      if (!isMemrefResultOp(user))
        continue;
      for (Value result : user->getResults())
        propagate(result);
    }
  }
}
 
struct ConvertVectorToXeGPUPass
    : public impl::ConvertVectorToXeGPUBase<ConvertVectorToXeGPUPass> {
  void runOnOperation() override {
    // Promote local allocations to SLM (address space 3) so that
    // load_matrix/store_matrix lowerings have well-typed memref operands.
    promoteAllocasToSLM(getOperation());
 
    // Undo any flattening of gather/scatter operands, so that the conversion
    // below sees the N-D shape the XeGPU layouts are expressed in, and rewrite
    // transfer_reads the conversion patterns cannot lower as they are.
    if (failed(legalizeForVectorToXeGPU(getOperation())))
      return signalPassFailure();
 
    RewritePatternSet patterns(&getContext());
    populateVectorToXeGPUConversionPatterns(patterns);
    populatePrepareVectorToMMAPatterns(patterns);
    if (failed(applyPatternsGreedily(getOperation(), std::move(patterns))))
      return signalPassFailure();
  }
};
} // namespace
 
void mlir::populateVectorToXeGPUConversionPatterns(
    RewritePatternSet &patterns) {
  MLIRContext *ctx = patterns.getContext();
  patterns
      .add<TransferReadLowering, TransferWriteLowering, LoadLowering,
           ScatterLowering, GatherLowering, StoreLowering, ContractionLowering>(
          ctx);
  // Tried ahead of TransferReadLowering, which would otherwise take a 1D
  // transfer on the scattered path.
  patterns.add<PromoteTransferReadToUnitDim>(ctx, /*benefit=*/2);
}
 
void mlir::populateVectorToXeGPULegalizationPatterns(
    RewritePatternSet &patterns) {
  MLIRContext *ctx = patterns.getContext();
  // Un-flatten every gather/scatter that was flattened to 1-D operands, so
  // that the conversion patterns and the XeGPU layouts downstream of them see
  // the N-D shape of the accessed data. The shape_cast canonicalizations clean
  // up the casts the un-flattening leaves behind.
  patterns.add<UnflattenGatherScatter<vector::GatherOp>,
               UnflattenGatherScatter<vector::ScatterOp>,
               TransposedBroadcastToStretch>(ctx);
  vector::ShapeCastOp::getCanonicalizationPatterns(patterns, ctx);
  // Rewrite transfer_reads with a broadcast permutation map into a transfer of
  // the data plus the vector ops that expand it.
  patterns.add<LegalizeBroadcastTransferRead>(ctx);
}