XeGPUUnroll.cpp

Go to the documentation of this file.

//===- XeGPUUnroll.cpp - patterns to do unrolling ---------------*- 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 contains patterns for unrolling XeGPU operations. It follows a
// similar concept and design as vector unroll patterns, serving as a complement
// to them.
//
//===----------------------------------------------------------------------===//
 
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/XeGPU/IR/XeGPU.h"
#include "mlir/Dialect/XeGPU/Transforms/Transforms.h"
#include "mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h"
#include "mlir/Dialect/XeGPU/Utils/XeGPUUtils.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/Support/DebugLog.h"
 
namespace mlir {
namespace xegpu {
#define GEN_PASS_DEF_XEGPUUNROLL
#include "mlir/Dialect/XeGPU/Transforms/Passes.h.inc"
} // namespace xegpu
} // namespace mlir
 
#define DEBUG_TYPE "xegpu-unroll"
 
using namespace mlir;
 
namespace {
 
// Forward declaration for use inside UnrollPattern below.
SmallVector<Value>
unrollByTile(SmallVector<OpFoldResult> mixedOffsets,
             xegpu::TensorDescType tdescTy, ArrayRef<int64_t> targetShape,
             const std::function<Value(SmallVector<OpFoldResult>)> &createOp,
             Location loc, PatternRewriter &rewriter);
 
template <typename SourceOp>
struct UnrollPattern : public OpRewritePattern<SourceOp> {
  UnrollPattern(MLIRContext *context, const xegpu::UnrollOptions &options,
                PatternBenefit benefit = 1)
      : OpRewritePattern<SourceOp>(context, benefit), options(options) {}
 
protected:
  /// Return the target shape for the given `op`. Return std::nullopt if the
  /// op shouldn't be or cannot be unrolled.
  std::optional<SmallVector<int64_t>> getTargetShape(Operation *op) const {
    LDBG() << "Get unroll shape for: " << *op;
 
    if (options.filterConstraint && failed(options.filterConstraint(op))) {
      LDBG() << "--no filter constraint -> BAIL";
      return std::nullopt;
    }
 
    assert(options.nativeShape &&
           "expects the native shape for native shape call back function.");
    auto nativeShape = options.nativeShape(op);
    return nativeShape;
  }
 
  SmallVector<Type> getUnrolledTypes(ShapedType type,
                                     ArrayRef<int64_t> tileShape) const {
    return options.getUnrolledTypes(type, tileShape);
  }
 
  /// Emulate the the unpack behavior using insert_strided_slice for VectorType
  /// values and unrealized_conversion_cast for TensorDescType values.
  Value unpack(ValueRange srcs, Type destTy, ArrayRef<int64_t> blockSize,
               Location loc, PatternRewriter &rewriter) const {
    if (auto vecTy = dyn_cast<VectorType>(destTy)) {
      auto shape = vecTy.getShape();
      return xegpu::createVectorWithShapeFromValues(rewriter, loc, srcs, shape);
    }
 
    if (isa<xegpu::TensorDescType>(destTy)) {
      auto attr = NamedAttribute(rewriter.getStringAttr(unpackAttrName),
                                 rewriter.getUnitAttr());
      auto blkAttr = NamedAttribute(rewriter.getStringAttr(blockAttrName),
                                    rewriter.getDenseI64ArrayAttr(blockSize));
      auto castOp = UnrealizedConversionCastOp::create(
          rewriter, loc, destTy, srcs,
          ArrayRef<NamedAttribute>({attr, blkAttr}));
      return castOp.getResult(0);
    }
 
    llvm_unreachable("Unexpected destTy.");
    return Value();
  }
 
  /// Emulate the the pack behavior using extract_strided_slice for VectorType
  /// values and unrealized_conversion_cast for TensorDescType values.
  SmallVector<Value> pack(Value src, TypeRange destTypes,
                          ArrayRef<int64_t> blockSize, Location loc,
                          PatternRewriter &rewriter) const {
    if (auto vecTy = dyn_cast<VectorType>(src.getType())) {
      return xegpu::extractVectorsWithShapeFromValue(rewriter, loc, src,
                                                     blockSize);
    }
 
    if (isa<xegpu::TensorDescType>(src.getType())) {
      auto attr = NamedAttribute(rewriter.getStringAttr(packAttrName),
                                 rewriter.getUnitAttr());
      auto blkAttr = NamedAttribute(rewriter.getStringAttr(blockAttrName),
                                    rewriter.getDenseI64ArrayAttr(blockSize));
      auto castOp = UnrealizedConversionCastOp::create(
          rewriter, loc, destTypes, src,
          ArrayRef<NamedAttribute>({attr, blkAttr}));
      return castOp.getResults();
    }
 
    llvm_unreachable("Unexpected src type.");
    return SmallVector<Value>();
  }
 
  /// Helper to pack operands for DPAS-like operations with early return if
  /// no unrolling is needed.
  SmallVector<Value> packOperandForDpas(Value operand,
                                        ArrayRef<int64_t> blockSize,
                                        Location loc,
                                        PatternRewriter &rewriter) const {
    auto vecType = cast<VectorType>(operand.getType());
    std::optional<SmallVector<int64_t>> grids =
        computeShapeRatio(vecType.getShape(), blockSize);
    assert(grids && "Expecting grids to be computed.");
    auto numNewOps = computeProduct(*grids);
    if (numNewOps == 1)
      return SmallVector<Value>({operand});
    VectorType newVecTy =
        vecType.cloneWith(blockSize, vecType.getElementType());
    SmallVector<Type> convertedTypes(numNewOps, newVecTy);
    return pack(operand, convertedTypes, blockSize, loc, rewriter);
  }
 
private:
  const char *const packAttrName = "__xegpu_blocking_pack__";
  const char *const unpackAttrName = "__xegpu_blocking_unpack__";
  const char *const blockAttrName = "__xegpu_blocking_tile_shape__";
 
