[SYCL][USM] Add support for APIs that query properties of USM pointers (#1016)

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

Author: jbrodman

sycl/include/CL/sycl/usm.hpp

@@ -145,5 +145,19 @@ T *aligned_alloc(size_t Alignment, size_t Count, const queue &Q,
                           Kind);
 }
 
+// Pointer queries
+/// Query the allocation type from a USM pointer
+///
+/// @param ptr is the USM pointer to query
+/// @param ctxt is the sycl context the ptr was allocated in
+usm::alloc get_pointer_type(const void *ptr, const context &ctxt);
+
+/// Queries the device against which the pointer was allocated
+/// Throws an invalid_object_error if ptr is a host allocation.
+///
+/// @param ptr is the USM pointer to query
+/// @param ctxt is the sycl context the ptr was allocated in
+device get_pointer_device(const void *ptr, const context &ctxt);
+
 } // namespace sycl
 } // namespace cl

sycl/source/detail/usm/usm_impl.cpp

@@ -231,5 +231,80 @@ void *aligned_alloc(size_t Alignment, size_t Size, const queue &Q, alloc Kind) {
   return aligned_alloc(Alignment, Size, Q.get_device(), Q.get_context(), Kind);
 }
 
+// Pointer queries
+/// Query the allocation type from a USM pointer
+/// Returns alloc::host for all pointers in a host context.
+///
+/// @param ptr is the USM pointer to query
+/// @param ctxt is the sycl context the ptr was allocated in
+alloc get_pointer_type(const void *Ptr, const context &Ctxt) {
+  // Everything on a host device is just system malloc so call it host
+  if (Ctxt.is_host())
+    return alloc::host;
+
+  std::shared_ptr<detail::context_impl> CtxImpl = detail::getSyclObjImpl(Ctxt);
+  pi_context PICtx = CtxImpl->getHandleRef();
+  pi_usm_type AllocTy;
+
+  // query type using PI function
+  PI_CALL(piextUSMGetMemAllocInfo)(PICtx, Ptr, PI_MEM_ALLOC_TYPE,
+                                   sizeof(pi_usm_type), &AllocTy, nullptr);
+
+  alloc ResultAlloc;
+  switch (AllocTy) {
+  case PI_MEM_TYPE_HOST:
+    ResultAlloc = alloc::host;
+    break;
+  case PI_MEM_TYPE_DEVICE:
+    ResultAlloc = alloc::device;
+    break;
+  case PI_MEM_TYPE_SHARED:
+    ResultAlloc = alloc::shared;
+    break;
+  default:
+    ResultAlloc = alloc::unknown;
+    break;
+  }
+
+  return ResultAlloc;
+}
+
+/// Queries the device against which the pointer was allocated
+///
+/// @param ptr is the USM pointer to query
+/// @param ctxt is the sycl context the ptr was allocated in
+device get_pointer_device(const void *Ptr, const context &Ctxt) {
+  // Just return the host device in the host context
+  if (Ctxt.is_host())
+    return Ctxt.get_devices()[0];
+
+  std::shared_ptr<detail::context_impl> CtxImpl = detail::getSyclObjImpl(Ctxt);
+
+  // Check if ptr is a host allocation
+  if (get_pointer_type(Ptr, Ctxt) == alloc::host) {
+    auto Devs = CtxImpl->getDevices();
+    if (Devs.size() == 0)
+      throw runtime_error("No devices in passed context!");
+
+    // Just return the first device in the context
+    return Devs[0];
+  }
+
+  pi_context PICtx = CtxImpl->getHandleRef();
+  pi_device DeviceId;
+
+  // query device using PI function
+  PI_CALL(piextUSMGetMemAllocInfo)(PICtx, Ptr, PI_MEM_ALLOC_DEVICE,
+                                   sizeof(pi_device), &DeviceId, nullptr);
+
+  for (const device &Dev : CtxImpl->getDevices()) {
+    // Try to find the real sycl device used in the context
+    if (detail::getSyclObjImpl(Dev)->getHandleRef() == DeviceId)
+      return Dev;
+  }
+
+  throw runtime_error("Cannot find device associated with USM allocation!");
+}
+
 } // namespace sycl
 } // namespace cl

sycl/test/usm/pointer_query.cpp

@@ -0,0 +1,96 @@
+// RUN: %clangxx -fsycl %s -o %t1.out
+// RUN: env SYCL_DEVICE_TYPE=HOST %t1.out
+
+//==-------------- pointer_query.cpp - Pointer Query 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;
+
+int main() {
+  int *array = nullptr;
+  const int N = 4;
+  queue q;
+  auto dev = q.get_device();
+  auto ctxt = q.get_context();
+
+  if (!(dev.get_info<info::device::usm_device_allocations>() &&
+        dev.get_info<info::device::usm_shared_allocations>() &&
+        dev.get_info<info::device::usm_host_allocations>()))
+    return 0;
+
+  usm::alloc Kind;
+  device D;
+
+  // Test device allocs
+  array = (int *)malloc_device(N * sizeof(int), q);
+  if (array == nullptr) {
+    return 1;
+  }
+  Kind = get_pointer_type(array, ctxt);
+  if (ctxt.is_host()) {
+    // for now, host device treats all allocations
+    // as host allocations
+    if (Kind != usm::alloc::host) {
+      return 2;
+    }
+  } else {
+    if (Kind != usm::alloc::device) {
+      return 3;
+    }
+  }
+  D = get_pointer_device(array, ctxt);
+  if (D != dev) {
+    return 4;
+  }
+  free(array, ctxt);
+
+  // Test shared allocs
+  array = (int *)malloc_shared(N * sizeof(int), q);
+  if (array == nullptr) {
+    return 5;
+  }
+  Kind = get_pointer_type(array, ctxt);
+  if (ctxt.is_host()) {
+    // for now, host device treats all allocations
+    // as host allocations
+    if (Kind != usm::alloc::host) {
+      return 6;
+    }
+  } else {
+    if (Kind != usm::alloc::shared) {
+      return 7;
+    }
+  }
+  D = get_pointer_device(array, ctxt);
+  if (D != dev) {
+    return 8;
+  }
+  free(array, ctxt);
+
+  // Test host allocs
+  array = (int *)malloc_host(N * sizeof(int), q);
+  if (array == nullptr) {
+    return 9;
+  }
+  Kind = get_pointer_type(array, ctxt);
+  if (Kind != usm::alloc::host) {
+    return 10;
+  }
+  D = get_pointer_device(array, ctxt);
+  auto Devs = ctxt.get_devices();
+  auto result = std::find(Devs.begin(), Devs.end(), D);
+  if (result == Devs.end()) {
+    // Returned device was not in queried context
+    return 11;
+  }
+  free(array, ctxt);
+
+  return 0;
+}