tt-mlir documentation

emitpy.assign_global (tt::emitpy::AssignGlobalOp)

Assign a value to a global variable

The emitpy.assign_global assignes the value to an existing named global variable.

Example:

emitpy.assign_global @value = %arg0 : !emitpy.opaque<"[ttnn.Tensor]">

Interfaces: SymbolUserOpInterface

Attributes:

Operands:

emitpy.assign (tt::emitpy::AssignOp)

Assign operation

Syntax:

operation ::= `emitpy.assign` $value attr-dict `:` functional-type(operands, results)

The emitpy.assign operation represents a Python variable assignment. This models new_var = old_var or var = constant.

Example:

%2 = emitpy.assign %1 : <!emitpy.opaque<"ttnn.Tensor">>

// Code emitted for the operation above.
v2 = v1;

Operands:

Results:

emitpy.call_opaque (tt::emitpy::CallOpaqueOp)

Opaque call operation

Syntax:

operation ::= `emitpy.call_opaque` $callee `(` $operands `)` attr-dict `:` functional-type($operands, results)

The emitpy.call_opaque operation represents a Python function call. The callee can be an arbitrary non-empty string.

Example:

%2 = emitpy.call_opaque "ttnn.add"(%0, %1) {args = [0 : index, 1 : index, #emitpy.opaque<"ttnn.DataType.BFLOAT16">, #emitpy.opaque<"ttnn.MemoryConfig(ttnn.TensorMemoryLayout.INTERLEAVED, ttnn.BufferType.DRAM, None)">], keyword_args = ["", "", "dtype", "memory_config"]} : (!emitpy.opaque<"ttnn.Tensor">, !emitpy.opaque<"ttnn.Tensor">) -> !emitpy.opaque<"ttnn.Tensor">

Interfaces: PyExpressionInterface

Attributes:

Operands:

Results:

emitpy.class (tt::emitpy::ClassOp)

Class definition

The emitpy.class operation represents a Python class definition. It is a symbol and contains a single-region body with class members.

Class verification enforces:

• At most one __init__ method.
• __init__ returns nothing and is an instance method.
• Methods can be annotated with emitpy.method_kind attribute on func.func using one of: instance, staticmethod, classmethod.
• Instance and class methods must have a receiver argument as the first argument; if named, it must be self (instance) or cls (classmethod).

Example:

emitpy.class @Model(#emitpy.opaque<"torch.nn.Module">) {
  func.func @__init__(%self: !emitpy.class<"Model">) {
    %w = emitpy.call_opaque "ttnn.load_weight"() : () -> !emitpy.opaque<"ttnn.Tensor">
    emitpy.set_attr %self, "weight", %w : (!emitpy.class<"Model">, !emitpy.opaque<"ttnn.Tensor">)
    return
  }
}

Traits: IsolatedFromAbove, NoTerminator, SingleBlock, SymbolTable

Interfaces: Symbol

Attributes:

emitpy.constant (tt::emitpy::ConstantOp)

Constant operation

The emitpy.constant operation produces an SSA value equal to some constant specified by an attribute. This can be used to form simple integer and floating point constants, as well as more exotic things like tensor constants.

Traits: ConstantLike

Attributes:

Results:

emitpy.create_dict (tt::emitpy::CreateDictOp)

Create dictionary operation

Syntax:

operation ::= `emitpy.create_dict` $dict_name (`(` $items^ `)`)? attr-dict `:` functional-type(operands, results)

Creates a Python dictionary with a specified name.

Can be created in two ways:

1. From key-value pairs (items operands) - pairs are alternating key, value
2. From a Python literal expression (literal_expr attribute)

When literal_expr is provided, items must be empty. When items are provided, they must be even count (key-value pairs).

Example:

// Empty dict using literal_expr
%dict1 = emitpy.create_dict "my_dict" {literal_expr = "{}"} : () -> !emitpy.dict

// Dict from key-value operands
%dict2 = emitpy.create_dict "cache" (%k0, %v0) : (index, !emitpy.opaque<"None">) -> !emitpy.dict

// Dict from literal expression (most compact for complex dicts)
%dict3 = emitpy.create_dict "_CONST_EVAL_CACHE" {literal_expr = "{i: None for i in range(133)}"} : () -> !emitpy.dict

// Code emitted for the operations above:
my_dict = {}
cache = {0: None}
_CONST_EVAL_CACHE = {i: None for i in range(133)}

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Attributes:

Operands:

Results:

emitpy.expression (tt::emitpy::ExpressionOp)

Expression operation to inline Python code

The emitpy.expression operation returns a single SSA value yielded by its single-basic-block region. The operations inside the region must form a DAG (directed acyclic graph) of computations that produces a single result value. The expression's operands are mapped to block arguments and can be used within the expression body.

When generating Python code, operations within the expression body are emitted inline without intermediate variable assignments, producing cleaner code.

The optional do_not_inline attribute can be used to force the translator to emit the expression as separate statements rather than inline.

Example:

%expression = emitpy.expression(%a, %b) : (!emitpy.opaque, !emitpy.opaque) -> !emitpy.opaque {
^bb0(%arg0: !emitpy.opaque, %arg1: !emitpy.opaque):
  %0 = emitpy.call_opaque "ttnn.mul"(%arg0, %arg1) : (!emitpy.opaque, !emitpy.opaque) -> !emitpy.opaque
  %1 = emitpy.call_opaque "ttnn.add"(%0, %arg1) : (!emitpy.opaque, !emitpy.opaque) -> !emitpy.opaque
  emitpy.yield %1 : !emitpy.opaque
}
return %expression : !emitpy.opaque

This would generate:

return ttnn.add(ttnn.mul(a, b), b)

If the do_not_inline attribute is set:

%expression = ttnn.add(ttnn.mul(a, b), b)
return %expression

Traits: HasOnlyGraphRegion, IsolatedFromAbove, SingleBlockImplicitTerminator<YieldOp>, SingleBlock

Attributes:

Operands:

Results:

emitpy.get_attr (tt::emitpy::GetAttrOp)

Get attribute from an object

Syntax:

operation ::= `emitpy.get_attr` $object `,` $attr_name attr-dict `:` functional-type(operands, results)

The emitpy.get_attr operation represents Python attribute access via the dot operator.

Example:

%w = emitpy.get_attr %self, "weight" : (!emitpy.class<"Model">) -> !emitpy.opaque<"ttnn.Tensor">

# Code emitted:
w = self.weight

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface), PyExpressionInterface

Effects: MemoryEffects::Effect{}

Attributes:

Operands:

Results:

emitpy.get_value_for_dict_key (tt::emitpy::GetValueForDictKeyOp)

Get dictionary value for a given key

Syntax:

operation ::= `emitpy.get_value_for_dict_key` $dict `[` $key `]` attr-dict `:` functional-type(operands, results)

Retrieves a value for a given key in a dictionary. The key must be either an index type or an opaque type representing a string.

Example:

%dict = emitpy.global_statement @_CONST_EVAL_CACHE : !emitpy.dict
%key = emitpy.literal "5" : index
%tensors = emitpy.get_value_for_dict_key %dict[%key] : (!emitpy.dict, index) -> !emitpy.opaque<"[ttnn.Tensor]">

// Code emitted:
tensors = _CONST_EVAL_CACHE[5]

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface)

