【翻译】Pytorch机制，源代码分析与内存管理调研

翻译自原文：Pytorch机制，源代码分析与内存管理调研。

这篇文章也主要参考了PyTorch internals : ezyang’s blog [PPT]

• 翻译1：一文搞懂 PyTorch 内部机制
• 翻译2：万字综述，核心开发者全面解读PyTorch内部机制

ezyang 的另一篇文章：Let’s talk about the PyTorch dispatcher : ezyang’s blog

• 翻译1：【译】聊聊Pytorch Dispatcher
• 翻译2：让我们谈谈 PyTorch dispatcher

1 Pytorch 发布版本构成

从 GitHub pytorch 克隆的软件仓库与使用pip install 或 conda install 下载的软件包不同。实际上，前者包含许多基于 C/C++ 的文件，构成了 PyTorch 的基础，而后者更加简洁，包含了编译后的库和 dll 文件。
在此，让我们首先讨论发行版本，或者安装包。该安装包有许多组件，我在这里只挑选一些最重要的部分进行说明。

1.1 nn

所有深度学习层的 Python 入口都位于这里。它们主要从输入的初始化参数中收集参数，并对输入数据进行一些修改。之后，它会将核心计算操作以及参数一起发送给基于 torch._C 的函数执行。

1.2 autograd

这部分包含一系列基础函数，用于实现反向传播。如果深入研究，可以发现核心实现仍然来自 C 库。Variable 的封装也在这里，但由于张量（tensor）和 变量（Variable）的合并，现在已经被省略掉了。

1.3 CUDA

主要的部分都包含在 cuda 文件夹中：Stream（流）、Event（事件）、Broadcast（广播）和 Random（随机）。

• CUDA 流是属于特定设备的线性执行序列，与其他流相互独立。
• CUDA 事件是同步标记，可用于监视设备的进度、准确测量时间并同步 CUDA 流。
• 广播相关的函数主要确保操作在不同的 GPU 上运行并正确地进行聚合。

1.4 optim

torch.optim 是一个实现各种优化算法的包。其中已经支持了最常用的方法，比如 adam、sgd 和 adagrad。

1.5 distributed

distributions 包包含了可参数化的概率分布和采样函数。这允许构建随机计算图和用于优化的随机梯度估计器。

1.6 onnx

torch.onnx 模块包含将模型导出为 ONNX IR 格式的函数。这些模型可以使用 ONNX 库加载，然后转换为在其他深度学习框架上运行的模型。

1.7 tensor

这里定义了最基本的张量类。它继承自 C 库中的一个超类，称为 torch._C._TensorBase。它附加了许多方法，如register_hook、resize 和 norm 等，这些方法最终调用基于 C 的库来执行操作。

1.8 lib

这个库存放了编译后的 C/C++ 文件。其中包括 .dll 文件和 .lib 文件。根据在谷歌上的bug报告，我相信 .dll 文件是专门为了与 Windows 兼容性而编译的，而 .lib 文件可以在 Linux 中使用，其中一些也可以在 Windows 中使用。（如果你找到更准确的解释，请告诉我。这些文件包括：_C.lib，c10.lib，torch.lib，c10_cuda.lib。

1.9 functional

与张量操作相关的函数都位于这里。实际上，它们再次是来自 C 库函数的封装。在这个文件中，你可以找到一些函数，比如 tensordot、unique、split 等。

1.10 utils

这里包含了各种实用工具代码。这包括与数据集相关的代码如 dataloader.py、dataset.py、sampler.py，还包括保存和输出相关的 checkpoint.py。一些 TensorBoard 支持的代码也可以在这里找到。

2 Pytorch 如何管理其内部资源

2.1 什么是张量（Tensor）？

在数学中，张量是一种代数对象，描述了从一组代数对象到另一组代数对象的线性映射。张量可以映射的对象包括但不限于向量和标量，甚至可以递归地映射到其他张量。在 PyTorch 中，张量是中心数据结构。它是一个 n 维数据结构，其中包含某种标量类型，例如浮点数、整数等。我们可以将张量视为由一些数据组成，并包含一些元数据，描述张量的大小、元素类型（dtype），张量存储在哪个设备上（CPU内存、CUDA内存）等信息。

