[RFC v2 PATCH 0/16] Optimize this_cpu_*() ops for non-x86 (ARM64 for this series)

[RFC v2 PATCH 0/16] Optimize this_cpu_*() ops for non-x86 (ARM64 for this series)

[Yang Shi]

Hi,

This is v2 RFC. In v2 a lot problems found out by Sashiko were fixed and more
feature gaps were closed (please see the below changelog for the details).
Although there are still some open issues, for example, it just can support
48 bits VA (for 4K and 64K) and 47 bits VA (for 16K), KPTI support has not
been solved yet, etc, but I think the delta should be big enough and worth
a new RFC to gather comments in order to make sure I'm on the right track.

Some more benchmarks were done, for example, some latency related benchmarks
that I mentioned at LSFMM because I thought responsiveness should be improved
due to the removal of preempt_disable. Collected more PMU counters as well.
Please refer to the benchmark section for more details. 

Look forward to comments.


Changelog
v2: * Added 3-level and 2-level page table support.
    * Tested with 16K and 64K page size. But we just support 48 bits VA (4K
      and 64K) and 47 bits VA with 16K for now.  Please refer to the below
      "known issue" section for the detail reason.
    * Added support for memory hotplug.
    * Added support for KASAN (generic).
    * Treated percpu and local percpu area address as vmalloc address.
    * Fixed build failure for x86.
    * Fixed build failure for !CONFIG_NUMA.
    * Added KASAN support for local percpu area.
    * Some other misc bug fixes found out by Sashiko.
    * More code refactor and cleanup.
    * Regorganized the patches.
    * More benchmarks, refer to benchmark section for more details.
    * Rebased to v7.2-rc1.


Introduction
============
This patch series implemented the LSFMM 2026 proposal for optimizing
this_cpu_*() ops on ARM64. For the details of the proposal, Please refer to:
https://lore.kernel.org/linux-mm/CAHbLzkpcN-T8MH6=W3jCxcFj1gVZp8fRqe231yzZT-rV_E_org@mail.gmail.com/
I didn't repeat it in the cover letter because there is no change to the
proposal.

The series is based on 7.1-rc1. It is basically minimum viable patches.
There are still a few hacks in this series and it may break something,
for example, KPTI, SMT machines which shared TLB, etc. But it shoule be
good enough for now to demonstrate the core idea. The main purpose of the
RFC is to gather feedback, figure out missing parts and risks, and make sure
we are on the right track, as well as hopefully it can help the discussion
for the upcoming LSFMM.

I broke the patches down to arch-dependent and arch-independent parts so that
hopefully the interested persons can do experiments on other architectures,
for example, S390, easier.

A new kernel config is introduced, HAVE_LOCAL_PER_CPU_MAP. The architectures
which can support this feature will select it. Allocating and freeing percpu
local mapping is protected by this config so that others won't pay the cost.

 
Known Issues
============
1. KPTI
-------
We need determine what CPU we are on, then switch to the right page table.
Currently arm64 kernel fetches tramp_pg_dir via swapper_pg_dir - fixed_offset,
and fetches swapper_pg_dir from ttbr1. But ttbr1 may not hold swapper_pg_dir
anymore except CPU #0. So we need to figure out the other way to handle it.
Switching to tramp_pg_dir should be easy, but the reverse seems harder because
tramp_pg_dir just maps the trampoline vectors.
Maybe we can do two steps switch. Switch to swapper_pg_dir at the first step,
then switch to per cpu page table (for entry) or tramp page table (for exit).
Nobody should call this_cpu_*() at either userspace -> kernel entry stage or
kernel -> userspace exit stage.

2. SW PAN
---------
Has the similar issue as KPTI. It installs reserved_pg_dir to TTBR0 when running
in kernel space, but fetching reserved_pg_dir via swapper_pg_dir - fixed_offset.
Maybe we can save the physical address of swapper_pg_dir in a variable, then load
it from that variable instead of ttbr1.

3. Shared TLB machines
----------------------
Some machines may share TLB between CPUs, for example, SMT machines may share
TLB between the two hardware threads in one core.
The per cpu page table just can't work with it. Maybe we need a new
cpufeature to indicate whether per cpu page table is allowed or not. Then
just enable it for not-shared-TLB machines.

4. Don't support all VA bits
----------------------------
We just support 48 bits VA (4K and 64K) and 47 bits VA (16K) for now. For 4K
and 64K, supporting other VA bits is not hard, we just need to determine the
size for percpu and local percpu area.
But it is harder for supporting 48 bits VA + 16K page size. We just have two
top level kernel page table entries with this configuration, but we assume we
just need to sync up kernel page table at the top level for now. We need to
sync up kernel page table at the second level in order to support it. I'm not
sure whether it is worth it or not.


Benchmark
=========
The benchmarks are done on 160 core AmpereOne machine. The baseline is
v7.2-rc1 kernel.

1. Reduction of kernel text size
--------------------------------
The patchset can reduce at least 11 instructions for this_cpu_*() ops. Both
preempt_disable() and preempt_enable() need 4 instructions to manipulate
the preempt count, and preempt_enable() needs more instructions (compare +
READ + compare) to determine whether reschedule is needed or not.
Because this_cpu_*() ops are inlined and called in a lot of places so we
can save a lot of instructions.

The size of kernel text is reduced by ~184KB with default Fedora kernel
config. This also helps reduce kernel icache miss rate and stalled frontend
cycles as kernel build benchmark result showed.

2. Kernel Build
---------------
Run kernel build (make -j160) with the default Fedora kernel config in a
memcg.
13% - 18% sys time improvment
3% - 7% wall time improvement

5% fewer kernel icache miss, 5% fewer executed kernel instructions and
15% fewer stalled frontend cycles for kernel.

3. stress-ng vm ops
-------------------
stress-ng --vm 160 --vm-bytes 128M --vm-ops 100000000
8.5% improvement

4. stress-ng vm ops + fork
--------------------------
stress-ng --mmapfork 160 --mmapfork-bytes 128M --mmapfork-ops 500
15% improvement

5. Specjbb
----------
The specjbb test latency curves showed the patched kernel has consistently
lower p99 latency (the lower the better) than the baseline.

2.5% improvement on max-jOPS and 4% - 5% improvement on critical-jOPS.
The specjbb benchmark is quite sensitive to latency and responsiveness, 
particularly critical-jOPS result. The patches are supposed to improve the
responsiveness due to the reduction of preempt-disabled critical sections.

6. MySQL
--------
1% - 2% gains on read-only test, 2% - 4% gains on write-only test. Also see
15% decrease on frontend cache stall.


Regression test
===============
1. memcg creation
-----------------
Create 10K memcgs. Each memcg creation needs to allocate multiple percpu
variables, for example, percpu refcnt, rstat and objcg percpu refcnt.

Consumed 2112K more virtual memory for percpu “local mapping” and a few
more mega bytes consumed by per cpu page tables.
No noticeable regression was found for elapsed time.

2. fork test
------------
stress-ng --fork 160 --fork-ops 10000000
fork() needs to allocate multiple percpu variables, for example, rss
counters and mm_cid_cpu.

Roughly 1% regression was found. However stress-ng fork test has quites
small address space, the real life workloads typically have much larger
address space and do more complicated works. The stress-ng mmapfork
benchmark saw 15% improvement.


The organization of patches
===========================
The refactor and prepatory patches (patch 1 - patch 4)
Percpu page table support patches (patch 5 - patch 8)
Local percpu area support patches (patch 7 - patch 15)
Use local percpu area for this_cpu ops (patch 16)


Yang Shi (16):
      drivers: arch_numa: move percpu set up code to arch
      arm64: kconfig: make percpu related configs not depend on NUMA
      mm: pgalloc: introduce {pud|pmd}_populate_sync()
      vmalloc: pass in pgd pointer for vmap{__vunmap}_range_noflush()
      arm64: mm: enable percpu kernel page table
      arm64: mm: defined {pud|pmd}_populate_sync()
      arm64: mm: sync percpu page table for memory hotplug/unplug
      arm64: kasan: sync up kasan shadow area page table
      arm64: mm: define percpu virtual space area
      mm: percpu: prepare to use dedicated percpu area
      arm64: mm: map local percpu first chunk
      mm: percpu: set up first chunk and reserve chunk
      arm64: mm: introduce __per_cpu_local_off
      mm: percpu: allocate and free local percpu vm area
      arm64: kconfig: select HAVE_LOCAL_PER_CPU_MAP
      arm64: percpu: use local percpu for this_cpu_*() APIs

 arch/arm64/Kconfig                   |  12 +++++++---
 arch/arm64/include/asm/mmu.h         |   5 ++++
 arch/arm64/include/asm/mmu_context.h |   9 +++++++-
 arch/arm64/include/asm/percpu.h      |  37 ++++++++++++++++++++++++++++-
 arch/arm64/include/asm/pgalloc.h     |  24 +++++++++++++++++++
 arch/arm64/include/asm/pgtable.h     |  37 ++++++++++++++++++++++++++---
 arch/arm64/kernel/setup.c            |   3 +++
 arch/arm64/kernel/smp.c              |  44 +++++++++++++++++++++++++++++++++++
 arch/arm64/mm/kasan_init.c           |  47 +++++++++++++++++++++++--------------
 arch/arm64/mm/mmu.c                  | 165 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++----------------------
 arch/arm64/mm/ptdump.c               |   4 ++++
 arch/riscv/kernel/smp.c              |  51 ++++++++++++++++++++++++++++++++++++++++
 drivers/base/arch_numa.c             |  51 +---------------------------------------
 include/linux/mm.h                   |  11 +++++++++
 include/linux/percpu.h               |   4 +++-
 include/linux/pgalloc.h              |  13 +++++++++++
 include/linux/vmalloc.h              |   3 +++
 mm/Kconfig                           |   9 ++++++++
 mm/internal.h                        |   5 +++-
 mm/kmsan/hooks.c                     |  14 +++++------
 mm/percpu-internal.h                 |  14 +++++++++++
 mm/percpu-vm.c                       |  94 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 mm/percpu.c                          |  58 +++++++++++++++++++++++++++++++++++++---------
 mm/sparse-vmemmap.c                  |   4 ++--
 mm/vmalloc.c                         | 138 +++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++++---------------
 25 files changed, 712 insertions(+), 144 deletions(-)


Thanks,
Yang

[PATCH 01/16] drivers: arch_numa: move percpu set up code to arch

[Yang Shi]

The following patches will need to modify it to set up local percpu area
on arm64 so the generic code won't work anymore. ARM64 and RISCV are the
only users, so move the code to their own kernel/smp.c file.

We can't leave the generic code in arch_numa.c because both ARM64 and
RISCV select CONFIG_GENERIC_ARCH_NUMA, so it will result in multiple
definition problem.

There is no functional change.

Signed-off-by: Yang Shi <yang@os.amperecomputing.com>
---
 arch/arm64/kernel/smp.c  | 51 ++++++++++++++++++++++++++++++++++++++++
 arch/riscv/kernel/smp.c  | 51 ++++++++++++++++++++++++++++++++++++++++
 drivers/base/arch_numa.c | 51 +---------------------------------------
 3 files changed, 103 insertions(+), 50 deletions(-)

diff --git a/arch/arm64/kernel/smp.c b/arch/arm64/kernel/smp.c
index d46022f72075..0dab4e9c97e4 100644
--- a/arch/arm64/kernel/smp.c
+++ b/arch/arm64/kernel/smp.c
@@ -819,6 +819,57 @@ void __init smp_prepare_cpus(unsigned int max_cpus)
 	}
 }
 
+#ifdef CONFIG_HAVE_SETUP_PER_CPU_AREA
+extern int cpu_to_node_map[NR_CPUS];
+
+unsigned long __per_cpu_offset[NR_CPUS] __read_mostly;
+EXPORT_SYMBOL(__per_cpu_offset);
+
+int early_cpu_to_node(int cpu)
+{
+	return cpu_to_node_map[cpu];
+}
+
+static int __init pcpu_cpu_distance(unsigned int from, unsigned int to)
+{
+	return node_distance(early_cpu_to_node(from), early_cpu_to_node(to));
+}
+
+void __init setup_per_cpu_areas(void)
+{
+	unsigned long delta;
+	unsigned int cpu;
+	int rc = -EINVAL;
+
+	if (pcpu_chosen_fc != PCPU_FC_PAGE) {
+		/*
+		 * Always reserve area for module percpu variables.  That's
+		 * what the legacy allocator did.
+		 */
+		rc = pcpu_embed_first_chunk(PERCPU_MODULE_RESERVE,
+					    PERCPU_DYNAMIC_RESERVE, PAGE_SIZE,
+					    pcpu_cpu_distance,
+					    early_cpu_to_node);
+#ifdef CONFIG_NEED_PER_CPU_PAGE_FIRST_CHUNK
+		if (rc < 0)
+			pr_warn("PERCPU: %s allocator failed (%d), falling back to page size\n",
+				   pcpu_fc_names[pcpu_chosen_fc], rc);
+#endif
+	}
+
+#ifdef CONFIG_NEED_PER_CPU_PAGE_FIRST_CHUNK
+	if (rc < 0)
+		rc = pcpu_page_first_chunk(PERCPU_MODULE_RESERVE, early_cpu_to_node);
+#endif
+	if (rc < 0)
+		panic("Failed to initialize percpu areas (err=%d).", rc);
+
+	delta = (unsigned long)pcpu_base_addr - (unsigned long)__per_cpu_start;
+	for_each_possible_cpu(cpu)
+		__per_cpu_offset[cpu] = delta + pcpu_unit_offsets[cpu];
+}
+#endif
+
 static const char *ipi_types[MAX_IPI] __tracepoint_string = {
 	[IPI_RESCHEDULE]	= "Rescheduling interrupts",
 	[IPI_CALL_FUNC]		= "Function call interrupts",
diff --git a/arch/riscv/kernel/smp.c b/arch/riscv/kernel/smp.c
index fa66f9c97d74..d420ff9dcf4f 100644
--- a/arch/riscv/kernel/smp.c
+++ b/arch/riscv/kernel/smp.c
@@ -364,3 +364,54 @@ void kgdb_roundup_cpus(void)
 	}
 }
 #endif
+
+#ifdef CONFIG_HAVE_SETUP_PER_CPU_AREA
+extern int cpu_to_node_map[NR_CPUS];
+
+unsigned long __per_cpu_offset[NR_CPUS] __read_mostly;
+EXPORT_SYMBOL(__per_cpu_offset);
+
+int early_cpu_to_node(int cpu)
+{
+	return cpu_to_node_map[cpu];
+}
+
+static int __init pcpu_cpu_distance(unsigned int from, unsigned int to)
+{
+	return node_distance(early_cpu_to_node(from), early_cpu_to_node(to));
+}
+
+void __init setup_per_cpu_areas(void)
+{
+	unsigned long delta;
+	unsigned int cpu;
+	int rc = -EINVAL;
+
+	if (pcpu_chosen_fc != PCPU_FC_PAGE) {
+		/*
+		 * Always reserve area for module percpu variables.  That's
+		 * what the legacy allocator did.
+		 */
+		rc = pcpu_embed_first_chunk(PERCPU_MODULE_RESERVE,
+					    PERCPU_DYNAMIC_RESERVE, PAGE_SIZE,
+					    pcpu_cpu_distance,
+					    early_cpu_to_node);
+#ifdef CONFIG_NEED_PER_CPU_PAGE_FIRST_CHUNK
+		if (rc < 0)
+			pr_warn("PERCPU: %s allocator failed (%d), falling back to page size\n",
+				   pcpu_fc_names[pcpu_chosen_fc], rc);
+#endif
+	}
+
+#ifdef CONFIG_NEED_PER_CPU_PAGE_FIRST_CHUNK
+	if (rc < 0)
+		rc = pcpu_page_first_chunk(PERCPU_MODULE_RESERVE, early_cpu_to_node);
+#endif
+	if (rc < 0)
+		panic("Failed to initialize percpu areas (err=%d).", rc);
+
+	delta = (unsigned long)pcpu_base_addr - (unsigned long)__per_cpu_start;
+	for_each_possible_cpu(cpu)
+		__per_cpu_offset[cpu] = delta + pcpu_unit_offsets[cpu];
+}
+#endif
diff --git a/drivers/base/arch_numa.c b/drivers/base/arch_numa.c
index c99f2ab105e5..b3b91ceed6a9 100644
--- a/drivers/base/arch_numa.c
+++ b/drivers/base/arch_numa.c
@@ -16,7 +16,7 @@
 
 #include <asm/sections.h>
 
