mlir.dialects.transform.xegpu

Classes

ConvertLayoutOp	Specialization for ConvertLayoutOp class.
ConvertLayoutOpAdaptor
GetLoadOp	Specialization for GetLoadOp class.
GetLoadOpAdaptor
InsertPrefetchOp	Specialization for InsertPrefetchOp class.
InsertPrefetchOpAdaptor
SetAnchorLayoutOp	Specialization for SetAnchorLayoutOp class.
SetAnchorLayoutOpAdaptor
SetGPULaunchThreadsOp	Specialization for SetGPULaunchThreadsOp class.
SetGPULaunchThreadsOpAdaptor
GetLoadOp	Specialization for GetLoadOp class.
SetAnchorLayoutOp	Specialization for SetAnchorLayoutOp class.
SetGPULaunchThreadsOp	Specialization for SetGPULaunchThreadsOp class.
InsertPrefetchOp	Specialization for InsertPrefetchOp class.
ConvertLayoutOp	Specialization for ConvertLayoutOp class.

Functions

xegpu_convert_layout(→ _ods_ir)
xegpu_get_load_op(→ _ods_ir)
xegpu_insert_prefetch(→ _ods_ir)
xegpu_set_anchor_layout(→ SetAnchorLayoutOp)
xegpu_set_gpu_launch_threads(→ SetGPULaunchThreadsOp)
get_load_op(→ OpResult)
set_anchor_layout(→ SetAnchorLayoutOp)
set_gpu_launch_threads(→ SetGPULaunchThreadsOp)
insert_prefetch(→ OpResult)
convert_layout(→ ConvertLayoutOp)

Module Contents

class mlir.dialects.transform.xegpu.ConvertLayoutOp(newConvertOp: _ods_ir, target: _ods_ir, input_sg_layout: Sequence[_ods_ir], input_sg_data: Sequence[_ods_ir], input_inst_data: Sequence[_ods_ir], target_sg_layout: Sequence[_ods_ir], target_sg_data: Sequence[_ods_ir], target_inst_data: Sequence[_ods_ir], *, input_order: Sequence[int] | _ods_ir | None = None, static_input_sg_layout: Sequence[int] | _ods_ir | None = None, static_input_sg_data: Sequence[int] | _ods_ir | None = None, static_input_inst_data: Sequence[int] | _ods_ir | None = None, static_target_sg_layout: Sequence[int] | _ods_ir | None = None, static_target_sg_data: Sequence[int] | _ods_ir | None = None, static_target_inst_data: Sequence[int] | _ods_ir | None = None, target_order: Sequence[int] | _ods_ir | None = None, loc: _ods_ir | None = None, ip: _ods_ir | None = None)

Bases: _ods_ir

Adds an xegpu.convert_layout op to convert the xegpu.layout attribute of a value before its first use. The input and target layouts are defined by the *sg_layout, *sg_data and optional *inst_data and *order attributes. Returns a handle to the emitted xegpu.convert_layout op.

OPERATION_NAME = 'transform.xegpu.convert_layout'

_ODS_OPERAND_SEGMENTS

_ODS_REGIONS = (0, True)

target() → _ods_ir

input_sg_layout() → _ods_ir

input_sg_data() → _ods_ir

input_inst_data() → _ods_ir

target_sg_layout() → _ods_ir

target_sg_data() → _ods_ir

target_inst_data() → _ods_ir

input_order() → _ods_ir | None

static_input_sg_layout() → _ods_ir | None

static_input_sg_data() → _ods_ir | None

static_input_inst_data() → _ods_ir | None

static_target_sg_layout() → _ods_ir | None

static_target_sg_data() → _ods_ir | None

static_target_inst_data() → _ods_ir | None

target_order() → _ods_ir | None

newConvertOp() → _ods_ir

class mlir.dialects.transform.xegpu.ConvertLayoutOpAdaptor

Bases: _ods_ir

OPERATION_NAME = 'transform.xegpu.convert_layout'

target() → _ods_ir

input_sg_layout() → _ods_ir

