deepspeed_py_aio.h 1.0 KB

12345678910111213141516171819202122232425262728293031
  1. // Copyright (c) Microsoft Corporation.
  2. // SPDX-License-Identifier: Apache-2.0
  3. // DeepSpeed Team
  4. /*
  5. Copyright 2020 The Microsoft DeepSpeed Team
  6. Licensed under the MIT license.
  7. Functionality for swapping optimizer tensors to/from (NVMe) storage devices.
  8. */
  9. #include <deepspeed_aio_common.h>
  10. #include <stdlib.h>
  11. #include <torch/extension.h>
  12. int deepspeed_py_aio_write(const torch::Tensor& buffer,
  13. const char* filename,
  14. const int block_size,
  15. const int queue_depth,
  16. const bool single_submit,
  17. const bool overlap_events,
  18. const bool validate);
  19. int deepspeed_py_aio_read(torch::Tensor& buffer,
  20. const char* filename,
  21. const int block_size,
  22. const int queue_depth,
  23. const bool single_submit,
  24. const bool overlap_events,
  25. const bool validate);