From 9e0c05abae2ab17eedcc0952f572f9c041c9c953 Mon Sep 17 00:00:00 2001
From: Ben Tracy <ben.tracy@codeplay.com>
Date: Thu, 28 Nov 2024 13:41:00 +0000
Subject: [PATCH] [SYCL][Graph] Implement ext_oneapi_weak_object for
 command_graph class

- Add owner_less and weak_object support for the command_graph class
- Unit test which creates objects for all graph states and tests using them in a map
---
 .../sycl/ext/oneapi/experimental/graph.hpp    |  7 ++--
 sycl/include/sycl/ext/oneapi/owner_less.hpp   | 33 +++++++++--------
 .../Extensions/CommandGraph/CommandGraph.cpp  | 36 +++++++++++++++++++
 3 files changed, 60 insertions(+), 16 deletions(-)

diff --git a/sycl/include/sycl/ext/oneapi/experimental/graph.hpp b/sycl/include/sycl/ext/oneapi/experimental/graph.hpp
index 9cca0ed2aa532..88385a6f5ba34 100644
--- a/sycl/include/sycl/ext/oneapi/experimental/graph.hpp
+++ b/sycl/include/sycl/ext/oneapi/experimental/graph.hpp
@@ -12,6 +12,7 @@
 #include <sycl/context.hpp>                // for context
 #include <sycl/detail/export.hpp>          // for __SYCL_EXPORT
 #include <sycl/detail/kernel_desc.hpp>     // for kernel_param_kind_t
+#include <sycl/detail/owner_less_base.hpp> // for OwnerLessBase
 #include <sycl/detail/property_helper.hpp> // for DataLessPropKind, PropWith...
 #ifdef __INTEL_PREVIEW_BREAKING_CHANGES
 #include <sycl/detail/string_view.hpp>
