window.swap = state->swap;
 
        for (i = 0; i < fb->format->num_planes; i++) {
-               struct tegra_bo *bo = tegra_fb_get_plane(fb, i);
-
-               window.base[i] = bo->iova + fb->offsets[i];
+               window.base[i] = state->iova[i] + fb->offsets[i];
 
                /*
                 * Tegra uses a shared stride for UV planes. Framebuffers are
 }
 
 static const struct drm_plane_helper_funcs tegra_plane_helper_funcs = {
+       .prepare_fb = tegra_plane_prepare_fb,
+       .cleanup_fb = tegra_plane_cleanup_fb,
        .atomic_check = tegra_plane_atomic_check,
        .atomic_disable = tegra_plane_atomic_disable,
        .atomic_update = tegra_plane_atomic_update,
 }
 
 static const struct drm_plane_helper_funcs tegra_cursor_plane_helper_funcs = {
+       .prepare_fb = tegra_plane_prepare_fb,
+       .cleanup_fb = tegra_plane_cleanup_fb,
        .atomic_check = tegra_cursor_atomic_check,
        .atomic_update = tegra_cursor_atomic_update,
        .atomic_disable = tegra_cursor_atomic_disable,
 
        unsigned int zpos = plane->state->normalized_zpos;
        struct drm_framebuffer *fb = plane->state->fb;
        struct tegra_plane *p = to_tegra_plane(plane);
-       struct tegra_bo *bo;
        dma_addr_t base;
        u32 value;
 
        /* disable compression */
        tegra_plane_writel(p, 0, DC_WINBUF_CDE_CONTROL);
 
-       bo = tegra_fb_get_plane(fb, 0);
-       base = bo->iova;
+       base = state->iova[0] + fb->offsets[0];
 
        tegra_plane_writel(p, state->format, DC_WIN_COLOR_DEPTH);
        tegra_plane_writel(p, 0, DC_WIN_PRECOMP_WGRP_PARAMS);
 }
 
 static const struct drm_plane_helper_funcs tegra_shared_plane_helper_funcs = {
+       .prepare_fb = tegra_plane_prepare_fb,
+       .cleanup_fb = tegra_plane_cleanup_fb,
        .atomic_check = tegra_shared_plane_atomic_check,
        .atomic_update = tegra_shared_plane_atomic_update,
        .atomic_disable = tegra_shared_plane_atomic_disable,
 
 #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 "dc.h"
 {
        struct tegra_plane *p = to_tegra_plane(plane);
        struct tegra_plane_state *state;
+       unsigned int i;
 
        if (plane->state)
                __drm_atomic_helper_plane_destroy_state(plane->state);
                plane->state->plane = plane;
                plane->state->zpos = p->index;
                plane->state->normalized_zpos = p->index;
+
+               for (i = 0; i < 3; i++)
+                       state->iova[i] = DMA_MAPPING_ERROR;
        }
 }
 
        for (i = 0; i < 2; i++)
                copy->blending[i] = state->blending[i];
 
+       for (i = 0; i < 3; i++) {
+               copy->iova[i] = DMA_MAPPING_ERROR;
+               copy->sgt[i] = NULL;
+       }
+
        return ©->base;
 }
 
        .format_mod_supported = tegra_plane_format_mod_supported,
 };
 
+static int tegra_dc_pin(struct tegra_dc *dc, struct tegra_plane_state *state)
+{
+       unsigned int i;
+       int err;
+
+       for (i = 0; i < state->base.fb->format->num_planes; i++) {
+               struct tegra_bo *bo = tegra_fb_get_plane(state->base.fb, i);
+
+               if (!dc->client.group) {
+                       struct sg_table *sgt;
+
+                       sgt = host1x_bo_pin(dc->dev, &bo->base, NULL);
+                       if (IS_ERR(sgt)) {
+                               err = PTR_ERR(sgt);
+                               goto unpin;
+                       }
+
+                       err = dma_map_sg(dc->dev, sgt->sgl, sgt->nents,
+                                        DMA_TO_DEVICE);
+                       if (err == 0) {
+                               err = -ENOMEM;
+                               goto unpin;
+                       }
+
+                       state->iova[i] = sg_dma_address(sgt->sgl);
+                       state->sgt[i] = sgt;
+               } else {
+                       state->iova[i] = bo->iova;
+               }
+       }
+
+       return 0;
+
+unpin:
+       dev_err(dc->dev, "failed to map plane %u: %d\n", i, err);
+
+       while (i--) {
+               struct tegra_bo *bo = tegra_fb_get_plane(state->base.fb, i);
+               struct sg_table *sgt = state->sgt[i];
+
+               dma_unmap_sg(dc->dev, sgt->sgl, sgt->nents, DMA_TO_DEVICE);
+               host1x_bo_unpin(dc->dev, &bo->base, sgt);
+
+               state->iova[i] = DMA_MAPPING_ERROR;
+               state->sgt[i] = NULL;
+       }
+
+       return err;
+}
+
+static void tegra_dc_unpin(struct tegra_dc *dc, struct tegra_plane_state *state)
+{
+       unsigned int i;
+
+       for (i = 0; i < state->base.fb->format->num_planes; i++) {
+               struct tegra_bo *bo = tegra_fb_get_plane(state->base.fb, i);
+
+               if (!dc->client.group) {
+                       struct sg_table *sgt = state->sgt[i];
+
+                       if (sgt) {
+                               dma_unmap_sg(dc->dev, sgt->sgl, sgt->nents,
+                                            DMA_TO_DEVICE);
+                               host1x_bo_unpin(dc->dev, &bo->base, sgt);
+                       }
+               }
+
+               state->iova[i] = DMA_MAPPING_ERROR;
+               state->sgt[i] = NULL;
+       }
+}
+
+int tegra_plane_prepare_fb(struct drm_plane *plane,
+                          struct drm_plane_state *state)
+{
+       struct tegra_dc *dc = to_tegra_dc(state->crtc);
+
+       if (!state->fb)
+               return 0;
+
+       drm_gem_fb_prepare_fb(plane, state);
+
+       return tegra_dc_pin(dc, to_tegra_plane_state(state));
+}
+
+void tegra_plane_cleanup_fb(struct drm_plane *plane,
+                           struct drm_plane_state *state)
+{
+       struct tegra_dc *dc = to_tegra_dc(state->crtc);
+
+       if (dc)
+               tegra_dc_unpin(dc, to_tegra_plane_state(state));
+}
+
 int tegra_plane_state_add(struct tegra_plane *plane,
                          struct drm_plane_state *state)
 {