[SYCL] Avoid device_impl shared_ptr copy in getMaxWorkGroups() (#17705)

Author: slawekptak

sycl/include/sycl/handler.hpp

@@ -213,6 +213,7 @@ class HandlerAccess;
 class HostTask;
 
 using EventImplPtr = std::shared_ptr<event_impl>;
+using DeviceImplPtr = std::shared_ptr<device_impl>;
 
 template <typename RetType, typename Func, typename Arg>
 static Arg member_ptr_helper(RetType (Func::*)(Arg) const);
@@ -249,6 +250,7 @@ template <typename Type> struct get_kernel_wrapper_name_t {
 };
 
 __SYCL_EXPORT device getDeviceFromHandler(handler &);
+const DeviceImplPtr &getDeviceImplFromHandler(handler &);
 
 // Checks if a device_global has any registered kernel usage.
 __SYCL_EXPORT bool isDeviceGlobalUsedInKernel(const void *DeviceGlobalPtr);
@@ -3481,6 +3483,8 @@ class __SYCL_EXPORT handler {
             typename PropertyListT>
   friend class accessor;
   friend device detail::getDeviceFromHandler(handler &);
+  friend const detail::DeviceImplPtr &
+  detail::getDeviceImplFromHandler(handler &);
 
   template <typename DataT, int Dimensions, access::mode AccessMode,
             access::target AccessTarget, access::placeholder IsPlaceholder>

sycl/source/detail/graph_impl.hpp

@@ -919,6 +919,12 @@ class graph_impl : public std::enable_shared_from_this<graph_impl> {
   /// @return Context associated with graph.
   sycl::context getContext() const { return MContext; }
 
+  /// Query for the device_impl tied to this graph.
+  /// @return device_impl shared ptr reference associated with graph.
+  const DeviceImplPtr &getDeviceImplPtr() const {
+    return getSyclObjImpl(MDevice);
+  }
+
   /// Query for the device tied to this graph.
   /// @return Device associated with graph.
   sycl::device getDevice() const { return MDevice; }

sycl/source/handler.cpp

@@ -43,6 +43,15 @@ inline namespace _V1 {
 
 namespace detail {
 
+const DeviceImplPtr &getDeviceImplFromHandler(handler &CGH) {
+  assert((CGH.MQueue || getSyclObjImpl(CGH)->MGraph) &&
+         "One of MQueue or MGraph should be nonnull!");
+  if (CGH.MQueue)
+    return CGH.MQueue->getDeviceImplPtr();
+
+  return getSyclObjImpl(CGH)->MGraph->getDeviceImplPtr();
+}
+
 bool isDeviceGlobalUsedInKernel(const void *DeviceGlobalPtr) {
   DeviceGlobalMapEntry *DGEntry =
       detail::ProgramManager::getInstance().getDeviceGlobalEntry(
@@ -2059,10 +2068,10 @@ void handler::setUserFacingNodeType(ext::oneapi::experimental::node_type Type) {
 }
 
 std::optional<std::array<size_t, 3>> handler::getMaxWorkGroups() {
-  auto Dev = detail::getSyclObjImpl(detail::getDeviceFromHandler(*this));
+  const auto &DeviceImpl = detail::getDeviceImplFromHandler(*this);
   std::array<size_t, 3> UrResult = {};
-  auto Ret = Dev->getAdapter()->call_nocheck<UrApiKind::urDeviceGetInfo>(
-      Dev->getHandleRef(),
+  auto Ret = DeviceImpl->getAdapter()->call_nocheck<UrApiKind::urDeviceGetInfo>(
+      DeviceImpl->getHandleRef(),
       UrInfoCode<
           ext::oneapi::experimental::info::device::max_work_groups<3>>::value,
       sizeof(UrResult), &UrResult, nullptr);