[SYCL] fix for tensorflow, linux only (#14153)

On Linux, the TensorFlow library (which uses SYCL .so) is having a
problem wherein its own releasing of resource is running afoul ours. The
solution is to stop deferring any new memory releases once we know we
are shutting down. To that end I've broken the shutdown into two parts
(early and late) and added a flag for whether its OK to defer memory or
not. This has been confirmed to fix the problem.
It is not entirely clear to me what, exactly, the underlying cause of
the problem is on the TensorFlow side. But I realized that while SYCL is
used by several intermediate libs we have no testing of this scenario,
so I've begun to add some.
Lastly, on Windows, I've got encouraging results that this approach may
help us unify shutdown so it's the same on both Linux and Win, but that
is a larger work and should be in its own PR.

Author: cperkinsintel

sycl/source/detail/global_handler.cpp

@@ -36,6 +36,11 @@ namespace detail {
 using LockGuard = std::lock_guard<SpinLock>;
 SpinLock GlobalHandler::MSyclGlobalHandlerProtector{};
 
+// forward decl
+void shutdown_win(); // TODO: win variant will go away soon
+void shutdown_early();
+void shutdown_late();
+
 // Utility class to track references on object.
 // Used for GlobalHandler now and created as thread_local object on the first
 // Scheduler usage. Origin idea is to track usage of Scheduler from main and
@@ -227,16 +232,25 @@ void GlobalHandler::releaseDefaultContexts() {
   MPlatformToDefaultContextCache.Inst.reset(nullptr);
 }
 
-struct DefaultContextReleaseHandler {
-  ~DefaultContextReleaseHandler() {
+struct EarlyShutdownHandler {
+  ~EarlyShutdownHandler() {
+#ifdef _WIN32
+    // on Windows we keep to the existing shutdown procedure
     GlobalHandler::instance().releaseDefaultContexts();
+#else
+    shutdown_early();
+#endif
   }
 };
 
-void GlobalHandler::registerDefaultContextReleaseHandler() {
-  static DefaultContextReleaseHandler handler{};
+void GlobalHandler::registerEarlyShutdownHandler() {
+  static EarlyShutdownHandler handler{};
 }
 
+bool GlobalHandler::isOkToDefer() const { return OkToDefer; }
+
+void GlobalHandler::endDeferredRelease() { OkToDefer = false; }
+
 // Note: Split from shutdown so it is available to the unittests for ensuring
 //       that the mock plugin is the lone plugin.
 void GlobalHandler::unloadPlugins() {
@@ -279,30 +293,37 @@ void GlobalHandler::drainThreadPool() {
 // itself is very aggressive about reclaiming memory. Thus,
 // we focus solely on unloading the plugins, so as to not
 // accidentally retain device handles. etc
-void shutdown() {
+void shutdown_win() {
   GlobalHandler *&Handler = GlobalHandler::getInstancePtr();
   Handler->unloadPlugins();
 }
 #else
-void shutdown() {
+void shutdown_early() {
   const LockGuard Lock{GlobalHandler::MSyclGlobalHandlerProtector};
   GlobalHandler *&Handler = GlobalHandler::getInstancePtr();
   if (!Handler)
     return;
 
+  // Now that we are shutting down, we will no longer defer MemObj releases.
+  Handler->endDeferredRelease();
+
   // Ensure neither host task is working so that no default context is accessed
   // upon its release
   Handler->prepareSchedulerToRelease(true);
 
   if (Handler->MHostTaskThreadPool.Inst)
     Handler->MHostTaskThreadPool.Inst->finishAndWait();
 
-  // If default contexts are requested after the first default contexts have
-  // been released there may be a new default context. These must be released
-  // prior to closing the plugins.
-  // Note: Releasing a default context here may cause failures in plugins with
-  // global state as the global state may have been released.
+  // This releases OUR reference to the default context, but
+  // other may yet have refs
   Handler->releaseDefaultContexts();
+}
+
+void shutdown_late() {
+  const LockGuard Lock{GlobalHandler::MSyclGlobalHandlerProtector};
+  GlobalHandler *&Handler = GlobalHandler::getInstancePtr();
+  if (!Handler)
+    return;
 
   // First, release resources, that may access plugins.
   Handler->MPlatformCache.Inst.reset(nullptr);
@@ -345,7 +366,7 @@ extern "C" __SYCL_EXPORT BOOL WINAPI DllMain(HINSTANCE hinstDLL,
                    // TODO: figure out what XPTI is doing that prevents release.
 #endif
 
-    shutdown();
+    shutdown_win();
     break;
   case DLL_PROCESS_ATTACH:
     if (PrintPiTrace)
@@ -363,7 +384,7 @@ extern "C" __SYCL_EXPORT BOOL WINAPI DllMain(HINSTANCE hinstDLL,
 // destructors. Priorities 0-100 are reserved by the compiler. The priority
 // value 110 allows SYCL users to run their destructors after runtime library
 // deinitialization.
-__attribute__((destructor(110))) static void syclUnload() { shutdown(); }
+__attribute__((destructor(110))) static void syclUnload() { shutdown_late(); }
 #endif
 } // namespace detail
 } // namespace _V1

sycl/source/detail/global_handler.hpp

@@ -73,8 +73,10 @@ class GlobalHandler {
   XPTIRegistry &getXPTIRegistry();
   ThreadPool &getHostTaskThreadPool();
 
-  static void registerDefaultContextReleaseHandler();
+  static void registerEarlyShutdownHandler();
 
+  bool isOkToDefer() const;
+  void endDeferredRelease();
   void unloadPlugins();
   void releaseDefaultContexts();
   void drainThreadPool();
@@ -91,7 +93,11 @@ class GlobalHandler {
   void *GSYCLCallEvent = nullptr;
 #endif
 
-  friend void shutdown();
+  bool OkToDefer = true;
+
+  friend void shutdown_win();
+  friend void shutdown_early();
+  friend void shutdown_late();
   friend class ObjectUsageCounter;
   static GlobalHandler *&getInstancePtr();
   static SpinLock MSyclGlobalHandlerProtector;

sycl/source/detail/platform_impl.cpp

@@ -193,15 +193,9 @@ std::vector<platform> platform_impl::get_platforms() {
     Platforms.push_back(Platform.first);
   }
 
-  // Register default context release handler after plugins have been loaded and
-  // after the first calls to each plugin. This initializes a function-local
-  // variable that should be destroyed before any global variables in the
-  // plugins are destroyed. This is done after the first call to the backends to
-  // ensure any lazy-loaded dependencies are loaded prior to the handler
-  // variable's initialization. Note: The default context release handler is not
-  // guaranteed to be destroyed before function-local static variables as they
-  // may be initialized after.
-  GlobalHandler::registerDefaultContextReleaseHandler();
+  // This initializes a function-local variable whose destructor is invoked as
+  // the SYCL shared library is first being unloaded.
+  GlobalHandler::registerEarlyShutdownHandler();
 
   return Platforms;
 }

sycl/source/detail/sycl_mem_obj_t.cpp

@@ -229,8 +229,11 @@ void SYCLMemObjT::detachMemoryObject(
       (MInteropContext && !MInteropContext->isOwnedByRuntime());
 
   if (MRecord && MRecord->MCurContext->isOwnedByRuntime() &&
-      !InteropObjectsUsed && (!MHostPtrProvided || MIsInternal))
-    Scheduler::getInstance().deferMemObjRelease(Self);
+      !InteropObjectsUsed && (!MHostPtrProvided || MIsInternal)) {
+    bool okToDefer = GlobalHandler::instance().isOkToDefer();
+    if (okToDefer)
+      Scheduler::getInstance().deferMemObjRelease(Self);
+  }
 }
 
 void SYCLMemObjT::handleWriteAccessorCreation() {

sycl/test-e2e/IntermediateLib/Inputs/simple_lib.cpp

@@ -0,0 +1,58 @@
+/*
+    // compile to static lib
+    clang++ -fsycl -c -fPIC -o simple_lib.o simple_lib.cpp
+
+    // compile to dynamic lib
+    clang++ -fsycl  -fPIC -shared -o simple_lib.so simple_lib.cpp
+
+*/
+
+#include "simple_lib.h"
+#include <sycl/detail/core.hpp>
+
+const size_t BUFF_SIZE = 1;
+
+class Delay {
+public:
+  std::shared_ptr<sycl::buffer<int, 1>> sharedBuffer;
+
+  void release() {
+    std::cout << "Delay.release()" << std::endl;
+    sharedBuffer.reset();
+  }
+
+  const sycl::buffer<int, 1> &getBuffer() {
+    if (!sharedBuffer) {
+      sharedBuffer = std::make_shared<sycl::buffer<int, 1>>(BUFF_SIZE);
+    }
+    return *sharedBuffer;
+  }
+
+  Delay() : sharedBuffer(nullptr) {}
+  ~Delay() { release(); }
+};
+
+#ifdef _WIN32
+static Delay theDelay;
+Delay *MyDelay = &theDelay;
+#else
+Delay *MyDelay = new Delay;
+
+__attribute__((destructor(101))) static void Unload101() {
+  std::cout << "lib unload - __attribute__((destructor(101)))" << std::endl;
+  delete MyDelay;
+}
+#endif
+
+EXPORTDECL int add_using_device(int a, int b) {
+  sycl::queue q;
+  sycl::buffer<int, 1> buf = MyDelay->getBuffer();
+  q.submit([&](sycl::handler &cgh) {
+     sycl::accessor acc(buf, cgh, sycl::write_only);
+
+     cgh.single_task([=] { acc[0] = a + b; });
+   }).wait();
+
+  sycl::host_accessor acc(buf);
+  return acc[0];
+}

sycl/test-e2e/IntermediateLib/Inputs/simple_lib.h

@@ -0,0 +1,12 @@
+#ifndef SIMPLE_SYCL_LIB_H
+#define SIMPLE_SYCL_LIB_H
+
+#ifdef _WIN32
+#define EXPORTDECL extern "C" __declspec(dllexport)
+#else
+#define EXPORTDECL extern "C"
+#endif
+
+EXPORTDECL int add_using_device(int a, int b);
+
+#endif

sycl/test-e2e/IntermediateLib/dynamic_app_linux.cpp

@@ -0,0 +1,68 @@
+// REQUIRES: level_zero && linux
+
+// build shared library
+// RUN: %clangxx -fsycl -fPIC -shared -o %T/simple_lib.so %S/Inputs/simple_lib.cpp
+
+// build app
+// RUN: %clangxx -DSO_PATH="%T/simple_lib.so" -o %t.out %s
+
+// RUN: %{run} %t.out
+// RUN: env UR_L0_LEAKS_DEBUG=1 %{run} %t.out
+
+// In these tests we are building an intermediate library which uses SYCL and an
+// app that employs that intermediate library, using both static and dynamic
+// linking, and delayed release. This is to test that release and shutdown are
+// working correctly.
+
+/*
+    //library
+    clang++ -fsycl  -fPIC -shared -o simple_lib.so Inputs/simple_lib.cpp
+
+    //app
+    clang++ -DSO_PATH="simple_lib.so" -o dynamic_app.bin dynamic_app_linux.cpp
+
+    UR_L0_LEAKS_DEBUG=1 ./dynamic_app.bin
+
+*/
+
+#include "Inputs/simple_lib.h"
+#include <assert.h>
+#include <dlfcn.h>
+#include <iostream>
+
+void *handle = nullptr;
+
+__attribute__((destructor(101))) static void Unload101() {
+  std::cout << "app unload - __attribute__((destructor(101)))" << std::endl;
+  if (handle) {
+    dlclose(handle);
+    handle = nullptr;
+  }
+}
+
+#define STRINGIFY_HELPER(A) #A
+#define STRINGIFY(A) STRINGIFY_HELPER(A)
+#define SO_FNAME "" STRINGIFY(SO_PATH) ""
+
+int main() {
+
+  handle = dlopen(SO_FNAME, RTLD_NOW);
+  if (!handle) {
+    std::cout << "failed to load" << std::endl;
+    return 1;
+  }
+
+  // Function pointer to the exported function
+  int (*add_using_device)(int, int) =
+      (int (*)(int, int))dlsym(handle, "add_using_device");
+  if (!add_using_device) {
+    std::cout << "failed to get function" << std::endl;
+    return 2;
+  }
+
+  int result = add_using_device(3, 4);
+  std::cout << "Result: " << result << std::endl;
+  assert(result == 7);
+
+  return 0;
+}

sycl/test-e2e/IntermediateLib/static_app.cpp

@@ -0,0 +1,37 @@
+// REQUIRES: level_zero && linux
+
+// DEFINE: %{fPIC_flag} =  %if windows %{%} %else %{-fPIC%}
+// build static library
+// RUN: %clangxx -fsycl -c  %{fPIC_flag} -o simple_lib.o %S/Inputs/simple_lib.cpp
+
+// build app
+// RUN:  %clangxx -fsycl -o %t.out %s simple_lib.o
+
+// RUN: %{run} %t.out
+// RUN: env UR_L0_LEAKS_DEBUG=1 %{run} %t.out
+
+// In these tests we are building an intermediate library which uses SYCL and an
+// app that employs that intermediate library, using both static and dynamic
+// linking, and delayed release. This is to test that release and shutdown are
+// working correctly.
+
+/*
+    //library
+    clang++ -fsycl -c -fPIC -o simple_lib.o Inputs/simple_lib.cpp
+
+    //app
+    clang++ -fsycl -o static_app.bin static_app.cpp simple_lib.o
+
+    UR_L0_LEAKS_DEBUG=1 ./simple_app.bin
+
+*/
+
+#include "Inputs/simple_lib.h"
+#include <assert.h>
+#include <iostream>
+
+int main() {
+  int result = add_using_device(3, 4);
+  std::cout << "result: " << result << std::endl;
+  assert(result == 7);
+}