Using ntuple() in generated function

General Usage Performance
,
1

I have found that using ntuple in generated functions produces better code (faster, and with fewer allocations), than constructing the tuple expression manually. Presumably the code that ntuple generates is more inferable. From a maintenance perspective it feels preferable to use it too.

However, I can only use ntuple on the compile time constants. I have been unable to get it work with variables. To show you what I mean, the following code throws an error:

@generated function myfunc1(x::T, ::Val{N}) where {T, N}
    ex = Expr(:tuple, [:(x * $i) for i in 1:N]...)
    return :(SVector{N, T}($ex))
end

@generated function myfunc2(x::T, ::Val{N}) where {T, N}
    ex = ntuple(i -> :(x * $i), N)
    return :(SVector{N, T}($ex))
end

display(myfunc1(2., Val(3)))
display(myfunc2(2., Val(3)))

Output:

3-element SVector{3, Float64} with indices SOneTo(3):
2.0
4.0
6.0
ERROR: MethodError: Cannot convert an object of type Expr to an object of type Float64

How can I fix myfunc2 to work like myfun1?

2

Don’t you want ex = :(ntuple(i -> x * i, $N))

1 Like
3

Thank you. But this gives the error:

ERROR: The function body AST defined by this @generated function is not pure. This likely means it contains a closure, a comprehension or a generator.

4

However, this works:

mult(a, b) = a * b
@generated function myfunc2(x::T, ::Val{N}) where {T, N}
    ex = :(ntuple(Base.Fix1(mult, x), $N))
    return :(SVector{N, T}($ex))
end

Many thanks for the nudge in the right direction.

5

You can go one step simpler with

@generated function myfunc2(x::T, ::Val{N}) where {T, N}
    ex = :(ntuple(Base.Fix1(*, x), $N))
    return :(SVector{N, T}($ex))
end
1 Like
6

Also, without resorting to Base.Fix1 (which is completely appropriate here):

nt(x,n) = ntuple(i->i*x,n)
@generated function myfunc3(x::T, ::Val{N}) where {T, N}
       ex = :(nt(x, $N))
       return :(SVector{N, T}($ex))
end

The non-pureness earlier comes from defining the anonymous function.
This example will work in other cases where Base.Fix1 will not fit.

1 Like