};
 EXPORT_SYMBOL_NS_GPL(ssd130x_variants, DRM_SSD130X);
 
+struct ssd130x_plane_state {
+       struct drm_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_plane_state *to_ssd130x_plane_state(struct drm_plane_state *state)
+{
+       return container_of(state, struct ssd130x_plane_state, base);
+}
+
 static inline struct ssd130x_device *drm_to_ssd130x(struct drm_device *drm)
 {
        return container_of(drm, struct ssd130x_device, drm);
                                 SSD130X_SET_ADDRESS_MODE_HORIZONTAL);
 }
 
-static int ssd130x_update_rect(struct ssd130x_device *ssd130x, struct drm_rect *rect)
+static int ssd130x_update_rect(struct ssd130x_device *ssd130x,
+                              struct ssd130x_plane_state *ssd130x_state,
+                              struct drm_rect *rect)
 {
        unsigned int x = rect->x1;
        unsigned int y = rect->y1;
-       u8 *buf = ssd130x->buffer;
-       u8 *data_array = ssd130x->data_array;
+       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)
+static void ssd130x_clear_screen(struct ssd130x_device *ssd130x,
+                                struct ssd130x_plane_state *ssd130x_state)
 {
        struct drm_rect fullscreen = {
                .x1 = 0,
                .y2 = ssd130x->height,
        };
 
-       ssd130x_update_rect(ssd130x, &fullscreen);
+       ssd130x_update_rect(ssd130x, ssd130x_state, &fullscreen);
 }
 
-static int ssd130x_fb_blit_rect(struct drm_framebuffer *fb, const struct iosys_map *vmap,
+static int ssd130x_fb_blit_rect(struct drm_plane_state *state,
+                               const struct iosys_map *vmap,
                                struct drm_rect *rect)
 {
+       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;
-       u8 *buf = ssd130x->buffer;
-
-       if (!buf)
-               return 0;
 
        /* Align y to display page boundaries */
        rect->y1 = round_down(rect->y1, page_height);
 
        drm_gem_fb_end_cpu_access(fb, DMA_FROM_DEVICE);
 
-       ssd130x_update_rect(ssd130x, rect);
+       ssd130x_update_rect(ssd130x, ssd130x_state, rect);
 
        return ret;
 }
 
+static int ssd130x_primary_plane_helper_atomic_check(struct drm_plane *plane,
+                                                    struct drm_atomic_state *state)
+{
+       struct drm_device *drm = plane->dev;
+       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);
+       const struct drm_format_info *fi;
+       unsigned int pitch;
+       int ret;
+
+       ret = drm_plane_helper_atomic_check(plane, state);
+       if (ret)
+               return ret;
+
+       fi = drm_format_info(DRM_FORMAT_R1);
+       if (!fi)
+               return -EINVAL;
+
+       pitch = drm_format_info_min_pitch(fi, 0, ssd130x->width);
+
+       ssd130x_state->buffer = kcalloc(pitch, ssd130x->height, GFP_KERNEL);
+       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;
+}
+
 static void ssd130x_primary_plane_helper_atomic_update(struct drm_plane *plane,
                                                       struct drm_atomic_state *state)
 {
                if (!drm_rect_intersect(&dst_clip, &damage))
                        continue;
 
-               ssd130x_fb_blit_rect(plane_state->fb, &shadow_plane_state->data[0], &dst_clip);
+               ssd130x_fb_blit_rect(plane_state, &shadow_plane_state->data[0], &dst_clip);
        }
 
        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);
        int idx;
 
        if (!drm_dev_enter(drm, &idx))
                return;
 
-       ssd130x_clear_screen(ssd130x);
+       ssd130x_clear_screen(ssd130x, ssd130x_state);
 
        drm_dev_exit(idx);
 }
 
