[SYCL][USM] Make host device return nullptr for bad mallocs (huge, 0, etc) (#988)

Host device wasn't properly handling alloc::unknown.
Also, aligned allocator wasn't failing when it really should, so check to throw an error.

Signed-off-by: James Brodman <[email protected]>

Author: jbrodman

sycl/include/CL/sycl/detail/aligned_allocator.hpp

@@ -53,6 +53,8 @@ template <typename T> class aligned_allocator {
   pointer allocate(size_t Size) {
     size_t NumBytes = Size * sizeof(value_type);
     NumBytes = ((NumBytes - 1) | (MAlignment - 1)) + 1;
+    if (Size > NumBytes)
+      throw std::bad_alloc();
 
     pointer Result = reinterpret_cast<pointer>(
         detail::OSUtil::alignedAlloc(MAlignment, NumBytes));

sycl/source/detail/usm/usm_impl.cpp

@@ -71,17 +71,21 @@ void *alignedAlloc(size_t Alignment, size_t Size, const context &Ctxt,
                    const device &Dev, alloc Kind) {
   void *RetVal = nullptr;
   if (Ctxt.is_host()) {
-    if (!Alignment) {
-      // worst case default
-      Alignment = 128;
-    }
-
-    aligned_allocator<char> Alloc(Alignment);
-    try {
-      RetVal = Alloc.allocate(Size);
-    } catch (const std::bad_alloc &) {
-      // Conform with Specification behavior
+    if (Kind == alloc::unknown) {
       RetVal = nullptr;
+    } else {
+      if (!Alignment) {
+        // worst case default
+        Alignment = 128;
+      }
+
+      aligned_allocator<char> Alloc(Alignment);
+      try {
+        RetVal = Alloc.allocate(Size);
+      } catch (const std::bad_alloc &) {
+        // Conform with Specification behavior
+        RetVal = nullptr;
+      }
     }
   } else {
     std::shared_ptr<context_impl> CtxImpl = detail::getSyclObjImpl(Ctxt);

sycl/test/usm/badmalloc.cpp

@@ -0,0 +1,75 @@
+// RUN: %clangxx -fsycl %s -o %t1.out
+// RUN: env SYCL_DEVICE_TYPE=HOST %t1.out
+// RUN: %CPU_RUN_PLACEHOLDER %t1.out
+// RUN: %GPU_RUN_PLACEHOLDER %t1.out
+
+// UNSUPPORTED: windows
+
+//==----------------- badmalloc.cpp - Bad Mallocs test ---------------------==//
+//
+// 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
+//
+//===----------------------------------------------------------------------===//
+
+
+// This test verifies that things fail in the proper way when they should.
+
+#include <CL/sycl.hpp>
+#include <iostream>
+
+using namespace cl::sycl;
+
+int main(int argc, char *argv[]) {
+  queue q;
+
+  // Good size, bad type
+  auto p = malloc(8, q, usm::alloc::unknown);
+  if (p != nullptr)
+    return 1;
+
+  // Bad size, host
+  p = malloc(-1, q, usm::alloc::host);
+  std::cout << "p = " << p << std::endl;
+  if (p != nullptr)
+    return 2;
+  p = malloc(-1, q, usm::alloc::device);
+  std::cout << "p = " << p << std::endl;
+  if (p != nullptr)
+    return 3;
+  p = malloc(-1, q, usm::alloc::shared);
+  std::cout << "p = " << p << std::endl;
+  if (p != nullptr)
+    return 4;
+  p = malloc(-1, q, usm::alloc::unknown);
+  std::cout << "p = " << p << std::endl;
+  if (p != nullptr)
+    return 5;
+
+  // Bad size, auto aligned
+  p = aligned_alloc(0, -1, q,  usm::alloc::host);
+  std::cout << "p = " << p << std::endl;
+  if (p != nullptr)
+    return 6;
+  p = aligned_alloc(0, -1, q,  usm::alloc::device);
+  std::cout << "p = " << p << std::endl;
+  if (p != nullptr)
+    return 7;
+  p = aligned_alloc(0, -1, q,  usm::alloc::shared);
+  std::cout << "p = " << p << std::endl;
+  if (p != nullptr)
+    return 8;
+  p = aligned_alloc(0, -1, q,  usm::alloc::unknown);
+  std::cout << "p = " << p << std::endl;
+  if (p != nullptr)
+    return 9;
+
+  // Allocs of 0 undefined, but bad type
+  p = aligned_alloc(4, 0, q,  usm::alloc::unknown);
+  std::cout << "p = " << p << std::endl;
+  if (p != nullptr)
+    return 10;
+
+  return 0;
+}