[SYCL] Fix weak_object for host_accessor and stream (#8903)

This commit makes a specialization of weak_object for stream, similar to
buffer, to help recreate a stream from its constituents. Additionally it
fixes a problem for host_accessor where the bases would conflict with
eachother as they defined the same functions.

---------

Signed-off-by: Larsen, Steffen <[email protected]>

Author: steffenlarsen

sycl/include/sycl/accessor.hpp

@@ -3060,9 +3060,7 @@ template <typename DataT, int Dimensions = 1,
           access_mode AccessMode = access_mode::read_write>
 class __SYCL_EBO host_accessor
     : public accessor<DataT, Dimensions, AccessMode, target::host_buffer,
-                      access::placeholder::false_t>,
-      public detail::OwnerLessBase<
-          host_accessor<DataT, Dimensions, AccessMode>> {
+                      access::placeholder::false_t> {
 protected:
   using AccessorT = accessor<DataT, Dimensions, AccessMode, target::host_buffer,
                              access::placeholder::false_t>;
@@ -3083,6 +3081,18 @@ class __SYCL_EBO host_accessor
     AccessorT::__init(Ptr, AccessRange, MemRange, Offset);
   }
 
+#ifndef __SYCL_DEVICE_ONLY__
+  host_accessor(const detail::AccessorImplPtr &Impl)
+      : accessor<DataT, Dimensions, AccessMode, target::host_buffer,
+                 access::placeholder::false_t>{Impl} {}
+
+  template <class Obj>
+  friend decltype(Obj::impl) getSyclObjImpl(const Obj &SyclObject);
+
+  template <class T>
+  friend T detail::createSyclObjFromImpl(decltype(T::impl) ImplObj);
+#endif // __SYCL_DEVICE_ONLY__
+
 public:
   host_accessor() : AccessorT() {}
 
@@ -3240,6 +3250,28 @@ class __SYCL_EBO host_accessor
     *AccessorT::getQualifiedPtr() = std::move(Other);
     return *this;
   }
+
+  // host_accessor needs to explicitly define the owner_before member functions
+  // as inheriting from OwnerLessBase causes base class conflicts.
+  // TODO: Once host_accessor is detached from accessor, inherit from
+  // OwnerLessBase instead.
+#ifndef __SYCL_DEVICE_ONLY__
+  bool ext_oneapi_owner_before(
+      const ext::oneapi::detail::weak_object_base<host_accessor> &Other)
+      const noexcept {
+    return this->impl.owner_before(
+        ext::oneapi::detail::getSyclWeakObjImpl(Other));
+  }
+
+  bool ext_oneapi_owner_before(const host_accessor &Other) const noexcept {
+    return this->impl.owner_before(Other.impl);
+  }
+#else
+  bool ext_oneapi_owner_before(
+      const ext::oneapi::detail::weak_object_base<host_accessor> &Other)
+      const noexcept;
+  bool ext_oneapi_owner_before(const host_accessor &Other) const noexcept;
+#endif
 };
 
 template <typename DataT, int Dimensions, typename AllocatorT>

sycl/include/sycl/buffer.hpp

@@ -722,6 +722,9 @@ class buffer : public detail::buffer_plain,
   detail::make_buffer_helper(pi_native_handle, const context &, event, bool);
   template <typename SYCLObjT> friend class ext::oneapi::weak_object;
 
+  // NOTE: These members are required for reconstructing the buffer, but are not
+  // part of the implementation class. If more members are added, they should
+  // also be added to the weak_object specialization for buffers.
   range<dimensions> Range;
   // Offset field specifies the origin of the sub buffer inside the parent
   // buffer

sycl/include/sycl/ext/oneapi/owner_less.hpp

@@ -61,6 +61,8 @@ struct owner_less<kernel_id> : public detail::owner_less_base<kernel_id> {};
 template <>
 struct owner_less<platform> : public detail::owner_less_base<platform> {};
 template <> struct owner_less<queue> : public detail::owner_less_base<queue> {};
+template <>
+struct owner_less<stream> : public detail::owner_less_base<stream> {};
 
 template <bundle_state State>
 struct owner_less<device_image<State>>
@@ -86,10 +88,6 @@ struct owner_less<host_accessor<DataT, Dimensions, AccessMode>>
     : public detail::owner_less_base<
           host_accessor<DataT, Dimensions, AccessMode>> {};
 
-template <typename DataT, int Dimensions>
-struct owner_less<host_accessor<DataT, Dimensions>>
-    : public detail::owner_less_base<host_accessor<DataT, Dimensions>> {};
-
 template <typename DataT, int Dimensions>
 struct owner_less<local_accessor<DataT, Dimensions>>
     : public detail::owner_less_base<local_accessor<DataT, Dimensions>> {};

