题图是陪伴我写作的一杯奶茶(茶话弄)。
Parameter Pack 可以允许在编译期借助可变长模板参数,实现可变长的函数参数(各自拥有不同的类型),或者可变的结构体定义。关键在于:
- 参数的数量以及类型是在编译期可推断的
- 每个参数是不同的类型,因此没法使用
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!