iommu: msm: Refactor code to separate compilation

Modify the code and update makefile and defconfig to allow us to
only compile one of the two IOMMU drivers that we have. Only one of
the drivers are needed at the same time. This will ensure we are not
compiling and linking code that is not needed.

Change-Id: I7db69f873245e57bddab6625e0b04e2ed48f44b6
Signed-off-by: Olav Haugan <ohaugan@codeaurora.org>

21 files changed, 170 insertions, 107 deletions

diff --git a/arch/arm/configs/apq8084_defconfig b/arch/arm/configs/apq8084_defconfig
index 549bc733f7ea..78eb35269962 100644
--- a/arch/arm/configs/apq8084_defconfig
+++ b/arch/arm/configs/apq8084_defconfig
@@ -353,7 +353,7 @@ CONFIG_SPS_SUPPORT_NDP_BAM=y
 CONFIG_QPNP_PWM=y
 CONFIG_QPNP_POWER_ON=y
 CONFIG_QPNP_CLKDIV=y
-CONFIG_MSM_IOMMU=y
+CONFIG_MSM_IOMMU_V1=y
 CONFIG_IOMMU_PGTABLES_L2=y
 CONFIG_EXT2_FS=y
 CONFIG_EXT2_FS_XATTR=y
diff --git a/arch/arm/configs/mpq8092_defconfig b/arch/arm/configs/mpq8092_defconfig
index 179750eb72bd..5d055968364b 100644
--- a/arch/arm/configs/mpq8092_defconfig
+++ b/arch/arm/configs/mpq8092_defconfig
@@ -332,7 +332,7 @@ CONFIG_SPS_SUPPORT_NDP_BAM=y
 CONFIG_QPNP_PWM=y
 CONFIG_QPNP_POWER_ON=y
 CONFIG_QPNP_CLKDIV=y
-CONFIG_MSM_IOMMU=y
+CONFIG_MSM_IOMMU_V1=y
 CONFIG_IOMMU_PGTABLES_L2=y
 CONFIG_EXT2_FS=y
 CONFIG_EXT2_FS_XATTR=y
diff --git a/arch/arm/configs/msm8226-perf_defconfig b/arch/arm/configs/msm8226-perf_defconfig
index 2f2b1e82d7ec..903ca5093aed 100644
--- a/arch/arm/configs/msm8226-perf_defconfig
+++ b/arch/arm/configs/msm8226-perf_defconfig
@@ -332,7 +332,7 @@ CONFIG_SPS_SUPPORT_NDP_BAM=y
 CONFIG_QPNP_PWM=y
 CONFIG_QPNP_POWER_ON=y
 CONFIG_QPNP_VIBRATOR=y
-CONFIG_MSM_IOMMU=y
+CONFIG_MSM_IOMMU_V1=y
 CONFIG_CORESIGHT=y
 CONFIG_CORESIGHT_TMC=y
 CONFIG_CORESIGHT_TPIU=y
diff --git a/arch/arm/configs/msm8226_defconfig b/arch/arm/configs/msm8226_defconfig
index 6eff5be65539..92be81237416 100644
--- a/arch/arm/configs/msm8226_defconfig
+++ b/arch/arm/configs/msm8226_defconfig
@@ -357,7 +357,7 @@ CONFIG_SPS_SUPPORT_NDP_BAM=y
 CONFIG_QPNP_PWM=y
 CONFIG_QPNP_POWER_ON=y
 CONFIG_QPNP_VIBRATOR=y
-CONFIG_MSM_IOMMU=y
+CONFIG_MSM_IOMMU_V1=y
 CONFIG_CORESIGHT=y
 CONFIG_CORESIGHT_TMC=y
 CONFIG_CORESIGHT_TPIU=y
diff --git a/arch/arm/configs/msm8610-perf_defconfig b/arch/arm/configs/msm8610-perf_defconfig
index 3c56bd252e81..4b6534140ee7 100644
--- a/arch/arm/configs/msm8610-perf_defconfig
+++ b/arch/arm/configs/msm8610-perf_defconfig
@@ -342,7 +342,7 @@ CONFIG_SPS_SUPPORT_NDP_BAM=y
 CONFIG_QPNP_PWM=y
 CONFIG_QPNP_POWER_ON=y
 CONFIG_QPNP_VIBRATOR=y