2.2 张量如何组织？

TH 库负责张量的计算、存储和内存管理。它将 Tensor 分为两个独立的部分：存储和访问/视图。
存储：THStorage。它管理张量的存储方式。
访问：THTensor。它为用户提供访问张量的接口。

2.2.1 数据存储

typedef struct THStorage    // Tensor Storage，PyTorch 中用于存储 Tensor 数据的结构体
{
 real *data;                
 ptrdiff_t size;            
 int refcount;              // 使用引用计数方法来进行自动垃圾回收。当引用计数变为 0 时，该结构将自动释放。
 char flag;
 THAllocator *allocator;
 void *allocatorContext;
 struct THStorage *view;
} THStorage;

CPU 中的所有 Tensor 实际上是指向内存中类似结构的 C 指针。它使用**引用计数（reference count）**来进行内存管理。

2.2.2 数据访问

typedef struct THTensor
{
 long *size;                // 包含所有维度的长度信息
 long *stride;              // 每个维度的大小
 int nDimension;            // 维度的数量
 // Attention: storage->size might be bigger than the size of tensor.
 THStorage *storage;        // 该数据结构的指针
 ptrdiff_t storageOffset;
 int refcount;              // 引用计数
 char flag;
} THTensor;

2.2.3 内存分配器

/c10/core/Allocator.h

通过定义 Allocator 抽象类和相关函数，PyTorch 提供了一种标准化的接口，使得不同的内存分配器可以实现自己的分配和释放逻辑，并与 PyTorch 的内存管理系统进行交互。这样，PyTorch 可以根据具体的场景和需求，选择合适的内存分配器来进行内存管理。

#include <memory>
struct C10_API Allocator {                        // Allocator 抽象类定义了用于分配和释放内存的接口。
  virtual ~Allocator() = default;                 // 虚析构函数
  // 纯虚函数，用于分配 n 字节大小的内存。
  // 派生类需要实现该函数来提供具体的内存分配逻辑，并返回一个指向已分配内存的 DataPtr。
  // DataPtr 是一个智能指针，用于管理已分配内存的生命周期。
  virtual DataPtr allocate(size_t n) const = 0;
  // 虚函数，默认实现为返回空指针。
  // 派生类可以选择性地实现该函数来提供具体的释放内存的逻辑。
  // DeleterFnPtr 是一个函数指针类型，用于释放内存。
  virtual DeleterFnPtr raw_deleter() const {
    return nullptr;
  }
  // 帮助函数，用于直接调用派生类的 allocate 函数来分配内存，并返回指向已分配内存的指针。
  void* raw_allocate(size_t n) {
    auto dptr = allocate(n);
    AT_ASSERT(dptr.get() == dptr.get_context());
    return dptr.release_context();
  }
  // 帮助函数，用于调用派生类的 raw_deleter()函数来释放内存。
  void raw_deallocate(void* ptr) {
    auto d = raw_deleter();
    AT_ASSERT(d);
    d(ptr);
  }
};

allocate 函数直接包含在头文件 memory 中。

/aten/src/TH/THAllocator.cpp

at::DataPtr THMapAllocator::makeDataPtr(const char 
filename, int flags, size_t size, size_t
 actual_size_out) {
  auto* context = new THMapAllocator(filename, flags, size);
  if (actual_size_out) *actual_size_out = context->size();
  return {context->data(), context, &deleteTHMapAllocator, at::DeviceType::CPU};
}

默认分配器是 malloc/free 分配器。在分配失败时，malloc 和 realloc 会引发错误（使用 THError）。

2.3 理解参数

理解步幅（stride）和存储偏移（storage offset）是困难且不够直观，因此让我们借用一些来自 PyTorch 内部开发者 ezyang 的图片来详细阐述这个问题。
张量是一个数学概念。但为了在计算机上表示它们，我们必须为它们定义一种物理表示方式。最常见的表示方式是在内存中连续地布置张量的每个元素（这就是连续性一词的含义），将每一行写入内存。

