583 lines · cpp
1//===---- XeGPUUtils.cpp - MLIR Utilities for XeGPUOps ------------------===//2//3// Part of the MLIR 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//9// This file implements utility methods for working with the XeGPU dialect.10//11//===----------------------------------------------------------------------===//12 13#include "mlir/Dialect/XeGPU/Utils/XeGPUUtils.h"14#include "mlir/Dialect/GPU/IR/GPUDialect.h"15#include "mlir/Dialect/Index/IR/IndexOps.h"16#include "mlir/Dialect/LLVMIR/XeVMDialect.h"17#include "mlir/Dialect/SCF/Transforms/Patterns.h"18#include "mlir/Dialect/Utils/IndexingUtils.h"19#include "mlir/Dialect/XeGPU/IR/XeGPU.h"20#include "mlir/IR/Builders.h"21#include "mlir/IR/Operation.h"22#include "mlir/IR/ValueRange.h"23#include "mlir/Interfaces/LoopLikeInterface.h"24#include "mlir/Transforms/DialectConversion.h"25#include "llvm/Support/FormatVariadic.h"26#include <cstdint>27#include <numeric>28 29using namespace mlir;30 31/// convert ArrayRef<ValueRange> into SmallVector<Value>32SmallVector<Value> xegpu::flattenValues(ArrayRef<ValueRange> values) {33 SmallVector<Value> result;34 for (const auto &vals : values)35 llvm::append_range(result, vals);36 return result;37}38 39FailureOr<VectorType>40mlir::xegpu::getDistributedVectorType(xegpu::TensorDescType tdescTy) {41 auto layout = llvm::dyn_cast_if_present<LayoutAttr>(tdescTy.getLayout());42 // It only works for subgroup level layout, which only has lane_layout43 // and lane_data, and is to distribute a SIMD code into SIMT code.44 if (!layout || !layout.isForSubgroup())45 return failure();46 47 SmallVector<int64_t> laneData(layout.getLaneData().asArrayRef());48 SmallVector<int64_t> laneLayout(layout.getLaneLayout().asArrayRef());49 auto tdescShape = tdescTy.getShape();50 auto elementType = tdescTy.getElementType();51 52 // compute sgSize by multiply elements of laneLayout53 // e.g. for 2D layout, sgSize = laneLayout[0] * laneLayout[1]54 // e.g. for 1D layout, sgSize = laneLayout[0]55 int64_t sgSize = llvm::product_of(laneLayout);56 57 // Case 1: regular loads/stores58 auto scatterAttr = tdescTy.getEncodingOfType<ScatterTensorDescAttr>();59 if (scatterAttr) {60 auto chunkSize = scatterAttr.getChunkSize().getInt();61 // Verify if the first dimension of the tensor descriptor shape is62 // distributable.63 assert(tdescShape[0] == laneLayout[0] &&64 "tensor descriptor shape is not distributable");65 return VectorType::get({chunkSize}, elementType);66 }67 68 // Case 2: block loads/stores69 // Check if the tensor descriptor shape is distributable.70 int64_t tensorSize = 1;71 for (auto [tdescDim, laneDim, laneDataDim] :72 llvm::zip_equal(tdescShape, laneLayout, laneData)) {73 assert((tdescDim % (laneDim * laneDataDim) == 0) &&74 "tensor descriptor shape is not distributable");75 tensorSize *= tdescDim;76 }77 // tensorSize must be adjusted for array_length.78 tensorSize *= tdescTy.getArrayLength();79 80 return VectorType::get({tensorSize / sgSize}, elementType);81}82 83FailureOr<VectorType>84mlir::xegpu::getDistributedVectorType(VectorType originalType,85 xegpu::LayoutAttr layout) {86 int64_t rank = originalType.getRank();87 // Distributed vector type is only supported for 1D, 2D and 3D vectors.88 if (rank < 1 || rank > 3)89 return failure();90 ArrayRef<int64_t> shape = originalType.getShape();91 // arrayLength is 1 for 1D and 2D vectors, and equal to the first dimension92 // of the 3D vector.93 int arrayLength = 1;94 if (rank == 3) {95 arrayLength = shape[0];96 shape = shape.drop_front();97 }98 auto helperTdescTy = xegpu::TensorDescType::get(99 shape, originalType.getElementType(), arrayLength,100 /*boundary_check=*/true,101 /*memory_space=*/xegpu::MemorySpace::Global, layout);102 return xegpu::getDistributedVectorType(helperTdescTy);103}104 105std::string xegpu::getLayoutName(const OpOperand &operand) {106 const StringRef prefix("layout_operand_");107 unsigned idx = const_cast<OpOperand &>(operand).getOperandNumber();108 return llvm::formatv("{0}{1}", prefix, idx).str();109}110 111std::string xegpu::getLayoutName(const OpResult result) {112 const StringRef prefix = "layout_result_";113 return llvm::formatv("{0}{1}", prefix, result.getResultNumber()).str();114}115 116xegpu::DistributeLayoutAttr xegpu::getDistributeLayoutAttr(const Value value) {117 if (!value)118 return nullptr;119 120 if (auto tdescTy =121 dyn_cast_if_present<xegpu::TensorDescType>(value.getType()))122 return tdescTy.getLayoutAttr();123 124 if (auto result = dyn_cast<OpResult>(value)) {125 Operation *defOp = result.getDefiningOp();126 assert(defOp && "result must have a defining op");127 128 // For ConvertLayoutOp, the layout is stored in the targetLayoutAttr129 if (auto convertOp = dyn_cast<xegpu::ConvertLayoutOp>(defOp))130 return convertOp.getTargetLayoutAttr();131 132 // for LoadNdOp, the layout is stored in the tensor descriptor133 if (auto loadNd = dyn_cast<xegpu::LoadNdOp>(defOp))134 return getDistributeLayoutAttr(loadNd.getTensorDesc());135 136 // for LoadMatrixOp, the layout is attached to the property of the op137 if (auto loadOp = dyn_cast<xegpu::LoadMatrixOp>(defOp))138 return loadOp.getLayoutAttr();139 140 // for StoreMatrixOp, the layout is attached to the property of the op141 if (auto storeOp = dyn_cast<xegpu::StoreMatrixOp>(defOp))142 return storeOp.getLayoutAttr();143 std::string layoutName = getLayoutName(result);144 if (defOp->hasAttr(layoutName))145 return defOp->getAttrOfType<xegpu::DistributeLayoutAttr>(layoutName);146 147 // check for "permament" layout only after "temporary" layout name lookup148 // for backward compatibility149 if (auto loadGatherOp = dyn_cast<xegpu::LoadGatherOp>(defOp))150 return loadGatherOp.getLayoutAttr();151 }152 153 if (auto arg = dyn_cast<BlockArgument>(value)) {154 auto *parentOp = arg.getOwner()->getParentOp();155 if (auto loop = dyn_cast<LoopLikeOpInterface>(parentOp)) {156 OpOperand *tiedInit = loop.getTiedLoopInit(arg);157 if (tiedInit)158 return getDistributeLayoutAttr(tiedInit->get());159 }160 }161 162 return nullptr;163}164 165xegpu::DistributeLayoutAttr166xegpu::getDistributeLayoutAttr(const OpOperand &opr) {167 Operation *op = opr.getOwner();168 169 if (auto loadOp = dyn_cast<xegpu::LoadMatrixOp>(op))170 return loadOp.getLayoutAttr();171 172 if (auto storeOp = dyn_cast<xegpu::StoreMatrixOp>(op))173 return storeOp.getLayoutAttr();174 175 std::string layoutName = xegpu::getLayoutName(opr);176 if (op->hasAttr(layoutName))177 return op->getAttrOfType<xegpu::DistributeLayoutAttr>(layoutName);178 179 // check for "permament" layout only after "temporary" layout name lookup180 if (auto storeScatterOp = dyn_cast<xegpu::StoreScatterOp>(op))181 if (auto layout = storeScatterOp.getLayoutAttr())182 return layout;183 184 return getDistributeLayoutAttr(opr.get());185}186 187// Returns the permanent layout attribute for the given result if it's188// available on the defining op. Otherwise returns the provided layout.189xegpu::DistributeLayoutAttr190maybePickPermanentLayout(xegpu::DistributeLayoutAttr layout,191 const OpResult &result, mlir::Operation *owner,192 const std::string &name) {193 xegpu::DistributeLayoutAttr candidate = layout;194 195 if (auto loadOp = dyn_cast<xegpu::LoadGatherOp>(owner)) {196 if (auto perm = loadOp.getLayoutAttr())197 candidate = perm;198 }199 200 return candidate;201}202 203// Returns the permanent layout attribute for the given operand if it's204// available on the defining op. Otherwise returns the provided layout.205xegpu::DistributeLayoutAttr206maybePickPermanentLayout(xegpu::DistributeLayoutAttr layout,207 const OpOperand &operand, mlir::Operation *owner,208 const std::string &name) {209 xegpu::DistributeLayoutAttr candidate = layout;210 unsigned idx = const_cast<OpOperand &>(operand).getOperandNumber();211 212 if (auto storeOp = dyn_cast<xegpu::StoreScatterOp>(owner)) {213 if (idx == 0) {214 if (auto perm = storeOp.getLayoutAttr())215 candidate = perm;216 }217 }218 219 return candidate;220}221 222template <typename T, typename>223void xegpu::setDistributeLayoutAttr(const T &operandOrResult,224 const DistributeLayoutAttr layout,225 bool respectPermLayout) {226 Operation *owner = operandOrResult.getOwner();227 std::string name = xegpu::getLayoutName(operandOrResult);228 229 if (owner->hasAttrOfType<DistributeLayoutAttr>(name))230 return;231 232 DistributeLayoutAttr candidate = layout;233 if (respectPermLayout)234 candidate = maybePickPermanentLayout(layout, operandOrResult, owner, name);235 236 if (candidate)237 owner->setAttr(name, candidate);238}239 240// Explicit instantiation for OpResult241template void xegpu::setDistributeLayoutAttr<mlir::OpResult>(242 const mlir::OpResult &result,243 const mlir::xegpu::DistributeLayoutAttr layout, bool respectPermLayout);244 245// Explicit instantiation for OpOperand246template void xegpu::setDistributeLayoutAttr<mlir::OpOperand>(247 const mlir::OpOperand &operand,248 const mlir::xegpu::DistributeLayoutAttr layout, bool respectPermLayout);249 250void xegpu::setDistributeLayoutAttrs(251 Operation *op, function_ref<DistributeLayoutAttr(Value)> getLayoutImpl) {252 op->walk([&](Operation *nestOp) {253 if (isa<xegpu::LoadMatrixOp, xegpu::StoreMatrixOp>(nestOp))254 return;255 256 for (OpOperand &opr : nestOp->getOpOperands()) {257 auto layout = getLayoutImpl(opr.get());258 setDistributeLayoutAttr(opr, layout);259 }260 for (OpResult result : nestOp->getOpResults()) {261 auto layout = getLayoutImpl(result);262 setDistributeLayoutAttr(result, layout);263 }264 });265}266 267template <typename T, typename>268void xegpu::removeLayoutAttr(const T &operandOrResult) {269 Operation *owner = operandOrResult.getOwner();270 std::string name = xegpu::getLayoutName(operandOrResult);271 if (owner->hasAttrOfType<DistributeLayoutAttr>(name))272 owner->removeAttr(name);273}274 275// Explicit instantiation for OpResult276template void277xegpu::removeLayoutAttr<mlir::OpResult>(const mlir::OpResult &result);278 279// Explicit instantiation for OpOperand280template void281xegpu::removeLayoutAttr<mlir::OpOperand>(const mlir::OpOperand &operand);282 283void xegpu::removeLayoutAttrs(Operation *op) {284 op->walk([&](Operation *nestOp) {285 for (OpOperand &opr : nestOp->getOpOperands())286 removeLayoutAttr(opr);287 for (OpResult result : nestOp->getOpResults())288 removeLayoutAttr(result);289 });290}291 292SmallVector<Value>293xegpu::extractVectorsWithShapeFromValue(OpBuilder &builder, Location loc,294 Value value, ArrayRef<int64_t> shape) {295 auto vecTy = dyn_cast<VectorType>(value.getType());296 if (!vecTy)297 return {value};298 299 ArrayRef<int64_t> srcShape = vecTy.getShape();300 if (!computeShapeRatio(srcShape, shape))301 return {value};302 303 int64_t srcShapeRank = srcShape.size();304 int64_t targetShapeRank = shape.size();305 306 SmallVector<int64_t> adjustedTargetShape(srcShape.size());307 int64_t rankDiff = srcShapeRank - targetShapeRank;308 std::fill(adjustedTargetShape.begin(), adjustedTargetShape.begin() + rankDiff,309 1);310 llvm::copy(shape, adjustedTargetShape.begin() + rankDiff);311 312 SmallVector<Value> result;313 for (SmallVector<int64_t> offsets :314 StaticTileOffsetRange(srcShape, adjustedTargetShape)) {315 SmallVector<int64_t> staticStrides(offsets.size(), 1);316 Value slice = vector::ExtractStridedSliceOp::create(317 builder, loc, value, offsets, adjustedTargetShape, staticStrides);318 319 // Reshape to remove leading unit dims if needed320 if (srcShapeRank > targetShapeRank) {321 auto targetTy = VectorType::get(shape, vecTy.getElementType());322 slice = vector::ShapeCastOp::create(builder, loc, targetTy, slice);323 }324 result.push_back(slice);325 }326 327 return result;328}329 330Value xegpu::createVectorWithShapeFromValues(OpBuilder &builder, Location loc,331 ValueRange values,332 ArrayRef<int64_t> shape) {333 VectorType inputTy = dyn_cast<VectorType>(values[0].getType());334 assert(llvm::all_of(values.getTypes(),335 [&](Type type) { return type == inputTy; }) &&336 "values must be of the same VectorType");337 338 Type elemTy = inputTy.getElementType();339 ArrayRef<int64_t> tileShape = inputTy.getShape();340 341 VectorType resultTy = VectorType::get(shape, elemTy);342 auto zeroAttr = builder.getZeroAttr(elemTy);343 Value result = arith::ConstantOp::create(344 builder, loc, resultTy, DenseElementsAttr::get(resultTy, zeroAttr));345 346 for (auto [src, offsets] :347 llvm::zip_equal(values, StaticTileOffsetRange(shape, tileShape))) {348 SmallVector<int64_t> staticStrides(tileShape.size(), 1);349 result = vector::InsertStridedSliceOp::create(builder, loc, src, result,350 offsets, staticStrides);351 }352 return result;353}354 355void xegpu::doSCFStructuralTypeConversionWithTensorType(356 Operation *op, TypeConverter converter) {357 MLIRContext *context = op->getContext();358 359 auto materializeCast = [](OpBuilder &builder, Type type, ValueRange inputs,360 Location loc) -> Value {361 return UnrealizedConversionCastOp::create(builder, loc, type, inputs)362 .getResult(0);363 };364 365 { // convert VectorType to RankedTensorType for SCF Structural ops366 TypeConverter converter;367 converter.addConversion([](Type type) -> Type { return type; });368 converter.addConversion([](VectorType type) -> Type {369 return RankedTensorType::get(type.getShape(), type.getElementType());370 });371 converter.addSourceMaterialization(materializeCast);372 converter.addTargetMaterialization(materializeCast);373 374 mlir::ConversionTarget target(*context);375 target.addLegalOp<UnrealizedConversionCastOp>();376 377 mlir::RewritePatternSet patterns(context);378 scf::populateSCFStructuralTypeConversionsAndLegality(converter, patterns,379 target);380 (void)mlir::applyPartialConversion(op, target, std::move(patterns));381 }382 383 { // propagate the layout attribute to RankedTensorType by checking384 // BuiltInUnrealizedCastOps385 // for VectorType to RankedTensorType cast.386 op->walk([](UnrealizedConversionCastOp castOp) {387 if (castOp.getNumOperands() != 1 || castOp.getNumResults() != 1)388 return WalkResult::skip();389 390 Value input = castOp.getInputs()[0];391 Value result = castOp.getResults()[0];392 auto inputTy = dyn_cast<VectorType>(input.getType());393 auto resultTy = dyn_cast<RankedTensorType>(result.getType());394 395 // Only look at ops casting from VectorType to RankedTensorType396 if (!inputTy || !resultTy)397 return WalkResult::skip();398 399 xegpu::DistributeLayoutAttr layout =400 xegpu::getDistributeLayoutAttr(input);401 if (!layout)402 return WalkResult::skip();403 404 RankedTensorType newTy = resultTy.cloneWithEncoding(layout);405 result.setType(newTy);406 407 // update the arguments if user is a LoopLike op.408 for (OpOperand &use : result.getUses()) {409 if (auto loop = dyn_cast<LoopLikeOpInterface>(use.getOwner())) {410 BlockArgument arg = loop.getTiedLoopRegionIterArg(&use);411 arg.setType(newTy);412 }413 // whileOp has two regions, the BlockArgument of the after region414 // is not exposed by LoopLikeOpInterface415 if (auto whileOp = dyn_cast<scf::WhileOp>(use.getOwner())) {416 unsigned idx = use.getOperandNumber();417 BlockArgument arg = whileOp.getAfterArguments()[idx];418 arg.setType(newTy);419 }420 }421 return WalkResult::advance();422 });423 424 // using yieldOp as anchor to update the result type of its ParentOp425 op->walk([](scf::YieldOp yieldOp) {426 Operation *parentOp = yieldOp->getParentOp();427 for (OpResult r : parentOp->getOpResults()) {428 unsigned idx = r.getResultNumber();429 Type resultTy = r.getType();430 Type yieldTy = yieldOp.getResults()[idx].getType();431 if (isa<RankedTensorType>(resultTy) && yieldTy != resultTy)432 r.setType(yieldTy);433 }434 });435 }436 437 { // perform the conversion from RankedTensorType to VectorType based on the438 // DistributeLayoutAttr439 440 // Handle the UnrealizedConversionCastOp introduced by the first step.441 // For vector->RankedTensorType, it will simply forward the inputs.442 // For RankedTensorType->vector, it will update the inputs with the443 // one from the adaptor.444 class UnrealizedConversionCastOpPattern445 : public OpConversionPattern<mlir::UnrealizedConversionCastOp> {446 using OpConversionPattern<447 mlir::UnrealizedConversionCastOp>::OpConversionPattern;448 449 mlir::LogicalResult450 matchAndRewrite(mlir::UnrealizedConversionCastOp op,451 OneToNOpAdaptor adaptor,452 ConversionPatternRewriter &rewriter) const override {453 auto inputs = op.getOperands();454 auto outputs = op.getOutputs();455 456 if (inputs.size() != 1 || outputs.size() != 1)457 return failure();458 459 auto inputTy = inputs[0].getType();460 auto outputTy = outputs[0].getType();461 462 if (isa<VectorType>(inputTy) && isa<RankedTensorType>(outputTy)) {463 rewriter.replaceOpWithMultiple(op, adaptor.getInputs());464 return success();465 }466 467 if (isa<RankedTensorType>(inputTy) && isa<VectorType>(outputTy)) {468 SmallVector<Value> values = xegpu::flattenValues(adaptor.getInputs());469 auto newOp = UnrealizedConversionCastOp::create(rewriter, op.getLoc(),470 outputTy, values);471 rewriter.replaceOp(op, newOp);472 return success();473 }474 return failure();475 }476 };477 478 converter.addSourceMaterialization(materializeCast);479 converter.addTargetMaterialization([&](OpBuilder &builder, TypeRange type,480 ValueRange inputs, Location loc) {481 return UnrealizedConversionCastOp::create(builder, loc, type, inputs)482 .getResults();483 });484 485 mlir::ConversionTarget target(*context);486 target.addDynamicallyLegalOp<UnrealizedConversionCastOp>(487 [](UnrealizedConversionCastOp op) {488 auto isTensorTy = [](Type type) {489 return isa<RankedTensorType>(type);490 };491 return llvm::none_of(op->getOperandTypes(), isTensorTy) &&492 llvm::none_of(op->getResultTypes(), isTensorTy);493 });494 mlir::RewritePatternSet patterns(context);495 patterns.insert<UnrealizedConversionCastOpPattern>(context);496 scf::populateSCFStructuralTypeConversionsAndLegality(converter, patterns,497 target);498 (void)mlir::applyPartialConversion(op, target, std::move(patterns));499 }500}501 502std::optional<std::string> xegpu::getChipStr(Operation *op) {503 auto gpuModuleOp = op->getParentOfType<gpu::GPUModuleOp>();504 505 if (!gpuModuleOp)506 return std::nullopt;507 508 auto targetAttrs = gpuModuleOp.getTargets();509 if (targetAttrs) {510 for (auto &attr : *targetAttrs) {511 auto xevmAttr = llvm::dyn_cast<xevm::XeVMTargetAttr>(attr);512 if (xevmAttr)513 return xevmAttr.getChip().str();514 }515 }516 517 return std::nullopt;518}519 520/// Generates element-wise addition ops of two arrays with same length.521SmallVector<OpFoldResult> xegpu::addElementwise(OpBuilder &builder,522 Location loc,523 ArrayRef<OpFoldResult> lhs,524 ArrayRef<OpFoldResult> rhs) {525 assert(lhs.size() == rhs.size() && "lhs and rhs must have the same size");526 SmallVector<OpFoldResult> results;527 for (auto [l, r] : llvm::zip_equal(lhs, rhs)) {528 auto lval = getValueOrCreateConstantIndexOp(builder, loc, l);529 auto rval = getValueOrCreateConstantIndexOp(builder, loc, r);530 results.push_back(builder.createOrFold<index::AddOp>(loc, lval, rval));531 }532 return results;533}534 535/// Generates element-wise addition ops of two arrays with automatic alignment.536/// When the input arrays have different sizes, the shorter array is537/// right-aligned with the longer array, and the unmatched leading elements from538/// the longer array are preserved unchanged. This is commonly used for offset539/// computation where higher-dimensional offsets need to be added to540/// lower-dimensional adjustments.541///542/// Example:543/// lhs = [l1, l2, l3], rhs = [r1, r2]544/// Result: [11, l2+r1, l3+r2]545SmallVector<OpFoldResult>546xegpu::addWithRightAligned(OpBuilder &builder, Location loc,547 ArrayRef<OpFoldResult> lhs,548 ArrayRef<OpFoldResult> rhs) {549 // ensure a is longer than b550 ArrayRef<OpFoldResult> a = lhs.size() >= rhs.size() ? lhs : rhs;551 ArrayRef<OpFoldResult> b = lhs.size() >= rhs.size() ? rhs : lhs;552 SmallVector<OpFoldResult> results(a.take_front(a.size() - b.size()));553 a = a.slice(a.size() - b.size());554 results.append(addElementwise(builder, loc, a, b));555 return results;556}557 558template <typename T>559int xegpu::getLargestDivisor(T dim, ArrayRef<T> candidates,560 ArrayRef<T> candidateMultiples) {561 static_assert(std::is_integral<T>::value, "T must be an integer type");562 int largest = -1;563 SmallVector<T> multiples = {1};564 if (!candidateMultiples.empty())565 multiples =566 SmallVector<T>(candidateMultiples.begin(), candidateMultiples.end());567 for (T candidate : candidates) {568 for (T multiple : multiples) {569 int value = static_cast<int>(candidate * multiple);570 if (value != 0 && dim % value == 0 && value > largest)571 largest = value;572 }573 }574 return largest;575}576 577/// Explicit instantiations578template int xegpu::getLargestDivisor<int>(int dim, ArrayRef<int> candidates,579 ArrayRef<int> candidateMultiples);580template int581xegpu::getLargestDivisor<unsigned>(unsigned dim, ArrayRef<unsigned> candidates,582 ArrayRef<unsigned> candidateMultiples);583