Skip to main content

Mojo struct

Device

Represents a logical instance of a device, for eg: CPU. This can be used to allocate and manage memory in a device's address space, and to compile and execute models and graphs on a device.

Implemented traits

AnyType, CollectionElement, Copyable, Movable, Stringable

Methods

__init__

__init__(inout self: Self)

Constructs a default initialized Device in a state that is only valid for deletion. Can be used to represent a 'moved from' state.

Use cpu_device() or cuda_device() to create a CPU or GPU Device.

__copyinit__

__copyinit__(inout self: Self, existing: Self)

Create a copy of the Device (bumping a refcount on the underlying Device).

Args:

  • existing (Self): Instance from which to copy.

__moveinit__

__moveinit__(inout self: Self, owned existing: Self)

Create a new Device and consume existing.

Args:

  • existing (Self): Instance from which to move from.

__del__

__del__(owned self: Self)

Destroys the device.

Note that any DeviceBuffer allocated on the Device will contain a reference to the Device, and the Device will only be de-allocated when all of its DeviceBuffers have also been destroyed.

__eq__

__eq__(self: Self, other: Self) -> Bool

Check if self and other point to the same underlying Device.

Args:

  • other (Self): Instance to compare against.

allocate

allocate(self: Self, spec: TensorSpec, name: Optional[String] = #kgen.none) -> DeviceTensor

Creates tensor allocated in the Device's address space.

Args:

  • spec (TensorSpec): TensorSpec descripting the shape and type of the tensor to allocate.
  • name (Optional[String]): An optional name for the DeviceTensor.

allocate(self: Self, bytecount: Int, name: Optional[String] = #kgen.none) -> DeviceMemory

Allocates a DeviceMemory object in the Device's address space.

Args:

  • bytecount (Int): The size of the memory to allocate in bytes.
  • name (Optional[String]): An optional name for the DeviceMemory.

Returns:

A DeviceMemory object allocated in the Device's address space.

__str__

__str__(self: Self) -> String

Returns a descriptor of the device.

Returns:

String representation of device.