#include "mlir/Dialect/XeGPU/IR/XeGPU.h"
#include "mlir/Dialect/XeGPU/Utils/XeGPUUtils.h"
#include "mlir/Dialect/XeGPU/uArch/uArchCommon.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/OpDefinition.h"
#include "mlir/Interfaces/ControlFlowInterfaces.h"
#include "llvm/ADT/STLFunctionalExtras.h"

Namespaces
namespace	mlir
	Include the generated interface declarations.
namespace	mlir::xegpu

Typedefs
using	mlir::xegpu::GetLayoutFnTy = llvm::function_ref<DistributeLayoutAttr(Value)>
	Callable returning the propagated layout for a given Value, used by the layout-propagation helpers below.

Functions
LogicalResult	mlir::xegpu::propagateLayouts (OpBuilder &builder, Operation *target, LayoutKind layoutKind, unsigned indexBitWidth, bool printOnly=false)
LogicalResult	mlir::xegpu::resolveLayoutConflicts (Operation *target)
LogicalResult	mlir::xegpu::propagateRegionArgsToInits (RegionBranchOpInterface regionOp, GetLayoutFnTy getLayoutOfValue)
	Propagate layouts from a region branch op's region entry block arguments back to its init operands.
bool	mlir::xegpu::recoverTemporaryLayouts (Operation *rootOp)
	Attach layout attributes to all vector-type operands of operations within the given operation's nested region.
template<typename T, typename = std::enable_if_t<std::is_same_v<T, OpOperand> \|\| std::is_same_v<T, OpResult>>>
void	mlir::xegpu::removeLayoutAttr (const T &operandOrResult)
	Removes the LayoutAttr for a given OpOperand or OpResult if it exists.
void	mlir::xegpu::removeLayoutAttrs (Operation *op)
	Removes the DistributeLayoutAttr for each OpOperand and OpResult of the given operation if they exist.
void	mlir::xegpu::removeTemporaryLayoutAttrs (Operation *op)
	Removes the temporary layout attributes for each OpOperand and OpResult of the given operation.
SmallVector< NamedAttribute >	mlir::xegpu::dropSgLayoutAndDataOnAttrs (ArrayRef< NamedAttribute > attrs)
	Updates the NamedAttribute sequence by dropping sg-layout and sg-data information from any DistributeLayoutAttr found.
SmallVector< NamedAttribute >	mlir::xegpu::dropInstDataOnAttrs (ArrayRef< NamedAttribute > attrs)
	Updates the NamedAttribute sequence by dropping inst-data information from any DistributeLayoutAttr found.
DistributeLayoutAttr	mlir::xegpu::inferBroadcastSourceLayout (DistributeLayoutAttr resLayout, ArrayRef< int64_t > resShape, ArrayRef< int64_t > srcShape)
	Infers the source layout attribute for a broadcast operation given the result layout attribute, result shape, and source shape.
DistributeLayoutAttr	mlir::xegpu::inferMultiReductionSourceLayout (DistributeLayoutAttr resLayout, SmallVector< int64_t > reduceDims)
	Infers the source layout attribute for a reduction operation given the result layout attribute and reduced dims.
DistributeLayoutAttr	mlir::xegpu::inferReductionSourceLayout (DistributeLayoutAttr resLayout)
	Infers the source layout attribute for a reduction operation given the result layout attribute and reduced dims.
DistributeLayoutAttr	mlir::xegpu::inferTransposeSourceLayout (DistributeLayoutAttr resLayout, ArrayRef< int64_t > permutation)
	Infers the source layout attribute for a transpose operation given the result layout attribute and permutation.
DistributeLayoutAttr	mlir::xegpu::inferBitCastSourceLayout (DistributeLayoutAttr resLayout, int resElemTyBitWidth, int srcElemTyBitWidth)
	Infers the source layout attribute for a bitcast operation given the result layout attribute, result element type bitwidth, and source element type bitwidth.
DistributeLayoutAttr	mlir::xegpu::inferInterleaveSourceLayout (DistributeLayoutAttr resLayout)
	Infers the source layout attribute for an interleave operation given the result layout attribute.
DistributeLayoutAttr	mlir::xegpu::inferDeinterleaveSourceLayout (DistributeLayoutAttr resLayout)
	Infers the source layout attribute for a deinterleave operation given the result layout attribute.
DistributeLayoutAttr	mlir::xegpu::inferShapeCastSourceLayout (DistributeLayoutAttr resLayout, ArrayRef< int64_t > resShape, ArrayRef< int64_t > srcShape)
	Infers the source layout attribute for a shape cast operation given the result layout attribute, result shape, and source shape.
