//===- XeGPUSgToWiDistributeExperimental.cpp - XeGPU SG to WI Pass --------===//
//
// 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/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/Index/IR/IndexDialect.h"
#include "mlir/Dialect/Math/IR/Math.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/SCF/Transforms/Patterns.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/IntelGpuXe2.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/BuiltinOps.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/MLIRContext.h"
#include "mlir/IR/Operation.h"
#include "mlir/IR/Value.h"
#include "mlir/IR/ValueRange.h"
#include "mlir/Transforms/DialectConversion.h"
#include "llvm/ADT/SetVector.h"
#include "llvm/Support/LogicalResult.h"
#include "llvm/Support/raw_ostream.h"
#include <optional>

namespace mlir {
namespace xegpu {
#define GEN_PASS_DEF_XEGPUSGTOWIDISTRIBUTEEXPERIMENTAL
#include "mlir/Dialect/XeGPU/Transforms/Passes.h.inc"
} // namespace xegpu
} // namespace mlir

using namespace mlir;

#define DEBUG_TYPE "xegpu-sg-to-wi-distribute-experimental"
#define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "]: ")

namespace {

/// Casts the given vector value `v` to the expected vector type `expectedTy`.
static Value castValueTo(ConversionPatternRewriter &rewriter,
                         TypedValue<VectorType> v, VectorType expectedTy) {
  // If the type matches, simply return the value itself.
  if (v.getType() == expectedTy)
    return v;
  // If only shape differs, use shape cast.
  if (isa<VectorType>(v.getType()) &&
      v.getType().getNumElements() == expectedTy.getNumElements())
    return vector::ShapeCastOp::create(rewriter, v.getLoc(), expectedTy, v);

  // Else create an unrealized cast.
  auto newOp = UnrealizedConversionCastOp::create(rewriter, v.getLoc(),
                                                  expectedTy, ValueRange{v});
  return newOp.getResult(0);
}

/// Checks if all XeGPU anchor ops and vector results have valid layouts.
static LogicalResult verifyLayouts(Operation *root) {
  auto walkResult = root->walk([&](Operation *nestedOp) -> WalkResult {
    if (auto anchorOp = dyn_cast<xegpu::AnchorLayoutInterface>(nestedOp)) {
      auto layout = anchorOp.getAnchorLayout();
      if (!layout) {
        nestedOp->emitError("expected anchor layout attribute on operation");
        return WalkResult::interrupt();
      }
      return WalkResult::advance();
    }
    // For each vector result, check if the op contains a result layout
    // attribute.
    for (OpResult result : nestedOp->getResults()) {
      if (isa<VectorType>(result.getType())) {
        auto layout = xegpu::getDistributeLayoutAttr(result);
        if (!layout) {
          nestedOp->emitError(
              "expected result layout attribute on vector result");
          return WalkResult::interrupt();
        }
      }
    }
    return WalkResult::advance();
  });
  return walkResult.wasInterrupted() ? failure() : success();
}

/// A vector::MultiDimReductionOp at subgroup level in expected form if, it has
/// exactly 1 reduction dimension, it had valid result layout attribute, and
/// result type can be distributed to lanes using the layout.
static bool isValidSubgroupMultiReductionOp(vector::MultiDimReductionOp op) {
  auto resLayout = xegpu::getTemporaryLayout(op->getOpResult(0));
  // If no layout, not valid.
  if (!resLayout || !resLayout.isForSubgroup())
    return false;
  VectorType resTy = dyn_cast<VectorType>(op.getType());
  if (!resTy)
    return false;
  // Compute the distributed result vector type based on the layout.
  FailureOr<VectorType> resDistTypeOrFailure =
      getDistVecTypeBasedOnLaneLayout(resLayout, resTy);
  if (failed(resDistTypeOrFailure))
    return false;
  return op.getReductionDims().size() == 1;
}

/// A vector::MultiDimReductionOp is doing lane-local reduction if each workitem
/// is doing its own local reduction. In this case the result layout ensures
/// that result vector is distributed to lanes, i.e. the result vector type is
/// different from the distributed result vector type.
static bool isReductionLaneLocal(vector::MultiDimReductionOp op) {
  // Must be valid MultiDimReductionOp.
  assert(isValidSubgroupMultiReductionOp(op) && "Expecting a valid subgroup "
                                                "MultiDimReductionOp");
  auto resLayout = xegpu::getTemporaryLayout(op->getOpResult(0));
  VectorType resTy = dyn_cast<VectorType>(op.getType());
  auto resDistTypeOrFailure = getDistVecTypeBasedOnLaneLayout(resLayout, resTy);
  return resTy != resDistTypeOrFailure.value();
}

/// Given a vector type and its distributed vector type, return the list of
/// dimensions that are distributed.
static SmallVector<int64_t> getDistributedDims(VectorType originalType,
                                               VectorType distributedType) {
  assert(originalType.getRank() == distributedType.getRank() &&
         "original and distributed vector types must have the same rank");
  SmallVector<int64_t> distributedDims;
  for (int64_t i = 0; i < originalType.getRank(); ++i) {
    if (distributedType.getDimSize(i) != originalType.getDimSize(i))
      distributedDims.push_back(i);
  }
  return distributedDims;
}

/// Distributes a subgroup-level CreateNdDesc op to workitem-level CreateNdDesc
/// op. This simply drops the layout attribute from the tensor descriptor type.
struct SgToWiCreateNdDesc : public OpConversionPattern<xegpu::CreateNdDescOp> {
  using OpConversionPattern<xegpu::CreateNdDescOp>::OpConversionPattern;

  LogicalResult
  matchAndRewrite(xegpu::CreateNdDescOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    xegpu::TensorDescType resultType = op.getType();
    // If no layout, nothing to do.
    if (!resultType.getLayout())
      return failure();

    auto newOp = xegpu::CreateNdDescOp::create(
        rewriter, op.getLoc(), resultType.dropLayouts(), op.getOperands(),
        op->getAttrs());
    rewriter.replaceOp(op, newOp.getResult());
    return success();
  }
};

/// Distributes a subgroup-level LoadNd op to workitem-level LoadNd op. Output
/// of workitem-level LoadNd op is 1D. ShapeCast is added to restore the
/// original rank.
struct SgToWiLoadNd : public OpConversionPattern<xegpu::LoadNdOp> {
  using OpConversionPattern<xegpu::LoadNdOp>::OpConversionPattern;

  LogicalResult
  matchAndRewrite(xegpu::LoadNdOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    xegpu::DistributeLayoutAttr layout = op.getAnchorLayout();
    // If no layout, nothing to do.
    if (!layout)
      return failure();
    // Check if the layout attached to the tensor descriptor is same as the
    // anchor layout. Otherwise, this is a conflict.
    if (op.getTensorDescType().getLayout() != layout)
      return rewriter.notifyMatchFailure(
          op, "conflicting layout attributes on tensor descriptor and anchor");
    auto uArch = getUArch(xegpu::getChipStr(op).value_or(""));
    if (!uArch)
      return rewriter.notifyMatchFailure(
          op, "xegpu::LoadNdOp require target attribute attached to "
              "determine transpose "
              "requirement");
    auto supportedWiResultTyOrFailure =
        xegpu::getDistributedVectorType(op.getTensorDescType());
    auto expectedWiResultTyOrFailure =
        xegpu::getDistVecTypeBasedOnLaneLayout(layout, op.getType());
    if (failed(supportedWiResultTyOrFailure))
      return rewriter.notifyMatchFailure(
          op, "unable to compute the workitem vector type for LoadNdOp");
    if (failed(expectedWiResultTyOrFailure))
      return rewriter.notifyMatchFailure(
          op,
          "unable to compute expected workitem vector type from lane layout");
    auto newOp = xegpu::LoadNdOp::create(
        rewriter, op.getLoc(), supportedWiResultTyOrFailure.value(),
        adaptor.getTensorDesc(), op.getMixedOffsets(), op.getPackedAttr(),
        op.getTransposeAttr(), op.getL1HintAttr(), op.getL2HintAttr(),
        op.getL3HintAttr(), /**layout**/ nullptr);
    // Set the packed attribute if the layout requires it.
    newOp.setPacked(xegpu::requirePacked(cast<xegpu::LayoutAttr>(layout)));
    // Set the transpose attribute if the layout requires it.
    if (xegpu::requireTranspose(cast<xegpu::LayoutAttr>(layout), uArch))
      newOp.setTranspose(DenseI64ArrayAttr::get(rewriter.getContext(), {1, 0}));
    rewriter.replaceOp(op, castValueTo(rewriter, newOp.getResult(),
                                       expectedWiResultTyOrFailure.value()));
    return success();
  }
};

/// Distributes a subgroup-level StoreNd op to workitem-level StoreNd op. Stored
/// value in workitem-level StoreNd op is 1D. ShapeCast is added to cast the
/// incoming value to 1D.
struct SgToWiStoreNd : public OpConversionPattern<xegpu::StoreNdOp> {
  using OpConversionPattern<xegpu::StoreNdOp>::OpConversionPattern;

