vulkan-playground/vaapi_encoder.h

#ifndef VAAPI_ENCODER_H
#define VAAPI_ENCODER_H

#include <cassert>
#include <cstdio>
#include <unistd.h>
#include <va/va.h>
#include <va/va_drm.h>
#include <va/va_drmcommon.h>
#include <pthread.h>
#include <fcntl.h>
#include "bitstream.h"

#define CHAIN_SIZE 4

struct VaapiEncoder
{
	VADisplay dpy = nullptr;
	VAContextID ctx;
	VASurfaceID inputFrames[CHAIN_SIZE];
	int width, height;
	int frame_count;
	int destroying;
	VASurfaceID reference_picture[3];
	pthread_t output_thread = 0;
	pthread_mutex_t mutex;
	pthread_cond_t output_cond;
	VABufferID output_buf;
	int output_fd = -1;
	VAConfigAttrib attrib[VAConfigAttribTypeMax];
	VAConfigID cfg;
	inline void Destroy()
	{
		if(output_fd >= 0)
			close(output_fd);
		output_fd = -1;


		if(output_thread)
		{
			pthread_mutex_lock(&mutex);
			destroying = 1;
			pthread_cond_signal(&output_cond);
			pthread_mutex_unlock(&mutex);
			pthread_join(output_thread, NULL);
			pthread_mutex_destroy(&mutex);
			pthread_cond_destroy(&output_cond);
		}
		if(dpy)
		{
			if(ctx != VA_INVALID_ID)
			{

				vaDestroyContext(dpy, ctx);
				vaDestroyConfig(dpy, cfg);
				cfg = VA_INVALID_ID;
				ctx = VA_INVALID_ID;
				vaDestroySurfaces(dpy, reference_picture, 3);
				vaDestroySurfaces(dpy, inputFrames, CHAIN_SIZE);
				if(output_buf != VA_INVALID_ID)
					vaDestroyBuffer(dpy, output_buf);
				output_buf = VA_INVALID_ID;
			}
			vaTerminate(dpy);
			dpy = nullptr;
		}
	}
	bool SetupVA(VAProfile profile, uint32_t format, uint32_t fourcc, int drm_fd, int width, int height, const char *filename, int *dmabuf_fd, uint64_t *mod, uint32_t *size, uint32_t *offset, uint32_t *pitch1, uint32_t *pitch2, uint64_t *modifiers, int modifierscount)
	{
		VASurfaceAttrib va_attribs[5];
		VASurfaceAttribExternalBuffers va_attrib_extbuf = {0};
		VADRMPRIMESurfaceDescriptor drmSurface = {0};
		VAStatus status;
		int major, minor;

		this->width = width;
		this->height = height;
		frame_count = 0;
		destroying = 0;
		dpy = vaGetDisplayDRM(drm_fd);
		output_fd = open(filename, O_WRONLY | O_CREAT | O_TRUNC | O_CLOEXEC, 0644);
		output_buf = VA_INVALID_ID;
		status = vaInitialize(dpy, &major, &minor);
		if(status != VA_STATUS_SUCCESS)
		{
			return false;
		}
		va_attrib_extbuf.pixel_format = fourcc;
		va_attrib_extbuf.width = width;
		va_attrib_extbuf.height = height;
		va_attrib_extbuf.flags = VA_SURFACE_EXTBUF_DESC_ENABLE_TILING;
		va_attrib_extbuf.private_data = NULL;
		VADRMFormatModifierList modList;
		modList.modifiers = modifiers;
		modList.num_modifiers = modifierscount;

		va_attribs[0].type = VASurfaceAttribMemoryType;
		va_attribs[0].flags = VA_SURFACE_ATTRIB_SETTABLE;
		va_attribs[0].value.type = VAGenericValueTypeInteger;
		va_attribs[0].value.value.i = VA_SURFACE_ATTRIB_MEM_TYPE_VA;
		va_attribs[1].type = VASurfaceAttribUsageHint;
		va_attribs[1].flags = VA_SURFACE_ATTRIB_SETTABLE;
		va_attribs[1].value.type = VAGenericValueTypeInteger;
		va_attribs[1].value.value.i = VA_SURFACE_ATTRIB_USAGE_HINT_EXPORT | VA_SURFACE_ATTRIB_USAGE_HINT_ENCODER;
		va_attribs[2].type = VASurfaceAttribPixelFormat;
		va_attribs[2].flags = VA_SURFACE_ATTRIB_SETTABLE;
		va_attribs[2].value.type = VAGenericValueTypeInteger;
		va_attribs[2].value.value.i = fourcc;
		va_attribs[3].type = VASurfaceAttribExternalBufferDescriptor;
		va_attribs[3].flags = VA_SURFACE_ATTRIB_SETTABLE;
		va_attribs[3].value.type = VAGenericValueTypePointer;
		va_attribs[3].value.value.p = &va_attrib_extbuf;
		va_attribs[4].type = VASurfaceAttribDRMFormatModifiers;
		va_attribs[4].flags = VA_SURFACE_ATTRIB_SETTABLE;
		va_attribs[4].value.type = VAGenericValueTypePointer;
		va_attribs[4].value.value.p = &modList;

		status = vaCreateSurfaces(dpy, format,
								  width, height, inputFrames, CHAIN_SIZE,
								  &va_attribs[0], 5);

		for(int i = 0; i < CHAIN_SIZE; i++)
		{
			status = vaExportSurfaceHandle(dpy, inputFrames[i], VA_SURFACE_ATTRIB_MEM_TYPE_DRM_PRIME_2, VA_EXPORT_SURFACE_WRITE_ONLY | VA_EXPORT_SURFACE_SEPARATE_LAYERS, &drmSurface );
			//printf("%d %d %llx\n", status, drmSurface.objects[0].fd, drmSurface.objects[0].drm_format_modifier);
			dmabuf_fd[i] = drmSurface.objects[0].fd;
			if(status != VA_STATUS_SUCCESS)
			{
				return false;
			}
		}
		*mod = drmSurface.objects[0].drm_format_modifier;
		*size = drmSurface.objects[0].size;
		*offset = drmSurface.layers[1].offset[0];
		*pitch1 = drmSurface.layers[0].pitch[0];
		*pitch2 = drmSurface.layers[1].pitch[0];
		pthread_mutex_init(&mutex, NULL);
		pthread_cond_init(&output_cond, NULL);
		pthread_create(&output_thread, NULL, &OutputThread, this);
		for (int i = 0; i < VAConfigAttribTypeMax; i++)
			attrib[i].type = (VAConfigAttribType)i;
		vaGetConfigAttributes(dpy, profile, VAEntrypointEncSlice, attrib, VAConfigAttribTypeMax);

		return true;
	}