sycl/include/sycl/ext/oneapi/weak_object.hpp

@@ -11,6 +11,7 @@
 #include <sycl/buffer.hpp>
 #include <sycl/detail/defines_elementary.hpp>
 #include <sycl/ext/oneapi/weak_object_base.hpp>
+#include <sycl/stream.hpp>
 
 #include <optional>
 
@@ -50,7 +51,8 @@ class weak_object : public detail::weak_object_base<SYCLObjT> {
 
   weak_object &operator=(const SYCLObjT &SYCLObj) noexcept {
     // Create weak_ptr from the shared_ptr to SYCLObj's implementation object.
-    this->MObjWeakPtr = GetWeakImpl(SYCLObj);
+    this->MObjWeakPtr =
+        detail::weak_object_base<SYCLObjT>::GetWeakImpl(SYCLObj);
     return *this;
   }
   weak_object &operator=(const weak_object &Other) noexcept = default;
@@ -103,7 +105,8 @@ class weak_object<buffer<T, Dimensions, AllocatorT>>
 
   weak_object &operator=(const buffer_type &SYCLObj) noexcept {
     // Create weak_ptr from the shared_ptr to SYCLObj's implementation object.
-    this->MObjWeakPtr = GetWeakImpl(SYCLObj);
+    this->MObjWeakPtr = detail::weak_object_base<
+        buffer<T, Dimensions, AllocatorT>>::GetWeakImpl(SYCLObj);
     this->MRange = SYCLObj.Range;
     this->MOffsetInBytes = SYCLObj.OffsetInBytes;
     this->MIsSubBuffer = SYCLObj.IsSubBuffer;
@@ -112,6 +115,13 @@ class weak_object<buffer<T, Dimensions, AllocatorT>>
   weak_object &operator=(const weak_object &Other) noexcept = default;
   weak_object &operator=(weak_object &&Other) noexcept = default;
 
+  void swap(weak_object &Other) noexcept {
+    this->MObjWeakPtr.swap(Other.MObjWeakPtr);
+    std::swap(MRange, Other.MRange);
+    std::swap(MOffsetInBytes, Other.MOffsetInBytes);
+    std::swap(MIsSubBuffer, Other.MIsSubBuffer);
+  }
+
 #ifndef __SYCL_DEVICE_ONLY__
   std::optional<buffer_type> try_lock() const noexcept {
     auto MObjImplPtr = this->MObjWeakPtr.lock();
@@ -141,6 +151,78 @@ class weak_object<buffer<T, Dimensions, AllocatorT>>
   bool MIsSubBuffer;
 };
 
+// Specialization of weak_object for stream as it needs additional members
+// to reconstruct the original stream.
+template <>
+class weak_object<stream> : public detail::weak_object_base<stream> {
+public:
+  using object_type = typename detail::weak_object_base<stream>::object_type;
+
+  constexpr weak_object() noexcept : detail::weak_object_base<stream>() {}
+  weak_object(const stream &SYCLObj) noexcept
+      : detail::weak_object_base<stream>(SYCLObj),
+        MWeakGlobalBuf{SYCLObj.GlobalBuf},
+        MWeakGlobalOffset{SYCLObj.GlobalOffset},
+        MWeakGlobalFlushBuf{SYCLObj.GlobalFlushBuf} {}
+  weak_object(const weak_object &Other) noexcept = default;
+  weak_object(weak_object &&Other) noexcept = default;
+
+  weak_object &operator=(const stream &SYCLObj) noexcept {
+    // Create weak_ptr from the shared_ptr to SYCLObj's implementation object.
+    this->MObjWeakPtr = detail::weak_object_base<stream>::GetWeakImpl(SYCLObj);
+    MWeakGlobalBuf = SYCLObj.GlobalBuf;
+    MWeakGlobalOffset = SYCLObj.GlobalOffset;
+    MWeakGlobalFlushBuf = SYCLObj.GlobalFlushBuf;
+    return *this;
+  }
+  weak_object &operator=(const weak_object &Other) noexcept = default;
+  weak_object &operator=(weak_object &&Other) noexcept = default;
+
+  void swap(weak_object &Other) noexcept {
+    this->MObjWeakPtr.swap(Other.MObjWeakPtr);
+    MWeakGlobalBuf.swap(Other.MWeakGlobalBuf);
+    MWeakGlobalOffset.swap(Other.MWeakGlobalOffset);
+    MWeakGlobalFlushBuf.swap(Other.MWeakGlobalFlushBuf);
+  }
+
+  void reset() noexcept {
+    this->MObjWeakPtr.reset();
+    MWeakGlobalBuf.reset();
+    MWeakGlobalOffset.reset();
+    MWeakGlobalFlushBuf.reset();
+  }
+
+#ifndef __SYCL_DEVICE_ONLY__
+  std::optional<stream> try_lock() const noexcept {
+    auto ObjImplPtr = this->MObjWeakPtr.lock();
+    auto GlobalBuf = MWeakGlobalBuf.try_lock();
+    auto GlobalOffset = MWeakGlobalOffset.try_lock();
+    auto GlobalFlushBuf = MWeakGlobalFlushBuf.try_lock();
+    if (!ObjImplPtr || !GlobalBuf || !GlobalOffset || !GlobalFlushBuf)
+      return std::nullopt;
+    return stream{ObjImplPtr, *GlobalBuf, *GlobalOffset, *GlobalFlushBuf};
+  }
+  stream lock() const {
+    std::optional<stream> OptionalObj = try_lock();
+    if (!OptionalObj)
+      throw sycl::exception(sycl::make_error_code(sycl::errc::invalid),
+                            "Referenced object has expired.");
+    return *OptionalObj;
+  }
+#else
+  // On device calls to these functions are disallowed, so declare them but
+  // don't define them to avoid compilation failures.
+  std::optional<stream> try_lock() const noexcept;
+  stream lock() const;
+#endif // __SYCL_DEVICE_ONLY__
+
+private:
+  // Additional members required for recreating stream.
+  weak_object<detail::GlobalBufAccessorT> MWeakGlobalBuf;
+  weak_object<detail::GlobalOffsetAccessorT> MWeakGlobalOffset;
+  weak_object<detail::GlobalBufAccessorT> MWeakGlobalFlushBuf;
+};
+
 } // namespace ext::oneapi
 } // __SYCL_INLINE_VER_NAMESPACE(_V1)
 } // namespace sycl

