Safeguard Memory for v1.10

Author: Qfl3x

src/Tracing.jl

@@ -1812,95 +1812,97 @@ Base.@nospecializeinfer function make_tracer(
     return res
 end
 
-
-Base.@nospecializeinfer function make_tracer(
-    seen,
-    @nospecialize(prev::Memory),
-    @nospecialize(path),
-    mode;
-    @nospecialize(track_numbers::Type = Union{}),
-    @nospecialize(sharding = Sharding.NoSharding()),
-    @nospecialize(runtime = nothing),
-    @nospecialize(device = nothing),
-    @nospecialize(client = nothing),
-    kwargs...,
-)
-    RT = Core.Typeof(prev)
-    # XXX: If someone wants to shard the same array with different shardings, we need to
-    #      somehow handle this correctly... Right now we just use the first sharding.
-    if mode != NoStopTracedTrack && haskey(seen, prev)
-        if mode == TracedToTypes
-            visited = seen[prev]
-            push!(path, visited)
-            return nothing
+if isdefined(Base, :Memory)
+    Base.@nospecializeinfer function make_tracer(
+        seen,
+        @nospecialize(prev::Memory),
+        @nospecialize(path),
+        mode;
+        @nospecialize(track_numbers::Type = Union{}),
+        @nospecialize(sharding = Sharding.NoSharding()),
+        @nospecialize(runtime = nothing),
+        @nospecialize(device = nothing),
+        @nospecialize(client = nothing),
+        kwargs...,
+    )
+        RT = Core.Typeof(prev)
+        # XXX: If someone wants to shard the same array with different shardings, we need to
+        #      somehow handle this correctly... Right now we just use the first sharding.
+        if mode != NoStopTracedTrack && haskey(seen, prev)
+            if mode == TracedToTypes
+                visited = seen[prev]
+                push!(path, visited)
+                return nothing
+            end
+            return seen[prev]
         end
-        return seen[prev]
-    end
-    if eltype(RT) <: ReactantPrimitive
-        if mode == ArrayToConcrete
-            runtime isa Val{:PJRT} &&
-                (return seen[prev] = ConcretePJRTArray(prev; sharding, device, client))
-            runtime isa Val{:IFRT} &&
-                (return seen[prev] = ConcreteIFRTArray(prev; sharding, device, client))
-            error("Unsupported runtime $runtime")
+        if eltype(RT) <: ReactantPrimitive
+            if mode == ArrayToConcrete
+                runtime isa Val{:PJRT} &&
+                    (return seen[prev] = ConcretePJRTArray(prev; sharding, device, client))
+                runtime isa Val{:IFRT} &&
+                    (return seen[prev] = ConcreteIFRTArray(prev; sharding, device, client))
+                error("Unsupported runtime $runtime")
+            elseif mode == TracedToTypes
+                # Original array can get mutated so we store a copy:
+                push!(path, copy(prev))
+                seen[prev] = VisitedObject(length(seen) + 1)
+                return nothing
+            end
         elseif mode == TracedToTypes
-            # Original array can get mutated so we store a copy:
-            push!(path, copy(prev))
-            seen[prev] = VisitedObject(length(seen) + 1)
+            push!(path, RT)
+            for I in eachindex(prev)
+                if isassigned(prev, I)
+                    pv = prev[I]
+                    make_tracer(
+                        seen,
+                        pv,
+                        path,
+                        mode;
+                        track_numbers,
+                        sharding,
+                        runtime,
+                        device,
+                        client,
+                        kwargs...,
+                    )
+                end
+            end
             return nothing
         end
-    elseif mode == TracedToTypes
-        push!(path, RT)
+        TT = traced_type(eltype(RT), Val(mode), track_numbers, sharding, runtime)
+        newa = Array{TT,ndims(RT)}(undef, size(prev))
+        seen[prev] = newa
+        same = true
         for I in eachindex(prev)
             if isassigned(prev, I)
                 pv = prev[I]
-                make_tracer(
+                nv = make_tracer(
                     seen,
                     pv,
-                    path,
+                    append_path(path, I),
                     mode;
                     track_numbers,
-                    sharding,
+                    sharding=Base.getproperty(sharding, I),
                     runtime,
                     device,
                     client,
                     kwargs...,
                 )
+                if pv !== nv
+                    same = false
+                end
+                @inbounds newa[I] = nv
             end
         end
-        return nothing
-    end
-    TT = traced_type(eltype(RT), Val(mode), track_numbers, sharding, runtime)
-    newa = Array{TT,ndims(RT)}(undef, size(prev))
-    seen[prev] = newa
-    same = true
-    for I in eachindex(prev)
-        if isassigned(prev, I)
-            pv = prev[I]
-            nv = make_tracer(
-                seen,
-                pv,
-                append_path(path, I),
-                mode;
-                track_numbers,
-                sharding=Base.getproperty(sharding, I),
-                runtime,
-                device,
-                client,
-                kwargs...,
-            )
-            if pv !== nv
-                same = false
-            end
-            @inbounds newa[I] = nv
+        if same
+            seen[prev] = prev
+            return prev
         end
+        return newa
     end
-    if same
-        seen[prev] = prev
-        return prev
-    end
-    return newa
 end
+
 Base.@nospecializeinfer function make_tracer(
     seen,
     @nospecialize(prev::Sharding.Mesh),

src/Types.jl

@@ -228,18 +228,20 @@ function ConcretePJRTArray(
     return ConcretePJRTArray{T,N,nsharded,typeof(shardinfo)}(sharded_data, shape, shardinfo)
 end
 
-function ConcretePJRTArray(
-    data::Memory{T};
-    client::Union{Nothing,XLA.PJRT.Client}=nothing,
-    idx::Union{Int,Nothing}=nothing,
-    device::Union{Nothing,XLA.PJRT.Device}=nothing,
-    sharding::Sharding.AbstractSharding=Sharding.NoSharding(),
-) where {T}
-    theclient, thedevice = _select_client_and_device(client, idx, device, sharding)
-    sharded_data, shardinfo = sharding(theclient, thedevice, data)
-    shape = size(data)
-    nsharded = length(sharded_data)
-    return ConcretePJRTArray{T,1,nsharded,typeof(shardinfo)}(sharded_data, shape, shardinfo)
+if isdefined(Base, :Memory)
+    function ConcretePJRTArray(
+        data::Memory{T};
+        client::Union{Nothing,XLA.PJRT.Client}=nothing,
+        idx::Union{Int,Nothing}=nothing,
+        device::Union{Nothing,XLA.PJRT.Device}=nothing,
+        sharding::Sharding.AbstractSharding=Sharding.NoSharding(),
+    ) where {T}
+        theclient, thedevice = _select_client_and_device(client, idx, device, sharding)
+        sharded_data, shardinfo = sharding(theclient, thedevice, data)
+        shape = size(data)
+        nsharded = length(sharded_data)
+        return ConcretePJRTArray{T,1,nsharded,typeof(shardinfo)}(sharded_data, shape, shardinfo)
+    end
 end
 
 Base.wait(x::Union{ConcretePJRTArray,ConcretePJRTNumber}) = foreach(wait, x.data)
@@ -370,17 +372,19 @@ function ConcreteIFRTArray(
     return ConcreteIFRTArray{T,N,typeof(shardinfo)}(sharded_data, shape, shardinfo, padding)
 end
 
-function ConcreteIFRTArray(
-    data::Memory{T};
-    client::Union{Nothing,XLA.IFRT.Client}=nothing,
-    idx::Union{Int,Nothing}=nothing,
-    device::Union{Nothing,XLA.IFRT.Device}=nothing,
-    sharding::Sharding.AbstractSharding=Sharding.NoSharding(),
-) where {T}
-    theclient, thedevice = _select_client_and_device(client, idx, device, sharding)
-    sharded_data, shardinfo, padding = sharding(theclient, nothing, data)
-    shape = size(data)
-    return ConcreteIFRTArray{T,1,typeof(shardinfo)}(sharded_data, shape, shardinfo)
+if isdefined(Base, :Memory)
+    function ConcreteIFRTArray(
+        data::Memory{T};
+        client::Union{Nothing,XLA.IFRT.Client}=nothing,
+        idx::Union{Int,Nothing}=nothing,
+        device::Union{Nothing,XLA.IFRT.Device}=nothing,
+        sharding::Sharding.AbstractSharding=Sharding.NoSharding(),
+    ) where {T}
+        theclient, thedevice = _select_client_and_device(client, idx, device, sharding)
+        sharded_data, shardinfo, padding = sharding(theclient, nothing, data)
+        shape = size(data)
+        return ConcreteIFRTArray{T,1,typeof(shardinfo)}(sharded_data, shape, shardinfo)
+    end
 end
 
 # Assemble data from multiple arrays. Needed in distributed setting where each process wont

src/xla/PJRT/Buffer.jl

@@ -21,20 +21,21 @@ function Buffer(client::Client, array::Array{T,N}, device::Device) where {T,N}
     return Buffer(buffer)
 end
 
-
-function Buffer(client::Client, memory::Memory{T}, device::Device) where {T}
-    sizear = collect(Int64, reverse(size(memory)))
-    buffer = GC.@preserve memory sizear begin
-        @ccall MLIR.API.mlir_c.ArrayFromHostBuffer(
-            client.client::Ptr{Cvoid},
-            pointer(memory)::Ptr{T},
-            XLA.primitive_type(T)::UInt64,
-            1::Csize_t,
-            pointer(sizear)::Ptr{Int64},
-            device.device::Ptr{Cvoid},
-        )::Ptr{Cvoid}
+if isdefined(Base, :Memory)
+    function Buffer(client::Client, memory::Memory{T}, device::Device) where {T}
+        sizear = collect(Int64, reverse(size(memory)))
+        buffer = GC.@preserve memory sizear begin
+            @ccall MLIR.API.mlir_c.ArrayFromHostBuffer(
+                client.client::Ptr{Cvoid},
+                pointer(memory)::Ptr{T},
+                XLA.primitive_type(T)::UInt64,
+                1::Csize_t,
+                pointer(sizear)::Ptr{Int64},
+                device.device::Ptr{Cvoid},
+            )::Ptr{Cvoid}
+        end
+        return Buffer(buffer)
     end
-    return Buffer(buffer)
 end
 
 function Base.similar(a::Buffer)