diff --git a/drivers/iommu/arm-smmu.c b/drivers/iommu/arm-smmu.c
index 065379c1397f..5f2a8bb3e3a3 100644
--- a/drivers/iommu/arm-smmu.c
+++ b/drivers/iommu/arm-smmu.c
@@ -367,6 +367,8 @@ struct arm_smmu_device {
 	u32				num_mapping_groups;
 	DECLARE_BITMAP(smr_map, ARM_SMMU_MAX_SMRS);
 
+	u32				ubs;
+
 	unsigned long			va_size;
 	unsigned long			ipa_size;
 	unsigned long			pa_size;
@@ -1754,6 +1756,7 @@ static int arm_smmu_init_domain_context(struct iommu_domain *domain,
 {
 	int irq, start, ret = 0;
 	unsigned long ias, oas;
+	int sep = 0;
 	struct io_pgtable_ops *pgtbl_ops;
 	enum io_pgtable_fmt fmt;
 	struct arm_smmu_domain *smmu_domain = to_smmu_domain(domain);
@@ -1795,9 +1798,27 @@ static int arm_smmu_init_domain_context(struct iommu_domain *domain,
 		start = smmu->num_s2_context_banks;
 		ias = smmu->va_size;
 		oas = smmu->ipa_size;
-		if (IS_ENABLED(CONFIG_64BIT))
+		if (IS_ENABLED(CONFIG_64BIT)) {
 			fmt = ARM_64_LPAE_S1;
-		else
+
+			if (quirks & IO_PGTABLE_QUIRK_ARM_TTBR1) {
+
+				/*
+				 * When the UBS id is 5 we know that the bus
+				 * size is 49 bits and that bit 48 is the fixed
+				 * sign extension bit.  For any other bus size
+				 * we need to specify the sign extension bit
+				 * and adjust the input size accordingly
+				 */
+
+				if (smmu->ubs == 5) {
+					sep = 48;
+				} else {
+					sep = ias - 1;
+					ias--;
+				}
+			}
+		} else
 			fmt = ARM_32_LPAE_S1;
 		break;
 	case ARM_SMMU_DOMAIN_NESTED:
@@ -1859,6 +1880,7 @@ static int arm_smmu_init_domain_context(struct iommu_domain *domain,
 			.pgsize_bitmap	= arm_smmu_ops.pgsize_bitmap,
 			.ias		= ias,
 			.oas		= oas,
+			.sep		= sep,
 			.tlb		= &arm_smmu_gather_ops,
 			.iommu_dev	= smmu->dev,
 			.iova_base	= domain->geometry.aperture_start,
@@ -3888,8 +3910,9 @@ static int arm_smmu_device_cfg_probe(struct arm_smmu_device *smmu)
 		smmu->va_size = smmu->ipa_size;
 		size = SZ_4K | SZ_2M | SZ_1G;
 	} else {
-		size = (id >> ID2_UBS_SHIFT) & ID2_UBS_MASK;
-		smmu->va_size = arm_smmu_id_size_to_bits(size);
+		smmu->ubs = (id >> ID2_UBS_SHIFT) & ID2_UBS_MASK;
+
+		smmu->va_size = arm_smmu_id_size_to_bits(smmu->ubs);
 #ifndef CONFIG_64BIT
 		smmu->va_size = min(32UL, smmu->va_size);
 #endif
diff --git a/drivers/iommu/io-pgtable-arm.c b/drivers/iommu/io-pgtable-arm.c
index 2d2583c78bdb..7651545e3f2e 100644
--- a/drivers/iommu/io-pgtable-arm.c
+++ b/drivers/iommu/io-pgtable-arm.c
@@ -550,9 +550,18 @@ static inline arm_lpae_iopte *arm_lpae_get_table(
 {
 	struct io_pgtable_cfg *cfg = &data->iop.cfg;
 
-	return ((cfg->quirks & IO_PGTABLE_QUIRK_ARM_TTBR1) &&
-		(iova & (1UL << (cfg->ias - 1)))) ?
-		data->pgd[1] : data->pgd[0];
+	/*
+	 * iovas for TTBR1 will have all the bits set between the input address
+	 * region and the sign extension bit
+	 */
+	if (unlikely(cfg->quirks & IO_PGTABLE_QUIRK_ARM_TTBR1)) {
+		unsigned long mask = GENMASK(cfg->sep, cfg->ias);
+
+		if ((iova & mask) == mask)
+			return data->pgd[1];
+	}
+
+	return data->pgd[0];
 }
 
 static int arm_lpae_map(struct io_pgtable_ops *ops, unsigned long iova,
@@ -1089,26 +1098,26 @@ static u64 arm64_lpae_setup_ttbr1(struct io_pgtable_cfg *cfg,
 	/* Set T1SZ */
 	reg |= (64ULL - cfg->ias) << ARM_LPAE_TCR_T1SZ_SHIFT;
 
-	/* Set the SEP bit based on the size */
-	switch (cfg->ias) {
-	case 32:
+	switch (cfg->sep) {
+	case 31:
 		reg |= (ARM_LPAE_TCR_SEP_31 << ARM_LPAE_TCR_SEP_SHIFT);
 		break;
-	case 36:
+	case 35:
 		reg |= (ARM_LPAE_TCR_SEP_35 << ARM_LPAE_TCR_SEP_SHIFT);
 		break;
-	case 40:
+	case 39:
 		reg |= (ARM_LPAE_TCR_SEP_39 << ARM_LPAE_TCR_SEP_SHIFT);
 		break;
-	case 42:
+	case 41:
 		reg |= (ARM_LPAE_TCR_SEP_41 << ARM_LPAE_TCR_SEP_SHIFT);
 		break;
-	case 44:
+	case 43:
 		reg |= (ARM_LPAE_TCR_SEP_43 << ARM_LPAE_TCR_SEP_SHIFT);
 		break;
-	case 48:
+	case 47:
 		reg |= (ARM_LPAE_TCR_SEP_47 << ARM_LPAE_TCR_SEP_SHIFT);
 		break;
+	case 48:
 	default:
 		reg |= (ARM_LPAE_TCR_SEP_UPSTREAM << ARM_LPAE_TCR_SEP_SHIFT);
 		break;
diff --git a/drivers/iommu/io-pgtable.h b/drivers/iommu/io-pgtable.h
index 2cf213514221..0326bb6a4afa 100644
--- a/drivers/iommu/io-pgtable.h
+++ b/drivers/iommu/io-pgtable.h
@@ -67,6 +67,7 @@ struct io_pgtable_cfg {
 	unsigned long			pgsize_bitmap;
 	unsigned int			ias;
 	unsigned int			oas;
+	int				sep;
 	const struct iommu_gather_ops	*tlb;
 	struct device			*iommu_dev;
 	dma_addr_t			iova_base;