Classes
struct	UnrollOptions
	Options to control the XeGPU unrolling. More...

Functions
void	populateXeGPUFoldAliasOpsPatterns (RewritePatternSet &patterns)
	Appends patterns for folding aliasing ops into XeGPU ops into `patterns`. More...

void	populateXeGPUSubgroupDistributePatterns (RewritePatternSet &patterns)
	Appends patterns for XeGPU SIMT distribution into `patterns`. More...

void	populateXeGPUWgToSgDistributePatterns (RewritePatternSet &patterns)

void	populateXeGPUUnrollPatterns (RewritePatternSet &patterns, const UnrollOptions &options)
	Collect a set of patterns to unroll xegpu operations to a smaller shapes. More...

SmallVector< Value >	flattenValues (ArrayRef< ValueRange > values)
	Flatten a set of ValueRange into a single SmallVector<Value> More...

FailureOr< VectorType >	getDistributedVectorType (xegpu::TensorDescType tdescTy)
	If tensor descriptor has a layout attribute it is used in SIMT mode. More...

FailureOr< VectorType >	getDistributedVectorType (VectorType originalType, LayoutAttr layout)
	Helper to get the distributed vector type for a given vector type according to a given LayoutAttr. More...

std::string	getLayoutName (const OpOperand &operand)
	Return the attribute name for the OpOperand to attach LayoutAttr. More...

std::string	getLayoutName (const OpResult result)
	Return the attribute name for the OpResult to attach LayoutAttr. More...

LayoutAttr	getLayoutAttr (const Value value)
	Retrieves the LayoutAttr associated with a given Value. More...

LayoutAttr	getLayoutAttr (const OpOperand &opr)
	Retrieves the LayoutAttr associated with a given OpOperand. More...

template<typename T , typename = std::enable_if_t<std::is_same_v<T, OpOperand> \|\| std::is_same_v<T, OpResult>>>
void	setLayoutAttr (const T &operandOrResult, const LayoutAttr layout)
	Sets the LayoutAttr for a given OpOperand or OpResult by attaching it to the owner's dictionary attributes. More...

void	setLayoutAttrs (Operation *op, function_ref< LayoutAttr(Value)> getLayoutImpl)
	Set the LayoutAttr for each OpOperand and OpResult of the given operation. More...

SmallVector< Value >	extractVectorsWithShapeFromValue (OpBuilder &builder, Location loc, Value value, ArrayRef< int64_t > shape)
	Extract a set of small vectors from a value with a given shape using vector.extract_stride_slice. More...

Value	createVectorWithShapeFromValues (OpBuilder &builder, Location loc, ValueRange values, ArrayRef< int64_t > shape)
	Create a vector of shape from a set of values using vector.insert_stride_slice. More...

void	doSCFStructuralTypeConversionWithTensorType (Operation *op, TypeConverter converter)
	Do type conversion for SCF structural ops, e.g., scf.for using SCF structure type convertion patterns. More...

static void	transpose (llvm::ArrayRef< int64_t > trans, SmallVector< int64_t > &shape)

template<typename T >
static std::string	makeString (T array, bool breakline=false)

static SmallVector< int64_t >	getShapeOf (Type type)

static int64_t	getRankOf (Value val)

static bool	isReadHintOrNone (const CachePolicyAttr &attr)

static bool	isWriteHintOrNone (const CachePolicyAttr &attr)

static LogicalResult	isValidGatherScatterParams (Type maskTy, VectorType valueTy, TensorDescType tdescTy, UnitAttr transposeAttr, function_ref< InFlightDiagnostic()> emitError)

Function Documentation

◆ createVectorWithShapeFromValues()

Value mlir::xegpu::createVectorWithShapeFromValues	(	OpBuilder &	builder,
		Location	loc,
		ValueRange	values,
		ArrayRef< int64_t >	shape
	)

Create a vector of shape from a set of values using vector.insert_stride_slice.

Definition at line 208 of file XeGPUUtils.cpp.

