arm64 kvm:implement basic lazy save and restore for FPSIMD registers

  Implement basic lazy save and restore for FPSIMD registers, which only
  restore FPSIMD state on el0_fpsimd_acc and save FPSIMD state in switch().

Signed-off-by: Robin Luk <lubin.lu@antgroup.com>

5 files changed, 104 insertions, 29 deletions

diff --git a/pkg/ring0/defs_arm64.go b/pkg/ring0/defs_arm64.go
index 4b2375b71..c372b02bb 100644
--- a/pkg/ring0/defs_arm64.go
+++ b/pkg/ring0/defs_arm64.go
@@ -55,6 +55,9 @@ type CPUArchState struct {
 	// faultAddr is the value of far_el1.
 	faultAddr uintptr
 
+	// el0Fp is the address of application's fpstate.
+	el0Fp uintptr
+
 	// ttbr0Kvm is the value of ttbr0_el1 for sentry.
 	ttbr0Kvm uintptr
 
diff --git a/pkg/ring0/entry_arm64.s b/pkg/ring0/entry_arm64.s
index 4a656de8e..f801b8e11 100644
--- a/pkg/ring0/entry_arm64.s
+++ b/pkg/ring0/entry_arm64.s
@@ -65,7 +65,7 @@
 //
 // The following registers are not saved: R18, R19.
 #define REGISTERS_SAVE(reg, offset) \
-STP (R0, R1), offset+PTRACE_R0(reg); \
+  STP (R0, R1), offset+PTRACE_R0(reg); \
   STP (R2, R3), offset+PTRACE_R2(reg); \
   STP (R4, R5), offset+PTRACE_R4(reg); \
   STP (R6, R7), offset+PTRACE_R6(reg); \
@@ -88,7 +88,7 @@ STP (R0, R1), offset+PTRACE_R0(reg); \
 //
 // The following registers are not loaded: R18, R19.
 #define REGISTERS_LOAD(reg, offset) \
-LDP offset+PTRACE_R0(reg), (R0, R1); \
+  LDP offset+PTRACE_R0(reg), (R0, R1); \
   LDP offset+PTRACE_R2(reg), (R2, R3); \
   LDP offset+PTRACE_R4(reg), (R4, R5); \
   LDP offset+PTRACE_R6(reg), (R6, R7); \
@@ -104,6 +104,32 @@ LDP offset+PTRACE_R0(reg), (R0, R1); \
   LDP offset+PTRACE_R28(reg), (g, R29); \
   MOVD offset+PTRACE_R30(reg), R30;
 
+// Loads the application's fpstate.
+#define FPSTATE_EL0_LOAD() \
+  MRS TPIDR_EL1, RSV_REG; \
+  MOVD CPU_FPSTATE_EL0(RSV_REG), RSV_REG; \
+  MOVD 0(RSV_REG), RSV_REG_APP; \
+  MOVD RSV_REG_APP, FPSR; \
+  MOVD 8(RSV_REG), RSV_REG_APP; \
+  MOVD RSV_REG_APP, FPCR; \
+  ADD $16, RSV_REG, RSV_REG; \
+  WORD $0xad400640; \ // ldp q0, q1, [x18]
+  WORD $0xad410e42; \
+  WORD $0xad421644; \
+  WORD $0xad431e46; \
+  WORD $0xad442648; \
+  WORD $0xad452e4a; \
+  WORD $0xad46364c; \
+  WORD $0xad473e4e; \
+  WORD $0xad484650; \
+  WORD $0xad494e52; \
+  WORD $0xad4a5654; \
+  WORD $0xad4b5e56; \
+  WORD $0xad4c6658; \
+  WORD $0xad4d6e5a; \
+  WORD $0xad4e765c; \
+  WORD $0xad4f7e5e;
+
 #define ESR_ELx_EC_UNKNOWN	(0x00)
 #define ESR_ELx_EC_WFx		(0x01)
 /* Unallocated EC: 0x02 */
@@ -258,26 +284,28 @@ LDP offset+PTRACE_R0(reg), (R0, R1); \
 	MSR RSV_REG, TTBR0_EL1; \
 	ISB $15;
 
-TEXT ·EnableVFP(SB),NOSPLIT,$0
+// FPSIMDDisableTrap disables the trap for accessing fpsimd.
+TEXT ·FPSIMDDisableTrap(SB),NOSPLIT,$0
 	MOVD $FPEN_ENABLE, R0
-	WORD $0xd5181040 //MSR R0, CPACR_EL1
+	MSR R0, CPACR_EL1
 	ISB $15
 	RET
 
-TEXT ·DisableVFP(SB),NOSPLIT,$0
-	MOVD $0, R0
-	WORD $0xd5181040 //MSR R0, CPACR_EL1
+// FPSIMDEnableTrap enables the trap for accessing fpsimd.
+TEXT ·FPSIMDEnableTrap(SB),NOSPLIT,$0
+	MSR $0, CPACR_EL1
 	ISB $15
 	RET
 
-#define VFP_ENABLE \
-	MOVD $FPEN_ENABLE, R0; \
-	WORD $0xd5181040; \ //MSR R0, CPACR_EL1
+// FPSIMD_DISABLE_TRAP disables the trap for accessing fpsimd.
+#define FPSIMD_DISABLE_TRAP(reg) \
+	MOVD $FPEN_ENABLE, reg; \
+	MSR reg, CPACR_EL1; \
 	ISB $15;
 
-#define VFP_DISABLE \
-	MOVD $0x0, R0; \
-	WORD $0xd5181040; \ //MSR R0, CPACR_EL1
+// FPSIMD_ENABLE_TRAP enables the trap for accessing fpsimd.
+#define FPSIMD_ENABLE_TRAP(reg) \
+	MSR $0, CPACR_EL1; \
 	ISB $15;
 
 // KERNEL_ENTRY_FROM_EL0 is the entry code of the vcpu from el0 to el1.
@@ -334,6 +362,14 @@ TEXT ·DisableVFP(SB),NOSPLIT,$0
 	MOVD R3, 8(RSP); \
 	B ·HaltEl1ExceptionAndResume(SB);
 
+// storeEl0Fpstate writes the address of application's fpstate.
+TEXT ·storeEl0Fpstate(SB),NOSPLIT,$0-8
+	MOVD value+0(FP), R1
+	ORR $0xffff000000000000, R1, R1
+	MRS  TPIDR_EL1, RSV_REG
+	MOVD R1, CPU_FPSTATE_EL0(RSV_REG)
+	RET
+
 // storeAppASID writes the application's asid value.
 TEXT ·storeAppASID(SB),NOSPLIT,$0-8
 	MOVD asid+0(FP), R1
@@ -348,7 +384,7 @@ TEXT ·Halt(SB),NOSPLIT,$0
 	MOVD R1, CPU_LAZY_VFP(RSV_REG)
 	DSB $15
 
-	VFP_DISABLE
+	FPSIMD_ENABLE_TRAP(RSV_REG)
 
 	// Trigger MMIO_EXIT/_KVM_HYPERCALL_VMEXIT.
 	//
@@ -511,6 +547,9 @@ TEXT ·Start(SB),NOSPLIT,$0
 	ORR $0xffff000000000000, RSV_REG, RSV_REG
 	WORD $0xd518d092        //MSR R18, TPIDR_EL1
 
+	// Enable trap for accessing fpsimd.
+	MSR $0, CPACR_EL1
+
 	// Init.
 	MOVD $SCTLR_EL1_DEFAULT, R1 // re-enable the mmu.
 	MSR R1, SCTLR_EL1
@@ -547,6 +586,10 @@ TEXT ·El1_sync(SB),NOSPLIT,$0
 	BEQ el1_da                        // data abort in EL1
 	CMP $ESR_ELx_EC_IABT_CUR, R24
 	BEQ el1_ia                        // instruction abort in EL1
+	CMP $ESR_ELx_EC_FP_ASIMD, R24
+	BEQ el1_fpsimd_acc                // FP/ASIMD access
+	CMP $ESR_ELx_EC_SVE, R24
+	BEQ el1_sve_acc                   // SVE access
 	CMP $ESR_ELx_EC_SP_ALIGN, R24
 	BEQ el1_sp_pc                     // stack alignment exception
 	CMP $ESR_ELx_EC_PC_ALIGN, R24
@@ -557,10 +600,6 @@ TEXT ·El1_sync(SB),NOSPLIT,$0
 	BEQ el1_svc                       // SVC in 64-bit state
 	CMP $ESR_ELx_EC_BREAKPT_CUR, R24
 	BEQ el1_dbg                       // debug exception in EL1
-	CMP $ESR_ELx_EC_FP_ASIMD, R24
-	BEQ el1_fpsimd_acc                // FP/ASIMD access
-	CMP $ESR_ELx_EC_SVE, R24
-	BEQ el1_sve_acc                   // SVE access
 	B el1_invalid
 
 el1_da:
@@ -577,8 +616,21 @@ el1_dbg:
 	EXCEPTION_EL1(El1SyncDbg)
 el1_fpsimd_acc:
 el1_sve_acc:
-	VFP_ENABLE
-	B ·kernelExitToEl1(SB)  // Resume.
+	FPSIMD_DISABLE_TRAP(RSV_REG)
+
+	// Restore context.
+	MRS TPIDR_EL1, RSV_REG
+
+	// Restore sp.
+	MOVD CPU_REGISTERS+PTRACE_SP(RSV_REG), R1
+	MOVD R1, RSP
+
+	// Restore common registers.
+	REGISTERS_LOAD(RSV_REG, CPU_REGISTERS)
+	MOVD CPU_REGISTERS+PTRACE_R19(RSV_REG), RSV_REG_APP
+
+	ERET()	// return to el1.
+
 el1_invalid:
 	EXCEPTION_EL1(El1SyncInv)
 
@@ -642,9 +694,20 @@ el0_da:
 el0_ia:
 	EXCEPTION_EL0(PageFault)
 el0_fpsimd_acc:
-	EXCEPTION_EL0(El0SyncFpsimdAcc)
 el0_sve_acc:
-	EXCEPTION_EL0(El0SyncSveAcc)
+	FPSIMD_DISABLE_TRAP(RSV_REG)
+	FPSTATE_EL0_LOAD()
+
+	// Restore context.
+	MRS TPIDR_EL1, RSV_REG
+	MOVD CPU_APP_ADDR(RSV_REG), RSV_REG_APP
+
+	// Restore R0-R30
+	REGISTERS_LOAD(RSV_REG_APP, 0)
+	MOVD PTRACE_R18(RSV_REG_APP), RSV_REG
+	MOVD PTRACE_R19(RSV_REG_APP), RSV_REG_APP
+
+	ERET()  // return to el0.
 el0_fpsimd_exc:
 	EXCEPTION_EL0(El0SyncFpsimdExc)
 el0_sp_pc:
diff --git a/pkg/ring0/kernel_arm64.go b/pkg/ring0/kernel_arm64.go
index c05284641..41909b3a0 100644
--- a/pkg/ring0/kernel_arm64.go
+++ b/pkg/ring0/kernel_arm64.go
@@ -62,6 +62,8 @@ func IsCanonical(addr uint64) bool {
 //go:nosplit
 func (c *CPU) SwitchToUser(switchOpts SwitchOpts) (vector Vector) {
 	storeAppASID(uintptr(switchOpts.UserASID))
+	storeEl0Fpstate(switchOpts.FloatingPointState)
+
 	if switchOpts.Flush {
 		FlushTlbByASID(uintptr(switchOpts.UserASID))
 	}
@@ -71,13 +73,17 @@ func (c *CPU) SwitchToUser(switchOpts SwitchOpts) (vector Vector) {
 	regs.Pstate &= ^uint64(PsrFlagsClear)
 	regs.Pstate |= UserFlagsSet
 
-	EnableVFP()
-	LoadFloatingPoint(switchOpts.FloatingPointState)
+	fpDisableTrap := CPACREL1()
+	if fpDisableTrap != 0 {
+		FPSIMDEnableTrap()
+	}
 
 	kernelExitToEl0()
 
-	SaveFloatingPoint(switchOpts.FloatingPointState)
-	DisableVFP()
+	fpDisableTrap = CPACREL1()
+	if fpDisableTrap != 0 {
+		SaveFloatingPoint(switchOpts.FloatingPointState)
+	}
 
 	vector = c.vecCode
 
diff --git a/pkg/ring0/lib_arm64.go b/pkg/ring0/lib_arm64.go
index a490bf3af..edf24eda3 100644
--- a/pkg/ring0/lib_arm64.go
+++ b/pkg/ring0/lib_arm64.go
@@ -16,6 +16,9 @@
 
 package ring0
 
+// storeEl0Fpstate writes the address of application's fpstate.
+func storeEl0Fpstate(value *byte)
+
 // storeAppASID writes the application's asid value.
 func storeAppASID(asid uintptr)
 
@@ -59,11 +62,10 @@ func LoadFloatingPoint(*byte)
 // SaveFloatingPoint saves floating point state.
 func SaveFloatingPoint(*byte)
 
-// EnableVFP enables fpsimd.
-func EnableVFP()
+func FPSIMDDisableTrap()
 
 // DisableVFP disables fpsimd.
-func DisableVFP()
+func FPSIMDEnableTrap()
 
 // Init sets function pointers based on architectural features.
 //
diff --git a/pkg/ring0/offsets_arm64.go b/pkg/ring0/offsets_arm64.go
index 164db6d5a..03adaa6b0 100644
--- a/pkg/ring0/offsets_arm64.go
+++ b/pkg/ring0/offsets_arm64.go
@@ -36,6 +36,7 @@ func Emit(w io.Writer) {
 	fmt.Fprintf(w, "#define CPU_ERROR_CODE       0x%02x\n", reflect.ValueOf(&c.errorCode).Pointer()-reflect.ValueOf(c).Pointer())
 	fmt.Fprintf(w, "#define CPU_ERROR_TYPE       0x%02x\n", reflect.ValueOf(&c.errorType).Pointer()-reflect.ValueOf(c).Pointer())
 	fmt.Fprintf(w, "#define CPU_FAULT_ADDR       0x%02x\n", reflect.ValueOf(&c.faultAddr).Pointer()-reflect.ValueOf(c).Pointer())
+	fmt.Fprintf(w, "#define CPU_FPSTATE_EL0      0x%02x\n", reflect.ValueOf(&c.el0Fp).Pointer()-reflect.ValueOf(c).Pointer())
 	fmt.Fprintf(w, "#define CPU_TTBR0_KVM	     0x%02x\n", reflect.ValueOf(&c.ttbr0Kvm).Pointer()-reflect.ValueOf(c).Pointer())
 	fmt.Fprintf(w, "#define CPU_TTBR0_APP        0x%02x\n", reflect.ValueOf(&c.ttbr0App).Pointer()-reflect.ValueOf(c).Pointer())
 	fmt.Fprintf(w, "#define CPU_VECTOR_CODE      0x%02x\n", reflect.ValueOf(&c.vecCode).Pointer()-reflect.ValueOf(c).Pointer())