[SYCL][Grapg] 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/include/sycl/ext/oneapi/virtual_mem/virtual_mem.hpp

@@ -28,7 +28,8 @@ enum class granularity_mode : char {
 
 __SYCL_EXPORT size_t
 get_mem_granularity(const device &SyclDevice, const context &SyclContext,
-                    granularity_mode Mode = granularity_mode::recommended);
+                    granularity_mode Mode = granularity_mode::recommended,
+                    size_t allocationSize = 1);
 
 __SYCL_EXPORT size_t
 get_mem_granularity(const context &SyclContext,

sycl/source/detail/graph/memory_pool.cpp

@@ -42,7 +42,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(
+        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/virtual_mem.cpp

@@ -24,7 +24,8 @@ namespace ext::oneapi::experimental {
 
 __SYCL_EXPORT size_t get_mem_granularity(const device &SyclDevice,
                                          const context &SyclContext,
-                                         granularity_mode Mode) {
+                                         granularity_mode Mode,
+                                         size_t allocationSize) {
   if (!SyclDevice.has(aspect::ext_oneapi_virtual_mem))
     throw sycl::exception(
         sycl::make_error_code(sycl::errc::feature_not_supported),
@@ -45,13 +46,15 @@ __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, 0, nullptr, &InfoOutputSize);
+      urCtx, urDevice, allocationSize, GranularityQuery, 0, 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),

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."

unified-runtime/include/ur_ddi.h

@@ -18,6 +18,7 @@
 
 UR_APIEXPORT ur_result_t UR_APICALL urVirtualMemGranularityGetInfo(
     ur_context_handle_t, ur_device_handle_t hDevice,
+    [[maybe_unused]] size_t allocationSize,
     ur_virtual_mem_granularity_info_t propName, size_t propSize,
     void *pPropValue, size_t *pPropSizeRet) {
   UrReturnHelper ReturnValue(propSize, pPropValue, pPropSizeRet);