Repository: llvm/llvm-project
File: mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToWiDistributeExperimental.cpp (C++)

Code before:
    if (!layout)
      return failure();

    VectorType valueTy = op.getValueType();
    if (!valueTy)
      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 = valueTy.getShape();
    if (llvm::any_of(shape.take_front(valueTy.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 expectedWiValueTyOrFailure =
        xegpu::getDistVecTypeBasedOnLaneLayout(layout, valueTy);
    if (failed(expectedWiValueTyOrFailure))
      return rewriter.notifyMatchFailure(
          op,
          "unable to compute expected workitem vector type from lane layout");

    VectorType expectedWiValueTy = expectedWiValueTyOrFailure.value();
    VectorType supportedWiValueTy =
        VectorType::get({expectedWiValueTy.getNumElements()},
                        expectedWiValueTy.getElementType());

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

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

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

    xegpu::StoreScatterOp::create(
        rewriter, op.getLoc(), adaptedValue, adaptor.getDest(), offsets, mask,

Review comment (bug):
Not sure about "expected" and "supported" prefix. why not just distValueTy and distValueTy1D? 

Somthing like 
    auto distValueTyOrFailure = xegpu::getDistVecTypeBasedOnLaneLayout(layout, origValueTy);
    VectorType distValueTy = distValueTyOrFailure.value();
    VectorType distValueTy1D = vectorType::get(distValueTy.getNumElements()} ...

Code after:

/// 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);