Skip to main content

/

Docs

Nightly
v24.5
Log in

Mojo struct

DeviceTensor

Fields​

  • ​spec (TensorSpec):

Implemented traits​

AnyType, CollectionElement, Copyable, DeviceBuffer, Movable, StringableRaising

Methods​

__init__​

__init__(inout self: Self, spec: TensorSpec, device: Device, name: Optional[String])

Allocates a DeviceTensor in the Device's address space.

Args:

  • ​spec (TensorSpec): TensorSpec describing the dtype and shape of the DeviceTensor.
  • ​device (Device): Device on which to perform the allocation.
  • ​name (Optional[String]): Optional name for the DeviceMemory.

__init__(inout self: Self)

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

__init__(inout self: Self, owned storage: DeviceMemory, spec: TensorSpec)

Constructs a DeviceTensor from an existing storage buffer and spec.

Args:

  • ​storage (DeviceMemory): The storage backing the DeviceTensor.
  • ​spec (TensorSpec): TensorSpec describing the type and shape of the DeviceTensor.

__copyinit__​

__copyinit__(inout self: Self, existing: Self)

__moveinit__​

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

copy_to​

copy_to(self: Self, dev: Device, name: Optional[String] = #kgen.none) -> Self

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

Args:

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

Returns:

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

Raises:

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

copy_into​

copy_into(self: Self, inout dst_tensor: Self)

Copies the contents of self into a preallocated DeviceTensor.

Args:

  • ​dst_tensor (Self): The destination DeviceTensor of the copy.

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:

A DeviceTensor 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 DeviceTensor's storage in device memory.

to_tensor​

to_tensor[type: DType, rank: Int](owned self: Self) -> Tensor[$0, $1]

Returns a Tensor created using the DeviceTensor's shape and storage.

device​

device(self: Self) -> Device

Returns the Device on which the DeviceTensor was allocated.

name​

name(self: Self) -> Optional[String]

Returns the name of the DeviceTensor.

__str__​

__str__(self: Self) -> String

Returns a descriptor for the DeviceTensor.

bytecount​

bytecount(self: Self) -> Int

Returns the number of bytes in the DeviceTensor.

take​

take(inout self: Self) -> Self

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