msm: mdss: delay iommu mapping of buffers until commit

Currently iommu map call is done in client's context during overlay
play, this requires iommu attachment to correctly map buffers.
This mapping can be delayed until the actual commit and avoid
unnecessary delays in client's context. Also, only unmap buffers if they
have been mapped, instead of checking if iommu is attached.

Modified to work one kernel 3.14, iommu changes.

Change-Id: I802ec83a5af52ce6f372e8bda3ceb2646fb9cc55
Signed-off-by: Terence Hampson <thampson@codeaurora.org>

8 files changed, 216 insertions, 122 deletions

diff --git a/drivers/video/fbdev/msm/mdss_fb.c b/drivers/video/fbdev/msm/mdss_fb.c
index 7c0c710b430d..9ee80c3520d3 100644
--- a/drivers/video/fbdev/msm/mdss_fb.c
+++ b/drivers/video/fbdev/msm/mdss_fb.c
@@ -1372,10 +1372,8 @@ void mdss_fb_free_fb_ion_memory(struct msm_fb_data_type *mfd)
 					DMA_BIDIRECTIONAL);
 		dma_buf_detach(mfd->fbmem_buf, mfd->fb_attachment);
 		dma_buf_put(mfd->fbmem_buf);
-
 	}
 
-	dma_buf_put(mfd->fbmem_buf);
 	ion_free(mfd->fb_ion_client, mfd->fb_ion_handle);
 	mfd->fb_ion_handle = NULL;
 }
@@ -1442,9 +1440,6 @@ int mdss_fb_alloc_fb_ion_memory(struct msm_fb_data_type *mfd, size_t fb_size)
 		goto fb_mmap_failed;
 	}
 