input_sg_data() → _ods_ir

input_inst_data() → _ods_ir

target_sg_layout() → _ods_ir

target_sg_data() → _ods_ir

target_inst_data() → _ods_ir

input_order() → _ods_ir | None

static_input_sg_layout() → _ods_ir | None

static_input_sg_data() → _ods_ir | None

static_input_inst_data() → _ods_ir | None

static_target_sg_layout() → _ods_ir | None

static_target_sg_data() → _ods_ir | None

static_target_inst_data() → _ods_ir | None

target_order() → _ods_ir | None

mlir.dialects.transform.xegpu.xegpu_convert_layout(new_convert_op: _ods_ir, target: _ods_ir, input_sg_layout: Sequence[_ods_ir], input_sg_data: Sequence[_ods_ir], input_inst_data: Sequence[_ods_ir], target_sg_layout: Sequence[_ods_ir], target_sg_data: Sequence[_ods_ir], target_inst_data: Sequence[_ods_ir], *, input_order: Sequence[int] | _ods_ir | None = None, static_input_sg_layout: Sequence[int] | _ods_ir | None = None, static_input_sg_data: Sequence[int] | _ods_ir | None = None, static_input_inst_data: Sequence[int] | _ods_ir | None = None, static_target_sg_layout: Sequence[int] | _ods_ir | None = None, static_target_sg_data: Sequence[int] | _ods_ir | None = None, static_target_inst_data: Sequence[int] | _ods_ir | None = None, target_order: Sequence[int] | _ods_ir | None = None, loc: _ods_ir | None = None, ip: _ods_ir | None = None) → _ods_ir

class mlir.dialects.transform.xegpu.GetLoadOp(loadNdHandle: _ods_ir, target: _ods_ir, *, loc: _ods_ir | None = None, ip: _ods_ir | None = None)

Bases: _ods_ir

Traces the producers of the given value until an xegpu.load_nd or xegpu.load op is found. Returns a handle to it. Currently traces producers by following only the first operand of producer ops.

OPERATION_NAME = 'transform.xegpu.get_load_op'

_ODS_REGIONS = (0, True)

target() → _ods_ir

loadNdHandle() → _ods_ir

class mlir.dialects.transform.xegpu.GetLoadOpAdaptor

Bases: _ods_ir

OPERATION_NAME = 'transform.xegpu.get_load_op'

target() → _ods_ir

mlir.dialects.transform.xegpu.xegpu_get_load_op(load_nd_handle: _ods_ir, target: _ods_ir, *, loc: _ods_ir | None = None, ip: _ods_ir | None = None) → _ods_ir

class mlir.dialects.transform.xegpu.InsertPrefetchOp(desc_op: _ods_ir, target: _ods_ir, *, dynamic_nb_prefetch: _ods_ir | None = None, static_nb_prefetch: int | _ods_ir | None = None, loc: _ods_ir | None = None, ip: _ods_ir | None = None)

Bases: _ods_ir

Inserts xegpu.prefetch_nd operations for the given xegpu.load_nd op. The load op must reside within the scf.for loop. Number of prefetch steps is set by the nb_prefetch argument (default value is 1). Returns a handle to the created xegpu.create_nd_desc op.

OPERATION_NAME = 'transform.xegpu.insert_prefetch'

_ODS_REGIONS = (0, True)

target() → _ods_ir

dynamic_nb_prefetch() → _ods_ir | None

static_nb_prefetch() → _ods_ir | None

desc_op() → _ods_ir

class mlir.dialects.transform.xegpu.InsertPrefetchOpAdaptor

Bases: _ods_ir

OPERATION_NAME = 'transform.xegpu.insert_prefetch'

target() → _ods_ir

dynamic_nb_prefetch() → _ods_ir | None

static_nb_prefetch() → _ods_ir | None

mlir.dialects.transform.xegpu.xegpu_insert_prefetch(desc_op: _ods_ir, target: _ods_ir, *, dynamic_nb_prefetch: _ods_ir | None = None, static_nb_prefetch: int | _ods_ir | None = None, loc: _ods_ir | None = None, ip: _ods_ir | None = None) → _ods_ir

