Skip to main content
Log in

Mojo struct

DeviceMemory

struct DeviceMemory

DeviceMemory is an owning buffer allocated on a (possibly non-CPU) Device.

Fields

  • name (Optional[String]):
  • num_bytes (Int):

Implemented traits

AnyType, CollectionElement, Copyable, DeviceBuffer, Movable, StringableRaising, UnknownDestructibility, Writable

Methods

__init__

__init__(out self)

Constructs a DeviceMemory object in a state that is only valid for deletion. Can be used to represent a moved from state.

__init__(out self, num_bytes: Int, device: Device, name: Optional[String] = Optional(None))

Allocates DeviceMemory from the Device's address space.

Args:

  • num_bytes (Int): Size of the DeviceMemory buffer to allocate in bytes.
  • device (Device): Device on which to perform the allocation.
  • name (Optional[String]): Optional name for the DeviceMemory.

__init__[type: DType, rank: Int](out self, owned tensor: Tensor[type, rank])

Creates a DeviceMemory from the existing tensor storage.

Args:

  • tensor (Tensor[type, rank]): Tensor whose storage to use.

__init__(out self, owned anytensor: AnyTensor)

Creates a device tensor the existing anytensor storage.

Args:

  • anytensor (AnyTensor): AnyTensor whose storage to use.

__copyinit__

__copyinit__(out self, existing: Self)

__moveinit__

__moveinit__(out self, owned existing: Self)

__del__

__del__(owned self)

De-allocate and destroy the DeviceMemory.

Note: this will also decrement the refcount on the Device used to allocate the DeviceMemory.

copy

copy(self) -> Self

Explicitly construct a copy of self.

Returns:

A copy of this value.

bytecount

bytecount(self) -> Int

Returns the number of bytes in the DeviceMemory.

get_device

get_device(self) -> Device

Returns the device on which the DeviceMemory was allocated.

device

device(self) -> Device

Returns the device on which the DeviceMemory was allocated.

__str__

__str__(self) -> String

Returns a description of the DeviceMemory.

write_to

write_to[W: Writer](self, mut writer: W)

Formats a description of the DeviceMemory to the provided Writer.

Parameters:

  • W (Writer): A type conforming to the Writable trait.

Args:

  • writer (W): The object to write to.

copy_into

copy_into(self, mut dst_memory: Self)

Copies the contents of self into preallocated DeviceMemory.

Args:

  • dst_memory (Self): The destination DeviceMemory of the copy.

copy_to

copy_to(self, dev: Device, name: Optional[String] = Optional(None)) -> Self

Copies the contents of self into DeviceMemory allocated on dev. Note: this function allocates memory on dev.

Args:

  • dev (Device): The Device on which to allocate the new DeviceMemory.
  • name (Optional[String]): Optional name of the new DeviceMemory.

Returns:

Newly allocated DeviceMemory containing a copy of self's contents.

Raises:

If the DeviceMemory is backed by the same Device object as dev.

move_to

move_to(owned self, dev: Device) -> Self

Returns self if already allocated on dev, otherwise copy the contents of self to dev.

Args:

  • dev (Device): The Device on which the returned buffer is allocated.

Returns:

DeviceMemory located in dev's address space.

unsafe_ptr

unsafe_ptr(self) -> UnsafePointer[SIMD[uint8, 1]]

Returns a pointer to the underlying device memory.

Note: The caller is responsible for ensuring that the returned pointer is not used after its owner is last used.

take

take(mut self) -> Self

Takes and returns the contents of self, leaving self in an empty but destructible state.