Replace IdDict for OrderedIdDict (#128)

Author: mofeing

Project.toml

@@ -8,6 +8,7 @@ CEnum = "fa961155-64e5-5f13-b03f-caf6b980ea82"
 Downloads = "f43a241f-c20a-4ad4-852c-f6b1247861c6"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 Libdl = "8f399da3-3557-5675-b5ff-fb832c97cbdb"
+OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
 Preferences = "21216c6a-2e73-6563-6e65-726566657250"
 Reactant_jll = "0192cb87-2b54-54ad-80e0-3be72ad8a3c0"
 Scratch = "6c6a2e73-6563-6170-7368-637461726353"

src/Compiler.jl

@@ -176,7 +176,7 @@ function compile_to_module!(mod, f, args; optimize=true)
         end
     end
 
-    concrete_seen = IdDict()
+    concrete_seen = OrderedIdDict()
 
     concrete_result = make_tracer(
         concrete_seen, traced_result, ("result",), TracedToConcrete

src/OrderedIdDict.jl

@@ -0,0 +1,59 @@
+using OrderedCollections: OrderedCollections
+
+struct OrderedIdDict{K,V} <: AbstractDict{K,V}
+    keys::OrderedCollections.OrderedDict{UInt,K}
+    values::OrderedCollections.OrderedDict{UInt,V}
+
+    function OrderedIdDict{K,V}(pairs) where {K,V}
+        return new(
+            OrderedCollections.OrderedDict{UInt,K}(objectid(k) => k for (k, _) in pairs),
+            OrderedCollections.OrderedDict{UInt,V}(objectid(k) => v for (k, v) in pairs),
+        )
+    end
+end
+
+OrderedIdDict() = OrderedIdDict{Any,Any}()
+OrderedIdDict{K,V}() where {K,V} = OrderedIdDict{K,V}(Pair{K,V}[])
+
+Base.show(io::IO, d::OrderedIdDict) = show(io, d.keys)
+
+OrderedCollections.isordered(::OrderedIdDict) = true
+
+Base.length(d::OrderedIdDict) = length(d.keys)
+Base.isempty(d::OrderedIdDict) = isempty(d.keys)
+
+function Base.getindex(d::OrderedIdDict, k)
+    return d.values[objectid(k)]
+end
+
+function Base.setindex!(d::OrderedIdDict, v, k)
+    d.keys[objectid(k)] = k
+    d.values[objectid(k)] = v
+    return d
+end
+
+function Base.haskey(d::OrderedIdDict, k)
+    return haskey(d.keys, objectid(k))
+end
+
+function Base.delete!(d::OrderedIdDict, k)
+    delete!(d.keys, objectid(k))
+    delete!(d.values, objectid(k))
+    return d
+end
+
+function Base.iterate(d::OrderedIdDict)
+    k = iterate(d.keys)
+    isnothing(k) && return nothing
+    ((_, k), _) = k
+    (_, v), state = iterate(d.values)
+    return k => v, state
+end
+
+function Base.iterate(d::OrderedIdDict, state)
+    k = iterate(d.keys, state)
+    isnothing(k) && return nothing
+    ((_, k), _) = k
+    (_, v), state = iterate(d.values, state)
+    return k => v, state
+end

src/Reactant.jl

@@ -1,5 +1,8 @@
 module Reactant
 
+# auxiliary types and functions
+include("OrderedIdDict.jl")
+
 using Enzyme
 
 abstract type RArray{T,N} <: AbstractArray{T,N} end

src/TracedRArray.jl

@@ -453,7 +453,7 @@ function elem_apply(f, args::Vararg{Any,Nargs}) where {Nargs}
         end
     end
 
-    seen_results = IdDict()
+    seen_results = OrderedIdDict()
     traced2_result = make_tracer(seen_results, result, (), TracedSetPath; tobatch=OutShape)
 
     func2.operation = MLIR.API.MlirOperation(C_NULL)

src/Tracing.jl

@@ -442,4 +442,6 @@ function make_tracer(seen, prev::Core.Box, @nospecialize(path), mode; kwargs...)
     return res
 end
 
-@inline to_rarray(@nospecialize(x)) = make_tracer(IdDict(), x, (), Reactant.ArrayToConcrete)
+@inline function to_rarray(@nospecialize(x))
+    return make_tracer(OrderedIdDict(), x, (), Reactant.ArrayToConcrete)
+end

src/utils.jl

@@ -27,7 +27,7 @@ function make_mlir_fn(f, args, kwargs, name="main", concretein=true; toscalar=fa
     end
 
     N = length(args)
-    seen_args = IdDict()
+    seen_args = OrderedIdDict()
     traced_args = ntuple(N) do i
         return make_tracer(
             seen_args,
@@ -101,7 +101,7 @@ function make_mlir_fn(f, args, kwargs, name="main", concretein=true; toscalar=fa
         end
     end
 
-    seen_results = IdDict()
+    seen_results = OrderedIdDict()
 
     traced_result = make_tracer(
         seen_results, result, (:result,), concretein ? TracedTrack : TracedSetPath

test/tracing.jl

@@ -82,7 +82,9 @@ using Test
                 (Val{0.5}, Val{0.5}),
                 (Val{:x}, Val{:x}),
             ]
-                tracedty = traced_type(origty, IdDict(), Val(ConcreteToTraced))
+                tracedty = traced_type(
+                    origty, Reactant.OrderedIdDict(), Val(ConcreteToTraced)
+                )
                 @test tracedty == targetty
             end
 
@@ -93,7 +95,7 @@ using Test
                 TracedRArray{Float64,3},
             ]
                 @test_throws Union{ErrorException,String} traced_type(
-                    type, IdDict(), Val(ConcreteToTraced)
+                    type, Reactant.OrderedIdDict(), Val(ConcreteToTraced)
                 )
             end
         end