crc_work);
        struct drm_crtc *crtc = crtc_state->base.crtc;
        struct vkms_output *out = drm_crtc_to_vkms_output(crtc);
-       struct vkms_device *vdev = container_of(out, struct vkms_device,
-                                               output);
        struct vkms_crc_data *primary_crc = NULL;
        struct vkms_crc_data *cursor_crc = NULL;
-       struct drm_plane *plane;
        u32 crc32 = 0;
        u64 frame_start, frame_end;
        bool crc_pending;
        if (!crc_pending)
                return;
 
-       drm_for_each_plane(plane, &vdev->drm) {
-               struct vkms_plane_state *vplane_state;
-               struct vkms_crc_data *crc_data;
+       if (crtc_state->num_active_planes >= 1)
+               primary_crc = crtc_state->active_planes[0]->crc_data;
 
-               vplane_state = to_vkms_plane_state(plane->state);
-               crc_data = vplane_state->crc_data;
-
-               if (drm_framebuffer_read_refcount(&crc_data->fb) == 0)
-                       continue;
-
-               if (plane->type == DRM_PLANE_TYPE_PRIMARY)
-                       primary_crc = crc_data;
-               else
-                       cursor_crc = crc_data;
-       }
+       if (crtc_state->num_active_planes == 2)
+               cursor_crc = crtc_state->active_planes[1]->crc_data;
 
        if (primary_crc)
                crc32 = _vkms_get_crc(primary_crc, cursor_crc);
 
 // SPDX-License-Identifier: GPL-2.0+
 
 #include "vkms_drv.h"
+#include <drm/drm_atomic.h>
 #include <drm/drm_atomic_helper.h>
 #include <drm/drm_probe_helper.h>
 
        struct vkms_output *output = container_of(timer, struct vkms_output,
                                                  vblank_hrtimer);
        struct drm_crtc *crtc = &output->crtc;
-       struct vkms_crtc_state *state = to_vkms_crtc_state(crtc->state);
+       struct vkms_crtc_state *state;
        u64 ret_overrun;
        bool ret;
 
        if (!ret)
                DRM_ERROR("vkms failure on handling vblank");
 
+       state = output->crc_state;
        if (state && output->crc_enabled) {
                u64 frame = drm_crtc_accurate_vblank_count(crtc);
 
 
        __drm_atomic_helper_crtc_destroy_state(state);
 
-       if (vkms_state) {
-               WARN_ON(work_pending(&vkms_state->crc_work));
-               kfree(vkms_state);
-       }
+       WARN_ON(work_pending(&vkms_state->crc_work));
+       kfree(vkms_state->active_planes);
+       kfree(vkms_state);
 }
 
 static void vkms_atomic_crtc_reset(struct drm_crtc *crtc)
        .verify_crc_source      = vkms_verify_crc_source,
 };
 
+static int vkms_crtc_atomic_check(struct drm_crtc *crtc,
+                                 struct drm_crtc_state *state)
+{
+       struct vkms_crtc_state *vkms_state = to_vkms_crtc_state(state);
+       struct drm_plane *plane;
+       struct drm_plane_state *plane_state;
+       int i = 0, ret;
+
+       if (vkms_state->active_planes)
+               return 0;
+
+       ret = drm_atomic_add_affected_planes(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,
+                                                                 plane);
+               WARN_ON(!plane_state);
+
+               if (!plane_state->visible)
+                       continue;
+
+               i++;
+       }
+
+       vkms_state->active_planes = kcalloc(i, sizeof(plane), GFP_KERNEL);
+       if (!vkms_state->active_planes)
+               return -ENOMEM;
+       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,
+                                                                 plane);
+
+               if (!plane_state->visible)
+                       continue;
+
+               vkms_state->active_planes[i++] =
+                       to_vkms_plane_state(plane_state);
+       }
+
+       return 0;
+}
+
 static void vkms_crtc_atomic_enable(struct drm_crtc *crtc,
                                    struct drm_crtc_state *old_state)
 {
                crtc->state->event = NULL;
        }
 
+       vkms_output->crc_state = to_vkms_crtc_state(crtc->state);
+
        spin_unlock_irq(&vkms_output->lock);
 }
 
 static const struct drm_crtc_helper_funcs vkms_crtc_helper_funcs = {
+       .atomic_check   = vkms_crtc_atomic_check,
        .atomic_begin   = vkms_crtc_atomic_begin,
        .atomic_flush   = vkms_crtc_atomic_flush,
        .atomic_enable  = vkms_crtc_atomic_enable,
 
        struct drm_crtc_state base;
        struct work_struct crc_work;
 
+       int num_active_planes;
+       /* stack of active planes for crc computation, should be in z order */
+       struct vkms_plane_state **active_planes;
+
        /* below three are protected by vkms_output.crc_lock */
        bool crc_pending;
        u64 frame_start;
        struct hrtimer vblank_hrtimer;
        ktime_t period_ns;
        struct drm_pending_vblank_event *event;
-       bool crc_enabled;
        /* ordered wq for crc_work */
        struct workqueue_struct *crc_workq;
        /* protects concurrent access to crc_data */
        spinlock_t lock;
 
+       /* protected by @lock */
+       bool crc_enabled;
+       struct vkms_crtc_state *crc_state;
+
        spinlock_t crc_lock;
 };