Working with the IR

Important

Loki is still under active development and has not yet seen a stable release. Interfaces can change at any time, objects may be renamed, or concepts may be re-thought. Make sure to sync your work to the current release frequently by rebasing feature branches and upstreaming more general applicable work in the form of pull requests.

The most important tool for working with Loki’s internal representation are utilities that traverse the IR to find specific nodes or patterns, to modify or replace subtrees, or to annotate the tree. In Loki there exist two types of tree traversal tools, depending on which level of the IR they operate on:

Visitors that traverse the tree of control flow nodes;
Mappers (following Pymbolic’s pymbolic.mapper naming convention) that traverse expression trees.

Visitors 

Loki’s visitors work by inspecting the type of each IR node they encounter and then selecting the best matching handler method for that node type. This allows implementing visitors that perform tasks either for very specific node types or generally applicable for any node type, depending on the handler’s name.

Loki includes a range of ready-to-use and configurable visitors for many common use cases, such as discovering certain node types, modifying or replacing nodes in the tree, or creating a string representation of the tree. For some use cases it may be easier to implement new visitors tailored to the task.

Searching the tree 

The first category of visitors traverses the IR and collects a list of results subject to certain criteria. In almost all cases FindNodes is the tool for that job with some bespoke variants for specific use cases.

`loki.ir.find.FindNodes`(match[, mode, greedy])	Find `Node` instances that match a given criterion.
`loki.ir.find.FindScopes`(match[, greedy])	Find all parent nodes for node `match`.
`loki.ir.find.SequenceFinder`(node_type)	Find repeated nodes of the same type in lists/tuples within a given tree.
`loki.ir.find.PatternFinder`(pattern)	Find a pattern of nodes given as tuple/list of types within a given tree.

A common pattern for using FindNodes is the following:

for loop in FindNodes((Loop, WhileLoop)).visit(routine.body):
    # ...do something with loop...

There are additional visitors that search all expression trees embedded in the control flow IR, which are explained further down.

Transforming the tree 

A core feature of Loki is the ability to transform the IR, which is done using the Transformer. It is a visitor that rebuilds the tree and replaces nodes according to a mapper.

`loki.ir.transformer.Transformer`([mapper, ...])	Visitor class to rebuild the tree and replace nodes according to a mapper.
`loki.ir.transformer.NestedTransformer`([...])	A `Transformer` that applies replacements in a depth-first fashion.
`loki.ir.transformer.MaskedTransformer`([...])	An enriched `Transformer` that can selectively include or exclude parts of the tree.
`loki.ir.transformer.NestedMaskedTransformer`([...])	A `MaskedTransformer` that retains parents for children that are included in the produced tree.

Transformer is commonly used in conjunction with FindNodes, with the latter being used to build the mapper for the first. The following example removes all loops over the horizontal dimension and replaces them by their body. This code snippet is a simplified version of a transformation used in ExtractSCATransformation:

routine = Subroutine(...)
horizontal = Dimension(...)

...

loop_map = {}
for loop in FindNodes(Loop).visit(routine.body):
    if loop.variable == horizontal.variable:
        loop_map[loop] = loop.body
routine.body = Transformer(loop_map).visit(routine.body)

Converting the tree to string 

The last step in a transformation pipeline is usually to write the transformed IR to a file. This is a task for Loki’s backends which themselves are subclasses of loki.visitors.pprint.Stringifier, yet another visitor. loki.visitors.pprint.Stringifier doubles as a pretty-printer for the IR that is useful for debugging.

`loki.ir.pprint.Stringifier`([depth, indent, ...])	Convert a given IR tree to a string representation.
`loki.ir.pprint`(ir[, stream])	Pretty-print the given IR using `Stringifier`.

Implementing new visitors 

Any new visitor should subclass Visitor (or any of its subclasses).

The common base class for all visitors is GenericVisitor, declared in loki.visitors that provides the basic functionality for matching objects to their handler methods. Derived from that is Visitor which adds a default handler visit_Node (for Node) and functionality to recurse for all items in a list or tuple and return the combined result.

To define handlers in new visitors, they should define visit_Foo methods for each class Foo they want to handle. If a specific method for a class Foo is not found, the MRO of the class is walked in order until a matching method is found (all the way until, for example, Visitor.visit_Node applies). The method signature is:

def visit_Foo(self, o, [*args, **kwargs]):
    pass

