MLIR 23.0.0git
Typedefs | Functions
XeGPULayoutImpl.cpp File Reference
#include "mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/LLVMIR/XeVMDialect.h"
#include "mlir/Dialect/SCF/Transforms/Patterns.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/XeGPU/IR/XeGPU.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/Operation.h"
#include "mlir/IR/ValueRange.h"
#include "mlir/Interfaces/ControlFlowInterfaces.h"
#include "mlir/Interfaces/LoopLikeInterface.h"
#include "mlir/Interfaces/SideEffectInterfaces.h"
#include "mlir/Transforms/DialectConversion.h"
#include "llvm/ADT/PostOrderIterator.h"
#include "llvm/Support/FormatVariadic.h"
#include <cstdint>
#include <numeric>

Go to the source code of this file.

Typedefs

using LayoutRepresentation = SmallVector<int64_t>

Functions

static void setTensorDescLayout (Value val, xegpu::DistributeLayoutAttr layout)
static void walkRegionBackward (Region &region, llvm::function_ref< void(Operation *)> visit)
static xegpu::DistributeLayoutAttr getLayoutFromUsePoints (Value result)
static void propagateResultsToRegularOperands (Operation *op)
static void propagateRegionResultsToYieldOperands (mlir::RegionBranchTerminatorOpInterface yieldOp)
template void xegpu::removeLayoutAttr< mlir::OpResult > (const mlir::OpResult &result)
template void xegpu::removeLayoutAttr< mlir::OpOperand > (const mlir::OpOperand &operand)
static bool leadingDimsAreUnit (ArrayRef< int64_t > shape, int numInnerDims)
 Returns true if every dimension of shape except the innermost numInnerDims is a unit (size-1) dimension.
static xegpu::LayoutAttr buildInstDataLayoutWithLane (mlir::MLIRContext *context, ArrayRef< int64_t > instData, ArrayRef< int64_t > laneLayout, ArrayRef< int64_t > laneData, DenseI32ArrayAttr orderAttr=nullptr)
static bool isValidLaneLayout (ArrayRef< int64_t > dataShape, ArrayRef< int64_t > laneLayout, ArrayRef< int64_t > laneData)
static xegpu::LayoutAttr buildLaneLayout (mlir::MLIRContext *context, ArrayRef< int64_t > laneLayout, ArrayRef< int64_t > laneData, DenseI32ArrayAttr orderAttr=nullptr)
static xegpu::LayoutAttr buildLayout (mlir::MLIRContext *context, ArrayRef< int64_t > sgLayout, ArrayRef< int64_t > sgData, ArrayRef< int64_t > instData, ArrayRef< int64_t > laneLayout, ArrayRef< int64_t > laneData, DenseI32ArrayAttr orderAttr=nullptr)
static xegpu::LayoutAttr buildSgLayout (mlir::MLIRContext *context, ArrayRef< int64_t > wgTileShape, ArrayRef< int64_t > sgLayout, int dimK=-1, DenseI32ArrayAttr orderAttr=nullptr)
static SmallVector< LayoutRepresentation > enumerateFactorizations (int64_t total, int64_t rank)
 Enumerates all ways to split total into rank factors whose product equals total.
static SmallVector< LayoutRepresentation > getSgLayoutCandidates (ArrayRef< int64_t > wgShape, ArrayRef< int64_t > instData, int64_t sgCount)
static std::optional< SmallVector< int64_t > > get2DBlockIOInstDataLayout (ArrayRef< int64_t > dataShape, Type elemTy, const xegpu::uArch::BlockIOInstructionInterface *uArchInstruction, bool transform=false, bool transpose=false)
 Helper function to compute inst_data vectors for DPAS operands A, B, and C/D.
static std::optional< std::tuple< SmallVector< int64_t >, SmallVector< int64_t >, SmallVector< int64_t > > > getDpasInstDataLayouts (VectorType aTy, VectorType bTy, VectorType cdTy, const xegpu::uArch::MMAInstructionInterface *uArchInstruction)
 Helper function to compute inst_data vectors for DPAS operands A, B, and C/D.
