Remove some literals from shared data (#343)

* Improve which literals get put in shared data

* Check that strings do not appear in shared data

* Remove redundant line space

* Tidy up testing

* Bump patch

Project.toml

@@ -1,7 +1,7 @@
 name = "Mooncake"
 uuid = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 authors = ["Will Tebbutt, Hong Ge, and contributors"]
-version = "0.4.32"
+version = "0.4.33"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

src/interpreter/s2s_reverse_mode_ad.jl

@@ -65,10 +65,12 @@ IDInstPair[
 =#
 function shared_data_stmts(p::SharedDataPairs)::Vector{IDInstPair}
     return map(enumerate(p.pairs)) do (n, p)
-        return (p[1], new_inst(Expr(:call, getfield, Argument(1), n)))
+        return (p[1], new_inst(Expr(:call, get_shared_data_field, Argument(1), n)))
     end
 end
 
+@inline get_shared_data_field(shared_data, n) = getfield(shared_data, n)
+
 #=
 The block stack is the stack used to keep track of which basic blocks are visited on the
 forwards pass, and therefore which blocks need to be visited on the reverse pass. There is
@@ -453,10 +455,13 @@ end
 # a bits type, then it is returned. If it is not, then the CoDual is put into shared data,
 # and the ID associated to it in the forwards- and reverse-passes returned.
 function const_codual(stmt, info::ADInfo)
-    x = uninit_fcodual(get_const_primal_value(stmt))
-    return isbitstype(_typeof(x)) ? x : add_data!(info, x)
+    v = get_const_primal_value(stmt)
+    x = uninit_fcodual(v)
+    return safe_for_literal(v) ? x : add_data!(info, x)
 end
 
+safe_for_literal(v) = v isa String || v isa Type || isbitstype(_typeof(v))
+
 inc_or_const(stmt, info::ADInfo) = is_active(stmt) ? __inc(stmt) : const_codual(stmt, info)
 
 # Get the value associated to `x`. For `GlobalRef`s, verify that `x` is indeed a constant.

test/interpreter/s2s_reverse_mode_ad.jl

@@ -315,4 +315,12 @@ end
         rule = Mooncake.build_rrule(f, 0.0)
         @benchmark Mooncake.value_and_gradient!!($rule, $f, $(Ref(0.0))[])
     end
+    @testset "literal Strings do not appear in shared data" begin
+        f() = "hello"
+        @test length(build_rrule(Tuple{typeof(f)}).fwds_oc.oc.captures) == 2
+    end
+    @testset "Literal Types do not appear in shared data" begin
+        f() = Float64
+        @test length(build_rrule(Tuple{typeof(f)}).fwds_oc.oc.captures) == 2
+    end
 end