11. Context manager

PopXL contains a number of context managers.

You can get the graph context with:

  • get_current_graph() (or gcg()): returns the current graph from the current context.

  • get_main_graph() (or gmg()): returns the main graph of the current context.

You can also use the following context managers in popxl:

  • ipu(): Sets the IPU id on ops created in this context. Internally, it uses the Poplar virtual graph id as the IPU id. This indicates which IPU the operations within the context should be mapped to. For instance, you can use with ipu(0) to include the operations you want to put on the first IPU requested.

    The usage of virtual graph in PopXL is similar to the virtual graph in Poplar, except that PopXL requires only one virtual computational graph per IPU.

  • in_sequence(): Forces operations created in its context to be executed in the same order as they are created. In Listing 11.1, in_sequence() guarantees host_restore is executed after the copy_var_update. copy_var_update_ updates the value of x in-place with b. So the value of x is 5. If in_sequence() is not used, this will cause ambiguity with the in-place operation between copy_var_update and host_store.

    Listing 11.1 Example of using in_sequence context manager
    13with main:
14    x = popxl.variable(1, popxl.float32)
15    b = popxl.constant(5, popxl.float32)
16
17    with popxl.in_sequence():
18        ops.var_updates.copy_var_update_(x, b)
19        # host store
20        o_d2h = popxl.d2h_stream(x.shape, x.dtype, name="output_stream")
21        ops.host_store(o_d2h, x)
22
23# run the model
24with popxl.Session(ir, "ipu_model") as session:
25    outputs = session.run()

    Download in_sequence.py

  • name_scope(): sets the name scope for the operations created within this context.

  • io_tiles(): executes for the operations created within this context on the I/O tiles of the current IPU.

You can also use the contexts merge_exchange() and io_tile_exchange() to do transforms on graphs. Section 10, Transforms contains details.