From 00c475e6465435939e2465e8847d97f6e807fcb3 Mon Sep 17 00:00:00 2001
From: columbarius <co1umbarius@protonmail.com>
Date: Thu, 10 Aug 2023 03:15:11 +0200
Subject: [PATCH] vulkan: download buffer to MemPtr

---
 spa/plugins/vulkan/vulkan-compute-utils.c | 27 +++++++++++++++++++++-
 spa/plugins/vulkan/vulkan-utils.c         | 28 +++++++++++++++++++++++
 spa/plugins/vulkan/vulkan-utils.h         | 10 ++++++++
 3 files changed, 64 insertions(+), 1 deletion(-)

diff --git a/spa/plugins/vulkan/vulkan-compute-utils.c b/spa/plugins/vulkan/vulkan-compute-utils.c
index dfd8f08d4..af7807e23 100644
--- a/spa/plugins/vulkan/vulkan-compute-utils.c
+++ b/spa/plugins/vulkan/vulkan-compute-utils.c
@@ -247,6 +247,30 @@ static int createCommandBuffer(struct vulkan_compute_state *s)
 	return 0;
 }
 
+static int runExportSHMBuffers(struct vulkan_compute_state *s) {
+	for (uint32_t i = 0; i < s->n_streams; i++) {
+		struct vulkan_stream *p = &s->streams[i];
+
+		if (p->direction == SPA_DIRECTION_INPUT)
+			continue;
+
+		if (p->spa_buffers[p->current_buffer_id]->datas[0].type == SPA_DATA_MemPtr) {
+			struct spa_buffer *spa_buf = p->spa_buffers[p->current_buffer_id];
+			struct vulkan_read_pixels_info readInfo = {
+				.data = spa_buf->datas[0].data,
+				.offset = spa_buf->datas[0].chunk->offset,
+				.stride = spa_buf->datas[0].chunk->stride,
+				.bytes_per_pixel = 16,
+				.size.width = s->constants.width,
+				.size.height = s->constants.height,
+			};
+			CHECK(vulkan_read_pixels(&s->base, &readInfo, &p->buffers[p->current_buffer_id]));
+		}
+	}
+
+	return 0;
+}
+
 /** runCommandBuffer
  *  The return value of this functions means the following:
  *  ret < 0: Error
@@ -610,6 +634,7 @@ int spa_vulkan_process(struct vulkan_compute_state *s)
 	CHECK(updateDescriptors(s));
 	CHECK(runCommandBuffer(s));
         VK_CHECK_RESULT(vkDeviceWaitIdle(s->base.device));
+	CHECK(runExportSHMBuffers(s));
 
 	return 0;
 }
@@ -620,7 +645,7 @@ int spa_vulkan_get_buffer_caps(struct vulkan_compute_state *s, enum spa_directio
 	case SPA_DIRECTION_INPUT:
 		return VULKAN_BUFFER_TYPE_CAP_DMABUF;
 	case SPA_DIRECTION_OUTPUT:
-		return VULKAN_BUFFER_TYPE_CAP_DMABUF;
+		return VULKAN_BUFFER_TYPE_CAP_DMABUF | VULKAN_BUFFER_TYPE_CAP_SHM;
 	}
 	return 0;
 }
diff --git a/spa/plugins/vulkan/vulkan-utils.c b/spa/plugins/vulkan/vulkan-utils.c
index 62e42eb84..3dfdc0d85 100644
--- a/spa/plugins/vulkan/vulkan-utils.c
+++ b/spa/plugins/vulkan/vulkan-utils.c
@@ -338,6 +338,34 @@ static void destroyFormatInfo(struct vulkan_base *s)
 	s->formatInfoCount = 0;
 }
 
+int vulkan_read_pixels(struct vulkan_base *s, struct vulkan_read_pixels_info *info, struct vulkan_buffer *vk_buf)
+{
+	VkImageSubresource img_sub_res = {
+		.aspectMask = VK_IMAGE_ASPECT_COLOR_BIT,
+		.arrayLayer = 0,
+		.mipLevel = 0,
+	};
+	VkSubresourceLayout img_sub_layout;
+	vkGetImageSubresourceLayout(s->device, vk_buf->image, &img_sub_res, &img_sub_layout);
+
+	void *v;
+	VK_CHECK_RESULT(vkMapMemory(s->device, vk_buf->memory, 0, VK_WHOLE_SIZE, 0, &v));
+
+	const char *d = (const char *)v + img_sub_layout.offset;
+	unsigned char *p = (unsigned char *)info->data + info->offset;
+	uint32_t bytes_per_pixel = 16;
+	uint32_t pack_stride = img_sub_layout.rowPitch;
+	if (pack_stride == info->stride) {
+		memcpy(p, d, info->stride * info->size.height);
+	} else {
+		for (uint32_t i = 0; i < info->size.height; i++) {
+			memcpy(p + i * info->stride, d + i * pack_stride, info->size.width * bytes_per_pixel);
+		}
+	}
+	vkUnmapMemory(s->device, vk_buf->memory);
+	return 0;
+}
+
 int vulkan_sync_foreign_dmabuf(struct vulkan_base *s, struct vulkan_buffer *vk_buf)
 {
 	VULKAN_INSTANCE_FUNCTION(vkImportSemaphoreFdKHR);
diff --git a/spa/plugins/vulkan/vulkan-utils.h b/spa/plugins/vulkan/vulkan-utils.h
index 7dac6bcbb..9f9611f95 100644
--- a/spa/plugins/vulkan/vulkan-utils.h
+++ b/spa/plugins/vulkan/vulkan-utils.h
@@ -46,6 +46,14 @@
 		return _res;								\
 }
 
+struct vulkan_read_pixels_info {
+	struct spa_rectangle size;
+	void *data;
+	uint32_t offset;
+	uint32_t stride;
+	uint32_t bytes_per_pixel;
+};
+
 struct dmabuf_fixation_info {
 	VkFormat format;
 	uint64_t modifierCount;
@@ -62,6 +70,8 @@ struct external_buffer_info {
 	struct spa_buffer *spa_buf;
 };
 
+int vulkan_read_pixels(struct vulkan_base *s, struct vulkan_read_pixels_info *info, struct vulkan_buffer *vk_buf);
+
 int vulkan_sync_foreign_dmabuf(struct vulkan_base *s, struct vulkan_buffer *vk_buf);
 bool vulkan_sync_export_dmabuf(struct vulkan_base *s, struct vulkan_buffer *vk_buf, int sync_file_fd);