NVIDIA · jbrasen · Apr 17, 2025
diff --git a/Platform/NVIDIA/NVIDIA.common.dsc.inc b/Platform/NVIDIA/NVIDIA.common.dsc.inc
@@ -441,6 +441,7 @@ CONFIG_LOGO_IMPLEMENTER = noimplementer.inf
   MemoryVerificationLib|Silicon/NVIDIA/Library/MemoryVerificationLib/MemoryVerificationLib.inf
 
   SystemContextLib|Silicon/NVIDIA/Library/SystemContextLib/SystemContextLib.inf
+  SystemFiberLib|Silicon/NVIDIA/Library/SystemFiberLib/SystemFiberLib.inf
   StatusRegLib|Silicon/NVIDIA/Library/StatusRegLib/StatusRegLib.inf
 
   OrderedCollectionLib|MdePkg/Library/BaseOrderedCollectionRedBlackTreeLib/BaseOrderedCollectionRedBlackTreeLib.inf

diff --git a/Silicon/NVIDIA/Include/Library/SystemFiberLib.h b/Silicon/NVIDIA/Include/Library/SystemFiberLib.h
@@ -0,0 +1,91 @@
+/** @file
+
+  SystemFiberLib - Library for system fiber operations
+
+  SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+  SPDX-License-Identifier: BSD-2-Clause-Patent
+
+**/
+
+#ifndef SYSTEM_FIBER_LIB_H__
+#define SYSTEM_FIBER_LIB_H__
+
+#include <Uefi/UefiBaseType.h>
+
+typedef VOID (*SYSTEM_FIBER_ENTRY_POINT)(
+  VOID  *Context
+  );
+typedef VOID *SYSTEM_FIBER;
+
+/**
+  Creates a new fiber.
+  Fiber will not be started until ResumeSystemFiber is called.
+
+  @param[in] EntryPoint  Entry point of the fiber
+  @param[in] Context     Context of the fiber
+  @param[in] StackSize   Stack size of the fiber
+  @param[out] Fiber     Fiber object
+
+  @return EFI_SUCCESS if the fiber was created successfully
+  @return EFI_OUT_OF_RESOURCES if the fiber could not be created
+  @return EFI_INVALID_PARAMETER if the entry point is NULL
+*/
+EFI_STATUS
+EFIAPI
+CreateSystemFiber (
+  IN SYSTEM_FIBER_ENTRY_POINT  EntryPoint,
+  IN VOID                      *Context,
+  IN UINTN                     StackSize,
+  OUT SYSTEM_FIBER             *Fiber
+  );
+
+/**
+  Destroys a fiber
+  If the fiber is currently running, it will yield before being destroyed
+
+  @param[in] Fiber  Fiber object
+
+  @return EFI_SUCCESS if the fiber was destroyed successfully
+  @return EFI_INVALID_PARAMETER if the fiber is NULL
+*/
+EFI_STATUS
+EFIAPI
+DestroySystemFiber (
+  IN SYSTEM_FIBER  Fiber
+  );
+
+/**
+  Resumes a fiber
+  Will not return until the fiber yields.
+
+  @param[in] Fiber  Fiber object
+
+  @return EFI_SUCCESS if the fiber was resumed successfully
+  @return EFI_INVALID_PARAMETER if the fiber is NULL
+  @return EFI_ALREADY_STARTED if the fiber is already running
+  @return EFI_ABORTED if the fiber is marked as destroyed
+*/
+EFI_STATUS
+EFIAPI
+ResumeSystemFiber (
+  IN SYSTEM_FIBER  Fiber
+  );
+
+/**
+  Yields the current fiber.
+  This function will return to the caller of ResumeSystemFiber.
+
+  @param[in] Fiber  Fiber object
+
+  @return EFI_SUCCESS if the fiber was yielded successfully
+  @return EFI_INVALID_PARAMETER if the fiber is NULL
+  @return EFI_NOT_STARTED if the fiber is not running
+*/
+EFI_STATUS
+EFIAPI
+YieldSystemFiber (
+  IN SYSTEM_FIBER  Fiber
+  );
+
+#endif
diff --git a/Silicon/NVIDIA/Library/DeviceDiscoveryDriverLib/DeviceDiscoveryDriverLib.c b/Silicon/NVIDIA/Library/DeviceDiscoveryDriverLib/DeviceDiscoveryDriverLib.c
@@ -2,7 +2,7 @@
 
   Device Discovery Driver Library
 
-  SPDX-FileCopyrightText: Copyright (c) 2018-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+  SPDX-FileCopyrightText: Copyright (c) 2018-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
   SPDX-License-Identifier: BSD-2-Clause-Patent
 
@@ -22,8 +22,8 @@
 #include <Library/DeviceTreeHelperLib.h>
 #include <Library/PlatformResourceLib.h>
 #include <Library/MemoryAllocationLib.h>
+#include <Library/SystemFiberLib.h>
 #include <Library/TimerLib.h>
