brintos / llvm-project-archived public Read only

1//===- XeGPUSubgroupDistribute.cpp - XeGPU Subgroup Distribute Pass -------===//2//3// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.4// See https://llvm.org/LICENSE.txt for license information.5// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception6//7//===----------------------------------------------------------------------===//8#include "mlir/Dialect/GPU/IR/GPUDialect.h"9#include "mlir/Dialect/GPU/Utils/DistributionUtils.h"10#include "mlir/Dialect/Index/IR/IndexDialect.h"11#include "mlir/Dialect/MemRef/IR/MemRef.h"12#include "mlir/Dialect/Vector/IR/VectorOps.h"13#include "mlir/Dialect/Vector/Transforms/VectorDistribution.h"14#include "mlir/Dialect/XeGPU/IR/XeGPU.h"15#include "mlir/Dialect/XeGPU/Transforms/Passes.h"16#include "mlir/Dialect/XeGPU/Transforms/Transforms.h"17#include "mlir/Dialect/XeGPU/Utils/XeGPUUtils.h"18#include "mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h"19#include "mlir/IR/AffineMap.h"20#include "mlir/IR/Attributes.h"21#include "mlir/IR/Builders.h"22#include "mlir/IR/BuiltinAttributes.h"23#include "mlir/IR/BuiltinOps.h"24#include "mlir/IR/BuiltinTypes.h"25#include "mlir/IR/Operation.h"26#include "mlir/IR/PatternMatch.h"27#include "mlir/IR/TypeRange.h"28#include "mlir/IR/Value.h"29#include "mlir/IR/Visitors.h"30#include "mlir/Interfaces/FunctionInterfaces.h"31#include "mlir/Support/LLVM.h"32#include "mlir/Transforms/DialectConversion.h"33#include "mlir/Transforms/GreedyPatternRewriteDriver.h"34#include "mlir/Transforms/InliningUtils.h"35#include "llvm/ADT/ArrayRef.h"36#include "llvm/ADT/STLExtras.h"37#include "llvm/ADT/SmallVector.h"38 39namespace mlir {40namespace xegpu {41#define GEN_PASS_DEF_XEGPUSUBGROUPDISTRIBUTE42#include "mlir/Dialect/XeGPU/Transforms/Passes.h.inc"43} // namespace xegpu44} // namespace mlir45 46#define DEBUG_TYPE "xegpu-subgroup-distribute"47#define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "]: ")48 49using namespace mlir;50 51static const char *const resolveSIMTTypeMismatch =52    "resolve_simt_type_mismatch"; // Attribute name for identifying53                                  // UnrelizedConversionCastOp added to resolve54                                  // SIMT type mismatches.55 56namespace {57 58//===----------------------------------------------------------------------===//59// SIMT Distribution Patterns60//===----------------------------------------------------------------------===//61 62/// In certain cases, we may need to favor XeGPU specific distribution patterns63/// over generic vector distribution patterns. In such cases, we can assign64/// priorities to patterns.65static constexpr unsigned regularPatternBenefit = 1;66static constexpr unsigned highPatternBenefit = 2;67 68/// Helper function to get  distributed vector type for a source vector type69/// according to the lane_layout. We simply divide each dimension of tensor70/// descriptor shape by corresponding lane_layout dimension. If71/// array_length > 1, that is appended to the front of the ditributed shape.72/// NOTE: This is the vector type that will be returned by the73/// gpu.warp_execute_on_lane0 op.74///75/// Examples:76/// | original vector shape | lane_layout | distributed vector shape |77/// |-----------------------|-------------|--------------------------|78/// | 32x16                 | [1, 16]     | 32x1                     |79/// | 32x16                 | [2, 8]      | 16x2                     |80/// | 2x32x16               | [1, 16]     | 2x32x1                   |81static FailureOr<VectorType>82getDistVecTypeBasedOnLaneLayout(xegpu::DistributeLayoutAttr layout,83                                VectorType originalType) {84  if (!layout)85    return failure();86  assert((isa<xegpu::LayoutAttr>(layout) || isa<xegpu::SliceAttr>(layout)) &&87         "Expecting a valid layout.");88  SmallVector<int64_t> effectiveLaneLayout =89      layout.getEffectiveLaneLayoutAsInt();90  assert(static_cast<size_t>(originalType.getRank()) >=91             effectiveLaneLayout.size() &&92         "Rank of the original vector type should be greater or equal to the "93         "size of the lane layout to distribute the vector type.");94  SmallVector<int64_t> distributedShape(originalType.getShape());95  // Only distribute the last `laneLayout.size()` dimensions. The remaining96  // dimensions are not distributed.97  unsigned distributionStart =98      originalType.getRank() - effectiveLaneLayout.size();99  for (auto [i, dim] : llvm::enumerate(originalType.getShape())) {100    if (i < distributionStart)101      continue;102 103    // Check if the dimension can be distributed evenly.104    if (dim % effectiveLaneLayout[i - distributionStart] != 0)105      return failure();106    distributedShape[i] = dim / effectiveLaneLayout[i - distributionStart];107  }108  return VectorType::get(distributedShape, originalType.getElementType());109}110 111/// Helper function to resolve types if the distributed type out of112/// gpu.warp_execute_on_lane0 is different from the expected xegpu SIMT type.113/// Example 1:114///   distributed type: vector<8x1xf32>115///   expected type: vector<8xf32>116///   resolved using,117///   %0 = vector.shape_cast %1 : vector<8x1xf32> to vector<8xf32>118/// Example 2:119///   distributed type: xegpu.tensor_desc<8x16xf32, #xegpu.layout<...>>120///   expected type: xegpu.tensor_desc<8x16xf32>121///   resolved using,122///   %0 = unrealized_conversion_cast %1 :123///      xegpu.tensor_desc<8x16xf32, #xegpu.layout<..>> ->124///      xegpu.tensor_desc<8x16xf32>125template <typename T>126static Value resolveDistributedTy(Value orig, T expected,127                                  PatternRewriter &rewriter) {128  // If orig and expected types are the same, return orig.129  if (orig.getType() == expected)130    return orig;131  // If orig is a vector type, create a shape cast op to reconcile the types.132  if (isa<VectorType>(orig.getType())) {133    auto castOp =134        vector::ShapeCastOp::create(rewriter, orig.getLoc(), expected, orig);135    return castOp.getResult();136  }137  // If orig is a tensor descriptor type, create an unrealized conversion cast138  // op to reconcile the types.139  if (isa<xegpu::TensorDescType>(orig.getType())) {140    auto castOp = UnrealizedConversionCastOp::create(rewriter, orig.getLoc(),141                                                     expected, orig);142    castOp->setAttr(resolveSIMTTypeMismatch, rewriter.getUnitAttr());143    return castOp.getResult(0);144  }145  llvm_unreachable("Unsupported type for reconciliation");146  return orig;147}148 149/// Helper function to check if the layout is packed. Layout is packed if it is150/// 2D and lane_data[0] != 1 (data packed from col dimension).151/// TODO: Move to target info.152static bool requirePacked(const xegpu::LayoutAttr layout) {153  if (!layout)154    return false;155  auto laneData = layout.getEffectiveLaneDataAsInt();156  if (laneData.size() != 2)157    return false;158  return laneData[0] != 1;159}160 161/// Helper function to check if the layout requires a transpose effect.162static bool requireTranspose(const xegpu::LayoutAttr layout,163                             const xegpu::uArch::uArch *uArch) {164  // Return false for unsupported targets.165  // TODO: Add more support or move to target info.166  if (uArch->getName().equals_insensitive("pvc") &&167      uArch->getName().equals_insensitive("bmg"))168    return false;169  if (!layout)170    return false;171  auto laneLayout = layout.getEffectiveLaneLayoutAsInt();172  if (laneLayout.size() != 2)173    return false;174  return laneLayout[0] == uArch->getSubgroupSize() && laneLayout[1] == 1;175}176 177/// Given a vector type and its distributed vector type, return the list of178/// dimensions that are distributed.179static SmallVector<int64_t> getDistributedDims(VectorType originalType,180                                               VectorType distributedType) {181  assert(originalType.getRank() == distributedType.getRank() &&182         "sequential and distributed vector types must have the same rank");183  SmallVector<int64_t> distributedDims;184  for (int64_t i = 0; i < originalType.getRank(); ++i) {185    if (distributedType.getDimSize(i) != originalType.getDimSize(i)) {186      distributedDims.push_back(i);187    }188  }189  return distributedDims;190}191 192/// Given a GPUFuncOp, this pattern creates a new GPUFuncOp and moves the body193/// of the original GPUFuncOp to the new GPUFuncOp such that entire body is194/// contained within a WarpExecuteOnLane0Op.195/// Example:196///197/// ```198///   gpu.func @foo(%arg0: memref<*xf16>) -> vector<8x16xf32> {199///     ...200///     ...201///     gpu.return %result: vector<8x16xf32>202///   }203/// ```204/// To205/// ```206///   gpu.func @foo(%arg0: memref<*xf16>) -> vector<8x16xf32> {207///     %laneid = gpu.lane_id : index208///     %0 = gpu.warp_execute_on_lane_0(%laneid) -> vector<8x16xf32> {209///       ...210///       ...211///       gpu.yield %result: vector<8x16xf32>212///     }213///     return %0214///   }215struct MoveFuncBodyToWarpOp : public OpRewritePattern<gpu::GPUFuncOp> {216  using OpRewritePattern<gpu::GPUFuncOp>::OpRewritePattern;217  LogicalResult matchAndRewrite(gpu::GPUFuncOp gpuFuncOp,218                                PatternRewriter &rewriter) const override {219    auto uArch = getUArch(xegpu::getChipStr(gpuFuncOp).value_or(""));220    if (!uArch)221      return rewriter.notifyMatchFailure(222          gpuFuncOp, "Subgroup distribution requires target attribute attached "223                     "to set the warp size");224    // If the function only contains a single void return, skip.225    if (llvm::all_of(gpuFuncOp.getBody().getOps(), [](Operation &op) {226          return isa<gpu::ReturnOp>(op) && !op.getNumOperands();227        }))228      return failure();229    // If the function already moved inside a warp_execute_on_lane0, skip.230    if (llvm::any_of(gpuFuncOp.getBody().getOps(), [](Operation &op) {231          return isa<gpu::WarpExecuteOnLane0Op>(op);232        }))233      return failure();234    // Create a new function with the same signature and same attributes.235    SmallVector<Type> workgroupAttributionsTypes =236        llvm::map_to_vector(gpuFuncOp.getWorkgroupAttributions(),237                            [](BlockArgument arg) { return arg.getType(); });238    SmallVector<Type> privateAttributionsTypes =239        llvm::map_to_vector(gpuFuncOp.getPrivateAttributions(),240                            [](BlockArgument arg) { return arg.getType(); });241    auto newGpuFunc = gpu::GPUFuncOp::create(242        rewriter, gpuFuncOp.getLoc(), gpuFuncOp.getName(),243        gpuFuncOp.getFunctionType(), workgroupAttributionsTypes,244        privateAttributionsTypes);245    newGpuFunc->setAttrs(gpuFuncOp->getAttrs());246    // Create a WarpExecuteOnLane0Op with same arguments and results as the247    // original gpuFuncOp.248    rewriter.setInsertionPointToEnd(&newGpuFunc.getFunctionBody().front());249    auto laneId = gpu::LaneIdOp::create(250        rewriter, newGpuFunc.getLoc(), rewriter.getIndexType(),251        /** upperBound = **/ mlir::IntegerAttr());252    ArrayRef<Type> gpuFuncResultType = gpuFuncOp.getFunctionType().getResults();253    auto warpOp = gpu::WarpExecuteOnLane0Op::create(254        rewriter, laneId.getLoc(), gpuFuncResultType, laneId,255        uArch->getSubgroupSize(), newGpuFunc.getArguments(),256        newGpuFunc.getArgumentTypes());257    Block &warpBodyBlock = warpOp.getBodyRegion().front();258    // Replace the ReturnOp of the original gpu function with a YieldOp.259    auto origRetunOp =260        cast<gpu::ReturnOp>(gpuFuncOp.getBlocks().back().getTerminator());261    rewriter.setInsertionPointAfter(origRetunOp);262    gpu::YieldOp::create(rewriter, origRetunOp.getLoc(),263                         origRetunOp.getOperands());264    rewriter.eraseOp(origRetunOp);265    // Move the original function body to the WarpExecuteOnLane0Op body.266    rewriter.inlineRegionBefore(gpuFuncOp.getBody(), warpOp.getBodyRegion(),267                                warpOp.getBodyRegion().begin());268    rewriter.eraseBlock(&warpBodyBlock);269    // Insert a new ReturnOp after the WarpExecuteOnLane0Op.270    rewriter.setInsertionPointAfter(warpOp);271    gpu::ReturnOp::create(rewriter, newGpuFunc.getLoc(), warpOp.getResults());272    rewriter.replaceOp(gpuFuncOp, newGpuFunc);273    return success();274  }275};276 277/// Distribute a create_nd_tdesc feeding into vector.yield op of the enclosing278/// `gpu.warp_execute_on_lane_0` region. After the sinking, the warp op will279/// still contain the original op that will not be used by the yield op (and280/// should be cleaned up later). The yield op will bypass the create_nd_tdesc's281/// arguments. Tensor descriptor shape is not distributed because it is a282/// uniform value across all work items within the subgroup. However, the283/// layout information is dropped in the new tensor descriptor type.284///285/// Example:286///287/// ```288///   #layout0 = #xegpu.layout<wi_layout = [1, 8], wi_data = [1, 1]>289///   %r = gpu.warp_execute_on_lane_0(%laneid) ->290///                   (!xegpu.tensor_desc<4x8xf32, #layout0>) {291///     ...292///     %td = xegpu.create_nd_tdesc %arg0293///               : memref<4x8xf32> -> !xegpu.tensor_desc<4x8xf32, #layout0>294///     vector.yield %td295///   }296/// ```297/// To298/// ```299///   %r:2 = gpu.warp_execute_on_lane_0(%laneid) -> (...) {300///     ...301///     %dead = xegpu.create_nd_tdesc %arg0302///               : memref<4x8xf32> -> !xegpu.tensor_desc<4x8xf32, #layout0>303///     vector.yield %arg0, %dead304///   }305///   %td = xegpu.create_nd_tdesc %r#0: memref<4x8xf32>306///                                 -> !xegpu.tensor_desc<4x8xf32>307///308/// ```309struct CreateNdDescDistribution final : public gpu::WarpDistributionPattern {310  using gpu::WarpDistributionPattern::WarpDistributionPattern;311  LogicalResult matchAndRewrite(gpu::WarpExecuteOnLane0Op warpOp,312                                PatternRewriter &rewriter) const override {313    OpOperand *operand =314        getWarpResult(warpOp, llvm::IsaPred<xegpu::CreateNdDescOp>);315    if (!operand)316      return rewriter.notifyMatchFailure(317          warpOp, "warp result is not a xegpu::CreateNdDesc op");318    auto descOp = operand->get().getDefiningOp<xegpu::CreateNdDescOp>();319    unsigned operandIdx = operand->getOperandNumber();320 321    xegpu::LayoutAttr layout = descOp.getType().getLayoutAttr();322    if (!layout)323      return rewriter.notifyMatchFailure(324          descOp, "the tensor descriptor lacks layout attribute");325    // CreateNdOp must not have offsets.326    if (descOp.getMixedOffsets().size())327      return rewriter.notifyMatchFailure(328          descOp, "xegpu::CreateNdDescOp must not have offsets");329 330    SmallVector<size_t> newRetIndices;331    rewriter.setInsertionPoint(warpOp);332    gpu::WarpExecuteOnLane0Op newWarpOp = moveRegionToNewWarpOpAndAppendReturns(333        rewriter, warpOp, /* new yieled values = */ descOp->getOperands(),334        /* new yielded types = */ descOp.getOperandTypes(), newRetIndices);335 336    SmallVector<Value> newDescOperands = llvm::map_to_vector(337        newRetIndices, [&](size_t i) { return newWarpOp.getResult(i); });338    rewriter.setInsertionPointAfter(newWarpOp);339    xegpu::TensorDescType distributedTensorDescTy =340        descOp.getType().dropLayouts(); // Distributed tensor descriptor type341                                        // does not contain layout info.342    Value newDescOp = xegpu::CreateNdDescOp::create(343        rewriter, newWarpOp.getLoc(), distributedTensorDescTy, newDescOperands,344        descOp->getAttrs());345 346    Value distributedVal = newWarpOp.getResult(operandIdx);347    // Resolve the distributed type to the expected type.348    newDescOp =349        resolveDistributedTy(newDescOp, distributedVal.getType(), rewriter);350    rewriter.replaceAllUsesWith(distributedVal, newDescOp);351    return success();352  }353};354 355/// Distribute a store_nd op at the end of enclosing356/// `gpu.warp_execute_on_lane_0`. In case arguments for the store are passed357/// through the warp op interface they would be propagated as returned values.358/// Source vector is distributed based on lane layout. Appropriate cast ops are359/// inserted if the distributed types does not match expected xegpu SIMT types.360///361/// Example:362///363/// ```364///   #layout0 = #xegpu.layout<wi_layout = [1, 8], wi_data = [1, 1]>365///   gpu.warp_execute_on_lane_0(%laneid) -> () {366///     ...367///     xegpu.store_nd %arg0, %arg1 [%x, %y]: vector<4x8xf32>,368///                                 !xegpu.tensor_desc<4x8xf32, #layout0>369///   }370/// ```371/// To372/// ```373///   %r:2 = gpu.warp_execute_on_lane_0(%laneid) -> (vector<4x1xf32>,374///   !xegpu.tensor_desc<4x8xf32, #layout0>, index, index) {375///     ...376///     gpu.yield %arg0, %arg1, %x, %y: vector<4x8xf32>,377///     !xegpu.tensor_desc<4x8xf32, #layout0>, index, index378///   }379///   %0 = vector.shape_cast %r#0: vector<4x1xf32> to vector<4xf32>380///   %1 = unrealized_conversion_cast %r#1: !xegpu.tensor_desc<4x8xf32,381///   #layout0>382///     -> !xegpu.tensor_desc<4x8xf32>383///   xegpu.store_nd %0, %1 [%r#2, %r#3]: vector<4xf32>,384///     !xegpu.tensor_desc<4x8xf32>385///386/// ```387struct StoreNdDistribution final : public gpu::WarpDistributionPattern {388  using gpu::WarpDistributionPattern::WarpDistributionPattern;389  LogicalResult matchAndRewrite(gpu::WarpExecuteOnLane0Op warpOp,390                                PatternRewriter &rewriter) const override {391    gpu::YieldOp yield = warpOp.getTerminator();392    Operation *lastNode = yield->getPrevNode();393    auto storeOp = dyn_cast_or_null<xegpu::StoreNdOp>(lastNode);394    if (!storeOp)395      return failure();396 397    SmallVector<OpFoldResult> offsets = storeOp.getMixedOffsets();398    // Expecting offsets to be present.399    if (offsets.empty())400      return rewriter.notifyMatchFailure(storeOp,401                                         "the store op must have offsets");402    SmallVector<Value> offsetsAsValues =403        vector::getAsValues(rewriter, storeOp.getLoc(), offsets);404    SmallVector<Type> offsetTypes = llvm::to_vector(405        llvm::map_range(offsetsAsValues, [](Value v) { return v.getType(); }));406    xegpu::TensorDescType tensorDescTy = storeOp.getTensorDescType();407    xegpu::LayoutAttr layout = tensorDescTy.getLayoutAttr();408    if (!layout)409      return rewriter.notifyMatchFailure(410          storeOp, "the source tensor descriptor lacks layout attribute");411 412    FailureOr<VectorType> distributedTypeByWarpOpOrFailure =413        getDistVecTypeBasedOnLaneLayout(layout, storeOp.getValueType());414    if (failed(distributedTypeByWarpOpOrFailure))415      return rewriter.notifyMatchFailure(storeOp,416                                         "Failed to distribute the type");417    VectorType distributedTypeByWarpOp =418        distributedTypeByWarpOpOrFailure.value();419 420    SmallVector<size_t> newRetIndices;421    SmallVector<Value> newYieldedValues = {storeOp.getValue(),422                                           storeOp.getTensorDesc()};423    SmallVector<Type> newYieldedTypes = {distributedTypeByWarpOp, tensorDescTy};424    newYieldedValues.append(offsetsAsValues.begin(), offsetsAsValues.end());425    newYieldedTypes.append(offsetTypes.begin(), offsetTypes.end());426    gpu::WarpExecuteOnLane0Op newWarpOp = moveRegionToNewWarpOpAndAppendReturns(427        rewriter, warpOp, newYieldedValues, newYieldedTypes, newRetIndices);428    // Create a new store op outside the warp op with the distributed vector429    // type. Tensor descriptor is not distributed.430    rewriter.setInsertionPointAfter(newWarpOp);431    SmallVector<Value> newStoreOperands;432 433    // For the value operand, there can be a mismatch between the vector type434    // distributed by the warp op and (xegpu-specific) distributed type435    // supported by the store op. Type mismatch must be resolved using436    // appropriate cast op.437    FailureOr<VectorType> storeNdDistributedValueTyOrFailure =438        xegpu::getDistributedVectorType(storeOp.getTensorDescType());439    if (failed(storeNdDistributedValueTyOrFailure))440      return rewriter.notifyMatchFailure(441          storeOp, "Failed to get distributed vector type for the store op");442    newStoreOperands.push_back(resolveDistributedTy(443        newWarpOp.getResult(newRetIndices[0]),444        storeNdDistributedValueTyOrFailure.value(), rewriter));445    // For the tensor descriptor operand, the layout attribute is dropped after446    // distribution. Types needs to be resolved in this case also.447    xegpu::TensorDescType distributedTensorDescTy =448        storeOp.getTensorDescType().dropLayouts();449    newStoreOperands.push_back(450        resolveDistributedTy(newWarpOp.getResult(newRetIndices[1]),451                             distributedTensorDescTy, rewriter));452    // Collect offsets.453    for (size_t i = 2; i < newRetIndices.size(); ++i)454      newStoreOperands.push_back(newWarpOp.getResult(newRetIndices[i]));455 456    auto newStoreOp =457        xegpu::StoreNdOp::create(rewriter, newWarpOp.getLoc(), TypeRange{},458                                 newStoreOperands, storeOp->getAttrs());459    xegpu::removeLayoutAttrs(newStoreOp);460    rewriter.eraseOp(storeOp);461    return success();462  }463};464 465/// Distribute a load_nd op feeding into vector.yield op for the enclosing466/// `gpu.warp_execute_on_lane_0` and put it after the warp op.467/// The warp op will still contain the original op that will not be used by468/// the yield op (and should be cleaned up later). The yield op will469/// bypass the load's arguments. Only the loaded vector is distributed470/// according to lane layout and, tensor descriptor types is not471/// distributed. Appropriate cast ops are inserted if the distributed types does472/// not match expected xegpu SIMT types.473///474/// Example:475///476/// ```477///   #layout0 = #xegpu.layout<wi_layout = [1, 8], wi_data = [1, 1]>478///   %r = gpu.warp_execute_on_lane_0(%laneid) ->479///                   (vector<4x1xf32>) {480///     ...481///     %ld = xegpu.load_nd %arg0, %arg1: !xegpu.tensor_desc<4x8xf32, #layout0>482///     ->483///       vector<4x8xf32>484///     gpu.yield %ld485///   }486/// ```487/// To488/// ```489///   %r:2 = gpu.warp_execute_on_lane_0(%laneid) -> (vector<4x1xf32>,490///   !xegpu.tensor_desc<4x8xf32, #layout0>) {491///     ...492///     %dead = xegpu.load_nd %arg0: !xegpu.tensor_desc<4x8xf32, #layout0> ->493///     vector<4x8xf32> gpu.yield %dead, %arg0494///   }495///   %0 = unrealized_conversion_cast %r#1: !xegpu.tensor_desc<4x8xf32,496///        #layout0> -> !xegpu.tensor_desc<4x8xf32>497///   %1 = xegpu.load_nd %0: !xegpu.tensor_desc<4x8xf32> -> vector<4xf32>498///   %2 = vector.shape_cast %r#0: vector<4xf32> to vector<4x1xf32>499///500/// ```501struct LoadNdDistribution final : public gpu::WarpDistributionPattern {502  using gpu::WarpDistributionPattern::WarpDistributionPattern;503  LogicalResult matchAndRewrite(gpu::WarpExecuteOnLane0Op warpOp,504                                PatternRewriter &rewriter) const override {505    OpOperand *operand = getWarpResult(warpOp, [&](Operation *op) {506      if (!isa<xegpu::LoadNdOp>(op))507        return false;508      // Make sure the same load op is the last operation in the warp op body.509      // This ensure that load op is not sinked earlier violating any barrier510      // synchronizations.511      gpu::YieldOp yield = warpOp.getTerminator();512      return yield->getPrevNode() == op;513    });514 515    if (!operand)516      return rewriter.notifyMatchFailure(517          warpOp, "warp result is not a xegpu::LoadNd op");518 519    auto loadOp = operand->get().getDefiningOp<xegpu::LoadNdOp>();520    auto uArch = getUArch(xegpu::getChipStr(loadOp).value_or(""));521    if (!uArch)522      return rewriter.notifyMatchFailure(523          loadOp, "xegpu::LoadNdOp require target attribute attached to "524                  "determine transpose "525                  "requirement");526    // Chip information is required to decide if the layout requires transpose527    // effect.528    // Expecting offsets to be present.529    SmallVector<OpFoldResult> offsets = loadOp.getMixedOffsets();530    if (offsets.empty())531      return rewriter.notifyMatchFailure(loadOp,532                                         "the load op must have offsets");533    SmallVector<Value> offsetsAsValues =534        vector::getAsValues(rewriter, loadOp.getLoc(), offsets);535    SmallVector<Type> offsetTypes = llvm::to_vector(536        llvm::map_range(offsetsAsValues, [](Value v) { return v.getType(); }));537 538    xegpu::TensorDescType tensorDescTy = loadOp.getTensorDescType();539    xegpu::LayoutAttr layout = tensorDescTy.getLayoutAttr();540    if (!layout)541      return rewriter.notifyMatchFailure(542          loadOp, "the source tensor descriptor lacks layout attribute");543 544    unsigned operandIdx = operand->getOperandNumber();545    VectorType distributedTypeByWarpOp =546        cast<VectorType>(warpOp.getResult(operandIdx).getType());547 548    SmallVector<size_t> newRetIndices;549    SmallVector<Value> newYieldedValues = {loadOp.getTensorDesc()};550    SmallVector<Type> newYieldedTypes = {tensorDescTy};551    newYieldedValues.append(offsetsAsValues.begin(), offsetsAsValues.end());552    newYieldedTypes.append(offsetTypes.begin(), offsetTypes.end());553    gpu::WarpExecuteOnLane0Op newWarpOp = moveRegionToNewWarpOpAndAppendReturns(554        rewriter, warpOp, newYieldedValues, newYieldedTypes, newRetIndices);555 556    // Create a new load op outside the warp op with the distributed vector557    // type.558    rewriter.setInsertionPointAfter(newWarpOp);559    FailureOr<VectorType> loadNdDistValueTyOrFailure =560        xegpu::getDistributedVectorType(loadOp.getTensorDescType());561    if (failed(loadNdDistValueTyOrFailure))562      return rewriter.notifyMatchFailure(563          loadOp, "Failed to get distributed vector type for the load op");564    xegpu::TensorDescType distributedTensorDescTy =565        loadOp.getTensorDescType().dropLayouts(); // Distributed tensor566                                                  // descriptor type does not567                                                  // contain layout info.568    SmallVector<Value> newLoadOperands{569        resolveDistributedTy(newWarpOp.getResult(newRetIndices[0]),570                             distributedTensorDescTy, rewriter)};571    // Collect offsets.572    for (size_t i = 1; i < newRetIndices.size(); ++i)573      newLoadOperands.push_back(newWarpOp.getResult(newRetIndices[i]));574    auto newLoadOp = xegpu::LoadNdOp::create(575        rewriter, newWarpOp.getLoc(), loadNdDistValueTyOrFailure.value(),576        newLoadOperands, loadOp->getAttrs());577    xegpu::removeLayoutAttrs(newLoadOp);578    // Set the packed attribute if the layout requires it.579    newLoadOp.setPacked(requirePacked(layout));580    // Set the transpose attribute if the layout requires it.581    if (requireTranspose(layout, uArch))582      newLoadOp.setTranspose(583          DenseI64ArrayAttr::get(rewriter.getContext(), {1, 0}));584    Value distributedVal = newWarpOp.getResult(operandIdx);585    // There can be a conflict between the vector type distributed by the586    // warp op and (xegpu-specific) distributed type supported by the load587    // op. Resolve these mismatches by inserting a cast.588    Value tyResolvedVal = resolveDistributedTy(589        newLoadOp.getResult(), distributedTypeByWarpOp, rewriter);590    rewriter.replaceAllUsesWith(distributedVal, tyResolvedVal);591    return success();592  }593};594 595/// Distribute a dpas op feeding into vector.yield op for the enclosing596/// `gpu.warp_execute_on_lane_0` and put it after the warp op.597/// The warp op will still contain the original op that will not be used by598/// the yield op (and should be cleaned up later). The yield op will599/// bypass the dpas's arguments. Appropriate cast ops are inserted if the600/// distributed types does not match expected xegpu SIMT types.601/// Example:602/// ```603///   #lo_a = #xegpu.layout<wi_layout = [1, 16], wi_data = [1, 1]>604///   #lo_b = #xegpu.layout<wi_layout = [1, 16], wi_data = [2, 1]>605///   #lo_c = #xegpu.layout<wi_layout = [1, 16], wi_data = [1, 1]>606///   %r = gpu.warp_execute_on_lane_0(%laneid) ->607///                   (vector<8x1xf32>) {608///     ...609///     %dpas = xegpu.dpas %arg0, %arg1: vector<8x16xf16>, vector<16x16xf16> ->610///       vector<8x16xf32>611///     gpu.yield %dpas612///   }613/// ```614/// To615/// ```616///   %r:2 = gpu.warp_execute_on_lane_0(%laneid) -> (vector<8x1xf32>,617///   vector<8x1xf16>, vector<16x1xf16>) {618///     ...619///     %dead = xegpu.dpas %arg0, %arg1: vector<8x16xf16>, vector<16x16xf16>620///       -> vector<8x16xf32>621///     gpu.yield %dead, %arg0, %arg1622///   }623///   %0 = vector.shape_cast %r#1: vector<8x1xf16> to vector<8xf16>624///   %1 = vector.shape_cast %r#2: vector<16x1xf16> to vector<16xf16>625///   %2 = xegpu.dpas %0, %1: vector<8xf16>, vector<16xf16> ->626///     vector<8xf32>627///   %dpas = vector.shape_cast %2: vector<8xf32> to vector<8x1xf32>628/// ```629struct DpasDistribution final : public gpu::WarpDistributionPattern {630  using gpu::WarpDistributionPattern::WarpDistributionPattern;631  LogicalResult matchAndRewrite(gpu::WarpExecuteOnLane0Op warpOp,632                                PatternRewriter &rewriter) const override {633    OpOperand *operand = getWarpResult(warpOp, llvm::IsaPred<xegpu::DpasOp>);634    if (!operand)635      return rewriter.notifyMatchFailure(warpOp,636                                         "warp result is not a xegpu::Dpas op");637 638    auto dpasOp = operand->get().getDefiningOp<xegpu::DpasOp>();639    unsigned operandIdx = operand->getOperandNumber();640    std::string layoutAName = xegpu::getLayoutName(dpasOp->getOpOperand(0));641    std::string layoutBName = xegpu::getLayoutName(dpasOp->getOpOperand(1));642    std::string layoutCName = xegpu::getLayoutName(dpasOp->getOpResult(0));643 644    xegpu::LayoutAttr layoutA =645        dpasOp->getAttrOfType<xegpu::LayoutAttr>(layoutAName);646    xegpu::LayoutAttr layoutB =647        dpasOp->getAttrOfType<xegpu::LayoutAttr>(layoutBName);648    xegpu::LayoutAttr layoutOut =649        dpasOp->getAttrOfType<xegpu::LayoutAttr>(layoutCName);650    if (!layoutA || !layoutB || !layoutOut)651      return rewriter.notifyMatchFailure(652          dpasOp,653          "the xegpu::Dpas op lacks layout attribute for A, B or output");654 655    FailureOr<VectorType> distLhsTypeByWarpOpOrFailure =656        getDistVecTypeBasedOnLaneLayout(layoutA, dpasOp.getLhsType());657    FailureOr<VectorType> distRhsTypeByWarpOpOrFailure =658        getDistVecTypeBasedOnLaneLayout(layoutB, dpasOp.getRhsType());659    FailureOr<VectorType> distResultTypeByWarpOpOrFailure =660        getDistVecTypeBasedOnLaneLayout(layoutOut, dpasOp.getResultType());661    if (failed(distLhsTypeByWarpOpOrFailure) ||662        failed(distRhsTypeByWarpOpOrFailure) ||663        failed(distResultTypeByWarpOpOrFailure))664      return rewriter.notifyMatchFailure(665          dpasOp,666          "Failed to distribute the A, B or output types in xegpu::Dpas op");667 668    llvm::SmallVector<Value, 3> newYieldValues{dpasOp.getLhs(),669                                               dpasOp.getRhs()};670    llvm::SmallVector<Type, 3> newYieldTypes{671        distLhsTypeByWarpOpOrFailure.value(),672        distRhsTypeByWarpOpOrFailure.value()};673    // Dpas acc operand is optional.674    if (dpasOp.getAcc()) {675      newYieldValues.push_back(dpasOp.getAcc());676      newYieldTypes.push_back(distResultTypeByWarpOpOrFailure.value());677    }678    // Create a new warp op without the dpas.679    SmallVector<size_t> newRetIndices;680    gpu::WarpExecuteOnLane0Op newWarpOp = moveRegionToNewWarpOpAndAppendReturns(681        rewriter, warpOp, newYieldValues, newYieldTypes, newRetIndices);682 683    FailureOr<VectorType> expectedDistLhsTyOrFailure =684        xegpu::getDistributedVectorType(dpasOp.getLhsType(), layoutA);685    FailureOr<VectorType> expectedDistRhsTyOrFailure =686        xegpu::getDistributedVectorType(dpasOp.getRhsType(), layoutB);687    FailureOr<VectorType> expectedDistResultTyOrFailure =688        xegpu::getDistributedVectorType(dpasOp.getResultType(), layoutOut);689    if (failed(expectedDistLhsTyOrFailure) ||690        failed(expectedDistRhsTyOrFailure) ||691        failed(expectedDistResultTyOrFailure))692      return rewriter.notifyMatchFailure(693          dpasOp,694          "Failed to get distributed vector type for the dpas operands.");695    // Create a new dpas op outside the warp op.696    rewriter.setInsertionPointAfter(newWarpOp);697    SmallVector<Value> newDpasOperands;698    SmallVector<VectorType> newDpasOperandExpectedTypes;699 700    // Resolve the distributed types with the original types.701    newDpasOperandExpectedTypes.push_back(expectedDistLhsTyOrFailure.value());702    newDpasOperandExpectedTypes.push_back(expectedDistRhsTyOrFailure.value());703    VectorType distributedResultTy = expectedDistResultTyOrFailure.value();704    if (dpasOp.getAcc())705      newDpasOperandExpectedTypes.push_back(distributedResultTy);706 707    for (unsigned i = 0; i < newRetIndices.size(); i++) {708      newDpasOperands.push_back(709          resolveDistributedTy(newWarpOp.getResult(newRetIndices[i]),710                               newDpasOperandExpectedTypes[i], rewriter));711    }712    auto newDpasOp = xegpu::DpasOp::create(rewriter, newWarpOp->getLoc(),713                                           distributedResultTy, newDpasOperands,714                                           dpasOp->getAttrs());715    xegpu::removeLayoutAttrs(newDpasOp);716    Value distributedVal = newWarpOp.getResult(operandIdx);717    // Resolve the output type.718    Value typeResolved =719        resolveDistributedTy(newDpasOp.getResult(),720                             distResultTypeByWarpOpOrFailure.value(), rewriter);721    rewriter.replaceAllUsesWith(distributedVal, typeResolved);722    return success();723  }724};725 726/// Distribute a prefetch_nd op at the end of enclosing727/// `gpu.warp_execute_on_lane_0`. In case arguments for the prefetch are passed728/// through the warp op interface they would be propagated as returned values.729/// Tensor descriptor shape is not distributed because it is a uniform value730/// across all work items within the subgroup. Appropriate cast ops are inserted731/// if the distributed types does not match expected xegpu SIMT types.732///733/// Example:734///735/// ```736///   #layout0 = #xegpu.layout<wi_layout = [1, 8], wi_data = [1, 1]>737///   gpu.warp_execute_on_lane_0(%laneid) -> () {738///     ...739///     xegpu.prefetch_nd %arg0 [%x, %y] : !xegpu.tensor_desc<4x8xf32, #layout0>740///   }741/// ```742/// To743/// ```744///   %r:1 = gpu.warp_execute_on_lane_0(%laneid) -> (745///    !xegpu.tensor_desc<4x8xf32, #layout0>, index, index) {746///     gpu.yield %arg0, %x, %y: !xegpu.tensor_desc<4x8xf32, #layout0>, index,747///     index748///   }749///   %1 = unrealized_conversion_cast %r#0: !xegpu.tensor_desc<4x8xf32,750///     #layout0> -> !xegpu.tensor_desc<4x8xf32>751///   xegpu.prefetch_nd %1 [%r#1, %r#2] : !xegpu.tensor_desc<4x8xf32>752///753/// ```754struct PrefetchNdDistribution final : public gpu::WarpDistributionPattern {755  using gpu::WarpDistributionPattern::WarpDistributionPattern;756  LogicalResult matchAndRewrite(gpu::WarpExecuteOnLane0Op warpOp,757                                PatternRewriter &rewriter) const override {758    gpu::YieldOp yield = warpOp.getTerminator();759    Operation *lastNode = yield->getPrevNode();760    auto prefetchOp = dyn_cast_or_null<xegpu::PrefetchNdOp>(lastNode);761    if (!prefetchOp)762      return failure();763 764    SmallVector<OpFoldResult> offsets = prefetchOp.getMixedOffsets();765    // PrefetchNdOp must have offsets.766    if (offsets.empty())767      return rewriter.notifyMatchFailure(prefetchOp,768                                         "the prefetch op must have offsets");769    SmallVector<Value> offsetsAsValues =770        vector::getAsValues(rewriter, prefetchOp.getLoc(), offsets);771    SmallVector<Type> offsetTypes = llvm::to_vector(772        llvm::map_range(offsetsAsValues, [](Value v) { return v.getType(); }));773 774    xegpu::LayoutAttr layout = prefetchOp.getTensorDescType().getLayoutAttr();775    if (!layout)776      return rewriter.notifyMatchFailure(777          prefetchOp, "the source tensor descriptor lacks layout attribute");778 779    SmallVector<Value> newYieldValues = {prefetchOp.getTensorDesc()};780    SmallVector<Type> newYieldTypes = {prefetchOp.getTensorDescType()};781    newYieldValues.append(offsetsAsValues.begin(), offsetsAsValues.end());782    newYieldTypes.append(offsetTypes.begin(), offsetTypes.end());783    SmallVector<size_t> newRetIndices;784    gpu::WarpExecuteOnLane0Op newWarpOp = moveRegionToNewWarpOpAndAppendReturns(785        rewriter, warpOp, newYieldValues, newYieldTypes, newRetIndices);786    // Create a new prefetch op outside the warp op with updated tensor787    // descriptor type. Source tensor descriptor require type resolution.788    xegpu::TensorDescType newTensorDescTy =789        prefetchOp.getTensorDescType().dropLayouts();790    rewriter.setInsertionPointAfter(newWarpOp);791    SmallVector<Value> newPrefetchOperands = {resolveDistributedTy(792        newWarpOp.getResult(newRetIndices[0]), newTensorDescTy, rewriter)};793    // Collect offsets.794    for (size_t i = 1; i < newRetIndices.size(); ++i)795      newPrefetchOperands.push_back(newWarpOp.getResult(newRetIndices[i]));796    xegpu::PrefetchNdOp::create(rewriter, newWarpOp.getLoc(), TypeRange{},797                                newPrefetchOperands, prefetchOp->getAttrs());798    xegpu::removeLayoutAttrs(prefetchOp);799    rewriter.eraseOp(prefetchOp);800    return success();801  }802};803 804/// Sink a gpu::BarrierOp at the end of enclosing `gpu.warp_execute_on_lane_0`805/// region. This will simply move the barrier op outside of the warp op.806struct GpuBarrierDistribution final : public gpu::WarpDistributionPattern {807  using gpu::WarpDistributionPattern::WarpDistributionPattern;808  LogicalResult matchAndRewrite(gpu::WarpExecuteOnLane0Op warpOp,809                                PatternRewriter &rewriter) const override {810    gpu::YieldOp yield = warpOp.getTerminator();811    Operation *lastNode = yield->getPrevNode();812    // The last node must be a gpu::BarrierOp.813    auto barrierOp = dyn_cast_or_null<gpu::BarrierOp>(lastNode);814    if (!barrierOp)815      return failure();816    // Move the barrier op outside of the warp op.817    rewriter.setInsertionPointAfter(warpOp);818    gpu::BarrierOp::create(rewriter, barrierOp.getLoc(),819                           barrierOp->getResultTypes(),820                           barrierOp->getOperands(), barrierOp->getAttrs());821    rewriter.eraseOp(barrierOp);822    return success();823  }824};825 826/// Distribute a scattered store op. The offsets argument is required.827/// Both offset and mask vectors must be 1D and have #subgroup_size elements.828/// The layouts are fixed and implicit: one offset/mask per lane.829/// The pass changes the offset/mask vector shapes to a830/// single-element vector, **it is assumed that their producer will also be831/// distributed**. The payload vector also has a fixed distribution:832///   no chunk size -> vector of one element.833///   chunk size    -> vector of the innermost dimension of the SG-payload.834/// Example 1 (no chunk size):835///    %mask = producer_op : vector<16xi1>836///    %offset = producer_op : vector<16xindex>837///    xegpu.store %payload, %src[%offset], %mask : vector<16xf16>,838///     memref<256xf16>, vector<16xindex>, vector<16xi1>839/// To840///    %mask = producer_op : vector<1xi1>841///    %offset = producer_op : vector<1xindex>842///    xegpu.store %payload, %src[%offset], %mask : vector<1xf16>,843///     memref<256xf16>, vector<1xindex>, vector<1xi1>844/// Example 2 (chunk size, same mask and offsets):845///    xegpu.store %payload, %src[%offset], %mask <{chunk_size=8}> :846///     vector<16x8xf16>, memref<256xf16>, vector<16xindex>, vector<16xi1>847/// To848///    xegpu.store %payload, %src[%offset], %mask <{chunk_size=8}> :849///     vector<8xf16>, memref<256xf16>, vector<1xindex>, vector<1xi1>850struct StoreDistribution final : public gpu::WarpDistributionPattern {851  using gpu::WarpDistributionPattern::WarpDistributionPattern;852  LogicalResult matchAndRewrite(gpu::WarpExecuteOnLane0Op warpOp,853                                PatternRewriter &rewriter) const override {854    Operation *lastNode = warpOp.getTerminator()->getPrevNode();855    auto storeScatterOp = dyn_cast_or_null<xegpu::StoreScatterOp>(lastNode);856    if (!storeScatterOp)857      return failure();858    auto offsets = storeScatterOp.getOffsets();859    if (!offsets || !isa<VectorType>(offsets.getType()))860      return rewriter.notifyMatchFailure(861          storeScatterOp, "Store op must have a vector of offsets argument");862    VectorType offsetsTy = cast<VectorType>(offsets.getType());863    VectorType maskTy = cast<VectorType>(storeScatterOp.getMask().getType());864    if (offsetsTy.getRank() != 1 || maskTy.getRank() != 1)865      return rewriter.notifyMatchFailure(storeScatterOp,866                                         "Expected 1D offsets and mask vector");867    VectorType storeVecTy = cast<VectorType>(storeScatterOp.getValueType());868    if (storeVecTy.getRank() > 2)869      return rewriter.notifyMatchFailure(870          storeScatterOp, "Expected at most 2D result at SG level");871 872    std::string layoutPayloadName =873        xegpu::getLayoutName(storeScatterOp->getOpOperand(0));874    std::string layoutOffsetsName =875        xegpu::getLayoutName(storeScatterOp->getOpOperand(2));876    std::string layoutMaskName =877        xegpu::getLayoutName(storeScatterOp->getOpOperand(3));878 879    xegpu::LayoutAttr layoutPayload =880        storeScatterOp->getAttrOfType<xegpu::LayoutAttr>(layoutPayloadName);881    xegpu::LayoutAttr layoutOffsets =882        storeScatterOp->getAttrOfType<xegpu::LayoutAttr>(layoutOffsetsName);883    xegpu::LayoutAttr layoutMask =884        storeScatterOp->getAttrOfType<xegpu::LayoutAttr>(layoutMaskName);885 886    FailureOr<VectorType> distStoreVecByWarpOpOrFailure =887        getDistVecTypeBasedOnLaneLayout(layoutPayload, storeVecTy);888    FailureOr<VectorType> distOffsetsByWarpOpOrFailure =889        getDistVecTypeBasedOnLaneLayout(layoutOffsets, offsetsTy);890    FailureOr<VectorType> distMaskByWarpOpOrFailure =891        getDistVecTypeBasedOnLaneLayout(layoutMask, maskTy);892    if (failed(distStoreVecByWarpOpOrFailure) ||893        failed(distOffsetsByWarpOpOrFailure) ||894        failed(distMaskByWarpOpOrFailure)) {895      return rewriter.notifyMatchFailure(896          storeScatterOp,897          "Some vector operands have no layouts, using defaults instead.");898    }899    // Distributed store payload type according to the lane layout.900    VectorType distPayloadTyByWarpOp = distStoreVecByWarpOpOrFailure.value();901    // Expected distributed payload type is always 1D.902    VectorType expectedPayloadTy =903        VectorType::get({distPayloadTyByWarpOp.getNumElements()},904                        distPayloadTyByWarpOp.getElementType());905 906    SmallVector<size_t> newRetIndices;907    SmallVector<Value> operands = storeScatterOp->getOperands();908    SmallVector<Type> operandTypesToYield = {909        distPayloadTyByWarpOp, operands[1].getType(),910        distOffsetsByWarpOpOrFailure.value(),911        distMaskByWarpOpOrFailure.value()};912 913    gpu::WarpExecuteOnLane0Op newWarpOp = moveRegionToNewWarpOpAndAppendReturns(914        rewriter, warpOp, operands, operandTypesToYield, newRetIndices);915    SmallVector<Value> newStoreScatterOpOperands = llvm::map_to_vector(916        newRetIndices, [&](size_t idx) { return newWarpOp.getResult(idx); });917    // The payload operand may need type adjustment due to mismatch between warp918    // distributed type and expected SIMT type.919    rewriter.setInsertionPointAfter(newWarpOp);920    newStoreScatterOpOperands[0] = resolveDistributedTy(921        newStoreScatterOpOperands[0], expectedPayloadTy, rewriter);922    xegpu::StoreScatterOp newOp = xegpu::StoreScatterOp::create(923        rewriter, newWarpOp.getLoc(), TypeRange{}, newStoreScatterOpOperands,924        storeScatterOp->getAttrs());925    xegpu::removeLayoutAttrs(newOp);926    rewriter.eraseOp(storeScatterOp);927    return success();928  }929};930 931static SmallVector<Value> computeDistributedCoordinatesForMatrixOp(932    PatternRewriter &rewriter, Location loc, xegpu::DistributeLayoutAttr layout,933    Value laneId, ArrayRef<int64_t> payloadShape, ValueRange origOffsets) {934  SmallVector<Value> newCoods;935  auto maybeCoords =936      layout.computeDistributedCoords(rewriter, loc, laneId, payloadShape);937  if (failed(maybeCoords))938    return {};939  assert(maybeCoords.value().size() == 1 &&940         "Expected one set of distributed offsets");941  SmallVector<OpFoldResult> ofrVec = xegpu::addWithRightAligned(942      rewriter, loc, getAsOpFoldResult(maybeCoords.value()[0]),943      getAsOpFoldResult(origOffsets));944  newCoods = llvm::map_to_vector(ofrVec, llvm::CastTo<Value>);945  return newCoods;946}947 948/// Pattern for distributing xegpu::LoadMatrixOp.949struct LoadMatrixDistribution final : public gpu::WarpDistributionPattern {950  using gpu::WarpDistributionPattern::WarpDistributionPattern;951  LogicalResult matchAndRewrite(gpu::WarpExecuteOnLane0Op warpOp,952                                PatternRewriter &rewriter) const override {953    gpu::YieldOp yield = warpOp.getTerminator();954    Operation *lastNode = yield->getPrevNode();955    auto matrixOp = dyn_cast_or_null<xegpu::LoadMatrixOp>(lastNode);956    if (!matrixOp)957      return failure();958 959    OpOperand *producedByLastLoad = getWarpResult(warpOp, [&](Operation *op) {960      return isa<xegpu::LoadMatrixOp>(op) && matrixOp == op;961    });962    if (!producedByLastLoad)963      return rewriter.notifyMatchFailure(964          warpOp, "The last op is not xegpu::LoadMatrixOp");965    const int operandIdx = producedByLastLoad->getOperandNumber();966 967    VectorType sgPayloadTy =968        dyn_cast<VectorType>(matrixOp.getResult().getType());969    VectorType warpResultTy =970        cast<VectorType>(warpOp.getResult(operandIdx).getType());971    if (!sgPayloadTy)972      return rewriter.notifyMatchFailure(973          matrixOp, "the matrix op payload must be a vector type");974 975    auto loc = matrixOp.getLoc();976    auto offsets = matrixOp.getMixedOffsets();977    if (offsets.empty())978      return rewriter.notifyMatchFailure(matrixOp,979                                         "the load op must have offsets");980    SmallVector<Value> offsetsAsValues =981        vector::getAsValues(rewriter, matrixOp.getLoc(), offsets);982 983    auto layout = matrixOp.getLayoutAttr();984    if (!layout)985      return rewriter.notifyMatchFailure(986          matrixOp, "the matrix operation lacks layout attribute");987 988    FailureOr<VectorType> distPayloadByWarpOpOrFailure =989        getDistVecTypeBasedOnLaneLayout(layout, sgPayloadTy);990    if (failed(distPayloadByWarpOpOrFailure))991      return rewriter.notifyMatchFailure(992          matrixOp, "Failed to distribute matrix op payload based on layout.");993 994    SmallVector<Value> operands = {matrixOp.getMemDesc()};995    const unsigned offsetsStartIdx = operands.size();996    operands.append(offsetsAsValues);997 998    SmallVector<Type> operandTypes = llvm::to_vector(999        llvm::map_range(operands, [](Value v) { return v.getType(); }));1000 1001    SmallVector<size_t> newRetIndices;1002    gpu::WarpExecuteOnLane0Op newWarpOp = moveRegionToNewWarpOpAndAppendReturns(1003        rewriter, warpOp, operands, operandTypes, newRetIndices);1004    SmallVector<Value> newOperands = llvm::map_to_vector(1005        newRetIndices, [&](size_t idx) { return newWarpOp.getResult(idx); });1006 1007    SmallVector<int64_t> newConstOffsets(matrixOp.getConstOffsets().size(),1008                                         ShapedType::kDynamic);1009    DenseI64ArrayAttr newConstOffsetsAttr =1010        rewriter.getDenseI64ArrayAttr(newConstOffsets);1011    ValueRange currentOffsets =1012        ValueRange(newOperands).drop_front(offsetsStartIdx);1013 1014    SmallVector<Value> newCoords = currentOffsets;1015    rewriter.setInsertionPointAfter(newWarpOp);1016 1017    if (!matrixOp.getSubgroupBlockIoAttr()) {1018      newCoords = computeDistributedCoordinatesForMatrixOp(1019          rewriter, loc, layout, newWarpOp.getLaneid(), sgPayloadTy.getShape(),1020          currentOffsets);1021    }1022    xegpu::LoadMatrixOp newOp = xegpu::LoadMatrixOp::create(1023        rewriter, newWarpOp.getLoc(), *distPayloadByWarpOpOrFailure,1024        newOperands[0], ValueRange(newCoords), newConstOffsetsAttr,1025        matrixOp.getSubgroupBlockIoAttr(), xegpu::DistributeLayoutAttr{});1026    // Resolve the output type and replace all uses.1027    rewriter.replaceAllUsesWith(1028        newWarpOp.getResult(operandIdx),1029        resolveDistributedTy(newOp.getResult(), warpResultTy, rewriter));1030    return success();1031  }1032};1033 1034/// Pattern for distributing xegpu::StoreMatrixOp.1035struct StoreMatrixDistribution final : public gpu::WarpDistributionPattern {1036  using gpu::WarpDistributionPattern::WarpDistributionPattern;1037  LogicalResult matchAndRewrite(gpu::WarpExecuteOnLane0Op warpOp,1038                                PatternRewriter &rewriter) const override {1039    gpu::YieldOp yield = warpOp.getTerminator();1040    Operation *lastNode = yield->getPrevNode();1041    auto matrixOp = dyn_cast_or_null<xegpu::StoreMatrixOp>(lastNode);1042    if (!matrixOp)1043      return failure();1044 1045    VectorType sgPayloadTy = dyn_cast<VectorType>(matrixOp.getData().getType());1046    if (!sgPayloadTy)1047      return rewriter.notifyMatchFailure(1048          matrixOp, "the matrix op payload must be a vector type");1049 1050    auto loc = matrixOp.getLoc();1051    auto offsets = matrixOp.getMixedOffsets();1052    if (offsets.empty())1053      return rewriter.notifyMatchFailure(matrixOp,1054                                         "the store op must have offsets");1055    SmallVector<Value> offsetsAsValues =1056        vector::getAsValues(rewriter, matrixOp.getLoc(), offsets);1057 1058    auto layout = matrixOp.getLayoutAttr();1059    if (!layout)1060      return rewriter.notifyMatchFailure(1061          matrixOp, "the matrix operation lacks layout attribute");1062 1063    FailureOr<VectorType> distPayloadByWarpOpOrFailure =1064        getDistVecTypeBasedOnLaneLayout(layout, sgPayloadTy);1065    if (failed(distPayloadByWarpOpOrFailure))1066      return rewriter.notifyMatchFailure(1067          matrixOp, "Failed to distribute matrix op payload based on layout.");1068 1069    SmallVector<Value> operands = {matrixOp.getData(), matrixOp.getMemDesc()};1070    const unsigned offsetsStartIdx = operands.size();1071    operands.append(offsetsAsValues);1072 1073    SmallVector<Type> operandTypes = llvm::to_vector(1074        llvm::map_range(operands, [](Value v) { return v.getType(); }));1075    operandTypes[0] = *distPayloadByWarpOpOrFailure;1076 1077    SmallVector<size_t> newRetIndices;1078    gpu::WarpExecuteOnLane0Op newWarpOp = moveRegionToNewWarpOpAndAppendReturns(1079        rewriter, warpOp, operands, operandTypes, newRetIndices);1080    SmallVector<Value> newOperands = llvm::map_to_vector(1081        newRetIndices, [&](size_t idx) { return newWarpOp.getResult(idx); });1082 1083    SmallVector<int64_t> newConstOffsets(matrixOp.getConstOffsets().size(),1084                                         ShapedType::kDynamic);1085    DenseI64ArrayAttr newConstOffsetsAttr =1086        rewriter.getDenseI64ArrayAttr(newConstOffsets);1087    ValueRange currentOffsets =1088        ValueRange(newOperands).drop_front(offsetsStartIdx);1089 1090    SmallVector<Value> newCoords = currentOffsets;1091    rewriter.setInsertionPointAfter(newWarpOp);1092 1093    if (!matrixOp.getSubgroupBlockIoAttr()) {1094      newCoords = computeDistributedCoordinatesForMatrixOp(1095          rewriter, loc, layout, newWarpOp.getLaneid(), sgPayloadTy.getShape(),1096          currentOffsets);1097    }1098 1099    xegpu::StoreMatrixOp::create(1100        rewriter, loc, TypeRange{}, newOperands[0], newOperands[1],1101        ValueRange(newCoords), newConstOffsetsAttr,1102        matrixOp.getSubgroupBlockIoAttr(), xegpu::DistributeLayoutAttr{});1103    rewriter.eraseOp(matrixOp);1104    return success();1105  }1106};1107 1108/// Distribute a scattered load op. The logic and requirements are the same as1109/// for the scattered store distribution. The warpOp's payload vector is1110/// expected to be distributed by the load's result consumer.1111/// Example 1 (no chunk size):1112///    %mask = producer_op : vector<16xi1>1113///    %offset = producer_op : vector<16xindex>1114///    %0 = xegpu.load %payload, %src[%offset], %mask : memref<256xf16>,1115///    vector<16xindex>, vector<16xi1> -> vector<16xf16>1116/// To1117///    %mask = producer_op : vector<1xi1>1118///    %offset = producer_op : vector<1xindex>1119///    %0 = xegpu.load %payload, %src[%offset], %mask : memref<256xf16>,1120///     vector<1xindex>, vector<1xi1> -> vector<1xf16>1121/// Example 2 (chunk size, same mask and offsets):1122///    %0 = xegpu.load %payload, %src[%offset], %mask <{chunk_size=8}> :1123///     memref<256xf16>, vector<16xindex>, vector<16xi1> -> vector<16x8xf16>1124/// To1125///    %0 = xegpu.load %payload, %src[%offset], %mask <{chunk_size=8}> :1126///     memref<256xf16>, vector<1xindex>, vector<1xi1> -> vector<8xf16>1127struct LoadDistribution final : public gpu::WarpDistributionPattern {1128  using gpu::WarpDistributionPattern::WarpDistributionPattern;1129  LogicalResult matchAndRewrite(gpu::WarpExecuteOnLane0Op warpOp,1130                                PatternRewriter &rewriter) const override {1131    OpOperand *producedByLastLoad = getWarpResult(warpOp, [&](Operation *op) {1132      // Check if the yield operand that was produced by the *last* scattered1133      // load op to avoid sinking it before barriers (maintain memory order).1134      return isa<xegpu::LoadGatherOp>(op) &&1135             warpOp.getTerminator()->getPrevNode() == op;1136    });1137    if (!producedByLastLoad)1138      return rewriter.notifyMatchFailure(1139          warpOp, "The last op is not xegpu::LoadGatherOp");1140 1141    auto loadGatherOp =1142        producedByLastLoad->get().getDefiningOp<xegpu::LoadGatherOp>();1143    auto offsets = loadGatherOp.getOffsets();1144    if (!offsets || !isa<VectorType>(offsets.getType()) ||1145        !isa<VectorType>(loadGatherOp.getMask().getType()))1146      return rewriter.notifyMatchFailure(1147          loadGatherOp,1148          "Load op must have a vector arguments for offsets and mask");1149    VectorType offsetsTy = cast<VectorType>(offsets.getType());1150    VectorType maskTy = cast<VectorType>(loadGatherOp.getMask().getType());1151    if (offsetsTy.getRank() != 1 || maskTy.getRank() != 1)1152      return rewriter.notifyMatchFailure(loadGatherOp,1153                                         "Expected 1D offsets and mask vector");1154    // Assume offset and mask producers will be distributed as well.1155    std::string layoutOffsetsName =1156        xegpu::getLayoutName(loadGatherOp->getOpOperand(1));1157    std::string layoutMaskName =1158        xegpu::getLayoutName(loadGatherOp->getOpOperand(2));1159 1160    xegpu::LayoutAttr layoutOffsets =1161        loadGatherOp->getAttrOfType<xegpu::LayoutAttr>(layoutOffsetsName);1162    xegpu::LayoutAttr layoutMask =1163        loadGatherOp->getAttrOfType<xegpu::LayoutAttr>(layoutMaskName);1164 1165    FailureOr<VectorType> distOffsetsByWarpOpOrFailure =1166        getDistVecTypeBasedOnLaneLayout(layoutOffsets, offsetsTy);1167    FailureOr<VectorType> distMaskByWarpOpOrFailure =1168        getDistVecTypeBasedOnLaneLayout(layoutMask, maskTy);1169    if (failed(distOffsetsByWarpOpOrFailure) ||1170        failed(distMaskByWarpOpOrFailure)) {1171      return rewriter.notifyMatchFailure(1172          loadGatherOp,1173          "Some vector operands have no layouts, using defaults instead.");1174    }1175 1176    SmallVector<size_t> newRetIndices;1177    SmallVector<Value> operands = loadGatherOp->getOperands();1178    SmallVector<Type> operandTypesToYield = {1179        operands[0].getType(), distOffsetsByWarpOpOrFailure.value(),1180        distMaskByWarpOpOrFailure.value()};1181 1182    const unsigned operandIdx = producedByLastLoad->getOperandNumber();1183    VectorType distResultTy =1184        cast<VectorType>(warpOp.getResult(operandIdx).getType());1185    // Distributed load op will always be 1D.1186    VectorType loadVecTy = VectorType::get({distResultTy.getNumElements()},1187                                           distResultTy.getElementType());1188 1189    gpu::WarpExecuteOnLane0Op newWarpOp = moveRegionToNewWarpOpAndAppendReturns(1190        rewriter, warpOp, operands, operandTypesToYield, newRetIndices);1191 1192    SmallVector<Value> newLoadGatherOperands = llvm::map_to_vector(1193        newRetIndices, [&](size_t idx) { return newWarpOp.getResult(idx); });1194 1195    rewriter.setInsertionPointAfter(newWarpOp);1196    xegpu::LoadGatherOp newOp = xegpu::LoadGatherOp::create(1197        rewriter, newWarpOp.getLoc(), loadVecTy, newLoadGatherOperands,1198        loadGatherOp->getAttrs());1199    xegpu::removeLayoutAttrs(newOp);1200    Value distributedVal = newWarpOp.getResult(operandIdx);1201    // Resolve the output type and replace all uses.1202    rewriter.replaceAllUsesWith(1203        distributedVal,1204        resolveDistributedTy(newOp.getResult(), distResultTy, rewriter));1205    return success();1206  }1207};1208 1209/// Helper to rewrite a 2D VectorMultiReductionOp into a sequence of 1D1210/// VectorReductionOps. We also insert layouts for the newly created ops.1211static Value lowerToVectorReductions(TypedValue<VectorType> src,1212                                     TypedValue<VectorType> acc,1213                                     vector::CombiningKind kind,1214                                     int64_t reductionDim, Location loc,1215                                     PatternRewriter &rewriter) {1216  // Expecting a 2D source vector.1217  assert(src.getType().getRank() == 2 && "expected a 2D source vector");1218  VectorType sourceType = src.getType();1219  int64_t sourceH = sourceType.getShape()[0];1220  int64_t sourceW = sourceType.getShape()[1];1221  int nSlices = (reductionDim == 0) ? sourceW : sourceH;1222  // Create a constant vector to hold the result of the reduction.1223  TypedAttr zeroAttr = rewriter.getZeroAttr(sourceType.getElementType());1224  Value reductionResult = arith::ConstantOp::create(1225      rewriter, loc, acc.getType(),1226      DenseElementsAttr::get(acc.getType(), zeroAttr));1227  // Reduction result should have the same layout as the accumulator.1228  xegpu::setDistributeLayoutAttr(cast<OpResult>(reductionResult),1229                                 xegpu::getDistributeLayoutAttr(acc));1230  // For each slice of the source, extract the slice vector, do a reduction1231  // and, insert the reduced value back to the result vector.1232  for (int i = 0; i < nSlices; ++i) {1233    SmallVector<int64_t, 2> sliceOffsets, sliceSizes;1234    if (reductionDim == 1) {1235      sliceOffsets = {i, 0};1236      sliceSizes = {1, sourceW};1237    } else {1238      sliceOffsets = {0, i};1239      sliceSizes = {sourceH, 1};1240    }1241    vector::ExtractStridedSliceOp extractOp =1242        vector::ExtractStridedSliceOp::create(rewriter, loc, src, sliceOffsets,1243                                              sliceSizes, {1, 1});1244    int64_t nSliceElements = extractOp.getResult().getType().getNumElements();1245    vector::ShapeCastOp slice = vector::ShapeCastOp::create(1246        rewriter, loc,1247        VectorType::get({nSliceElements}, sourceType.getElementType()),1248        extractOp.getResult());1249    // Shape cast is currently handled in xegpu side. So layouts must be1250    // retained during lowering. Shape cast output has the same layout as the1251    // accumulator. Shape cast source has the same layout as the original1252    // reduction source.1253    // TODO: other ops generated here may also need layout attributes.1254    xegpu::setDistributeLayoutAttr(slice->getOpOperand(0),1255                                   xegpu::getDistributeLayoutAttr(src));1256    xegpu::setDistributeLayoutAttr(slice->getOpResult(0),1257                                   xegpu::getDistributeLayoutAttr(acc));1258    // Extract and reduction results in scalars, so no result layout is needed.1259    Value accExtract = vector::ExtractOp::create(rewriter, loc, acc, i);1260    Value reduction = vector::ReductionOp::create(1261        rewriter, loc, kind, slice.getResult(), accExtract);1262    reductionResult =1263        vector::InsertOp::create(rewriter, loc, reduction, reductionResult, i);1264  }1265  return reductionResult;1266}1267 1268/// This patterns distribute the `vector.multi_reduction` operation across1269/// lanes in a warp. Currently only 2D to 1D reductions are supported. Given1270/// layouts for the source and accumulator vectors,1271/// * If the reduction dimension is distributed across lanes, the reduction is1272///   non-lane-local and the reduction is done using warp shuffles. Here we1273///   simply rewrite the MultiDimReductionOp to a sequence of ReductionOps in1274///   the warp op body.1275/// * If the reduction dimension is not distributed across lanes, the reduction1276///   is lane-local. In this case, we yield the source and accumulator vectors1277///   from the warp op and perform the lane-local reduction outside the warp op1278///   using a sequence of ReductionOps.1279/// Example 1 (Reduction is lane-local):1280/// ```1281/// %r = gpu.warp_execute_on_lane_0(%laneid)[32] -> (vector<1xf32>) {1282///   %0 = "some_def"() : () -> (vector<16x32xf32>)1283///   %acc = "some_def"() : () -> (vector<32xf32>)1284///   %1 = vector.multi_reduction <add>, %0, %acc [0] : vector<16x32xf32> to1285///   vector<32xf32> gpu.yield %1 : vector<32xf32>1286/// }1287/// ```1288/// is lowered to:1289/// ```1290/// %r:2 = gpu.warp_execute_on_lane_0(%laneid)[32] -> (vector<16x1xf32>,1291/// vector<1xf32>) {1292///   %0 = "some_def"() : () -> (vector<16x32xf32>)1293///   %acc = "some_def"() : () -> (vector<32xf32>)1294///   gpu.yield %0, %acc : vector<16x32xf32>, vector<32xf32>1295/// }1296/// %c = arith.constant dense<0.0> : vector<1xf32>1297/// %1 = vector.shape_cast %r#0 : vector<16x1xf32> to vector<16xf32>1298/// %2 = vector.reduction <add>, %1, %r#1 : vector<16xf32> to f321299/// %3 = vector.insert %2, %c[0] : f32 into vector<1xf32>1300/// ```1301/// Example 2 (Reduction is non-lane-local):1302/// ```1303/// %r = gpu.warp_execute_on_lane_0(%laneid)[32] -> (vector<2xf32>) {1304///   %0 = "some_def"() : () -> (vector<2x32xf32>)1305///   %acc = "some_def"() : () -> (vector<2xf32>)1306///   %1 = vector.multi_reduction <add>, %0, %acc [1] : vector<2x32xf32> to1307///   vector<2xf32>1308///   gpu.yield %1 : vector<2xf32>1309/// }1310/// ```1311/// is lowered to:1312/// ```1313/// %r = gpu.warp_execute_on_lane_0(%laneid)[32] -> (vector<2xf32>) {1314///   %0 = "some_def"() : () -> (vector<2x32xf32>)1315///   %acc = "some_def"() : () -> (vector<2xf32>)1316///   %1 = arith.constant dense<0.0> : vector<2xf32>1317///   %2 = vector.extract %0[0] : vector<32xf32> from <vector<2x32xf32>>1318///   %3 = ("warp.reduction %2") : f321319///   %4 = vector.insert %3, %1[0] : f32 into vector<2xf32>1320///   ... repeat for row 11321///   gpu.yield %1 : vector<2xf32>1322/// }1323struct VectorMultiReductionDistribution : public gpu::WarpDistributionPattern {1324  using gpu::WarpDistributionPattern::WarpDistributionPattern;1325  LogicalResult matchAndRewrite(gpu::WarpExecuteOnLane0Op warpOp,1326                                PatternRewriter &rewriter) const override {1327    OpOperand *yieldOperand =1328        getWarpResult(warpOp, llvm::IsaPred<vector::MultiDimReductionOp>);1329    if (!yieldOperand)1330      return failure();1331    auto reductionOp =1332        cast<vector::MultiDimReductionOp>(yieldOperand->get().getDefiningOp());1333    unsigned operandIdx = yieldOperand->getOperandNumber();1334    VectorType sourceType = reductionOp.getSourceVectorType();1335    // Only 2D vectors are supported.1336    if (sourceType.getRank() != 2)1337      return rewriter.notifyMatchFailure(warpOp,1338                                         "Only 2D reductions are supported.");1339    ArrayRef<int64_t> reductionDims = reductionOp.getReductionDims();1340    // Only 1 reduction dimension supported. This also ensures that the result1341    // is vector type.1342    if (reductionDims.size() != 1)1343      return rewriter.notifyMatchFailure(1344          warpOp, "Only 1 reduction dimension is supported.");1345    int64_t reductionDim = reductionDims[0];1346    VectorType distributedResultType =1347        cast<VectorType>(warpOp.getResult(operandIdx).getType());1348    VectorType resultType = cast<VectorType>(reductionOp.getType());1349    xegpu::DistributeLayoutAttr sourceLayout =1350        xegpu::getDistributeLayoutAttr(reductionOp.getSource());1351 1352    FailureOr<VectorType> sourceDistTypeOrFailure =1353        getDistVecTypeBasedOnLaneLayout(sourceLayout, sourceType);1354    if (failed(sourceDistTypeOrFailure))1355      return rewriter.notifyMatchFailure(1356          warpOp, "Failed to distribute the source vector type.");1357    VectorType sourceDistType = sourceDistTypeOrFailure.value();1358    // Only single dimension distribution is supported.1359    bool dim0Distributed =1360        sourceDistType.getShape()[0] != sourceType.getShape()[0];1361    bool dim1Distributed =1362        sourceDistType.getShape()[1] != sourceType.getShape()[1];1363    if (dim0Distributed && dim1Distributed)1364      return rewriter.notifyMatchFailure(1365          warpOp, "Expecting source to be distributed in a single dimension.");1366    int64_t sourceDistDim = dim0Distributed ? 0 : (dim1Distributed ? 1 : -1);1367    if (sourceDistDim == -1)1368      return rewriter.notifyMatchFailure(1369          warpOp, "Expecting a distributed source vector.");1370    bool resultDistributed =1371        distributedResultType.getNumElements() < resultType.getNumElements();1372    // If the lane owns all the data required for reduction (i.e. reduction is1373    // fully parallel accross lanes), then each lane owns part of the result1374    // (i.e. result is distributed). If the reduction require cross-lane1375    // shuffling, then the result is shared among all lanes (broadcasted).1376    // Therefore we expect following cases:1377    //1378    // | Source vector        | Reduction dim  | Result vector  |1379    // |----------------------|----------------|----------------|1380    // |  dim-0 distributed   |       0        | broadcasted    |1381    // |  dim-0 distributed   |       1        | distributed    |1382    // |  dim-1 distributed   |       0        | distributed    |1383    // |  dim-1 distributed   |       1        | broadcasted    |1384 1385    bool isReductionLaneLocal = (sourceDistDim == 0 && reductionDim == 1) ||1386                                (sourceDistDim == 1 && reductionDim == 0);1387    if (isReductionLaneLocal && !resultDistributed)1388      return rewriter.notifyMatchFailure(1389          warpOp, "Expecting a distributed result for lane-local reduction.");1390 1391    if (!isReductionLaneLocal && resultDistributed)1392      return rewriter.notifyMatchFailure(1393          warpOp,1394          "Expecting a broadcasted result for non-lane-local reduction.");1395 1396    // Handle lane-local reduction case. In this case we fully distribute the1397    // reduction result.1398    if (isReductionLaneLocal) {1399      // Yield the source and acc vectors from the WarpOp.1400      SmallVector<size_t> newRetIndices;1401      auto newWarpOp = moveRegionToNewWarpOpAndAppendReturns(1402          rewriter, warpOp, {reductionOp.getSource(), reductionOp.getAcc()},1403          {sourceDistType, distributedResultType}, newRetIndices);1404      rewriter.setInsertionPointAfter(newWarpOp);1405      Value result = lowerToVectorReductions(1406          cast<TypedValue<VectorType>>(newWarpOp->getResult(newRetIndices[0])),1407          cast<TypedValue<VectorType>>(newWarpOp->getResult(newRetIndices[1])),1408          reductionOp.getKind(), reductionDim, reductionOp.getLoc(), rewriter);1409      // Replace the warp op result with the final result.1410      rewriter.replaceAllUsesWith(newWarpOp.getResult(operandIdx), result);1411      return success();1412    }1413    // For non-lane-local case, we simply rewrite the MultiReductionOp in terms1414    // of multiple ReductionOps. Actual distribution is done by the1415    // WarpOpReduction pattern.1416    rewriter.setInsertionPointAfter(reductionOp);1417    Value result = lowerToVectorReductions(1418        cast<TypedValue<VectorType>>(reductionOp.getSource()),1419        cast<TypedValue<VectorType>>(reductionOp.getAcc()),1420        reductionOp.getKind(), reductionDim, reductionOp.getLoc(), rewriter);1421    // Replace the warp op result with the final result.1422    rewriter.replaceAllUsesWith(reductionOp.getResult(), result);1423    return success();1424  }1425};1426 1427/// Distribute a `vector.shape_cast` op feeding into yield op of an enclosing1428/// `gpu.warp_execute_on_lane_0` region.1429struct VectorShapeCastDistribution : public gpu::WarpDistributionPattern {1430  using gpu::WarpDistributionPattern::WarpDistributionPattern;1431  LogicalResult matchAndRewrite(gpu::WarpExecuteOnLane0Op warpOp,1432                                PatternRewriter &rewriter) const override {1433    OpOperand *yieldOperand =1434        getWarpResult(warpOp, llvm::IsaPred<vector::ShapeCastOp>);1435    if (!yieldOperand)1436      return failure();1437    auto shapeCastOp =1438        cast<vector::ShapeCastOp>(yieldOperand->get().getDefiningOp());1439    unsigned operandNumber = yieldOperand->getOperandNumber();1440    auto resultDistTy =1441        cast<VectorType>(warpOp.getResult(operandNumber).getType());1442    xegpu::DistributeLayoutAttr sourceLayout =1443        xegpu::getDistributeLayoutAttr(shapeCastOp->getOpOperand(0));1444    xegpu::DistributeLayoutAttr resultLayout =1445        xegpu::getDistributeLayoutAttr(shapeCastOp.getResult());1446    if (!sourceLayout || !resultLayout)1447      return rewriter.notifyMatchFailure(1448          warpOp,1449          "the source or result of shape_cast op lacks distribution layout");1450 1451    // For rank reducing or increasing shape_cast ops, the lower rank layout1452    // must be a slice of higher rank layout.1453    int64_t sourceRank = shapeCastOp.getSourceVectorType().getRank();1454    int64_t resultRank = shapeCastOp.getResultVectorType().getRank();1455    if (sourceRank < resultRank && !sourceLayout.isSliceOf(resultLayout))1456      return rewriter.notifyMatchFailure(1457          warpOp, "shape_cast is rank reducing but source layout is not a "1458                  "slice of result layout");1459    if (sourceRank > resultRank && !resultLayout.isSliceOf(sourceLayout))1460      return rewriter.notifyMatchFailure(1461          warpOp, "shape_cast is rank increasing but result layout is not a "1462                  "slice of source layout");1463 1464    FailureOr<VectorType> sourceDistTypeOrFailure =1465        getDistVecTypeBasedOnLaneLayout(sourceLayout,1466                                        shapeCastOp.getSourceVectorType());1467    if (failed(sourceDistTypeOrFailure))1468      return rewriter.notifyMatchFailure(1469          warpOp, "failed to get distributed vector type for source");1470    VectorType sourceDistType = sourceDistTypeOrFailure.value();1471    // Create a new warp op that yields the source of the shape_cast op.1472    SmallVector<size_t> newRetIndices;1473    auto newWarpOp = moveRegionToNewWarpOpAndAppendReturns(1474        rewriter, warpOp, {shapeCastOp.getSource()}, {sourceDistType},1475        newRetIndices);1476    rewriter.setInsertionPointAfter(newWarpOp);1477    Value source = newWarpOp.getResult(newRetIndices[0]);1478    // Create a new shape_cast op outside the warp op.1479    Value newShapeCast = vector::ShapeCastOp::create(1480        rewriter, shapeCastOp.getLoc(), resultDistTy, source);1481    rewriter.replaceAllUsesWith(newWarpOp.getResult(operandNumber),1482                                newShapeCast);1483    return success();1484  }1485};1486 1487// Distribute a `vector.extract_strided_slice` op feeding into yield op of an1488// enclosing `gpu.warp_execute_on_lane_0` region. This pattern covers1489// advanced cases where the distributed dimension is partially extracted and1490// currently not supported by the generic vector distribution patterns.1491struct VectorExtractStridedSliceDistribution1492    : public gpu::WarpDistributionPattern {1493  using gpu::WarpDistributionPattern::WarpDistributionPattern;1494  LogicalResult matchAndRewrite(gpu::WarpExecuteOnLane0Op warpOp,1495                                PatternRewriter &rewriter) const override {1496    OpOperand *operand =1497        getWarpResult(warpOp, llvm::IsaPred<vector::ExtractStridedSliceOp>);1498    if (!operand)1499      return failure();1500    auto extractOp =1501        cast<vector::ExtractStridedSliceOp>(operand->get().getDefiningOp());1502    unsigned operandIdx = operand->getOperandNumber();1503    auto distributedType =1504        cast<VectorType>(warpOp.getResult(operandIdx).getType());1505    // Find the distributed dimensions.1506    auto extractResultType = cast<VectorType>(operand->get().getType());1507    auto distributedDims =1508        getDistributedDims(extractResultType, distributedType);1509    // Collect updated source type, sizes and offsets. They may be adjusted1510    // later if the data is distributed to lanes (as opposed to being owned by1511    // all lanes uniformly).1512    VectorType updatedSourceType = extractOp.getSourceVectorType();1513    SmallVector<Attribute> updatedSizes = llvm::map_to_vector(1514        extractOp.getSizes(), [](Attribute attr) { return attr; });1515    SmallVector<Attribute> updatedOffsets = llvm::map_to_vector(1516        extractOp.getOffsets(), [](Attribute attr) { return attr; });1517    // If the result is distributed, it must be distributed in exactly one1518    // dimension. In this case, we adjust the sourceDistType, distributedSizes1519    // and distributedOffsets accordingly.1520    if (distributedDims.size() > 0) {1521      if (distributedDims.size() != 1)1522        return rewriter.notifyMatchFailure(1523            warpOp, "Source can not be distributed in multiple dimensions.");1524      int64_t distributedDim = distributedDims[0];1525      int sourceDistrDimSize =1526          extractOp.getSourceVectorType().getShape()[distributedDim];1527      auto sourceLayout =1528          xegpu::getDistributeLayoutAttr(extractOp->getOpOperand(0));1529      if (!sourceLayout || sourceLayout.getEffectiveLaneLayoutAsInt().empty())1530        return rewriter.notifyMatchFailure(1531            warpOp, "the source of extract_strided_slice op lacks distribution "1532                    "layout");1533      auto sourceLaneLayout = sourceLayout.getEffectiveLaneLayoutAsInt();1534      // Because only single dimension distribution is supported, lane layout1535      // size at the distributed dim must be the subgroup size.1536      int subgroupSize = sourceLaneLayout[distributedDim];1537      // Check if the source size in the distributed dimension is a multiple of1538      // subgroup size.1539      if (sourceDistrDimSize % subgroupSize != 0)1540        return rewriter.notifyMatchFailure(1541            warpOp,1542            "Source size along distributed dimension is not a multiple of "1543            "subgroup size.");1544      auto sourceLaneData = sourceLayout.getEffectiveLaneDataAsInt();1545      // We expect lane data to be all ones in this case.1546      if (!llvm::all_of(sourceLaneData, [](int64_t v) { return v == 1; }))1547        return rewriter.notifyMatchFailure(1548            warpOp, "Expecting unit lane data in source layout");1549      // The offsets in the distributed dimention must be a multiple of subgroup1550      // size.1551      int64_t distrDimOffset =1552          cast<IntegerAttr>(extractOp.getOffsets()[distributedDim]).getInt();1553      if (distrDimOffset % subgroupSize != 0)1554        return rewriter.notifyMatchFailure(1555            warpOp, "Offset along distributed dimension "1556                    "is not a multiple of subgroup size.");1557      updatedSourceType = getDistVecTypeBasedOnLaneLayout(1558                              sourceLayout, extractOp.getSourceVectorType())1559                              .value();1560      // Update the distributed sizes to match the distributed type.1561      updatedSizes[distributedDim] = rewriter.getI64IntegerAttr(1562          distributedType.getDimSize(distributedDim));1563      // Update the distributed offsets to match round robin distribution (i.e.1564      // each lane owns data at `subgroupSize` stride given unit lane data).1565      updatedOffsets[distributedDim] =1566          rewriter.getI64IntegerAttr(distrDimOffset / subgroupSize);1567    }1568    // Do the distribution by yielding the source of the extract op from1569    // the warp op and creating a new extract op outside the warp op.1570    SmallVector<size_t> newRetIndices;1571    auto newWarpOp = moveRegionToNewWarpOpAndAppendReturns(1572        rewriter, warpOp, {extractOp.getSource()}, {updatedSourceType},1573        newRetIndices);1574    rewriter.setInsertionPointAfter(newWarpOp);1575    Value source = newWarpOp.getResult(newRetIndices[0]);1576    // Create a new extract op outside the warp op.1577    Value newExtractOp = vector::ExtractStridedSliceOp::create(1578        rewriter, extractOp.getLoc(), distributedType, source,1579        ArrayAttr::get(rewriter.getContext(), updatedOffsets),1580        ArrayAttr::get(rewriter.getContext(), updatedSizes),1581        extractOp.getStrides());1582    rewriter.replaceAllUsesWith(newWarpOp.getResult(operandIdx), newExtractOp);1583    return success();1584  }1585};1586 1587/// Distribute a `vector.insert_strided_slice` op feeding into yield op of an1588/// enclosing `gpu.warp_execute_on_lane_0` region. This pattern covers1589/// advanced cases where the distributed dimension is partially inserted and1590/// currently not supported by the generic vector distribution patterns.1591struct VectorInsertStridedSliceDistribution1592    : public gpu::WarpDistributionPattern {1593  using gpu::WarpDistributionPattern::WarpDistributionPattern;1594  LogicalResult matchAndRewrite(gpu::WarpExecuteOnLane0Op warpOp,1595                                PatternRewriter &rewriter) const override {1596    OpOperand *operand =1597        getWarpResult(warpOp, llvm::IsaPred<vector::InsertStridedSliceOp>);1598    if (!operand)1599      return failure();1600    unsigned int operandNumber = operand->getOperandNumber();1601    auto insertOp =1602        operand->get().getDefiningOp<vector::InsertStridedSliceOp>();1603    auto distributedType =1604        cast<VectorType>(warpOp.getResult(operandNumber).getType());1605    // Find the distributed dimensions of the dest vector.1606    auto insertResultType = cast<VectorType>(operand->get().getType());1607    auto destDistributedDims =1608        getDistributedDims(insertResultType, distributedType);1609    // Collect updated offsets, source type and dest type. They may be adjusted1610    // later if the data is distributed to lanes (as opposed to being owned by1611    // all lanes uniformly).1612    SmallVector<Attribute> updatedOffsets = llvm::map_to_vector(1613        insertOp.getOffsets(), [](Attribute attr) { return attr; });1614    VectorType updatedSourceType = insertOp.getSourceVectorType();1615    VectorType updatedDestType = insertOp.getDestVectorType();1616    if (destDistributedDims.size() > 0) {1617      // Only single dimension distribution is supported.1618      if (destDistributedDims.size() != 1)1619        return rewriter.notifyMatchFailure(1620            warpOp,1621            "Expecting source to be distributed in a single dimension.");1622      int64_t destDistributedDim = destDistributedDims[0];1623 1624      VectorType srcType = insertOp.getSourceVectorType();1625      VectorType destType = insertOp.getDestVectorType();1626      // Currently we require that both source (kD) and dest (nD) vectors are1627      // distributed. This requires that distributedDim (d) is contained in the1628      // last k dims of the dest vector (d >= n - k).1629      int64_t sourceDistributedDim =1630          destDistributedDim - (destType.getRank() - srcType.getRank());1631      if (sourceDistributedDim < 0)1632        return rewriter.notifyMatchFailure(1633            insertOp,1634            "distributed dimension must be in the last k (i.e. source "1635            "rank) dims of dest vector");1636      int64_t srcDistrDimSize = srcType.getDimSize(sourceDistributedDim);1637      // Obtain the source and dest layouts.1638      auto destLayout =1639          xegpu::getDistributeLayoutAttr(insertOp->getOpOperand(1));1640      auto sourceLayout =1641          xegpu::getDistributeLayoutAttr(insertOp->getOpOperand(0));1642      if (!destLayout || !sourceLayout ||1643          destLayout.getEffectiveLaneLayoutAsInt().empty() ||1644          sourceLayout.getEffectiveLaneLayoutAsInt().empty())1645        return rewriter.notifyMatchFailure(1646            warpOp, "the source or dest of insert_strided_slice op lacks "1647                    "distribution layout");1648      // Because only single dimension distribution is supported, lane layout1649      // size at the distributed dim must be the subgroup size.1650      int subgroupSize =1651          destLayout.getEffectiveLaneLayoutAsInt()[destDistributedDim];1652      // We require that source and dest lane data are all ones to ensure1653      // uniform round robin distribution.1654      auto destLaneData = destLayout.getEffectiveLaneDataAsInt();1655      auto sourceLaneData = sourceLayout.getEffectiveLaneDataAsInt();1656      if (!llvm::all_of(destLaneData, [](int64_t v) { return v == 1; }) ||1657          !llvm::all_of(sourceLaneData, [](int64_t v) { return v == 1; }))1658        return rewriter.notifyMatchFailure(1659            warpOp, "Expecting unit lane data in source and dest layouts");1660      // Source distributed dim size must be multiples of subgroup size.1661      if (srcDistrDimSize % subgroupSize != 0)1662        return rewriter.notifyMatchFailure(1663            warpOp, "Distributed dimension size in source is not a multiple of "1664                    "subgroup size.");1665      // Offsets in the distributed dimension must be multiples of subgroup1666      // size.1667      int64_t destDistrDimOffset =1668          cast<IntegerAttr>(insertOp.getOffsets()[destDistributedDim]).getInt();1669      if (destDistrDimOffset % subgroupSize != 0)1670        return rewriter.notifyMatchFailure(1671            warpOp,1672            "Offset along distributed dimension in dest is not a multiple of "1673            "subgroup size.");1674      // Update the source and dest types based on their layouts.1675      updatedSourceType = getDistVecTypeBasedOnLaneLayout(1676                              sourceLayout, insertOp.getSourceVectorType())1677                              .value();1678      updatedDestType = getDistVecTypeBasedOnLaneLayout(1679                            destLayout, insertOp.getDestVectorType())1680                            .value();1681      // Update the distributed offsets to match round robin distribution (i.e.1682      // each lane owns data at `subgroupSize` stride given unit lane data).1683      updatedOffsets[destDistributedDim] =1684          rewriter.getI64IntegerAttr(destDistrDimOffset / subgroupSize);1685    }1686    // Do the distribution by yielding the source and dest of the insert op1687    // from the warp op and creating a new insert op outside the warp op.1688    SmallVector<size_t> newRetIndices;1689    auto newWarpOp = moveRegionToNewWarpOpAndAppendReturns(1690        rewriter, warpOp, {insertOp.getValueToStore(), insertOp.getDest()},1691        {updatedSourceType, updatedDestType}, newRetIndices);1692    rewriter.setInsertionPointAfter(newWarpOp);1693 1694    Value valueToStore = newWarpOp.getResult(newRetIndices[0]);1695    Value dest = newWarpOp.getResult(newRetIndices[1]);1696    // Create a new insert op outside the warp op.1697    Value newInsertOp = vector::InsertStridedSliceOp::create(1698        rewriter, insertOp.getLoc(), updatedDestType, valueToStore, dest,1699        ArrayAttr::get(rewriter.getContext(), updatedOffsets),1700        insertOp.getStrides());1701    rewriter.replaceAllUsesWith(newWarpOp.getResult(operandNumber),1702                                newInsertOp);1703    return success();1704  }1705};1706 1707/// Sink a memref::ExtractAlignedPointerAsIndex op feeding into yield op of an1708/// enclosing `gpu.warp_execute_on_lane_0` region. This will simply move the op1709/// outside of the warp op.1710struct MemrefExtractAlignedPointerAsIndexDistribution final1711    : public gpu::WarpDistributionPattern {1712  using gpu::WarpDistributionPattern::WarpDistributionPattern;1713  LogicalResult matchAndRewrite(gpu::WarpExecuteOnLane0Op warpOp,1714                                PatternRewriter &rewriter) const override {1715    OpOperand *operand = getWarpResult(1716        warpOp, llvm::IsaPred<memref::ExtractAlignedPointerAsIndexOp>);1717    if (!operand)1718      return rewriter.notifyMatchFailure(1719          warpOp,1720          "warp result is not a memref::MemrefExtractAlignedPointerAsIndex op");1721    auto extractOp =1722        operand->get().getDefiningOp<memref::ExtractAlignedPointerAsIndexOp>();1723    unsigned operandIdx = operand->getOperandNumber();1724    SmallVector<size_t> newRetIndices;1725    gpu::WarpExecuteOnLane0Op newWarpOp = moveRegionToNewWarpOpAndAppendReturns(1726        rewriter, warpOp, extractOp.getSource(),1727        TypeRange{extractOp.getSource().getType()}, newRetIndices);1728    rewriter.setInsertionPointAfter(newWarpOp);1729    auto newExtractOp = memref::ExtractAlignedPointerAsIndexOp::create(1730        rewriter, newWarpOp.getLoc(), extractOp.getType(),1731        newWarpOp.getResult(newRetIndices[0]));1732    Value distributedVal = newWarpOp.getResult(operandIdx);1733    rewriter.replaceAllUsesWith(distributedVal, newExtractOp.getResult());1734    return success();1735  }1736};1737 1738/// Distribute a vector::BitCastOp feeding into yield op of an enclosing1739/// `gpu.warp_execute_on_lane_0` region. Bitcast only impacts the innermost1740/// diemension of the source/result vectors. Equivalent vector::BitCastOp is1741/// created outside of the warp op with distributed source vector type (computed1742/// using assigned layout).1743struct VectorBitcastDistribution final : public gpu::WarpDistributionPattern {1744  using gpu::WarpDistributionPattern::WarpDistributionPattern;1745  LogicalResult matchAndRewrite(gpu::WarpExecuteOnLane0Op warpOp,1746                                PatternRewriter &rewriter) const override {1747    OpOperand *operand =1748        getWarpResult(warpOp, llvm::IsaPred<vector::BitCastOp>);1749    if (!operand)1750      return rewriter.notifyMatchFailure(1751          warpOp, "warp result is not a vector::BitCast op");1752    auto bitcastOp = operand->get().getDefiningOp<vector::BitCastOp>();1753    unsigned operandIdx = operand->getOperandNumber();1754    VectorType distributedSourceType =1755        getDistVecTypeBasedOnLaneLayout(1756            xegpu::getDistributeLayoutAttr(bitcastOp.getSource()),1757            bitcastOp.getSourceVectorType())1758            .value_or(VectorType());1759    if (!distributedSourceType)1760      return rewriter.notifyMatchFailure(1761          bitcastOp, "Failed to distribute the source vector type in "1762                     "vector::BitCast op");1763    VectorType distributedResultType =1764        cast<VectorType>(warpOp.getResult(operandIdx).getType());1765    SmallVector<size_t> newRetIndices;1766    gpu::WarpExecuteOnLane0Op newWarpOp = moveRegionToNewWarpOpAndAppendReturns(1767        rewriter, warpOp, bitcastOp.getSource(),1768        TypeRange{distributedSourceType}, newRetIndices);1769    rewriter.setInsertionPointAfter(newWarpOp);1770    auto newBitcastOp = vector::BitCastOp::create(1771        rewriter, newWarpOp.getLoc(), distributedResultType,1772        newWarpOp.getResult(newRetIndices[0]));1773    Value distributedVal = newWarpOp.getResult(operandIdx);1774    rewriter.replaceAllUsesWith(distributedVal, newBitcastOp.getResult());1775    return success();1776  }1777};1778 1779/// Distribute a vector::TransposeOp feeding into yield op of an enclosing1780/// `gpu.warp_execute_on_lane_0` region. Currently only 2D transposes are1781/// supported. In most cases, transpose is a no op because it is entirely1782/// handled using the layouts (e.g. 16x1 -> 1x16). However, if each lane owns1783/// multiple slices of data after distribution (e.g. 16x2 -> 2x16), a lane-local1784/// transpose (i.e. shuffle) is needed. Therefore, we create an equivalent1785/// vector::TransposeOp outside of the warp op with distributed source vector1786/// type (computed using assigned layout).1787struct VectorTransposeDistribution final : public gpu::WarpDistributionPattern {1788  using gpu::WarpDistributionPattern::WarpDistributionPattern;1789  LogicalResult matchAndRewrite(gpu::WarpExecuteOnLane0Op warpOp,1790                                PatternRewriter &rewriter) const override {1791    OpOperand *operand =1792        getWarpResult(warpOp, llvm::IsaPred<vector::TransposeOp>);1793    if (!operand)1794      return rewriter.notifyMatchFailure(1795          warpOp, "warp result is not a vector::Transpose op");1796    auto transposeOp = operand->get().getDefiningOp<vector::TransposeOp>();1797    unsigned operandIdx = operand->getOperandNumber();1798    xegpu::DistributeLayoutAttr sourceLayout =1799        xegpu::getDistributeLayoutAttr(transposeOp.getVector());1800    xegpu::DistributeLayoutAttr resultLayout =1801        xegpu::getDistributeLayoutAttr(transposeOp.getResult());1802    if (!sourceLayout || !resultLayout)1803      return rewriter.notifyMatchFailure(1804          transposeOp,1805          "the source or result vector of the transpose op lacks layout "1806          "attribute");1807    int64_t sourceRank = transposeOp.getSourceVectorType().getRank();1808    int64_t resultRank = transposeOp.getResultVectorType().getRank();1809    // Only 2D transposes are supported for now.1810    // TODO: Support nD transposes.1811    if (sourceRank != 2 || resultRank != 2)1812      return rewriter.notifyMatchFailure(1813          transposeOp, "the source or result vector of the transpose op "1814                       "does not have 2D layout");1815    ArrayRef<int64_t> perm = transposeOp.getPermutation();1816    // Result layout must be a transpose of source layout.1817    if (!resultLayout.isTransposeOf(sourceLayout, perm))1818      return rewriter.notifyMatchFailure(1819          transposeOp,1820          "the source or result vector layouts must be 2D transposes of each "1821          "other");1822    FailureOr<VectorType> distributedSourceTypeOrFailure =1823        getDistVecTypeBasedOnLaneLayout(sourceLayout,1824                                        transposeOp.getSourceVectorType());1825    if (failed(distributedSourceTypeOrFailure))1826      return rewriter.notifyMatchFailure(1827          transposeOp, "Failed to distribute the source vector type in "1828                       "vector::Transpose op");1829    SmallVector<size_t> newRetIndices;1830    gpu::WarpExecuteOnLane0Op newWarpOp = moveRegionToNewWarpOpAndAppendReturns(1831        rewriter, warpOp, transposeOp.getVector(),1832        TypeRange{distributedSourceTypeOrFailure.value()}, newRetIndices);1833    rewriter.setInsertionPointAfter(newWarpOp);1834    auto newTransposeOp = vector::TransposeOp::create(1835        rewriter, newWarpOp.getLoc(), newWarpOp.getResult(newRetIndices[0]),1836        perm);1837    Value distributedVal = newWarpOp.getResult(operandIdx);1838    rewriter.replaceAllUsesWith(distributedVal, newTransposeOp.getResult());1839    return success();1840  }1841};1842 1843} // namespace1844 1845namespace {1846struct XeGPUSubgroupDistributePass final1847    : public xegpu::impl::XeGPUSubgroupDistributeBase<1848          XeGPUSubgroupDistributePass> {1849  void runOnOperation() override;1850};1851} // namespace1852 1853void xegpu::populateXeGPUSubgroupDistributePatterns(1854    RewritePatternSet &patterns) {1855  patterns.add<CreateNdDescDistribution, StoreNdDistribution,1856               LoadNdDistribution, DpasDistribution, PrefetchNdDistribution,1857               GpuBarrierDistribution, VectorMultiReductionDistribution,1858               LoadDistribution, StoreDistribution, VectorTransposeDistribution,1859               VectorBitcastDistribution, LoadMatrixDistribution,1860               StoreMatrixDistribution,1861               MemrefExtractAlignedPointerAsIndexDistribution>(1862      patterns.getContext(),1863      /*pattern benefit=*/regularPatternBenefit);1864  // For following patterns, we need to override the regular vector distribution1865  // patterns. Therefore, assign higher benefit.1866  patterns1867      .add<VectorShapeCastDistribution, VectorExtractStridedSliceDistribution,1868           VectorInsertStridedSliceDistribution>(1869          patterns.getContext(),1870          /*pattern benefit=*/highPatternBenefit);1871}1872 1873void xegpu::populateXeGPUMoveFuncBodyToWarpOpPatterns(1874    RewritePatternSet &patterns) {1875  patterns.add<MoveFuncBodyToWarpOp>(patterns.getContext());1876}1877 1878void XeGPUSubgroupDistributePass::runOnOperation() {1879  // Step 1: Attach layouts to op operands.1880  // TODO: Following assumptions are made:1881  // 1) It is assumed that there are no layout conflicts.1882  // 2) Any existing layout attributes attached to the operands are ignored.1883  Operation *op = getOperation();1884  op->walk([&](Operation *op) {1885    for (OpOperand &operand : op->getOpOperands()) {1886      // Layouts are needed for vector type only.1887      if (!isa<VectorType>(operand.get().getType()))1888        continue;1889      if (isa<xegpu::LoadMatrixOp, xegpu::StoreMatrixOp>(op))1890        continue;1891 1892      auto layout = xegpu::getDistributeLayoutAttr(operand.get());1893      if (!layout) {1894        op->emitError("Could not find layout attribute for operand ")1895            << operand.getOperandNumber() << " of operation " << op->getName();1896        signalPassFailure();1897        return;1898      }1899      xegpu::setDistributeLayoutAttr(operand, layout);1900    }1901  });1902  // Step 2: Move all operations of a GPU function inside1903  // gpu.warp_execute_on_lane_0 operation.1904  {1905    RewritePatternSet patterns(&getContext());1906    xegpu::populateXeGPUMoveFuncBodyToWarpOpPatterns(patterns);1907 1908    if (failed(applyPatternsGreedily(getOperation(), std::move(patterns)))) {1909      signalPassFailure();1910      return;1911    }1912    // At this point, we have moved the entire function body inside the1913    // warpOp. Now move any scalar uniform code outside of the warpOp (like1914    // GPU index ops, scalar constants, etc.). This will simplify the1915    // later lowering and avoid custom patterns for these ops.1916    getOperation()->walk([&](Operation *op) {1917      if (auto warpOp = dyn_cast<gpu::WarpExecuteOnLane0Op>(op))1918        vector::moveScalarUniformCode(warpOp);1919    });1920  }1921  // Step 3: Apply subgroup to workitem distribution patterns.1922  RewritePatternSet patterns(&getContext());1923  xegpu::populateXeGPUSubgroupDistributePatterns(patterns);1924  // distributionFn is used by vector distribution patterns to determine the1925  // distributed vector type for a given vector value. In XeGPU subgroup1926  // distribution context, we compute this based on lane layout.1927  auto distributionFn = [](Value val) {1928    VectorType vecType = dyn_cast<VectorType>(val.getType());1929    int64_t vecRank = vecType ? vecType.getRank() : 0;1930    if (vecRank == 0)1931      return AffineMap::get(val.getContext());1932    // Get the layout of the vector type.1933    xegpu::DistributeLayoutAttr layout = xegpu::getDistributeLayoutAttr(val);1934    // If no layout is specified, that means no distribution.1935    if (!layout)1936      return AffineMap::getMultiDimMapWithTargets(vecRank, {},1937                                                  val.getContext());1938    // Expecting vector and layout rank to match.1939    assert(layout.getRank() == vecRank &&1940           "Expecting vector and layout rank to match");1941    // A dimension is distributed only if layout suggests there are1942    // multiple lanes assigned for this dimension and the shape can be evenly1943    // distributed to those lanes.1944    SmallVector<unsigned int> distributedDims;1945    for (auto [i, v] : llvm::enumerate(layout.getEffectiveLaneLayoutAsInt())) {1946      if (v > 1 && vecType.getShape()[i] % v == 0)1947        distributedDims.push_back(i);1948    }1949    return AffineMap::getMultiDimMapWithTargets(vecRank, distributedDims,1950                                                val.getContext());1951  };1952  // TODO: shuffleFn is not used.1953  auto shuffleFn = [](Location loc, OpBuilder &builder, Value val, Value srcIdx,1954                      int64_t warpSz) { return Value(); };1955 1956  auto warpReduction = [](Location loc, OpBuilder &builder, Value input,1957                          vector::CombiningKind kind, uint32_t size) {1958    // First reduce on a single thread to get per lane reduction value.1959    Value laneVal = vector::ReductionOp::create(builder, loc, kind, input);1960    // Parallel reduction using butterfly shuffles.1961    for (uint64_t i = 1; i < size; i <<= 1) {1962      Value shuffled = gpu::ShuffleOp::create(builder, loc, laneVal, i,1963                                              /*width=*/size,1964                                              /*mode=*/gpu::ShuffleMode::XOR)1965                           .getShuffleResult();1966      laneVal = makeArithReduction(builder, loc, kind, laneVal, shuffled);1967    }1968    return laneVal;1969  };1970 1971  vector::populateDistributeReduction(1972      patterns, warpReduction,1973      /*pattern benefit=*/regularPatternBenefit);1974 1975  vector::populatePropagateWarpVectorDistributionPatterns(1976      patterns, distributionFn, shuffleFn,1977      /*pattern benefit=*/regularPatternBenefit);1978  if (failed(applyPatternsGreedily(getOperation(), std::move(patterns)))) {1979    signalPassFailure();1980    return;1981  }1982 1983  // Step 4: Finally, clean up UnrealizedConversionCastOps that were inserted1984  // due to tensor desc type mismatches created by using upstream distribution1985  // patterns (scf.for). This cleanup should only be done if all the ops are1986  // distributed successfully, if some ops are still not distributed and remains1987  // inside any WarpExecuteOnLane0Op we avoid this simplication step to avoid1988  // breaking the IR.1989  bool foundWarpOp = false;1990  getOperation()->walk([&](gpu::WarpExecuteOnLane0Op warpOp) {1991    // Look for WarpOps that are not trivially dead.1992    if (isOpTriviallyDead(warpOp))1993      return WalkResult::advance();1994    foundWarpOp = true;1995    return WalkResult::interrupt();1996  });1997  if (foundWarpOp)1998    return;1999 2000  getOperation()->walk([&](mlir::UnrealizedConversionCastOp op) {2001    // We are only interested in UnrealizedConversionCastOps there were added2002    // for resolving SIMT type mismatches.2003    if (!op->getAttr(resolveSIMTTypeMismatch))2004      return WalkResult::skip();2005 2006    Value input = op.getOperand(0);2007    Value output = op.getResult(0);2008 2009    // Both input and output must have tensor descriptor types.2010    xegpu::TensorDescType inputDescType =2011        mlir::dyn_cast<xegpu::TensorDescType>(input.getType());2012    xegpu::TensorDescType outputDescType =2013        mlir::dyn_cast<xegpu::TensorDescType>(output.getType());2014    assert(inputDescType && outputDescType &&2015           "Unrealized conversion cast must have tensor descriptor types");2016 2017    // tensor_desc<shape, layout> -> tensor_desc<shape> Type of conversions.2018    // This occurs inside scf.for body to resolve the block argument type to2019    // SIMT type.2020    if (inputDescType.getLayout()) {2021      auto argument = mlir::dyn_cast<mlir::BlockArgument>(input);2022      if (argument) {2023        argument.setType(output.getType());2024        output.replaceAllUsesWith(argument);2025        if (auto loopOp = mlir::dyn_cast<mlir::LoopLikeOpInterface>(2026                argument.getOwner()->getParentOp())) {2027          auto result = loopOp.getTiedLoopResult(argument);2028          result.setType(output.getType());2029        }2030      }2031    }2032 2033    // tensor_desc<shape> -> tensor_desc<shape, layout> Type of2034    // conversions. This occurs at the yield op of scf.for body to go back2035    // from SIMT type to original type.2036    if (outputDescType.getLayout())2037      output.replaceAllUsesWith(input);2038 2039    if (op->use_empty())2040      op->erase();2041    return WalkResult::advance();2042  });2043}2044