3 ms·
One thing that's important to note about memory management with XLA, is that inside a compiled program there's no user-exposed "malloc"/"free". The memory usag
by alevskaya 5y ago
One thing that's important to note about memory management with XLA, is that inside a compiled program there's no user-exposed "malloc"/"free". The memory usage and schedule of a given program/graph is statically optimized by the compiler (thus JAX requires static shapes when jitting). When running in op-by-op/eager mode the allocated buffers are coupled to the lifetime of the python array object, and are free'd when this array is garbage-collected.
- tubby12345 5y ago>The memory usage and schedule of a given program/graph is statically optimized by the compiler Is it really though? The only thing I see is https://github.com/tensorflow/tensorflow/blob/95cdeaa8c848fd4b89087232c23be44f0a160c41/tensorflow/compiler/tf2xla/xla_compiled_cpu_function.cc#L41 https://github.com/tensorflow/tensorflow/blob/95cdeaa8c848fd... which traces back to https://github.com/tensorflow/tensorflow/blob/54a8a3b373918bed29d7cfbdb8fd0a6c6065967c/tensorflow/compiler/xla/service/cpu/buffer_info_util.cc#L24 https://github.com/tensorflow/tensorflow/blob/54a8a3b373918b... which doesn't anything smart that i can tell.
- alevskaya 5y agoThat first call is from tf2xla which JAX doesn't use. But I think those calls just set up the global buffers used by the launched program inputs/temporaries. The entire point of XLA's static shape discipline is allowing for a highly optimized memory schedule inside compiled programs.
- tubby12345 5y agoI believe you but I'm wondering what that memory scheduling actually looks like (e.g. solving the full MIP?). do you have a code pointer to that scheduling routine?