deepspeed_py_aio.h 954 B

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