T0SZ controls the input address range for TTBR0, so use the input
address range rather than the output address range for the calculation.
For stage-2, this means using the output size of stage-1.
Signed-off-by: Will Deacon <will.deacon@arm.com>
                        reg = TTBCR_TG0_64K;
 
                if (!stage1) {
+                       reg |= (64 - smmu->s1_output_size) << TTBCR_T0SZ_SHIFT;
+
                        switch (smmu->s2_output_size) {
                        case 32:
                                reg |= (TTBCR2_ADDR_32 << TTBCR_PASIZE_SHIFT);
                                break;
                        }
                } else {
-                       reg |= (64 - smmu->s1_output_size) << TTBCR_T0SZ_SHIFT;
+                       reg |= (64 - smmu->input_size) << TTBCR_T0SZ_SHIFT;
                }
        } else {
                reg = 0;