};
 EXPORT_SYMBOL_NS_GPL(ssd130x_variants, DRM_SSD130X);
 
+struct ssd130x_crtc_state {
+       struct drm_crtc_state base;
+       /* Buffer to store pixels in HW format and written to the panel */
+       u8 *data_array;
+};
+
 struct ssd130x_plane_state {
        struct drm_shadow_plane_state base;
        /* Intermediate buffer to convert pixels from XRGB8888 to HW format */
        u8 *buffer;
-       /* Buffer to store pixels in HW format and written to the panel */
-       u8 *data_array;
 };
 
+static inline struct ssd130x_crtc_state *to_ssd130x_crtc_state(struct drm_crtc_state *state)
+{
+       return container_of(state, struct ssd130x_crtc_state, base);
+}
+
 static inline struct ssd130x_plane_state *to_ssd130x_plane_state(struct drm_plane_state *state)
 {
        return container_of(state, struct ssd130x_plane_state, base.base);
 }
 
 static int ssd130x_update_rect(struct ssd130x_device *ssd130x,
-                              struct ssd130x_plane_state *ssd130x_state,
-                              struct drm_rect *rect)
+                              struct drm_rect *rect, u8 *buf,
+                              u8 *data_array)
 {
        unsigned int x = rect->x1;
        unsigned int y = rect->y1;
-       u8 *buf = ssd130x_state->buffer;
-       u8 *data_array = ssd130x_state->data_array;
        unsigned int width = drm_rect_width(rect);
        unsigned int height = drm_rect_height(rect);
        unsigned int line_length = DIV_ROUND_UP(width, 8);
        return ret;
 }
 
-static void ssd130x_clear_screen(struct ssd130x_device *ssd130x,
-                                struct ssd130x_plane_state *ssd130x_state)
+static void ssd130x_clear_screen(struct ssd130x_device *ssd130x, u8 *data_array)
 {
        unsigned int page_height = ssd130x->device_info->page_height;
        unsigned int pages = DIV_ROUND_UP(ssd130x->height, page_height);
-       u8 *data_array = ssd130x_state->data_array;
        unsigned int width = ssd130x->width;
        int ret, i;
 
        }
 }
 
-static int ssd130x_fb_blit_rect(struct drm_plane_state *state,
+static int ssd130x_fb_blit_rect(struct drm_framebuffer *fb,
                                const struct iosys_map *vmap,
-                               struct drm_rect *rect)
+                               struct drm_rect *rect,
+                               u8 *buf, u8 *data_array)
 {
-       struct drm_framebuffer *fb = state->fb;
        struct ssd130x_device *ssd130x = drm_to_ssd130x(fb->dev);
        unsigned int page_height = ssd130x->device_info->page_height;
-       struct ssd130x_plane_state *ssd130x_state = to_ssd130x_plane_state(state);
-       u8 *buf = ssd130x_state->buffer;
        struct iosys_map dst;
        unsigned int dst_pitch;
        int ret = 0;
 
        drm_gem_fb_end_cpu_access(fb, DMA_FROM_DEVICE);
 
-       ssd130x_update_rect(ssd130x, ssd130x_state, rect);
+       ssd130x_update_rect(ssd130x, rect, buf, data_array);
 
        return ret;
 }
        struct ssd130x_device *ssd130x = drm_to_ssd130x(drm);
        struct drm_plane_state *plane_state = drm_atomic_get_new_plane_state(state, plane);
        struct ssd130x_plane_state *ssd130x_state = to_ssd130x_plane_state(plane_state);
-       unsigned int page_height = ssd130x->device_info->page_height;
-       unsigned int pages = DIV_ROUND_UP(ssd130x->height, page_height);
+       struct drm_crtc *crtc = plane_state->crtc;
+       struct drm_crtc_state *crtc_state;
        const struct drm_format_info *fi;
        unsigned int pitch;
        int ret;
 
