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 for the rank > 2 case shared by Load/Store/PrefetchNd unroll
  /// patterns. The matching CreateNdDesc unroll pattern produces one tdesc
  /// per batch tile (the batch offset is baked into its base pointer via
  /// memref.subview), so here we only need to iterate the inner 2D offsets
  /// for each batch tdesc.
  ///
  /// Packs `srcTdesc` into one tdesc per batch tile, then iterates the inner
  /// 2D tile offsets for each batch tdesc and invokes `createOp` with
  /// (batchTdesc, fullOffsets), where fullOffsets is `batchRank` zeros
  /// followed by the inner offsets. Returns the values produced by createOp,
  /// flattened across (batch, inner) iteration order.
  SmallVector<Value> unrollNdBatch(
      Value srcTdesc, xegpu::TensorDescType tdescTy,
      ArrayRef<int64_t> targetShape, ArrayRef<OpFoldResult> mixedOffsets,
      int64_t batchRank,
      llvm::function_ref<Value(Value, SmallVector<OpFoldResult>)> createOp,
      Location loc, PatternRewriter &rewriter) const {
    ArrayRef<int64_t> shape = tdescTy.getShape();
    SmallVector<int64_t> innerShape(shape.begin() + batchRank, shape.end());
    SmallVector<int64_t> innerTarget(targetShape.begin() + batchRank,
                                     targetShape.end());
 
    SmallVector<Type> batchTdescTypes = getUnrolledTypes(tdescTy, targetShape);
    SmallVector<Value> batchTdescs =
        pack(srcTdesc, batchTdescTypes, targetShape, loc, rewriter);
 
    auto innerTdescTy = xegpu::TensorDescType::get(
        tdescTy.getContext(), innerShape, tdescTy.getElementType(),
        tdescTy.getEncoding(), /*layout=*/nullptr);
 
    SmallVector<OpFoldResult> innerOffsets(mixedOffsets.begin() + batchRank,
                                           mixedOffsets.end());
 
    SmallVector<Value> newOps;
    for (Value batchTdesc : batchTdescs) {
      auto wrappedCreate = [&](SmallVector<OpFoldResult> offsets) -> Value {
        SmallVector<OpFoldResult> fullOffsets(batchRank,
                                              rewriter.getIndexAttr(0));
        fullOffsets.append(offsets.begin(), offsets.end());
        return createOp(batchTdesc, fullOffsets);
      };
      auto perBatch = unrollByTile(innerOffsets, innerTdescTy, innerTarget,
                                   wrappedCreate, loc, rewriter);
      newOps.append(perBatch.begin(), perBatch.end());
    }
    return newOps;
  }
 
  /// 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();
 
    int64_t rank = tdescTy.getRank();
    int64_t batchRank = rank - 2;
 
    // For rank <= 2 or non-memref source: existing single-tdesc behavior.
    if (batchRank <= 0 || !isa<MemRefType>(op.getSourceType())) {
      SmallVector<Value> newOps;
      auto newTdescTy = getUnrolledTypes(tdescTy, *targetShape)[0];
      auto newOp = xegpu::CreateNdDescOp::create(
          rewriter, loc, newTdescTy, op.getSource(), op.getMixedSizes(),
          op.getMixedStrides());
      newOps.push_back(newOp);
      Value castOp = unpack(newOps, tdescTy, *targetShape, loc, rewriter);
      rewriter.replaceOp(op, castOp);
      return success();
    }
 
    // For rank > 2 with memref source: create one tdesc per batch tile via
    // memref.subview. Each subview slices the batch dimensions, so the
    // resulting tdesc has the batch offset baked into its base pointer.
    // The inner dimensions remain full-size for reuse across multiple
    // load/store operations with different offsets.
    ArrayRef<int64_t> shape = tdescTy.getShape();
    SmallVector<int64_t> batchBlockSize(targetShape->begin(),
                                        targetShape->begin() + batchRank);
    batchBlockSize.append(shape.begin() + batchRank, shape.end());
 
    auto newTdescTy =
        cast<xegpu::TensorDescType>(getUnrolledTypes(tdescTy, *targetShape)[0]);
 
    SmallVector<Value> newOps;
    for (SmallVector<int64_t> batchOffsets :
         StaticTileOffsetRange(shape, batchBlockSize)) {
      // Build memref.subview operands. The subview slices contiguously along
      // each batch dimension (no gaps), so the subview's element stride is 1
      // for every dim. This is unrelated to the source memref's strides, which
      // describe the layout of the original buffer and are propagated by the
      // SubViewOp builder onto the resulting memref type.
      SmallVector<OpFoldResult> subviewOffsets;
      for (int64_t off : batchOffsets)
        subviewOffsets.push_back(rewriter.getIndexAttr(off));
 
      SmallVector<OpFoldResult> subviewSizes;
      for (int64_t d : batchBlockSize)
        subviewSizes.push_back(rewriter.getIndexAttr(d));
 
      SmallVector<OpFoldResult> subviewStrides(rank, rewriter.getIndexAttr(1));
 
      auto subview = memref::SubViewOp::create(rewriter, loc, op.getSource(),
                                               subviewOffsets, subviewSizes,
                                               subviewStrides);
 
      auto newOp = xegpu::CreateNdDescOp::create(
          rewriter, loc, newTdescTy,
          cast<TypedValue<MemRefType>>(subview.getResult()));
      newOps.push_back(newOp);
    }
 
    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();
 
    int64_t rank = tdescTy.getRank();
    int64_t batchRank = rank - 2;
 
    if (batchRank <= 0) {
      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);
    } else {
      // Rank > 2: batch tdescs cover [batchTarget..., innerShape...].
      // Each batch tdesc is reused for multiple inner prefetches via offsets.
      auto createPrefetch =
          [&](Value tdesc, SmallVector<OpFoldResult> fullOffsets) -> Value {
        xegpu::PrefetchNdOp::create(rewriter, loc, tdesc, fullOffsets,
                                    op.getL1HintAttr(), op.getL2HintAttr(),
                                    op.getL3HintAttr(), layout);
        return nullptr;
      };
      this->unrollNdBatch(op.getTensorDesc(), tdescTy, *targetShape,
                          op.getMixedOffsets(), batchRank, 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);
 
    int64_t rank = tdescTy.getRank();
    int64_t batchRank = rank - 2;
    SmallVector<Value> newOps;
 
    if (batchRank <= 0) {
      // Rank <= 2: original behavior with single tdesc.
      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);
    } else {
      // Rank > 2: batch tdescs cover [batchTarget..., innerShape...].
      // Each batch tdesc is reused for multiple inner loads via offsets.
      auto createLoad = [&](Value tdesc,
                            SmallVector<OpFoldResult> fullOffsets) -> Value {
        return xegpu::LoadNdOp::create(
            rewriter, loc, newValueTy, tdesc, fullOffsets, op.getPackedAttr(),
            op.getTransposeAttr(), op.getL1HintAttr(), op.getL2HintAttr(),
            op.getL3HintAttr(), layout);
      };
      newOps = this->unrollNdBatch(op.getTensorDesc(), tdescTy, *targetShape,
                                   op.getMixedOffsets(), batchRank, 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);
 
    int64_t rank = tdescTy.getRank();
    int64_t batchRank = rank - 2;
    size_t valueIndex = 0;
 
    if (batchRank <= 0) {
      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);
    } else {
      // Rank > 2: batch tdescs cover [batchTarget..., innerShape...].
      // Each batch tdesc is reused for multiple inner stores via offsets.
      // valueIndex advances across (batch, inner) iterations in the same
      // order unrollNdBatch invokes the callback, so it stays in sync with
      // the pre-packed convertedValues.
      auto createStore = [&](Value tdesc,
                             SmallVector<OpFoldResult> fullOffsets) -> Value {
        xegpu::StoreNdOp::create(
            rewriter, loc, convertedValues[valueIndex++], tdesc, fullOffsets,
            op.getL1HintAttr(), op.getL2HintAttr(), op.getL3HintAttr(), layout);
        return nullptr;
      };
      this->unrollNdBatch(op.getTensorDesc(), tdescTy, *targetShape,
                          op.getMixedOffsets(), batchRank, 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);
 
            tmpC = xegpu::DpasOp::create(
                rewriter, loc, vecTy, operands,
                xegpu::dropInstDataOnAttrs(op->getAttrs()));
          }
          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, vecTy, operands,
                xegpu::dropInstDataOnAttrs(op->getAttrs()));
            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 = 1;
    if (auto chunkSizeAttr = op->getAttr("chunk_size")) {
      if (auto intAttr = llvm::dyn_cast<IntegerAttr>(chunkSizeAttr))
        chunkSize = intAttr.getInt();
    }
 
    // 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);
      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 = 1;
    if (auto chunkSizeAttr = op->getAttr("chunk_size")) {
      if (auto intAttr = llvm::dyn_cast<IntegerAttr>(chunkSizeAttr))
        chunkSize = intAttr.getInt();
    }
 
    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);
    }
 
    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.
