diff --git a/include/linux/memcontrol.h b/include/linux/memcontrol.h
index 0219e670a436c6c9a78a4ea8be46ca167fd97302..24794b56eaafd5084d3a75c53097a46761a539d4 100644
--- a/include/linux/memcontrol.h
+++ b/include/linux/memcontrol.h
@@ -437,6 +437,11 @@ void memcg_print_bad_task(struct oom_control *oc);
 
 #ifdef CONFIG_MEMCG_SWAP_QOS
 DECLARE_STATIC_KEY_FALSE(memcg_swap_qos_key);
+
+#define MEMCG_SWAP_STAT_DISABLE		0
+#define MEMCG_SWAP_STAT_ALL		1
+#define MEMCG_SWAP_STAT_NONE		2
+#define MAX_MEMCG_SWAP_TYPE		MEMCG_SWAP_STAT_NONE
 #endif
 
 /*
diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index 9603b6dff9aae20818bfb34308baa2924711ea1b..cd578c2401013927a86c090081fed510633ee596 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -4064,14 +4064,15 @@ DEFINE_STATIC_KEY_FALSE(memcg_swap_qos_key);
 
 #ifdef CONFIG_SYSCTL
 static int sysctl_memcg_swap_qos_stat;
+static int swap_qos_type_max = MAX_MEMCG_SWAP_TYPE;
 
-static void memcg_swap_qos_reset(void)
+static void memcg_swap_qos_reset(int type)
 {
 	struct mem_cgroup *memcg;
 
 	for_each_mem_cgroup(memcg) {
 		WRITE_ONCE(memcg->swap_dev->max, PAGE_COUNTER_MAX);
-		WRITE_ONCE(memcg->swap_dev->type, SWAP_TYPE_ALL);
+		WRITE_ONCE(memcg->swap_dev->type, type);
 	}
 }
 
@@ -4079,21 +4080,39 @@ static int sysctl_memcg_swap_qos_handler(struct ctl_table *table, int write,
 			void __user *buffer, size_t *length, loff_t *ppos)
 {
 	int ret;
+	int qos_stat_old = sysctl_memcg_swap_qos_stat;
+	int swap_type;
 
 	ret = proc_dointvec_minmax(table, write, buffer, length, ppos);
 	if (ret)
 		return ret;
+
+	if (qos_stat_old == sysctl_memcg_swap_qos_stat)
+		return 0;
+
 	if (write) {
-		if (sysctl_memcg_swap_qos_stat &&
-				!static_branch_likely(&memcg_swap_qos_key)) {
-			memcg_swap_qos_reset();
+
+		switch (sysctl_memcg_swap_qos_stat) {
+		case MEMCG_SWAP_STAT_DISABLE:
+			static_branch_disable(&memcg_swap_qos_key);
+			return 0;
+		case MEMCG_SWAP_STAT_ALL:
+			swap_type = SWAP_TYPE_ALL;
+			break;
+		case MEMCG_SWAP_STAT_NONE:
+			swap_type = SWAP_TYPE_NONE;
+			break;
+		}
+
+		if (!qos_stat_old) {
+			memcg_swap_qos_reset(swap_type);
 			static_branch_enable(&memcg_swap_qos_key);
 			enable_swap_slots_cache_max();
-		} else if (!sysctl_memcg_swap_qos_stat &&
-				static_branch_likely(&memcg_swap_qos_key)) {
-			static_branch_disable(&memcg_swap_qos_key);
+		} else {
+			return -EINVAL;
 		}
 	}
+
 	return 0;
 }
 
@@ -4105,7 +4124,7 @@ static struct ctl_table memcg_swap_qos_sysctls[] = {
 		.mode		= 0644,
 		.proc_handler	= sysctl_memcg_swap_qos_handler,
 		.extra1		= SYSCTL_ZERO,
-		.extra2		= SYSCTL_ONE,
+		.extra2		= &swap_qos_type_max,
 	},
 	{ }
 };