  xegpu::UnrollOptions options;
};
 
// Walks tile offsets within the tensor descriptor shape and emits one op per
// tile by calling `createOp` with the per-tile offsets. Used by LoadNd,
// StoreNd, CreateNdDesc, and PrefetchNd unrollers, which all need to adjust
// their explicit offsets for each unrolled tile.
SmallVector<Value>
unrollByTile(SmallVector<OpFoldResult> mixedOffsets,
             xegpu::TensorDescType tdescTy, ArrayRef<int64_t> targetShape,
             const std::function<Value(SmallVector<OpFoldResult>)> &createOp,
             Location loc, PatternRewriter &rewriter) {
  int64_t rank = tdescTy.getRank();
  ArrayRef<int64_t> shape = tdescTy.getShape();
 
  auto addi = [&](OpFoldResult a, int64_t b) -> Value {
    std::optional<int64_t> maybeInt = getConstantIntValue(a);
    if (maybeInt) {
      return arith::ConstantIndexOp::create(rewriter, loc, *maybeInt + b);
    } else {
      auto aV = llvm::cast<Value>(a);
      auto bV = arith::ConstantIndexOp::create(rewriter, loc, b);
      return rewriter.createOrFold<arith::AddIOp>(loc, aV, bV);
    }
  };
 
  SmallVector<OpFoldResult> oldOffsets = llvm::to_vector(
      llvm::drop_begin(mixedOffsets, mixedOffsets.size() - rank));
  auto validIdxes =
      llvm::seq<int64_t>(mixedOffsets.size() - rank, mixedOffsets.size());
 
  SmallVector<Value> newOps;
  for (SmallVector<int64_t> offsets :
       StaticTileOffsetRange(shape, targetShape)) {
 
    for (auto [idx, oldOff, offset] :
         llvm::zip(validIdxes, oldOffsets, offsets))
      mixedOffsets[idx] = addi(oldOff, offset);
 
    auto newOp = createOp(mixedOffsets);
    newOps.push_back(newOp);
  }
  return newOps;
}
 
struct UnrollCreateNdOp : public UnrollPattern<xegpu::CreateNdDescOp> {
  using UnrollPattern<xegpu::CreateNdDescOp>::UnrollPattern;
  LogicalResult matchAndRewrite(xegpu::CreateNdDescOp op,
                                PatternRewriter &rewriter) const override {
    Location loc = op.getLoc();
    xegpu::TensorDescType tdescTy = op.getType();
 
    std::optional<SmallVector<int64_t>> targetShape = getTargetShape(op);
    if (!targetShape)
      return failure();
 
    // Keep the high-D source; only the tile shape shrinks.
    Value src = op.getSource();
    auto makeCreateNd = [&](Type tdesc) -> Value {
      auto ndTy = cast<xegpu::TensorDescType>(tdesc);
      if (isa<MemRefType>(src.getType()))
        return xegpu::CreateNdDescOp::create(rewriter, loc, ndTy,
                                             cast<TypedValue<MemRefType>>(src));
      return xegpu::CreateNdDescOp::create(
          rewriter, loc, ndTy, src, op.getMixedSizes(), op.getMixedStrides());
    };
 
    SmallVector<Type> newTdescTys = getUnrolledTypes(tdescTy, *targetShape);
    SmallVector<Value> newOps;
    if (tdescTy.getRank() <= 2) {
      // 2D: one tdesc, broadcast across tiles by pack/unpack.
      newOps.push_back(makeCreateNd(newTdescTys[0]));
    } else {
      // >2D: one tdesc per tile, so the source count matches the pack count.
      for (Type t : newTdescTys)
        newOps.push_back(makeCreateNd(t));
    }
    Value castOp = unpack(newOps, tdescTy, *targetShape, loc, rewriter);
    rewriter.replaceOp(op, castOp);
    return success();
  }
};
 
struct UnrollPrefetchNdOp : public UnrollPattern<xegpu::PrefetchNdOp> {
  using UnrollPattern<xegpu::PrefetchNdOp>::UnrollPattern;
  LogicalResult matchAndRewrite(xegpu::PrefetchNdOp op,
                                PatternRewriter &rewriter) const override {
    Location loc = op.getLoc();
    xegpu::TensorDescType tdescTy = op.getTensorDescType();
 
    std::optional<SmallVector<int64_t>> targetShape = getTargetShape(op);
    if (!targetShape)
      return failure();
 
    xegpu::DistributeLayoutAttr layout = op.getLayoutAttr();
    if (layout)
      layout = layout.dropInstData();
 
    // Batch (leading) dims unroll to unit tiles; one tdesc serves all.
    SmallVector<Type> convertedTdescTypes =
        getUnrolledTypes(tdescTy, *targetShape);
    SmallVector<Value> convertedTdesc = pack(
        op.getTensorDesc(), convertedTdescTypes, *targetShape, loc, rewriter);
 
    auto createPrefetch = [&](SmallVector<OpFoldResult> offsets) -> Value {
      xegpu::PrefetchNdOp::create(rewriter, loc, convertedTdesc[0], offsets,
                                  op.getL1HintAttr(), op.getL2HintAttr(),
                                  op.getL3HintAttr(), layout);
      return nullptr;
    };
    unrollByTile(op.getMixedOffsets(), tdescTy, *targetShape, createPrefetch,
                 loc, rewriter);
 
    rewriter.eraseOp(op);
    return success();
  }
};
 
struct UnrollLoadNdOp : public UnrollPattern<xegpu::LoadNdOp> {
  using UnrollPattern<xegpu::LoadNdOp>::UnrollPattern;
  LogicalResult matchAndRewrite(xegpu::LoadNdOp op,
                                PatternRewriter &rewriter) const override {
 
    Location loc = op.getLoc();
    VectorType valueTy = op.getType();
    xegpu::TensorDescType tdescTy = op.getTensorDescType();
 
    std::optional<SmallVector<int64_t>> targetShape = getTargetShape(op);
    if (!targetShape)
      return failure();
 
    xegpu::DistributeLayoutAttr layout = op.getLayoutAttr();
    if (layout)
      layout = layout.dropInstData();
 
    Type elemTy = tdescTy.getElementType();
    VectorType newValueTy = valueTy.cloneWith(*targetShape, elemTy);
 
    SmallVector<Value> newOps;
 
    // Batch (leading) dims unroll to unit tiles; one tdesc serves all.
    SmallVector<Type> convertedTdescTypes =
        getUnrolledTypes(tdescTy, *targetShape);
    SmallVector<Value> convertedTdescs = pack(
        op.getTensorDesc(), convertedTdescTypes, *targetShape, loc, rewriter);
 
    auto createLoad = [&](SmallVector<OpFoldResult> offsets) -> Value {
      return xegpu::LoadNdOp::create(
          rewriter, loc, newValueTy, convertedTdescs[0], offsets,
          op.getPackedAttr(), op.getTransposeAttr(), op.getL1HintAttr(),
          op.getL2HintAttr(), op.getL3HintAttr(), layout);
    };
    newOps = unrollByTile(op.getMixedOffsets(), tdescTy, *targetShape,
                          createLoad, loc, rewriter);
 
    Value castOp = unpack(newOps, op.getType(), *targetShape, loc, rewriter);
    rewriter.replaceOp(op, castOp);
    return success();
  }
};
 
struct UnrollStoreNdOp : public UnrollPattern<xegpu::StoreNdOp> {
  using UnrollPattern<xegpu::StoreNdOp>::UnrollPattern;
  LogicalResult matchAndRewrite(xegpu::StoreNdOp op,
                                PatternRewriter &rewriter) const override {
    Location loc = op.getLoc();
    VectorType valueTy = op.getValueType();
    xegpu::TensorDescType tdescTy = op.getTensorDescType();
 
    std::optional<SmallVector<int64_t>> targetShape = getTargetShape(op);
    if (!targetShape)
      return failure();
 
    xegpu::DistributeLayoutAttr layout = op.getLayoutAttr();
    if (layout)
      layout = layout.dropInstData();
 
    SmallVector<Type> convertedValTypes =
        getUnrolledTypes(valueTy, *targetShape);
 
    SmallVector<Value> convertedValues =
        pack(op.getValue(), convertedValTypes, *targetShape, loc, rewriter);
 
    size_t valueIndex = 0;
 
    // Batch (leading) dims unroll to unit tiles like any other dim. valueIndex
    // advances in unrollByTile's tile order, staying in sync with the packed
    // values.
    SmallVector<Type> convertedTdescTypes =
        getUnrolledTypes(tdescTy, *targetShape);
    SmallVector<Value> convertedTdescs = pack(
        op.getTensorDesc(), convertedTdescTypes, *targetShape, loc, rewriter);
 
    auto createStore = [&](SmallVector<OpFoldResult> offsets) {
      xegpu::StoreNdOp::create(rewriter, loc, convertedValues[valueIndex++],
                               convertedTdescs[0], offsets, op.getL1HintAttr(),
                               op.getL2HintAttr(), op.getL3HintAttr(), layout);
      return (Value) nullptr;
    };
    unrollByTile(op.getMixedOffsets(), tdescTy, *targetShape, createStore, loc,
                 rewriter);
 
    rewriter.eraseOp(op);
    return success();
  }
};
 
struct UnrollDpasOp : public UnrollPattern<xegpu::DpasOp> {
  using UnrollPattern<xegpu::DpasOp>::UnrollPattern;
  LogicalResult matchAndRewrite(xegpu::DpasOp op,
                                PatternRewriter &rewriter) const override {
    Location loc = op.getLoc();
 
    std::optional<SmallVector<int64_t>> targetShape = getTargetShape(op);
    if (!targetShape || targetShape->size() < 3)
      return failure();
 
    // targetShape is [batch..., M, K, N]
    int64_t tsRank = targetShape->size();
    auto M = (*targetShape)[tsRank - 3];
    auto K = (*targetShape)[tsRank - 2];
    auto N = (*targetShape)[tsRank - 1];
    ArrayRef<int64_t> batchDims(targetShape->data(), tsRank - 3);
 
    // Build block sizes including batch dimensions.
    SmallVector<int64_t> aBlockSize(batchDims);
    aBlockSize.push_back(M);
    aBlockSize.push_back(K);
    SmallVector<int64_t> bBlockSize(batchDims);
    bBlockSize.push_back(K);
    bBlockSize.push_back(N);
    SmallVector<int64_t> cBlockSize(batchDims);
    cBlockSize.push_back(M);
    cBlockSize.push_back(N);
 
    auto a = op.getLhs();
    auto b = op.getRhs();
    auto c = op.getAcc();
 
    SmallVector<Value> aVals = packOperandForDpas(a, aBlockSize, loc, rewriter);
    SmallVector<Value> bVals = packOperandForDpas(b, bBlockSize, loc, rewriter);
    SmallVector<Value> cVals;
    if (c)
      cVals = packOperandForDpas(c, cBlockSize, loc, rewriter);
 
    auto ranges = c ? SmallVector<ValueRange>({aVals, bVals, cVals})
                    : SmallVector<ValueRange>({aVals, bVals});
    if (llvm::any_of(ranges, [](auto &v) { return v.size() == 0; }) ||
        llvm::all_of(ranges, [](auto &v) { return v.size() == 1; }))
      return failure();
 
    VectorType resultTy = op.getResult().getType();
    auto vecTy = VectorType::get(cBlockSize, resultTy.getElementType());
 
    auto aShape = a.getType().getShape();
    auto bShape = b.getType().getShape();
 
    // Compute iteration counts. Batch dims only iterate over M and N (not
    // K-reduction), so compute batch iterations from the C block size.
    int64_t batchRank = batchDims.size();
    int64_t mIters = aShape[batchRank] / M;
    int64_t kIters = aShape[batchRank + 1] / K;
    int64_t nIters = bShape[batchRank + 1] / N;
 
    // Compute batch iterations (product of batch dim ratios).
    int64_t batchIters = 1;
    for (int64_t d = 0; d < batchRank; ++d)
      batchIters *= aShape[d] / batchDims[d];
 
    SmallVector<Value> newOps;
    for (int64_t batch = 0; batch < batchIters; ++batch) {
      for (int64_t i = 0; i < mIters; ++i) {
        for (int64_t j = 0; j < nIters; ++j) {
          Value tmpC;
          if (c)
            tmpC = cVals[batch * (mIters * nIters) + i * nIters + j];
 
          for (int64_t k = 0; k < kIters; ++k) {
            Value aVec = aVals[batch * (mIters * kIters) + i * kIters + k];
            Value bVec = bVals[batch * (kIters * nIters) + k * nIters + j];
            SmallVector<Value> operands({aVec, bVec});
            if (tmpC)
              operands.push_back(tmpC);
 
            auto newDpasOp = xegpu::DpasOp::create(
                rewriter, loc, TypeRange{vecTy}, operands, op.getProperties(),
                xegpu::dropInstDataOnAttrs(
                    op->getDiscardableAttrDictionary().getValue()));
            xegpu::dropInstDataOnInherentAttrs(newDpasOp);
            tmpC = newDpasOp.getResult();
          }
          newOps.push_back(tmpC);
        }
      }
    }
    Value castOp = unpack(newOps, resultTy, cBlockSize, loc, rewriter);
    rewriter.replaceOp(op, castOp);
    return success();
  }
};
 
struct UnrollDpasMxOp : public UnrollPattern<xegpu::DpasMxOp> {
  using UnrollPattern<xegpu::DpasMxOp>::UnrollPattern;
  LogicalResult matchAndRewrite(xegpu::DpasMxOp op,
                                PatternRewriter &rewriter) const override {
    Location loc = op.getLoc();
 
    std::optional<SmallVector<int64_t>> targetShape = getTargetShape(op);
    if (!targetShape || targetShape->size() < 4)
      return failure();
 
    // targetShape is [batch..., M, K, N, S]
    int64_t tsRank = targetShape->size();
    auto M = (*targetShape)[tsRank - 4];
    auto K = (*targetShape)[tsRank - 3];
    auto N = (*targetShape)[tsRank - 2];
    auto S = (*targetShape)[tsRank - 1];
    ArrayRef<int64_t> batchDims(targetShape->data(), tsRank - 4);
 
    SmallVector<int64_t> aBlockSize(batchDims);
    aBlockSize.push_back(M);
    aBlockSize.push_back(K);
    SmallVector<int64_t> bBlockSize(batchDims);
    bBlockSize.push_back(K);
    bBlockSize.push_back(N);
    SmallVector<int64_t> cBlockSize(batchDims);
    cBlockSize.push_back(M);
    cBlockSize.push_back(N);
    SmallVector<int64_t> aScaleBlockSize(batchDims);
    aScaleBlockSize.push_back(M);
    aScaleBlockSize.push_back(S);
    SmallVector<int64_t> bScaleBlockSize(batchDims);
    bScaleBlockSize.push_back(S);
    bScaleBlockSize.push_back(N);
 
    auto a = op.getA();
    auto b = op.getB();
    auto c = op.getAcc();
    auto ascale = dyn_cast<TypedValue<VectorType>>(op.getScaleA());
    auto bscale = dyn_cast<TypedValue<VectorType>>(op.getScaleB());
 
    SmallVector<Value> aVals = packOperandForDpas(a, aBlockSize, loc, rewriter);
    SmallVector<Value> bVals = packOperandForDpas(b, bBlockSize, loc, rewriter);
    SmallVector<Value> cVals;
    if (c)
      cVals = packOperandForDpas(c, cBlockSize, loc, rewriter);
    SmallVector<Value> aScaleVals;
    if (ascale)
      aScaleVals = packOperandForDpas(ascale, aScaleBlockSize, loc, rewriter);
    SmallVector<Value> bScaleVals;
    if (bscale)
      bScaleVals = packOperandForDpas(bscale, bScaleBlockSize, loc, rewriter);
 
    VectorType resultTy = op.getResult().getType();
    auto vecTy = VectorType::get(cBlockSize, resultTy.getElementType());
 
    auto aShape = a.getType().getShape();
    auto bShape = b.getType().getShape();
    int64_t batchRank = batchDims.size();
    int64_t mIters = aShape[batchRank] / M;
    int64_t kIters = aShape[batchRank + 1] / K;
    int64_t nIters = bShape[batchRank + 1] / N;
 
    int64_t batchIters = 1;
    for (int64_t d = 0; d < batchRank; ++d)
      batchIters *= aShape[d] / batchDims[d];
 
    SmallVector<Value> newOps;
    xegpu::DpasMxOp newDpasMxOp;
    for (int64_t batch = 0; batch < batchIters; ++batch) {
      for (int64_t i = 0; i < mIters; ++i) {
        for (int64_t j = 0; j < nIters; ++j) {
          Value tmpC;
          if (c)
            tmpC = cVals[batch * (mIters * nIters) + i * nIters + j];
 
          for (int64_t k = 0; k < kIters; ++k) {
            Value aVec = aVals[batch * (mIters * kIters) + i * kIters + k];
            Value bVec = bVals[batch * (kIters * nIters) + k * nIters + j];
            SmallVector<Value> operands({aVec, bVec});
            if (tmpC)
              operands.push_back(tmpC);
            if (ascale)
              operands.push_back(
                  aScaleVals[batch * (mIters * kIters) + i * kIters + k]);
            if (bscale)
              operands.push_back(
                  bScaleVals[batch * (kIters * nIters) + k * nIters + j]);
 
            newDpasMxOp = xegpu::DpasMxOp::create(
                rewriter, loc, TypeRange{vecTy}, operands, op.getProperties(),
                xegpu::dropInstDataOnAttrs(
                    op->getDiscardableAttrDictionary().getValue()));
            xegpu::dropInstDataOnInherentAttrs(newDpasMxOp);
            tmpC = newDpasMxOp.getResult();
          }
          newOps.push_back(newDpasMxOp);
        }
      }
    }
    Value castOp = unpack(newOps, resultTy, cBlockSize, loc, rewriter);
    rewriter.replaceOp(op, castOp);
    return success();
  }
};
 
/// This pattern handles the unrolling of LoadGatherOp with offsets (gathered
/// load).
/// It unrolls the offsets and mask operands accordingly, and creates multiple
/// LoadGatherOp with the unrolled operands.
struct UnrollLoadGatherOp : public UnrollPattern<xegpu::LoadGatherOp> {
  using UnrollPattern<xegpu::LoadGatherOp>::UnrollPattern;
  LogicalResult matchAndRewrite(xegpu::LoadGatherOp op,
                                PatternRewriter &rewriter) const override {
    Location loc = op.getLoc();
    VectorType valueTy = llvm::dyn_cast<VectorType>(op.getType());
    Value offsets = op.getOffsets();
    Value mask = op.getMask();
 
    std::optional<SmallVector<int64_t>> targetShape = getTargetShape(op);
    if (!targetShape)
      return failure();
 
    SmallVector<int64_t> targetMaskShape(*targetShape);
    int64_t chunkSize = op.getChunkSize().value_or(1);
 
    // Unroll mask and offsets with correct shape
    VectorType maskTy = llvm::dyn_cast<VectorType>(mask.getType());
    VectorType offsetsTy = llvm::dyn_cast<VectorType>(offsets.getType());
    Type elemTy = valueTy.getElementType();
    VectorType newValueTy = VectorType::get(*targetShape, elemTy);
 
    SmallVector<Type> convertedMaskTypes;
    SmallVector<Value> convertedMasks;
    SmallVector<Type> convertedOffsetTypes;
    SmallVector<Value> convertedOffsets;
 
    if (chunkSize > 1) {
      // For chunked loads, mask and offsets have one less dimension
      targetMaskShape.pop_back();
      int64_t blockedChunkSize = targetShape->back();
      int64_t numNewChunks = chunkSize / blockedChunkSize;
      chunkSize = blockedChunkSize;
 
      convertedMaskTypes = getUnrolledTypes(maskTy, targetMaskShape);
      convertedOffsetTypes = getUnrolledTypes(offsetsTy, targetMaskShape);
 
      SmallVector<Value> convertedMasksBase =
          pack(mask, convertedMaskTypes, targetMaskShape, loc, rewriter);
      SmallVector<Value> convertedOffsetsBase =
          pack(offsets, convertedOffsetTypes, targetMaskShape, loc, rewriter);
 
      for (auto maskVal : convertedMasksBase)
        convertedMasks.append(numNewChunks, maskVal);
 
      for (auto [baseOffset, offsetType] :
           llvm::zip(convertedOffsetsBase, convertedOffsetTypes)) {
        for (int64_t i = 0; i < numNewChunks; ++i) {
          Value inc = arith::ConstantIndexOp::create(rewriter, loc,
                                                     i * blockedChunkSize);
          Value incVec =
              vector::BroadcastOp::create(rewriter, loc, offsetType, inc);
          Value offsetVal =
              arith::AddIOp::create(rewriter, loc, baseOffset, incVec);
          convertedOffsets.push_back(offsetVal);
        }
      }
    } else {
      convertedMaskTypes = getUnrolledTypes(maskTy, targetMaskShape);
      convertedMasks =
          pack(mask, convertedMaskTypes, targetMaskShape, loc, rewriter);
 
      convertedOffsetTypes = getUnrolledTypes(offsetsTy, *targetShape);
      convertedOffsets =
          pack(offsets, convertedOffsetTypes, *targetShape, loc, rewriter);
    }
 
    auto layout = op.getLayoutAttr();
    if (layout)
      layout = layout.dropInstData();
 
    SmallVector<Value> newOps;
    for (auto [o, m] : llvm::zip(convertedOffsets, convertedMasks)) {
      auto newOp = xegpu::LoadGatherOp::create(
          rewriter, loc, newValueTy, op.getSource(), o, m,
          rewriter.getI64IntegerAttr(chunkSize), op.getL1HintAttr(),
          op.getL2HintAttr(), op.getL3HintAttr(), layout,
          /*contiguity=*/nullptr);
      newOps.push_back(newOp);
    }
 
    Value castOp = unpack(newOps, op.getType(), *targetShape, loc, rewriter);
    rewriter.replaceOp(op, castOp);
    return success();
  }
};
 
/// This pattern handles the unrolling of StoreScatterOp with offsets (scattered
/// store).
/// It unrolls the offsets and mask operands accordingly, and creates multiple
/// StoreScatterOp with the unrolled operands.
struct UnrollStoreScatterOp : public UnrollPattern<xegpu::StoreScatterOp> {
  using UnrollPattern<xegpu::StoreScatterOp>::UnrollPattern;
  LogicalResult matchAndRewrite(xegpu::StoreScatterOp op,
                                PatternRewriter &rewriter) const override {
    Location loc = op.getLoc();
    VectorType valueTy = llvm::dyn_cast<VectorType>(op.getValue().getType());
    Value offsets = op.getOffsets();
    Value mask = op.getMask();
 
    std::optional<SmallVector<int64_t>> targetShape = getTargetShape(op);
    if (!targetShape)
      return failure();
 
    int64_t chunkSize = op.getChunkSize().value_or(1);
 
    SmallVector<int64_t> targetMaskShape(*targetShape);
    VectorType maskTy = llvm::dyn_cast<VectorType>(mask.getType());
    VectorType offsetsTy = llvm::dyn_cast<VectorType>(offsets.getType());
 
    SmallVector<Type> convertedMaskTypes;
    SmallVector<Value> convertedMasks;
    SmallVector<Type> convertedOffsetTypes;
    SmallVector<Value> convertedOffsets;
 
    if (chunkSize > 1) {
      targetMaskShape.pop_back();
      int64_t blockedChunkSize = targetShape->back();
      int64_t numNewChunks = chunkSize / blockedChunkSize;
      chunkSize = blockedChunkSize;
 
      convertedMaskTypes = getUnrolledTypes(maskTy, targetMaskShape);
      convertedOffsetTypes = getUnrolledTypes(offsetsTy, targetMaskShape);
 
      SmallVector<Value> convertedMasksBase =
          pack(mask, convertedMaskTypes, targetMaskShape, loc, rewriter);
      SmallVector<Value> convertedOffsetsBase =
          pack(offsets, convertedOffsetTypes, targetMaskShape, loc, rewriter);
 
      for (auto maskVal : convertedMasksBase)
        convertedMasks.append(numNewChunks, maskVal);
 
      for (auto [baseOffset, offsetType] :
           llvm::zip(convertedOffsetsBase, convertedOffsetTypes)) {
        for (int64_t i = 0; i < numNewChunks; ++i) {
          Value inc = arith::ConstantIndexOp::create(rewriter, loc,
                                                     i * blockedChunkSize);
          Value incVec =
              vector::BroadcastOp::create(rewriter, loc, offsetType, inc);
          Value offsetVal =
              arith::AddIOp::create(rewriter, loc, baseOffset, incVec);
          convertedOffsets.push_back(offsetVal);
        }
      }
    } else {
      convertedMaskTypes = getUnrolledTypes(maskTy, targetMaskShape);
      convertedMasks =
          pack(mask, convertedMaskTypes, targetMaskShape, loc, rewriter);
 
      convertedOffsetTypes = getUnrolledTypes(offsetsTy, *targetShape);
      convertedOffsets =
          pack(offsets, convertedOffsetTypes, *targetShape, loc, rewriter);
    }
 
    SmallVector<Type> convertedValTypes =
        getUnrolledTypes(valueTy, *targetShape);
    SmallVector<Value> convertedValues =
        pack(op.getValue(), convertedValTypes, *targetShape, loc, rewriter);
 
    auto layout = op.getLayoutAttr();
    if (layout)
      layout = layout.dropInstData();
 
    for (auto [v, o, m] :
         llvm::zip(convertedValues, convertedOffsets, convertedMasks)) {
      xegpu::StoreScatterOp::create(rewriter, loc, v, op.getDest(), o, m,
                                    rewriter.getI64IntegerAttr(chunkSize),
                                    op.getL1HintAttr(), op.getL2HintAttr(),
                                    op.getL3HintAttr(), layout,
                                    /*contiguity=*/nullptr);
    }
 
    rewriter.eraseOp(op);
    return success();
  }
};
 
struct UnrollLoadMatrixOp : public UnrollPattern<xegpu::LoadMatrixOp> {
  using UnrollPattern<xegpu::LoadMatrixOp>::UnrollPattern;
  LogicalResult matchAndRewrite(xegpu::LoadMatrixOp op,
                                PatternRewriter &rewriter) const override {
    Location loc = op.getLoc();
    VectorType valueTy = llvm::dyn_cast<VectorType>(op.getType());
    assert(valueTy && "the value type must be vector type!");
 
    std::optional<SmallVector<int64_t>> targetShape = getTargetShape(op);
    if (!targetShape || targetShape->size() != (size_t)valueTy.getRank())
      return failure();
 
    Type elemTy = valueTy.getElementType();
    ArrayRef<int64_t> shape = valueTy.getShape();
    xegpu::DistributeLayoutAttr layout = op.getLayoutAttr();
 
    VectorType newValueTy = valueTy.cloneWith(*targetShape, elemTy);
 
    SmallVector<OpFoldResult> mixedOffsets = op.getMixedOffsets();
    SmallVector<SmallVector<OpFoldResult>> offsetsList;
    for (SmallVector<int64_t> offsets :
         StaticTileOffsetRange(shape, *targetShape)) {
      auto adds = xegpu::addElementwise(
          rewriter, loc, mixedOffsets,
          getAsIndexOpFoldResult(op.getContext(), offsets));
      offsetsList.push_back(adds);
    }
 
    SmallVector<Value> newOps;
    if (layout)
      layout = layout.dropInstData();
    for (SmallVector<OpFoldResult> offsets : offsetsList) {
      auto newOp = xegpu::LoadMatrixOp::create(
          rewriter, op.getLoc(), newValueTy, op.getMemDesc(), offsets, layout);
      newOps.push_back(newOp);
    }
    Value castOp = unpack(newOps, op.getType(), *targetShape, loc, rewriter);
    rewriter.replaceOp(op, castOp);
    return success();
  }
};
 
struct UnrollStoreMatrixOp : public UnrollPattern<xegpu::StoreMatrixOp> {
  using UnrollPattern<xegpu::StoreMatrixOp>::UnrollPattern;
  LogicalResult matchAndRewrite(xegpu::StoreMatrixOp op,
                                PatternRewriter &rewriter) const override {
    std::optional<SmallVector<int64_t>> targetShape = getTargetShape(op);
    if (!targetShape)
      return failure();
 
    Location loc = op.getLoc();
    VectorType valueTy = llvm::dyn_cast<VectorType>(op.getData().getType());
    assert(valueTy && "the value type must be vector type!");
    ArrayRef<int64_t> shape = valueTy.getShape();
    xegpu::DistributeLayoutAttr layout = op.getLayoutAttr();
    if (layout)
      layout = layout.dropInstData();
 
    SmallVector<Type> convertedValTypes =
        getUnrolledTypes(valueTy, *targetShape);
    SmallVector<Value> convertedValues =
        pack(op.getData(), convertedValTypes, *targetShape, loc, rewriter);
 
    SmallVector<OpFoldResult> mixedOffsets = op.getMixedOffsets();
    SmallVector<SmallVector<OpFoldResult>> offsetsList;
    for (SmallVector<int64_t> offsets :
         StaticTileOffsetRange(shape, *targetShape)) {
      auto adds = xegpu::addElementwise(
          rewriter, loc, mixedOffsets,
          getAsIndexOpFoldResult(op.getContext(), offsets));
      offsetsList.push_back(adds);
    }
 
    for (auto [v, offsets] : llvm::zip_equal(convertedValues, offsetsList))
      xegpu::StoreMatrixOp::create(rewriter, loc, v, op.getMemDesc(), offsets,
                                   layout);
 
    rewriter.eraseOp(op);
    return success();
  }
};
 
/// UnrollConvertLayoutOp pattern for unrolling xegpu::ConvertLayoutOp
/// operations. It first check whether the convert layout op has valid layouts
/// after inst_data stripped. If it does, it will unroll the vector into
/// multiple smaller vectors according to the target shape, and create multiple
/// ConvertLayoutOp with the unrolled vectors and the stripped layouts.
///
/// When the input and target layouts have different inst_data, the source is
/// extracted at the input inst_data granularity and the result is inserted at
/// the target inst_data granularity, enabling slice cancellation during
/// canonicalization.
struct UnrollConvertLayoutOp : public UnrollPattern<xegpu::ConvertLayoutOp> {
  using UnrollPattern<xegpu::ConvertLayoutOp>::UnrollPattern;
 
