#include <linux/memblock.h>
 #include <linux/vmalloc.h>
 #include <linux/memory.h>
+#include <linux/memory_hotplug.h>
 
 #include <asm/firmware.h>
 #include <asm/machdep.h>
 }
 
 #ifdef CONFIG_MEMORY_HOTREMOVE
-static int pseries_remove_memblock(unsigned long base, unsigned int memblock_size)
+static int pseries_remove_memory(u64 start, u64 size)
 {
-       unsigned long start, start_pfn;
-       struct zone *zone;
        int ret;
-       unsigned long section;
-       unsigned long sections_to_remove;
+
+       /* Remove htab bolted mappings for this section of memory */
+       start = (unsigned long)__va(start);
+       ret = remove_section_mapping(start, start + size);
+
+       /* Ensure all vmalloc mappings are flushed in case they also
+        * hit that section of memory
+        */
+       vm_unmap_aliases();
+
+       return ret;
+}
+
+static int pseries_remove_memblock(unsigned long base, unsigned int memblock_size)
+{
+       unsigned long block_sz, start_pfn;
+       int sections_per_block;
+       int i, nid;
 
        start_pfn = base >> PAGE_SHIFT;
 
                return 0;
        }
 
-       zone = page_zone(pfn_to_page(start_pfn));
+       block_sz = memory_block_size_bytes();
+       sections_per_block = block_sz / MIN_MEMORY_BLOCK_SIZE;
+       nid = memory_add_physaddr_to_nid(base);
 
-       /*
-        * Remove section mappings and sysfs entries for the
-        * section of the memory we are removing.
-        *
-        * NOTE: Ideally, this should be done in generic code like
-        * remove_memory(). But remove_memory() gets called by writing
-        * to sysfs "state" file and we can't remove sysfs entries
-        * while writing to it. So we have to defer it to here.
-        */
-       sections_to_remove = (memblock_size >> PAGE_SHIFT) / PAGES_PER_SECTION;
-       for (section = 0; section < sections_to_remove; section++) {
-               unsigned long pfn = start_pfn + section * PAGES_PER_SECTION;
-               ret = __remove_pages(zone, pfn, PAGES_PER_SECTION);
-               if (ret)
-                       return ret;
+       for (i = 0; i < sections_per_block; i++) {
+               remove_memory(nid, base, MIN_MEMORY_BLOCK_SIZE);
+               base += MIN_MEMORY_BLOCK_SIZE;
        }
 
-       /*
-        * Update memory regions for memory remove
-        */
+       /* Update memory regions for memory remove */
        memblock_remove(base, memblock_size);
-
-       /*
-        * Remove htab bolted mappings for this section of memory
-        */
-       start = (unsigned long)__va(base);
-       ret = remove_section_mapping(start, start + memblock_size);
-
-       /* Ensure all vmalloc mappings are flushed in case they also
-        * hit that section of memory
-        */
-       vm_unmap_aliases();
-
-       return ret;
+       return 0;
 }
 
-static int pseries_remove_memory(struct device_node *np)
+static int pseries_remove_mem_node(struct device_node *np)
 {
        const char *type;
        const unsigned int *regs;
        base = *(unsigned long *)regs;
        lmb_size = regs[3];
 
-       ret = pseries_remove_memblock(base, lmb_size);
-       return ret;
+       pseries_remove_memblock(base, lmb_size);
+       return 0;
 }
 #else
 static inline int pseries_remove_memblock(unsigned long base,
 {
        return -EOPNOTSUPP;
 }
-static inline int pseries_remove_memory(struct device_node *np)
+static inline int pseries_remove_mem_node(struct device_node *np)
 {
        return -EOPNOTSUPP;
 }
 #endif /* CONFIG_MEMORY_HOTREMOVE */
 
-static int pseries_add_memory(struct device_node *np)
+static int pseries_add_mem_node(struct device_node *np)
 {
        const char *type;
        const unsigned int *regs;
 
        switch (action) {
        case OF_RECONFIG_ATTACH_NODE:
-               err = pseries_add_memory(node);
+               err = pseries_add_mem_node(node);
                break;
        case OF_RECONFIG_DETACH_NODE:
-               err = pseries_remove_memory(node);
+               err = pseries_remove_mem_node(node);
                break;
        case OF_RECONFIG_UPDATE_PROPERTY:
                pr = (struct of_prop_reconfig *)node;
        if (firmware_has_feature(FW_FEATURE_LPAR))
                of_reconfig_notifier_register(&pseries_mem_nb);
 
+#ifdef CONFIG_MEMORY_HOTREMOVE
+       ppc_md.remove_memory = pseries_remove_memory;
+#endif
+
        return 0;
 }
 machine_device_initcall(pseries, pseries_memory_hotplug_init);