Ticket #370: nvenc-dualbuffers.patch

xpra/codecs/nvenc/encoder.pyx

@@ -1187 +1187 @@
     cdef NV_ENCODE_API_FUNCTION_LIST functionList               #@DuplicatedSignature
     cdef void *context
     cdef NV_ENC_REGISTERED_PTR inputHandle
-    cdef CUdeviceptr cudaBuffer
+    cdef CUdeviceptr cudaInputBuffer
+    cdef CUdeviceptr cudaNV12Buffer
+    cdef size_t inputPitch
+    cdef size_t NV12Pitch
     cdef void *inputBuffer
-    cdef size_t pitch
     cdef void *bitstreamBuffer
     cdef NV_ENC_BUFFER_FORMAT bufferFmt
     cdef object codec_name
@@ -1242 +1244 @@
         raiseCUDA(cuCtxCreate(&self.cuda_context, 0, cuda_device))
         debug("cuCtxCreate: device_id=%s, cuda_device=%s, cuda_context=%s", device_id, cuda_device, hex(<long> self.cuda_context))
         #allocate CUDA input buffer (on device):
-        raiseCUDA(cuMemAllocPitch(&self.cudaBuffer, &self.pitch, self.encoder_width, self.encoder_height*3/2, 16), "allocating CUDA input buffer on device")
-        debug("cudaBuffer=%s, pitch=%s", hex(<long> self.cudaBuffer), self.pitch)
+        raiseCUDA(cuMemAllocPitch(&self.cudaInputBuffer, &self.inputPitch, self.encoder_width, self.encoder_height*3/2, 16), "allocating CUDA input buffer on device")
+        debug("CUDA Input Buffer=%s, pitch=%s", hex(<long> self.cudaInputBuffer), self.inputPitch)
+        #allocate CUDA NV12 buffer (on device):
+        raiseCUDA(cuMemAllocPitch(&self.cudaNV12Buffer, &self.NV12Pitch, self.encoder_width, self.encoder_height*3/2, 16), "allocating CUDA NV12 buffer on device")
+        debug("CUDA NV12 Buffer=%s, pitch=%s", hex(<long> self.cudaNV12Buffer), self.NV12Pitch)
         #allocate buffer on host:
-        raiseCUDA(cuMemAllocHost(&self.inputBuffer, self.pitch*self.encoder_height*3/2), "allocating CUDA input buffer on host")
+        raiseCUDA(cuMemAllocHost(&self.inputBuffer, self.inputPitch*self.encoder_height*3/2), "allocating CUDA input buffer on host")
         debug("inputBuffer=%s", hex(<long> self.inputBuffer))
 
         self.init_nvenc()
@@ -1297 +1302 @@
             memset(&registerResource, 0, sizeof(NV_ENC_REGISTER_RESOURCE))
             registerResource.version = NV_ENC_REGISTER_RESOURCE_VER
             registerResource.resourceType = NV_ENC_INPUT_RESOURCE_TYPE_CUDADEVICEPTR
-            registerResource.resourceToRegister = <void *> self.cudaBuffer
+            registerResource.resourceToRegister = <void *> self.cudaNV12Buffer
             registerResource.width = self.encoder_width
             registerResource.height = self.encoder_height
-            registerResource.pitch = self.pitch
+            registerResource.pitch = self.NV12Pitch
             raiseNVENC(self.functionList.nvEncRegisterResource(self.context, &registerResource), "registering CUDA input buffer")
             self.inputHandle = registerResource.registeredResource
             debug("input handle for CUDA buffer: %s", hex(<long> self.inputHandle))
@@ -1347 +1352 @@
                 self.cuda_context = NULL
 
     def cuda_clean(self):
-        if self.inputHandle!=NULL:
+        if self.inputHandle!=NULL and self.context!=NULL:
             debug("clean() unregistering %s", hex(<long> self.inputHandle))
             raiseNVENC(self.functionList.nvEncUnregisterResource(self.context, self.inputHandle), "unregistering CUDA input buffer")
             self.inputHandle = NULL
@@ -1355 +1360 @@
             debug("clean() freeing CUDA host buffer %s", hex(<long> self.inputBuffer))
             raiseCUDA(cuMemFreeHost(self.inputBuffer), "freeing host buffer")
             self.inputBuffer = NULL
