如何对std::index_sequence进行前缀求和

How to do a prefix sum on std::index_sequence?

本文关键字:前缀 求和 sequence index std      更新时间:2023-10-16

我想计算std::index_sequence索引的独占前缀和(扫描),但我不确定从哪里开始。我已经研究了std::make_index_sequence的实现,以寻求一个泛化,但它对我来说很神秘

如何实现下面的exclusive_scan_index_sequence以使程序成功?

#include <type_traits>
#include <cassert>
#include <cstddef>
// define something like std::index_sequence
template<size_t... Indices>
struct index_sequence {};
// compute the exclusive scan of IndexSequence
// initializing the first value in the result sequence to Init
template<size_t Init, class IndexSequence>
struct exclusive_scan_index_sequence;
template<size_t Init, size_t... Indices>
struct exclusive_scan_index_sequence<Init,index_sequence<Indices...>>
{
  // what goes here?
};
int main()
{
  using ones = index_sequence<1,1,1,1,1>;
  using scanned = exclusive_scan_index_sequence<0,ones>;
  using expected = index_sequence<0,1,2,3,4>;
  assert((std::is_same<expected,scanned>::value));
  return 0;
}

我认为以下操作虽然仍然是递归的,但更容易:

template<typename Current, size_t Next, class IndexSequence>
struct exclusive_scan_index_sequence_impl
{
    using type = Current;
};
template<size_t... Current, size_t Next, size_t First, size_t... Indices>
struct exclusive_scan_index_sequence_impl<
           index_sequence<Current...>,
           Next,
           index_sequence<First,Indices...>
       >
    : exclusive_scan_index_sequence_impl<
          index_sequence<Current...,Next>,
          Next+First,
          index_sequence<Indices...>
      >
{ };
template<size_t Init, class IndexSequence>
using exclusive_scan_index_sequence =
    typename exclusive_scan_index_sequence_impl<
        index_sequence<>,
        Init,
        IndexSequence
    >::type;

实时示例

这里有一个递归解决方案,它连接index_sequences:

template<class IndexSequence1, class IndexSequence2>
struct index_sequence_cat_impl;
template<size_t... Indices1, size_t... Indices2>
struct index_sequence_cat_impl<index_sequence<Indices1...>,index_sequence<Indices2...>>
{
  using type = index_sequence<Indices1...,Indices2...>;
};
template<class IndexSequence1, class IndexSequence2>
using index_sequence_cat = typename index_sequence_cat_impl<IndexSequence1,IndexSequence2>::type;
// compute the exclusive scan of IndexSequence
// initializing the first value in the sequence to Init
template<size_t Init, class IndexSequence>
struct exclusive_scan_index_sequence_impl;
template<size_t Init, size_t Index0, size_t... Indices>
struct exclusive_scan_index_sequence_impl<Init,index_sequence<Index0, Indices...>>
{
  using rest = typename exclusive_scan_index_sequence_impl<Init + Index0, index_sequence<Indices...>>::type; 
  using type = index_sequence_cat<index_sequence<Init>, rest>;
};
template<size_t Init, size_t Index0>
struct exclusive_scan_index_sequence_impl<Init,index_sequence<Index0>>
{
  using type = index_sequence<Init>;
};
template<size_t Init, class IndexSequence>
using exclusive_scan_index_sequence = typename exclusive_scan_index_sequence_impl<Init,IndexSequence>::type;

也许有可能以某种方式反复进行。