MLIR 23.0.0git
XeGPUUtils.h File Reference
#include "mlir/Dialect/XeGPU/IR/XeGPU.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/OpDefinition.h"
#include "llvm/ADT/SetVector.h"
#include <functional>

Go to the source code of this file.

Namespaces

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

Typedefs

using mlir::xegpu::SubShapeAndCountFn
 Callback type for computing sub-shape and count for 1:N (or 1:1 shape-changing) VectorType conversion.

Functions

SmallVector< Valuemlir::xegpu::flattenValues (ArrayRef< ValueRange > values)
 Flatten a set of ValueRange into a single SmallVector<Value>
FailureOr< VectorType > mlir::xegpu::getDistributedVectorType (xegpu::TensorDescType tdescTy)
 If tensor descriptor has a layout attribute it is used in SIMT mode.
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.
FailureOr< VectorType > mlir::xegpu::getDistVecTypeBasedOnLaneLayout (DistributeLayoutAttr layout, VectorType originalType)
 Helper function to get distributed vector type for a source vector type according to the lane_layout.
SmallVector< Valuemlir::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.
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.
std::optional< std::string > mlir::xegpu::getChipStr (Operation *op)
 Retrieves the chip string from the XeVM target attribute of the parent GPU module operation.
SmallVector< OpFoldResultmlir::xegpu::addElementwise (OpBuilder &builder, Location loc, ArrayRef< OpFoldResult > lhs, ArrayRef< OpFoldResult > rhs)
 Generates element-wise addition ops of two arrays with same length.
SmallVector< OpFoldResultmlir::xegpu::addWithRightAligned (OpBuilder &builder, Location loc, ArrayRef< OpFoldResult > lhs, ArrayRef< OpFoldResult > rhs)
 Generates element-wise addition ops of two arrays with automatic alignment.
Value mlir::xegpu::subgroupReduction (Location loc, OpBuilder &builder, Value input, vector::CombiningKind kind, uint32_t size)
 Given an input value representing per-lane data, this function returns the result after performing a reduction on the input over all lanes (number of lanes given by size).
Value mlir::xegpu::lowerToVectorReductions (TypedValue< VectorType > src, TypedValue< VectorType > acc, vector::CombiningKind kind, int64_t reductionDim, Location loc, PatternRewriter &rewriter)
 Given a src and an acc argumments from a vector::MultiDimReductionOp, lower to a set of vector::ReductionOp ops over 1D slices extracted from src.
Value mlir::xegpu::createReductionNeutralValue (OpBuilder &builder, Location loc, Type type, vector::CombiningKind kind)
 Creates a constant filled with the neutral (identity) value for the given reduction kind.
Value mlir::xegpu::lowerCrossLaneReductionToShuffles (TypedValue< VectorType > src, TypedValue< VectorType > acc, vector::CombiningKind kind, int64_t reductionDim, int64_t reductionSize, Location loc, PatternRewriter &rewriter)
 Lowers cross-lane reductions to shuffle operations on a 2D vector.
template<typename T>
int mlir::xegpu::getLargestDivisor (T dim, ArrayRef< T > candidates, ArrayRef< T > candidateMultiples={})
 Helper Function to find a proper instruction multiple for the user-supplied sg-level data shape (diven by dim).
DistributeLayoutAttr mlir::xegpu::getDistributeLayoutAttr (const Value value)
 Retrieves the DistributeLayoutAttr associated with a given Value.
DistributeLayoutAttr mlir::xegpu::getDistributeLayoutAttr (const OpOperand &opr)
 Retrieves the DistributeLayoutAttr associated with a given OpOperand.
void mlir::xegpu::setDistributeLayoutAttr (const OpResult &Result, const DistributeLayoutAttr layout)
 [to-be-deprecated] Sets the DistributeLayoutAttr for a given OpResult user should use setAnchorLayout instead
void mlir::xegpu::setDistributeLayoutAttr (const OpOperand &opr, const DistributeLayoutAttr layout)
 [to-be-deprecated] Sets the DistributeLayoutAttr for a given OpOperand user should use setAnchorLayout instead
std::string mlir::xegpu::getTemporaryLayoutName (const OpOperand &operand)
 Return the attribute name for the OpOperand to attach DistributeLayoutAttr.
std::string mlir::xegpu::getTemporaryLayoutName (const OpResult result)
 Return the attribute name for the OpResult to attach DistributeLayoutAttr.
template<typename T, typename = std::enable_if_t<std::is_same_v<T, OpOperand> || std::is_same_v<T, OpResult>>>
DistributeLayoutAttr mlir::xegpu::getTemporaryLayout (const T &operandOrResult)
 get and set distribute layout attribute for non-anchor operations (and offsets/masks of load/store ops before we get rid of their temp attrs)
template<typename T, typename = std::enable_if_t<std::is_same_v<T, OpOperand> || std::is_same_v<T, OpResult>>>
void mlir::xegpu::setTemporaryLayout (const T &operandOrResult, const DistributeLayoutAttr layout)
bool mlir::xegpu::requirePacked (const DistributeLayoutAttr layout)
 Helper function to check if the layout is packed.
bool mlir::xegpu::requireTranspose (const DistributeLayoutAttr layout, const uArch::uArch *uArch)
 Helper function to check if the layout requires a transpose effect.
bool mlir::xegpu::matchUnitDimExpansion (ArrayRef< int64_t > src, ArrayRef< int64_t > dst, SmallVector< int64_t > &expandedUnitDims)
bool mlir::xegpu::matchSplitDimExpansion (ArrayRef< int64_t > src, ArrayRef< int64_t > dst, SmallVector< SmallVector< int64_t > > &splitDimGroups)
DenseMap< Value, SmallVector< Type > > mlir::xegpu::precomputeLoopBlockArgTypes (Operation *topLevelOp, SubShapeAndCountFn getSubShapeAndCount)
 Pre-computes distributed VectorType mappings for every value carried through an SCF loop under topLevelOp (1:1 shape-changing or 1:N): the region block args (scf.while before/after args, scf.for iter_args), the loop results, and the terminator operands feeding them.
void mlir::xegpu::addVectorTypeConversion (TypeConverter &converter, SubShapeAndCountFn getSubShapeAndCount, DenseMap< Value, SmallVector< Type > > loopArgTypes)
 Adds a context-aware VectorType conversion to converter (1:1 shape-changing or 1:N, depending on getSubShapeAndCount's returned count).
void mlir::xegpu::cleanupUnrealizedConversionCasts (Operation *root, const llvm::SmallSetVector< UnrealizedConversionCastOp, 8 > &existingCasts)
 Cleans up UnrealizedConversionCastOps inserted during SCF structural type conversion and/or XeGPU unrolling.
bool mlir::xegpu::matchDimCollapse (ArrayRef< int64_t > src, ArrayRef< int64_t > dst, SmallVector< SmallVector< int64_t > > &collapseDims)