How to efficiently create a struct to store parametric structs in tuples

New to Julia
1

Hi!, I’m working with a PDE and I want suggestions on how properly store components inside a struct with no type instabilities considering that components are parametric structs too.

For a MWE imagine something minimal like this (This behaves like what I need / want but don’t make any sense nor is correct.):

using BenchmarkTools

struct Offset{M}
    offsets::NTuple{M, Int}
end

Base.length(o::Offset{M}) where{M}  = M

struct SymmetricOffsets{N}
    offsets::NTuple{N,Offset}
end

Base.length(o::SymmetricOffsets) = sum(length,o.offsets)

function apply_offsets(val::NTuple{N,Int},offsets::SymmetricOffsets{N}) where N

    res = fill(val,length(offsets))
    
    idx = 1
    dim = 1
    for offset in offsets.offsets
        for oidx in offset.offsets
            tmp = val[dim] + oidx
            res[idx] = Base.setindex(res[idx],tmp,dim)
            idx += 1
        end
        dim += 1
    end
    res
end


struct SymmetricOffsets2{M1,M2,M3}
    offsets::Tuple{Offset{M1},Offset{M2},Offset{M3}}
end


Base.length(o::SymmetricOffsets2{M1,M2,M3}) where {M1,M2,M3} = M1+M2+M3

function apply_offsets(val::NTuple{N,Int},offsets::SymmetricOffsets2{M1,M2,M3}) where {N,M1,M2,M3}

    res = fill(val,length(offsets))
    
    idx = 1
    for dim in 1:3
        for oidx in offsets.offsets[dim].offsets
            tmp = val[dim] + oidx
            res[idx] = Base.setindex(res[idx],tmp,dim)
            idx += 1
        end
    end

    res
end

struct SymmetricOffsets3{M1,M2,M3}
    o1::Offset{M1}
    o2::Offset{M2}
    o3::Offset{M3}
end

Base.length(o::SymmetricOffsets3{M1,M2,M3}) where {M1,M2,M3} = M1+M2+M3


function apply_offsets(val::NTuple{N,Int},offsets::SymmetricOffsets3{M1,M2,M3}) where {N,M1,M2,M3}

    res = fill(val,length(offsets))
    
    idx = 1

    for oidx in offsets.o1.offsets
        tmp = val[1] + oidx
        res[idx] = Base.setindex(res[idx],tmp,1)
        idx += 1
    end

    for oidx in offsets.o2.offsets
        tmp = val[2] + oidx
        res[idx] = Base.setindex(res[idx],tmp,2)
        idx += 1
    end

    for oidx in offsets.o3.offsets
        tmp = val[3] + oidx
        res[idx] = Base.setindex(res[idx],tmp,3)
        idx += 1
    end

    res
end




offset1 = Offset((1,2,3))
offset2 = Offset((4,5,6,7))
offset3 = Offset((8,9,10,11,12))

symmetric_offsets = SymmetricOffsets((offset1,offset2,offset3))
symmetric_offsets2 = SymmetricOffsets2((offset1,offset2,offset3))
symmetric_offsets3 = SymmetricOffsets3(offset1,offset2,offset3)


@btime apply_offsets($(1,2,3),$symmetric_offsets)
@btime apply_offsets($(1,2,3),$symmetric_offsets2)
@btime apply_offsets($(1,2,3),$symmetric_offsets3)

@code_warntype apply_offsets((1,2,3),symmetric_offsets)
@code_warntype apply_offsets((1,2,3),symmetric_offsets2)
@code_warntype apply_offsets((1,2,3),symmetric_offsets3)

Note that if I unroll the for loop in SymmetricOffsets2 then I get similar results to SymmetricOffsets3.

Results:

  773.333 ns (34 allocations: 1.61 KiB)
  77.403 ns (7 allocations: 832 bytes)
  25.251 ns (1 allocation: 368 bytes)