	void PushOutput(VABufferID buf)
	{
		pthread_mutex_lock(&mutex);

		/* The mutex is never released while encoding, so this point should
		 * never be reached if input.valid is true. */
		assert(output_buf == VA_INVALID_ID);

		output_buf = buf;
		pthread_cond_signal(&output_cond);
		pthread_mutex_unlock(&mutex);
	}

	static void *OutputThread(void *data)
	{
		VaapiEncoder *r = (VaapiEncoder*)data;
		pthread_mutex_lock(&r->mutex);

		while (!r->destroying) {
			if (r->output_buf == VA_INVALID_ID)
				pthread_cond_wait(&r->output_cond, &r->mutex);

			/* If the thread is awaken by destroy_worker_thread(),
			 * there might not be valid input */
			if (r->output_buf == VA_INVALID_ID)
				continue;

			vaSyncBuffer(r->dpy, r->output_buf, UINT64_MAX);
			r->WriteOutput(r->output_buf);
			vaDestroyBuffer(r->dpy, r->output_buf);
			r->output_buf = VA_INVALID_ID;
		}

		pthread_mutex_unlock(&r->mutex);
		return NULL;
	}

	bool WriteOutput(VABufferID buf)
	{
		VACodedBufferSegment *segment;
		VAStatus status;
		int count = 0;

		status = vaMapBuffer(dpy, buf, (void **) &segment);
		if (status != VA_STATUS_SUCCESS)
			return false;

		do
		{
#if 0
			if (segment->status & VA_CODED_BUF_STATUS_SLICE_OVERFLOW_MASK) {
				vaUnmapBuffer(dpy, buf);
				printf("overflow!\n");

				// todo: warning or mark for restart in lq?
				// some mesa versions seems to cause false-overflow here!
				return false;
			}
#endif
			assert(segment->size);
			count += write(output_fd, segment->buf, segment->size);
			segment = (VACodedBufferSegment *)segment->next;
		}
		while(segment);

		vaUnmapBuffer(dpy, buf);

		if (count < 0)
			return false;

		return true;
	}

