Python API Reference

Python API is used to control & access SPU, for example, to do data infeed/outfeed, to compile an XLA program to PPHlo, or to fire a PPHlo on an SPU runtime.

Runtime Setup

class spu.Runtime(link: Context, config: RuntimeConfig)

The SPU Virtual Machine Slice.

clear() → None

Delete all SPU values.

del_var(name: str) → None

Delete an SPU value.

Args:

name (str): Id of the value.

get_var(name: str) → bytes

Get an SPU value.

Args:

name (str): Id of value.

Returns:

spu_pb2.ValueProto: Data data.

get_var_meta(name: str) → ValueMeta

Get an SPU value without content.

Args:

name (str): Id of value.

Returns:

spu_pb2.ValueProto: Data with out content.

run(executable: ExecutableProto) → None

Run an SPU executable.

Args:

executable (spu_pb2.ExecutableProto): executable.

set_var(name: str, value: bytes) → None

Set an SPU value.

Args:

name (str): Id of value. value (spu_pb2.ValueProto): value data.

Runtime IO

class spu.Io(world_size: int, config: RuntimeConfig)

The SPU IO interface.

make_shares(x: np.ndarray, vtype: spu_pb2.Visibility, owner_rank: int = -1) → List[bytes]

Convert from NumPy array to list of SPU value(s).

Args:

x (np.ndarray): input. vtype (spu_pb2.Visibility): visibility. owner_rank (int): the index of the trusted piece. if >= 0, colocation optimization may be applied.

Returns:

[spu_pb2.ValueProto]: output.

reconstruct(str_shares: List[bytes]) → np.ndarray

Convert from list of SPU value(s) to NumPy array.

Args:

xs (spu_pb2.ValueProto]): input.

Returns:

np.ndarray: output.

Compiler

spu.compile(ir_text: str, ir_type: str, vis: List[spu_pb2.Visibility]) → str

Compile from textual HLO/MHLO IR to SPU bytecode.

Args:

ir_text (str): textual HLO/MHLO IR protobuf binary format. ir_type (str): “hlo” or “mhlo”. vtype (spu_pb2.Visibility): Visbilities .

Returns:

[spu_pb2.ValueProto]: output.

Simulation