叁 - 搭建向量库

Part 1: NDArray 类

在开始这个作业之前,你应该首先熟悉我们提供的 NDArray.py 类,作为作业的起点。这段代码相当简洁(大约 500 行,但其中很多是为你将要实现的函数提供的注释)。

在核心部分,NDArray 类是一个处理通用 n 维数组操作的 Python 包装器。回想一下,几乎任何这样的数组都会在内部以浮点值向量的形式存储,即:

float data[size];

然后数组不同维度的实际访问都由额外的字段处理(例如数组形状、步长等),这些字段指示了这个“扁平”数组如何映射到 n 维结构。为了获得任何合理的速度,原始操作(如加法、二进制操作,以及更结构化的操作如矩阵乘法等)都需要在某个级别用一些本地语言如 C++ 编写(包括例如进行 CUDA 调用)。但是很多操作,如转置、广播、矩阵子集等,都可以通过调整数组的高级结构,如步长,来处理。

NDArray 类的理念是我们希望所有处理数组结构的逻辑都写在 Python 中。只有真正执行对扁平向量的底层操作的“真正”低级代码(以及管理这些扁平向量的代码,因为它们可能需要在 GPU 上分配),才会用 C++ 编写。这种分离的具体性质在你完成作业时可能会变得更加清晰,但一般来说,所有可以在 Python 中完成的事情都是用 Python 完成的;通常我们会很慷慨地调用 .compact()(复制内存),以便使底层的 C++ 实现更简单。

更详细地说,NDArray 类中有五个字段你需要熟悉(注意,所有这些字段的真实类成员都以下划线开头,例如 _handle_strides 等,其中一些会作为公共属性暴露出来... 对于你的所有代码,使用内部带下划线的版本是可以的):

  1. device - 一个 BackendDevice 类型的对象,它是一个简单的包装器,包含了指向底层设备后端(例如 CPU 或 CUDA)的链接。
  2. handle - 一个类对象,用于存储数组的底层内存。这被分配为 device.Array() 类的一个实例,虽然这个分配都是在提供的代码中完成的(具体是在 NDArray.make 函数中),你不需要担心自己调用它。
  3. shape - 一个元组,指定数组中每个维度的大小。
  4. strides - 一个元组,指定数组中每个维度的步长。
  5. offset - 一个整数,指示数组实际从底层 device.Array 内存的哪个位置开始(这样做很方便,因为我们可以更容易地管理指向现有内存的指针,而不必跟踪分配)。

通过操作这些字段,甚至纯 Python 代码都可以执行数组的许多必要操作,例如置换维度(即转置)、广播等。然后,对于原始数组操作调用,device 类有许多方法(在 C++ 中实现)包含了必要的实现。

有几点需要注意:

def make(shape, strides=None, device=None, handle=None, offset=0):

它创建一个具有给定形状、步长、设备、句柄和偏移量的新 NDArray。如果没有指定 handle(即没有引用预先存在的内存),则调用会分配所需的内存,但如果指定了 handle,则不会分配新的内存,而是新的 NDArray 指向与旧 NDArray 相同的内存。对于尽可能多地 分配新内存的有效实现非常重要,因此在许多情况下你会希望使用此调用来实现。

Part 2: CPU 后端 - Compact 和 setitem

ndarray_backend_cpu.cc 中实现以下函数:

为什么这些函数都属于同一类别?让我们来考虑一下 Compact() 函数的实现。回想一下,如果矩阵在内存中按照 " 行优先 " 的形式顺序排列(实际上是对行优先的一般化,适用于更高维的数组),那么矩阵被认为是紧凑的,即最后一个维度排在第一位,然后是倒数第二个维度,以此类推,直到第一个维度。在我们的实现中,我们还要求分配的后端数组的总大小等于数组的大小(即底层数组在数组数据之前或之后也不能有任何数据,这意味着,例如,.offset 字段等于零)。

现在,让我们考虑一个三维数组作为例子,看看紧凑调用可能是如何工作的。这里的 shapestrides 是被紧凑的矩阵的形状和步长(即在我们进行紧凑之前的形状和步长)。

cnt = 0;

for (size_t i = 0; i < shape[0]; i++)
    for (size_t j = 0; j < shape[1]; j++)
        for (size_t k = 0; k < shape[2]; k++)
            out[cnt++] = in[strides[0]*i + strides[1]*j + strides[2]*k];

换句话说,我们正在从基于步长的矩阵表示转换为纯粹的顺序表示。

