Type instability for `SArray` wrapped in `struct`

General Usage
,
1

I’m struggling with an apparent type instability when wrapping a StaticArray in a struct. This also leads to a ~3x slowdown. Here’s an MWE running on version 0.6.0:

struct Foo
    x::SMatrix{4, 4, Float64}
end

function bar(f, M)
    B = M' * f.x * M
    return B
end
function baz(x, M)
    B = M' * x * M
    return B
end

S = @SMatrix rand(4, 4)
f = Foo(@SMatrix rand(4, 4))
x = f.x

Timings:

jl-0.6> @btime bar($f, $S);
  140.712 ns (4 allocations: 576 bytes)

jl-0.6> @btime baz($x, $S);
  47.050 ns (0 allocations: 0 bytes)

The output from @code_warntype is very long, so I just include some snippets:

jl-0.6> @code_warntype bar(f, S)
Variables:
  #self#::#bar
  f::Foo
  M::StaticArrays.SArray{Tuple{4,4},Float64,2,16}
  B::Any
  #temp#::StaticArrays.SArray{Tuple{4,4},Float64,2,16}

Body:
  begin
      $(Expr(:inbounds, false))
      [snip]
      # meta: location operators.jl * 424
      SSAValue(0) = ((StaticArrays._A_mul_B)($(QuoteNode(Size(4, 4))), $(QuoteNode(Size(4, 4))), #temp#::StaticArrays.SArray{Tuple{4,4},Float64,2,16}, (Core.getfield)(f::Foo, :x)::SMatrix{4,4,Float64})::Any * M::StaticArrays.SArray{Tuple{4,4},Float64,2,16})::Any
      # meta: pop location
      $(Expr(:inbounds, :pop)) # line 3:
      return SSAValue(0)
  end::Any

Note that B is inferred as Any, and the return type is also Any.

jl-0.6> @code_warntype baz(x, S)
Variables:
  #self#::#baz
  x::StaticArrays.SArray{Tuple{4,4},Float64,2,16}
  M::StaticArrays.SArray{Tuple{4,4},Float64,2,16}
  B::StaticArrays.SArray{Tuple{4,4},Float64,2,16}
  #temp#@_5::StaticArrays.SArray{Tuple{4,4},Float64,2,16}
  T@_6::Any
  #temp#@_7::StaticArrays.SArray{Tuple{4,4},Float64,2,16}
  T@_8::Any
  #temp#@_9::StaticArrays.SArray{Tuple{4,4},Float64,2,16}

Body:
  begin
      $(Expr(:inbounds, false))
      [snip]
      $(Expr(:inbounds, :pop)) # line 3:
      return #temp#@_9::StaticArrays.SArray{Tuple{4,4},Float64,2,16}
  end::StaticArrays.SArray{Tuple{4,4},Float64,2,16}

B is correctly inferred, as is the return type, though there are a couple of temps that are Any.

What’s going on here?

2

The type has an extra parameter for the total storage. Though planned as a future feature, right now julia can’t compute that on the fly, so you’ll have to pass it explicitly:

struct Foo
    x::SMatrix{4, 4, Float64, 4*4}
end
1 Like
3

Oh, that’s great! I was aware of that parameter, but I had the impression that is was not needed. I’m pretty sure I’ve never used it before. Is it only needed field types, not for dispatch?

4

Yes, dispatch is separate from specialization and the compiler will specialize on the subtype that has the variable set when the function called.

5

A follow-up question. How do I do this for arbitrary dimensions?

struct Foo{M, N}
    x::SMatrix{M, N, Float64, M*N}
end

doesn’t work.

6 
struct Foo{M, N, L}
  x::SMatrix{M, N, Float64, L}
end
1 Like
7

You can also use GitHub - vtjnash/ComputedFieldTypes.jl: Build types in Julia where some fields have computed types to compute L automatically

1 Like