diff --git a/include/linux/memcontrol.h b/include/linux/memcontrol.h
index 924177502479..a1395b584947 100644
--- a/include/linux/memcontrol.h
+++ b/include/linux/memcontrol.h
@@ -1460,21 +1460,48 @@ static inline void mem_cgroup_flush_foreign(struct bdi_writeback *wb)
 #endif	/* CONFIG_CGROUP_WRITEBACK */
 
 struct sock;
-bool mem_cgroup_charge_skmem(struct mem_cgroup *memcg, unsigned int nr_pages);
-void mem_cgroup_uncharge_skmem(struct mem_cgroup *memcg, unsigned int nr_pages);
 #ifdef CONFIG_MEMCG
+bool __mem_cgroup_charge_skmem(struct sock *sk, unsigned int nr_pages);
+void __mem_cgroup_uncharge_skmem(struct sock *sk, unsigned int nr_pages);
+
 extern struct static_key_false memcg_sockets_enabled_key;
-#define mem_cgroup_sockets_enabled static_branch_unlikely(&memcg_sockets_enabled_key)
-void mem_cgroup_sk_alloc(struct sock *sk);
+
+static inline bool mem_cgroup_sockets_enabled(void)
+{
+	return static_branch_unlikely(&memcg_sockets_enabled_key);
+}
+
+static inline bool mem_cgroup_charge_skmem(struct sock *sk,
+					   unsigned int nr_pages)
+{
+	if (mem_cgroup_sockets_enabled())
+		return __mem_cgroup_charge_skmem(sk, nr_pages);
+	return true;
+}
+
+static inline void mem_cgroup_uncharge_skmem(struct sock *sk,
+					     unsigned int nr_pages)
+{
+	if (mem_cgroup_sockets_enabled())
+		__mem_cgroup_uncharge_skmem(sk, nr_pages);
+}
+
+void __mem_cgroup_sk_alloc(struct sock *sk);
+
+static inline void mem_cgroup_sk_alloc(struct sock *sk)
+{
+	if (mem_cgroup_sockets_enabled())
+		__mem_cgroup_sk_alloc(sk);
+}
+
 void mem_cgroup_sk_free(struct sock *sk);
-static inline bool mem_cgroup_under_socket_pressure(struct mem_cgroup *memcg)
+
+bool __mem_cgroup_under_socket_pressure(const struct sock *sk);
+
+static inline bool mem_cgroup_under_socket_pressure(const struct sock *sk)
 {
-	if (!cgroup_subsys_on_dfl(memory_cgrp_subsys) && memcg->tcpmem_pressure)
-		return true;
-	do {
-		if (time_before(jiffies, memcg->socket_pressure))
-			return true;
-	} while ((memcg = parent_mem_cgroup(memcg)));
+	if (mem_cgroup_sockets_enabled())
+		return __mem_cgroup_under_socket_pressure(sk);
 	return false;
 }
 
@@ -1483,10 +1510,31 @@ extern int memcg_expand_shrinker_maps(int new_id);
 extern void memcg_set_shrinker_bit(struct mem_cgroup *memcg,
 				   int nid, int shrinker_id);
 #else
-#define mem_cgroup_sockets_enabled 0
-static inline void mem_cgroup_sk_alloc(struct sock *sk) { };
-static inline void mem_cgroup_sk_free(struct sock *sk) { };
-static inline bool mem_cgroup_under_socket_pressure(struct mem_cgroup *memcg)
+static inline bool mem_cgroup_sockets_enabled(void)
+{
+	return false;
+}
+
+static inline bool mem_cgroup_charge_skmem(struct sock *sk,
+					   unsigned int nr_pages)
+{
+	return true;
+}
+
+static inline void mem_cgroup_uncharge_skmem(struct sock *sk,
+					     unsigned int nr_pages)
+{
+}
+
+static inline void mem_cgroup_sk_alloc(struct sock *sk)
+{
+}
+
+static inline void mem_cgroup_sk_free(struct sock *sk)
+{
+}
+
+static inline bool mem_cgroup_under_socket_pressure(const struct sock *sk)
 {
 	return false;
 }