现在,实现 Compact() 的挑战在于,你希望该方法适用于任意数量的输入维度。为不同的固定维度大小数组专门化很容易,但对于通用实现,你需要考虑如何执行相同的操作,其中你实际上需要一个“可变数量的 for 循环”。一个提示是,一种做法是维护一个索引向量(大小等于维度数),然后在循环中手动递增它们(包括在任何达到其最大大小时的“进位”操作)。

然而,如果你在这方面遇到困难,你始终可以利用这样一个事实,即我们可能不会要求你处理超过 6 个维度的矩阵(尽管我们将使用 6 个维度,用于我们在课堂上讨论的 im2col 操作)。

与 setitem 的关联

虽然设置项功能看起来非常不同,但实际上与 Compact() 密切相关。__setitem()__ 是在设置对象的某些元素时调用的 Python 函数,例如,

A[::2,4:5,9] = 0 # or = some_other_array

你该如何实现这个功能?在上面的 __getitem()__ 调用中,你已经实现了一种方法来获取数组的子集而不复制它(只修改步长)。但你如何实际上去设置这个数组的元素?在这份作业的几乎所有其他设置中,我们在将项设置到输出数组之前调用了 .compact(),但在这种情况下这样做是行不通的:调用 .compact() 会将子集数组复制到某个新的内存中,但 __setitem__() 调用的整个目的是我们想要修改现有的内存。

如果你仔细考虑了一段时间,你会意识到答案看起来与 .compact() 的逆操作非常相似。如果我们想要将右侧矩阵(已经是紧凑的)赋给 __getitem()__ 的结果,那么我们需要在这里像 shapestride 一样是输出矩阵的字段。然后我们可以如下实现 setitem 调用:

cnt = 0;
for (size_t i = 0; i < shape[0]; i++)
    for (size_t j = 0; j < shape[1]; j++)
        for (size_t k = 0; k < shape[2]; k++)
            out[strides[0]*i + strides[1]*j + strides[2]*k] = in[cnt++]; // or "= val;"

由于这种相似性,如果你以模块化的方式实现了你的索引策略,你将能够在 Compact() 调用和 EwiseSetitem()ScalarSetitem() 调用之间重用它。

Part 3: CUDA 后端 - Compact 和 setitem

ndarray_backend_cuda.cu 中实现以下函数:

在这部分中,您将在 CUDA 后端实现紧凑和 setitem 调用。这与 C++ 版本非常相似,但是,取决于您如何实现该函数,可能也会有一些重大差异。我们特别想要强调 C++ 和 CUDA 实现之间的一些差异。

首先,与在 CUDA 后端代码中实现的示例函数一样,对于上述所有函数,您实际上将要实现两个函数:上面列出的基本函数,您将从 Python 调用这些函数,以及实际执行计算的相应 CUDA 内核。在大多数情况下,我们仅在 ndarray_backend_cuda.cu 文件中提供“基本”函数的原型,并且您需要自己定义和实现内核函数。但是,为了了解这些函数的工作原理,对于 Compact() 调用,我们提供了完整的 Compact() 调用,以及 CompactKernel() 调用的函数原型。

您可能注意到的一件事是似乎奇怪地使用了 CudaVec 结构,这是一个用于传递形状/步幅参数的结构。在 C++ 版本中,我们使用 STL std::vector 变量来存储这些输入(在基本 Compact() 调用中也是如此,但是 CUDA 内核无法操作 STL 向量,因此需要其他东西)。此外,尽管我们可以将向量转换为普通的 CUDA 数组,但这将非常繁琐,因为我们需要调用 cudaMalloc(),将参数作为整数指针传递,然后在调用之后释放它们。当然,这样的内存管理对于数组中的实际底层数据是必需的,但是对于仅传递变量大小的形状/步幅值的小元组来说,这似乎有些过度。解决方案是创建一个具有数组可能具有的“最大化”大小的结构,然后只将实际的形状/步幅数据存储在这些字段的第一个条目中。包含的 CudaVec 结构和 VecToCuda() 函数已经完成了所有这些工作,您可以按照提供的方式将它们用于所有需要将形状/步幅传递给内核本身的 CUDA 内核。

C++ 和 CUDA 实现 Compact() 的另一个(更概念性的)主要区别是,在 C++ 中,您通常会顺序遍历非紧凑数组的元素,这使得您可以对计算紧凑和非紧凑数组之间对应索引的操作进行一些优化。在 CUDA 中,您无法这样做,需要实现能够直接将紧凑数组中的索引映射到分步数组中的索引的代码。

与以前一样,我们建议您以一种能够轻松在 Compact()Setitem() 调用之间重复使用的方式实现代码。简短地注意,如果您想要从内核代码调用(单独的、非内核)函数,则需要将其定义为 __device__ 函数。