class mlir.dialects.transform.xegpu.SetAnchorLayoutOp(target: _ods_ir, sg_layout: Sequence[_ods_ir], sg_data: Sequence[_ods_ir], inst_data: Sequence[_ods_ir], *, index: int | _ods_ir | None = None, static_sg_layout: Sequence[int] | _ods_ir | None = None, static_sg_data: Sequence[int] | _ods_ir | None = None, static_inst_data: Sequence[int] | _ods_ir | None = None, order: Sequence[int] | _ods_ir | None = None, slice_dims: Sequence[int] | _ods_ir | None = None, loc: _ods_ir | None = None, ip: _ods_ir | None = None)

Bases: _ods_ir

Sets the xegpu.layout anchor layout for XeGPU ops that support it. The target operand value can be set by the index argument (currently only applicable to a DPAS op). The layout is defined by the sg_layout, sg_data and optional inst_data and order attributes. If slice_dims is provided, the xegpu.layout attribute is wrapped in an xegpu.slice<..., dims=slice_dims> attribute. Emits a silenceable failure if the target op does not support anchor layouts.

OPERATION_NAME = 'transform.xegpu.set_anchor_layout'

_ODS_OPERAND_SEGMENTS

_ODS_REGIONS = (0, True)

target() → _ods_ir

sg_layout() → _ods_ir

sg_data() → _ods_ir

inst_data() → _ods_ir

index() → _ods_ir | None

static_sg_layout() → _ods_ir | None

static_sg_data() → _ods_ir | None

static_inst_data() → _ods_ir | None

order() → _ods_ir | None

slice_dims() → _ods_ir | None

class mlir.dialects.transform.xegpu.SetAnchorLayoutOpAdaptor

Bases: _ods_ir

OPERATION_NAME = 'transform.xegpu.set_anchor_layout'

target() → _ods_ir

sg_layout() → _ods_ir

sg_data() → _ods_ir

inst_data() → _ods_ir

index() → _ods_ir | None

static_sg_layout() → _ods_ir | None

static_sg_data() → _ods_ir | None

static_inst_data() → _ods_ir | None

order() → _ods_ir | None

slice_dims() → _ods_ir | None

mlir.dialects.transform.xegpu.xegpu_set_anchor_layout(target: _ods_ir, sg_layout: Sequence[_ods_ir], sg_data: Sequence[_ods_ir], inst_data: Sequence[_ods_ir], *, index: int | _ods_ir | None = None, static_sg_layout: Sequence[int] | _ods_ir | None = None, static_sg_data: Sequence[int] | _ods_ir | None = None, static_inst_data: Sequence[int] | _ods_ir | None = None, order: Sequence[int] | _ods_ir | None = None, slice_dims: Sequence[int] | _ods_ir | None = None, loc: _ods_ir | None = None, ip: _ods_ir | None = None) → SetAnchorLayoutOp

class mlir.dialects.transform.xegpu.SetGPULaunchThreadsOp(target: _ods_ir, threads: Sequence[_ods_ir], *, static_threads: Sequence[int] | _ods_ir | None = None, loc: _ods_ir | None = None, ip: _ods_ir | None = None)

Bases: _ods_ir

Overrides the x,y,z threads operands of a given gpu.launch operation in-place.

OPERATION_NAME = 'transform.xegpu.set_gpu_launch_threads'

_ODS_REGIONS = (0, True)

target() → _ods_ir

threads() → _ods_ir

static_threads() → _ods_ir | None

class mlir.dialects.transform.xegpu.SetGPULaunchThreadsOpAdaptor

Bases: _ods_ir

OPERATION_NAME = 'transform.xegpu.set_gpu_launch_threads'

target() → _ods_ir

threads() → _ods_ir

static_threads() → _ods_ir | None

