题图是陪伴我写作的一杯奶茶(茶话弄)。

Parameter Pack 可以允许在编译期借助可变长模板参数,实现可变长的函数参数(各自拥有不同的类型),或者可变的结构体定义。关键在于:

  1. 参数的数量以及类型是在编译期可推断的
  2. 每个参数是不同的类型,因此没法使用vector<T>形式的模板

例如我们要编写一个函数,对一系列不同类型的对象都增加一个相同的 bias,待处理的对象数量和类型有多种组合:

void my_op(const float* bias, float* tensor_1, int64_t* tensor_2, int32_t* tensor_3) {
    // for each tensor, modify with bias
    *tensor_1 += bias;
    if (*tensor_1 < 0) {
    	*tensor_1 = 0;
    }
}

单独看上面的函数,其实拆分成一次只对一个 tensor 进行操作,并通过模板特化成不同类型,也是一种实现的方法。但例如在 CUDA 编程里,一个 kernel 大概率会比多个 kernel 要更好。

Basic

那么对于上述函数,我们可以借助 Parameter Pack 定义一个函数模板:

template <typename P, typename... Ts>
void my_op_basic(const P* bias, Ts... *tensors) {
    //...?
}

// instance
template my_op_basic<float, double>(const float*, double*);
template my_op_basic<float, float, int64_t, int32_t>(const float*, float*, int64_t*, int32_t*);

对于这样的一个函数,在内部我们必然需要对 tensors 里包含的每一个实参进行遍历。我们期待编译器能发现 tensors 的长度是已知的,并且其中的类型也是已知的,因此在 tensors 上的 for-loop 也是可以直接展开的……因此顺着我们的直觉,写下以下代码:

template <typename P, typename... Ts>
void my_op_basic(const P* bias, Ts... *tensors) {
    for (int i = 0; i < sizeof...(tensors); ++i) {
    	auto ts = std::get<i>(std::forward_as_tuple(tensors));
        // ERROR: i is not a constexpr
    }
}

但是很不幸,i是一个无法逾越的鸿沟,毕竟我们总是可以随意篡改 i 的值,进而使得整个循环不是那么“静态”。但这也让我们意识到,std::get<size_t>std::forward_as_tuple可能会对解决这个问题很有帮助。
换个思路,如果我们使用一个 for-in-loop 呢?

template <typename P, typename... Ts>
void my_op_basic(const P* bias, Ts... *tensors) {
    for (const auto& ts: {...tensors}) {
    	// this will work, only when all tensors have the same type
    }
}

同样不幸的是,一个 for-in-loop 总是需要一个可以迭代的容器作为目标(而 tuple 不可迭代)。因此参数列表必须是相同的类型(那为什么我们不用 vector<T>呢)。
几番尝试结束,我们意识到这件事情可能并没有原先设想的那么简单。事实上整个社区也直到 cpp-17 的出现,才有了比较完美的解法。

Fold Expression

通过折叠表达式 ...语法可以将 parameter pack 对象按顺序展开并应用到同一个运算符上,这个运算符可以是一元或者二元的。例如:

template <typename... Ts>
auto sumup(Ts... t) {
	return (t + ...);
}

sumup(1, 2, 5.5, 15); // 23.5

那么我们可以定义一个 lambda 来负责对每个元素进行处理,并使用 fold expression 将参数展开并依次调用。最妙的是这里展开的运算符是逗号表达式。

/* https://godbolt.org/z/4GhrGPqbT */
template <typename ... Ts>
void Foo (Ts && ... multi_inputs)
{
    int i = 0;

    auto processOne = [&] (auto & input)
    {
        // Do things in your "loop" lambda.
        ++i;
        std::cout << "input " << i << " = " << input << std::endl;
    };

    // Call lambda for each input (using fold expression over comma operator).
    (processOne(multi_inputs), ...);
}

更进一步简化语法之后,我们的 my_op 已经呼之欲出:

template <typename P, typename... Ts>
void my_op_basic(const P* bias, Ts... *tensors) {
    (
       [&]() {
           auto& t = tensors;
           *t += *bias;
           if (*t < 0) {
            	*t = static_cast<typename std::decay<decltype(t)>::type>(0.0f);
           }
           // NOTE: it's a pointer extracted from tensors, not the parameter pack itself.
       } (), 
    ...);
}

Advance

此时我们已经基本完成了任务,但是这个函数的接口仍然不够优雅,如果可以用一个容器来将需要处理的对象统一包裹一下,接口会更简洁一些。此时我们就需要一个结构体模板来容纳这些可变的参数。

template <typename P, typename... Ts>
void my_op_advance(const P* bias, OpPack<Ts...> tensor_pack) {}

可变结构体模板

Tuple 这个时候变得重要了:我们通过在结构体中定义一个 tuple 以存储可变长的对象列表。

#include <tuple>

template <typename... Ts>
struct OpPack {
	std::tuple<Ts...> tensors;
};

// instance
std::tuple<float, int, double> t {1.5, 114514, 3.1415926};
OpPack<float, int, double> args {
	.tensors = t,
};

而为了继续使用我们的 fold expression 魔法,我们需要一种办法来将结构体里的 tuple 直接展开到参数列表里面。我们已经知道了 std::get可以从 tuple 里摘取对应 index 的值,只要 index 是一个编译期常量。因此我们用一个编译期生成的 index 序列去将 tuple 进行展开。

template <typename P, typename... Ts, size_t... Is>
void my_op_proxy(P* bias, OpPack<Ts...> tensor_pack, std::index_sequence<Is...>) {
	my_op_basic(bias, std::get<Is>(tensor_pack.tensors)...);
}

template <typename P, typename... Ts>
void my_op_advance(const P* bias, OpPack<Ts...> tensor_pack) {
    my_op_proxy(bias, tensor_pack, std::index_sequence_for(tensor_pack.tensors));
}

Vola!