1# -*- coding: utf-8 -*-
2# This file is part of Xpra.
3# Copyright (C) 2017-2021 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
7#cython: wraparound=False
8
9import os
10import sys
11from time import monotonic
12
13from xpra.os_util import WIN32
14from xpra.util import csv, roundup
15from xpra.codecs.codec_constants import TransientCodecException, CodecStateException
16from xpra.codecs.image_wrapper import ImageWrapper
17from xpra.codecs.nvfbc.cuda_image_wrapper import CUDAImageWrapper
18from xpra.codecs.nv_util import get_nvidia_module_version, get_cards, get_license_keys, parse_nvfbc_hex_key
19
20from xpra.log import Logger
21log = Logger("encoder", "nvfbc")
22
23try:
24    from pycuda import driver
25    from xpra.codecs.cuda_common.cuda_context import CUDA_ERRORS_INFO, select_device, device_info
26except ImportError:
27    raise
28except Exception as e:
29    log.error("Error: NvFBC requires CUDA", exc_info=True)
30    raise ImportError("NvFBC requires CUDA: %s" % e) from None
31
32import ctypes
33from ctypes import wintypes
34
35from libc.stdint cimport uintptr_t, uint8_t, int64_t  #pylint: disable=syntax-error
36from libc.string cimport memset, memcpy
37from xpra.buffers.membuf cimport padbuf, MemBuf
38
39DEFAULT_PIXEL_FORMAT = os.environ.get("XPRA_NVFBC_DEFAULT_PIXEL_FORMAT", "RGB")
40CLIENT_KEYS_STRS = get_license_keys(basefilename="nvfbc")
41
42
43ctypedef unsigned long DWORD
44ctypedef int BOOL
45
46
47cdef extern from "NvFBC/nvFBC.h":
48    ctypedef int NVFBCRESULT
49    ctypedef unsigned long NvU32
50
51    int NVFBC_DLL_VERSION
52    int NVFBC_GLOBAL_FLAGS_NONE
53    #deprecated:
54    #int NVFBC_GLOBAL_FLAGS_STEREO_BUFFER
55    int NVFBC_GLOBAL_FLAGS_NO_INITIAL_REFRESH
56    int NVFBC_GLOBAL_FLAGS_NO_DEVICE_RESET_TOGGLE
57
58    int NVFBC_CREATE_PARAMS_VER
59    int NVFBC_STATUS_VER
60    int NVFBC_CURSOR_CAPTURE_PARAMS_VER
61
62    NVFBCRESULT NVFBC_SUCCESS
63    NVFBCRESULT NVFBC_ERROR_GENERIC                     # Unexpected failure in NVFBC.
64    NVFBCRESULT NVFBC_ERROR_INVALID_PARAM               # One or more of the paramteres passed to NvFBC are invalid [This include NULL pointers].
65    NVFBCRESULT NVFBC_ERROR_INVALIDATED_SESSION         # NvFBC session is invalid. Client needs to recreate session.
66    NVFBCRESULT NVFBC_ERROR_PROTECTED_CONTENT           # Protected content detected. Capture failed.
67    NVFBCRESULT NVFBC_ERROR_DRIVER_FAILURE              # GPU driver returned failure to process NvFBC command.
68    NVFBCRESULT NVFBC_ERROR_CUDA_FAILURE                # CUDA driver returned failure to process NvFBC command.
69    NVFBCRESULT NVFBC_ERROR_UNSUPPORTED                 # API Unsupported on this version of NvFBC.
70    NVFBCRESULT NVFBC_ERROR_HW_ENC_FAILURE              # HW Encoder returned failure to process NVFBC command.
71    NVFBCRESULT NVFBC_ERROR_INCOMPATIBLE_DRIVER         # NVFBC is not compatible with this version of the GPU driver.
72    NVFBCRESULT NVFBC_ERROR_UNSUPPORTED_PLATFORM        # NVFBC is not supported on this platform.
73    NVFBCRESULT NVFBC_ERROR_OUT_OF_MEMORY               # Failed to allocate memory.
74    NVFBCRESULT NVFBC_ERROR_INVALID_PTR                 # A NULL pointer was passed.
75    NVFBCRESULT NVFBC_ERROR_INCOMPATIBLE_VERSION        # An API was called with a parameter struct that has an incompatible version. Check dwVersion field of paramter struct.
76    NVFBCRESULT NVFBC_ERROR_OPT_CAPTURE_FAILURE         # Desktop Capture failed.
77    NVFBCRESULT NVFBC_ERROR_INSUFFICIENT_PRIVILEGES     # User doesn't have appropriate previlages.
78    NVFBCRESULT NVFBC_ERROR_INVALID_CALL                # NVFBC APIs called in wrong sequence.
79    NVFBCRESULT NVFBC_ERROR_SYSTEM_ERROR                # Win32 error.
80    NVFBCRESULT NVFBC_ERROR_INVALID_TARGET              # The target adapter idx can not be used for NVFBC capture. It may not correspond to an NVIDIA GPU, or may not be attached to desktop.
81    NVFBCRESULT NVFBC_ERROR_NVAPI_FAILURE               # NvAPI Error
82    NVFBCRESULT NVFBC_ERROR_DYNAMIC_DISABLE             # NvFBC is dynamically disabled. Cannot continue to capture
83    NVFBCRESULT NVFBC_ERROR_IPC_FAILURE                 # NVFBC encountered an error in state management
84    NVFBCRESULT NVFBC_ERROR_CURSOR_CAPTURE_FAILURE      # Hardware cursor capture failed
85
86
87    ctypedef int NVFBC_STATE
88    NVFBC_STATE NVFBC_STATE_DISABLE
89    NVFBC_STATE NVFBC_STATE_ENABLE
90
91    # Defines parameters that describe the grabbed data,
92    # and provides detailed information about status of the NVFBC session.
93    ctypedef struct NvFBCFrameGrabInfo:
94        DWORD   dwWidth                 #[out] Indicates the current width of captured buffer.
95        DWORD   dwHeight                #[out] Indicates the current height of captured buffer.
96        DWORD   dwBufferWidth           #[out] Indicates the current width of the pixel buffer(padded width).
97        DWORD   dwReserved              #[out] Reserved, do not use.
98        BOOL    bOverlayActive          #[out] Is set to 1 if overlay was active.
99        BOOL    bMustRecreate           #[out] Is set to 1 if the compressor must call NvBFC_Create again.
100        BOOL    bFirstBuffer            #[out] Is set to 1 is this was the first capture call, or first call after a desktop mode change.
101                                        # Relevant only for XOR and diff modes supported by NVFBCToSys interface.
102        BOOL    bHWMouseVisible         #[out] Is set to 1 if HW cursor was enabled by OS at the time of the grab.
103        BOOL    bProtectedContent       #[out] Is set to 1 if protected content was active (DXVA encryption Session).
104        DWORD   dwDriverInternalError   #[out] To be used as diagnostic info if Grab() fails. Status is non-fatal if Grab() returns success.
105                                        # Indicates the status code from lower layers. 0 or 0xFBCA11F9 indicates no error was returned.
106        BOOL    bStereoOn               #[out] Is set to 1 if stereo was on.
107        BOOL    bIGPUCapture            #[out] Is set to 1 if the captured frame is from iGPU. 0 if capture fails or if captured from dGPU*/
108        DWORD   dwSourcePID             #[out] Indicates which process caused the last screen update that got grabbed*/
109        DWORD   dwReserved3             #[out] Reserved, do not use.
110        DWORD   bIsHDR                  #[out] Is set to 1 if grabbed content is in HDR format.
111        #DWORD   bReservedBit1           #[out] Reserved, do not use.
112        #DWORD   bReservedBits           #[out] Reserved, do not use.
113        DWORD   dwWaitModeUsed          #[out] The mode used for this Grab operation (blocking or non-blocking), based on the grab flags passed by the application.
114                                        # Actual blocking mode can differ from application's request if incorrect grab flags are passed.
115        #NvU32   dwReserved2[11]         #[out] Resereved, should be set to 0.
116
117    # Defines the parameters to be used with NvFBC_GetStatusEx API
118    ctypedef struct NvFBCStatusEx:
119        NvU32  dwVersion                #[in]  Struct version. Set to NVFBC_STATUS_VER.
120        NvU32  bIsCapturePossible       #[out] Indicates if NvFBC feature is enabled.
121        NvU32  bCurrentlyCapturing      #[out] Indicates if NVFBC is currently capturing for the Adapter ordinal specified in dwAdapterIdx.
122        NvU32  bCanCreateNow            #[out] Deprecated. Do not use.
123        NvU32  bSupportMultiHead        #[out] MultiHead grab supported.
124        NvU32  bSupportConfigurableDiffMap     #[out] Difference map with configurable blocksize supported. Supported sizes 16x16, 32x32, 64x64, 128x128(default)
125        NvU32  bSupportImageClassification     #[out] Generation of 'classification map' demarkating high frequency content in the captured image is supported
126        #NvU32  bReservedBits            #[in]  Reserved, do not use.
127        NvU32  dwNvFBCVersion           #[out] Indicates the highest NvFBC interface version supported by the loaded NVFBC library.
128        NvU32  dwAdapterIdx             #[in]  Adapter Ordinal corresponding to the display to be grabbed. IGNORED if bCapturePID is set
129        void*  pPrivateData             #[in]  optional **/
130        NvU32  dwPrivateDataSize        #[in]  optional **/
131        NvU32  dwReserved[59]           #[in]  Reserved. Should be set to 0.
132        void*  pReserved[31]            #[in]  Reserved. Should be set to NULL.
133
134    # Defines the parameters to be used with NvFBC_CreateEx API
135    ctypedef struct NvFBCCreateParams:
136        NvU32  dwVersion                #[in]  Struct version. Set to NVFBC_CREATE_PARAMS_VER.
137        NvU32  dwInterfaceType          #[in]  ID of the NVFBC interface Type being requested.
138        NvU32  dwMaxDisplayWidth        #[out] Max. display width allowed.
139        NvU32  dwMaxDisplayHeight       #[out] Max. display height allowed.
140        void*  pDevice                  #[in]  Device pointer.
141        void*  pPrivateData             #[in]  Private data [optional].
142        NvU32  dwPrivateDataSize        #[in]  Size of private data.
143        NvU32  dwInterfaceVersion       #[in]  Version of the capture interface.
144        void*  pNvFBC                   #[out] A pointer to the requested NVFBC object.
145        NvU32  dwAdapterIdx             #[in]  Adapter Ordinal corresponding to the display to be grabbed. If pDevice is set, this parameter is ignored.
146        NvU32  dwNvFBCVersion           #[out] Indicates the highest NvFBC interface version supported by the loaded NVFBC library.
147        void*  cudaCtx                  #[in]  CUDA context created using cuD3D9CtxCreate with the D3D9 device passed as pDevice. Only used for NvFBCCuda interface.
148                                        # It is mandatory to pass a valid D3D9 device if cudaCtx is passed. The call will fail otherwise.
149                                        # Client must release NvFBCCuda object before destroying the cudaCtx.
150        void*  pPrivateData2            #[in]  Private data [optional].
151        NvU32  dwPrivateData2Size       #[in]  Size of private data.
152        #NvU32  dwReserved[55]           #[in]  Reserved. Should be set to 0.
153        #void*  pReserved[27]            #[in]  Reserved. Should be set to NULL.
154
155    # Defines parameters for a Grab\Capture call to get HW cursor data in the NVFBCToSys capture session
156    ctypedef struct NVFBC_CURSOR_CAPTURE_PARAMS:
157        NvU32 dwVersion                 #[in]:  Struct version. Set to NVFBC_MOUSE_GRAB_INFO_VER
158        NvU32 dwWidth                   #[out]: Width of mouse glyph captured
159        NvU32 dwHeight                  #[out]: Height of mouse glyph captured
160        NvU32 dwPitch                   #[out]: Pitch of mouse glyph captured
161        NvU32 bIsHwCursor               #[out]: Tells if cursor is HW cursor or SW cursor. If set to 0, ignore height, width, pitch and pBits
162        #NvU32 bReserved : 32           #[in]:  Reserved
163        NvU32 dwPointerFlags            #[out]: Maps to DXGK_POINTERFLAGS::Value
164        NvU32 dwXHotSpot                #[out]: Maps to DXGKARG_SETPOINTERSHAPE::XHot
165        NvU32 dwYHotSpot                #[out]: Maps to DXGKARG_SETPOINTERSHAPE::YHot
166        NvU32 dwUpdateCounter           #[out]: Cursor update Counter.
167        NvU32 dwBufferSize              #[out]: Size of the buffer contaiing the captured cursor glyph.
168        void * pBits                    #[out]: pointer to buffer containing the captured cursor glyph
169        NvU32 dwReservedA[22]           #[in]:  Reserved. Set to 0
170        void * pReserved[15]            #[in]:  Reserved. Set to 0
171
172    # NVFBC API to set global overrides
173    # param [in] dwFlags Global overrides for NVFBC. Use ::NVFBC_GLOBAL_FLAGS value.
174    void NvFBC_SetGlobalFlags(DWORD dwFlags)
175
176    # NVFBC API to create an NVFBC capture session.
177    # Instantiates an interface identified by NvFBCCreateParams::dwInterfaceType.
178    # param [inout] pCreateParams Pointer to a struct of type ::NvFBCCreateParams, typecast to void*
179    # return An applicable ::NVFBCRESULT value.
180    NVFBCRESULT NvFBC_CreateEx(void * pCreateParams)
181
182    # NVFBC API to query Current NVFBC status.
183    # Queries the status for the adapter pointed to by the NvFBCStatusEx::dwAdapterIdx parameter.
184    # [inout] pCreateParams Pointer to a struct of type ::NvFBCStatusEx.
185    # return An applicable ::NVFBCRESULT value.
186    NVFBCRESULT NvFBC_GetStatusEx(NvFBCStatusEx *pNvFBCStatusEx)
187
188    # NVFBC API to enable \ disable NVFBC feature.
189    # param [in] nvFBCState Refer ::NVFBC_STATE
190    # return An applicable ::NVFBCRESULT value.
191    NVFBCRESULT NvFBC_Enable(NVFBC_STATE nvFBCState)
192
193    # NVFBC API to query highest GRID SDK version supported by the loaded NVFBC library.
194    # param [out] pVersion Pointer to a 32-bit integer to hold the supported GRID SDK version.
195    # return An applicable ::NVFBCRESULT value.
196    NVFBCRESULT NvFBC_GetSDKVersion(NvU32 * pVersion)
197
198
199cdef extern from "NvFBC/nvFBCToSys.h":
200    int NVFBC_TO_SYS
201    int NVFBC_SHARED_CUDA
202    int NVFBC_TOSYS_SETUP_PARAMS_VER
203    int NVFBC_TOSYS_GRAB_FRAME_PARAMS_VER
204
205    ctypedef int NVFBCToSysBufferFormat
206    NVFBCToSysBufferFormat NVFBC_TOSYS_ARGB         # 32bpp, one byte per channel.
207    NVFBCToSysBufferFormat NVFBC_TOSYS_RGB          # 24bpp, one byte per channel.
208    NVFBCToSysBufferFormat NVFBC_TOSYS_YYYYUV420p   # 12bpp, the Y' channel at full resolution, U channel at half resolution (1 byte for four pixels), V channel at half resolution.
209    NVFBCToSysBufferFormat NVFBC_TOSYS_RGB_PLANAR   # 24bpp, stored sequentially in memory as complete red channel, complete green channel, complete blue channel.
210    NVFBCToSysBufferFormat NVFBC_TOSYS_XOR          # RGB format: 24bpp XOR�d with the prior frame.
211    NVFBCToSysBufferFormat NVFBC_TOSYS_YUV444p      # Output Pixels in YUV444 planar format, i.e. separate 8-bpp Y, U, V planes with no subsampling.
212    NVFBCToSysBufferFormat NVFBC_TOSYS_ARGB10       # RGB 10 bit format: A2B10G10R10, 32bpp.
213
214    ctypedef int  NVFBCToSysGrabMode
215    NVFBCToSysGrabMode NVFBC_TOSYS_SOURCEMODE_FULL  # Grab full res
216    NVFBCToSysGrabMode NVFBC_TOSYS_SOURCEMODE_SCALE # Will convert current res to supplied resolution (dwTargetWidth and dwTargetHeight)
217    NVFBCToSysGrabMode NVFBC_TOSYS_SOURCEMODE_CROP  # Native res, crops a subwindow, of dwTargetWidth and dwTargetHeight sizes, starting at dwStartX and dwStartY
218
219    ctypedef int NVFBC_TOSYS_GRAB_FLAGS
220    NVFBC_TOSYS_GRAB_FLAGS NVFBC_TOSYS_NOFLAGS      # Default (no flags set). Grabbing will wait for a new frame or HW mouse move.
221    NVFBC_TOSYS_GRAB_FLAGS NVFBC_TOSYS_NOWAIT       # Grabbing will not wait for a new frame nor a HW cursor move.
222    NVFBC_TOSYS_GRAB_FLAGS NVFBC_TOSYS_WAIT_WITH_TIMEOUT # Grabbing will wait for a new frame or HW mouse move with a maximum wait time of NVFBC_TOSYS_GRAB_FRAME_PARAMS::dwWaitTime millisecond
223
224    ctypedef struct NVFBC_TOSYS_SETUP_PARAMS_V2:
225        NvU32 dwVersion                             #[in]: Struct version. Set to NVFBC_TOSYS_SETUP_PARAMS_VER
226        NvU32 bWithHWCursor                         #[in]: The client should set this to 1 if it requires the HW cursor to be composited on the captured image
227        NvU32 bDiffMap                              #[in]: The client should set this to use the DiffMap feature
228        NvU32 bEnableSeparateCursorCapture          #[in]: The client should set this to 1 if it wants to enable mouse capture in separate stream
229        NvU32 bHDRRequest                           #[in]: The client should set this to 1 to request HDR capture
230        NvU32 b16x16DiffMap                         #[in]: Valid only if bDiffMap is set. The client should set this to 1 it it wants to request 16x16 Diffmap, set it to 0 if it wants 128x128 Diffmap
231        #NvU32 bReservedBits :27                     #[in]: Reserved. Set to 0
232        NVFBCToSysBufferFormat eMode                #[in]: Output image format
233        #NvU32 dwReserved1                           #[in]: Reserved. Set to 0
234        void **ppBuffer                             #[out]: Container to hold NvFBC output buffers
235        void **ppDiffMap                            #[out]: Container to hold NvFBC output diffmap buffers
236        void  *hCursorCaptureEvent                  #[out]: Client should wait for mouseEventHandle event before calling MouseGrab function. */
237        #NvU32 dwReserved[58]                        #[in]: Reserved. Set to 0
238        #void *pReserved[29]                         #[in]: Reserved. Set to 0
239    ctypedef NVFBC_TOSYS_SETUP_PARAMS_V2 NVFBC_TOSYS_SETUP_PARAMS
240
241    ctypedef struct NVFBC_TOSYS_GRAB_FRAME_PARAMS_V1:
242        NvU32 dwVersion                             #[in]: Struct version. Set to NVFBC_TOSYS_GRAB_FRAME_PARAMS_VER.
243        NvU32 dwFlags                               #[in]: Special grabbing requests. This should be a bit-mask of NVFBC_TOSYS_GRAB_FLAGS values.
244        NvU32 dwTargetWidth                         #[in]: Target image width. NvFBC will scale the captured image to fit taret width and height. Used with NVFBC_TOSYS_SOURCEMODE_SCALE and NVFBC_TOSYS_SOURCEMODE_CROP.
245        NvU32 dwTargetHeight                        #[in]: Target image height. NvFBC will scale the captured image to fit taret width and height. Used with NVFBC_TOSYS_SOURCEMODE_SCALE and NVFBC_TOSYS_SOURCEMODE_CROP.
246        NvU32 dwStartX                              #[in]: x-coordinate of starting pixel for cropping. Used with NVFBC_TOSYS_SOURCEMODE_CROP.
247        NvU32 dwStartY                              #[in]: y-coordinate of starting pixel for cropping. Used with NVFBC_TOSYS_SOURCEMODE_CROP.
248        NVFBCToSysGrabMode eGMode                   #[in]: Frame grab mode.
249        NvU32 dwWaitTime                            #[in]: Time limit for NvFBCToSysGrabFrame() to wait until a new frame is available or a HW mouse moves. Use with NVFBC_TOSYS_WAIT_WITH_TIMEOUT
250        NvFBCFrameGrabInfo *pNvFBCFrameGrabInfo     #[in/out]: Frame grab information and feedback from NvFBC driver.
251        NvU32 dwReserved[56]                        #[in]: Reserved. Set to 0.
252        void *pReserved[31]                         #[in]: Reserved. Set to NULL.
253    ctypedef NVFBC_TOSYS_GRAB_FRAME_PARAMS_V1 NVFBC_TOSYS_GRAB_FRAME_PARAMS
254
255    # Sets up NVFBC System Memory capture according to the provided parameters.
256    # [in] pParam Pointer to a struct of type ::NVFBC_TOSYS_SETUP_PARAMS.
257    ctypedef NVFBCRESULT (*NVFBCTOSYSSETUP) (NVFBC_TOSYS_SETUP_PARAMS *pParam) nogil
258    # Captures the desktop and dumps the captured data to a System memory buffer.
259    # If the API returns a failure, the client should check the return codes and
260    # ::NvFBCFrameGrabInfo output fields to determine if the session needs to be re-created.
261    # [inout] pParam Pointer to a struct of type ::NVFBC_TOSYS_GRAB_FRAME_PARAMS.
262    ctypedef NVFBCRESULT (*NVFBCTOSYSGRABFRAME) (NVFBC_TOSYS_GRAB_FRAME_PARAMS *pParam) nogil
263    # Captures HW cursor data whenever shape of mouse is changed
264    # [inout] pParam Pointer to a struct of type ::NVFBC_CURSOR_CAPTURE_PARAMS
265    ctypedef NVFBCRESULT (*NVFBCTOSYSCURSORCAPTURE) (NVFBC_CURSOR_CAPTURE_PARAMS *pParam) nogil
266    # A high precision implementation of Sleep().
267    # Can provide sub quantum (usually 16ms) sleep that does not burn CPU cycles.
268    # [in] qwMicroSeconds The number of microseconds that the thread should sleep for.
269    ctypedef NVFBCRESULT (*NVFBCTOSYSGPUBASEDCPUSLEEP) (int64_t qwMicroSeconds) nogil
270    # Destroys the NVFBCToSys capture session.
271    ctypedef NVFBCRESULT (*NVFBCTOSYSRELEASE) () nogil
272
273    ctypedef struct NvFBCToSys:
274        NVFBCTOSYSSETUP NvFBCToSysSetUp
275        NVFBCTOSYSGRABFRAME NvFBCToSysGrabFrame
276        NVFBCTOSYSCURSORCAPTURE NvFBCToSysCursorCapture
277        NVFBCTOSYSGPUBASEDCPUSLEEP NvFBCToSysGPUBasedCPUSleep
278        NVFBCTOSYSRELEASE NvFBCToSysRelease
279
280cdef extern from "NvFBC/nvFBCCuda.h":
281    int NVFBC_TOCUDA_NOFLAGS            # Default (no flags set). Grabbing will wait for a new frame or HW mouse move
282    int NVFBC_TOCUDA_NOWAIT             # Grabbing will not wait for a new frame nor a HW cursor move.
283    int NVFBC_TOCUDA_CPU_SYNC           # Does a cpu event signal when grab is complete
284    int NVFBC_TOCUDA_WITH_HWCURSOR      # Grabs the HW cursor if any visible
285    int NVFBC_TOCUDA_RESERVED_A         # reserved
286    int NVFBC_TOCUDA_WAIT_WITH_TIMEOUT  # Grabbing will wait for a new frame or HW mouse move with a maximum wait time of NVFBC_CUDA_GRAB_FRAME_PARAMS::dwWaitTime millisecond
287
288    ctypedef int NVFBCToCUDABufferFormat
289    NVFBCToCUDABufferFormat NVFBC_TOCUDA_ARGB       # Output in 32-bit packed ARGB format
290    NVFBCToCUDABufferFormat NVFBC_TOCUDA_ARGB10     # Output in 32-bit packed ARGB10 format (A2B10G10R10)
291
292    ctypedef struct NVFBC_CUDA_SETUP_PARAMS_V1:
293        NvU32 dwVersion                     # [in]: Struct version. Set to NVFBC_CUDA_SETUP_PARMS_VER
294        NvU32 bEnableSeparateCursorCapture  # [in]: The client should set this to 1 if it wants to enable mouse capture separately from Grab()
295        NvU32 bHDRRequest                   # [in]: The client should set this to 1 if it wants to request HDR capture
296        #NvU32 bReserved                     # [in]: Reserved. Seto to 0
297        void *hCursorCaptureEvent           # [out]: Event handle to be signalled when there is an update to the HW cursor state.
298        NVFBCToCUDABufferFormat eFormat     # [in]: Output image format
299        #NvU32 dwReserved[61]                # [in]: Reserved. Set to 0
300        #void *pReserved[31]                 # [in]: Reserved. Set to NULL
301    int NVFBC_CUDA_SETUP_PARAMS_V1_VER
302    ctypedef NVFBC_CUDA_SETUP_PARAMS_V1 NVFBC_CUDA_SETUP_PARAMS
303
304    ctypedef struct NVFBC_CUDA_GRAB_FRAME_PARAMS_V1:
305        NvU32 dwVersion                     # [in]: Struct version. Set to NVFBC_CUDA_GRAB_FRAME_PARAMS_V1_VER
306        NvU32 dwFlags                       # [in]: Flags for grab frame
307        void *pCUDADeviceBuffer             # [in]: Output buffer
308        NvFBCFrameGrabInfo *pNvFBCFrameGrabInfo # [in/out]: Frame grab configuration and feedback from NvFBC driver
309        NvU32 dwWaitTime                    # [in] Time limit in millisecond to wait for a new frame or HW mouse move. Use with NVFBC_TOCUDA_WAIT_WITH_TIMEOUT
310        #NvU32 dwReserved[61]                # [in]: Reserved. Set to 0
311        #void *pReserved[30]                 # [in]: Reserved. Set to NULL
312    int NVFBC_CUDA_GRAB_FRAME_PARAMS_V1_VER
313    ctypedef NVFBC_CUDA_GRAB_FRAME_PARAMS_V1 NVFBC_CUDA_GRAB_FRAME_PARAMS
314
315    # Returns the maximum buffer size, in bytes for allocating a CUDA buffer to hold output data generated by the NvFBCCuda interface
316    # [out] pdwMaxBufSize Pointer to a 32-bit unsigned integer
317    ctypedef NVFBCRESULT (*NVFBCCUDAGETMAXBUFFERSIZE) (NvU32 *pdwMaxBufSize) nogil
318    #Performs initial setup
319    # [in] pParams Pointer to a struct of type ::NVFBC_CUDA_SETUP_PARAMS
320    ctypedef NVFBCRESULT (*NVFBCCUDASETUP) (NVFBC_CUDA_SETUP_PARAMS *pParams) nogil
321    # Captures the desktop and dumps captured data to a CUDA buffer provided by the client
322    # If the API returns a failure, the client should check the return codes and ::NvFBCFrameGrabInfo output fields to determine if the session needs to be re-created
323    # [inout] pParams Pointer to a struct of type ::NVFBC_CUDA_GRAB_FRAME_PARAMS
324    ctypedef NVFBCRESULT (*NVFBCCUDAGRABFRAME) (NVFBC_CUDA_GRAB_FRAME_PARAMS *pParams) nogil
325    # A high precision implementation of Sleep()
326    # Can provide sub quantum (usually 16ms) sleep that does not burn CPU cycles
327    # [in] qwMicroSeconds The number of microseconds that the thread should sleep for.
328    ctypedef NVFBCRESULT (*NVFBCCUDAGPUBASEDCPUSLEEP) (int64_t qwMicroSeconds) nogil
329    # Captures HW cursor data whenever shape of mouse is changed
330    # [inout] pParam Pointer to a struct of type ::NVFBC_TOSYS_GRAB_MOUSE_PARAMS
331    ctypedef NVFBCRESULT (*NVFBCCUDACURSORCAPTURE) (NVFBC_CURSOR_CAPTURE_PARAMS *pParam) nogil
332    # Destroys the NvFBCCuda capture session.
333    ctypedef NVFBCRESULT (*NVFBCCUDARELEASE) ()
334
335
336    ctypedef struct NvFBCCuda:
337        NVFBCCUDAGETMAXBUFFERSIZE NvFBCCudaGetMaxBufferSize
338        NVFBCCUDASETUP NvFBCCudaSetup
339        NVFBCCUDAGRABFRAME NvFBCCudaGrabFrame
340        NVFBCCUDAGPUBASEDCPUSLEEP NvFBCCudaGPUBasedCPUSleep
341        NVFBCCUDACURSORCAPTURE NvFBCCudaCursorCapture
342        NVFBCCUDARELEASE NvFBCCudaRelease
343
344
345ERRORS = {
346    NVFBC_SUCCESS                       : "SUCCESS",
347    NVFBC_ERROR_GENERIC                 : "GENERIC",
348    NVFBC_ERROR_INVALID_PARAM           : "INVALID_PARAM",
349    NVFBC_ERROR_INVALIDATED_SESSION     : "INVALIDATED_SESSION",
350    NVFBC_ERROR_PROTECTED_CONTENT       : "PROTECTED_CONTENT",
351    NVFBC_ERROR_DRIVER_FAILURE          : "DRIVER_FAILURE",
352    NVFBC_ERROR_CUDA_FAILURE            : "CUDA_FAILURE",
353    NVFBC_ERROR_UNSUPPORTED             : "UNSUPPORTED",
354    NVFBC_ERROR_HW_ENC_FAILURE          : "HW_ENC_FAILURE",
355    NVFBC_ERROR_INCOMPATIBLE_DRIVER     : "INCOMPATIBLE_DRIVER",
356    NVFBC_ERROR_UNSUPPORTED_PLATFORM    : "UNSUPPORTED_PLATFORM",
357    NVFBC_ERROR_OUT_OF_MEMORY           : "OUT_OF_MEMORY",
358    NVFBC_ERROR_INVALID_PTR             : "INVALID_PTR",
359    NVFBC_ERROR_INCOMPATIBLE_VERSION    : "INCOMPATIBLE_VERSION",
360    NVFBC_ERROR_OPT_CAPTURE_FAILURE     : "OPT_CAPTURE_FAILURE",
361    NVFBC_ERROR_INSUFFICIENT_PRIVILEGES : "INSUFFICIENT_PRIVILEGES",
362    NVFBC_ERROR_INVALID_CALL            : "INVALID_CALL",
363    NVFBC_ERROR_SYSTEM_ERROR            : "SYSTEM_ERROR",
364    NVFBC_ERROR_INVALID_TARGET          : "INVALID_TARGET",
365    NVFBC_ERROR_NVAPI_FAILURE           : "NVAPI_FAILURE",
366    NVFBC_ERROR_DYNAMIC_DISABLE         : "DYNAMIC_DISABLE",
367    NVFBC_ERROR_IPC_FAILURE             : "IPC_FAILURE",
368    NVFBC_ERROR_CURSOR_CAPTURE_FAILURE  : "CURSOR_CAPTURE_FAILURE",
369    }
370
371
372cdef inline cvp(val):
373    return ctypes.cast(<uintptr_t> val, ctypes.c_void_p)
374
375
376class NvFBCException(Exception):
377    def __init__(self, code, fn):
378        self.function = fn
379        self.code = code
380        msg = "%s - returned %s" % (fn, ERRORS.get(code, code))
381        super().__init__(msg)
382
383cdef inline raiseNvFBC(NVFBCRESULT ret, msg):
384    if ret!=0:
385        raise NvFBCException(ret, msg)
386
387
388NvFBC = None
389def init_nvfbc_library():
390    global NvFBC
391    if NvFBC is not None:
392        return NvFBC
393    if not WIN32:
394        NvFBC = False
395        raise Exception("nvfbc is not supported on %s" % sys.platform)
396    load = ctypes.WinDLL
397    #we only support 64-bit:
398    nvfbc_libname = "NvFBC64.dll"
399    log("init_nvfbc_library() will try to load %s", nvfbc_libname)
400    try:
401        NvFBC = load(nvfbc_libname)
402        log("init_nvfbc_library() %s(%s)=%s", load, nvfbc_libname, NvFBC)
403    except Exception as e:
404        NvFBC = False
405        log("failed to load '%s'", nvfbc_libname, exc_info=True)
406        raise ImportError("nvfbc: the required library %s cannot be loaded: %s" % (nvfbc_libname, e)) from None
407    NvFBC.NvFBC_GetSDKVersion.argtypes = [ctypes.c_void_p]
408    NvFBC.NvFBC_GetSDKVersion.restype = wintypes.INT
409    NvFBC.NvFBC_GetStatusEx.argtypes = [ctypes.c_void_p]
410    NvFBC.NvFBC_GetStatusEx.restype = wintypes.INT
411    NvFBC.NvFBC_SetGlobalFlags.argtypes = [wintypes.DWORD]
412    NvFBC.NvFBC_SetGlobalFlags.restype = wintypes.INT
413    NvFBC.NvFBC_Enable.argtypes = [wintypes.INT]
414    NvFBC.NvFBC_Enable.restype = wintypes.INT
415    return NvFBC
416
417def unload_library():
418    global NvFBC
419    NvFBC = None
420
421
422def get_status(int adapter=0):
423    global NvFBC
424    assert NvFBC
425    cdef NvFBCStatusEx status
426    memset(&status, 0, sizeof(NvFBCStatusEx))
427    status.dwVersion = NVFBC_STATUS_VER
428    status.dwAdapterIdx = adapter
429    cdef NVFBCRESULT res = NvFBC.NvFBC_GetStatusEx(cvp(<uintptr_t> &status))
430    log("NvFBC_GetStatusEx()=%i", res)
431    raiseNvFBC(res, "NvFBC_GetStatusEx")
432    s = {
433        "capture-possible"      : bool(status.bIsCapturePossible),
434        "currently-capturing"   : bool(status.bCurrentlyCapturing),
435        "can-create-now"        : bool(status.bCanCreateNow),
436        "support-multihead"     : bool(status.bSupportMultiHead),
437        "support-diffmap"       : bool(status.bSupportConfigurableDiffMap),
438        "version"               : int(status.dwNvFBCVersion),
439        "adapter"               : int(status.dwAdapterIdx),
440        }
441    log("get_status()=%s", s)
442    return s
443
444def check_status():
445    status = get_status()
446    if not status.get("capture-possible"):
447        raise Exception("NvFBC status error: capture is not possible")
448    if status.get("currently-capturing"):
449        raise TransientCodecException("NvFBC status error: currently capturing")
450    if not status.get("can-create-now"):
451        raise TransientCodecException("NvFBC status error: cannot create now")
452
453def set_global_flags(DWORD flags):
454    global NvFBC
455    assert NvFBC
456    cdef NVFBCRESULT res = NvFBC.NvFBC_SetGlobalFlags(flags)
457    log("NvFBC_SetGlobalFlags(%i)=%i", flags, res)
458    raiseNvFBC(res, "NvFBC_SetGlobalFlags")
459
460def create_context(int width=-1, int height=-1, interface_type=NVFBC_TO_SYS):
461    log("create_context(%i, %i, %s)", width, height, {NVFBC_TO_SYS : "SYS", NVFBC_SHARED_CUDA : "CUDA"}.get(interface_type))
462    check_status()
463    cdef NvFBCCreateParams create
464    cdef NVFBCRESULT res = <NVFBCRESULT> 0
465    cdef char* ckey
466    keys = CLIENT_KEYS_STRS or [None]
467    log("create_context() will try with keys: %s", csv(keys))
468    assert len(keys)>0
469    for key in keys:
470        memset(&create, 0, sizeof(NvFBCCreateParams))
471        create.dwVersion = NVFBC_CREATE_PARAMS_VER
472        create.dwInterfaceType = interface_type
473        create.dwMaxDisplayWidth = width
474        create.dwMaxDisplayHeight = height
475        #create.pDevice = 0
476        create.dwInterfaceVersion = NVFBC_DLL_VERSION
477        if key:
478            binkey = parse_nvfbc_hex_key(key)
479            ckey = binkey
480            create.pPrivateData = <void*> ckey
481            create.dwPrivateDataSize = len(ckey)
482            log("create_context() key data=%#x, size=%i", <uintptr_t> ckey, len(ckey))
483        res = NvFBC.NvFBC_CreateEx(cvp(<uintptr_t> &create))
484        log("create_context() NvFBC_CreateEx()=%i for key=%s", res, key)
485        if res==0:
486            #success!
487            break
488    log("NvFBC_CreateEx(%#x)=%i", <uintptr_t> &create, res)
489    raiseNvFBC(res, "NvFBC_CreateEx")
490    info = {
491        "max-display-width"     : create.dwMaxDisplayWidth,
492        "max-display-height"    : create.dwMaxDisplayHeight,
493        "version"               : create.dwNvFBCVersion,
494        "context"               : <uintptr_t> create.pNvFBC,
495        }
496    log("NvFBC_CreateEx: %s", info)
497    return info
498
499cdef get_frame_grab_info(NvFBCFrameGrabInfo *grab_info):
500    return {
501        "width"             : int(grab_info.dwWidth),
502        "height"            : int(grab_info.dwHeight),
503        "stride"            : int(grab_info.dwBufferWidth),
504        "overlay-active"    : bool(grab_info.bOverlayActive),
505        "first-buffer"      : bool(grab_info.bFirstBuffer),
506        "hw-mouse-visible"  : bool(grab_info.bHWMouseVisible),
507        "protected-content" : bool(grab_info.bProtectedContent),
508        "stereo"            : bool(grab_info.bStereoOn),
509        "IGPU-capture"      : bool(grab_info.bIGPUCapture),
510        "source-pid"        : int(grab_info.dwSourcePID),
511        "HDR"               : bool(grab_info.bIsHDR),
512        "wait-mode"         : int(grab_info.dwWaitModeUsed),
513        }
514
515def get_version():
516    global NvFBC
517    assert NvFBC
518    cdef NvU32 version = 0
519    cdef NVFBCRESULT res = NvFBC.NvFBC_GetSDKVersion(cvp(<uintptr_t> &version))
520    log("NvFBC_GetSDKVersion()=%i version=%i", res, version)
521    raiseNvFBC(res, "NvFBC_GetSDKVersion")
522    return version
523
524def get_type():
525    return "nvfbc"
526
527def get_info():
528    info = {
529            "type"              : "nvfbc",
530            "version"           : get_version(),
531            }
532    cards = get_cards()
533    if cards:
534        info["cards"] = cards
535    #only show the version if we have it already (don't probe now)
536    v = get_nvidia_module_version(False)
537    if v:
538        info["kernel_module_version"] = v
539    return info
540
541
542SYS_PIXEL_FORMAT_CONST = {
543    "BGRX"      : NVFBC_TOSYS_ARGB,
544    "RGB"       : NVFBC_TOSYS_RGB,
545    #"YUV420P"   : NVFBC_TOSYS_YYYYUV420p,
546    #"RGBP"      : NVFBC_TOSYS_RGB_PLANAR,
547    #NVFBC_TOSYS_XOR,
548    #"YUV444P"   : NVFBC_TOSYS_YUV444p,
549    "r210"      : NVFBC_TOSYS_ARGB10,
550    }
551
552
553cdef class NvFBC_SysCapture:
554    cdef NvFBCToSys *context
555    cdef uint8_t *framebuffer
556    cdef uint8_t setup
557    cdef object pixel_format
558    cdef NvFBCFrameGrabInfo grab_info
559    cdef NVFBC_TOSYS_GRAB_FRAME_PARAMS grab
560
561    cdef object __weakref__
562
563    def init_context(self, int width=-1, int height=-1, pixel_format=DEFAULT_PIXEL_FORMAT):
564        log("init_context(%i, %i, %s)", width, height, pixel_format)
565        global SYS_PIXEL_FORMAT_CONST
566        if pixel_format not in SYS_PIXEL_FORMAT_CONST:
567            raise Exception("unsupported pixel format '%s'" % pixel_format)
568        self.pixel_format = pixel_format
569        self.framebuffer = NULL
570        info = create_context(-1, -1, NVFBC_TO_SYS)
571        maxw = info["max-display-width"]
572        maxh = info["max-display-height"]
573        assert width<=maxw and height<=maxh, "display dimension %ix%i is too large, the maximum supported by this card and driver is %ix%i" % (width, height, maxw, maxh)
574        self.context = <NvFBCToSys*> (<uintptr_t> info["context"])
575        assert self.context!=NULL
576        cdef NVFBC_TOSYS_SETUP_PARAMS params
577        memset(&params, 0, sizeof(NVFBC_TOSYS_SETUP_PARAMS))
578        params.dwVersion = NVFBC_TOSYS_SETUP_PARAMS_VER
579        params.eMode = SYS_PIXEL_FORMAT_CONST[pixel_format]
580        params.bWithHWCursor = False
581        params.bDiffMap = False
582        params.ppBuffer = <void**> &self.framebuffer
583        params.ppDiffMap = NULL
584        cdef NVFBCRESULT res = self.context.NvFBCToSysSetUp(&params)
585        raiseNvFBC(res, "NvFBCToSysSetUp")
586        self.setup = True
587
588    def get_info(self) -> dict:
589        info = get_info()
590        info["pixel-format"] = self.pixel_format
591        return info
592
593    def get_type(self):
594        return  "nvfbc-sys"
595
596    def __repr__(self):
597        return "NvFBC_SysCapture(%#x)" % (<uintptr_t> self.context)
598
599    def __dealloc__(self):
600        self.clean()
601
602    def refresh(self):
603        assert self.context
604        cdef double start = monotonic()
605        memset(&self.grab_info, 0, sizeof(NvFBCFrameGrabInfo))
606        memset(&self.grab, 0, sizeof(NVFBC_TOSYS_GRAB_FRAME_PARAMS))
607        self.grab.dwVersion = NVFBC_TOSYS_GRAB_FRAME_PARAMS_VER
608        self.grab.dwFlags = NVFBC_TOSYS_NOWAIT
609        self.grab.dwTargetWidth = 0  #width
610        self.grab.dwTargetHeight = 0 #height
611        self.grab.dwStartX = 0
612        self.grab.dwStartY = 0
613        self.grab.eGMode = NVFBC_TOSYS_SOURCEMODE_FULL
614        self.grab.pNvFBCFrameGrabInfo = &self.grab_info
615        cdef NVFBCRESULT res
616        with nogil:
617            res = self.context.NvFBCToSysGrabFrame(&self.grab)
618        if res!=0 and self.grab_info.dwDriverInternalError:
619            raise CodecStateException("NvFBC driver internal error")
620        if res==NVFBC_ERROR_DYNAMIC_DISABLE:
621            raise CodecStateException("NvFBC capture has been disabled")
622        if (res!=0 and self.grab_info.bMustRecreate) or res==NVFBC_ERROR_INVALIDATED_SESSION:
623            raise TransientCodecException("NvFBC context invalidated")
624        raiseNvFBC(res, "NvFBCToSysGrabFrame")
625        log("NvFBCToSysGrabFrame() info=%s", get_frame_grab_info(&self.grab_info))
626        return True
627
628    def get_image(self, unsigned int x=0, unsigned int y=0, unsigned int width=0, unsigned int height=0):
629        assert self.context
630        log("nvfbc sys get_image%s", (x, y, width, height))
631        if width==0:
632            width = self.grab_info.dwWidth
633        if height==0:
634            height = self.grab_info.dwHeight
635        assert x==0 and y==0 and width>0 and height>0
636        assert x+width<=self.grab_info.dwWidth, "invalid capture width: %i+%i, capture size is only %i" % (x, width, self.grab_info.dwWidth)
637        assert y+height<=self.grab_info.dwHeight, "invalid capture height: %i+%i, capture size is only %i" % (y, height, self.grab_info.dwHeight)
638        cdef double start = monotonic()
639        #TODO: only copy when the next frame is going to overwrite the buffer,
640        #or when closing the context
641        cdef unsigned int Bpp = len(self.pixel_format)    # ie: "BGR" -> 3
642        cdef unsigned int grab_stride = self.grab_info.dwWidth*Bpp
643        cdef unsigned int stride = grab_stride
644        cdef size_t size
645        cdef MemBuf buf
646        cdef uintptr_t buf_ptr = 0
647        cdef uintptr_t grab_ptr = <uintptr_t> (self.framebuffer+x*Bpp+y*grab_stride)
648        if x>0 or y>0 or self.grab_info.dwWidth-width>16 or self.grab_info.dwHeight-height>16:
649            #copy sub-image with smaller stride:
650            stride = roundup(width*Bpp, 16)
651            size = stride*height
652            buf = padbuf(size, stride)
653            buf_ptr = <uintptr_t> buf.get_mem()
654            with nogil:
655                for _ in range(height):
656                    memcpy(<void *> buf_ptr, <void *> grab_ptr, width*Bpp)
657                    grab_ptr += grab_stride
658                    buf_ptr += stride
659        else:
660            #copy whole:
661            size = self.grab_info.dwBufferWidth*self.grab_info.dwHeight*Bpp
662            buf = padbuf(size, stride)
663            buf_ptr = <uintptr_t> buf.get_mem()
664            with nogil:
665                memcpy(<void *> buf_ptr, <void *> grab_ptr, size)
666        image = ImageWrapper(0, 0, width, height, memoryview(buf), self.pixel_format, Bpp*8, stride, Bpp)
667        end = monotonic()
668        log("image=%s buffer size=%i, (copy took %ims)", image, size, int((end-start)*1000))
669        return image
670
671    def clean(self):
672        log("clean()")
673        if self.setup:
674            self.setup = False
675            ctx = self.context
676            if ctx:
677                self.context = NULL
678                ctx.NvFBCToSysRelease()
679
680
681cdef class NvFBC_CUDACapture:
682    cdef NvFBCCuda *context
683    cdef uint8_t setup
684    cdef object pixel_format
685    cdef NvU32 max_buffer_size
686    cdef int cuda_device_id
687    cdef object cuda_device
688    cdef object cuda_context
689    cdef object cuda_device_buffer
690    cdef NvFBCFrameGrabInfo grab_info
691    cdef NVFBC_CUDA_GRAB_FRAME_PARAMS grab
692
693    cdef object __weakref__
694
695    def init_context(self, int width=-1, int height=-1, pixel_format="BGRX"):
696        log("init_context(%i, %i, %s)", width, height, pixel_format)
697        if pixel_format not in ("BGRX", "r210"):
698            raise Exception("unsupported pixel format '%s'" % pixel_format)
699        self.pixel_format = pixel_format
700        #CUDA init:
701        self.cuda_device_id, self.cuda_device = select_device()
702        if not self.cuda_device:
703            raise Exception("no valid CUDA device")
704        d = self.cuda_device
705        self.cuda_context = d.make_context(flags=driver.ctx_flags.SCHED_AUTO | driver.ctx_flags.MAP_HOST)
706        assert self.cuda_context, "failed to create a CUDA context for device %s" % device_info(d)
707        self.cuda_context.pop()
708        self.cuda_context.push()
709        #NvFBC init:
710        info = create_context(-1, -1, NVFBC_SHARED_CUDA)
711        maxw = info["max-display-width"]
712        maxh = info["max-display-height"]
713        assert width<=maxw and height<=maxh, "display dimension %ix%i is too large, maximum supported by this card and driver is %ix%i" % (width, height, maxw, maxh)
714        self.context = <NvFBCCuda*> (<uintptr_t> info["context"])
715        assert self.context!=NULL
716        self.context.NvFBCCudaGetMaxBufferSize(&self.max_buffer_size)
717        log("NvFBCCudaGetMaxBufferSize: %#x", self.max_buffer_size)
718        cdef NVFBC_CUDA_SETUP_PARAMS params
719        memset(&params, 0, sizeof(NVFBC_CUDA_SETUP_PARAMS))
720        params.dwVersion = NVFBC_CUDA_SETUP_PARAMS_V1_VER
721        params.bEnableSeparateCursorCapture = 1
722        params.bHDRRequest = 0
723        if pixel_format=="BGRX":
724            params.eFormat = NVFBC_TOCUDA_ARGB
725        else:
726            params.eFormat = NVFBC_TOCUDA_ARGB10
727        cdef NVFBCRESULT res = self.context.NvFBCCudaSetup(&params)
728        raiseNvFBC(res, "NvFBCCudaSetup")
729        self.setup = True
730
731    def get_info(self) -> dict:
732        info = get_info()
733        info["pixel-format"] = self.pixel_format
734        return info
735
736    def get_type(self):
737        return  "nvfbc-cuda"
738
739    def __repr__(self):
740        return "NvFBC_CUDACapture(%#x)" % (<uintptr_t> self.context)
741
742    def __dealloc__(self):
743        self.clean()
744
745    def refresh(self):
746        return True
747
748    def get_image(self, unsigned int x=0, unsigned int y=0, unsigned int width=0, unsigned int height=0):
749        assert self.context
750        log("nvfbc cuda get_image%s", (x, y, width, height))
751        if width==0:
752            width = self.grab_info.dwWidth
753        if height==0:
754            height = self.grab_info.dwHeight
755        assert x==0 and y==0 and width>0 and height>0
756        cdef double start = monotonic()
757        #allocate CUDA device memory:
758        if not self.cuda_device_buffer:
759            #TODO: choose a better size
760            self.cuda_device_buffer = driver.mem_alloc(self.max_buffer_size)
761            log("max_buffer_size=%#x, cuda device buffer=%s", self.max_buffer_size, self.cuda_device_buffer)
762        #cuda_device_buffer, stride = self.cuda_device.mem_alloc_pitch(4096, 2160, 16)
763        memset(&self.grab_info, 0, sizeof(NvFBCFrameGrabInfo))
764        memset(&self.grab, 0, sizeof(NVFBC_CUDA_GRAB_FRAME_PARAMS))
765        self.grab.dwVersion = NVFBC_CUDA_GRAB_FRAME_PARAMS_V1_VER
766        ptr = <uintptr_t> int(self.cuda_device_buffer)
767        self.grab.pCUDADeviceBuffer = <void*> ptr
768        self.grab.pNvFBCFrameGrabInfo = &self.grab_info
769        self.grab.dwFlags = NVFBC_TOCUDA_NOWAIT
770        cdef NVFBCRESULT res
771        with nogil:
772            res = self.context.NvFBCCudaGrabFrame(&self.grab)
773        if res<0:
774            raiseNvFBC(res, "NvFBCToSysGrabFrame")
775        elif res!=0:
776            raise Exception("CUDA Grab Frame failed: %s" % CUDA_ERRORS_INFO.get(res, res))
777        cdef double end = monotonic()
778        log("NvFBCCudaGrabFrame: info=%s, elapsed=%ims", get_frame_grab_info(&self.grab_info), int((end-start)*1000))
779        assert x==0 and y==0 and width>0 and height>0
780        assert x+width<=self.grab_info.dwWidth, "invalid capture width: %i+%i, capture size is only %i" % (x, width, self.grab_info.dwWidth)
781        assert y+height<=self.grab_info.dwHeight, "invalid capture height: %i+%i, capture size is only %i" % (y, height, self.grab_info.dwHeight)
782        Bpp = len(self.pixel_format)    # ie: "BGR" -> 3
783        image = SharedCUDAImageWrapper(0, 0, width, height, None, self.pixel_format, Bpp*8, int(self.grab_info.dwBufferWidth*Bpp), Bpp, False, None)
784        image.cuda_device_buffer = self.cuda_device_buffer
785        image.cuda_context = self.cuda_context
786        image.buffer_size = self.max_buffer_size
787        return image
788
789    def clean(self):
790        log("clean()")
791        cuda_context = self.cuda_context
792        self.cuda_context = None
793        if self.setup:
794            self.setup = False
795            if self.context:
796                self.context.NvFBCCudaRelease()
797                self.context = NULL
798        if cuda_context:
799            try:
800                cuda_context.pop()
801                cuda_context.detach()
802            except:
803                log("%s.pop() or detach()", cuda_context, exc_info=True)
804        #don't free it - an imagewrapper may still use it:
805        #TODO: we should invalidate it
806        self.cuda_device_buffer = None
807
808
809class SharedCUDAImageWrapper(CUDAImageWrapper):
810
811    def free_cuda_device_buffer(self):
812        #override so we only clear the reference,
813        #the buffer is going to be re-used so we cannot free it
814        self.cuda_device_buffer = None
815
816
817def init_module():
818    log("nvfbc.init_module()")
819    init_nvfbc_library()
820
821def cleanup_module():
822    log("nvfbc.cleanup_module()")
823    unload_library()
824
825def selftest(full=False):
826    pass
827