Device

class Device

Compute backend a Tensor’s storage lives on.

One instance per physical device (CPU, or a specific CUDA or HIP GPU); obtained through cpu_device(), cuda_device() or hip_device() rather than constructed directly. Every device-dependent tensor operation (allocation, copy, the elementary kernels used to train the network globals) is dispatched through this interface.

Public Functions

virtual ~Device() = default

virtual std::pair<void*, Tensor> allocate(std::size_t size, AllocHint hint) const = 0

Allocate size bytes for hint; the returned Tensor, if not empty, is the pool allocation this storage was carved out of.

virtual void free(void *ptr) const = 0

Free a pointer returned by allocate.

inline virtual void free_on_stream(void *ptr, void *stream) const

Free a pointer returned by allocate after the work queued on stream.

inline virtual void *default_free_stream() const

Stream to free stream-ordered storage on if it has no known stream.

inline virtual void order_streams(void *from, void *to) const

Make work queued on to from now on wait for the work queued on from.

virtual void memcpy(void *to, void *from, std::size_t size) const = 0

Copy size bytes from from to to, both on this device.

virtual void tensor_copy(const Tensor &source, Tensor &target) const = 0

Copy the elements of source into target.

virtual void tensor_zero(Tensor &tensor) const = 0

Set every element of tensor to zero.

virtual void tensor_add(const Tensor &source, Tensor &target) const = 0

Add source into target element-wise.

virtual void tensor_cpu(const Tensor &source, Tensor &target) const = 0

Copy source into target, which lives on the CPU.

virtual const Device *device_ptr() const = 0

This device, as the pointer type used elsewhere in the API.

inline virtual void sync_barrier() const

Block until every operation queued on this device has completed.

virtual DeviceType device_type() const = 0

The kind of device this is.

inline virtual int device_index() const

Index of the device among those of its type, 0 for the CPU.

virtual void activate() const = 0

Make this device current for the calling thread.

virtual void adam_step(const Tensor &gradient, Tensor &parameter, Tensor &exp_avg, Tensor &exp_avg_sq, double step_size, double beta1, double beta2, double eps, double bias_corr2_sqrt, double weight_decay) const = 0

One Adam optimizer update step; see AdamOptimizer.

Public Static Attributes

static constexpr bool stream_ordered_alloc