7 files changed, 83 insertions, 23 deletions
diff --git a/pkg/sentry/platform/kvm/bluepill_arm64.go b/pkg/sentry/platform/kvm/bluepill_arm64.go
index 552341721..c215d443c 100644
--- a/pkg/sentry/platform/kvm/bluepill_arm64.go
+++ b/pkg/sentry/platform/kvm/bluepill_arm64.go
@@ -54,6 +54,14 @@ func bluepillArchExit(c *vCPU, context *arch.SignalContext64) {
 	context.Pstate = regs.Pstate
 	context.Pstate &^= uint64(ring0.UserFlagsClear)
 	context.Pstate |= ring0.UserFlagsSet
+
+	lazyVfp := c.GetLazyVFP()
+	if lazyVfp != 0 {
+		fpsimd := fpsimdPtr((*byte)(c.floatingPointState))
+		context.Fpsimd64.Fpsr = fpsimd.Fpsr
+		context.Fpsimd64.Fpcr = fpsimd.Fpcr
+		context.Fpsimd64.Vregs = fpsimd.Vregs
+	}
 }
 
 // KernelSyscall handles kernel syscalls.
@@ -64,6 +72,17 @@ func (c *vCPU) KernelSyscall() {
 	if regs.Regs[8] != ^uint64(0) {
 		regs.Pc -= 4 // Rewind.
 	}
+
+	vfpEnable := ring0.CPACREL1()
+	if vfpEnable != 0 {
+		fpsimd := fpsimdPtr((*byte)(c.floatingPointState))
+		fpcr := ring0.GetFPCR()
+		fpsr := ring0.GetFPSR()
+		fpsimd.Fpcr = uint32(fpcr)
+		fpsimd.Fpsr = uint32(fpsr)
+		ring0.SaveVRegs((*byte)(c.floatingPointState))
+	}
+
 	ring0.Halt()
 }
 
@@ -75,5 +94,16 @@ func (c *vCPU) KernelException(vector ring0.Vector) {
 	if vector == ring0.Vector(bounce) {
 		regs.Pc = 0
 	}
+
+	vfpEnable := ring0.CPACREL1()
+	if vfpEnable != 0 {
+		fpsimd := fpsimdPtr((*byte)(c.floatingPointState))
+		fpcr := ring0.GetFPCR()
+		fpsr := ring0.GetFPSR()
+		fpsimd.Fpcr = uint32(fpcr)
+		fpsimd.Fpsr = uint32(fpsr)
+		ring0.SaveVRegs((*byte)(c.floatingPointState))
+	}
+
 	ring0.Halt()
 }
diff --git a/pkg/sentry/platform/kvm/bluepill_arm64_unsafe.go b/pkg/sentry/platform/kvm/bluepill_arm64_unsafe.go
index 195331383..4ca2b7717 100644
--- a/pkg/sentry/platform/kvm/bluepill_arm64_unsafe.go
+++ b/pkg/sentry/platform/kvm/bluepill_arm64_unsafe.go
@@ -17,10 +17,19 @@
 package kvm
 
 import (
+	"unsafe"
+
 	"gvisor.dev/gvisor/pkg/sentry/arch"
 	"gvisor.dev/gvisor/pkg/sentry/platform/ring0"
 )
 
+// fpsimdPtr returns a fpsimd64 for the given address.
+//
+//go:nosplit
+func fpsimdPtr(addr *byte) *arch.FpsimdContext {
+	return (*arch.FpsimdContext)(unsafe.Pointer(addr))
+}
+
 // dieArchSetup initialies the state for dieTrampoline.
 //
 // The arm64 dieTrampoline requires the vCPU to be set in R1, and the last PC
@@ -45,3 +54,10 @@ func dieArchSetup(c *vCPU, context *arch.SignalContext64, guestRegs *userRegs) {
 	context.Regs[1] = uint64(uintptr(unsafe.Pointer(c)))
 	context.Pc = uint64(dieTrampolineAddr)
 }
