Help with type instability needed

General Usage
1

I would like to write a function grideval(f,x) that evaluates a function f on the cartesian grid described by the tuple x of one-dimensional grid coordinates. This is what I have so far:

"""
    cartesian(x::AbstractVector...)
    cartesian(x::NTuple{N,AbstractVector})

Cartesian product.

# Examples
``
julia> cartesian([1,2],[3,4])
2×2 ApproxTools.TensorGrid{...}:
 (1, 3)  (1, 4)
 (2, 3)  (2, 4)
``
"""
cartesian(x::AbstractVector...) = cartesian(x)
cartesian(x::NTuple{N,AbstractVector}) where {N} = TensorGrid{Tuple{eltype.(x)...},length(x),typeof(x)}(x)
struct TensorGrid{T,N,F} <: AbstractArray{T,N}
    factors::F
end
Base.size(x::TensorGrid) = length.(x.factors)
Base.getindex(x::TensorGrid{<:Any,N,<:Any}, i::Vararg{Int,N}) where {N} = map(getindex,x.factors,i)

"""
    grideval(f, x::AbstractVector...)
    grideval(f, x::NTuple{N,AbstractVector})

Evaluate `f` on the grid spanned by the `x`.

# Examples
``
julia> grideval(*, ([1,2],[3,4]))
2×2 Array{Int64,2}:
 3  4
 6  8
``
"""
grideval(f, x::AbstractVector...) = grideval(f,x)
grideval(f, x::NTuple{N,AbstractVector}) where {N} = (x->f(x...)).(cartesian(x))

Unfortunately, grideval turns out to have strange type-stability properties. For example, this is fine:

julia> function f1()
           x = linspace(0,1,10)
           grideval((x1,x2)->2*x1*x2, (x,x))
       end
       @code_warntype f1()
Variables:
  #self# <optimized out>
  #11::##11#12
  x::StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}

Body:
  begin 
      x::StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}} = $(Expr(:invoke, MethodInstance for linspace(::Type{Float64}, ::Int64, ::Int64, ::Int64, ::Int64), :(Base.linspace), :(Base.Float64), 0, 1, 10, 1)) # line 3:
      #11::##11#12 = $(Expr(:new, :(Main.##11#12)))
      SSAValue(1) = (Core.tuple)(x::StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}, x::StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}})::Tuple{StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}},StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}
      return $(Expr(:invoke, MethodInstance for grideval(::Function, ::Tuple{StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}},StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}), :(Main.grideval), :(#11), SSAValue(1)))
  end::Array{Float64,2}

But this isn’t:

julia> function f2()
           a = 2
           x = linspace(0,1,10)
           grideval((x1,x2)->a*x1*x2, (x,x))
       end
       @code_warntype f2()
Variables:
  #self# <optimized out>
  #13::##13#14{Int64}
  a <optimized out>
  x::StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}

Body:
  begin  # line 3:
      x::StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}} = $(Expr(:invoke, MethodInstance for linspace(::Type{Float64}, ::Int64, ::Int64, ::Int64, ::Int64), :(Base.linspace), :(Base.Float64), 0, 1, 10, 1)) # line 4:
      #13::##13#14{Int64} = $(Expr(:new, ##13#14{Int64}, 2))
      SSAValue(1) = (Core.tuple)(x::StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}, x::StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}})::Tuple{StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}},StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}
      return $(Expr(:invoke, MethodInstance for grideval(::Function, ::Tuple{StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}},StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}), :(Main.grideval), :(#13), SSAValue(1)))
  end::Any

Is there a way to make the second version type-stable as well?

2

This probably works:

function f2()
           a = 2
           x = linspace(0,1,10)
           let a=a
           grideval((x1,x2)->a*x1*x2, (x,x))
       end
end

It’s the dreaded performance of captured variables in closures · Issue #15276 · JuliaLang/julia · GitHub

1 Like
3

Unfortunately it doesn’t:

julia> @code_warntype f2()
Variables:
  #self# <optimized out>
  #3::##3#4{Int64}
  a@_3 <optimized out>
  a@_4 <optimized out>
  x::StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}

Body:
  begin  # line 3:
      x::StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}} = $(Expr(:invoke, MethodInstance for linspace(::Type{Float64}, ::Int64, ::Int64, ::Int64, ::Int64), :(Base.linspace), :(Base.Float64), 0, 1, 10, 1)) # line 5:
      #3::##3#4{Int64} = $(Expr(:new, ##3#4{Int64}, 2))
      SSAValue(2) = (Core.tuple)(x::StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}, x::StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}})::Tuple{StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}},StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}
      return $(Expr(:invoke, MethodInstance for grideval(::Function, ::Tuple{StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}},StepRangeLen{Float64,Base.TwicePrecision{Float64},Base.TwicePrecision{Float64}}}), :(Main.grideval), :(#3), SSAValue(2)))
  end::Any

I am aware of the issues with closures capturing variables, but as far as I can tell they shouldn’t apply here: even without the let block, a is just a simple local integer variable, and the compiler claims it optimised it away.

4

It is type stable in Julia 0.7.

For 0.6 you can try this workaround:

grideval(f, ::Type{T}, x::NTuple{N,AbstractVector}) where {T,N} = (x->f(x...)::T).(cartesian(x))
function f2()
    a = 2
    x = linspace(0,1,10)
    grideval((x1,x2)->a*x1*x2, Float64, (x,x))
end
2 Likes
6

For completeness, here’s an alternative approach to grideval that is type-stable on 0.6 as well:

# Utility function for grideval
# x is only used for its length. It would be cleaner to pass Val(N) instead, but unlike
# Base.tail, Val(N-1) is not type-stable. 
shape(x::NTuple{N},i,l) where {N} = (shape(Base.tail(x),i,l)..., i == N ? l : 1)
shape(::NTuple{0},i,l) = ()
7

For completeness, here’s an alternative approach to grideval that is type-stable on 0.6 as well:

"""
    grideval(f, x::AbstractVector...)
    grideval(f, x::NTuple{N,AbstractVector})

Evaluate `f` on the grid spanned by the `x`.

# Examples
``
julia> grideval(*, ([1,2],[3,4]))
2×2 Array{Int64,2}:
 3  4
 6  8
``
"""
grideval(f, x::AbstractVector...) = grideval(f,x)
@generated grideval(f, x::NTuple{N,AbstractVector}) where {N} = :(Base.Cartesian.@ncall($N,broadcast,f,i->reshape(x[i],shape(x,i))))

# Utility function for grideval
# x is only used for its length. It would be cleaner to pass Val(N) instead, but unlike
# Base.tail, Val(N-1) is not type-stable.
shape(x::NTuple{N,Any},i) where {N} = (shape(Base.tail(x),i)..., i == N ? length(x[i]) : 1)
shape(::NTuple{0},i) = ()