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.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!
If you'd like to share more information, please report an issue on GitHub
π What went wrong?