1# -*- coding: utf-8 -*- 2# This file is part of Xpra. 3# Copyright (C) 2018 Antoine Martin <antoine@xpra.org> 4# Xpra is released under the terms of the GNU GPL v2, or, at your option, any 5# later version. See the file COPYING for details. 6 7from time import monotonic 8import numpy 9from pycuda import driver #@UnresolvedImport 10 11from xpra.codecs.image_wrapper import ImageWrapper 12from xpra.log import Logger 13 14log = Logger("cuda", "nvfbc") 15 16 17class CUDAImageWrapper(ImageWrapper): 18 19 def __init__(self, *args): 20 super().__init__(*args) 21 self.stream = None 22 self.cuda_device_buffer = None 23 self.cuda_context = None 24 self.buffer_size = 0 25 26 def wait_for_stream(self): 27 s = self.stream 28 if s and not s.is_done(): 29 self.stream.synchronize() 30 31 32 def may_download(self): 33 ctx = self.cuda_context 34 if self.pixels is not None or not ctx or self.freed: 35 return 36 assert self.cuda_device_buffer, "bug: no device buffer" 37 start = monotonic() 38 ctx.push() 39 host_buffer = driver.pagelocked_empty(self.buffer_size, dtype=numpy.byte) #pylint: disable=no-member 40 driver.memcpy_dtoh_async(host_buffer, self.cuda_device_buffer, self.stream) #pylint: disable=no-member 41 self.wait_for_stream() 42 self.pixels = host_buffer.tobytes() 43 elapsed = monotonic()-start 44 log("may_download() from %#x to %s, size=%s, elapsed=%ims - %iMB/s", 45 int(self.cuda_device_buffer), host_buffer, self.buffer_size, 46 int(1000*elapsed), self.buffer_size/elapsed/1024/1024) 47 self.free_cuda() 48 ctx.pop() 49 50 def freeze(self): 51 #this image is already a copy when we get it 52 return True 53 54 def get_gpu_buffer(self): 55 self.wait_for_stream() 56 return self.cuda_device_buffer 57 58 def has_pixels(self): 59 return self.pixels is not None 60 61 def get_pixels(self): 62 self.may_download() 63 return super().get_pixels() 64 65 def clone_pixel_data(self): 66 self.may_download() 67 return super().clone_pixel_data() 68 69 def get_sub_image(self, x, y, w, h): 70 self.may_download() 71 return super().get_sub_image(x, y, w, h) 72 73 def free_cuda_device_buffer(self): 74 cdb = self.cuda_device_buffer 75 if not cdb: 76 return 77 log("%s.free_cuda() cuda_device_buffer=%#x", self, int(cdb or 0)) 78 self.cuda_device_buffer = None 79 cdb.free() 80 81 def free_cuda(self): 82 self.free_cuda_device_buffer() 83 self.stream = None 84 self.cuda_context = None 85 self.buffer_size = 0 86 87 def free(self): 88 self.free_cuda() 89 return super().free() 90 91 def clean(self): 92 try: 93 self.wait_for_stream() 94 except driver.LogicError: #pylint: disable=no-member 95 log("%s.clean()", self, exc_info=True) 96 self.free() 97