static std::pair< SmallVector< int64_t >, SmallVector< int64_t > > computeScatterIOLaneLayoutAndData (ArrayRef< int64_t > instShape, int64_t subgroupSize, int64_t maxChunkSize)
 Computes lane_layout and lane_data for scatter-style store anchor layouts (store scatter, store matrix).
static std::pair< SmallVector< int64_t >, SmallVector< int64_t > > compute2DBlockIOLaneLayoutAndData (ArrayRef< int64_t > instShape, int64_t subgroupSize, int64_t bitwidth, int64_t packingSize, bool transform=false)
static std::pair< SmallVector< int64_t >, SmallVector< int64_t > > computeReductionLaneLayoutAndData (ArrayRef< int64_t > srcShape, ArrayRef< int64_t > reductionDims, int subgroupSize, int64_t maxReduceVectorSize, bool verticalLaneLayout=false)
 Computes the (lane_layout, lane_data) for a multi-reduction's source layout.
static std::optional< std::tuple< xegpu::DistributeLayoutAttr, xegpu::DistributeLayoutAttr, xegpu::DistributeLayoutAttr > > getDpasSubgroupLayouts (mlir::MLIRContext *context, VectorType aTy, VectorType bTy, VectorType cdTy, xegpu::DistributeLayoutAttr consumerLayout, int numSg, std::tuple< SmallVector< int64_t >, SmallVector< int64_t >, SmallVector< int64_t > > instDataVecs)
 Helper function to set up subgroup layouts for DPAS operands A, B, and C/D.
static xegpu::DistributeLayoutAttr createScaleLayout (mlir::MLIRContext *context, VectorType matrixTy, VectorType scaleTy, xegpu::DistributeLayoutAttr matrixLayout, bool isBScale, const xegpu::uArch::uArch *uArch)
 Helper to create a scale layout derived from a matrix operand layout.
static xegpu::DistributeLayoutAttr setupGenericLoadAnchorLayout (xegpu::LayoutKind layoutKind, mlir::MLIRContext *context, xegpu::DistributeLayoutAttr consumerLayout, int maxChunkSize, ArrayRef< int64_t > resShape, int subgroupSize)
 Sets up the anchor layout for load gather and load matrix operation.
static xegpu::DistributeLayoutAttr getStoreSubgroupLayouts (mlir::MLIRContext *context, ArrayRef< int64_t > wgShape, ArrayRef< int64_t > instData, int numSg)
 Picks the subgroup layout for a scatter-style store (store_scatter / store_matrix): the most balanced numSg factorization that divides wgShape with sg_data a multiple of instData.
static xegpu::DistributeLayoutAttr setupGenericStoreAnchorLayout (xegpu::LayoutKind layoutKind, mlir::MLIRContext *context, int maxChunkSize, ArrayRef< int64_t > srcShape, int subgroupSize, int numSg)
 Sets up the anchor layout for store scatter and store matrix operation, which share the same logic.
static xegpu::DistributeLayoutAttr adjustInnermostDimForDivisibility (xegpu::DistributeLayoutAttr consumerLayout, xegpu::LayoutKind layoutKind, size_t innerMostDim, int ratio, int64_t bound, const xegpu::uArch::uArch *uArch)
 Adjusts consumerLayout's innermost-dim data field selected by layoutKind so that the source layout can be safely inferred by dividing that value by ratio.

Typedef Documentation

◆ LayoutRepresentation

using LayoutRepresentation = SmallVector<int64_t>

Definition at line 856 of file XeGPULayoutImpl.cpp.

Function Documentation

◆ adjustInnermostDimForDivisibility()

xegpu::DistributeLayoutAttr adjustInnermostDimForDivisibility ( xegpu::DistributeLayoutAttr consumerLayout,
xegpu::LayoutKind layoutKind,
size_t innerMostDim,
int ratio,
int64_t bound,
const xegpu::uArch::uArch * uArch )
static

Adjusts consumerLayout's innermost-dim data field selected by layoutKind so that the source layout can be safely inferred by dividing that value by ratio.

Doubles the value until the divisibility constraint is met, bounded above by bound like result-shape.