  LogicalResult
  matchAndRewrite(xegpu::StoreNdOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    xegpu::DistributeLayoutAttr layout = op.getAnchorLayout();
    // If no layout, nothing to do.
    if (!layout)
      return failure();
    // Check if the layout attached to the tensor descriptor and value layout is
    // same as the anchor layout. Otherwise, this is a conflict.
    if (op.getTensorDescType().getLayout() != layout)
      return rewriter.notifyMatchFailure(
          op, "conflicting layout attributes on tensor descriptor and anchor");
    auto valueLayout = xegpu::getDistributeLayoutAttr(op->getOpOperand(0));
    if (valueLayout != layout)
      return rewriter.notifyMatchFailure(
          op, "conflicting layout attributes on value and anchor");
    auto supportedWiValueTyOrFailure =
        xegpu::getDistributedVectorType(op.getTensorDescType());
    if (failed(supportedWiValueTyOrFailure))
      return rewriter.notifyMatchFailure(
          op,
          "unable to compute wi vector type for StoreNdOp value from tensor "
          "descriptor");

    xegpu::StoreNdOp::create(
        rewriter, op.getLoc(),
        castValueTo(rewriter, cast<TypedValue<VectorType>>(adaptor.getValue()),
                    supportedWiValueTyOrFailure.value()),
        adaptor.getTensorDesc(), op.getMixedOffsets(), op.getL1HintAttr(),
        op.getL2HintAttr(), op.getL3HintAttr(), /**layout**/ nullptr);
    rewriter.eraseOp(op);
    return success();
  }
};

/// Distributes a subgroup-level Dpas op to workitem-level Dpas op. All inpputs
/// and output of workitem-level Dpas op are 1D. Necessary casts are added to
/// convert the inputs and output to/from 1D.
struct SgToWiDpas : public OpConversionPattern<xegpu::DpasOp> {
  using OpConversionPattern<xegpu::DpasOp>::OpConversionPattern;

  LogicalResult
  matchAndRewrite(xegpu::DpasOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    // llvm::errs() << "DpasOpPattern matchAndRewrite called\n";
    // Check if the op has A, B and CD layouts attached.
    auto layoutA = cast<xegpu::LayoutAttr>(op.getLayoutAAttr());
    auto layoutB = cast<xegpu::LayoutAttr>(op.getLayoutBAttr());
    auto layoutCd = cast<xegpu::LayoutAttr>(op.getLayoutCdAttr());
    if (!layoutA || !layoutB || !layoutCd)
      return failure();
    // llvm::errs() << "tryning to calculate wi types for dpas op\n";
    auto wiResultTyOrFailure =
        xegpu::getDistributedVectorType(op.getType(), layoutCd);
    auto wiATypeOrFailure =
        xegpu::getDistributedVectorType(op.getLhs().getType(), layoutA);
    auto wiBTypeOrFailure =
        xegpu::getDistributedVectorType(op.getRhs().getType(), layoutB);
    auto expectedWiResultTyOrFailure =
        xegpu::getDistVecTypeBasedOnLaneLayout(layoutCd, op.getType());
    if (failed(wiResultTyOrFailure) || failed(wiATypeOrFailure) ||
        failed(wiBTypeOrFailure))
      return rewriter.notifyMatchFailure(
          op, "failed to calculate supported workitem vector types for DpasOp "
              "from layouts");
    if (failed(expectedWiResultTyOrFailure))
      return rewriter.notifyMatchFailure(
          op, "unable to compute expected workitem vector type for DpasOp from "
              "lane layout");
    auto newOp = xegpu::DpasOp::create(
        rewriter, op->getLoc(), wiResultTyOrFailure.value(),
        castValueTo(rewriter, cast<TypedValue<VectorType>>(adaptor.getLhs()),
                    wiATypeOrFailure.value()),
        castValueTo(rewriter, cast<TypedValue<VectorType>>(adaptor.getRhs()),
                    wiBTypeOrFailure.value()),
        castValueTo(rewriter, cast<TypedValue<VectorType>>(adaptor.getAcc()),
                    wiResultTyOrFailure.value()),
        /** layoutA**/ nullptr,
        /** layoutB**/ nullptr, /** layoutCd**/ nullptr);
    // Explicitly set the new types to enable correct type materializations.
    rewriter.replaceOp(op, castValueTo(rewriter, newOp.getResult(),
                                       expectedWiResultTyOrFailure.value()));
    return success();
  }
};

/// Distributes elementwise ops to workitem-level elementwise ops. This
/// currently handles elementwise ops with single result only.
struct SgToWiElementWise : public ConversionPattern {
  SgToWiElementWise(TypeConverter &typeConverter, MLIRContext *ctx)
      : ConversionPattern(MatchAnyOpTypeTag(), /*benefit=*/1, ctx) {}

  LogicalResult
  matchAndRewrite(Operation *op, ArrayRef<Value> operands,
                  ConversionPatternRewriter &rewriter) const override {
    // Only match ops with elementwise trait and single result.
    if (!OpTrait::hasElementwiseMappableTraits(op) || op->getNumResults() != 1)
      return failure();

    auto resultType = dyn_cast<VectorType>(op->getResult(0).getType());
    if (!resultType)
      return rewriter.notifyMatchFailure(
          op, "operation result is not a vector type");

    xegpu::DistributeLayoutAttr layout =
        xegpu::getTemporaryLayout(llvm::cast<OpResult>(op->getResult(0)));
    if (!layout || !layout.isForSubgroup())
      return rewriter.notifyMatchFailure(
          op, "operation result does not have subgroup distribute layout");

    auto wiShapeOrFailure =
        xegpu::getDistVecTypeBasedOnLaneLayout(layout, resultType);

    if (failed(wiShapeOrFailure))
      return rewriter.notifyMatchFailure(
          op, "unable to compute workitem vector type from the layout");

    VectorType newResultType = wiShapeOrFailure.value();
    OperationState state(op->getLoc(), op->getName());
    state.addOperands(operands);
    state.addTypes(newResultType);
    // Copy all attributes except for DistributeLayoutAttr.
    for (auto attr : op->getAttrs()) {
      if (!isa<xegpu::DistributeLayoutAttr>(attr.getValue()))
        state.addAttribute(attr.getName(), attr.getValue());
    }
    Operation *newOp = rewriter.create(state);

    rewriter.replaceOp(op, newOp->getResult(0));
    return success();
  }
};

/// Distributes a subgroup-level arith ConstantOp to workitem-level arith
/// ConstantOp.
struct SgToWiArithConstant : public OpConversionPattern<arith::ConstantOp> {
  using OpConversionPattern<arith::ConstantOp>::OpConversionPattern;

  LogicalResult
  matchAndRewrite(arith::ConstantOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto resultType = dyn_cast<VectorType>(op.getType());
    if (!resultType)
      return failure();

    // Only handle dense vector constants
    auto dense = dyn_cast<SplatElementsAttr>(op.getValue());
    if (!dense)
      return rewriter.notifyMatchFailure(
          op, "only dense splat vector constants are supported");

    xegpu::DistributeLayoutAttr layout =
        xegpu::getTemporaryLayout(llvm::cast<OpResult>(op.getResult()));
    if (!layout || !layout.isForSubgroup())
      return rewriter.notifyMatchFailure(
          op, "operation result does not have subgroup distribute layout");

    auto wiShapeOrFailure =
        xegpu::getDistVecTypeBasedOnLaneLayout(layout, resultType);

    if (failed(wiShapeOrFailure))
      return rewriter.notifyMatchFailure(
          op, "unable to compute workitem vector type from the layout");

    VectorType newResultType = wiShapeOrFailure.value();
    auto sclarValue = dense.getSplatValue<Attribute>();
    auto newDenseAttr = DenseElementsAttr::get(newResultType, sclarValue);

    auto newOp = arith::ConstantOp::create(rewriter, op.getLoc(), newResultType,
                                           newDenseAttr);
    rewriter.replaceOp(op, newOp.getResult());
    return success();
  }
};

/// Distributes a subgroup-level PrefetchNd op to workitem-level PrefetchNd op.
struct SgToWiPrefetchNd : public OpConversionPattern<xegpu::PrefetchNdOp> {
  using OpConversionPattern<xegpu::PrefetchNdOp>::OpConversionPattern;

  LogicalResult
  matchAndRewrite(xegpu::PrefetchNdOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    xegpu::DistributeLayoutAttr layout = op.getAnchorLayout();
    // If no layout, nothing to do.
    if (!layout)
      return failure();

    xegpu::PrefetchNdOp::create(rewriter, op.getLoc(), adaptor.getTensorDesc(),
                                op.getMixedOffsets(), op.getL1HintAttr(),
                                op.getL2HintAttr(), op.getL3HintAttr(),
                                /**layout**/ nullptr);
    rewriter.eraseOp(op);
    return success();
  }
};

/// Distributes a subgroup-level LoadGather (xegpu.load) op to workitem-level.
///
/// Example 1 (1D, no chunk size):
///   layout = #xegpu.layout<lane_layout = [16], lane_data = [1]>
///   %mask = producer_op : vector<16xi1>
///   %offset = producer_op : vector<16xindex>
///   %0 = xegpu.load %src[%offset], %mask : memref<256xf16>,
///     vector<16xindex>, vector<16xi1> -> vector<16xf16>
/// Distributed to:
///   %mask = producer_op : vector<1xi1>
///   %offset = producer_op : vector<1xindex>
///   %0 = xegpu.load %src[%offset], %mask : memref<256xf16>,
///     vector<1xindex>, vector<1xi1> -> vector<1xf16>
///
/// Example 2 (2D with chunk size, same mask & offset):
///   layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>
///   %0 = xegpu.load %src[%offset], %mask <{chunk_size=8}> :
///     memref<256xf16>, vector<16xindex>, vector<16xi1> -> vector<16x8xf16>
/// Distributed to:
///   %0 = xegpu.load %src[%offset], %mask <{chunk_size=8}> :
///     memref<256xf16>, vector<1xindex>, vector<1xi1> -> vector<8xf16>
///
/// Example 3 (3D with leading unit dims):
///   layout = #xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 1, 1]>
///   %mask = producer_op : vector<1x1x16xi1>
///   %offset = producer_op : vector<1x1x16xindex>
///   %0 = xegpu.load %src[%offset], %mask : memref<256xf16>,
///     vector<1x1x16xindex>, vector<1x1x16xi1> -> vector<1x1x16xf16>
/// Distributed to:
///   %mask = producer_op : vector<1x1x1xi1>
///   %offset = producer_op : vector<1x1x1xindex>
///   %0 = xegpu.load %src[%offset], %mask : memref<256xf16>,
///     vector<1xindex>, vector<1xi1> -> vector<1xf16>
struct SgToWiLoadGather : public OpConversionPattern<xegpu::LoadGatherOp> {
  using OpConversionPattern<xegpu::LoadGatherOp>::OpConversionPattern;

