XeGPUPeepHoleOptimizer.cpp

Go to the documentation of this file.

//===- XeGPUPeepHoleOptimizer.cpp - XeGPU optimize block loads -*- 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
//
//===----------------------------------------------------------------------===//
 
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/Transforms/Patterns.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/XeGPU/IR/XeGPU.h"
#include "mlir/Dialect/XeGPU/Transforms/Passes.h"
#include "mlir/Dialect/XeGPU/Transforms/Transforms.h"
#include "mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h"
#include "mlir/Dialect/XeGPU/Utils/XeGPUUtils.h"
#include "mlir/Dialect/XeGPU/uArch/uArchBase.h"
#include "mlir/Dialect/XeGPU/uArch/uArchCommon.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/OpDefinition.h"
#include "mlir/IR/Types.h"
#include "mlir/IR/Value.h"
#include "mlir/Transforms/DialectConversion.h"
#include "mlir/Transforms/GreedyPatternRewriteDriver.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include <optional>
 
namespace mlir {
namespace xegpu {
#define GEN_PASS_DEF_XEGPUPEEPHOLEOPTIMIZER
#include "mlir/Dialect/XeGPU/Transforms/Passes.h.inc"
} // namespace xegpu
} // namespace mlir
 
#define DEBUG_TYPE "xegpu-optimize-peephole"
#define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "]: ")
 
using namespace mlir;
 
