2 files changed, 7 insertions, 31 deletions

diff --git a/arch/x86/kernel/pci-calgary_64.c b/arch/x86/kernel/pci-calgary_64.c
index 0b96b5589f08..078d4ec1a9d9 100644
--- a/arch/x86/kernel/pci-calgary_64.c
+++ b/arch/x86/kernel/pci-calgary_64.c
@@ -110,7 +110,7 @@ int use_calgary __read_mostly = 0;
  * x3950 (PCIE): 8 chassis, 32 PHBs per chassis   = 256
  * x3950 (PCIX): 8 chassis, 16 PHBs per chassis   = 128
  */
-#define MAX_PHB_BUS_NUM         384
+#define MAX_PHB_BUS_NUM         256
 
 #define PHBS_PER_CALGARY          4
 
@@ -1056,8 +1056,6 @@ static int __init calgary_init_one(struct pci_dev *dev)
         struct iommu_table *tbl;
         int ret;
 
-        BUG_ON(dev->bus->number >= MAX_PHB_BUS_NUM);
-
         bbar = busno_to_bbar(dev->bus->number);
         ret = calgary_setup_tar(dev, bbar);
         if (ret)
diff --git a/arch/x86/mm/pat_rbtree.c b/arch/x86/mm/pat_rbtree.c
index f20eeec85a86..8acaddd0fb21 100644
--- a/arch/x86/mm/pat_rbtree.c
+++ b/arch/x86/mm/pat_rbtree.c
@@ -34,8 +34,7 @@
  * memtype_lock protects the rbtree.
  */
 
-static void memtype_rb_augment_cb(struct rb_node *node);
-static struct rb_root memtype_rbroot = RB_AUGMENT_ROOT(&memtype_rb_augment_cb);
+static struct rb_root memtype_rbroot = RB_ROOT;
 
 static int is_node_overlap(struct memtype *node, u64 start, u64 end)
 {
@@ -56,7 +55,7 @@ static u64 get_subtree_max_end(struct rb_node *node)
 }
 
 /* Update 'subtree_max_end' for a node, based on node and its children */
-static void update_node_max_end(struct rb_node *node)
+static void memtype_rb_augment_cb(struct rb_node *node, void *__unused)
 {
         struct memtype *data;
         u64 max_end, child_max_end;
@@ -78,25 +77,6 @@ static void update_node_max_end(struct rb_node *node)
         data->subtree_max_end = max_end;
 }
 
-/* Update 'subtree_max_end' for a node and all its ancestors */
-static void update_path_max_end(struct rb_node *node)
-{
-        u64 old_max_end, new_max_end;
-
-        while (node) {
-                struct memtype *data = container_of(node, struct memtype, rb);
-
-                old_max_end = data->subtree_max_end;
-                update_node_max_end(node);
-                new_max_end = data->subtree_max_end;
-
-                if (new_max_end == old_max_end)
-                        break;
-
-                node = rb_parent(node);
-        }
-}
-
 /* Find the first (lowest start addr) overlapping range from rb tree */
 static struct memtype *memtype_rb_lowest_match(struct rb_root *root,
                                 u64 start, u64 end)
@@ -190,12 +170,6 @@ failure:
         return -EBUSY;
 }
 
-static void memtype_rb_augment_cb(struct rb_node *node)
-{
-        if (node)
-                update_path_max_end(node);
-}
-
 static void memtype_rb_insert(struct rb_root *root, struct memtype *newdata)
 {
         struct rb_node **node = &(root->rb_node);
@@ -213,6 +187,7 @@ static void memtype_rb_insert(struct rb_root *root, struct memtype *newdata)
 
         rb_link_node(&newdata->rb, parent, node);
         rb_insert_color(&newdata->rb, root);
+        rb_augment_insert(&newdata->rb, memtype_rb_augment_cb, NULL);
 }
 
 int rbt_memtype_check_insert(struct memtype *new, unsigned long *ret_type)
@@ -234,13 +209,16 @@ int rbt_memtype_check_insert(struct memtype *new, unsigned long *ret_type)
 
 struct memtype *rbt_memtype_erase(u64 start, u64 end)
 {
+        struct rb_node *deepest;
         struct memtype *data;
 
         data = memtype_rb_exact_match(&memtype_rbroot, start, end);
         if (!data)
                 goto out;
 
+        deepest = rb_augment_erase_begin(&data->rb);
         rb_erase(&data->rb, &memtype_rbroot);
+        rb_augment_erase_end(deepest, memtype_rb_augment_cb, NULL);
 out:
         return data;
 }