Used by ops whose source relates to the result by a fixed factor along the innermost dim (e.g., bitcast: bitwidth ratio; interleave: 2x).

Divisibility constraints per LayoutKind:

  • Subgroup: sgData[innermost] % ratio == 0
  • InstData: instData[innermost] % (laneLayout[innermost] * ratio) == 0 (laneLayout falls back to subgroupSize if absent)
  • Lane: laneData[innermost] % ratio == 0

Definition at line 2380 of file XeGPULayoutImpl.cpp.

References mlir::xegpu::uArch::uArch::getSubgroupSize(), mlir::xegpu::InstData, mlir::xegpu::Lane, and mlir::xegpu::Subgroup.

Referenced by mlir::xegpu::setupBitCastResultLayout(), and mlir::xegpu::setupInterleaveResultLayout().

◆ buildInstDataLayoutWithLane()

xegpu::LayoutAttr buildInstDataLayoutWithLane ( mlir::MLIRContext * context,
ArrayRef< int64_t > instData,
ArrayRef< int64_t > laneLayout,
ArrayRef< int64_t > laneData,
DenseI32ArrayAttr orderAttr = nullptr )
static

Definition at line 384 of file XeGPULayoutImpl.cpp.

References mlir::detail::DenseArrayAttrImpl< int32_t >::get().

Referenced by setupGenericLoadAnchorLayout(), setupGenericStoreAnchorLayout(), mlir::xegpu::setupMultiReductionResultLayout(), and mlir::xegpu::setupStoreNdAnchorLayout().

◆ buildLaneLayout()

xegpu::LayoutAttr buildLaneLayout ( mlir::MLIRContext * context,
ArrayRef< int64_t > laneLayout,
ArrayRef< int64_t > laneData,
DenseI32ArrayAttr orderAttr = nullptr )
static

Definition at line 407 of file XeGPULayoutImpl.cpp.

References mlir::detail::DenseArrayAttrImpl< int32_t >::get().

Referenced by setupGenericLoadAnchorLayout(), setupGenericStoreAnchorLayout(), mlir::xegpu::setupMultiReductionResultLayout(), mlir::xegpu::setupReductionResultLayout(), and mlir::xegpu::setupStoreNdAnchorLayout().

◆ buildLayout()

xegpu::LayoutAttr buildLayout ( mlir::MLIRContext * context,
ArrayRef< int64_t > sgLayout,
ArrayRef< int64_t > sgData,
ArrayRef< int64_t > instData,
ArrayRef< int64_t > laneLayout,
ArrayRef< int64_t > laneData,
DenseI32ArrayAttr orderAttr = nullptr )
static

Definition at line 421 of file XeGPULayoutImpl.cpp.

References mlir::detail::DenseArrayAttrImpl< int32_t >::get().

Referenced by buildSgLayout(), createScaleLayout(), and mlir::xegpu::setupMultiReductionResultLayout().

◆ buildSgLayout()

xegpu::LayoutAttr buildSgLayout ( mlir::MLIRContext * context,
ArrayRef< int64_t > wgTileShape,
ArrayRef< int64_t > sgLayout,
int dimK = -1,
DenseI32ArrayAttr orderAttr = nullptr )
static

Definition at line 437 of file XeGPULayoutImpl.cpp.

References buildLayout().

Referenced by getDpasSubgroupLayouts(), getStoreSubgroupLayouts(), and mlir::xegpu::setupStoreNdAnchorLayout().

◆ compute2DBlockIOLaneLayoutAndData()

std::pair< SmallVector< int64_t >, SmallVector< int64_t > > compute2DBlockIOLaneLayoutAndData ( ArrayRef< int64_t > instShape,
int64_t subgroupSize,
int64_t bitwidth,
int64_t packingSize,
bool transform = false )
static

Definition at line 1043 of file XeGPULayoutImpl.cpp.

Referenced by mlir::xegpu::setupStoreNdAnchorLayout().

◆ computeReductionLaneLayoutAndData()

