[PATCH 6/6] arm64: use alternatie callback to patch TPIDR_EL1 accesses

[Ada Couprie Diaz <ada.coupriediaz@arm.com>]

ARM64_HAS_VIRT_HOST_EXTN represents about 3k alternatives, most of them
are found in per_cpu calls switching between TPIDR_EL1 and TPIDR_EL2.

Replace the current alternative instructions with a callback that
changes the target system register, saving a potential 10kB.

`aarch64_insn_extract_system_reg()` is safe to use as it was made
`noinstr` in the previous patch.

Signed-off-by: Ada Couprie Diaz <ada.coupriediaz@arm.com>
---
 arch/arm64/include/asm/alternative.h |  3 +++
 arch/arm64/include/asm/assembler.h   | 12 ++++--------
 arch/arm64/include/asm/percpu.h      | 12 ++++++------
 arch/arm64/kernel/alternative.c      | 25 +++++++++++++++++++++++++
 arch/arm64/kernel/image-vars.h       |  1 +
 5 files changed, 39 insertions(+), 14 deletions(-)

diff --git a/arch/arm64/include/asm/alternative.h b/arch/arm64/include/asm/alternative.h
index 621aa8550174f..e3842a28f88ba 100644
--- a/arch/arm64/include/asm/alternative.h
+++ b/arch/arm64/include/asm/alternative.h
@@ -37,5 +37,8 @@ static inline int apply_alternatives_module(void *start, size_t length)
 void alt_cb_patch_nops(struct alt_instr *alt, __le32 *origptr,
 		       __le32 *updptr, int nr_inst);
 
+void alt_cb_patch_tpidr_el2(struct alt_instr *alt, __le32 *origptr,
+			    __le32 *updptr, int nr_inst);
+
 #endif /* __ASSEMBLER__ */
 #endif /* __ASM_ALTERNATIVE_H */
diff --git a/arch/arm64/include/asm/assembler.h b/arch/arm64/include/asm/assembler.h
index 0b58b550e8dce..6cbdfefe45382 100644
--- a/arch/arm64/include/asm/assembler.h
+++ b/arch/arm64/include/asm/assembler.h
@@ -210,19 +210,15 @@ lr	.req	x30		// link register
 	.endm
 #else
 	.macro	get_this_cpu_offset, dst
-alternative_if_not ARM64_HAS_VIRT_HOST_EXTN
+alternative_cb ARM64_HAS_VIRT_HOST_EXTN, alt_cb_patch_tpidr_el2
 	mrs	\dst, tpidr_el1
-alternative_else
-	mrs	\dst, tpidr_el2
-alternative_endif
+alternative_cb_end
 	.endm
 
 	.macro	set_this_cpu_offset, src
-alternative_if_not ARM64_HAS_VIRT_HOST_EXTN
+alternative_cb ARM64_HAS_VIRT_HOST_EXTN, alt_cb_patch_tpidr_el2
 	msr	tpidr_el1, \src
-alternative_else
-	msr	tpidr_el2, \src
-alternative_endif
+alternative_cb_end
 	.endm
 #endif
 
diff --git a/arch/arm64/include/asm/percpu.h b/arch/arm64/include/asm/percpu.h
index 8cf4068ce1b56..6999ae86a7e9f 100644
--- a/arch/arm64/include/asm/percpu.h
+++ b/arch/arm64/include/asm/percpu.h
@@ -14,9 +14,9 @@
 
 static inline void set_my_cpu_offset(unsigned long off)
 {
-	asm volatile(ALTERNATIVE("msr tpidr_el1, %0",
-				 "msr tpidr_el2, %0",
-				 ARM64_HAS_VIRT_HOST_EXTN)
+	asm volatile(ALTERNATIVE_CB("msr tpidr_el1, %0",
+				    ARM64_HAS_VIRT_HOST_EXTN,
+				    alt_cb_patch_tpidr_el2)
 			:: "r" (off) : "memory");
 }
 
@@ -37,9 +37,9 @@ static inline unsigned long __kern_my_cpu_offset(void)
 	 * We want to allow caching the value, so avoid using volatile and
 	 * instead use a fake stack read to hazard against barrier().
 	 */
-	asm(ALTERNATIVE("mrs %0, tpidr_el1",
-			"mrs %0, tpidr_el2",
-			ARM64_HAS_VIRT_HOST_EXTN)
+	asm(ALTERNATIVE_CB("mrs %0, tpidr_el1",
+			   ARM64_HAS_VIRT_HOST_EXTN,
+			   alt_cb_patch_tpidr_el2)
 		: "=r" (off) :
 		"Q" (*(const unsigned long *)current_stack_pointer));
 
diff --git a/arch/arm64/kernel/alternative.c b/arch/arm64/kernel/alternative.c
index f5ec7e7c1d3fd..d853eda04ed23 100644
--- a/arch/arm64/kernel/alternative.c
+++ b/arch/arm64/kernel/alternative.c
@@ -303,3 +303,28 @@ noinstr void alt_cb_patch_nops(struct alt_instr *alt, __le32 *origptr,
 		updptr[i] = cpu_to_le32(aarch64_insn_gen_nop());
 }
 EXPORT_SYMBOL(alt_cb_patch_nops);
+
+/*
+ * Patch MSR and MRS instructions to switch from TPIDR_EL1 to TPIDR_EL2.
+ * MSR (register) and MRS have the same target system register encoding,
+ * so we can handle both here.
+ */
+noinstr void alt_cb_patch_tpidr_el2(struct alt_instr *alt, __le32 *origptr,
+				    __le32 *updptr, int nr_inst)
+{
+	u32 orinst, altinst;
+
+	BUG_ON(nr_inst != 1);
+
+	orinst = le32_to_cpu(origptr[0]);
+	BUG_ON(!aarch64_insn_is_msr_reg(orinst) &&
+	       !aarch64_insn_is_mrs(orinst));
+	BUG_ON(aarch64_insn_extract_system_reg(orinst) !=
+	       AARCH64_INSN_SYSREG_TPIDR_EL1);
+
+	altinst = orinst & ~AARCH64_INSN_SYSREG_MASK;
+	altinst |= AARCH64_INSN_SYSREG_TPIDR_EL2 << AARCH64_INSN_SYSREG_OFFSET;
+
+	updptr[0] = cpu_to_le32(altinst);
+}
+EXPORT_SYMBOL(alt_cb_patch_tpidr_el2);
diff --git a/arch/arm64/kernel/image-vars.h b/arch/arm64/kernel/image-vars.h
index 615b77ef403fe..49d9984d7cd89 100644
--- a/arch/arm64/kernel/image-vars.h
+++ b/arch/arm64/kernel/image-vars.h
@@ -96,6 +96,7 @@ KVM_NVHE_ALIAS(kvm_compute_ich_hcr_trap_bits);
 KVM_NVHE_ALIAS(kvm_patch_ich_vtr_el2);
 KVM_NVHE_ALIAS(__io_nvidia_olympus_patch_dmb);
 KVM_NVHE_ALIAS(__io_arm_a57_patch_ladr);
+KVM_NVHE_ALIAS(alt_cb_patch_tpidr_el2);
 
 /* Global kernel state accessed by nVHE hyp code. */
 KVM_NVHE_ALIAS(kvm_vgic_global_state);
-- 
2.43.0