diff --git a/drivers/iommu/arm/arm-smmu-v3/pkvm/arm-smmu-v3.c b/drivers/iommu/arm/arm-smmu-v3/pkvm/arm-smmu-v3.c
index fbc0c9cdcde9..0e462aa3c601 100644
--- a/drivers/iommu/arm/arm-smmu-v3/pkvm/arm-smmu-v3.c
+++ b/drivers/iommu/arm/arm-smmu-v3/pkvm/arm-smmu-v3.c
@@ -908,25 +908,8 @@ static void smmu_tlb_flush_walk(unsigned long iova, size_t size,
 				size_t granule, void *cookie)
 {
 	struct kvm_hyp_iommu_domain *domain = cookie;
-	struct hyp_arm_smmu_v3_domain *smmu_domain = domain->priv;
-	struct io_pgtable *pgtable = smmu_domain->pgtable;
-	struct arm_lpae_io_pgtable *data = io_pgtable_to_data(pgtable);
-	struct arm_lpae_io_pgtable_walk_data wd = {
-		.cookie = data,
-	};
-	struct io_pgtable_walk_common walk_data = {
-		.visit_leaf = smmu_unmap_visit_leaf,
-		.data = &wd,
-	};
 
 	smmu_tlb_inv_range(domain, iova, size, granule, false);
-
-	/* idmapped domains doesn't elevate refcounts. */
-	if (data->idmapped)
-		return;
-
-	/* We need to walk the table to make sure all leafs un-tracked. */
-	pgtable->ops.pgtable_walk(&pgtable->ops, iova, size, &walk_data);
 }
 
 static void smmu_tlb_add_page(struct iommu_iotlb_gather *gather,
@@ -939,10 +922,16 @@ static void smmu_tlb_add_page(struct iommu_iotlb_gather *gather,
 		smmu_tlb_inv_range(cookie, iova, granule, granule, true);
 }
 
+static void smmu_free_leaf(unsigned long phys, size_t granule, void *cookie)
+{
+	WARN_ON(iommu_pkvm_unuse_dma(phys, granule));
+}
+
 static const struct iommu_flush_ops smmu_tlb_ops = {
 	.tlb_flush_all	= smmu_tlb_flush_all,
 	.tlb_flush_walk = smmu_tlb_flush_walk,
 	.tlb_add_page	= smmu_tlb_add_page,
+	.free_leaf	= smmu_free_leaf,
 };
 
 static void smmu_iotlb_sync(struct kvm_hyp_iommu_domain *domain,
diff --git a/drivers/iommu/io-pgtable-arm-common.c b/drivers/iommu/io-pgtable-arm-common.c
index 4fee3985c647..21c94d2b84fc 100644
--- a/drivers/iommu/io-pgtable-arm-common.c
+++ b/drivers/iommu/io-pgtable-arm-common.c
@@ -351,20 +351,20 @@ void __arm_lpae_free_pgtable(struct arm_lpae_io_pgtable *data, int lvl,
 		table_size = ARM_LPAE_GRANULE(data);
 
 	start = ptep;
-
-	/* Only leaf entries at the last level */
-	if (lvl == ARM_LPAE_MAX_LEVELS - 1)
-		end = ptep;
-	else
-		end = (void *)ptep + table_size;
+	end = (void *)ptep + table_size;
 
 	while (ptep != end) {
 		arm_lpae_iopte pte = *ptep++;
 
-		if (!iopte_valid(pte) || iopte_leaf(pte, lvl, data->iop.fmt))
+		if (iopte_leaf(pte, lvl, data->iop.fmt)) {
+			io_pgtable_free_leaf(&data->iop, iopte_to_paddr(pte, data),
+					     ARM_LPAE_BLOCK_SIZE(lvl, data));
 			continue;
+		}
 
-		__arm_lpae_free_pgtable(data, lvl + 1, iopte_deref(pte, data));
+		/* Only leaf entries at the last level */
+		if ((lvl != ARM_LPAE_MAX_LEVELS - 1) && iopte_valid(pte))
+			__arm_lpae_free_pgtable(data, lvl + 1, iopte_deref(pte, data));
 	}
 
 	__arm_lpae_free_pages(start, table_size, &data->iop.cfg, data->iop.cookie);
diff --git a/include/linux/io-pgtable.h b/include/linux/io-pgtable.h
index 7ddf23404988..c8e88d46e5e8 100644
--- a/include/linux/io-pgtable.h
+++ b/include/linux/io-pgtable.h
@@ -35,6 +35,7 @@ enum io_pgtable_fmt {
  *                  operations efficiently will typically issue them here, but
  *                  others may decide to update the iommu_iotlb_gather structure
  *                  and defer the invalidation until iommu_iotlb_sync() instead.
+ * @free_leaf:      Called when a valid leaf is removed when page table is freed.
  *
  * Note that these can all be called in atomic context and must therefore
  * not block.
@@ -45,6 +46,7 @@ struct iommu_flush_ops {
 			       void *cookie);
 	void (*tlb_add_page)(struct iommu_iotlb_gather *gather,
 			     unsigned long iova, size_t granule, void *cookie);
+	void (*free_leaf)(unsigned long phys, size_t granule, void *cookie);
 };
 
 /**
@@ -322,6 +324,13 @@ io_pgtable_tlb_add_page(struct io_pgtable *iop,
 		iop->cfg.tlb->tlb_add_page(gather, iova, granule, iop->cookie);
 }
 
+static inline void
+io_pgtable_free_leaf(struct io_pgtable *iop, unsigned long phys, size_t granule)
+{
+	if (iop->cfg.tlb && iop->cfg.tlb->free_leaf)
+		iop->cfg.tlb->free_leaf(phys, granule, iop->cookie);
+}
+
 /**
  * enum io_pgtable_caps - IO page table backend capabilities.
  */