-static int cpu_to_node_map[NR_CPUS] = { [0 ... NR_CPUS-1] = NUMA_NO_NODE };
+int cpu_to_node_map[NR_CPUS] = { [0 ... NR_CPUS-1] = NUMA_NO_NODE };
 
 bool numa_off;
 
@@ -140,55 +140,6 @@ void __init early_map_cpu_to_node(unsigned int cpu, int nid)
 		set_cpu_numa_node(cpu, nid);
 }
 
-#ifdef CONFIG_HAVE_SETUP_PER_CPU_AREA
-unsigned long __per_cpu_offset[NR_CPUS] __read_mostly;
-EXPORT_SYMBOL(__per_cpu_offset);
-
-int early_cpu_to_node(int cpu)
-{
-	return cpu_to_node_map[cpu];
-}
-
-static int __init pcpu_cpu_distance(unsigned int from, unsigned int to)
-{
-	return node_distance(early_cpu_to_node(from), early_cpu_to_node(to));
-}
-
-void __init setup_per_cpu_areas(void)
-{
-	unsigned long delta;
-	unsigned int cpu;
-	int rc = -EINVAL;
-
-	if (pcpu_chosen_fc != PCPU_FC_PAGE) {
-		/*
-		 * Always reserve area for module percpu variables.  That's
-		 * what the legacy allocator did.
-		 */
-		rc = pcpu_embed_first_chunk(PERCPU_MODULE_RESERVE,
-					    PERCPU_DYNAMIC_RESERVE, PAGE_SIZE,
-					    pcpu_cpu_distance,
-					    early_cpu_to_node);
-#ifdef CONFIG_NEED_PER_CPU_PAGE_FIRST_CHUNK
-		if (rc < 0)
-			pr_warn("PERCPU: %s allocator failed (%d), falling back to page size\n",
-				   pcpu_fc_names[pcpu_chosen_fc], rc);
-#endif
-	}
-
-#ifdef CONFIG_NEED_PER_CPU_PAGE_FIRST_CHUNK
-	if (rc < 0)
-		rc = pcpu_page_first_chunk(PERCPU_MODULE_RESERVE, early_cpu_to_node);
-#endif
-	if (rc < 0)
-		panic("Failed to initialize percpu areas (err=%d).", rc);
-
-	delta = (unsigned long)pcpu_base_addr - (unsigned long)__per_cpu_start;
-	for_each_possible_cpu(cpu)
-		__per_cpu_offset[cpu] = delta + pcpu_unit_offsets[cpu];
-}
-#endif
-
 /*
  * Initialize NODE_DATA for a node on the local memory
  */
-- 
2.47.0

[PATCH 02/16] arm64: kconfig: make percpu related configs not depend on NUMA

[Yang Shi]

The percpu set up code has been moved out of arch_numa.c, so it is
unnecessary to have percpu related kernel configs depend on NUMA.

Signed-off-by: Yang Shi <yang@os.amperecomputing.com>
---
 arch/arm64/Kconfig      |  6 +++---
 arch/arm64/kernel/smp.c | 10 ++++++----
 2 files changed, 9 insertions(+), 7 deletions(-)

diff --git a/arch/arm64/Kconfig b/arch/arm64/Kconfig
index b3afe0688919..2e1bd79a579c 100644
--- a/arch/arm64/Kconfig
+++ b/arch/arm64/Kconfig
@@ -224,6 +224,7 @@ config ARM64
 	select MMU_GATHER_RCU_TABLE_FREE
 	select HAVE_RSEQ
 	select HAVE_RUST if RUSTC_SUPPORTS_ARM64
+	select HAVE_SETUP_PER_CPU_AREA
 	select HAVE_STACKPROTECTOR
 	select HAVE_STATIC_CALL if CFI
 	select HAVE_SYSCALL_TRACEPOINTS
@@ -239,6 +240,8 @@ config ARM64
 	select LOCK_MM_AND_FIND_VMA
 	select MODULES_USE_ELF_RELA
 	select NEED_DMA_MAP_STATE
+	select NEED_PER_CPU_EMBED_FIRST_CHUNK
+	select NEED_PER_CPU_PAGE_FIRST_CHUNK
 	select NEED_SG_DMA_LENGTH
 	select OF
 	select OF_EARLY_FLATTREE
@@ -1581,9 +1584,6 @@ config NUMA
 	bool "NUMA Memory Allocation and Scheduler Support"
 	select GENERIC_ARCH_NUMA
 	select OF_NUMA
-	select HAVE_SETUP_PER_CPU_AREA
-	select NEED_PER_CPU_EMBED_FIRST_CHUNK
-	select NEED_PER_CPU_PAGE_FIRST_CHUNK
 	select USE_PERCPU_NUMA_NODE_ID
 	help
 	  Enable NUMA (Non-Uniform Memory Access) support.
diff --git a/arch/arm64/kernel/smp.c b/arch/arm64/kernel/smp.c
index 0dab4e9c97e4..28ff889658d3 100644
--- a/arch/arm64/kernel/smp.c
+++ b/arch/arm64/kernel/smp.c
@@ -819,16 +819,18 @@ void __init smp_prepare_cpus(unsigned int max_cpus)
 	}
 }
 
-#ifdef CONFIG_HAVE_SETUP_PER_CPU_AREA
+#ifdef CONFIG_NUMA
 extern int cpu_to_node_map[NR_CPUS];
 
-unsigned long __per_cpu_offset[NR_CPUS] __read_mostly;
-EXPORT_SYMBOL(__per_cpu_offset);
-
 int early_cpu_to_node(int cpu)
 {
 	return cpu_to_node_map[cpu];
 }
+#endif
+
+#ifdef CONFIG_HAVE_SETUP_PER_CPU_AREA
+unsigned long __per_cpu_offset[NR_CPUS] __read_mostly;
+EXPORT_SYMBOL(__per_cpu_offset);
 
 static int __init pcpu_cpu_distance(unsigned int from, unsigned int to)
 {
-- 
2.47.0

[PATCH 03/16] mm: pgalloc: introduce {pud|pmd}_populate_sync()

[Yang Shi]

pud_populate() and pmd_populate_kernel() are called by percpu and
vmemmap to populate PUD and PMD.  They need to sync up kernel page table
when percpu page table support for ARM64 is added in the later patch.
So introduce {pud|pmd}_populate_sync() and protect them with
ARCH_HAS_{PUD|PMD}_POPULATE_SYNC kernel configs.  ARM64 specific
implementation will be defined in the later patch.

We can change pud_populate() and pmd_populate_kernel() to take address
as a new parameter, but this change is needed for all architectures and
not all callsites of them need sync page table.  So introducing new APIs
for this usecase seems much more simple.

This is a preparation patch, no functional change.

Signed-off-by: Yang Shi <yang@os.amperecomputing.com>
---
 include/linux/pgalloc.h | 13 +++++++++++++
 mm/Kconfig              |  6 ++++++
 mm/percpu.c             |  4 ++--
 mm/sparse-vmemmap.c     |  4 ++--
 4 files changed, 23 insertions(+), 4 deletions(-)

diff --git a/include/linux/pgalloc.h b/include/linux/pgalloc.h
index 9174fa59bbc5..062c64a5d4bd 100644
--- a/include/linux/pgalloc.h
+++ b/include/linux/pgalloc.h
@@ -26,4 +26,17 @@
 			arch_sync_kernel_mappings(addr, addr);		\
 	} while (0)
 
+#ifndef CONFIG_ARCH_HAS_PUD_POPULATE_SYNC
+#define pud_populate_sync(addr, pud, pmd)				\
+	do {								\
+		pud_populate(&init_mm, pud, pmd);			\
+	} while (0)
+#endif
+
+#ifndef CONFIG_ARCH_HAS_PMD_POPULATE_SYNC
+#define pmd_populate_sync(addr, pmd, pte)				\
+	do {								\
+		pmd_populate_kernel(&init_mm, pmd, pte);		\
+	} while (0)
+#endif
 #endif /* _LINUX_PGALLOC_H */
diff --git a/mm/Kconfig b/mm/Kconfig
index 9e0ca4824905..13ea4b29ee7b 100644
--- a/mm/Kconfig
+++ b/mm/Kconfig
@@ -1509,6 +1509,12 @@ config LAZY_MMU_MODE_KUNIT_TEST
 
 	  If unsure, say N.
 
+config ARCH_HAS_PMD_POPULATE_SYNC
+	bool
+
+config ARCH_HAS_PUD_POPULATE_SYNC
+	bool
+
 source "mm/damon/Kconfig"
 
 endmenu
diff --git a/mm/percpu.c b/mm/percpu.c
index b0676b8054ed..45c30c6531c5 100644
--- a/mm/percpu.c
+++ b/mm/percpu.c
@@ -3163,7 +3163,7 @@ void __init __weak pcpu_populate_pte(unsigned long addr)
 	pud = pud_offset(p4d, addr);
 	if (pud_none(*pud)) {
 		pmd = memblock_alloc_or_panic(PMD_TABLE_SIZE, PMD_TABLE_SIZE);
-		pud_populate(&init_mm, pud, pmd);
+		pud_populate_sync(addr, pud, pmd);
 	}
 
 	pmd = pmd_offset(pud, addr);
@@ -3171,7 +3171,7 @@ void __init __weak pcpu_populate_pte(unsigned long addr)
 		pte_t *new;
 
 		new = memblock_alloc_or_panic(PTE_TABLE_SIZE, PTE_TABLE_SIZE);
-		pmd_populate_kernel(&init_mm, pmd, new);
+		pmd_populate_sync(addr, pmd, new);
 	}
 
 	return;
diff --git a/mm/sparse-vmemmap.c b/mm/sparse-vmemmap.c
index 99e2be39671b..463c2f8b611c 100644
--- a/mm/sparse-vmemmap.c
+++ b/mm/sparse-vmemmap.c
@@ -198,7 +198,7 @@ static pmd_t * __meminit vmemmap_pmd_populate(pud_t *pud, unsigned long addr, in
 		if (!p)
 			return NULL;
 		kernel_pte_init(p);
-		pmd_populate_kernel(&init_mm, pmd, p);
+		pmd_populate_sync(addr, pmd, p);
 	}
 	return pmd;
 }
@@ -211,7 +211,7 @@ static pud_t * __meminit vmemmap_pud_populate(p4d_t *p4d, unsigned long addr, in
 		if (!p)
 			return NULL;
 		pmd_init(p);
-		pud_populate(&init_mm, pud, p);
+		pud_populate_sync(addr, pud, p);
 	}
 	return pud;
 }
-- 
2.47.0

[PATCH 04/16] vmalloc: pass in pgd pointer for vmap{__vunmap}_range_noflush()

[Yang Shi]

vmap{__vunmap}_range_noflush() assume manipulate init_mm pgd.  The
following patch will map percpu local mapping into percpu page table by
calling them, so the assumption will no longer stand.  Make them take
pgd pointer as an parameter.

Also make vmap_range_noflush() non static, it will be called outside
vmalloc in the following patch.

There is no functional change.

Signed-off-by: Yang Shi <yang@os.amperecomputing.com>
---
 mm/internal.h    |  5 ++++-
 mm/kmsan/hooks.c | 14 +++++++-------
 mm/vmalloc.c     | 25 +++++++++++++------------
 3 files changed, 24 insertions(+), 20 deletions(-)

diff --git a/mm/internal.h b/mm/internal.h
index 181e79f1d6a2..9d9f08cb26a1 100644
--- a/mm/internal.h
+++ b/mm/internal.h
@@ -1552,10 +1552,13 @@ void clear_vm_uninitialized_flag(struct vm_struct *vm);
 int __must_check __vmap_pages_range_noflush(unsigned long addr,
 			       unsigned long end, pgprot_t prot,
 			       struct page **pages, unsigned int page_shift);
+int __must_check vmap_range_noflush(pgd_t *pgdir, unsigned long addr,
+			unsigned long end, phys_addr_t phys_addr,
+			pgprot_t prot, unsigned int max_page_shift);
 
 void vunmap_range_noflush(unsigned long start, unsigned long end);
 