sycl/include/sycl/stream.hpp

@@ -743,6 +743,17 @@ inline __width_manipulator__ setw(int Width) {
 /// \ingroup sycl_api
 class __SYCL_EXPORT __SYCL_SPECIAL_CLASS __SYCL_TYPE(stream) stream
     : public detail::OwnerLessBase<stream> {
+private:
+#ifndef __SYCL_DEVICE_ONLY__
+  // Constructor for recreating a stream.
+  stream(std::shared_ptr<detail::stream_impl> Impl,
+         detail::GlobalBufAccessorT GlobalBuf,
+         detail::GlobalOffsetAccessorT GlobalOffset,
+         detail::GlobalBufAccessorT GlobalFlushBuf)
+      : impl{Impl}, GlobalBuf{GlobalBuf}, GlobalOffset{GlobalOffset},
+        GlobalFlushBuf{GlobalFlushBuf} {}
+#endif
+
 public:
 #ifdef __SYCL_DEVICE_ONLY__
   // Default constructor for objects later initialized with __init member.
@@ -811,6 +822,10 @@ class __SYCL_EXPORT __SYCL_SPECIAL_CLASS __SYCL_TYPE(stream) stream
   friend decltype(Obj::impl) detail::getSyclObjImpl(const Obj &SyclObject);
 #endif
 
+  // NOTE: Some members are required for reconstructing the stream, but are not
+  // part of the implementation class. If more members are added, they should
+  // also be added to the weak_object specialization for streams.
+
   // Accessor to the global stream buffer. Global buffer contains all output
   // from the kernel.
   mutable detail::GlobalBufAccessorT GlobalBuf;
@@ -942,6 +957,8 @@ class __SYCL_EXPORT __SYCL_SPECIAL_CLASS __SYCL_TYPE(stream) stream
 
   friend class handler;
 
+  template <typename SYCLObjT> friend class ext::oneapi::weak_object;
+
   friend const stream &operator<<(const stream &, const char);
   friend const stream &operator<<(const stream &, const char *);
   template <typename ValueType>

sycl/test-e2e/WeakObject/weak_object_utils.hpp

@@ -37,6 +37,9 @@ template <template <typename> typename CallableT> void runTest(sycl::queue Q) {
   sycl::accessor<int, 3, sycl::access::mode::read_write,
                  sycl::access::target::host_buffer>
       HAcc3D;
+  sycl::host_accessor<int, 1> HAcc1D_2020;
+  sycl::host_accessor<int, 2> HAcc2D_2020;
+  sycl::host_accessor<int, 3> HAcc3D_2020;
 
   CallableT<decltype(Plt)>()(Plt);
   CallableT<decltype(Dev)>()(Dev);
@@ -54,6 +57,9 @@ template <template <typename> typename CallableT> void runTest(sycl::queue Q) {
   CallableT<decltype(HAcc1D)>()(HAcc1D);
   CallableT<decltype(HAcc2D)>()(HAcc2D);
   CallableT<decltype(HAcc3D)>()(HAcc3D);
+  CallableT<decltype(HAcc1D_2020)>()(HAcc1D_2020);
+  CallableT<decltype(HAcc2D_2020)>()(HAcc2D_2020);
+  CallableT<decltype(HAcc3D_2020)>()(HAcc3D_2020);
 
   Q.submit([&](sycl::handler &CGH) {
     sycl::accessor DAcc1D{Buf1D, CGH, sycl::read_only};
@@ -62,13 +68,15 @@ template <template <typename> typename CallableT> void runTest(sycl::queue Q) {
     sycl::local_accessor<int> LAcc1D{1, CGH};
     sycl::local_accessor<int, 2> LAcc2D{sycl::range<2>{1, 2}, CGH};
     sycl::local_accessor<int, 3> LAcc3D{sycl::range<3>{1, 2, 3}, CGH};
+    sycl::stream Stream{1024, 32, CGH};
 
     CallableT<decltype(DAcc1D)>()(DAcc1D);
     CallableT<decltype(DAcc2D)>()(DAcc2D);
     CallableT<decltype(DAcc3D)>()(DAcc3D);
     CallableT<decltype(LAcc1D)>()(LAcc1D);
     CallableT<decltype(LAcc2D)>()(LAcc2D);
     CallableT<decltype(LAcc3D)>()(LAcc3D);
+    CallableT<decltype(Stream)>()(Stream);
   });
 }
 
@@ -120,6 +128,12 @@ void runTestMulti(sycl::queue Q1) {
   sycl::accessor<int, 3, sycl::access::mode::read_write,
                  sycl::access::target::host_buffer>
       HAcc3D2;
+  sycl::host_accessor<int, 1> HAcc1D1_2020;
+  sycl::host_accessor<int, 2> HAcc2D1_2020;
+  sycl::host_accessor<int, 3> HAcc3D1_2020;
+  sycl::host_accessor<int, 1> HAcc1D2_2020;
+  sycl::host_accessor<int, 2> HAcc2D2_2020;
+  sycl::host_accessor<int, 3> HAcc3D2_2020;
 
   CallableT<decltype(Ctx1)>()(Ctx1, Ctx2);
   CallableT<decltype(Q1)>()(Q1, Q2);
@@ -135,6 +149,9 @@ void runTestMulti(sycl::queue Q1) {
   CallableT<decltype(HAcc1D1)>()(HAcc1D1, HAcc1D2);
   CallableT<decltype(HAcc2D1)>()(HAcc2D1, HAcc2D2);
   CallableT<decltype(HAcc3D1)>()(HAcc3D1, HAcc3D2);
+  CallableT<decltype(HAcc1D1_2020)>()(HAcc1D1_2020, HAcc1D2_2020);
+  CallableT<decltype(HAcc2D1_2020)>()(HAcc2D1_2020, HAcc2D2_2020);
+  CallableT<decltype(HAcc3D1_2020)>()(HAcc3D1_2020, HAcc3D2_2020);
 
   Q1.submit([&](sycl::handler &CGH) {
     sycl::accessor DAcc1D1{Buf1D1, CGH, sycl::read_only};
@@ -149,12 +166,15 @@ void runTestMulti(sycl::queue Q1) {
     sycl::local_accessor<int, 2> LAcc2D2{sycl::range<2>{1, 2}, CGH};
     sycl::local_accessor<int, 3> LAcc3D1{sycl::range<3>{1, 2, 3}, CGH};
     sycl::local_accessor<int, 3> LAcc3D2{sycl::range<3>{1, 2, 3}, CGH};
+    sycl::stream Stream1{1024, 32, CGH};
+    sycl::stream Stream2{1024, 32, CGH};
 
     CallableT<decltype(DAcc1D1)>()(DAcc1D1, DAcc1D2);
     CallableT<decltype(DAcc2D1)>()(DAcc2D1, DAcc2D2);
     CallableT<decltype(DAcc3D1)>()(DAcc3D1, DAcc3D2);
     CallableT<decltype(LAcc1D1)>()(LAcc1D1, LAcc1D2);
     CallableT<decltype(LAcc2D1)>()(LAcc2D1, LAcc2D2);
     CallableT<decltype(LAcc3D1)>()(LAcc3D1, LAcc3D2);
+    CallableT<decltype(Stream1)>()(Stream1, Stream2);
   });
 }