Effects: MemoryEffects::Effect{}

Operands:

Results:

emitpy.global (tt::emitpy::GlobalOp)

A global variable

Syntax:

operation ::= `emitpy.global` $sym_name custom<EmitPyGlobalOpInitialValue>($initial_value) attr-dict

The emitpy.global operation defines a named global variable.

Example:

emitpy.global @x : #emitpy.opaque<"[]">

Interfaces: Symbol

Attributes:

emitpy.global_statement (tt::emitpy::GlobalStatementOp)

Marks a variable as global for the current function scope

Syntax:

operation ::= `emitpy.global_statement` $name `:` type($result) attr-dict

The emitpy.global_statement operation represents a Python global statement.

It marks an identifier, specified by its symbol @name, as referring to a variable in the global (module) scope, rather than a new local variable.

Example:

%0 = emitpy.global_statement @x : !emitpy.opaque<"[ttnn.Tensor]">

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface), SymbolUserOpInterface

Effects: MemoryEffects::Effect{}

Attributes:

Results:

emitpy.import (tt::emitpy::ImportOp)

Import operation

The emitpy.import operation allows to define a Python module import via various forms of the import statement.

Example:

 emitpy.import import "ttnn"

Attributes:

emitpy.literal (tt::emitpy::LiteralOp)