std::pair< SmallVector< int64_t >, SmallVector< int64_t > > computeReductionLaneLayoutAndData ( ArrayRef< int64_t > srcShape,
ArrayRef< int64_t > reductionDims,
int subgroupSize,
int64_t maxReduceVectorSize,
bool verticalLaneLayout = false )
static

Computes the (lane_layout, lane_data) for a multi-reduction's source layout.

Only the innermost two dims are distributed; leading dims are assumed unit. subgroupSize lanes go on one dim; up to maxReduceVectorSize elements are packed into lane_data on the other. To minimize cross-lane reduction, lanes are spread across a non-reduction dim when possible so the reduction happens within a lane. inst_data is the element-wise product lane_layout * lane_data.

e.g. with srcShape=[32, 128], subgroupSize=16, maxReduceVectorSize=2:

  • Switch: reductionDims=[1] and consumerReductionDims=[] -> lanes move to the non-reduction dim 0: lane_layout=[16, 1], lane_data=[1, 2].
  • Default: reductionDims=[0, 1] (both reduced) -> lanes stay on the innermost dim: lane_layout=[1, 16], lane_data=[2, 1].

Definition at line 1078 of file XeGPULayoutImpl.cpp.

Referenced by mlir::xegpu::setupMultiReductionResultLayout().

◆ computeScatterIOLaneLayoutAndData()

std::pair< SmallVector< int64_t >, SmallVector< int64_t > > computeScatterIOLaneLayoutAndData ( ArrayRef< int64_t > instShape,
int64_t subgroupSize,
int64_t maxChunkSize )
static

Computes lane_layout and lane_data for scatter-style store anchor layouts (store scatter, store matrix).

Lanes and the per-lane vector both live on the innermost dim:

  • laneLayout[innermost] = min(subgroupSize, srcShape[innermost])
  • laneData[innermost] = min(srcShape[innermost] / laneLayout[innermost], maxChunkSize) All other entries are 1.

Definition at line 1028 of file XeGPULayoutImpl.cpp.

Referenced by setupGenericStoreAnchorLayout().

◆ createScaleLayout()

xegpu::DistributeLayoutAttr createScaleLayout ( mlir::MLIRContext * context,
VectorType matrixTy,
VectorType scaleTy,
xegpu::DistributeLayoutAttr matrixLayout,
bool isBScale,
const xegpu::uArch::uArch * uArch )
static

Helper to create a scale layout derived from a matrix operand layout.

The scale layout is computed by mapping each dimension of the matrix layout to the corresponding scale tensor dimension using the ratio between the matrix and scale shapes.

Definition at line 1285 of file XeGPULayoutImpl.cpp.

References buildLayout(), mlir::xegpu::uArch::uArch::getInstruction(), and mlir::xegpu::uArch::SubgroupScaledMatrixMultiplyAcc.

◆ enumerateFactorizations()

SmallVector< LayoutRepresentation > enumerateFactorizations ( int64_t total,
int64_t rank )
static

Enumerates all ways to split total into rank factors whose product equals total.

Returns the list of all such factorizations.

Definition at line 860 of file XeGPULayoutImpl.cpp.

Referenced by getSgLayoutCandidates().

◆ get2DBlockIOInstDataLayout()

std::optional< SmallVector< int64_t > > get2DBlockIOInstDataLayout ( ArrayRef< int64_t > dataShape,
Type elemTy,
const xegpu::uArch::BlockIOInstructionInterface * uArchInstruction,
bool transform = false,
bool transpose = false )
static

Helper function to compute inst_data vectors for DPAS operands A, B, and C/D.

Definition at line 950 of file XeGPULayoutImpl.cpp.

References mlir::xegpu::uArch::BlockIOInstructionInterface::getBlockWidthHeightCount(), and mlir::xegpu::getLargestDivisor().

Referenced by mlir::xegpu::setupStoreNdAnchorLayout().

◆ getDpasInstDataLayouts()

std::optional< std::tuple< SmallVector< int64_t >, SmallVector< int64_t >, SmallVector< int64_t > > > getDpasInstDataLayouts ( VectorType aTy,
VectorType bTy,
VectorType cdTy,
const xegpu::uArch::MMAInstructionInterface * uArchInstruction )
static