请注意尺寸（sizes）和步幅（strides）之间的关系。如果我们得到一个尺寸为（D，H，W）的张量，并且这个张量是用户直接定义的，而不是某个切片或操作的结果，那么它的步幅将是（H*W，W，1）。你可以进行比较并得出结论。在某个维度上，每个步幅元素都是其后面所有维度的乘积，并且最后一个维度的步幅为1。

在物理上，步幅表示计算机在获取下一个对应维度的起始位置时需要跳过多少块内存。如果我们使用一个公式来计算 $Tensor[i,j,k]$ 的内存位置，它将是 $storageOffset+i\times stride[0]+j\times stride[1]+k\times stride[2]$。其中，storageOffset 表示起始偏移量，stride 表示步幅。

在上面的示例图像中，我指定了张量包含32位整数，因此您可以看到每个整数位于物理地址上，彼此相隔四个字节的偏移量。为了记住张量的实际维度，我们还必须记录尺寸作为额外的元数据。

然后是内存偏移（memory offset）。这是什么意思？正如我们之前提到的，张量存储可以支持多个张量视图，如果我们对前 N 个元素进行切片，那么我们将从第 N+1 个内存位置开始。下面的示例将进一步解释这个概念。

在左侧的示例中，可以看到我们从第三个元素块开始，这意味着我们跳过了两个块，因此偏移量为2。由于切片的存在，二维张量变成了一维张量，连续的元素在物理存储中是连续的，这意味着步幅为[1]。在这种情况下，尺寸表示元素的数量，即2。

在右侧的示例中，连续的元素不是连续的，但是它们确实从开头开始，所以步幅为[2]，偏移量为0。总共仍然有两个元素，所以尺寸没有改变。

另外，如果您仍然发现有些难以理解，您可以尝试使用这个网站来调整这些参数并查看动态过程。

2.4 张量实现调度

正如我们所知，尽管在 Python 中，您可以根据需要使用任何类型的数据，因为解释器会处理其余的事情。然而，由于基本内核是用 C/C++ 编写的，因此需要将 Python 中的函数派发到具有不同输入和设备类型的相同函数中。对于 C/C++ 函数而言，某个函数不能同时将 int 和 float Tensor 作为相同的 x 进行处理，它们需要分别实现。

2.5 如何调度

正如我们上面所讨论的，基本的 C/C++ 实现需要根据数据和设备类型进行派发。那么在代码中，实际上该如何完成这项工作呢？

基本上有三种方法。

1. 将这些函数根据不同的数据和设备类型分别手动编写。
2. 使用模板函数在编译时构建这些派发函数。但这种方法只适用于 C++，而 PyTorch 中的许多代码仍然是用 C 编写的。
3. 应用宏技巧。通过将函数名定义为宏，并接受一个或多个参数，比如数据类型名，我们可以通过多次使用 #define 和 #undef，在函数名宏中设置变量为不同的类型名，将函数编译成支持不同类型的多个副本。

这里是一个简化的示例：

File structure:
.
├── add.c # Used to extend generic/add.c
├── add.h # Used to extend generic/add.h
├── general.h # Including other header files
└── generic
 ├── add.c # Definition of generic function add
 └── add.h # Definition of generic type Vector

add.h

// add.h
#include "general.h"
#define CONCAT_2_EXPAND(A, B) A ## B                // 宏定义，将两个参数进行连接
#define CONCAT_2(A, B) CONCAT_2_EXPAND(A, B)        // 将两个参数传递给前一个宏定义进行展开
#define CONCAT_3_EXPAND(A, B, C) A ## B ## C        // 将三个参数进行连接
#define CONCAT_3(A, B, C) CONCAT_3_EXPAND(A, B, C)  // 将三个参数传递给前一个宏定义进行展开
#define Vector_(NAME) CONCAT_3(Num, Vector_, NAME)  // 使用CONCAT_3宏来将Num、Vector_和NAME三个参数进行连接
#define Vector CONCAT_2(Num, Vector)                // 使用CONCAT_2宏来将Num和Vector两个参数进行连接
#define num float        // 将num替换为float，使用num作为特定类型(float)的别名
#define Num Float        // 将Num替换为Float，使用Num作为特定类型(Float)的别名
#include "generic/add.h"
// 取消了之前定义的num和Num的宏定义
#undef num
#undef Num
#define num double
#define Num Double
#include "generic/add.h"
#undef num
#undef Num

