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

Question

给定具有形状的多维数组[A][B][C][D]但存储为具有长度的1-dim数组[A*B*C*D].我想使用模板元编程来简化索引计算.指数(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

Answer 1

是的，可以获得相同的组装（证明）。我通过在索引对象的构造函数中“计算”每个维度的间距并“初始化”非静态数组数据成员来到达那里。

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...);
}