[SYCL][Graph] async_malloc use allocation size for zeVirtualMemQueryPageSize

Author: lslusarczyk

sycl/cmake/modules/FetchUnifiedRuntime.cmake

@@ -28,6 +28,9 @@ set(UR_BUILD_TESTS "${SYCL_UR_BUILD_TESTS}" CACHE BOOL "" FORCE)
 # UR tests require the examples to be built
 set(UR_BUILD_EXAMPLES "${SYCL_UR_BUILD_TESTS}" CACHE BOOL "" FORCE)
 
+option(SYCL_UR_FORMAT_CPP_STYLE "Format code style of UR C++ sources" OFF)
+set(UR_FORMAT_CPP_STYLE "${SYCL_UR_FORMAT_CPP_STYLE}" CACHE BOOL "" FORCE)
+
 # Here we override the defaults to unified-runtime
 set(UR_BUILD_XPTI_LIBS OFF CACHE BOOL "")
 set(UR_ENABLE_SYMBOLIZER ON CACHE BOOL "Enable symbolizer for sanitizer layer.")

sycl/source/detail/graph/memory_pool.cpp

@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "memory_pool.hpp"
+#include "detail/virtual_mem.hpp"
 #include "graph_impl.hpp"
 
 #include <optional>
@@ -42,7 +43,8 @@ void *graph_mem_pool::malloc(size_t Size, usm::alloc AllocType,
     context_impl &CtxImpl = *getSyclObjImpl(MContext);
     adapter_impl &Adapter = CtxImpl.getAdapter();
 
-    size_t Granularity = get_mem_granularity(MDevice, MContext);
+    const size_t Granularity = get_mem_granularity_for_allocation_size(
+        MDevice, MContext, granularity_mode::recommended, Size);
     uintptr_t StartPtr = 0;
     size_t AlignedSize = alignByteSize(Size, Granularity);
     // See if we can find an allocation to reuse

sycl/source/detail/graph/memory_pool.hpp

@@ -12,7 +12,7 @@
 #include <detail/physical_mem_impl.hpp>                // For physical_mem_impl
 #include <sycl/context.hpp>                            // For context
 #include <sycl/device.hpp>                             // For device
-#include <sycl/ext/oneapi/virtual_mem/virtual_mem.hpp> // For get_mem_granularity
+#include <sycl/ext/oneapi/virtual_mem/virtual_mem.hpp> // For unmap
 
 namespace sycl {
 inline namespace _V1 {

sycl/source/detail/virtual_mem.hpp

@@ -0,0 +1,24 @@
+//==---------------- virtual_mem.hpp ---------------------------------------==//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+#pragma once
+
+#include <sycl/ext/oneapi/virtual_mem/virtual_mem.hpp>
+
+namespace sycl {
+inline namespace _V1 {
+namespace ext::oneapi::experimental {
+
+size_t get_mem_granularity_for_allocation_size(const device &SyclDevice,
+                                               const context &SyclContext,
+                                               granularity_mode Mode,
+                                               size_t AllocationSize);
+
+} // namespace ext::oneapi::experimental
+} // namespace _V1
+} // namespace sycl

sycl/source/virtual_mem.cpp

@@ -9,7 +9,7 @@
 #include <detail/context_impl.hpp>
 #include <detail/device_impl.hpp>
 #include <detail/physical_mem_impl.hpp>
-#include <sycl/ext/oneapi/virtual_mem/virtual_mem.hpp>
+#include <detail/virtual_mem.hpp>
 
 // System headers for querying page-size.
 #ifdef _WIN32
@@ -22,9 +22,10 @@ namespace sycl {
 inline namespace _V1 {
 namespace ext::oneapi::experimental {
 
-__SYCL_EXPORT size_t get_mem_granularity(const device &SyclDevice,
-                                         const context &SyclContext,
-                                         granularity_mode Mode) {
+size_t get_mem_granularity_for_allocation_size(const device &SyclDevice,
+                                               const context &SyclContext,
+                                               granularity_mode Mode,
+                                               const size_t AllocationSize) {
   if (!SyclDevice.has(aspect::ext_oneapi_virtual_mem))
     throw sycl::exception(
         sycl::make_error_code(sycl::errc::feature_not_supported),
@@ -45,20 +46,29 @@ __SYCL_EXPORT size_t get_mem_granularity(const device &SyclDevice,
 #ifndef NDEBUG
   size_t InfoOutputSize = 0;
   Adapter->call<sycl::detail::UrApiKind::urVirtualMemGranularityGetInfo>(
-      urCtx, urDevice, GranularityQuery, 0u, nullptr, &InfoOutputSize);
+      urCtx, urDevice, AllocationSize, GranularityQuery, 0u, nullptr,
+      &InfoOutputSize);
   assert(InfoOutputSize == sizeof(size_t) &&
          "Unexpected output size of granularity info query.");
 #endif // NDEBUG
   size_t Granularity = 0;
   Adapter->call<sycl::detail::UrApiKind::urVirtualMemGranularityGetInfo>(
-      urCtx, urDevice, GranularityQuery, sizeof(size_t), &Granularity, nullptr);
+      urCtx, urDevice, AllocationSize, GranularityQuery, sizeof(size_t),
+      &Granularity, nullptr);
   if (Granularity == 0)
     throw sycl::exception(
         sycl::make_error_code(sycl::errc::invalid),
         "Unexpected granularity result: memory granularity shouldn't be 0.");
   return Granularity;
 }
 
+__SYCL_EXPORT size_t get_mem_granularity(const device &SyclDevice,
+                                         const context &SyclContext,
+                                         granularity_mode Mode) {
+  return get_mem_granularity_for_allocation_size(SyclDevice, SyclContext, Mode,
+                                                 1);
+}
+
 __SYCL_EXPORT size_t get_mem_granularity(const context &SyclContext,
                                          granularity_mode Mode) {
   const std::vector<device> Devices = SyclContext.get_devices();

sycl/test-e2e/Graph/AsyncAlloc/Inputs/async_alloc_different_sizes.cpp

@@ -30,5 +30,5 @@ int main() {
   asyncAllocWorksWithSize(131);
   asyncAllocWorksWithSize(10071);
   asyncAllocWorksWithSize(1007177);
-  // asyncAllocWorksWithSize(191439360); // BUG
+  asyncAllocWorksWithSize(191439360);
 }

unified-runtime/include/ur_api.h

@@ -41,6 +41,9 @@ params:
           [in][optional] is the device to get the granularity from, if the
           device is null then the granularity is suitable for all devices in
           context.
+    - type: size_t
+      name: allocationSize
+      desc: "[in] size in bytes of allocation size which granurality we search for."
     - type: $x_virtual_mem_granularity_info_t
       name: propName
       desc: "[in] type of the info to query."