From bd88f59069d7bc7d52f78f6b5fd4d3487094ac79 Mon Sep 17 00:00:00 2001
From: Andrei Elovikov <andrei.elovikov@intel.com>
Date: Tue, 24 Jun 2025 14:45:04 -0700
Subject: [PATCH] [NFC][SYCL] Return raw `context_impl *` from
 `Command::getWorkerContext`

Continuation of the refactoring in
https://github.com/intel/llvm/pull/18795
https://github.com/intel/llvm/pull/18877
https://github.com/intel/llvm/pull/18966
https://github.com/intel/llvm/pull/18979
https://github.com/intel/llvm/pull/18980
https://github.com/intel/llvm/pull/18981
https://github.com/intel/llvm/pull/19007
https://github.com/intel/llvm/pull/19030
---
 sycl/source/detail/scheduler/commands.cpp      | 16 ++++++++--------
 sycl/source/detail/scheduler/commands.hpp      |  6 +++---
 sycl/source/detail/scheduler/graph_builder.cpp |  2 +-
 3 files changed, 12 insertions(+), 12 deletions(-)

diff --git a/sycl/source/detail/scheduler/commands.cpp b/sycl/source/detail/scheduler/commands.cpp
index 40246b26d4b9f..835fdc535f6b7 100644
--- a/sycl/source/detail/scheduler/commands.cpp
+++ b/sycl/source/detail/scheduler/commands.cpp
@@ -759,7 +759,6 @@ void Command::makeTraceEventEpilog() {
 
 Command *Command::processDepEvent(EventImplPtr DepEvent, const DepDesc &Dep,
                                   std::vector<Command *> &ToCleanUp) {
-  const ContextImplPtr &WorkerContext = getWorkerContext();
 
   // 1. Non-host events can be ignored if they are not fully initialized.
   // 2. Some types of commands do not produce UR events after they are
@@ -780,8 +779,9 @@ Command *Command::processDepEvent(EventImplPtr DepEvent, const DepDesc &Dep,
   Command *ConnectionCmd = nullptr;
 
   context_impl &DepEventContext = DepEvent->getContextImpl();
+  context_impl *WorkerContext = getWorkerContext();
   // If contexts don't match we'll connect them using host task
-  if (&DepEventContext != WorkerContext.get() && WorkerContext) {
+  if (&DepEventContext != WorkerContext && WorkerContext) {
     Scheduler::GraphBuilder &GB = Scheduler::getInstance().MGraphBuilder;
     ConnectionCmd = GB.connectDepEvent(this, DepEvent, Dep, ToCleanUp);
   } else
@@ -790,10 +790,10 @@ Command *Command::processDepEvent(EventImplPtr DepEvent, const DepDesc &Dep,
   return ConnectionCmd;
 }
 
-ContextImplPtr Command::getWorkerContext() const {
+context_impl *Command::getWorkerContext() const {
   if (!MQueue)
     return nullptr;
-  return MQueue->getContextImplPtr();
+  return &MQueue->getContextImpl();
 }
 
 bool Command::producesPiEvent() const { return true; }
@@ -1547,10 +1547,10 @@ void MemCpyCommand::emitInstrumentationData() {
 #endif
 }
 
-ContextImplPtr MemCpyCommand::getWorkerContext() const {
+context_impl *MemCpyCommand::getWorkerContext() const {
   if (!MWorkerQueue)
     return nullptr;
-  return MWorkerQueue->getContextImplPtr();
+  return &MWorkerQueue->getContextImpl();
 }
 
 bool MemCpyCommand::producesPiEvent() const {
@@ -1720,10 +1720,10 @@ void MemCpyCommandHost::emitInstrumentationData() {
 #endif
 }
 
-ContextImplPtr MemCpyCommandHost::getWorkerContext() const {
+context_impl *MemCpyCommandHost::getWorkerContext() const {
   if (!MWorkerQueue)
     return nullptr;
-  return MWorkerQueue->getContextImplPtr();
+  return &MWorkerQueue->getContextImpl();
 }
 
 ur_result_t MemCpyCommandHost::enqueueImp() {
diff --git a/sycl/source/detail/scheduler/commands.hpp b/sycl/source/detail/scheduler/commands.hpp
index 8141f7635eda6..a54fc7e231284 100644
--- a/sycl/source/detail/scheduler/commands.hpp
+++ b/sycl/source/detail/scheduler/commands.hpp
@@ -221,7 +221,7 @@ class Command {
 
   /// Get the context of the queue this command will be submitted to. Could
   /// differ from the context of MQueue for memory copy commands.
-  virtual ContextImplPtr getWorkerContext() const;
+  virtual context_impl *getWorkerContext() const;
 
   /// Returns true iff the command produces a UR event on non-host devices.
   virtual bool producesPiEvent() const;
@@ -584,7 +584,7 @@ class MemCpyCommand : public Command {
   void printDot(std::ostream &Stream) const final;
   const Requirement *getRequirement() const final { return &MDstReq; }
   void emitInstrumentationData() final;
-  ContextImplPtr getWorkerContext() const final;
+  context_impl *getWorkerContext() const final;
   bool producesPiEvent() const final;
 
 private:
@@ -608,7 +608,7 @@ class MemCpyCommandHost : public Command {
   void printDot(std::ostream &Stream) const final;
   const Requirement *getRequirement() const final { return &MDstReq; }
   void emitInstrumentationData() final;
-  ContextImplPtr getWorkerContext() const final;
+  context_impl *getWorkerContext() const final;
 
 private:
   ur_result_t enqueueImp() final;
diff --git a/sycl/source/detail/scheduler/graph_builder.cpp b/sycl/source/detail/scheduler/graph_builder.cpp
index 8f7bfb29fb109..77f28f95169ff 100644
--- a/sycl/source/detail/scheduler/graph_builder.cpp
+++ b/sycl/source/detail/scheduler/graph_builder.cpp
@@ -1221,7 +1221,7 @@ void Scheduler::GraphBuilder::removeRecordForMemObj(SYCLMemObjI *MemObject) {
 Command *Scheduler::GraphBuilder::connectDepEvent(
     Command *const Cmd, const EventImplPtr &DepEvent, const DepDesc &Dep,
     std::vector<Command *> &ToCleanUp) {
-  assert(Cmd->getWorkerContext().get() != &DepEvent->getContextImpl());
+  assert(Cmd->getWorkerContext() != &DepEvent->getContextImpl());
 
   // construct Host Task type command manually and make it depend on DepEvent
   ExecCGCommand *ConnectCmd = nullptr;