Type stability using reduce

General Usage
, ,
1

I’d like to better understand how to ensure type stability when applying a reduce operation on a collection. In the example below, the CrossProd type maintains a collection of objects that represent a “cross product”, which are assembled using the cross method:

abstract type Atom end
struct Cd <: Atom
    w::Int64
end
struct CrossProd{T<:Tuple{Vararg{Atom}}} <: Atom
    a::T
end
cross(A::Atom, B::Atom) = CrossProd(tuple(A,B))
cross(A::CrossProd, B::Atom) = CrossProd(tuple(A.a..., B))
cross(A::Atom, B::CrossProd) = CrossProd(tuple(A, B.a...))
cross(A::CrossProd, B::CrossProd) = CrossProd(tuple(A.a..., B.a...))

If a fixed number of atoms is collected into a tuple, the result is type stable:

function testcross()
    nc = (Cd(1), Cd(2))
    nc2 = reduce(cross, nc)
end
julia> @code_warntype testcross()

Variables
  #self#::Core.Compiler.Const(testcross, false)
  nc::Tuple{Cd,Cd}
  nc2::CrossProd{Tuple{Cd,Cd}}

Body::CrossProd{Tuple{Cd,Cd}}
1 ─ %1 = Main.Cd(1)::Core.Compiler.Const(Cd(1), false)
│   %2 = Main.Cd(2)::Core.Compiler.Const(Cd(2), false)
│        (nc = Core.tuple(%1, %2))
│   %4 = Main.reduce(Main.cross, nc::Core.Compiler.Const((Cd(1), Cd(2)), false))::Core.Compiler.Const(CrossProd{Tuple{Cd,Cd}}((Cd(1), Cd(2))), false)
│        (nc2 = %4)
└──      return %4

Now suppose we wish to collect an arbitrary number of atoms. I might do this as follows, which isn’t type stable:

function testcross(n)
    nc = ntuple(i->Cd(i), n)
    nc2 = reduce(cross, nc)
end
julia> @code_warntype testcross(2)

Variables
  #self#::Core.Compiler.Const(testcross, false)
  n::Int64
  #3::var"#3#4"
  nc::Tuple{Vararg{Cd,N} where N}
  nc2::Union{Cd, CrossProd}

Body::Union{Cd, CrossProd}
1 ─      (#3 = %new(Main.:(var"#3#4")))
│   %2 = #3::Core.Compiler.Const(var"#3#4"(), false)
│        (nc = Main.ntuple(%2, n))
│   %4 = Main.reduce(Main.cross, nc)::Union{Cd, CrossProd}
│        (nc2 = %4)
└──      return %4

A few questions:

  • Why is the type of variable nc indeterminate? Couldn’t the compiler infer the length of the tuple from the function argument?
  • Why does nc2 have the type Union{Cd, CrossProd}? I had expected it to be of type CrossProd{Tuple{Cd,Cd}}, as in the first example.
  • Can the definition of cross be changed so that reduce can infer the right type?

Thanks!

2

This is because you get Cd(1) from testcross(1). Note that the compiler only knows the type of n and not the value of it. If you want to check if const propagation works, you can wrap it in another function:

julia> f() = testcross(2)
f (generic function with 1 method)

julia> @code_warntype f()
Variables
  #self#::Core.Compiler.Const(f, false)

Body::CrossProd{Tuple{Cd,Cd}}
1 ─ %1 = Main.testcross(2)::Core.Compiler.Const(CrossProd{Tuple{Cd,Cd}}((Cd(1), Cd(2))), false)
└──      return %1

In this particular case, you can pass Val(2) to testcross as ntuple accepts Val types:

julia> @code_warntype testcross(Val(2))
Variables
  #self#::Core.Compiler.Const(testcross, false)
  n::Core.Compiler.Const(Val{2}(), false)
  #3::var"#3#4"
  nc::Tuple{Cd,Cd}
  nc2::CrossProd{Tuple{Cd,Cd}}

Body::CrossProd{Tuple{Cd,Cd}}
1 ─      (#3 = %new(Main.:(var"#3#4")))
│   %2 = #3::Core.Compiler.Const(var"#3#4"(), false)
│        (nc = Main.ntuple(%2, n))
│   %4 = Main.reduce(Main.cross, nc::Core.Compiler.Const((Cd(1), Cd(2)), false))::Core.Compiler.Const(CrossProd{Tuple{Cd,Cd}}((Cd(1), Cd(2))), false)
│        (nc2 = %4)
└──      return %4
2 Likes
3

Thanks for the explanation, @tkf. In the context I’m working in, there are two additional methods

Base.:(*)(t::Real, A::CrossProd) = CrossProd(map(a->t*a, A.a))
Base.:(*)(t::Real, A::Cd) = Cd(t*A.w)

which will be called repeatedly within a loop. My worry is that calls to the CrossProd variant will be dynamically dispatched and incur unnecessary overhead. In the code below, I replaced the call to reduce with a call directly to the CrossProd constructor:

function testcross(n)
    nc = ntuple(i->Cd(i),n)
    nc2 = CrossProd(nc)
    nc3 = 2*nc2
end

julia> @code_warntype testcross(3)
Variables
  #self#::Core.Compiler.Const(testcross, false)
  n::Int64
  #17::var"#17#18"
  nc::Tuple{Vararg{Cd,N} where N}
  nc2::CrossProd
  nc3::CrossProd

Body::CrossProd
1 ─      (#17 = %new(Main.:(var"#17#18")))
│   %2 = #17::Core.Compiler.Const(var"#17#18"(), false)
│        (nc = Main.ntuple(%2, n))
│        (nc2 = Main.CrossProd(nc))
│   %5 = (2 * nc2)::CrossProd
│        (nc3 = %5)
└──      return %5

The output of @code_warntype highlights the type for nc3 in red, even though it seems to correctly infer the type. Presumably that’s because of the dynamic dispatch.

Is there a way to reparameterize these types to avoid the type inference issue? Also, should I avoid using a tuple in the definition of CrossProd, and instead use an array?

4

Julia’s type inference is always correct (unless you hit by a bug) although it may not be concrete (see, e.g., ?isconcretetype). For example, CrossProd{Tuple{Cd,Cd}} is concrete but CrossProd is not (you don’t know how many Atoms it has). You’d want to have concrete types if you want speed (although small Union is OK). That’s why @code_warntype shows them in red.

It depends on your usecase. If you only have a few Atoms (or maybe a few tens of them if you are OK with long compilation time), I think using tuples is a good choice although you need to program carefully to make things type stable (e.g., you’d need to write things in “functional” style rather than using loops). If you have a lot of Atoms, I think using arrays is a better solution.

1 Like