#include "mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/LLVMIR/XeVMDialect.h"
#include "mlir/Dialect/SCF/Transforms/Patterns.h"
#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Dialect/XeGPU/IR/XeGPU.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/Operation.h"
#include "mlir/IR/ValueRange.h"
#include "mlir/Interfaces/LoopLikeInterface.h"
#include "mlir/Transforms/DialectConversion.h"
#include "llvm/Support/FormatVariadic.h"
#include <cstdint>
#include <numeric>

Typedefs
using	LayoutRepresentation = std::pair<int64_t, int64_t>

Functions
template void	xegpu::removeLayoutAttr< mlir::OpResult > (const mlir::OpResult &result)
template void	xegpu::removeLayoutAttr< mlir::OpOperand > (const mlir::OpOperand &operand)
static xegpu::DistributeLayoutAttr	setupGenericLoadAnchorLayout (xegpu::LayoutKind layoutKind, mlir::MLIRContext *context, xegpu::DistributeLayoutAttr consumerLayout, bool isChunkedLoad, int maxChunkSize, int valShapeSize, int subgroupSize)
	Sets up the anchor layout for load gather and load matrix operation.
static xegpu::DistributeLayoutAttr	setupGenericStoreAnchorLayout (xegpu::LayoutKind layoutKind, mlir::MLIRContext *context, bool isChunkedStore, int maxChunkSize, ArrayRef< int64_t > srcShape, int subgroupSize)
	Sets up the anchor layout for store scatter and store matrix operation.
template<typename RankedTy>
static xegpu::LayoutAttr	getDefaultLaneLayout2DBlockIo (RankedTy ty, const xegpu::uArch::uArch *uArch, std::optional< unsigned > packingSize=std::nullopt, bool vnni=false)
static SmallVector< LayoutRepresentation >	getValidLayouts (ArrayRef< int64_t > wgShape, ArrayRef< int64_t > instData, int64_t sgCount)

Typedef Documentation

◆ LayoutRepresentation

using LayoutRepresentation = std::pair<int64_t, int64_t>

Definition at line 893 of file XeGPULayoutImpl.cpp.

Function Documentation

◆ getDefaultLaneLayout2DBlockIo()

template<typename RankedTy>

xegpu::LayoutAttr getDefaultLaneLayout2DBlockIo	(	RankedTy	ty,
		const xegpu::uArch::uArch *	uArch,
		std::optional< unsigned >	packingSize = std::nullopt,
		bool	vnni = false )

static

Definition at line 863 of file XeGPULayoutImpl.cpp.

References mlir::xegpu::uArch::uArch::getSubgroupSize().

◆ getValidLayouts()

SmallVector< LayoutRepresentation > getValidLayouts	(	ArrayRef< int64_t >	wgShape,
		ArrayRef< int64_t >	instData,
		int64_t	sgCount )

static

Definition at line 895 of file XeGPULayoutImpl.cpp.

References lhs, and rhs.

◆ setupGenericLoadAnchorLayout()

xegpu::DistributeLayoutAttr setupGenericLoadAnchorLayout	(	xegpu::LayoutKind	layoutKind,
		mlir::MLIRContext *	context,
		xegpu::DistributeLayoutAttr	consumerLayout,
		bool	isChunkedLoad,
		int	maxChunkSize,
		int	valShapeSize,
		int	subgroupSize )

static

Sets up the anchor layout for load gather and load matrix operation.

load matrix lowers to load gather and 1d block load. All of them share the same layout setup logic. For Subgroup layout, uses the consumer layout directly. non-chunked loads: InstData = {1, ..., min(consumer, maxLaneLoadSize * subgroupSize)} LaneLayout = {1, ..., subgroupSize} lane_data = {1, ..., min(consumer, maxLaneLoadSize)} chunked loads: InstData = {subgroupSize, min(consumer, maxLaneLoadSize)} LaneLayout = {subgroupSize, 1} lane_data={1,min(consumer, maxLaneLoadSize)}

Definition at line 681 of file XeGPULayoutImpl.cpp.

References mlir::xegpu::InstData, mlir::xegpu::Lane, and mlir::xegpu::Subgroup.

◆ setupGenericStoreAnchorLayout()

xegpu::DistributeLayoutAttr setupGenericStoreAnchorLayout	(	xegpu::LayoutKind	layoutKind,
		mlir::MLIRContext *	context,
		bool	isChunkedStore,
		int	maxChunkSize,
		ArrayRef< int64_t >	srcShape,
		int	subgroupSize )

static

Sets up the anchor layout for store scatter and store matrix operation.

store matrix lowers to store scatter and 1d block store. All of them share the same layout setup logic. For Subgroup layout, not support yet. non-chunked stores: InstData = {1, ..., subgroupSize} LaneLayout = {1, ..., subgroupSize} lane_data = {1, ..., 1} chunked stores: InstData = {subgroupSize, min(srcVec, maxLaneStoreSize)} LaneLayout = {subgroupSize, 1} lane_data={1,min(srcVec, maxLaneStoreSize)}

Definition at line 780 of file XeGPULayoutImpl.cpp.

References mlir::xegpu::InstData, mlir::xegpu::Lane, and mlir::xegpu::Subgroup.

Referenced by mlir::xegpu::setupStoreMatrixAnchorLayout(), and mlir::xegpu::setupStoreScatterAnchorLayout().