Devices

Declaring devices

A program declares once where it runs. A .ss file declares it with a line at the top:

devices cuda[0]

The devices are cpu, mps (Apple GPUs) and cuda[...], with one or more different device numbers in the brackets, as in cuda[0] or cuda[0 1 2]. cpu[...], with two or more different numbers, is that many worker processes on the CPU, which stand in for GPUs: devices cpu[0 1 2] runs a program just as devices cuda[0 1 2] would, with three CPU processes in place of three GPUs, to test a program for several GPUs on a machine without them (Several devices). Anything else, including cpu[0], is E1001. Indented lines under the declaration are options, for running on several devices: shard params and split NAME, one per line (Options for several devices; anything else is E1002, and an option given twice E1903). With one device they do nothing, so a program keeps them when its devices line changes. A file has one devices line (E1003), before all of its declarations (E1004); comments and blank lines may come first.

From Python, shapescript.devices("cuda[0]") declares the devices, with the same text as the line. A notebook declares them with a devices line in a cell or from Python; a cell may repeat the devices an earlier cell declared, but not change them (E1003).

Devices at run time

There is no default device. A module with a devices line declares its devices when it is imported, and every ShapeScript function asks for the declared device when it is called, so calling one before anything has declared devices raises shapescript.DeviceError, which names both ways to declare them. Declaring the same devices again does nothing. Declaring different ones, whether from two modules or from a module and Python, raises DeviceError naming both declarations, and so does declaring a device the machine doesn’t have. With several devices, a module runs as Several devices says; a function compiled for one device, from a module without the same devices line, raises DeviceError when it’s called, rather than running on one device.

Every tensor ShapeScript code makes, such as causal[n m], a number that becomes a tensor, or a model’s field, is made on the declared device. A tensor argument must already be there: nothing is moved for you. With shape checks on (Runtime checks), each argument that is a tensor is checked, and one on another device raises DeviceError naming the argument and the device.

Building for a device

shapescript build --devices "cuda[0]" declares devices for a file without a devices line, as the line would; a file that has one must say the same (E1006). shapescript.compile(src, path, devices="cuda[0]") does the same from Python. A build with --standalone can’t ask the shapescript package which devices the program declared, so it needs a devices line or --devices (E1005), and it can’t be for several devices, whose runtime is the package’s (E1805). --devices "cuda[0 1 2]" compiles a file without a devices line for several devices, as the line would.