drm/msm: get an iova from the address space instead of an id

In the future we won't have a fixed set of addresses spaces.
Instead of going through the effort of assigning a ID for each
address space just use the address space itself as a token for
getting / putting an iova.

This forces a few changes in the gem object however: instead
of using a simple index into a list of domains, we need to
maintain a list of them. Luckily the list will be pretty small;
even with dynamic address spaces we wouldn't ever see more than
two or three.

CRs-Fixed: 2050484
Change-Id: Ic0dedbad4495f02a21135217f3605b93f8b8dfea
Signed-off-by: Jordan Crouse <jcrouse@codeaurora.org>
Signed-off-by: Abhijit Kulkarni <kabhijit@codeaurora.org>
diff --git a/drivers/gpu/drm/msm/sde/sde_connector.c b/drivers/gpu/drm/msm/sde/sde_connector.c
index 2970b28..c3c5a13 100644
--- a/drivers/gpu/drm/msm/sde/sde_connector.c
+++ b/drivers/gpu/drm/msm/sde/sde_connector.c
@@ -327,8 +327,7 @@
 		return;
 	}
 
-	msm_framebuffer_cleanup(c_state->out_fb,
-			c_state->mmu_id);
+	msm_framebuffer_cleanup(c_state->out_fb, c_state->aspace);
 	drm_framebuffer_unreference(c_state->out_fb);
 	c_state->out_fb = NULL;
 
@@ -432,7 +431,7 @@
 	if (c_state->out_fb) {
 		drm_framebuffer_reference(c_state->out_fb);
 		rc = msm_framebuffer_prepare(c_state->out_fb,
-				c_state->mmu_id);
+				c_state->aspace);
 		if (rc)
 			SDE_ERROR("failed to prepare fb, %d\n", rc);
 	}
@@ -652,14 +651,14 @@
 					c_conn->fb_kmap);
 
 			if (c_state->out_fb->flags & DRM_MODE_FB_SECURE)
-				c_state->mmu_id =
-				c_conn->mmu_id[SDE_IOMMU_DOMAIN_SECURE];
+				c_state->aspace =
+				c_conn->aspace[SDE_IOMMU_DOMAIN_SECURE];
 			else
-				c_state->mmu_id =
-				c_conn->mmu_id[SDE_IOMMU_DOMAIN_UNSECURE];
+				c_state->aspace =
+				c_conn->aspace[SDE_IOMMU_DOMAIN_UNSECURE];
 
 			rc = msm_framebuffer_prepare(c_state->out_fb,
-					c_state->mmu_id);
+					c_state->aspace);
 			if (rc)
 				SDE_ERROR("prep fb failed, %d\n", rc);
 		}
@@ -1010,18 +1009,17 @@
 	c_conn->lp_mode = 0;
 	c_conn->last_panel_power_mode = SDE_MODE_DPMS_ON;
 
-	/* cache mmu_id's for later */
 	sde_kms = to_sde_kms(priv->kms);
 	if (sde_kms->vbif[VBIF_NRT]) {
-		c_conn->mmu_id[SDE_IOMMU_DOMAIN_UNSECURE] =
-			sde_kms->mmu_id[MSM_SMMU_DOMAIN_NRT_UNSECURE];
-		c_conn->mmu_id[SDE_IOMMU_DOMAIN_SECURE] =
-			sde_kms->mmu_id[MSM_SMMU_DOMAIN_NRT_SECURE];
+		c_conn->aspace[SDE_IOMMU_DOMAIN_UNSECURE] =
+			sde_kms->aspace[MSM_SMMU_DOMAIN_NRT_UNSECURE];
+		c_conn->aspace[SDE_IOMMU_DOMAIN_SECURE] =
+			sde_kms->aspace[MSM_SMMU_DOMAIN_NRT_SECURE];
 	} else {
-		c_conn->mmu_id[SDE_IOMMU_DOMAIN_UNSECURE] =
-			sde_kms->mmu_id[MSM_SMMU_DOMAIN_UNSECURE];
-		c_conn->mmu_id[SDE_IOMMU_DOMAIN_SECURE] =
-			sde_kms->mmu_id[MSM_SMMU_DOMAIN_SECURE];
+		c_conn->aspace[SDE_IOMMU_DOMAIN_UNSECURE] =
+			sde_kms->aspace[MSM_SMMU_DOMAIN_UNSECURE];
+		c_conn->aspace[SDE_IOMMU_DOMAIN_SECURE] =
+			sde_kms->aspace[MSM_SMMU_DOMAIN_SECURE];
 	}
 
 	if (ops)
