CINN 算子单测问题总结

[zzk0]

简介

在编写了一段时间 CINN 算子单测之后，发现 CINN 存在如下的问题。

1. External Call 机制（正确性问题、性能问题、数据类型问题）
2. 编译时间过长
3. 报错信息位置
4. 显存暴涨
5. 调度代码冗余
6. LowerVec 返回多于一个 LoweredFunc

问题

External Call 机制

正确性问题

将指针传递给 External Function，导致了 inline compute 的错误和算子融合的错误。

1. inline compute 优化的时候无法感知到指针被 External Function 使用，inline compute 检测 IR 中 Store 和 Load 来判断变量是否被使用，然而指针的使用未被考虑在内，于是 inline 消除了相关指针变量。inline compute 的错误已经由 [BUGFIX] inline compute skip #1329 解决。
2. 算子融合可以减少数据的搬运开销，例如有两个算子进行融合，第一个算子的计算结果可以在 kernel 里面创建共享内存进行存储，第二个算子则从共享内存读取使用即可。但是，这块共享内存可能作为指针传递给 External Function，在 Compute 函数中不知道这个 Tensor 最后是一个 shared buffer，误以为可以访问 Tensor 所有的元素，进而出现越界访问。算子融合的错误出现在 Op mapper scatter #1301 中，临时采用了将 scatter_add 标记为 kNonfusible 来解决。因为 scatter_add 这个算子和任何 kInjective 算子进行融合，都会发生越界访问的错误。
3. 生成的 CUDA kernel 中声明了过长的数组，超过了 CUDA kernel stack 的限制，例如 sort 算子排序 >32768 个元素会报错。op unittest for sort & enhance test helper #1411

性能问题

部分生成的 CUDA kernel 只使用了一个 block，几乎是单线程模式。

以 sort 算子为例子：#1411

对于 128 个元素的数组，生成的 CUDA kernel 如下。这种问题并非 External Call 机制所直接造成，kernel 限定使用一个 block 恰恰是 External Call 机制对数据依赖分析正确的结果。算子编写人员在使用 External Call 机制未能意识到最终生成的 kernel 无法充分利用 GPU 特性，算法设计上存在缺陷，调度优化也无能为力。

__global__
void __launch_bounds__(1) fn_sort_0_kernel(const int64_t* __restrict__ x, int64_t* __restrict__ var_0)
{
  int32_t _var_0_index_temp_buffer [ 128 ];
  int32_t _var_0_index_temp_temp_buffer [ 128 ];
  int32_t* var_0_index = _var_0_index_temp_buffer;
  int32_t* var_0_index_temp = _var_0_index_temp_temp_buffer;
  for (int32_t i = 0; i < 128; i += 1) {
    var_0_index_temp[i] = cinn_nvgpu_lt_num_int64(x, 128, x[i], 0, 1);
  };
  for (int32_t i = 0; i < 128; i += 1) {
    var_0_index[i] = cinn_nvgpu_next_smallest_int32(var_0_index_temp, 128, i, 0, 1);
  };
  for (int32_t i = 0; i < 128; i += 1) {
    var_0[i] = x[var_0_index[i]];
  };
}

数据类型支持不完整

数据类型支持不完整，相关的修改有很多，这里就简单列举几个：

#1301
#1312
#1500

运行过程会抛出以下错误：

F0328 08:07:48.080608 588260 nvrtc_util.cc:107] Check failed: compile_res == NVRTC_SUCCESS (6 vs. 0)
default_program(21): error: argument of type "const double *" is incompatible with parameter of type "const float *"

其修改方法比较简单：

1. 在 cinn/runtime/cuda/cinn_cuda_runtime_source.cuh 编写并实现 kernel
2. 在 cinn/runtime/cuda/cuda_intrinsics.cc 注册新的 Extern 函数
3. 如果有 float16 类型，在 cinn/runtime/cuda/cuda_instrinsics_float16.cc 注册新的 Extern 函数

编译时间过长的问题

IR 构建过于复杂造成编译时间过程。

例如，在 cinn/hlir/pe/elementwise.h 中存在下面的代码片段，IR 语法树的复杂度和元素的个数成正比。

template <typename T>
ir::Tensor AssignValue(const std::vector<T>& values,
                       const common::Type& type       = common::type_of<T>(),
                       const std::string& output_name = "T_assign_value_out") {
  CHECK(!values.empty()) << "The input of pe::AssignValue should not empty! Please check.";

  auto out = lang::Compute(
      {ir::Expr(static_cast<int>(values.size()))},
      [=](const std::vector<ir::Expr>& indice) {
        auto init_value =
            (type == common::type_of<T>()) ? ir::Expr(values[0]) : common::cast(ir::Expr(values[0]), type);
        ir::Expr previous = ir::Select::Make(ir::EQ::Make(indice[0], ir::Expr(0)), init_value, lang::Zero(type));

        for (int i = 1; i < values.size(); ++i) {
          auto val = (type == common::type_of<T>()) ? ir::Expr(values[i]) : common::cast(ir::Expr(values[i]), type);
          previous = ir::Select::Make(ir::EQ::Make(indice[0], ir::Expr(i)), val, previous);
        }
        return previous;
      },
      output_name);

  return out;
}

存在这个问题的算子有：

• constant: op unittest for cbrt/ceil/cholesky/concat/constant/fill_constant #1495
• split: op unittest for split #1453

报错信息位置

以 scatter_add 算子为例子：#1500

在编写单测的过程中，由于单测编写人员疏忽，设置错误的 index，超过了合理的范围。报错信息将在 cuda module 里面抛出，未能及时发现是单测本身存在的问题导致数组越界的发生，单测编写人员甚至不知道 unspecified launch failure 是什么原因造成的。

# core dumped: cuda_module.cc:118] RAW: The error `CUDA_ERROR_LAUNCH_FAILED` occurs
# while compiling the ptx! And its message is `unspecified launch failure`.

调度代码冗余问题

cinn/hlir/op/op_util.cc 提供了常见的调度函数实现方式，在此之前编写的算子存在代码冗余。在 cinn/hlir/op/contrib 中翻一翻，可以发现来自社区贡献的代码，重复编写了调度函数代码。

LowerVec 返回多于一个 LoweredFunc

以 scatter 算子为例子：#1500

scatter 算子的实现中，会调用 pe::Transpose 生成一个转置后的张量，如果对多个张量创建 Stage，那么在 LowerVec 后将返回两个 LoweredFunc，而后续过程只允许一个 LoweredFunc。

目前对 IRSchedule 的过程理解不够深刻，未能理解其本质。

显存暴涨的问题

#1411

对于 sort 算子，在输入数组的大小大于 32K/64K 的情况下，kernel 中的 temp buffer 将会超过 stack size，抛出 CUDA ERROR。此外，该计算过程分配过多内存，输入数组越大，中间将占用越多显存。

该问题的原因不详，需要跟进分析。

未来工作

为了更好进行 CINN 单测，我认为有如下可改进的点：

1. 显存监控，监控算子实现合理性。对于部分算子，占用过多显存导致 CUDA 显存分配失败，建立相关机制可保障 CI 执行。
2. 超时机制，控制 CINN 算子的编译时长、运行时长在合理的范围内。
3. 单测冗余，当前的单测中存在大量重复代码，比如 shape、dtype，可以编写相关 util 函数，提供建议测试的 shape.