Skip to main content

/

Docs

Log in

Mojo 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

Methods​

__init__​

__init__(inout self: Self)

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

__init__(inout self: Self, num_bytes: Int, device: Device, name: Optional[String] = #kgen.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](inout self: 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__(inout self: Self, owned anytensor: AnyTensor)

Creates a device tensor the existing anytensor storage.

Args:

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

__copyinit__​

__copyinit__(inout self: Self, existing: Self)

__moveinit__​

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

__del__​

__del__(owned self: Self)

De-allocate and destroy the DeviceMemory.

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

bytecount​

bytecount(self: Self) -> Int

Returns the number of bytes in the DeviceMemory.

get_device​

get_device(self: Self) -> Device

Returns the device on which the DeviceMemory was allocated.

device​

device(self: Self) -> Device

Returns the device on which the DeviceMemory was allocated.

__str__​

__str__(self: Self) -> String

Returns a description of the DeviceMemory.

copy_into​

copy_into(self: Self, inout 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: Self, dev: Device, name: Optional[String] = #kgen.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: 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: Self) -> UnsafePointer[SIMD[uint8, 1], 0, 0, alignof[::AnyType,__mlir_type.!kgen.target]() if triple_is_nvidia_cuda() else 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(inout self: Self) -> Self

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

Was this page helpful?