8 files changed, 51 insertions, 329 deletions

diff --git a/drivers/gpu/drm/vkms/Makefile b/drivers/gpu/drm/vkms/Makefile
index 333d3cead0e3..72f779cbfedd 100644
--- a/drivers/gpu/drm/vkms/Makefile
+++ b/drivers/gpu/drm/vkms/Makefile
@@ -4,7 +4,6 @@ vkms-y := \
 	vkms_plane.o \
 	vkms_output.o \
 	vkms_crtc.o \
-	vkms_gem.o \
 	vkms_composer.o \
 	vkms_writeback.o
 
diff --git a/drivers/gpu/drm/vkms/vkms_composer.c b/drivers/gpu/drm/vkms/vkms_composer.c
index 33c031f27c2c..66c6842d70db 100644
--- a/drivers/gpu/drm/vkms/vkms_composer.c
+++ b/drivers/gpu/drm/vkms/vkms_composer.c
@@ -5,6 +5,7 @@
 #include <drm/drm_atomic.h>
 #include <drm/drm_atomic_helper.h>
 #include <drm/drm_gem_framebuffer_helper.h>
+#include <drm/drm_gem_shmem_helper.h>
 #include <drm/drm_vblank.h>
 
 #include "vkms_drv.h"
@@ -129,15 +130,15 @@ static void compose_cursor(struct vkms_composer *cursor_composer,
 			   void *vaddr_out)
 {
 	struct drm_gem_object *cursor_obj;
-	struct vkms_gem_object *cursor_vkms_obj;
+	struct drm_gem_shmem_object *cursor_shmem_obj;
 
 	cursor_obj = drm_gem_fb_get_obj(&cursor_composer->fb, 0);
-	cursor_vkms_obj = drm_gem_to_vkms_gem(cursor_obj);
+	cursor_shmem_obj = to_drm_gem_shmem_obj(cursor_obj);
 
-	if (WARN_ON(!cursor_vkms_obj->vaddr))
+	if (WARN_ON(!cursor_shmem_obj->vaddr))
 		return;
 
-	blend(vaddr_out, cursor_vkms_obj->vaddr,
+	blend(vaddr_out, cursor_shmem_obj->vaddr,
 	      primary_composer, cursor_composer);
 }
 
@@ -147,20 +148,20 @@ static int compose_planes(void **vaddr_out,
 {
 	struct drm_framebuffer *fb = &primary_composer->fb;
 	struct drm_gem_object *gem_obj = drm_gem_fb_get_obj(fb, 0);
-	struct vkms_gem_object *vkms_obj = drm_gem_to_vkms_gem(gem_obj);
+	struct drm_gem_shmem_object *shmem_obj = to_drm_gem_shmem_obj(gem_obj);
 
 	if (!*vaddr_out) {
-		*vaddr_out = kzalloc(vkms_obj->gem.size, GFP_KERNEL);
+		*vaddr_out = kzalloc(shmem_obj->base.size, GFP_KERNEL);
 		if (!*vaddr_out) {
 			DRM_ERROR("Cannot allocate memory for output frame.");
 			return -ENOMEM;
 		}
 	}
 
-	if (WARN_ON(!vkms_obj->vaddr))
+	if (WARN_ON(!shmem_obj->vaddr))
 		return -EINVAL;
 
-	memcpy(*vaddr_out, vkms_obj->vaddr, vkms_obj->gem.size);
+	memcpy(*vaddr_out, shmem_obj->vaddr, shmem_obj->base.size);
 
 	if (cursor_composer)
 		compose_cursor(cursor_composer, primary_composer, *vaddr_out);
diff --git a/drivers/gpu/drm/vkms/vkms_crtc.c b/drivers/gpu/drm/vkms/vkms_crtc.c
index 09c012d54d58..0443b7deeaef 100644
--- a/drivers/gpu/drm/vkms/vkms_crtc.c
+++ b/drivers/gpu/drm/vkms/vkms_crtc.c
@@ -168,9 +168,11 @@ static const struct drm_crtc_funcs vkms_crtc_funcs = {
 };
 
 static int vkms_crtc_atomic_check(struct drm_crtc *crtc,
-				  struct drm_crtc_state *state)
+				  struct drm_atomic_state *state)
 {
-	struct vkms_crtc_state *vkms_state = to_vkms_crtc_state(state);
+	struct drm_crtc_state *crtc_state = drm_atomic_get_new_crtc_state(state,
+									  crtc);
+	struct vkms_crtc_state *vkms_state = to_vkms_crtc_state(crtc_state);
 	struct drm_plane *plane;
 	struct drm_plane_state *plane_state;
 	int i = 0, ret;
@@ -178,12 +180,12 @@ static int vkms_crtc_atomic_check(struct drm_crtc *crtc,
 	if (vkms_state->active_planes)
 		return 0;
 
-	ret = drm_atomic_add_affected_planes(state->state, crtc);
+	ret = drm_atomic_add_affected_planes(crtc_state->state, crtc);
 	if (ret < 0)
 		return ret;
 
-	drm_for_each_plane_mask(plane, crtc->dev, state->plane_mask) {
-		plane_state = drm_atomic_get_existing_plane_state(state->state,
+	drm_for_each_plane_mask(plane, crtc->dev, crtc_state->plane_mask) {
+		plane_state = drm_atomic_get_existing_plane_state(crtc_state->state,
 								  plane);
 		WARN_ON(!plane_state);
 
@@ -199,8 +201,8 @@ static int vkms_crtc_atomic_check(struct drm_crtc *crtc,
 	vkms_state->num_active_planes = i;
 
 	i = 0;
-	drm_for_each_plane_mask(plane, crtc->dev, state->plane_mask) {
-		plane_state = drm_atomic_get_existing_plane_state(state->state,
+	drm_for_each_plane_mask(plane, crtc->dev, crtc_state->plane_mask) {
+		plane_state = drm_atomic_get_existing_plane_state(crtc_state->state,
 								  plane);
 
 		if (!plane_state->visible)
@@ -214,19 +216,19 @@ static int vkms_crtc_atomic_check(struct drm_crtc *crtc,
 }
 
 static void vkms_crtc_atomic_enable(struct drm_crtc *crtc,
-				    struct drm_crtc_state *old_state)
+				    struct drm_atomic_state *state)
 {
 	drm_crtc_vblank_on(crtc);
 }
 
 static void vkms_crtc_atomic_disable(struct drm_crtc *crtc,
-				     struct drm_crtc_state *old_state)
+				     struct drm_atomic_state *state)
 {
 	drm_crtc_vblank_off(crtc);
 }
 
 static void vkms_crtc_atomic_begin(struct drm_crtc *crtc,
-				   struct drm_crtc_state *old_crtc_state)
+				   struct drm_atomic_state *state)
 {
 	struct vkms_output *vkms_output = drm_crtc_to_vkms_output(crtc);
 
@@ -237,7 +239,7 @@ static void vkms_crtc_atomic_begin(struct drm_crtc *crtc,
 }
 
 static void vkms_crtc_atomic_flush(struct drm_crtc *crtc,
-				   struct drm_crtc_state *old_crtc_state)
+				   struct drm_atomic_state *state)
 {
 	struct vkms_output *vkms_output = drm_crtc_to_vkms_output(crtc);
 
diff --git a/drivers/gpu/drm/vkms/vkms_drv.c b/drivers/gpu/drm/vkms/vkms_drv.c
index cb0b6230c22c..1a1b5bc8e121 100644
--- a/drivers/gpu/drm/vkms/vkms_drv.c
+++ b/drivers/gpu/drm/vkms/vkms_drv.c
@@ -23,6 +23,7 @@
 #include <drm/drm_ioctl.h>
 #include <drm/drm_managed.h>
 #include <drm/drm_probe_helper.h>
+#include <drm/drm_gem_shmem_helper.h>
 #include <drm/drm_vblank.h>
 
 #include "vkms_drv.h"
@@ -39,23 +40,7 @@ bool enable_cursor = true;
 module_param_named(enable_cursor, enable_cursor, bool, 0444);
 MODULE_PARM_DESC(enable_cursor, "Enable/Disable cursor support");
 
-static const struct file_operations vkms_driver_fops = {
-	.owner		= THIS_MODULE,
-	.open		= drm_open,
-	.mmap		= drm_gem_mmap,
-	.unlocked_ioctl	= drm_ioctl,
-	.compat_ioctl	= drm_compat_ioctl,
-	.poll		= drm_poll,
-	.read		= drm_read,
-	.llseek		= no_llseek,
-	.release	= drm_release,
-};
-
-static const struct vm_operations_struct vkms_gem_vm_ops = {
-	.fault = vkms_gem_fault,
-	.open = drm_gem_vm_open,
-	.close = drm_gem_vm_close,
-};
+DEFINE_DRM_GEM_FOPS(vkms_driver_fops);
 
 static void vkms_release(struct drm_device *dev)
 {
@@ -93,15 +78,12 @@ static void vkms_atomic_commit_tail(struct drm_atomic_state *old_state)
 	drm_atomic_helper_cleanup_planes(dev, old_state);
 }
 
-static struct drm_driver vkms_driver = {
+static const struct drm_driver vkms_driver = {
 	.driver_features	= DRIVER_MODESET | DRIVER_ATOMIC | DRIVER_GEM,
 	.release		= vkms_release,
 	.fops			= &vkms_driver_fops,
-	.dumb_create		= vkms_dumb_create,
-	.gem_vm_ops		= &vkms_gem_vm_ops,
-	.gem_free_object_unlocked = vkms_gem_free_object,
-	.prime_fd_to_handle	= drm_gem_prime_fd_to_handle,
-	.gem_prime_import_sg_table = vkms_prime_import_sg_table,
+	.gem_create_object = drm_gem_shmem_create_object_cached,
+	DRM_GEM_SHMEM_DRIVER_OPS,
 
 	.name			= DRIVER_NAME,
 	.desc			= DRIVER_DESC,
@@ -132,7 +114,7 @@ static int vkms_modeset_init(struct vkms_device *vkmsdev)
 	dev->mode_config.max_height = YRES_MAX;
 	dev->mode_config.cursor_width = 512;
 	dev->mode_config.cursor_height = 512;
-	dev->mode_config.preferred_depth = 24;
+	dev->mode_config.preferred_depth = 32;
 	dev->mode_config.helper_private = &vkms_mode_config_helpers;
 
 	return vkms_output_init(vkmsdev, 0);
@@ -184,6 +166,8 @@ static int __init vkms_init(void)
 	if (ret)
 		goto out_devres;
 
+	drm_fbdev_generic_setup(&vkms_device->drm, 0);
+
 	return 0;
 
 out_devres:
diff --git a/drivers/gpu/drm/vkms/vkms_drv.h b/drivers/gpu/drm/vkms/vkms_drv.h
index 380a8f27e156..5ed91ff08cb3 100644
--- a/drivers/gpu/drm/vkms/vkms_drv.h
+++ b/drivers/gpu/drm/vkms/vkms_drv.h
@@ -88,23 +88,12 @@ struct vkms_device {
 	struct vkms_output output;
 };
 
-struct vkms_gem_object {
-	struct drm_gem_object gem;
-	struct mutex pages_lock; /* Page lock used in page fault handler */
-	struct page **pages;
-	unsigned int vmap_count;
-	void *vaddr;
-};
-
 #define drm_crtc_to_vkms_output(target) \
 	container_of(target, struct vkms_output, crtc)
 
 #define drm_device_to_vkms_device(target) \
 	container_of(target, struct vkms_device, drm)
 
-#define drm_gem_to_vkms_gem(target)\
-	container_of(target, struct vkms_gem_object, gem)
-
 #define to_vkms_crtc_state(target)\
 	container_of(target, struct vkms_crtc_state, base)
 
@@ -120,24 +109,6 @@ int vkms_output_init(struct vkms_device *vkmsdev, int index);
 struct drm_plane *vkms_plane_init(struct vkms_device *vkmsdev,
 				  enum drm_plane_type type, int index);
 
-/* Gem stuff */
-vm_fault_t vkms_gem_fault(struct vm_fault *vmf);
-
-int vkms_dumb_create(struct drm_file *file, struct drm_device *dev,
-		     struct drm_mode_create_dumb *args);
-
-void vkms_gem_free_object(struct drm_gem_object *obj);
-
-int vkms_gem_vmap(struct drm_gem_object *obj);
-
-void vkms_gem_vunmap(struct drm_gem_object *obj);
-
-/* Prime */
-struct drm_gem_object *
-vkms_prime_import_sg_table(struct drm_device *dev,
-			   struct dma_buf_attachment *attach,
-			   struct sg_table *sg);
-
 /* CRC Support */
 const char *const *vkms_get_crc_sources(struct drm_crtc *crtc,
 					size_t *count);
diff --git a/drivers/gpu/drm/vkms/vkms_gem.c b/drivers/gpu/drm/vkms/vkms_gem.c
deleted file mode 100644
index a017fc59905e..000000000000
--- a/drivers/gpu/drm/vkms/vkms_gem.c
+++ /dev/null
@@ -1,248 +0,0 @@
-// SPDX-License-Identifier: GPL-2.0+
-
-#include <linux/dma-buf.h>
-#include <linux/shmem_fs.h>
-#include <linux/vmalloc.h>
-#include <drm/drm_prime.h>
-
-#include "vkms_drv.h"
-
-static struct vkms_gem_object *__vkms_gem_create(struct drm_device *dev,
-						 u64 size)
-{
-	struct vkms_gem_object *obj;
-	int ret;
-
-	obj = kzalloc(sizeof(*obj), GFP_KERNEL);
-	if (!obj)
-		return ERR_PTR(-ENOMEM);
-
-	size = roundup(size, PAGE_SIZE);
-	ret = drm_gem_object_init(dev, &obj->gem, size);
-	if (ret) {
-		kfree(obj);
-		return ERR_PTR(ret);
-	}
-
-	mutex_init(&obj->pages_lock);
-
-	return obj;
-}
-
-void vkms_gem_free_object(struct drm_gem_object *obj)
-{
-	struct vkms_gem_object *gem = container_of(obj, struct vkms_gem_object,
-						   gem);
-
-	WARN_ON(gem->pages);
-	WARN_ON(gem->vaddr);
-
-	mutex_destroy(&gem->pages_lock);
-	drm_gem_object_release(obj);
-	kfree(gem);
-}
-
-vm_fault_t vkms_gem_fault(struct vm_fault *vmf)
-{
-	struct vm_area_struct *vma = vmf->vma;
-	struct vkms_gem_object *obj = vma->vm_private_data;
-	unsigned long vaddr = vmf->address;
-	pgoff_t page_offset;
-	loff_t num_pages;
-	vm_fault_t ret = VM_FAULT_SIGBUS;
-
-	page_offset = (vaddr - vma->vm_start) >> PAGE_SHIFT;
-	num_pages = DIV_ROUND_UP(obj->gem.size, PAGE_SIZE);
-
-	if (page_offset > num_pages)
-		return VM_FAULT_SIGBUS;
-
-	mutex_lock(&obj->pages_lock);
-	if (obj->pages) {
-		get_page(obj->pages[page_offset]);
-		vmf->page = obj->pages[page_offset];
-		ret = 0;
-	}
-	mutex_unlock(&obj->pages_lock);
-	if (ret) {
-		struct page *page;
-		struct address_space *mapping;
-
-		mapping = file_inode(obj->gem.filp)->i_mapping;
-		page = shmem_read_mapping_page(mapping, page_offset);
-
-		if (!IS_ERR(page)) {
-			vmf->page = page;
-			ret = 0;
-		} else {
-			switch (PTR_ERR(page)) {
-			case -ENOSPC:
-			case -ENOMEM:
-				ret = VM_FAULT_OOM;
-				break;
-			case -EBUSY:
-				ret = VM_FAULT_RETRY;
-				break;
-			case -EFAULT:
-			case -EINVAL:
-				ret = VM_FAULT_SIGBUS;
-				break;
-			default:
-				WARN_ON(PTR_ERR(page));
-				ret = VM_FAULT_SIGBUS;
-				break;
-			}
-		}
-	}
-	return ret;
-}
-
-static struct drm_gem_object *vkms_gem_create(struct drm_device *dev,
-					      struct drm_file *file,
-					      u32 *handle,
-					      u64 size)
-{
-	struct vkms_gem_object *obj;
-	int ret;
-
-	if (!file || !dev || !handle)
-		return ERR_PTR(-EINVAL);
-
-	obj = __vkms_gem_create(dev, size);
-	if (IS_ERR(obj))
-		return ERR_CAST(obj);
-
-	ret = drm_gem_handle_create(file, &obj->gem, handle);
-	if (ret)
-		return ERR_PTR(ret);
-
-	return &obj->gem;
-}
-
-int vkms_dumb_create(struct drm_file *file, struct drm_device *dev,
-		     struct drm_mode_create_dumb *args)
-{
-	struct drm_gem_object *gem_obj;
-	u64 pitch, size;
-
-	if (!args || !dev || !file)
-		return -EINVAL;
-
-	pitch = args->width * DIV_ROUND_UP(args->bpp, 8);
-	size = pitch * args->height;
-
-	if (!size)
-		return -EINVAL;
-
-	gem_obj = vkms_gem_create(dev, file, &args->handle, size);
-	if (IS_ERR(gem_obj))
-		return PTR_ERR(gem_obj);
-
-	args->size = gem_obj->size;
-	args->pitch = pitch;
-
-	drm_gem_object_put(gem_obj);
-
-	DRM_DEBUG_DRIVER("Created object of size %lld\n", size);
-
-	return 0;
-}
-
-static struct page **_get_pages(struct vkms_gem_object *vkms_obj)
-{
-	struct drm_gem_object *gem_obj = &vkms_obj->gem;
-
-	if (!vkms_obj->pages) {
-		struct page **pages = drm_gem_get_pages(gem_obj);
-
-		if (IS_ERR(pages))
-			return pages;
-
-		if (cmpxchg(&vkms_obj->pages, NULL, pages))
-			drm_gem_put_pages(gem_obj, pages, false, true);
-	}
-
-	return vkms_obj->pages;
-}
-
-void vkms_gem_vunmap(struct drm_gem_object *obj)
-{
-	struct vkms_gem_object *vkms_obj = drm_gem_to_vkms_gem(obj);
-
-	mutex_lock(&vkms_obj->pages_lock);
-	if (vkms_obj->vmap_count < 1) {
-		WARN_ON(vkms_obj->vaddr);
-		WARN_ON(vkms_obj->pages);
-		mutex_unlock(&vkms_obj->pages_lock);
-		return;
-	}
-
-	vkms_obj->vmap_count--;
-
-	if (vkms_obj->vmap_count == 0) {
-		vunmap(vkms_obj->vaddr);
-		vkms_obj->vaddr = NULL;
-		drm_gem_put_pages(obj, vkms_obj->pages, false, true);
-		vkms_obj->pages = NULL;
-	}
-
-	mutex_unlock(&vkms_obj->pages_lock);
-}
-
-int vkms_gem_vmap(struct drm_gem_object *obj)
-{
-	struct vkms_gem_object *vkms_obj = drm_gem_to_vkms_gem(obj);
-	int ret = 0;
-
-	mutex_lock(&vkms_obj->pages_lock);
-
-	if (!vkms_obj->vaddr) {
-		unsigned int n_pages = obj->size >> PAGE_SHIFT;
-		struct page **pages = _get_pages(vkms_obj);
-
-		if (IS_ERR(pages)) {
-			ret = PTR_ERR(pages);
-			goto out;
-		}
-
-		vkms_obj->vaddr = vmap(pages, n_pages, VM_MAP, PAGE_KERNEL);
-		if (!vkms_obj->vaddr)
-			goto err_vmap;
-	}
-
-	vkms_obj->vmap_count++;
-	goto out;
-
-err_vmap:
-	ret = -ENOMEM;
-	drm_gem_put_pages(obj, vkms_obj->pages, false, true);
-	vkms_obj->pages = NULL;
-out:
-	mutex_unlock(&vkms_obj->pages_lock);
-	return ret;
-}
-
-struct drm_gem_object *
-vkms_prime_import_sg_table(struct drm_device *dev,
-			   struct dma_buf_attachment *attach,
-			   struct sg_table *sg)
-{
-	struct vkms_gem_object *obj;
-	int npages;
-
-	obj = __vkms_gem_create(dev, attach->dmabuf->size);
-	if (IS_ERR(obj))
-		return ERR_CAST(obj);
-
-	npages = PAGE_ALIGN(attach->dmabuf->size) / PAGE_SIZE;
-	DRM_DEBUG_PRIME("Importing %d pages\n", npages);
-
-	obj->pages = kvmalloc_array(npages, sizeof(struct page *), GFP_KERNEL);
-	if (!obj->pages) {
-		vkms_gem_free_object(&obj->gem);
-		return ERR_PTR(-ENOMEM);
-	}
-
-	drm_prime_sg_to_page_addr_arrays(sg, obj->pages, NULL, npages);
-	return &obj->gem;
-}
diff --git a/drivers/gpu/drm/vkms/vkms_plane.c b/drivers/gpu/drm/vkms/vkms_plane.c
index 6d31265a2ab7..0824327cc860 100644
--- a/drivers/gpu/drm/vkms/vkms_plane.c
+++ b/drivers/gpu/drm/vkms/vkms_plane.c
@@ -1,10 +1,13 @@
 // SPDX-License-Identifier: GPL-2.0+
 
+#include <linux/dma-buf-map.h>
+
 #include <drm/drm_atomic.h>
 #include <drm/drm_atomic_helper.h>
 #include <drm/drm_fourcc.h>
 #include <drm/drm_gem_framebuffer_helper.h>
 #include <drm/drm_plane_helper.h>
+#include <drm/drm_gem_shmem_helper.h>
 
 #include "vkms_drv.h"
 
@@ -145,13 +148,14 @@ static int vkms_prepare_fb(struct drm_plane *plane,
 			   struct drm_plane_state *state)
 {
 	struct drm_gem_object *gem_obj;
+	struct dma_buf_map map;
 	int ret;
 
 	if (!state->fb)
 		return 0;
 
 	gem_obj = drm_gem_fb_get_obj(state->fb, 0);
-	ret = vkms_gem_vmap(gem_obj);
+	ret = drm_gem_shmem_vmap(gem_obj, &map);
 	if (ret)
 		DRM_ERROR("vmap failed: %d\n", ret);
 
@@ -162,12 +166,16 @@ static void vkms_cleanup_fb(struct drm_plane *plane,
 			    struct drm_plane_state *old_state)
 {
 	struct drm_gem_object *gem_obj;
+	struct drm_gem_shmem_object *shmem_obj;
+	struct dma_buf_map map;
 
 	if (!old_state->fb)
 		return;
 
 	gem_obj = drm_gem_fb_get_obj(old_state->fb, 0);
-	vkms_gem_vunmap(gem_obj);
+	shmem_obj = to_drm_gem_shmem_obj(drm_gem_fb_get_obj(old_state->fb, 0));
+	dma_buf_map_set_vaddr(&map, shmem_obj->vaddr);
+	drm_gem_shmem_vunmap(gem_obj, &map);
 }
 
 static const struct drm_plane_helper_funcs vkms_primary_helper_funcs = {
diff --git a/drivers/gpu/drm/vkms/vkms_writeback.c b/drivers/gpu/drm/vkms/vkms_writeback.c
index 094fa4aa061d..67f80ab1e85f 100644
--- a/drivers/gpu/drm/vkms/vkms_writeback.c
+++ b/drivers/gpu/drm/vkms/vkms_writeback.c
@@ -1,11 +1,15 @@
 // SPDX-License-Identifier: GPL-2.0+
 
-#include "vkms_drv.h"
+#include <linux/dma-buf-map.h>
+
 #include <drm/drm_fourcc.h>
 #include <drm/drm_writeback.h>
 #include <drm/drm_probe_helper.h>
 #include <drm/drm_atomic_helper.h>
 #include <drm/drm_gem_framebuffer_helper.h>
+#include <drm/drm_gem_shmem_helper.h>
+
+#include "vkms_drv.h"
 
 static const u32 vkms_wb_formats[] = {
 	DRM_FORMAT_XRGB8888,
@@ -63,22 +67,21 @@ static int vkms_wb_connector_get_modes(struct drm_connector *connector)
 static int vkms_wb_prepare_job(struct drm_writeback_connector *wb_connector,
 			       struct drm_writeback_job *job)
 {
-	struct vkms_gem_object *vkms_obj;
 	struct drm_gem_object *gem_obj;
+	struct dma_buf_map map;
 	int ret;
 
 	if (!job->fb)
 		return 0;
 
 	gem_obj = drm_gem_fb_get_obj(job->fb, 0);
-	ret = vkms_gem_vmap(gem_obj);
+	ret = drm_gem_shmem_vmap(gem_obj, &map);
 	if (ret) {
 		DRM_ERROR("vmap failed: %d\n", ret);
 		return ret;
 	}
 
-	vkms_obj = drm_gem_to_vkms_gem(gem_obj);
-	job->priv = vkms_obj->vaddr;
+	job->priv = map.vaddr;
 
 	return 0;
 }
@@ -88,12 +91,14 @@ static void vkms_wb_cleanup_job(struct drm_writeback_connector *connector,
 {
 	struct drm_gem_object *gem_obj;
 	struct vkms_device *vkmsdev;
+	struct dma_buf_map map;
 
 	if (!job->fb)
 		return;
 
 	gem_obj = drm_gem_fb_get_obj(job->fb, 0);
-	vkms_gem_vunmap(gem_obj);
+	dma_buf_map_set_vaddr(&map, job->priv);
+	drm_gem_shmem_vunmap(gem_obj, &map);
 
 	vkmsdev = drm_device_to_vkms_device(gem_obj->dev);
 	vkms_set_composer(&vkmsdev->output, false);