mlir.dialects.transform.xegpu.xegpu_set_gpu_launch_threads(target: _ods_ir, threads: Sequence[_ods_ir], *, static_threads: Sequence[int] | _ods_ir | None = None, loc: _ods_ir | None = None, ip: _ods_ir | None = None) → SetGPULaunchThreadsOp

class mlir.dialects.transform.xegpu.GetLoadOp(target: Value, *, loc=None, ip=None)

Bases: GetLoadOp

Specialization for GetLoadOp class.

mlir.dialects.transform.xegpu.get_load_op(target: Value, *, loc=None, ip=None) → OpResult

class mlir.dialects.transform.xegpu.SetAnchorLayoutOp(target: Operation | Value, sg_layout: mlir.dialects._ods_common.MixedValues, sg_data: mlir.dialects._ods_common.MixedValues, *, inst_data: mlir.dialects._ods_common.MixedValues | None = None, order: mlir.dialects._ods_common.MixedInt | None = None, slice_dims: mlir.dialects._ods_common.MixedInt | None = None, index: int | Attribute | None = None, loc=None, ip=None)

Bases: SetAnchorLayoutOp

Specialization for SetAnchorLayoutOp class.

mlir.dialects.transform.xegpu.set_anchor_layout(target: Operation | Value, sg_layout: mlir.dialects._ods_common.MixedValues, sg_data: mlir.dialects._ods_common.MixedValues, *, inst_data: mlir.dialects._ods_common.MixedValues | None = None, order: mlir.dialects._ods_common.MixedInt | None = None, slice_dims: mlir.dialects._ods_common.MixedInt | None = None, index: int | Attribute | None = None, loc=None, ip=None) → SetAnchorLayoutOp

class mlir.dialects.transform.xegpu.SetGPULaunchThreadsOp(launch_op: Operation | Value, threads: mlir.dialects._ods_common.MixedValues, *, loc=None, ip=None)

Bases: SetGPULaunchThreadsOp

Specialization for SetGPULaunchThreadsOp class.

mlir.dialects.transform.xegpu.set_gpu_launch_threads(launch_op: Operation | Value, threads: mlir.dialects._ods_common.MixedValues, *, loc=None, ip=None) → SetGPULaunchThreadsOp

class mlir.dialects.transform.xegpu.InsertPrefetchOp(target: Operation | Value, *, nb_prefetch: mlir.dialects._ods_common.MixedInt | None = 1, loc=None, ip=None)

Bases: InsertPrefetchOp

Specialization for InsertPrefetchOp class.

mlir.dialects.transform.xegpu.insert_prefetch(target: Operation | Value, *, nb_prefetch: mlir.dialects._ods_common.MixedInt | None = 1, loc=None, ip=None) → OpResult

class mlir.dialects.transform.xegpu.ConvertLayoutOp(target: Value, input_sg_layout: mlir.dialects._ods_common.MixedValues, input_sg_data: mlir.dialects._ods_common.MixedValues, target_sg_layout: mlir.dialects._ods_common.MixedValues, target_sg_data: mlir.dialects._ods_common.MixedValues, *, input_inst_data: mlir.dialects._ods_common.MixedValues | None = None, target_inst_data: mlir.dialects._ods_common.MixedValues | None = None, input_order: mlir.dialects._ods_common.MixedInt | None = None, target_order: mlir.dialects._ods_common.MixedInt | None = None, loc=None, ip=None)

Bases: ConvertLayoutOp

Specialization for ConvertLayoutOp class.

mlir.dialects.transform.xegpu.convert_layout(target: Value, input_sg_layout: mlir.dialects._ods_common.MixedValues, input_sg_data: mlir.dialects._ods_common.MixedValues, target_sg_layout: mlir.dialects._ods_common.MixedValues, target_sg_data: mlir.dialects._ods_common.MixedValues, *, input_inst_data: mlir.dialects._ods_common.MixedValues | None = None, target_inst_data: mlir.dialects._ods_common.MixedValues | None = None, input_order: mlir.dialects._ods_common.MixedInt | None = None, target_order: mlir.dialects._ods_common.MixedInt | None = None, loc=None, ip=None) → ConvertLayoutOp