+       if (!crtc)
+               return -EINVAL;
+
+       crtc_state = drm_atomic_get_crtc_state(state, crtc);
+       if (IS_ERR(crtc_state))
+               return PTR_ERR(crtc_state);
+
        ret = drm_plane_helper_atomic_check(plane, state);
        if (ret)
                return ret;
        if (!ssd130x_state->buffer)
                return -ENOMEM;
 
-       ssd130x_state->data_array = kcalloc(ssd130x->width, pages, GFP_KERNEL);
-       if (!ssd130x_state->data_array) {
-               kfree(ssd130x_state->buffer);
-               /* Set to prevent a double free in .atomic_destroy_state() */
-               ssd130x_state->buffer = NULL;
-               return -ENOMEM;
-       }
-
        return 0;
 }
 
        struct drm_plane_state *plane_state = drm_atomic_get_new_plane_state(state, plane);
        struct drm_plane_state *old_plane_state = drm_atomic_get_old_plane_state(state, plane);
        struct drm_shadow_plane_state *shadow_plane_state = to_drm_shadow_plane_state(plane_state);
+       struct drm_crtc_state *crtc_state = drm_atomic_get_new_crtc_state(state, plane_state->crtc);
+       struct ssd130x_crtc_state *ssd130x_crtc_state =  to_ssd130x_crtc_state(crtc_state);
+       struct ssd130x_plane_state *ssd130x_plane_state = to_ssd130x_plane_state(plane_state);
+       struct drm_framebuffer *fb = plane_state->fb;
        struct drm_atomic_helper_damage_iter iter;
        struct drm_device *drm = plane->dev;
        struct drm_rect dst_clip;
                if (!drm_rect_intersect(&dst_clip, &damage))
                        continue;
 
-               ssd130x_fb_blit_rect(plane_state, &shadow_plane_state->data[0], &dst_clip);
+               ssd130x_fb_blit_rect(fb, &shadow_plane_state->data[0], &dst_clip,
+                                    ssd130x_plane_state->buffer,
+                                    ssd130x_crtc_state->data_array);
        }
 
        drm_dev_exit(idx);
 {
        struct drm_device *drm = plane->dev;
        struct ssd130x_device *ssd130x = drm_to_ssd130x(drm);
-       struct ssd130x_plane_state *ssd130x_state = to_ssd130x_plane_state(plane->state);
+       struct drm_plane_state *plane_state = drm_atomic_get_new_plane_state(state, plane);
+       struct drm_crtc_state *crtc_state;
+       struct ssd130x_crtc_state *ssd130x_crtc_state;
        int idx;
 
+       if (!plane_state->crtc)
+               return;
+
+       crtc_state = drm_atomic_get_new_crtc_state(state, plane_state->crtc);
+       ssd130x_crtc_state = to_ssd130x_crtc_state(crtc_state);
+
        if (!drm_dev_enter(drm, &idx))
                return;
 
-       ssd130x_clear_screen(ssd130x, ssd130x_state);
+       ssd130x_clear_screen(ssd130x, ssd130x_crtc_state->data_array);
 
        drm_dev_exit(idx);
 }
        if (!ssd130x_state)
                return NULL;
 
-       /* The buffers are not duplicated and are allocated in .atomic_check */
+       /* The buffer is not duplicated and is allocated in .atomic_check */
        ssd130x_state->buffer = NULL;
-       ssd130x_state->data_array = NULL;
 
        new_shadow_plane_state = &ssd130x_state->base;
 
 {
        struct ssd130x_plane_state *ssd130x_state = to_ssd130x_plane_state(state);
 
-       kfree(ssd130x_state->data_array);
        kfree(ssd130x_state->buffer);
 
        __drm_gem_destroy_shadow_plane_state(&ssd130x_state->base);
        return MODE_OK;
 }
 