The handler is responsible for visiting the children (if any) of the node o. *args and **kwargs may be used to pass information up and down the call stack. You can also pass named keyword arguments, e.g.:

def visit_Foo(self, o, parent=None, *args, **kwargs):
    pass

Mappers 

Mappers are visitors that traverse expression trees.

They are built upon pymbolic.mapper classes and for that reason use a slightly different way of determining the handler methods: each expression tree node (pymbolic.primitives.Expression) holds a class attribute mapper_method with the name of the relevant method.

Loki provides, similarly to control flow tree visitors, ready-to-use mappers for searching or transforming expression trees, all of which are implemented in loki.expression.mappers. In addition, loki.expression.expr_visitors provides visitors that apply the same mapper to all expression trees in the IR.

Searching in expression trees 

The equivalent to FindNodes for expression trees is ExpressionRetriever. Using a generic function handle, (almost) arbitrary conditions can be used as a query that decides whether to include a given node into the list of results.

loki.expression.mappers.ExpressionRetriever(query)

A mapper for the expression tree that looks for entries specified by a query.

Note that mappers operate only on expression trees, i.e. using them directly is only useful when working with a single property of a control flow node, such as loki.ir.Assignment.rhs. If one wanted to search for expression nodes in all expression trees in the IR, e.g. to find all variables, bespoke visitors exist that apply ExpressionRetriever to all expression trees.

`loki.expression.expr_visitors.ExpressionFinder`([...])	Base class visitor to collect specific sub-expressions, eg.
`loki.expression.expr_visitors.FindExpressions`([...])	A visitor to collect all expression tree nodes (i.e., `pymbolic.primitives.Expression`) in an IR tree.
`loki.expression.expr_visitors.FindTypedSymbols`([...])	A visitor to collect all `TypedSymbol` used in an IR tree.
`loki.expression.expr_visitors.FindVariables`([...])	A visitor to collect all variables used in an IR tree
`loki.expression.expr_visitors.FindInlineCalls`([...])	A visitor to collect all `InlineCall` symbols used in an IR tree.
`loki.expression.expr_visitors.FindLiterals`([...])	A visitor to collect all literals (which includes `FloatLiteral`, `IntLiteral`, `LogicLiteral`, `StringLiteral`, and `IntrinsicLiteral`) used in an IR tree.

For example, the following finds all function calls embedded in expressions (InlineCall, as opposed to subroutine calls in CallStatement):

for call in FindInlineCalls().visit(routine.body):
    # ...do something with call...

Transforming expression trees 

Transformations of the expression tree are done very similar to Transformer, using the mapper SubstituteExpressionsMapper that is given a map to replace matching expression nodes.

`loki.expression.mappers.LokiIdentityMapper`([...])	A visitor to traverse and transform an expression tree
`loki.expression.mappers.SubstituteExpressionsMapper`(...)	A Pymbolic expression mapper (i.e., a visitor for the expression tree) that defines on-the-fly handlers from a given substitution map.

In the same way that searching can be done on all expression trees in the IR, transformations can be applied to all expression trees at the same time using SubstituteExpressions:

loki.expression.expr_visitors.SubstituteExpressions(...)

A dedicated visitor to perform expression substitution in all IR nodes

The following example shows how expression node discovery and substitution can be combined to replace all occurences of intrinsic function calls. (The code snippet is taken from replace_intrinsics, where two dict, function_map and symbol_map, provide the mapping to rename certain function calls that appear in routine.)

from loki.expression import symbols as sym

callmap = {}
for c in FindInlineCalls(unique=False).visit(routine.body):
    cname = c.name.lower()

    if cname in symbol_map:
        callmap[c] = sym.Variable(name=symbol_map[cname], scope=routine.scope)

    if cname in function_map:
        fct_symbol = sym.ProcedureSymbol(function_map[cname], scope=routine.scope)
        callmap[c] = sym.InlineCall(fct_symbol, parameters=c.parameters,
                                    kw_parameters=c.kw_parameters)

routine.body = SubstituteExpressions(callmap).visit(routine.body)

Converting expressions to string 

Every backend has their own mapper to convert expressions to a source code string, according to the corresponding language specification. All build on a common base class LokiStringifyMapper, which is also called automatically when converting any expression node to string.

loki.expression.mappers.LokiStringifyMapper(...)

A class derived from the default StringifyMapper that adds mappings for nodes of the expression tree that we added ourselves.