MLIR

Multi-Level IR Compiler Framework

'ml_program' Dialect

The MLProgram dialect contains structural operations and types for defining a compiled Machine-Learning program, as created from common ML frameworks, such as TensorFlow, PyTorch, JAX, etc. It does not itself define computation ops common to such frameworks but establishes a common programming model for establishing modules, functions, globals and memory model components appropriate for such an abstract level of detail.

This dialect is under active development, and while stability is an eventual goal, it is not guaranteed at this juncture. Given the early state, it is recommended to inquire further prior to using this dialect.

Operations 

source

ml_program.func (ml_program::FuncOp) 

Function containing a single SSACFG region

This simple function container represents callables in an ML program where the body is an SSACFG region. It must be terminated by a return op which yields values with the same arity and types as the FunctionType results of the containing func.

This op is a Symbol but does not introduce a new SymbolTable. As such, it cannot represent nested symbols.

Example:

ml_program.func private @some_extern(i32) -> i32
ml_program.func @compute(%arg0 : i32) -> i32 {
  ml_program.return %arg0 : i32
}

Traits: IsolatedFromAbove

Interfaces: CallableOpInterface, FunctionOpInterface, RegionKindInterface, Symbol

Attributes: 

AttributeMLIR TypeDescription
sym_name::mlir::StringAttrstring attribute
function_type::mlir::TypeAttrtype attribute of function type
arg_attrs::mlir::ArrayAttrArray of dictionary attributes
res_attrs::mlir::ArrayAttrArray of dictionary attributes
sym_visibility::mlir::StringAttrstring attribute

ml_program.global (ml_program::GlobalOp) 

Module level declaration of a global variable

Syntax:

operation ::= `ml_program.global` custom<SymbolVisibility>($sym_visibility)
              (`mutable` $is_mutable^)?
              $sym_name ``
              custom<TypedInitialValue>($type, $value)
              attr-dict

Declares a named global variable (or constant).

A global contains a value of a specified type which can be accessed at runtime via appropriate load/store operations. It can be mutable or constant, optionally taking an initial value or declared as extern (in which case, the initial value is found in external storage by symbol name).

Generally, the type of the global and the type of the initial value will be the same. However, for type hierarchies which can have a more generalized bounding type that can be assigned from a narrow type, this is allowed (but not verified).

Examples:

// Constant global.
ml_program.global @foobar(dense<4> : tensor<4xi32>) : tensor<?xi32>

// Constant with external linkage.
ml_program.global mutable @foobar(#ml_program.extern<tensor<4xi32>>)
  : tensor<?xi32>

// Mutable global with an undefined initial value.
ml_program.global mutable @foobar : tensor<?xi32>

Interfaces: Symbol

Attributes: 

AttributeMLIR TypeDescription
sym_name::mlir::StringAttrstring attribute
type::mlir::TypeAttrany type attribute
is_mutable::mlir::UnitAttrunit attribute
value::mlir::Attributeany attribute
sym_visibility::mlir::StringAttrstring attribute

ml_program.global_load (ml_program::GlobalLoadOp) 

Direct load of a mutable value from a global

Syntax:

operation ::= `ml_program.global_load` $global `:` type($result) attr-dict

Performs a non-atomic, non-volatile, non-synchronized load from a global that may be mutable.

It is fully expected that these constraints are not suitable for all situations, and alternative ops should be defined and used for more advanced cases.

This op is side effecting and may not be valid to use in graph regions without additional consideration to evaluation order constraints. See global_load_graph for op which allows for explicit ordering constraints.

Example:

%0 = ml_program.global_load @foobar : tensor<?xi32>

Interfaces: OpAsmOpInterface, SymbolUserOpInterface

Attributes: 

AttributeMLIR TypeDescription
global::mlir::SymbolRefAttrsymbol reference attribute

Results: 

ResultDescription
resultany type

ml_program.global_load_const (ml_program::GlobalLoadConstOp) 

Direct load a constant value from a global

Syntax:

operation ::= `ml_program.global_load_const` $global `:` type($result) attr-dict

Loads a constant (immutable) value from a global directly by symbol.

This op is only legal for globals that are not mutable and exists because such a load can be considered to have no side effects.

Example:

%0 = ml_program.global_load_const @foobar : tensor<?xi32>

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface), OpAsmOpInterface, SymbolUserOpInterface

Effects: MemoryEffects::Effect{}

Attributes: 

AttributeMLIR TypeDescription
global::mlir::SymbolRefAttrsymbol reference attribute

Results: 

ResultDescription
resultany type

ml_program.global_load_graph (ml_program::GlobalLoadGraphOp) 

Direct load of a mutable value from a global in Graph region

Syntax:

operation ::= `ml_program.global_load_graph` $global `` custom<TokenOrdering>($consumeTokens, type($produceToken)) `:` type($result) attr-dict

Performs a non-atomic, non-volatile, non-synchronized load from a global that may be mutable.

It is fully expected that these constraints are not suitable for all situations, and alternative ops should be defined and used for more advanced cases.

This op is side effecting and may not be valid to use in graph regions without additional consideration to evaluation order constraints.

Example:

%0, %cstr = ml_program.global_load_graph @foobar
  ordering (%token -> !ml_program.token) : tensor<?xi32>

Interfaces: SymbolUserOpInterface

Attributes: 