Helper function to compute inst_data vectors for DPAS operands A, B, and C/D.

Look up the uArch table and search for the largest supported block size that divides the data shape

Definition at line 981 of file XeGPULayoutImpl.cpp.

References mlir::xegpu::getLargestDivisor(), mlir::xegpu::uArch::MMAInstructionInterface::getSupportedK(), mlir::xegpu::uArch::MMAInstructionInterface::getSupportedM(), and mlir::xegpu::uArch::MMAInstructionInterface::getSupportedN().

◆ getDpasSubgroupLayouts()

std::optional< std::tuple< xegpu::DistributeLayoutAttr, xegpu::DistributeLayoutAttr, xegpu::DistributeLayoutAttr > > getDpasSubgroupLayouts ( mlir::MLIRContext * context,
VectorType aTy,
VectorType bTy,
VectorType cdTy,
xegpu::DistributeLayoutAttr consumerLayout,
int numSg,
std::tuple< SmallVector< int64_t >, SmallVector< int64_t >, SmallVector< int64_t > > instDataVecs )
static

Helper function to set up subgroup layouts for DPAS operands A, B, and C/D.

Compute subgroup layout candidates based on wgtile and instData, and then pick the best one that satisfies all operands and the consumer (if specified).

Definition at line 1168 of file XeGPULayoutImpl.cpp.

References buildSgLayout(), and getSgLayoutCandidates().

◆ getLayoutFromUsePoints()

xegpu::DistributeLayoutAttr getLayoutFromUsePoints ( Value result)
static

Definition at line 116 of file XeGPULayoutImpl.cpp.

References mlir::xegpu::getDistributeLayoutAttr(), and result.

Referenced by propagateRegionResultsToYieldOperands(), propagateResultsToRegularOperands(), and mlir::xegpu::recoverTemporaryLayouts().

◆ getSgLayoutCandidates()

SmallVector< LayoutRepresentation > getSgLayoutCandidates ( ArrayRef< int64_t > wgShape,
ArrayRef< int64_t > instData,
int64_t sgCount )
static

Definition at line 907 of file XeGPULayoutImpl.cpp.

References enumerateFactorizations(), lhs, and rhs.

Referenced by getDpasSubgroupLayouts(), getStoreSubgroupLayouts(), and mlir::xegpu::setupStoreNdAnchorLayout().

◆ getStoreSubgroupLayouts()

xegpu::DistributeLayoutAttr getStoreSubgroupLayouts ( mlir::MLIRContext * context,
ArrayRef< int64_t > wgShape,
ArrayRef< int64_t > instData,
int numSg )
static

Picks the subgroup layout for a scatter-style store (store_scatter / store_matrix): the most balanced numSg factorization that divides wgShape with sg_data a multiple of instData.

A store has no consumer.

Definition at line 1759 of file XeGPULayoutImpl.cpp.

References buildSgLayout(), and getSgLayoutCandidates().

Referenced by setupGenericStoreAnchorLayout().

◆ isValidLaneLayout()

bool isValidLaneLayout ( ArrayRef< int64_t > dataShape,
ArrayRef< int64_t > laneLayout,
ArrayRef< int64_t > laneData )
static

Definition at line 398 of file XeGPULayoutImpl.cpp.

Referenced by mlir::xegpu::setupStoreNdAnchorLayout().

◆ leadingDimsAreUnit()

bool leadingDimsAreUnit ( ArrayRef< int64_t > shape,
int numInnerDims )
static

Returns true if every dimension of shape except the innermost numInnerDims is a unit (size-1) dimension.

Definition at line 375 of file XeGPULayoutImpl.cpp.

Referenced by mlir::xegpu::setupMultiReductionResultLayout().

◆ propagateRegionResultsToYieldOperands()

void propagateRegionResultsToYieldOperands ( mlir::RegionBranchTerminatorOpInterface yieldOp)
static

Definition at line 188 of file XeGPULayoutImpl.cpp.

References mlir::OperandRange::getBeginOperandIndex(), getLayoutFromUsePoints(), mlir::OperandRange::getType(), mlir::xegpu::setTemporaryLayout(), and setTensorDescLayout().