add.c

// add.c
#include "add.h"
#define num float
#define Num Float
#include "generic/add.c"
#undef num
#undef Num
#define num double
#define Num Double
#include "generic/add.c"
#undef num
#undef Num

generic/add.h

// generic/add.h
typedef struct Vector
{
num *data;
int n;
} Vector;
extern void Vector_(add)(Vector *C, Vector *A, Vector *B);

generic/add.c

// generic/add.c
void Vector_(add)(Vector *C, Vector *A, Vector *B)
{
int i, n;
n = C->n;
for(i=0;i<n;i++)
{
C->data[i] = A->data[i] + B->data[i];
}
}

3 一个查找 THStorage 的示例

我试图找到 THStorage 的定义，因为它将使我们对 PyTorch 的文件管理结构有一个简要了解，并且我们还可以对这些宏和包含的形式构建这个庞大项目有一个基本的概念。我们从 torch/csrc/Storage.cpp 开始，逐步检查所包含的文件。

Storage.cpp                 ->
#include <TH/TH.h>          ->
#include <TH/THStorageFunction.h>   ->
#include <TH/generic/THStorage.h>   ->
#include <c10/core/StorageImpl.h>

在 TH/generic/THStorage.h 中查找宏定义：

#define THStorage at::StorageImpl    // 将 THStorage 替换为 at::StorageImpl

在 c10/core/StorageImpl.h 中查找结构体定义：

namespace c10 {    // 位于 c10 命名空间下
struct C10_API StorageImpl final : public c10::intrusive_ptr_target { // StorageImpl 结构体继承自 c10::intrusive_ptr_target 类
...
private:
  caffe2::TypeMeta  data_type_;  // Data type
  DataPtr data_ptr_;             // Data pointer
  int64_t numel_;                // Data number
  bool resizable_;
  bool received_cuda_;
  Allocator* allocator_;         // Data's allocator
};
}

因此，THWStorage 的实际类型是 at::StorageImpl，它是数据存储的实现。首先，让我们查看 THPStorage_(pynew) 的定义，在没有提供 cdata 的情况下，它需要使用函数 THWStorage_(NAME) 来创建类 THWStorage 的实现，其中 NAME 的值可能是：

new                // New a THStorage, if size not specified, size=0, that means using default Allocator
free
size
get
set
data
newWithSize        // New THStorage，specify size but use default Allocator
newWithAllocator   // New THStorage，specify size and Allocator
copy_functions
copyByte
...
copyCudaByte
...

还有一些宏定义：

这些宏定义用于生成具体的函数和方法名称，方便使用不同的类型和操作来定义相应的存储函数。

#define THWStorage_(NAME) THStorage_(NAME)     // torch/csrc/THP.h
#define THStorage_(NAME) TH_CONCAT_4(TH,Real,Storage_,NAME)   // TH/THStorageFunctions.h

THStorage_(NAME) 函数的声明位于 TH/generic/THStorage.h、TH/generic/THStorageCopy.h 文件中，而实现部分则位于相应的 cpp 文件中。
（顺便提一下，如果使用 CUDA，#define THWStorage_(NAME) 的声明位于 THC/generic/THCStorage.h 和 THC/generic/THCStorageCopy.h 中）
以 THStorage_(newWithSize) 函数为例，查看 TH/generic/THStorage.cpp 文件，我们可以找到如下定义：