struct UnrollConvertLayoutOp : public UnrollPattern<xegpu::ConvertLayoutOp> {
  using UnrollPattern<xegpu::ConvertLayoutOp>::UnrollPattern;
  LogicalResult matchAndRewrite(xegpu::ConvertLayoutOp op,
                                PatternRewriter &rewriter) const override {
    Location loc = op.getLoc();
    Type valType = op.getType();
 
    xegpu::DistributeLayoutAttr inputLayout = op.getInputLayoutAttr();
    xegpu::DistributeLayoutAttr targetLayout = op.getTargetLayoutAttr();
    if (!inputLayout || !targetLayout)
      return rewriter.notifyMatchFailure(op, "missing layout attributes.");
 
    if (valType.isIntOrFloat()) {
      rewriter.replaceOp(op, op.getSource());
      assert(!inputLayout.dropInstData() && !targetLayout.dropInstData() &&
             "unexpected layout attributes for scalar type");
      return success();
    }
 
    if (inputLayout.getEffectiveInstDataAsInt().empty() ||
        targetLayout.getEffectiveInstDataAsInt().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();
 
    Value newSource = op.getSource();
    SmallVector<Value> newOps;
    if (inputLayout && targetLayout) {
      SmallVector<Type> convertedValTypes =
          getUnrolledTypes(valueTy, *targetShape);
      SmallVector<Value> convertedValues =
          pack(op.getOperand(), convertedValTypes, *targetShape, loc, rewriter);
      for (auto [v, t] : llvm::zip(convertedValues, convertedValTypes)) {
        auto newOp = xegpu::ConvertLayoutOp::create(rewriter, loc, t, v,
                                                    inputLayout, targetLayout);
        newOps.push_back(newOp);
      }
      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);
}