MethodInstance for apply_offsets(::Tuple{Int64, Int64, Int64}, ::SymmetricOffsets{3})
  from apply_offsets(val::Tuple{Vararg{Int64, N}}, offsets::SymmetricOffsets{N}) where N @ Main 
Static Parameters
  N = 3
Arguments
  #self#::Core.Const(apply_offsets)
  val::Tuple{Int64, Int64, Int64}
  offsets::SymmetricOffsets{3}
Locals
  @_4::Union{Nothing, Tuple{Offset, Int64}}
  dim::Int64
  idx::Int64
  res::Any
  @_8::Union{Nothing, Tuple{Int64, Int64}}
  offset::Offset
  oidx::Int64
  tmp::Int64
Body::Any
1 ─ %1  = Main.length(offsets)::Any
│         (res = Main.fill(val, %1))
│         (idx = 1)
│         (dim = 1)
│   %5  = Base.getproperty(offsets, :offsets)::Tuple{Offset, Offset, Offset}
│         (@_4 = Base.iterate(%5))
│   %7  = (@_4::Core.PartialStruct(Tuple{Offset, Int64}, Any[Offset, Core.Const(2)]) === nothing)::Core.Const(false)
│   %8  = Base.not_int(%7)::Core.Const(true)
└──       goto #7 if not %8
2 ┄ %10 = @_4::Tuple{Offset, Int64}
│         (offset = Core.getfield(%10, 1))
│   %12 = Core.getfield(%10, 2)::Int64
│   %13 = Base.getproperty(offset, :offsets)::Tuple{Vararg{Int64, M}} where M
│         (@_8 = Base.iterate(%13))
│   %15 = (@_8 === nothing)::Bool
│   %16 = Base.not_int(%15)::Bool
└──       goto #5 if not %16
3 ┄ %18 = @_8::Tuple{Int64, Int64}
│         (oidx = Core.getfield(%18, 1))
│   %20 = Core.getfield(%18, 2)::Int64
│   %21 = Base.getindex(val, dim)::Int64
│         (tmp = %21 + oidx)
│   %23 = Base.setindex::Core.Const(Base.setindex)
│   %24 = Base.getindex(res, idx)::Any
│   %25 = tmp::Int64
│   %26 = (%23)(%24, %25, dim)::Any
│         Base.setindex!(res, %26, idx)
│         (idx = idx + 1)
│         (@_8 = Base.iterate(%13, %20))
│   %30 = (@_8 === nothing)::Bool
│   %31 = Base.not_int(%30)::Bool
└──       goto #5 if not %31
4 ─       goto #3
5 ┄       (dim = dim + 1)
│         (@_4 = Base.iterate(%5, %12))
│   %36 = (@_4 === nothing)::Bool
│   %37 = Base.not_int(%36)::Bool
└──       goto #7 if not %37
6 ─       goto #2
7 ┄       return res

MethodInstance for apply_offsets(::Tuple{Int64, Int64, Int64}, ::SymmetricOffsets2{3, 4, 5})
  from apply_offsets(val::Tuple{Vararg{Int64, N}}, offsets::SymmetricOffsets2{M1, M2, M3}) where {N, M1, M2, M3} @ Main 
Static Parameters
  N = 3
  M1 = 3
  M2 = 4
  M3 = 5
Arguments
  #self#::Core.Const(apply_offsets)
  val::Tuple{Int64, Int64, Int64}
  offsets::SymmetricOffsets2{3, 4, 5}
Locals
  @_4::Union{Nothing, Tuple{Int64, Int64}}
  idx::Int64
  res::Vector{Tuple{Int64, Int64, Int64}}
  @_7::Union{Nothing, Tuple{Int64, Int64}}
  dim::Int64
  oidx::Int64
  tmp::Int64