References mlir::OpBuilder::create(), mlir::get(), mlir::DenseElementsAttr::get(), mlir::getType(), mlir::ValueRange::getTypes(), and mlir::Builder::getZeroAttr().

◆ doSCFStructuralTypeConversionWithTensorType()

void mlir::xegpu::doSCFStructuralTypeConversionWithTensorType	(	Operation *	op,
		TypeConverter	converter
	)

Do type conversion for SCF structural ops, e.g., scf.for using SCF structure type convertion patterns.

Since VectorType cannot carry the layout attribute, which is needed to guide the type conversion for XeGPU, they are first converted into RankedTensorType, where the layout attribute can be attached. And then upstream SCF structural type conversion patterns are applied with the provided converter. TODO: This is a temporary solution. We should refactor it when context-aware type conversion is available.

Definition at line 233 of file XeGPUUtils.cpp.

References mlir::TypeConverter::addConversion(), mlir::ConversionTarget::addLegalOp(), mlir::TypeConverter::addSourceMaterialization(), mlir::TypeConverter::addTargetMaterialization(), mlir::WalkResult::advance(), mlir::applyPartialConversion(), mlir::OpBuilder::create(), mlir::get(), mlir::Operation::getContext(), getLayoutAttr(), mlir::Operation::getParentOp(), mlir::Value::getType(), mlir::Value::getUses(), mlir::patterns, mlir::scf::populateSCFStructuralTypeConversionsAndLegality(), mlir::Value::setType(), mlir::WalkResult::skip(), and mlir::Operation::walk().

◆ extractVectorsWithShapeFromValue()

SmallVector< Value > mlir::xegpu::extractVectorsWithShapeFromValue	(	OpBuilder &	builder,
		Location	loc,
		Value	value,
		ArrayRef< int64_t >	shape
	)

Extract a set of small vectors from a value with a given shape using vector.extract_stride_slice.

Definition at line 188 of file XeGPUUtils.cpp.

References mlir::computeShapeRatio(), mlir::OpBuilder::create(), and mlir::Value::getType().

◆ flattenValues()

SmallVector< Value > mlir::xegpu::flattenValues ( ArrayRef< ValueRange > values )

Flatten a set of ValueRange into a single SmallVector<Value>

convert ArrayRef<ValueRange> into SmallVector<Value>

Definition at line 30 of file XeGPUUtils.cpp.

◆ getDistributedVectorType() [1/2]

FailureOr<VectorType> mlir::xegpu::getDistributedVectorType	(	VectorType	originalType,
		LayoutAttr	layout
	)

Helper to get the distributed vector type for a given vector type according to a given LayoutAttr.

◆ getDistributedVectorType() [2/2]

FailureOr< VectorType > mlir::xegpu::getDistributedVectorType ( xegpu::TensorDescType tdescTy )

If tensor descriptor has a layout attribute it is used in SIMT mode.

In this mode, the distributed vector shape is determined as follows: Definitions: lane_data_size = lane_data[0] × lane_data[1] subgroup_size = lane_layout[0] × lane_layout[1] distribution_unit_size = subgroup_size × lane_data_size

Case 1: Regular loads/stores. The following conditions must be met:

tensor_desc[0] == lane_layout[0] Distributed vector is a 1D vector with shape: [chunk_size]

Case 2: Block loads/stores Additional definitions: tensor_size = tensor_desc[0] * .. * tensor_desc[r-1] * array_length n_distribution_units = tensor_size / distribution_unit_size fragment_size = n_distribution_units * lane_data_size Given above definitions, the following conditions must be met:

tensor_desc[0] % (lane_layout[0] × lane_data[0]) == 0
tensor_desc[1] % (lane_layout[1] × lane_data[1]) == 0 Distributed vector is a 1D vector with shape: [fragment_size]

Definition at line 38 of file XeGPUUtils.cpp.

References mlir::get().

◆ getLayoutAttr() [1/2]

