MLIR  22.0.0git
Namespaces | Classes | Functions
mlir::xegpu Namespace Reference

Namespaces

 targetinfo
 HW dependent constants.
 

Classes

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

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< ValueflattenValues (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 DistributeLayoutAttr. More...
 
std::string getLayoutName (const OpResult result)
 Return the attribute name for the OpResult to attach DistributeLayoutAttr. More...
 
DistributeLayoutAttr getDistributeLayoutAttr (const Value value)
 Retrieves the DistributeLayoutAttr associated with a given Value. More...
 
template<typename AttrTy >
AttrTy getDistributeLayoutAttrOfType (const Value value)
 
DistributeLayoutAttr getDistributeLayoutAttr (const OpOperand &opr)
 Retrieves the DistributeLayoutAttr associated with a given OpOperand. More...
 
template<typename AttrTy >
AttrTy getDistributeLayoutAttrOfType (const OpOperand &opr)
 
template<typename T , typename = std::enable_if_t<std::is_same_v<T, OpOperand> || std::is_same_v<T, OpResult>>>
void removeLayoutAttr (const T &operandOrResult)
 Removes the LayoutAttr for a given OpOperand or OpResult if it exists. More...
 
void removeLayoutAttrs (Operation *op)
 Removes the DistributeLayoutAttr for each OpOperand and OpResult of the given operation if they exist. More...
 
template<typename T , typename = std::enable_if_t<std::is_same_v<T, OpOperand> || std::is_same_v<T, OpResult>>>
void setDistributeLayoutAttr (const T &operandOrResult, const DistributeLayoutAttr layout)
 Sets the DistributeLayoutAttr for a given OpOperand or OpResult by attaching it to the owner's dictionary attributes. More...
 
void setDistributeLayoutAttrs (Operation *op, function_ref< DistributeLayoutAttr(Value)> getLayoutImpl)
 Set the DistributeLayoutAttr for each OpOperand and OpResult of the given operation. More...
 
SmallVector< ValueextractVectorsWithShapeFromValue (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...
 
std::optional< std::string > getChipStr (Operation *op)
 Retrieves the chip string from the XeVM target attribute of the parent GPU module operation. More...
 
SmallVector< OpFoldResultaddWithRightAligned (OpBuilder &builder, Location loc, ArrayRef< OpFoldResult > lhs, ArrayRef< OpFoldResult > rhs)
 Generates element-wise addition ops of two arrays with automatic alignment. More...
 
static SmallVector< SmallVector< Value > > genOffsetsComputingInsts (OpBuilder &builder, Location loc, SmallVector< Value > sgId, ArrayRef< int64_t > sgLayout, ArrayRef< int64_t > sizePerSg, ArrayRef< int64_t > sizePerWg)
 Generates instructions to compute offsets for a subgroup identified by its multidimensional indices (sgId), using the specified subgroup layout (sgLayout), subgroup data dimensions (sizePerSg), and the overall data dimensions (sizePerWg). More...
 
bool isSharedMemory (const MemRefType &memrefTy)
 
template<typename T >
static std::string makeString (T array, bool breakline=false)
 
static SmallVector< int64_t > getShapeOf (Type type)
 
static bool isReadHintOrNone (const CachePolicyAttr &attr)
 
static bool isWriteHintOrNone (const CachePolicyAttr &attr)
 
static LogicalResult isValidGatherScatterParams (Type maskTy, VectorType valueTy, TensorDescType tdescTy, function_ref< InFlightDiagnostic()> emitError)
 
static LogicalResult isValidGatherScatterBufferParams (Type offsetsTy, Type maskTy, VectorType valueTy, int64_t chunkSize, function_ref< InFlightDiagnostic()> emitError)
 
ParseResult parseOptionalDynamicIndexList (OpAsmParser &parser, SmallVectorImpl< OpAsmParser::UnresolvedOperand > &values, DenseI64ArrayAttr &integers, SmallVectorImpl< Type > *valueTypes=nullptr, AsmParser::Delimiter delimiter=AsmParser::Delimiter::Square)
 
void printOptionalDynamicIndexList (OpAsmPrinter &printer, Operation *op, OperandRange values, DenseI64ArrayAttr integers)
 

Function Documentation

◆ addWithRightAligned()

SmallVector< OpFoldResult > mlir::xegpu::addWithRightAligned ( OpBuilder builder,
Location  loc,
ArrayRef< OpFoldResult lhs,
ArrayRef< OpFoldResult rhs 
)

Generates element-wise addition ops of two arrays with automatic alignment.

When the input arrays have different sizes, the shorter array is right-aligned with the longer array, and the unmatched leading elements from the longer array are preserved unchanged. This is commonly used for offset computation where higher-dimensional offsets need to be added to lower-dimensional adjustments.

Example: lhs = [l1, l2, l3], rhs = [r1, r2] Result: [11, l2+r1, l3+r2]

Definition at line 443 of file XeGPUUtils.cpp.

References mlir::OpBuilder::createOrFold(), and mlir::getValueOrCreateConstantIndexOp().

◆ 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 242 of file XeGPUUtils.cpp.

References 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 267 of file XeGPUUtils.cpp.

References mlir::TypeConverter::addConversion(), mlir::ConversionTarget::addDynamicallyLegalOp(), mlir::ConversionTarget::addLegalOp(), mlir::TypeConverter::addSourceMaterialization(), mlir::TypeConverter::addTargetMaterialization(), mlir::WalkResult::advance(), mlir::applyPartialConversion(), flattenValues(), mlir::get(), mlir::Operation::getContext(), getDistributeLayoutAttr(), mlir::Operation::getOperandTypes(), mlir::Operation::getOpResults(), mlir::Operation::getParentOp(), mlir::Operation::getResultTypes(), mlir::Value::getType(), mlir::Value::getUses(), mlir::patterns, mlir::scf::populateSCFStructuralTypeConversionsAndLegality(), mlir::ConversionPatternRewriter::replaceOp(), mlir::ConversionPatternRewriter::replaceOpWithMultiple(), 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 222 of file XeGPUUtils.cpp.

References mlir::computeShapeRatio(), 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 32 of file XeGPUUtils.cpp.

Referenced by doSCFStructuralTypeConversionWithTensorType().

◆ genOffsetsComputingInsts()

static SmallVector<SmallVector<Value> > mlir::xegpu::genOffsetsComputingInsts ( OpBuilder builder,
Location  loc,
SmallVector< Value sgId,
ArrayRef< int64_t >  sgLayout,
ArrayRef< int64_t >  sizePerSg,
ArrayRef< int64_t >  sizePerWg 
)
static

Generates instructions to compute offsets for a subgroup identified by its multidimensional indices (sgId), using the specified subgroup layout (sgLayout), subgroup data dimensions (sizePerSg), and the overall data dimensions (sizePerWg).

Definition at line 45 of file XeGPUDialect.cpp.

References mlir::computeElementwiseMul(), mlir::arith::ConstantIndexOp::create(), mlir::OpBuilder::createOrFold(), and min().

◆ getChipStr()

std::optional< std::string > mlir::xegpu::getChipStr ( Operation op)

Retrieves the chip string from the XeVM target attribute of the parent GPU module operation.

Returns the chip identifier if found, or nullopt if no GPU module parent or XeVM target attribute exists.

Definition at line 414 of file XeGPUUtils.cpp.

References mlir::Operation::getParentOfType().

◆ 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 40 of file XeGPUUtils.cpp.

References mlir::get().

◆ getDistributeLayoutAttr() [1/2]

xegpu::DistributeLayoutAttr mlir::xegpu::getDistributeLayoutAttr ( const OpOperand opr)

Retrieves the DistributeLayoutAttr 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 155 of file XeGPUUtils.cpp.

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

◆ getDistributeLayoutAttr() [2/2]

xegpu::DistributeLayoutAttr mlir::xegpu::getDistributeLayoutAttr ( const Value  value)

Retrieves the DistributeLayoutAttr associated with a given Value.

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

Definition at line 117 of file XeGPUUtils.cpp.

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

Referenced by doSCFStructuralTypeConversionWithTensorType(), getDistributeLayoutAttr(), and getDistributeLayoutAttrOfType().

◆ getDistributeLayoutAttrOfType() [1/2]

template<typename AttrTy >
AttrTy mlir::xegpu::getDistributeLayoutAttrOfType ( const OpOperand opr)

Definition at line 89 of file XeGPUUtils.h.

References getDistributeLayoutAttr().

◆ getDistributeLayoutAttrOfType() [2/2]

template<typename AttrTy >
AttrTy mlir::xegpu::getDistributeLayoutAttrOfType ( const Value  value)

Definition at line 79 of file XeGPUUtils.h.

References getDistributeLayoutAttr().

◆ getLayoutName() [1/2]

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

Return the attribute name for the OpOperand to attach DistributeLayoutAttr.

Definition at line 106 of file XeGPUUtils.cpp.

Referenced by getDistributeLayoutAttr(), removeLayoutAttr(), and setDistributeLayoutAttr().

◆ getLayoutName() [2/2]

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

Return the attribute name for the OpResult to attach DistributeLayoutAttr.

Definition at line 112 of file XeGPUUtils.cpp.

References mlir::OpResult::getResultNumber().

◆ getShapeOf()

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

Definition at line 52 of file XeGPUOps.cpp.

Referenced by isValidGatherScatterBufferParams(), and isValidGatherScatterParams().

◆ isReadHintOrNone()

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

Definition at line 61 of file XeGPUOps.cpp.

References kind.

◆ isSharedMemory()

bool mlir::xegpu::isSharedMemory ( const MemRefType &  memrefTy)

Definition at line 26 of file XeGPUOps.cpp.

◆ isValidGatherScatterBufferParams()

static LogicalResult mlir::xegpu::isValidGatherScatterBufferParams ( Type  offsetsTy,
Type  maskTy,
VectorType  valueTy,
int64_t  chunkSize,
function_ref< InFlightDiagnostic()>  emitError 
)
static

Definition at line 125 of file XeGPUOps.cpp.

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

◆ isValidGatherScatterParams()

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

Definition at line 78 of file XeGPUOps.cpp.

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

◆ 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 ( array,
bool  breakline = false 
)
static

Definition at line 38 of file XeGPUOps.cpp.

Referenced by isValidGatherScatterParams().

◆ parseOptionalDynamicIndexList()

ParseResult mlir::xegpu::parseOptionalDynamicIndexList ( OpAsmParser parser,
SmallVectorImpl< OpAsmParser::UnresolvedOperand > &  values,
DenseI64ArrayAttr integers,
SmallVectorImpl< Type > *  valueTypes = nullptr,
AsmParser::Delimiter  delimiter = AsmParser::Delimiter::Square 
)

◆ populateXeGPUFoldAliasOpsPatterns()

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

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

Definition at line 63 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 820 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:

  1. the unrolled type unrolledType and number of unrolled instances numUnrolledInstances are computed from the targetShape.
  2. pack each operand. ExtractStridedSlice are created to break-up the vector operands. And BuiltinUnrealizedCastop are created to break-up the TensorDesc operands.
  3. the original op is cloned numUnrolledInstances times, once for each result.
  4. 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 687 of file XeGPUUnroll.cpp.

References options, and mlir::patterns.

◆ populateXeGPUWgToSgDistributePatterns()

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

Definition at line 822 of file XeGPUWgToSgDistribute.cpp.

References mlir::patterns.

◆ printOptionalDynamicIndexList()

void mlir::xegpu::printOptionalDynamicIndexList ( OpAsmPrinter printer,
Operation op,
OperandRange  values,
DenseI64ArrayAttr  integers 
)

Definition at line 381 of file XeGPUOps.cpp.

References mlir::printDynamicIndexList(), and mlir::AsmParser::Square.

◆ removeLayoutAttr()

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.

Definition at line 197 of file XeGPUUtils.cpp.

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

Referenced by removeLayoutAttrs().

◆ removeLayoutAttrs()

void mlir::xegpu::removeLayoutAttrs ( Operation op)

Removes the DistributeLayoutAttr for each OpOperand and OpResult of the given operation if they exist.

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

Definition at line 212 of file XeGPUUtils.cpp.

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

◆ setDistributeLayoutAttr()

template<typename T , typename = std::enable_if_t<std::is_same_v<T, OpOperand> || std::is_same_v<T, OpResult>>>
void mlir::xegpu::setDistributeLayoutAttr ( const T &  operandOrResult,
const DistributeLayoutAttr  layout 
)

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

Definition at line 164 of file XeGPUUtils.cpp.

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

Referenced by updateControlFlowOps(), and updateOp().

◆ setDistributeLayoutAttrs()

void mlir::xegpu::setDistributeLayoutAttrs ( Operation op,
function_ref< DistributeLayoutAttr(Value)>  getLayoutImpl 
)

Set the DistributeLayoutAttr 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 182 of file XeGPUUtils.cpp.

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