123456789101112131415161718192021222324 |
- '''
- Copyright 2020 The Microsoft DeepSpeed Team
- '''
- import cupy
- from torch.utils.dlpack import to_dlpack
- from torch.utils.dlpack import from_dlpack
- class CupyBackend(object):
- def __init__(self):
- pass
- def torch2cupy(self, tensor):
- return cupy.fromDlpack(to_dlpack(tensor))
- def cupy2torch(self, cupy_tensor):
- return from_dlpack(cupy_tensor.toDlpack())
- def compress_by_chunk(self, cupy_bool_tensor, num_chunks):
- packed_sign = cupy.packbits(cupy_bool_tensor)
- sign_list_packed = cupy.split(packed_sign, num_chunks)
- cupy.cuda.get_current_stream().synchronize()
- return sign_list_packed
|