xegpu::LayoutAttr mlir::xegpu::getLayoutAttr ( const OpOperand & opr )

Retrieves the LayoutAttr associated with a given OpOperand.

It will first check the operand_layout_{id} of the owner operation. If not found, it will check the operand itself and its defining op.

Definition at line 147 of file XeGPUUtils.cpp.

References mlir::IROperand< DerivedT, IRValueT >::get(), mlir::Operation::getAttrOfType(), getLayoutAttr(), getLayoutName(), mlir::detail::IROperandBase::getOwner(), and mlir::Operation::hasAttr().

◆ getLayoutAttr() [2/2]

xegpu::LayoutAttr mlir::xegpu::getLayoutAttr ( const Value value )

Retrieves the LayoutAttr associated with a given Value.

For TensorDescType values, the LayoutAttr is extracted from the TensorDescType itself. For other values, it is obtained from the attributes of the defining operation. Returns nullptr if no LayoutAttr is found.

Definition at line 115 of file XeGPUUtils.cpp.

References mlir::IROperand< DerivedT, IRValueT >::get(), mlir::Operation::getAttrOfType(), getLayoutName(), mlir::Value::getType(), and mlir::Operation::hasAttr().

Referenced by doSCFStructuralTypeConversionWithTensorType(), and getLayoutAttr().

◆ getLayoutName() [1/2]

std::string mlir::xegpu::getLayoutName ( const OpOperand & operand )

Return the attribute name for the OpOperand to attach LayoutAttr.

Definition at line 104 of file XeGPUUtils.cpp.

Referenced by getLayoutAttr(), and setLayoutAttr().

◆ getLayoutName() [2/2]

std::string mlir::xegpu::getLayoutName ( const OpResult result )

Return the attribute name for the OpResult to attach LayoutAttr.

Definition at line 110 of file XeGPUUtils.cpp.

References mlir::OpResult::getResultNumber().

◆ getRankOf()

static int64_t mlir::xegpu::getRankOf ( Value val )

static

Definition at line 54 of file XeGPUOps.cpp.

References mlir::Value::getType().

◆ getShapeOf()

static SmallVector<int64_t> mlir::xegpu::getShapeOf ( Type type )

static

Definition at line 45 of file XeGPUOps.cpp.

Referenced by isValidGatherScatterParams().

◆ isReadHintOrNone()

static bool mlir::xegpu::isReadHintOrNone ( const CachePolicyAttr & attr )

static

Definition at line 61 of file XeGPUOps.cpp.

References kind.

◆ isValidGatherScatterParams()

static LogicalResult mlir::xegpu::isValidGatherScatterParams	(	Type	maskTy,
		VectorType	valueTy,
		TensorDescType	tdescTy,
		UnitAttr	transposeAttr,
		function_ref< InFlightDiagnostic()>	emitError
	)

static

Definition at line 78 of file XeGPUOps.cpp.

References mlir::emitError(), getShapeOf(), makeString(), and transpose().

◆ isWriteHintOrNone()

static bool mlir::xegpu::isWriteHintOrNone ( const CachePolicyAttr & attr )

static

Definition at line 69 of file XeGPUOps.cpp.

References kind.

◆ makeString()

template<typename T >

static std::string mlir::xegpu::makeString	(	T	array,
		bool	breakline = `false`
	)

static

Definition at line 31 of file XeGPUOps.cpp.

Referenced by isValidGatherScatterParams().

◆ populateXeGPUFoldAliasOpsPatterns()

void mlir::xegpu::populateXeGPUFoldAliasOpsPatterns ( RewritePatternSet & patterns )

Appends patterns for folding aliasing ops into XeGPU ops into patterns.

Definition at line 65 of file XeGPUFoldAliasOps.cpp.

References mlir::patterns.

◆ populateXeGPUSubgroupDistributePatterns()

void mlir::xegpu::populateXeGPUSubgroupDistributePatterns ( RewritePatternSet & patterns )

Appends patterns for XeGPU SIMT distribution into patterns.

