Brian Chen (Feb 11 2023 at 16:38): Brian Chen (Feb 11 2023 at 16:38):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
Sukera (Feb 11 2023 at 17:07):I don't think we have an interface for querying whether two fields of a parametrized struct are constrained to be the same type
Sukera (Feb 11 2023 at 17:08):after all, you can do S3{Any} and still put whatever you want into it, you just lose type stability
Mason Protter (Feb 11 2023 at 18:08):Yeah I don’t believe we have a way to do this.
Brian Chen (Feb 11 2023 at 18:13):The fieldtypes trick is a good one. I didn't realize it defaulted to Any for parametric types
Sukera (Feb 11 2023 at 18:15):what else should it default to? All field type limitations due to type parameters are just refinements of Any after all
Brian Chen (Feb 11 2023 at 18:17):Some custom marker perhaps. Either way, less important since it doesn't cover every case
Brian Chen (Feb 11 2023 at 18:17):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
Sukera (Feb 11 2023 at 18:18):more efficient for doing what?
Brian Chen (Feb 11 2023 at 18:18):Storing accumulated values over time
Sukera (Feb 11 2023 at 18:19):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)
Sukera (Feb 11 2023 at 18:19):except less convenient, due to not being able to set individual fields
Sukera (Feb 11 2023 at 18:19):but storing things continously over time really screams "struct of vectors" to me
Brian Chen (Feb 11 2023 at 18:20):No history is required, would just be pure memory overhead. This is for accumulating gradients of arbitrary mutable structs
Brian Chen (Feb 11 2023 at 18:22):My main knowledge gap right now is knowing how each type of value is stored. Specifically what and how much is boxed
Sukera (Feb 11 2023 at 18:25):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) Sukera (Feb 11 2023 at 18:25):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
Brian Chen (Feb 11 2023 at 18:26):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
Brian Chen (Feb 11 2023 at 18:27):*struct type. We can assume only mutable structs are being considered here
Brenhin Keller (Feb 11 2023 at 19:05):What's the application for this?
Brian Chen (Feb 11 2023 at 22:07):@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.
Brian Chen (Feb 11 2023 at 22:11):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).
Brian Chen (Feb 11 2023 at 22:13):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
Brenhin Keller (Feb 11 2023 at 22:26):Oh interesting
Brenhin Keller (Feb 11 2023 at 22:27):Yeah that sounds better
Mason Protter (Feb 11 2023 at 23:05):Why not just do something like
mutable struct AdjointStruct{T, fields, FieldTangentTypes}
fields::NamedTuple{fields, FieldTangentTypes}
end
?
Mason Protter (Feb 11 2023 at 23:07):I don't get why you'd want to wrap it in a Ref{Any}
Brian Chen (Feb 11 2023 at 23:26):That's basically what I'm proposing
Brian Chen (Feb 11 2023 at 23:27):The Ref{Any} is how things are done now
Brian Chen (Feb 19 2023 at 22:11):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 :(
Brian Chen (Feb 19 2023 at 22:30):Maybe O(fields) small allocations using a MutableNamedTuples-like struct wouldn't be too bad if they happen to be pooled with good locality.
Mason Protter (Feb 20 2023 at 00:13):@Brian Chen You can avoid that using something like Accessors.jl
Mason Protter (Feb 20 2023 at 00:14):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.
Brian Chen (Feb 20 2023 at 00:32):(this is the case for using a type stable Ref as well)
Brian Chen (Feb 20 2023 at 00:35):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.
Mason Protter (Feb 20 2023 at 00:40):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
Brian Chen (Feb 20 2023 at 01:04):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
Brian Chen (Feb 20 2023 at 01:08):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 :)
Mason Protter (Feb 20 2023 at 01:09):It'd be nice if we had @generated structs or something.
Notification Bot (Feb 20 2023 at 01:11):Brian Chen has marked this topic as resolved.
Brian Chen (Feb 20 2023 at 01:11):Very much so! For now, looks like compiler smarts are enough here.
Mason Protter (Feb 20 2023 at 02:47):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)
Brian Chen (Feb 20 2023 at 15:27):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.
Mason Protter (Mar 21 2023 at 23:58):@Brian Chen are you still thinking about this?
Brian Chen (Mar 22 2023 at 00:00):Every once in a while, though it's not as high on my priority list as I'd like :)
Mason Protter (Mar 22 2023 at 00:00):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?
Mason Protter (Mar 22 2023 at 00:05):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
Mason Protter (Mar 22 2023 at 00:05):Brian Chen said:
Interesting. Is this more or less identical to
mergewhenxis a NamedTuple?
Yep
Mason Protter (Mar 22 2023 at 00:08):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
Mason Protter (Mar 22 2023 at 00:11):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
aplavin (Mar 22 2023 at 11:28):Isn't copy_with_changes the same as setproperties from ConstructionBase?
Mason Protter (Mar 22 2023 at 15:56):Not quite because it bypasses constructors, it’s really about the underlying data fields instead of properties
Mason Protter (Mar 22 2023 at 15:57):But it also doesnt do what I wanted because it cant handle uninitalized fields :frown:
aplavin (Mar 22 2023 at 16:09):Hm, can you give some examples of these two issues?
Comparing copy_with_changes vs setproperties.
Brian Chen (Oct 22 2023 at 18:04):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
Mason Protter (Oct 22 2023 at 18:09):Ooh very nice
Last updated: Dec 28 2024 at 04:38 UTC