1571 lines · cpp
1//===- XeGPUWgToSgDistribute.cpp - XeGPU Workgroup to Subgroup 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/XeGPU/Transforms/Passes.h"9 10#include "mlir/Dialect/Affine/Utils.h"11#include "mlir/Dialect/Arith/IR/Arith.h"12#include "mlir/Dialect/Arith/Utils/Utils.h"13#include "mlir/Dialect/GPU/IR/GPUDialect.h"14#include "mlir/Dialect/Index/IR/IndexDialect.h"15#include "mlir/Dialect/Index/IR/IndexOps.h"16#include "mlir/Dialect/Math/IR/Math.h"17#include "mlir/Dialect/MemRef/IR/MemRef.h"18#include "mlir/Dialect/SCF/Transforms/Patterns.h"19#include "mlir/Dialect/Utils/IndexingUtils.h"20#include "mlir/Dialect/XeGPU/IR/XeGPU.h"21#include "mlir/Dialect/XeGPU/Transforms/Transforms.h"22#include "mlir/Dialect/XeGPU/Utils/XeGPUUtils.h"23#include "mlir/Transforms/DialectConversion.h"24#include <optional>25 26namespace mlir {27namespace xegpu {28#define GEN_PASS_DEF_XEGPUWGTOSGDISTRIBUTE29#include "mlir/Dialect/XeGPU/Transforms/Passes.h.inc"30} // namespace xegpu31} // namespace mlir32 33using namespace mlir;34 35namespace {36 37// Retrieve the RangeAttr if it is specified.38static xegpu::RangeAttr getRangeSpecAttr(Operation *op) {39 Operation *parent = op->getParentOfType<scf::IfOp>();40 while (parent) {41 if (auto attr = llvm::dyn_cast_if_present<xegpu::RangeAttr>(42 parent->getAttr("sg_id_range")))43 return attr;44 parent = parent->getParentOfType<scf::IfOp>();45 }46 return {};47}48 49static std::pair<SmallVector<int64_t>, int>50getSgShapeAndCount(ArrayRef<int64_t> shape,51 xegpu::DistributeLayoutAttr layout) {52 int count = 1;53 SmallVector<int64_t> sgShape(shape);54 if (layout && layout.isForWorkgroup()) {55 SmallVector<int64_t> sgLayout = layout.getEffectiveSgLayoutAsInt();56 if (!layout.getEffectiveSgDataAsInt().empty())57 sgShape = layout.getEffectiveSgDataAsInt();58 else if (auto maybeDerivedSgData = computeShapeRatio(shape, sgLayout))59 sgShape = *maybeDerivedSgData;60 SmallVector<int64_t> distUnit = computeElementwiseMul(sgLayout, sgShape);61 // Clamp distUnit to the original shape to handle cases where data is62 // shared among subgroups, which may cause distUnit to exceed the original63 // shape.64 for (size_t i = 0; i < distUnit.size(); ++i)65 distUnit[i] = std::min(shape[i], distUnit[i]);66 count = computeProduct(shape) / computeProduct(distUnit);67 }68 return std::make_pair(sgShape, count);69}70 71/// Utility helper for deriving a list of offsets for each sub-TensorDescs72/// or sub-MemDescs to be accessed by current subgroup (sgId) based on the73/// associated distribute layout attribute, the shape, subgroup id and the74/// original offsets of the op75template <76 typename OpType,77 typename = std::enable_if_t<llvm::is_one_of<78 OpType, xegpu::CreateNdDescOp, xegpu::LoadNdOp, xegpu::StoreNdOp,79 xegpu::PrefetchNdOp, xegpu::LoadMatrixOp, xegpu::StoreMatrixOp>::value>>80static LogicalResult81genOffsetsList(ConversionPatternRewriter &rewriter, OpType op,82 SmallVector<SmallVector<OpFoldResult>> &offsetsList) {83 Location loc = op.getLoc();84 SmallVector<OpFoldResult> origOffsets = op.getMixedOffsets();85 // not applicable to ops without offsets operands.86 if (origOffsets.empty())87 return failure();88 89 // if op is xegpu::CreateNdDescOp, call op.getDescLayoutAttr()90 xegpu::DistributeLayoutAttr layout;91 if constexpr (std::is_same_v<OpType, xegpu::LoadMatrixOp> ||92 std::is_same_v<OpType, xegpu::StoreMatrixOp>) {93 layout = op.getLayoutAttr();94 } else {95 layout = op.getDescLayoutAttr();96 }97 98 // not applicable to ops without workgroup layout attributes99 if (!layout || !layout.isForWorkgroup())100 return failure();101 102 Value sgId =103 gpu::SubgroupIdOp::create(rewriter, loc, /*upper_bound=*/nullptr);104 105 // verify and adjust the sgId if the range specifier is present106 xegpu::RangeAttr sgIdRange = getRangeSpecAttr(op);107 if (sgIdRange) {108 int64_t startOfRange = sgIdRange.getStart().getInt();109 int64_t endOfRange = sgIdRange.getEnd().getInt();110 // verify the RangeAttr against the layout attribute111 if (layout.getNumSubgroups() != endOfRange - startOfRange)112 return rewriter.notifyMatchFailure(113 op, "sg_layout size must match the sg_id_range");114 // adjust the sgId if necessary115 if (startOfRange > 0) {116 Value startOfRangeVal =117 arith::ConstantIndexOp::create(rewriter, loc, startOfRange);118 sgId = index::SubOp::create(rewriter, loc, sgId, startOfRangeVal);119 }120 }121 122 // Compute the list of subgroup-relative offsets for sub-tensors or sub-memory123 // descriptors to be accessed, based on the layout information.124 ArrayRef<int64_t> wgShape = op.getDataShape();125 auto maybeDescOffsets =126 layout.computeDistributedCoords(rewriter, loc, sgId, wgShape);127 if (failed(maybeDescOffsets))128 return failure();129 130 // Compute the final global offsets for each accessed sub-tensor131 // or sub-memory descriptor.132 for (const auto &sgOffsets : *maybeDescOffsets) {133 SmallVector<OpFoldResult> newOffsets = xegpu::addWithRightAligned(134 rewriter, loc, getAsOpFoldResult(sgOffsets), origOffsets);135 offsetsList.push_back(std::move(newOffsets));136 }137 138 // callback(offsetsList);139 return success();140}141 142/// This pattern transforms the CreateNdDescOp to create a subgroup descriptor143/// from a workgroup descriptor. It replaces the offsets and sizes with144/// appropriate values for the subgroup.145/// It uses round-robin assignment to distribute the work to the subgroups.146/// Following create_nd_desc operation:,147/// %tdesc = xegpu.create_nd_tdesc %src[0, 0] : memref<24x24xf32>148/// -> !xegpu.tensor_desc<24x24xf32, #xegpu.layout<sg_layout = [4, 4],149/// sg_data = [2, 2], lane_layout = [2, 2], lane_data = [1, 1]>>150/// is converted to 9 subgroup level operations based on the sg_layout &151/// sg_data:152/// %tdesc = xegpu.create_nd_tdesc %src[off1, off2] : memref<24x24xf32> ->153/// !xegpu.tensor_desc<2x2xf32, #xegpu.layout<lane_layout = [2, 2],154/// lane_data = [1, 1]>>155///156/// The sg_layout and sg_data attributes are dropped after the pass as they are157/// no longer needed.158///159/// 24x24 matrix distribution example:160/// sg_layout = [4, 4], sg_data = [2, 2]161/// Each 8x8 matrix within the 24x24 matrix is called a distribution unit.162/// dist_unit_shape = [8, 8] --> sg_layout[i] * sg_data[i]163///164/// +------------------------+165/// | 8x8 | 8x8 | 8x8 | <- 3 tiles across166/// |-----+-----+-----|167/// | 8x8 | 8x8 | 8x8 | <- 3 tiles down168/// |-----+-----+-----|169/// | 8x8 | 8x8 | 8x8 |170/// +------------------------+171///172/// Each 8x8 tile is further subdivided among subgroups:173/// +------------------------+174/// | 2x2 2x2 2x2 2x2 | <- 4 subgroups across (each handles 2 columns)175/// | 2x2 2x2 2x2 2x2 | <- 4 subgroups down (each handles 2 rows)176/// | 2x2 2x2 2x2 2x2 |177/// | 2x2 2x2 2x2 2x2 |178/// +------------------------+179///180/// Since the 24x24 matrix is divided into 8x8 distribution units, there will be181/// 9 distribution units (3x3) in total. Hence the 9 subgroup level operations.182 183/// The pass currently has entire distribution logic in the WgToSgCreateNdOp184/// pattern and all the other ops just follow.185/// TODO: Decouple the distribution logic from WgToSgCreateNdOp for all the186/// ops in the pass.187struct WgToSgCreateNdOp : public OpConversionPattern<xegpu::CreateNdDescOp> {188 using OpConversionPattern<xegpu::CreateNdDescOp>::OpConversionPattern;189 190 LogicalResult191 matchAndRewrite(xegpu::CreateNdDescOp op, OneToNOpAdaptor adaptor,192 ConversionPatternRewriter &rewriter) const override {193 SmallVector<SmallVector<OpFoldResult>> offsetsList;194 if (failed(genOffsetsList(rewriter, op, offsetsList)))195 return failure();196 197 MLIRContext *ctx = op.getContext();198 xegpu::TensorDescType tdescTy = op.getType();199 ArrayRef<int64_t> wgShape = tdescTy.getShape();200 Type elemTy = tdescTy.getElementType();201 xegpu::DistributeLayoutAttr layout = tdescTy.getLayoutAttr();202 SmallVector<int64_t> sgShape = getSgShapeAndCount(wgShape, layout).first;203 auto newTdescTy =204 xegpu::TensorDescType::get(ctx, sgShape, elemTy, tdescTy.getEncoding(),205 layout.dropSgLayoutAndData());206 207 SmallVector<Value> newOps;208 for (auto offsets : offsetsList) {209 auto newOp = xegpu::CreateNdDescOp::create(210 rewriter, op.getLoc(), newTdescTy, op.getSource(), offsets,211 op.getMixedSizes(), op.getMixedStrides());212 213 newOps.push_back(newOp);214 }215 rewriter.replaceOpWithMultiple(op, {newOps});216 217 return success();218 }219};220 221// This pattern transforms the CreateNdDescOp without offsets to create a222// subgroup descriptor from a workgroup descriptor223struct WgToSgCreateNdOpNoOffset224 : public OpConversionPattern<xegpu::CreateNdDescOp> {225 using OpConversionPattern<xegpu::CreateNdDescOp>::OpConversionPattern;226 227 LogicalResult228 matchAndRewrite(xegpu::CreateNdDescOp op, OneToNOpAdaptor adaptor,229 ConversionPatternRewriter &rewriter) const override {230 231 // Check no offsets are specified.232 if (!op.getMixedOffsets().empty())233 return failure();234 235 Location loc = op.getLoc();236 MLIRContext *ctx = op.getContext();237 xegpu::TensorDescType tdescTy = op.getType();238 auto layout = dyn_cast<xegpu::LayoutAttr>(tdescTy.getLayout());239 if (!layout || !layout.isForWorkgroup())240 return failure();241 242 Type elemTy = tdescTy.getElementType();243 ArrayRef<int64_t> wgShape = tdescTy.getShape();244 245 SmallVector<int64_t> sgShape;246 int count;247 std::tie(sgShape, count) = getSgShapeAndCount(wgShape, layout);248 xegpu::TensorDescType newTdescTy =249 xegpu::TensorDescType::get(ctx, sgShape, elemTy, tdescTy.getEncoding(),250 layout.dropSgLayoutAndData());251 252 SmallVector<Value> newCreateNdOps(count);253 std::generate(newCreateNdOps.begin(), newCreateNdOps.end(), [&]() {254 return xegpu::CreateNdDescOp::create(rewriter, loc, newTdescTy,255 op.getSource(), op.getMixedSizes(),256 op.getMixedStrides());257 });258 259 rewriter.replaceOpWithMultiple(op, {newCreateNdOps});260 return success();261 }262};263 264/// This pattern transforms the LoadNdOp to load subgroup data.265struct WgToSgLoadNdOp : public OpConversionPattern<xegpu::LoadNdOp> {266 using OpConversionPattern<xegpu::LoadNdOp>::OpConversionPattern;267 LogicalResult268 matchAndRewrite(xegpu::LoadNdOp op, OneToNOpAdaptor adaptor,269 ConversionPatternRewriter &rewriter) const override {270 if (!op.getMixedOffsets().empty())271 return failure();272 273 SmallVector<Value> newLoadOps;274 for (auto src : adaptor.getTensorDesc()) {275 xegpu::TensorDescType tdescTy =276 dyn_cast<xegpu::TensorDescType>(src.getType());277 ArrayRef<int64_t> srcShape = tdescTy.getShape();278 VectorType newResTy = VectorType::get(srcShape, tdescTy.getElementType());279 auto newLoadOp = xegpu::LoadNdOp::create(rewriter, op.getLoc(), newResTy,280 src, op->getAttrs());281 newLoadOps.push_back(newLoadOp);282 }283 rewriter.replaceOpWithMultiple(op, {newLoadOps});284 return mlir::success();285 }286};287 288/// This pattern transforms the StoreNdOp to store to a subgroup descriptor289/// It creates a StoreNdOp op to store the updated values to the new subgroup290/// src tensor descriptors.291struct WgToSgStoreNdOp : public OpConversionPattern<xegpu::StoreNdOp> {292 using OpConversionPattern<xegpu::StoreNdOp>::OpConversionPattern;293 LogicalResult294 matchAndRewrite(xegpu::StoreNdOp op, OneToNOpAdaptor adaptor,295 ConversionPatternRewriter &rewriter) const override {296 if (!op.getMixedOffsets().empty())297 return failure();298 299 for (auto [v, t] : llvm::zip(adaptor.getValue(), adaptor.getTensorDesc()))300 xegpu::StoreNdOp::create(rewriter, op.getLoc(), v, t, op.getL1HintAttr(),301 op.getL2HintAttr(), op.getL3HintAttr());302 303 rewriter.eraseOp(op);304 return success();305 }306};307 308// This pattern transforms the LoadNdOp with explicit offsets to load309// subgroup data.310struct WgToSgLoadNdOpWithOffset : public OpConversionPattern<xegpu::LoadNdOp> {311 using OpConversionPattern<xegpu::LoadNdOp>::OpConversionPattern;312 LogicalResult313 matchAndRewrite(xegpu::LoadNdOp op, OneToNOpAdaptor adaptor,314 ConversionPatternRewriter &rewriter) const override {315 316 SmallVector<SmallVector<OpFoldResult>> offsetsList;317 if (failed(genOffsetsList(rewriter, op, offsetsList)))318 return failure();319 320 SmallVector<Value> newOps;321 for (auto [tdesc, offsets] :322 llvm::zip(adaptor.getTensorDesc(), offsetsList)) {323 auto tdescTy = dyn_cast<xegpu::TensorDescType>(tdesc.getType());324 VectorType newResTy =325 VectorType::get(tdescTy.getShape(), tdescTy.getElementType());326 auto newOp = xegpu::LoadNdOp::create(327 rewriter, op.getLoc(), newResTy, tdesc, offsets,328 /*packed = */ nullptr, /*transpose = */ nullptr, op.getL1HintAttr(),329 op.getL2HintAttr(), op.getL3HintAttr());330 newOps.push_back(newOp);331 }332 rewriter.replaceOpWithMultiple(op, {newOps});333 334 return success();335 }336};337 338// This pattern transforms the StoreNdOp with explicit offsets to store339// subgroup data.340struct WgToSgStoreNdOpWithOffset341 : public OpConversionPattern<xegpu::StoreNdOp> {342 using OpConversionPattern<xegpu::StoreNdOp>::OpConversionPattern;343 LogicalResult344 matchAndRewrite(xegpu::StoreNdOp op, OneToNOpAdaptor adaptor,345 ConversionPatternRewriter &rewriter) const override {346 SmallVector<SmallVector<OpFoldResult>> offsetsList;347 if (failed(genOffsetsList(rewriter, op, offsetsList)))348 return failure();349 350 for (auto [v, tdesc, offsets] :351 llvm::zip(adaptor.getValue(), adaptor.getTensorDesc(), offsetsList)) {352 xegpu::StoreNdOp::create(rewriter, op.getLoc(), v, tdesc, offsets,353 op.getL1HintAttr(), op.getL2HintAttr(),354 op.getL3HintAttr());355 }356 rewriter.eraseOp(op);357 358 return success();359 }360};361 362// This pattern transforms the PrefetchNdOp with explicit offsets to prefetch363// subgroup data.364struct WgToSgPrefetchNdOpWithOffset365 : public OpConversionPattern<xegpu::PrefetchNdOp> {366 using OpConversionPattern<xegpu::PrefetchNdOp>::OpConversionPattern;367 LogicalResult368 matchAndRewrite(xegpu::PrefetchNdOp op, OneToNOpAdaptor adaptor,369 ConversionPatternRewriter &rewriter) const override {370 SmallVector<SmallVector<OpFoldResult>> offsetsList;371 if (failed(genOffsetsList(rewriter, op, offsetsList)))372 return failure();373 374 for (auto [tdesc, offsets] :375 llvm::zip(adaptor.getTensorDesc(), offsetsList)) {376 xegpu::PrefetchNdOp::create(rewriter, op.getLoc(), tdesc, offsets,377 op.getL1HintAttr(), op.getL2HintAttr(),378 op.getL3HintAttr());379 }380 rewriter.eraseOp(op);381 382 return success();383 }384};385 386/// This pattern transforms the UpdateNdOffsetOp to update the offsets of a387/// subgroup descriptor. It creates an UpdateNdOffsetOp op to update the388/// offsets of the new subgroup src tensor descriptors.389struct WgToSgUpdateNdOffsetOp390 : public OpConversionPattern<xegpu::UpdateNdOffsetOp> {391 using OpConversionPattern<xegpu::UpdateNdOffsetOp>::OpConversionPattern;392 LogicalResult393 matchAndRewrite(xegpu::UpdateNdOffsetOp op, OneToNOpAdaptor adaptor,394 ConversionPatternRewriter &rewriter) const override {395 llvm::SmallVector<Value> newUpdateTileOffsetOps;396 for (auto tDesc : adaptor.getTensorDesc()) {397 auto newUpdateTileOffsetOp = xegpu::UpdateNdOffsetOp::create(398 rewriter, op.getLoc(), tDesc.getType(), tDesc, op.getOffsets(),399 op.getConstOffsets());400 newUpdateTileOffsetOps.push_back(newUpdateTileOffsetOp);401 }402 403 rewriter.replaceOpWithMultiple(op, {newUpdateTileOffsetOps});404 return success();405 }406};407 408/// This pattern transforms the DpasOp to work at subgroup level.409struct WgToSgDpasOp : public OpConversionPattern<xegpu::DpasOp> {410 using OpConversionPattern<xegpu::DpasOp>::OpConversionPattern;411 LogicalResult412 matchAndRewrite(xegpu::DpasOp op, OneToNOpAdaptor adaptor,413 ConversionPatternRewriter &rewriter) const override {414 Location loc = op.getLoc();415 VectorType resultTy = op.getResult().getType();416 if (resultTy.getRank() != 2)417 return failure();418 419 auto originalLayout = xegpu::getDistributeLayoutAttr(op.getResult());420 if (!originalLayout)421 return failure();422 423 size_t i = 0;424 SmallVector<Value> newDpasOps;425 for (auto aVec : adaptor.getLhs()) {426 for (auto bVec : adaptor.getRhs()) {427 428 llvm::SmallVector<Value> operands({aVec, bVec});429 Value tmpC;430 if (op.getAcc()) {431 tmpC = adaptor.getAcc()[i++];432 operands.push_back(tmpC);433 }434 435 ArrayRef<int64_t> aVecShape =436 llvm::cast<VectorType>(aVec.getType()).getShape();437 ArrayRef<int64_t> bVecShape =438 llvm::cast<VectorType>(bVec.getType()).getShape();439 VectorType resTy = VectorType::get({aVecShape[0], bVecShape[1]},440 resultTy.getElementType());441 tmpC = xegpu::DpasOp::create(rewriter, loc, resTy, operands);442 xegpu::setDistributeLayoutAttr(cast<OpResult>(tmpC),443 originalLayout.dropSgLayoutAndData());444 445 newDpasOps.push_back(tmpC);446 }447 }448 rewriter.replaceOpWithMultiple(op, {newDpasOps});449 return success();450 }451};452 453/// This pattern transforms the PrefetchNdOp to prefetch the subgroup data.454struct WgToSgPrefetchNdOp : public OpConversionPattern<xegpu::PrefetchNdOp> {455 using OpConversionPattern<xegpu::PrefetchNdOp>::OpConversionPattern;456 LogicalResult457 matchAndRewrite(xegpu::PrefetchNdOp op, OneToNOpAdaptor adaptor,458 ConversionPatternRewriter &rewriter) const override {459 460 int64_t offsetSize = static_cast<int64_t>(op.getOffsets().size());461 if ((offsetSize != 0) || op.getConstOffsetsAttr())462 return failure();463 464 for (auto src : adaptor.getTensorDesc())465 xegpu::PrefetchNdOp::create(rewriter, op.getLoc(), TypeRange(), src,466 op->getAttrs());467 rewriter.eraseOp(op);468 return success();469 }470};471 472/// This pattern transforms vector.broadcast ops to work at subgroup level.473struct WgToSgVectorBroadcastOp474 : public OpConversionPattern<vector::BroadcastOp> {475 using OpConversionPattern<vector::BroadcastOp>::OpConversionPattern;476 477 LogicalResult478 matchAndRewrite(vector::BroadcastOp op, OneToNOpAdaptor adaptor,479 ConversionPatternRewriter &rewriter) const override {480 481 VectorType resultType = op.getResult().getType();482 ArrayRef<int64_t> wgShape = resultType.getShape();483 484 xegpu::DistributeLayoutAttr layout =485 xegpu::getDistributeLayoutAttr(op.getResult());486 if (!layout || !layout.isForWorkgroup())487 return failure();488 489 SmallVector<int64_t> sgShape = getSgShapeAndCount(wgShape, layout).first;490 VectorType newResultType =491 VectorType::get(sgShape, resultType.getElementType());492 493 if (!xegpu::XeGPUDialect::isEvenlyDistributable(wgShape, layout))494 return failure();495 496 SmallVector<Value> newBroadcastOps;497 for (auto operand : adaptor.getOperands().front()) {498 auto newBroadcast = vector::BroadcastOp::create(rewriter, op.getLoc(),499 newResultType, operand);500 xegpu::setDistributeLayoutAttr(newBroadcast->getResult(0),501 layout.dropSgLayoutAndData());502 503 newBroadcastOps.push_back(newBroadcast.getResult());504 }505 rewriter.replaceOpWithMultiple(op, {newBroadcastOps});506 return success();507 }508};509 510// This pattern transforms elementwise ops to work at subgroup level.511struct WgToSgElementwiseOp : public ConversionPattern {512 WgToSgElementwiseOp(MLIRContext *ctx)513 : ConversionPattern(MatchAnyOpTypeTag(), /*benefit=*/1, ctx) {}514 515 LogicalResult516 matchAndRewrite(Operation *op, ArrayRef<ValueRange> operands,517 ConversionPatternRewriter &rewriter) const override {518 // Only match ops with elementwise trait and single result.519 if (!OpTrait::hasElementwiseMappableTraits(op) || op->getNumResults() != 1)520 return failure();521 522 auto resultType = dyn_cast<VectorType>(op->getResult(0).getType());523 assert(resultType && "Expected result to be a VectorType");524 525 ArrayRef<int64_t> wgShape = resultType.getShape();526 527 xegpu::DistributeLayoutAttr layout =528 xegpu::getDistributeLayoutAttr(op->getResult(0));529 if (!layout || !layout.isForWorkgroup())530 return failure();531 532 SmallVector<int64_t> sgShape = getSgShapeAndCount(wgShape, layout).first;533 534 size_t numVariants = operands.empty() ? 0 : operands.front().size();535 536 if (llvm::any_of(operands, [&](const ValueRange &operandVec) {537 return operandVec.size() != numVariants;538 }))539 return failure();540 541 SmallVector<Value> newResults;542 VectorType newResultType =543 VectorType::get(sgShape, resultType.getElementType());544 545 for (size_t i = 0; i < numVariants; ++i) {546 SmallVector<Value> opOperands;547 for (auto &operandVec : operands)548 opOperands.push_back(operandVec[i]);549 550 OperationState state(op->getLoc(), op->getName());551 state.addOperands(opOperands);552 state.addTypes(newResultType);553 // Copy all attributes, but update "layout_result_0" to drop554 // sgLayout/sgData555 for (auto attr : op->getAttrs()) {556 if (auto layout =557 dyn_cast<xegpu::DistributeLayoutAttr>(attr.getValue())) {558 if (!layout.getEffectiveLaneLayoutAsInt().empty() ||559 !layout.getEffectiveInstDataAsInt().empty())560 state.addAttribute(attr.getName(), layout.dropSgLayoutAndData());561 } else {562 state.addAttribute(attr.getName(), attr.getValue());563 }564 }565 Operation *newOp = rewriter.create(state);566 newResults.push_back(newOp->getResult(0));567 }568 569 rewriter.replaceOpWithMultiple(op, {newResults});570 return success();571 }572};573 574// clang-format off575// Pattern for lowering ConvertLayoutOp based on sg_layout and sg_data.576// If input_layout and target_layout have identical sg_layout and sg_data,577// the op is rewritten to a subgroup-level ConvertLayoutOp with these fields578// dropped. For example:579// #a = #xegpu.layout<sg_layout = [2, 2], sg_data = [16, 16], inst_data = [16, 16]>580// #b = #xegpu.layout<sg_layout = [2, 2], sg_data = [16, 16], inst_data = [8, 16]>581// xegpu.convert_layout %1 <{input_layout = #a, target_layout = #b}> : vector<32x64xf32>582// becomes:583// #a = #xegpu.layout<inst_data = [16, 16]>584// #b = #xegpu.layout<inst_data = [8, 16]>585// xegpu.convert_layout %1 <{input_layout = #a, target_layout = #b}> : vector<16x16xf32>586// (vector<16x16xf32> is determined by sg_data = [16, 16])587//588// If sg_layout or sg_data differ, SLM is used to redistribute data across subgroups.589// For example:590// #a = #xegpu.layout<sg_layout = [1, 4], sg_data = [32, 16], inst_data = [16, 16]>591// #b = #xegpu.layout<sg_layout = [2, 2], sg_data = [16, 32], inst_data = [8, 16]>592// xegpu.convert_layout %1 <{input_layout = #a, target_layout = #b}> : vector<32x64xf32>593// is lowered to:594// #a = #xegpu.layout<inst_data = [16, 16]>595// #b = #xegpu.layout<inst_data = [8, 16]>596// store_matrix %1, %slm <{layout_input_0 = #a}> : vector<32x16>, mem_desc<32x64xf32>597// %d = load_matrix %slm <{layout_result_0 = #a}> : mem_desc<32x64xf32> -> vector<16x32xf32>598// xegpu.convert_layout %d <{input_layout = #a, target_layout = #b}> : vector<16x32xf32>599// clang-format on600struct WgToSgConvertLayoutOp601 : public OpConversionPattern<xegpu::ConvertLayoutOp> {602 using OpConversionPattern<xegpu::ConvertLayoutOp>::OpConversionPattern;603 LogicalResult604 matchAndRewrite(xegpu::ConvertLayoutOp op, OneToNOpAdaptor adaptor,605 ConversionPatternRewriter &rewriter) const override {606 // TODO: currently, we only support LayoutAttr607 auto input = dyn_cast<xegpu::LayoutAttr>(op.getInputLayout());608 auto target = dyn_cast<xegpu::LayoutAttr>(op.getTargetLayout());609 610 if (!input || !target || !input.isForWorkgroup() ||611 !target.isForWorkgroup())612 return rewriter.notifyMatchFailure(613 op, "Input and target layouts must have subgroup layout");614 615 DenseI32ArrayAttr inputSgLayout = input.getSgLayout();616 DenseI32ArrayAttr inputSgData = input.getSgData();617 DenseI32ArrayAttr inputOrder = input.getOrder();618 DenseI32ArrayAttr targetSgLayout = target.getSgLayout();619 DenseI32ArrayAttr targetSgData = target.getSgData();620 DenseI32ArrayAttr targetOrder = target.getOrder();621 622 // TODO: currently we only support for optimal case, where input and623 // output has the same sg_layout and sg_data, so SLM is not involved.624 if (inputSgLayout != targetSgLayout || inputSgData != targetSgData ||625 inputOrder != targetOrder)626 return failure();627 628 input = input.dropSgLayoutAndData();629 target = target.dropSgLayoutAndData();630 631 SmallVector<Value> newOps(adaptor.getSource());632 if (input && target) {633 // keep the ConvertLayoutOp for rest fields, e.g., inst_data.634 for (auto [i, src] : llvm::enumerate(adaptor.getSource())) {635 auto newOp = xegpu::ConvertLayoutOp::create(636 rewriter, op.getLoc(), src.getType(), src, input, target);637 newOps[i] = newOp;638 }639 }640 rewriter.replaceOpWithMultiple(op, {newOps});641 return success();642 }643};644 645// Handles UnrealizedConversionCastOp generated during646// SCFStructuralTypeConversions (step 1). This op may appear as either a647// target or source materialization for Vector values, e.g.:648// 1. unrealized_cast %1 : vector<256xf32> to vector<16xf32>, ...649// 2. unrealized_cast %1 : vector<16xf32>, ... to vector<256xf32>650// it could be either 1:N or N:1 cast. In both cases, the pattern651// simply forwards the inputs to the outputs using 1:1 or 1:N interface.652// for example, the following scf::forOp653// ```654// %for = scf.for ... iter_args(%arg1 = %0)->(vector<128x128xf16>) {655// %n = use(%arg1): vector<128x128xf16>656// scf.yield %n : vector<128x128xf16>657// }658// ```659// Could be converted to:660// ```661// %1 = unrealized_conversion_cast %0662// : vector<128x128xf16> to vector<16x16xf16>, vector<16x16xf16>663// %for:2 = scf.for ... iter_args(%arg1 = %1#1, %arg2 = %1#2)664// -> (vector<16x16xf16>, vector<16x16xf16) {665// %m = unrealized_conversion_cast %arg1, %arg2666// : vector<16x16xf16>, vector<16x16xf16> to vector<128x128xf16>667// %n = use(%m): vector<128x128xf16>668// %b = unrealized_conversion_cast %n669// : vector<128x128xf16> to vector<16x16xf16>, vector<16x16xf16>670// scf.yield %b#1, %b#2 : vector<16x16xf16>, vector<16x16xf16>671// }672// %cast = unrealized_conversion_cast %for:2673// : vector<16x16xf16>, vector<16x16xf16> to vector<128x128xf16>674// ```675// TODO: remove it when context-aware type converter is ready.676struct UnrealizedConversionCastOpPattern677 : public OpConversionPattern<mlir::UnrealizedConversionCastOp> {678 using OpConversionPattern<679 mlir::UnrealizedConversionCastOp>::OpConversionPattern;680 681 mlir::LogicalResult682 matchAndRewrite(mlir::UnrealizedConversionCastOp op, OneToNOpAdaptor adaptor,683 ConversionPatternRewriter &rewriter) const override {684 SmallVector<Value> inputs = xegpu::flattenValues(adaptor.getInputs());685 686 auto inputTy = dyn_cast<VectorType>(inputs[0].getType());687 auto outputTy = dyn_cast<VectorType>(op->getOpResult(0).getType());688 689 if (!inputTy || !outputTy || !llvm::all_equal(op->getResultTypes()) ||690 !llvm::all_equal(ValueRange(inputs).getTypes()))691 return failure();692 693 // Handles the case "cast %1 : vector<256xf32> to vector<16xf32>, ...".694 // It is generated by source materialization (e.g., inits to scf forOp).695 // The input values provided by the adaptor should already be distributed,696 // and their types should correspond exactly to the result types of the697 // operation.698 if (op.getNumOperands() == 1 &&699 llvm::equal(ValueRange(inputs).getTypes(), op->getResultTypes())) {700 rewriter.replaceOp(op, inputs);701 return success();702 }703 704 // Handles the case "cast %1 : vector<16xf32>, ... to vector<256xf32>".705 // It is generated by target materialization (e.g., arguments/results706 // of scf forOp). All input values must have the same vector type, and707 // their shape must be evenly divisible by the output vector's shape708 // (determined by the nature of the workgroup to subgroup distribution).709 // TODO: it is not safe to do such forward, since such N:1 cast could be710 // from others.711 if (op.getNumResults() == 1 &&712 computeShapeRatio(outputTy.getShape(), inputTy.getShape())) {713 rewriter.replaceOpWithMultiple(op, {inputs});714 return success();715 }716 717 return mlir::failure();718 }719};720 721// This pattern distributes arith.constant op into subgroup-level constants722struct WgToSgArithConstantOp : public OpConversionPattern<arith::ConstantOp> {723 using OpConversionPattern<arith::ConstantOp>::OpConversionPattern;724 725 LogicalResult726 matchAndRewrite(arith::ConstantOp op, OneToNOpAdaptor adaptor,727 ConversionPatternRewriter &rewriter) const override {728 auto vecAttr = dyn_cast<DenseElementsAttr>(op.getValue());729 auto vecType = dyn_cast<VectorType>(op.getType());730 if (!vecAttr || !vecType)731 return failure();732 733 xegpu::DistributeLayoutAttr layout =734 xegpu::getDistributeLayoutAttr(op.getResult());735 if (!layout || !layout.isForWorkgroup())736 return failure();737 738 ArrayRef<int64_t> wgShape = vecType.getShape();739 SmallVector<int64_t> sgShape;740 int count;741 std::tie(sgShape, count) = getSgShapeAndCount(wgShape, layout);742 743 auto newType = VectorType::get(sgShape, vecType.getElementType());744 Location loc = op.getLoc();745 auto eltType = vecType.getElementType();746 747 auto setLayout = [&](Value val) {748 xegpu::setDistributeLayoutAttr(llvm::dyn_cast<OpResult>(val),749 layout.dropSgLayoutAndData());750 };751 752 if (vecAttr.isSplat()) {753 // Splat: single value for all subgroups754 Attribute singleVal = vecAttr.getSplatValue<Attribute>();755 auto sgAttr = DenseElementsAttr::get(newType, singleVal);756 auto cstOp = arith::ConstantOp::create(rewriter, loc, newType, sgAttr);757 setLayout(cstOp->getResult(0));758 rewriter.replaceOp(op, cstOp);759 return success();760 } else if (sgShape == wgShape) { // if the entire vector is shared by all761 // subgroups, don't distribute762 auto newConstOp =763 arith::ConstantOp::create(rewriter, op.getLoc(), vecType, vecAttr);764 setLayout(newConstOp->getResult(0));765 rewriter.replaceOp(op, newConstOp);766 return success();767 } else {768 // Non-splat constant769 // Only supports 1D & 2D770 // TODO: support other cases that require SLM access771 if (!eltType.isIndex())772 return rewriter.notifyMatchFailure(773 op, "Unsupported element type for non-splat constant op.");774 775 if (wgShape.size() > 2)776 return rewriter.notifyMatchFailure(777 op, "Only 1D & 2D vector constant supported");778 779 SmallVector<Attribute> values(vecAttr.getValues<Attribute>());780 int64_t rowStride = 0, colStride = 0;781 int64_t rows = wgShape.size() == 1 ? 1 : wgShape[0];782 int64_t cols = wgShape.size() == 1 ? wgShape[0] : wgShape[1];783 784 // Compute colStride and rowStride, and check for constant strides.785 if (cols > 1) {786 colStride = cast<IntegerAttr>(values[1]).getInt() -787 cast<IntegerAttr>(values[0]).getInt();788 }789 if (rows > 1) {790 rowStride = cast<IntegerAttr>(values[cols]).getInt() -791 cast<IntegerAttr>(values[0]).getInt();792 }793 794 for (int64_t r = 0; r < rows; ++r) {795 for (int64_t c = 0; c < cols; ++c) {796 int64_t idx = r * cols + c;797 // Check column stride798 if (c > 0 && cols > 1) {799 int64_t prevIdx = r * cols + (c - 1);800 int64_t diff = cast<IntegerAttr>(values[idx]).getInt() -801 cast<IntegerAttr>(values[prevIdx]).getInt();802 if (diff != colStride)803 return rewriter.notifyMatchFailure(804 op, "Non-constant column stride in constant op.");805 }806 // Check row stride807 if (r > 0 && rows > 1) {808 int64_t prevIdx = (r - 1) * cols + c;809 int64_t diff = cast<IntegerAttr>(values[idx]).getInt() -810 cast<IntegerAttr>(values[prevIdx]).getInt();811 if (diff != rowStride)812 return rewriter.notifyMatchFailure(813 op, "Non-constant row stride in constant op.");814 }815 }816 }817 818 // Create a constant for the base tile.819 // For 2D case, extract the top-left sgShape[0] x sgShape[1] submatrix.820 // For 1D case, extract the first sgShape[0] elements.821 SmallVector<Attribute> baseTileValues;822 int baseTileCols = sgShape[sgShape.size() - 1];823 int64_t baseTileRows = sgShape.size() == 1 ? 1 : sgShape[0];824 for (int64_t r = 0; r < baseTileRows; ++r) {825 for (int64_t c = 0; c < baseTileCols; ++c) {826 baseTileValues.push_back(values[r * cols + c]);827 }828 }829 830 auto tileAttr = DenseElementsAttr::get(VectorType::get(sgShape, eltType),831 baseTileValues);832 auto baseConstVec = arith::ConstantOp::create(rewriter, loc, tileAttr);833 834 // Get subgroup id835 Value sgId =836 gpu::SubgroupIdOp::create(rewriter, loc, /*upper_bound=*/nullptr);837 auto sgOffsets =838 layout.computeDistributedCoords(rewriter, loc, sgId, wgShape);839 if (failed(sgOffsets))840 return failure();841 842 SmallVector<Value, 2> strideConsts;843 strideConsts.push_back(844 arith::ConstantIndexOp::create(rewriter, loc, colStride));845 if (rows > 1)846 strideConsts.insert(847 strideConsts.begin(),848 arith::ConstantIndexOp::create(rewriter, loc, rowStride));849 850 SmallVector<Value> newConstOps;851 for (auto offsets : *sgOffsets) {852 // Multiply offset with stride, broadcast it and add to baseConstVec853 Value mulOffset = arith::ConstantIndexOp::create(rewriter, loc, 0);854 for (size_t i = 0; i < strideConsts.size(); ++i) {855 Value mul =856 arith::MulIOp::create(rewriter, loc, rewriter.getIndexType(),857 offsets[i], strideConsts[i]);858 mulOffset = arith::AddIOp::create(859 rewriter, loc, rewriter.getIndexType(), mulOffset, mul);860 }861 // Broadcast to baseConstVec size862 auto bcastOffset = vector::BroadcastOp::create(863 rewriter, loc, baseConstVec.getType(), mulOffset);864 auto finalConst =865 arith::AddIOp::create(rewriter, loc, baseConstVec, bcastOffset);866 setLayout(baseConstVec);867 setLayout(bcastOffset);868 setLayout(finalConst);869 newConstOps.push_back(finalConst);870 }871 rewriter.replaceOpWithMultiple(op, {newConstOps});872 return success();873 }874 }875};876 877// This pattern transforms the LoadGatherOp with explicit offsets to load878// subgroup data879struct WgToSgLoadGatherOpWithOffset880 : public OpConversionPattern<xegpu::LoadGatherOp> {881 using OpConversionPattern<xegpu::LoadGatherOp>::OpConversionPattern;882 LogicalResult883 matchAndRewrite(xegpu::LoadGatherOp op, OneToNOpAdaptor adaptor,884 ConversionPatternRewriter &rewriter) const override {885 886 if (!op.getOffsets())887 return failure();888 889 Location loc = op.getLoc();890 VectorType resultType = dyn_cast<VectorType>(op.getResult().getType());891 if (!resultType)892 return failure();893 ArrayRef<int64_t> wgShape = resultType.getShape();894 895 xegpu::DistributeLayoutAttr layout =896 xegpu::getDistributeLayoutAttr(op.getResult());897 if (!layout || !layout.isForWorkgroup())898 return failure();899 900 SmallVector<int64_t> sgShape = getSgShapeAndCount(wgShape, layout).first;901 902 // The offsets need to be distributed903 auto offsetsVecType =904 dyn_cast<VectorType>(adaptor.getOffsets().front().getType());905 auto maskVecType =906 dyn_cast<VectorType>(adaptor.getMask().front().getType());907 if (!offsetsVecType || !maskVecType ||908 offsetsVecType.getShape() != maskVecType.getShape()) {909 return rewriter.notifyMatchFailure(op,910 "offsets have not been distributed");911 }912 913 SmallVector<Value> newLoadOps;914 auto chunkSizeAttr =915 rewriter.getI64IntegerAttr(op.getChunkSize().value_or(1));916 VectorType newTy = VectorType::get(sgShape, resultType.getElementType());917 for (auto [offsets, mask] :918 llvm::zip(adaptor.getOffsets(), adaptor.getMask())) {919 auto newLayout = layout.dropSgLayoutAndData();920 auto newLoadOp = xegpu::LoadGatherOp::create(921 rewriter, loc, newTy, op.getSource(), offsets, mask, chunkSizeAttr,922 op.getL1HintAttr(), op.getL2HintAttr(), op.getL3HintAttr(),923 newLayout);924 xegpu::setDistributeLayoutAttr(newLoadOp->getResult(0), newLayout);925 newLoadOps.push_back(newLoadOp);926 }927 rewriter.replaceOpWithMultiple(op, {newLoadOps});928 return success();929 }930};931 932// This pattern transforms the StoreScatterOp with explicit offsets to store933// subgroup data934struct WgToSgStoreScatterOpWithOffset935 : public OpConversionPattern<xegpu::StoreScatterOp> {936 using OpConversionPattern<xegpu::StoreScatterOp>::OpConversionPattern;937 LogicalResult938 matchAndRewrite(xegpu::StoreScatterOp op, OneToNOpAdaptor adaptor,939 ConversionPatternRewriter &rewriter) const override {940 941 if (!op.getOffsets())942 return failure();943 944 Location loc = op.getLoc();945 VectorType valueType = dyn_cast<VectorType>(op.getValue().getType());946 if (!valueType)947 return failure();948 949 xegpu::DistributeLayoutAttr layout =950 xegpu::getDistributeLayoutAttr(op.getOperand(0));951 if (!layout || !layout.isForWorkgroup())952 return failure();953 954 // The offsets need to be distributed955 auto offsetsVecType =956 dyn_cast<VectorType>(adaptor.getOffsets().front().getType());957 auto maskVecType =958 dyn_cast<VectorType>(adaptor.getMask().front().getType());959 if (!offsetsVecType || !maskVecType ||960 offsetsVecType.getShape() != maskVecType.getShape()) {961 return rewriter.notifyMatchFailure(op,962 "offsets have not been distributed");963 }964 965 auto chunkSizeOpt = op.getChunkSize();966 int64_t chunkSize = chunkSizeOpt ? static_cast<int64_t>(*chunkSizeOpt) : 1;967 auto chunkSizeAttr = rewriter.getI64IntegerAttr(chunkSize);968 for (auto [val, offs, mask] : llvm::zip(969 adaptor.getValue(), adaptor.getOffsets(), adaptor.getMask())) {970 auto store = xegpu::StoreScatterOp::create(971 rewriter, loc, val, op.getDest(), offs, mask, chunkSizeAttr,972 op.getL1HintAttr(), op.getL2HintAttr(), op.getL3HintAttr(),973 layout.dropSgLayoutAndData());974 // Update the layout attribute to drop sg_layout and sg_data.975 for (OpOperand &operand : store->getOpOperands()) {976 // Skip for operand one (memref)977 if (operand.getOperandNumber() == 1)978 continue;979 xegpu::setDistributeLayoutAttr(operand, layout.dropSgLayoutAndData());980 }981 }982 rewriter.eraseOp(op);983 return success();984 }985};986 987struct WgToSgLoadMatrixOp : public OpConversionPattern<xegpu::LoadMatrixOp> {988 using OpConversionPattern<xegpu::LoadMatrixOp>::OpConversionPattern;989 LogicalResult990 matchAndRewrite(xegpu::LoadMatrixOp op, OneToNOpAdaptor adaptor,991 ConversionPatternRewriter &rewriter) const override {992 993 SmallVector<SmallVector<OpFoldResult>> offsetsList;994 if (failed(genOffsetsList(rewriter, op, offsetsList)))995 return failure();996 997 ArrayRef<int64_t> wgShape = op.getDataShape();998 VectorType valueTy = llvm::dyn_cast<VectorType>(op.getRes().getType());999 assert(valueTy && "the value type must be vector type!");1000 Type elemTy = valueTy.getElementType();1001 1002 xegpu::DistributeLayoutAttr layout = op.getLayoutAttr();1003 SmallVector<int64_t> sgShape = getSgShapeAndCount(wgShape, layout).first;1004 VectorType newResTy = VectorType::get(sgShape, elemTy);1005 SmallVector<Value> newOps;1006 for (auto offsets : offsetsList) {1007 auto newOp = xegpu::LoadMatrixOp::create(rewriter, op.getLoc(), newResTy,1008 op.getMemDesc(), offsets,1009 layout.dropSgLayoutAndData());1010 newOps.push_back(newOp);1011 }1012 rewriter.replaceOpWithMultiple(op, {newOps});1013 1014 return success();1015 }1016};1017 1018struct WgToSgStoreMatrixOp : public OpConversionPattern<xegpu::StoreMatrixOp> {1019 using OpConversionPattern<xegpu::StoreMatrixOp>::OpConversionPattern;1020 LogicalResult1021 matchAndRewrite(xegpu::StoreMatrixOp op, OneToNOpAdaptor adaptor,1022 ConversionPatternRewriter &rewriter) const override {1023 1024 SmallVector<SmallVector<OpFoldResult>> offsetsList;1025 if (failed(genOffsetsList(rewriter, op, offsetsList)))1026 return failure();1027 1028 xegpu::DistributeLayoutAttr layout = op.getLayoutAttr();1029 for (auto [v, offsets] : llvm::zip(adaptor.getData(), offsetsList))1030 xegpu::StoreMatrixOp::create(rewriter, op.getLoc(), v, op.getMemDesc(),1031 offsets, layout.dropSgLayoutAndData());1032 rewriter.eraseOp(op);1033 return success();1034 }1035};1036 1037// This pattern distributes the vector.step ops to work at subgroup level1038struct WgToSgVectorStepOp : public OpConversionPattern<vector::StepOp> {1039 using OpConversionPattern<vector::StepOp>::OpConversionPattern;1040 LogicalResult1041 matchAndRewrite(vector::StepOp op, OneToNOpAdaptor adaptor,1042 ConversionPatternRewriter &rewriter) const override {1043 xegpu::DistributeLayoutAttr layout =1044 xegpu::getDistributeLayoutAttr(op.getResult());1045 if (!layout || !layout.isForWorkgroup())1046 return failure();1047 1048 Location loc = op.getLoc();1049 VectorType type = op.getResult().getType();1050 auto wgShape = type.getShape();1051 std::optional<SmallVector<int64_t>> sgShape =1052 getSgShapeAndCount(wgShape, layout).first;1053 if (!sgShape)1054 return failure();1055 1056 Value sgId =1057 gpu::SubgroupIdOp::create(rewriter, loc, /*upper_bound=*/nullptr);1058 auto sgOffsets =1059 layout.computeDistributedCoords(rewriter, loc, sgId, wgShape);1060 if (failed(sgOffsets))1061 return failure();1062 1063 VectorType newTy = type.cloneWith(*sgShape, type.getElementType());1064 auto steps = vector::StepOp::create(rewriter, loc, newTy);1065 SmallVector<Value> newOps;1066 for (auto offsets : *sgOffsets) {1067 // Broadcast the offset scalar to a vector & add to the base steps1068 auto bcastOffset =1069 vector::BroadcastOp::create(rewriter, loc, newTy, offsets[0]);1070 auto finalSteps =1071 arith::AddIOp::create(rewriter, loc, steps, bcastOffset);1072 xegpu::setDistributeLayoutAttr(steps->getResult(0),1073 layout.dropSgLayoutAndData());1074 xegpu::setDistributeLayoutAttr(bcastOffset->getResult(0),1075 layout.dropSgLayoutAndData());1076 xegpu::setDistributeLayoutAttr(finalSteps->getResult(0),1077 layout.dropSgLayoutAndData());1078 newOps.push_back(finalSteps);1079 }1080 1081 rewriter.replaceOpWithMultiple(op, {newOps});1082 return success();1083 }1084};1085 1086// This pattern transforms vector.shape_cast ops to work at subgroup level.1087struct WgToSgVectorShapeCastOp1088 : public OpConversionPattern<vector::ShapeCastOp> {1089 using OpConversionPattern<vector::ShapeCastOp>::OpConversionPattern;1090 1091 LogicalResult1092 matchAndRewrite(vector::ShapeCastOp op, OneToNOpAdaptor adaptor,1093 ConversionPatternRewriter &rewriter) const override {1094 1095 VectorType resultType = dyn_cast<VectorType>(op.getResult().getType());1096 if (!resultType)1097 return failure();1098 1099 ArrayRef<int64_t> wgShape = resultType.getShape();1100 xegpu::DistributeLayoutAttr layout =1101 xegpu::getDistributeLayoutAttr(op.getResult());1102 if (!layout || !layout.isForWorkgroup())1103 return failure();1104 1105 SmallVector<int64_t> sgShape = getSgShapeAndCount(wgShape, layout).first;1106 VectorType newResultType =1107 VectorType::get(sgShape, resultType.getElementType());1108 1109 // TODO: Add check for compatible layouts in layout attr.1110 auto srcType = dyn_cast<VectorType>(adaptor.getSource()[0].getType());1111 if (!srcType)1112 return failure();1113 1114 // Check that shape_cast only adds/removes unit dimensions,1115 auto onlyUnitDims = [](ArrayRef<int64_t> src, ArrayRef<int64_t> dst) {1116 // Remove all 1s from both shapes and compare the rest.1117 SmallVector<int64_t> srcNonUnit, dstNonUnit;1118 for (int64_t d : src)1119 if (d != 1)1120 srcNonUnit.push_back(d);1121 for (int64_t d : dst)1122 if (d != 1)1123 dstNonUnit.push_back(d);1124 return srcNonUnit == dstNonUnit;1125 };1126 1127 if (!onlyUnitDims(srcType.getShape(), sgShape))1128 return failure();1129 1130 // For rank reducing or increasing shape_cast ops, the lower rank layout1131 // must be a slice of higher rank layout.1132 int64_t sourceRank = srcType.getRank();1133 int64_t resultRank = sgShape.size();1134 xegpu::DistributeLayoutAttr sourceLayout =1135 xegpu::getDistributeLayoutAttr(op.getSource());1136 if (sourceRank < resultRank && !sourceLayout.isSliceOf(layout))1137 return failure();1138 if (sourceRank > resultRank && !layout.isSliceOf(sourceLayout))1139 return failure();1140 1141 SmallVector<Value> newShapeCastOps;1142 for (auto src : adaptor.getSource()) {1143 auto newShapeCast = vector::ShapeCastOp::create(rewriter, op.getLoc(),1144 newResultType, src);1145 xegpu::setDistributeLayoutAttr(newShapeCast->getResult(0),1146 layout.dropSgLayoutAndData());1147 newShapeCastOps.push_back(newShapeCast.getResult());1148 }1149 1150 rewriter.replaceOpWithMultiple(op, {newShapeCastOps});1151 return success();1152 }1153};1154 1155/// Pattern for lowering vector.multi_reduction op to subgroup level.1156/// Current limitation: the sg_layout in the reduced dimension being 11157/// so that reduction is local to subgroup & no cross-subgroup communication is1158/// needed.1159/// TODO: Add cases to handle more general situations which require SLM access.1160struct WgToSgMultiDimReductionOp1161 : public OpConversionPattern<vector::MultiDimReductionOp> {1162 using OpConversionPattern<vector::MultiDimReductionOp>::OpConversionPattern;1163 1164 LogicalResult1165 matchAndRewrite(vector::MultiDimReductionOp op, OneToNOpAdaptor adaptor,1166 ConversionPatternRewriter &rewriter) const override {1167 VectorType srcType = op.getSourceVectorType();1168 VectorType dstType = dyn_cast<VectorType>(op.getResult().getType());1169 if (!dstType)1170 return failure();1171 1172 auto srcShape = srcType.getShape();1173 xegpu::DistributeLayoutAttr layout =1174 xegpu::getDistributeLayoutAttr(op.getResult());1175 if (!layout || !layout.isForWorkgroup())1176 return failure();1177 1178 auto reductionDims = llvm::to_vector(op.getReductionDims());1179 1180 SmallVector<int64_t> sgLayout = llvm::cast<xegpu::SliceAttr>(layout)1181 .getParent()1182 .getEffectiveSgLayoutAsInt();1183 SmallVector<int64_t> sgData = llvm::cast<xegpu::SliceAttr>(layout)1184 .getParent()1185 .getEffectiveSgDataAsInt();1186 1187 // Check that the sgLayout in the reduced dimension is 1 and1188 // each sg gets the entire slice to reduce.1189 for (int64_t dim : reductionDims) {1190 if (sgLayout[dim] != 1 || sgData[dim] != srcShape[dim])1191 return rewriter.notifyMatchFailure(1192 op,1193 "sgLayout in each reduced dimension must be 1 and sgData in the "1194 "reduced dim must match srcShape in that dim");1195 }1196 1197 SmallVector<int64_t> sgShape = getSgShapeAndCount(srcShape, layout).first;1198 1199 VectorType newDstType =1200 VectorType::get({sgShape}, dstType.getElementType());1201 1202 SmallVector<Value> newReductions;1203 for (auto sgSrc : adaptor.getSource()) {1204 auto newOp = vector::MultiDimReductionOp::create(1205 rewriter, op.getLoc(), newDstType, op.getKind(), sgSrc,1206 adaptor.getAcc()[0], op.getReductionDims());1207 xegpu::setDistributeLayoutAttr(newOp->getResult(0),1208 layout.dropSgLayoutAndData());1209 newReductions.push_back(newOp.getResult());1210 }1211 1212 rewriter.replaceOpWithMultiple(op, {newReductions});1213 return success();1214 }1215};1216 1217// This pattern transforms vector.transpose ops to work at subgroup level.1218struct WgToSgVectorTransposeOp1219 : public OpConversionPattern<vector::TransposeOp> {1220 using OpConversionPattern<vector::TransposeOp>::OpConversionPattern;1221 1222 LogicalResult1223 matchAndRewrite(vector::TransposeOp op, OneToNOpAdaptor adaptor,1224 ConversionPatternRewriter &rewriter) const override {1225 VectorType resultType = op.getResultVectorType();1226 1227 ArrayRef<int64_t> wgShape = resultType.getShape();1228 xegpu::DistributeLayoutAttr layout =1229 xegpu::getDistributeLayoutAttr(op.getResult());1230 if (!layout || !layout.isForWorkgroup())1231 return failure();1232 1233 xegpu::DistributeLayoutAttr sourceLayout =1234 xegpu::getDistributeLayoutAttr(op.getVector());1235 if (!sourceLayout || !sourceLayout.isForWorkgroup())1236 return failure();1237 1238 SmallVector<int64_t> sourceSgLayout =1239 sourceLayout.getEffectiveSgLayoutAsInt();1240 SmallVector<int64_t> resultSgLayout = layout.getEffectiveSgLayoutAsInt();1241 DenseI32ArrayAttr sourceOrder = sourceLayout.getOrder();1242 DenseI32ArrayAttr resultOrder = layout.getOrder();1243 1244 if (!sourceOrder || !resultOrder) {1245 return rewriter.notifyMatchFailure(1246 op, "Both source and result must have order attributes");1247 }1248 1249 ArrayRef<int64_t> permutation = op.getPermutation();1250 size_t permutationSize = permutation.size();1251 if (sourceSgLayout.size() != permutationSize ||1252 resultSgLayout.size() != permutationSize) {1253 return rewriter.notifyMatchFailure(1254 op, "Layouts and permutation must have the same rank");1255 }1256 1257 // Check that sgLayout, sgData & order are properly transposed for source1258 // and result1259 if (!layout.isTransposeOf(sourceLayout, permutation))1260 return rewriter.notifyMatchFailure(1261 op, "Result layout is not a valid transpose of source layout "1262 "according to permutation");1263 1264 SmallVector<int64_t> sgShape = getSgShapeAndCount(wgShape, layout).first;1265 VectorType newResultType =1266 VectorType::get(sgShape, resultType.getElementType());1267 SmallVector<Value> newTransposeOps;1268 for (auto src : adaptor.getVector()) {1269 auto newTranspose = vector::TransposeOp::create(1270 rewriter, op.getLoc(), newResultType, src, permutation);1271 xegpu::setDistributeLayoutAttr(newTranspose->getResult(0),1272 layout.dropSgLayoutAndData());1273 newTransposeOps.push_back(newTranspose.getResult());1274 }1275 1276 rewriter.replaceOpWithMultiple(op, {newTransposeOps});1277 return success();1278 }1279};1280 1281// This pattern distributes the vector.constant_mask ops to work at subgroup1282// level.1283struct WgToSgVectorConstantMaskOp1284 : public OpConversionPattern<vector::ConstantMaskOp> {1285 using OpConversionPattern<vector::ConstantMaskOp>::OpConversionPattern;1286 1287 LogicalResult1288 matchAndRewrite(vector::ConstantMaskOp op, OneToNOpAdaptor adaptor,1289 ConversionPatternRewriter &rewriter) const override {1290 xegpu::DistributeLayoutAttr layout =1291 xegpu::getDistributeLayoutAttr(op.getResult());1292 if (!layout || !layout.isForWorkgroup())1293 return failure();1294 1295 Location loc = op.getLoc();1296 VectorType type = op.getResult().getType();1297 auto wgShape = type.getShape();1298 1299 ArrayRef<int64_t> wgMaskDimSizes = op.getMaskDimSizes();1300 1301 // Get subgroup ID.1302 Value sgId =1303 gpu::SubgroupIdOp::create(rewriter, loc, /*upper_bound=*/nullptr);1304 auto sgOffsets =1305 layout.computeDistributedCoords(rewriter, loc, sgId, wgShape);1306 if (failed(sgOffsets))1307 return failure();1308 1309 SmallVector<int64_t> sgShape = getSgShapeAndCount(wgShape, layout).first;1310 VectorType resultType = VectorType::get(sgShape, type.getElementType());1311 1312 // In each dimension, each subgroup computes its local mask size as:1313 // min(max(wgMaskSize[d] - offset[d], 0), sgDimSize[d])1314 SmallVector<Value> newCreateMaskOps;1315 for (auto offsetSet : *sgOffsets) {1316 SmallVector<Value> maskOperands;1317 1318 for (auto [i, wgMaskSize] : llvm::enumerate(wgMaskDimSizes)) {1319 Value wgMaskSizeVal =1320 arith::ConstantIndexOp::create(rewriter, loc, wgMaskSize);1321 Value dimSizeVal =1322 arith::ConstantIndexOp::create(rewriter, loc, sgShape[i]);1323 Value offset = offsetSet[i];1324 Value adjustedMaskSize =1325 arith::SubIOp::create(rewriter, loc, wgMaskSizeVal, offset);1326 Value zero = arith::ConstantIndexOp::create(rewriter, loc, 0);1327 Value nonNegative =1328 arith::MaxSIOp::create(rewriter, loc, adjustedMaskSize, zero);1329 Value sgMaskSize =1330 arith::MinSIOp::create(rewriter, loc, nonNegative, dimSizeVal);1331 maskOperands.push_back(sgMaskSize);1332 }1333 1334 auto newCreateMaskOp =1335 vector::CreateMaskOp::create(rewriter, loc, resultType, maskOperands);1336 xegpu::setDistributeLayoutAttr(newCreateMaskOp->getResult(0),1337 layout.dropSgLayoutAndData());1338 newCreateMaskOps.push_back(newCreateMaskOp.getResult());1339 }1340 1341 rewriter.replaceOpWithMultiple(op, {newCreateMaskOps});1342 return success();1343 }1344};1345 1346} // namespace1347 1348namespace mlir {1349namespace xegpu {1350void populateXeGPUWgToSgDistributePatterns(RewritePatternSet &patterns) {1351 patterns1352 .add<WgToSgCreateNdOp, WgToSgCreateNdOpNoOffset, WgToSgLoadNdOp,1353 WgToSgLoadNdOpWithOffset, WgToSgStoreNdOp, WgToSgStoreNdOpWithOffset,1354 WgToSgUpdateNdOffsetOp, WgToSgDpasOp, WgToSgPrefetchNdOp,1355 WgToSgPrefetchNdOpWithOffset, UnrealizedConversionCastOpPattern,1356 WgToSgElementwiseOp, WgToSgVectorBroadcastOp, WgToSgConvertLayoutOp,1357 WgToSgArithConstantOp, WgToSgLoadGatherOpWithOffset,1358 WgToSgStoreScatterOpWithOffset, WgToSgLoadMatrixOp,1359 WgToSgStoreMatrixOp, WgToSgVectorStepOp, WgToSgVectorShapeCastOp,1360 WgToSgMultiDimReductionOp, WgToSgVectorTransposeOp,1361 WgToSgVectorConstantMaskOp>(patterns.getContext());1362}1363} // namespace xegpu1364} // namespace mlir1365 1366namespace {1367struct XeGPUWgToSgDistributePass1368 : public xegpu::impl::XeGPUWgToSgDistributeBase<XeGPUWgToSgDistributePass> {1369 void runOnOperation() override;1370};1371} // namespace1372 1373void XeGPUWgToSgDistributePass::runOnOperation() {1374 // Track existing UnrealizedConversionCastOps1375 SmallVector<Operation *> existingCastOps;1376 getOperation()->walk([&](UnrealizedConversionCastOp castOp) {1377 existingCastOps.push_back(castOp.getOperation());1378 });1379 1380 {1381 // Step 1: Apply SCFStructuralTypeConversions to SCF operations with1382 // VectorType operands. This first converts such operands to1383 // RankedTensorType, propagates the layout attribute into the encoding1384 // attribute, and finally converts the RankedTensorType to VectorType based1385 // on the encoding.1386 1387 TypeConverter converter;1388 converter.addConversion([&](Type type) -> Type { return type; });1389 converter.addConversion(1390 [&](RankedTensorType type,1391 SmallVectorImpl<Type> &result) -> std::optional<LogicalResult> {1392 Type elemTy = type.getElementType();1393 ArrayRef<int64_t> shape = type.getShape();1394 1395 int count;1396 SmallVector<int64_t> subShape;1397 std::tie(subShape, count) = getSgShapeAndCount(1398 shape,1399 dyn_cast_if_present<xegpu::LayoutAttr>(type.getEncoding()));1400 1401 auto newTy = VectorType::get(subShape, elemTy);1402 result.append(count, newTy);1403 return success();1404 });1405 1406 xegpu::doSCFStructuralTypeConversionWithTensorType(getOperation(),1407 converter);1408 }1409 1410 // Step 2: Perform workgroup to subgroup distribution for TensorDesc values,1411 // as well as XeGPU, Arith, and Vector operations.1412 MLIRContext *ctx = &getContext();1413 RewritePatternSet patterns(ctx);1414 ConversionTarget target(*ctx);1415 TypeConverter converter;1416 converter.addConversion([&](Type type) -> Type { return type; });1417 converter.addConversion(1418 [&](xegpu::TensorDescType type,1419 SmallVectorImpl<Type> &result) -> std::optional<LogicalResult> {1420 Type elemTy = type.getElementType();1421 ArrayRef<int64_t> shape = type.getShape();1422 1423 int count;1424 SmallVector<int64_t> subShape;1425 xegpu::LayoutAttr layout = type.getLayoutAttr();1426 std::tie(subShape, count) = getSgShapeAndCount(shape, layout);1427 1428 if (layout)1429 layout = layout.dropSgLayoutAndData();1430 1431 auto newTy = xegpu::TensorDescType::get(1432 type.getContext(), subShape, elemTy, type.getEncoding(), layout);1433 result.append(count, newTy);1434 return success();1435 });1436 1437 auto getTensorDescType = [](Operation *op) -> xegpu::TensorDescType {1438 if (auto createOp = dyn_cast<xegpu::CreateNdDescOp>(op))1439 return createOp.getType();1440 if (auto loadOp = dyn_cast<xegpu::LoadNdOp>(op))1441 return loadOp.getTensorDescType();1442 if (auto storeOp = dyn_cast<xegpu::StoreNdOp>(op))1443 return storeOp.getTensorDescType();1444 if (auto updateOp = dyn_cast<xegpu::UpdateNdOffsetOp>(op))1445 return updateOp.getType();1446 if (auto prefetchOp = dyn_cast<xegpu::PrefetchNdOp>(op))1447 return prefetchOp.getTensorDescType();1448 return xegpu::TensorDescType();1449 };1450 1451 auto isLegal = [&](xegpu::DistributeLayoutAttr layout) -> bool {1452 return !layout || !layout.isForWorkgroup();1453 };1454 1455 target.addDynamicallyLegalOp<xegpu::CreateNdDescOp, xegpu::LoadNdOp,1456 xegpu::StoreNdOp, xegpu::UpdateNdOffsetOp,1457 xegpu::PrefetchNdOp>([=](Operation *op) -> bool {1458 auto tdescTy = getTensorDescType(op);1459 auto layout = dyn_cast_if_present<xegpu::LayoutAttr>(tdescTy.getLayout());1460 return isLegal(layout);1461 });1462 1463 target.addDynamicallyLegalOp<xegpu::DpasOp>([=](xegpu::DpasOp op) -> bool {1464 auto layout = xegpu::getDistributeLayoutAttr(op.getResult());1465 return isLegal(layout);1466 });1467 1468 target.addDynamicallyLegalOp<xegpu::LoadMatrixOp>(1469 [=](xegpu::LoadMatrixOp op) -> bool {1470 return isLegal(op.getLayoutAttr());1471 });1472 1473 target.addDynamicallyLegalOp<xegpu::StoreMatrixOp>(1474 [=](xegpu::StoreMatrixOp op) -> bool {1475 return isLegal(op.getLayoutAttr());1476 });1477 1478 target.addDynamicallyLegalOp<arith::ConstantOp>(1479 [=](arith::ConstantOp op) -> bool {1480 auto vecType = dyn_cast<VectorType>(op.getType());1481 if (!vecType)1482 return true;1483 1484 auto layout = xegpu::getDistributeLayoutAttr(op.getResult());1485 return isLegal(layout);1486 });1487 1488 target.addDynamicallyLegalOp<1489 vector::ShapeCastOp, vector::StepOp, vector::TransposeOp,1490 vector::BroadcastOp, vector::MultiDimReductionOp, vector::ConstantMaskOp>(1491 [=](Operation *op) -> bool {1492 // Check for either a SliceAttr or LayoutAttr on the result.1493 auto layout = xegpu::getDistributeLayoutAttr(op->getResult(0));1494 return isLegal(layout);1495 });1496 1497 target.addDynamicallyLegalOp<xegpu::LoadGatherOp>(1498 [=](xegpu::LoadGatherOp op) -> bool {1499 auto layout = xegpu::getDistributeLayoutAttr(op.getResult());1500 return isLegal(layout);1501 });1502 1503 target.addDynamicallyLegalOp<xegpu::StoreScatterOp>(1504 [=](xegpu::StoreScatterOp op) -> bool {1505 auto layout = xegpu::getDistributeLayoutAttr(op.getOperand(0));1506 return isLegal(layout);1507 });1508 1509 target.addDynamicallyLegalOp<xegpu::ConvertLayoutOp>(1510 [=](xegpu::ConvertLayoutOp op) -> bool {1511 return isLegal(op.getInputLayout()) && isLegal(op.getTargetLayout());1512 });1513 1514 target.addDynamicallyLegalDialect<math::MathDialect, arith::ArithDialect>(1515 [=](Operation *op) -> std::optional<bool> {1516 // Only handle elementwise mappable ops1517 if (!OpTrait::hasElementwiseMappableTraits(op))1518 return true;1519 1520 VectorType resultType =1521 dyn_cast<VectorType>(op->getResult(0).getType());1522 if (!resultType)1523 return true;1524 1525 // Check if all operands are vectors of the same shape1526 // TODO: Support other types.1527 for (Value operand : op->getOperands()) {1528 VectorType operandType = dyn_cast<VectorType>(operand.getType());1529 if (!operandType || operandType.getShape() != resultType.getShape()) {1530 return true;1531 }1532 }1533 1534 xegpu::DistributeLayoutAttr layout =1535 xegpu::getDistributeLayoutAttr(op->getResult(0));1536 return isLegal(layout);1537 });1538 1539 target.addDynamicallyLegalOp<UnrealizedConversionCastOp>(1540 [=](UnrealizedConversionCastOp op) {1541 return llvm::is_contained(existingCastOps, op.getOperation());1542 });1543 1544 target.markUnknownOpDynamicallyLegal([](Operation *) { return true; });1545 1546 scf::populateSCFStructuralTypeConversionsAndLegality(converter, patterns,1547 target);1548 xegpu::populateXeGPUWgToSgDistributePatterns(patterns);1549 if (failed(1550 applyPartialConversion(getOperation(), target, std::move(patterns))))1551 return signalPassFailure();1552 1553 // Remove sg_layout and sg_data attributes from the Layout1554 // attribute for each VectorType result of the operation.1555 // For Structured Control Flow ops, the layout is simply removed,1556 // since in 1:N case, the layout for new results are missing.1557 // Layout propagation pass will activated.1558 getOperation()->walk([](Operation *op) {1559 for (OpResult result : op->getOpResults()) {1560 std::string name = xegpu::getLayoutName(result);1561 if (auto layout = op->getAttrOfType<xegpu::LayoutAttr>(name)) {1562 op->removeAttr(name);1563 if (!isa<scf::IfOp, scf::ForOp, scf::WhileOp, scf::ConditionOp>(op)) {1564 if (auto newLayout = layout.dropSgLayoutAndData())1565 op->setAttr(name, newLayout);1566 }1567 }1568 }1569 });1570}1571