1356 lines · cpp
1//===- XeGPUPropagateLayout.cpp - XeGPU Layout Propagation ------*- C++ -*-===//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 9#include "mlir/Analysis/DataFlow/DeadCodeAnalysis.h"10#include "mlir/Analysis/DataFlow/SparseAnalysis.h"11#include "mlir/Analysis/DataFlow/Utils.h"12#include "mlir/Analysis/DataFlowFramework.h"13#include "mlir/Dialect/GPU/IR/GPUDialect.h"14#include "mlir/Dialect/MemRef/IR/MemRef.h"15#include "mlir/Dialect/Vector/IR/VectorOps.h"16#include "mlir/Dialect/XeGPU/IR/XeGPU.h"17#include "mlir/Dialect/XeGPU/Transforms/Passes.h"18#include "mlir/Dialect/XeGPU/Utils/XeGPUUtils.h"19#include "mlir/IR/Attributes.h"20#include "mlir/IR/Builders.h"21#include "mlir/IR/BuiltinAttributes.h"22#include "mlir/IR/BuiltinTypes.h"23#include "mlir/IR/Operation.h"24#include "mlir/IR/Value.h"25#include "mlir/IR/Visitors.h"26#include "mlir/Interfaces/ControlFlowInterfaces.h"27#include "mlir/Interfaces/FunctionInterfaces.h"28#include "mlir/Support/LLVM.h"29#include "llvm/ADT/ArrayRef.h"30#include "llvm/ADT/STLExtras.h"31#include "llvm/ADT/SmallSet.h"32#include "llvm/ADT/SmallVector.h"33#include "llvm/ADT/TypeSwitch.h"34#include "llvm/Support/Casting.h"35#include "llvm/Support/Debug.h"36#include "llvm/Support/LogicalResult.h"37#include "llvm/Support/raw_ostream.h"38 39#include "mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h"40 41namespace mlir {42namespace xegpu {43#define GEN_PASS_DEF_XEGPUPROPAGATELAYOUT44#include "mlir/Dialect/XeGPU/Transforms/Passes.h.inc"45} // namespace xegpu46} // namespace mlir47 48#define DEBUG_TYPE "xegpu-propagate-layout"49#define DBGS() (llvm::dbgs() << "[" DEBUG_TYPE "]: ")50 51using namespace mlir;52using namespace mlir::dataflow;53 54namespace {55 56enum class LayoutKind { Lane, InstData };57 58//===----------------------------------------------------------------------===//59// LayoutInfo60//===----------------------------------------------------------------------===//61 62/// Helper class for tracking the analysis state of an mlir value. For layout63/// propagation, the analysis state is simply the distribution layout of64/// each value. The distribution layout information is encapsulated using65/// xegpu::DistributeLayoutAttr class which can hold information about any type66/// of distribution layout that XeGPU dialect supports. Purpose of this analysis67/// to propagate some unique distribution layout for each value in the program68/// starting from a set of anchor operations (like DPAS, StoreNd, etc.). Note69/// that analysis will reach a fixed point when all values are reached some70/// layout and, analysis does not try to modify any already assigned layouts.71///72/// Given this, LayoutInfo satisifies the following properties:73/// 1) A LayoutInfo value can be in one of two states - `assigned` or `not74/// assigned`.75/// 2) Two LayoutInfo values are equal if they are both assigned or76/// both not assigned. The concrete value of assigned state does not matter.77/// 3) The meet operator works as follows:78/// - If current state is assigned, return the current state. (already79/// a unique layout is assigned. don't change it)80/// - Otherwise, return the other state.81 82struct LayoutInfo {83private:84 xegpu::DistributeLayoutAttr storage = nullptr;85 86public:87 LayoutInfo() = default;88 LayoutInfo(const xegpu::DistributeLayoutAttr &layout) : storage(layout) {}89 90 // Two lattice values are equal if they have `some` layout. The actual91 // content of the layout does not matter.92 bool operator==(const LayoutInfo &other) const {93 return this->isAssigned() == other.isAssigned();94 }95 96 static LayoutInfo meet(const LayoutInfo &lhs, const LayoutInfo &rhs);97 98 static LayoutInfo join(const LayoutInfo &lhs, const LayoutInfo &rhs);99 100 void print(raw_ostream &os) const;101 102 bool isAssigned() const { return storage != nullptr; }103 104 LayoutInfo transpose(ArrayRef<int64_t> permutation) const;105 106 SmallVector<int> getLaneLayout() const;107 108 SmallVector<int> getLaneData() const;109 110 SmallVector<int> getInstData() const;111 112 bool isSliceLayout() const {113 if (!isAssigned())114 return false;115 return isa<xegpu::SliceAttr>(storage);116 }117 118 int64_t getRank() const {119 if (!isAssigned())120 return -1;121 return storage.getRank();122 }123 124 Attribute get() { return storage; }125};126 127SmallVector<int> LayoutInfo::getLaneLayout() const {128 if (!isAssigned())129 return {};130 assert(storage.getEffectiveLaneLayoutAsInt().size() &&131 "Expected lane layout to be assigned");132 return llvm::map_to_vector(storage.getEffectiveLaneLayoutAsInt(),133 [](int64_t val) { return static_cast<int>(val); });134}135 136SmallVector<int> LayoutInfo::getLaneData() const {137 if (!isAssigned())138 return {};139 assert(storage.getEffectiveLaneDataAsInt().size() &&140 "Expected lane data to be assigned");141 return llvm::map_to_vector(storage.getEffectiveLaneDataAsInt(),142 [](int64_t val) { return static_cast<int>(val); });143}144 145SmallVector<int> LayoutInfo::getInstData() const {146 if (!isAssigned())147 return {};148 return llvm::map_to_vector(storage.getEffectiveInstDataAsInt(),149 [](int64_t val) { return static_cast<int>(val); });150}151 152void LayoutInfo::print(raw_ostream &os) const {153 if (isAssigned()) {154 os << storage;155 } else {156 os << "Not assigned.";157 }158}159 160LayoutInfo LayoutInfo::meet(const LayoutInfo &lhs, const LayoutInfo &rhs) {161 if (!lhs.isAssigned())162 return rhs;163 return lhs;164}165 166/// Since this is a backward analysis, join method is not used.167LayoutInfo LayoutInfo::join(const LayoutInfo &lhs, const LayoutInfo &rhs) {168 llvm_unreachable("Join should not be triggered by layout propagation.");169}170 171/// Construct a new layout with the transposed inst_data or lane_layout,172/// lane_data.173LayoutInfo LayoutInfo::transpose(ArrayRef<int64_t> permutation) const {174 if (!isAssigned())175 return {};176 // Check if the permutation is valid.177 llvm::SmallSet<int64_t, 4> seen(permutation.begin(), permutation.end());178 bool hasDuplicates = seen.size() != permutation.size();179 bool withinRange = llvm::all_of(permutation, [&](int64_t idx) {180 return idx >= 0 && idx < static_cast<int64_t>(permutation.size());181 });182 183 if (!withinRange || hasDuplicates) {184 assert(false && "Invalid permutation for transpose.");185 return {};186 }187 188 SmallVector<int32_t> laneLayout;189 SmallVector<int32_t> laneData;190 SmallVector<int32_t> instData;191 for (int64_t idx : permutation) {192 if (getLaneLayout().size()) {193 laneLayout.push_back(static_cast<int32_t>(getLaneLayout()[idx]));194 laneData.push_back(static_cast<int32_t>(getLaneData()[idx]));195 }196 if (getInstData().size())197 instData.push_back(static_cast<int32_t>(getInstData()[idx]));198 }199 xegpu::LayoutAttr layoutAttr;200 if (getLaneLayout().size())201 layoutAttr =202 xegpu::LayoutAttr::get(storage.getContext(), laneLayout, laneData);203 if (getInstData().size())204 layoutAttr = xegpu::LayoutAttr::get(storage.getContext(), instData);205 return LayoutInfo(layoutAttr);206}207 208//===----------------------------------------------------------------------===//209// LayoutInfoLattice210//===----------------------------------------------------------------------===//211 212/// Lattice holding the LayoutInfo for each value.213struct LayoutInfoLattice : public Lattice<LayoutInfo> {214 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(LayoutInfoLattice)215 using Lattice::Lattice;216};217 218/// Helper Functions to get default layouts. A `default layout` is a layout that219/// is assigned to a value when the layout is not fixed by some anchor operation220/// (like DPAS).221 222/// Helper Function to get the default layout for uniform values like constants.223/// For 1D vector, lane_layout is [subgroupSize] and lane_data is [1].224/// For 2D vector, lane_layout is [1, subgroupSize] and lane_data is [1, 1].225static LayoutInfo getDefaultSIMTLayoutInfo(mlir::MLIRContext *ctx,226 unsigned rank,227 const xegpu::uArch::uArch *uArch) {228 assert((rank == 1 || rank == 2) && "Expected 1D or 2D vector.");229 if (rank == 1) {230 return LayoutInfo(231 xegpu::LayoutAttr::get(ctx, {uArch->getSubgroupSize()}, {1}));232 }233 return LayoutInfo(234 xegpu::LayoutAttr::get(ctx, {1, uArch->getSubgroupSize()}, {1, 1}));235}236 237static LayoutInfo getDefaultSIMTLayoutInfo(mlir::MLIRContext *ctx,238 unsigned rank, int subgroupSize) {239 assert((rank == 1 || rank == 2) && "Expected 1D or 2D vector.");240 if (rank == 1) {241 return LayoutInfo(xegpu::LayoutAttr::get(ctx, {subgroupSize}, {1}));242 }243 return LayoutInfo(xegpu::LayoutAttr::get(ctx, {1, subgroupSize}, {1, 1}));244}245 246/// Helper to get the default layout for a vector type.247static LayoutInfo getDefaultSIMTLayoutInfo(VectorType vectorTy,248 const xegpu::uArch::uArch *uArch,249 unsigned packingSize,250 bool isScattered = false) {251 // Expecting a 1D or 2D vector.252 assert((vectorTy.getRank() == 1 || vectorTy.getRank() == 2) &&253 "Expected 1D or 2D vector.");254 // Expecting int or float element type.255 assert(vectorTy.getElementType().isIntOrFloat() &&256 "Expected int or float element type.");257 // If the rank is 1, then return default layout for 1D vector.258 if (vectorTy.getRank() == 1)259 return getDefaultSIMTLayoutInfo(vectorTy.getContext(), 1, uArch);260 // Packing factor is determined by the element type bitwidth.261 unsigned bitwidth = vectorTy.getElementType().getIntOrFloatBitWidth();262 int packingFactor = bitwidth < packingSize ? packingSize / bitwidth : 1;263 if (isScattered) {264 return LayoutInfo(xegpu::LayoutAttr::get(vectorTy.getContext(),265 {uArch->getSubgroupSize(), 1},266 {1, packingFactor}));267 }268 return LayoutInfo(xegpu::LayoutAttr::get(vectorTy.getContext(),269 {1, uArch->getSubgroupSize()},270 {1, packingFactor}));271}272 273/// Helper to get the default layout for a vector type.274static LayoutInfo getDefaultSIMTLayoutInfo(xegpu::TensorDescType tdescTy,275 const xegpu::uArch::uArch *uArch,276 unsigned packingSize,277 bool isScattered = false) {278 // Expecting a 1D or 2D vector.279 assert((tdescTy.getRank() == 1 || tdescTy.getRank() == 2) &&280 "Expected 1D or 2D TensorDesc.");281 // Expecting int or float element type.282 assert(tdescTy.getElementType().isIntOrFloat() &&283 "Expected int or float element type.");284 // If the rank is 1, then return default layout for 1D vector.285 if (tdescTy.getRank() == 1)286 return getDefaultSIMTLayoutInfo(tdescTy.getContext(), 1, uArch);287 // Packing factor is determined by the element type bitwidth.288 unsigned bitwidth = tdescTy.getElementType().getIntOrFloatBitWidth();289 int subgroupSize = uArch->getSubgroupSize();290 int packingFactor = bitwidth < packingSize ? packingSize / bitwidth : 1;291 if (isScattered) {292 return LayoutInfo(xegpu::LayoutAttr::get(293 tdescTy.getContext(), {subgroupSize, 1}, {1, packingFactor}));294 }295 296 return LayoutInfo(xegpu::LayoutAttr::get(297 tdescTy.getContext(), {1, subgroupSize}, {1, packingFactor}));298}299 300/// Helper Function to get the expected layouts for DPAS operands. `lane_data`301/// is set according to the following criteria:302/// * For A operand, the data must be packed in minimum303/// `packedSizeInBitsForDefault`304/// * For B operand, the data must be packed in minimum305/// `packedSizeInBitsForDpasB`306static LayoutInfo307getSIMTLayoutInfoForDPASOperand(VectorType vectorTy, unsigned operandNum,308 const xegpu::uArch::uArch *uArch,309 unsigned packingSize) {310 Type elementTy = vectorTy.getElementType();311 assert(elementTy.isIntOrFloat() &&312 "Expected int or float type in DPAS operands");313 SmallVector<int32_t, 2> layout({1, uArch->getSubgroupSize()});314 // For B operand, data must be packed in minimum `packedDpasBSizeInBits` and315 // must have the VNNI format.316 if (operandNum == 1 && elementTy.getIntOrFloatBitWidth() < packingSize) {317 SmallVector<int32_t, 2> data(318 {static_cast<int32_t>(packingSize / elementTy.getIntOrFloatBitWidth()),319 1});320 return LayoutInfo(321 xegpu::LayoutAttr::get(vectorTy.getContext(), layout, data));322 }323 // Otherwise, return the default layout for the vector type.324 return getDefaultSIMTLayoutInfo(vectorTy, uArch, packingSize);325}326 327//===----------------------------------------------------------------------===//328// LayoutInfoPropagation329//===----------------------------------------------------------------------===//330 331/// Backward data flow analysis to propagate the lane_layout and lane_data of332/// each value in the program. Currently, the layouts for operands DPAS,333/// StoreNd, and StoreScatter are fixed (known before propagation). Purpose of334/// this analysis is to propagate those known layouts to all their producers and335/// (other) consumers.336class LayoutInfoPropagation337 : public SparseBackwardDataFlowAnalysis<LayoutInfoLattice> {338private:339 LayoutKind layoutKind;340 void visitDpasOp(xegpu::DpasOp dpas, ArrayRef<LayoutInfoLattice *> operands,341 ArrayRef<const LayoutInfoLattice *> results);342 343 void visitStoreNdOp(xegpu::StoreNdOp store,344 ArrayRef<LayoutInfoLattice *> operands,345 ArrayRef<const LayoutInfoLattice *> results);346 347 void visitStoreScatterOp(xegpu::StoreScatterOp storeScatter,348 ArrayRef<LayoutInfoLattice *> operands,349 ArrayRef<const LayoutInfoLattice *> results);350 351 void visitLoadNdOp(xegpu::LoadNdOp load,352 ArrayRef<LayoutInfoLattice *> operands,353 ArrayRef<const LayoutInfoLattice *> results);354 355 void visitLoadGatherOp(xegpu::LoadGatherOp load,356 ArrayRef<LayoutInfoLattice *> operands,357 ArrayRef<const LayoutInfoLattice *> results);358 359 void visitTransposeOp(vector::TransposeOp transpose,360 ArrayRef<LayoutInfoLattice *> operands,361 ArrayRef<const LayoutInfoLattice *> results);362 363 void visitVectorBitcastOp(vector::BitCastOp bitcast,364 ArrayRef<LayoutInfoLattice *> operands,365 ArrayRef<const LayoutInfoLattice *> results);366 367 void visitCreateDescOp(xegpu::CreateDescOp createDesc,368 ArrayRef<LayoutInfoLattice *> operands,369 ArrayRef<const LayoutInfoLattice *> results);370 371 void visitUpdateNdOffsetOp(xegpu::UpdateNdOffsetOp updateNdOffset,372 ArrayRef<LayoutInfoLattice *> operands,373 ArrayRef<const LayoutInfoLattice *> results);374 375 void visitPrefetchNdOp(xegpu::PrefetchNdOp prefetch,376 ArrayRef<LayoutInfoLattice *> operands,377 ArrayRef<const LayoutInfoLattice *> results);378 379 void visitVectorMultiReductionOp(vector::MultiDimReductionOp reduction,380 ArrayRef<LayoutInfoLattice *> operands,381 ArrayRef<const LayoutInfoLattice *> results);382 383 void visitVectorBroadCastOp(vector::BroadcastOp broadcast,384 ArrayRef<LayoutInfoLattice *> operands,385 ArrayRef<const LayoutInfoLattice *> results);386 void visitShapeCastOp(vector::ShapeCastOp shapeCast,387 ArrayRef<LayoutInfoLattice *> operands,388 ArrayRef<const LayoutInfoLattice *> results);389 390 bool hasParamsOfLayoutKind(xegpu::DistributeLayoutAttr anchorLayout);391 392public:393 LayoutInfoPropagation(DataFlowSolver &solver,394 SymbolTableCollection &symbolTable,395 LayoutKind layoutKind)396 : SparseBackwardDataFlowAnalysis(solver, symbolTable),397 layoutKind(layoutKind) {}398 using SparseBackwardDataFlowAnalysis::SparseBackwardDataFlowAnalysis;399 400 LogicalResult401 visitOperation(Operation *op, ArrayRef<LayoutInfoLattice *> operands,402 ArrayRef<const LayoutInfoLattice *> results) override;403 404 void visitBranchOperand(OpOperand &operand) override {};405 406 void visitCallOperand(OpOperand &operand) override {};407 408 void visitExternalCall(CallOpInterface call,409 ArrayRef<LayoutInfoLattice *> operands,410 ArrayRef<const LayoutInfoLattice *> results) override {411 };412 413 void setToExitState(LayoutInfoLattice *lattice) override {414 (void)lattice->meet(LayoutInfo());415 }416};417} // namespace418 419LogicalResult LayoutInfoPropagation::visitOperation(420 Operation *op, ArrayRef<LayoutInfoLattice *> operands,421 ArrayRef<const LayoutInfoLattice *> results) {422 TypeSwitch<Operation *>(op)423 .Case<xegpu::DpasOp>(424 [&](auto dpasOp) { visitDpasOp(dpasOp, operands, results); })425 .Case<xegpu::StoreNdOp>(426 [&](auto storeNdOp) { visitStoreNdOp(storeNdOp, operands, results); })427 .Case<xegpu::StoreScatterOp>([&](auto storeScatterOp) {428 visitStoreScatterOp(storeScatterOp, operands, results);429 })430 .Case<xegpu::LoadNdOp>(431 [&](auto loadNdOp) { visitLoadNdOp(loadNdOp, operands, results); })432 .Case<xegpu::LoadGatherOp>([&](auto loadGatherOp) {433 visitLoadGatherOp(loadGatherOp, operands, results);434 })435 .Case<xegpu::CreateDescOp>([&](auto createDescOp) {436 visitCreateDescOp(createDescOp, operands, results);437 })438 .Case<xegpu::UpdateNdOffsetOp>([&](auto updateNdOffsetOp) {439 visitUpdateNdOffsetOp(updateNdOffsetOp, operands, results);440 })441 .Case<xegpu::PrefetchNdOp>([&](auto prefetchNdOp) {442 visitPrefetchNdOp(prefetchNdOp, operands, results);443 })444 .Case<vector::TransposeOp>([&](auto transposeOp) {445 visitTransposeOp(transposeOp, operands, results);446 })447 .Case<vector::BitCastOp>([&](auto bitcastOp) {448 visitVectorBitcastOp(bitcastOp, operands, results);449 })450 .Case<vector::MultiDimReductionOp>([&](auto reductionOp) {451 visitVectorMultiReductionOp(reductionOp, operands, results);452 })453 .Case<vector::BroadcastOp>([&](auto broadcastOp) {454 visitVectorBroadCastOp(broadcastOp, operands, results);455 })456 .Case<vector::ShapeCastOp>([&](auto shapeCastOp) {457 visitShapeCastOp(shapeCastOp, operands, results);458 })459 // All other ops.460 .Default([&](Operation *op) {461 for (const LayoutInfoLattice *resultInfo : results) {462 if (!resultInfo->getValue().isAssigned())463 continue;464 for (auto [operandInfo, operand] :465 llvm::zip(operands, op->getOpOperands())) {466 // If the operand type is not a vector or tensor descriptor, skip467 // it.468 if (!isa<xegpu::TensorDescType, VectorType>(469 operand.get().getType()))470 continue;471 // Propagate the result layout to the operand.472 meet(operandInfo, *resultInfo);473 }474 }475 });476 477 return success();478}479 480bool LayoutInfoPropagation::hasParamsOfLayoutKind(481 xegpu::DistributeLayoutAttr anchorLayout) {482 if (anchorLayout == nullptr) {483 return false;484 }485 if (layoutKind == LayoutKind::InstData) {486 return !(anchorLayout.getEffectiveInstDataAsInt().empty());487 } else if (layoutKind == LayoutKind::Lane) {488 return !(anchorLayout.getEffectiveLaneLayoutAsInt().empty() ||489 anchorLayout.getEffectiveLaneDataAsInt().empty());490 }491 return false;492}493 494void LayoutInfoPropagation::visitPrefetchNdOp(495 xegpu::PrefetchNdOp prefetch, ArrayRef<LayoutInfoLattice *> operands,496 ArrayRef<const LayoutInfoLattice *> results) {497 498 LayoutInfo prefetchLayout;499 xegpu::DistributeLayoutAttr anchorLayout = prefetch.getLayoutAttr();500 if (hasParamsOfLayoutKind(anchorLayout)) {501 prefetchLayout = LayoutInfo(anchorLayout);502 } else {503 // Here we assign the default layout to the tensor descriptor operand of504 // prefetch.505 auto tdescTy = prefetch.getTensorDescType();506 507 auto uArch = getUArch(getChipStr(prefetch).value_or(""));508 const auto *uArchInstruction =509 dyn_cast<xegpu::uArch::Subgroup2DBlockPrefetchInstruction>(510 uArch->getInstruction(511 xegpu::uArch::InstructionKind::Subgroup2DBlockPrefetch));512 513 auto blockWHC =514 uArchInstruction->getBlockWidthHeightCount(tdescTy.getElementType());515 if (!blockWHC)516 prefetch.emitWarning("No known block params found for the element type.");517 auto [bWidth, bHeight, bCount] = blockWHC.value();518 SmallVector<int> instData;519 int instWidth = xegpu::getLargestDivisor(520 static_cast<int>(tdescTy.getDimSize(tdescTy.getRank() - 1)), bWidth,521 bCount);522 if (instWidth == -1)523 prefetch.emitWarning(524 "No suitable instruction multiple found for the given shape.");525 if (tdescTy.getRank() == 1)526 instData = {instWidth};527 else {528 int instHeight = xegpu::getLargestDivisor(529 static_cast<int>(tdescTy.getDimSize(tdescTy.getRank() - 2)), bHeight);530 if (instHeight == -1)531 prefetch.emitWarning(532 "No suitable instruction multiple found for the given shape.");533 instData = {instHeight, instWidth};534 }535 536 if (layoutKind == LayoutKind::InstData)537 prefetchLayout =538 LayoutInfo(xegpu::LayoutAttr::get(tdescTy.getContext(), instData));539 else540 prefetchLayout = getDefaultSIMTLayoutInfo(541 tdescTy, uArch, uArchInstruction->getPackedFormatBitSize());542 543 prefetch.setLayoutAttr(544 dyn_cast<xegpu::DistributeLayoutAttr>(prefetchLayout.get()));545 }546 // Propagate the layout to the source tensor descriptor.547 propagateIfChanged(operands[0], operands[0]->meet(prefetchLayout));548}549 550void LayoutInfoPropagation::visitVectorMultiReductionOp(551 vector::MultiDimReductionOp reduction,552 ArrayRef<LayoutInfoLattice *> operands,553 ArrayRef<const LayoutInfoLattice *> results) {554 // The layout of the result must be present.555 LayoutInfo resultLayout = results[0]->getValue();556 if (!resultLayout.isAssigned())557 return;558 // We only consider 2D -> 1D reductions at this point.559 VectorType resultTy = llvm::dyn_cast<VectorType>(reduction.getDestType());560 if (!resultTy || resultTy.getRank() != 1) {561 reduction.emitWarning("Expecting output type to be 1D vector.");562 return;563 }564 auto uArch = getUArch(xegpu::getChipStr(reduction).value_or(""));565 // Given that the result is 1D, the layout of the operand should be 2D with566 // default layout.567 LayoutInfo operandLayout = getDefaultSIMTLayoutInfo(568 reduction->getContext(), 2, uArch->getSubgroupSize());569 propagateIfChanged(operands[0], operands[0]->meet(operandLayout));570 // Accumulator should have the same layout as the result.571 propagateIfChanged(operands[1], operands[1]->meet(resultLayout));572}573 574void LayoutInfoPropagation::visitVectorBroadCastOp(575 vector::BroadcastOp broadcast, ArrayRef<LayoutInfoLattice *> operands,576 ArrayRef<const LayoutInfoLattice *> results) {577 // The layout of the result must be present.578 LayoutInfo resultLayout = results[0]->getValue();579 if (!resultLayout.isAssigned())580 return;581 // Only consider vector to vector broadcasts for now.582 VectorType resultTy = broadcast.getResultVectorType();583 VectorType sourceTy = dyn_cast<VectorType>(broadcast.getSourceType());584 if (!sourceTy) {585 broadcast.emitWarning("Expecting source type to be a vector type.");586 return;587 }588 589 // Only consider nD -> nD broadcast.590 if (sourceTy.getRank() != resultTy.getRank()) {591 broadcast.emitWarning("Expecting source and result to have same rank.");592 return;593 }594 SetVector<int64_t> broadcastUnitDims = broadcast.computeBroadcastedUnitDims();595 if (broadcastUnitDims.size() != 1) {596 broadcast.emitWarning("Expecting source type to be nD vector only with "597 "one broadcasted dimension.");598 return;599 }600 // Propagate the result layout to the source operand.601 propagateIfChanged(operands[0], operands[0]->meet(resultLayout));602}603 604void LayoutInfoPropagation::visitShapeCastOp(605 vector::ShapeCastOp shapeCast, ArrayRef<LayoutInfoLattice *> operands,606 ArrayRef<const LayoutInfoLattice *> results) {607 // The layout of the result must be present.608 LayoutInfo resultLayout = results[0]->getValue();609 if (!resultLayout.isAssigned())610 return;611 VectorType sourceTy = shapeCast.getSourceVectorType();612 VectorType resultTy = shapeCast.getResultVectorType();613 // Shape cast layout propagation only supports 1D -> 2D shape casts.614 // TODO: Support kD -> nD shape casts (k < n, n >= 2) where expanded dims are615 // unit dimensions and non-unit dims match.616 if (sourceTy.getRank() != 1 || resultTy.getRank() != 2) {617 shapeCast.emitWarning("Expecting shape cast to be 1D -> 2D.");618 return;619 }620 int64_t slicedDim = resultTy.getShape()[0] == 1 ? 0 : 1;621 xegpu::SliceAttr sliceLayout = xegpu::SliceAttr::get(622 shapeCast->getContext(), cast<xegpu::LayoutAttr>(resultLayout.get()),623 DenseI64ArrayAttr::get(shapeCast->getContext(), {slicedDim}));624 propagateIfChanged(operands[0], operands[0]->meet(LayoutInfo(sliceLayout)));625}626 627/// Propagate the layout of the result tensor to the source tensor descriptor628/// in UpdateNdOffsetOp.629void LayoutInfoPropagation::visitUpdateNdOffsetOp(630 xegpu::UpdateNdOffsetOp updateNdOffset,631 ArrayRef<LayoutInfoLattice *> operands,632 ArrayRef<const LayoutInfoLattice *> results) {633 // The layout of the result must be present.634 LayoutInfo resultLayout = results[0]->getValue();635 if (!resultLayout.isAssigned())636 return;637 // Propagate the layout to the source operand.638 propagateIfChanged(operands[0], operands[0]->meet(resultLayout));639}640 641/// Set the layouts for DPAS A, B, and C operands.642void LayoutInfoPropagation::visitDpasOp(643 xegpu::DpasOp dpas, ArrayRef<LayoutInfoLattice *> operands,644 ArrayRef<const LayoutInfoLattice *> results) {645 646 LayoutInfo dpasALayout;647 LayoutInfo dpasBLayout;648 LayoutInfo dpasCDLayout;649 650 xegpu::DistributeLayoutAttr anchorLayoutCD = dpas.getLayoutCdAttr();651 if (hasParamsOfLayoutKind(anchorLayoutCD)) {652 xegpu::DistributeLayoutAttr anchorLayoutA = dpas.getLayoutAAttr();653 xegpu::DistributeLayoutAttr anchorLayoutB = dpas.getLayoutBAttr();654 assert(hasParamsOfLayoutKind(anchorLayoutA) &&655 "Expected anchor layout for DPAS A operand.");656 assert(hasParamsOfLayoutKind(anchorLayoutB) &&657 "Expected anchor layout for DPAS B operand.");658 dpasALayout = LayoutInfo(anchorLayoutA);659 dpasBLayout = LayoutInfo(anchorLayoutB);660 dpasCDLayout = LayoutInfo(anchorLayoutCD);661 662 } else {663 664 VectorType aTy = dpas.getLhsType();665 VectorType bTy = dpas.getRhsType();666 667 auto uArch = getUArch(getChipStr(dpas).value_or(""));668 const int subgroupSize = uArch->getSubgroupSize();669 const auto *uArchInstruction =670 dyn_cast<xegpu::uArch::SubgroupMatrixMultiplyAcc>(uArch->getInstruction(671 xegpu::uArch::InstructionKind::SubgroupMatrixMultiplyAcc));672 673 const unsigned dataALen = aTy.getShape().front();674 auto supportedALen = uArchInstruction->getSupportedM(aTy.getElementType());675 const int maxALen =676 xegpu::getLargestDivisor(dataALen, ArrayRef<unsigned>(supportedALen));677 if (maxALen == -1)678 dpas.emitWarning(679 "No suitable instruction multiple found for the given shape.");680 681 const unsigned dataBLen = bTy.getShape().back();682 auto supportedBLen = uArchInstruction->getSupportedN(bTy.getElementType());683 684 const int maxBLen =685 xegpu::getLargestDivisor(dataBLen, ArrayRef<unsigned>(supportedBLen));686 687 if (maxBLen == -1)688 dpas.emitWarning(689 "No suitable instruction multiple found for the given shape.");690 SmallVector<int> instDataA = {maxALen, subgroupSize};691 SmallVector<int> instDataB = {subgroupSize, maxBLen};692 693 if (layoutKind == LayoutKind::InstData) {694 dpasALayout =695 LayoutInfo(xegpu::LayoutAttr::get(dpas.getContext(), instDataA));696 dpasBLayout =697 LayoutInfo(xegpu::LayoutAttr::get(dpas.getContext(), instDataB));698 } else {699 dpasALayout = getSIMTLayoutInfoForDPASOperand(700 aTy, 0, uArch, uArchInstruction->getPackedFormatBitSizeA());701 dpasBLayout = getSIMTLayoutInfoForDPASOperand(702 bTy, 1, uArch, uArchInstruction->getPackedFormatBitSizeB());703 }704 705 if (operands.size() > 2) {706 VectorType cTy = dpas.getAccType();707 if (layoutKind == LayoutKind::InstData) {708 const unsigned dataCLen = bTy.getShape().back();709 auto supportedCLen =710 uArchInstruction->getSupportedN(bTy.getElementType());711 const int maxCLen = xegpu::getLargestDivisor(712 dataCLen, ArrayRef<unsigned>(supportedCLen));713 if (maxCLen == -1)714 dpas.emitWarning(715 "No suitable instruction multiple found for the given shape.");716 SmallVector<int> instDataC = {maxALen, maxCLen};717 dpasCDLayout =718 LayoutInfo(xegpu::LayoutAttr::get(dpas.getContext(), instDataC));719 } else720 dpasCDLayout = getSIMTLayoutInfoForDPASOperand(721 cTy, 2, uArch, uArchInstruction->getPackedFormatBitSizeB());722 723 dpas.setLayoutCdAttr(724 dyn_cast<xegpu::DistributeLayoutAttr>(dpasCDLayout.get()));725 }726 dpas.setLayoutAAttr(727 dyn_cast<xegpu::DistributeLayoutAttr>(dpasALayout.get()));728 dpas.setLayoutBAttr(729 dyn_cast<xegpu::DistributeLayoutAttr>(dpasBLayout.get()));730 }731 732 propagateIfChanged(operands[0], operands[0]->meet(dpasALayout));733 propagateIfChanged(operands[1], operands[1]->meet(dpasBLayout));734 if (operands.size() > 2) {735 propagateIfChanged(operands[2], operands[2]->meet(dpasCDLayout));736 }737}738 739/// Set the layout for the value and tensor descriptor operands in StoreNdOp.740void LayoutInfoPropagation::visitStoreNdOp(741 xegpu::StoreNdOp store, ArrayRef<LayoutInfoLattice *> operands,742 ArrayRef<const LayoutInfoLattice *> results) {743 744 LayoutInfo storeLayout;745 xegpu::DistributeLayoutAttr anchorLayout = store.getLayoutAttr();746 if (hasParamsOfLayoutKind(anchorLayout)) {747 storeLayout = LayoutInfo(anchorLayout);748 } else {749 auto uArch = getUArch(getChipStr(store).value_or(""));750 const auto *uArchInstruction =751 dyn_cast<xegpu::uArch::Subgroup2DBlockStoreInstruction>(752 uArch->getInstruction(753 xegpu::uArch::InstructionKind::Subgroup2DBlockStore));754 VectorType dataTy = store.getValueType();755 auto blockWHC = uArchInstruction->getBlockWidthHeightCount(756 store.getValueType().getElementType());757 if (!blockWHC)758 store.emitWarning("No known block params found for the element type.");759 auto [bWidth, bHeight, bCount] = blockWHC.value();760 SmallVector<int> instData;761 int instWidth = xegpu::getLargestDivisor(762 static_cast<int>(dataTy.getDimSize(dataTy.getRank() - 1)), bWidth,763 bCount);764 if (instWidth == -1)765 store.emitWarning(766 "No suitable instruction multiple found for the given shape.");767 if (dataTy.getRank() == 1)768 instData = {instWidth};769 else {770 int instHeight = xegpu::getLargestDivisor(771 static_cast<int>(dataTy.getDimSize(dataTy.getRank() - 2)), bHeight);772 if (instHeight == -1)773 store.emitWarning(774 "No suitable instruction multiple found for the given shape.");775 instData = {instHeight, instWidth};776 }777 778 if (layoutKind == LayoutKind::InstData)779 storeLayout =780 LayoutInfo(xegpu::LayoutAttr::get(dataTy.getContext(), instData));781 else782 storeLayout =783 getDefaultSIMTLayoutInfo(store.getValueType(), uArch,784 uArchInstruction->getPackedFormatBitSize());785 store.setLayoutAttr(786 dyn_cast<xegpu::DistributeLayoutAttr>(storeLayout.get()));787 }788 // Propagate the layout to the value operand.789 // Both operands should have the same layout790 for (LayoutInfoLattice *operand : operands)791 propagateIfChanged(operand, operand->meet(storeLayout));792}793 794/// Propagate the layout of the value to the tensor descriptor operand in795/// LoadNdOp.796void LayoutInfoPropagation::visitLoadNdOp(797 xegpu::LoadNdOp load, ArrayRef<LayoutInfoLattice *> operands,798 ArrayRef<const LayoutInfoLattice *> results) {799 800 LayoutInfo loadLayout;801 xegpu::DistributeLayoutAttr anchorLayout = load.getLayoutAttr();802 if (hasParamsOfLayoutKind(anchorLayout)) {803 loadLayout = LayoutInfo(anchorLayout);804 } else {805 806 LayoutInfo valueLayout = results[0]->getValue();807 // Need the layout of the value to propagate to the tensor descriptor.808 if (!valueLayout.isAssigned())809 return;810 loadLayout = valueLayout;811 // LoadNdOp has the transpose effect. However, at the stage of this analysis812 // this effect is not expected and should be abstracted away. Emit a813 // warning.814 if (auto transpose = load.getTranspose()) {815 load.emitWarning("Transpose effect is not expected for LoadNdOp at "816 "LayoutInfoPropagation stage.");817 loadLayout = valueLayout.transpose(transpose.value());818 }819 load.setLayoutAttr(dyn_cast<xegpu::DistributeLayoutAttr>(loadLayout.get()));820 }821 // Propagate the new layout to the tensor descriptor operand.822 propagateIfChanged(operands[0], operands[0]->meet(loadLayout));823}824 825/// For vector::TransposeOp, the layout of the result is transposed and826/// propagated to the operand.827void LayoutInfoPropagation::visitTransposeOp(828 vector::TransposeOp transpose, ArrayRef<LayoutInfoLattice *> operands,829 ArrayRef<const LayoutInfoLattice *> results) {830 // Need the layout of transpose result to propagate to the operands.831 LayoutInfo resultLayout = results[0]->getValue();832 if (!resultLayout.isAssigned())833 return;834 LayoutInfo newLayout = resultLayout.transpose(transpose.getPermutation());835 // Propagate the new layout to the vector operand.836 propagateIfChanged(operands[0], operands[0]->meet(newLayout));837}838 839/// For vector::BitCastOp, the lane_data of the source layout is changed based840/// on the bit width of the source and result types.841void LayoutInfoPropagation::visitVectorBitcastOp(842 vector::BitCastOp bitcast, ArrayRef<LayoutInfoLattice *> operands,843 ArrayRef<const LayoutInfoLattice *> results) {844 // Need the layout of bitcast result to propagate to the operands.845 LayoutInfo resultLayout = results[0]->getValue();846 if (!resultLayout.isAssigned())847 return;848 int inElemTyBitWidth =849 bitcast.getSourceVectorType().getElementType().getIntOrFloatBitWidth();850 int outElemTyBitWidth =851 bitcast.getResultVectorType().getElementType().getIntOrFloatBitWidth();852 // If the element bit widths are the same, then the layout does not change.853 if (inElemTyBitWidth == outElemTyBitWidth) {854 propagateIfChanged(operands[0], operands[0]->meet(resultLayout));855 return;856 }857 // Check if the result layout is valid. i.e. result vector can be distributed.858 auto resultLaneLayout = resultLayout.getLaneLayout();859 auto resultLaneData = resultLayout.getLaneData();860 if (failed(xegpu::getDistributedVectorType(861 bitcast.getResultVectorType(),862 xegpu::LayoutAttr::get(bitcast->getContext(), resultLaneLayout,863 resultLaneData)))) {864 bitcast.emitWarning(865 "Result vector type can not be evenly distributed across lanes.");866 return;867 }868 int64_t rank = bitcast.getSourceVectorType().getRank();869 // Bitcast is a `narrowing` if the input element type bit width larger than870 // the output element type bit width. eg. f32 -> f16 is a narrowing bitcast.871 bool isNarrowing = inElemTyBitWidth > outElemTyBitWidth;872 int bitCastRatio = isNarrowing ? inElemTyBitWidth / outElemTyBitWidth873 : outElemTyBitWidth / inElemTyBitWidth;874 SmallVector<int> sourceLaneLayout =875 resultLayout.getLaneLayout(); // Lane layout does not change for bitcast.876 SmallVector<int> outData = resultLayout.getLaneData();877 878 // TODO: Currently we assume that bitcasts does not require cross lane879 // communication. So each lane must own the required number of elements to880 // perform the bitcast locally without cross-lane communication.881 int outInnerBitsPerLane = outData[rank - 1] * outElemTyBitWidth;882 if (outInnerBitsPerLane < inElemTyBitWidth) {883 bitcast.emitWarning(884 "Narrowing bitcast with cross lane communication is not supported.");885 return;886 }887 // Check if each lane owns a single element in all dimensions except the888 // innermost dimension.889 SmallVector<int> sourceLaneData(outData.begin(), outData.end() - 1);890 if (llvm::any_of(sourceLaneData, [](int64_t d) { return d != 1; })) {891 bitcast.emitWarning("Each lane must not own multiple elements in any "892 "dimension other than "893 "the innermost dimension.");894 return;895 }896 // Decide lane data based on whether the bitcast is narrowing or widening.897 int64_t innerMostLaneData = isNarrowing ? outData[rank - 1] / bitCastRatio898 : outData[rank - 1] * bitCastRatio;899 sourceLaneData.push_back(innerMostLaneData);900 901 propagateIfChanged(902 operands[0],903 operands[0]->meet(LayoutInfo(xegpu::LayoutAttr::get(904 bitcast->getContext(), sourceLaneLayout, sourceLaneData))));905}906 907/// Propagate the layout of the result to the tensor descriptor, mask and offset908/// operands in LoadGatherOp.909void LayoutInfoPropagation::visitLoadGatherOp(910 xegpu::LoadGatherOp load, ArrayRef<LayoutInfoLattice *> operands,911 ArrayRef<const LayoutInfoLattice *> results) {912 913 LayoutInfo loadLayout;914 LayoutInfo maskLayout;915 xegpu::DistributeLayoutAttr anchorLayout = load.getLayoutAttr();916 if (hasParamsOfLayoutKind(anchorLayout)) {917 loadLayout = LayoutInfo(anchorLayout);918 maskLayout = loadLayout;919 } else {920 921 // The layout is strictly determined by the payload type.922 auto payloadTy = dyn_cast<VectorType>(load.getValueType());923 if (!payloadTy) {924 load.emitWarning("Not propagating, non-vector payload supplied.");925 return;926 }927 auto uArch = getUArch(getChipStr(load).value_or(""));928 const int subgroupSize = uArch->getSubgroupSize();929 SmallVector<int> instData{subgroupSize};930 if (auto chunkSize = load.getChunkSize().value_or(0); chunkSize > 1)931 instData.push_back(chunkSize);932 else if (auto srcTdescTy =933 dyn_cast<xegpu::TensorDescType>(load.getSourceType())) {934 if (srcTdescTy.getChunkSizeAsInt() > 1)935 instData.push_back(chunkSize);936 }937 938 if (layoutKind == LayoutKind::InstData)939 loadLayout =940 LayoutInfo(xegpu::LayoutAttr::get(load.getContext(), instData));941 else942 loadLayout = getDefaultSIMTLayoutInfo(943 payloadTy, uArch, uArch->getGeneralPackedFormatBitSize(),944 /*scattered*/ true);945 946 // Mask operand should have 1D default layout.947 maskLayout = getDefaultSIMTLayoutInfo(load->getContext(), 1, subgroupSize);948 949 load.setLayoutAttr(dyn_cast<xegpu::DistributeLayoutAttr>(loadLayout.get()));950 }951 // Propagate the new layout to the tensor descriptor operand.952 if (isa<xegpu::TensorDescType>(load.getSourceType()))953 propagateIfChanged(operands[0], operands[0]->meet(loadLayout));954 // Propagate the new layout to the mask and optional offset operand.955 propagateIfChanged(operands[1], operands[1]->meet(maskLayout));956 if (load.getOffsets())957 propagateIfChanged(operands[2], operands[2]->meet(maskLayout));958}959 960/// Propagate the layout of the descriptor to the vector offset operand in961/// CreateDescOp.962void LayoutInfoPropagation::visitCreateDescOp(963 xegpu::CreateDescOp createDesc, ArrayRef<LayoutInfoLattice *> operands,964 ArrayRef<const LayoutInfoLattice *> results) {965 LayoutInfo descLayout = results[0]->getValue();966 // Need the layout of the descriptor to propagate to the operands.967 if (!descLayout.isAssigned())968 return;969 auto uArch = getUArch(getChipStr(createDesc).value_or(""));970 // For offset operand propagate 1D default layout.971 LayoutInfo layout = getDefaultSIMTLayoutInfo(createDesc->getContext(), 1,972 uArch->getSubgroupSize());973 propagateIfChanged(operands[1], operands[1]->meet(layout));974}975 976/// Set the layout for the value, tensor descriptor, offset and mask operands in977/// the StoreScatterOp.978void LayoutInfoPropagation::visitStoreScatterOp(979 xegpu::StoreScatterOp storeScatter, ArrayRef<LayoutInfoLattice *> operands,980 ArrayRef<const LayoutInfoLattice *> results) {981 982 LayoutInfo payloadLayout;983 LayoutInfo maskLayout;984 xegpu::DistributeLayoutAttr anchorLayout = storeScatter.getLayoutAttr();985 if (hasParamsOfLayoutKind(anchorLayout)) {986 payloadLayout = LayoutInfo(anchorLayout);987 maskLayout = payloadLayout;988 } else {989 // Currently, for 2D StoreScatterOp we expect that the height dimension of990 // the tensor descriptor is equal to the subgroup size. This is ensured by991 // the op verifier.992 auto payloadTy = dyn_cast<VectorType>(storeScatter.getValueType());993 if (!payloadTy) {994 storeScatter.emitWarning("Not propagating, non-vector payload supplied.");995 return;996 }997 998 auto uArch = getUArch(getChipStr(storeScatter).value_or(""));999 const int subgroupSize = uArch->getSubgroupSize();1000 1001 if (layoutKind == LayoutKind::InstData) {1002 SmallVector<int> instData{subgroupSize};1003 if (auto chunkSize = storeScatter.getChunkSize().value_or(0);1004 chunkSize > 1)1005 instData.push_back(chunkSize);1006 else if (auto dstTdescTy = dyn_cast<xegpu::TensorDescType>(1007 storeScatter.getDestType())) {1008 if (dstTdescTy.getChunkSizeAsInt() > 1)1009 instData.push_back(chunkSize);1010 }1011 payloadLayout = LayoutInfo(1012 xegpu::LayoutAttr::get(storeScatter.getContext(), instData));1013 } else {1014 auto payloadShape = payloadTy.getShape();1015 if (payloadShape.size() > 1)1016 assert(payloadShape[0] == subgroupSize &&1017 "Expected the first dimension of 2D tensor descriptor to be "1018 "equal to "1019 "subgroup size.");1020 payloadLayout = getDefaultSIMTLayoutInfo(1021 payloadTy, uArch, uArch->getGeneralPackedFormatBitSize(),1022 /*scattered=*/true);1023 }1024 1025 maskLayout =1026 getDefaultSIMTLayoutInfo(storeScatter->getContext(), 1, subgroupSize);1027 1028 storeScatter.setLayoutAttr(1029 dyn_cast<xegpu::DistributeLayoutAttr>(payloadLayout.get()));1030 }1031 // Propagate the payload operand layout1032 propagateIfChanged(operands[0], operands[0]->meet(payloadLayout));1033 // Propagate the destination (if tdesc) operand layout1034 if (isa<xegpu::TensorDescType>(storeScatter.getDestType()))1035 propagateIfChanged(operands[1], operands[1]->meet(payloadLayout));1036 // Propagate the new layout to the mask and optional offset operand.1037 propagateIfChanged(operands[2], operands[2]->meet(maskLayout));1038 if (storeScatter.getOffsets())1039 propagateIfChanged(operands[3], operands[3]->meet(maskLayout));1040}1041 1042namespace {1043//===----------------------------------------------------------------------===//1044// RunLayoutInfoPropagation1045//===----------------------------------------------------------------------===//1046 1047/// Driver class for running the LayoutInfoPropagation analysis.1048class RunLayoutInfoPropagation {1049public:1050 MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(RunLayoutInfoPropagation)1051 1052 RunLayoutInfoPropagation(Operation *op, LayoutKind layoutKind) : target(op) {1053 SymbolTableCollection symbolTable;1054 loadBaselineAnalyses(solver);1055 solver.load<LayoutInfoPropagation>(symbolTable, layoutKind);1056 (void)solver.initializeAndRun(op);1057 }1058 1059 LayoutInfo getLayoutInfo(Value val);1060 1061 void printAnalysisResult(llvm::raw_ostream &os);1062 1063private:1064 DataFlowSolver solver;1065 const Operation *target;1066};1067} // namespace1068 1069LayoutInfo RunLayoutInfoPropagation::getLayoutInfo(Value val) {1070 auto *state = solver.lookupState<LayoutInfoLattice>(val);1071 if (!state)1072 return {};1073 return state->getValue();1074}1075 1076// Print the analysis result for debugging purposes.1077void RunLayoutInfoPropagation::printAnalysisResult(llvm::raw_ostream &os) {1078 auto printFunctionResult = [&](FunctionOpInterface funcOp) {1079 os << "function: " << funcOp.getName() << ":\n";1080 // Function arguments1081 for (BlockArgument arg : funcOp.getArguments()) {1082 LayoutInfo layout = getLayoutInfo(arg);1083 os << "argument: " << arg << "\n";1084 os << "layout : ";1085 layout.print(os);1086 os << "\n";1087 }1088 // Function ops1089 funcOp.walk([&](Operation *op) {1090 // Skip ops that do not have results1091 if (op->getResults().empty())1092 return;1093 os << "op : ";1094 // For control-flow ops, print the op name only.1095 if (isa<BranchOpInterface>(op) || isa<RegionBranchOpInterface>(op))1096 os << op->getName();1097 else1098 op->print(os);1099 os << "\n";1100 // Print the layout for each result.1101 for (auto [i, r] : llvm::enumerate(op->getResults())) {1102 LayoutInfo layout = getLayoutInfo(r);1103 os << "layout for result #" << i << ": ";1104 layout.print(os);1105 os << "\n";1106 }1107 });1108 };1109 1110 SmallVector<FunctionOpInterface> funcOps;1111 if (auto modOp = dyn_cast<ModuleOp>(target)) {1112 for (auto funcOp : modOp.getOps<FunctionOpInterface>())1113 funcOps.push_back(funcOp);1114 1115 // Collect all GpuFuncOps in the module.1116 for (auto gpuModOp : modOp.getOps<gpu::GPUModuleOp>()) {1117 for (auto gpuFuncOp : gpuModOp.getOps<FunctionOpInterface>())1118 funcOps.push_back(gpuFuncOp);1119 }1120 }1121 // Print the analysis result for each function.1122 for (FunctionOpInterface funcOp : funcOps)1123 printFunctionResult(funcOp);1124}1125 1126using GetLayoutFnTy = function_ref<xegpu::DistributeLayoutAttr(Value)>;1127/// Update an operation with the layout of its results. If the result type is1128/// a vector type, a temporary layout attribute is added to the operation. If1129/// the result type is a tensor descriptor type, the type is updated with the1130/// layout attribute. The users of the result are also updated with the layout1131/// attribute.1132static LogicalResult updateOp(mlir::OpBuilder &builder, mlir::Operation *op,1133 GetLayoutFnTy getLayoutOfValue) {1134 // Region ops (like scf.for) are already handled by the1135 // updateControlFlowOps.1136 if (mlir::isa<mlir::RegionBranchOpInterface>(op))1137 return success();1138 1139 // Iterate over all the results.1140 for (OpResult result : op->getResults()) {1141 Type resultType = result.getType();1142 // Layouts are needed only for vector and tensor descriptor types.1143 if (!isa<VectorType, xegpu::TensorDescType>(resultType))1144 continue;1145 // If the result has no layout but has users, emit a warning and continue.1146 xegpu::DistributeLayoutAttr layout = getLayoutOfValue(result);1147 if (!layout && result.getNumUses() > 0) {1148 op->emitWarning("op has users but no layout assigned for its result");1149 continue;1150 }1151 // If the result is a tensor descriptor type, update the tensor desc type1152 // with layout.1153 if (auto tensorDescTy = dyn_cast<xegpu::TensorDescType>(resultType)) {1154 auto typeWithLayout = xegpu::TensorDescType::get(1155 tensorDescTy.getContext(), tensorDescTy.getShape(),1156 tensorDescTy.getElementType(), tensorDescTy.getEncoding(), layout);1157 result.setType(typeWithLayout);1158 continue;1159 }1160 // If the result is a vector type, add a temporary layout attribute to the1161 // op.1162 xegpu::setDistributeLayoutAttr(result, layout, /*respectPermLayout*/ true);1163 }1164 return success();1165}1166 1167/// Region ops like scf.for need special handling because they have blocks1168/// inside. If the blocks have tensor descriptor type as block arguments,1169/// thier types must be updated. Also region op can have results that may not1170/// have any users (e.g. A and B tiles). They are not assigned a layout by1171/// layout analysis because they have no users. However inside the region op1172/// corresponding block arguments for these results do have layouts.1173/// Therefore, in this case we still need to update the result types with the1174/// layout attribute. This function function updates the internal block1175/// arguments and the result types of the region op with the assigned layouts.1176/// clang-format off1177/// Example: scf.for ... iter_args(...) -> (out types) {1178/// ^bb0(block types):1179/// ...1180/// scf.yield ... : (yield types)1181/// }1182/// clang-format on1183/// In this example, at scf.yield, control-flow can transfer to two successor1184/// regions. One is the ^bb0 (for loop body) and the other is the scf.for op1185/// itself (yield the results). So we update both the block arguments of the1186/// successor region (i.e. block types) and the result types of the scf.for op1187/// (i.e. out types). Note that yield types are updated by respective1188/// producers inside bb0.1189static LogicalResult1190updateControlFlowOps(mlir::OpBuilder &builder,1191 mlir::RegionBranchTerminatorOpInterface terminator,1192 GetLayoutFnTy getLayoutOfValue) {1193 // Only process if the terminator is inside a region branch op.1194 if (!mlir::isa<mlir::RegionBranchOpInterface>(terminator->getParentOp()))1195 return success();1196 1197 llvm::SmallVector<mlir::RegionSuccessor> successors;1198 llvm::SmallVector<mlir::Attribute> operands(terminator->getNumOperands(),1199 nullptr);1200 terminator.getSuccessorRegions(operands, successors);1201 1202 for (mlir::RegionSuccessor &successor : successors) {1203 mlir::OperandRange successorOperands =1204 terminator.getSuccessorOperands(successor);1205 mlir::ValueRange successorInputs = successor.getSuccessorInputs();1206 for (auto [successorOperand, successorInput] :1207 llvm::zip(successorOperands, successorInputs)) {1208 Type inputType = successorInput.getType();1209 // We only need to operate on tensor descriptor or vector types.1210 if (!isa<xegpu::TensorDescType, VectorType>(inputType))1211 continue;1212 xegpu::DistributeLayoutAttr successorInputLayout =1213 getLayoutOfValue(successorInput);1214 xegpu::DistributeLayoutAttr successorOperandLayout =1215 getLayoutOfValue(successorOperand);1216 1217 // If either of the layouts is not assigned, we cannot proceed.1218 if (!successorOperandLayout) {1219 LLVM_DEBUG(DBGS() << "No layout assigned for forwarded operand in "1220 "branch terminator: "1221 << successorOperand << "\n");1222 return failure();1223 }1224 // We expect the layouts to match.1225 if (successorInputLayout &&1226 successorInputLayout != successorOperandLayout) {1227 LLVM_DEBUG(DBGS() << "Conflicting layouts for region argument and "1228 "operand forwarded as the argument: "1229 << successorInputLayout << " vs "1230 << successorOperandLayout << "\n");1231 return failure();1232 }1233 // Get tensor descriptor type with the layout.1234 if (auto tdescTy = dyn_cast<xegpu::TensorDescType>(inputType)) {1235 auto newTdescTy = xegpu::TensorDescType::get(1236 tdescTy.getContext(), tdescTy.getShape(), tdescTy.getElementType(),1237 tdescTy.getEncoding(), successorOperandLayout);1238 successorInput.setType(newTdescTy);1239 continue;1240 }1241 // If the type is a vector type and this region argument is an OpResult,1242 // set the layout attribute on the OpResult.1243 if (auto result = dyn_cast<OpResult>(successorInput))1244 xegpu::setDistributeLayoutAttr(result, successorOperandLayout);1245 }1246 }1247 return success();1248}1249 1250/// Update the function arguments and results with the layouts.1251static LogicalResult updateFunctionOpInterface(mlir::OpBuilder &builder,1252 mlir::FunctionOpInterface funcOp,1253 GetLayoutFnTy getLayoutOfValue) {1254 SmallVector<Type> newArgTypes;1255 // Update the function arguments.1256 for (BlockArgument arg : funcOp.getArguments()) {1257 Type argType = arg.getType();1258 newArgTypes.push_back(argType);1259 if (!isa<VectorType, xegpu::TensorDescType>(argType))1260 continue;1261 xegpu::DistributeLayoutAttr layout = getLayoutOfValue(arg);1262 if (!layout) {1263 LLVM_DEBUG(DBGS() << "Expecting layout for function argument: " << arg1264 << " but got none.\n");1265 return failure();1266 }1267 if (auto tensorDescTy = dyn_cast<xegpu::TensorDescType>(argType)) {1268 auto newTdescTy = xegpu::TensorDescType::get(1269 tensorDescTy.getContext(), tensorDescTy.getShape(),1270 tensorDescTy.getElementType(), tensorDescTy.getEncoding(), layout);1271 arg.setType(newTdescTy);1272 newArgTypes.back() = newTdescTy;1273 }1274 }1275 // Update the function type with the new argument types.1276 // NOTE: We assume that function results are not expected to have layouts.1277 funcOp.setType(FunctionType::get(funcOp.getContext(), newArgTypes,1278 funcOp.getResultTypes()));1279 return success();1280}1281 1282namespace {1283struct XeGPUPropagateLayoutPass final1284 : public xegpu::impl::XeGPUPropagateLayoutBase<XeGPUPropagateLayoutPass> {1285 XeGPUPropagateLayoutPass() = default;1286 XeGPUPropagateLayoutPass(const XeGPUPropagateLayoutPass &other) = default;1287 XeGPUPropagateLayoutPass(xegpu::XeGPUPropagateLayoutOptions options)1288 : XeGPUPropagateLayoutBase(options) {}1289 void runOnOperation() override;1290};1291 1292} // namespace1293 1294void XeGPUPropagateLayoutPass::runOnOperation() {1295 LayoutKind layoutKind;1296 if (this->layoutKind == "lane") {1297 layoutKind = LayoutKind::Lane;1298 } else if (this->layoutKind == "inst") {1299 layoutKind = LayoutKind::InstData;1300 } else {1301 getOperation()->emitError("Unsupported layout kind option: " +1302 this->layoutKind);1303 signalPassFailure();1304 return;1305 }1306 RunLayoutInfoPropagation analysis(getOperation(), layoutKind);1307 // Print the analysis result and exit. (for debugging purposes)1308 if (printOnly) {1309 auto &os = llvm::outs();1310 analysis.printAnalysisResult(os);1311 return;1312 }1313 // Helper to convert LayoutInfo to xegpu::LayoutAttr.1314 auto getXeGPULayoutForValue = [&](Value val) -> xegpu::DistributeLayoutAttr {1315 LayoutInfo layout = analysis.getLayoutInfo(val);1316 if (!layout.isAssigned())1317 return {};1318 xegpu::DistributeLayoutAttr layoutAttr =1319 cast<xegpu::DistributeLayoutAttr>(layout.get());1320 if (layout.isSliceLayout())1321 return cast<xegpu::SliceAttr>(layoutAttr);1322 return cast<xegpu::LayoutAttr>(layoutAttr);1323 };1324 1325 mlir::OpBuilder builder(&getContext());1326 Operation *op = getOperation();1327 auto walkResult = op->walk([&](mlir::Block *block) -> WalkResult {1328 for (mlir::Operation &op : llvm::reverse(block->getOperations())) {1329 LogicalResult r = success();1330 TypeSwitch<Operation *>(&op)1331 .Case<mlir::RegionBranchTerminatorOpInterface>(1332 [&](mlir::RegionBranchTerminatorOpInterface branchTermOp) {1333 r = updateControlFlowOps(builder, branchTermOp,1334 getXeGPULayoutForValue);1335 })1336 .Case<mlir::FunctionOpInterface>(1337 [&](mlir::FunctionOpInterface funcOp) {1338 r = updateFunctionOpInterface(builder, funcOp,1339 getXeGPULayoutForValue);1340 })1341 .Default([&](Operation *op) {1342 r = updateOp(builder, op, getXeGPULayoutForValue);1343 });1344 if (failed(r)) {1345 op.emitError("Failed to update operation with the layout.");1346 return WalkResult::interrupt();1347 }1348 }1349 return WalkResult::advance();1350 });1351 if (walkResult.wasInterrupted()) {1352 signalPassFailure();1353 return;1354 }1355}1356