DistributeLayoutAttr	mlir::xegpu::inferInsertStridedSliceSourceLayout (DistributeLayoutAttr resLayout, ArrayRef< int64_t > resShape, ArrayRef< int64_t > srcShape)
	Infers the source layout attribute for an insert strided slice operation given the result layout attribute, result shape, and source shape.
DistributeLayoutAttr	mlir::xegpu::inferInsertSourceLayout (DistributeLayoutAttr resLayout, ArrayRef< int64_t > resShape, ArrayRef< int64_t > srcShape)
	Infers the source layout attribute for an insert operation.
DistributeLayoutAttr	mlir::xegpu::inferExtractSourceLayout (DistributeLayoutAttr resLayout, ArrayRef< int64_t > resShape, ArrayRef< int64_t > srcShape)
	Infers the source layout attribute for an extract operation.
DistributeLayoutAttr	mlir::xegpu::inferMaskOffsetLayoutForScatterIO (DistributeLayoutAttr payloadLayout, int chunkSize)
	Infers the layout attribute for mask and offset operand for Chunked load and store, given the anchor layout attribute for the value being load/store.
DistributeLayoutAttr	mlir::xegpu::inferSourceLayoutFromResultForNonAnchorOp (OpOperand &operand, DistributeLayoutAttr resLayout)
	Infers the source layout attribute for an operand using result layout attribute.
SliceAttr	mlir::xegpu::setupMultiReductionResultLayout (LayoutKind layoutKind, VectorType srcVectorTy, DistributeLayoutAttr consumerLayout, SmallVector< int64_t > reductionDims, int numSg, const uArch::uArch *uArch)
	Note on the consumerLayout argument used by the consumer-driven setup* / complete* helpers below:
SliceAttr	mlir::xegpu::setupReductionResultLayout (LayoutKind layoutKind, VectorType srcVectorTy, const uArch::uArch *uArch)
	Sets up layout for Reduction operations by creating a SliceAttr for the result.
DistributeLayoutAttr	mlir::xegpu::setupBitCastResultLayout (LayoutKind layoutKind, VectorType srcVectorTy, VectorType resVectorTy, DistributeLayoutAttr consumerLayout, const uArch::uArch *uArch)
	Setup the result layout attribute for a bitcast operation based on element type bitwidths.
DistributeLayoutAttr	mlir::xegpu::setupInterleaveResultLayout (LayoutKind layoutKind, VectorType srcVectorTy, VectorType resVectorTy, DistributeLayoutAttr consumerLayout, const uArch::uArch *uArch)
	Sets up the result layout for an interleave operation to ensure the source layout can be safely derived.
DistributeLayoutAttr	mlir::xegpu::setupInsertStridedSliceResultLayout (LayoutKind layoutKind, VectorType srcVectorTy, VectorType resVectorTy, DistributeLayoutAttr consumerLayout, const uArch::uArch *uArch)
	Sets up the result layout for an insert strided slice operation.
DistributeLayoutAttr	mlir::xegpu::setupLoadGatherAnchorLayout (LayoutKind layoutKind, VectorType vectorTy, int contigChunkSize, DistributeLayoutAttr consumerLayout, const uArch::uArch *uArch)
	Sets up the anchor layout for a load gather operation.
DistributeLayoutAttr	mlir::xegpu::setupLoadMatrixAnchorLayout (LayoutKind layoutKind, VectorType vectorTy, int contigChunkSize, DistributeLayoutAttr consumerLayout, const uArch::uArch *uArch)
	Sets up the anchor layout for load matrix operation.
DistributeLayoutAttr	mlir::xegpu::setupStoreScatterAnchorLayout (LayoutKind layoutKind, VectorType vectorTy, int contigChunkSize, const uArch::uArch *uArch)
	Sets up the anchor layout for a store scatter operation.
DistributeLayoutAttr	mlir::xegpu::setupStoreMatrixAnchorLayout (LayoutKind layoutKind, VectorType vectorTy, int contigChunkSize, const uArch::uArch *uArch)
	Sets up the anchor layout for a store matrix operation.
std::optional< DistributeLayoutAttr >	mlir::xegpu::completeScatterLoadLaneLayoutFromInstData (DistributeLayoutAttr userSpecifiedLayout, DistributeLayoutAttr consumerLayout, Type elemTy, const xegpu::uArch::LoadGatherInstruction *uArchInstruction, const int subgroupSize)
	If the consumer layout has only inst_data (no lane_layout/lane_data), completes it by running the corresponding scatter-style Lane-kind setup rule with inst_data as the destination shape.
