Fixup AMX workaround for ptrace.

SETREGSET/GETREGSET expect AMX portions of fpstate to always be used.
For this reason we need to allocate enough memory for this to happen,
even if we never populate the AMX portions within initX86FPState.

PiperOrigin-RevId: 599702181

pkg/cpuid/cpuid_amd64.go

@@ -384,17 +384,13 @@ const (
 // Extended state includes floating point registers, and other cpu state that's
 // not associated with the normal task context.
 //
-// We do not support enabling AMX within gVisor, therefore always exclude
-// AMXExtendedStateSize.
-// TODO(gvisor.dev/issues/9896): Implement AMX Support.
-//
 // Note: the return value matches the size of signal FP state frames.
 // Look at check_xstate_in_sigframe() in the kernel sources for more details.
 //
 //go:nosplit
 func (fs FeatureSet) ExtendedStateSize() (size, align uint) {
 	if fs.UseXsave() {
-		return uint(xsaveSize) - fs.AMXExtendedStateSize(), 64
+		return uint(xsaveSize), 64
 	}
 
 	// If we don't support xsave, we fall back to fxsave, which requires
@@ -423,7 +419,7 @@ func (fs FeatureSet) ValidXCR0Mask() uint64 {
 		return 0
 	}
 	ax, _, _, dx := fs.query(xSaveInfo)
-	return (uint64(dx)<<32 | uint64(ax)) ^ XCR0AMXMask
+	return (uint64(dx)<<32 | uint64(ax)) &^ XCR0AMXMask
 }
 
 // UseXsave returns the choice of fp state saving instruction.

pkg/ring0/entry_amd64.s

@@ -233,9 +233,7 @@ TEXT ·doSwitchToUser(SB),NOSPLIT,$16-48
 	MOVB ·hasXSAVE(SB), BX
 	TESTB BX, BX
 	JZ no_xrstor
-	// Use xrstor to restore all available fp state. For now, we restore
-	// everything unconditionally by setting the implicit operand edx:eax
-	// (the "requested feature bitmap") to all 1's.
+	// Use xrstor to restore all available fp state.
 	MOVL $XCR0_EAX, AX
 	MOVL $XCR0_EDX, DX
 	BYTE $0x48; BYTE $0x0f; BYTE $0xae; BYTE $0x2f // XRSTOR64 0(DI)
@@ -281,7 +279,6 @@ done_sysret_or_iret:
 	TESTB BX, BX
 	JZ no_xsave
 	// Use xsave/xsaveopt to save all extended state.
-	// We save everything unconditionally by setting RFBM to all 1's.
 	MOVL $XCR0_EAX, AX
 	MOVL $XCR0_EDX, DX
 	TESTB CX, CX
@@ -519,7 +516,6 @@ kernel:
 	TESTB BX, BX
 	JZ no_xsave
 	// Use xsave/xsaveopt to save all extended state.
-	// We save everything unconditionally by setting RFBM to all 1's.
 	MOVL $XCR0_EAX, AX
 	MOVL $XCR0_EDX, DX
 	TESTB CX, CX
@@ -651,7 +647,6 @@ kernel:
 	TESTB BX, BX
 	JZ no_xsave
 	// Use xsave/xsaveopt to save all extended state.
-	// We save everything unconditionally by setting RFBM to all 1's.
 	MOVL $XCR0_EAX, AX
 	MOVL $XCR0_EDX, DX
 	TESTB CX, CX

pkg/ring0/lib_amd64.go

@@ -119,7 +119,8 @@ func Init(fs cpuid.FeatureSet) {
 	hasFSGSBASE = fs.HasFeature(cpuid.X86FeatureFSGSBase)
 	validXCR0Mask = uintptr(fs.ValidXCR0Mask())
 	if hasXSAVE {
-		localXCR0 = xgetbv(0)
+		XCR0AMXMask := uintptr((1 << 17) | (1 << 18))
+		localXCR0 = xgetbv(0) &^ XCR0AMXMask
 	}
 }
 

pkg/sentry/arch/fpu/fpu_amd64.go

@@ -70,8 +70,13 @@ const (
 func initX86FPState(data *byte, useXsave bool)
 
 func newX86FPStateSlice() State {
-	size, align := cpuid.HostFeatureSet().ExtendedStateSize()
-	capacity := size + FP_XSTATE_MAGIC2_SIZE
+	maxsize, align := cpuid.HostFeatureSet().ExtendedStateSize()
+	// We need capacity to be large enough to hold AMX bytes because of
+	// ptrace. PTRACE_SETREGSET/GETREGSET assume that AMX portions should
+	// always be used.
+	// TODO(gvisor.dev/issues/9896): Implement AMX Support.
+	capacity := maxsize + FP_XSTATE_MAGIC2_SIZE
+	size := maxsize - cpuid.HostFeatureSet().AMXExtendedStateSize()
 	// Always use at least 4096 bytes.
 	//
 	// For the KVM platform, this state is a fixed 4096 bytes, so make sure
@@ -129,6 +134,8 @@ func InitHostState() {
 		hostXCR0Mask = featureSet.ValidXCR0Mask()
 		hostUseXsave = featureSet.UseXsave()
 		hostFPSize, _ = featureSet.ExtendedStateSize()
+		// TODO(gvisor.dev/issues/9896): Implement AMX Support.
+		hostFPSize = hostFPSize - featureSet.AMXExtendedStateSize()
 	})
 }
 

pkg/sentry/platform/systrap/sysmsg/sysmsg_amd64.go

@@ -52,7 +52,8 @@ func (s *ArchState) Init() {
 	fs := cpuid.HostFeatureSet()
 
 	fpLenUint, _ := fs.ExtendedStateSize()
-	s.fpLen = uint32(fpLenUint)
+	// TODO(gvisor.dev/issues/9896): Implement AMX Support.
+	s.fpLen = uint32(fpLenUint - fs.AMXExtendedStateSize())
 	if fs.UseXsaveopt() {
 		s.xsaveMode = xsaveopt
 	} else if fs.UseXsave() {