By which I mean it should be possible to construct a new instance with each field set to any type I want. So these would pass the test:
struct S1{A, B}
a::A
b::B
end
struct S2{A}
a::A
b::Any
end
But these would not because they put additional constraints on field types:
struct S3{A}
a::A
b::A
end
struct S4{A, B<:Vector{A}}
a::A
b::B
end
all(==(Any), fieldtypes(foo))
can at least prevent S4
I don't think we have an interface for querying whether two fields of a parametrized struct are constrained to be the same type
after all, you can do S3{Any}
and still put whatever you want into it, you just lose type stability
Yeah I don’t believe we have a way to do this.
The fieldtypes
trick is a good one. I didn't realize it defaulted to Any
for parametric types
what else should it default to? All field type limitations due to type parameters are just refinements of Any
after all
Some custom marker perhaps. Either way, less important since it doesn't cover every case
Context here is that I'm trying to figure out whether a Ref{Any} containing a NamedTuple, a NamedTuple of Refs or a mutable struct which allows all fields to be independently typed is more efficient
more efficient for doing what?
Storing accumulated values over time
my gut says that the raw NamedTuple
will be the same as a mutable struct with all types parametrized (that's more or less what it is already)
except less convenient, due to not being able to set individual fields
but storing things continously over time really screams "struct of vectors" to me
No history is required, would just be pure memory overhead. This is for accumulating gradients of arbitrary mutable structs
My main knowledge gap right now is knowing how each type of value is stored. Specifically what and how much is boxed
Any
)isbitstype
(or otherwise the object has a size that's not statically inferrable from the type) -> it's a pointer (also with type safety)making this work in a generic way is difficult, and with a struct on your end doubly so, since they can't hold a variable number of fields
That's why you repurpose the existing struct. But to do that you need to know whether you can plug the types you want into its fields, hence the original question
*struct type. We can assume only mutable structs are being considered here
What's the application for this?
@Brenhin Keller I was thinking about whether Zygote's current method for tracking gradients of mutable structs could be improved meaningfully. Right now the mechanism is a NamedTuple wrapped in a Ref{Any}
. Any update has to deref the ref, update the NamedTuple out-of-place and write it back. Rampant type instabilities aside, this seems pretty inefficient.
The problem is that most of the going alternatives seem to have tradeoffs. Using a NamedTuple of Refs like https://github.com/MasonProtter/MutableNamedTuples.jl means O(fields)
allocations even as we're trying to get rid of those by removing type instabilities. Creating a copy of the original struct with tangent types instead of primal ones would be great, but that requires knowing if it's safe to do so (hence this thread).
Currently I'm leaning towards defining an internal type which works like MArray
/ https://github.com/JuliaLang/julia/pull/35453. More complexity but should only require one allocation per instance and not incur double indirects on field reads/writes
Oh interesting
Yeah that sounds better
Why not just do something like
mutable struct AdjointStruct{T, fields, FieldTangentTypes}
fields::NamedTuple{fields, FieldTangentTypes}
end
?
I don't get why you'd want to wrap it in a Ref{Any}
That's basically what I'm proposing
The Ref{Any} is how things are done now
After looking into this more, it may be dead end because of https://github.com/JuliaArrays/StaticArrays.jl/blob/d6e0fde34e5b3f02a7399d89cbed9c8b6e036831/src/MArray.jl#L37-L39 :(
Maybe O(fields)
small allocations using a MutableNamedTuples-like struct wouldn't be too bad if they happen to be pooled with good locality.
@Brian Chen You can avoid that using something like Accessors.jl
using Accessors, ConstructionBase
mutable struct AdjointStruct{T, fields, FieldTangentTypes}
fields::NamedTuple{fields, FieldTangentTypes}
end
function Base.setproperty!(s::AdjointStruct{T, fieldnames, FTs}, name::Symbol, x) where {T, fieldnames, FTs}
fields = s.fields
newfields = set(fields, Accessors.opticcompose((PropertyLens){name}()), x)
setfield!(s, :fields, newfields)
end
ConstructionBase.constructorof(::Type{T}) where {T<: AdjointStruct} = T
julia> let s = AdjointStruct{Foo, (:x, :y), Tuple{String, Symbol}}((;x="hi", y=:bye))
@btime $s.x = "bye"
s
end
1.940 ns (0 allocations: 0 bytes)
AdjointStruct{Foo, (:x, :y), Tuple{String, Symbol}}((x = "bye", y = :bye))
Isn't that more or less the Ref approach but with a nicer interface? At least looking at @code_llvm
, I see unpacking of every field in s.fields
and then re-packing, all for what should be a write to a single field.
(this is the case for using a type stable Ref as well)
Perhaps that is unavoidable, but boy would it be nice to have the compiler synthesize an appropriate type like https://github.com/apple/swift/blob/main/docs/DifferentiableProgramming.md#synthesis-conditions.
Hm, well it does at least do the right thing for isbits types at least:
julia> let s = AdjointStruct{Complex{Int}, (:re, :im), Tuple{Int, Int}}((;re=1, im=2))
code_llvm((typeof(s),)) do s
s.re = 2
end
end
; @ REPL[31]:3 within `#7`
define i64 @"julia_#7_731"({}* noundef nonnull align 8 dereferenceable(16) %0) #0 {
top:
; ┌ @ REPL[27]:4 within `setproperty!`
%1 = bitcast {}* %0 to i64*
store i64 2, i64* %1, align 8
; └
ret i64 2
}
I think unpacking and repacking should be relatively cheap compared to O(N)
allocations
Are you on nightly? This is what I see on 1.8.5:
julia> let s3 = AdjointStruct{Complex{Int}, (:re, :im), Tuple{Int, Int}}((;re=1, im=2))
code_llvm((typeof(s3),)) do s
s.re = 2
end
end
; @ REPL[34]:3 within `#5`
define i64 @"julia_#5_1177"({}* nonnull align 8 dereferenceable(16) %0) #0 {
top:
; ┌ @ REPL[32]:2 within `setproperty!`
; │┌ @ Base.jl:38 within `getproperty`
%1 = bitcast {}* %0 to [2 x i64]*
%.elt3 = getelementptr inbounds [2 x i64], [2 x i64]* %1, i64 0, i64 1
%.unpack4 = load i64, i64* %.elt3, align 8
; │└
; │ @ REPL[32]:4 within `setproperty!`
%.repack = bitcast {}* %0 to i64*
store i64 2, i64* %.repack, align 8
%.repack5 = getelementptr inbounds [2 x i64], [2 x i64]* %1, i64 0, i64 1
store i64 %.unpack4, i64* %.repack5, align 8
; └
ret i64 2
}
that was 1.9.0-beta4
Mason Protter said:
I think unpacking and repacking should be relatively cheap compared to
O(N)
allocations
Most likely, though I've seen some monster types out of e.g. certain SciML libraries. At the very least it wouldn't be any worse than the status quo, but I was really hoping to be greedy here and find a way to write directly into fields of the tangent type :)
It'd be nice if we had @generated struct
s or something.
Brian Chen has marked this topic as resolved.
Very much so! For now, looks like compiler smarts are enough here.
I think also, if we're talking about performance sensitive applications, then probably people aren't putting non-isbits types inside of mutable containers (at least I'd hope so)
I see you're not acquainted with Flux.Recur
and Flux.BatchNorm
:laughter_tears:. But yes, the main objective here is to try to improve type stability (which directly affects TTFG in Zygote) while not regressing anywhere else.
@Brian Chen are you still thinking about this?
Every once in a while, though it's not as high on my priority list as I'd like :)
I was playing with something related, and realized there's a nice way to update many fields without unpacking and repacking multiple times at least. Basically just
@generated function copy_with_changes(x::T, vals::NamedTuple{keys}) where {T, keys}
args = map(fieldnames(T)) do key
key ∈ keys ? :(vals[$(QuoteNode(key))]) : :(getfield(x, $(QuoteNode(key))))
end
Expr(:new, T, args...)
end
Interesting. Is this more or less identical to merge
when x
is a NamedTuple?
E.g.
julia> struct Foo
a::Int
b::Any
c::String
d::Float64
end
julia> code_llvm((Foo,)) do x
copy_with_changes(x, (a=2, d=3.0))
end
; @ REPL[11]:2 within `#10`
define void @"julia_#10_981"({ i64, {}*, {}*, double }* noalias nocapture sret({ i64, {}*, {}*, double }) %0, [2 x {}*]* noalias nocapture %1, { i64, {}*, {}*, double }* nocapture nonnull readonly align 8 dereferenceable(32) %2) #0 {
top:
; ┌ @ REPL[2]:1 within `copy_with_changes`
; │┌ @ REPL[2]:1 within `macro expansion`
%3 = getelementptr inbounds { i64, {}*, {}*, double }, { i64, {}*, {}*, double }* %2, i64 0, i32 1
%4 = load atomic {}*, {}** %3 unordered, align 8
%5 = getelementptr inbounds { i64, {}*, {}*, double }, { i64, {}*, {}*, double }* %2, i64 0, i32 2
%6 = load atomic {}*, {}** %5 unordered, align 8
; └└
%7 = getelementptr inbounds [2 x {}*], [2 x {}*]* %1, i64 0, i64 0
store {}* %4, {}** %7, align 8
%8 = getelementptr inbounds [2 x {}*], [2 x {}*]* %1, i64 0, i64 1
store {}* %6, {}** %8, align 8
%.repack = getelementptr inbounds { i64, {}*, {}*, double }, { i64, {}*, {}*, double }* %0, i64 0, i32 0
store i64 2, i64* %.repack, align 8
%.repack1 = getelementptr inbounds { i64, {}*, {}*, double }, { i64, {}*, {}*, double }* %0, i64 0, i32 1
store {}* %4, {}** %.repack1, align 8
%.repack3 = getelementptr inbounds { i64, {}*, {}*, double }, { i64, {}*, {}*, double }* %0, i64 0, i32 2
store {}* %6, {}** %.repack3, align 8
%.repack5 = getelementptr inbounds { i64, {}*, {}*, double }, { i64, {}*, {}*, double }* %0, i64 0, i32 3
store double 3.000000e+00, double* %.repack5, align 8
ret void
}
Unfortunately, in the presence of non-isbits stuff you still need to repack the whole struct, but this would avoid doing it multiple times if you need to change multiple fields
Brian Chen said:
Interesting. Is this more or less identical to
merge
whenx
is a NamedTuple?
Yep
The repacking is quite fast:
julia> let f = Foo(1, 2, "hi", 4), a= Ref(1), d=Ref(3.0)
@btime copy_with_changes($f, (;a=$a[], d=$d[]))
end
2.478 ns (0 allocations: 0 bytes)
Foo(1, 2, "hi", 3.0)
As long as you don't have any massive Tuples stored in the struct, I imagine it'll be super quick
Even then it's pretty tolerable:
julia> let f = Foo2(ntuple(identity, 100), 1, "hi", 4.0), b = Ref(1), d=Ref(3.0)
@btime copy_with_changes($f, (;b=$b[], d=$d[]))
end;
22.746 ns (0 allocations: 0 bytes)
julia> let f = Foo2(ntuple(identity, 100), 1, "hi", 4.0), a = Ref(f.a), d=Ref(3.0)
@btime copy_with_changes($f, (;a=$a[], d=$d[]))
end;
20.242 ns (0 allocations: 0 bytes)
I think something soon I'll make a SumTypes 2.0 which tries to address the issues raised in https://github.com/JuliaLang/julia/discussions/48883, partially using this
Isn't copy_with_changes
the same as setproperties
from ConstructionBase
?
Not quite because it bypasses constructors, it’s really about the underlying data fields instead of properties
But it also doesnt do what I wanted because it cant handle uninitalized fields :frown:
Hm, can you give some examples of these two issues?
Comparing copy_with_changes vs setproperties.
Saw https://github.com/JuliaLang/julia/pull/51748 floating around on Slack today and thought of this thread. If something which works with non-isbits types does come out of this, I think you'd be interested @Mason Protter
Ooh very nice
Last updated: Nov 06 2024 at 04:40 UTC