| 47 | |
| 48 | |
| 49 | class CUBLAS_Context: |
| 50 | _instance = None |
| 51 | |
| 52 | def __init__(self): |
| 53 | raise RuntimeError("Call get_instance() instead") |
| 54 | |
| 55 | def initialize(self): |
| 56 | self.context = {} |
| 57 | |
| 58 | @classmethod |
| 59 | def get_instance(cls): |
| 60 | if cls._instance is None: |
| 61 | cls._instance = cls.__new__(cls) |
| 62 | cls._instance.initialize() |
| 63 | return cls._instance |
| 64 | |
| 65 | def get_context(self, device): |
| 66 | if device.index not in self.context: |
| 67 | prev_device = torch.cuda.current_device() |
| 68 | torch.cuda.set_device(device) |
| 69 | self.context[device.index] = ct.c_void_p(lib.get_context()) |
| 70 | torch.cuda.set_device(prev_device) |
| 71 | return self.context[device.index] |
| 72 | |
| 73 | |
| 74 | FIRST_CUDA_DEVICE = torch.device("cuda", index=0) |
nothing calls this directly
no outgoing calls
no test coverage detected