// 创建一个具有指定大小的 THStorage 对象
THStorage* THStorage_(newWithSize)(ptrdiff_t size)
{
  // 使用 c10::make_intrusive 函数创建一个 at::StorageImpl 对象，
  // 并将其存储在 THStorage* 类型的 storage 指针中
  THStorage* storage = c10::make_instrusive
[at::StorageImpl](at::StorageImpl)
(
#ifdef THQUANTIZED
    caffe2::TypeMeta::Make<quantized_t>(),
#else
    caffe2::TypeMeta::Make<scalar_t>(),        // New a scalar_t type
#endif
    size,                        // 指定创建的 THStorage 对象的大小
    getTHDefaultAllocator(),     // 获取默认的分配器，用于为存储对象分配内存
    true).release();             // 指定创建的 THStorage 对象是否需要自动释放内存。
                                 // release()：释放 make_intrusive 返回的所有权
  return storage;
}

从这段代码块可以推断出，它使用 new 关键字创建了一个 StorageImpl 对象，并且使用一个侵入式指针指向其中一个，最后返回一个指向 StorageImpl 的指针，然后销毁了这个侵入式指针。宏 THStorage 等同于 at::StorageImpl，因此这个方法简单地创建了一个 StorageImpl 对象并返回一个指向它的指针。根据 c10::make_intrusive 的定义，这个工作实际上是由 StorageImpl 的构造函数完成的，具体的构造函数如下所示：

StorageImpl(
    caffe2::TypeMeta data_type,
    int64_4 numel,
    at::Allocator* allocator,
    bool resizable)
...

在这里我们只追踪到此处，展示了 PyTorch 内部代码如何调用和实现这些方法的一个典型示例。

4 Autograd（自动微分）

自动微分（Autograd）是一种支持自动计算梯度的方法，它在反向传播中被使用。Autograd 直接依赖于计算图（Computational Graph）。计算图被用于定义模型的流程。它将函数和变量组合起来，并展示它们之间的连接关系。

计算图是一个具有以下特性的有向图：

• 边（Edge）：表示函数或函数的依赖关系
• 带有输入边的节点：表示函数（或操作符）
• 带有输出边的节点：表示变量

计算图有两种主要类型，动态计算图和静态计算图。TensorFlow使用静态计算图，它具有以下特点：

• 首先定义图的结构，然后为叶子节点赋值（这是占位符的原始目的）
• 然后根据叶子节点的赋值进行前向计算

另一方面，PyTorch使用动态图。图的结构与前向计算同时建立，因此不需要使用占位符。

下面是 autograd 的内部代码示例。

我们将详细解释涉及到这个过程中的这些参数。

• Data（数据）：它是一个变量所持有的数据。
• requires_grad（是否需要计算梯度）：如果设置为 True，则开始跟踪所有操作的历史记录，并形成反向图用于梯度计算。
• grad（梯度值）：grad 存储了梯度的值。
  • 如果 requires_grad 为 False，则它将持有 None 值。
  • 即使 requires_grad 为 True，除非从其他节点调用了 .backward() 函数，否则它也将持有 None 值。
• grad_fn（梯度计算的反向函数）：用于计算梯度的反向函数。
• is_leaf（节点是否为叶节点）：一个节点是叶节点的条件：
  1. 它由一些函数显式初始化，比如 x = torch.tensor(1.0) 或 x = torch.randn(1, 1)（基本上是本文开头讨论的所有张量初始化方法之一）。
  2. 它在 requires_grad = False 的张量上进行的操作后创建。
  3. 它是在某个张量上调用了 .detach() 方法创建的。

5 Pytorch 源代码的构成

由于支持不同的数据类型和设备，并且 Python 代码调用了 C/C++ 代码，源代码结构不容易理解。下面是根目录中最重要的部分。

提供更详细的目录说明和解释如下。

★5.1 重要文件夹的解释

5.1.1 C10

Caffe Tensor Library：最基础的张量库。这里的代码可以部署到移动设备和服务器上。它包含了 PyTorch 的核心抽象，包括张量（Tensor）和储存（Storage）数据结构的实际实现。

5.1.2 ATen