Body::Vector{Tuple{Int64, Int64, Int64}}
1 ─ %1  = Main.length(offsets)::Core.Const(12)
│         (res = Main.fill(val, %1))
│         (idx = 1)
│   %4  = (1:3)::Core.Const(1:3)
│         (@_4 = Base.iterate(%4))
│   %6  = (@_4::Core.Const((1, 1)) === nothing)::Core.Const(false)
│   %7  = Base.not_int(%6)::Core.Const(true)
└──       goto #7 if not %7
2 ┄ %9  = @_4::Tuple{Int64, Int64}
│         (dim = Core.getfield(%9, 1))
│   %11 = Core.getfield(%9, 2)::Int64
│   %12 = Base.getproperty(offsets, :offsets)::Tuple{Offset{3}, Offset{4}, Offset{5}}
│   %13 = Base.getindex(%12, dim)::Union{Offset{3}, Offset{4}, Offset{5}}
│   %14 = Base.getproperty(%13, :offsets)::Union{Tuple{Int64, Int64, Int64}, NTuple{4, Int64}, NTuple{5, Int64}}
│         (@_7 = Base.iterate(%14))
│   %16 = (@_7::Core.PartialStruct(Tuple{Int64, Int64}, Any[Int64, Core.Const(2)]) === nothing)::Core.Const(false)
│   %17 = Base.not_int(%16)::Core.Const(true)
└──       goto #5 if not %17
3 ┄ %19 = @_7::Tuple{Int64, Int64}
│         (oidx = Core.getfield(%19, 1))
│   %21 = Core.getfield(%19, 2)::Int64
│   %22 = Base.getindex(val, dim)::Int64
│         (tmp = %22 + oidx)
│   %24 = Base.setindex::Core.Const(Base.setindex)
│   %25 = Base.getindex(res, idx)::Tuple{Int64, Int64, Int64}
│   %26 = tmp::Int64
│   %27 = (%24)(%25, %26, dim)::Tuple{Int64, Int64, Int64}
│         Base.setindex!(res, %27, idx)
│         (idx = idx + 1)
│         (@_7 = Base.iterate(%14, %21))
│   %31 = (@_7 === nothing)::Bool
│   %32 = Base.not_int(%31)::Bool
└──       goto #5 if not %32
4 ─       goto #3
5 ┄       (@_4 = Base.iterate(%4, %11))
│   %36 = (@_4 === nothing)::Bool
│   %37 = Base.not_int(%36)::Bool
└──       goto #7 if not %37
6 ─       goto #2
7 ┄       return res

MethodInstance for apply_offsets(::Tuple{Int64, Int64, Int64}, ::SymmetricOffsets3{3, 4, 5})
  from apply_offsets(val::Tuple{Vararg{Int64, N}}, offsets::SymmetricOffsets3{M1, M2, M3}) where {N, M1, M2, M3} @ Main 
Static Parameters
  N = 3
  M1 = 3
  M2 = 4
  M3 = 5
Arguments
  #self#::Core.Const(apply_offsets)
  val::Tuple{Int64, Int64, Int64}
  offsets::SymmetricOffsets3{3, 4, 5}
Locals
  @_4::Union{Nothing, Tuple{Int64, Int64}}
  @_5::Union{Nothing, Tuple{Int64, Int64}}
  @_6::Union{Nothing, Tuple{Int64, Int64}}
  idx::Int64
  res::Vector{Tuple{Int64, Int64, Int64}}
  oidx@_9::Int64
  tmp@_10::Int64
  oidx@_11::Int64
  tmp@_12::Int64
  oidx@_13::Int64
  tmp@_14::Int64