  /// Extracts source in `inTile` slices, regroups into `convTile`-sized
  /// ConvertLayoutOps, and inserts the result in `outTile` slices.
  /// Returns failure if the tiles do not evenly divide.
  LogicalResult
  rewriteWithRegrouping(xegpu::ConvertLayoutOp op, VectorType valueTy,
                        ArrayRef<int64_t> convTile, ArrayRef<int64_t> inTile,
                        ArrayRef<int64_t> outTile,
                        xegpu::DistributeLayoutAttr inputLayout,
                        xegpu::DistributeLayoutAttr targetLayout, Location loc,
                        PatternRewriter &rewriter) const {
    ArrayRef<int64_t> vecShape = valueTy.getShape();
    if (!computeShapeRatio(vecShape, convTile) ||
        !computeShapeRatio(convTile, inTile) ||
        !computeShapeRatio(convTile, outTile))
      return failure();
 
    Type elemTy = valueTy.getElementType();
    int64_t rank = valueTy.getRank();
    VectorType convTy = VectorType::get(convTile, elemTy);
    SmallVector<int64_t> strides(rank, 1);
 
    Value source = op.getSource();
    auto zeroOf = [&](VectorType ty) -> Value {
      return arith::ConstantOp::create(
          rewriter, loc, ty,
          DenseElementsAttr::get(ty, rewriter.getZeroAttr(elemTy)));
    };
    auto addOffsets = [](ArrayRef<int64_t> a,
                         ArrayRef<int64_t> b) -> SmallVector<int64_t> {
      SmallVector<int64_t> res(a);
      for (auto [r, v] : llvm::zip_equal(res, b))
        r += v;
      return res;
    };
 
    Value result = zeroOf(valueTy);
    for (SmallVector<int64_t> convOff :
         StaticTileOffsetRange(vecShape, convTile)) {
      // Build the convert tile from inTile-sized slices of the source.
      Value conv;
      if (convTile == inTile) {
        conv = vector::ExtractStridedSliceOp::create(
            rewriter, loc, source, convOff, convTile, strides);
      } else {
        conv = zeroOf(convTy);
        for (SmallVector<int64_t> inLocal :
             StaticTileOffsetRange(convTile, inTile)) {
          Value piece = vector::ExtractStridedSliceOp::create(
              rewriter, loc, source, addOffsets(convOff, inLocal), inTile,
              strides);
          conv = vector::InsertStridedSliceOp::create(rewriter, loc, piece,
                                                      conv, inLocal, strides);
        }
      }
 
      conv = xegpu::ConvertLayoutOp::create(rewriter, loc, convTy, conv,
                                            inputLayout, targetLayout);
 
      // Write the converted tile into the result as outTile-sized slices.
      if (convTile == outTile) {
        result = vector::InsertStridedSliceOp::create(rewriter, loc, conv,
                                                      result, convOff, strides);
      } else {
        for (SmallVector<int64_t> outLocal :
             StaticTileOffsetRange(convTile, outTile)) {
          Value piece = vector::ExtractStridedSliceOp::create(
              rewriter, loc, conv, outLocal, outTile, strides);
          result = vector::InsertStridedSliceOp::create(
              rewriter, loc, piece, result, addOffsets(convOff, outLocal),
              strides);
        }
      }
    }
 
    rewriter.replaceOp(op, result);
    return success();
  }
 
