Type Stable Iteration over Static Parameters

General Usage
,
1

I have a use case where I need to iterate over a set of static parameters. I’m not sure how to do that in a type stable way. A MWE is below:

function foo(::Val{N}) where N
    n = 0
    for i in 1:N
        n += sum(ntuple(identity, Val(i)))
    end 
    return n
end

I’ve tried hacking together something like the following, but it’s still not type stable

foo(::Val{N}) where N = foo(Val.(ntuple(identity, N))...)

function foo(inds::Vararg{<:Val})
    n = 0
    for i in inds
        n += sum(ntuple(identity, i))
    end 
    return n
end

Is there a good way to make something like this type stable?

2

As far as I can tell, this should be stable because of constant propagation. How are you measuring type stability here?

3

I’m seeing problems with @code_warntype as well as @btime telling me things are allocating.

4

But is that how loop compilation works?

I assume it’s not unrolled even though everything is known to the compiler, so each Val is still unstable and boxed. @code_llvm suggests that’s whats happening.

So @bmit If you want this to compile away, you need to use recursion somehow or other rather than a loop:

julia> function foo2(::Val{N}) where N
           reduce(ntuple(identity, N); init=0) do acc, i
               acc + sum(ntuple(identity, Val{i}()))
           end
       end

julia> @btime foo2(Val{10}())
  1.072 ns (0 allocations: 0 bytes)
220

But this also has type stability limits, and wont compile away for all N. If you use N = 11 it could be a few orders of magnitude slower.

Edit: by specifing this is foldable you can make it compile away for all N

Base.@assume_effects foldable function foo2(::Val{N}) where N
   reduce(ntuple(identity, N); init=0) do acc, i
       acc + sum(ntuple(identity, Val{i}()))
   end
end

But, if you make N=1000 compilation may hang.

1 Like
5

An example that more faithfully/realistically represents your problem may help. In particular, is the final result supposed to get completely constant folded in your actual problem?

6

I think I’ve found the issue is actually more with Generators and Iterators.flatten.

I’ve posted a better MWE here if you’re interested.

1 Like