#include "mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h"
#include "mlir/Dialect/Arith/IR/Arith.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/IRMapping.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/DenseSet.h"
#include "llvm/ADT/PostOrderIterator.h"
#include "llvm/Support/FormatVariadic.h"
#include <cstdint>
#include <functional>
#include <numeric>

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 bool	isLayoutPreservingVectorProducer (Operation *op)
	Accept pure vector generators and elementwise ops that preserve vector shape.
static bool	collectRematerializableChain (Operation root, SmallVectorImpl< Operation > &chain)
	Collect each eligible producer once, producers first, for cloning with an IRMapping.
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 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, int64_t broadcastDim=-1)
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 xegpu::DistributeLayoutAttr	get2DBlockIOSubgroupLayout (MLIRContext context, ArrayRef< int64_t > dataShape, Type elemTy, const xegpu::uArch::BlockIOInstructionInterface uArchInstruction, int numSg)
	Picks the subgroup layout for a 2D block IO op (store_nd / prefetch_nd).
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::tuple< SmallVector< int64_t >, SmallVector< int64_t >, SmallVector< int64_t > >	compute2DBlockIOLaneLayout (ArrayRef< int64_t > instShape, int64_t subgroupSize, int64_t bitwidth, int64_t packingSize, bool transform=false, bool transpose=false)
static std::tuple< SmallVector< int64_t >, SmallVector< int64_t >, SmallVector< int64_t > >	computeReductionLaneLayoutAndData (ArrayRef< int64_t > srcShape, ArrayRef< int64_t > reductionDims, int subgroupSize, int64_t maxReduceVectorSize, xegpu::DistributeLayoutAttr consumerLayout=nullptr)
	Derives (lane_layout, lane_data, inst_data) for a vector.multi_reduction source from the layout requested by its result's consumer.
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, const xegpu::uArch::MMAInstructionInterface uArchInstruction)
	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.
static bool	isContiguousTiling (ArrayRef< int64_t > dimGroup, ArrayRef< int64_t > shape, ArrayRef< int64_t > counts, ArrayRef< int64_t > tileShape, int64_t expectedStride)
	Checks that the tiles a split group is cut into walk the collapsed source dim with one constant stride, which is what collapseDims assumes when it multiplies a layout field across the group.
static xegpu::DistributeLayoutAttr	getLoopCarriedLayoutForYieldOperand (RegionBranchTerminatorOpInterface terminator, OpOperand &operand)
static xegpu::DistributeLayoutAttr	getParentResultLayoutForYieldOperand (RegionBranchTerminatorOpInterface terminator, OpOperand &operand)

Variables
static constexpr unsigned	kMaxRematerializableChainOps = 16

Typedef Documentation

◆ LayoutRepresentation

using LayoutRepresentation = SmallVector<int64_t>

Definition at line 1085 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 2810 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 501 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 524 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 538 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 554 of file XeGPULayoutImpl.cpp.

References buildLayout().

Referenced by get2DBlockIOSubgroupLayout(), getDpasSubgroupLayouts(), and getStoreSubgroupLayouts().

◆ collectRematerializableChain()

bool collectRematerializableChain	(	Operation *	root,
		SmallVectorImpl< Operation * > &	chain )

static

Collect each eligible producer once, producers first, for cloning with an IRMapping.

Reject chains exceeding the duplication limit or reaching an unsupported producer or vector block argument.

Definition at line 163 of file XeGPULayoutImpl.cpp.

References isLayoutPreservingVectorProducer(), and kMaxRematerializableChainOps.

Referenced by mlir::xegpu::isTriviallyRematerializable(), and mlir::xegpu::rematerialize().

◆ compute2DBlockIOLaneLayout()

std::tuple< SmallVector< int64_t >, SmallVector< int64_t >, SmallVector< int64_t > > compute2DBlockIOLaneLayout	(	ArrayRef< int64_t >	instShape,
		int64_t	subgroupSize,
		int64_t	bitwidth,
		int64_t	packingSize,
		bool	transform = false,
		bool	transpose = false )

static

Definition at line 1343 of file XeGPULayoutImpl.cpp.

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

◆ computeReductionLaneLayoutAndData()

std::tuple< SmallVector< int64_t >, SmallVector< int64_t >, SmallVector< int64_t > > computeReductionLaneLayoutAndData	(	ArrayRef< int64_t >	srcShape,
		ArrayRef< int64_t >	reductionDims,
		int	subgroupSize,
		int64_t	maxReduceVectorSize,
		xegpu::DistributeLayoutAttr	consumerLayout = nullptr )

static

Derives (lane_layout, lane_data, inst_data) for a vector.multi_reduction source from the layout requested by its result's consumer.

Each returned vector has one entry per source dimension.