diff --git a/include/net/sock.h b/include/net/sock.h
index 7dd3051551fb..f486a121fb99 100644
--- a/include/net/sock.h
+++ b/include/net/sock.h
@@ -502,7 +502,9 @@ struct sock {
 	void			*sk_security;
 #endif
 	struct sock_cgroup_data	sk_cgrp_data;
+#ifdef CONFIG_MEMCG
 	struct mem_cgroup	*sk_memcg;
+#endif
 	void			(*sk_state_change)(struct sock *sk);
 	void			(*sk_data_ready)(struct sock *sk);
 	void			(*sk_write_space)(struct sock *sk);
@@ -1311,8 +1313,7 @@ static inline bool sk_under_memory_pressure(const struct sock *sk)
 	if (!sk->sk_prot->memory_pressure)
 		return false;
 
-	if (mem_cgroup_sockets_enabled && sk->sk_memcg &&
-	    mem_cgroup_under_socket_pressure(sk->sk_memcg))
+	if (mem_cgroup_under_socket_pressure(sk))
 		return true;
 
 	return !!*sk->sk_prot->memory_pressure;
diff --git a/include/net/tcp.h b/include/net/tcp.h
index e85d564446c6..c2efcf3c1495 100644
--- a/include/net/tcp.h
+++ b/include/net/tcp.h
@@ -257,8 +257,7 @@ extern unsigned long tcp_memory_pressure;
 /* optimized version of sk_under_memory_pressure() for TCP sockets */
 static inline bool tcp_under_memory_pressure(const struct sock *sk)
 {
-	if (mem_cgroup_sockets_enabled && sk->sk_memcg &&
-	    mem_cgroup_under_socket_pressure(sk->sk_memcg))
+	if (mem_cgroup_under_socket_pressure(sk))
 		return true;
 
 	return READ_ONCE(tcp_memory_pressure);
diff --git a/mm/memcontrol.c b/mm/memcontrol.c
index 4d722f4eaba8..331d7cd92e3d 100644
--- a/mm/memcontrol.c
+++ b/mm/memcontrol.c
@@ -7012,13 +7012,10 @@ void mem_cgroup_migrate(struct page *oldpage, struct page *newpage)
 DEFINE_STATIC_KEY_FALSE(memcg_sockets_enabled_key);
 EXPORT_SYMBOL(memcg_sockets_enabled_key);
 
-void mem_cgroup_sk_alloc(struct sock *sk)
+void __mem_cgroup_sk_alloc(struct sock *sk)
 {
 	struct mem_cgroup *memcg;
 
-	if (!mem_cgroup_sockets_enabled)
-		return;
-
 	/* Do not associate the sock with unrelated interrupted task's memcg. */
 	if (in_interrupt())
 		return;
@@ -7041,18 +7038,39 @@ void mem_cgroup_sk_free(struct sock *sk)
 		css_put(&sk->sk_memcg->css);
 }
 
+bool __mem_cgroup_under_socket_pressure(const struct sock *sk)
+{
+	struct mem_cgroup *memcg = sk->sk_memcg;
+
+	if (!memcg)
+		return false;
+
+	if (!cgroup_subsys_on_dfl(memory_cgrp_subsys) && memcg->tcpmem_pressure)
+		return true;
+	do {
+		if (time_before(jiffies, memcg->socket_pressure))
+			return true;
+	} while ((memcg = parent_mem_cgroup(memcg)));
+	return false;
+}
+EXPORT_SYMBOL(__mem_cgroup_under_socket_pressure);
+
 /**
- * mem_cgroup_charge_skmem - charge socket memory
- * @memcg: memcg to charge
+ * __mem_cgroup_charge_skmem - charge socket memory
+ * @sk: socket to charge
  * @nr_pages: number of pages to charge
  *
  * Charges @nr_pages to @memcg. Returns %true if the charge fit within
  * @memcg's configured limit, %false if the charge had to be forced.
  */
-bool mem_cgroup_charge_skmem(struct mem_cgroup *memcg, unsigned int nr_pages)
+bool __mem_cgroup_charge_skmem(struct sock *sk, unsigned int nr_pages)
 {
+	struct mem_cgroup *memcg = sk->sk_memcg;
 	gfp_t gfp_mask = GFP_KERNEL;
 
+	if (!memcg)
+		return true;
+
 	if (!cgroup_subsys_on_dfl(memory_cgrp_subsys)) {
 		struct page_counter *fail;
 
@@ -7079,12 +7097,17 @@ bool mem_cgroup_charge_skmem(struct mem_cgroup *memcg, unsigned int nr_pages)
 }
 
 /**
- * mem_cgroup_uncharge_skmem - uncharge socket memory
- * @memcg: memcg to uncharge
+ * __mem_cgroup_uncharge_skmem - uncharge socket memory
+ * @sk: socket to uncharge
  * @nr_pages: number of pages to uncharge
  */
-void mem_cgroup_uncharge_skmem(struct mem_cgroup *memcg, unsigned int nr_pages)
+void __mem_cgroup_uncharge_skmem(struct sock *sk, unsigned int nr_pages)
 {
+	struct mem_cgroup *memcg = sk->sk_memcg;
+
+	if (!memcg)
+		return;
+
 	if (!cgroup_subsys_on_dfl(memory_cgrp_subsys)) {
 		page_counter_uncharge(&memcg->tcpmem, nr_pages);
 		return;
diff --git a/net/core/sock.c b/net/core/sock.c
index ba9e7d91e2ef..5972d26f03ae 100644
--- a/net/core/sock.c
+++ b/net/core/sock.c
@@ -1901,9 +1901,10 @@ struct sock *sk_clone_lock(const struct sock *sk, const gfp_t priority)
 
 		sock_reset_flag(newsk, SOCK_DONE);
 
+#ifdef CONFIG_MEMCG
 		/* sk->sk_memcg will be populated at accept() time */
 		newsk->sk_memcg = NULL;
-
+#endif
 		cgroup_sk_clone(&newsk->sk_cgrp_data);
 
 		rcu_read_lock();
@@ -2587,10 +2588,10 @@ int __sk_mem_raise_allocated(struct sock *sk, int size, int amt, int kind)
 {
 	struct proto *prot = sk->sk_prot;
 	long allocated = sk_memory_allocated_add(sk, amt);
-	bool charged = true;
+	bool charged;
 
-	if (mem_cgroup_sockets_enabled && sk->sk_memcg &&
-	    !(charged = mem_cgroup_charge_skmem(sk->sk_memcg, amt)))
+	charged = mem_cgroup_charge_skmem(sk, amt);
+	if (!charged)
 		goto suppress_allocation;
 
 	/* Under limit. */
@@ -2653,8 +2654,7 @@ int __sk_mem_raise_allocated(struct sock *sk, int size, int amt, int kind)
 
 	sk_memory_allocated_sub(sk, amt);
 
-	if (mem_cgroup_sockets_enabled && sk->sk_memcg)
-		mem_cgroup_uncharge_skmem(sk->sk_memcg, amt);
+	mem_cgroup_uncharge_skmem(sk, amt);
 
 	return 0;
 }
@@ -2693,8 +2693,7 @@ void __sk_mem_reduce_allocated(struct sock *sk, int amount)
 {
 	sk_memory_allocated_sub(sk, amount);
 
-	if (mem_cgroup_sockets_enabled && sk->sk_memcg)
-		mem_cgroup_uncharge_skmem(sk->sk_memcg, amount);
+	mem_cgroup_uncharge_skmem(sk, amount);
 
 	if (sk_under_memory_pressure(sk) &&
 	    (sk_memory_allocated(sk) < sk_prot_mem_limits(sk, 0)))
diff --git a/net/ipv4/inet_connection_sock.c b/net/ipv4/inet_connection_sock.c
index b457dd2d6c75..06959f4ec70e 100644
--- a/net/ipv4/inet_connection_sock.c
+++ b/net/ipv4/inet_connection_sock.c
@@ -511,7 +511,7 @@ struct sock *inet_csk_accept(struct sock *sk, int flags, int *err, bool kern)
 
 out:
 	release_sock(sk);
-	if (newsk && mem_cgroup_sockets_enabled) {
+	if (newsk && mem_cgroup_sockets_enabled()) {
 		int amt;
 
 		/* atomically get the memory usage, set and charge the
@@ -525,8 +525,8 @@ struct sock *inet_csk_accept(struct sock *sk, int flags, int *err, bool kern)
 		amt = sk_mem_pages(newsk->sk_forward_alloc +
 				   atomic_read(&newsk->sk_rmem_alloc));
 		mem_cgroup_sk_alloc(newsk);
-		if (newsk->sk_memcg && amt)
-			mem_cgroup_charge_skmem(newsk->sk_memcg, amt);
+		if (amt)
+			mem_cgroup_charge_skmem(newsk, amt);
 
 		release_sock(newsk);
 	}
diff --git a/net/ipv4/tcp_output.c b/net/ipv4/tcp_output.c
index ab79d36ed07f..ce235cc02b7e 100644
--- a/net/ipv4/tcp_output.c
+++ b/net/ipv4/tcp_output.c
@@ -3369,8 +3369,7 @@ void sk_forced_mem_schedule(struct sock *sk, int size)
 	sk->sk_forward_alloc += amt * SK_MEM_QUANTUM;
 	sk_memory_allocated_add(sk, amt);
 
-	if (mem_cgroup_sockets_enabled && sk->sk_memcg)
-		mem_cgroup_charge_skmem(sk->sk_memcg, amt);
+	mem_cgroup_charge_skmem(sk, amt);
 }
 
 /* Send a FIN. The caller locks the socket for us.
