Nested loops unrolling using metaprogramming

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.

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

Recommended topics

Hot tags