mlir.dialects._xegpu_transform_ops_gen¶

Attributes¶

_ods_ir

Classes¶

`ConvertLayoutOp`	Adds an `xegpu.convert_layout` op to convert the `xegpu.layout` attribute
`GetDescOp`	Traces the producers of the given value until an `xegpu.create_nd_tdesc`
`InsertPrefetchOp`	Given a target value (e.g., `vector`) residing in a `scf.for` loop, this
`SetDescLayoutOp`	Given an `xegpu.create_nd_desc` operation, this transform adds `xegpu.layout`
`SetGPULaunchThreadsOp`	Overrides the x,y,z threads operands of a given `gpu.launch` operation in-place.
`SetOpLayoutAttrOp`	Sets the `xegpu.layout` attribute of an op. If `result=true`, sets the

Functions¶

`xegpu_convert_layout`(→ _ods_ir)
`xegpu_get_desc_op`(→ _ods_ir)
`xegpu_insert_prefetch`(→ _ods_ir)
`xegpu_set_desc_layout`(→ _ods_ir)
`xegpu_set_gpu_launch_threads`(→ SetGPULaunchThreadsOp)
`xegpu_set_op_layout_attr`(→ SetOpLayoutAttrOp)

Module Contents¶

mlir.dialects._xegpu_transform_ops_gen._ods_ir¶

class mlir.dialects._xegpu_transform_ops_gen.ConvertLayoutOp(newConvertOp, target, input_sg_layout, input_sg_data, input_inst_data, target_sg_layout, target_sg_data, target_inst_data, *, static_input_sg_layout=None, static_input_sg_data=None, static_input_inst_data=None, static_target_sg_layout=None, static_target_sg_data=None, static_target_inst_data=None, loc=None, ip=None)¶

Bases: _ods_ir

Adds an xegpu.convert_layout op to convert the xegpu.layout attribute of a value. The input and target layouts are defined by the *sg_layout, *sg_data and optional *inst_data 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¶

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¶

newConvertOp() → _ods_ir¶

mlir.dialects._xegpu_transform_ops_gen.xegpu_convert_layout(new_convert_op, target, input_sg_layout, input_sg_data, input_inst_data, target_sg_layout, target_sg_data, target_inst_data, *, static_input_sg_layout=None, static_input_sg_data=None, static_input_inst_data=None, static_target_sg_layout=None, static_target_sg_data=None, static_target_inst_data=None, loc=None, ip=None) → _ods_ir¶

class mlir.dialects._xegpu_transform_ops_gen.GetDescOp(descHandle, target, *, loc=None, ip=None)¶

Bases: _ods_ir

Traces the producers of the given value until an xegpu.create_nd_tdesc descriptor 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_desc_op'¶

_ODS_REGIONS = (0, True)¶

target() → _ods_ir¶

descHandle() → _ods_ir¶

mlir.dialects._xegpu_transform_ops_gen.xegpu_get_desc_op(desc_handle, target, *, loc=None, ip=None) → _ods_ir¶

class mlir.dialects._xegpu_transform_ops_gen.InsertPrefetchOp(desc_op, target, *, dynamic_nb_prefetch=None, static_nb_prefetch=None, loc=None, ip=None)¶

Bases: _ods_ir

Given a target value (e.g., vector) residing in a scf.for loop, this transform finds the corresponding xegpu.load_nd op and inserts xegpu.prefetch_nd operations for the tile. 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¶

mlir.dialects._xegpu_transform_ops_gen.xegpu_insert_prefetch(desc_op, target, *, dynamic_nb_prefetch=None, static_nb_prefetch=None, loc=None, ip=None) → _ods_ir¶

class mlir.dialects._xegpu_transform_ops_gen.SetDescLayoutOp(transformed, target, sg_layout, sg_data, inst_data, *, static_sg_layout=None, static_sg_data=None, static_inst_data=None, loc=None, ip=None)¶

Bases: _ods_ir

Given an xegpu.create_nd_desc operation, this transform adds xegpu.layout attribute to the result tensor descriptor. The layout is defined by the sg_layout, and sg_data and optional inst_data attributes. Returns a handle to the transformed op.

OPERATION_NAME = 'transform.xegpu.set_desc_layout'¶

_ODS_OPERAND_SEGMENTS¶

_ODS_REGIONS = (0, True)¶

target() → _ods_ir¶

sg_layout() → _ods_ir¶

sg_data() → _ods_ir¶

inst_data() → _ods_ir¶

static_sg_layout() → _ods_ir | None¶

static_sg_data() → _ods_ir | None¶

static_inst_data() → _ods_ir | None¶

transformed() → _ods_ir¶

mlir.dialects._xegpu_transform_ops_gen.xegpu_set_desc_layout(transformed, target, sg_layout, sg_data, inst_data, *, static_sg_layout=None, static_sg_data=None, static_inst_data=None, loc=None, ip=None) → _ods_ir¶

class mlir.dialects._xegpu_transform_ops_gen.SetGPULaunchThreadsOp(target, threads, *, static_threads=None, loc=None, ip=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¶

mlir.dialects._xegpu_transform_ops_gen.xegpu_set_gpu_launch_threads(target, threads, *, static_threads=None, loc=None, ip=None) → SetGPULaunchThreadsOp¶

class mlir.dialects._xegpu_transform_ops_gen.SetOpLayoutAttrOp(target, sg_layout, sg_data, inst_data, *, index=None, static_sg_layout=None, static_sg_data=None, static_inst_data=None, result=None, loc=None, ip=None)¶

Bases: _ods_ir

Sets the xegpu.layout attribute of an op. If result=true, sets the layout_result_{index}, otherwise layout_operand_{index} attribute. The target operand/result value is defined by the index argument. The layout is defined by the sg_layout, sg_data and optional inst_data attributes.

OPERATION_NAME = 'transform.xegpu.set_op_layout_attr'¶

_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¶

result() → bool¶: Shortcut to get an op result if it has only one (throws an error otherwise).

mlir.dialects._xegpu_transform_ops_gen.xegpu_set_op_layout_attr(target, sg_layout, sg_data, inst_data, *, index=None, static_sg_layout=None, static_sg_data=None, static_inst_data=None, result=None, loc=None, ip=None) → SetOpLayoutAttrOp¶