  LogicalResult
  matchAndRewrite(xegpu::LoadGatherOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    xegpu::DistributeLayoutAttr layout = op.getAnchorLayout();
    if (!layout)
      return failure();

    VectorType origResultTy = op.getValueType();
    if (!origResultTy)
      return failure();

    // Check that leading dimensions are unit.
    int chunkSize = op.getChunkSize().value_or(1);
    int effectiveVecRank = (chunkSize == 1) ? 1 : 2;
    ArrayRef<int64_t> shape = origResultTy.getShape();
    if (llvm::any_of(
            shape.take_front(origResultTy.getRank() - effectiveVecRank),
            [](int64_t d) { return d != 1; }))
      return rewriter.notifyMatchFailure(
          op, "Only unit dimensions allowed for the leading "
              "dimensions of the load vector!");

    auto distResultTyOrFailure =
        xegpu::getDistVecTypeBasedOnLaneLayout(layout, origResultTy);
    if (failed(distResultTyOrFailure))
      return rewriter.notifyMatchFailure(
          op,
          "unable to compute expected workitem vector type from lane layout");

    VectorType distResultTy = distResultTyOrFailure.value();
    VectorType distResultTy1D = VectorType::get({distResultTy.getNumElements()},
                                                distResultTy.getElementType());

    // Flatten offsets and mask to 1D to match the 1D result type.
    Value distOffsets = adaptor.getOffsets();
    auto distOffsetsTy = cast<VectorType>(distOffsets.getType());
    VectorType offsetsTy1D = VectorType::get({distOffsetsTy.getNumElements()},
                                             distOffsetsTy.getElementType());
    distOffsets = castValueTo(
        rewriter, cast<TypedValue<VectorType>>(distOffsets), offsetsTy1D);

    Value distMask = adaptor.getMask();
    auto distMaskTy = cast<VectorType>(distMask.getType());
    VectorType maskTy1D = VectorType::get({distMaskTy.getNumElements()},
                                          distMaskTy.getElementType());
    distMask =
        castValueTo(rewriter, cast<TypedValue<VectorType>>(distMask), maskTy1D);

    Value distSource = adaptor.getSource();
    auto newOp = xegpu::LoadGatherOp::create(
        rewriter, op.getLoc(), distResultTy1D, distSource, distOffsets,
        distMask, op.getChunkSizeAttr(), op.getL1HintAttr(), op.getL2HintAttr(),
        op.getL3HintAttr(), /*layout=*/nullptr);

    Value result = newOp->getResult(0);
    if (distResultTy1D != distResultTy)
      result = castValueTo(rewriter, cast<TypedValue<VectorType>>(result),
                           distResultTy);
    rewriter.replaceOp(op, result);
    return success();
  }
};

/// This pattern distributes a subgroup-level vector.reduction op to
/// workitem-level. This require shuffling the data across the workitems (using
/// gpu::ShuffleOp) and reducing in stages until all workitems have the final
/// result.
struct SgToWiVectorReduction : public OpConversionPattern<vector::ReductionOp> {
  using OpConversionPattern<vector::ReductionOp>::OpConversionPattern;

  LogicalResult
  matchAndRewrite(vector::ReductionOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto layout = xegpu::getDistributeLayoutAttr(op.getVector());

    // If no layout, nothing to do.
    if (!layout || !layout.isForSubgroup())
      return failure();

    VectorType srcVecType = op.getSourceVectorType();
    // Only rank 1 vectors supported.
    if (srcVecType.getRank() != 1)
      return rewriter.notifyMatchFailure(
          op, "Only rank 1 reductions can be distributed.");
    // Lane layout must have the same rank as the vector.
    if (layout.getRank() != srcVecType.getRank())
      return rewriter.notifyMatchFailure(
          op, "Layout rank does not match vector rank.");

    // Get the subgroup size from the layout.
    int64_t sgSize = layout.getEffectiveLaneLayoutAsInt()[0];
    const uArch *uArch = getUArch(xegpu::getChipStr(op).value_or(""));
    if (!uArch)
      return rewriter.notifyMatchFailure(
          op, "xegpu::ReductionOp require target attribute attached to "
              "determine subgroup size");

    // Only subgroup-sized vectors supported.
    if (sgSize != uArch->getSubgroupSize() ||
        srcVecType.getShape()[0] % sgSize != 0)
      return rewriter.notifyMatchFailure(op,
                                         "Invalid layout or reduction vector "
                                         "dimension must match subgroup size.");

    if (!op.getType().isIntOrFloat())
      return rewriter.notifyMatchFailure(
          op, "Reduction distribution currently only supports floats and "
              "integer types.");

    // Get the distributed vector (per work-item portion).
    Value laneValVec = adaptor.getVector();

    // Distribute and reduce across work-items in the subgroup.
    Value fullReduce = xegpu::subgroupReduction(
        op.getLoc(), rewriter, laneValVec, op.getKind(), sgSize);

    // If there's an accumulator, combine it with the reduced value.
    if (adaptor.getAcc())
      fullReduce = vector::makeArithReduction(
          rewriter, op.getLoc(), op.getKind(), fullReduce, adaptor.getAcc());

    rewriter.replaceOp(op, fullReduce);
    return success();
  }
};

/// This pattern distributes a subgroup-level vector.multi_reduction op to
/// workitem-level only if the reduction is lane-local. This means that
/// reduction dimension is not distributed to lanes and each lane does its own
/// local reduction.
struct SgToWiMultiDimReduction
    : public OpConversionPattern<vector::MultiDimReductionOp> {
  using OpConversionPattern<vector::MultiDimReductionOp>::OpConversionPattern;

  LogicalResult
  matchAndRewrite(vector::MultiDimReductionOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    Value result;
    ArrayRef<int64_t> reductionDims = op.getReductionDims();
    assert(reductionDims.size() == 1 &&
           "Expecting single reduction dimension for subgroup multi "
           "reduction op");
    // For rank > 2, ensure leading dimensions are unit.
    VectorType sourceType = op.getSourceVectorType();
    int64_t rank = sourceType.getRank();
    if (rank > 2) {
      ArrayRef<int64_t> shape = sourceType.getShape();
      if (llvm::any_of(shape.take_front(rank - 2),
                       [](int64_t d) { return d != 1; }))
        return rewriter.notifyMatchFailure(
            op, "only unit leading dimensions are supported for "
                "multi_reduction with rank > 2");
    }
    if (isReductionLaneLocal(op)) {
      auto resLayout = xegpu::getTemporaryLayout(op->getOpResult(0));
      VectorType resVecTy = dyn_cast<VectorType>(op.getType());
      auto resDistVecTyOrFailure =
          getDistVecTypeBasedOnLaneLayout(resLayout, resVecTy);
      // For lane local reduction, simply create a new MultiDimReductionOp using
      // adaptor operands and the new result type.
      result = vector::MultiDimReductionOp::create(
          rewriter, op.getLoc(), resDistVecTyOrFailure.value(), op.getKind(),
          adaptor.getSource(), adaptor.getAcc(), op.getReductionDims());
    } else {
      auto reductionDim = reductionDims[0];
      VectorType sourceType = op.getSourceVectorType();
      int64_t reductionDimSize = sourceType.getShape()[reductionDim];
      result = xegpu::lowerCrossLaneReductionToShuffles(
          cast<TypedValue<VectorType>>(adaptor.getSource()),
          cast<TypedValue<VectorType>>(adaptor.getAcc()), op.getKind(),
          reductionDim, reductionDimSize, op.getLoc(), rewriter);
    }
    rewriter.replaceOp(op, result);
    return success();
  }
};

/// Helper to compute distributed coordinates for matrix ops.
/// When not using subgroup_block_io, each workitem computes its own
/// coordinates based on the layout and lane ID.
static SmallVector<Value> computeDistributedCoordsForMatrixOp(
    ConversionPatternRewriter &rewriter, Location loc,
    xegpu::DistributeLayoutAttr layout, ArrayRef<int64_t> payloadShape,
    ValueRange origOffsets) {
  Value laneId = gpu::LaneIdOp::create(rewriter, loc, rewriter.getIndexType(),
                                       /*upperBound=*/mlir::IntegerAttr());
  auto maybeCoords =
      layout.computeDistributedCoords(rewriter, loc, laneId, payloadShape);
  if (failed(maybeCoords))
    return {};
  assert(maybeCoords.value().size() == 1 &&
         "Expected one set of distributed offsets");
  SmallVector<OpFoldResult> ofrVec = xegpu::addWithRightAligned(
      rewriter, loc, getAsOpFoldResult(maybeCoords.value()[0]),
      getAsOpFoldResult(origOffsets));
  return llvm::map_to_vector(ofrVec, llvm::CastTo<Value>);
}

/// This pattern distributes a subgroup-level LoadMatrix op to workitem-level.
struct SgToWiLoadMatrix : public OpConversionPattern<xegpu::LoadMatrixOp> {
  using OpConversionPattern<xegpu::LoadMatrixOp>::OpConversionPattern;

