'io_parameters' Dialect

External parameter resource management APIs.

Parameters are externalized storage for resources that are asynchronously accessible and device-aware. Parameters can be read or written on the same device timelines as the operations that consume or produce them and with locality pinning to ensure memory doesn't need to move. Parameters are referenced by a scope and a key, with the scope being optional but strongly recommended as a way to distinguish sets of parameters that may exist when multiple model parts are compiled together and would otherwise collide.

Parameters are provided by a few operations implementing a virtual interface and can support shared parameters (same storage used in multiple contexts, or outliving a single instantiation in a context), in-memory caches, memory-mapped files (including directly using the mapped memory for execution when devices support it), iree_hal_file_t usage for device-supported I/O, and parameter subsetting for things like runtime sharding.

Alongside read(+load) and write operations gather and scatter allow for batching of large numbers of reads and writes into/from single buffers. For parameter providers that can batch operations this allows for a handful (~1-4) of calls out to perform many more operations (~thousands). Modeling the gather/scatter also gives us a point where we could extract the mapping and use it to repack files/defrag memory in the future.

See io_parameters.imports.mlir for the full list of exported functions.

• 'io_parameters' Dialect
  • Operations
    • Parameter I/O ops
      • io_parameters.gather (IO::Parameters::GatherOp)
      • io_parameters.load (IO::Parameters::LoadOp)
      • io_parameters.scatter (IO::Parameters::ScatterOp)
  • Enums
    • AccessScopeBitfield
      • Cases:
    • BufferUsageBitfield
      • Cases:
    • CallingConvention
      • Cases:
    • CollectiveElementType
      • Cases:
    • CollectiveKind
      • Cases:
    • CollectiveReductionOp
      • Cases:
    • CommandBufferModeBitfield
      • Cases:
    • CommandCategoryBitfield
      • Cases:
    • DescriptorFlags
      • Cases:
    • DispatchFlags
      • Cases:
    • ExecutionBarrierFlagBitfield
      • Cases:
    • ExecutionStageBitfield
      • Cases:
    • FenceFlagBitfield
      • Cases:
    • MemoryAccessBitfield
      • Cases:
    • MemoryModel
      • Cases:
    • MemoryTypeBitfield
      • Cases:
    • PipelineLayoutFlags
      • Cases:
    • DescriptorType
      • Cases:

Operations

Parameter I/O ops

Ops parameter I/O.

io_parameters.gather (IO::Parameters::GatherOp)

Gathers multiple parameters from a parameter scope

Syntax:

operation ::= `io_parameters.gather` `<` $device `:` type($device) `>`
              `affinity` `(` $queue_affinity `)`
              `wait` `(` $wait_fence `)`
              `signal` `(` $signal_fence `)`
              `{`
              custom<ParameterGatherOperations>(
              $source_scope, $source_keys, $source_offsets,
              $target_buffer, type($target_buffer), $target_offsets, $target_lengths)
              `}`
              attr-dict-with-keyword

Asynchronously gathers one or more parameters into a single target buffer. This is equivalent to one read per parameter but allows implementations that can batch operations to do so without additional overhead.

Traits: AttrSizedOperandSegments

Attributes:

Attribute	MLIR Type	Description
source_scope	::mlir::StringAttr	string attribute
source_keys	::mlir::ArrayAttr	string array attribute

Operands:

Operand	Description
device	device
queue_affinity	64-bit signless integer
wait_fence	fence
signal_fence	fence
source_offsets	variadic of 64-bit signless integer
target_buffer	buffer
target_offsets	variadic of index
target_lengths	variadic of index

io_parameters.load (IO::Parameters::LoadOp)

Reads one or more parameters from a parameter scope

Syntax:

