LLM Systems Assignment 1: CUDA Programming

本次作业主要涉及两个文件：

• llmsys_hw1/minitorch/cuda_kernel_ops.py：Python 侧调用 kernel 编译成的 .so 接口
• llmsys_hw1/src/combine.cu：CUDA kernel 实现

完成本次作业的主要预备知识要求是掌握基本的 CUDA 编程能力（如：按照课程 readings 要求完成 PMPP 这本书前 6 章的学习），在此基础上，需要理解课程提供的骨架代码中几个关键的辅助函数。

本文目的就是介绍完成 homework 1 需要理解的几个辅助函数，并不会给出作业具体 solution 代码。

shape & strides tensor 表示

homework 1 对 tensor 的表示基于 shape & strides 布局表示，这一点 homework 1 的官方链接已经有基本解释，网上相关的介绍也非常多，对此不赘述。

这里主要介绍几个完成本次作业需要理解的辅助函数。

index_to_position

/**
 * Index: multi-dim coordinate.
 * Position: a int number means the offset in 1-dim storage.
 */
__device__ int index_to_position(const int* index, const int* strides, int num_dims) {
  /**
   * Converts a multidimensional tensor index into a single-dimensional position in storage
   * based on strides.
   * Args:
   *    index: index tuple of ints
   *    strides: tensor strides
   *    num_dims: number of dimensions in the tensor, e.g. shape/strides of [2, 3, 4] has 3 dimensions
   *
   * Returns:
   *    int - position in storage
   */
  int position = 0;
  for (int i = 0; i < num_dims; ++i) {
    position += index[i] * strides[i];
  }
  return position;
}

笔者已经在代码上补充了注释，函数功能和参数说明如下： - 函数功能：将多维 index 坐标表示，基于 strides，计算得到该 index 元素在一维存储空间的偏移量。 - index：与 tensor shape 对应的多维坐标 - position：根据 index & strides 计算得到的存储空间 offset

to_index

/**
 * Convert a global index to a multi-dim coordinate.
 * For example:
 *   - ordinal = 10
 *   - shape = [3, 4]
 *   - num_dims = 2
 *
 * According to the code, the compute process is as follows:
 *   - Init: cur_ord = 10
 *   - Loop 0: i=1, out_index[1] = cur_ord % shape[1] = 10 % 4 = 2, cur_ord = 10/4 = 2
 *   - Loop 1: i=0, out_index[0] = cur_ord % shape[0] = 2 % 3 = 2
 *
 * So the final out_index is [2, 2], which is corresponding to a row-major strides [4, 1], where `2*4 + 2*1 = 10`.
 */
__device__ void to_index(int ordinal, const int* shape, int* out_index, int num_dims) {
  /**
   * Convert an ordinal to an index in the shape. Should ensure that enumerating position 0 ... size of
   * a tensor produces every index exactly once. It may not be the inverse of index_to_position.
   * Args:
   *    ordinal: ordinal position to convert
   *    shape: tensor shape
   *    out_index: return index corresponding to position
   *    num_dims: number of dimensions in the tensor
   *
   * Returns:
   *    None (Fills in out_index)
   */
  int cur_ord = ordinal;
  for (int i = num_dims - 1; i >= 0; --i) {
    int sh = shape[i];
    out_index[i] = cur_ord % sh;
    cur_ord /= sh;
  }
}

同样的，上面贴出的代码已经补充了相关注释。解释如下： - 函数功能：给定一个 ordinal （通常由 kernel 内每个 thread 根据 blockDim & threadIdx 计算得到的 global thread id）作为函数入参以及 tensor shape & dims 作为入参，将 flatten 后的 ordinal 转换为与 shape 的维度对应的多维 index。 - 输入的 shape 数组元素存储顺序为最高维度到最低维度，对于普通的 tensor（直接创建，没有 transpose/permute 等操作），最低维度是内存连续存储的维度（strides[-1] = 1）。 - 代码开头给了一个例子和计算过程，实际上就是 row-major 将元素的 global index 映射为多维的 index。 - 这个函数通常用于根据 kernel 内某个 thread 的 global id 计算得到其负责处理的元素的多维 index。

broadcast_index

/**
 * For a broadcast result coordinate(the parameter `big_index`), compute its corresponding input coordinate according
 * to the params `shape`, `num_dims_big` and `num_dims`. The result is written to `out_index`.
 *
 * Example 1:
 *   - big_index: [1, 3]
 *   - big_shape: [3, 4]
 *   - shape: [4]
 *   - num_dims_big: 2
 *   - num_dims: 1
 *
 * The compute process is as follows:
 *   - Loop 0: i=0, shape[0] = 4 (>1), out_index[0] = big_index[0 + (2 - 1)] = 3
 *
 * So the `out_index` is [3].
 *
 * Example 2:
 *   - big_index: [1, 3]
 *   - big_shape: [3, 4]
 *   - shape: [1, 4] (also 2 dims, but with a dim is 1.)
 *   - num_dims_big: 2
 *   - num_dims: 2
 *
 * The compute process is as follows:
 *   - Loop 0: i=0, shape[0] = 1, out_index[0] = 0
 *   - Loop 1: i=1, shape[1] = 4, out_index[1] = big_index[1 + (2 - 2)] = 3
 *
 * So the `out_index` is [0, 3].
 */
__device__ void broadcast_index(const int* big_index, const int* big_shape, const int* shape, int* out_index,
                                int num_dims_big, int num_dims) {
  /**
   * Convert a big_index into big_shape to a smaller out_index into shape following broadcasting rules.
   * In this case it may be larger or with more dimensions than the shape given.
   * Additional dimensions may need to be mapped to 0 or removed.
   *
   * Args:
   *    big_index: multidimensional index of bigger tensor
   *    big_shape: tensor shape of bigger tensor
   *    shape: tensor shape of smaller tensor
   *    out_index: multidimensional index of smaller tensor
   *    num_dims_big: number of dimensions in bigger tensor
   *    num_dims: number of dimensions in smaller tensor
   *
   * Returns:
   *    None (Fills in out_index)
   */
  for (int i = 0; i < num_dims; ++i) {
    if (shape[i] > 1) {
      out_index[i] = big_index[i + (num_dims_big - num_dims)];
    } else {
      out_index[i] = 0;
    }
  }
}

这个函数的行为 & 例子同样也已经在注释中详细说明。具体行为请参照上面注释理解。这里简单说明下什么时候该使用这个函数：

• elementwise kernel 一般处理流程
  • 根据 blockDim & threadIdx 计算得到 global thread id
  • 根据 global thread id 调用 to_index 得到当前 thread 负责处理的 output tensor index
  • 由于可能涉及到广播操作，需要根据 output tensor index 使用广播规则计算得到对应的 input tensor index，也就是通过调用 broadcast_index 这个函数实现
  • 有了 input tensor index & output tensor index 后，调用 index_to_position 接口得到 offset，完成 elementwise 计算写入 output tensor。