#include "internal.h"
 #include "shuffle.h"
 
+#ifdef CONFIG_MHP_MEMMAP_ON_MEMORY
+static int memmap_on_memory_set(const char *val, const struct kernel_param *kp)
+{
+       if (hugetlb_optimize_vmemmap_enabled())
+               return 0;
+       return param_set_bool(val, kp);
+}
+
+static const struct kernel_param_ops memmap_on_memory_ops = {
+       .flags  = KERNEL_PARAM_OPS_FL_NOARG,
+       .set    = memmap_on_memory_set,
+       .get    = param_get_bool,
+};
 
 /*
  * memory_hotplug.memmap_on_memory parameter
  */
 static bool memmap_on_memory __ro_after_init;
-#ifdef CONFIG_MHP_MEMMAP_ON_MEMORY
-module_param(memmap_on_memory, bool, 0444);
+module_param_cb(memmap_on_memory, &memmap_on_memory_ops, &memmap_on_memory, 0444);
 MODULE_PARM_DESC(memmap_on_memory, "Enable memmap on memory for memory hotplug");
+
+static inline bool mhp_memmap_on_memory(void)
+{
+       return memmap_on_memory;
+}
+#else
+static inline bool mhp_memmap_on_memory(void)
+{
+       return false;
+}
 #endif
 
 enum {
         *       altmap as an alternative source of memory, and we do not exactly
         *       populate a single PMD.
         */
-       return memmap_on_memory &&
-              !hugetlb_optimize_vmemmap_enabled() &&
-              IS_ENABLED(CONFIG_MHP_MEMMAP_ON_MEMORY) &&
+       return mhp_memmap_on_memory() &&
               size == memory_block_size_bytes() &&
               IS_ALIGNED(vmemmap_size, PMD_SIZE) &&
               IS_ALIGNED(remaining_size, (pageblock_nr_pages << PAGE_SHIFT));
         * We only support removing memory added with MHP_MEMMAP_ON_MEMORY in
         * the same granularity it was added - a single memory block.
         */
-       if (memmap_on_memory) {
+       if (mhp_memmap_on_memory()) {
                nr_vmemmap_pages = walk_memory_blocks(start, size, NULL,
                                                      get_nr_vmemmap_pages_cb);
                if (nr_vmemmap_pages) {