+/* Called during init to allocate the plane's atomic state. */
+static void ssd130x_primary_plane_reset(struct drm_plane *plane)
+{
+       struct ssd130x_plane_state *ssd130x_state;
+
+       WARN_ON(plane->state);
+
+       ssd130x_state = kzalloc(sizeof(*ssd130x_state), GFP_KERNEL);
+       if (!ssd130x_state)
+               return;
+
+       __drm_atomic_helper_plane_reset(plane, &ssd130x_state->base);
+}
+
+static struct drm_plane_state *ssd130x_primary_plane_duplicate_state(struct drm_plane *plane)
+{
+       struct ssd130x_plane_state *old_ssd130x_state;
+       struct ssd130x_plane_state *ssd130x_state;
+
+       if (WARN_ON(!plane->state))
+               return NULL;
+
+       old_ssd130x_state = to_ssd130x_plane_state(plane->state);
+       ssd130x_state = kmemdup(old_ssd130x_state, sizeof(*ssd130x_state), GFP_KERNEL);
+       if (!ssd130x_state)
+               return NULL;
+
+       /* The buffers are not duplicated and are allocated in .atomic_check */
+       ssd130x_state->buffer = NULL;
+       ssd130x_state->data_array = NULL;
+
+       __drm_atomic_helper_plane_duplicate_state(plane, &ssd130x_state->base);
+
+       return &ssd130x_state->base;
+}
+
+static void ssd130x_primary_plane_destroy_state(struct drm_plane *plane,
+                                               struct drm_plane_state *state)
+{
+       struct ssd130x_plane_state *ssd130x_state = to_ssd130x_plane_state(state);
+
+       kfree(ssd130x_state->data_array);
+       kfree(ssd130x_state->buffer);
+
+       __drm_atomic_helper_plane_destroy_state(&ssd130x_state->base);
+
+       kfree(ssd130x_state);
+}
+
 static const struct drm_plane_helper_funcs ssd130x_primary_plane_helper_funcs = {
        DRM_GEM_SHADOW_PLANE_HELPER_FUNCS,
-       .atomic_check = drm_plane_helper_atomic_check,
+       .atomic_check = ssd130x_primary_plane_helper_atomic_check,
        .atomic_update = ssd130x_primary_plane_helper_atomic_update,
        .atomic_disable = ssd130x_primary_plane_helper_atomic_disable,
 };
 static const struct drm_plane_funcs ssd130x_primary_plane_funcs = {
        .update_plane = drm_atomic_helper_update_plane,
        .disable_plane = drm_atomic_helper_disable_plane,
+       .reset = ssd130x_primary_plane_reset,
+       .atomic_duplicate_state = ssd130x_primary_plane_duplicate_state,
+       .atomic_destroy_state = ssd130x_primary_plane_destroy_state,
        .destroy = drm_plane_cleanup,
        DRM_GEM_SHADOW_PLANE_FUNCS,
 };
 {
        struct drm_device *drm = encoder->dev;
        struct ssd130x_device *ssd130x = drm_to_ssd130x(drm);
-       unsigned int page_height = ssd130x->device_info->page_height;
-       unsigned int pages = DIV_ROUND_UP(ssd130x->height, page_height);
-       const struct drm_format_info *fi;
-       unsigned int pitch;
        int ret;
 
        ret = ssd130x_power_on(ssd130x);
        if (ret)
                goto power_off;
 
-       fi = drm_format_info(DRM_FORMAT_R1);
-       if (!fi)
-               goto power_off;
-
-       pitch = drm_format_info_min_pitch(fi, 0, ssd130x->width);
-
-       ssd130x->buffer = kcalloc(pitch, ssd130x->height, GFP_KERNEL);
-       if (!ssd130x->buffer)
-               goto power_off;
-
-       ssd130x->data_array = kcalloc(ssd130x->width, pages, GFP_KERNEL);
-       if (!ssd130x->data_array) {
-               kfree(ssd130x->buffer);
-               goto power_off;
-       }
-
        ssd130x_write_cmd(ssd130x, 1, SSD130X_DISPLAY_ON);
 
        backlight_enable(ssd130x->bl_dev);
 
        ssd130x_write_cmd(ssd130x, 1, SSD130X_DISPLAY_OFF);
 
-       kfree(ssd130x->data_array);
-       kfree(ssd130x->buffer);
-
        ssd130x_power_off(ssd130x);
 }