  LogicalResult
  matchAndRewrite(xegpu::LoadMatrixOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto layout = op.getLayoutAttr();
    // If no layout, nothing to do.
    if (!layout)
      return failure();

    VectorType sgPayloadTy = dyn_cast<VectorType>(op.getResult().getType());
    if (!sgPayloadTy)
      return rewriter.notifyMatchFailure(
          op, "the matrix op payload must be a vector type");

    auto loc = op.getLoc();
    auto offsets = op.getMixedOffsets();
    if (offsets.empty())
      return rewriter.notifyMatchFailure(op, "the load op must have offsets");

    FailureOr<VectorType> distPayloadTyOrFailure =
        getDistVecTypeBasedOnLaneLayout(layout, sgPayloadTy);
    if (failed(distPayloadTyOrFailure))
      return rewriter.notifyMatchFailure(
          op, "Failed to distribute matrix op payload based on layout.");

    SmallVector<Value> offsetsAsValues =
        vector::getAsValues(rewriter, loc, offsets);

    SmallVector<Value> newCoords = offsetsAsValues;
    if (!op.getSubgroupBlockIoAttr()) {
      newCoords = computeDistributedCoordsForMatrixOp(
          rewriter, loc, layout, sgPayloadTy.getShape(), offsetsAsValues);
      if (newCoords.empty())
        return rewriter.notifyMatchFailure(
            op, "Failed to compute distributed coordinates.");
    }

    SmallVector<int64_t> newConstOffsets(op.getConstOffsets().size(),
                                         ShapedType::kDynamic);
    DenseI64ArrayAttr newConstOffsetsAttr =
        rewriter.getDenseI64ArrayAttr(newConstOffsets);

    auto newOp = xegpu::LoadMatrixOp::create(
        rewriter, loc, *distPayloadTyOrFailure, adaptor.getMemDesc(),
        ValueRange(newCoords), newConstOffsetsAttr, op.getSubgroupBlockIoAttr(),
        xegpu::DistributeLayoutAttr{});
    rewriter.replaceOp(op, newOp.getResult());
    return success();
  }
};

/// Distributes a subgroup-level vector.transpose op to workitem-level.
struct SgToWiVectorTranspose : public OpConversionPattern<vector::TransposeOp> {
  using OpConversionPattern<vector::TransposeOp>::OpConversionPattern;

  LogicalResult
  matchAndRewrite(vector::TransposeOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    xegpu::DistributeLayoutAttr sourceLayout =
        xegpu::getTemporaryLayout(op->getOpOperand(0));
    xegpu::DistributeLayoutAttr resultLayout =
        xegpu::getTemporaryLayout(op->getOpResult(0));
    if (!sourceLayout || !resultLayout)
      return rewriter.notifyMatchFailure(
          op, "the source or result vector of the transpose op lacks layout "
              "attribute");
    ArrayRef<int64_t> perm = op.getPermutation();
    // Result layout must be a transpose of source layout.
    if (!resultLayout.isTransposeOf(sourceLayout, perm,
                                    xegpu::LayoutKind::Lane))
      return rewriter.notifyMatchFailure(
          op, "the source or result vector layouts must be transposes of "
              "each other");
    FailureOr<VectorType> distributedResultTypeOrFailure =
        getDistVecTypeBasedOnLaneLayout(resultLayout, op.getResultVectorType());
    if (failed(distributedResultTypeOrFailure))
      return rewriter.notifyMatchFailure(
          op, "Failed to distribute the result vector type in "
              "vector::Transpose op");
    auto newOp = vector::TransposeOp::create(rewriter, op.getLoc(),
                                             adaptor.getVector(), perm);
    rewriter.replaceOp(op, castValueTo(rewriter, newOp.getResult(),
                                       distributedResultTypeOrFailure.value()));
    return success();
  }
};

/// Distributes a subgroup-level vector.bitcast op to workitem-level.
/// Bitcast only impacts the innermost dimension of the source/result vectors.
struct SgToWiVectorBitcast : public OpConversionPattern<vector::BitCastOp> {
  using OpConversionPattern<vector::BitCastOp>::OpConversionPattern;

  LogicalResult
  matchAndRewrite(vector::BitCastOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    xegpu::DistributeLayoutAttr resultLayout =
        xegpu::getTemporaryLayout(op->getOpResult(0));
    if (!resultLayout)
      return rewriter.notifyMatchFailure(
          op, "result vector of the bitcast op lacks layout attribute");
    FailureOr<VectorType> distributedResultTypeOrFailure =
        getDistVecTypeBasedOnLaneLayout(resultLayout, op.getResultVectorType());
    if (failed(distributedResultTypeOrFailure))
      return rewriter.notifyMatchFailure(
          op, "Failed to distribute the result vector type in "
              "vector::BitCast op");
    auto newOp = vector::BitCastOp::create(
        rewriter, op.getLoc(), distributedResultTypeOrFailure.value(),
        adaptor.getSource());
    rewriter.replaceOp(op, newOp.getResult());
    return success();
  }
};

/// Distributes a subgroup-level vector.create_mask or vector.constant_mask op
/// to workitem-level. Uses `computeDistributedCoords()` to obtain the
/// coordinates each workitem owns, then compares each coordinate against the
/// original mask bounds using `arith.cmpi slt`. The per-element boolean
/// results are assembled into the distributed mask vector.
///
/// For multi-dimensional masks, the element is in-bounds when ALL dimensions
/// satisfy `coord[i] < bound[i]`.
///
/// Example (1D):
///   layout = #xegpu.layout<lane_layout = [16], lane_data = [1]>
///   %mask = vector.create_mask %m0 : vector<16xi1>
/// For lane k, computeDistributedCoords gives coord = [k], so:
///   %in_bounds = arith.cmpi slt, %coord, %m0  →  i1
///   %mask = vector.broadcast %in_bounds : i1 to vector<1xi1>
///
/// Example (2D):
///   layout = #xegpu.layout<lane_layout = [8, 2], lane_data = [1, 1]>
///   %mask = vector.create_mask %m0, %m1 : vector<8x4xi1>
/// Each WI owns a 1x2 slice. computeDistributedCoords returns 2 coords:
///   [[r0, c0], [r0, c1]]
/// For each coord: in_bounds = (r < m0) && (c < m1)
///   %mask = vector.from_elements %bit0, %bit1 : vector<1x2xi1>
template <typename OpType,
          typename = std::enable_if_t<llvm::is_one_of<
              OpType, vector::CreateMaskOp, vector::ConstantMaskOp>::value>>
struct SgToWiCreateMask : public OpConversionPattern<OpType> {
  using OpConversionPattern<OpType>::OpConversionPattern;

  LogicalResult
  matchAndRewrite(OpType op, typename OpType::Adaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    xegpu::DistributeLayoutAttr layout =
        xegpu::getTemporaryLayout(op->getOpResult(0));
    if (!layout || !layout.isForSubgroup())
      return rewriter.notifyMatchFailure(
          op, "operation result does not have subgroup distribute layout");

    VectorType origType = op.getType();
    FailureOr<VectorType> distTypeOrFailure =
        getDistVecTypeBasedOnLaneLayout(layout, origType);
    if (failed(distTypeOrFailure))
      return rewriter.notifyMatchFailure(
          op, "unable to compute workitem vector type from the layout");

    VectorType distType = distTypeOrFailure.value();
    Location loc = op.getLoc();

    // Materialize the original mask bounds as Values.
    SmallVector<Value> origBounds;
    if constexpr (std::is_same_v<OpType, vector::CreateMaskOp>) {
      origBounds.append(op.getOperands().begin(), op.getOperands().end());
    } else {
      auto dimSizes = op.getMaskDimSizesAttr().asArrayRef();
      for (auto dimSize : dimSizes)
        origBounds.push_back(
            arith::ConstantIndexOp::create(rewriter, loc, dimSize).getResult());
    }

    ArrayRef<int64_t> origShape = origType.getShape();

    // Use computeDistributedCoords to get the coordinates each WI owns.
    Value laneId = gpu::LaneIdOp::create(rewriter, loc, rewriter.getIndexType(),
                                         /*upperBound=*/mlir::IntegerAttr());
    auto maybeCoordsVec =
        layout.computeDistributedCoords(rewriter, loc, laneId, origShape);
    if (failed(maybeCoordsVec))
      return rewriter.notifyMatchFailure(
          op, "failed to compute distributed coordinates from layout");

    SmallVector<SmallVector<Value>> coordsVec = maybeCoordsVec.value();
    int64_t numElements = distType.getNumElements();
    assert(static_cast<int64_t>(coordsVec.size()) == numElements &&
           "number of coordinate sets must match number of distributed "
           "elements");

    // For each element, compare all coordinates against bounds.
    Value trueVal =
        arith::ConstantIntOp::create(rewriter, loc, /*value=*/1, /*width=*/1);
    SmallVector<Value> maskBits;
    for (auto &coords : coordsVec) {
      Value inBounds = trueVal;
      for (size_t i = 0; i < coords.size(); ++i) {
        Value cmp = arith::CmpIOp::create(
            rewriter, loc, arith::CmpIPredicate::slt, coords[i], origBounds[i]);
        inBounds = arith::AndIOp::create(rewriter, loc, inBounds, cmp);
      }
      maskBits.push_back(inBounds);
    }

    // Build the distributed mask vector.
    Value result;
    if (numElements == 1) {
      result =
          vector::BroadcastOp::create(rewriter, loc, distType, maskBits[0]);
    } else {
      result =
          vector::FromElementsOp::create(rewriter, loc, distType, maskBits);
    }
    rewriter.replaceOp(op, result);
    return success();
  }
};

/// This pattern distributes a subgroup-level StoreMatrix op to workitem-level.
struct SgToWiStoreMatrix : public OpConversionPattern<xegpu::StoreMatrixOp> {
  using OpConversionPattern<xegpu::StoreMatrixOp>::OpConversionPattern;

