'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)

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