obs-ffmpeg: Cache registered CUDA graphics resources

obsproject · Jan 31, 2024 · a5d154a · a5d154a
1 parent 646fc7b
commit a5d154a
Showing 1 changed file with 56 additions and 13 deletions.
diff --git a/plugins/obs-ffmpeg/obs-nvenc.c b/plugins/obs-ffmpeg/obs-nvenc.c
@@ -66,13 +66,17 @@
 struct nv_bitstream;
 struct nv_texture;
 
-#ifdef _WIN32
 struct handle_tex {
+#ifdef _WIN32
 	uint32_t handle;
 	ID3D11Texture2D *tex;
 	IDXGIKeyedMutex *km;
-};
+#else
+	GLuint tex_id;
+	CUgraphicsResource res_y;
+	CUgraphicsResource res_uv;
 #endif
+};
 
 /* ------------------------------------------------------------------------- */
 /* Main Implementation Structure                                             */
@@ -118,6 +122,7 @@ struct nvenc_data {
 	bool fallback;
 	int32_t bframes;
 
+	DARRAY(struct handle_tex) input_textures;
 	DARRAY(struct nv_bitstream) bitstreams;
 	DARRAY(struct nv_cuda_surface) surfaces;
 	NV_ENC_BUFFER_FORMAT surface_format;
@@ -129,7 +134,6 @@ struct nvenc_data {
 
 #ifdef _WIN32
 	DARRAY(struct nv_texture) textures;
-	DARRAY(struct handle_tex) input_textures;
 	ID3D11Device *device;
 	ID3D11DeviceContext *context;
 #endif
@@ -1608,6 +1612,13 @@ static void nvenc_destroy(void *data)
 	if (enc->device) {
 		enc->device->lpVtbl->Release(enc->device);
 	}
+#else
+	for (size_t i = 0; i < enc->input_textures.num; i++) {
+		CUgraphicsResource res_y = enc->input_textures.array[i].res_y;
+		CUgraphicsResource res_uv = enc->input_textures.array[i].res_uv;
+		cu->cuGraphicsUnregisterResource(res_y);
+		cu->cuGraphicsUnregisterResource(res_uv);
+	}
 #endif
 	if (enc->cu_ctx) {
 		cu->cuCtxPopCurrent(NULL);
@@ -1618,10 +1629,10 @@ static void nvenc_destroy(void *data)
 	bfree(enc->sei);
 	deque_free(&enc->dts_list);
 	da_free(enc->surfaces);
+	da_free(enc->input_textures);
 	da_free(enc->bitstreams);
 #ifdef _WIN32
 	da_free(enc->textures);
-	da_free(enc->input_textures);
 #endif
 	da_free(enc->packet_data);
 	bfree(enc->roi_map);
@@ -2011,19 +2022,53 @@ static bool nvenc_encode_tex(void *data, uint32_t handle, int64_t pts,
 		}                                                            \
 	}
 
-static inline bool copy_tex_cuda(struct nvenc_data *enc, const bool p010,
-				 GLuint tex[2], struct nv_cuda_surface *surf)
+static inline bool get_res_for_tex_ids(struct nvenc_data *enc, GLuint tex_id_y,
+				       GLuint tex_id_uv,
+				       CUgraphicsResource *tex_y,
+				       CUgraphicsResource *tex_uv)
 {
 	bool success = true;
-	CUgraphicsResource mapped_tex[2] = {0};
-	CUarray mapped_cuda;
+
+	for (size_t idx = 0; idx < enc->input_textures.num; idx++) {
+		struct handle_tex *ht = &enc->input_textures.array[idx];
+		if (ht->tex_id != tex_id_y)
+			continue;
+
+		*tex_y = ht->res_y;
+		*tex_uv = ht->res_uv;
+		return success;
+	}
 
 	CU_CHECK(cu->cuGraphicsGLRegisterImage(
-		&mapped_tex[0], tex[0], GL_TEXTURE_2D,
+		tex_y, tex_id_y, GL_TEXTURE_2D,
 		CU_GRAPHICS_REGISTER_FLAGS_READ_ONLY))
 	CU_CHECK(cu->cuGraphicsGLRegisterImage(
-		&mapped_tex[1], tex[1], GL_TEXTURE_2D,
+		tex_uv, tex_id_uv, GL_TEXTURE_2D,
 		CU_GRAPHICS_REGISTER_FLAGS_READ_ONLY))
+
+	struct handle_tex ht = {tex_id_y, *tex_y, *tex_uv};
+	da_push_back(enc->input_textures, &ht);
+
+unmap:
+	if (!success) {
+		cu->cuGraphicsUnregisterResource(*tex_y);
+		cu->cuGraphicsUnregisterResource(*tex_uv);
+	}
+
+	return success;
+}
+
+static inline bool copy_tex_cuda(struct nvenc_data *enc, const bool p010,
+				 GLuint tex[2], struct nv_cuda_surface *surf)
+{
+	bool success = true;
+	CUgraphicsResource mapped_tex[2] = {0};
+	CUarray mapped_cuda;
+
+	if (!get_res_for_tex_ids(enc, tex[0], tex[1], &mapped_tex[0],
+				 &mapped_tex[1]))
+		return false;
+
 	CU_CHECK(cu->cuGraphicsMapResources(2, mapped_tex, 0))
 
 	CUDA_MEMCPY2D m = {0};
@@ -2039,7 +2084,7 @@ static inline bool copy_tex_cuda(struct nvenc_data *enc, const bool p010,
 	m.srcArray = mapped_cuda;
 	CU_CHECK(cu->cuMemcpy2D(&m))
 
-	// Map anc copy UV texture
+	// Map and copy UV texture
 	CU_CHECK(cu->cuGraphicsSubResourceGetMappedArray(&mapped_cuda,
 							 mapped_tex[1], 0, 0))
 	m.srcArray = mapped_cuda;
@@ -2050,8 +2095,6 @@ static inline bool copy_tex_cuda(struct nvenc_data *enc, const bool p010,
 
 unmap:
 	cu->cuGraphicsUnmapResources(2, mapped_tex, 0);
-	cu->cuGraphicsUnregisterResource(mapped_tex[0]);
-	cu->cuGraphicsUnregisterResource(mapped_tex[1]);
 
 	return success;
 }