  LogicalResult
  matchAndRewrite(xegpu::StoreMatrixOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto layout = op.getLayoutAttr();
    // If no layout, nothing to do.
    if (!layout)
      return failure();

    VectorType sgPayloadTy = dyn_cast<VectorType>(op.getData().getType());
    if (!sgPayloadTy)
      return rewriter.notifyMatchFailure(
          op, "the matrix op payload must be a vector type");

    auto loc = op.getLoc();
    auto offsets = op.getMixedOffsets();
    if (offsets.empty())
      return rewriter.notifyMatchFailure(op, "the store op must have offsets");

    FailureOr<VectorType> distPayloadTyOrFailure =
        getDistVecTypeBasedOnLaneLayout(layout, sgPayloadTy);
    if (failed(distPayloadTyOrFailure))
      return rewriter.notifyMatchFailure(
          op, "Failed to distribute matrix op payload based on layout.");

    SmallVector<Value> offsetsAsValues =
        vector::getAsValues(rewriter, loc, offsets);

    SmallVector<Value> newCoords = offsetsAsValues;
    if (!op.getSubgroupBlockIoAttr()) {
      newCoords = computeDistributedCoordsForMatrixOp(
          rewriter, loc, layout, sgPayloadTy.getShape(), offsetsAsValues);
      if (newCoords.empty())
        return rewriter.notifyMatchFailure(
            op, "Failed to compute distributed coordinates.");
    }

    SmallVector<int64_t> newConstOffsets(op.getConstOffsets().size(),
                                         ShapedType::kDynamic);
    DenseI64ArrayAttr newConstOffsetsAttr =
        rewriter.getDenseI64ArrayAttr(newConstOffsets);

    xegpu::StoreMatrixOp::create(
        rewriter, loc, TypeRange{},
        castValueTo(rewriter, cast<TypedValue<VectorType>>(adaptor.getData()),
                    distPayloadTyOrFailure.value()),
        adaptor.getMemDesc(), ValueRange(newCoords), newConstOffsetsAttr,
        op.getSubgroupBlockIoAttr(), xegpu::DistributeLayoutAttr{});
    rewriter.eraseOp(op);
    return success();
  }
};

/// Distributes a subgroup-level StoreScatter (xegpu.store) op to
/// workitem-level.
///
/// Example 1 (1D, no chunk size):
///   layout = #xegpu.layout<lane_layout = [16], lane_data = [1]>
///   %mask = producer_op : vector<16xi1>
///   %offset = producer_op : vector<16xindex>
///   xegpu.store %payload, %src[%offset], %mask : vector<16xf16>,
///     memref<256xf16>, vector<16xindex>, vector<16xi1>
/// Distributed to:
///   %mask = producer_op : vector<1xi1>
///   %offset = producer_op : vector<1xindex>
///   xegpu.store %payload, %src[%offset], %mask : vector<1xf16>,
///     memref<256xf16>, vector<1xindex>, vector<1xi1>
///
/// Example 2 (2D with chunk size, same mask & offset):
///   layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>
///   xegpu.store %payload, %src[%offset], %mask <{chunk_size=8}> :
///     vector<16x8xf16>, memref<256xf16>, vector<16xindex>, vector<16xi1>
/// Distributed to:
///   xegpu.store %payload, %src[%offset], %mask <{chunk_size=8}> :
///     vector<8xf16>, memref<256xf16>, vector<1xindex>, vector<1xi1>
///
/// Example 3 (3D with leading unit dims):
///   layout = #xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 1, 1]>
///   %mask = producer_op : vector<1x1x16xi1>
///   %offset = producer_op : vector<1x1x16xindex>
///   xegpu.store %payload, %src[%offset], %mask : vector<1x1x16xf16>,
///     memref<256xf16>, vector<1x1x16xindex>, vector<1x1x16xi1>
/// Distributed to:
///   %mask = producer_op : vector<1x1x1xi1>
///   %offset = producer_op : vector<1x1x1xindex>
///   xegpu.store %payload, %src[%offset], %mask : vector<1xf16>,
///     memref<256xf16>, vector<1xindex>, vector<1xi1>
struct SgToWiStoreScatter : public OpConversionPattern<xegpu::StoreScatterOp> {
  using OpConversionPattern<xegpu::StoreScatterOp>::OpConversionPattern;

  LogicalResult
  matchAndRewrite(xegpu::StoreScatterOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    xegpu::DistributeLayoutAttr layout = op.getAnchorLayout();
    if (!layout)
      return failure();

    VectorType origValueTy = op.getValueType();
    if (!origValueTy)
      return failure();

    // Check that all leading dimensions are unit dimensions.
    int chunkSize = op.getChunkSize().value_or(1);
    int effectiveVecRank = (chunkSize == 1) ? 1 : 2;
    ArrayRef<int64_t> shape = origValueTy.getShape();
    if (llvm::any_of(shape.take_front(origValueTy.getRank() - effectiveVecRank),
                     [](int64_t d) { return d != 1; }))
      return rewriter.notifyMatchFailure(
          op, "Only unit dimensions allowed for the leading "
              "dimensions of the store vector!");

    auto distValueTyOrFailure =
        xegpu::getDistVecTypeBasedOnLaneLayout(layout, origValueTy);
    if (failed(distValueTyOrFailure))
      return rewriter.notifyMatchFailure(
          op,
          "unable to compute expected workitem vector type from lane layout");

    VectorType distValueTy = distValueTyOrFailure.value();
    VectorType distValueTy1D = VectorType::get({distValueTy.getNumElements()},
                                               distValueTy.getElementType());

    Value distValue = adaptor.getValue();
    if (distValue.getType() != distValueTy1D)
      distValue = castValueTo(rewriter, cast<TypedValue<VectorType>>(distValue),
                              distValueTy1D);

    // Flatten offsets and mask to 1D to match the 1D value type.
    Value distOffsets = adaptor.getOffsets();
    auto distOffsetsTy = cast<VectorType>(distOffsets.getType());
    VectorType offsetsTy1D = VectorType::get({distOffsetsTy.getNumElements()},
                                             distOffsetsTy.getElementType());
    distOffsets = castValueTo(
        rewriter, cast<TypedValue<VectorType>>(distOffsets), offsetsTy1D);

    Value distMask = adaptor.getMask();
    auto distMaskTy = cast<VectorType>(distMask.getType());
    VectorType maskTy1D = VectorType::get({distMaskTy.getNumElements()},
                                          distMaskTy.getElementType());
    distMask =
        castValueTo(rewriter, cast<TypedValue<VectorType>>(distMask), maskTy1D);

    Value distDest = adaptor.getDest();
    xegpu::StoreScatterOp::create(rewriter, op.getLoc(), distValue, distDest,
                                  distOffsets, distMask, op.getChunkSizeAttr(),
                                  op.getL1HintAttr(), op.getL2HintAttr(),
                                  op.getL3HintAttr(), /*layout=*/nullptr);
    rewriter.eraseOp(op);
    return success();
  }
};

/// Distribute a vector::StepOp to workitem-level.
/// The layout must have exactly 1 effective lane dimension.
/// We completely resolve the vector::StepOp by computing the lane_data-sized
/// subranges.
struct SgToWiVectorStep : public OpConversionPattern<vector::StepOp> {
  using OpConversionPattern<vector::StepOp>::OpConversionPattern;

  LogicalResult
  matchAndRewrite(vector::StepOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    xegpu::DistributeLayoutAttr resultLayout =
        xegpu::getTemporaryLayout(op->getResult(0));
    if (!resultLayout || !resultLayout.isForSubgroup())
      return rewriter.notifyMatchFailure(
          op, "the result vector of the step op lacks subgroup layout");

    auto loc = op.getLoc();
    auto stepResultVecTy = op.getResult().getType();
    auto wiShapeOrFailure =
        xegpu::getDistVecTypeBasedOnLaneLayout(resultLayout, stepResultVecTy);
    if (failed(wiShapeOrFailure))
      return rewriter.notifyMatchFailure(
          op, "unable to compute workitem vector type from the layout");
    VectorType newVecTy = wiShapeOrFailure.value();

    Value laneId = gpu::LaneIdOp::create(rewriter, loc, rewriter.getIndexType(),
                                         /*upperBound=*/mlir::IntegerAttr());
    auto laneDataBlockCoords = resultLayout.computeDistributedCoords(
        rewriter, loc, laneId, stepResultVecTy.getShape());
    if (failed(laneDataBlockCoords))
      return rewriter.notifyMatchFailure(
          op, "failed to compute lane data block coordinates");

    auto laneDataBlockCoordsVec = laneDataBlockCoords.value();
    auto laneDataBlockLength = resultLayout.getEffectiveLaneDataAsInt()[0];
    assert(static_cast<int64_t>(laneDataBlockCoordsVec.size()) ==
           newVecTy.getNumElements() / laneDataBlockLength);
    SmallVector<Value> stepVals;
    // For each lane_data block, reconstruct its sub-range
    // from the range of SG-level vector.step.Example: vector.step
    // {slice<layout<lane_layout=[2,4,2], lane_data=[1,2,1]>, dims=[0,2]>} :
    // vector<16xindex>
    // Each logical lane holds 4 elements as 2 blocks of 2 elements each.
    // The blocks are round-robin distributed, so logical lane id 0
    // holds values [0,1, 8,9].
    for (auto &laneDataBlockCoords : laneDataBlockCoordsVec) {
      auto laneDataBlockStartCoord = laneDataBlockCoords[0];
      stepVals.push_back(laneDataBlockStartCoord);
      for (int i = 1; i < laneDataBlockLength; ++i) {
        auto offset = arith::ConstantIndexOp::create(rewriter, loc, i);
        stepVals.push_back(arith::AddIOp::create(
            rewriter, loc, laneDataBlockStartCoord, offset));
      }
    }
    assert(static_cast<int64_t>(stepVals.size()) == newVecTy.getNumElements() &&
           "Expecting the number of step values to match the number of "
           "elements in the vector");
    auto stepOpVal =
        vector::FromElementsOp::create(rewriter, loc, newVecTy, stepVals);
    rewriter.replaceOp(op, stepOpVal);
    return success();
  }
};

/// Distributes a subgroup-level vector.extract op to workitem-level. Only
/// handles sub-vector extraction (result is VectorType, not scalar).
struct SgToWiVectorExtract : public OpConversionPattern<vector::ExtractOp> {
  using OpConversionPattern<vector::ExtractOp>::OpConversionPattern;