std::optional< DistributeLayoutAttr >	mlir::xegpu::completeScatterStoreLaneLayoutFromInstData (DistributeLayoutAttr specifiedLayout, Type elemTy, const xegpu::uArch::StoreScatterInstruction *uArchInstruction, const int subgroupSize)
	Like completeScatterLoadLaneLayoutFromInstData, but for scatter stores (store_scatter / store_matrix).
std::optional< DistributeLayoutAttr >	mlir::xegpu::completeBlockStoreLaneLayoutFromInstData (DistributeLayoutAttr specifiedLayout, Type elemTy, const xegpu::uArch::BlockIOInstructionInterface *uArchInstruction, const int subgroupSize)
	Completes a user-provided 2D-block store_nd / prefetch_nd anchor that has only inst_data.
std::optional< DistributeLayoutAttr >	mlir::xegpu::completeBlockLoadLaneLayoutFromInstData (DistributeLayoutAttr specifiedLayout, DistributeLayoutAttr consumerLayout, Type elemTy, const xegpu::uArch::BlockIOInstructionInterface *uArchInstruction, const int subgroupSize)
	Like completeBlockStoreLaneLayoutFromInstData, but for load_nd.
DistributeLayoutAttr	mlir::xegpu::setupStoreNdAnchorLayout (LayoutKind layoutKind, VectorType vectorTy, int numSg, const uArch::uArch *uArch)
	Sets up the anchor layout for a store_nd operation.
DistributeLayoutAttr	mlir::xegpu::setupPrefetchNdAnchorLayout (LayoutKind layoutKind, TensorDescType tdescTy, int numSg, const uArch::uArch *uArch)
	Sets up the anchor layout for a prefetch_nd operation.
DistributeLayoutAttr	mlir::xegpu::setupLoadNdAnchorLayout (LayoutKind layoutKind, VectorType vectorTy, DistributeLayoutAttr consumerLayout, int numSg, const uArch::uArch *uArch)
	Sets up the anchor layout for a load_nd operation.
std::optional< std::tuple< DistributeLayoutAttr, DistributeLayoutAttr, DistributeLayoutAttr > >	mlir::xegpu::setupDpasLayout (LayoutKind layoutKind, VectorType aTy, VectorType bTy, VectorType cdTy, DistributeLayoutAttr consumerLayout, int numSg, const uArch::uArch *uArch)
	Sets up the anchor layouts for a dpas operands (A, B, and C/D).
std::optional< std::tuple< DistributeLayoutAttr, DistributeLayoutAttr, DistributeLayoutAttr, DistributeLayoutAttr, DistributeLayoutAttr > >	mlir::xegpu::setupDpasMxLayout (LayoutKind layoutKind, VectorType aTy, VectorType bTy, VectorType cdTy, VectorType aScaleTy, VectorType bScaleTy, DistributeLayoutAttr consumerLayout, int numSg, const uArch::uArch *uArch)
	Sets up the anchor layouts for dpas_mx operands (A, B, C/D, A_scale, and B_scale).
std::optional< std::tuple< DistributeLayoutAttr, DistributeLayoutAttr, DistributeLayoutAttr > >	mlir::xegpu::completeDpasLaneLayoutFromInstData (DistributeLayoutAttr aLayout, DistributeLayoutAttr bLayout, DistributeLayoutAttr cdLayout, VectorType aTy, VectorType bTy, VectorType cdTy, const uArch::uArch *uArch)
	Completes user-provided DPAS A/B/C-D anchors that carry only inst_data by filling in lane_layout / lane_data derived from the operand shapes (mirrors the InstData branch of setupDpasLayout).
std::optional< std::tuple< DistributeLayoutAttr, DistributeLayoutAttr, DistributeLayoutAttr, DistributeLayoutAttr, DistributeLayoutAttr > >	mlir::xegpu::completeDpasMxLaneLayoutFromInstData (DistributeLayoutAttr aLayout, DistributeLayoutAttr bLayout, DistributeLayoutAttr cdLayout, VectorType aTy, VectorType bTy, VectorType cdTy, VectorType aScaleTy, VectorType bScaleTy, const uArch::uArch *uArch)
	Like completeDpasLaneLayoutFromInstData, but for dpas_mx: additionally re-derives the A_scale / B_scale layouts from the completed A / B layouts.
DistributeLayoutAttr	mlir::xegpu::getConsumerLayoutAt (OpOperand &operand)
	Gets the expected layout for a given consumer operand.
bool	mlir::xegpu::isTriviallyRematerializable (Operation *op)
	Returns true if op is safe and cheap to clone: it has no side effects, no regions, and all of its operands are themselves trivially rematerializable (e.g.

Namespaces

Typedefs

Functions