Nested loops unrolling using metaprogramming

2019-06-16 12:52发布

I have a number of nested loops with small sizes I, J, ... known at compile time, e.g.

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

I need to unroll the loops using the sizes I, J, ... in such a way that I can use each coordinate combination at compile time.

To clarify, consider the following structure and take 2 nested loops with sizes I = 2, J = 3.

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

I can not use the indices i, j (similar to above) to index the structure C since they are not known at compile time. However what I would like to generate is exactly what would have been the case had I been allowed to use the indices, e.g.

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

I am not particularly concerned with the order of call generations as long as all combinations are produced. The generation mechanism should generalize to arbitrary number of nested loops.

1条回答
戒情不戒烟
2楼-- · 2019-06-16 13:51

You can do this by instantiating templates in a tree-like manner, keeping track of the nodes currently visited.

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

Usage is pretty simple:

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

On Clang this prints:

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 = <>]

Live Demo


You could make this more generic so that you can inject the loop body into the class through a lambda, but the above solution should do if you only want to do this once and don't want to pull in other dependencies like boost::hana. Here's a possible implementation of the more generic version (you could improve it with perfect forwarding and the like):

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

You would use this like so:

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

Live Demo


Here's a quick version I mocked up with boost::hana. There are likely better ways to do this, but this should give you an idea of what can be done.

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...); });
                   });
}
查看更多
登录 后发表回答