+
+// bluepillArchFpContext returns the arch-specific fpsimd context.
+//
+//go:nosplit
+func bluepillArchFpContext(context unsafe.Pointer) *arch.FpsimdContext {
+	return &((*arch.SignalContext64)(context).Fpsimd64)
+}
diff --git a/pkg/sentry/platform/kvm/kvm_const_arm64.go b/pkg/sentry/platform/kvm/kvm_const_arm64.go
index 5a74c6e36..531ae8b1e 100644
--- a/pkg/sentry/platform/kvm/kvm_const_arm64.go
+++ b/pkg/sentry/platform/kvm/kvm_const_arm64.go
@@ -19,30 +19,31 @@ const (
 	_KVM_GET_ONE_REG = 0x4010aeab
 	_KVM_SET_ONE_REG = 0x4010aeac
 
-	_KVM_ARM_PREFERRED_TARGET = 0x8020aeaf
-	_KVM_ARM_VCPU_INIT        = 0x4020aeae
-	_KVM_ARM64_REGS_PSTATE    = 0x6030000000100042
-	_KVM_ARM64_REGS_SP_EL1    = 0x6030000000100044
-	_KVM_ARM64_REGS_R0        = 0x6030000000100000
-	_KVM_ARM64_REGS_R1        = 0x6030000000100002
-	_KVM_ARM64_REGS_R2        = 0x6030000000100004
-	_KVM_ARM64_REGS_R3        = 0x6030000000100006
-	_KVM_ARM64_REGS_R8        = 0x6030000000100010
-	_KVM_ARM64_REGS_R18       = 0x6030000000100024
-	_KVM_ARM64_REGS_PC        = 0x6030000000100040
-	_KVM_ARM64_REGS_MAIR_EL1  = 0x603000000013c510
-	_KVM_ARM64_REGS_TCR_EL1   = 0x603000000013c102
-	_KVM_ARM64_REGS_TTBR0_EL1 = 0x603000000013c100
-	_KVM_ARM64_REGS_TTBR1_EL1 = 0x603000000013c101
-	_KVM_ARM64_REGS_SCTLR_EL1 = 0x603000000013c080
-	_KVM_ARM64_REGS_CPACR_EL1 = 0x603000000013c082
-	_KVM_ARM64_REGS_VBAR_EL1  = 0x603000000013c600
+	_KVM_ARM_TARGET_GENERIC_V8 = 5
+	_KVM_ARM_PREFERRED_TARGET  = 0x8020aeaf
+	_KVM_ARM_VCPU_INIT         = 0x4020aeae
+	_KVM_ARM64_REGS_PSTATE     = 0x6030000000100042
+	_KVM_ARM64_REGS_SP_EL1     = 0x6030000000100044
+	_KVM_ARM64_REGS_R0         = 0x6030000000100000
+	_KVM_ARM64_REGS_R1         = 0x6030000000100002
+	_KVM_ARM64_REGS_R2         = 0x6030000000100004
+	_KVM_ARM64_REGS_R3         = 0x6030000000100006
+	_KVM_ARM64_REGS_R8         = 0x6030000000100010
+	_KVM_ARM64_REGS_R18        = 0x6030000000100024
+	_KVM_ARM64_REGS_PC         = 0x6030000000100040
+	_KVM_ARM64_REGS_MAIR_EL1   = 0x603000000013c510
+	_KVM_ARM64_REGS_TCR_EL1    = 0x603000000013c102
+	_KVM_ARM64_REGS_TTBR0_EL1  = 0x603000000013c100
+	_KVM_ARM64_REGS_TTBR1_EL1  = 0x603000000013c101
+	_KVM_ARM64_REGS_SCTLR_EL1  = 0x603000000013c080
+	_KVM_ARM64_REGS_CPACR_EL1  = 0x603000000013c082
+	_KVM_ARM64_REGS_VBAR_EL1   = 0x603000000013c600
 )
 
 // Arm64: Architectural Feature Access Control Register EL1.
 const (
-	_FPEN_NOTRAP = 0x3
-	_FPEN_SHIFT  = 0x20
+	_FPEN_NOTRAP = 3
+	_FPEN_SHIFT  = 20
 )
 
 // Arm64: System Control Register EL1.