namespace {
 
/// Get the 2D lane data from a tensor desc type if it exists.
static std::optional<SmallVector<int64_t>>
getMaybeLaneData(xegpu::TensorDescType tdescType) {
  auto layout = tdescType.getLayoutAttr();
  if (!layout)
    return std::nullopt;
  return xegpu::getInner2DIfUnitLeadingDims(layout.getEffectiveLaneDataAsInt());
}
 
/// Get the 2D lane layout from a tensor desc type if it exists.
static std::optional<SmallVector<int64_t>>
getMaybeLaneLayout(xegpu::TensorDescType tdescType) {
  auto layout = tdescType.getLayoutAttr();
  if (!layout)
    return std::nullopt;
  return xegpu::getInner2DIfUnitLeadingDims(
      layout.getEffectiveLaneLayoutAsInt());
}
 
/// A layout can be optimized if its lane layout is transposed (lane[0] != 1 &&
/// lane[1] == 1), but inner lane data is not equal to [1, 1].
/// Example:
///     !xegpu.tensor_desc<16x16xf16,
///         #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>>
/// In this case, lane layout is transposed (from the usual [1, SG_SIZE] form)
/// indicating that this is a load that requires transpose effect. However,
/// lane data is [1, 2], meaning that each lane must grab 2 f16 elements from
/// the inner dimension. We convert this to a optimized form by converting the
/// tensor_desc to i32 type such that lane data becomes [1, 1]. This makes the
/// later lowering easily use the load with transpose instruction.
static bool canBeOptimizedForTranspose(ArrayRef<int64_t> laneLayout,
                                       ArrayRef<int64_t> laneData) {
  if (laneLayout.size() != 2 || laneData.size() != 2)
    return false;
  if (laneLayout[0] == 1 || laneLayout[1] != 1)
    return false;
  if (laneData[0] != 1 || laneData[1] == 1)
    return false;
  return true;
}
 
/// A tensor desc type can be optimized if its element type is less than 32 bits
/// and its layout can be optimized.
static bool canBeOptimizedForTranspose(xegpu::TensorDescType tdescType) {
  // If the dtype is greater or equal to 32 bits, layout must be valid.
  int elementTyBitwidth = tdescType.getElementType().getIntOrFloatBitWidth();
  if (elementTyBitwidth >= 32)
    return false;
  auto maybeLaneLayout = getMaybeLaneLayout(tdescType);
  auto maybeLaneData = getMaybeLaneData(tdescType);
  if (!maybeLaneData || !maybeLaneLayout)
    return false;
  return canBeOptimizedForTranspose(*maybeLaneLayout, *maybeLaneData);
}
 
/// Check if a tensor desc type can be optimized for transpose, if so return the
/// new optimized tensor desc type with a valid transpose layout.
static xegpu::TensorDescType
tryOptimize(xegpu::TensorDescType tdescType,
            const xegpu::uArch::uArch *targetuArch) {
  if (!canBeOptimizedForTranspose(tdescType))
    return tdescType;
  auto laneData = getMaybeLaneData(tdescType)
                      .value(); // Lane data must exist if we reach here.
  int64_t innerLaneData = laneData[1];
  int elementTyBitwidth = tdescType.getElementType().getIntOrFloatBitWidth();
  // Required shape is total shape of the vector result that this tensor desc
  // must eventually load after adjusting for the new bitwidth and array
  // length.
  SmallVector<int64_t> requiredShape(tdescType.getShape());
  requiredShape.back() =
      requiredShape.back() * tdescType.getArrayLength() / innerLaneData;
  int newBitWidth = elementTyBitwidth * innerLaneData;
  Type newElemTy = IntegerType::get(tdescType.getContext(), newBitWidth);
  // Supported shape is the max transpose shape that can be supported by
  // hardware that is less than or equal to required shape.
  auto *blockLoadTarget =
      dyn_cast<xegpu::uArch::Subgroup2DBlockLoadInstruction>(
          targetuArch->getInstruction(
              xegpu::uArch::InstructionKind::Subgroup2DBlockLoad));
  auto maybeHWParams = blockLoadTarget->getBlockWidthHeightCount(
      newElemTy, /** has transform */ false, /** has transpose */ true);
  // If no HW params found, return the original type.
  if (!maybeHWParams)
    return tdescType;
  auto [widths, heights, counts] = maybeHWParams.value();
  // TODO: Currently we expect array length to be 1 for transpose case.
  if (counts.size() != 1 || counts[0] != 1)
    return tdescType;
  int arrayLen = counts[0];
  // The transpose acts on the innermost 2 dims; any leading dims are unit
  // batch.
  int64_t rank = requiredShape.size();
  int supportedHeight = xegpu::getLargestDivisor(
      static_cast<int>(requiredShape[rank - 2]), heights);
  int supportedWidth = xegpu::getLargestDivisor(
      static_cast<int>(requiredShape[rank - 1]), widths);
  // If no supported height or width found, return the original type.
  if (supportedHeight == -1 || supportedWidth == -1)
    return tdescType;
 
  // Preserve leading (unit) batch dims; only the inner 2 dims are reshaped.
  SmallVector<int64_t> supportedShape(requiredShape.begin(),
                                      requiredShape.end() - 2);
  supportedShape.push_back(supportedHeight);
  supportedShape.push_back(supportedWidth);
  auto ctx = tdescType.getContext();
  auto origLayout = tdescType.getLayoutAttr();
  auto laneLayoutI64 = origLayout.getEffectiveLaneLayoutAsInt();
  SmallVector<int32_t> laneLayoutI32(laneLayoutI64.begin(),
                                     laneLayoutI64.end());
 
  xegpu::LayoutAttr newLayout = xegpu::LayoutAttr::get(
      ctx, /*lane_layout=*/DenseI32ArrayAttr::get(ctx, laneLayoutI32),
      /*lane_data=*/
      DenseI32ArrayAttr::get(ctx,
                             SmallVector<int32_t>(laneLayoutI32.size(), 1)),
      /*order=*/origLayout.getOrder());
 
  // Array length can not be larger than 1 for transpose case.
  return xegpu::TensorDescType::get(supportedShape, newElemTy, arrayLen,
                                    tdescType.getBoundaryCheck(),
                                    tdescType.getMemorySpace(), newLayout);
}
 
/// Helper to convert an OpFoldResult to Value.
static Value convertToValue(ConversionPatternRewriter &rewriter, Location loc,
                            OpFoldResult ofr) {
  std::optional<int64_t> mayBeInt = getConstantIntValue(ofr);
  if (mayBeInt)
    return arith::ConstantIndexOp::create(rewriter, loc, *mayBeInt).getResult();
  return llvm::cast<Value>(ofr);
}
 
/// Helper to divide a Value by a constant integer.
static Value divideByConstant(ConversionPatternRewriter &rewriter, Location loc,
                              Value val, int64_t constant) {
  // If the constant is a power of 2, use right shift for division.
  if (llvm::isPowerOf2_64(constant)) {
    int64_t shiftAmount = llvm::Log2_64(constant);
    return arith::ShRUIOp::create(
               rewriter, loc, val,
               arith::ConstantIndexOp::create(rewriter, loc, shiftAmount)
                   .getResult())
        .getResult();
  }
  auto constantOp =
      arith::ConstantIndexOp::create(rewriter, loc, constant).getResult();
  return arith::DivUIOp::create(rewriter, loc, val, constantOp).getResult();
}
 
/// This function takes a larger register block `data` and generates multiple
/// smaller loads (size given by `newTensorDesc`) to fill in the `data` block
/// starting from `offsets`.
static Value generateLoads(ConversionPatternRewriter &rewriter,
                           TypedValue<VectorType> data,
                           SmallVector<OpFoldResult> offsets,
                           TypedValue<xegpu::TensorDescType> newTensorDesc,
                           xegpu::LoadNdOp origLoadOp) {
  Location loc = data.getLoc();
  assert(offsets.size() >= 2 && "Expecting at least 2 offsets for 2D LoadNdOp");
  int64_t rank = data.getType().getRank();
  Value offsetDim0 = convertToValue(rewriter, loc, offsets[offsets.size() - 2]);
  Value offsetDim1 = convertToValue(rewriter, loc, offsets[offsets.size() - 1]);
  SmallVector<int64_t> supportedShape(newTensorDesc.getType().getShape());
  // Compute the ratio between original shape and supported shape. We need to
  // generate loads in this ratio arrangement.
  auto shapeRatio = computeShapeRatio(data.getType().getShape(),
                                      supportedShape)
                        .value(); // `ratio` must be defined if we reach here.
  // The loop below only walks the last 2 entries of `shapeRatio`; leading
  // (batch) dims are unit, so their ratio is 1 and their offsets pass through
  // unchanged.
  int64_t suppDim0 = supportedShape[rank - 2];
  int64_t suppDim1 = supportedShape[rank - 1];
  for (int64_t h = 0; h < shapeRatio[rank - 2]; ++h) {
    for (int64_t w = 0; w < shapeRatio[rank - 1]; ++w) {
      int64_t localOffsetDim0 = h * suppDim0;
      int64_t localOffsetDim1 = w * suppDim1;
      Value loadOffsetX = arith::AddIOp::create(
          rewriter, loc, offsetDim0,
          arith::ConstantIndexOp::create(rewriter, loc, localOffsetDim0)
              .getResult());
      Value loadOffsetY = arith::AddIOp::create(
          rewriter, loc, offsetDim1,
          arith::ConstantIndexOp::create(rewriter, loc, localOffsetDim1)
              .getResult());
      // Keep the leading (batch) offsets; replace only the inner 2.
      SmallVector<OpFoldResult> loadOffsets(offsets.begin(), offsets.end());
      loadOffsets[loadOffsets.size() - 2] = loadOffsetX;
      loadOffsets[loadOffsets.size() - 1] = loadOffsetY;
      auto loadOp = xegpu::LoadNdOp::create(
          rewriter, loc,
          VectorType::get(supportedShape, data.getType().getElementType()),
          newTensorDesc, loadOffsets, origLoadOp.getPackedAttr(),
          origLoadOp.getTransposeAttr(), origLoadOp.getL1HintAttr(),
          origLoadOp.getL2HintAttr(), origLoadOp.getL3HintAttr(),
          origLoadOp.getLayoutAttr());
      // Set the layout for the loadOp.
      auto layoutAttr = newTensorDesc.getType().getLayoutAttr();
      loadOp.setAnchorLayout(layoutAttr);
      // Insert the loaded block into the right position in data (leading dims
      // at 0, inner 2 dims at the local tile offset).
      SmallVector<int64_t> insertPos(rank, 0);
      insertPos[rank - 2] = localOffsetDim0;
      insertPos[rank - 1] = localOffsetDim1;
      SmallVector<int64_t> insertStrides(rank, 1);
      auto insertOp = vector::InsertStridedSliceOp::create(
          rewriter, loc, loadOp.getResult(), data, insertPos, insertStrides);
      // InsertOp must have the same layout as newTensorDesc.
      xegpu::setTemporaryLayout(insertOp->getOpResult(0), layoutAttr);
      data = insertOp.getResult();
    }
  }
  return data;
}
 
/// Checks if a CreateNdDescOp can be optimized for transpose, if so creates a
/// new CreateNdDescOp with optimized tensor desc type. This involves extracting
/// the base pointer from the original memory source and adjusting the shape and
/// strides of the tensor desc to fit with the new optimized transpose layout.
class XeGPUCreateNdDescOpPattern final
    : public OpConversionPattern<xegpu::CreateNdDescOp> {
public:
  using OpConversionPattern<xegpu::CreateNdDescOp>::OpConversionPattern;
  LogicalResult
  matchAndRewrite(xegpu::CreateNdDescOp createNdOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto tdescTy = createNdOp.getType();
    // Get the target uArch info.
    auto chipStr = xegpu::getChipStr(createNdOp);
    // Check if the chip is supported.
    assert(chipStr &&
           (chipStr.value() == "pvc" || chipStr.value() == "bmg" ||
            chipStr.value() == "cri") &&
           "Expecting target chip to be pvc, bmg or cri for transpose "
           "optimization.");
    const auto *targetuArch = xegpu::uArch::getUArch(chipStr.value());
 
    auto convertType = tryOptimize(tdescTy, targetuArch);
    if (convertType == tdescTy)
      return failure();
    Location loc = createNdOp.getLoc();
    Value source = createNdOp.getSource();
    auto memrefType = dyn_cast<MemRefType>(source.getType());
 
    SmallVector<OpFoldResult> mixedSizes;
    SmallVector<OpFoldResult> mixedStrides;
    memref::ExtractStridedMetadataOp meta;
    if (memrefType) {
      // extract_aligned_pointer_as_index returns the base pointer without the
      // memref offset, so always extract strided metadata to fold it in.
      meta = memref::ExtractStridedMetadataOp::create(rewriter, loc, source);
      mixedSizes = meta.getConstifiedMixedSizes();
      mixedStrides = meta.getConstifiedMixedStrides();
    } else {
      mixedSizes = createNdOp.getMixedSizes();
      mixedStrides = createNdOp.getMixedStrides();
    }
 
    auto maybeConstInnerStride = getConstantIntValue(mixedStrides.back());
    // Only row-major memrefs are expected for now.
    if (!maybeConstInnerStride || *maybeConstInnerStride != 1)
      return rewriter.notifyMatchFailure(
          createNdOp, "Expecting row-major memref for transpose optimization.");
    auto optionalLaneData = getMaybeLaneData(tdescTy);
    assert(optionalLaneData && "Expected 2D lane data");
    auto laneData = optionalLaneData.value();
    int64_t innerLaneData = laneData[1];
    // Inner dimension of the shape must be adjusted based on innerLaneData.
    SmallVector<OpFoldResult> modifiedShape(mixedSizes);
    modifiedShape.back() = divideByConstant(
        rewriter, loc, convertToValue(rewriter, loc, modifiedShape.back()),
        innerLaneData);
    // Repacking to a wider element rescales every stride but the innermost.
    assert(mixedStrides.size() >= 2 &&
           "Expected at least 2 strides for CreateNdDescOp");
    SmallVector<OpFoldResult> modifiedStrides(mixedStrides);
    for (size_t i = 0; i + 1 < modifiedStrides.size(); ++i)
      modifiedStrides[i] = divideByConstant(
          rewriter, loc, convertToValue(rewriter, loc, modifiedStrides[i]),
          innerLaneData);
 
    if (memrefType) {
      // Compute base = aligned pointer + offset (in bytes).
      Value alignedPtr = memref::ExtractAlignedPointerAsIndexOp::create(
          rewriter, loc, meta.getBaseBuffer());
      Value elemBytes = arith::ConstantIndexOp::create(
          rewriter, loc, memrefType.getElementTypeBitWidth() / 8);
      Value offsetBytes =
          arith::MulIOp::create(rewriter, loc, meta.getOffset(), elemBytes);
      Value baseIdx =
          arith::AddIOp::create(rewriter, loc, alignedPtr, offsetBytes);
      source = arith::IndexCastOp::create(rewriter, loc, rewriter.getI64Type(),
                                          baseIdx);
    }
    // Create a new CreateNdDescOp with the modified shape and converted type.
    auto newCreateNdDescOp = xegpu::CreateNdDescOp::create(
        rewriter, loc, convertType, source, modifiedShape, modifiedStrides);
    rewriter.replaceOp(createNdOp, newCreateNdDescOp.getResult());
    return success();
  }
};
 
/// Checks if a LoadNdOp consumes a tensor desc type that was rewritten for
/// tranpose optimization. If so, rewrites the LoadNdOp to to align with the
/// adjusted tensor desc type. This can result in multiple LoadNdOps being
/// generated to fill in the original load shape.
class XeGPULoadNdDescOpPattern final
    : public OpConversionPattern<xegpu::LoadNdOp> {
public:
  using OpConversionPattern<xegpu::LoadNdOp>::OpConversionPattern;
  LogicalResult
  matchAndRewrite(xegpu::LoadNdOp loadNdOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto origTensorDescType = loadNdOp.getTensorDescType();
    auto adaptorType =
        cast<xegpu::TensorDescType>(adaptor.getTensorDesc().getType());
    if (adaptorType == origTensorDescType)
      return failure();
    // Offsets must be adjusted based on innerLaneData.
    auto laneData = getMaybeLaneData(loadNdOp.getTensorDescType()).value();
    int64_t innerLaneData = laneData[1];
    auto offsets = loadNdOp.getMixedOffsets();
    if (offsets.empty())
      return rewriter.notifyMatchFailure(loadNdOp,
                                         "Expecting offsets in LoadNd");
    SmallVector<OpFoldResult> modifiedOffsets(offsets);
    modifiedOffsets.back() = divideByConstant(
        rewriter, loadNdOp.getLoc(),
        convertToValue(rewriter, loadNdOp.getLoc(), modifiedOffsets.back()),
        innerLaneData);
    // Get the 2D data shape of this loadNdOp in its original type including
    // array length.
    SmallVector<int64_t> origDataShape(origTensorDescType.getShape());
    // Adjust the data shape based on innerLaneData.
    origDataShape.back() /= innerLaneData;
    // HW supported shape is the new tensor desc shape after conversion.
    SmallVector<int64_t> hwSupportedShape(adaptorType.getShape());
    VectorType origVectorType =
        VectorType::get(origDataShape, adaptorType.getElementType());
    Value data;
    // Orig data shape is 3D for the array length case.
    if (origTensorDescType.getArrayLength() > 1) {
      SmallVector<Value> arraySlices;
      for (int64_t i = 0; i < origTensorDescType.getArrayLength(); ++i) {
        Value slice = arith::ConstantOp::create(
            rewriter, loadNdOp->getLoc(), origVectorType,
            rewriter.getZeroAttr(origVectorType));
        // Increase the Y offset for each array slice.
        Value offsetY = convertToValue(rewriter, loadNdOp->getLoc(),
                                       modifiedOffsets.back());
        modifiedOffsets.back() =
            arith::AddIOp::create(
                rewriter, loadNdOp->getLoc(), offsetY,
                arith::ConstantIndexOp::create(rewriter, loadNdOp->getLoc(),
                                               i * origDataShape[1])
                    .getResult())
                .getResult();
        slice = generateLoads(
            rewriter, cast<TypedValue<VectorType>>(slice), modifiedOffsets,
            cast<TypedValue<xegpu::TensorDescType>>(adaptor.getTensorDesc()),
            loadNdOp);
        // BitCast back to original load shape without array length.
        auto bitcastType = VectorType::get(origTensorDescType.getShape(),
                                           origTensorDescType.getElementType());
        auto bitCastOp = vector::BitCastOp::create(rewriter, loadNdOp->getLoc(),
                                                   bitcastType, slice);
        // BitCastOp must have the same layout as the original loadNdOp.
        xegpu::setTemporaryLayout(bitCastOp->getOpResult(0),
                                  origTensorDescType.getLayoutAttr());
        arraySlices.push_back(bitCastOp.getResult());
      }
      rewriter.replaceOpWithMultiple(loadNdOp, {arraySlices});
      return success();
    }
    data = arith::ConstantOp::create(
        rewriter, loadNdOp->getLoc(),
        VectorType::get(origDataShape, adaptorType.getElementType()),
        rewriter.getZeroAttr(origVectorType));
    data = generateLoads(
        rewriter, cast<TypedValue<VectorType>>(data), modifiedOffsets,
        cast<TypedValue<xegpu::TensorDescType>>(adaptor.getTensorDesc()),
        loadNdOp);
    auto bitCastOp = vector::BitCastOp::create(rewriter, loadNdOp->getLoc(),
                                               loadNdOp.getType(), data);
    // BitCastOp must have the same layout as the original loadNdOp.
    xegpu::setTemporaryLayout(bitCastOp->getOpResult(0),
                              origTensorDescType.getLayoutAttr());
    rewriter.replaceOp(loadNdOp, bitCastOp);
    return success();
  }
};
 
/// Vector ExtractOp must be processed if the original tensor desc type has
/// array length greater than 1. In this case, the LoadNdOp is replaced with
/// multiple LoadNdOps for each array slice making the extraction unnecessary.
/// In this case, we simply remove the ExtractOp.
class VectorExtractOpPattern final
    : public OpConversionPattern<vector::ExtractOp> {
public:
  using OpConversionPattern<vector::ExtractOp>::OpConversionPattern;
  LogicalResult
  matchAndRewrite(vector::ExtractOp extractOp, OneToNOpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    // Check if the source of the extraction is split to multiple values.
    if (adaptor.getSource().size() == 1)
      return failure();
    auto mixedPos = extractOp.getMixedPosition();
    if (mixedPos.size() != 1)
      return failure();
    auto mayBeInt = getConstantIntValue(mixedPos[0]);
    if (!mayBeInt)
      return failure();
    rewriter.replaceOp(extractOp, adaptor.getSource()[*mayBeInt]);
    return success();
  }
};
 
/// Performs a reduction over 2 dimensions by decomposing it into two 1D
/// reductions ordered based on layout to minimize cross-lane communication.
class MultiRed2dOpPattern
    : public OpConversionPattern<vector::MultiDimReductionOp> {
  using OpConversionPattern::OpConversionPattern;
  LogicalResult
  matchAndRewrite(vector::MultiDimReductionOp reductionOp, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto sourceVecType = reductionOp.getSourceVectorType();
    if (reductionOp.getReductionDims().size() != 2)
      return rewriter.notifyMatchFailure(reductionOp, "Expected 2D reduction");
    auto resLayout = xegpu::getDistributeLayoutAttr(reductionOp.getResult());
    // Retrieve and order dims for 1D decomposition (prefer intra-lane first).
    auto dims = llvm::to_vector(reductionOp.getReductionDims());
    auto [intraLaneDim, crossLaneDim] = getReductionDimOrder(dims, resLayout);
    // Order does not matter
    if (intraLaneDim == -1 || crossLaneDim == -1) {
      intraLaneDim = dims[0];
      crossLaneDim = dims[1];
    }
    auto loc = reductionOp.getLoc();
    auto acc = reductionOp.getAcc();
 
    // The decomposition below splits the 2D reduction into an intra-lane
    // then a cross-lane 1D reduction. The natural result layout of the
    // decomposed sequence (a doubly-sliced layout) differs from the
    // original 2D reduction's result layout that the rest of the IR was
    // written/propagated against. To keep the post-peephole IR
    // self-consistent without depending on a follow-up layout
    // propagation pass, we always insert a bridge xegpu.convert_layout
    // from the natural post-decomposition layout to the original
    // reduction's result layout. Trivial bridges fold away in
    // canonicalization.
    xegpu::DistributeLayoutAttr postDecompLayout;
    if (resLayout) {
      // Derive the source vector's layout.
      xegpu::DistributeLayoutAttr srcLayoutForCvt;
      if (auto resSlice = dyn_cast_if_present<xegpu::SliceAttr>(resLayout))
        srcLayoutForCvt = resSlice.getParent();
      if (!srcLayoutForCvt)
        srcLayoutForCvt =
            xegpu::getDistributeLayoutAttr(reductionOp.getSource());
      if (srcLayoutForCvt) {
        // The natural layout of the post-decomposition reduction result
        // is a nested SliceAttr: REDUCE_1 (reduces `intraLaneDim` from
        // the source) yields `slice<src, [intraLaneDim]>`; REDUCE_2
        // then reduces `adjCrossLaneDim` from that intermediate, giving
        // `slice<slice<src, [intraLaneDim]>, [adjCrossLaneDim]>`.
        MLIRContext *ctx = reductionOp.getContext();
        int64_t adjCrossLaneDim =
            crossLaneDim > intraLaneDim ? crossLaneDim - 1 : crossLaneDim;
        auto intermediateLayout = xegpu::SliceAttr::get(
            ctx, srcLayoutForCvt, DenseI64ArrayAttr::get(ctx, {intraLaneDim}));
        postDecompLayout = xegpu::SliceAttr::get(
            ctx, intermediateLayout,
            DenseI64ArrayAttr::get(ctx, {adjCrossLaneDim}));
      }
    }
 
    SmallVector<int64_t> accShape(sourceVecType.getShape());
    accShape.erase(accShape.begin() + intraLaneDim);
    Type eTy = sourceVecType.getElementType();
    Value constNeutralVal = xegpu::createReductionNeutralValue(
        rewriter, loc, VectorType::get(accShape, eTy), reductionOp.getKind());
 
    Value intraLaneReduced = vector::MultiDimReductionOp::create(
        rewriter, loc, reductionOp.getKind(), reductionOp.getSource(),
        constNeutralVal, ArrayRef<int64_t>(intraLaneDim));
 
    // Adjust crossLaneDim after the first reduction.
    if (crossLaneDim > intraLaneDim)
      crossLaneDim -= 1;
    Value crossLaneReduced = vector::MultiDimReductionOp::create(
        rewriter, loc, reductionOp.getKind(), intraLaneReduced, acc,
        ArrayRef<int64_t>(crossLaneDim));
    assert(crossLaneReduced.getType() == reductionOp.getResult().getType() &&
           "Type mismatch");
 
    Value replacement = crossLaneReduced;
    if (resLayout && postDecompLayout) {
      // Bridge from the natural post-decomposition layout to the
      // original reduction's result layout. This preserves the contract
      // any consumer (convert_layout, anchor op, or otherwise) was
      // written against, so the rewrite is correct independent of
      // whether layout propagation runs afterwards.
      auto bridgeOp = xegpu::ConvertLayoutOp::create(
          rewriter, loc, crossLaneReduced.getType(), crossLaneReduced,
          postDecompLayout, resLayout);
      replacement = bridgeOp.getResult();
    }
 
    rewriter.replaceOp(reductionOp, replacement);
    return success();
  }
 
private:
  std::pair<int64_t, int64_t>
  getReductionDimOrder(ArrayRef<int64_t> reductionDims,
                       xegpu::DistributeLayoutAttr layout) const {
    assert(layout.isForSubgroup() && "Must know the lane layout");
    assert(reductionDims.size() == 2 && "Expected 2D reduction");
    int64_t intra = -1, cross = -1;
    xegpu::LayoutAttr layoutAttr = dyn_cast<xegpu::LayoutAttr>(layout);
    if (auto layoutSliceAttr = dyn_cast<xegpu::SliceAttr>(layout))
      layoutAttr =
          dyn_cast<xegpu::LayoutAttr>(layoutSliceAttr.flatten().getParent());
    assert(layoutAttr);
    SmallVector<int64_t> laneLayout = layoutAttr.getEffectiveLaneLayoutAsInt();
 
    assert(laneLayout.size() && "Expected a non-empty layout");
    // try to pick a dim that does not communicate
    for (auto dim : reductionDims) {
      if (laneLayout[dim] == 1)
        intra = dim;
      else
        cross = dim;
    }
    return {intra, cross};
  }
};
 
} // namespace
 
