diff --git a/drivers/iommu/sun50i-iommu.c b/drivers/iommu/sun50i-iommu.c
index 8d8f11854676..7a51364ca77d 100644
--- a/drivers/iommu/sun50i-iommu.c
+++ b/drivers/iommu/sun50i-iommu.c
@@ -256,7 +256,12 @@ enum sun50i_iommu_aci {
 
 static phys_addr_t sun50i_pte_get_page_address(u32 pte)
 {
-	return (phys_addr_t)pte & SUN50I_PTE_PAGE_ADDRESS_MASK;
+	phys_addr_t paddr = pte & SUN50I_PTE_PAGE_ADDRESS_MASK;
+
+	if (paddr < PHYS_OFFSET)
+		paddr |= SZ_4G;
+
+	return paddr;
 }
 
 static enum sun50i_iommu_aci sun50i_get_pte_aci(u32 pte)
@@ -602,14 +607,6 @@ static int sun50i_iommu_map(struct iommu_domain *domain, unsigned long iova,
 	u32 *page_table, *pte_addr;
 	int ret = 0;
 
-	/* the IOMMU can only handle 32-bit addresses, both input and output */
-	if ((uint64_t)paddr >> 32) {
-		ret = -EINVAL;
-		dev_warn_once(iommu->dev,
-			      "attempt to map address beyond 4GB\n");
-		goto out;
-	}
-
 	page_table = sun50i_dte_get_page_table(sun50i_domain, iova, gfp);
 	if (IS_ERR(page_table)) {
 		ret = PTR_ERR(page_table);