  LogicalResult
  matchAndRewrite(vector::ExtractOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    // Only handle vector results (not scalar extraction).
    auto resultType = dyn_cast<VectorType>(op.getType());
    if (!resultType)
      return rewriter.notifyMatchFailure(op, "scalar extract not supported");

    xegpu::DistributeLayoutAttr layout =
        xegpu::getTemporaryLayout(op->getOpResult(0));
    if (!layout || !layout.isForSubgroup())
      return failure();

    // This implementation assumes distribution only happens on the innermost
    // dimension. Verify that lane_layout[0...n-2] are all unit.
    auto laneLayout = layout.getEffectiveLaneLayoutAsInt();
    if (llvm::any_of(ArrayRef<int64_t>(laneLayout).drop_back(1),
                     [](int64_t v) { return v != 1; }))
      return rewriter.notifyMatchFailure(
          op, "only innermost dimension distribution is supported for "
              "vector.extract");

    auto newOp = vector::ExtractOp::create(
        rewriter, op.getLoc(), adaptor.getSource(), op.getMixedPosition());
    rewriter.replaceOp(op, newOp.getResult());
    return success();
  }
};

/// This pattern distributes a subgroup-level ShapeCast op to workitem-level.
struct SgToWiVectorShapeCast : public OpConversionPattern<vector::ShapeCastOp> {
  using OpConversionPattern<vector::ShapeCastOp>::OpConversionPattern;

  LogicalResult
  matchAndRewrite(vector::ShapeCastOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    xegpu::DistributeLayoutAttr resultLayout =
        xegpu::getTemporaryLayout(op->getOpResult(0));
    if (!resultLayout || !resultLayout.isForSubgroup())
      return rewriter.notifyMatchFailure(
          op, "the result vector of the shape_cast op lacks subgroup layout");

    auto resultDistTypeOrFailure = xegpu::getDistVecTypeBasedOnLaneLayout(
        resultLayout, op.getResultVectorType());
    if (failed(resultDistTypeOrFailure))
      return rewriter.notifyMatchFailure(
          op, "failed to get distributed vector type for result");

    Value source = adaptor.getSource();
    auto newShapeCast = vector::ShapeCastOp::create(
        rewriter, op.getLoc(), resultDistTypeOrFailure.value(), source);
    rewriter.replaceOp(op, newShapeCast);
    return success();
  }
};

/// Distributes a subgroup-level vector.extract_strided_slice op to
/// workitem-level. If the result is distributed, the offsets and sizes are
/// adjusted to match the distributed types.
struct SgToWiVectorExtractStridedSlice
    : public OpConversionPattern<vector::ExtractStridedSliceOp> {
  using OpConversionPattern<vector::ExtractStridedSliceOp>::OpConversionPattern;

  LogicalResult
  matchAndRewrite(vector::ExtractStridedSliceOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    xegpu::DistributeLayoutAttr resultLayout =
        xegpu::getTemporaryLayout(op->getOpResult(0));
    if (!resultLayout || !resultLayout.isForSubgroup())
      return failure();

    VectorType resultType = op.getType();
    auto distResultTyOrFailure =
        xegpu::getDistVecTypeBasedOnLaneLayout(resultLayout, resultType);
    if (failed(distResultTyOrFailure))
      return rewriter.notifyMatchFailure(
          op, "unable to compute distributed vector type from lane layout");
    VectorType distResultTy = *distResultTyOrFailure;

    SmallVector<int64_t> distributedDims =
        getDistributedDims(resultType, distResultTy);

    // Collect updated sizes, offsets, strides. Pad to full source rank.
    int64_t sourceRank = op.getSourceVectorType().getRank();
    SmallVector<Attribute> updatedSizes =
        llvm::map_to_vector(op.getSizes(), [](Attribute attr) { return attr; });
    SmallVector<Attribute> updatedOffsets = llvm::map_to_vector(
        op.getOffsets(), [](Attribute attr) { return attr; });
    SmallVector<Attribute> updatedStrides = llvm::map_to_vector(
        op.getStrides(), [](Attribute attr) { return attr; });
    for (int64_t i = op.getSizes().size(); i < sourceRank; ++i) {
      updatedSizes.push_back(
          rewriter.getI64IntegerAttr(op.getSourceVectorType().getDimSize(i)));
      updatedOffsets.push_back(rewriter.getI64IntegerAttr(0));
      updatedStrides.push_back(rewriter.getI64IntegerAttr(1));
    }

    // If the result is distributed, adjust offsets and sizes in the
    // distributed dimension.
    if (!distributedDims.empty()) {
      if (distributedDims.size() != 1)
        return rewriter.notifyMatchFailure(
            op, "only single dimension distribution is supported");
      int64_t distDim = distributedDims[0];
      const uArch *uArch = getUArch(xegpu::getChipStr(op).value_or(""));
      if (!uArch)
        return rewriter.notifyMatchFailure(
            op, "target attribute required to determine subgroup size");
      int subgroupSize = uArch->getSubgroupSize();
      auto sourceLayout = xegpu::getTemporaryLayout(op->getOpOperand(0));
      if (!sourceLayout || sourceLayout.getEffectiveLaneLayoutAsInt().empty())
        return rewriter.notifyMatchFailure(
            op, "source of extract_strided_slice lacks distribution layout");
      int sourceDistrDimSize = op.getSourceVectorType().getShape()[distDim];
      if (sourceDistrDimSize % subgroupSize != 0)
        return rewriter.notifyMatchFailure(
            op, "source size along distributed dim is not a multiple of "
                "subgroup size");
      auto sourceLaneData = sourceLayout.getEffectiveLaneDataAsInt();
      // Only check lane_data for the distributed dimension. Non-distributed
      // dimensions may have non-unit lane_data (e.g., packed layouts).
      if (distDim < static_cast<int64_t>(sourceLaneData.size()) &&
          sourceLaneData[distDim] != 1)
        return rewriter.notifyMatchFailure(
            op, "expecting unit lane data along the distributed dimension");
      int64_t distrDimOffset =
          cast<IntegerAttr>(updatedOffsets[distDim]).getInt();
      if (distrDimOffset % subgroupSize != 0)
        return rewriter.notifyMatchFailure(
            op, "offset along distributed dim is not a multiple of "
                "subgroup size");
      // Adjust sizes and offsets for the distributed dimension.
      updatedSizes[distDim] =
          rewriter.getI64IntegerAttr(distResultTy.getDimSize(distDim));
      updatedOffsets[distDim] =
          rewriter.getI64IntegerAttr(distrDimOffset / subgroupSize);
    }

    auto newOp = vector::ExtractStridedSliceOp::create(
        rewriter, op.getLoc(), distResultTy, adaptor.getSource(),
        ArrayAttr::get(rewriter.getContext(), updatedOffsets),
        ArrayAttr::get(rewriter.getContext(), updatedSizes),
        ArrayAttr::get(rewriter.getContext(), updatedStrides));
    rewriter.replaceOp(op, newOp.getResult());
    return success();
  }
};

/// This pattern distributes a subgroup-level `vector.broadcast` op to
/// workitem-level. The pattern supports three cases:
///
/// 1) Broadcast a low-rank vector to high-rank vector: The low-rank input
///    vector must have a slice layout of the result. If the distributed source
///    and target vector types are identical, this lowers to a no-op; otherwise,
///    it remains a broadcast but operates on distributed vectors.
///
/// 2) Broadcast a same-rank vector with identical layouts for source and
///    target: The source vector must have unit dimensions, and lane_data must
///    be unit size for those unit dims. This always lowers to a no-op.
///
/// 3) Broadcast a scalar with no layout: This always lowers to a broadcast
///    from scalar to distributed result type.
///
/// Example 1 (low-rank to high-rank broadcast):
/// ```
///   %0 = "some_op"() {layout_result_0 =
///     #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,
///     dims = [0]>} : () -> vector<16xf16>
///   %1 = vector.broadcast %0 {layout_result_0 =
///     #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
///     : vector<16xf16> to vector<16x16xf16>
/// ```
/// is distributed to:
/// ```
///   %0 = "some_op"() : () -> vector<1xf16>
///   %1 = vector.broadcast %0 : vector<1xf16> to vector<16x1xf16>
/// ```
///
/// Example 2 (same-rank broadcast, no-op):
/// ```
///   %0 = "some_op"() {layout_result_0 =
///     #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
///     : () -> vector<16x1xf16>
///   %1 = vector.broadcast %0 {layout_result_0 =
///     #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
///     : vector<16x1xf16> to vector<16x16xf16>
/// ```
/// is distributed to (no-op, source already matches distributed result type):
/// ```
///   %0 = "some_op"() : () -> vector<16x1xf16>
///   // broadcast is eliminated, %0 is used directly
/// ```
///
/// Example 3 (scalar to vector broadcast):
/// ```
///   %0 = "some_op"() : () -> f16
///   %1 = vector.broadcast %0 {layout_result_0 =
///     #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
///     : f16 to vector<16x16xf16>
/// ```
/// is distributed to:
/// ```
///   %0 = "some_op"() : f16
///   %1 = vector.broadcast %0 : f16 to vector<16x1xf16>
/// ```
struct SgToWiBroadcast : public OpConversionPattern<vector::BroadcastOp> {
  using OpConversionPattern<vector::BroadcastOp>::OpConversionPattern;