A TENsor library 是一个针对 C11 的张量库，也是 PyTorch 的 C 张量库。它是一个 C++ 库，用于实现张量的操作。如果你正在寻找某些内核代码的位置，那么它很可能在 ATen 中。ATen 本身分为两个操作符的区域：有现代化的、C++ 实现的“本地”操作符，以及有遗留的、C 实现的“遗留”操作符（TH、THC、THNN、THCUNN）。遗留操作符被认为是较差的区域，如果可能的话，尽量不要花太多时间在那里。

5.1.3 Caffe2

这部分来自原始的Caffe2。在PyTorch和Caffe2合并之后，Caffe2成为了PyTorch中的一种后端。

5.1.4 Torch

这部分通常由用户在使用PyTorch进行模型训练或测试时调用。它包含了你最熟悉的内容：你导入和使用的实际的Python模块。

5.1.5 Torch/csrc

这部分是实现了你可以称之为 PyTorch 前端的 C++ 代码。更具体地说，它实现了在 Python 和 C++ 之间进行转换的绑定代码，还包含了 PyTorch 的一些非常重要的部分，例如自动求导引擎和 JIT 编译器。它还包含了 C++ 前端代码。

5.2 简单调用中的机制

调用像 torch.add 这样的函数。

1. 我们必须从 Python 领域转换到 C++ 领域（Python 参数解析）。
2. 我们处理变量调度 （VariableType一Type）。
3. 我们处理设备类型/布局调度（类型）。
4. 我们有实际的内核，它要么是现代的原生函数，要么是传统的 TH 函数。

6 Pytorch 内存管理基本情况

每个张量在初始化时都会被分配一个分配器。

c10/core/Allocator.h：在这里定义了PyTorch的默认分配器类。其中的一些策略：

• DataPtr 是一个唯一指针（附带了一个删除器和一些删除器的上下文），指向一些内存，还记录了其数据所在的设备。即使是空指针的 DataPtr 也可以具有唯一的设备；这使得我们可以将大小为零的分配与非零的分配一并处理。
• 这里选择 CPU 是任意的；如果有一个“未定义”的设备，我们也可以使用它。
• 使用函数 compare_exchange_deleter 可以在运行时更改删除器。
• 这个上下文被用于生成具有任意 std::function 删除器的 DataPtr。在一些面向用户的函数中，我们提供了一个（用户友好的）接口，用于从外部数据构建张量，其中接受任意 std::function 删除器。搜索 InefficientStdFunctionContext 可以找到这些出现的地方。

  这个上下文是低效的，因为除了 std::function 本身的动态分配之外，还需要在 InefficientStdFunctionContext 上进行一次动态分配。

在 Aten（aten/src/ATen/CPUFixedAllocator.h）中有一个伪分配器，如果实际使用了某些 CPU 固定操作（如 cpu_fixed_malloc、cpu_fixed_realloc、cpu_fixed_free），它会抛出异常。

c10/core/CPUAllocator.cpp 包含以下函数：alloc_cpu、free_cpu、memset_junk、alloc_cpu 甚至还处理了 NUMA 机器的代码。而且还有一个名为 MemoryAllocationReporter 的类，用于报告 C10 的内存分配和释放状态。

c10/core/Allocator.cpp：为不同的设备类型设置和获取分配器。

DeviceType::CPU
DeviceType::CUDA
DeviceType::OPENGL
DeviceType::OPENCL
DeviceType::MKLDNN
DeviceType::IDEEP
DeviceType::HIP
DeviceType::FPGA
DeviceType::MSNPU
DeviceType::XLA

c10/core/StorageImpl.h 和 c10/core/Storage.h：主要使用给定的分配器分配内存缓冲区，并创建一个带有该内存的存储。

c10/cuda/CUDACachingAllocator.cpp 是一个用于 CUDA 的缓存分配器。它具有以下描述：

另一个适用于 CUDA 设备分配的缓存分配器。