diff --git a/pkg/sentry/platform/kvm/machine_arm64.go b/pkg/sentry/platform/kvm/machine_arm64.go
index 09552837a..e42505542 100644
--- a/pkg/sentry/platform/kvm/machine_arm64.go
+++ b/pkg/sentry/platform/kvm/machine_arm64.go
@@ -28,6 +28,10 @@ type vCPUArchState struct {
 	//
 	// This starts above fixedKernelPCID.
 	PCIDs *pagetables.PCIDs
+
+	// floatingPointState is the floating point state buffer used in guest
+	// to host transitions. See usage in bluepill_arm64.go.
+	floatingPointState *arch.FloatingPointData
 }
 
 const (
diff --git a/pkg/sentry/platform/kvm/machine_arm64_unsafe.go b/pkg/sentry/platform/kvm/machine_arm64_unsafe.go
index b531f2f85..00801dee6 100644
--- a/pkg/sentry/platform/kvm/machine_arm64_unsafe.go
+++ b/pkg/sentry/platform/kvm/machine_arm64_unsafe.go
@@ -123,6 +123,7 @@ func (c *vCPU) initArchState() error {
 	reg.addr = uint64(reflect.ValueOf(&data).Pointer())
 	regGet.addr = uint64(reflect.ValueOf(&dataGet).Pointer())
 
+	vcpuInit.target = _KVM_ARM_TARGET_GENERIC_V8
 	vcpuInit.features[0] |= (1 << _KVM_ARM_VCPU_PSCI_0_2)
 	if _, _, errno := syscall.RawSyscall(
 		syscall.SYS_IOCTL,
@@ -134,7 +135,8 @@ func (c *vCPU) initArchState() error {
 
 	// cpacr_el1
 	reg.id = _KVM_ARM64_REGS_CPACR_EL1
-	data = (_FPEN_NOTRAP << _FPEN_SHIFT)
+	// It is off by default, and it is turned on only when in use.
+	data = 0 // Disable fpsimd.
 	if err := c.setOneRegister(&reg); err != nil {
 		return err
 	}
@@ -226,6 +228,7 @@ func (c *vCPU) initArchState() error {
 		return err
 	}
 
+	c.floatingPointState = arch.NewFloatingPointData()
 	return nil
 }
 
diff --git a/pkg/sentry/platform/ring0/defs_arm64.go b/pkg/sentry/platform/ring0/defs_arm64.go
index 1583dda12..0e2ab716c 100644
--- a/pkg/sentry/platform/ring0/defs_arm64.go
+++ b/pkg/sentry/platform/ring0/defs_arm64.go
@@ -124,6 +124,12 @@ func (c *CPU) SetAppAddr(value uintptr) {
 	c.appAddr = value
 }
 
+// GetLazyVFP returns the value of cpacr_el1.
+//go:nosplit
+func (c *CPU) GetLazyVFP() (value uintptr) {
+	return c.lazyVFP
+}
+
 // SwitchArchOpts are embedded in SwitchOpts.
 type SwitchArchOpts struct {
 	// UserASID indicates that the application ASID to be used on switch,
diff --git a/pkg/sentry/platform/ring0/lib_arm64.go b/pkg/sentry/platform/ring0/lib_arm64.go
index af075aae4..80922f43d 100644
--- a/pkg/sentry/platform/ring0/lib_arm64.go
+++ b/pkg/sentry/platform/ring0/lib_arm64.go
@@ -20,13 +20,13 @@ package ring0
 func CPACREL1() (value uintptr)
 
 // FPCR returns the value of FPCR register.
-func FPCR() (value uintptr)
+func GetFPCR() (value uintptr)
 
 // SetFPCR writes the FPCR value.
 func SetFPCR(value uintptr)
 
 // FPSR returns the value of FPSR register.
-func FPSR() (value uintptr)
+func GetFPSR() (value uintptr)
 
 // SetFPSR writes the FPSR value.
 func SetFPSR(value uintptr)