Add gpu_headers_lib for use by op kernels. Expose the device event mgr to
the event manager. Change: 142660218
This commit is contained in:
parent
22cbeb7e0b
commit
d2f25ad5cf
@ -969,6 +969,14 @@ cc_library(
|
||||
],
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "gpu_headers_lib",
|
||||
hdrs = [
|
||||
"common_runtime/gpu/gpu_event_mgr.h",
|
||||
],
|
||||
visibility = ["//visibility:public"],
|
||||
)
|
||||
|
||||
cc_library(
|
||||
name = "cuda",
|
||||
visibility = ["//visibility:public"],
|
||||
|
@ -80,6 +80,14 @@ class BaseGPUDevice : public LocalDevice {
|
||||
void ReinitializeGpuDevice(OpKernelContext* context, PerOpGpuDevice* device,
|
||||
DeviceContext* dc, Allocator* allocator) override;
|
||||
|
||||
// Returns the id of this device within the native driver system; e.g., for
|
||||
// CUDA this is the ordinal of the GPU within the system.
|
||||
int gpu_id() const { return gpu_id_; }
|
||||
|
||||
// The executor that provides control for the device; e.g., for CUDA this
|
||||
// corresponds to the cuda context.
|
||||
gpu::StreamExecutor* executor() const { return executor_; }
|
||||
|
||||
protected:
|
||||
Allocator* gpu_allocator_; // not owned
|
||||
Allocator* cpu_allocator_; // not owned
|
||||
|
Loading…
Reference in New Issue
Block a user