Literal operation

Syntax:

operation ::= `emitpy.literal` $value attr-dict `:` type($result)

The emitpy.literal operation produces an SSA value equal to some constant specified by an attribute.

Example:

%0 = emitpy.literal "0" : index

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface), PyExpressionInterface

Effects: MemoryEffects::Effect{}

Attributes:

Results:

emitpy.set_attr (tt::emitpy::SetAttrOp)

Set attribute on an object

Syntax:

operation ::= `emitpy.set_attr` $object `,` $attr_name `,` $value attr-dict `:` `(` type($object) `,` type($value) `)`

The emitpy.set_attr operation represents Python attribute assignment via the dot operator.

Example:

emitpy.set_attr %self, "weight", %w : (!emitpy.opaque<"Model">, !emitpy.opaque<"ttnn.Tensor">)

# Code emitted:
self.weight = w

Attributes:

Operands:

emitpy.set_value_for_dict_key (tt::emitpy::SetValueForDictKeyOp)

Set dictionary value for a given key

Syntax:

operation ::= `emitpy.set_value_for_dict_key` $dict `[` $key `]` `=` $value attr-dict `:` `(` type($dict) `,` type($key) `,` type($value) `)`

Changes or inserts a value in the dictionary for a given key. The key must be either an index type or an opaque type representing a string.

Example:

%dict = emitpy.global_statement @_CONST_EVAL_CACHE : !emitpy.dict
%key = emitpy.literal "5" : index
emitpy.set_value_for_dict_key %dict[%key] = %tensors : (!emitpy.dict, index, !emitpy.opaque<"[ttnn.Tensor]">)

// Code emitted:
_CONST_EVAL_CACHE[5] = tensors

Operands:

emitpy.subscript (tt::emitpy::SubscriptOp)

Subscript operation

Syntax:

operation ::= `emitpy.subscript` $value `[` $index `]` attr-dict `:` functional-type(operands, results)

With the emitpy.subscript operation the subscript operator [] can be applied to variables or arguments of opaque type.

Example:

%0 = emitpy.literal "0" : index
%1 = emitpy.subscript %arg0[%0] : (!emitpy.opaque<"[ttnn.Tensor]">, index) -> !emitpy.opaque<"ttnn.Tensor">

Traits: AlwaysSpeculatableImplTrait

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface), PyExpressionInterface

Effects: MemoryEffects::Effect{}

Operands:

Results:

emitpy.verbatim (tt::emitpy::VerbatimOp)

Verbatim operation

Syntax:

operation ::= `emitpy.verbatim` $value (`args` $fmtArgs^ `:` type($fmtArgs))? attr-dict

The emitpy.verbatim operation produces no results and the value is emitted as is followed by a line break ('\n' character) during translation.

This operation can be used in situations where a more suitable operation is not yet implemented in the dialect.

Note: Use with caution. This operation can have arbitrary effects on the semantics of the emitted code. Use semantically more meaningful operations whenever possible. Additionally this op is NOT intended to be used to inject large snippets of code.

Attributes:

Operands:

emitpy.yield (tt::emitpy::YieldOp)

Yield operation for expression termination

Syntax:

operation ::= `emitpy.yield` $result attr-dict `:` type($result)

The emitpy.yield terminates its parent EmitPy operation's region, optionally yielding an SSA value. The semantics of how the values are yielded is defined by the parent operation.

Example:

%0 = emitpy.expression(%arg0, %arg1) : (!emitpy.opaque, !emitpy.opaque) -> !emitpy.opaque {
^bb0(%a: !emitpy.opaque, %b: !emitpy.opaque):
  %1 = emitpy.call_opaque "foo"(%a, %b) : (!emitpy.opaque, !emitpy.opaque) -> !emitpy.opaque
  emitpy.yield %1 : !emitpy.opaque
}

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

Interfaces: ConditionallySpeculatable, NoMemoryEffect (MemoryEffectOpInterface), RegionBranchTerminatorOpInterface

Effects: MemoryEffects::Effect{}

Operands: