[UR][L0] Diagnose errors in host memory registration API

unified-runtime/source/adapters/level_zero/common.hpp

@@ -350,3 +350,14 @@ extern thread_local int32_t ErrorAdapterNativeCode;
 // Utility function for setting a message and warning
 [[maybe_unused]] void setErrorMessage(const char *pMessage, int32_t ErrorCode,
                                       int32_t AdapterErrorCode);
+
+// Returns the host memory page size in bytes.
+inline size_t getHostPageSize() {
+#ifdef _WIN32
+  SYSTEM_INFO SystemInfo;
+  GetSystemInfo(&SystemInfo);
+  return static_cast<size_t>(SystemInfo.dwPageSize);
+#else
+  return static_cast<size_t>(sysconf(_SC_PAGESIZE));
+#endif
+}

unified-runtime/source/adapters/level_zero/device.cpp

@@ -1277,6 +1277,16 @@ ur_result_t urDeviceGetInfo(
     // L0 supports importing external memory.
     return ReturnValue(true);
   }
+  case UR_DEVICE_INFO_USM_HOST_ALLOC_REGISTER_SUPPORT_EXP: {
+#if defined(UR_ADAPTER_LEVEL_ZERO_V2) && defined(__linux__)
+    // Registering existing host memory as a USM host allocation relies on the
+    // external system memory mapping extension being supported by the driver.
+    return ReturnValue(
+        Device->Platform->ZeExternalMemoryMappingExtensionSupported);
+#else
+    return ReturnValue(false);
+#endif
+  }
   case UR_DEVICE_INFO_EXTERNAL_SEMAPHORE_IMPORT_SUPPORT_EXP: {
     return ReturnValue(Device->Platform->ZeExternalSemaphoreExt.Supported);
   }

unified-runtime/source/adapters/level_zero/v2/usm.cpp

@@ -907,6 +907,23 @@ ur_result_t urUSMHostAllocRegisterExp(
   if (!hContext->getPlatform()->ZeExternalMemoryMappingExtensionSupported) {
     return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
   }
+  if (pHostMem == nullptr || size == 0) {
+    return UR_RESULT_ERROR_INVALID_VALUE;
+  }
+
+  // Check that the half-open address range [pHostMem, pHostMem + size)
+  // does not wrap around the host address space.
+  uintptr_t Begin = reinterpret_cast<uintptr_t>(pHostMem);
+  if (size > std::numeric_limits<uintptr_t>::max() - Begin) {
+    return UR_RESULT_ERROR_INVALID_VALUE;
+  }
+
+  // The pointer and size must be aligned to the host page size.
+  const size_t PageSize = getHostPageSize();
+  if (reinterpret_cast<uintptr_t>(pHostMem) % PageSize != 0 ||
+      size % PageSize != 0) {
+    return UR_RESULT_ERROR_INVALID_VALUE;
+  }
 
   ze_external_memmap_sysmem_ext_desc_t sysMemDesc = {
       ZE_STRUCTURE_TYPE_EXTERNAL_MEMMAP_SYSMEM_EXT_DESC, nullptr, pHostMem,
@@ -918,7 +935,16 @@ ur_result_t urUSMHostAllocRegisterExp(
   void *mappedMem = nullptr;
   ZE2UR_CALL(zeMemAllocHost,
              (hContext->getZeHandle(), &hostDesc, size, 1, &mappedMem));
-  assert(mappedMem == pHostMem);
+
+  // This extension maps the existing host memory in place, so the mapped
+  // device virtual address must match the host pointer that was registered.
+  // If the driver ever returns a different address we cannot honor the
+  // contract, so undo the mapping and report a failure rather than silently
+  // handing back an unusable registration.
+  if (mappedMem != pHostMem) {
+    ZE_CALL_NOCHECK(zeMemFree, (hContext->getZeHandle(), mappedMem));
+    return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+  }
 
   return UR_RESULT_SUCCESS;
 }
@@ -928,6 +954,9 @@ ur_result_t urUSMHostAllocUnregisterExp(ur_context_handle_t hContext,
   if (!hContext->getPlatform()->ZeExternalMemoryMappingExtensionSupported) {
     return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
   }
+  if (pHostMem == nullptr) {
+    return UR_RESULT_ERROR_INVALID_VALUE;
+  }
 
   ZE2UR_CALL(zeMemFree, (hContext->getZeHandle(), pHostMem));