@@ -236,7 +237,8 @@ class __SYCL_EXPORT dynamic_command_group {
 
 namespace detail {
 // Templateless modifiable command-graph base class.
-class __SYCL_EXPORT modifiable_command_graph {
+class __SYCL_EXPORT modifiable_command_graph
+    : public sycl::detail::OwnerLessBase<modifiable_command_graph> {
 public:
   /// Constructor.
   /// @param SyclContext Context to use for graph.
@@ -390,7 +392,8 @@ class __SYCL_EXPORT modifiable_command_graph {
 };
 
 // Templateless executable command-graph base class.
-class __SYCL_EXPORT executable_command_graph {
+class __SYCL_EXPORT executable_command_graph
+    : public sycl::detail::OwnerLessBase<executable_command_graph> {
 public:
   /// An executable command-graph is not user constructable.
   executable_command_graph() = delete;
diff --git a/sycl/include/sycl/ext/oneapi/owner_less.hpp b/sycl/include/sycl/ext/oneapi/owner_less.hpp
index a69072652c5b7..265761a42f3c6 100644
--- a/sycl/include/sycl/ext/oneapi/owner_less.hpp
+++ b/sycl/include/sycl/ext/oneapi/owner_less.hpp
@@ -8,20 +8,21 @@
 
 #pragma once
 
-#include <sycl/access/access.hpp>               // for access_mode
-#include <sycl/accessor.hpp>                    // for host_acce...
-#include <sycl/accessor.hpp>                    // for accessor
-#include <sycl/context.hpp>                     // for context
-#include <sycl/device.hpp>                      // for device
-#include <sycl/event.hpp>                       // for event
-#include <sycl/ext/oneapi/weak_object.hpp>      // for weak_object
-#include <sycl/kernel.hpp>                      // for kernel
-#include <sycl/kernel_bundle.hpp>               // for kernel_id
-#include <sycl/kernel_bundle_enums.hpp>         // for bundle_state
-#include <sycl/platform.hpp>                    // for platform
-#include <sycl/properties/image_properties.hpp> // for sampled_i...
-#include <sycl/queue.hpp>                       // for queue
-#include <sycl/stream.hpp>                      // for stream
+#include <sycl/access/access.hpp>                 // for access_mode
+#include <sycl/accessor.hpp>                      // for host_acce...
+#include <sycl/accessor.hpp>                      // for accessor
+#include <sycl/context.hpp>                       // for context
+#include <sycl/device.hpp>                        // for device
+#include <sycl/event.hpp>                         // for event
+#include <sycl/ext/oneapi/experimental/graph.hpp> // for command_graph
+#include <sycl/ext/oneapi/weak_object.hpp>        // for weak_object
+#include <sycl/kernel.hpp>                        // for kernel
+#include <sycl/kernel_bundle.hpp>                 // for kernel_id
+#include <sycl/kernel_bundle_enums.hpp>           // for bundle_state
+#include <sycl/platform.hpp>                      // for platform
+#include <sycl/properties/image_properties.hpp>   // for sampled_i...
+#include <sycl/queue.hpp>                         // for queue
+#include <sycl/stream.hpp>                        // for stream
 
 namespace sycl {
 inline namespace _V1 {
@@ -126,6 +127,10 @@ struct owner_less<host_sampled_image_accessor<DataT, Dimensions>>
     : public detail::owner_less_base<
           host_sampled_image_accessor<DataT, Dimensions>> {};
 
+template <experimental::graph_state State>
+struct owner_less<experimental::command_graph<State>>
+    : public detail::owner_less_base<experimental::command_graph<State>> {};
+
 } // namespace ext::oneapi
 } // namespace _V1
 } // namespace sycl
diff --git a/sycl/unittests/Extensions/CommandGraph/CommandGraph.cpp b/sycl/unittests/Extensions/CommandGraph/CommandGraph.cpp
index 6ac141528a589..df3c3ba81c221 100644
--- a/sycl/unittests/Extensions/CommandGraph/CommandGraph.cpp
+++ b/sycl/unittests/Extensions/CommandGraph/CommandGraph.cpp
@@ -7,9 +7,45 @@
 //===----------------------------------------------------------------------===//
 #include "Common.hpp"
 
+#include <map>
+
 using namespace sycl;
 using namespace sycl::ext::oneapi;
 
+// Test creating and using ext::oneapi::weak_object and owner_less for
+// command_graph class in a map
+TEST_F(CommandGraphTest, OwnerLessGraph) {
+
+  using ModifiableGraphT =
+      experimental::command_graph<experimental::graph_state::modifiable>;
+  using ExecutableGraphT =
+      experimental::command_graph<experimental::graph_state::executable>;
+
+  // Test graph objects using the default template parameter
+  weak_object<experimental::command_graph<>> DefaultWeakGraph = Graph;
+  std::map<weak_object<experimental::command_graph<>>, int,
+           owner_less<experimental::command_graph<>>>
+      DefaultGraphMap;
+
+  ASSERT_NO_THROW(DefaultGraphMap.insert({DefaultWeakGraph, 1}));
+
+  // Test graph objects in the modifiable state
+  weak_object<ModifiableGraphT> WeakGraph = Graph;
+  std::map<weak_object<ModifiableGraphT>, int, owner_less<ModifiableGraphT>>
+      ModifiableGraphMap;
+
+  ASSERT_NO_THROW(ModifiableGraphMap.insert({WeakGraph, 1}));
+
+  // Test graph objects in the executable state
+  auto ExecGraph = Graph.finalize();
+
+  weak_object<ExecutableGraphT> WeakGraphExec = ExecGraph;
+  std::map<weak_object<ExecutableGraphT>, int, owner_less<ExecutableGraphT>>
+      ExecGraphMap;
+
+  ASSERT_NO_THROW(ExecGraphMap.insert({WeakGraphExec, 1}));
+}
+
 TEST_F(CommandGraphTest, AddNode) {
   auto GraphImpl = sycl::detail::getSyclObjImpl(Graph);