getting memory footprint of new kernel

index d0b8458..3a317b6 100644
@@ -203,6 +203,11 @@ int crash_shrink_memory(unsigned long new_size);
 size_t crash_get_memory_size(void);
 void crash_free_reserved_phys_range(unsigned long begin, unsigned long end);
 
+unsigned long kexec_newkernel_get_phys(void);
+void kexec_newkernel_set_phys(unsigned long phys);
+unsigned long kexec_newkernel_get_memsz(void);
+void kexec_newkernel_set_memsz(unsigned long memsz);
+
 #else /* !CONFIG_KEXEC */
 struct pt_regs;
 struct task_struct;

index 5e4bd78..9e421a8 100644
@@ -38,6 +38,9 @@
 #include <asm/io.h>
 #include <asm/sections.h>
 
+static unsigned long kexec_newkernel_phys;
+static unsigned long kexec_newkernel_memsz;
+
 /* Per cpu memory for storing cpu states in case of system crash. */
 note_buf_t __percpu *crash_notes;
 

@@ -62,6 +65,27 @@ int kexec_should_crash(struct task_struct *p)
 	return 0;
 }
 
+unsigned long kexec_newkernel_get_phys(void)
+{
+	return kexec_newkernel_phys;
+}
+
+void kexec_newkernel_set_phys(unsigned long phys)
+{
+	printk("%s: kexec_newkernel_phys %lx -> %lx\n", __func__, kexec_newkernel_phys, phys);
+	kexec_newkernel_phys = phys;
+}
+
+unsigned long kexec_newkernel_get_memsz(void)
+{
+	return kexec_newkernel_memsz;
+}
+
+void kexec_newkernel_set_memsz(unsigned long memsz)
+{
+	printk("%s: kexec_newkernel_mensz %lx -> %lx\n", __func__, kexec_newkernel_memsz, memsz);
+	kexec_newkernel_memsz = memsz;
+}
 /*
  * When kexec transitions to the new kernel there is a one-to-one
  * mapping between physical and virtual addresses.  On processors

index 6ada93c..0660e40 100644
@@ -121,6 +121,44 @@ static ssize_t kexec_crash_size_store(struct kobject *kobj,
 }
 KERNEL_ATTR_RW(kexec_crash_size);
 
+static ssize_t kexec_newkernel_phys_show(struct kobject *kobj,
+				       struct kobj_attribute *attr, char *buf)
+{
+	return sprintf(buf, "%zu\n", kexec_newkernel_get_phys());
+}
+static ssize_t kexec_newkernel_phys_store(struct kobject *kobj,
+				   struct kobj_attribute *attr,
+				   const char *buf, size_t count)
+{
+	unsigned long phys;
+
+	if (strict_strtoul(buf, 0, &phys))
+		return -EINVAL;
+
+	kexec_newkernel_set_phys(phys);
+	return count;
+}
+KERNEL_ATTR_RW(kexec_newkernel_phys);
+
+static ssize_t kexec_newkernel_memsz_show(struct kobject *kobj,
+				       struct kobj_attribute *attr, char *buf)
+{
+	return sprintf(buf, "%zu\n", kexec_newkernel_get_memsz());
+}
+static ssize_t kexec_newkernel_memsz_store(struct kobject *kobj,
+				   struct kobj_attribute *attr,
+				   const char *buf, size_t count)
+{
+	unsigned long memsz;
+
+	if (strict_strtoul(buf, 0, &memsz))
+		return -EINVAL;
+
+	kexec_newkernel_set_memsz(memsz);
+	return count;
+}
+KERNEL_ATTR_RW(kexec_newkernel_memsz);
+
 static ssize_t vmcoreinfo_show(struct kobject *kobj,
 			       struct kobj_attribute *attr, char *buf)
 {

@@ -194,6 +232,8 @@ static struct attribute * kernel_attrs[] = {
 	&kexec_loaded_attr.attr,
 	&kexec_crash_loaded_attr.attr,
 	&kexec_crash_size_attr.attr,
+	&kexec_newkernel_phys_attr.attr,
+	&kexec_newkernel_memsz_attr.attr,
 	&vmcoreinfo_attr.attr,
 #endif
 	&rcu_expedited_attr.attr,

index 3a317b6..b46f9c6 100644
@@ -207,11 +207,13 @@ unsigned long kexec_newkernel_get_phys(void);
 void kexec_newkernel_set_phys(unsigned long phys);
 unsigned long kexec_newkernel_get_memsz(void);
 void kexec_newkernel_set_memsz(unsigned long memsz);
+void register_balloon_populate_range(void (*fn)(unsigned long pfn, unsigned long num));
 
 #else /* !CONFIG_KEXEC */
 struct pt_regs;
 struct task_struct;
 static inline void crash_kexec(struct pt_regs *regs) { }
 static inline int kexec_should_crash(struct task_struct *p) { return 0; }
+static inline void register_balloon_populate_range(void (*fn)(unsigned long pfn, unsigned long num)) { }
 #endif /* CONFIG_KEXEC */
 #endif /* LINUX_KEXEC_H */

index 9e421a8..22a3881 100644
@@ -32,6 +32,7 @@
 #include <linux/vmalloc.h>
 #include <linux/swap.h>
 #include <linux/syscore_ops.h>
+#include <linux/export.h>
 
 #include <asm/page.h>
 #include <asm/uaccess.h>

@@ -86,6 +87,41 @@ void kexec_newkernel_set_memsz(unsigned long memsz)
 	printk("%s: kexec_newkernel_mensz %lx -> %lx\n", __func__, kexec_newkernel_memsz, memsz);
 	kexec_newkernel_memsz = memsz;
 }
+
+static void (*balloon_populate_range)(unsigned long pfn, unsigned long num);
+
+void register_balloon_populate_range(void (*fn)(unsigned long pfn, unsigned long num))
+{
+	if (!balloon_populate_range)
+		balloon_populate_range = fn;
+}
+EXPORT_SYMBOL_GPL(register_balloon_populate_range);
+
+static void populate_range(unsigned long mem, size_t memsz)
+{
+	unsigned long pfn, num;
+
+
+
+	pfn = PFN_DOWN(mem);
+	num = PFN_UP(memsz);
+	if (pfn && num)
+		balloon_populate_range(pfn, num);
+}
+
+static void kimage_populate_ranges(struct kimage *image)
+{
+	unsigned long i;
+
+	if (!balloon_populate_range)
+		return;
+
+	for (i = 0; i < image->nr_segments; i++)
+		populate_range(image->segment[i].mem, image->segment[i].memsz);
+
+	populate_range(kexec_newkernel_phys, kexec_newkernel_memsz);
+}
+
 /*
  * When kexec transitions to the new kernel there is a one-to-one
  * mapping between physical and virtual addresses.  On processors

@@ -1598,6 +1634,8 @@ int kernel_kexec(void)
 		machine_shutdown();
 	}
 
+	kimage_populate_ranges(kexec_image);
+
 	machine_kexec(kexec_image);
 
 #ifdef CONFIG_KEXEC_JUMP