Refactor kms_vaapi and kms_cuda

Also fixes color metadata, color range and hdr on nvidia wayland.

10 files changed, 129 insertions, 36 deletions

diff --git a/include/capture/capture.h b/include/capture/capture.h
index 82a9555..2ca57cc 100644
--- a/include/capture/capture.h
+++ b/include/capture/capture.h
@@ -1,17 +1,24 @@
 #ifndef GSR_CAPTURE_CAPTURE_H
 #define GSR_CAPTURE_CAPTURE_H
 
+#include "../color_conversion.h"
 #include <stdbool.h>
 
 typedef struct AVCodecContext AVCodecContext;
 typedef struct AVFrame AVFrame;
+typedef void* VADisplay;
+typedef struct _VADRMPRIMESurfaceDescriptor VADRMPRIMESurfaceDescriptor;
+typedef struct gsr_cuda gsr_cuda;
+typedef struct AVFrame AVFrame;
+typedef struct CUgraphicsResource_st *CUgraphicsResource;
+typedef struct CUarray_st *CUarray;
 
 typedef struct gsr_capture gsr_capture;
 
 struct gsr_capture {
     /* These methods should not be called manually. Call gsr_capture_* instead */
-    int (*start)(gsr_capture *cap, AVCodecContext *video_codec_context);
-    void (*tick)(gsr_capture *cap, AVCodecContext *video_codec_context, AVFrame **frame); /* can be NULL */
+    int (*start)(gsr_capture *cap, AVCodecContext *video_codec_context, AVFrame *frame);
+    void (*tick)(gsr_capture *cap, AVCodecContext *video_codec_context); /* can be NULL */
     bool (*should_stop)(gsr_capture *cap, bool *err); /* can be NULL */
     int (*capture)(gsr_capture *cap, AVFrame *frame);
     void (*capture_end)(gsr_capture *cap, AVFrame *frame); /* can be NULL */
@@ -21,12 +28,34 @@ struct gsr_capture {
     bool started;
 };
 
-int gsr_capture_start(gsr_capture *cap, AVCodecContext *video_codec_context);
-void gsr_capture_tick(gsr_capture *cap, AVCodecContext *video_codec_context, AVFrame **frame);
+typedef struct gsr_capture_base gsr_capture_base;
+
+struct gsr_capture_base {
+    unsigned int input_texture;
+    unsigned int target_textures[2];
+    unsigned int cursor_texture;
+
+    gsr_color_conversion color_conversion;
+
+    AVCodecContext *video_codec_context;
+};
+
+typedef struct {
+    gsr_cuda *cuda;
+    CUgraphicsResource *cuda_graphics_resources;
+    CUarray *mapped_arrays;
+} gsr_cuda_context;
+
+int gsr_capture_start(gsr_capture *cap, AVCodecContext *video_codec_context, AVFrame *frame);
+void gsr_capture_tick(gsr_capture *cap, AVCodecContext *video_codec_context);
 bool gsr_capture_should_stop(gsr_capture *cap, bool *err);
 int gsr_capture_capture(gsr_capture *cap, AVFrame *frame);
 void gsr_capture_end(gsr_capture *cap, AVFrame *frame);
 /* Calls |gsr_capture_stop| as well */
 void gsr_capture_destroy(gsr_capture *cap, AVCodecContext *video_codec_context);
 
+bool gsr_capture_base_setup_vaapi_textures(gsr_capture_base *self, AVFrame *frame, gsr_egl *egl, VADisplay va_dpy, VADRMPRIMESurfaceDescriptor *prime, gsr_color_range color_range);
+bool gsr_capture_base_setup_cuda_textures(gsr_capture_base *base, AVFrame *frame, gsr_cuda_context *cuda_context, gsr_egl *egl, gsr_color_range color_range, bool hdr);
+void gsr_capture_base_stop(gsr_capture_base *self, gsr_egl *egl);
+
 #endif /* GSR_CAPTURE_CAPTURE_H */
diff --git a/include/capture/kms.h b/include/capture/kms.h
new file mode 100644
index 0000000..ae53f15
--- /dev/null
+++ b/include/capture/kms.h
@@ -0,0 +1,48 @@
+#ifndef GSR_CAPTURE_KMS_H
+#define GSR_CAPTURE_KMS_H
+
+#include "../../kms/client/kms_client.h"
+#include "../color_conversion.h"
+#include "../vec2.h"
+#include "../defs.h"
+#include <stdbool.h>
+
+typedef struct gsr_capture_base gsr_capture_base;
+typedef struct AVCodecContext AVCodecContext;
+typedef struct AVMasteringDisplayMetadata AVMasteringDisplayMetadata;
+typedef struct AVContentLightMetadata AVContentLightMetadata;
+typedef struct gsr_capture_kms gsr_capture_kms;
+typedef struct gsr_egl gsr_egl;
+typedef struct AVFrame AVFrame;
+
+#define MAX_CONNECTOR_IDS 32
+
+typedef struct {
+    uint32_t connector_ids[MAX_CONNECTOR_IDS];
+    int num_connector_ids;
+} MonitorId;
+
+struct gsr_capture_kms {
+    bool should_stop;
+    bool stop_is_error;
+    
+    gsr_kms_client kms_client;
+    gsr_kms_response kms_response;
+
+    vec2i capture_pos;
+    vec2i capture_size;
+    MonitorId monitor_id;
+
+    AVMasteringDisplayMetadata *mastering_display_metadata;
+    AVContentLightMetadata *light_metadata;
+
+    gsr_monitor_rotation monitor_rotation;
+};
+
+/* Returns 0 on success */
+int gsr_capture_kms_start(gsr_capture_kms *self, gsr_capture_base *base, const char *display_to_capture, gsr_egl *egl, AVCodecContext *video_codec_context);
+void gsr_capture_kms_stop(gsr_capture_kms *self);
+bool gsr_capture_kms_capture(gsr_capture_kms *self, gsr_capture_base *base, AVFrame *frame, gsr_egl *egl, bool hdr, bool screen_plane_use_modifiers, bool cursor_texture_is_external);
+void gsr_capture_kms_cleanup_kms_fds(gsr_capture_kms *self);
+
+#endif /* GSR_CAPTURE_KMS_H */
diff --git a/include/capture/kms_cuda.h b/include/capture/kms_cuda.h
index 5691f61..ce3e8ce 100644
--- a/include/capture/kms_cuda.h
+++ b/include/capture/kms_cuda.h
@@ -3,6 +3,7 @@
 
 #include "../vec2.h"
 #include "../utils.h"
+#include "../color_conversion.h"
 #include "capture.h"
 
 typedef struct _XDisplay Display;
@@ -12,6 +13,7 @@ typedef struct {
     const char *display_to_capture; /* if this is "screen", then the first monitor is captured. A copy is made of this */
     gsr_gpu_info gpu_inf;
     bool hdr;
+    gsr_color_range color_range;
 } gsr_capture_kms_cuda_params;
 
 gsr_capture* gsr_capture_kms_cuda_create(const gsr_capture_kms_cuda_params *params);
diff --git a/include/capture/kms_vaapi.h b/include/capture/kms_vaapi.h
index b670a0d..54d0bad 100644
--- a/include/capture/kms_vaapi.h
+++ b/include/capture/kms_vaapi.h
@@ -12,7 +12,6 @@ typedef struct {
     gsr_egl *egl;
     const char *display_to_capture; /* if this is "screen", then the first monitor is captured. A copy is made of this */
     gsr_gpu_info gpu_inf;
-    bool wayland;
     bool hdr;
     gsr_color_range color_range;
 } gsr_capture_kms_vaapi_params;
diff --git a/include/color_conversion.h b/include/color_conversion.h
index 8bfe599..222be4b 100644
--- a/include/color_conversion.h
+++ b/include/color_conversion.h
@@ -3,6 +3,7 @@
 
 #include "shader.h"
 #include "vec2.h"
+#include <stdbool.h>
 
 typedef enum {
     GSR_COLOR_RANGE_LIMITED,
diff --git a/include/cuda.h b/include/cuda.h
index 41fe15b..fd1f9f9 100644
--- a/include/cuda.h
+++ b/include/cuda.h
@@ -73,7 +73,8 @@ typedef CUDA_MEMCPY2D_v2 CUDA_MEMCPY2D;
 
 typedef struct CUgraphicsResource_st *CUgraphicsResource;
 
-typedef struct {
+typedef struct gsr_cuda gsr_cuda;
+struct gsr_cuda {
     gsr_overclock overclock;
     bool do_overclock;
 
@@ -88,8 +89,9 @@ typedef struct {
     CUresult (*cuCtxPushCurrent_v2)(CUcontext ctx);
     CUresult (*cuCtxPopCurrent_v2)(CUcontext *pctx);
     CUresult (*cuGetErrorString)(CUresult error, const char **pStr);
-    CUresult (*cuMemsetD8_v2)(CUdeviceptr dstDevice, unsigned char uc, size_t N);
     CUresult (*cuMemcpy2D_v2)(const CUDA_MEMCPY2D *pCopy);
+    CUresult (*cuMemcpy2DAsync_v2)(const CUDA_MEMCPY2D *pcopy, CUstream hStream);
+    CUresult (*cuStreamSynchronize)(CUstream hStream);
 
     CUresult (*cuGraphicsGLRegisterImage)(CUgraphicsResource *pCudaResource, unsigned int image, unsigned int target, unsigned int flags);
     CUresult (*cuGraphicsEGLRegisterImage)(CUgraphicsResource *pCudaResource, void *image, unsigned int flags);
@@ -98,7 +100,7 @@ typedef struct {
     CUresult (*cuGraphicsUnmapResources)(unsigned int count, CUgraphicsResource *resources, CUstream hStream);
     CUresult (*cuGraphicsUnregisterResource)(CUgraphicsResource resource);
     CUresult (*cuGraphicsSubResourceGetMappedArray)(CUarray *pArray, CUgraphicsResource resource, unsigned int arrayIndex, unsigned int mipLevel);
-} gsr_cuda;
+};
 
 bool gsr_cuda_load(gsr_cuda *self, Display *display, bool overclock);
 void gsr_cuda_unload(gsr_cuda *self);
diff --git a/include/defs.h b/include/defs.h
new file mode 100644
index 0000000..473583c
--- /dev/null
+++ b/include/defs.h
@@ -0,0 +1,28 @@
+#ifndef GSR_DEFS_H
+#define GSR_DEFS_H
+
+typedef enum {
+    GSR_GPU_VENDOR_AMD,
+    GSR_GPU_VENDOR_INTEL,
+    GSR_GPU_VENDOR_NVIDIA
+} gsr_gpu_vendor;
+
+typedef struct {
+    gsr_gpu_vendor vendor;
+    int gpu_version; /* 0 if unknown */
+} gsr_gpu_info;
+
+typedef enum {
+    GSR_MONITOR_ROT_0,
+    GSR_MONITOR_ROT_90,
+    GSR_MONITOR_ROT_180,
+    GSR_MONITOR_ROT_270
+} gsr_monitor_rotation;
+
+typedef enum {
+    GSR_CONNECTION_X11,
+    GSR_CONNECTION_WAYLAND,
+    GSR_CONNECTION_DRM
+} gsr_connection_type;
+
+#endif /* GSR_DEFS_H */
diff --git a/include/egl.h b/include/egl.h
index 4705742..6f0c11b 100644
--- a/include/egl.h
+++ b/include/egl.h
@@ -69,10 +69,16 @@ typedef void (*__eglMustCastToProperFunctionPointerType)(void);
 #define GL_TRUE                                 1
 #define GL_TRIANGLES                            0x0004
 #define GL_TEXTURE_2D                           0x0DE1
-#define GL_TEXTURE_EXTERNAL_OES                 0x8D65 // TODO: Use this where applicable
+#define GL_TEXTURE_EXTERNAL_OES                 0x8D65
+#define GL_RED                                  0x1903
+#define GL_RG                                   0x8227
 #define GL_RGB                                  0x1907
 #define GL_RGBA                                 0x1908
 #define GL_RGBA8                                0x8058
+#define GL_R8                                   0x8229
+#define GL_RG8                                  0x822B
+#define GL_R16                                  0x822A
+#define GL_RG16                                 0x822C
 #define GL_UNSIGNED_BYTE                        0x1401
 #define GL_COLOR_BUFFER_BIT                     0x00004000
 #define GL_TEXTURE_WRAP_S                       0x2802
@@ -133,7 +139,8 @@ typedef struct {
     int num_outputs;
 } gsr_wayland;
 
-typedef struct {
+typedef struct gsr_egl gsr_egl;
+struct gsr_egl {
     void *egl_library;
     void *gl_library;
 
@@ -217,7 +224,7 @@ typedef struct {
     int (*glGetUniformLocation)(unsigned int program, const char *name);
     void (*glUniform1f)(int location, float v0);
     void (*glUniform2f)(int location, float v0, float v1);
-} gsr_egl;
+};
 
 bool gsr_egl_load(gsr_egl *self, Display *dpy, bool wayland);
 void gsr_egl_unload(gsr_egl *self);
diff --git a/include/shader.h b/include/shader.h
index 37f4c09..57d1096 100644
--- a/include/shader.h
+++ b/include/shader.h
@@ -1,7 +1,7 @@
 #ifndef GSR_SHADER_H
 #define GSR_SHADER_H
 
-#include "egl.h"
+typedef struct gsr_egl gsr_egl;
 
 typedef struct {
     gsr_egl *egl;
diff --git a/include/utils.h b/include/utils.h
index 898f375..cb3bd6b 100644
--- a/include/utils.h
+++ b/include/utils.h
@@ -3,28 +3,11 @@
 
 #include "vec2.h"
 #include "../include/egl.h"
+#include "../include/defs.h"
 #include <stdbool.h>
 #include <stdint.h>
 #include <X11/extensions/Xrandr.h>
 
-typedef enum {
-    GSR_GPU_VENDOR_AMD,
-    GSR_GPU_VENDOR_INTEL,
-    GSR_GPU_VENDOR_NVIDIA
-} gsr_gpu_vendor;
-
-typedef struct {
-    gsr_gpu_vendor vendor;
-    int gpu_version; /* 0 if unknown */
-} gsr_gpu_info;
-
-typedef enum {
-    GSR_MONITOR_ROT_0,
-    GSR_MONITOR_ROT_90,
-    GSR_MONITOR_ROT_180,
-    GSR_MONITOR_ROT_270
-} gsr_monitor_rotation;
-
 typedef struct {
     const char *name;
     int name_len;
@@ -36,12 +19,6 @@ typedef struct {
     uint32_t monitor_identifier; /* Only on drm and wayland */
 } gsr_monitor;
 
-typedef enum {
-    GSR_CONNECTION_X11,
-    GSR_CONNECTION_WAYLAND,
-    GSR_CONNECTION_DRM
-} gsr_connection_type;
-
 typedef struct {
     const char *name;
     int name_len;