+static int ssd130x_crtc_helper_atomic_check(struct drm_crtc *crtc,
+                                           struct drm_atomic_state *state)
+{
+       struct drm_device *drm = crtc->dev;
+       struct ssd130x_device *ssd130x = drm_to_ssd130x(drm);
+       struct drm_crtc_state *crtc_state = drm_atomic_get_new_crtc_state(state, crtc);
+       struct ssd130x_crtc_state *ssd130x_state = to_ssd130x_crtc_state(crtc_state);
+       unsigned int page_height = ssd130x->device_info->page_height;
+       unsigned int pages = DIV_ROUND_UP(ssd130x->height, page_height);
+       int ret;
+
+       ret = drm_crtc_helper_atomic_check(crtc, state);
+       if (ret)
+               return ret;
+
+       ssd130x_state->data_array = kmalloc(ssd130x->width * pages, GFP_KERNEL);
+       if (!ssd130x_state->data_array)
+               return -ENOMEM;
+
+       return 0;
+}
+
+/* Called during init to allocate the CRTC's atomic state. */
+static void ssd130x_crtc_reset(struct drm_crtc *crtc)
+{
+       struct ssd130x_crtc_state *ssd130x_state;
+
+       WARN_ON(crtc->state);
+
+       ssd130x_state = kzalloc(sizeof(*ssd130x_state), GFP_KERNEL);
+       if (!ssd130x_state)
+               return;
+
+       __drm_atomic_helper_crtc_reset(crtc, &ssd130x_state->base);
+}
+
+static struct drm_crtc_state *ssd130x_crtc_duplicate_state(struct drm_crtc *crtc)
+{
+       struct ssd130x_crtc_state *old_ssd130x_state;
+       struct ssd130x_crtc_state *ssd130x_state;
+
+       if (WARN_ON(!crtc->state))
+               return NULL;
+
+       old_ssd130x_state = to_ssd130x_crtc_state(crtc->state);
+       ssd130x_state = kmemdup(old_ssd130x_state, sizeof(*ssd130x_state), GFP_KERNEL);
+       if (!ssd130x_state)
+               return NULL;
+
+       /* The buffer is not duplicated and is allocated in .atomic_check */
+       ssd130x_state->data_array = NULL;
+
+       __drm_atomic_helper_crtc_duplicate_state(crtc, &ssd130x_state->base);
+
+       return &ssd130x_state->base;
+}
+
+static void ssd130x_crtc_destroy_state(struct drm_crtc *crtc,
+                                      struct drm_crtc_state *state)
+{
+       struct ssd130x_crtc_state *ssd130x_state = to_ssd130x_crtc_state(state);
+
+       kfree(ssd130x_state->data_array);
+
+       __drm_atomic_helper_crtc_destroy_state(state);
+
+       kfree(ssd130x_state);
+}
+
 /*
  * The CRTC is always enabled. Screen updates are performed by
  * the primary plane's atomic_update function. Disabling clears
  */
 static const struct drm_crtc_helper_funcs ssd130x_crtc_helper_funcs = {
        .mode_valid = ssd130x_crtc_helper_mode_valid,
-       .atomic_check = drm_crtc_helper_atomic_check,
+       .atomic_check = ssd130x_crtc_helper_atomic_check,
 };
 
 static const struct drm_crtc_funcs ssd130x_crtc_funcs = {
-       .reset = drm_atomic_helper_crtc_reset,
+       .reset = ssd130x_crtc_reset,
        .destroy = drm_crtc_cleanup,
        .set_config = drm_atomic_helper_set_config,
        .page_flip = drm_atomic_helper_page_flip,
-       .atomic_duplicate_state = drm_atomic_helper_crtc_duplicate_state,
-       .atomic_destroy_state = drm_atomic_helper_crtc_destroy_state,
+       .atomic_duplicate_state = ssd130x_crtc_duplicate_state,
+       .atomic_destroy_state = ssd130x_crtc_destroy_state,
 };
 
 static void ssd130x_encoder_helper_atomic_enable(struct drm_encoder *encoder,