XeGPUPropagateLayout.cpp File Reference

#include "mlir/Analysis/DataFlow/DeadCodeAnalysis.h"
#include "mlir/Analysis/DataFlow/SparseAnalysis.h"
#include "mlir/Analysis/DataFlow/Utils.h"
#include "mlir/Analysis/DataFlowFramework.h"
#include "mlir/Dialect/GPU/IR/GPUDialect.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/XeGPU/IR/XeGPU.h"
#include "mlir/Dialect/XeGPU/Transforms/Passes.h"
#include "mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h"
#include "mlir/Dialect/XeGPU/Utils/XeGPUUtils.h"
#include "mlir/Dialect/XeGPU/uArch/IntelGpuXe2.h"
#include "mlir/IR/Attributes.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/BuiltinAttributes.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/Operation.h"
#include "mlir/IR/Value.h"
#include "mlir/IR/Visitors.h"
#include "mlir/Interfaces/ControlFlowInterfaces.h"
#include "mlir/Interfaces/FunctionInterfaces.h"
#include "mlir/Interfaces/LoopLikeInterface.h"
#include "mlir/Support/LLVM.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallSet.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/TypeSwitch.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/LogicalResult.h"
#include "llvm/Support/raw_ostream.h"
#include "mlir/Dialect/XeGPU/Transforms/Passes.h.inc"

Go to the source code of this file.

Classes
class	mlir::xegpu::impl::XeGPUPropagateLayoutBase< DerivedT >

Namespaces
namespace	mlir
	Include the generated interface declarations.
namespace	mlir::xegpu
namespace	mlir::xegpu::impl

Macros
#define	GEN_PASS_DEF_XEGPUPROPAGATELAYOUT
#define	DEBUG_TYPE   "xegpu-propagate-layout"
#define	DBGS()

Typedefs
using	GetLayoutFnTy = function_ref<xegpu::DistributeLayoutAttr(Value)>

Functions
std::unique_ptr<::mlir::Pass >	mlir::xegpu::impl::createXeGPUPropagateLayout ()
std::unique_ptr<::mlir::Pass >	mlir::xegpu::impl::createXeGPUPropagateLayout (XeGPUPropagateLayoutOptions options)
std::unique_ptr<::mlir::Pass >	mlir::xegpu::createXeGPUPropagateLayout ()
std::unique_ptr<::mlir::Pass >	mlir::xegpu::createXeGPUPropagateLayout (XeGPUPropagateLayoutOptions options)
SmallVector< std::pair< int, int > >	getValidLayouts (ArrayRef< int64_t > wgShape, ArrayRef< int > instData, int64_t sgCount)
FailureOr< int64_t >	getNumSg (Operation *op, const int sgSize)
static xegpu::CreateNdDescOp	getDefiningCreateNdDescOp (Value tdescValue)
	Helper to get the defining CreateNdDescOp of a tensor descriptor value.
static LogicalResult	updateOp (mlir::OpBuilder &builder, mlir::Operation *op, GetLayoutFnTy getLayoutOfValue)
	Update an operation with the layout of its results.
static LogicalResult	updateControlFlowOps (mlir::OpBuilder &builder, mlir::RegionBranchTerminatorOpInterface terminator, GetLayoutFnTy getLayoutOfValue)
	Region ops like scf.for need special handling because they have blocks inside.
static LogicalResult	updateFunctionOpInterface (mlir::OpBuilder &builder, mlir::FunctionOpInterface funcOp, GetLayoutFnTy getLayoutOfValue)
	Update the function arguments and results with the layouts.

Macro Definition Documentation

DBGS

#define DBGS

(

)

Value:

(llvm::dbgs() << "[" DEBUG_TYPE "]: ")

Definition at line 50 of file XeGPUPropagateLayout.cpp.

DEBUG_TYPE

#define DEBUG_TYPE   "xegpu-propagate-layout"

Definition at line 49 of file XeGPUPropagateLayout.cpp.