-        if (<void *> self.cudaBuffer)!=NULL:
-            debug("clean() freeing CUDA device buffer %s", hex(<long> self.cudaBuffer))
-            raiseCUDA(cuMemFree(self.cudaBuffer), "freeing CUDA buffer")
-            self.cudaBuffer = <CUdeviceptr> NULL
-        if self.bitstreamBuffer!=NULL:
-            debug("clean() destroying bitstream buffer %s", hex(<long> self.bitstreamBuffer))
-            raiseNVENC(self.functionList.nvEncDestroyBitstreamBuffer(self.context, self.bitstreamBuffer), "destroying output buffer")
-            self.bitstreamBuffer = NULL
-        raiseNVENC(self.functionList.nvEncDestroyEncoder(self.context), "destroying context")
+        if (<void *> self.cudaInputBuffer)!=NULL:
+            debug("clean() freeing CUDA input buffer %s", hex(<long> self.cudaInputBuffer))
+            raiseCUDA(cuMemFree(self.cudaInputBuffer), "freeing CUDA input buffer")
+            self.cudaInputBuffer = <CUdeviceptr> NULL
+        if (<void *> self.cudaNV12Buffer)!=NULL:
+            debug("clean() freeing CUDA NV12 buffer %s", hex(<long> self.cudaNV12Buffer))
+            raiseCUDA(cuMemFree(self.cudaNV12Buffer), "freeing CUDA NV12 buffer")
+            self.cudaNV12Buffer = <CUdeviceptr> NULL
+        if self.context!=NULL:
+            if self.bitstreamBuffer!=NULL:
+                debug("clean() destroying bitstream buffer %s", hex(<long> self.bitstreamBuffer))
+                raiseNVENC(self.functionList.nvEncDestroyBitstreamBuffer(self.context, self.bitstreamBuffer), "destroying output buffer")
+                self.bitstreamBuffer = NULL
+            debug("clean() destroying encoder %s", hex(<long> self.context))
+            raiseNVENC(self.functionList.nvEncDestroyEncoder(self.context), "destroying context")
 
     def get_width(self):
         return self.width
@@ -1431 +1442 @@
         debug("compress_image(..) pixels=%s", type(pixels))
 
         #copy to input buffer:
-        size = self.pitch * self.encoder_height * 3/2
+        size = self.inputPitch * self.encoder_height * 3/2
         memset(self.inputBuffer, 0, size)
         #copy luma:
         assert PyObject_AsReadBuffer(pixels[0], &Y, &Y_len)==0
@@ -1439 +1450 @@
         assert PyObject_AsReadBuffer(pixels[2], &Cr, &Cr_len)==0
         stride = strides[0]
         for y in range(h):
-            memcpy(self.inputBuffer + y*self.pitch, Y + stride*y, w)
+            memcpy(self.inputBuffer + y*self.inputPitch, Y + stride*y, w)
         #copy chroma packed:
         assert strides[1]==strides[2], "U and V strides differ: %s vs %s" % (strides[1], strides[2])
         stride = strides[1]
         for y in range(h/2):
-            offset = (self.encoder_height + y) * self.pitch
+            offset = (self.encoder_height + y) * self.inputPitch
             for x in range(w/2):
                 (<char*> self.inputBuffer)[offset + (x*2)] = (<char *> Cb)[stride*y + x]
                 (<char*> self.inputBuffer)[offset + (x*2)+1] = (<char *> Cr)[stride*y + x]
 
         #copy input buffer to CUDA buffer:
-        raiseCUDA(cuMemcpyHtoD(self.cudaBuffer, self.inputBuffer, size), "copy from host to device")
+        raiseCUDA(cuMemcpyHtoD(self.cudaInputBuffer, self.inputBuffer, size), "copy from host to device input buffer")
+        #FIXME: just for testing:
+        raiseCUDA(cuMemcpyHtoD(self.cudaNV12Buffer, self.inputBuffer, size), "TEMPORARY")
         debug("compress_image(..) input buffer copied to device")
 
         #map buffer so nvenc can access it:
@@ -1467 +1480 @@
             picParams.pictureStruct = NV_ENC_PIC_STRUCT_FRAME
             picParams.inputWidth = self.encoder_width
             picParams.inputHeight = self.encoder_height
-            picParams.inputPitch = self.pitch
+            picParams.inputPitch = self.NV12Pitch
             picParams.inputBuffer = mapInputResource.mappedResource
             picParams.outputBitstream = self.bitstreamBuffer
             #picParams.pictureType: required when enablePTD is disabled