XeGPUUtils.h

Go to the documentation of this file.

//===- XeGPUUtils.h - Vector Utilities --------------------------*- C++ -*-===//
//
// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
//
//===----------------------------------------------------------------------===//
 
#ifndef MLIR_DIALECT_XEGPU_UTILS_XEGPUUTILS_H_
#define MLIR_DIALECT_XEGPU_UTILS_XEGPUUTILS_H_
 
#include "mlir/Dialect/XeGPU/IR/XeGPU.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/OpDefinition.h"
#include "llvm/ADT/SetVector.h"
#include <functional>
#include <optional>
 
namespace mlir {
class UnrealizedConversionCastOp;
class VectorType;
class OpOperand;
class OpResult;
class OpBuilder;
class ValueRange;
class TypeConverter;
class OpFoldResult;
 
namespace xegpu {
class DistributeLayoutAttr;
class LayoutAttr;
class TensorDescType;
 
namespace uArch {
struct uArch;
} // namespace uArch
} // namespace xegpu
 
namespace xegpu {
/// Flatten a set of ValueRange into a single SmallVector<Value>
SmallVector<Value> flattenValues(ArrayRef<ValueRange> values);
 
/// 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]
FailureOr<VectorType> getDistributedVectorType(xegpu::TensorDescType tdescTy);
 
/// Helper to get the distributed vector type for a given vector type according
/// to a given LayoutAttr.
FailureOr<VectorType> getDistributedVectorType(VectorType originalType,
                                               LayoutAttr layout);
 
/// Helper function to get distributed vector type for a source vector type
/// according to the lane_layout. We simply divide each dimension of tensor
/// descriptor shape by corresponding lane_layout dimension. If
/// array_length > 1, that is appended to the front of the distributed shape.
///
/// Examples:
/// | original vector shape | lane_layout | distributed vector shape |
/// |-----------------------|-------------|--------------------------|
/// | 32x16                 | [1, 16]     | 32x1                     |
/// | 32x16                 | [2, 8]      | 16x2                     |
/// | 2x32x16               | [1, 16]     | 2x32x1                   |
FailureOr<VectorType>
getDistVecTypeBasedOnLaneLayout(DistributeLayoutAttr layout,
                                VectorType originalType);
 
/// Extract a set of small vectors from a value with a given shape using
/// vector.extract_stride_slice
SmallVector<Value> extractVectorsWithShapeFromValue(OpBuilder &builder,
                                                    Location loc, Value value,
                                                    ArrayRef<int64_t> shape);
 
/// Create a vector of shape from a set of values using
/// vector.insert_stride_slice.
Value createVectorWithShapeFromValues(OpBuilder &builder, Location loc,
                                      ValueRange values,
                                      ArrayRef<int64_t> shape);
 
/// 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.
std::optional<std::string> getChipStr(Operation *op);
 
/// Returns the number of subgroups the kernel enclosing `op` runs, derived from
/// the `known_block_size` of its parent `gpu.func`. Fails when there is no such
/// parent, when the attribute is absent, when a block dimension is not a
/// positive power of two, or when the block does not cover a whole subgroup.
FailureOr<int64_t> getNumSubgroupsFromBlockSize(Operation *op,
                                                int64_t subgroupSize);
 
/// Generates element-wise addition ops of two arrays with same length.
SmallVector<OpFoldResult> addElementwise(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]
SmallVector<OpFoldResult> addWithRightAligned(OpBuilder &builder, Location loc,
                                              ArrayRef<OpFoldResult> lhs,
                                              ArrayRef<OpFoldResult> rhs);
 
/// 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`). This uses butterfly shuffles to perform the
/// reduction in a log2(size) number of steps.
/// NOTE: Implementation taken from TestVectorTransforms.cpp
/// For float kinds, `fastmath` is set on every combining operation.
Value subgroupReduction(Location loc, OpBuilder &builder, Value input,
                        vector::CombiningKind kind, uint32_t size,
                        arith::FastMathFlagsAttr fastmath = nullptr);
 
/// 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`. The reduction is performed along `reductionDim`. The result is a
/// vector with the same shape as `acc`. `fastmath` holds the fast-math flags
/// of the vector::MultiDimReductionOp. Each new vector::ReductionOp gets them.
/// TODO: Only 2D to 1D reduction is supported for now.
Value lowerToVectorReductions(TypedValue<VectorType> src,
                              TypedValue<VectorType> acc,
                              vector::CombiningKind kind, int64_t reductionDim,
                              Location loc, PatternRewriter &rewriter,
                              arith::FastMathFlagsAttr fastmath = nullptr);
 
/// Creates a constant filled with the neutral (identity) value for the
/// given reduction kind. For example: 0 for ADD/OR/XOR, 1 for MUL, all ones
/// for AND, and max/min signed/unsigned int for MINSI/MINUI/MAXSI/MAXUI.
/// For float min/max kinds, the value must fit the element type and be
/// allowed by \p fastmath:
/// - A quiet NaN for the kinds that ignore NaN operands, if the type has NaN
///   and `nnan` is not set.
/// - Otherwise the infinity that never wins, if the type has it and `ninf`
///   is not set.
/// - Otherwise the lowest (for max) or highest (for min) finite value.
/// If \p type is a VectorType, returns a splat vector constant; otherwise
/// returns a scalar constant. Returns nullptr if the element type is
/// incompatible with the requested reduction kind.
Value createReductionNeutralValue(OpBuilder &builder, Location loc, Type type,
                                  vector::CombiningKind kind,
                                  arith::FastMathFlagsAttr fastmath = nullptr);
 
/// Lowers cross-lane reductions to shuffle operations on a 2D vector.
/// Extracts slices along the reduction dimension, performs subgroup reductions
/// with shuffles across reductionSize work-items, and inserts the results back
/// into an accumulator vector. `fastmath` holds the fast-math flags of the
/// reduction being lowered. For float kinds, they are set on every combining
/// operation: the shuffle reductions (see subgroupReduction) and the combine
/// with each accumulator element.
Value lowerCrossLaneReductionToShuffles(
    TypedValue<VectorType> src, TypedValue<VectorType> acc,
    vector::CombiningKind kind, int64_t reductionDim, int64_t reductionSize,
    Location loc, PatternRewriter &rewriter,
    arith::FastMathFlagsAttr fastmath = nullptr);
 
/// Helper Function to find a proper instruction multiple for the user-supplied
/// sg-level data shape (diven by `dim`). `candidates` are uArch allowed shapes.
/// `candidateMultiples` are uArch multiples of such shapes (i.e. block count or
/// array length).
template <typename T>
int getLargestDivisor(T dim, ArrayRef<T> candidates,
                      ArrayRef<T> candidateMultiples = {});
 
/// Retrieves the DistributeLayoutAttr associated with a given Value, or nullptr
/// if none is found. For TensorDescType values it is extracted from the
/// TensorDescType itself; for an op result, from the attributes of the defining
/// operation.
///
/// A block argument carries no attribute of its own, so it is resolved through
/// the operand that feeds it:
///  - an iter_arg of a loop, including scf.while's "before" arguments, resolves
///    through its tied init operand;
///  - an scf.while "after" argument is fed by scf.condition rather than by an
///    init operand. Only a pass-through before region is handled: the forwarded
///    value must be a "before" argument, and the layout comes from that
///    argument's tied init operand. If the before region forwards anything
///    computed, nullptr is returned.
DistributeLayoutAttr getDistributeLayoutAttr(const Value value);
 
/// 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.
DistributeLayoutAttr getDistributeLayoutAttr(const OpOperand &opr);
 
/// [to-be-deprecated] Sets the DistributeLayoutAttr for a given OpResult
/// user should use setAnchorLayout instead
void setDistributeLayoutAttr(const OpResult &Result,
                             const DistributeLayoutAttr layout);
 
/// [to-be-deprecated] Sets the DistributeLayoutAttr for a given OpOperand
/// user should use setAnchorLayout instead
void setDistributeLayoutAttr(const OpOperand &opr,
                             const DistributeLayoutAttr layout);
 
/// Return the attribute name for the OpOperand to attach DistributeLayoutAttr
std::string getTemporaryLayoutName(const OpOperand &operand);
 
/// Return the attribute name for the OpResult to attach DistributeLayoutAttr
std::string getTemporaryLayoutName(const OpResult result);
 
/// 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>>>
DistributeLayoutAttr getTemporaryLayout(const T &operandOrResult);
 
template <typename T,
          typename = std::enable_if_t<std::is_same_v<T, OpOperand> ||
                                      std::is_same_v<T, OpResult>>>
void setTemporaryLayout(const T &operandOrResult,
                        const DistributeLayoutAttr layout);
 
/// Returns the innermost 2 entries of `vals` if it is at least 2D and all of
/// its leading entries are unit; std::nullopt otherwise.
std::optional<SmallVector<int64_t>>
getInner2DIfUnitLeadingDims(ArrayRef<int64_t> vals);
 
/// Helper function to check if the layout is packed. Layout is packed if
/// lane_data[rank-2] != 1 (data packed from col dimension).
/// TODO: Move to target info.
bool requirePacked(const DistributeLayoutAttr layout);
 
/// Helper function to check if the layout requires a transpose effect.
bool requireTranspose(const DistributeLayoutAttr layout,
                      const uArch::uArch *uArch);
 
/// Returns true if `type` has a static shape and static strides.
bool hasStaticShapeAndStrides(MemRefType type);
 
/// Returns `baseAddr + elemOffset * bitWidth / 8`, i.e. advances the byte
/// address `baseAddr` by `elemOffset` elements of `memrefTy`'s element type.
/// Sub-byte element types (e.g. i4) are supported. `baseAddr` and `elemOffset`
/// must have the same index or integer type, which is also the result type.
Value addElemOffsetInBytes(OpBuilder &builder, Location loc, Value baseAddr,
                           Value elemOffset, MemRefType memrefTy);
 
// Check if dst shape is an expansion of src shape by inserting unit dimensions.
bool matchUnitDimExpansion(ArrayRef<int64_t> src, ArrayRef<int64_t> dst,
                           SmallVector<int64_t> &expandedUnitDims);
 
// Checks if dst shape is a non-unit expansion of src shape where each
// dimension in src is split into one or more consecutive dimensions in dst.
bool matchSplitDimExpansion(ArrayRef<int64_t> src, ArrayRef<int64_t> dst,
                            SmallVector<SmallVector<int64_t>> &splitDimGroups);
 
/// Callback type for computing sub-shape and count for 1:N (or 1:1
/// shape-changing) VectorType conversion. Given a VectorType and its
/// DistributeLayoutAttr, returns (subShape, count). A count <= 0 signals
/// "no conversion needed"; count == 1 is a 1:1 shape-changing conversion;
/// count > 1 produces `count` copies of `subShape`.
using SubShapeAndCountFn = std::function<std::pair<SmallVector<int64_t>, int>(
    VectorType, DistributeLayoutAttr)>;
 
/// 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. Each is derived from
/// a single source -- the layout of the feeding value (loop init or
/// `scf.condition` operand) -- and keyed by `Value`, because the SCF converters
/// detach/replace the loop body mid-conversion, after which a layout query on a
/// block arg returns null. Recording results and terminator operands lets a 1:N
/// pass resolve them from the map after stripping the loop op's transient
/// layout attrs. `scf.if` has no loop-carried block args and needs no entry.
DenseMap<Value, SmallVector<Type>>
precomputeLoopBlockArgTypes(Operation *topLevelOp,
                            SubShapeAndCountFn getSubShapeAndCount);
 
/// Adds a context-aware VectorType conversion to `converter` (1:1
/// shape-changing or 1:N, depending on `getSubShapeAndCount`'s returned
/// count). `getSubShapeAndCount` computes (subShape, count) for a VectorType
/// and its layout; count <= 0 means no conversion needed. `loopArgTypes`
/// (typically obtained from `precomputeLoopBlockArgTypes`) provides the
/// pre-computed types for SCF loop block arguments (`scf.while`,
/// `scf.for`); pass an empty map if the IR has no such loops.
void addVectorTypeConversion(TypeConverter &converter,
                             SubShapeAndCountFn getSubShapeAndCount,
                             DenseMap<Value, SmallVector<Type>> loopArgTypes);
 
/// Cleans up UnrealizedConversionCastOps inserted during SCF structural type
/// conversion and/or XeGPU unrolling. Folds cancelling N:1->1:N and 1:N->N:1
/// cast chains (inserting vector.shape_cast when shapes differ but element
/// counts match). Unpaired pack (1:N) and unpack (N:1) casts between a single
/// large VectorType and N identically-typed smaller VectorTypes are lowered
/// to vector.extract_strided_slice / vector.insert_strided_slice. Dead casts
/// are erased. Casts in `existingCasts` are preserved.
void cleanupUnrealizedConversionCasts(
    Operation *root,
    const llvm::SmallSetVector<UnrealizedConversionCastOp, 8> &existingCasts);
 
// Checks if dst shape is a collapse of src shape where each dimension in dst is
// produced by one or more consecutive dimensions in src whose product equals
// the dst dimension. Populates collapseDims with groups of src indices that are
// collapsed into each dst dimension. Leading or trailing unit dst dimensions
// (with no backing src dim) result in empty groups. Example: src=[8,16,32],
// dst=[1,4096] -> true, collapseDims=[[],[0,1,2]].
bool matchDimCollapse(ArrayRef<int64_t> src, ArrayRef<int64_t> dst,
                      SmallVector<SmallVector<int64_t>> &collapseDims);
} // namespace xegpu
} // namespace mlir
 
#endif // MLIR_DIALECT_XEGPU_UTILS_XEGPUUTILS_H_