diff --git a/drivers/gpu/drm/msm/sde/sde_connector.h b/drivers/gpu/drm/msm/sde/sde_connector.h
index 497d0db..2318756 100644
--- a/drivers/gpu/drm/msm/sde/sde_connector.h
+++ b/drivers/gpu/drm/msm/sde/sde_connector.h
@@ -240,7 +240,7 @@
 	struct drm_panel *panel;
 	void *display;
 
-	int mmu_id[SDE_IOMMU_DOMAIN_MAX];
+	struct msm_gem_address_space *aspace[SDE_IOMMU_DOMAIN_MAX];
 
 	char name[SDE_CONNECTOR_NAME_SIZE];
 
@@ -304,14 +304,14 @@
  * struct sde_connector_state - private connector status structure
  * @base: Base drm connector structure
  * @out_fb: Pointer to output frame buffer, if applicable
- * @mmu_id: MMU ID for accessing frame buffer objects, if applicable
+ * @aspace: Address space for accessing frame buffer objects, if applicable
  * @property_values: Local cache of current connector property values
  * @rois: Regions of interest structure for mapping CRTC to Connector output
  */
 struct sde_connector_state {
 	struct drm_connector_state base;
 	struct drm_framebuffer *out_fb;
-	int mmu_id;
+	struct msm_gem_address_space *aspace;
 	uint64_t property_values[CONNECTOR_PROP_COUNT];
 
 	struct msm_roi_list rois;
diff --git a/drivers/gpu/drm/msm/sde/sde_encoder_phys.h b/drivers/gpu/drm/msm/sde/sde_encoder_phys.h
index b173876..4b12651 100644
--- a/drivers/gpu/drm/msm/sde/sde_encoder_phys.h
+++ b/drivers/gpu/drm/msm/sde/sde_encoder_phys.h
@@ -333,7 +333,7 @@
  * @wb_fmt:		Writeback pixel format
  * @frame_count:	Counter of completed writeback operations
  * @kickoff_count:	Counter of issued writeback operations
- * @mmu_id:		mmu identifier for non-secure/secure domain
+ * @aspace:		address space identifier for non-secure/secure domain
  * @wb_dev:		Pointer to writeback device
  * @start_time:		Start time of writeback latest request
  * @end_time:		End time of writeback latest request
@@ -355,7 +355,7 @@
 	const struct sde_format *wb_fmt;
 	u32 frame_count;
 	u32 kickoff_count;
-	int mmu_id[SDE_IOMMU_DOMAIN_MAX];
+	struct msm_gem_address_space *aspace[SDE_IOMMU_DOMAIN_MAX];
 	struct sde_wb_device *wb_dev;
 	ktime_t start_time;
 	ktime_t end_time;
diff --git a/drivers/gpu/drm/msm/sde/sde_encoder_phys_wb.c b/drivers/gpu/drm/msm/sde/sde_encoder_phys_wb.c
index 54c1397..875d99d 100644
--- a/drivers/gpu/drm/msm/sde/sde_encoder_phys_wb.c
+++ b/drivers/gpu/drm/msm/sde/sde_encoder_phys_wb.c
@@ -250,7 +250,8 @@
 	struct sde_hw_wb_cfg *wb_cfg;
 	struct sde_hw_wb_cdp_cfg *cdp_cfg;
 	const struct msm_format *format;
-	int ret, mmu_id;
+	int ret;
+	struct msm_gem_address_space *aspace;
 
 	if (!phys_enc || !phys_enc->sde_kms || !phys_enc->sde_kms->catalog) {
 		SDE_ERROR("invalid encoder\n");
@@ -264,9 +265,9 @@
 
 	wb_cfg->intf_mode = phys_enc->intf_mode;
 	wb_cfg->is_secure = (fb->flags & DRM_MODE_FB_SECURE) ? true : false;
-	mmu_id = (wb_cfg->is_secure) ?
-			wb_enc->mmu_id[SDE_IOMMU_DOMAIN_SECURE] :
-			wb_enc->mmu_id[SDE_IOMMU_DOMAIN_UNSECURE];
+	aspace = (wb_cfg->is_secure) ?
+			wb_enc->aspace[SDE_IOMMU_DOMAIN_SECURE] :
+			wb_enc->aspace[SDE_IOMMU_DOMAIN_UNSECURE];
 
 	SDE_DEBUG("[fb_secure:%d]\n", wb_cfg->is_secure);
 
@@ -288,7 +289,7 @@
 	wb_cfg->roi = *wb_roi;
 
 	if (hw_wb->caps->features & BIT(SDE_WB_XY_ROI_OFFSET)) {
-		ret = sde_format_populate_layout(mmu_id, fb, &wb_cfg->dest);
+		ret = sde_format_populate_layout(aspace, fb, &wb_cfg->dest);
 		if (ret) {
 			SDE_DEBUG("failed to populate layout %d\n", ret);
 			return;
@@ -297,7 +298,7 @@
 		wb_cfg->dest.height = fb->height;
 		wb_cfg->dest.num_planes = wb_cfg->dest.format->num_planes;
 	} else {
-		ret = sde_format_populate_layout_with_roi(mmu_id, fb, wb_roi,
+		ret = sde_format_populate_layout_with_roi(aspace, fb, wb_roi,
 			&wb_cfg->dest);
 		if (ret) {
 			/* this error should be detected during atomic_check */
@@ -914,12 +915,19 @@
 	struct drm_mode_fb_cmd2 mode_cmd;
 	uint32_t size;
 	int nplanes, i, ret;
+	struct msm_gem_address_space *aspace;
 
 	if (!wb_enc || !wb_enc->base.parent || !wb_enc->base.sde_kms) {
 		SDE_ERROR("invalid params\n");
 		return -EINVAL;
 	}
 
+	aspace = wb_enc->base.sde_kms->aspace[SDE_IOMMU_DOMAIN_UNSECURE];
+	if (!aspace) {
+		SDE_ERROR("invalid address space\n");
+		return -EINVAL;
+	}
+
 	dev = wb_enc->base.sde_kms->dev;
 	if (!dev) {
 		SDE_ERROR("invalid dev\n");
@@ -974,8 +982,7 @@
 	}
 
 	/* prepare the backing buffer now so that it's available later */
-	ret = msm_framebuffer_prepare(fb,
-			wb_enc->mmu_id[SDE_IOMMU_DOMAIN_UNSECURE]);
+	ret = msm_framebuffer_prepare(fb, aspace);
 	if (!ret)
 		wb_enc->fb_disable = fb;
 	return ret;
@@ -1234,15 +1241,15 @@
 	phys_enc = &wb_enc->base;
 
 	if (p->sde_kms->vbif[VBIF_NRT]) {
-		wb_enc->mmu_id[SDE_IOMMU_DOMAIN_UNSECURE] =
-			p->sde_kms->mmu_id[MSM_SMMU_DOMAIN_NRT_UNSECURE];
-		wb_enc->mmu_id[SDE_IOMMU_DOMAIN_SECURE] =
-			p->sde_kms->mmu_id[MSM_SMMU_DOMAIN_NRT_SECURE];
+		wb_enc->aspace[SDE_IOMMU_DOMAIN_UNSECURE] =
+			p->sde_kms->aspace[MSM_SMMU_DOMAIN_NRT_UNSECURE];
+		wb_enc->aspace[SDE_IOMMU_DOMAIN_SECURE] =
+			p->sde_kms->aspace[MSM_SMMU_DOMAIN_NRT_SECURE];
 	} else {
-		wb_enc->mmu_id[SDE_IOMMU_DOMAIN_UNSECURE] =
-			p->sde_kms->mmu_id[MSM_SMMU_DOMAIN_UNSECURE];
-		wb_enc->mmu_id[SDE_IOMMU_DOMAIN_SECURE] =
-			p->sde_kms->mmu_id[MSM_SMMU_DOMAIN_SECURE];
+		wb_enc->aspace[SDE_IOMMU_DOMAIN_UNSECURE] =
+			p->sde_kms->aspace[MSM_SMMU_DOMAIN_UNSECURE];
+		wb_enc->aspace[SDE_IOMMU_DOMAIN_SECURE] =
+			p->sde_kms->aspace[MSM_SMMU_DOMAIN_SECURE];
 	}
 
 	hw_mdp = sde_rm_get_mdp(&p->sde_kms->rm);
diff --git a/drivers/gpu/drm/msm/sde/sde_formats.c b/drivers/gpu/drm/msm/sde/sde_formats.c
index c3477b5..04c9e79 100644
--- a/drivers/gpu/drm/msm/sde/sde_formats.c
+++ b/drivers/gpu/drm/msm/sde/sde_formats.c
@@ -818,7 +818,7 @@
 }
 
 static int _sde_format_populate_addrs_ubwc(
-		int mmu_id,
+		struct msm_gem_address_space *aspace,
 		struct drm_framebuffer *fb,
 		struct sde_hw_fmt_layout *layout)
 {
@@ -830,7 +830,7 @@
 		return -EINVAL;
 	}
 
-	base_addr = msm_framebuffer_iova(fb, mmu_id, 0);
+	base_addr = msm_framebuffer_iova(fb, aspace, 0);
 	if (!base_addr) {
 		DRM_ERROR("failed to retrieve base addr\n");
 		return -EFAULT;
@@ -909,7 +909,7 @@
 }
 
 static int _sde_format_populate_addrs_linear(
-		int mmu_id,
+		struct msm_gem_address_space *aspace,
 		struct drm_framebuffer *fb,
 		struct sde_hw_fmt_layout *layout)
 {
@@ -926,7 +926,7 @@
 
 	/* Populate addresses for simple formats here */
 	for (i = 0; i < layout->num_planes; ++i) {
-		layout->plane_addr[i] = msm_framebuffer_iova(fb, mmu_id, i);
+		layout->plane_addr[i] = msm_framebuffer_iova(fb, aspace, i);
 		if (!layout->plane_addr[i]) {
 			DRM_ERROR("failed to retrieve base addr\n");
 			return -EFAULT;
@@ -937,7 +937,7 @@
 }
 
 int sde_format_populate_layout(
-		int mmu_id,
+		struct msm_gem_address_space *aspace,
 		struct drm_framebuffer *fb,
 		struct sde_hw_fmt_layout *layout)
 {
@@ -969,9 +969,9 @@
 	/* Populate the addresses given the fb */
 	if (SDE_FORMAT_IS_UBWC(layout->format) ||
 			SDE_FORMAT_IS_TILE(layout->format))
-		ret = _sde_format_populate_addrs_ubwc(mmu_id, fb, layout);
+		ret = _sde_format_populate_addrs_ubwc(aspace, fb, layout);
 	else
-		ret = _sde_format_populate_addrs_linear(mmu_id, fb, layout);
+		ret = _sde_format_populate_addrs_linear(aspace, fb, layout);
 
 	/* check if anything changed */
 	if (!ret && !memcmp(plane_addr, layout->plane_addr, sizeof(plane_addr)))
@@ -1013,14 +1013,14 @@
 }
 
 int sde_format_populate_layout_with_roi(
-		int mmu_id,
+		struct msm_gem_address_space *aspace,
 		struct drm_framebuffer *fb,
 		struct sde_rect *roi,
 		struct sde_hw_fmt_layout *layout)
 {
 	int ret;
 
-	ret = sde_format_populate_layout(mmu_id, fb, layout);
+	ret = sde_format_populate_layout(aspace, fb, layout);
 	if (ret || !roi)
 		return ret;
 
diff --git a/drivers/gpu/drm/msm/sde/sde_formats.h b/drivers/gpu/drm/msm/sde/sde_formats.h
index 40aab22..2333a72 100644
--- a/drivers/gpu/drm/msm/sde/sde_formats.h
+++ b/drivers/gpu/drm/msm/sde/sde_formats.h
@@ -14,6 +14,7 @@
 #define _SDE_FORMATS_H
 
 #include <drm/drm_fourcc.h>
+#include "msm_gem.h"
 #include "sde_hw_mdss.h"
 
 /**
@@ -103,7 +104,7 @@
 /**
  * sde_format_populate_layout - populate the given format layout based on
  *                     mmu, fb, and format found in the fb
- * @mmu_id:            mmu id handle
+ * @aspace:            address space pointer
  * @fb:                framebuffer pointer
  * @fmtl:              format layout structure to populate
  *
@@ -111,14 +112,14 @@
  *         are the same as before or 0 if new addresses were populated
  */
 int sde_format_populate_layout(
-		int mmu_id,
+		struct msm_gem_address_space *aspace,
 		struct drm_framebuffer *fb,
 		struct sde_hw_fmt_layout *fmtl);
 
 /**
  * sde_format_populate_layout_with_roi - populate the given format layout
  *                     based on mmu, fb, roi, and format found in the fb
- * @mmu_id:            mmu id handle
+ * @aspace:            address space pointer
  * @fb:                framebuffer pointer
  * @roi:               region of interest (optional)
  * @fmtl:              format layout structure to populate
@@ -126,7 +127,7 @@
  * Return: error code on failure, 0 on success
  */
 int sde_format_populate_layout_with_roi(
-		int mmu_id,
+		struct msm_gem_address_space *aspace,
 		struct drm_framebuffer *fb,
 		struct sde_rect *roi,
 		struct sde_hw_fmt_layout *fmtl);
diff --git a/drivers/gpu/drm/msm/sde/sde_hw_reg_dma_v1.c b/drivers/gpu/drm/msm/sde/sde_hw_reg_dma_v1.c
index dbd435b..9bc9837 100644
--- a/drivers/gpu/drm/msm/sde/sde_hw_reg_dma_v1.c
+++ b/drivers/gpu/drm/msm/sde/sde_hw_reg_dma_v1.c
@@ -13,6 +13,7 @@
 #include "sde_hw_ctl.h"
 #include "sde_hw_reg_dma_v1.h"
 #include "msm_drv.h"
+#include "msm_mmu.h"
 
 #define GUARD_BYTES (BIT(8) - 1)
 #define ALIGNED_OFFSET (U32_MAX & ~(GUARD_BYTES))
@@ -582,6 +583,7 @@
 	struct sde_reg_dma_buffer *dma_buf = NULL;
 	u32 iova_aligned, offset;
 	u32 rsize = size + GUARD_BYTES;
+	struct msm_gem_address_space *aspace = NULL;
 	int rc = 0;
 
 	if (!size || SIZE_DWORD(size) > MAX_DWORDS_SZ) {
@@ -602,7 +604,15 @@
 		goto fail;
 	}
 
-	rc = msm_gem_get_iova(dma_buf->buf, 0, &dma_buf->iova);
+	aspace = msm_gem_smmu_address_space_get(reg_dma->drm_dev,
+			MSM_SMMU_DOMAIN_UNSECURE);
+	if (!aspace) {
+		DRM_ERROR("failed to get aspace\n");
+		rc = -EINVAL;
+		goto free_gem;
+	}
+
+	rc = msm_gem_get_iova(dma_buf->buf, aspace, &dma_buf->iova);
 	if (rc) {
 		DRM_ERROR("failed to get the iova rc %d\n", rc);
 		goto free_gem;
@@ -625,7 +635,7 @@
 	return dma_buf;
 
 put_iova:
-	msm_gem_put_iova(dma_buf->buf, 0);
+	msm_gem_put_iova(dma_buf->buf, aspace);
 free_gem:
 	msm_gem_free_object(dma_buf->buf);
 fail:
diff --git a/drivers/gpu/drm/msm/sde/sde_kms.c b/drivers/gpu/drm/msm/sde/sde_kms.c
index 78ea685..abb378d 100644
--- a/drivers/gpu/drm/msm/sde/sde_kms.c
+++ b/drivers/gpu/drm/msm/sde/sde_kms.c
@@ -1414,6 +1414,29 @@
 	return drm_atomic_helper_check(dev, state);
 }
 
+static struct msm_gem_address_space*
+_sde_kms_get_address_space(struct msm_kms *kms,
+		unsigned int domain)
+{
+	struct sde_kms *sde_kms;
+
+	if (!kms) {
+		SDE_ERROR("invalid kms\n");
+		return  NULL;
+	}
+
+	sde_kms = to_sde_kms(kms);
+	if (!sde_kms) {
+		SDE_ERROR("invalid sde_kms\n");
+		return NULL;
+	}
+
+	if (domain >= MSM_SMMU_DOMAIN_MAX)
+		return NULL;
+
+	return sde_kms->aspace[domain];
+}
+
 static const struct msm_kms_funcs kms_funcs = {
 	.hw_init         = sde_kms_hw_init,
 	.postinit        = sde_kms_postinit,
@@ -1436,6 +1459,7 @@
 	.round_pixclk    = sde_kms_round_pixclk,
 	.destroy         = sde_kms_destroy,
 	.register_events = _sde_kms_register_events,
+	.get_address_space = _sde_kms_get_address_space,
 };
 
 /* the caller api needs to turn on clock before calling it */
@@ -1449,17 +1473,17 @@
 	struct msm_mmu *mmu;
 	int i;
 
-	for (i = ARRAY_SIZE(sde_kms->mmu_id) - 1; i >= 0; i--) {
-		mmu = sde_kms->aspace[i]->mmu;
-
-		if (!mmu)
+	for (i = ARRAY_SIZE(sde_kms->aspace) - 1; i >= 0; i--) {
+		if (!sde_kms->aspace[i])
 			continue;
 
+		mmu = sde_kms->aspace[i]->mmu;
+
 		mmu->funcs->detach(mmu, (const char **)iommu_ports,
 				ARRAY_SIZE(iommu_ports));
 		msm_gem_address_space_destroy(sde_kms->aspace[i]);
 
-		sde_kms->mmu_id[i] = 0;
+		sde_kms->aspace[i] = NULL;
 	}
 
 	return 0;
@@ -1499,17 +1523,6 @@
 			goto fail;
 		}
 
-		sde_kms->mmu_id[i] = msm_register_address_space(sde_kms->dev,
-			aspace);
-		if (sde_kms->mmu_id[i] < 0) {
-			ret = sde_kms->mmu_id[i];
-			SDE_ERROR("failed to register sde iommu %d: %d\n",
-					i, ret);
-			mmu->funcs->detach(mmu, (const char **)iommu_ports,
-					ARRAY_SIZE(iommu_ports));
-			msm_gem_address_space_destroy(aspace);
-			goto fail;
-		}
 	}
 
 	return 0;
diff --git a/drivers/gpu/drm/msm/sde/sde_kms.h b/drivers/gpu/drm/msm/sde/sde_kms.h
index 0c5c286..d818fdf 100644
--- a/drivers/gpu/drm/msm/sde/sde_kms.h
+++ b/drivers/gpu/drm/msm/sde/sde_kms.h
@@ -160,7 +160,6 @@
 	struct sde_mdss_cfg *catalog;
 
 	struct msm_gem_address_space *aspace[MSM_SMMU_DOMAIN_MAX];
-	int mmu_id[MSM_SMMU_DOMAIN_MAX];
 	struct sde_power_client *core_client;
 
 	struct ion_client *iclient;
diff --git a/drivers/gpu/drm/msm/sde/sde_plane.c b/drivers/gpu/drm/msm/sde/sde_plane.c
index 2a98af4..fb3523d 100644
--- a/drivers/gpu/drm/msm/sde/sde_plane.c
+++ b/drivers/gpu/drm/msm/sde/sde_plane.c
@@ -115,6 +115,7 @@
 
 /*
  * struct sde_plane - local sde plane structure
+ * @aspace: address space pointer
  * @csc_cfg: Decoded user configuration for csc
  * @csc_usr_ptr: Points to csc_cfg if valid user config available
  * @csc_ptr: Points to sde_csc_cfg structure to use for current
@@ -129,7 +130,7 @@
 struct sde_plane {
 	struct drm_plane base;
 
-	int mmu_id;
+	struct msm_gem_address_space *aspace;
 
 	struct mutex lock;
 
@@ -888,7 +889,7 @@
 		return;
 	}
 
-	ret = sde_format_populate_layout(psde->mmu_id, fb, &pipe_cfg->layout);
+	ret = sde_format_populate_layout(psde->aspace, fb, &pipe_cfg->layout);
 	if (ret == -EAGAIN)
 		SDE_DEBUG_PLANE(psde, "not updating same src addrs\n");
 	else if (ret)
@@ -1801,7 +1802,7 @@
 		struct sde_hw_fmt_layout layout;
 
 		memset(&layout, 0, sizeof(struct sde_hw_fmt_layout));
-		sde_format_populate_layout(rstate->mmu_id, state->fb,
+		sde_format_populate_layout(rstate->aspace, state->fb,
 				&layout);
 		for (i = 0; i < ARRAY_SIZE(rot_cmd->src_iova); i++) {
 			rot_cmd->src_iova[i] = layout.plane_addr[i];
@@ -1810,7 +1811,7 @@
 		rot_cmd->src_planes = layout.num_planes;
 
 		memset(&layout, 0, sizeof(struct sde_hw_fmt_layout));
-		sde_format_populate_layout(rstate->mmu_id, rstate->out_fb,
+		sde_format_populate_layout(rstate->aspace, rstate->out_fb,
 				&layout);
 		for (i = 0; i < ARRAY_SIZE(rot_cmd->dst_iova); i++) {
 			rot_cmd->dst_iova[i] = layout.plane_addr[i];
@@ -1950,6 +1951,7 @@
 	struct sde_plane_state *new_pstate = to_sde_plane_state(new_state);
 	struct sde_plane_rot_state *new_rstate = &new_pstate->rot;
 	struct drm_crtc_state *cstate;
+	struct sde_kms *kms = _sde_plane_get_kms(plane);
 	int ret;
 
 	SDE_DEBUG("plane%d.%d FB[%u] sbuf:%d rot:%d crtc:%d\n",
@@ -1958,6 +1960,9 @@
 			!!new_rstate->out_sbuf, !!new_rstate->rot_hw,
 			sde_plane_crtc_enabled(new_state));
 
+	if (!kms)
+		return -EINVAL;
+
 	if (!new_rstate->out_sbuf || !new_rstate->rot_hw)
 		return 0;
 
@@ -2005,9 +2010,11 @@
 				new_rstate->sequence_id);
 
 		if (new_state->fb->flags & DRM_MODE_FB_SECURE)
-			new_rstate->mmu_id = MSM_SMMU_DOMAIN_SECURE;
+			new_rstate->aspace =
+				kms->aspace[MSM_SMMU_DOMAIN_SECURE];
 		else
-			new_rstate->mmu_id = MSM_SMMU_DOMAIN_UNSECURE;
+			new_rstate->aspace =
+				kms->aspace[MSM_SMMU_DOMAIN_UNSECURE];
 
 		/* check if out_fb is already attached to rotator */
 		new_rstate->out_fbo = sde_kms_fbo_alloc(plane->dev, fb_w, fb_h,
@@ -2046,7 +2053,7 @@
 	}
 
 	/* prepare rotator input buffer */
-	ret = msm_framebuffer_prepare(new_state->fb, new_rstate->mmu_id);
+	ret = msm_framebuffer_prepare(new_state->fb, new_rstate->aspace);
 	if (ret) {
 		SDE_ERROR("failed to prepare input framebuffer\n");
 		goto error_prepare_input_buffer;
@@ -2058,7 +2065,7 @@
 				new_rstate->sequence_id);
 
 		ret = msm_framebuffer_prepare(new_rstate->out_fb,
-				new_rstate->mmu_id);
+				new_rstate->aspace);
 		if (ret) {
 			SDE_ERROR("failed to prepare inline framebuffer\n");
 			goto error_prepare_output_buffer;
@@ -2068,7 +2075,7 @@
 	return 0;
 
 error_prepare_output_buffer:
-	msm_framebuffer_cleanup(new_state->fb, new_rstate->mmu_id);
+	msm_framebuffer_cleanup(new_state->fb, new_rstate->aspace);
 error_prepare_input_buffer:
 	sde_crtc_res_put(cstate, SDE_CRTC_RES_ROT_OUT_FB,
 			(u64) &new_rstate->rot_hw->base);
@@ -2124,7 +2131,7 @@
 	if (sde_plane_enabled(old_state)) {
 		if (old_rstate->out_fb) {
 			msm_framebuffer_cleanup(old_rstate->out_fb,
-					old_rstate->mmu_id);
+					old_rstate->aspace);
 			sde_crtc_res_put(cstate, SDE_CRTC_RES_ROT_OUT_FB,
 					(u64) &old_rstate->rot_hw->base);
 			old_rstate->out_fb = NULL;
@@ -2133,7 +2140,7 @@
 			old_rstate->out_fbo = NULL;
 		}
 
-		msm_framebuffer_cleanup(old_state->fb, old_rstate->mmu_id);
+		msm_framebuffer_cleanup(old_state->fb, old_rstate->aspace);
 	}
 }
 
@@ -2163,6 +2170,7 @@
 	old_pstate = to_sde_plane_state(plane->state);
 	rstate = &pstate->rot;
 	old_rstate = &old_pstate->rot;
+	rstate->aspace = psde->aspace;
 
 	/* cstate will be null if crtc is disconnected from plane */
 	cstate = _sde_plane_get_crtc_state(state);
@@ -2657,14 +2665,14 @@
 
 	new_rstate = &to_sde_plane_state(new_state)->rot;
 
-	ret = msm_framebuffer_prepare(new_rstate->out_fb, new_rstate->mmu_id);
+	ret = msm_framebuffer_prepare(new_rstate->out_fb, psde->aspace);
 	if (ret) {
 		SDE_ERROR("failed to prepare framebuffer\n");
 		return ret;
 	}
 
 	/* validate framebuffer layout before commit */
-	ret = sde_format_populate_layout(new_rstate->mmu_id,
+	ret = sde_format_populate_layout(psde->aspace,
 			new_rstate->out_fb, &layout);
 	if (ret) {
 		SDE_ERROR_PLANE(psde, "failed to get format layout, %d\n", ret);
@@ -2687,7 +2695,7 @@
 
 	old_rstate = &to_sde_plane_state(old_state)->rot;
 
-	msm_framebuffer_cleanup(old_rstate->out_fb, old_rstate->mmu_id);
+	msm_framebuffer_cleanup(old_rstate->out_fb, old_rstate->aspace);
 
 	sde_plane_rot_cleanup_fb(plane, old_state);
 }
@@ -4481,7 +4489,7 @@
 	/* cache local stuff for later */
 	plane = &psde->base;
 	psde->pipe = pipe;
-	psde->mmu_id = kms->mmu_id[MSM_SMMU_DOMAIN_UNSECURE];
+	psde->aspace = kms->aspace[MSM_SMMU_DOMAIN_UNSECURE];
 	psde->is_virtual = (master_plane_id != 0);
 	psde->scaler_check_state = SDE_PLANE_SCLCHECK_NONE;
 	INIT_LIST_HEAD(&psde->mplane_list);
diff --git a/drivers/gpu/drm/msm/sde/sde_plane.h b/drivers/gpu/drm/msm/sde/sde_plane.h
index f83a891..46784e7 100644
--- a/drivers/gpu/drm/msm/sde/sde_plane.h
+++ b/drivers/gpu/drm/msm/sde/sde_plane.h
@@ -34,7 +34,7 @@
  * @rot90: true if rotation of 90 degree is required
  * @hflip: true if horizontal flip is required
  * @vflip: true if vertical flip is required
- * @mmu_id: iommu identifier for input/output buffers
+ * @aspace:  pointer address space for input/output buffers
  * @rot_cmd: rotator configuration command
  * @nplane: total number of drm plane attached to rotator
  * @in_fb: input fb attached to rotator
@@ -64,7 +64,7 @@
 	bool rot90;
 	bool hflip;
 	bool vflip;
-	u32 mmu_id;
+	struct msm_gem_address_space *aspace;
 	struct sde_hw_rot_cmd rot_cmd;
 	int nplane;
 	/* input */