MCPcopy Create free account
hub / github.com/deepspeedai/DeepSpeed / _alloc_bounce_buffer

Method _alloc_bounce_buffer

csrc/aio/py_lib/deepspeed_cpu_op.cpp:85–99  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

83}
84
85void cpu_op_desc_t::_alloc_bounce_buffer()
86{
87 auto options = torch::TensorOptions()
88 .dtype(_buffer.dtype())
89 .layout(_buffer.layout())
90 .device(torch::kCPU)
91 .requires_grad(false);
92
93#if defined(__CUDA_ARCH__)
94 _cpu_buffer = torch::empty(_buffer.numel(), options).pin_memory();
95#else
96 _is_managed_bounce_buffer = true;
97 _cpu_buffer = _pinned_tensor_mgr->alloc(_buffer.numel(), options);
98#endif
99}
100
101void cpu_op_desc_t::_free_bounce_buffer()
102{

Callers

nothing calls this directly

Calls 5

numelMethod · 0.80
allocMethod · 0.80
deviceMethod · 0.45
dtypeMethod · 0.45
pin_memoryMethod · 0.45

Tested by

no test coverage detected