Referenced by mlir::xegpu::recoverTemporaryLayouts().

◆ propagateResultsToRegularOperands()

void propagateResultsToRegularOperands ( Operation * op)
static

Definition at line 149 of file XeGPULayoutImpl.cpp.

References getLayoutFromUsePoints(), mlir::Operation::getNumResults(), mlir::Operation::getOpOperands(), mlir::Operation::getResult(), mlir::xegpu::inferSourceLayoutFromResultForNonAnchorOp(), result, mlir::xegpu::setTemporaryLayout(), and setTensorDescLayout().

Referenced by mlir::xegpu::recoverTemporaryLayouts().

◆ setTensorDescLayout()

void setTensorDescLayout ( Value val,
xegpu::DistributeLayoutAttr layout )
static

Definition at line 75 of file XeGPULayoutImpl.cpp.

References mlir::Value::getType(), and mlir::Value::setType().

Referenced by propagateRegionResultsToYieldOperands(), and propagateResultsToRegularOperands().

◆ setupGenericLoadAnchorLayout()

xegpu::DistributeLayoutAttr setupGenericLoadAnchorLayout ( xegpu::LayoutKind layoutKind,
mlir::MLIRContext * context,
xegpu::DistributeLayoutAttr consumerLayout,
int maxChunkSize,
ArrayRef< int64_t > resShape,
int subgroupSize )
static

Sets up the anchor layout for load gather and load matrix operation.

load matrix lowers to load gather and 1d block load. All of them share the same layout setup logic.

For Subgroup layout, uses the consumer layout directly.

For InstData layout, takes consumer's inst_data as-is. lane_layout and lane_data are taken from the consumer when present; otherwise the helper derives the standard scatter-style default (subgroupSize lanes on the innermost dim, per-lane vector capped by maxChunkSize).

For Lane layout, lane_layout/lane_data are taken from the consumer when present; otherwise derived from the same default.

Definition at line 1674 of file XeGPULayoutImpl.cpp.

References buildInstDataLayoutWithLane(), buildLaneLayout(), mlir::xegpu::InstData, mlir::xegpu::Lane, and mlir::xegpu::Subgroup.

◆ setupGenericStoreAnchorLayout()

xegpu::DistributeLayoutAttr setupGenericStoreAnchorLayout ( xegpu::LayoutKind layoutKind,
mlir::MLIRContext * context,
int maxChunkSize,
ArrayRef< int64_t > srcShape,
int subgroupSize,
int numSg )
static

Sets up the anchor layout for store scatter and store matrix operation, which share the same logic.

Lane layout comes from computeScatterIOLaneLayoutAndData; inst_data is lane_layout * lane_data.

Definition at line 1771 of file XeGPULayoutImpl.cpp.

References buildInstDataLayoutWithLane(), buildLaneLayout(), computeScatterIOLaneLayoutAndData(), getStoreSubgroupLayouts(), mlir::xegpu::InstData, mlir::xegpu::Lane, and mlir::xegpu::Subgroup.

Referenced by mlir::xegpu::setupStoreMatrixAnchorLayout(), and mlir::xegpu::setupStoreScatterAnchorLayout().

◆ walkRegionBackward()

void walkRegionBackward ( Region & region,
llvm::function_ref< void(Operation *)> visit )
static

Definition at line 88 of file XeGPULayoutImpl.cpp.

References mlir::Region::empty(), visit(), and walkRegionBackward().

Referenced by mlir::xegpu::recoverTemporaryLayouts(), and walkRegionBackward().

◆ xegpu::removeLayoutAttr< mlir::OpOperand >()

template void xegpu::removeLayoutAttr< mlir::OpOperand > ( const mlir::OpOperand & operand)

References mlir::xegpu::removeLayoutAttr().

◆ xegpu::removeLayoutAttr< mlir::OpResult >()

template void xegpu::removeLayoutAttr< mlir::OpResult > ( const mlir::OpResult & result)

References mlir::xegpu::removeLayoutAttr(), and result.

Generated on for MLIR by doxygen 1.14.0