- 分配与流相关联。一旦释放，块可以在同一流上重新分配，但不能在任何其他流上分配。
- 分配器尝试找到最小的可用块以适应所请求的大小。如果块大于请求的大小，可能会被拆分。如果找不到块，则分配器将委托给 cudaMalloc。
- 如果 cudaMalloc 失败，分配器将释放所有未拆分的缓存块，并重新尝试分配。
- 大型（>1MB）和小型分配存储在不同的池中。小型请求被打包到 2MB 的缓冲区中。大型请求将使用最小的可用空闲块或使用 cudaMalloc 分配新的块。为了减少碎片化，如果没有足够大小的空闲块，则介于 1MB 和 10MB 之间的请求将分配并拆分一个 20MB 的块。

使用此分配器，分配和释放应逻辑上被视为与流相关联的内存段的"使用"，就像内核启动一样。如果内存段在多个流中使用，程序员必须插入适当的同步。
该库提供了一个 recordStream() 函数，以帮助在多个流上使用分配时插入正确的同步。这将确保在每个记录的流完成工作之前，块不会被重新使用。

7 Python 如何与 C/C++ 交互

7.1 将 C 程序编译为 .so 库，并在 Python 中调用它

7.1.1 以共享库形式编译

1. 完成你的 C 代码编写。
2. 将其编译为一个 *.so 文件。
3. 在 Python 文件中导入 ctypes 模块。
4. 在 Python 文件中加载 *.so 文件。
5. 定义 C 函数的输入类型。
6. 调用 *.so 文件中的函数。

function.c

// 判断给定的数字是否是2的幂
int myFunction(int num)
{
    if (num == 0){
        return 0;
    }
    else{
      if ((num & (num - 1)) == 0)
          return 1;
      else
          return 0;
    }
}

Compile

将名为 "function.c" 的 C 代码编译为一个共享库 libfun.so
gcc -fPIC -shared -o libfun.so function.c

fPIC：生成位置独立的代码（Position Independent Code），这是为了使得共享库可以在内存中的任意位置加载，提高可移植性和共享性。

function.py

在 Python 中使用 ctypes 模块导入共享库 "libfun.so"
import ctypes 
NUM = 16      
fun = ctypes.CDLL(libfun.so)   
fun.myFunction.argtypes=[ctypes.c_int] 
returnVale = fun.myFunction(NUM)

7.1.2 在 C++ 文件中添加包装器

如果这是一个 C++ 文件，你需要在一个 extern “C” 包装器中暴露你想使用的函数。

#include <iostream>
class Foo{
    public:
        void bar(){
            std::cout << "Hello" << std::endl;
        }
};
// Since ctypes can only talk to C functions, you need 
// to provide those declaring them as extern "C"
// 使用 extern "C" 包装器来在 C++ 代码中声明可以供 ctypes 调用的函数
extern "C" {
    Foo* Foo_new(){ return new Foo(); }    // 返回指向 Foo 对象的指针
    void Foo_bar(Foo* foo){ foo->bar(); }  // 函数 Foo_bar() 接受指向 Foo 对象的指针作为参数
}

然后进行编译：

# 将源文件 foo.cpp 编译为目标文件 foo.o
g++ -c -fPIC foo.cpp -o foo.o
# 将目标文件 foo.o 链接为共享库文件 libfoo.so
g++ -shared -Wl,-install_name,libfoo.so -o libfoo.so foo.o

后续，在 Python 代码中的某些事物与 C 语言中的类似。

from ctypes import cdll
lib = cdll.LoadLibrary('./libfoo.so')

class Foo(object):
    def __init__(self):
        self.obj = lib.Foo_new()

    def bar(self):
        lib.Foo_bar(self.obj)
# Once you have that you can call it like

f = Foo()
f.bar() #and you will see "Hello" on the screen

7.2 C++ 文件包含模块并公开

在 BOOST_PYTHON_MODULE 中包含 <boost/python.hpp> 函数。
通过编写一个 Boost.Python 的包装器，可以将 C++ 函数暴露给 Python。

