使用元编程进行高效的索引计算

efficient index computation using meta programming

本文关键字：高效 索引 计算 编程      更新时间：2023-10-16

给定一个形状为[A][B][C][D]但存储为长度为[A*B*C*D]的 1-dim 数组的多维数组。我想使用模板元编程来简化索引计算。索引(a,b,c,d)应位于位置

a*B*C*D + b*C*D + c*D + d

我目前使用

#include <iostream>
#include <cstdlib>
#include <array>

template<size_t start, size_t AXES>
struct prod_func
{
constexpr inline size_t operator()(const std::array<const size_t, AXES> arr) const
{
return arr[start] * prod_func < start + 1, AXES > ()(arr);
}
} ;
template<size_t AXES>
struct prod_func<AXES, AXES>
{
constexpr inline size_t operator()(const std::array<const size_t, AXES> arr) const
{
return 1;
}
} ;

template<int AXES>
class index
{
const std::array<const size_t, AXES> shapes;
public:
index(std::array<const size_t, AXES> s) : shapes(s) {}
template <typename... Dims>
constexpr inline size_t operator()(int off, Dims... dims) const {
return off * (prod_func < AXES - (sizeof...(Dims)), AXES > ()(shapes)) + operator()(dims...);
}
constexpr inline size_t operator()(int t) const {
return t;
}

};

int main()
{
size_t A=2, B=3, C=6, D=7;
auto idx = index<4>({A,B,C,D});
int a=1, b=1, c=1, d=1;
std::cin >> a;
std::cin >> b;
std::cin >> c;
std::cin >> d;
asm ("nop");
size_t result =  idx(a,b,c,d);
asm ("nop"); 
std::cout << result << std::endl;
asm ("nop"); 
result = (a*B*C*D + b*C*D + c*D + d);
asm ("nop");
std::cout << result << std::endl;
return 0;
}

cin只是为了确保运行时值。检查g++ -O2 -S ../main.cpp -std=c++11给出的组件

imull   $105, 8(%rsp), %edx
imull   $35, 12(%rsp), %eax
movl    $_ZSt4cout, %edi
addl    %edx, %eax
movl    16(%rsp), %edx
leal    (%rax,%rdx,8), %esi
subl    %edx, %esi
addl    20(%rsp), %esi

对于(a*B*C*D + b*C*D + c*D + d)部分。这是我对编译器的期望。但是对于索引类，它会产生更多的操作：

movslq  8(%rsp), %rax
movl    $_ZSt4cout, %edi
leaq    (%rax,%rax,2), %rdx
leaq    (%rax,%rdx,4), %rdx
leaq    (%rax,%rdx,8), %rcx
movslq  12(%rsp), %rax
leaq    (%rax,%rax,4), %rdx
leaq    (%rcx,%rdx,8), %rax
subq    %rdx, %rax
movslq  20(%rsp), %rdx
addq    %rdx, %rax
movslq  16(%rsp), %rdx
leaq    (%rax,%rdx,8), %rsi
subq    %rdx, %rsi

并且没有得到优化B*C*D=105. 有没有办法获得类似的组装？我想包装一些 CUDA 代码，所以它确实需要是相同的代码(在 C++11 中(。需要明确的是，在编译时只知道轴的数量。 或者任何其他写法？

编辑：虽然我现在确信它具有相同的效率，但我仍然希望获得相同的程序集：https://godbolt.org/g/RHwBV6

是的，可以获得相同的程序集(证明(。我通过"计算"索引对象构造函数中每个维度的音高并"初始化"非静态数组数据成员来达到目标。

template<size_t Nd>
struct Index {
static_assert(Nd >= 1, "");
size_t extents_[Nd];
size_t pitches_[Nd];
public:
template<class... Ts>
constexpr Index(size_t e0, Ts... es) noexcept
: Index{MakeIndSeq<Nd>{}, e0, size_t(es)...}
{}
private:
template<size_t... ds, class... Ts>
constexpr Index(IndSeq<ds...>, size_t e0, Ts... es) noexcept
: extents_{e0, es...}
, pitches_{extents2pitch<ds>(e0, es...)...}
{}
public:
template<class... Ts>
constexpr size_t operator()(size_t i0, Ts... is) const {
return operator()(MakeIndSeq<Nd>{}, i0, is...);
}
private:
template<size_t... ds, class... Ts>
constexpr size_t operator()(IndSeq<ds...>, Ts... is) const {
return sum((is*pitches_[ds])...);
}
};

extents2pitch看起来像什么样

子
template<size_t d, size_t... ds, class... Ts>
constexpr size_t extents2pitch_impl(IndSeq<ds...>, size_t N0, Ts... Ns) {
return product<size_t>(
Array<size_t, size_t(1)+sizeof...(Ns)>{N0, Ns...}[sizeof...(Ns)-ds]...
);
}
template<size_t d, class... Ts>
constexpr size_t extents2pitch(size_t N0, Ts... Ns) {
return extents2pitch_impl<d>(MakeIndSeq<sizeof...(Ns)-d>{}, N0, Ns...);
}
相关文章：