Body::Vector{Tuple{Int64, Int64, Int64}}
1 ──       Core.NewvarNode(:(@_4))
│          Core.NewvarNode(:(@_5))
│    %3  = Main.length(offsets)::Core.Const(12)
│          (res = Main.fill(val, %3))
│          (idx = 1)
│    %6  = Base.getproperty(offsets, :o1)::Offset{3}
│    %7  = Base.getproperty(%6, :offsets)::Tuple{Int64, Int64, Int64}
│          (@_6 = Base.iterate(%7))
│    %9  = (@_6::Core.PartialStruct(Tuple{Int64, Int64}, Any[Int64, Core.Const(2)]) === nothing)::Core.Const(false)
│    %10 = Base.not_int(%9)::Core.Const(true)
└───       goto #4 if not %10
2 ┄─ %12 = @_6::Tuple{Int64, Int64}
│          (oidx@_9 = Core.getfield(%12, 1))
│    %14 = Core.getfield(%12, 2)::Int64
│    %15 = Base.getindex(val, 1)::Int64
│          (tmp@_10 = %15 + oidx@_9)
│    %17 = Base.setindex::Core.Const(Base.setindex)
│    %18 = Base.getindex(res, idx)::Tuple{Int64, Int64, Int64}
│    %19 = tmp@_10::Int64
│    %20 = (%17)(%18, %19, 1)::Tuple{Int64, Int64, Int64}
│          Base.setindex!(res, %20, idx)
│          (idx = idx + 1)
│          (@_6 = Base.iterate(%7, %14))
│    %24 = (@_6 === nothing)::Bool
│    %25 = Base.not_int(%24)::Bool
└───       goto #4 if not %25
3 ──       goto #2
4 ┄─ %28 = Base.getproperty(offsets, :o2)::Offset{4}
│    %29 = Base.getproperty(%28, :offsets)::NTuple{4, Int64}
│          (@_5 = Base.iterate(%29))
│    %31 = (@_5::Core.PartialStruct(Tuple{Int64, Int64}, Any[Int64, Core.Const(2)]) === nothing)::Core.Const(false)
│    %32 = Base.not_int(%31)::Core.Const(true)
└───       goto #7 if not %32
5 ┄─ %34 = @_5::Tuple{Int64, Int64}
│          (oidx@_11 = Core.getfield(%34, 1))
│    %36 = Core.getfield(%34, 2)::Int64
│    %37 = Base.getindex(val, 2)::Int64
│          (tmp@_12 = %37 + oidx@_11)
│    %39 = Base.setindex::Core.Const(Base.setindex)
│    %40 = Base.getindex(res, idx)::Tuple{Int64, Int64, Int64}
│    %41 = tmp@_12::Int64
│    %42 = (%39)(%40, %41, 2)::Tuple{Int64, Int64, Int64}
│          Base.setindex!(res, %42, idx)
│          (idx = idx + 1)
│          (@_5 = Base.iterate(%29, %36))
│    %46 = (@_5 === nothing)::Bool
│    %47 = Base.not_int(%46)::Bool
└───       goto #7 if not %47
6 ──       goto #5
7 ┄─ %50 = Base.getproperty(offsets, :o3)::Offset{5}
│    %51 = Base.getproperty(%50, :offsets)::NTuple{5, Int64}
│          (@_4 = Base.iterate(%51))
│    %53 = (@_4::Core.PartialStruct(Tuple{Int64, Int64}, Any[Int64, Core.Const(2)]) === nothing)::Core.Const(false)
│    %54 = Base.not_int(%53)::Core.Const(true)
└───       goto #10 if not %54
8 ┄─ %56 = @_4::Tuple{Int64, Int64}
│          (oidx@_13 = Core.getfield(%56, 1))
│    %58 = Core.getfield(%56, 2)::Int64
│    %59 = Base.getindex(val, 3)::Int64
│          (tmp@_14 = %59 + oidx@_13)
│    %61 = Base.setindex::Core.Const(Base.setindex)
│    %62 = Base.getindex(res, idx)::Tuple{Int64, Int64, Int64}
│    %63 = tmp@_14::Int64
│    %64 = (%61)(%62, %63, 3)::Tuple{Int64, Int64, Int64}
│          Base.setindex!(res, %64, idx)
│          (idx = idx + 1)
│          (@_4 = Base.iterate(%51, %58))
│    %68 = (@_4 === nothing)::Bool
│    %69 = Base.not_int(%68)::Bool
└───       goto #10 if not %69
9 ──       goto #8
10 ┄       return res