  LogicalResult matchAndRewrite(xegpu::ConvertLayoutOp op,
                                PatternRewriter &rewriter) const override {
    Location loc = op.getLoc();
    Type valType = op.getType();
 
    xegpu::DistributeLayoutAttr inputLayout = op.getEffectiveInputLayout();
    xegpu::DistributeLayoutAttr targetLayout = op.getTargetLayoutAttr();
    if (!inputLayout || !targetLayout)
      return rewriter.notifyMatchFailure(op, "missing layout attributes.");
 
    if (valType.isIntOrFloat()) {
      rewriter.replaceOp(op, op.getSource());
      return success();
    }
 
    // Capture inst_data granularities before stripping them.
    SmallVector<int64_t> inTile = inputLayout.getEffectiveInstDataAsInt();
    SmallVector<int64_t> outTile = targetLayout.getEffectiveInstDataAsInt();
    if (inTile.empty() || outTile.empty())
      return rewriter.notifyMatchFailure(op, "Not a target ConvertLayoutOp.");
 
    inputLayout = inputLayout.dropInstData();
    targetLayout = targetLayout.dropInstData();
 
    VectorType valueTy = llvm::dyn_cast<VectorType>(op.getType());
    assert(valueTy && "the value type must be vector type!");
 
    std::optional<SmallVector<int64_t>> targetShape = getTargetShape(op);
    if (!targetShape || targetShape->size() != (size_t)valueTy.getRank())
      return failure();
 
    // Nothing to convert if layouts match after stripping inst_data.
    if (!inputLayout || !targetLayout || inputLayout.isEqualTo(targetLayout)) {
      rewriter.replaceOp(op, op.getSource());
      return success();
    }
 
    // Try regrouping: extract at inTile, convert, insert at outTile.
    if (succeeded(rewriteWithRegrouping(op, valueTy, *targetShape, inTile,
                                        outTile, inputLayout, targetLayout, loc,
                                        rewriter)))
      return success();
 
    // Fallback: pack/unpack at the convert tile granularity.
    SmallVector<Type> convertedValTypes =
        getUnrolledTypes(valueTy, *targetShape);
    SmallVector<Value> convertedValues =
        pack(op.getOperand(), convertedValTypes, *targetShape, loc, rewriter);
    SmallVector<Value> newOps;
    for (auto [v, t] : llvm::zip(convertedValues, convertedValTypes)) {
      auto newOp = xegpu::ConvertLayoutOp::create(rewriter, loc, t, v,
                                                  inputLayout, targetLayout);
      newOps.push_back(newOp);
    }
    Value newSource = unpack(newOps, op.getType(), *targetShape, loc, rewriter);
 
    rewriter.replaceOp(op, newSource);
    return success();
  }
};
 
/// Unrolls vector.multi_reduction by sequentially reducing tiles with
/// elementwise arith operations first, then a single multi_reduction
/// per non-reduced tile position. This avoids generating long chains of
/// multi_reduction ops (as the upstream pattern does) and is more efficient.
///
/// Example:
/// vector.multi_reduction <32x64xf16> to <32xf16> (tile_shape=32, 32)
/// -- Upstream pattern generates:
/// %tmp1 = vector.multi_reduction %tile0, %zero_acc <32x32xf16> to <32xf16>
/// %res = vector.multi_reduction %tmp1, %tile1 <32x32xf16> to <32xf16>
/// -- This pattern generates:
/// %tmp1 = arith.reduction %tile0, %tile1 <32x32xf16> -> <32x32xf16> //
/// elementwise %res = vector.multi_reduction %tmp1, %zero_acc <32x32xf16> to
/// <32xf16>
struct UnrollMultiReductionOp
    : public UnrollPattern<vector::MultiDimReductionOp> {
  UnrollMultiReductionOp(MLIRContext *context,
                         const xegpu::UnrollOptions &options,
                         PatternBenefit benefit = 2)
      : UnrollPattern<vector::MultiDimReductionOp>(context, options, benefit) {}
 
  LogicalResult matchAndRewrite(vector::MultiDimReductionOp reductionOp,
                                PatternRewriter &rewriter) const override {
    VectorType srcTy = reductionOp.getSourceVectorType();
    ArrayRef<int64_t> srcShape = srcTy.getShape();
    int64_t srcRank = srcTy.getRank();
 
    Location loc = reductionOp.getLoc();
    Value source = reductionOp.getSource();
    Value acc = reductionOp.getAcc();
    vector::CombiningKind kind = reductionOp.getKind();
 
    // Result must be a vector (not scalar).
    auto resultType = dyn_cast<VectorType>(reductionOp.getDestType());
    if (!resultType)
      return failure();
 
    std::optional<SmallVector<int64_t>> targetShapeOpt =
        getTargetShape(reductionOp);
    if (!targetShapeOpt ||
        static_cast<int64_t>(targetShapeOpt->size()) != srcRank)
      return failure();
 
    SmallVector<int64_t> targetShape = *targetShapeOpt;
 
    // Check divisibility for all dimensions.
    for (int64_t i = 0; i < srcRank; ++i) {
      if (srcShape[i] % targetShape[i] != 0)
        return failure();
    }
 
    SmallVector<bool> reductionMask = reductionOp.getReductionMask();
    // Identify reduced and kept dimensions from the reduction mask.
    SmallVector<int64_t> reducedDims, keptDims;
    for (int64_t i = 0; i < srcRank; ++i) {
      if (reductionMask[i])
        reducedDims.push_back(i);
      else
        keptDims.push_back(i);
    }
 
    // Compute the number of tiles along each reduced dimension and their
    // product
    SmallVector<int64_t> numReducedTilesPerDim;
    for (int64_t d : reducedDims)
      numReducedTilesPerDim.push_back(srcShape[d] / targetShape[d]);
 
    // Build kept shapes for iterating over non-reduced dimensions.
    SmallVector<int64_t> keptShape, keptTileShape;
    for (int64_t d : keptDims) {
      keptShape.push_back(srcShape[d]);
      keptTileShape.push_back(targetShape[d]);
    }
 
    // Initialize the result vector for assembly.
    Value result = arith::ConstantOp::create(rewriter, loc, resultType,
                                             rewriter.getZeroAttr(resultType));
 
    // Iterate over all tile positions in the kept dimensions.
    // Ex: [off0, off1, _ _ off4]
    // blanks are offsets for the reduced dims, they will be
    // generated in the inner loop below
    for (SmallVector<int64_t> keptOffsets :
         StaticTileOffsetRange(keptShape, keptTileShape)) {
 
      // Reconstruct full-rank base offsets with 0 for reduced dims.
      // Ex: [off0, off1, 0, 0, off4]
      SmallVector<int64_t> baseOffsets(srcRank, 0);
      for (auto [idx, dim] : llvm::enumerate(keptDims))
        baseOffsets[dim] = keptOffsets[idx];
 
      // Generate the full tile indices for the reduced dimensions.
      // Ex: if reduceDimShapes = [32, 64] and
      // reducedDimTargetShapes = [16, 16], then reducedTileCoords:
      // [(0, 0), (0, 1), (0, 2), (0, 3),
      //  (1, 0), (1, 1), (1, 2), (1, 3)]
      auto reducedTileCoords = StaticTileOffsetRange(
          numReducedTilesPerDim, SmallVector<int64_t>(reducedDims.size(), 1));
 
      // Step 1: Fill "blanks" in the offsets for the reduced dimensions
      // using 'reducedTileCoords' and extract according tiles.
      // Ex: tiles = [source[off0, off1, off2_red, off3_red, off4], ...]
      SmallVector<Value> tiles;
      for (SmallVector<int64_t> reducedTileIdx : reducedTileCoords) {
        SmallVector<int64_t> offsets(baseOffsets);
        for (auto [idx, dim] : llvm::enumerate(reducedDims))
          offsets[dim] = reducedTileIdx[idx] * targetShape[dim];
        SmallVector<int64_t> strides(srcRank, 1);
        Value tile = vector::ExtractStridedSliceOp::create(
            rewriter, loc, source, offsets, targetShape, strides);
        tiles.push_back(tile);
      }
 
      // Step 2: Sequentially reduce tiles using elementwise arith operations.
      Value reduced = tiles[0];
      for (size_t i = 1; i < tiles.size(); ++i)
        reduced =
            vector::makeArithReduction(rewriter, loc, kind, reduced, tiles[i]);
 
      // Step 3: Perform a single multi_reduction with the accumulator slice.
      SmallVector<int64_t> accStrides(keptTileShape.size(), 1);
      Value accSlice = vector::ExtractStridedSliceOp::create(
          rewriter, loc, acc, keptOffsets, keptTileShape, accStrides);
 
      auto newReduction = vector::MultiDimReductionOp::create(
          rewriter, loc, reduced, accSlice, reductionMask, kind);
 
      // Step 4: Insert the reduced result into the output vector.
      SmallVector<int64_t> dstStrides(keptTileShape.size(), 1);
      result = vector::InsertStridedSliceOp::create(
          rewriter, loc, newReduction, result, keptOffsets, dstStrides);
    }
 
    rewriter.replaceOp(reductionOp, result);
    return success();
  }
};
 
} // namespace
 
void mlir::xegpu::populateXeGPUUnrollPatterns(
    RewritePatternSet &patterns, const xegpu::UnrollOptions &options) {
  patterns
      .add<UnrollCreateNdOp, UnrollPrefetchNdOp, UnrollLoadNdOp,
           UnrollStoreNdOp, UnrollDpasOp, UnrollDpasMxOp, UnrollLoadMatrixOp,
           UnrollStoreMatrixOp, UnrollLoadGatherOp, UnrollStoreScatterOp,
           UnrollConvertLayoutOp, UnrollMultiReductionOp>(patterns.getContext(),
                                                          options);
}