@annarev Thanks for the PR!
Do you mean the PluggableDeviceSimpleAllocator will directly call the stream_exec_->AllocateArray(), which calls SP_StreamExecutor::allocate(const SP_Device* device, uint64_t size, int64_t memory_space,SP_DeviceMemoryBase* mem ), and plugin authors decides whether allocate is build on top of a custom allocator? if it is, I think it is better to add an field such as void* custom_allocator in SP_Device? since the allocator should be per-device? Thanks

Thanks @annarev for the PR. SimpleAllocator can be something straightforward calling into AllocateRaw and that would be fine for us. We had brief discussion of this during the design phase. To keep it simple we can leverage SubAllocator currently. After the pluggable device implementation merges we can use the use_bfc_allocator flag to pass this information in device_process_state..

There are few places in pluggable_device_process_state like AllocatorParts will require changes to accomodate the SimpleAllocator design, currently it assumes the Device memory allocator is BFC. I can upload a PR which I have tested locally works with our design. Please let me know.

@annarev Thanks for the PR!
Do you mean the PluggableDeviceSimpleAllocator will directly call the stream_exec_->AllocateArray(), which calls SP_StreamExecutor::allocate(const SP_Device* device, uint64_t size, int64_t memory_space,SP_DeviceMemoryBase* mem ), and plugin authors decides whether allocate is build on top of a custom allocator? if it is, I think it is better to add an field such as void* custom_allocator in SP_Device? since the allocator should be per-device? Thanks

@jzhoulon Sounds about right. use_bfc_allocator will decide whether to use a simple allocator or a BFC allocator. So, a plugin that wants to use their own custom allocation strategy and doesn't want to wrap with BFC, can set this to true. Do you mean add use_bfc_allocator parameter per device? I was just thinking having one setting for all device types is easier to pass through on TF end. Do we have a usecase for using or not using bfc allocator per device?

Thanks @annarev for the PR. SimpleAllocator can be something straightforward calling into AllocateRaw and that would be fine for us. We had brief discussion of this during the design phase. To keep it simple we can leverage SubAllocator currently. After the pluggable device implementation merges we can use the use_bfc_allocator flag to pass this information in device_process_state..

@kulinseth Yep that makes sense. I misread it back then as custom allocator wrapper itself being something customized by the plugin.

There are few places in pluggable_device_process_state like AllocatorParts will require changes to accomodate the SimpleAllocator design, currently it assumes the Device memory allocator is BFC. I can upload a PR which I have tested locally works with our design. Please let me know.

Sounds good, you can upload the PR and we can comment on it if anything needs changing.

@annarev Thanks for the PR!
Do you mean the PluggableDeviceSimpleAllocator will directly call the stream_exec_->AllocateArray(), which calls SP_StreamExecutor::allocate(const SP_Device* device, uint64_t size, int64_t memory_space,SP_DeviceMemoryBase* mem ), and plugin authors decides whether allocate is build on top of a custom allocator? if it is, I think it is better to add an field such as void* custom_allocator in SP_Device? since the allocator should be per-device? Thanks

@jzhoulon Sounds about right. use_bfc_allocator will decide whether to use a simple allocator or a BFC allocator. So, a plugin that wants to use their own custom allocation strategy and doesn't want to wrap with BFC, can set this to true. Do you mean add use_bfc_allocator parameter per device? I was just thinking having one setting for all device types is easier to pass through on TF end. Do we have a usecase for using or not using bfc allocator per device?

No, I mean an opaque handle (void* custom_allocator) in SP_Device. See the following code example:

void create_device (SP_Platform platform, SE_CreateDeviceParams params, status) {
 params->custom_allocator = new CustomAllocator()
}

void allocate(SP_Device device, unit64_t size, ...) {
  CustomAllocator* allocator = static_cast<CustomAllocator*>(device->custom_allocate);
  allocator->allocate_raw(size, ...)
} 

@annarev Thanks for the PR!
Do you mean the PluggableDeviceSimpleAllocator will directly call the stream_exec_->AllocateArray(), which calls SP_StreamExecutor::allocate(const SP_Device* device, uint64_t size, int64_t memory_space,SP_DeviceMemoryBase* mem ), and plugin authors decides whether allocate is build on top of a custom allocator? if it is, I think it is better to add an field such as void* custom_allocator in SP_Device? since the allocator should be per-device? Thanks

@jzhoulon Sounds about right. use_bfc_allocator will decide whether to use a simple allocator or a BFC allocator. So, a plugin that wants to use their own custom allocation strategy and doesn't want to wrap with BFC, can set this to true. Do you mean add use_bfc_allocator parameter per device? I was just thinking having one setting for all device types is easier to pass through on TF end. Do we have a usecase for using or not using bfc allocator per device?

No, I mean an opaque handle (void* custom_allocator) in SP_Device. See the following code example:

void create_device (SP_Platform platform, SE_CreateDeviceParams params, status) {
 params->custom_allocator = new CustomAllocator()
}

void allocate(SP_Device device, unit64_t size, ...) {
  CustomAllocator* allocator = static_cast<CustomAllocator*>(device->custom_allocate);
  allocator->allocate_raw(size, ...)
} 

I think setting a custom allocator makes things more complicated. I think we should do it only if we think of a usecase that cannot be covered by just implementing allocate/deallocate in the StreamExecutor API: https://cs.opensource.google/tensorflow/tensorflow/+/master:tensorflow/c/experimental/stream_executor/stream_executor.h;l=245.

I think setting a custom allocator makes things more complicated. I think we should do it only if we think of a usecase that cannot be covered by just implementing allocate/deallocate in the StreamExecutor API: https://cs.opensource.google/tensorflow/tensorflow/+/master:tensorflow/c/experimental/stream_executor/stream_executor.h;l=245.

Agreed, I can queue up the PR after #45784 merges.

Agreed, I can queue up the PR after #45784 merges.

That would be awesome. Thank you very much, Kulin!