diff --git a/drivers/iommu/arm-smmu.c b/drivers/iommu/arm-smmu.c
index bc83596a20a3..aaf9cf59dca9 100644
--- a/drivers/iommu/arm-smmu.c
+++ b/drivers/iommu/arm-smmu.c
@@ -1820,23 +1820,18 @@ static int arm_smmu_init_domain_context(struct iommu_domain *domain,
 	struct arm_smmu_cfg *cfg = &smmu_domain->cfg;
 	unsigned long quirks = 0;
 	struct iommu_group *group;
+	struct io_pgtable *iop;
 
 	mutex_lock(&smmu_domain->init_mutex);
 	if (smmu_domain->smmu)
 		goto out_unlock;
 
-	group = iommu_group_get(dev);
-	ret = iommu_logger_register(&smmu_domain->logger, domain, group);
-	iommu_group_put(group);
-	if (ret)
-		goto out_unlock;
-
 	if (domain->type == IOMMU_DOMAIN_DMA) {
 		ret = arm_smmu_setup_default_domain(dev, domain);
 		if (ret) {
 			dev_err(dev, "%s: default domain setup failed\n",
 				__func__);
-			goto out_logger;
+			goto out_unlock;
 		}
 	}
 
@@ -1893,7 +1888,7 @@ static int arm_smmu_init_domain_context(struct iommu_domain *domain,
 
 	if (cfg->fmt == ARM_SMMU_CTX_FMT_NONE) {
 		ret = -EINVAL;
-		goto out_logger;
+		goto out_unlock;
 	}
 
 	switch (smmu_domain->stage) {
@@ -1941,7 +1936,7 @@ static int arm_smmu_init_domain_context(struct iommu_domain *domain,
 		break;
 	default:
 		ret = -EINVAL;
-		goto out_logger;
+		goto out_unlock;
 	}
 
 #ifdef CONFIG_IOMMU_IO_PGTABLE_FAST
@@ -1955,7 +1950,7 @@ static int arm_smmu_init_domain_context(struct iommu_domain *domain,
 
 	ret = arm_smmu_alloc_cb(domain, smmu, dev);
 	if (ret < 0)
-		goto out_logger;
+		goto out_unlock;
 
 	cfg->cbndx = ret;
 
@@ -1985,6 +1980,13 @@ static int arm_smmu_init_domain_context(struct iommu_domain *domain,
 		goto out_clear_smmu;
 	}
 
+	group = iommu_group_get_for_dev(smmu_domain->dev);
+	iop = container_of(smmu_domain->pgtbl_ops, struct io_pgtable, ops);
+	ret = iommu_logger_register(&smmu_domain->logger, domain, group, iop);
+	iommu_group_put(group);
+	if (ret)
+		goto out_clear_smmu;
+
 	/*
 	 * assign any page table memory that might have been allocated
 	 * during alloc_io_pgtable_ops
@@ -1998,23 +2000,23 @@ static int arm_smmu_init_domain_context(struct iommu_domain *domain,
 	domain->geometry.aperture_end = (1UL << ias) - 1;
 	ret = arm_smmu_adjust_domain_geometry(dev, domain);
 	if (ret)
-		goto out_clear_smmu;
+		goto out_logger;
 	domain->geometry.force_aperture = true;
 
 	if (domain->type == IOMMU_DOMAIN_DMA) {
 		ret = arm_smmu_get_dma_cookie(dev, smmu_domain);
 		if (ret)
-			goto out_clear_smmu;
+			goto out_logger;
 	}
 
 	/* Assign an asid */
 	ret = arm_smmu_init_asid(domain, smmu);
 	if (ret)
-		goto out_clear_smmu;
+		goto out_logger;
 
 	ret = arm_smmu_setup_context_bank(smmu_domain, smmu, dev);
 	if (ret)
-		goto out_clear_smmu;
+		goto out_logger;
 
 	strlcpy(smmu_domain->domain.name, dev_name(dev),
 		sizeof(smmu_domain->domain.name));
@@ -2022,12 +2024,12 @@ static int arm_smmu_init_domain_context(struct iommu_domain *domain,
 
 	return 0;
 
-out_clear_smmu:
-	arm_smmu_destroy_domain_context(domain);
-	smmu_domain->smmu = NULL;
 out_logger:
 	iommu_logger_unregister(smmu_domain->logger);
 	smmu_domain->logger = NULL;
+out_clear_smmu:
+	arm_smmu_destroy_domain_context(domain);
+	smmu_domain->smmu = NULL;
 out_unlock:
 	mutex_unlock(&smmu_domain->init_mutex);
 	return ret;
diff --git a/drivers/iommu/iommu-logger.c b/drivers/iommu/iommu-logger.c
index 483c9c4be897..190b2604f48f 100644
--- a/drivers/iommu/iommu-logger.c
+++ b/drivers/iommu/iommu-logger.c
@@ -3,6 +3,7 @@
  * Copyright (c) 2020, The Linux Foundation. All rights reserved.
  */
 
+#include <linux/bitfield.h>
 #include <linux/module.h>
 #include <linux/iommu.h>
 #include <linux/slab.h>
@@ -12,47 +13,126 @@
 static DEFINE_MUTEX(iommu_debug_attachments_lock);
 static LIST_HEAD(iommu_debug_attachments);
 
-/*
- * Each group may have more than one domain; but each domain may
- * only have one group.
- * Used by debug tools to display the name of the device(s) associated
- * with a particular domain.
- */
-struct iommu_debug_attachment {
-	struct iommu_domain *domain;
-	struct iommu_group *group;
-	struct list_head list;
-};
+static unsigned int iommu_logger_pgtable_levels(struct iommu_domain *domain,
+						struct io_pgtable *iop)
+{
+	unsigned int va_bits, pte_size, bits_per_level, pg_shift;
+	unsigned long ias = iop->cfg.ias;
+
+	switch (iop->fmt) {
+	case ARM_32_LPAE_S1:
+	case ARM_64_LPAE_S1:
+#ifdef CONFIG_IOMMU_IO_PGTABLE_FAST
+	case ARM_V8L_FAST:
+#endif
+		pte_size = sizeof(u64);
+		break;
+	default:
+		return 0;
+	}
+
+	pg_shift = __ffs(domain->pgsize_bitmap);
+	bits_per_level = pg_shift - ilog2(pte_size);
+	va_bits = ias - pg_shift;
+	return DIV_ROUND_UP(va_bits, bits_per_level);
+}
+
+static enum iommu_logger_pgtable_fmt iommu_logger_pgtable_fmt_lut(
+							enum io_pgtable_fmt fmt)
+{
+	switch (fmt) {
+	case ARM_32_LPAE_S1:
+		return IOMMU_LOGGER_ARM_32_LPAE_S1;
+	case ARM_64_LPAE_S1:
+#ifdef CONFIG_IOMMU_IO_PGTABLE_FAST
+	case ARM_V8L_FAST:
+#endif
+		return IOMMU_LOGGER_ARM_64_LPAE_S1;
+	default:
+		return IOMMU_LOGGER_MAX_PGTABLE_FMTS;
+	}
+}
+
+static int iommu_logger_domain_ttbrs(struct io_pgtable *iop, void **ttbr0_ptr,
+				     void **ttbr1_ptr)
+{
+	int ret;
+	u64 ttbr0, ttbr1;
+
+	switch (iop->fmt) {
+	case ARM_32_LPAE_S1:
+	case ARM_64_LPAE_S1:
+		ttbr0 = iop->cfg.arm_lpae_s1_cfg.ttbr[0];
+		ttbr1 = iop->cfg.arm_lpae_s1_cfg.ttbr[1];
+		ret = 0;
+		break;
+#ifdef CONFIG_IOMMU_IO_PGTABLE_FAST
+	case ARM_V8L_FAST:
+		ttbr0 = iop->cfg.av8l_fast_cfg.ttbr[0];
+		ttbr1 = iop->cfg.av8l_fast_cfg.ttbr[1];
+		ret = 0;
+		break;
+#endif
+	default:
+		ret = -EINVAL;
+	}
+
+	if (!ret) {
+		*ttbr0_ptr = phys_to_virt(ttbr0);
+		*ttbr1_ptr = ttbr1 ? phys_to_virt(ttbr1) : NULL;
+	}
+
+	return ret;
+}
 
 static struct iommu_debug_attachment *iommu_logger_init(
 						struct iommu_domain *domain,
-						struct iommu_group *group)
+						struct iommu_group *group,
+						struct io_pgtable *iop)
 {
 	struct iommu_debug_attachment *logger;
+	unsigned int levels = iommu_logger_pgtable_levels(domain, iop);
+	enum iommu_logger_pgtable_fmt fmt = iommu_logger_pgtable_fmt_lut(
+								iop->fmt);
+	void *ttbr0, *ttbr1;
+	int ret;
+
+	if (!levels || fmt == IOMMU_LOGGER_MAX_PGTABLE_FMTS)
+		return ERR_PTR(-EINVAL);
+
+	ret = iommu_logger_domain_ttbrs(iop, &ttbr0, &ttbr1);
+	if (ret)
+		return ERR_PTR(ret);
 
 	logger = kzalloc(sizeof(*logger), GFP_KERNEL);
 	if (!logger)
-		return NULL;
+		return ERR_PTR(-ENOMEM);
 
 	INIT_LIST_HEAD(&logger->list);
 	logger->domain = domain;
 	logger->group = group;
+	logger->fmt = fmt;
+	logger->levels = levels;
+	logger->ttbr0 = ttbr0;
+	logger->ttbr1 = ttbr1;
 
 	return logger;
 }
 
 int iommu_logger_register(struct iommu_debug_attachment **logger_out,
 			  struct iommu_domain *domain,
-			  struct iommu_group *group)
+			  struct iommu_group *group,
+			  struct io_pgtable *iop)
 {
 	struct iommu_debug_attachment *logger;
 
-	if (!logger_out)
+	if (!logger_out || !domain || !group || !iop ||
+	    iop->fmt >= IO_PGTABLE_NUM_FMTS)
 		return -EINVAL;
 
-	logger = iommu_logger_init(domain, group);
-	if (!logger)
-		return -ENOMEM;
+	logger = iommu_logger_init(domain, group, iop);
+	if (IS_ERR(logger))
+		return PTR_ERR(logger);
 
 	mutex_lock(&iommu_debug_attachments_lock);
 	list_add(&logger->list, &iommu_debug_attachments);
diff --git a/drivers/iommu/iommu-logger.h b/drivers/iommu/iommu-logger.h
index 3afe1872a201..5e3d9dd6f194 100644
--- a/drivers/iommu/iommu-logger.h
+++ b/drivers/iommu/iommu-logger.h
@@ -4,18 +4,46 @@
 #ifndef __LINUX_QTI_IOMMU_LOGGER_H
 #define __LINUX_QTI_IOMMU_LOGGER_H
 
-struct iommu_debug_attachment;
+#include <linux/io-pgtable.h>
+
+enum iommu_logger_pgtable_fmt {
+	IOMMU_LOGGER_ARM_32_LPAE_S1,
+	IOMMU_LOGGER_ARM_64_LPAE_S1,
+	IOMMU_LOGGER_MAX_PGTABLE_FMTS,
+};
+
+/*
+ * Each group may have more than one domain; but each domain may
+ * only have one group.
+ * Used by debug tools to display the name of the device(s) associated
+ * with a particular domain.
+ */
+struct iommu_debug_attachment {
+	struct iommu_domain *domain;
+	struct iommu_group *group;
+	enum iommu_logger_pgtable_fmt fmt;
+	unsigned int levels;
+	/*
+	 * Virtual addresses of the top-level page tables are stored here,
+	 * as they are more useful for debug tools than physical addresses.
+	 */
+	void *ttbr0;
+	void *ttbr1;
+	struct list_head list;
+};
 
 #if IS_ENABLED(CONFIG_QTI_IOMMU_SUPPORT)
 
 int iommu_logger_register(struct iommu_debug_attachment **a,
 			  struct iommu_domain *domain,
-			  struct iommu_group *group);
+			  struct iommu_group *group,
+			  struct io_pgtable *iop);
 void iommu_logger_unregister(struct iommu_debug_attachment *a);
 #else
 static inline int iommu_logger_register(struct iommu_debug_attachment **a,
 					struct iommu_domain *domain,
-					struct iommu_group *group)
+					struct iommu_group *group,
+					struct io_pgtable *iop)
 {
 	return 0;
 }