From cd50dbda9b72c387adaaeadd74010b87e2d9cf25 Mon Sep 17 00:00:00 2001
From: Sherry Yuan <sherry.yuan@intel.com>
Date: Thu, 16 Dec 2021 10:34:24 -0500
Subject: [PATCH] Add USM buffer location property for clDeviceMemAlloc

---
 src/acl_usm.cpp       | 16 ++++++++++++-
 test/acl_usm_test.cpp | 56 +++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 71 insertions(+), 1 deletion(-)

diff --git a/src/acl_usm.cpp b/src/acl_usm.cpp
index 3621b8fd..94b8c6b4 100644
--- a/src/acl_usm.cpp
+++ b/src/acl_usm.cpp
@@ -239,11 +239,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;
+  cl_uint mem_id = acl_get_default_memory(device->def);
   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 = (cl_uint) * (properties + 1);
+    } break;
     default: {
       UNLOCK_BAIL_INFO(CL_INVALID_DEVICE, context, "Invalid properties");
     }
@@ -254,8 +258,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");
   }
@@ -605,6 +611,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(0);
+    }
+  } break;
+
   case CL_MEM_ALLOC_BASE_PTR_INTEL: {
     void *base_ptr = NULL;
     if (usm_alloc) {
diff --git a/test/acl_usm_test.cpp b/test/acl_usm_test.cpp
index 1e4578ae..12c0720f 100644
--- a/test/acl_usm_test.cpp
+++ b/test/acl_usm_test.cpp
@@ -388,6 +388,62 @@ 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);
+
+  // Check alloc information
+  cl_uint read_mem_id = 4;
+  size_t ret_size = 9;
+  status = clGetMemAllocInfoINTEL(m_context, test_ptr,
+                                  CL_MEM_ALLOC_BUFFER_LOCATION_INTEL,
+                                  sizeof(cl_uint), &read_mem_id, &ret_size);
+
+  CHECK_EQUAL(CL_SUCCESS, status);
+  CHECK_EQUAL(0, read_mem_id);
+  CHECK_EQUAL(sizeof(cl_uint), ret_size);
+
+  status = clMemFreeINTEL(m_context, test_ptr);
+  ACL_LOCKED(CHECK_EQUAL(acl_usm_ptr_belongs_to_context(m_context, test_ptr),
+                         CL_FALSE));
+  CHECK(m_context->usm_allocation.empty());
+
+  // Check when given pointer is already freed
+  status = clGetMemAllocInfoINTEL(m_context, test_ptr,
+                                  CL_MEM_ALLOC_BUFFER_LOCATION_INTEL,
+                                  sizeof(cl_uint), &read_mem_id, &ret_size);
+
+  CHECK_EQUAL(CL_SUCCESS, status);
+  CHECK_EQUAL(0, read_mem_id);
+  CHECK_EQUAL(sizeof(cl_uint), ret_size);
+  CHECK_EQUAL(status, CL_SUCCESS);
+
+  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;