Definition at line 1576 of file XeGPUSubgroupDistribute.cpp.

References mlir::patterns.

◆ populateXeGPUUnrollPatterns()

void mlir::xegpu::populateXeGPUUnrollPatterns	(	RewritePatternSet &	patterns,
		const UnrollOptions &	options
	)

Collect a set of patterns to unroll xegpu operations to a smaller shapes.

Users can control whether an operation to be unrolled or not, as well as its target shape via options structure. (via setting filterConstraint and nativeShape respectively, both of them are function refs taking op as input). An op is unrolled to the targetShape as follows, for each of its operands:

the unrolled type unrolledType and number of unrolled instances numUnrolledInstances are computed from the targetShape.
pack each operand. ExtractStridedSlice are created to break-up the vector operands. And BuiltinUnrealizedCastop are created to break-up the TensorDesc operands.
the original op is cloned numUnrolledInstances times, once for each result.
unpack the results. InsertStridedSlice are inserted for VectorType result, and BuiltinUnrealizedCastOp are inserted for TensorDescType result to re-assemble the slices into the original shape.

Definition at line 602 of file XeGPUUnroll.cpp.

References options, and mlir::patterns.

◆ populateXeGPUWgToSgDistributePatterns()

void mlir::xegpu::populateXeGPUWgToSgDistributePatterns ( RewritePatternSet & patterns )

Definition at line 411 of file XeGPUWgToSgDistribute.cpp.

References mlir::patterns.

◆ setLayoutAttr()

template<typename T , typename = std::enable_if_t<std::is_same_v<T, OpOperand> || std::is_same_v<T, OpResult>>>

void mlir::xegpu::setLayoutAttr	(	const T &	operandOrResult,
		const LayoutAttr	layout
	)

Sets the LayoutAttr for a given OpOperand or OpResult by attaching it to the owner's dictionary attributes.

Definition at line 156 of file XeGPUUtils.cpp.

References getLayoutName(), mlir::Operation::hasAttrOfType(), and mlir::Operation::setAttr().

◆ setLayoutAttrs()

void mlir::xegpu::setLayoutAttrs	(	Operation *	op,
		function_ref< LayoutAttr(Value)>	getLayoutImpl
	)

Set the LayoutAttr for each OpOperand and OpResult of the given operation.

If the operation contains regions, it is also applied recursively to the contained operations

Definition at line 173 of file XeGPUUtils.cpp.

References mlir::Operation::getOpOperands(), mlir::Operation::getOpResults(), and mlir::Operation::walk().

◆ transpose()

static void mlir::xegpu::transpose	(	llvm::ArrayRef< int64_t >	trans,
		SmallVector< int64_t > &	shape
	)

static

Definition at line 23 of file XeGPUOps.cpp.

Referenced by creatLdMatrixCompatibleLoads(), mlir::nvgpu::getLdMatrixParams(), isValidGatherScatterParams(), propagatesCapture(), and pruneNonTransposedDims().

Classes

Functions

Function Documentation

◆ createVectorWithShapeFromValues()

◆ doSCFStructuralTypeConversionWithTensorType()

◆ extractVectorsWithShapeFromValue()

◆ flattenValues()

◆ getDistributedVectorType() [1/2]

◆ getDistributedVectorType() [2/2]

◆ getLayoutAttr() [1/2]

◆ getLayoutAttr() [2/2]

◆ getLayoutName() [1/2]

◆ getLayoutName() [2/2]

◆ getRankOf()

◆ getShapeOf()

◆ isReadHintOrNone()

◆ isValidGatherScatterParams()

◆ isWriteHintOrNone()

◆ makeString()

◆ populateXeGPUFoldAliasOpsPatterns()

◆ populateXeGPUSubgroupDistributePatterns()

◆ populateXeGPUUnrollPatterns()

◆ populateXeGPUWgToSgDistributePatterns()

◆ setLayoutAttr()

◆ setLayoutAttrs()

◆ transpose()