diff --git a/drivers/gpu/drm/vkms/vkms_composer.c b/drivers/gpu/drm/vkms/vkms_composer.c
index 3ade0df173d25efb4aef6d5b232f00ea1470c99c..ead8fff81f30362ca7b3c05a96d5f4f686628188 100644
--- a/drivers/gpu/drm/vkms/vkms_composer.c
+++ b/drivers/gpu/drm/vkms/vkms_composer.c
@@ -6,7 +6,6 @@
 #include <drm/drm_atomic_helper.h>
 #include <drm/drm_fourcc.h>
 #include <drm/drm_gem_framebuffer_helper.h>
-#include <drm/drm_gem_shmem_helper.h>
 #include <drm/drm_vblank.h>
 
 #include "vkms_drv.h"
@@ -154,24 +153,21 @@ static void compose_plane(struct vkms_composer *primary_composer,
 			  struct vkms_composer *plane_composer,
 			  void *vaddr_out)
 {
-	struct drm_gem_object *plane_obj;
-	struct drm_gem_shmem_object *plane_shmem_obj;
 	struct drm_framebuffer *fb = &plane_composer->fb;
+	void *vaddr;
 	void (*pixel_blend)(const u8 *p_src, u8 *p_dst);
 
-	plane_obj = drm_gem_fb_get_obj(&plane_composer->fb, 0);
-	plane_shmem_obj = to_drm_gem_shmem_obj(plane_obj);
-
-	if (WARN_ON(!plane_shmem_obj->vaddr))
+	if (WARN_ON(dma_buf_map_is_null(&primary_composer->map[0])))
 		return;
 
+	vaddr = plane_composer->map[0].vaddr;
+
 	if (fb->format->format == DRM_FORMAT_ARGB8888)
 		pixel_blend = &alpha_blend;
 	else
 		pixel_blend = &x_blend;
 
-	blend(vaddr_out, plane_shmem_obj->vaddr, primary_composer,
-	      plane_composer, pixel_blend);
+	blend(vaddr_out, vaddr, primary_composer, plane_composer, pixel_blend);
 }
 
 static int compose_active_planes(void **vaddr_out,
@@ -180,21 +176,23 @@ static int compose_active_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 drm_gem_shmem_object *shmem_obj = to_drm_gem_shmem_obj(gem_obj);
+	const void *vaddr;
 	int i;
 
 	if (!*vaddr_out) {
-		*vaddr_out = kzalloc(shmem_obj->base.size, GFP_KERNEL);
+		*vaddr_out = kzalloc(gem_obj->size, GFP_KERNEL);
 		if (!*vaddr_out) {
 			DRM_ERROR("Cannot allocate memory for output frame.");
 			return -ENOMEM;
 		}
 	}
 
-	if (WARN_ON(!shmem_obj->vaddr))
+	if (WARN_ON(dma_buf_map_is_null(&primary_composer->map[0])))
 		return -EINVAL;
 
-	memcpy(*vaddr_out, shmem_obj->vaddr, shmem_obj->base.size);
+	vaddr = primary_composer->map[0].vaddr;
+
+	memcpy(*vaddr_out, vaddr, gem_obj->size);
 
 	/* If there are other planes besides primary, we consider the active
 	 * planes should be in z-order and compose them associatively:
diff --git a/drivers/gpu/drm/vkms/vkms_drv.h b/drivers/gpu/drm/vkms/vkms_drv.h
index 7a76dccd7316bdfe3185454f1f0a430bd48ef388..8c731b6dcba711101cfbcda9c85ec07d5ee0ca3b 100644
--- a/drivers/gpu/drm/vkms/vkms_drv.h
+++ b/drivers/gpu/drm/vkms/vkms_drv.h
@@ -23,6 +23,7 @@
 struct vkms_composer {
 	struct drm_framebuffer fb;
 	struct drm_rect src, dst;
+	struct dma_buf_map map[4];
 	unsigned int offset;
 	unsigned int pitch;
 	unsigned int cpp;
diff --git a/drivers/gpu/drm/vkms/vkms_plane.c b/drivers/gpu/drm/vkms/vkms_plane.c
index 8fbbd148163df3c9cdad709f4db5ae70897ee685..8a56fbf572b0c4f59c3c27b901dea4eac734328e 100644
--- a/drivers/gpu/drm/vkms/vkms_plane.c
+++ b/drivers/gpu/drm/vkms/vkms_plane.c
@@ -97,6 +97,7 @@ static void vkms_plane_atomic_update(struct drm_plane *plane,
 	struct drm_plane_state *new_state = drm_atomic_get_new_plane_state(state,
 									   plane);
 	struct vkms_plane_state *vkms_plane_state;
+	struct drm_shadow_plane_state *shadow_plane_state;
 	struct drm_framebuffer *fb = new_state->fb;
 	struct vkms_composer *composer;
 
@@ -104,11 +105,13 @@ static void vkms_plane_atomic_update(struct drm_plane *plane,
 		return;
 
 	vkms_plane_state = to_vkms_plane_state(new_state);
+	shadow_plane_state = &vkms_plane_state->base;
 
 	composer = vkms_plane_state->composer;
 	memcpy(&composer->src, &new_state->src, sizeof(struct drm_rect));
 	memcpy(&composer->dst, &new_state->dst, sizeof(struct drm_rect));
 	memcpy(&composer->fb, fb, sizeof(struct drm_framebuffer));
+	memcpy(&composer->map, &shadow_plane_state->map, sizeof(composer->map));
 	drm_framebuffer_get(&composer->fb);
 	composer->offset = fb->offsets[0];
 	composer->pitch = fb->pitches[0];