GEN_PASS_DEF_XEGPUPROPAGATELAYOUT

#define GEN_PASS_DEF_XEGPUPROPAGATELAYOUT

Definition at line 44 of file XeGPUPropagateLayout.cpp.

Typedef Documentation

GetLayoutFnTy

using GetLayoutFnTy = function_ref<xegpu::DistributeLayoutAttr(Value)>

Definition at line 1345 of file XeGPUPropagateLayout.cpp.

Function Documentation

getDefiningCreateNdDescOp()

xegpu::CreateNdDescOp getDefiningCreateNdDescOp ( Value tdescValue )

static

Helper to get the defining CreateNdDescOp of a tensor descriptor value.

This function tries to find the defining CreateNdDescOp recursively accross control-flow boundaries.

Definition at line 1276 of file XeGPUPropagateLayout.cpp.

References mlir::IROperand< DerivedT, IRValueT >::get(), getDefiningCreateNdDescOp(), and mlir::Value::getDefiningOp().

Referenced by getDefiningCreateNdDescOp().

getNumSg()

FailureOr< int64_t > getNumSg	(	Operation *	op,
		const int	sgSize )

Definition at line 552 of file XeGPUPropagateLayout.cpp.

References mlir::Operation::getParentOfType().

getValidLayouts()

SmallVector< std::pair< int, int > > getValidLayouts	(	ArrayRef< int64_t >	wgShape,
		ArrayRef< int >	instData,
		int64_t	sgCount )

Definition at line 523 of file XeGPUPropagateLayout.cpp.

References lhs, and rhs.

updateControlFlowOps()

LogicalResult updateControlFlowOps ( mlir::OpBuilder & builder, mlir::RegionBranchTerminatorOpInterface terminator, GetLayoutFnTy getLayoutOfValue )

static

Region ops like scf.for need special handling because they have blocks inside.

If the blocks have tensor descriptor type as block arguments, thier types must be updated. Also region op can have results that may not have any users (e.g. A and B tiles). They are not assigned a layout by layout analysis because they have no users. However inside the region op corresponding block arguments for these results do have layouts. Therefore, in this case we still need to update the result types with the layout attribute. This function function updates the internal block arguments and the result types of the region op with the assigned layouts. clang-format off Example: scf.for ... iter_args(...) -> (out types) { ^bb0(block types): ... scf.yield ... : (yield types) } clang-format on In this example, at scf.yield, control-flow can transfer to two successor regions. One is the ^bb0 (for loop body) and the other is the scf.for op itself (yield the results). So we update both the block arguments of the successor region (i.e. block types) and the result types of the scf.for op (i.e. out types). Note that yield types are updated by respective producers inside bb0.

Definition at line 1409 of file XeGPUPropagateLayout.cpp.

References DBGS, result, mlir::xegpu::setDistributeLayoutAttr(), and success().

Referenced by mlir::xegpu::propagateLayouts().

updateFunctionOpInterface()

LogicalResult updateFunctionOpInterface ( mlir::OpBuilder & builder, mlir::FunctionOpInterface funcOp, GetLayoutFnTy getLayoutOfValue )

static

Update the function arguments and results with the layouts.

Definition at line 1465 of file XeGPUPropagateLayout.cpp.

References DBGS, and success().

Referenced by mlir::xegpu::propagateLayouts().

updateOp()

LogicalResult updateOp ( mlir::OpBuilder & builder, mlir::Operation * op, GetLayoutFnTy getLayoutOfValue )

static

Update an operation with the layout of its results.

If the result type is a vector type, a temporary layout attribute is added to the operation. If the result type is a tensor descriptor type, the type is updated with the layout attribute. The users of the result are also updated with the layout attribute.

Definition at line 1351 of file XeGPUPropagateLayout.cpp.

References mlir::Operation::emitWarning(), mlir::Operation::getResults(), result, mlir::xegpu::setDistributeLayoutAttr(), and success().

Referenced by mlir::xegpu::propagateLayouts().