Stream: helpdesk (published)

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

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 Alexander (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 Alexander (Mar 22 2023 at 16:09):

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

Last updated: Oct 02 2023 at 04:34 UTC