[SYCL][USM] Fix bug with malloc(..., kind) impl and host allocations (#691)

Also rename host-only alignedAlloc to be less ambiguous

Signed-off-by: James Brodman <[email protected]>

Author: jbrodman

sycl/include/CL/sycl/detail/usm_impl.hpp

@@ -19,8 +19,8 @@ namespace usm {
 void *alignedAlloc(size_t Alignment, size_t Bytes, const context &Ctxt,
                    const device &Dev, cl::sycl::usm::alloc Kind);
 
-void *alignedAlloc(size_t Alignment, size_t Bytes, const context &Ctxt,
-                   cl::sycl::usm::alloc Kind);
+void *alignedAllocHost(size_t Alignment, size_t Bytes, const context &Ctxt,
+                       cl::sycl::usm::alloc Kind);
 
 void free(void *Ptr, const context &Ctxt);
 

sycl/include/CL/sycl/usm/usm_allocator.hpp

@@ -109,7 +109,7 @@ class usm_allocator {
       usm::alloc AllocT = AllocKind,
       typename std::enable_if<AllocT == usm::alloc::host, int>::type = 0>
   pointer allocate(size_t Size) {
-    auto Result = reinterpret_cast<pointer>(detail::usm::alignedAlloc(
+    auto Result = reinterpret_cast<pointer>(detail::usm::alignedAllocHost(
         getAlignment(), Size * sizeof(value_type), mContext, AllocKind));
     if (!Result) {
       throw memory_allocation_error();

sycl/source/detail/usm/usm_impl.cpp

@@ -23,8 +23,8 @@ using alloc = cl::sycl::usm::alloc;
 namespace detail {
 namespace usm {
 
-void *alignedAlloc(size_t Alignment, size_t Size, const context &Ctxt,
-                   alloc Kind) {
+void *alignedAllocHost(size_t Alignment, size_t Size, const context &Ctxt,
+                       alloc Kind) {
   void *RetVal = nullptr;
   if (Ctxt.is_host()) {
     if (!Alignment) {
@@ -118,7 +118,7 @@ void *alignedAlloc(size_t Alignment, size_t Size, const context &Ctxt,
   }
   return RetVal;
 }
-  
+
 void free(void *Ptr, const context &Ctxt) {
   if (Ctxt.is_host()) {
     // need to use alignedFree here for Windows
@@ -153,15 +153,15 @@ void free(void *ptr, const context &Ctxt) {
 // Restricted USM
 ///
 void *malloc_host(size_t Size, const context &Ctxt) {
-  return detail::usm::alignedAlloc(0, Size, Ctxt, alloc::host);
+  return detail::usm::alignedAllocHost(0, Size, Ctxt, alloc::host);
 }
 
 void *malloc_shared(size_t Size, const device &Dev, const context &Ctxt) {
   return detail::usm::alignedAlloc(0, Size, Ctxt, Dev, alloc::shared);
 }
 
 void *aligned_alloc_host(size_t Alignment, size_t Size, const context &Ctxt) {
-  return detail::usm::alignedAlloc(Alignment, Size, Ctxt, alloc::host);
+  return detail::usm::alignedAllocHost(Alignment, Size, Ctxt, alloc::host);
 }
 
 void *aligned_alloc_shared(size_t Alignment, size_t Size, const device &Dev,
@@ -172,12 +172,28 @@ void *aligned_alloc_shared(size_t Alignment, size_t Size, const device &Dev,
 // single form
 
 void *malloc(size_t Size, const device &Dev, const context &Ctxt, alloc Kind) {
-  return detail::usm::alignedAlloc(0, Size, Ctxt, Dev, Kind);
+  void *RetVal = nullptr;
+
+  if (Kind == alloc::host) {
+    RetVal = detail::usm::alignedAllocHost(0, Size, Ctxt, Kind);
+  } else {
+    RetVal = detail::usm::alignedAlloc(0, Size, Ctxt, Dev, Kind);
+  }
+
+  return RetVal;
 }
 
 void *aligned_alloc(size_t Alignment, size_t Size, const device &Dev,
                     const context &Ctxt, alloc Kind) {
-  return detail::usm::alignedAlloc(Alignment, Size, Ctxt, Dev, Kind);
+  void *RetVal = nullptr;
+
+  if (Kind == alloc::host) {
+    RetVal = detail::usm::alignedAllocHost(Alignment, Size, Ctxt, Kind);
+  } else {
+    RetVal = detail::usm::alignedAlloc(Alignment, Size, Ctxt, Dev, Kind);
+  }
+
+  return RetVal;
 }
 
 } // namespace sycl

sycl/test/usm/mixed2.cpp

@@ -0,0 +1,72 @@
+// RUN: %clangxx -fsycl %s -o %t1.out
+// RUN: env SYCL_DEVICE_TYPE=HOST %t1.out
+// RUN: %CPU_RUN_PLACEHOLDER %t1.out
+// RUN: %GPU_RUN_PLACEHOLDER %t1.out
+
+//==------------------- mixed2.cpp - Mixed Memory test ---------------------==//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include <CL/sycl.hpp>
+
+using namespace cl::sycl;
+
+class foo;
+int main() {
+  int *darray = nullptr;
+  int *sarray = nullptr;
+  int *harray = nullptr;
+  const int N = 4;
+  const int MAGIC_NUM = 42;
+
+  queue q;
+  auto dev = q.get_device();
+  auto ctxt = q.get_context();
+
+  darray = (int *)malloc(N * sizeof(int), dev, ctxt, usm::alloc::device);
+  if (darray == nullptr) {
+    return -1;
+  }
+  sarray = (int *)malloc(N * sizeof(int), dev, ctxt, usm::alloc::shared);
+
+  if (sarray == nullptr) {
+    return -1;
+  }
+
+  harray = (int *)malloc(N * sizeof(int), dev, ctxt, usm::alloc::host);
+  if (harray == nullptr) {
+    return -1;
+  }
+  for (int i = 0; i < N; i++) {
+    sarray[i] = MAGIC_NUM - 1;
+    harray[i] = 1;
+  }
+
+  auto e0 = q.memset(darray, 0, N * sizeof(int));
+  e0.wait();
+
+  auto e1 = q.submit([=](handler &cgh) {
+    cgh.single_task<class foo>([=]() {
+      for (int i = 0; i < N; i++) {
+        sarray[i] += darray[i] + harray[i];
+      }
+    });
+  });
+
+  e1.wait();
+
+  for (int i = 0; i < N; i++) {
+    if (sarray[i] != MAGIC_NUM) {
+      return -1;
+    }
+  }
+  free(darray, ctxt);
+  free(sarray, ctxt);
+  free(harray, ctxt);
+
+  return 0;
+}