使用元编程展开嵌套循环

Nested loops unrolling using metaprogramming

本文关键字：嵌套循环 编程      更新时间：2023-10-16

我有许多小尺寸的嵌套循环I，J编译时已知，例如

for(int i = 0; i < I; ++i) {
    for(int j = 0; j < J; ++j) {
        // ...
        // do sth with (i,j,...)
    }
}

我需要使用I、J，这样我就可以在编译时使用每个坐标组合。

为了澄清，请考虑以下结构，并取两个大小I=2，J=3的嵌套循环。

template<int... I>
struct C {
     static void f() {
          // do sth
     }
};

我不能使用索引I，j（类似于上面）来索引结构C，因为它们在编译时是未知的。然而，如果我被允许使用这些指数，我想生成的正是情况，例如

C<0,0>::f();
C<0,1>::f();
C<0,2>::f();
C<1,0>::f();
C<1,1>::f();
C<1,2>::f();

只要产生了所有的组合，我就不特别关心调用生成的顺序。生成机制应该推广到任意数量的嵌套循环。

您可以通过以树状方式实例化模板来实现这一点，并跟踪当前访问的节点。

namespace detail{
    //This is used to store the visited nodes
    template<int...> struct int_pack;
    //Primary template
    template<typename, int... I>
    struct C;
    //This is the leaf node
    template<int... Is>
    struct C<int_pack<Is...>> {
        //The loop body goes here
        static void f() {
            std::cout << __PRETTY_FUNCTION__ << 'n';
        }
    };
    //This is the recursive case
    template <int I, int... Is, int... PIs>
    struct C<int_pack<PIs...>, I,Is...> {
        template <std::size_t... Idx>
        static void f_help (std::index_sequence<Idx...>) {
            //Store the current node in the pack 
            //and call `C::f` for each loop iteration
            (void)std::initializer_list<int> {
                (C<int_pack<PIs...,Idx>,Is...>::f(), 0)... 
            };   
        }
        //Use tag dispatching to generate the loop iterations
        static void f() {
            f_help(std::make_index_sequence<I>{});
        }
    };
}
//Helper alias
template<int... Is>
using C = detail::C<detail::int_pack<>, Is...>;

用法非常简单：

C<2,3>::f();

在Clang上打印：

static void detail::C<detail::int_pack<0, 0>>::f() [I = <>]
static void detail::C<detail::int_pack<0, 1>>::f() [I = <>]
static void detail::C<detail::int_pack<0, 2>>::f() [I = <>]
static void detail::C<detail::int_pack<1, 0>>::f() [I = <>]
static void detail::C<detail::int_pack<1, 1>>::f() [I = <>]
static void detail::C<detail::int_pack<1, 2>>::f() [I = <>]

实时演示

您可以使其更通用，这样您就可以通过lambda将循环体注入到类中，但如果您只想这样做一次，而不想引入其他依赖项（如boost::hana），则应该使用上述解决方案。以下是一个更通用版本的可能实现（您可以通过完美的转发等来改进它）：

namespace detail{
    template<int...> struct int_pack;
    template<typename, int... I>
    struct C;
    template<int... Is>
    struct C<int_pack<Is...>> {
        template <typename Func>
        static void f(const Func& func) {
            func(Is...);
        }
    };
    template <int I, int... Is, int... PIs>
    struct C<int_pack<PIs...>, I,Is...> {
        template <std::size_t... Idx, typename Func>
        static void f_help (std::index_sequence<Idx...>, const Func& func) {
            (void)std::initializer_list<int>{ (C<int_pack<PIs...,Idx>,Is...>::f(func), 0)... };   
        }
        template <typename Func>
        static void f(const Func& func) {
            f_help(std::make_index_sequence<I>{}, func);
        }
    };
}

你会这样使用：

C<2,3>::f([](int i, int j){
    std::cout << "i " << i << " j " << j << 'n';
});

实时演示

这是我用boost::hana模拟的一个快速版本。可能有更好的方法可以做到这一点，但这应该让你知道可以做什么。

template <typename Func>
void unroll (const Func& func) {
    func();
}
template <std::size_t I1, std::size_t... Is, typename Func>
void unroll (const Func& func) {
    hana::for_each(hana::range_c<std::size_t, 0, I1>,
                   [&](auto x) {
                       unroll<Is...>([x, &func] (auto... xs) { func(x,xs...); });
                   });
}
相关文章：