Add USM buffer location property for clDeviceMemAlloc

Author: sherry-yuan

src/acl_usm.cpp

@@ -240,11 +240,15 @@ clDeviceMemAllocINTEL(cl_context context, cl_device_id device,
   // Iterate over properties.
   // The end of the properties list is specified with a zero.
   cl_mem_alloc_flags_intel alloc_flags = 0;
+  unsigned int mem_id = 0;
   while (properties != NULL && *properties != 0) {
     switch (*properties) {
     case CL_MEM_ALLOC_FLAGS_INTEL: {
       alloc_flags = *(properties + 1);
     } break;
+    case CL_MEM_ALLOC_BUFFER_LOCATION_INTEL: {
+      mem_id = (unsigned int)*(properties + 1);
+    } break;
     default: {
       UNLOCK_BAIL_INFO(CL_INVALID_DEVICE, context, "Invalid properties");
     }
@@ -255,8 +259,10 @@ clDeviceMemAllocINTEL(cl_context context, cl_device_id device,
   cl_int status;
 
   // Use cl_mem for convenience
+  cl_mem_properties_intel props[] = {CL_MEM_ALLOC_BUFFER_LOCATION_INTEL, mem_id,
+                                     0};
   cl_mem usm_device_buffer = clCreateBufferWithPropertiesINTEL(
-      context, NULL, CL_MEM_READ_WRITE, size, NULL, &status);
+      context, props, CL_MEM_READ_WRITE, size, NULL, &status);
   if (status != CL_SUCCESS) {
     UNLOCK_BAIL_INFO(status, context, "Failed to allocate device memory");
   }
@@ -606,6 +612,14 @@ CL_API_ENTRY cl_int CL_API_CALL clGetMemAllocInfoINTEL(
     }
   } break;
 
+  case CL_MEM_ALLOC_BUFFER_LOCATION_INTEL: {
+    if (usm_alloc) {
+      RESULT_UINT(usm_alloc->mem->mem_id);
+    } else {
+      RESULT_UINT(CL_MEM_TYPE_UNKNOWN_INTEL);
+    }
+  } break;
+
   case CL_MEM_ALLOC_BASE_PTR_INTEL: {
     void *base_ptr = NULL;
     if (usm_alloc) {

test/acl_usm_test.cpp

@@ -388,6 +388,42 @@ MT_TEST(acl_usm, alloc_and_free_device_usm) {
   ACL_LOCKED(acl_print_debug_msg("end alloc_and_free_device_usm\n"));
 }
 
+MT_TEST(acl_usm, buffer_location_usm) {
+  ACL_LOCKED(acl_print_debug_msg("begin buffer_location_usm\n"));
+  const int alignment = ACL_MEM_ALIGN;
+  cl_int status;
+  this->yeah = true;
+
+  acl_usm_allocation_t *test_device_alloc;
+
+  cl_mem_properties_intel good_property[3] = {
+      CL_MEM_ALLOC_BUFFER_LOCATION_INTEL, 0, 0};
+
+  ACL_LOCKED(CHECK(acl_context_is_valid(m_context)));
+
+  // Correct USM buffer allocation
+  void *test_ptr = clDeviceMemAllocINTEL(
+      m_context, m_device[0], &(good_property[0]), 8, alignment, &status);
+  CHECK_EQUAL(status, CL_SUCCESS);
+  CHECK(ACL_DEVICE_ALLOCATION(test_ptr));
+  CHECK(test_ptr != NULL);
+  CHECK(!m_context->usm_allocation.empty());
+  ACL_LOCKED(CHECK_EQUAL(acl_usm_ptr_belongs_to_context(m_context, test_ptr),
+                         CL_TRUE));
+  ACL_LOCKED(test_device_alloc =
+                 acl_get_usm_alloc_from_ptr(m_context, test_ptr));
+  assert(test_device_alloc);
+  CHECK_EQUAL(test_device_alloc->range.begin, test_ptr);
+
+  status = clMemFreeINTEL(m_context, test_ptr);
+  CHECK_EQUAL(status, CL_SUCCESS);
+  ACL_LOCKED(CHECK_EQUAL(acl_usm_ptr_belongs_to_context(m_context, test_ptr),
+                         CL_FALSE));
+  CHECK(m_context->usm_allocation.empty());
+
+  ACL_LOCKED(acl_print_debug_msg("end buffer_location_usm\n"));
+}
+
 MT_TEST(acl_usm, alloc_and_free_shared_usm) {
   ACL_LOCKED(acl_print_debug_msg("begin alloc_and_free_shared_usm\n"));
   const int alignment = 16;