-	mfd->fbmem_buf = ion_share_dma_buf(mfd->fb_ion_client,
-			mfd->fb_ion_handle);
-
 	vaddr  = ion_map_kernel(mfd->fb_ion_client, mfd->fb_ion_handle);
 	if (IS_ERR_OR_NULL(vaddr)) {
 		pr_err("ION memory mapping failed - %ld\n", PTR_ERR(vaddr));
diff --git a/drivers/video/fbdev/msm/mdss_mdp.c b/drivers/video/fbdev/msm/mdss_mdp.c
index 654d02db3ec7..4076e71035c7 100644
--- a/drivers/video/fbdev/msm/mdss_mdp.c
+++ b/drivers/video/fbdev/msm/mdss_mdp.c
@@ -641,8 +641,11 @@ int mdss_iommu_ctrl(int enable)
 		__builtin_return_address(0), enable, mdata->iommu_ref_cnt);
 
 	if (enable) {
-
-		if (mdata->iommu_ref_cnt == 0)
+		/*
+		 * delay iommu attach until continous splash screen has
+		 * finished handoff, as it may still be working with phys addr
+		 */
+		if (!mdata->iommu_attached && !mdata->handoff_pending)
 			rc = mdss_iommu_attach(mdata);
 		mdata->iommu_ref_cnt++;
 	} else {
diff --git a/drivers/video/fbdev/msm/mdss_mdp.h b/drivers/video/fbdev/msm/mdss_mdp.h
index 20eda2b9d2d8..2ec029684b47 100644
--- a/drivers/video/fbdev/msm/mdss_mdp.h
+++ b/drivers/video/fbdev/msm/mdss_mdp.h
@@ -311,7 +311,9 @@ struct mdss_mdp_plane_sizes {
 struct mdss_mdp_img_data {
 	dma_addr_t addr;
 	unsigned long len;
+	u32 offset;
 	u32 flags;
+	bool mapped;
 	struct fd srcp_f;
 	struct dma_buf *srcp_dma_buf;
 	struct dma_buf_attachment *srcp_attachment;
@@ -759,11 +761,6 @@ int mdss_mdp_overlay_req_check(struct msm_fb_data_type *mfd,
 			       struct mdp_overlay *req,
 			       struct mdss_mdp_format_params *fmt);
 int mdss_mdp_overlay_vsync_ctrl(struct msm_fb_data_type *mfd, int en);
-int mdss_mdp_overlay_get_buf(struct msm_fb_data_type *mfd,
-			     struct mdss_mdp_data *data,
-			     struct msmfb_data *planes,
-			     int num_planes,
-			     u32 flags);
 int mdss_mdp_overlay_pipe_setup(struct msm_fb_data_type *mfd,
 	struct mdp_overlay *req, struct mdss_mdp_pipe **ppipe,
 	struct mdss_mdp_pipe *left_blend_pipe, bool is_single_layer);
@@ -927,10 +924,11 @@ int mdss_mdp_get_rau_strides(u32 w, u32 h, struct mdss_mdp_format_params *fmt,
 void mdss_mdp_data_calc_offset(struct mdss_mdp_data *data, u16 x, u16 y,
 	struct mdss_mdp_plane_sizes *ps, struct mdss_mdp_format_params *fmt);
 struct mdss_mdp_format_params *mdss_mdp_get_format_params(u32 format);
-int mdss_mdp_put_img(struct mdss_mdp_img_data *data);
-int mdss_mdp_get_img(struct msmfb_data *img, struct mdss_mdp_img_data *data,
-			struct device *dev);
-int mdss_mdp_overlay_free_buf(struct mdss_mdp_data *data);
+int mdss_mdp_data_get(struct mdss_mdp_data *data, struct msmfb_data *planes,
+		int num_planes, u32 flags, struct device *dev);
+int mdss_mdp_data_map(struct mdss_mdp_data *data);
+void mdss_mdp_data_free(struct mdss_mdp_data *data);
+
 u32 mdss_get_panel_framerate(struct msm_fb_data_type *mfd);
 int mdss_mdp_calc_phase_step(u32 src, u32 dst, u32 *out_phase);
 void mdss_mdp_intersect_rect(struct mdss_rect *res_rect,
diff --git a/drivers/video/fbdev/msm/mdss_mdp_overlay.c b/drivers/video/fbdev/msm/mdss_mdp_overlay.c
index 614407246022..c42ee755160f 100644
--- a/drivers/video/fbdev/msm/mdss_mdp_overlay.c
+++ b/drivers/video/fbdev/msm/mdss_mdp_overlay.c
@@ -1038,61 +1038,6 @@ static int mdss_mdp_overlay_set(struct msm_fb_data_type *mfd,
 	return ret;
 }
 
-int mdss_mdp_overlay_get_buf(struct msm_fb_data_type *mfd,
-					   struct mdss_mdp_data *data,
-					   struct msmfb_data *planes,
-					   int num_planes,
-					   u32 flags)
-{
-	int i, rc;
-
-	if ((num_planes <= 0) || (num_planes > MAX_PLANES))
-		return -EINVAL;
-
-	rc = mdss_iommu_ctrl(1);
-	if (IS_ERR_VALUE(rc)) {
-		pr_err("Iommu attach failed\n");
-		goto end;
-	}
-
-	memset(data, 0, sizeof(*data));
-	for (i = 0; i < num_planes; i++) {
-		data->p[i].flags = flags;
-		rc = mdss_mdp_get_img(&planes[i], &data->p[i], &mfd->pdev->dev);
-		if (rc) {
-			pr_err("failed to map buf p=%d flags=%x\n", i, flags);
-			while (i > 0) {
-				i--;
-				mdss_mdp_put_img(&data->p[i]);
-			}
-			break;
-		}
-	}
-
-	mdss_iommu_ctrl(0);
-	data->num_planes = i;
-end:
-	return rc;
-}
-
-int mdss_mdp_overlay_free_buf(struct mdss_mdp_data *data)
-{
-	int i, rc;
-
-	rc = mdss_iommu_ctrl(1);
-	if (IS_ERR_VALUE(rc)) {
-		pr_err("Iommu attach failed\n");
-		return rc;
-	}
-
-	for (i = 0; i < data->num_planes && data->p[i].len; i++)
-		mdss_mdp_put_img(&data->p[i]);
-
-	mdss_iommu_ctrl(0);
-	data->num_planes = 0;
-	return 0;
-}
-
 /**
  * __mdss_mdp_overlay_free_list_purge() - clear free list of buffers
  * @mfd:	Msm frame buffer data structure for the associated fb
@@ -1106,7 +1051,7 @@ static void __mdss_mdp_overlay_free_list_purge(struct msm_fb_data_type *mfd)
 
 	pr_debug("purging fb%d free list\n", mfd->index);
 	for (i = 0; i < mdp5_data->free_list_size; i++)
-		mdss_mdp_overlay_free_buf(&mdp5_data->free_list[i]);
+		mdss_mdp_data_free(&mdp5_data->free_list[i]);
 	mdp5_data->free_list_size = 0;
 }
 
@@ -1184,10 +1129,10 @@ static void mdss_mdp_overlay_cleanup(struct msm_fb_data_type *mfd)
 		 * soon as session is closed.
 		 */
 		if (pipe->flags & MDP_SECURE_DISPLAY_OVERLAY_SESSION)
-			mdss_mdp_overlay_free_buf(&pipe->front_buf);
+			mdss_mdp_data_free(&pipe->front_buf);
 		else
 			__mdss_mdp_overlay_free_list_add(mfd, &pipe->front_buf);
-		mdss_mdp_overlay_free_buf(&pipe->back_buf);
+		mdss_mdp_data_free(&pipe->back_buf);
 		list_del_init(&pipe->list);
 		if (recovery_mode) {
 			mdss_mdp_mixer_pipe_unstage(pipe, pipe->mixer_left);
@@ -1255,6 +1200,15 @@ int mdss_mdp_overlay_start(struct msm_fb_data_type *mfd)
 			mdss_mdp_batfet_ctrl(mdp5_data->mdata, true);
 		mdss_mdp_release_splash_pipe(mfd);
 		return 0;
+	} else if (mfd->panel_info->cont_splash_enabled) {
+		mutex_lock(&mdp5_data->list_lock);
+		rc = list_empty(&mdp5_data->pipes_used);
+		mutex_unlock(&mdp5_data->list_lock);
+		if (rc) {
+			pr_debug("empty kickoff on fb%d during cont splash\n",
+					mfd->index);
+			return 0;
+		}
 	}
 
 	pr_debug("starting fb%d overlay\n", mfd->index);
@@ -1390,6 +1344,8 @@ static int __overlay_queue_pipes(struct msm_fb_data_type *mfd)
 
 		if (pipe->back_buf.num_planes) {
 			buf = &pipe->back_buf;
+
+			ret = mdss_mdp_data_map(buf);
 		} else if (!pipe->params_changed) {
 			continue;
 		} else if (pipe->front_buf.num_planes) {
@@ -1405,7 +1361,9 @@ static int __overlay_queue_pipes(struct msm_fb_data_type *mfd)
 			buf = NULL;
 		}
 
-		ret = mdss_mdp_pipe_queue_data(pipe, buf);
+		if (!IS_ERR_VALUE(ret))
+			ret = mdss_mdp_pipe_queue_data(pipe, buf);
+
 		if (IS_ERR_VALUE(ret)) {
 			pr_warn("Unable to queue data for pnum=%d\n",
 					pipe->num);
@@ -1477,6 +1435,23 @@ int mdss_mdp_overlay_kickoff(struct msm_fb_data_type *mfd,
 	}
 
 	mutex_lock(&mdp5_data->ov_lock);
+	ret = mdss_mdp_overlay_start(mfd);
+	if (ret) {
+		pr_err("unable to start overlay %d (%d)\n", mfd->index, ret);
+		mutex_unlock(&mdp5_data->ov_lock);
+		if (ctl->shared_lock)
+			mutex_unlock(ctl->shared_lock);
+		return ret;
+	}
+
+	ret = mdss_iommu_ctrl(1);
+	if (IS_ERR_VALUE(ret)) {
+		pr_err("iommu attach failed rc=%d\n", ret);
+		mutex_unlock(&mdp5_data->ov_lock);
+		if (ctl->shared_lock)
+			mutex_unlock(ctl->shared_lock);
+		return ret;
+	}
 	mutex_lock(&mdp5_data->list_lock);
 
 	/*
@@ -1597,7 +1572,9 @@ commit_fail:
 	mutex_unlock(&mdp5_data->ov_lock);
 	if (ctl->shared_lock)
 		mutex_unlock(ctl->shared_lock);
+	mdss_iommu_ctrl(0);
 	ATRACE_END(__func__);
+
 	return ret;
 }
 
@@ -1796,13 +1773,13 @@ static int mdss_mdp_overlay_queue(struct msm_fb_data_type *mfd,
 	if (src_data->num_planes) {
 		pr_warn("dropped buffer pnum=%d play=%d addr=0x%pa\n",
 			pipe->num, pipe->play_cnt, &src_data->p[0].addr);
-		mdss_mdp_overlay_free_buf(src_data);
+		mdss_mdp_data_free(src_data);
 	}
 
-	ret = mdss_mdp_overlay_get_buf(mfd, src_data, &req->data, 1, flags);
-	if (IS_ERR_VALUE(ret)) {
+	ret = mdss_mdp_data_get(src_data, &req->data, 1, flags,
+		&mfd->pdev->dev);
+	if (IS_ERR_VALUE(ret))
 		pr_err("src_data pmem error\n");
-	}
 
 	mdss_mdp_pipe_unmap(pipe);
 
@@ -1826,12 +1803,6 @@ static int mdss_mdp_overlay_play(struct msm_fb_data_type *mfd,
 		goto done;
 	}
 
-	ret = mdss_mdp_overlay_start(mfd);
-	if (ret) {
-		pr_err("unable to start overlay %d (%d)\n", mfd->index, ret);
-		goto done;
-	}
-
 	if (req->id & MDSS_MDP_ROT_SESSION_MASK) {
 		ret = mdss_mdp_rotator_play(mfd, req);
 	} else if (req->id == BORDERFILL_NDX) {
diff --git a/drivers/video/fbdev/msm/mdss_mdp_rotator.c b/drivers/video/fbdev/msm/mdss_mdp_rotator.c
index 1cf1fba24d0c..33144c6ea3e6 100644
--- a/drivers/video/fbdev/msm/mdss_mdp_rotator.c
+++ b/drivers/video/fbdev/msm/mdss_mdp_rotator.c
@@ -731,33 +731,50 @@ int mdss_mdp_rotator_play(struct msm_fb_data_type *mfd,
 	if (!rot) {
 		pr_err("invalid session id=%x\n", req->id);
 		ret = -ENOENT;
-		goto dst_buf_fail;
+		goto session_fail;
 	}
 
 	memset(&src_buf, 0, sizeof(struct mdss_mdp_data));
 
 	flgs = rot->flags & MDP_SECURE_OVERLAY_SESSION;
 
+	mdss_iommu_ctrl(1);
 	ret = mdss_mdp_rotator_busy_wait_ex(rot);
 	if (ret) {
 		pr_err("rotator busy wait error\n");
 		goto dst_buf_fail;
 	}
 
-	ret = mdss_mdp_overlay_get_buf(mfd, &src_buf, &req->data, 1, flgs);
+	ret = mdss_mdp_data_get(&src_buf, &req->data, 1, flgs,
+		&mfd->pdev->dev);
 	if (ret) {
 		pr_err("src_data pmem error\n");
-		mdss_mdp_overlay_free_buf(&rot->src_buf);
 		goto dst_buf_fail;
 	}
-	mdss_mdp_overlay_free_buf(&rot->src_buf);
+
+	ret = mdss_mdp_data_map(&src_buf);
+	if (ret) {
+		pr_err("unable to map source buffer\n");
+		mdss_mdp_data_free(&src_buf);
+		goto dst_buf_fail;
+	}
+	mdss_mdp_data_free(&rot->src_buf);
 	memcpy(&rot->src_buf, &src_buf, sizeof(struct mdss_mdp_data));
 
-	mdss_mdp_overlay_free_buf(&rot->dst_buf);
-	ret = mdss_mdp_overlay_get_buf(mfd, &rot->dst_buf,
-			&req->dst_data, 1, flgs);
+	mdss_mdp_data_free(&rot->dst_buf);
+	ret = mdss_mdp_data_get(&rot->dst_buf, &req->dst_data, 1, flgs,
+		&mfd->pdev->dev);
 	if (ret) {
 		pr_err("dst_data pmem error\n");
+		mdss_mdp_data_free(&rot->src_buf);
+		goto dst_buf_fail;
+	}
+
+	ret = mdss_mdp_data_map(&rot->dst_buf);
+	if (ret) {
+		pr_err("unable to map destination buffer\n");
+		mdss_mdp_data_free(&rot->dst_buf);
+		mdss_mdp_data_free(&rot->src_buf);
 		goto dst_buf_fail;
 	}
 
@@ -767,6 +784,8 @@ int mdss_mdp_rotator_play(struct msm_fb_data_type *mfd,
 		pr_err("rotator queue error session id=%x\n", req->id);
 
 dst_buf_fail:
+	mdss_iommu_ctrl(0);
+session_fail:
 	mutex_unlock(&rotator_lock);
 	return ret;
 }
@@ -778,8 +797,8 @@ int mdss_mdp_rotator_unset(int ndx)
 	mutex_lock(&rotator_lock);
 	rot = mdss_mdp_rotator_session_get(ndx);
 	if (rot) {
-		mdss_mdp_overlay_free_buf(&rot->src_buf);
-		mdss_mdp_overlay_free_buf(&rot->dst_buf);
+		mdss_mdp_data_free(&rot->src_buf);
+		mdss_mdp_data_free(&rot->dst_buf);
 
 		rot->pid = 0;
 		ret = mdss_mdp_rotator_finish(rot);
diff --git a/drivers/video/fbdev/msm/mdss_mdp_splash_logo.c b/drivers/video/fbdev/msm/mdss_mdp_splash_logo.c
index c8f036211d3c..18a90006a31e 100644
--- a/drivers/video/fbdev/msm/mdss_mdp_splash_logo.c
+++ b/drivers/video/fbdev/msm/mdss_mdp_splash_logo.c
@@ -53,7 +53,7 @@ static int mdss_mdp_splash_alloc_memory(struct msm_fb_data_type *mfd,
 
 	handle = ion_alloc(mdata->iclient, size, SZ_4K,
 				ION_HEAP(ION_SYSTEM_HEAP_ID), 0);
-	if (IS_ERR(handle)) {
+	if (IS_ERR_OR_NULL(handle)) {
 		pr_err("ion memory allocation failed\n");
 		rc = PTR_RET(handle);
 		goto end;
@@ -102,7 +102,6 @@ static int mdss_mdp_splash_alloc_memory(struct msm_fb_data_type *mfd,
 	ion_free(mdata->iclient, handle);
 
 	return rc;
-
 kmap_err:
 	msm_unmap_dma_buf(sinfo->table,
 		mdss_get_iommu_domain(MDSS_IOMMU_DOMAIN_UNSECURE), 0);
diff --git a/drivers/video/fbdev/msm/mdss_mdp_util.c b/drivers/video/fbdev/msm/mdss_mdp_util.c
index 660622712875..e7f1a5ea3bd6 100644
--- a/drivers/video/fbdev/msm/mdss_mdp_util.c
+++ b/drivers/video/fbdev/msm/mdss_mdp_util.c
@@ -660,7 +660,7 @@ void mdss_mdp_data_calc_offset(struct mdss_mdp_data *data, u16 x, u16 y,
 	}
 }
 
-int mdss_mdp_put_img(struct mdss_mdp_img_data *data)
+static int mdss_mdp_put_img(struct mdss_mdp_img_data *data)
 {
 	struct ion_client *iclient = mdss_get_ionclient();
 	if (data->flags & MDP_MEMORY_ID_TYPE_FB) {
@@ -677,7 +677,7 @@ int mdss_mdp_put_img(struct mdss_mdp_img_data *data)
 			pr_err("invalid ion client\n");
 			return -ENOMEM;
 		} else {
-			if (is_mdss_iommu_attached()) {
+			if (data->mapped) {
 				int domain;
 				if (data->flags & MDP_SECURE_OVERLAY_SESSION)
 					domain = MDSS_IOMMU_DOMAIN_SECURE;
@@ -687,6 +687,7 @@ int mdss_mdp_put_img(struct mdss_mdp_img_data *data)
 				msm_unmap_dma_buf(data->srcp_table,
 					mdss_get_iommu_domain(domain), 0);
 
+				data->mapped = false;
 			}
 
 			dma_buf_unmap_attachment(data->srcp_attachment,
@@ -706,8 +707,8 @@ int mdss_mdp_put_img(struct mdss_mdp_img_data *data)
 	return 0;
 }
 
-int mdss_mdp_get_img(struct msmfb_data *img, struct mdss_mdp_img_data *data,
-			struct device *dev)
+static int mdss_mdp_get_img(struct msmfb_data *img,
+		struct mdss_mdp_img_data *data, struct device *dev)
 {
 	struct fd f;
 	int ret = -EINVAL;
@@ -719,7 +720,7 @@ int mdss_mdp_get_img(struct msmfb_data *img, struct mdss_mdp_img_data *data,
 	start = &data->addr;
 	len = &data->len;
 	data->flags |= img->flags;
-
+	data->offset = img->offset;
 	if (img->flags & MDP_MEMORY_ID_TYPE_FB) {
 		f = fdget(img->memory_id);
 		if (f.file == NULL) {
@@ -761,6 +762,47 @@ int mdss_mdp_get_img(struct msmfb_data *img, struct mdss_mdp_img_data *data,
 			goto err_detach;
 		}
 
+		data->addr = 0;
+		data->len = 0;
+		data->mapped = false;
+		/* return early, mapping will be done later */
+
+		return 0;
+	}
+
+	if (!*start) {
+		pr_err("start address is zero!\n");
+		mdss_mdp_put_img(data);
+		return -ENOMEM;
+	}
+
+	if (!ret && (data->offset < data->len)) {
+		data->addr += data->offset;
+		data->len -= data->offset;
+
+		pr_debug("mem=%d ihdl=%p buf=0x%pa len=0x%lu\n", img->memory_id,
+			 data->srcp_dma_buf, &data->addr, data->len);
+	} else {
+		mdss_mdp_put_img(data);
+		return ret ? : -EOVERFLOW;
+	}
+
+	return ret;
+err_detach:
+	dma_buf_detach(data->srcp_dma_buf, data->srcp_attachment);
+err_put:
+	dma_buf_put(data->srcp_dma_buf);
+	return ret;
+}
+
+static int mdss_mdp_map_buffer(struct mdss_mdp_img_data *data)
+{
+	int ret = -EINVAL;
+
+	if (data->addr && data->len)
+		return 0;
+
+	if (!IS_ERR_OR_NULL(data->srcp_dma_buf)) {
 		if (is_mdss_iommu_attached()) {
 			int domain;
 			if (data->flags & MDP_SECURE_OVERLAY_SESSION)
@@ -771,10 +813,12 @@ int mdss_mdp_get_img(struct msmfb_data *img, struct mdss_mdp_img_data *data,
 			ret = msm_map_dma_buf(data->srcp_dma_buf,
 					    data->srcp_table,
 					    mdss_get_iommu_domain(domain),
-					    0, SZ_4K, 0, start, len, 0, 0);
+					    0, SZ_4K, 0, &data->addr,
+					    &data->len, 0, 0);
+			data->mapped = true;
 		} else {
-			*start = sg_phys(data->srcp_table->sgl);
-			*len = data->srcp_table->sgl->length;
+			data->addr = sg_phys(data->srcp_table->sgl);
+			data->len = data->srcp_table->sgl->length;
 			ret = 0;
 		}
 
@@ -784,17 +828,17 @@ int mdss_mdp_get_img(struct msmfb_data *img, struct mdss_mdp_img_data *data,
 		}
 	}
 
-	if (!*start) {
+	if (!data->addr) {
 		pr_err("start address is zero!\n");
 		mdss_mdp_put_img(data);
 		return -ENOMEM;
 	}
 
-	if (!ret && (img->offset < data->len)) {
-		data->addr += img->offset;
-		data->len -= img->offset;
+	if (!ret && (data->offset < data->len)) {
+		data->addr += data->offset;
+		data->len -= data->offset;
 
-		pr_debug("mem=%d ihdl=%p buf=0x%pa len=%lu\n", img->memory_id,
+		pr_debug("ihdl=%p buf=0x%pa len=0x%lu\n",
 			 data->srcp_dma_buf, &data->addr, data->len);
 	} else {
 		mdss_mdp_put_img(data);
@@ -806,13 +850,69 @@ int mdss_mdp_get_img(struct msmfb_data *img, struct mdss_mdp_img_data *data,
 err_unmap:
 	dma_buf_unmap_attachment(data->srcp_attachment, data->srcp_table,
 				DMA_BIDIRECTIONAL);
-err_detach:
 	dma_buf_detach(data->srcp_dma_buf, data->srcp_attachment);
-err_put:
 	dma_buf_put(data->srcp_dma_buf);
 	return ret;
 }
 
+int mdss_mdp_data_get(struct mdss_mdp_data *data, struct msmfb_data *planes,
+		int num_planes, u32 flags, struct device *dev)
+{
+	int i, rc = 0;
+
+	if ((num_planes <= 0) || (num_planes > MAX_PLANES))
+		return -EINVAL;
+
+	for (i = 0; i < num_planes; i++) {
+		data->p[i].flags = flags;
+		rc = mdss_mdp_get_img(&planes[i], &data->p[i], dev);
+		if (rc) {
+			pr_err("failed to get buf p=%d flags=%x\n", i, flags);
+			while (i > 0) {
+				i--;
+				mdss_mdp_put_img(&data->p[i]);
+			}
+			break;
+		}
+	}
+
+	data->num_planes = i;
+
+	return rc;
+}
+
+int mdss_mdp_data_map(struct mdss_mdp_data *data)
+{
+	int i, rc = 0;
+
+	if (!data || !data->num_planes)
+		return -EINVAL;
+
+	for (i = 0; i < data->num_planes; i++) {
+		rc = mdss_mdp_map_buffer(&data->p[i]);
+		if (rc) {
+			pr_err("failed to map buf p=%d\n", i);
+			while (i > 0) {
+				i--;
+				mdss_mdp_put_img(&data->p[i]);
+			}
+			break;
+		}
+	}
+
+	return rc;
+}
+
+void mdss_mdp_data_free(struct mdss_mdp_data *data)
+{
+	int i;
+
+	for (i = 0; i < data->num_planes && data->p[i].len; i++)
+		mdss_mdp_put_img(&data->p[i]);
+
+	data->num_planes = 0;
+}
+
 int mdss_mdp_calc_phase_step(u32 src, u32 dst, u32 *out_phase)
 {
 	u32 unit, residue, result;
diff --git a/drivers/video/fbdev/msm/mdss_mdp_wb.c b/drivers/video/fbdev/msm/mdss_mdp_wb.c
index aa8ce7cb088f..6efbf97518a5 100644
--- a/drivers/video/fbdev/msm/mdss_mdp_wb.c
+++ b/drivers/video/fbdev/msm/mdss_mdp_wb.c
@@ -413,6 +413,7 @@ static struct mdss_mdp_wb_data *get_user_node(struct msm_fb_data_type *mfd,
 	struct mdss_mdp_wb *wb = mfd_to_wb(mfd);
 	struct mdss_mdp_wb_data *node;
 	struct mdss_mdp_img_data *buf;
+	u32 flags = 0;
 	int ret;
 
 	if (!list_empty(&wb->register_queue)) {
@@ -433,22 +434,29 @@ static struct mdss_mdp_wb_data *get_user_node(struct msm_fb_data_type *mfd,
 	}
 
 	node->user_alloc = true;
-	node->buf_data.num_planes = 1;
-	buf = &node->buf_data.p[0];
 	if (wb->is_secure)
-		buf->flags |= MDP_SECURE_OVERLAY_SESSION;
+		flags |= MDP_SECURE_OVERLAY_SESSION;
+
+	ret = mdss_mdp_data_get(&node->buf_data, data, 1, flags,
+		&mfd->pdev->dev);
+	if (IS_ERR_VALUE(ret)) {
+		pr_err("error getting buffer info\n");
+		goto register_fail;
+	}
 
 	ret = mdss_iommu_ctrl(1);
 	if (IS_ERR_VALUE(ret)) {
 		pr_err("IOMMU attach failed\n");
-		goto register_fail;
+		goto fail_freebuf;
 	}
-	ret = mdss_mdp_get_img(data, buf, &mfd->pdev->dev);
+
+	ret = mdss_mdp_data_map(&node->buf_data);
 	if (IS_ERR_VALUE(ret)) {
-		pr_err("error getting buffer info\n");
+		pr_err("error mapping buffer\n");
 		mdss_iommu_ctrl(0);
-		goto register_fail;
+		goto fail_freebuf;
 	}
+
 	mdss_iommu_ctrl(0);
 
 	memcpy(&node->buf_info, data, sizeof(*data));
@@ -456,14 +464,16 @@ static struct mdss_mdp_wb_data *get_user_node(struct msm_fb_data_type *mfd,
 	ret = mdss_mdp_wb_register_node(wb, node);
 	if (IS_ERR_VALUE(ret)) {
 		pr_err("error registering wb node\n");
-		goto register_fail;
+		goto fail_freebuf;
 	}
 
+	buf = &node->buf_data.p[0];
 	pr_debug("register node mem_id=%d offset=%u addr=0x%pa len=%lu\n",
 		 data->memory_id, data->offset, &buf->addr, buf->len);
 
 	return node;
-
+fail_freebuf:
+	mdss_mdp_data_free(&node->buf_data);
 register_fail:
 	kfree(node);
 	return NULL;
@@ -480,7 +490,7 @@ static void mdss_mdp_wb_free_node(struct mdss_mdp_wb_data *node)
 				node->buf_info.offset,
 				&buf->addr);
 
-		mdss_mdp_put_img(&node->buf_data.p[0]);
+		mdss_mdp_data_free(&node->buf_data);
 		node->user_alloc = false;
 	}
 }
@@ -765,7 +775,6 @@ int mdss_mdp_wb_ioctl_handler(struct msm_fb_data_type *mfd, u32 cmd,
 		break;
 	case MSMFB_WRITEBACK_TERMINATE:
 		ret = mdss_mdp_wb_terminate(mfd);
-		mdss_iommu_ctrl(0);
 		break;
 	case MSMFB_WRITEBACK_SET_MIRRORING_HINT:
 		if (!copy_from_user(&hint, arg, sizeof(hint))) {