mlir.dialects.transform.xegpu¶

Classes¶

`ConvertLayoutOp`	Specialization for ConvertLayoutOp class.
`GetDescOp`	Specialization for GetDescOp class.
`InsertPrefetchOp`	Specialization for InsertPrefetchOp class.
`SetDescLayoutOp`	Specialization for SetDescLayoutOp class.
`SetGPULaunchThreadsOp`	Specialization for SetGPULaunchThreadsOp class.
`SetOpLayoutAttrOp`	Specialization for SetOpLayoutAttrOp class.
`GetDescOp`	Specialization for GetDescOp class.
`SetDescLayoutOp`	Specialization for SetDescLayoutOp class.
`SetOpLayoutAttrOp`	Specialization for SetOpLayoutAttrOp class.
`SetGPULaunchThreadsOp`	Specialization for SetGPULaunchThreadsOp class.
`InsertPrefetchOp`	Specialization for InsertPrefetchOp class.
`ConvertLayoutOp`	Specialization for ConvertLayoutOp class.

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)
`get_desc_op`(→ OpResult)
`set_desc_layout`(→ OpResult)
`set_op_layout_attr`(→ SetOpLayoutAttrOp)
`set_gpu_launch_threads`(→ SetGPULaunchThreadsOp)
`insert_prefetch`(→ OpResult)
`convert_layout`(→ ConvertLayoutOp)

Module Contents¶

class mlir.dialects.transform.xegpu.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.transform.xegpu.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.transform.xegpu.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.transform.xegpu.xegpu_get_desc_op(desc_handle, target, *, loc=None, ip=None) → _ods_ir¶

class mlir.dialects.transform.xegpu.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.transform.xegpu.xegpu_insert_prefetch(desc_op, target, *, dynamic_nb_prefetch=None, static_nb_prefetch=None, loc=None, ip=None) → _ods_ir¶

class mlir.dialects.transform.xegpu.SetDescLayoutOp(transformed, target, sg_layout, sg_data, inst_data, *, static_sg_layout=None, static_sg_data=None, static_inst_data=None, slice_dims=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. If slice_dims is provided, the xegpu.layout attribute is wrapped in an xegpu.slice<..., dims=slice_dims> attribute. 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¶

slice_dims() → _ods_ir | None¶

transformed() → _ods_ir¶

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

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

class mlir.dialects.transform.xegpu.SetOpLayoutAttrOp(target, sg_layout, sg_data, inst_data, *, index=None, static_sg_layout=None, static_sg_data=None, static_inst_data=None, slice_dims=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. If slice_dims is provided, the xegpu.layout attribute is wrapped in an xegpu.slice<..., dims=slice_dims> attribute.

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¶

slice_dims() → _ods_ir | None¶

result() → bool¶

mlir.dialects.transform.xegpu.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, slice_dims=None, result=None, loc=None, ip=None) → SetOpLayoutAttrOp¶

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

Bases: GetDescOp

Specialization for GetDescOp class.

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

class mlir.dialects.transform.xegpu.SetDescLayoutOp(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, slice_dims: mlir.dialects._ods_common.MixedInt | None = None, loc=None, ip=None)¶

Bases: SetDescLayoutOp

Specialization for SetDescLayoutOp class.

mlir.dialects.transform.xegpu.set_desc_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, slice_dims: mlir.dialects._ods_common.MixedInt | None = None, loc=None, ip=None) → OpResult¶

class mlir.dialects.transform.xegpu.SetOpLayoutAttrOp(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, slice_dims: mlir.dialects._ods_common.MixedInt | None = None, index: int | Attribute | None = None, result: bool | Attribute | None = None, loc=None, ip=None)¶

Bases: SetOpLayoutAttrOp

Specialization for SetOpLayoutAttrOp class.

mlir.dialects.transform.xegpu.set_op_layout_attr(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, slice_dims: mlir.dialects._ods_common.MixedInt | None = None, index: int | Attribute | None = None, result: bool | Attribute | None = None, loc=None, ip=None) → SetOpLayoutAttrOp¶

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: 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: 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, 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, loc=None, ip=None) → ConvertLayoutOp¶