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