	template <int maxSize>
	bool CreatePackedBuffer(VABufferID &par, VABufferID &dat, VAEncPackedHeaderType type, const BaseBitstream<maxSize> &buf )
	{
		VAEncPackedHeaderParameterBuffer packed_header;
		VAStatus status;

		packed_header.type = type;
		packed_header.bit_length = buf.bit_offset;
		packed_header.has_emulation_bytes = 0;

		status = vaCreateBuffer(dpy, ctx,
								VAEncPackedHeaderParameterBufferType,
								sizeof packed_header, 1, &packed_header,
								&par);
		if (status != VA_STATUS_SUCCESS)
			return false;

		status = vaCreateBuffer(dpy, ctx,
								VAEncPackedHeaderDataBufferType,
								(buf.bit_offset + 7) / 8, 1, (void*)buf.buffer, &dat);
		if (status != VA_STATUS_SUCCESS) {
			vaDestroyBuffer(dpy, par);
			return false;
		}

		return true;
	}
	bool CreateContext(VAProfile profile, uint32_t format, uint32_t fourcc, uint32_t rc)
	{
		VASurfaceAttrib attrs[2] = { {VASurfaceAttribMemoryType, VA_SURFACE_ATTRIB_SETTABLE,{VAGenericValueTypeInteger, 0},},
									{VASurfaceAttribPixelFormat,VA_SURFACE_ATTRIB_SETTABLE,{VAGenericValueTypeInteger, 0},} };
		VAConfigAttrib cfg_attrib[2];
		VAStatus status;
		attrs[0].value.value.i = VA_SURFACE_ATTRIB_MEM_TYPE_VA;
		attrs[1].value.value.i = fourcc;
		status = vaCreateSurfaces(dpy, format,
								  width, height,
								  reference_picture, 3,
								  attrs, 2);
		if(status != VA_STATUS_SUCCESS)
			return false;
		cfg_attrib[0].type = VAConfigAttribRTFormat;
		cfg_attrib[0].value = format;

		cfg_attrib[1].type = VAConfigAttribRateControl;
		cfg_attrib[1].value = rc;
		status = vaCreateConfig(dpy, profile,
								VAEntrypointEncSlice, cfg_attrib, 2,
								&cfg);
		if(status != VA_STATUS_SUCCESS)
			return false;

		status = vaCreateContext(dpy, cfg, width, height, VA_PROGRESSIVE, NULL, 0, &ctx);
		if(status != VA_STATUS_SUCCESS)
		{
			vaDestroyConfig(dpy, cfg);
			cfg = VA_INVALID_ID;
			return false;
		}
		return true;
	}
	VABufferID CreateOutputBuf(int output_size)
	{
		VABufferID output_buf;
		VAStatus status;

		status = vaCreateBuffer(dpy, ctx,
								VAEncCodedBufferType, output_size,
								1, NULL, &output_buf);
		if (status == VA_STATUS_SUCCESS)
			return output_buf;
		abort();
		return VA_INVALID_ID;
	}

	template <typename T>
	VABufferID CreateMiscParameterBuffer(VAEncMiscParameterType tp, const T &data)
	{
		struct d{
			VAEncMiscParameterType type;
			union{
			T data;
			uint32_t pad[1];
			} d;
		} buf;
		buf.type = tp;
		buf.d.data = data;
		assert(sizeof(buf) == sizeof(VAEncMiscParameterBuffer) + sizeof(T));
		VABufferID buffer;
		if( vaCreateBuffer(dpy, ctx, VAEncMiscParameterBufferType, sizeof(buf), 1, &buf, &buffer) == VA_STATUS_SUCCESS)
			return buffer;
		abort();
		return VA_INVALID_ID;
	}
	template <typename T>
	VABufferID CreateParamererBuffer(VABufferType tp, const T &data)
	{
		VABufferID buffer;
		if( vaCreateBuffer(dpy, ctx, tp, sizeof(data), 1, (void*)&data, &buffer) == VA_STATUS_SUCCESS)
			return buffer;
		abort();
		return VA_INVALID_ID;
	}
	~VaapiEncoder()
	{
		Destroy();
	}
};

#endif // VAAPI_ENCODER_H