Reuse the consumer's slice parent when the slice dimensions exactly match reductionDims, the parent's effective lane_layout and lane_data each have one entry per source dimension, and the lane layout uses all subgroupSize lanes. Clamp reused lane counts to srcShape[i] / lane_data[i], since the sliced dimensions may have different extents in the reduction source.

If the consumer's slice parent cannot be reused, initialize lane_layout and lane_data to 1, then copy available consumer entries onto non-reduced source dimensions in order, with the same clamp: a consumer that broadcasts along a unit dimension may place lanes on it.

For a derived (not reused) layout, assign any unused lanes to the innermost dimension if it is reduced, up to the same bound. If that dimension still has only one lane, then the lane can reduce a vector limited by maxReduceVectorSize and the source dimension's size. The vector budget (maxReduceVectorSize) applies to the product of lane_data across all dimensions. If that product already meets or exceeds the budget, leave lane_data unchanged.

Set inst_data[i] to lane_layout[i] * lane_data[i]. On non-reduced dimensions, preserve a reused parent's larger inst_data[i].

Source layout examples with subgroupSize = maxReduceVectorSize = 16:

// Reuse a matching slice parent's lane layout and data.
srcShape=[32, 64], reductionDims=[1], consumer=slice(parent, dims=[1])
parent: lane_layout=[1, 16], lane_data=[1, 1]
source: lane_layout=[1, 16], lane_data=[1, 1], inst_data=[1, 16]
 
// Preserve dim 0 blocking and limit dim 1 lanes to the source size.
srcShape=[16, 4], reductionDims=[1], consumer=slice(parent, dims=[1])
parent: lane_layout=[1, 16], lane_data=[1, 1], inst_data=[8, 16]
source: lane_layout=[1, 4], lane_data=[1, 1], inst_data=[8, 4]
 
// Assign spare lanes to the innermost dimension when it is reduced.
srcShape=[16, 32, 32], reductionDims=[2]
consumer: lane_layout=[1, 2], lane_data=[1, 1]
source:   lane_layout=[1, 2, 8], lane_data=[1, 1, 1], inst_data=[1, 2, 8]
 
// Keep the consumer's lanes on the non-reduced dimension.
srcShape=[32, 64], reductionDims=[0]
consumer: lane_layout=[16], lane_data=[1]
source:   lane_layout=[1, 16], lane_data=[1, 1], inst_data=[1, 16]
 
// For a scalar reduction, put lanes on the innermost dimension.
srcShape=[32, 64], reductionDims=[0, 1]
consumer: none
source:   lane_layout=[1, 16], lane_data=[1, 1], inst_data=[1, 16]

Definition at line 1428 of file XeGPULayoutImpl.cpp.

References mlir::computeProduct().

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 1322 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 1727 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 1089 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 1191 of file XeGPULayoutImpl.cpp.

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

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

◆ get2DBlockIOSubgroupLayout()

xegpu::DistributeLayoutAttr get2DBlockIOSubgroupLayout	(	MLIRContext *	context,
		ArrayRef< int64_t >	dataShape,
		Type	elemTy,
		const xegpu::uArch::BlockIOInstructionInterface *	uArchInstruction,
		int	numSg )

static

Picks the subgroup layout for a 2D block IO op (store_nd / prefetch_nd).

Every supported block (height, width) dividing the shape is tried, largest area first and the more square block first among equal areas, until one admits a subgroup layout.

Definition at line 1225 of file XeGPULayoutImpl.cpp.

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

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 1275 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,
		const xegpu::uArch::MMAInstructionInterface *	uArchInstruction )

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 1593 of file XeGPULayoutImpl.cpp.

References buildSgLayout(), getSgLayoutCandidates(), and mlir::xegpu::uArch::MMAInstructionInterface::getSupportedM().

◆ getLayoutFromUsePoints()

xegpu::DistributeLayoutAttr getLayoutFromUsePoints ( Value result )

static

Definition at line 127 of file XeGPULayoutImpl.cpp.

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

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

◆ getLoopCarriedLayoutForYieldOperand()

xegpu::DistributeLayoutAttr getLoopCarriedLayoutForYieldOperand	(	RegionBranchTerminatorOpInterface	terminator,
		OpOperand &	operand )

static

Definition at line 3208 of file XeGPULayoutImpl.cpp.

References mlir::xegpu::getDistributeLayoutAttr().

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

◆ getParentResultLayoutForYieldOperand()

xegpu::DistributeLayoutAttr getParentResultLayoutForYieldOperand	(	RegionBranchTerminatorOpInterface	terminator,
		OpOperand &	operand )

static

Definition at line 3232 of file XeGPULayoutImpl.cpp.

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

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

◆ getSgLayoutCandidates()