-void __vunmap_range_noflush(unsigned long start, unsigned long end);
+void __vunmap_range_noflush(pgd_t *pgdir, unsigned long start, unsigned long end);
 
 static inline bool vma_is_single_threaded_private(struct vm_area_struct *vma)
 {
diff --git a/mm/kmsan/hooks.c b/mm/kmsan/hooks.c
index 8f22d1f22981..e2a0faf344b9 100644
--- a/mm/kmsan/hooks.c
+++ b/mm/kmsan/hooks.c
@@ -135,8 +135,8 @@ static unsigned long vmalloc_origin(unsigned long addr)
 
 void kmsan_vunmap_range_noflush(unsigned long start, unsigned long end)
 {
-	__vunmap_range_noflush(vmalloc_shadow(start), vmalloc_shadow(end));
-	__vunmap_range_noflush(vmalloc_origin(start), vmalloc_origin(end));
+	__vunmap_range_noflush(init_mm.pgd, vmalloc_shadow(start), vmalloc_shadow(end));
+	__vunmap_range_noflush(init_mm.pgd, vmalloc_origin(start), vmalloc_origin(end));
 	flush_cache_vmap(vmalloc_shadow(start), vmalloc_shadow(end));
 	flush_cache_vmap(vmalloc_origin(start), vmalloc_origin(end));
 }
@@ -181,7 +181,7 @@ int kmsan_ioremap_page_range(unsigned long start, unsigned long end,
 			vmalloc_origin(start + off + PAGE_SIZE), prot, &origin,
 			PAGE_SHIFT);
 		if (mapped) {
-			__vunmap_range_noflush(
+			__vunmap_range_noflush(init_mm.pgd,
 				vmalloc_shadow(start + off),
 				vmalloc_shadow(start + off + PAGE_SIZE));
 			err = mapped;
@@ -203,10 +203,10 @@ int kmsan_ioremap_page_range(unsigned long start, unsigned long end,
 			__free_pages(shadow, 1);
 		if (origin)
 			__free_pages(origin, 1);
-		__vunmap_range_noflush(
+		__vunmap_range_noflush(init_mm.pgd,
 			vmalloc_shadow(start),
 			vmalloc_shadow(start + clean * PAGE_SIZE));
-		__vunmap_range_noflush(
+		__vunmap_range_noflush(init_mm.pgd,
 			vmalloc_origin(start),
 			vmalloc_origin(start + clean * PAGE_SIZE));
 	}
@@ -233,8 +233,8 @@ void kmsan_iounmap_page_range(unsigned long start, unsigned long end)
 	     i++, v_shadow += PAGE_SIZE, v_origin += PAGE_SIZE) {
 		shadow = kmsan_vmalloc_to_page_or_null((void *)v_shadow);
 		origin = kmsan_vmalloc_to_page_or_null((void *)v_origin);
-		__vunmap_range_noflush(v_shadow, vmalloc_shadow(end));
-		__vunmap_range_noflush(v_origin, vmalloc_origin(end));
+		__vunmap_range_noflush(init_mm.pgd, v_shadow, vmalloc_shadow(end));
+		__vunmap_range_noflush(init_mm.pgd, v_origin, vmalloc_origin(end));
 		if (shadow)
 			__free_pages(shadow, 1);
 		if (origin)
diff --git a/mm/vmalloc.c b/mm/vmalloc.c
index 1afca3568b9b..17c26e9796b2 100644
--- a/mm/vmalloc.c
+++ b/mm/vmalloc.c
@@ -295,9 +295,9 @@ static int vmap_p4d_range(pgd_t *pgd, unsigned long addr, unsigned long end,
 	return err;
 }
 
-static int vmap_range_noflush(unsigned long addr, unsigned long end,
-			phys_addr_t phys_addr, pgprot_t prot,
-			unsigned int max_page_shift)
+int vmap_range_noflush(pgd_t *pgdir, unsigned long addr, unsigned long end,
+		       phys_addr_t phys_addr, pgprot_t prot,
+		       unsigned int max_page_shift)
 {
 	pgd_t *pgd;
 	unsigned long start;
@@ -314,7 +314,7 @@ static int vmap_range_noflush(unsigned long addr, unsigned long end,
 	BUG_ON(addr >= end);
 
 	start = addr;
-	pgd = pgd_offset_k(addr);
+	pgd = pgd_offset_pgd(pgdir, addr);
 	do {
 		next = pgd_addr_end(addr, end);
 		err = vmap_p4d_range(pgd, addr, next, phys_addr, prot,
@@ -334,8 +334,8 @@ int vmap_page_range(unsigned long addr, unsigned long end,
 {
 	int err;
 
-	err = vmap_range_noflush(addr, end, phys_addr, pgprot_nx(prot),
-				 ioremap_max_page_shift);
+	err = vmap_range_noflush(init_mm.pgd, addr, end, phys_addr,
+				 pgprot_nx(prot), ioremap_max_page_shift);
 	flush_cache_vmap(addr, end);
 	if (!err)
 		err = kmsan_ioremap_page_range(addr, end, phys_addr, prot,
@@ -478,7 +478,7 @@ static void vunmap_p4d_range(pgd_t *pgd, unsigned long addr, unsigned long end,
  *
  * This is an internal function only. Do not use outside mm/.
  */
-void __vunmap_range_noflush(unsigned long start, unsigned long end)
+void __vunmap_range_noflush(pgd_t *pgdir, unsigned long start, unsigned long end)
 {
 	unsigned long next;
 	pgd_t *pgd;
@@ -486,7 +486,7 @@ void __vunmap_range_noflush(unsigned long start, unsigned long end)
 	pgtbl_mod_mask mask = 0;
 
 	BUG_ON(addr >= end);
-	pgd = pgd_offset_k(addr);
+	pgd = pgd_offset_pgd(pgdir, addr);
 	do {
 		next = pgd_addr_end(addr, end);
 		if (pgd_bad(*pgd))
@@ -503,7 +503,7 @@ void __vunmap_range_noflush(unsigned long start, unsigned long end)
 void vunmap_range_noflush(unsigned long start, unsigned long end)
 {
 	kmsan_vunmap_range_noflush(start, end);
-	__vunmap_range_noflush(start, end);
+	__vunmap_range_noflush(init_mm.pgd, start, end);
 }
 
 /**
@@ -670,9 +670,10 @@ int __vmap_pages_range_noflush(unsigned long addr, unsigned long end,
 	for (i = 0; i < nr; i += 1U << (page_shift - PAGE_SHIFT)) {
 		int err;
 
-		err = vmap_range_noflush(addr, addr + (1UL << page_shift),
-					page_to_phys(pages[i]), prot,
-					page_shift);
+		err = vmap_range_noflush(init_mm.pgd, addr,
+					 addr + (1UL << page_shift),
+					 page_to_phys(pages[i]), prot,
+					 page_shift);
 		if (err)
 			return err;
 
-- 
2.47.0

[PATCH 05/16] arm64: mm: enable percpu kernel page table

[Yang Shi]

Currently all cpus share the same kernel page table (swapper_pg_dir),
this patch creates kernel page table for each cpu and each cpu uses its
own kernel page table.  The cpu 0 keeps using swapper_pg_dir.  All the
kernel page tables share the same content.  So we don't have to
duplicate the whole page table for all cpus, we just need to have
different pgd page for each cpu.  All kernel page table modification (split,
creation, deletion, etc) actually still happens on swapper_pg_dir, the
modification needs to be synchronized to other cpu's page tables when
the top level is modified.

The percpu page table can't be shared across cores, so clear cnp bit
too even though CNP is supported.

Some features may not work with it for now, for example, memory hotplug,
KASAN, etc.  They will be fixed in the following patches.

Signed-off-by: Yang Shi <yang@os.amperecomputing.com>
---
 arch/arm64/include/asm/mmu.h         |  1 +
 arch/arm64/include/asm/mmu_context.h |  9 ++++-
 arch/arm64/include/asm/pgtable.h     | 10 ++++++
 arch/arm64/kernel/setup.c            |  3 ++
 arch/arm64/kernel/smp.c              |  8 +++++
 arch/arm64/mm/mmu.c                  | 53 ++++++++++++++++++++++++++++
 6 files changed, 83 insertions(+), 1 deletion(-)

diff --git a/arch/arm64/include/asm/mmu.h b/arch/arm64/include/asm/mmu.h
index 5e1211c540ab..8ed3b5f3cf84 100644
--- a/arch/arm64/include/asm/mmu.h
+++ b/arch/arm64/include/asm/mmu.h
@@ -63,6 +63,7 @@ static inline bool arm64_kernel_unmapped_at_el0(void)
 extern void arm64_memblock_init(void);
 extern void paging_init(void);
 extern void bootmem_init(void);
+extern void setup_percpu_pgd(void);
 extern void create_mapping_noalloc(phys_addr_t phys, unsigned long virt,
 				   phys_addr_t size, pgprot_t prot);
 extern void create_pgd_mapping(struct mm_struct *mm, phys_addr_t phys,
diff --git a/arch/arm64/include/asm/mmu_context.h b/arch/arm64/include/asm/mmu_context.h
index 803b68758152..0e991b98ed9e 100644
--- a/arch/arm64/include/asm/mmu_context.h
+++ b/arch/arm64/include/asm/mmu_context.h
@@ -27,6 +27,7 @@
 #include <asm/tlbflush.h>
 
 extern bool rodata_full;
+extern pgd_t *percpu_pgd[NR_CPUS];
 
 static inline void contextidr_thread_switch(struct task_struct *next)
 {
@@ -138,7 +139,13 @@ void __cpu_replace_ttbr1(pgd_t *pgdp, bool cnp);
 
 static inline void cpu_enable_swapper_cnp(void)
 {
-	__cpu_replace_ttbr1(lm_alias(swapper_pg_dir), true);
+	unsigned int cpu = smp_processor_id();
+	pgd_t *ttbr1 = percpu_pgd[cpu];
+
+	if (cpu == 0)
+		ttbr1 = lm_alias(swapper_pg_dir);
+
+	__cpu_replace_ttbr1(ttbr1, false);
 }
 
 static inline void cpu_replace_ttbr1(pgd_t *pgdp)
diff --git a/arch/arm64/include/asm/pgtable.h b/arch/arm64/include/asm/pgtable.h
index 27689c62bd25..13e069b01600 100644
--- a/arch/arm64/include/asm/pgtable.h
+++ b/arch/arm64/include/asm/pgtable.h
@@ -1216,6 +1216,16 @@ p4d_t *p4d_offset_lockless_folded(pgd_t *pgdp, pgd_t pgd, unsigned long addr)
 
 #endif  /* CONFIG_PGTABLE_LEVELS > 4 */
 
+#if CONFIG_PGTABLE_LEVELS == 2
+#define ARCH_PAGE_TABLE_SYNC_MASK      PGTBL_PMD_MODIFIED
+#elif CONFIG_PGTABLE_LEVELS == 3
+#define ARCH_PAGE_TABLE_SYNC_MASK      PGTBL_PUD_MODIFIED
+#elif CONFIG_PGTABLE_LEVELS == 4
+#define ARCH_PAGE_TABLE_SYNC_MASK      PGTBL_P4D_MODIFIED
+#elif CONFIG_PGTABLE_LEVELS == 5
+#define ARCH_PAGE_TABLE_SYNC_MASK      PGTBL_PGD_MODIFIED
+#endif
+
 #define pgd_ERROR(e)	\
 	pr_err("%s:%d: bad pgd %016llx.\n", __FILE__, __LINE__, pgd_val(e))
 
diff --git a/arch/arm64/kernel/setup.c b/arch/arm64/kernel/setup.c
index 23c05dc7a8f2..6d420ad59af4 100644
--- a/arch/arm64/kernel/setup.c
+++ b/arch/arm64/kernel/setup.c
@@ -360,6 +360,9 @@ void __init __no_sanitize_address setup_arch(char **cmdline_p)
 	smp_init_cpus();
 	smp_build_mpidr_hash();
 
+	/* Must be called after smp_init_cpus */
+	setup_percpu_pgd();
+
 #ifdef CONFIG_ARM64_SW_TTBR0_PAN
 	/*
 	 * Make sure init_thread_info.ttbr0 always generates translation
diff --git a/arch/arm64/kernel/smp.c b/arch/arm64/kernel/smp.c
index 28ff889658d3..49524a7176b2 100644
--- a/arch/arm64/kernel/smp.c
+++ b/arch/arm64/kernel/smp.c
@@ -55,6 +55,8 @@
 
 #include <trace/events/ipi.h>
 
+extern void idmap_cpu_replace_ttbr1(phys_addr_t pgdir);
+
 /*
  * as from 2.5, kernels no longer have an init_tasks structure
  * so we need some other way of telling a new secondary core
@@ -198,6 +200,12 @@ asmlinkage notrace void secondary_start_kernel(void)
 	struct mm_struct *mm = &init_mm;
 	const struct cpu_operations *ops;
 	unsigned int cpu = smp_processor_id();
+	typedef void (ttbr_replace_func)(phys_addr_t);
+	ttbr_replace_func *replace_ttbr;
+
+	phys_addr_t ttbr1 = phys_to_ttbr(virt_to_phys(percpu_pgd[cpu]));
+	replace_ttbr = (void *)__pa_symbol(idmap_cpu_replace_ttbr1);
+	replace_ttbr(ttbr1);
 
 	/*
 	 * All kernel threads share the same mm context; grab a
diff --git a/arch/arm64/mm/mmu.c b/arch/arm64/mm/mmu.c
index f2be501468ce..5a235491f268 100644
--- a/arch/arm64/mm/mmu.c
+++ b/arch/arm64/mm/mmu.c
@@ -67,6 +67,59 @@ long __section(".mmuoff.data.write") __early_cpu_boot_status;
 static DEFINE_SPINLOCK(swapper_pgdir_lock);
 static DEFINE_MUTEX(fixmap_lock);
 
+pgd_t *percpu_pgd[NR_CPUS] __ro_after_init;
+bool percpu_pgd_setup_done __ro_after_init = false;
+
+void __init setup_percpu_pgd(void)
+{
+	int cpu;
+
+	for_each_possible_cpu(cpu) {
+		void *addr;
+
+		if (cpu == 0) {
+			percpu_pgd[cpu] = swapper_pg_dir;
+			continue;
+		}
+
+		addr = memblock_alloc(PAGE_SIZE, PAGE_SIZE);
+		if (!addr)
+			panic("Can't alloc percpu pgd\n");
+
+		memcpy(addr, (void *)swapper_pg_dir, PAGE_SIZE);
+		percpu_pgd[cpu] = (pgd_t *)addr;
+	}
+
+	dsb(ishst);
+
+	percpu_pgd_setup_done = true;
+}
+
+void arch_sync_kernel_mappings(unsigned long start, unsigned long end)
+{
+	unsigned long addr, next;
+	int cpu;
+	pgd_t *pgdp = pgd_offset_k(start);
+	pgd_t pgd;
+	unsigned int index = pgd_index(start);
+
+	BUG_ON(start > end);
+
+	if (!percpu_pgd_setup_done)
+		return;
+
+	addr = start;
+	do {
+		pgd = READ_ONCE(*pgdp);
+		next = pgd_addr_end(addr, end);
+		for_each_possible_cpu(cpu) {
+			if (cpu == 0)
+				continue;
+			set_pgd(percpu_pgd[cpu] + index, pgd);
+		}
+	} while (pgdp++, index++, addr = next, addr != end);
+}
+
 void noinstr set_swapper_pgd(pgd_t *pgdp, pgd_t pgd)
 {
 	pgd_t *fixmap_pgdp;
-- 
2.47.0

[PATCH 06/16] arm64: mm: defined {pud|pmd}_populate_sync()

[Yang Shi]

Implement ARM64 specific {pud|pmd}_populate_sync(), which sync up page
table.

Signed-off-by: Yang Shi <yang@os.amperecomputing.com>
---
 arch/arm64/Kconfig               |  2 ++
 arch/arm64/include/asm/pgalloc.h | 24 ++++++++++++++++++++++++
 2 files changed, 26 insertions(+)

diff --git a/arch/arm64/Kconfig b/arch/arm64/Kconfig
index 2e1bd79a579c..23cad9fc7046 100644
--- a/arch/arm64/Kconfig
+++ b/arch/arm64/Kconfig
@@ -41,9 +41,11 @@ config ARM64
 	select ARCH_HAS_NMI_SAFE_THIS_CPU_OPS
 	select ARCH_HAS_NON_OVERLAPPING_ADDRESS_SPACE
 	select ARCH_HAS_NONLEAF_PMD_YOUNG if ARM64_HAFT
+	select ARCH_HAS_PMD_POPULATE_SYNC
 	select ARCH_HAS_PREEMPT_LAZY
 	select ARCH_HAS_PTDUMP
 	select ARCH_HAS_PTE_SPECIAL
+	select ARCH_HAS_PUD_POPULATE_SYNC
 	select ARCH_HAS_HW_PTE_YOUNG
 	select ARCH_HAS_SETUP_DMA_OPS
 	select ARCH_HAS_SET_DIRECT_MAP
diff --git a/arch/arm64/include/asm/pgalloc.h b/arch/arm64/include/asm/pgalloc.h
index 1b4509d3382c..dd1ac84a1160 100644
--- a/arch/arm64/include/asm/pgalloc.h
+++ b/arch/arm64/include/asm/pgalloc.h
@@ -33,11 +33,27 @@ static inline void pud_populate(struct mm_struct *mm, pud_t *pudp, pmd_t *pmdp)
 	pudval |= (mm == &init_mm) ? PUD_TABLE_UXN : PUD_TABLE_PXN;
 	__pud_populate(pudp, __pa(pmdp), pudval);
 }
+
+#ifdef CONFIG_ARCH_HAS_PUD_POPULATE_SYNC
+static inline void pud_populate_sync(unsigned long addr, pud_t *pudp, pmd_t *pmdp)
+{
+	pud_populate(&init_mm, pudp, pmdp);
+	if (ARCH_PAGE_TABLE_SYNC_MASK & PGTBL_PUD_MODIFIED)
+		arch_sync_kernel_mappings(addr, addr);
+}
+#endif
 #else
 static inline void __pud_populate(pud_t *pudp, phys_addr_t pmdp, pudval_t prot)
 {
 	BUILD_BUG();
 }
+
+#ifdef CONFIG_ARCH_HAS_PUD_POPULATE_SYNC
+static inline void pud_populate_sync(unsigned long addr, pud_t *pudp, pmd_t *pmdp)
+{
+	BUILD_BUG();
+}
+#endif
 #endif	/* CONFIG_PGTABLE_LEVELS > 2 */
 
 #if CONFIG_PGTABLE_LEVELS > 3
@@ -121,4 +137,12 @@ pmd_populate(struct mm_struct *mm, pmd_t *pmdp, pgtable_t ptep)
 		       PMD_TYPE_TABLE | PMD_TABLE_AF | PMD_TABLE_PXN);
 }
 
+#ifdef CONFIG_ARCH_HAS_PMD_POPULATE_SYNC
+static inline void pmd_populate_sync(unsigned long addr, pmd_t *pmdp, pte_t *ptep)
+{
+	pmd_populate_kernel(&init_mm, pmdp, ptep);
+	if (ARCH_PAGE_TABLE_SYNC_MASK & PGTBL_PMD_MODIFIED)
+		arch_sync_kernel_mappings(addr, addr);
+}
+#endif
 #endif
-- 
2.47.0

[PATCH 07/16] arm64: mm: sync percpu page table for memory hotplug/unplug

[Yang Shi]

Need to sync kernel percpu page table for memory hotplug/unplug,
including both linear mapping and vmemmap.

Signed-off-by: Yang Shi <yang@os.amperecomputing.com>
---
 arch/arm64/mm/mmu.c | 85 ++++++++++++++++++++++++++++++---------------
 1 file changed, 57 insertions(+), 28 deletions(-)

diff --git a/arch/arm64/mm/mmu.c b/arch/arm64/mm/mmu.c
index 5a235491f268..e1290ccaa5be 100644
--- a/arch/arm64/mm/mmu.c
+++ b/arch/arm64/mm/mmu.c
@@ -214,7 +214,7 @@ static void init_clear_pgtable(void *table)
 }
 
 static void init_pte(pte_t *ptep, unsigned long addr, unsigned long end,
-		     phys_addr_t phys, pgprot_t prot)
+		     phys_addr_t phys, pgprot_t prot, pgtbl_mod_mask *mask)
 {
 	do {
 		pte_t old_pte = __ptep_get(ptep);
@@ -224,6 +224,7 @@ static void init_pte(pte_t *ptep, unsigned long addr, unsigned long end,
 		 * are deferred to the end of alloc_init_cont_pte().
 		 */
 		__set_pte_nosync(ptep, pfn_pte(__phys_to_pfn(phys), prot));
+		*mask |= PGTBL_PTE_MODIFIED;
 
 		/*
 		 * After the PTE entry has been populated once, we
@@ -251,7 +252,7 @@ static int alloc_init_cont_pte(pmd_t *pmdp, unsigned long addr,
 			       unsigned long end, phys_addr_t phys,
 			       pgprot_t prot,
 			       phys_addr_t (*pgtable_alloc)(enum pgtable_level),
-			       int flags)
+			       int flags, pgtbl_mod_mask *mask)
 {
 	unsigned long next;
 	pmd_t pmd = READ_ONCE(*pmdp);
@@ -272,6 +273,7 @@ static int alloc_init_cont_pte(pmd_t *pmdp, unsigned long addr,
 		init_clear_pgtable(ptep);
 		ptep += pte_index(addr);
 		__pmd_populate(pmdp, pte_phys, pmdval);
+		*mask |= PGTBL_PMD_MODIFIED;
 	} else {
 		BUG_ON(pmd_bad(pmd));
 		ptep = pte_set_fixmap_offset(pmdp, addr);
@@ -288,7 +290,7 @@ static int alloc_init_cont_pte(pmd_t *pmdp, unsigned long addr,
 		    !pte_range_has_valid_noncont(ptep))
 			__prot = __pgprot(pgprot_val(prot) | PTE_CONT);
 
-		init_pte(ptep, addr, next, phys, __prot);
+		init_pte(ptep, addr, next, phys, __prot, mask);
 
 		ptep += pte_index(next) - pte_index(addr);
 		phys += next - addr;
@@ -306,7 +308,8 @@ static int alloc_init_cont_pte(pmd_t *pmdp, unsigned long addr,
 
 static int init_pmd(pmd_t *pmdp, unsigned long addr, unsigned long end,
 		    phys_addr_t phys, pgprot_t prot,
-		    phys_addr_t (*pgtable_alloc)(enum pgtable_level), int flags)
+		    phys_addr_t (*pgtable_alloc)(enum pgtable_level), int flags,
+		    pgtbl_mod_mask *mask)
 {
 	unsigned long next;
 
@@ -319,7 +322,10 @@ static int init_pmd(pmd_t *pmdp, unsigned long addr, unsigned long end,
 		if (((addr | next | phys) & ~PMD_MASK) == 0 &&
 		    (flags & NO_BLOCK_MAPPINGS) == 0 &&
 		    !pmd_table(old_pmd)) {
-			WARN_ON(!pmd_set_huge(pmdp, phys, prot));
+			if (pmd_set_huge(pmdp, phys, prot))
+				*mask |= PGTBL_PMD_MODIFIED;
+			else
+				WARN_ON(1);
 
 			/*
 			 * After the PMD entry has been populated once, we
@@ -331,7 +337,7 @@ static int init_pmd(pmd_t *pmdp, unsigned long addr, unsigned long end,
 			int ret;
 
 			ret = alloc_init_cont_pte(pmdp, addr, next, phys, prot,
-						  pgtable_alloc, flags);
+						  pgtable_alloc, flags, mask);
 			if (ret)
 				return ret;
 
@@ -359,7 +365,7 @@ static int alloc_init_cont_pmd(pud_t *pudp, unsigned long addr,
 			       unsigned long end, phys_addr_t phys,
 			       pgprot_t prot,
 			       phys_addr_t (*pgtable_alloc)(enum pgtable_level),
-			       int flags)
+			       int flags, pgtbl_mod_mask *mask)
 {
 	int ret;
 	unsigned long next;
@@ -384,6 +390,7 @@ static int alloc_init_cont_pmd(pud_t *pudp, unsigned long addr,
 		init_clear_pgtable(pmdp);
 		pmdp += pmd_index(addr);
 		__pud_populate(pudp, pmd_phys, pudval);
+		*mask |= PGTBL_PUD_MODIFIED;
 	} else {
 		BUG_ON(pud_bad(pud));
 		pmdp = pmd_set_fixmap_offset(pudp, addr);
@@ -400,7 +407,8 @@ static int alloc_init_cont_pmd(pud_t *pudp, unsigned long addr,
 		    !pmd_range_has_valid_noncont(pmdp))
 			__prot = __pgprot(pgprot_val(prot) | PTE_CONT);
 
-		ret = init_pmd(pmdp, addr, next, phys, __prot, pgtable_alloc, flags);
+		ret = init_pmd(pmdp, addr, next, phys, __prot, pgtable_alloc, flags,
+			       mask);
 		if (ret)
 			goto out;
 
@@ -417,7 +425,7 @@ static int alloc_init_cont_pmd(pud_t *pudp, unsigned long addr,
 static int alloc_init_pud(p4d_t *p4dp, unsigned long addr, unsigned long end,
 			  phys_addr_t phys, pgprot_t prot,
 			  phys_addr_t (*pgtable_alloc)(enum pgtable_level),
-			  int flags)
+			  int flags, pgtbl_mod_mask *mask)
 {
 	int ret = 0;
 	unsigned long next;
@@ -438,6 +446,7 @@ static int alloc_init_pud(p4d_t *p4dp, unsigned long addr, unsigned long end,
 		init_clear_pgtable(pudp);
 		pudp += pud_index(addr);
 		__p4d_populate(p4dp, pud_phys, p4dval);
+		*mask |= PGTBL_P4D_MODIFIED;
 	} else {
 		BUG_ON(p4d_bad(p4d));
 		pudp = pud_set_fixmap_offset(p4dp, addr);
@@ -455,7 +464,10 @@ static int alloc_init_pud(p4d_t *p4dp, unsigned long addr, unsigned long end,
 		   ((addr | next | phys) & ~PUD_MASK) == 0 &&
 		    (flags & NO_BLOCK_MAPPINGS) == 0 &&
 		    !pud_table(old_pud)) {
-			WARN_ON(!pud_set_huge(pudp, phys, prot));
+			if (pud_set_huge(pudp, phys, prot))
+				*mask |= PGTBL_PUD_MODIFIED;
+			else
+				WARN_ON(1);
 
 			/*
 			 * After the PUD entry has been populated once, we
@@ -465,7 +477,7 @@ static int alloc_init_pud(p4d_t *p4dp, unsigned long addr, unsigned long end,
 						      READ_ONCE(pud_val(*pudp))));
 		} else {
 			ret = alloc_init_cont_pmd(pudp, addr, next, phys, prot,
-						  pgtable_alloc, flags);
+						  pgtable_alloc, flags, mask);
 			if (ret)
 				goto out;
 
@@ -484,7 +496,7 @@ static int alloc_init_pud(p4d_t *p4dp, unsigned long addr, unsigned long end,
 static int alloc_init_p4d(pgd_t *pgdp, unsigned long addr, unsigned long end,
 			  phys_addr_t phys, pgprot_t prot,
 			  phys_addr_t (*pgtable_alloc)(enum pgtable_level),
-			  int flags)
+			  int flags, pgtbl_mod_mask *mask)
 {
 	int ret;
 	unsigned long next;
@@ -505,6 +517,7 @@ static int alloc_init_p4d(pgd_t *pgdp, unsigned long addr, unsigned long end,
 		init_clear_pgtable(p4dp);
 		p4dp += p4d_index(addr);
 		__pgd_populate(pgdp, p4d_phys, pgdval);
+		*mask |= PGTBL_PGD_MODIFIED;
 	} else {
 		BUG_ON(pgd_bad(pgd));
 		p4dp = p4d_set_fixmap_offset(pgdp, addr);
@@ -516,7 +529,7 @@ static int alloc_init_p4d(pgd_t *pgdp, unsigned long addr, unsigned long end,
 		next = p4d_addr_end(addr, end);
 
 		ret = alloc_init_pud(p4dp, addr, next, phys, prot,
-				     pgtable_alloc, flags);
+				     pgtable_alloc, flags, mask);
 		if (ret)
 			goto out;
 
@@ -536,7 +549,7 @@ static int __create_pgd_mapping_locked(pgd_t *pgdir, phys_addr_t phys,
 				       unsigned long virt, phys_addr_t size,
 				       pgprot_t prot,
 				       phys_addr_t (*pgtable_alloc)(enum pgtable_level),
-				       int flags)
+				       int flags, pgtbl_mod_mask *mask)
 {
 	int ret;
 	unsigned long addr, end, next;
@@ -556,7 +569,7 @@ static int __create_pgd_mapping_locked(pgd_t *pgdir, phys_addr_t phys,
 	do {
 		next = pgd_addr_end(addr, end);
 		ret = alloc_init_p4d(pgdp, addr, next, phys, prot, pgtable_alloc,
-				     flags);
+				     flags, mask);
 		if (ret)
 			return ret;
 		phys += next - addr;
@@ -569,13 +582,13 @@ static int __create_pgd_mapping(pgd_t *pgdir, phys_addr_t phys,
 				unsigned long virt, phys_addr_t size,
 				pgprot_t prot,
 				phys_addr_t (*pgtable_alloc)(enum pgtable_level),
-				int flags)
+				int flags, pgtbl_mod_mask *mask)
 {
 	int ret;
 
 	mutex_lock(&fixmap_lock);
 	ret = __create_pgd_mapping_locked(pgdir, phys, virt, size, prot,
-					  pgtable_alloc, flags);
+					  pgtable_alloc, flags, mask);
 	mutex_unlock(&fixmap_lock);
 
 	return ret;
@@ -588,9 +601,10 @@ static void early_create_pgd_mapping(pgd_t *pgdir, phys_addr_t phys,
 				     int flags)
 {
 	int ret;
+	pgtbl_mod_mask mask = 0;
 
 	ret = __create_pgd_mapping(pgdir, phys, virt, size, prot, pgtable_alloc,
-				   flags);
+				   flags, &mask);
 	if (ret)
 		panic("Failed to create page tables\n");
 }
@@ -1332,6 +1346,7 @@ static int __init __kpti_install_ng_mappings(void *__unused)
 
 	if (!cpu) {
 		int ret;
+		pgtbl_mod_mask mask = 0;
 
 		alloc = __get_free_pages(GFP_ATOMIC | __GFP_ZERO, order);
 		kpti_ng_temp_pgd = (pgd_t *)(alloc + (levels - 1) * PAGE_SIZE);
@@ -1355,7 +1370,7 @@ static int __init __kpti_install_ng_mappings(void *__unused)
 		//
 		ret = __create_pgd_mapping_locked(kpti_ng_temp_pgd, __pa(alloc),
 						  KPTI_NG_TEMP_VA, PAGE_SIZE, PAGE_KERNEL,
-						  kpti_ng_pgd_alloc, 0);
+						  kpti_ng_pgd_alloc, 0, &mask);
 		if (ret)
 			panic("Failed to create page tables\n");
 	}
@@ -1661,7 +1676,7 @@ static void unmap_hotplug_range(unsigned long addr, unsigned long end,
 
 static void free_empty_pte_table(pmd_t *pmdp, unsigned long addr,
 				 unsigned long end, unsigned long floor,
-				 unsigned long ceiling)
+				 unsigned long ceiling, pgtbl_mod_mask *mask)
 {
 	pte_t *ptep, pte;
 	unsigned long i, start = addr;
@@ -1677,6 +1692,7 @@ static void free_empty_pte_table(pmd_t *pmdp, unsigned long addr,
 		WARN_ON(!pte_none(pte));
 	} while (addr += PAGE_SIZE, addr < end);
 
+	*mask |= PGTBL_PTE_MODIFIED;
 	if (!pgtable_range_aligned(start, end, floor, ceiling, PMD_MASK))
 		return;
 
@@ -1692,13 +1708,14 @@ static void free_empty_pte_table(pmd_t *pmdp, unsigned long addr,
 	}
 
 	pmd_clear(pmdp);
+	*mask |= PGTBL_PMD_MODIFIED;
 	__flush_tlb_kernel_pgtable(start);
 	free_hotplug_pgtable_page(virt_to_page(ptep));
 }
 
 static void free_empty_pmd_table(pud_t *pudp, unsigned long addr,
 				 unsigned long end, unsigned long floor,
-				 unsigned long ceiling)
+				 unsigned long ceiling, pgtbl_mod_mask *mask)
 {
 	pmd_t *pmdp, pmd;
 	unsigned long i, next, start = addr;
@@ -1711,7 +1728,7 @@ static void free_empty_pmd_table(pud_t *pudp, unsigned long addr,
 			continue;
 
 		WARN_ON(!pmd_present(pmd) || !pmd_table(pmd));
-		free_empty_pte_table(pmdp, addr, next, floor, ceiling);
+		free_empty_pte_table(pmdp, addr, next, floor, ceiling, mask);
 	} while (addr = next, addr < end);
 
 	if (CONFIG_PGTABLE_LEVELS <= 2)
@@ -1732,13 +1749,14 @@ static void free_empty_pmd_table(pud_t *pudp, unsigned long addr,
 	}
 
 	pud_clear(pudp);
+	*mask |= PGTBL_PUD_MODIFIED;
 	__flush_tlb_kernel_pgtable(start);
 	free_hotplug_pgtable_page(virt_to_page(pmdp));
 }
 
 static void free_empty_pud_table(p4d_t *p4dp, unsigned long addr,
 				 unsigned long end, unsigned long floor,
-				 unsigned long ceiling)
+				 unsigned long ceiling, pgtbl_mod_mask *mask)
 {
 	pud_t *pudp, pud;
 	unsigned long i, next, start = addr;
@@ -1751,7 +1769,7 @@ static void free_empty_pud_table(p4d_t *p4dp, unsigned long addr,
 			continue;
 
 		WARN_ON(!pud_present(pud) || !pud_table(pud));
-		free_empty_pmd_table(pudp, addr, next, floor, ceiling);
+		free_empty_pmd_table(pudp, addr, next, floor, ceiling, mask);
 	} while (addr = next, addr < end);
 
 	if (!pgtable_l4_enabled())
@@ -1772,13 +1790,14 @@ static void free_empty_pud_table(p4d_t *p4dp, unsigned long addr,
 	}
 
 	p4d_clear(p4dp);
+	*mask |= PGTBL_P4D_MODIFIED;
 	__flush_tlb_kernel_pgtable(start);
 	free_hotplug_pgtable_page(virt_to_page(pudp));
 }
 
 static void free_empty_p4d_table(pgd_t *pgdp, unsigned long addr,
 				 unsigned long end, unsigned long floor,
-				 unsigned long ceiling)
+				 unsigned long ceiling, pgtbl_mod_mask *mask)
 {
 	p4d_t *p4dp, p4d;
 	unsigned long i, next, start = addr;
@@ -1791,7 +1810,7 @@ static void free_empty_p4d_table(pgd_t *pgdp, unsigned long addr,
 			continue;
 
 		WARN_ON(!p4d_present(p4d));
-		free_empty_pud_table(p4dp, addr, next, floor, ceiling);
+		free_empty_pud_table(p4dp, addr, next, floor, ceiling, mask);
 	} while (addr = next, addr < end);
 
 	if (!pgtable_l5_enabled())
@@ -1812,6 +1831,7 @@ static void free_empty_p4d_table(pgd_t *pgdp, unsigned long addr,
 	}
 
 	pgd_clear(pgdp);
+	*mask |= PGTBL_PGD_MODIFIED;
 	__flush_tlb_kernel_pgtable(start);
 	free_hotplug_pgtable_page(virt_to_page(p4dp));
 }
@@ -1821,6 +1841,8 @@ static void free_empty_tables(unsigned long addr, unsigned long end,
 {
 	unsigned long next;
 	pgd_t *pgdp, pgd;
+	pgtbl_mod_mask mask = 0;
+	unsigned long start = addr;
 
 	do {
 		next = pgd_addr_end(addr, end);
@@ -1830,8 +1852,11 @@ static void free_empty_tables(unsigned long addr, unsigned long end,
 			continue;
 
 		WARN_ON(!pgd_present(pgd));
-		free_empty_p4d_table(pgdp, addr, next, floor, ceiling);
+		free_empty_p4d_table(pgdp, addr, next, floor, ceiling, &mask);
 	} while (addr = next, addr < end);
+
+	if (mask & ARCH_PAGE_TABLE_SYNC_MASK)
+		arch_sync_kernel_mappings(start, end);
 }
 #endif
 
@@ -2039,6 +2064,7 @@ int arch_add_memory(int nid, u64 start, u64 size,
 		    struct mhp_params *params)
 {
 	int ret, flags = NO_EXEC_MAPPINGS;
+	pgtbl_mod_mask mask = 0;
 
 	VM_BUG_ON(!mhp_range_allowed(start, size, true));
 
@@ -2047,10 +2073,13 @@ int arch_add_memory(int nid, u64 start, u64 size,
 
 	ret = __create_pgd_mapping(swapper_pg_dir, start, __phys_to_virt(start),
 				   size, params->pgprot, pgd_pgtable_alloc_init_mm,
-				   flags);
+				   flags, &mask);
 	if (ret)
 		goto err;
 
+	if (mask & ARCH_PAGE_TABLE_SYNC_MASK)
+		arch_sync_kernel_mappings(start, start + size);
+
 	memblock_clear_nomap(start, size);
 
 	ret = __add_pages(nid, start >> PAGE_SHIFT, size >> PAGE_SHIFT,
-- 
2.47.0

[PATCH 08/16] arm64: kasan: sync up kasan shadow area page table

[Yang Shi]

Allocating vmalloc and percpu needs to populate kasan shadow area page
table, needs to sync up to percpu page table.

Signed-off-by: Yang Shi <yang@os.amperecomputing.com>
---
 arch/arm64/mm/kasan_init.c | 45 ++++++++++++++++++++++++--------------
 1 file changed, 29 insertions(+), 16 deletions(-)

diff --git a/arch/arm64/mm/kasan_init.c b/arch/arm64/mm/kasan_init.c
index 3fcad956fdf7..960020343760 100644
--- a/arch/arm64/mm/kasan_init.c
+++ b/arch/arm64/mm/kasan_init.c
@@ -60,13 +60,14 @@ static phys_addr_t __init kasan_alloc_raw_page(int node)
 }
 
 static pte_t *__init kasan_pte_offset(pmd_t *pmdp, unsigned long addr, int node,
-				      bool early)
+				      bool early, pgtbl_mod_mask *mask)
 {
 	if (pmd_none(READ_ONCE(*pmdp))) {
 		phys_addr_t pte_phys = early ?
 				__pa_symbol(kasan_early_shadow_pte)
 					: kasan_alloc_zeroed_page(node);
 		__pmd_populate(pmdp, pte_phys, PMD_TYPE_TABLE);
+		*mask |= PGTBL_PMD_MODIFIED;
 	}
 
 	return early ? pte_offset_kimg(pmdp, addr)
@@ -74,49 +75,53 @@ static pte_t *__init kasan_pte_offset(pmd_t *pmdp, unsigned long addr, int node,
 }
 
 static pmd_t *__init kasan_pmd_offset(pud_t *pudp, unsigned long addr, int node,
-				      bool early)
+				      bool early, pgtbl_mod_mask *mask)
 {
 	if (pud_none(READ_ONCE(*pudp))) {
 		phys_addr_t pmd_phys = early ?
 				__pa_symbol(kasan_early_shadow_pmd)
 					: kasan_alloc_zeroed_page(node);
 		__pud_populate(pudp, pmd_phys, PUD_TYPE_TABLE);
+		*mask |= PGTBL_PUD_MODIFIED;
 	}
 
 	return early ? pmd_offset_kimg(pudp, addr) : pmd_offset(pudp, addr);
 }
 
 static pud_t *__init kasan_pud_offset(p4d_t *p4dp, unsigned long addr, int node,
-				      bool early)
+				      bool early, pgtbl_mod_mask *mask)
 {
 	if (p4d_none(READ_ONCE(*p4dp))) {
 		phys_addr_t pud_phys = early ?
 				__pa_symbol(kasan_early_shadow_pud)
 					: kasan_alloc_zeroed_page(node);
 		__p4d_populate(p4dp, pud_phys, P4D_TYPE_TABLE);
+		*mask |= PGTBL_P4D_MODIFIED;
 	}
 
 	return early ? pud_offset_kimg(p4dp, addr) : pud_offset(p4dp, addr);
 }
 
 static p4d_t *__init kasan_p4d_offset(pgd_t *pgdp, unsigned long addr, int node,
-				      bool early)
+				      bool early, pgtbl_mod_mask *mask)
 {
 	if (pgd_none(READ_ONCE(*pgdp))) {
 		phys_addr_t p4d_phys = early ?
 				__pa_symbol(kasan_early_shadow_p4d)
 					: kasan_alloc_zeroed_page(node);
 		__pgd_populate(pgdp, p4d_phys, PGD_TYPE_TABLE);
+		*mask |= PGTBL_PGD_MODIFIED;
 	}
 
 	return early ? p4d_offset_kimg(pgdp, addr) : p4d_offset(pgdp, addr);
 }
 
 static void __init kasan_pte_populate(pmd_t *pmdp, unsigned long addr,
-				      unsigned long end, int node, bool early)
+				      unsigned long end, int node, bool early,
+				      pgtbl_mod_mask *mask)
 {
 	unsigned long next;
-	pte_t *ptep = kasan_pte_offset(pmdp, addr, node, early);
+	pte_t *ptep = kasan_pte_offset(pmdp, addr, node, early, mask);
 
 	do {
 		phys_addr_t page_phys = early ?
@@ -130,38 +135,41 @@ static void __init kasan_pte_populate(pmd_t *pmdp, unsigned long addr,
 }
 
 static void __init kasan_pmd_populate(pud_t *pudp, unsigned long addr,
-				      unsigned long end, int node, bool early)
+				      unsigned long end, int node, bool early,
+				      pgtbl_mod_mask *mask)
 {
 	unsigned long next;
-	pmd_t *pmdp = kasan_pmd_offset(pudp, addr, node, early);
+	pmd_t *pmdp = kasan_pmd_offset(pudp, addr, node, early, mask);
 
 	do {
 		next = pmd_addr_end(addr, end);
-		kasan_pte_populate(pmdp, addr, next, node, early);
+		kasan_pte_populate(pmdp, addr, next, node, early, mask);
 	} while (pmdp++, addr = next, addr != end && pmd_none(READ_ONCE(*pmdp)));
 }
 
 static void __init kasan_pud_populate(p4d_t *p4dp, unsigned long addr,
-				      unsigned long end, int node, bool early)
+				      unsigned long end, int node, bool early,
+				      pgtbl_mod_mask *mask)
 {
 	unsigned long next;
-	pud_t *pudp = kasan_pud_offset(p4dp, addr, node, early);
+	pud_t *pudp = kasan_pud_offset(p4dp, addr, node, early, mask);
 
 	do {
 		next = pud_addr_end(addr, end);
-		kasan_pmd_populate(pudp, addr, next, node, early);
+		kasan_pmd_populate(pudp, addr, next, node, early, mask);
 	} while (pudp++, addr = next, addr != end && pud_none(READ_ONCE(*pudp)));
 }
 
 static void __init kasan_p4d_populate(pgd_t *pgdp, unsigned long addr,
-				      unsigned long end, int node, bool early)
+				      unsigned long end, int node, bool early,
+				      pgtbl_mod_mask *mask)
 {
 	unsigned long next;
-	p4d_t *p4dp = kasan_p4d_offset(pgdp, addr, node, early);
+	p4d_t *p4dp = kasan_p4d_offset(pgdp, addr, node, early, mask);
 
 	do {
 		next = p4d_addr_end(addr, end);
-		kasan_pud_populate(p4dp, addr, next, node, early);
+		kasan_pud_populate(p4dp, addr, next, node, early, mask);
 	} while (p4dp++, addr = next, addr != end && p4d_none(READ_ONCE(*p4dp)));
 }
 
@@ -170,12 +178,17 @@ static void __init kasan_pgd_populate(unsigned long addr, unsigned long end,
 {
 	unsigned long next;
 	pgd_t *pgdp;
+	pgtbl_mod_mask mask = 0;
+	unsigned long start = addr;
 
 	pgdp = pgd_offset_k(addr);
 	do {
 		next = pgd_addr_end(addr, end);
-		kasan_p4d_populate(pgdp, addr, next, node, early);
+		kasan_p4d_populate(pgdp, addr, next, node, early, &mask);
 	} while (pgdp++, addr = next, addr != end);
+
+	if (mask & ARCH_PAGE_TABLE_SYNC_MASK)
+		arch_sync_kernel_mappings(start, end);
 }
 
 #if defined(CONFIG_ARM64_64K_PAGES) || CONFIG_PGTABLE_LEVELS > 4
-- 
2.47.0

[PATCH 09/16] arm64: mm: define percpu virtual space area

[Yang Shi]

The percpu allocator returns offset from percpu base address.  The
percpu base address is determined by the first chunk which is typically
in the low address of vmalloc space, however percpu varriables are
typically allocated from the high address of vmalloc space.  So the offset
could be quite big.  It may be the whole size of vmalloc space.  To support
local percpu mapping in order to optimize this_cpu_*() ops, the percpu
allocator needs to allocate memory from local percpu area too in the following
patch and the offset to local percpu base address must be same because
the offset returned by percpu allocator must be used to access both
global percpu and local percpu.

We can half vmalloc space to have either half dedicated to local percpu,
but it wastes too much address space.  So carve out dedicated global percpu
and local percpu areas.  Each area size is 2 * PGDIR_SIZE if page size is 4K.
It is 1TB with 4K page size, should be big enough for percpu.  It has different
size for 16K and 64K page size respectively.

The percpu areas are aligned to top level page table entry size in order to
just need to sync percpu page table at pgd level to minimize page table sync
overhead.

The kernel virtual address space layout now looks like:

+-----------------+
|  Linear mapping |
+-----------------+
|  Modules        |
+-----------------+
|  Vmalloc        |
+-----------------+
|  Global Percpu  |
+-----------------+
|  Local Percpu   |
+-----------------+
|  Vmemap         |
+-----------------+
|  PCI I/O        |
+-----------------+
|  Fixed map      |
+-----------------+

Signed-off-by: Yang Shi <yang@os.amperecomputing.com>
---
 arch/arm64/include/asm/pgtable.h | 27 ++++++++++++++++++++++++---
 arch/arm64/mm/kasan_init.c       |  2 +-
 arch/arm64/mm/mmu.c              |  4 ++++
 arch/arm64/mm/ptdump.c           |  4 ++++
 4 files changed, 33 insertions(+), 4 deletions(-)

diff --git a/arch/arm64/include/asm/pgtable.h b/arch/arm64/include/asm/pgtable.h
index 13e069b01600..83cd65610d7f 100644
--- a/arch/arm64/include/asm/pgtable.h
+++ b/arch/arm64/include/asm/pgtable.h
@@ -18,14 +18,35 @@
  * VMALLOC range.
  *
  * VMALLOC_START: beginning of the kernel vmalloc space
- * VMALLOC_END: extends to the available space below vmemmap
+ * VMALLOC_END: extends to the space below global percpu area
  */
 #define VMALLOC_START		(MODULES_END)
+#define VMALLOC_END		(PERCPU_START - SZ_8M)
+
+/*
+ * PERCPU range
+ *
+ * PERCPU_START: beginning of global percpu area
+ * PERCPU_END: end of global percpu area
+ * LOCAL_PERCPU_START: beginning of local percpu area
+ * LOCAL_PERCPU_END: end of local percpu area, extend to the available
+ *                   space below vmemmap
+ */
+#if defined(CONFIG_ARM64_4K_PAGES)
+#define PERCPU_SIZE		(2 * PGDIR_SIZE)
+#elif defined(CONFIG_ARM64_16K_PAGES)
+#define PERCPU_SIZE            (16 * PUD_SIZE)
+#elif defined(CONFIG_ARM64_64K_PAGES)
+#define PERCPU_SIZE            (PGDIR_SIZE)
+#endif
+#define PERCPU_START		(PERCPU_END - PERCPU_SIZE)
+#define PERCPU_END		(LOCAL_PERCPU_START)
+#define LOCAL_PERCPU_START	(LOCAL_PERCPU_END - PERCPU_SIZE)
 #if VA_BITS == VA_BITS_MIN
-#define VMALLOC_END		(VMEMMAP_START - SZ_8M)
+#define LOCAL_PERCPU_END	(ALIGN_DOWN(VMEMMAP_START, PERCPU_SIZE))
 #else
 #define VMEMMAP_UNUSED_NPAGES	((_PAGE_OFFSET(vabits_actual) - PAGE_OFFSET) >> PAGE_SHIFT)
-#define VMALLOC_END		(VMEMMAP_START + VMEMMAP_UNUSED_NPAGES * sizeof(struct page) - SZ_8M)
+#define LOCAL_PERCPU_END	(ALIGN_DOWN((VMEMMAP_START + VMEMMAP_UNUSED_NPAGES * sizeof(struct page)), PERCPU_SIZE))
 #endif
 
 #define vmemmap			((struct page *)VMEMMAP_START - (memstart_addr >> PAGE_SHIFT))
diff --git a/arch/arm64/mm/kasan_init.c b/arch/arm64/mm/kasan_init.c
index 960020343760..a8e272bfefe1 100644
--- a/arch/arm64/mm/kasan_init.c
+++ b/arch/arm64/mm/kasan_init.c
@@ -324,7 +324,7 @@ static void __init kasan_init_shadow(void)
 
 	mod_shadow_start = (u64)kasan_mem_to_shadow((void *)MODULES_VADDR);
 
-	vmalloc_shadow_end = (u64)kasan_mem_to_shadow((void *)VMALLOC_END);
+	vmalloc_shadow_end = (u64)kasan_mem_to_shadow((void *)LOCAL_PERCPU_END);
 
 	/*
 	 * We are going to perform proper setup of shadow memory.
diff --git a/arch/arm64/mm/mmu.c b/arch/arm64/mm/mmu.c
index e1290ccaa5be..f14ba512fd8f 100644
--- a/arch/arm64/mm/mmu.c
+++ b/arch/arm64/mm/mmu.c
@@ -108,6 +108,10 @@ void arch_sync_kernel_mappings(unsigned long start, unsigned long end)
 	if (!percpu_pgd_setup_done)
 		return;
 
+	/* Don't sync local percpu area page table */
+	if (start >= LOCAL_PERCPU_START && end < LOCAL_PERCPU_END)
+		return;
+
 	addr = start;
 	do {
 		pgd = READ_ONCE(*pgdp);
diff --git a/arch/arm64/mm/ptdump.c b/arch/arm64/mm/ptdump.c
index 1c20144700d7..dd293319823e 100644
--- a/arch/arm64/mm/ptdump.c
+++ b/arch/arm64/mm/ptdump.c
@@ -389,6 +389,10 @@ static int __init ptdump_init(void)
 		{ MODULES_END,		"Modules end" },
 		{ VMALLOC_START,	"vmalloc() area" },
 		{ VMALLOC_END,		"vmalloc() end" },
+		{ PERCPU_START,		"Global percpu start" },
+		{ PERCPU_END,		"Global percpu end" },
+		{ LOCAL_PERCPU_START,	"Local percpu start" },
+		{ LOCAL_PERCPU_END,	"Local percpu end" },
 		{ vmemmap_start,	"vmemmap start" },
 		{ VMEMMAP_END,		"vmemmap end" },
 		{ PCI_IO_START,		"PCI I/O start" },
-- 
2.47.0

[PATCH 10/16] mm: percpu: prepare to use dedicated percpu area

[Yang Shi]

The percpu variables are allocated from vmalloc area by default.  The
architectures which support local percpu map will allocate percpu
variables from dedicated percpu area, for example, ARM64.

Introduce a new kernel config, CONFIG_HAVE_LOCAL_PER_CPU_MAP.  The
architectures which support local percpu map will need to select it.  If
it is enabled, allocate percpu variables from the dedicated percpu area.

But still treat percpu area address as vmalloc address.

Signed-off-by: Yang Shi <yang@os.amperecomputing.com>
---
 include/linux/mm.h | 11 +++++++++++
 mm/Kconfig         |  3 +++
 mm/percpu.c        |  8 ++++++++
 mm/vmalloc.c       | 23 ++++++++++++++++++++---
 4 files changed, 42 insertions(+), 3 deletions(-)

diff --git a/include/linux/mm.h b/include/linux/mm.h
index 485df9c2dbdd..c8c3800fbd9f 100644
--- a/include/linux/mm.h
+++ b/include/linux/mm.h
@@ -1858,6 +1858,17 @@ unsigned long vmalloc_to_pfn(const void *addr);
 #ifdef CONFIG_MMU
 extern bool is_vmalloc_addr(const void *x);
 extern int is_vmalloc_or_module_addr(const void *x);
+#ifdef CONFIG_HAVE_LOCAL_PER_CPU_MAP
+static inline bool is_percpu_addr(unsigned long addr)
+{
+	return addr >= PERCPU_START && addr < LOCAL_PERCPU_END;
+}
+#else
+static inline bool is_percpu_addr(unsigned long addr)
+{
+	return false;
+}
+#endif
 #else
 static inline bool is_vmalloc_addr(const void *x)
 {
diff --git a/mm/Kconfig b/mm/Kconfig
index 13ea4b29ee7b..e1910c332638 100644
--- a/mm/Kconfig
+++ b/mm/Kconfig
@@ -1061,6 +1061,9 @@ config NEED_PER_CPU_PAGE_FIRST_CHUNK
 config USE_PERCPU_NUMA_NODE_ID
 	bool
 
+config HAVE_LOCAL_PER_CPU_MAP
+	bool
+
 config HAVE_SETUP_PER_CPU_AREA
 	bool
 
diff --git a/mm/percpu.c b/mm/percpu.c
index 45c30c6531c5..b27828602c32 100644
--- a/mm/percpu.c
+++ b/mm/percpu.c
@@ -3243,9 +3243,17 @@ int __init pcpu_page_first_chunk(size_t reserved_size, pcpu_fc_cpu_to_node_fn_t
 	}
 
 	/* allocate vm area, map the pages and copy static data */
+#ifdef CONFIG_HAVE_LOCAL_PER_CPU_MAP
+	vm.flags = VM_ALLOC;
+	vm.addr = (void *)ALIGN(PERCPU_START, PAGE_SIZE);
+	vm.size = num_possible_cpus() * ai->unit_size;
+	vm_area_add_early(&vm);
+	kasan_populate_early_vm_area_shadow(vm.addr, vm.size);
+#else
 	vm.flags = VM_ALLOC;
 	vm.size = num_possible_cpus() * ai->unit_size;
 	vm_area_register_early(&vm, PAGE_SIZE);
+#endif
 
 	for (unit = 0; unit < num_possible_cpus(); unit++) {
 		unsigned long unit_addr =
diff --git a/mm/vmalloc.c b/mm/vmalloc.c
index 17c26e9796b2..9ec401918736 100644
--- a/mm/vmalloc.c
+++ b/mm/vmalloc.c
@@ -80,6 +80,9 @@ bool is_vmalloc_addr(const void *x)
 {
 	unsigned long addr = (unsigned long)kasan_reset_tag(x);
 
+	if (is_percpu_addr(addr))
+		return true;
+
 	return addr >= VMALLOC_START && addr < VMALLOC_END;
 }
 EXPORT_SYMBOL(is_vmalloc_addr);
@@ -4926,9 +4929,15 @@ pvm_find_va_enclose_addr(unsigned long addr)
 static unsigned long
 pvm_determine_end_from_reverse(struct vmap_area **va, unsigned long align)
 {
-	unsigned long vmalloc_end = VMALLOC_END & ~(align - 1);
+	unsigned long vmalloc_end;
 	unsigned long addr;
 
+#ifdef CONFIG_HAVE_LOCAL_PER_CPU_MAP
+	vmalloc_end = PERCPU_END & ~(align - 1);
+#else
+	vmalloc_end = VMALLOC_END & ~(align - 1);
+#endif
+
 	if (likely(*va)) {
 		list_for_each_entry_from_reverse((*va),
 				&free_vmap_area_list, list) {
@@ -4969,14 +4978,22 @@ struct vm_struct **pcpu_get_vm_areas(const unsigned long *offsets,
 				     const size_t *sizes, int nr_vms,
 				     size_t align)
 {
-	const unsigned long vmalloc_start = ALIGN(VMALLOC_START, align);
-	const unsigned long vmalloc_end = VMALLOC_END & ~(align - 1);
+	unsigned long vmalloc_start;
+	unsigned long vmalloc_end;
 	struct vmap_area **vas, *va;
 	struct vm_struct **vms;
 	int area, area2, last_area, term_area;
 	unsigned long base, start, size, end, last_end, orig_start, orig_end;
 	bool purged = false;
 
+#ifdef CONFIG_HAVE_LOCAL_PER_CPU_MAP
+	vmalloc_start = ALIGN(PERCPU_START, align);
+	vmalloc_end = PERCPU_END & ~(align - 1);
+#else
+	vmalloc_start = ALIGN(VMALLOC_START, align);
+	vmalloc_end = VMALLOC_END & ~(align - 1);
+#endif
+
 	/* verify parameters and allocate data structures */
 	BUG_ON(offset_in_page(align) || !is_power_of_2(align));
 	for (last_area = 0, area = 0; area < nr_vms; area++) {
-- 
2.47.0

[PATCH 11/16] arm64: mm: map local percpu first chunk

[Yang Shi]

Allocate local percpu area and map to percpu page table for the first
chunk.

It doesn't work for PCPU_FC_EMBED because the percpu base adddress may
be in linear mapping space in this case, it will result in returning huge
offset for percpu allocator.  So percpu local map just can work with
PCPU_FC_PAGE which allocates percpu variables from vmalloc area or the
dedicated percpu area.  So unselect NEED_PER_CPU_EMBED_FIRST_CHUNK if
the architectures support percpu local map.

Signed-off-by: Yang Shi <yang@os.amperecomputing.com>
---
 arch/arm64/Kconfig           |  1 -
 arch/arm64/include/asm/mmu.h |  4 ++++
 arch/arm64/kernel/smp.c      | 25 ++-----------------------
 arch/arm64/mm/mmu.c          | 23 +++++++++++++++++++++++
 mm/percpu-internal.h         | 11 +++++++++++
 mm/percpu.c                  | 16 ++++++++++++++++
 6 files changed, 56 insertions(+), 24 deletions(-)

diff --git a/arch/arm64/Kconfig b/arch/arm64/Kconfig
index 23cad9fc7046..df55e6bdf81f 100644
--- a/arch/arm64/Kconfig
+++ b/arch/arm64/Kconfig
@@ -242,7 +242,6 @@ config ARM64
 	select LOCK_MM_AND_FIND_VMA
 	select MODULES_USE_ELF_RELA
 	select NEED_DMA_MAP_STATE
-	select NEED_PER_CPU_EMBED_FIRST_CHUNK
 	select NEED_PER_CPU_PAGE_FIRST_CHUNK
 	select NEED_SG_DMA_LENGTH
 	select OF
diff --git a/arch/arm64/include/asm/mmu.h b/arch/arm64/include/asm/mmu.h
index 8ed3b5f3cf84..aa0818f7ec58 100644
--- a/arch/arm64/include/asm/mmu.h
+++ b/arch/arm64/include/asm/mmu.h
@@ -73,6 +73,10 @@ extern void *fixmap_remap_fdt(phys_addr_t dt_phys, int *size, pgprot_t prot);
 extern void mark_linear_text_alias_ro(void);
 extern int split_kernel_leaf_mapping(unsigned long start, unsigned long end);
 extern void linear_map_maybe_split_to_ptes(void);
+#ifdef CONFIG_HAVE_LOCAL_PER_CPU_MAP
+extern void map_local_percpu_first_chunk(unsigned int cpu, unsigned long virt,
+				struct page **pages, unsigned int nr);
+#endif
 
 /*
  * This check is triggered during the early boot before the cpufeature
diff --git a/arch/arm64/kernel/smp.c b/arch/arm64/kernel/smp.c
index 49524a7176b2..4c112bab8356 100644
--- a/arch/arm64/kernel/smp.c
+++ b/arch/arm64/kernel/smp.c
@@ -840,36 +840,15 @@ int early_cpu_to_node(int cpu)
 unsigned long __per_cpu_offset[NR_CPUS] __read_mostly;
 EXPORT_SYMBOL(__per_cpu_offset);
 
-static int __init pcpu_cpu_distance(unsigned int from, unsigned int to)
-{
-	return node_distance(early_cpu_to_node(from), early_cpu_to_node(to));
-}
-
 void __init setup_per_cpu_areas(void)
 {
 	unsigned long delta;
 	unsigned int cpu;
 	int rc = -EINVAL;
 
-	if (pcpu_chosen_fc != PCPU_FC_PAGE) {
-		/*
-		 * Always reserve area for module percpu variables.  That's
-		 * what the legacy allocator did.
-		 */
-		rc = pcpu_embed_first_chunk(PERCPU_MODULE_RESERVE,
-					    PERCPU_DYNAMIC_RESERVE, PAGE_SIZE,
-					    pcpu_cpu_distance,
-					    early_cpu_to_node);
-#ifdef CONFIG_NEED_PER_CPU_PAGE_FIRST_CHUNK
-		if (rc < 0)
-			pr_warn("PERCPU: %s allocator failed (%d), falling back to page size\n",
-				   pcpu_fc_names[pcpu_chosen_fc], rc);
-#endif
-	}
-
 #ifdef CONFIG_NEED_PER_CPU_PAGE_FIRST_CHUNK
-	if (rc < 0)
-		rc = pcpu_page_first_chunk(PERCPU_MODULE_RESERVE, early_cpu_to_node);
+	/* PCPU page table just can support PCPU_FC_PAGE */
+	rc = pcpu_page_first_chunk(PERCPU_MODULE_RESERVE, early_cpu_to_node);
 #endif
 	if (rc < 0)
 		panic("Failed to initialize percpu areas (err=%d).", rc);
diff --git a/arch/arm64/mm/mmu.c b/arch/arm64/mm/mmu.c
index f14ba512fd8f..ee4ab031ebf8 100644
--- a/arch/arm64/mm/mmu.c
+++ b/arch/arm64/mm/mmu.c
@@ -1080,6 +1080,29 @@ void __init linear_map_maybe_split_to_ptes(void)
 	}
 }
 
+#ifdef CONFIG_HAVE_LOCAL_PER_CPU_MAP
+void __init map_local_percpu_first_chunk(unsigned int cpu, unsigned long virt,
+			    struct page **pages, unsigned int nr)
+{
+	int i, ret;
+	pgtbl_mod_mask mask = 0;
+
+	arch_enter_lazy_mmu_mode();
+
+	for (i = 0; i < nr; i++) {
+		phys_addr_t phys = page_to_phys(pages[i]);
+		ret = __create_pgd_mapping(percpu_pgd[cpu], phys, virt, PAGE_SIZE, PAGE_KERNEL,
+					   early_pgtable_alloc, NO_EXEC_MAPPINGS, &mask);
+		if (ret)
+			panic("Failed to create local percpu first chunk\n");
+
+		virt += PAGE_SIZE;
+	}
+
+	arch_leave_lazy_mmu_mode();
+}
+#endif
+
 /*
  * This function can only be used to modify existing table entries,
  * without allocating new levels of table. Note that this permits the
diff --git a/mm/percpu-internal.h b/mm/percpu-internal.h
index 8cbe039bf847..6d0542fb9e49 100644
--- a/mm/percpu-internal.h
+++ b/mm/percpu-internal.h
@@ -162,6 +162,17 @@ static inline size_t pcpu_obj_full_size(size_t size)
 	return size * num_possible_cpus() + extra_size;
 }
 
+#ifdef CONFIG_HAVE_LOCAL_PER_CPU_MAP
+extern void __init map_local_percpu_first_chunk(unsigned int cpu, unsigned long virt,
+                            struct page **pages, unsigned int nr);
+#else
+static inline void __init map_local_percpu_first_chunk(unsigned int cpu, unsigned long virt,
+                            struct page **pages, unsigned int nr)
+{
+	return;
+}
+#endif
+
 #ifdef CONFIG_PERCPU_STATS
 
 #include <linux/spinlock.h>
diff --git a/mm/percpu.c b/mm/percpu.c
index b27828602c32..df2daa57309b 100644
--- a/mm/percpu.c
+++ b/mm/percpu.c
@@ -3194,6 +3194,7 @@ void __init __weak pcpu_populate_pte(unsigned long addr)
 int __init pcpu_page_first_chunk(size_t reserved_size, pcpu_fc_cpu_to_node_fn_t cpu_to_nd_fn)
 {
 	static struct vm_struct vm;
+	static struct vm_struct pcpu_vm;
 	struct pcpu_alloc_info *ai;
 	char psize_str[16];
 	int unit_pages;
@@ -3249,6 +3250,12 @@ int __init pcpu_page_first_chunk(size_t reserved_size, pcpu_fc_cpu_to_node_fn_t
 	vm.size = num_possible_cpus() * ai->unit_size;
 	vm_area_add_early(&vm);
 	kasan_populate_early_vm_area_shadow(vm.addr, vm.size);
+
+	pcpu_vm.flags = VM_ALLOC;
+	pcpu_vm.addr = (void *)ALIGN(LOCAL_PERCPU_START, PAGE_SIZE);
+	pcpu_vm.size = ai->unit_size;
+	vm_area_add_early(&pcpu_vm);
+	kasan_populate_early_vm_area_shadow(pcpu_vm.addr, pcpu_vm.size);
 #else
 	vm.flags = VM_ALLOC;
 	vm.size = num_possible_cpus() * ai->unit_size;
@@ -3256,6 +3263,7 @@ int __init pcpu_page_first_chunk(size_t reserved_size, pcpu_fc_cpu_to_node_fn_t
 #endif
 
 	for (unit = 0; unit < num_possible_cpus(); unit++) {
+		unsigned int cpu = ai->groups[0].cpu_map[unit];
 		unsigned long unit_addr =
 			(unsigned long)vm.addr + unit * ai->unit_size;
 
@@ -3272,6 +3280,14 @@ int __init pcpu_page_first_chunk(size_t reserved_size, pcpu_fc_cpu_to_node_fn_t
 
 		/* copy static data */
 		memcpy((void *)unit_addr, __per_cpu_start, ai->static_size);
+
+		/*
+		 * Map percpu data to PERCPU map.
+		 *
+		 * PCPU_FC_EMBED can't support it.
+		 */
+		map_local_percpu_first_chunk(cpu, (unsigned long)pcpu_vm.addr,
+				&pages[unit * unit_pages], unit_pages);
 	}
 
 	/* we're ready, commit */
-- 
2.47.0

[PATCH 12/16] mm: percpu: set up first chunk and reserve chunk

[Yang Shi]

Set up the first chunk and reserve chunk with local percpu map.

Signed-off-by: Yang Shi <yang@os.amperecomputing.com>
---
 include/linux/percpu.h |  2 +-
 mm/percpu-internal.h   |  2 ++
 mm/percpu.c            | 27 ++++++++++++++++++---------
 3 files changed, 21 insertions(+), 10 deletions(-)

diff --git a/include/linux/percpu.h b/include/linux/percpu.h
index 4f36b8585995..2f0e190aca29 100644
--- a/include/linux/percpu.h
+++ b/include/linux/percpu.h
@@ -114,7 +114,7 @@ extern struct pcpu_alloc_info * __init pcpu_alloc_alloc_info(int nr_groups,
 extern void __init pcpu_free_alloc_info(struct pcpu_alloc_info *ai);
 
 extern void __init pcpu_setup_first_chunk(const struct pcpu_alloc_info *ai,
-					 void *base_addr);
+					 void *base_addr, void *local_base);
 
 extern int __init pcpu_embed_first_chunk(size_t reserved_size, size_t dyn_size,
 				size_t atom_size,
diff --git a/mm/percpu-internal.h b/mm/percpu-internal.h
index 6d0542fb9e49..ec2e55a2f836 100644
--- a/mm/percpu-internal.h
+++ b/mm/percpu-internal.h
@@ -63,6 +63,8 @@ struct pcpu_chunk {
 	 * chunk_md.
 	 */
 	void			*base_addr ____cacheline_aligned_in_smp;
+	/* percpu local base address of the chunk */
+	void                    *local_base;
 
 	unsigned long		*alloc_map;	/* allocation map */
 	struct pcpu_block_md	*md_blocks;	/* metadata blocks */
diff --git a/mm/percpu.c b/mm/percpu.c
index df2daa57309b..15303611ecaf 100644
--- a/mm/percpu.c
+++ b/mm/percpu.c
@@ -1338,6 +1338,7 @@ static void pcpu_init_md_blocks(struct pcpu_chunk *chunk)
 /**
  * pcpu_alloc_first_chunk - creates chunks that serve the first chunk
  * @tmp_addr: the start of the region served
+ * @local_tmp: the start of the local percpu region served
  * @map_size: size of the region served
  *
  * This is responsible for creating the chunks that serve the first chunk.  The
@@ -1349,15 +1350,16 @@ static void pcpu_init_md_blocks(struct pcpu_chunk *chunk)
  * Chunk serving the region at @tmp_addr of @map_size.
  */
 static struct pcpu_chunk * __init pcpu_alloc_first_chunk(unsigned long tmp_addr,
-							 int map_size)
+							 unsigned long local_tmp, int map_size)
 {
 	struct pcpu_chunk *chunk;
-	unsigned long aligned_addr;
+	unsigned long aligned_addr, aligned_local;
 	int start_offset, offset_bits, region_size, region_bits;
 	size_t alloc_size;
 
 	/* region calculations */
 	aligned_addr = tmp_addr & PAGE_MASK;
+	aligned_local = local_tmp & PAGE_MASK;
 
 	start_offset = tmp_addr - aligned_addr;
 	region_size = ALIGN(start_offset + map_size, PAGE_SIZE);
@@ -1370,6 +1372,7 @@ static struct pcpu_chunk * __init pcpu_alloc_first_chunk(unsigned long tmp_addr,
 	INIT_LIST_HEAD(&chunk->list);
 
 	chunk->base_addr = (void *)aligned_addr;
+	chunk->local_base = (void *)aligned_local;
 	chunk->start_offset = start_offset;
 	chunk->end_offset = region_size - chunk->start_offset - map_size;
 
@@ -2562,7 +2565,7 @@ static void pcpu_dump_alloc_info(const char *lvl,
  * and available for dynamic allocation like any other chunk.
  */
 void __init pcpu_setup_first_chunk(const struct pcpu_alloc_info *ai,
-				   void *base_addr)
+				   void *base_addr, void *local_base)
 {
 	size_t size_sum = ai->static_size + ai->reserved_size + ai->dyn_size;
 	size_t static_size, dyn_size;
@@ -2572,7 +2575,7 @@ void __init pcpu_setup_first_chunk(const struct pcpu_alloc_info *ai,
 	unsigned int cpu;
 	int *unit_map;
 	int group, unit, i;
-	unsigned long tmp_addr;
+	unsigned long tmp_addr, local_tmp = 0;
 	size_t alloc_size;
 
 #define PCPU_SETUP_BUG_ON(cond)	do {					\
@@ -2713,11 +2716,15 @@ void __init pcpu_setup_first_chunk(const struct pcpu_alloc_info *ai,
 	 *   chunk.
 	 */
 	tmp_addr = (unsigned long)base_addr + static_size;
+	if (local_base)
+		local_tmp = (unsigned long)local_base + static_size;
 	if (ai->reserved_size)
-		pcpu_reserved_chunk = pcpu_alloc_first_chunk(tmp_addr,
+		pcpu_reserved_chunk = pcpu_alloc_first_chunk(tmp_addr, local_tmp,
 						ai->reserved_size);
 	tmp_addr = (unsigned long)base_addr + static_size + ai->reserved_size;
-	pcpu_first_chunk = pcpu_alloc_first_chunk(tmp_addr, dyn_size);
+	if (local_base)
+		local_tmp = (unsigned long)local_base + static_size + ai->reserved_size;
+	pcpu_first_chunk = pcpu_alloc_first_chunk(tmp_addr, local_tmp, dyn_size);
 
 	pcpu_nr_empty_pop_pages = pcpu_first_chunk->nr_empty_pop_pages;
 	pcpu_chunk_relocate(pcpu_first_chunk, -1);
@@ -3108,7 +3115,7 @@ int __init pcpu_embed_first_chunk(size_t reserved_size, size_t dyn_size,
 		PFN_DOWN(size_sum), ai->static_size, ai->reserved_size,
 		ai->dyn_size, ai->unit_size);
 
-	pcpu_setup_first_chunk(ai, base);
+	pcpu_setup_first_chunk(ai, base, NULL);
 	goto out_free;
 
 out_free_areas:
@@ -3260,6 +3267,8 @@ int __init pcpu_page_first_chunk(size_t reserved_size, pcpu_fc_cpu_to_node_fn_t
 	vm.flags = VM_ALLOC;
 	vm.size = num_possible_cpus() * ai->unit_size;
 	vm_area_register_early(&vm, PAGE_SIZE);
+
+	pcpu_vm.addr = NULL;
 #endif
 
 	for (unit = 0; unit < num_possible_cpus(); unit++) {
@@ -3295,7 +3304,7 @@ int __init pcpu_page_first_chunk(size_t reserved_size, pcpu_fc_cpu_to_node_fn_t
 		unit_pages, psize_str, ai->static_size,
 		ai->reserved_size, ai->dyn_size);
 
-	pcpu_setup_first_chunk(ai, vm.addr);
+	pcpu_setup_first_chunk(ai, vm.addr, pcpu_vm.addr);
 	goto out_free_ar;
 
 enomem:
@@ -3377,7 +3386,7 @@ void __init setup_per_cpu_areas(void)
 	ai->groups[0].nr_units = 1;
 	ai->groups[0].cpu_map[0] = 0;
 
-	pcpu_setup_first_chunk(ai, fc);
+	pcpu_setup_first_chunk(ai, fc, NULL);
 	pcpu_free_alloc_info(ai);
 }
 
-- 
2.47.0

[PATCH 13/16] arm64: mm: introduce __per_cpu_local_off

[Yang Shi]

this_cpu_*() ops will use it to get the local percpu address.  It has
the same value for all CPUs.

Also introduce pcpu_local_base, which is the base address of local
percpu map.

Signed-off-by: Yang Shi <yang@os.amperecomputing.com>
---
 arch/arm64/kernel/smp.c | 4 ++++
 include/linux/percpu.h  | 2 ++
 mm/percpu.c             | 3 +++
 3 files changed, 9 insertions(+)

diff --git a/arch/arm64/kernel/smp.c b/arch/arm64/kernel/smp.c
index 4c112bab8356..4a811379c191 100644
--- a/arch/arm64/kernel/smp.c
+++ b/arch/arm64/kernel/smp.c
@@ -839,6 +839,9 @@ int early_cpu_to_node(int cpu)
 #ifdef CONFIG_HAVE_SETUP_PER_CPU_AREA
 unsigned long __per_cpu_offset[NR_CPUS] __read_mostly;
 EXPORT_SYMBOL(__per_cpu_offset);
+/* Used to calculate pcpu local address, the offset is same for all CPUs */
+unsigned long __per_cpu_local_off __ro_after_init;
+EXPORT_SYMBOL(__per_cpu_local_off);
 
 void __init setup_per_cpu_areas(void)
 {
@@ -856,6 +859,7 @@ void __init setup_per_cpu_areas(void)
 	delta = (unsigned long)pcpu_base_addr - (unsigned long)__per_cpu_start;
 	for_each_possible_cpu(cpu)
 		__per_cpu_offset[cpu] = delta + pcpu_unit_offsets[cpu];
+	__per_cpu_local_off = (unsigned long)pcpu_local_base - (unsigned long)__per_cpu_start;
 }
 #endif
 
diff --git a/include/linux/percpu.h b/include/linux/percpu.h
index 2f0e190aca29..5a27031dff1a 100644
--- a/include/linux/percpu.h
+++ b/include/linux/percpu.h
@@ -75,6 +75,8 @@
 
 extern void *pcpu_base_addr;
 extern const unsigned long *pcpu_unit_offsets;
+/* percpu local mapping base */
+extern void *pcpu_local_base;
 
 struct pcpu_group_info {
 	int			nr_units;	/* aligned # of units */
diff --git a/mm/percpu.c b/mm/percpu.c
index 15303611ecaf..8b273ed8e5a8 100644
--- a/mm/percpu.c
+++ b/mm/percpu.c
@@ -145,6 +145,8 @@ static unsigned int pcpu_high_unit_cpu __ro_after_init;
 
 /* the address of the first chunk which starts with the kernel static area */
 void *pcpu_base_addr __ro_after_init;
+/* The address of the first chunk local mapping */
+void *pcpu_local_base __ro_after_init;
 
 static const int *pcpu_unit_map __ro_after_init;		/* cpu -> unit */
 const unsigned long *pcpu_unit_offsets __ro_after_init;	/* cpu -> unit offset */
@@ -3305,6 +3307,7 @@ int __init pcpu_page_first_chunk(size_t reserved_size, pcpu_fc_cpu_to_node_fn_t
 		ai->reserved_size, ai->dyn_size);
 
 	pcpu_setup_first_chunk(ai, vm.addr, pcpu_vm.addr);
+	pcpu_local_base = pcpu_vm.addr;
 	goto out_free_ar;
 
 enomem:
-- 
2.47.0

[PATCH 14/16] mm: percpu: allocate and free local percpu vm area

[Yang Shi]

Allocate local percpu vm area.  The delta between the allocated addr
(chunk local base) and pcpu_local_base must be same with the delta
between chunk base and pcpu_base_addr.  Each CPU's local percpu area
will be mapped to its own page table.  This section of page table is not
shared between CPUs.

And free local percpu vm area.  Also unmap from percpu page table.

Signed-off-by: Yang Shi <yang@os.amperecomputing.com>
---
 include/linux/vmalloc.h |  3 ++
 mm/percpu-internal.h    |  1 +
 mm/percpu-vm.c          | 94 +++++++++++++++++++++++++++++++++++++++++
 mm/vmalloc.c            | 92 +++++++++++++++++++++++++++++++++++++---
 4 files changed, 185 insertions(+), 5 deletions(-)

diff --git a/include/linux/vmalloc.h b/include/linux/vmalloc.h
index d87dc7f77f4e..daf9965a111d 100644
--- a/include/linux/vmalloc.h
+++ b/include/linux/vmalloc.h
@@ -313,6 +313,9 @@ struct vm_struct **pcpu_get_vm_areas(const unsigned long *offsets,
 				     size_t align);
 
 void pcpu_free_vm_areas(struct vm_struct **vms, int nr_vms);
+struct vm_struct *pcpu_get_local_vm_area(unsigned long hint,
+				     int unit_size, size_t align);
+
 # else
 static inline struct vm_struct **
 pcpu_get_vm_areas(const unsigned long *offsets,
diff --git a/mm/percpu-internal.h b/mm/percpu-internal.h
index ec2e55a2f836..4d461305243b 100644
--- a/mm/percpu-internal.h
+++ b/mm/percpu-internal.h
@@ -70,6 +70,7 @@ struct pcpu_chunk {
 	struct pcpu_block_md	*md_blocks;	/* metadata blocks */
 
 	void			*data;		/* chunk data */
+	void			*local_data;	/* chunk local vm */
 	bool			immutable;	/* no [de]population allowed */
 	bool			isolated;	/* isolated from active chunk
 						   slots */
diff --git a/mm/percpu-vm.c b/mm/percpu-vm.c
index 4f5937090590..5decaad7640f 100644
--- a/mm/percpu-vm.c
+++ b/mm/percpu-vm.c
@@ -112,6 +112,10 @@ static int pcpu_alloc_pages(struct pcpu_chunk *chunk,
 	return -ENOMEM;
 }
 
+#ifdef CONFIG_HAVE_LOCAL_PER_CPU_MAP
+extern pgd_t *percpu_pgd[NR_CPUS];
+#endif
+
 /**
  * pcpu_pre_unmap_flush - flush cache prior to unmapping
  * @chunk: chunk the regions to be flushed belongs to
@@ -130,6 +134,11 @@ static void pcpu_pre_unmap_flush(struct pcpu_chunk *chunk,
 	flush_cache_vunmap(
 		pcpu_chunk_addr(chunk, pcpu_low_unit_cpu, page_start),
 		pcpu_chunk_addr(chunk, pcpu_high_unit_cpu, page_end));
+
+#ifdef CONFIG_HAVE_LOCAL_PER_CPU_MAP
+	flush_cache_vunmap((unsigned long)chunk->local_base + (page_start << PAGE_SHIFT),
+			(unsigned long)chunk->local_base + (page_end << PAGE_SHIFT));
+#endif
 }
 
 static void __pcpu_unmap_pages(unsigned long addr, int nr_pages)
@@ -137,6 +146,20 @@ static void __pcpu_unmap_pages(unsigned long addr, int nr_pages)
 	vunmap_range_noflush(addr, addr + (nr_pages << PAGE_SHIFT));
 }
 
+#ifdef CONFIG_HAVE_LOCAL_PER_CPU_MAP
+static void __pcpu_unmap_pages_local(unsigned int cpu, unsigned long virt,
+			    int nr_pages)
+{
+	__vunmap_range_noflush(percpu_pgd[cpu], virt, virt + (nr_pages << PAGE_SHIFT));
+}
+#else
+static void __pcpu_unmap_pages_local(unsigned int cpu, unsigned long virt,
+			    int nr_pages)
+{
+	return;
+}
+#endif
+
 /**
  * pcpu_unmap_pages - unmap pages out of a pcpu_chunk
  * @chunk: chunk of interest
@@ -166,6 +189,10 @@ static void pcpu_unmap_pages(struct pcpu_chunk *chunk,
 		}
 		__pcpu_unmap_pages(pcpu_chunk_addr(chunk, cpu, page_start),
 				   page_end - page_start);
+
+		__pcpu_unmap_pages_local(cpu,
+					(unsigned long)chunk->local_base + (page_start << PAGE_SHIFT),
+					page_end - page_start);
 	}
 }
 
@@ -188,6 +215,12 @@ static void pcpu_post_unmap_tlb_flush(struct pcpu_chunk *chunk,
 	flush_tlb_kernel_range(
 		pcpu_chunk_addr(chunk, pcpu_low_unit_cpu, page_start),
 		pcpu_chunk_addr(chunk, pcpu_high_unit_cpu, page_end));
+
+#ifdef CONFIG_HAVE_LOCAL_PER_CPU_MAP
+	flush_tlb_kernel_range(
+		(unsigned long)chunk->local_base + (page_start << PAGE_SHIFT),
+		(unsigned long)chunk->local_base + (page_end << PAGE_SHIFT));
+#endif
 }
 
 static int __pcpu_map_pages(unsigned long addr, struct page **pages,
@@ -197,6 +230,32 @@ static int __pcpu_map_pages(unsigned long addr, struct page **pages,
 			PAGE_KERNEL, pages, PAGE_SHIFT, GFP_KERNEL);
 }
 
+#ifdef CONFIG_HAVE_LOCAL_PER_CPU_MAP
+static int __pcpu_map_pages_local(unsigned int cpu, unsigned long virt, struct page **pages,
+			    int nr_pages)
+{
+	unsigned int i;
+	int err = 0;
+
+	for (i = 0; i < nr_pages; i++) {
+		err = vmap_range_noflush(percpu_pgd[cpu], virt, virt + PAGE_SIZE,
+				page_to_phys(pages[i]), PAGE_KERNEL, PAGE_SHIFT);
+		if (err)
+			return err;
+
+		virt += PAGE_SIZE;
+	}
+
+	return err;
+}
+#else
+static int __pcpu_map_pages_local(unsigned int cpu, unsigned long virt, struct page **pages,
+			    int nr_pages)
+{
+	return 0;
+}
+#endif
+
 /**
  * pcpu_map_pages - map pages into a pcpu_chunk
  * @chunk: chunk of interest
@@ -224,6 +283,13 @@ static int pcpu_map_pages(struct pcpu_chunk *chunk,
 		if (err < 0)
 			goto err;
 
+		err = __pcpu_map_pages_local(cpu,
+					(unsigned long)chunk->local_base + (page_start << PAGE_SHIFT),
+					&pages[pcpu_page_idx(cpu, page_start)],
+					page_end - page_start);
+		if (err < 0)
+			goto err;
+
 		for (i = page_start; i < page_end; i++)
 			pcpu_set_page_chunk(pages[pcpu_page_idx(cpu, i)],
 					    chunk);
@@ -233,6 +299,9 @@ static int pcpu_map_pages(struct pcpu_chunk *chunk,
 	for_each_possible_cpu(tcpu) {
 		__pcpu_unmap_pages(pcpu_chunk_addr(chunk, tcpu, page_start),
 				   page_end - page_start);
+		__pcpu_unmap_pages_local(tcpu,
+					(unsigned long)chunk->local_base + (page_start << PAGE_SHIFT),
+					page_end - page_start);
 		if (tcpu == cpu)
 			break;
 	}
@@ -258,6 +327,11 @@ static void pcpu_post_map_flush(struct pcpu_chunk *chunk,
 	flush_cache_vmap(
 		pcpu_chunk_addr(chunk, pcpu_low_unit_cpu, page_start),
 		pcpu_chunk_addr(chunk, pcpu_high_unit_cpu, page_end));
+
+#ifdef CONFIG_HAVE_LOCAL_PER_CPU_MAP
+	flush_cache_vmap((unsigned long)chunk->local_base + (page_start << PAGE_SHIFT),
+			 (unsigned long)chunk->local_base + (page_end << PAGE_SHIFT));
+#endif
 }
 
 /**
@@ -349,6 +423,24 @@ static struct pcpu_chunk *pcpu_create_chunk(gfp_t gfp)
 	chunk->data = vms;
 	chunk->base_addr = vms[0]->addr - pcpu_group_offsets[0];
 
+#ifdef CONFIG_HAVE_LOCAL_PER_CPU_MAP
+	unsigned long delta = (unsigned long)chunk->base_addr - (unsigned long)pcpu_base_addr;
+	unsigned long hint = delta + (unsigned long)pcpu_local_base;
+	struct vm_struct *local_vm = pcpu_get_local_vm_area(hint,
+					pcpu_unit_size, pcpu_atom_size);
+	if (!local_vm) {
+		pcpu_free_vm_areas(vms, pcpu_nr_groups);
+		pcpu_free_chunk(chunk);
+		return NULL;
+	}
+
+	chunk->local_base = local_vm->addr;
+	chunk->local_data = (void *)local_vm;
+#else
+	chunk->local_base = 0;
+	chunk->local_data = NULL;
+#endif
+
 	pcpu_stats_chunk_alloc();
 	trace_percpu_create_chunk(chunk->base_addr);
 
@@ -365,6 +457,8 @@ static void pcpu_destroy_chunk(struct pcpu_chunk *chunk)
 
 	if (chunk->data)
 		pcpu_free_vm_areas(chunk->data, pcpu_nr_groups);
+	if (chunk->local_data)
+		free_vm_area((struct vm_struct *)chunk->local_data);
 	pcpu_free_chunk(chunk);
 }
 
diff --git a/mm/vmalloc.c b/mm/vmalloc.c
index 9ec401918736..3d94ca4bbd17 100644
--- a/mm/vmalloc.c
+++ b/mm/vmalloc.c
@@ -4923,17 +4923,21 @@ pvm_find_va_enclose_addr(unsigned long addr)
  *   in - the VA we start the search(reverse order);
  *   out - the VA with the highest aligned end address.
  * @align: alignment for required highest address
+ * @pcpu: whether request allocation from local percpu area
  *
  * Returns: determined end address within vmap_area
  */
 static unsigned long
-pvm_determine_end_from_reverse(struct vmap_area **va, unsigned long align)
+pvm_determine_end_from_reverse(struct vmap_area **va, unsigned long align, bool pcpu)
 {
 	unsigned long vmalloc_end;
 	unsigned long addr;
 
 #ifdef CONFIG_HAVE_LOCAL_PER_CPU_MAP
-	vmalloc_end = PERCPU_END & ~(align - 1);
+	if (pcpu)
+		vmalloc_end = LOCAL_PERCPU_END & ~(align - 1);
+	else
+		vmalloc_end = PERCPU_END & ~(align - 1);
 #else
 	vmalloc_end = VMALLOC_END & ~(align - 1);
 #endif
@@ -5042,7 +5046,7 @@ struct vm_struct **pcpu_get_vm_areas(const unsigned long *offsets,
 	end = start + sizes[area];
 
 	va = pvm_find_va_enclose_addr(vmalloc_end);
-	base = pvm_determine_end_from_reverse(&va, align) - end;
+	base = pvm_determine_end_from_reverse(&va, align, false) - end;
 
 	while (true) {
 		/*
@@ -5063,7 +5067,7 @@ struct vm_struct **pcpu_get_vm_areas(const unsigned long *offsets,
 		 * base downwards and then recheck.
 		 */
 		if (base + end > va->va_end) {
-			base = pvm_determine_end_from_reverse(&va, align) - end;
+			base = pvm_determine_end_from_reverse(&va, align, false) - end;
 			term_area = area;
 			continue;
 		}
@@ -5073,7 +5077,7 @@ struct vm_struct **pcpu_get_vm_areas(const unsigned long *offsets,
 		 */
 		if (base + start < va->va_start) {
 			va = node_to_va(rb_prev(&va->rb_node));
-			base = pvm_determine_end_from_reverse(&va, align) - end;
+			base = pvm_determine_end_from_reverse(&va, align, false) - end;
 			term_area = area;
 			continue;
 		}
@@ -5236,6 +5240,84 @@ void pcpu_free_vm_areas(struct vm_struct **vms, int nr_vms)
 		free_vm_area(vms[i]);
 	kfree(vms);
 }
+
+#ifdef CONFIG_HAVE_LOCAL_PER_CPU_MAP
+/* Find free vm area starts from hint */
+struct vm_struct *pcpu_get_local_vm_area(unsigned long hint,
+				     int unit_size, size_t align)
+{
+	struct vmap_area *tmp_va, *va;
+	struct vm_struct *vm;
+	struct vmap_node *vn;
+	unsigned long end;
+	int ret;
+
+	va = kmem_cache_zalloc(vmap_area_cachep, GFP_KERNEL);
+	vm = kzalloc(sizeof(struct vm_struct), GFP_KERNEL);
+	if (!va || !vm)
+		goto err_free;
+
+	spin_lock(&free_vmap_area_lock);
+
+	tmp_va = pvm_find_va_enclose_addr(hint);
+	if (!tmp_va) {
+		spin_unlock(&free_vmap_area_lock);
+		goto err_free;
+	}
+
+	end = pvm_determine_end_from_reverse(&tmp_va, align, true);
+
+	if (hint + unit_size > end) {
+		spin_unlock(&free_vmap_area_lock);
+		goto err_free;
+	}
+
+	ret = va_clip(&free_vmap_area_root,
+			&free_vmap_area_list, tmp_va, hint, unit_size);
+	if (ret) {
+		spin_unlock(&free_vmap_area_lock);
+		goto err_free;
+	}
+
+	va->va_start = hint;
+	va->va_end = hint + unit_size;
+
+	spin_unlock(&free_vmap_area_lock);
+
+	if (kasan_populate_vmalloc(va->va_start, unit_size, GFP_KERNEL))
+		goto err_free_shadow;
+
+	vn = addr_to_node(va->va_start);
+
+	spin_lock(&vn->busy.lock);
+	insert_vmap_area(va, &vn->busy.root, &vn->busy.head);
+	setup_vmalloc_vm(vm, va, VM_ALLOC,
+			 pcpu_get_local_vm_area);
+	spin_unlock(&vn->busy.lock);
+
+	kasan_unpoison_vmap_areas(&vm, 1, KASAN_VMALLOC_PROT_NORMAL);
+
+	return vm;
+
+err_free:
+	kmem_cache_free(vmap_area_cachep, va);
+	kfree(vm);
+
+	return NULL;
+
+err_free_shadow:
+	spin_lock(&free_vmap_area_lock);
+	va = merge_or_add_vmap_area_augment(va, &free_vmap_area_root,
+			&free_vmap_area_list);
+	if (va)
+		kasan_release_vmalloc(va->va_start, va->va_end, va->va_start, va->va_end,
+			       KASAN_VMALLOC_PAGE_RANGE | KASAN_VMALLOC_TLB_FLUSH);
+	kfree(vm);
+	spin_unlock(&free_vmap_area_lock);
+	return NULL;
+
+}
+#endif
 #endif	/* CONFIG_SMP */
 
 #ifdef CONFIG_PRINTK
-- 
2.47.0

[PATCH 15/16] arm64: kconfig: select HAVE_LOCAL_PER_CPU_MAP

[Yang Shi]

ARM64 supports local percpu map, so select HAVE_LOCAL_PER_CPU_MAP. But
it depends on VA bits and page size for now.

Signed-off-by: Yang Shi <yang@os.amperecomputing.com>
---
 arch/arm64/Kconfig | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/arch/arm64/Kconfig b/arch/arm64/Kconfig
index df55e6bdf81f..95ffd151187f 100644
--- a/arch/arm64/Kconfig
+++ b/arch/arm64/Kconfig
@@ -415,6 +415,11 @@ config KASAN_SHADOW_OFFSET
 config UNWIND_TABLES
 	bool
 
+config HAVE_LOCAL_PER_CPU_MAP
+	def_bool y
+	depends on (ARM64_VA_BITS_48 && (ARM64_4K_PAGES || ARM64_64K_PAGES)) || \
+		   (ARM64_VA_BITS_47 && ARM64_16K_PAGES)
+
 source "arch/arm64/Kconfig.platforms"
 
 menu "Kernel Features"
-- 
2.47.0

[PATCH 16/16] arm64: percpu: use local percpu for this_cpu_*() APIs

[Yang Shi]

Use local percpu address for this_cpu_*() APIs.  Because the percpu
variable is mapped to the same virtual address, their address can be
calculated by using __per_cpu_local_off which has same value for all
CPUs.  So preempt_disable/preempt_enable is not needed anymore.  This
optimization can improve the performance for this_cpu_*() operations.

Kernel build test on AmpereOne (160 cores) with default Fedora kernel
config in a memcg roughly showed 13% - 15% sys time improvement.

Signed-off-by: Yang Shi <yang@os.amperecomputing.com>
---
 arch/arm64/include/asm/percpu.h | 37 ++++++++++++++++++++++++++++++++-
 1 file changed, 36 insertions(+), 1 deletion(-)

diff --git a/arch/arm64/include/asm/percpu.h b/arch/arm64/include/asm/percpu.h
index b57b2bb00967..03d149f9177c 100644
--- a/arch/arm64/include/asm/percpu.h
+++ b/arch/arm64/include/asm/percpu.h
@@ -142,6 +142,27 @@ PERCPU_RET_OP(add, add, ldadd)
 #undef PERCPU_OP
 #undef PERCPU_RET_OP
 
+#ifdef CONFIG_HAVE_LOCAL_PER_CPU_MAP
+extern unsigned long __per_cpu_local_off;
+
+#define local_cpu_ptr(ptr)						\
+({									\
+	__verify_pcpu_ptr(ptr);						\
+	SHIFT_PERCPU_PTR(ptr, __per_cpu_local_off);			\
+})
+
+#define _pcp_protect(op, pcp, ...)					\
+({									\
+	op(local_cpu_ptr(&(pcp)), __VA_ARGS__);				\
+})
+
+#define _pcp_protect_return(op, pcp, args...)				\
+({									\
+	typeof(pcp) __retval;						\
+	__retval = (typeof(pcp))op(local_cpu_ptr(&(pcp)), ##args);	\
+	__retval;							\
+})
+#else
 /*
  * It would be nice to avoid the conditional call into the scheduler when
  * re-enabling preemption for preemptible kernels, but doing that in a way
@@ -152,7 +173,6 @@ PERCPU_RET_OP(add, add, ldadd)
  * preemption point when TIF_NEED_RESCHED gets set while preemption is
  * disabled.
  */
-
 #define _pcp_protect(op, pcp, ...)					\
 ({									\
 	preempt_disable_notrace();					\
@@ -168,6 +188,7 @@ PERCPU_RET_OP(add, add, ldadd)
 	preempt_enable_notrace();					\
 	__retval;							\
 })
+#endif
 
 #define this_cpu_read_1(pcp)		\
 	_pcp_protect_return(__percpu_read_8, pcp)
@@ -243,6 +264,19 @@ PERCPU_RET_OP(add, add, ldadd)
 
 #define this_cpu_cmpxchg64(pcp, o, n)	this_cpu_cmpxchg_8(pcp, o, n)
 
+#ifdef CONFIG_HAVE_LOCAL_PER_CPU_MAP
+#define this_cpu_cmpxchg128(pcp, o, n)					\
+({									\
+	typedef typeof(pcp) pcp_op_T__;					\
+	u128 old__, new__, ret__;					\
+	pcp_op_T__ *ptr__;						\
+	old__ = o;							\
+	new__ = n;							\
+	ptr__ = local_cpu_ptr(&(pcp));					\
+	ret__ = cmpxchg128_local((void *)ptr__, old__, new__);		\
+	ret__;								\
+})
+#else
 #define this_cpu_cmpxchg128(pcp, o, n)					\
 ({									\
 	typedef typeof(pcp) pcp_op_T__;					\
@@ -256,6 +290,7 @@ PERCPU_RET_OP(add, add, ldadd)
 	preempt_enable_notrace();					\
 	ret__;								\
 })
+#endif
 
 #ifdef __KVM_NVHE_HYPERVISOR__
 extern unsigned long __hyp_per_cpu_offset(unsigned int cpu);
-- 
2.47.0