  LogicalResult
  matchAndRewrite(vector::BroadcastOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    xegpu::DistributeLayoutAttr resultLayout =
        xegpu::getTemporaryLayout(cast<OpResult>(op.getResult()));
    if (!resultLayout || !resultLayout.isForSubgroup())
      return rewriter.notifyMatchFailure(
          op, "result does not have subgroup distribute layout");

    VectorType destType = op.getResultVectorType();
    VectorType sourceType = dyn_cast<VectorType>(op.getSourceType());

    xegpu::DistributeLayoutAttr sourceLayout =
        xegpu::getTemporaryLayout(op->getOpOperand(0));

    if (sourceType) {
      int64_t rankDiff = destType.getRank() - sourceType.getRank();
      if (rankDiff > 0) {
        // Case 1: Low-rank to high-rank broadcast.
        if (!sourceLayout || !sourceLayout.isSliceOf(resultLayout))
          op.emitWarning(
              "broadcast source layout must be a slice of result layout");
      } else if (rankDiff == 0) {
        // Case 2: Same-rank broadcast.
        auto broadcastUnitDimsSet = op.computeBroadcastedUnitDims();
        SmallVector<int64_t> broadcastUnitDims(broadcastUnitDimsSet.begin(),
                                               broadcastUnitDimsSet.end());
        assert(sourceLayout.isEqualTo(
                   sourceLayout.setUnitDimData(broadcastUnitDims)) &&
               "The sg_data for unit dimensions should be set as 1");
        sourceLayout = sourceLayout.setUnitDimLayout(broadcastUnitDims);
      }
    } else {
      // Case 3: Scalar to vector broadcast.
      if (sourceLayout)
        return rewriter.notifyMatchFailure(
            op, "broadcast from scalar must not have a layout attribute");
    }

    auto destDistType =
        xegpu::getDistVecTypeBasedOnLaneLayout(resultLayout, destType);
    if (failed(destDistType))
      return rewriter.notifyMatchFailure(
          op, "failed to distribute the result vector type");

    Value source = adaptor.getSource();
    // If the adapted source already matches the dest dist type, it's a no-op.
    if (source.getType() == destDistType.value()) {
      rewriter.replaceOp(op, source);
      return success();
    }

    auto newOp = vector::BroadcastOp::create(rewriter, op.getLoc(),
                                             destDistType.value(), source);
    rewriter.replaceOp(op, newOp);
    return success();
  }
};

/// Distributes a subgroup-level vector.insert_strided_slice op to
/// workitem-level. If the dest is distributed, the offsets are adjusted to
/// match the distributed types.
struct SgToWiVectorInsertStridedSlice
    : public OpConversionPattern<vector::InsertStridedSliceOp> {
  using OpConversionPattern<vector::InsertStridedSliceOp>::OpConversionPattern;

  LogicalResult
  matchAndRewrite(vector::InsertStridedSliceOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    xegpu::DistributeLayoutAttr resultLayout =
        xegpu::getTemporaryLayout(op->getOpResult(0));
    if (!resultLayout || !resultLayout.isForSubgroup())
      return failure();

    VectorType destType = op.getDestVectorType();
    auto distDestTyOrFailure =
        xegpu::getDistVecTypeBasedOnLaneLayout(resultLayout, destType);
    if (failed(distDestTyOrFailure))
      return rewriter.notifyMatchFailure(
          op, "unable to compute distributed vector type from lane layout");
    VectorType distDestTy = *distDestTyOrFailure;

    SmallVector<int64_t> destDistributedDims =
        getDistributedDims(destType, distDestTy);

    SmallVector<Attribute> updatedOffsets = llvm::map_to_vector(
        op.getOffsets(), [](Attribute attr) { return attr; });

    if (!destDistributedDims.empty()) {
      if (destDistributedDims.size() != 1)
        return rewriter.notifyMatchFailure(
            op, "only single dimension distribution is supported");
      int64_t destDistDim = destDistributedDims[0];

      const uArch *uArch = getUArch(xegpu::getChipStr(op).value_or(""));
      if (!uArch)
        return rewriter.notifyMatchFailure(
            op, "target attribute required to determine subgroup size");
      int subgroupSize = uArch->getSubgroupSize();

      VectorType srcType = op.getSourceVectorType();
      // The distributed dim must be in the last k (source rank) dims of dest.
      int64_t sourceDistDim =
          destDistDim - (destType.getRank() - srcType.getRank());
      if (sourceDistDim < 0)
        return rewriter.notifyMatchFailure(
            op, "distributed dimension must be in the last k dims of dest");

      auto destLayout = xegpu::getTemporaryLayout(op->getOpOperand(1));
      auto sourceLayout = xegpu::getTemporaryLayout(op->getOpOperand(0));
      if (!destLayout || !sourceLayout ||
          destLayout.getEffectiveLaneLayoutAsInt().empty() ||
          sourceLayout.getEffectiveLaneLayoutAsInt().empty())
        return rewriter.notifyMatchFailure(
            op, "source or dest of insert_strided_slice lacks distribution "
                "layout");

      auto destLaneData = destLayout.getEffectiveLaneDataAsInt();
      auto sourceLaneData = sourceLayout.getEffectiveLaneDataAsInt();
      // Only check lane_data for the distributed dimension. Non-distributed
      // dimensions may have non-unit lane_data (e.g., packed layouts).
      if ((destDistDim < static_cast<int64_t>(destLaneData.size()) &&
           destLaneData[destDistDim] != 1) ||
          (sourceDistDim < static_cast<int64_t>(sourceLaneData.size()) &&
           sourceLaneData[sourceDistDim] != 1))
        return rewriter.notifyMatchFailure(
            op, "expecting unit lane data along the distributed dimension");

      int64_t srcDistrDimSize = srcType.getDimSize(sourceDistDim);
      if (srcDistrDimSize % subgroupSize != 0)
        return rewriter.notifyMatchFailure(
            op, "source distributed dim size is not a multiple of "
                "subgroup size");

      int64_t destDistrDimOffset =
          cast<IntegerAttr>(op.getOffsets()[destDistDim]).getInt();
      if (destDistrDimOffset % subgroupSize != 0)
        return rewriter.notifyMatchFailure(
            op, "offset along distributed dim is not a multiple of "
                "subgroup size");
      // Adjust offset for the distributed dimension.
      updatedOffsets[destDistDim] =
          rewriter.getI64IntegerAttr(destDistrDimOffset / subgroupSize);
    }

    auto newOp = vector::InsertStridedSliceOp::create(
        rewriter, op.getLoc(), distDestTy, adaptor.getValueToStore(),
        adaptor.getDest(),
        ArrayAttr::get(rewriter.getContext(), updatedOffsets), op.getStrides());
    rewriter.replaceOp(op, newOp.getResult());
    return success();
  }
};

/// Distributes a subgroup-level vector.insert op to workitem-level. Only
/// handles sub-vector insertion (value to store is VectorType, not scalar).
struct SgToWiVectorInsert : public OpConversionPattern<vector::InsertOp> {
  using OpConversionPattern<vector::InsertOp>::OpConversionPattern;

  LogicalResult
  matchAndRewrite(vector::InsertOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    // Only handle vector value-to-store (not scalar insertion).
    auto valueType = dyn_cast<VectorType>(op.getValueToStoreType());
    if (!valueType)
      return rewriter.notifyMatchFailure(op, "scalar insert not supported");

    xegpu::DistributeLayoutAttr layout =
        xegpu::getTemporaryLayout(op->getOpResult(0));
    if (!layout || !layout.isForSubgroup())
      return failure();

    // verify that the outer k dimensions (for offsets)
    // don't have non-unit lane_layout.
    auto laneLayout = layout.getEffectiveLaneLayoutAsInt();
    if (llvm::any_of(ArrayRef<int64_t>(laneLayout).drop_back(1),
                     [](int64_t v) { return v != 1; }))
      return rewriter.notifyMatchFailure(
          op, "only innermost dimension distribution is supported for "
              "vector.insert");

    auto newOp = vector::InsertOp::create(
        rewriter, op.getLoc(), adaptor.getValueToStore(), adaptor.getDest(),
        op.getMixedPosition());
    rewriter.replaceOp(op, newOp.getResult());
    return success();
  }
};

/// Folds a subgroup-level ConvertLayout op with compatible lane layouts.
struct SgToWiConvertLayout
    : public OpConversionPattern<xegpu::ConvertLayoutOp> {
  using OpConversionPattern<xegpu::ConvertLayoutOp>::OpConversionPattern;

  LogicalResult
  matchAndRewrite(xegpu::ConvertLayoutOp op, OpAdaptor adaptor,
                  ConversionPatternRewriter &rewriter) const override {
    auto inputLayout = op.getInputLayoutAttr();
    auto targetLayout = op.getTargetLayoutAttr();
    Type valType = op.getResult().getType();

    if (valType.isIntOrFloat()) {
      rewriter.replaceOp(op, op.getSource());
      return success();
    }

    auto resShape = cast<VectorType>(valType).getShape();
    SmallVector<int64_t> resShapeVec(resShape.begin(), resShape.end());
    if (!inputLayout.isCompatibleWith(targetLayout, resShapeVec,
                                      xegpu::LayoutKind::Lane)) {
      return rewriter.notifyMatchFailure(
          op, "lowering incompatible convert_layout not yet supported");
    }

    rewriter.replaceOp(op, adaptor.getSource());
    return success();
  }
};

struct XeGPUSgToWiDistributeExperimentalPass
    : public xegpu::impl::XeGPUSgToWiDistributeExperimentalBase<
          XeGPUSgToWiDistributeExperimentalPass> {
  void runOnOperation() override;
};

} // namespace