SmallVector< LayoutRepresentation > getSgLayoutCandidates	(	ArrayRef< int64_t >	wgShape,
		ArrayRef< int64_t >	instData,
		int64_t	sgCount,
		int64_t	broadcastDim = -1 )

static

Definition at line 1141 of file XeGPULayoutImpl.cpp.

References enumerateFactorizations().

Referenced by get2DBlockIOSubgroupLayout(), getDpasSubgroupLayouts(), and getStoreSubgroupLayouts().

◆ 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 2233 of file XeGPULayoutImpl.cpp.

References buildSgLayout(), and getSgLayoutCandidates().

Referenced by setupGenericStoreAnchorLayout().

◆ isContiguousTiling()

bool isContiguousTiling	(	ArrayRef< int64_t >	dimGroup,
		ArrayRef< int64_t >	shape,
		ArrayRef< int64_t >	counts,
		ArrayRef< int64_t >	tileShape,
		int64_t	expectedStride )

static

Checks that the tiles a split group is cut into walk the collapsed source dim with one constant stride, which is what collapseDims assumes when it multiplies a layout field across the group.

Along each dim the group is cut into counts[dim] tiles of tileShape[dim] elements (empty for single-element tiles) taken from shape[dim], which gives the row-major strides. Walking from the innermost dim out, a dim cut into several tiles steps by tileShape[dim] * stride(dim) and has to pick up where the inner dims left off, expectedStride; a dim cut into one tile does not step at all, so it only contributes its size.

Definition at line 2941 of file XeGPULayoutImpl.cpp.

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

◆ isLayoutPreservingVectorProducer()

bool isLayoutPreservingVectorProducer ( Operation * op )

static

Accept pure vector generators and elementwise ops that preserve vector shape.

Generators have no vector operands; elementwise ops share their result layout with their vector operands, so the chain needs no layout inference.

Definition at line 145 of file XeGPULayoutImpl.cpp.

References broadcast(), mlir::Operation::getNumRegions(), mlir::Operation::getNumResults(), mlir::Operation::getResult(), mlir::Value::getType(), mlir::OpTrait::hasElementwiseMappableTraits(), and mlir::isMemoryEffectFree().

Referenced by collectRematerializableChain().

◆ isValidLaneLayout()

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

static

Definition at line 515 of file XeGPULayoutImpl.cpp.

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

◆ propagateRegionResultsToYieldOperands()

void propagateRegionResultsToYieldOperands ( mlir::RegionBranchTerminatorOpInterface yieldOp )

static

Definition at line 258 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 219 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 86 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.

For Lane layout, lane_layout/lane_data are taken from the consumer.

A consumer layout that carries lane_layout and lane_data is required. maxChunkSize is not read yet: the consumer's lane_data already fixes the per-lane chunk.

TODO: derive lane_layout/lane_data here when the consumer has none, capped by maxChunkSize, the way setupGenericStoreAnchorLayout does via computeScatterIOLaneLayoutAndData. That path is missing today, so the assert below stands in for it.

Definition at line 2152 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 2245 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 99 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.

Variable Documentation

◆ kMaxRematerializableChainOps

unsigned kMaxRematerializableChainOps = 16

staticconstexpr

Definition at line 139 of file XeGPULayoutImpl.cpp.

Referenced by collectRematerializableChain().

Typedefs

Functions

Variables

Typedef Documentation

◆ LayoutRepresentation

Function Documentation

◆ adjustInnermostDimForDivisibility()

◆ buildInstDataLayoutWithLane()

◆ buildLaneLayout()

◆ buildLayout()

◆ buildSgLayout()

◆ collectRematerializableChain()

◆ compute2DBlockIOLaneLayout()

◆ computeReductionLaneLayoutAndData()

◆ computeScatterIOLaneLayoutAndData()

◆ createScaleLayout()

◆ enumerateFactorizations()

◆ get2DBlockIOInstDataLayout()

◆ get2DBlockIOSubgroupLayout()

◆ getDpasInstDataLayouts()

◆ getDpasSubgroupLayouts()

◆ getLayoutFromUsePoints()

◆ getLoopCarriedLayoutForYieldOperand()

◆ getParentResultLayoutForYieldOperand()

◆ getSgLayoutCandidates()

◆ getStoreSubgroupLayouts()

◆ isContiguousTiling()

◆ isLayoutPreservingVectorProducer()

◆ isValidLaneLayout()

◆ propagateRegionResultsToYieldOperands()

◆ propagateResultsToRegularOperands()

◆ setTensorDescLayout()

◆ setupGenericLoadAnchorLayout()

◆ setupGenericStoreAnchorLayout()

◆ walkRegionBackward()

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

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

Variable Documentation

◆ kMaxRematerializableChainOps