Когда я читал код XLA в TF, я встречал следующие фрагменты кода:
https://github.com/tensorflow/tensorflow/blob/master/tensorflow/compiler/jit/xla_launch_util.h#L119
// Adapter class that wraps a Tensorflow allocator as an XLA allocator.
// Assumes that the Tensorflow allocator permits asynchronous deallocation:
// see comment on `AllowsAsynchronousDeallocation()`.
class XlaAllocator : public xla::DeviceMemoryAllocator {
public:
XlaAllocator(const se::Platform* platform, Allocator* wrapped);
~XlaAllocator() override;
xla::StatusOr<xla::OwningDeviceMemory> Allocate(
int device_ordinal, uint64 size, bool retry_on_failure) override;
Status Deallocate(int device_ordinal, se::DeviceMemoryBase mem) override;
// The Tensorflow BFC allocator used on GPU allows host-side deallocation
// before GPU execution takes place. Tensorflow uses the ordering of the main
// compute stream to enforce a happens-before relationship between a memory
// allocation and code that reuses the same memory. If Tensorflow adds
// support for multiple GPU streams or allocators with different ordering
// requirements, this code may need to change.
// (This attribute has no effect on CPU.)
bool AllowsAsynchronousDeallocation() const override { return true; }
private:
Allocator* wrapped_;
};
Как показано в комментарии, в нем говорилось, что распределитель Tensorflow BFCиспользование в GPU позволяет освободить хост-сторону перед выполнением GPU.
Это действительно смутило меня.Окунувшись в код BFCAllocator, я не нашел ничего, что могло бы это поддержать.
Итак, мой вопрос:
- Действительно ли Tensorflow BFC-распределитель действительно поддерживает асинхронное освобождение со стороны хоста?
- Если это правда, то чего мне не хватает?
Спасибо!