-#include <Library/SystemContextLib.h>
 
 #include <Protocol/AsyncDriverStatus.h>
 #include <Protocol/NonDiscoverableDevice.h>
@@ -40,7 +40,6 @@
 SCMI_CLOCK2_PROTOCOL                           *gScmiClockProtocol    = NULL;
 NVIDIA_CLOCK_PARENTS_PROTOCOL                  *gClockParentsProtocol = NULL;
 STATIC EFI_HANDLE                              mImageHandle           = NULL;
-STATIC EFI_SYSTEM_CONTEXT_AARCH64              MainContext            = { 0 };
 STATIC UINTN                                   SubThreadsRunning      = 0;
 STATIC NVIDIA_ASYNC_DRIVER_STATUS_PROTOCOL     *AsyncProtocol         = NULL;
 STATIC NVIDIA_DEVICE_DISCOVERY_THREAD_CONTEXT  *CurrentThread         = NULL;
@@ -208,10 +207,7 @@ DeviceDiscoveryThreadCallback (
   )
 {
   CurrentThread = (NVIDIA_DEVICE_DISCOVERY_THREAD_CONTEXT *)Context;
-  SwapSystemContext (
-    (EFI_SYSTEM_CONTEXT)&MainContext,
-    (EFI_SYSTEM_CONTEXT)&CurrentThread->Context
-    );
+  ResumeSystemFiber (CurrentThread->Fiber);
 }
 
 /**
@@ -239,10 +235,7 @@ DeviceDiscoveryThreadMicroSecondDelay (
   ThreadContext = CurrentThread;
   gBS->SetTimer (ThreadContext->Timer, TimerRelative, MicroSeconds*10);
   CurrentThread = NULL;
-  SwapSystemContext (
-    (EFI_SYSTEM_CONTEXT)&ThreadContext->Context,
-    (EFI_SYSTEM_CONTEXT)&MainContext
-    );
+  YieldSystemFiber (ThreadContext->Fiber);
   return MicroSeconds;
 }
 
@@ -255,12 +248,14 @@ DeviceDiscoveryThreadMicroSecondDelay (
 STATIC
 VOID
 DeviceThreadMain (
-  IN NVIDIA_DEVICE_DISCOVERY_THREAD_CONTEXT  *ThreadContext
+  IN VOID  *Context
   )
 {
-  EFI_STATUS  Status;
+  EFI_STATUS                              Status;
+  NVIDIA_DEVICE_DISCOVERY_THREAD_CONTEXT  *ThreadContext;
+
+  ThreadContext = (NVIDIA_DEVICE_DISCOVERY_THREAD_CONTEXT *)Context;
 
-  CurrentThread = ThreadContext;
   SubThreadsRunning++;
 
   Status = DeviceDiscoveryNotify (
@@ -284,16 +279,10 @@ DeviceThreadMain (
                NULL,
                NULL
                );
+    if (EFI_ERROR (Status)) {
+      DEBUG ((DEBUG_ERROR, "%a, driver returned %r to start notification\r\n", __FUNCTION__, Status));
+    }
   }
-
-  SwapSystemContext (
-    (EFI_SYSTEM_CONTEXT)&ThreadContext->Context,
-    (EFI_SYSTEM_CONTEXT)&MainContext
-    );
-
-  // Should never get here
-  ASSERT (FALSE);
-  CpuDeadLoop ();
 }
 
 /**
@@ -349,17 +338,15 @@ ThreadedDeviceStart (
   NewContext->DriverHandle = DriverHandle;
   NewContext->Node         = Node;
 
-  // Don't change the special registers
-  GetSystemContext ((EFI_SYSTEM_CONTEXT)&MainContext);
-  NewContext->Context.ELR  = MainContext.ELR;
-  NewContext->Context.SPSR = MainContext.SPSR;
-  NewContext->Context.FPSR = MainContext.FPSR;
-  NewContext->Context.ESR  = MainContext.ESR;
-  NewContext->Context.FAR  = MainContext.FAR;
-
-  NewContext->Context.LR = (UINT64)DeviceThreadMain;
-  NewContext->Context.SP = NewContext->StackBase + THREAD_STACK_SIZE;
-  NewContext->Context.X0 = (UINT64)NewContext;
+  Status = CreateSystemFiber (
+             DeviceThreadMain,
+             NewContext,
+             THREAD_STACK_SIZE,
+             &NewContext->Fiber
+             );
+  if (EFI_ERROR (Status)) {
+    goto ErrorExit;
+  }
 
   if (AsyncProtocol == NULL) {
     AsyncProtocol = (NVIDIA_ASYNC_DRIVER_STATUS_PROTOCOL *)AllocatePool (sizeof (NVIDIA_ASYNC_DRIVER_STATUS_PROTOCOL));
@@ -373,10 +360,7 @@ ThreadedDeviceStart (
   }
 
   OldTpl = gBS->RaiseTPL (TPL_CALLBACK);
-  SwapSystemContext (
-    (EFI_SYSTEM_CONTEXT)&MainContext,
-    (EFI_SYSTEM_CONTEXT)&NewContext->Context
-    );
+  ResumeSystemFiber (NewContext->Fiber);
   gBS->RestoreTPL (OldTpl);
 
 ErrorExit:

diff --git a/Silicon/NVIDIA/Library/DeviceDiscoveryDriverLib/DeviceDiscoveryDriverLib.inf b/Silicon/NVIDIA/Library/DeviceDiscoveryDriverLib/DeviceDiscoveryDriverLib.inf
@@ -2,7 +2,7 @@
 #
 #  Device discovery driver library
 #
-#  SPDX-FileCopyrightText: Copyright (c) 2018-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+#  SPDX-FileCopyrightText: Copyright (c) 2018-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 #
 #  SPDX-License-Identifier: BSD-2-Clause-Patent
 #
@@ -37,7 +37,7 @@
   IoLib
   DeviceDiscoveryLib
   DeviceTreeHelperLib
-  SystemContextLib
+  SystemFiberLib
 
 [Protocols]
   gEdkiiNonDiscoverableDeviceProtocolGuid

diff --git a/Silicon/NVIDIA/Library/DeviceDiscoveryDriverLib/DeviceDiscoveryDriverLibPrivate.h b/Silicon/NVIDIA/Library/DeviceDiscoveryDriverLib/DeviceDiscoveryDriverLibPrivate.h
@@ -2,7 +2,7 @@
 
   Device Discovery Driver Library private structures
 
-  SPDX-FileCopyrightText: Copyright (c) 2018-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+  SPDX-FileCopyrightText: Copyright (c) 2018-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
   SPDX-License-Identifier: BSD-2-Clause-Patent
 
@@ -14,6 +14,7 @@
 #include <PiDxe.h>
 #include <Protocol/ArmScmiClock2Protocol.h>
 #include <Protocol/ClockParents.h>
+#include <Library/SystemFiberLib.h>
 
 extern SCMI_CLOCK2_PROTOCOL           *gScmiClockProtocol;
 extern NVIDIA_CLOCK_PARENTS_PROTOCOL  *gClockParentsProtocol;
@@ -27,7 +28,7 @@ typedef struct {
 typedef struct {
   EFI_PHYSICAL_ADDRESS                   StackBase;
   EFI_EVENT                              Timer;
-  EFI_SYSTEM_CONTEXT_AARCH64             Context;
+  SYSTEM_FIBER                           Fiber;
   EFI_HANDLE                             DriverHandle;
   EFI_HANDLE                             Controller;
   IN NVIDIA_DEVICE_TREE_NODE_PROTOCOL    *Node;

diff --git a/Silicon/NVIDIA/Library/SystemFiberLib/AArch64/SystemFiberLibAArch64.c b/Silicon/NVIDIA/Library/SystemFiberLib/AArch64/SystemFiberLibAArch64.c
@@ -0,0 +1,46 @@
+/** @file
+
+  SystemFiberLibPrivate - AArch64 functions for system fiber operations
+
+  SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+  SPDX-License-Identifier: BSD-2-Clause-Patent
+
+**/
+
+#include "SystemFiberLibPrivate.h"
+#include <Library/BaseMemoryLib.h>
+#include <Library/DebugLib.h>
+
+/**
+  Initialze SystemContext */
+EFI_STATUS
+EFIAPI
+InitializeSystemContext (
+  IN OUT EFI_SYSTEM_CONTEXT  *SystemContext,
+  SYSTEM_FIBER_ENTRY_POINT   EntryPoint,
+  VOID                       *Context,
+  UINT8                      *Stack,
+  UINTN                      StackSize
+  )
+{
+  EFI_SYSTEM_CONTEXT          CurrentSystemContext;
+  EFI_SYSTEM_CONTEXT_AARCH64  CurrentSystemContextAArch64;
+
+  CurrentSystemContext.SystemContextAArch64 = &CurrentSystemContextAArch64;
+
+  ZeroMem (SystemContext->SystemContextAArch64, sizeof (EFI_SYSTEM_CONTEXT_AARCH64));
+
+  GetSystemContext (CurrentSystemContext);
+  SystemContext->SystemContextAArch64->ELR  = CurrentSystemContextAArch64.ELR;
+  SystemContext->SystemContextAArch64->SPSR = CurrentSystemContextAArch64.SPSR;
+  SystemContext->SystemContextAArch64->FPSR = CurrentSystemContextAArch64.FPSR;
+  SystemContext->SystemContextAArch64->ESR  = CurrentSystemContextAArch64.ESR;
+  SystemContext->SystemContextAArch64->FAR  = CurrentSystemContextAArch64.FAR;
+
+  SystemContext->SystemContextAArch64->LR = (UINT64)EntryPoint;
+  SystemContext->SystemContextAArch64->SP = (UINT64)Stack + StackSize;
+  SystemContext->SystemContextAArch64->X0 = (UINT64)Context;
+
+  return EFI_SUCCESS;
+}