My questions are:

What is the most performant and easy way to have different parametric components stored inside a struct ?

Now, I’m thinking that this is only achievable with metaprogramming and generating combinations of parameters for different cases (N). Is that right?

For more context, I’m solving a PDE with any amount of components (Usually between 2 and 8) and for every component I need to store a pair of functions and some default values.

"PDE component"
struct SchrodingerPDEComponent{Tv,Fn,Hn}
    σ::Tv #Dispersion coefficient
    f::Fn 
    ψ::Hn #Initial condition
end

"Generic handler for non polynomic potentials"
struct SchrodingerPDENonPolynomic{N,M,Tv,Comp<:SchrodingerPDEComponent{Tv},Potential} <: SchrodingerPDE
    boundaries::NTuple{N,NTuple{2,Tv}}
    components::NTuple{M,Comp}
    F::Potential
    T::Tv
end

However, this kind of problem of store parametric structs inside other structs happens to me in a very wide varierty of parts of my solver and I want some insights on how to do that in a proper and Julian way. I don’t want to generate my structs and functions with macros but I’m open to it if that is the most efficient way to do that.

2

Here’s an issue. Offset is not fully parameterized here. This means ANY Offset can be in here.

Instead, you could parameterize the entire tuple.

struct SymmetricOffsets{T<:Tuple}
    offsets::T
end
3 Likes
3

Hi!

The general rule of thumb when defining parametric structs is that every field must have strictly defined type when all parameters are set. Let’s take a look at your example:

struct Offset{M}
    offsets::NTuple{M, Int}
end

struct SymmetricOffsets{N}
    offsets::NTuple{N,Offset}
end

The Offset types follows this rule - the offsets field of Offset{2} will be always NTuple{2,Int}, or, in other words, Tuple{Int,Int}. However, with SymmetricOffset{2} the type of offsets is not fully defined: it can be Tuple{Offset{1}, Offset{2}}, or Tuple{Offset{3}, Offset{4}} - you get the point.

To avoid this, you can make the type of offsets a parameter - like this:

struct SymmetricOffsets{TupleT}
    offsets::TupleT
end

This may look a bit redundant, but it makes the compiler’s work a lot easier. To enforce some type-checks for convenience, define an inner constructor:

struct SymmetricOffsets{TupleT}
    offsets::TupleT
    offsets(offsets::T) where {T<:Tuple{Vararg{Offset}}} = new{T}(offsets)
end

This will make the performance of SymmetricOffsets and SymmetricOffsets2 identical.

To make it as fast as SymmetricOffsets3, you need to avoid iterating over a tuple with undefined eltype (e. g. in your outer loop). This can be done by replacing it with recursion. Since we never iterate over the tuple explicitly, there will be no type instability.

const SymmetricOffsetsN{N} = SymmetricOffsets{<:Tuple{Vararg{Any,N}}}
@inline function apply_offsets!(res, val, idx, dim, offsets)
    isempty(offsets) && return
    offset = first(offsets)
    for oidx in offset.offsets
        tmp = val[dim] + oidx
        res[idx] = Base.setindex(res[idx],tmp,dim)
        idx += 1
    end
    apply_offsets!(res, val, idx, dim+1, Base.tail(offsets))
end
function apply_offsets(val::NTuple{N,Int},offsets::SymmetricOffsetsN{N}) where N
    res = fill(val,length(offsets))
    apply_offsets!(res, val, 1, 1, offsets.offsets)
    res
end

These are the benchmarks I got:

  37.462 ns (1 allocation: 368 bytes)
  132.453 ns (7 allocations: 832 bytes)
  38.066 ns (1 allocation: 368 bytes)

Also note that if all Offsets were of equal lengths, we would not have needed all this trickery

1 Like