[SYCL] Fixed copying subbuffer between different contexts (#1005)

Fixed an issue when no alloca command for subbuffer was found while performing
copying of memory between different contexts.
Now if an allocation command for a subbuffer is not found for the source context
graph builder uses an allocation command for the whole buffer.

Signed-off-by: Ivan Karachun <[email protected]>

Author: Ivan Karachun

sycl/source/detail/scheduler/graph_builder.cpp

@@ -232,16 +232,31 @@ Command *Scheduler::GraphBuilder::insertMemoryMove(MemObjRecord *Record,
 
   AllocaCommandBase *AllocaCmdSrc =
       findAllocaForReq(Record, Req, Record->MCurContext);
+  if (!AllocaCmdSrc && IsSuitableSubReq(Req)) {
+    // Since no alloca command for the sub buffer requirement was found in the
+    // current context, need to find a parent alloca command for it (it must be
+    // there)
+    auto IsSuitableAlloca = [Record, Req](AllocaCommandBase *AllocaCmd) {
+      bool Res = sameCtx(AllocaCmd->getQueue()->get_context_impl(),
+                         Record->MCurContext) &&
+                 // Looking for a parent buffer alloca command
+                 AllocaCmd->getType() == Command::CommandType::ALLOCA;
+      return Res;
+    };
+    const auto It =
+        std::find_if(Record->MAllocaCommands.begin(),
+                     Record->MAllocaCommands.end(), IsSuitableAlloca);
+    AllocaCmdSrc = (Record->MAllocaCommands.end() != It) ? *It : nullptr;
+  }
   if (!AllocaCmdSrc)
     throw runtime_error("Cannot find buffer allocation");
   // Get parent allocation of sub buffer to perform full copy of whole buffer
   if (IsSuitableSubReq(Req)) {
     if (AllocaCmdSrc->getType() == Command::CommandType::ALLOCA_SUB_BUF)
       AllocaCmdSrc =
           static_cast<AllocaSubBufCommand *>(AllocaCmdSrc)->getParentAlloca();
-    else
-      assert(
-          !"Inappropriate alloca command. AllocaSubBufCommand was expected.");
+    else if (AllocaCmdSrc->getSYCLMemObj() != Req->MSYCLMemObj)
+      assert(!"Inappropriate alloca command.");
   }
 
   Command *NewCmd = nullptr;

sycl/test/basic_tests/buffer/subbuffer.cpp

@@ -247,11 +247,37 @@ void copyBlock() {
   }
 }
 
+void checkMultipleContexts() {
+  constexpr int N = 64;
+  int a[N] = {0};
+  {
+    sycl::queue queue1;
+    sycl::queue queue2;
+    sycl::buffer<int, 1> buf(a, sycl::range<1>(N));
+    sycl::buffer<int, 1> subbuf1(buf, sycl::id<1>(0), sycl::range<1>(N / 2));
+    sycl::buffer<int, 1> subbuf2(buf, sycl::id<1>(N / 2),
+                                 sycl::range<1>(N / 2));
+    queue1.submit([&](sycl::handler &cgh) {
+      auto bufacc = subbuf1.get_access<sycl::access::mode::read_write>(cgh);
+      cgh.parallel_for<class sub_buffer_1>(
+          sycl::range<1>(N / 2), [=](sycl::id<1> idx) { bufacc[idx[0]] = 1; });
+    });
+
+    queue2.submit([&](sycl::handler &cgh) {
+      auto bufacc = subbuf2.get_access<sycl::access::mode::read_write>(cgh);
+      cgh.parallel_for<class sub_buffer_2>(
+          sycl::range<1>(N / 2), [=](sycl::id<1> idx) { bufacc[idx[0]] = 2; });
+    });
+  }
+  assert(a[N / 2 - 1] == 1 && a[N / 2] == 2 && "Sub buffer data loss");
+}
+
 int main() {
   cl::sycl::queue q;
   check1DSubBuffer(q);
   checkHostAccessor(q);
   checkExceptions();
   copyBlock();
+  checkMultipleContexts();
   return 0;
 }