void xegpu::populateXeGPUPeepHoleOptimizerPatterns(
    RewritePatternSet &patterns) {
  patterns.add<XeGPUCreateNdDescOpPattern, XeGPULoadNdDescOpPattern,
               VectorExtractOpPattern, MultiRed2dOpPattern>(
      patterns.getContext());
}
 
namespace {
 
struct XeGPUPeepHoleOptimizerPass final
    : public xegpu::impl::XeGPUPeepHoleOptimizerBase<
          XeGPUPeepHoleOptimizerPass> {
  void runOnOperation() override {
    MLIRContext &context = getContext();
    TypeConverter converter;
    RewritePatternSet patterns(&context);
    ConversionTarget target(context);
 
    // This pass is only meant for PVC, BMG or CRI targets. If unsupported
    // target is found, exit early.
    bool isTargetSupported = false;
    getOperation()->walk([&](gpu::GPUFuncOp funcOp) {
      auto chipStr = xegpu::getChipStr(funcOp);
      if (chipStr && (chipStr.value() == "pvc" || chipStr.value() == "bmg" ||
                      chipStr.value() == "cri"))
        isTargetSupported = true;
    });
 
    if (!isTargetSupported) {
      DBGS() << "XeGPUPeepHoleOptimizerPass only supports PVC, BMG targets."
             << "\n";
      return;
    }
 
    // Run array length optimization patterns first so that subsequent transpose
    // peephole patterns operate on the array-length-optimized tensor descs.
    {
      RewritePatternSet arrayLenPatterns(&context);
      xegpu::populateXeGPUArrayLengthOptimizationPatterns(arrayLenPatterns);
      if (failed(applyPatternsGreedily(getOperation(),
                                       std::move(arrayLenPatterns)))) {
        DBGS() << "Array length optimization patterns failed.\n";
        return signalPassFailure();
      }
    }
 
    // CreateNdDescOp and LoadNdOp with optimizable tensor desc types must be
    // converted.
    target.addDynamicallyLegalOp<xegpu::CreateNdDescOp>(
        [&](xegpu::CreateNdDescOp createNdOp) {
          return !canBeOptimizedForTranspose(createNdOp.getType());
        });
    target.addDynamicallyLegalOp<xegpu::LoadNdOp>(
        [&](xegpu::LoadNdOp loadNdOp) {
          return !canBeOptimizedForTranspose(loadNdOp.getTensorDescType());
        });
    // Vector ExtractOps can have optimizable layouts if they extract from
    // LoadNdOps with array length greater than 1. These ExtractOps must be
    // converted.
    target.addDynamicallyLegalOp<vector::ExtractOp>(
        [&](vector::ExtractOp extractOp) {
          auto layout = xegpu::getTemporaryLayout(
              dyn_cast<OpResult>(extractOp.getResult()));
          if (!layout)
            return true;
          auto laneLayout = layout.getEffectiveLaneLayoutAsInt();
          auto laneData = layout.getEffectiveLaneDataAsInt();
          return !canBeOptimizedForTranspose(laneLayout, laneData);
        });
 
    target.addDynamicallyLegalOp<vector::MultiDimReductionOp>(
        [=](Operation *op) -> bool {
          auto layout = xegpu::getDistributeLayoutAttr(op->getResult(0));
          if (!layout || !layout.isForSubgroup())
            return true;
          if (auto reductionOp = dyn_cast<vector::MultiDimReductionOp>(op))
            return reductionOp.getReductionDims().size() != 2;
          return true;
        });
 
    converter.addConversion([](Type type) { return type; });
 
    target.addLegalDialect<arith::ArithDialect, memref::MemRefDialect,
                           vector::VectorDialect>();
    // xegpu.convert_layout is left untouched by this pass; mark it legal
    // so in-place updates don't trigger re-legalization failures.
    target.addLegalOp<xegpu::ConvertLayoutOp>();
    scf::populateSCFStructuralTypeConversionsAndLegality(converter, patterns,
                                                         target);
    xegpu::populateXeGPUPeepHoleOptimizerPatterns(patterns);
    if (failed(applyPartialConversion(getOperation(), target,
                                      std::move(patterns)))) {
      DBGS() << "Optimize block loads pass failed.\n";
      return signalPassFailure();
    }
 
    // Apply folding for cleaning up IR.
    MLIRContext *ctx = &getContext();
    RewritePatternSet emptyPatterns(ctx);
    (void)applyPatternsGreedily(getOperation(), std::move(emptyPatterns));
 
    xegpu::removeTemporaryLayoutAttrs(getOperation());
  }
};
 
} // namespace