Add get_memory_allocated() and get_memory_allocated_gb() functions

[x66ccff]

This PR will solve #503

You can use the functions like this to see how much memory is used by Reactant.jl!

## How to check the actual memory allocated by Reactant.jl?

julia> using Reactant

julia> a = rand(2, 100)
julia> ar = Reactant.to_rarray(a)

julia> get_memory_allocated()        # We can now check the memory usage like this
1600                                 # 200 elements * 8 bytes = 1600 bytes

julia> b = rand(Float32, 2, 100)
julia> br = Reactant.to_rarray(b)
julia> get_memory_allocated()
2400                                 # 200 elements * 4 bytes = 800 bytes, so we get 1600 + 800 = 2400 bytes

julia> ar = nothing                  # Now, lets release ar's memory
julia> GC.gc()                       # Use Julia's built-in GC

julia> get_memory_allocated()        # Then we check again, we indeed only have `br` in our device!
800

julia> get_memory_allocated_gb()     # Also you can use this to get memory usage in GBs.
7.450580596923828e-7

[Pangoraw]

If xla allocates memory from the start, then this count will be off?

Maybe we could leverage the xla APIs for this:

https://github.com/openxla/xla/blob/0cb8b51b761603f83e91d61be1d513aff77f1f82/xla/pjrt/pjrt_client.h#L204-L209

It seems we can get allocation stats per device.

[x66ccff]

I apologize, but I'm actually not very familiar with Julia and XLA, and I'm not sure how to proceed with improvements. I'm basically clueless when it comes to Julia and C++.

[x66ccff]

usage:

julia> using Reactant

julia> v = []
Any[]

julia> vr = []
Any[]

julia> for i=1:10
       push!(v,rand(1,1000000*2^i))
       end

julia> for i=1:10
       r = Reactant.to_rarray(v[i])
       push!(vr,r)
       end

julia> get_memory_allocated()
16368000000

julia> print_largest_arrays()
Top 5 arrays by memory usage:
1. Type: Float64, Shape: (1, 1024000000), Size: 7812.5 MB, Client: Reactant.XLA.Client(Ptr{Nothing} @0x0000000000ea8a80)
2. Type: Float64, Shape: (1, 512000000), Size: 3906.25 MB, Client: Reactant.XLA.Client(Ptr{Nothing} @0x0000000000ea8a80)
3. Type: Float64, Shape: (1, 256000000), Size: 1953.12 MB, Client: Reactant.XLA.Client(Ptr{Nothing} @0x0000000000ea8a80)
4. Type: Float64, Shape: (1, 128000000), Size: 976.56 MB, Client: Reactant.XLA.Client(Ptr{Nothing} @0x0000000000ea8a80)
5. Type: Float64, Shape: (1, 64000000), Size: 488.28 MB, Client: Reactant.XLA.Client(Ptr{Nothing} @0x0000000000ea8a80)

julia> get_memory_allocated_gb()
15.243887901306152

julia> print_largest_arrays(3)
Top 3 arrays by memory usage:
1. Type: Float64, Shape: (1, 1024000000), Size: 7812.5 MB, Client: Reactant.XLA.Client(Ptr{Nothing} @0x0000000000ea8a80)
2. Type: Float64, Shape: (1, 512000000), Size: 3906.25 MB, Client: Reactant.XLA.Client(Ptr{Nothing} @0x0000000000ea8a80)
3. Type: Float64, Shape: (1, 256000000), Size: 1953.12 MB, Client: Reactant.XLA.Client(Ptr{Nothing} @0x0000000000ea8a80)

[Pangoraw]

I apologize, but I'm actually not very familiar with Julia and XLA, and I'm not sure how to proceed with improvements.

No problems, I was just suggesting another possible option. I tried to kickstart the needed c++ in #517 but this might be of limited use since for example the stats are not implemented for the CPU device (I have yet to try on GPU).

[mofeing]

... the stats are not implemented for the CPU device (I have yet to try on GPU)

not sure about this, but i checked out the internals a couple of weeks ago. could it be that it works for GPU because it uses a memory pool for that, while for CPU it just allocates as requested?

[wsmoses]

closed in favor of #517