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 

7import numpy 

8from pycuda import driver #@UnresolvedImport 

9 

10from xpra.os_util import monotonic_time 

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_time() 

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_time()-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 ImageWrapper.get_pixels(self) 

64 

65 def clone_pixel_data(self): 

66 self.may_download() 

67 return ImageWrapper.clone_pixel_data(self) 

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 ImageWrapper.free(self) 

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()