void XeGPUSgToWiDistributeExperimentalPass::runOnOperation() {

  // Recover temporary operand layouts for usage in patterns.
  Operation *root = getOperation();
  if (!xegpu::recoverTemporaryLayouts(root)) {
    signalPassFailure();
    return;
  }

  // Verify if all XeGPU anchor ops and vector ops have result layouts.
  // TODO: This can be removed once the full layout refactoring is done.
  if (failed(verifyLayouts(root))) {
    LLVM_DEBUG(DBGS() << "XeGPUSgToWiDistributeExperimentalPass: layout "
                         "verification failed\n");
    signalPassFailure();
    return;
  }
  // Collect existing UnrealizedConversionCastOps. These must be preserved.
  llvm::SmallSetVector<UnrealizedConversionCastOp, 8> existingCasts;
  root->walk(
      [&](UnrealizedConversionCastOp castOp) { existingCasts.insert(castOp); });
  // Perform a structural type conversion to convert structural ops to have WI
  // types. This will insert UnrealizedConversionCastOps to make the IR
  // valid.
  auto materializeCast = [&](mlir::OpBuilder &builder, mlir::Type type,
                             mlir::ValueRange inputs,
                             mlir::Location loc) -> mlir::Value {
    UnrealizedConversionCastOp castOp =
        UnrealizedConversionCastOp::create(builder, loc, type, inputs);
    return castOp.getResult(0);
  };
  {
    ConversionTarget target(getContext());
    TypeConverter typeConverter;
    RewritePatternSet patterns(&getContext());
    typeConverter.addSourceMaterialization(materializeCast);
    typeConverter.addTargetMaterialization(materializeCast);
    xegpu::populateXeGPUSgToWiDistributeTypeConversions(typeConverter);
    scf::populateSCFStructuralTypeConversionsAndLegality(typeConverter,
                                                         patterns, target);
    xegpu::populateXeGPUSgToWiDistributeTypeConversionAndLegality(
        typeConverter, patterns, target);
    target.addLegalOp<UnrealizedConversionCastOp>();
    (void)applyPartialConversion(root, target, std::move(patterns));
  }
  // Structural type conversion can generate some redundant
  // UnrealizedConversionCastOps to materialize the SG type from type converted
  // WI type. These are redundant at this point and can be eliminated by
  // inserting shape casts instead.
  // Example:
  // %1 = UnrealizedConversionCastOp %0 : vector<16x1xf32> to vector<16x16xf32>
  // %2 = UnrealizedConversionCastOp %1 : vector<16x16xf32> to vector<16xf32>
  // This can be replaced with:
  // %2 = vector.shape_cast %0 : vector<16x1xf32> to vector<16xf32>
  OpBuilder builder(root);
  root->walk([&](UnrealizedConversionCastOp op) {
    // If this op existed before, nothing to do.
    if (existingCasts.contains(op))
      return;
    // number of inputs and outputs must be 1.
    if (op.getNumOperands() != 1 || op.getNumResults() != 1)
      return;
    // Both input and output types must be vector types.
    auto singleInput = op.getInputs()[0];
    auto inputTy = dyn_cast<VectorType>(singleInput.getType());
    auto outputTy = dyn_cast<VectorType>(op.getResult(0).getType());
    if (!inputTy || !outputTy)
      return;

    // Check if the defining op of the input is also an
    // UnrealizedConversionCastOp and it has a single user (which is this
    // op).
    auto definingOp = singleInput.getDefiningOp<UnrealizedConversionCastOp>();
    if (!definingOp || !definingOp->hasOneUse())
      return;
    auto inputOfDefiningOp = definingOp.getInputs()[0];
    // If the input of the defining op and output type are both vector types
    // have same number of elements, insert a shape cast.
    auto inputOfDefiningOpTy =
        dyn_cast<VectorType>(inputOfDefiningOp.getType());
    if (inputOfDefiningOpTy &&
        inputOfDefiningOpTy.getNumElements() == outputTy.getNumElements()) {
      builder.setInsertionPoint(op);
      auto shapeCast = vector::ShapeCastOp::create(builder, op.getLoc(),
                                                   outputTy, inputOfDefiningOp);
      op.replaceAllUsesWith(ValueRange{shapeCast.getResult()});
      return;
    }
  });
  // At this point, we will have some dead UnrealizedConversionCastOps. Just
  // erase them.
  bool changed = true;
  while (changed) {
    changed = false;
    root->walk([&](UnrealizedConversionCastOp op) {
      // Skip existing casts.
      if (existingCasts.contains(op))
        return;
      if (op.use_empty()) {
        op.erase();
        changed = true;
      }
    });
  }
}

void xegpu::populateXeGPUSgToWiDistributeTypeConversions(
    TypeConverter &typeConverter) {
  // Any type other than TensorDescType and VectorType are legal as is.
  typeConverter.addConversion([](Type type) -> std::optional<Type> {
    if (!isa<TensorDescType, VectorType>(type))
      return type;
    return std::nullopt;
  });
  // For TensorDescType, drop the layout attribute if any.
  typeConverter.addConversion([](TensorDescType type) -> Type {
    if (type.getLayoutAttr()) {
      return type.dropLayouts();
    }
    return type;
  });
  // For VectorType, check if there is a distribute layout attribute on the
  // value. If so, convert to the distributed vector type based on the layout.
  typeConverter.addConversion([](Value v) -> std::optional<Type> {
    auto type = v.getType();
    // If value is not vector type, nothing to do.
    if (!isa<VectorType>(type))
      return std::nullopt;
    auto layout = xegpu::getDistributeLayoutAttr(v);
    if (!layout || !layout.isForSubgroup())
      return type;
    // Vector type is distributed based on lane layout.
    auto newTyOrFailure =
        getDistVecTypeBasedOnLaneLayout(layout, cast<VectorType>(type));
    if (failed(newTyOrFailure))
      return type;
    return *newTyOrFailure;
  });
}

void xegpu::populateXeGPUSgToWiDistributeTypeConversionAndLegality(
    TypeConverter &typeConverter, RewritePatternSet &patterns,
    ConversionTarget &target) {
  populateXeGPUSgToWiDistributeTypeConversions(typeConverter);
  // CreateNdDescOp is legal only if its result type has no layout attribute.
  target.addDynamicallyLegalOp<xegpu::CreateNdDescOp>(
      [&](xegpu::CreateNdDescOp op) { return !op.getType().getLayoutAttr(); });
  // Any anchor XeGPU op is legal only if it has no anchor layout.
  target.addDynamicallyLegalDialect<xegpu::XeGPUDialect>([](Operation *op) {
    auto anchorOp = dyn_cast<AnchorLayoutInterface>(op);
    if (!anchorOp)
      return true;
    return !anchorOp.getAnchorLayout();
  });
  // Arith constants are legal only if they have no temporary layout attribute.
  target.addDynamicallyLegalOp<arith::ConstantOp>(
      [=](arith::ConstantOp op) -> bool {
        // If the result type is not a vector, it's legal.
        if (!isa<VectorType>(op.getResult().getType()))
          return true;
        return !xegpu::getTemporaryLayout(dyn_cast<OpResult>(op.getResult()));
      });
  // In math and arith dialects, only handle elementwise ops with a single
  // result and with a result layout attribute.
  target.addDynamicallyLegalDialect<math::MathDialect, arith::ArithDialect>(
      [=](Operation *op) -> std::optional<bool> {
        // Only handle elementwise mappable ops
        if (!OpTrait::hasElementwiseMappableTraits(op))
          return true;
        // Only handle ops with single vector result
        if (op->getNumResults() != 1)
          return true;

        VectorType resultType =
            dyn_cast<VectorType>(op->getResult(0).getType());
        if (!resultType)
          return true;

        // Check if all operands are vectors of the same shape
        for (Value operand : op->getOperands()) {
          VectorType operandType = dyn_cast<VectorType>(operand.getType());
          if (!operandType || operandType.getShape() != resultType.getShape()) {
            return true;
          }
        }
        return !xegpu::getTemporaryLayout(dyn_cast<OpResult>(op->getResult(0)));
      });
  // vector::ReductionOp is legal only if its source has no distribute layout
  // attribute.
  target.addDynamicallyLegalOp<vector::ReductionOp>(
      [=](vector::ReductionOp op) -> bool {
        auto layout = xegpu::getDistributeLayoutAttr(op.getVector());
        return !layout;
      });
  // vector::MultiDimReductionOp op legality.
  target.addDynamicallyLegalOp<vector::MultiDimReductionOp>(
      [=](vector::MultiDimReductionOp op) -> bool {
        return !isValidSubgroupMultiReductionOp(op);
      });
  target.addDynamicallyLegalOp<vector::CreateMaskOp, vector::ConstantMaskOp,
                               vector::TransposeOp, vector::BitCastOp,
                               vector::ShapeCastOp, vector::StepOp,
                               vector::BroadcastOp>([=](Operation *op) -> bool {
    return !xegpu::getTemporaryLayout(op->getOpResult(0));
  });
  target.addDynamicallyLegalOp<vector::ExtractOp>(
      [=](vector::ExtractOp op) -> bool {
        if (!isa<VectorType>(op.getType()))
          return true;
        return !xegpu::getTemporaryLayout(op->getOpResult(0));
      });
  target.addDynamicallyLegalOp<vector::InsertOp>(
      [=](vector::InsertOp op) -> bool {
        return !xegpu::getTemporaryLayout(op->getOpResult(0));
      });
  target.addDynamicallyLegalOp<vector::ExtractStridedSliceOp>(
      [=](vector::ExtractStridedSliceOp op) -> bool {
        return !xegpu::getTemporaryLayout(op->getOpResult(0));
      });
  target.addDynamicallyLegalOp<vector::InsertStridedSliceOp>(
      [=](vector::InsertStridedSliceOp op) -> bool {
        return !xegpu::getTemporaryLayout(op->getOpResult(0));
      });
  target.markUnknownOpDynamicallyLegal([](Operation *op) { return true; });
  patterns.add<SgToWiCreateNdDesc, SgToWiLoadNd, SgToWiStoreNd, SgToWiDpas,
               SgToWiElementWise, SgToWiArithConstant, SgToWiPrefetchNd,
               SgToWiLoadGather, SgToWiStoreScatter, SgToWiVectorReduction,
               SgToWiMultiDimReduction, SgToWiVectorExtract, SgToWiVectorInsert,
               SgToWiVectorExtractStridedSlice, SgToWiVectorInsertStridedSlice,
               SgToWiLoadMatrix, SgToWiStoreMatrix, SgToWiConvertLayout,
               SgToWiVectorTranspose, SgToWiVectorBitcast, SgToWiVectorStep,
               SgToWiVectorShapeCast, SgToWiBroadcast,
               SgToWiCreateMask<vector::CreateMaskOp>,
               SgToWiCreateMask<vector::ConstantMaskOp>>(typeConverter,
                                                         patterns.getContext());
}