operation ::= `io_parameters.load` `<` $device `:` type($device) `>`
              `affinity` `(` $queue_affinity `)`
              `wait` `(` $wait_fence `)`
              `signal` `(` $signal_fence `)`
              `type` `(` $memory_types `)`
              `usage` `(` $buffer_usage `)`
              `{`
              custom<ParameterLoadOperations>(
              $source_scope, $source_keys, $source_offsets,
              type($results), $lengths)
              `}`
              attr-dict-with-keyword

Asynchronously reads one or more parameters from an external parameter provider and returns the resulting buffers. Depending on the parameter and buffer types this may alias existing cached storage or be directly mapped to the parameter origin or result in a copy as if an allocate + read had been used.

Traits: AttrSizedOperandSegments

Interfaces: Util_SizeAwareOp

Attributes:

Attribute	MLIR Type	Description
source_scope	::mlir::StringAttr	string attribute
source_keys	::mlir::ArrayAttr	string array attribute
memory_types	mlir::iree_compiler::IREE::HAL::MemoryTypeBitfieldAttr	valid MemoryType
buffer_usage	mlir::iree_compiler::IREE::HAL::BufferUsageBitfieldAttr	valid BufferUsage

Operands:

Operand	Description
device	device
queue_affinity	64-bit signless integer
wait_fence	fence
signal_fence	fence
source_offsets	variadic of 64-bit signless integer
lengths	variadic of index

Results:

Result	Description
results	variadic of buffer

io_parameters.scatter (IO::Parameters::ScatterOp)

Scatters multiple parameters to a parameter scope

Syntax:

operation ::= `io_parameters.scatter` `<` $device `:` type($device) `>`
              `affinity` `(` $queue_affinity `)`
              `wait` `(` $wait_fence `)`
              `signal` `(` $signal_fence `)`
              `{`
              custom<ParameterScatterOperations>(
              $source_buffer, type($source_buffer), $source_offsets, $source_lengths,
              $target_scope, $target_keys, $target_offsets)
              `}`
              attr-dict-with-keyword

Asynchronously scatters one or more parameters from a single source buffer into one or more parameters. This is equivalent to one write per parameter but allows implementations that can batch operations to do so without additional overhead.

Traits: AttrSizedOperandSegments

Attributes:

Attribute	MLIR Type	Description
target_scope	::mlir::StringAttr	string attribute
target_keys	::mlir::ArrayAttr	string array attribute

Operands:

Operand	Description
device	device
queue_affinity	64-bit signless integer
wait_fence	fence
signal_fence	fence
source_buffer	buffer
source_offsets	variadic of index
source_lengths	variadic of index
target_offsets	variadic of 64-bit signless integer

Enums

AccessScopeBitfield

valid AccessScope

Cases:

Symbol	Value	String
None	0	None
IndirectCommandRead	1	IndirectCommandRead
ConstantRead	2	ConstantRead
DispatchRead	4	DispatchRead
DispatchWrite	8	DispatchWrite
TransferRead	16	TransferRead
TransferWrite	32	TransferWrite
HostRead	64	HostRead
HostWrite	128	HostWrite
MemoryRead	256	MemoryRead
MemoryWrite	512	MemoryWrite

BufferUsageBitfield

valid BufferUsage

Cases:

Symbol	Value	String
None	0	None
TransferSource	1	TransferSource
TransferTarget	2	TransferTarget
Transfer	3	Transfer
DispatchIndirectParams	256	DispatchIndirectParams
DispatchUniformRead	512	DispatchUniformRead
DispatchStorageRead	1024	DispatchStorageRead
DispatchStorageWrite	2048	DispatchStorageWrite
DispatchStorage	3072	DispatchStorage
DispatchImageRead	4096	DispatchImageRead
DispatchImageWrite	8192	DispatchImageWrite
DispatchImage	12288	DispatchImage
SharingExport	65536	SharingExport
SharingReplicate	131072	SharingReplicate
SharingConcurrent	262144	SharingConcurrent
SharingImmutable	524288	SharingImmutable
MappingScoped	16777216	MappingScoped
MappingPersistent	33554432	MappingPersistent
MappingOptional	67108864	MappingOptional
MappingAccessRandom	134217728	MappingAccessRandom
MappingAccessSequentialWrite	268435456	MappingAccessSequentialWrite
Mapping	150994944	Mapping