-CONFIG_MSM_IOMMU=y
+CONFIG_MSM_IOMMU_V0=y
 CONFIG_EXT2_FS=y
 CONFIG_EXT2_FS_XATTR=y
 CONFIG_EXT3_FS=y
diff --git a/arch/arm/configs/msm8610_defconfig b/arch/arm/configs/msm8610_defconfig
index eab374a40057..841f71a35c9f 100644
--- a/arch/arm/configs/msm8610_defconfig
+++ b/arch/arm/configs/msm8610_defconfig
@@ -361,7 +361,7 @@ CONFIG_SPS_SUPPORT_NDP_BAM=y
 CONFIG_QPNP_PWM=y
 CONFIG_QPNP_POWER_ON=y
 CONFIG_QPNP_VIBRATOR=y
-CONFIG_MSM_IOMMU=y
+CONFIG_MSM_IOMMU_V0=y
 CONFIG_CORESIGHT=y
 CONFIG_CORESIGHT_TMC=y
 CONFIG_CORESIGHT_TPIU=y
diff --git a/arch/arm/configs/msm8960-perf_defconfig b/arch/arm/configs/msm8960-perf_defconfig
index 80520e67206d..aefc7c9b3d7c 100644
--- a/arch/arm/configs/msm8960-perf_defconfig
+++ b/arch/arm/configs/msm8960-perf_defconfig
@@ -505,7 +505,7 @@ CONFIG_MSM_SSBI=y
 CONFIG_SPS=y
 CONFIG_SPS_SUPPORT_BAMDMA=y
 CONFIG_MSM_AVTIMER=y
-CONFIG_MSM_IOMMU=y
+CONFIG_MSM_IOMMU_V0=y
 CONFIG_IOMMU_PGTABLES_L2=y
 CONFIG_MOBICORE_SUPPORT=m
 CONFIG_MOBICORE_API=m
diff --git a/arch/arm/configs/msm8960_defconfig b/arch/arm/configs/msm8960_defconfig
index 77615e7daff8..fbad97aedaed 100644
--- a/arch/arm/configs/msm8960_defconfig
+++ b/arch/arm/configs/msm8960_defconfig
@@ -505,7 +505,7 @@ CONFIG_ANDROID_LOW_MEMORY_KILLER=y
 CONFIG_MSM_SSBI=y
 CONFIG_SPS=y
 CONFIG_SPS_SUPPORT_BAMDMA=y
-CONFIG_MSM_IOMMU=y
+CONFIG_MSM_IOMMU_V0=y
 CONFIG_IOMMU_PGTABLES_L2=y
 CONFIG_MOBICORE_SUPPORT=m
 CONFIG_MOBICORE_API=m
diff --git a/arch/arm/configs/msm8974-perf_defconfig b/arch/arm/configs/msm8974-perf_defconfig
index 64a67f5a9c8d..154a36868055 100644
--- a/arch/arm/configs/msm8974-perf_defconfig
+++ b/arch/arm/configs/msm8974-perf_defconfig
@@ -440,7 +440,7 @@ CONFIG_QPNP_POWER_ON=y
 CONFIG_QPNP_CLKDIV=y
 CONFIG_QPNP_REVID=y
 CONFIG_QPNP_COINCELL=y
-CONFIG_MSM_IOMMU=y
+CONFIG_MSM_IOMMU_V1=y
 CONFIG_IOMMU_PGTABLES_L2=y
 CONFIG_MOBICORE_SUPPORT=m
 CONFIG_MOBICORE_API=m
diff --git a/arch/arm/configs/msm8974_defconfig b/arch/arm/configs/msm8974_defconfig
index b0c128454bc3..2605234faa67 100644
--- a/arch/arm/configs/msm8974_defconfig
+++ b/arch/arm/configs/msm8974_defconfig
@@ -446,7 +446,7 @@ CONFIG_QPNP_POWER_ON=y
 CONFIG_QPNP_CLKDIV=y
 CONFIG_QPNP_REVID=y
 CONFIG_QPNP_COINCELL=y
-CONFIG_MSM_IOMMU=y
+CONFIG_MSM_IOMMU_V1=y
 CONFIG_MSM_IOMMU_PMON=y
 CONFIG_IOMMU_PGTABLES_L2=y
 CONFIG_MOBICORE_SUPPORT=m