AttributeMLIR TypeDescription
global::mlir::SymbolRefAttrsymbol reference attribute

Operands: 

OperandDescription
consumeTokensvariadic of Token for establishing execution ordering in a graph

Results: 

ResultDescription
resultany type
produceTokenToken for establishing execution ordering in a graph

ml_program.global_store (ml_program::GlobalStoreOp) 

Direct store of a value into a mutable global

Syntax:

operation ::= `ml_program.global_store` $global `=` $value `:` type($value) attr-dict

Performs a non-atomic, non-volatile, non-synchronized store to a mutable global.

It is fully expected that these constraints are not suitable for all situations, and alternative ops should be defined and used for more advanced cases.

This op is side effecting and may not be valid to use in graph regions without additional consideration to evaluation order constraints. See global_store_graph for op which allows for explicit ordering constraints.

Example:

ml_program.global_store @foobar = %0 : tensor<?xi32>

Interfaces: SymbolUserOpInterface

Attributes: 

AttributeMLIR TypeDescription
global::mlir::SymbolRefAttrsymbol reference attribute

Operands: 

OperandDescription
valueany type

ml_program.global_store_graph (ml_program::GlobalStoreGraphOp) 

Direct store of a value into a mutable global

Syntax:

operation ::= `ml_program.global_store_graph` $global `=` $value `` custom<TokenOrdering>($consumeTokens, type($produceToken)) `:` type($value) attr-dict

Performs a non-atomic, non-volatile, non-synchronized store to a mutable global.

It is fully expected that these constraints are not suitable for all situations, and alternative ops should be defined and used for more advanced cases.

This op is side effecting and may not be valid to use in graph regions without additional consideration to evaluation order constraints.

Example:

%token = ml_program.global_store @foobar = %0 : tensor<?xi32>
  ordering (%in_token -> !ml_program.token) : tensor<?xi32>

Interfaces: SymbolUserOpInterface

Attributes: 

AttributeMLIR TypeDescription
global::mlir::SymbolRefAttrsymbol reference attribute

Operands: 

OperandDescription
valueany type
consumeTokensvariadic of Token for establishing execution ordering in a graph

Results: 

ResultDescription
produceTokenToken for establishing execution ordering in a graph

ml_program.output (ml_program::OutputOp) 

Outputs values from a subgraph function

Syntax:

operation ::= `ml_program.output` attr-dict ($operands^ `:` type($operands))?

The output operation terminates a subgraph by yielding values to the caller. The operation takes variable number of operands and produces no results. The operand number and types must match the signature of the function that contains the operation.

Traits: AlwaysSpeculatableImplTrait, HasParent<SubgraphOp>, ReturnLike, Terminator

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface), RegionBranchTerminatorOpInterface

Effects: MemoryEffects::Effect{}

Operands: 

OperandDescription
operandsvariadic of any type

ml_program.return (ml_program::ReturnOp) 

Returns values from a func function

Syntax:

operation ::= `ml_program.return` attr-dict ($operands^ `:` type($operands))?

The return operation terminates a func function by yielding values to the caller. The operation takes variable number of operands and produces no results. The operand number and types must match the signature of the function that contains the operation.

Traits: AlwaysSpeculatableImplTrait, HasParent<FuncOp>, ReturnLike, Terminator

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface), RegionBranchTerminatorOpInterface

Effects: MemoryEffects::Effect{}

Operands: 

OperandDescription
operandsvariadic of any type

ml_program.subgraph (ml_program::SubgraphOp) 

An function containing a single Graph region

This simple function container represents callables in an ML program where the body is a Graph region containing a single block. It must be terminated by an output op which yields values with the same arity and types as the FunctionType results of the containing subgraph.

This op is a Symbol but does not introduce a new SymbolTable. As such, it cannot represented nested symbols.

Example:

ml_program.subgraph private @some_extern(i32) -> i32
ml_program.subgraph @compute(%arg0 : i32) -> i32 {
  ml_program.output %arg0 : i32
}

Traits: HasOnlyGraphRegion, IsolatedFromAbove, SingleBlock

Interfaces: CallableOpInterface, FunctionOpInterface, RegionKindInterface, Symbol

Attributes: 

AttributeMLIR TypeDescription
sym_name::mlir::StringAttrstring attribute
function_type::mlir::TypeAttrtype attribute of function type
arg_attrs::mlir::ArrayAttrArray of dictionary attributes
res_attrs::mlir::ArrayAttrArray of dictionary attributes
sym_visibility::mlir::StringAttrstring attribute

ml_program.token (ml_program::TokenOp) 

Produces a new token value

Syntax:

operation ::= `ml_program.token` attr-dict

Token values are used to chain side effecting ops in a graph so as to establish an execution order. This op produces a token.

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Results: 

ResultDescription
tokenToken for establishing execution ordering in a graph

Attributes 

ExternAttr 

Value used for a global signalling external resolution

Syntax:

#ml_program.extern<
  ::mlir::Type   # type
>

When used as the value for a GlobalOp, this indicates that the actual value should be resolved externally in an implementation defined manner. The sym_name of the global is the key for locating the value.

Examples:

extern : tensor<4xi32>

Parameters: 

ParameterC++ typeDescription
type::mlir::Type

Types 

TokenType 

Token for establishing execution ordering in a graph

Syntax: !ml_program.token