Dizzixx
Dizzixx

Reputation: 328

How to overload Base.view() for a custom struct of arrays type in Julia or extend StructArrays type with overloaded base.Push!() method

I should preface this by stating I am very new to the Julia language and I think I am missing something fundamental.

I am trying to setup a numerical simulation code and ultimately want it to work on both CPU and GPU. The method i am attempting is the Material Point Method (MPM) and it has been reported (here for instance) that generally the better memory footprints for this type of problem are SoA and AoSoA (Struct of Array and Array of Struct of Array) as they allow for coalesced memory access. As a large portion of MPM falls into the category of "embarrassingly parallel" this is of great benefit.

I have a pretty good idea of how I might approach this just using AoS (Array of Structs) as it follows the OOP paradigms I am more familiar with. I also have some idea of what a SoA should (I think...) look like when converting from AoS but am having difficulties realizing things in Julia (as I said the language is new to me)

I realize that the Julia package StructArrays gets close to what I am after but does not have push built in and I wanted to get more familiar with the language by doing things more manually.

What I have so far

struct Node{S,V,M}
    mass::S
    velocity::V
    affine::M
end

function Node()::Node
    return Node(0.,zeros(3),zeros(3,3))
end

struct Nodes{S,V,M}
    mass::Array{S}
    velocity::Array{V}
    affine::Array{M}
end

function Nodes(S,V,M)::Nodes
    return Nodes{S,V,M}(Array{S}(undef,0),Array{V}(undef,0),Array{M}(undef,0))
end

function Base.push!(nodes::Nodes)
    push!(nodes.mass,0.0)
    push!(nodes.velocity,zeros(3))
    push!(nodes.affine,zeros(3,3))
end

function Base.push!(nodes::Nodes,m,v,a)
    push!(nodes.mass,m)
    push!(nodes.velocity,v)
    push!(nodes.affine,a)
end

function Base.getindex(nodes::Nodes,idx::Int)::Node
    return Node(nodes.mass[idx],nodes.velocity[idx],nodes.affine[idx])
end

Which works like it should, e.g. -

julia> nodes = Nodes(Float64,SVector{3,Float64},SMatrix{3,3,Float64})

Nodes{Float64, SVector{3, Float64}, SMatrix{3, 3, Float64, L} where L}(Float64[], SVector{3, Float64}[], SMatrix{3, 3, Float64, L} where L[])

julia> @benchmark push!(nodes)

BenchmarkTools.Trial: 10000 samples with 952 evaluations.
Range (min … max): 96.954 ns … 409.226 μs ┊ GC (min … max): 0.00% … 84.19%
Time (median): 166.807 ns ┊ GC (median): 0.00%
Time (mean ± σ): 480.813 ns ± 7.516 μs ┊ GC (mean ± σ): 64.80 ± 6.64%

Memory estimate: 352 bytes, allocs estimate: 3.

julia> nodes.mass

10020502-element Vector{Float64}:

but I wanted to then extend Nodes to be compatible with views and slices but:

julia> view(nodes,1:1) ERROR: MethodError: no method matching view(::Nodes{Float64, SVector{3, Float64}, SMatrix{3, 3, Float64, L} where L}, ::UnitRange{Int64})

Which seemed fairly straight forward as I have not defined a method having that (or a compatible) signature. I thought I would look at one of the closest candidates (which just so happens to be from Struct Arrays) below:

Closest candidates are: view(::ChainRulesCore.AbstractZero, ::Any...) at C:\Users\DIi.julia\packages\ChainRulesCore\a4mIA\src\tangent_types\abstract_zero.jl:39 view(::StructArray{T, N, C, I} where I, ::Any...) where {T, N, C} at C:\Users\DIi.julia\packages\StructArrays\F5fDf\src\structarray.jl:361 view(::AbstractUnitRange, ::AbstractUnitRange{var"#s77"} where var"#s77"<:Integer) at subarray.jl:186 ...

Which I went and looked at

@inline function Base.view(s::StructArray{T, N, C}, I...) where {T, N, C}
    @boundscheck checkbounds(s, I...)
    StructArray{T}(map(v -> @inbounds(view(v, I...)), components(s)))
end

And I seem to know relatively little of what is going on...

  1. In the above what is "I"? I know T is the type of the structure, N is number of elements, C is a tuple argument which allows forming a StructArray using a NamedTuple (and is also used by the constructor function using new() to package the field components by name ...), but I have no idea what the I is other than in the base library it is a UniformScaling (as in Identity matrix, vector, value 1., etc.) but what is it doing in the call signature? Note, the docstring of structarray.jl explains T,N, and C as well as 'components' below but makes no mention of I.
  2. Why is "I" followed by ellipses?
  3. What does the where {T,C,N} mean? Is it just asking that all three parametric types where supplied as an assertion before continuing?
  4. What is "StructArray{T}(map(v -> @inbounds(view(v, I...)), components(s)))" doing? I get what @inbounds macro does and map i think is supposed to be applying view to each of the elements of the StructArray and then is calling its own constructor to instantiate that view of the existing allocation but what is with the "v,I..." and for that matter what is "v->" doing?

Lastly, more questions:

  1. Is there a more julia way to define the structs and methods of Nodes, Node, etc? I think I should be returning something more general like AbstractVector?

  2. How should I write a view method for the Nodes struct?

  3. How would I add a push!() method to extend use of StructArray as I think I will have better luck using that than rolling my own. I can use

    nodes = StructArray{Node}(undef)

which does allow me to define the first element as in

nodes[1]=Node() 

but then BoundErrors if I try to add another

nodes[2]=Node()

I could pre-define some length

StructArray{Node}(undef,N) 

but then don't know how to grow beyond that size and would like to be able to arbitrarily shrink or grow the Nodes arrays as I intend to use a hierarchal grid structure and get some kind of automatic refinement working.

If you got this far, thank you for your patience.

Upvotes: 1

Views: 232

Answers (1)

Dizzixx
Dizzixx

Reputation: 328

Turns out I got myself confused. While messing around with push!() on the individual arrays of a StructArray I ended up leaving them in a state that was incompatible with push!().

Turns out StructArray is compatible with psuh!() right out of the box so there is no reason to re-invent the implementation.

As an aside there is additionally memory overhead associated with use of StructArray but indexing in is faster so do with that what you will.

Upvotes: 0

Related Questions