快捷方式

CUDA 内存操作

Tensor new_managed_tensor(const Tensor &self, const std::vector<std::int64_t> &sizes)

使用统一托管内存 (UVM) 分配 at::Tensor。然后将其首选存储位置设置为 CPU(主机内存),并在 CUDA 设备上建立到主机内存的映射。

参数:
  • self – 输入张量

  • sizes – 目标张量维度

返回:

由 UVM 支持的新张量

Tensor new_managed_tensor_meta(const Tensor &self, const std::vector<std::int64_t> &sizes)

用于 Meta 调度键的占位符操作。

参数:
  • self – 输入张量

  • sizes – 目标张量维度

返回:

一个新的空张量

Tensor new_host_mapped_tensor(const Tensor &self, const std::vector<std::int64_t> &sizes)

使用主机映射内存分配 at::Tensor

参数:
  • self – 输入张量

  • sizes – 目标张量维度

返回:

由主机映射内存支持的新张量

Tensor new_unified_tensor(const Tensor &self, const std::vector<std::int64_t> &sizes, bool is_host_mapped)

使用统一托管内存 (UVM) 或主机映射内存分配 at::Tensor

参数:
  • self – 输入张量

  • sizes – 目标张量维度

  • is_host_mapped – 是否分配 UVM 或主机映射内存

返回:

由 UVM 或主机映射内存支持的新张量,具体取决于 is_host_mapped 的值

Tensor new_unified_tensor_meta(const Tensor &self, const std::vector<std::int64_t> &sizes, bool is_host_mapped)

用于 new_unified_tensor 的 Meta 调度键的占位符操作

参数:
  • self – 输入张量

  • sizes – 目标张量维度

  • is_host_mapped – 是否分配 UVM 或主机映射内存

返回:

由 UVM 或主机映射内存支持的新张量,具体取决于 is_host_mapped 的值

Tensor new_vanilla_managed_tensor(const Tensor &self, const std::vector<std::int64_t> &sizes)

使用统一托管内存 (UVM) 分配 at::Tensor,但允许自动管理其首选存储位置。

参数:
  • self – 输入张量

  • sizes – 目标张量维度

返回:

由 UVM 支持的新张量

bool uvm_storage(const Tensor &self)

检查张量是否使用 UVM 分配(CPU 或 GPU 张量)。

参数:

self – 输入张量

返回:

true 如果张量使用 UVM 分配,否则为 false

bool is_uvm_tensor(const Tensor &self)

检查张量是否使用 UVM 分配,但不是 CPU 张量。

参数:

self – 输入张量

返回:

true 如果张量是使用 UVM 分配的非 CPU 张量,否则为 false

Tensor uvm_to_cpu(const Tensor &self)

将 UVM 张量转换为 CPU 张量。

参数:

self – 输入张量

返回:

一个新张量,它实际上是将输入从 UVM 移动到 CPU

Tensor uvm_to_device(const Tensor &self, const Tensor &prototype)

创建一个新的 UVM 张量,该张量与 prototype 共享相同的设备和 UVM 存储。

参数:
  • self – 输入张量

  • prototype – 目标张量,其设备和 UVM 存储将与新张量共享

返回:

一个新张量,它与 prototype 共享相同的设备和 UVM 存储。

void uvm_cuda_mem_advise(const Tensor &self, int64_t cuda_memory_advise)

在 UVM 张量的存储上调用 cudaMemAdvise()cudaMemoryAdvise 枚举在 Python 端的 fbgemm_gpu.uvm 命名空间中可用;有关有效值,请参阅那里的文档。

另请参阅

有关 cudaMemoryAdvise 枚举的更多信息,请参阅此处

参数:
  • self – 输入张量

  • cuda_memory_advisecudaMemoryAdvise 枚举值,为整数

void uvm_cuda_mem_prefetch_async(const Tensor &self, std::optional<Tensor> device_t)

在 UVM 张量的存储上调用 cudaMemPrefetchAsync(),以将内存预取到目标设备。

另请参阅

有关 cudaMemPrefetchAsync() 的更多信息,请参阅此处

参数:
  • self – 输入张量

  • device_t[可选] 张量,其设备将作为预取目标

void uvm_mem_advice_dont_fork(const Tensor &self)

在 UVM 张量的存储上调用 madvise(...MADV_DONTFORK)。这是为了解决一个问题,即 UVM 内核驱动程序在 fork 时从页表中取消映射 UVM 存储页,从而导致下次从 CPU 访问时速度减慢。

另请参阅

有关 madvise() 的更多信息,请参阅此处

参数:

self – 输入张量

Tensor uvm_to_cpu_clone(const Tensor &self)

将 UVM 张量的连续存储(uvm_storage(t) 为 true)复制到新的 CPU 张量中。复制操作使用单线程 memcpy()

参数:

self – 输入张量

返回:

包含从 UVM 张量复制的数据的新 CPU 张量

文档

访问 PyTorch 的全面开发者文档

查看文档

教程

获取面向初学者和高级开发者的深入教程

查看教程

资源

查找开发资源并获得问题解答

查看资源