diff --git a/arch/arm/mach-msm/include/mach/iommu.h b/arch/arm/mach-msm/include/mach/iommu.h
index 70c696c04318..68cfb2ec8701 100644
--- a/arch/arm/mach-msm/include/mach/iommu.h
+++ b/arch/arm/mach-msm/include/mach/iommu.h
@@ -21,6 +21,8 @@
 
 extern pgprot_t     pgprot_kernel;
 extern struct bus_type msm_iommu_sec_bus_type;
+extern struct iommu_access_ops iommu_access_ops_v0;
+extern struct iommu_access_ops iommu_access_ops_v1;
 
 /* Domain attributes */
 #define MSM_IOMMU_DOMAIN_PT_CACHEABLE	0x1
@@ -229,6 +231,8 @@ struct remote_iommu_petersons_spinlock {
 void *msm_iommu_lock_initialize(void);
 void msm_iommu_mutex_lock(void);
 void msm_iommu_mutex_unlock(void);
+void msm_set_iommu_access_ops(struct iommu_access_ops *ops);
+struct iommu_access_ops *msm_get_iommu_access_ops(void);
 #else
 static inline void *msm_iommu_lock_initialize(void)
 {
@@ -236,6 +240,14 @@ static inline void *msm_iommu_lock_initialize(void)
 }
 static inline void msm_iommu_mutex_lock(void) { }
 static inline void msm_iommu_mutex_unlock(void) { }
+static inline void msm_set_iommu_access_ops(struct iommu_access_ops *ops)
+{
+
+}
+static inline struct iommu_access_ops *msm_get_iommu_access_ops(void)
+{
+	return NULL;
+}
 #endif
 
 #ifdef CONFIG_MSM_IOMMU_GPU_SYNC
diff --git a/arch/arm/mach-msm/include/mach/iommu_perfmon.h b/arch/arm/mach-msm/include/mach/iommu_perfmon.h
index dc4671cf3695..45683f4ebd88 100644
--- a/arch/arm/mach-msm/include/mach/iommu_perfmon.h
+++ b/arch/arm/mach-msm/include/mach/iommu_perfmon.h
@@ -144,8 +144,6 @@ struct iommu_pm_hw_ops {
 	unsigned int (*read_counter)(struct iommu_pmon_counter *);
 };
 
-extern struct iommu_access_ops iommu_access_ops_v0;
-extern struct iommu_access_ops iommu_access_ops_v1;
 #define MSM_IOMMU_PMU_NO_EVENT_CLASS -1
 
 #ifdef CONFIG_MSM_IOMMU_PMON
diff --git a/drivers/gpu/msm/kgsl_iommu.c b/drivers/gpu/msm/kgsl_iommu.c
index b86e0e1a7e63..eb85900199d6 100644
--- a/drivers/gpu/msm/kgsl_iommu.c
+++ b/drivers/gpu/msm/kgsl_iommu.c
@@ -919,7 +919,7 @@ static int kgsl_iommu_init_sync_lock(struct kgsl_mmu *mmu)
 	if (iommu->sync_lock_initialized)
 		return status;
 
-	iommu_access_ops = get_iommu_access_ops_v0();
+	iommu_access_ops = msm_get_iommu_access_ops();
 
 	if (iommu_access_ops && iommu_access_ops->iommu_lock_initialize) {
 		lock_phy_addr = (uint32_t)
diff --git a/drivers/gpu/msm/kgsl_iommu.h b/drivers/gpu/msm/kgsl_iommu.h
index 53b39464f2bc..5c4b17ebdad8 100644
--- a/drivers/gpu/msm/kgsl_iommu.h
+++ b/drivers/gpu/msm/kgsl_iommu.h
@@ -83,20 +83,6 @@ struct kgsl_iommu_register_list {
 	int ctx_reg;
 };
 
-#ifdef CONFIG_MSM_IOMMU
-extern struct iommu_access_ops iommu_access_ops_v0;
-
-static inline struct iommu_access_ops *get_iommu_access_ops_v0(void)
-{
-	return &iommu_access_ops_v0;
-}
-#else
-static inline struct iommu_access_ops *get_iommu_access_ops_v0(void)
-{
-	return NULL;
-}
-#endif
-
 /*
  * Max number of iommu units that the gpu core can have
  * On APQ8064, KGSL can control a maximum of 2 IOMMU units.
diff --git a/drivers/iommu/Kconfig b/drivers/iommu/Kconfig
index 8fea1f3a9350..d2416dbbedd9 100644
--- a/drivers/iommu/Kconfig
+++ b/drivers/iommu/Kconfig
@@ -18,12 +18,32 @@ config OF_IOMMU
        depends on OF
 
 # MSM IOMMU support
+
+# MSM_IOMMU always gets selected by whoever wants it.
 config MSM_IOMMU
-	bool "MSM IOMMU Support"
-	depends on ARCH_MSM8X60 || ARCH_MSM8960 || ARCH_APQ8064 || ARCH_MSM8974 || ARCH_MPQ8092 || ARCH_MSM8610 || ARCH_MSM8226 || ARCH_APQ8084
+	bool
+
+# MSM IOMMUv0 support
+config MSM_IOMMU_V0
+	bool "MSM IOMMUv0 Support"
+	depends on ARCH_MSM8X60 || ARCH_MSM8960 || ARCH_APQ8064 || ARCH_MSM8610
+	select IOMMU_API
+	select MSM_IOMMU
+	help
+	  Support for the IOMMUs (v0) found on certain Qualcomm SOCs.
+	  These IOMMUs allow virtualization of the address space used by most
+	  cores within the multimedia subsystem.
+
+	  If unsure, say N here.
+
+# MSM IOMMUv1 support
+config MSM_IOMMU_V1
+	bool "MSM IOMMUv1 Support"
+	depends on ARCH_MSM8974 || ARCH_MPQ8092 || ARCH_MSM8226 || ARCH_APQ8084
 	select IOMMU_API
+	select MSM_IOMMU
 	help
-	  Support for the IOMMUs found on certain Qualcomm SOCs.
+	  Support for the IOMMUs (v1) found on certain Qualcomm SOCs.
 	  These IOMMUs allow virtualization of the address space used by most
 	  cores within the multimedia subsystem.
 
@@ -32,7 +52,7 @@ config MSM_IOMMU
 # MSM IOMMU CPU-GPU sync programming support
 config MSM_IOMMU_GPU_SYNC
 	bool "MSM IOMMU CPU-GPU Sync Support"
-	depends on (ARCH_MSM8X60 || ARCH_MSM8960 || ARCH_APQ8064 || ARCH_MSM8930) && MSM_IOMMU && MSM_REMOTE_SPINLOCK_SFPB
+	depends on (ARCH_MSM8X60 || ARCH_MSM8960 || ARCH_APQ8064 || ARCH_MSM8930) && MSM_IOMMU_V0 && MSM_REMOTE_SPINLOCK_SFPB
 	help
 	  Say Y here if you want to synchronize access to IOMMU configuration
 	  port between CPU and GPU. CPU will grab a remote spinlock before
diff --git a/drivers/iommu/Makefile b/drivers/iommu/Makefile
index 0cb35f976160..325a692335d8 100644
--- a/drivers/iommu/Makefile
+++ b/drivers/iommu/Makefile
@@ -1,10 +1,15 @@
 obj-$(CONFIG_IOMMU_API) += iommu.o
 obj-$(CONFIG_OF_IOMMU)	+= of_iommu.o
-obj-$(CONFIG_MSM_IOMMU) += msm_iommu-v0.o msm_iommu_dev-v0.o
-ifdef CONFIG_OF
-obj-$(CONFIG_MSM_IOMMU) += msm_iommu-v1.o msm_iommu_dev-v1.o msm_iommu_pagetable.o msm_iommu_sec.o
+obj-$(CONFIG_MSM_IOMMU) += msm_iommu.o
+obj-$(CONFIG_MSM_IOMMU_V0) += msm_iommu-v0.o msm_iommu_dev-v0.o
+obj-$(CONFIG_MSM_IOMMU_V1) += msm_iommu-v1.o msm_iommu_dev-v1.o msm_iommu_pagetable.o msm_iommu_sec.o
+obj-$(CONFIG_MSM_IOMMU_PMON) += msm_iommu_perfmon.o
+ifdef CONFIG_MSM_IOMMU_V0
+obj-$(CONFIG_MSM_IOMMU_PMON) += msm_iommu_perfmon-v0.o
+endif
+ifdef CONFIG_MSM_IOMMU_V1
+obj-$(CONFIG_MSM_IOMMU_PMON) += msm_iommu_perfmon-v1.o
 endif
-obj-$(CONFIG_MSM_IOMMU_PMON) += msm_iommu_perfmon.o msm_iommu_perfmon-v0.o msm_iommu_perfmon-v1.o
 obj-$(CONFIG_AMD_IOMMU) += amd_iommu.o amd_iommu_init.o
 obj-$(CONFIG_AMD_IOMMU_V2) += amd_iommu_v2.o
 obj-$(CONFIG_DMAR_TABLE) += dmar.o
diff --git a/drivers/iommu/msm_iommu-v0.c b/drivers/iommu/msm_iommu-v0.c
index 1dd520240e8a..3f018bafaf5f 100644
--- a/drivers/iommu/msm_iommu-v0.c
+++ b/drivers/iommu/msm_iommu-v0.c
@@ -52,10 +52,6 @@ __asm__ __volatile__ (							\
 #define MSM_IOMMU_ATTR_CACHED_WB_NWA	0x2
 #define MSM_IOMMU_ATTR_CACHED_WT	0x3
 
-struct bus_type msm_iommu_sec_bus_type = {
-	.name = "msm_iommu_sec_bus",
-};
-
 static int msm_iommu_unmap_range(struct iommu_domain *domain, unsigned int va,
 				 unsigned int len);
 
@@ -224,7 +220,6 @@ struct iommu_access_ops iommu_access_ops_v0 = {
 	.iommu_lock_acquire = _iommu_lock_acquire,
 	.iommu_lock_release = _iommu_lock_release,
 };
-EXPORT_SYMBOL(iommu_access_ops_v0);
 
 static int __flush_iotlb_va(struct iommu_domain *domain, unsigned int va)
 {
diff --git a/drivers/iommu/msm_iommu-v1.c b/drivers/iommu/msm_iommu-v1.c
index d92111ab46ff..5879d5465438 100644
--- a/drivers/iommu/msm_iommu-v1.c
+++ b/drivers/iommu/msm_iommu-v1.c
@@ -141,7 +141,6 @@ struct iommu_access_ops iommu_access_ops_v1 = {
 	.iommu_lock_acquire = _iommu_lock_acquire,
 	.iommu_lock_release = _iommu_lock_release,
 };
-EXPORT_SYMBOL(iommu_access_ops_v1);
 
 void iommu_halt(const struct msm_iommu_drvdata *iommu_drvdata)
 {
diff --git a/drivers/iommu/msm_iommu.c b/drivers/iommu/msm_iommu.c
new file mode 100644
index 000000000000..ebecd114e3cb
--- /dev/null
+++ b/drivers/iommu/msm_iommu.c
@@ -0,0 +1,97 @@
+/* Copyright (c) 2013, The Linux Foundation. All rights reserved.
+ *
+ * This program is free software; you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License version 2 and
+ * only version 2 as published by the Free Software Foundation.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ */
+
+#define pr_fmt(fmt)	KBUILD_MODNAME ": " fmt
+#include <linux/errno.h>
+#include <linux/err.h>
+#include <linux/list.h>
+#include <linux/mutex.h>
+#include <linux/platform_device.h>
+#include <linux/export.h>
+#include <linux/iommu.h>
+#include <mach/iommu.h>
+
+static DEFINE_MUTEX(iommu_list_lock);
+static LIST_HEAD(iommu_list);
+
+static struct iommu_access_ops *iommu_access_ops;
+
+struct bus_type msm_iommu_sec_bus_type = {
+	.name = "msm_iommu_sec_bus",
+};
+
+void msm_set_iommu_access_ops(struct iommu_access_ops *ops)
+{
+	iommu_access_ops = ops;
+}
+
+struct iommu_access_ops *msm_get_iommu_access_ops()
+{
+	BUG_ON(iommu_access_ops == NULL);
+	return iommu_access_ops;
+}
+EXPORT_SYMBOL(msm_get_iommu_access_ops);
+
+void msm_iommu_add_drv(struct msm_iommu_drvdata *drv)
+{
+	mutex_lock(&iommu_list_lock);
+	list_add(&drv->list, &iommu_list);
+	mutex_unlock(&iommu_list_lock);
+}
+
+void msm_iommu_remove_drv(struct msm_iommu_drvdata *drv)
+{
+	mutex_lock(&iommu_list_lock);
+	list_del(&drv->list);
+	mutex_unlock(&iommu_list_lock);
+}
+
+static int find_iommu_ctx(struct device *dev, void *data)
+{
+	struct msm_iommu_ctx_drvdata *c;
+
+	c = dev_get_drvdata(dev);
+	if (!c || !c->name)
+		return 0;
+
+	return !strcmp(data, c->name);
+}
+
+static struct device *find_context(struct device *dev, const char *name)
+{
+	return device_find_child(dev, (void *)name, find_iommu_ctx);
+}
+
+struct device *msm_iommu_get_ctx(const char *ctx_name)
+{
+	struct msm_iommu_drvdata *drv;
+	struct device *dev = NULL;
+
+	mutex_lock(&iommu_list_lock);
+	list_for_each_entry(drv, &iommu_list, list) {
+		dev = find_context(drv->dev, ctx_name);
+		if (dev)
+			break;
+	}
+	mutex_unlock(&iommu_list_lock);
+
+	put_device(dev);
+
+	if (!dev || !dev_get_drvdata(dev)) {
+		pr_debug("Could not find context <%s>\n", ctx_name);
+		dev = ERR_PTR(-EPROBE_DEFER);
+	}
+
+	return dev;
+}
+EXPORT_SYMBOL(msm_iommu_get_ctx);
+
diff --git a/drivers/iommu/msm_iommu_dev-v0.c b/drivers/iommu/msm_iommu_dev-v0.c
index f5a39bfa79f3..dc5582d1bb06 100644
--- a/drivers/iommu/msm_iommu_dev-v0.c
+++ b/drivers/iommu/msm_iommu_dev-v0.c
@@ -32,64 +32,8 @@
 #include <mach/iommu.h>
 #include <mach/msm_bus.h>
 
-static DEFINE_MUTEX(iommu_list_lock);
-static LIST_HEAD(iommu_list);
-
 static struct of_device_id msm_iommu_v0_ctx_match_table[];
-
-void msm_iommu_add_drv(struct msm_iommu_drvdata *drv)
-{
-	mutex_lock(&iommu_list_lock);
-	list_add(&drv->list, &iommu_list);
-	mutex_unlock(&iommu_list_lock);
-}
-
-void msm_iommu_remove_drv(struct msm_iommu_drvdata *drv)
-{
-	mutex_lock(&iommu_list_lock);
-	list_del(&drv->list);
-	mutex_unlock(&iommu_list_lock);
-}
-
-static int find_iommu_ctx(struct device *dev, void *data)
-{
-	struct msm_iommu_ctx_drvdata *c;
-
-	c = dev_get_drvdata(dev);
-	if (!c || !c->name)
-		return 0;
-
-	return !strcmp(data, c->name);
-}
-
-static struct device *find_context(struct device *dev, const char *name)
-{
-	return device_find_child(dev, (void *)name, find_iommu_ctx);
-}
-
-struct device *msm_iommu_get_ctx(const char *ctx_name)
-{
-	struct msm_iommu_drvdata *drv;
-	struct device *dev = NULL;
-
-	mutex_lock(&iommu_list_lock);
-	list_for_each_entry(drv, &iommu_list, list) {
-		dev = find_context(drv->dev, ctx_name);
-		if (dev)
-			break;
-	}
-	mutex_unlock(&iommu_list_lock);
-
-	put_device(dev);
-
-	if (!dev || !dev_get_drvdata(dev)) {
-		pr_debug("Could not find context <%s>\n", ctx_name);
-		dev = ERR_PTR(-EPROBE_DEFER);
-	}
-
-	return dev;
-}
-EXPORT_SYMBOL(msm_iommu_get_ctx);
+static struct iommu_access_ops *msm_iommu_access_ops;
 
 static void msm_iommu_reset(void __iomem *base, void __iomem *glb_base, int ncb)
 {
@@ -454,7 +398,7 @@ static int msm_iommu_probe(struct platform_device *pdev)
 
 	drvdata->dev = &pdev->dev;
 
-	iommu_access_ops_v0.iommu_clk_on(drvdata);
+	msm_iommu_access_ops->iommu_clk_on(drvdata);
 
 	msm_iommu_reset(drvdata->base, drvdata->glb_base, drvdata->ncb);
 
@@ -462,7 +406,7 @@ static int msm_iommu_probe(struct platform_device *pdev)
 	if (ret)
 		goto fail_clk;
 
-	iommu_access_ops_v0.iommu_clk_off(drvdata);
+	msm_iommu_access_ops->iommu_clk_off(drvdata);
 
 	pr_info("device %s mapped at %p, with %d ctx banks\n",
 		drvdata->name, drvdata->base, drvdata->ncb);
@@ -478,7 +422,7 @@ static int msm_iommu_probe(struct platform_device *pdev)
 			pr_info("%s: pmon not available.\n", drvdata->name);
 		} else {
 			pmon_info->iommu.base = drvdata->base;
-			pmon_info->iommu.ops = &iommu_access_ops_v0;
+			pmon_info->iommu.ops = msm_iommu_access_ops;
 			pmon_info->iommu.hw_ops = iommu_pm_get_hw_ops_v0();
 			pmon_info->iommu.iommu_name = drvdata->name;
 			pmon_info->iommu.always_on = 1;
@@ -497,7 +441,7 @@ static int msm_iommu_probe(struct platform_device *pdev)
 	return 0;
 
 fail_clk:
-	iommu_access_ops_v0.iommu_clk_off(drvdata);
+	msm_iommu_access_ops->iommu_clk_off(drvdata);
 fail:
 	__put_bus_vote_client(drvdata);
 fail_mem:
@@ -697,9 +641,9 @@ static int msm_iommu_ctx_probe(struct platform_device *pdev)
 		goto fail;
 	}
 
-	iommu_access_ops_v0.iommu_clk_on(drvdata);
+	msm_iommu_access_ops->iommu_clk_on(drvdata);
 	__program_m2v_tables(drvdata, ctx_drvdata);
-	iommu_access_ops_v0.iommu_clk_off(drvdata);
+	msm_iommu_access_ops->iommu_clk_off(drvdata);
 
 	dev_info(&pdev->dev, "context %s using bank %d\n", ctx_drvdata->name,
 							   ctx_drvdata->num);
@@ -746,6 +690,11 @@ static struct platform_driver msm_iommu_ctx_driver = {
 static int __init msm_iommu_driver_init(void)
 {
 	int ret;
+
+	if (msm_soc_version_supports_iommu_v0()) {
+		msm_set_iommu_access_ops(&iommu_access_ops_v0);
+		msm_iommu_access_ops = msm_get_iommu_access_ops();
+	}
 	ret = platform_driver_register(&msm_iommu_driver);
 	if (ret != 0) {
 		pr_err("Failed to register IOMMU driver\n");
diff --git a/drivers/iommu/msm_iommu_dev-v1.c b/drivers/iommu/msm_iommu_dev-v1.c
index 384f003e0536..a229bd0e4c01 100644
--- a/drivers/iommu/msm_iommu_dev-v1.c
+++ b/drivers/iommu/msm_iommu_dev-v1.c
@@ -312,8 +312,6 @@ static int msm_iommu_probe(struct platform_device *pdev)
 
 	platform_set_drvdata(pdev, drvdata);
 
-	msm_iommu_sec_set_access_ops(&iommu_access_ops_v1);
-
 	pmon_info = msm_iommu_pm_alloc(&pdev->dev);
 	if (pmon_info != NULL) {
 		ret = msm_iommu_pmon_parse_dt(pdev, pmon_info);
@@ -322,7 +320,7 @@ static int msm_iommu_probe(struct platform_device *pdev)
 			pr_info("%s: pmon not available.\n", drvdata->name);
 		} else {
 			pmon_info->iommu.base = drvdata->base;
-			pmon_info->iommu.ops = &iommu_access_ops_v1;
+			pmon_info->iommu.ops = msm_get_iommu_access_ops();
 			pmon_info->iommu.hw_ops = iommu_pm_get_hw_ops_v1();
 			pmon_info->iommu.iommu_name = drvdata->name;
 			ret = msm_iommu_pm_iommu_register(pmon_info);
@@ -501,6 +499,10 @@ static int __init msm_iommu_driver_init(void)
 {
 	int ret;
 
+	if (!msm_soc_version_supports_iommu_v0()) {
+		msm_set_iommu_access_ops(&iommu_access_ops_v1);
+		msm_iommu_sec_set_access_ops(&iommu_access_ops_v1);
+	}
 	ret = platform_driver_register(&msm_iommu_driver);
 	if (ret != 0) {
 		pr_err("Failed to register IOMMU driver\n");