u64 proto_tce;
        __be64 *tcep;
        u64 rpn;
+       const unsigned long tceshift = tbl->it_page_shift;
+       const unsigned long pagesize = IOMMU_PAGE_SIZE(tbl);
 
        proto_tce = TCE_PCI_READ; // Read allowed
 
 
        while (npages--) {
                /* can't move this out since we might cross MEMBLOCK boundary */
-               rpn = __pa(uaddr) >> TCE_SHIFT;
-               *tcep = cpu_to_be64(proto_tce | (rpn & TCE_RPN_MASK) << TCE_RPN_SHIFT);
+               rpn = __pa(uaddr) >> tceshift;
+               *tcep = cpu_to_be64(proto_tce | rpn << tceshift);
 
-               uaddr += TCE_PAGE_SIZE;
+               uaddr += pagesize;
                tcep++;
        }
        return 0;
        return be64_to_cpu(*tcep);
 }
 
-static void tce_free_pSeriesLP(unsigned long liobn, long, long);
+static void tce_free_pSeriesLP(unsigned long liobn, long, long, long);
 static void tce_freemulti_pSeriesLP(struct iommu_table*, long, long);
 
 static int tce_build_pSeriesLP(unsigned long liobn, long tcenum, long tceshift,
                proto_tce |= TCE_PCI_WRITE;
 
        while (npages--) {
-               tce = proto_tce | (rpn & TCE_RPN_MASK) << tceshift;
+               tce = proto_tce | rpn << tceshift;
                rc = plpar_tce_put((u64)liobn, (u64)tcenum << tceshift, tce);
 
                if (unlikely(rc == H_NOT_ENOUGH_RESOURCES)) {
                        ret = (int)rc;
-                       tce_free_pSeriesLP(liobn, tcenum_start,
+                       tce_free_pSeriesLP(liobn, tcenum_start, tceshift,
                                           (npages_start - (npages + 1)));
                        break;
                }
        long tcenum_start = tcenum, npages_start = npages;
        int ret = 0;
        unsigned long flags;
+       const unsigned long tceshift = tbl->it_page_shift;
 
        if ((npages == 1) || !firmware_has_feature(FW_FEATURE_PUT_TCE_IND)) {
                return tce_build_pSeriesLP(tbl->it_index, tcenum,
-                                          tbl->it_page_shift, npages, uaddr,
+                                          tceshift, npages, uaddr,
                                           direction, attrs);
        }
 
                if (!tcep) {
                        local_irq_restore(flags);
                        return tce_build_pSeriesLP(tbl->it_index, tcenum,
-                                       tbl->it_page_shift,
+                                       tceshift,
                                        npages, uaddr, direction, attrs);
                }
                __this_cpu_write(tce_page, tcep);
        }
 
-       rpn = __pa(uaddr) >> TCE_SHIFT;
+       rpn = __pa(uaddr) >> tceshift;
        proto_tce = TCE_PCI_READ;
        if (direction != DMA_TO_DEVICE)
                proto_tce |= TCE_PCI_WRITE;
                limit = min_t(long, npages, 4096/TCE_ENTRY_SIZE);
 
                for (l = 0; l < limit; l++) {
-                       tcep[l] = cpu_to_be64(proto_tce | (rpn & TCE_RPN_MASK) << TCE_RPN_SHIFT);
+                       tcep[l] = cpu_to_be64(proto_tce | rpn << tceshift);
                        rpn++;
                }
 
                rc = plpar_tce_put_indirect((u64)tbl->it_index,
-                                           (u64)tcenum << 12,
+                                           (u64)tcenum << tceshift,
                                            (u64)__pa(tcep),
                                            limit);
 
        return ret;
 }
 
-static void tce_free_pSeriesLP(unsigned long liobn, long tcenum, long npages)
+static void tce_free_pSeriesLP(unsigned long liobn, long tcenum, long tceshift,
+                              long npages)
 {
        u64 rc;
 
        while (npages--) {
-               rc = plpar_tce_put((u64)liobn, (u64)tcenum << 12, 0);
+               rc = plpar_tce_put((u64)liobn, (u64)tcenum << tceshift, 0);
 
                if (rc && printk_ratelimit()) {
                        printk("tce_free_pSeriesLP: plpar_tce_put failed. rc=%lld\n", rc);
        u64 rc;
 
        if (!firmware_has_feature(FW_FEATURE_STUFF_TCE))
-               return tce_free_pSeriesLP(tbl->it_index, tcenum, npages);
+               return tce_free_pSeriesLP(tbl->it_index, tcenum,
+                                         tbl->it_page_shift, npages);
 
-       rc = plpar_tce_stuff((u64)tbl->it_index, (u64)tcenum << 12, 0, npages);
+       rc = plpar_tce_stuff((u64)tbl->it_index,
+                            (u64)tcenum << tbl->it_page_shift, 0, npages);
 
        if (rc && printk_ratelimit()) {
                printk("tce_freemulti_pSeriesLP: plpar_tce_stuff failed\n");
        u64 rc;
        unsigned long tce_ret;
 
-       rc = plpar_tce_get((u64)tbl->it_index, (u64)tcenum << 12, &tce_ret);
+       rc = plpar_tce_get((u64)tbl->it_index,
+                          (u64)tcenum << tbl->it_page_shift, &tce_ret);
 
        if (rc && printk_ratelimit()) {
                printk("tce_get_pSeriesLP: plpar_tce_get failed. rc=%lld\n", rc);