Stream: helpdesk (published)

Topic: ✔ Check if a struct type is "fully generic"


view this post on Zulip 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

view this post on Zulip Sukera (Feb 11 2023 at 17:06):

all(==(Any), fieldtypes(foo)) can at least prevent S4

view this post on Zulip 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

view this post on Zulip 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

view this post on Zulip Mason Protter (Feb 11 2023 at 18:08):

Yeah I don’t believe we have a way to do this.

view this post on Zulip 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

view this post on Zulip 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

view this post on Zulip Brian Chen (Feb 11 2023 at 18:17):

Some custom marker perhaps. Either way, less important since it doesn't cover every case

view this post on Zulip 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

view this post on Zulip Sukera (Feb 11 2023 at 18:18):

more efficient for doing what?

view this post on Zulip Brian Chen (Feb 11 2023 at 18:18):

Storing accumulated values over time

view this post on Zulip 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)

view this post on Zulip Sukera (Feb 11 2023 at 18:19):

except less convenient, due to not being able to set individual fields

view this post on Zulip Sukera (Feb 11 2023 at 18:19):

but storing things continously over time really screams "struct of vectors" to me

view this post on Zulip 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

view this post on Zulip 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

view this post on Zulip Sukera (Feb 11 2023 at 18:25):

view this post on Zulip 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

view this post on Zulip 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

view this post on Zulip Brian Chen (Feb 11 2023 at 18:27):

*struct type. We can assume only mutable structs are being considered here

view this post on Zulip Brenhin Keller (Feb 11 2023 at 19:05):

What's the application for this?

view this post on Zulip 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.

view this post on Zulip 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).

view this post on Zulip 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

view this post on Zulip Brenhin Keller (Feb 11 2023 at 22:26):

Oh interesting

view this post on Zulip Brenhin Keller (Feb 11 2023 at 22:27):

Yeah that sounds better

view this post on Zulip 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

?

view this post on Zulip Mason Protter (Feb 11 2023 at 23:07):

I don't get why you'd want to wrap it in a Ref{Any}

view this post on Zulip Brian Chen (Feb 11 2023 at 23:26):

That's basically what I'm proposing

view this post on Zulip Brian Chen (Feb 11 2023 at 23:27):

The Ref{Any} is how things are done now

view this post on Zulip 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 :(

view this post on Zulip 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.

view this post on Zulip Mason Protter (Feb 20 2023 at 00:13):

@Brian Chen You can avoid that using something like Accessors.jl

view this post on Zulip 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

view this post on Zulip Mason Protter (Feb 20 2023 at 00:14):

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))

view this post on Zulip Brian Chen (Feb 20 2023 at 00:31):

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.

view this post on Zulip Brian Chen (Feb 20 2023 at 00:32):

(this is the case for using a type stable Ref as well)

view this post on Zulip 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.

view this post on Zulip 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
}

view this post on Zulip Mason Protter (Feb 20 2023 at 00:43):

I think unpacking and repacking should be relatively cheap compared to O(N) allocations

view this post on Zulip 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
}

view this post on Zulip Mason Protter (Feb 20 2023 at 01:07):

that was 1.9.0-beta4

view this post on Zulip 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 :)

view this post on Zulip Mason Protter (Feb 20 2023 at 01:09):

It'd be nice if we had @generated structs or something.

view this post on Zulip Notification Bot (Feb 20 2023 at 01:11):

Brian Chen has marked this topic as resolved.

view this post on Zulip Brian Chen (Feb 20 2023 at 01:11):

Very much so! For now, looks like compiler smarts are enough here.

view this post on Zulip 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)

view this post on Zulip 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.

view this post on Zulip Mason Protter (Mar 21 2023 at 23:58):

@Brian Chen are you still thinking about this?

view this post on Zulip 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 :)

view this post on Zulip 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

view this post on Zulip Brian Chen (Mar 22 2023 at 00:04):

Interesting. Is this more or less identical to merge when x is a NamedTuple?

view this post on Zulip 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

view this post on Zulip Mason Protter (Mar 22 2023 at 00:05):

Brian Chen said:

Interesting. Is this more or less identical to merge when x is a NamedTuple?

Yep

view this post on Zulip 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)

view this post on Zulip Brian Chen (Mar 22 2023 at 00:09):

As long as you don't have any massive Tuples stored in the struct, I imagine it'll be super quick

view this post on Zulip 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)

view this post on Zulip Mason Protter (Mar 22 2023 at 00:13):

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

view this post on Zulip aplavin (Mar 22 2023 at 11:28):

Isn't copy_with_changes the same as setproperties from ConstructionBase?

view this post on Zulip 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

view this post on Zulip Mason Protter (Mar 22 2023 at 15:57):

But it also doesnt do what I wanted because it cant handle uninitalized fields :frown:

view this post on Zulip aplavin (Mar 22 2023 at 16:09):

Hm, can you give some examples of these two issues?
Comparing copy_with_changes vs setproperties.

view this post on Zulip 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

view this post on Zulip Mason Protter (Oct 22 2023 at 18:09):

Ooh very nice


Last updated: Nov 06 2024 at 04:40 UTC