// 将 C++ 函数 greet 通过 Boost.Python 暴露给 Python，
// 并将其作为名为 hello_ext 的模块可用。
#include <boost/python.hpp>
char const* greet()
{
   return "hello, world";
}
BOOST_PYTHON_MODULE(hello_ext)
{
    using namespace boost::python;
    def("greet", greet);
}

就是这样。我们完成了。现在我们可以将其构建为一个共享库。生成的 DLL 现在对 Python 可见。下面是一个示例的 Python 会话：

import hello_ext
print hello_ext.greet()
# hello, world

8 将 C++/CUDA 操作与 PyTorch 集成

当我们想要构建自定义的方法或模块时，我们可以选择是在 Python 中构建还是在 C++ 中构建。前者更简单，但 C++ 版本更快且更高效，特别是当我们想要构建一个经常使用或耗时的模块时。下面是对此的解释。

8.1 CPU集成

除了在 Python 中集成 C++ 文件并在 PyTorch 中使用它之外，PyTorch 本身还为我们提供了两种非常直接的方式来完成此任务。它们是使用 setuptools 构建和使用 JIT 编译扩展。

对于“提前编译”的方式，我们通过编写一个 setup.py 脚本来使用 setuptools 编译我们的 C++ 代码来构建 C++ 扩展。对于 LLTM 而言，它看起来就像这样简单：

# 构建一个名为 "lltm_cpp" 的 C++ 扩展模块
from setuptools import setup, Extension
from torch.utils import cpp_extension
setup(name='lltm_cpp',    # 指定模块的名称为 "lltm_cpp"
      # 指定模块的名称为 "lltm_cpp"，以及 C++ 源文件的路径为 "lltm.cpp"
      ext_modules=[cpp_extension.CppExtension('lltm_cpp', ['lltm.cpp'])],
      cmdclass={'build_ext': cpp_extension.BuildExtension})

JIT 编译机制通过调用 PyTorch API 中的一个简单函数 torch.utils.cpp_extension.load()，为您提供了一种即时编译和加载扩展的方式。对于 LLTM，使用 JIT 编译机制的代码如下所示：

from torch.utils.cpp_extension import load
lltm_cpp = load(name="lltm_cpp", sources=["lltm.cpp"])

8.2 CUDA集成

将支持 CUDA 的操作与 PyTorch 集成也非常简单。如果你想编写一个 setup.py 脚本，它可能如下所示：

from setuptools import setup
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
setup(
    name='lltm',
    ext_modules=[
        CUDAExtension('lltm_cuda', [
            'lltm_cuda.cpp',
            'lltm_cuda_kernel.cu',
        ])
    ],
    cmdclass={
        'build_ext': BuildExtension
    })

与 CppExtension() 不同，我们现在使用 CUDAExtension()。我们只需要指定 .cu 文件以及 .cpp 文件 - 库将负责处理其中涉及的所有麻烦。JIT 机制甚至更简单：

from torch.utils.cpp_extension import load
lltm = load(name='lltm', sources=['lltm_cuda.cpp', 'lltm_cuda_kernel.cu'])

9 总结

• PyTorch 的 Python 部分对内存管理没有特殊的关注，它的工作方式与标准的 Python 程序相同。
• 当前的 PyTorch 源代码包含来自多个来源的代码，其中一些是纯粹的遗留代码，一些来自于 Caffe2，一些作为基础代码存在，一些被打包成动态链接库以供 Python 使用。此外，CPU 和 CUDA 的代码也有所不同，如果想要进行优化，需要关注正确的部分。
• 几乎所有 PyTorch 的核心模块和函数都是用 C++ 编写的，因此效率更高。
• 每个张量都附带一个内存分配器，它不仅可以完成分配和释放的工作，还可以记录其所在的设备。不同数据类型可以通过输入参数传递不同类型的分配器，这使得代码更加兼容。
• PyTorch 结合了多种代码调度方法，它们对 C 和 C++ 代码都能很好地工作。
• Python 可以使用 ctypes 调用编译的 C 文件，但 PyTorch 提供了一套工具集，使这个过程更加简单。