CallingConvention

Calling conversions for linked functions

Cases:

Symbol	Value	String
Default	0	Default

CollectiveElementType

valid CollectiveElementType

Cases:

Symbol	Value	String
Sint8	0	si8
Uint8	1	ui8
Sint16	2	si16
Uint16	3	ui16
Sint32	4	si32
Uint32	5	ui32
Sint64	6	si64
Uint64	7	ui64
Float16	8	f16
Float32	9	f32
Float64	10	f64
BFloat16	11	bf16
Float8E5M2	12	f8E5M2
Float8E4M3	13	f8E4M3
Float8E5M2FNUZ	14	f8E5M2FNUZ
Float8E4M3FNUZ	15	f8E4M3FNUZ

CollectiveKind

valid CollectiveKind

Cases:

Symbol	Value	String
AllGather	0	all_gather
AllReduce	1	all_reduce
AllToAll	2	all_to_all
Broadcast	3	broadcast
Reduce	4	reduce
ReduceScatter	5	reduce_scatter
Send	6	send
Recv	7	recv
SendRecv	8	send_recv

CollectiveReductionOp

valid CollectiveReductionOp

Cases:

Symbol	Value	String
None	0	none
ReductionSum	1	sum
ReductionProduct	2	product
ReductionMinimum	3	minimum
ReductionMaximum	4	maximum
ReductionAverage	5	average

CommandBufferModeBitfield

valid CommandBufferMode

Cases:

Symbol	Value	String
None	0	None
OneShot	1	OneShot
AllowInlineExecution	16	AllowInlineExecution

CommandCategoryBitfield

valid CommandCategory

Cases:

Symbol	Value	String
None	0	None
Transfer	1	Transfer
Dispatch	2	Dispatch

DescriptorFlags

valid Descriptor flags

Cases:

Symbol	Value	String
None	0	None
ReadOnly	1	ReadOnly
Indirect	2	Indirect

DispatchFlags

valid dispatch flags

Cases:

Symbol	Value	String
None	0	None
Reserved	1	Reserved

ExecutionBarrierFlagBitfield

valid ExecutionBarrierFlag

Cases:

Symbol	Value	String
None	0	None
Reserved	1	Reserved

ExecutionStageBitfield

valid ExecutionStage

Cases:

Symbol	Value	String
None	0	None
CommandIssue	1	CommandIssue
CommandProcess	2	CommandProcess
Dispatch	4	Dispatch
Transfer	8	Transfer
CommandRetire	16	CommandRetire
Host	32	Host

FenceFlagBitfield

valid FenceFlag

Cases:

Symbol	Value	String
None	0	None
Reserved	1	Reserved

MemoryAccessBitfield

valid MemoryAccess

Cases:

Symbol	Value	String
None	0	None
Read	1	Read
Write	2	Write
Discard	4	Discard
MayAlias	8	MayAlias
Unaligned	16	Unaligned
Any	32	Any

MemoryModel

IREE HAL MemoryModel

Cases:

Symbol	Value	String
Unified	0	Unified
Discrete	1	Discrete

MemoryTypeBitfield

valid MemoryType

Cases:

Symbol	Value	String
None	0	None
Optimal	1	Optimal
HostVisible	2	HostVisible
HostCoherent	4	HostCoherent
HostCached	8	HostCached
HostLocal	70	HostLocal
DeviceVisible	16	DeviceVisible
DeviceLocal	48	DeviceLocal

PipelineLayoutFlags

valid PipelineLayout flags

Cases:

Symbol	Value	String
None	0	None
Indirect	1	Indirect

DescriptorType

valid DescriptorType

Cases:

Symbol	Value	String
UniformBuffer	6	uniform_buffer
StorageBuffer	7	storage_buffer