[SYCL] Add logics for aligned_alloc_xxx<T> to deal with unsupported Alignment argument

sycl/include/sycl/usm.hpp

@@ -179,18 +179,26 @@ T *aligned_alloc_device(
     size_t Alignment, size_t Count, const device &Dev, const context &Ctxt,
     const property_list &PropList = {},
     const detail::code_location &CodeLoc = detail::code_location::current()) {
-  return static_cast<T *>(aligned_alloc_device(max(Alignment, alignof(T)),
-                                               Count * sizeof(T), Dev, Ctxt,
-                                               PropList, CodeLoc));
+  if (!(Alignment & (Alignment - 1))) {
+    return static_cast<T *>(aligned_alloc_device(max(Alignment, alignof(T)),
+                                                 Count * sizeof(T), Dev, Ctxt,
+                                                 PropList, CodeLoc));
+  } else {
+    return nullptr;
+  }
 }
 
 template <typename T>
 T *aligned_alloc_device(
     size_t Alignment, size_t Count, const queue &Q,
     const property_list &PropList = {},
     const detail::code_location &CodeLoc = detail::code_location::current()) {
-  return aligned_alloc_device<T>(Alignment, Count, Q.get_device(),
-                                 Q.get_context(), PropList, CodeLoc);
+  if (!(Alignment & (Alignment - 1))) {
+    return aligned_alloc_device<T>(Alignment, Count, Q.get_device(),
+                                   Q.get_context(), PropList, CodeLoc);
+  } else {
+    return nullptr;
+  }
 }
 
 template <typename T>
@@ -230,37 +238,53 @@ T *aligned_alloc_host(
     size_t Alignment, size_t Count, const context &Ctxt,
     const property_list &PropList = {},
     const detail::code_location &CodeLoc = detail::code_location::current()) {
-  return static_cast<T *>(aligned_alloc_host(std ::max(Alignment, alignof(T)),
-                                             Count * sizeof(T), Ctxt, PropList,
-                                             CodeLoc));
+  if (!(Alignment & (Alignment - 1))) {
+    return static_cast<T *>(aligned_alloc_host(std ::max(Alignment, alignof(T)),
+                                               Count * sizeof(T), Ctxt,
+                                               PropList, CodeLoc));
+  } else {
+    return nullptr;
+  }
 }
 
 template <typename T>
 T *aligned_alloc_host(
     size_t Alignment, size_t Count, const queue &Q,
     const property_list &PropList = {},
     const detail::code_location &CodeLoc = detail::code_location::current()) {
-  return aligned_alloc_host<T>(Alignment, Count, Q.get_context(), PropList,
-                               CodeLoc);
+  if (!(Alignment & (Alignment - 1))) {
+    return aligned_alloc_host<T>(Alignment, Count, Q.get_context(), PropList,
+                                 CodeLoc);
+  } else {
+    return nullptr;
+  }
 }
 
 template <typename T>
 T *aligned_alloc_shared(
     size_t Alignment, size_t Count, const device &Dev, const context &Ctxt,
     const property_list &PropList = {},
     const detail::code_location &CodeLoc = detail::code_location::current()) {
-  return static_cast<T *>(aligned_alloc_shared(max(Alignment, alignof(T)),
-                                               Count * sizeof(T), Dev, Ctxt,
-                                               PropList, CodeLoc));
+  if (!(Alignment & (Alignment - 1))) {
+    return static_cast<T *>(aligned_alloc_shared(max(Alignment, alignof(T)),
+                                                 Count * sizeof(T), Dev, Ctxt,
+                                                 PropList, CodeLoc));
+  } else {
+    return nullptr;
+  }
 }
 
 template <typename T>
 T *aligned_alloc_shared(
     size_t Alignment, size_t Count, const queue &Q,
     const property_list &PropList = {},
     const detail::code_location &CodeLoc = detail::code_location::current()) {
-  return aligned_alloc_shared<T>(Alignment, Count, Q.get_device(),
-                                 Q.get_context(), PropList, CodeLoc);
+  if (!(Alignment & (Alignment - 1))) {
+    return aligned_alloc_shared<T>(Alignment, Count, Q.get_device(),
+                                   Q.get_context(), PropList, CodeLoc);
+  } else {
+    return nullptr;
+  }
 }
 
 template <typename T>
@@ -286,18 +310,26 @@ T *aligned_alloc(
     size_t Alignment, size_t Count, const device &Dev, const context &Ctxt,
     usm::alloc Kind, const property_list &PropList = {},
     const detail::code_location &CodeLoc = detail::code_location::current()) {
-  return static_cast<T *>(aligned_alloc(max(Alignment, alignof(T)),
-                                        Count * sizeof(T), Dev, Ctxt, Kind,
-                                        PropList, CodeLoc));
+  if (!(Alignment & (Alignment - 1))) {
+    return static_cast<T *>(aligned_alloc(max(Alignment, alignof(T)),
+                                          Count * sizeof(T), Dev, Ctxt, Kind,
+                                          PropList, CodeLoc));
+  } else {
+    return nullptr;
+  }
 }
 
 template <typename T>
 T *aligned_alloc(
     size_t Alignment, size_t Count, const queue &Q, usm::alloc Kind,
     const property_list &PropList = {},
     const detail::code_location &CodeLoc = detail::code_location::current()) {
-  return aligned_alloc<T>(Alignment, Count, Q.get_device(), Q.get_context(),
-                          Kind, PropList, CodeLoc);
+  if (!(Alignment & (Alignment - 1))) {
+    return aligned_alloc<T>(Alignment, Count, Q.get_device(), Q.get_context(),
+                            Kind, PropList, CodeLoc);
+  } else {
+    return nullptr;
+  }
 }
 
 // Device copy enhancement APIs, prepare_for and release_from USM.

sycl/test-e2e/USM/align.cpp

@@ -0,0 +1,108 @@
+// RUN: %{build} -o %t.out
+// RUN: %{run} %t.out
+
+// UNSUPPORTED: gpu
+
+// E2E tests for annotated USM allocation functions with alignment arguments
+// that are not powers of 2. Note this test does not work on gpu because some
+// tests expect to return nullptr, e.g. when the alignment argument is not a
+// power of 2, while the gpu runtime has different behavior
+
+#include <sycl/sycl.hpp>
+
+#include <complex>
+#include <numeric>
+
+// clang-format on
+using namespace sycl;
+using namespace ext::oneapi::experimental;
+using namespace ext::intel::experimental;
+using alloc = usm::alloc;
+
+template <typename T> void testAlign(sycl::queue &q, unsigned align) {
+  const sycl::context &Ctx = q.get_context();
+  auto dev = q.get_device();
+
+  constexpr int N = 10;
+  assert(align > 0 || (align & (align - 1)) == 0);
+
+  auto ADevice = [&](size_t align, auto... args) {
+    return aligned_alloc_device(align, N, args...);
+  };
+  auto AHost = [&](size_t align, auto... args) {
+    return aligned_alloc_host(align, N, args...);
+  };
+  auto AShared = [&](size_t align, auto... args) {
+    return aligned_alloc_shared(align, N, args...);
+  };
+  auto AAnnotated = [&](size_t align, auto... args) {
+    return aligned_alloc(align, N, args...);
+  };
+
+  auto ATDevice = [&](size_t align, auto... args) {
+    return aligned_alloc_device<T>(align, N, args...);
+  };
+  auto ATHost = [&](size_t align, auto... args) {
+    return aligned_alloc_host<T>(align, N, args...);
+  };
+  auto ATShared = [&](size_t align, auto... args) {
+    return aligned_alloc_shared<T>(align, N, args...);
+  };
+  auto ATAnnotated = [&](size_t align, auto... args) {
+    return aligned_alloc<T>(align, N, args...);
+  };
+
+  // Test cases that are expected to return null
+  auto check_null = [&q](auto AllocFn, int Line = __builtin_LINE(),
+                         int Case = 0) {
+    decltype(AllocFn()) Ptr = AllocFn();
+    auto v = reinterpret_cast<uintptr_t>(Ptr);
+    if (v != 0) {
+      free(Ptr, q);
+      std::cout << "Failed at line " << Line << ", case " << Case << std::endl;
+      assert(false && "The return is not null!");
+    }
+  };
+
+  auto CheckNullAll = [&](auto Funcs, int Line = __builtin_LINE()) {
+    std::apply(
+        [&](auto... Fs) {
+          int Case = 0;
+          (void)std::initializer_list<int>{
+              (check_null(Fs, Line, Case++), 0)...};
+        },
+        Funcs);
+  };
+
+  CheckNullAll(std::tuple{
+      // Case: aligned_alloc_xxx with no alignment property, and the alignment
+      // argument is not a power of 2, the result is nullptr
+      [&]() { return ADevice(3, q); }, [&]() { return ADevice(5, dev, Ctx); },
+      [&]() { return AHost(7, q); }, [&]() { return AHost(9, Ctx); },
+      [&]() { return AShared(114, q); },
+      [&]() { return AShared(1023, dev, Ctx); },
+      [&]() { return AAnnotated(15, q, alloc::device); },
+      [&]() { return AAnnotated(17, dev, Ctx, alloc::host); }
+      // Case: aligned_alloc_xxx<T> with no alignment property, and the
+      // alignment
+      // argument is not a power of 2, the result is nullptr
+      ,
+      [&]() { return ATDevice(3, q); }, [&]() { return ATDevice(5, dev, Ctx); },
+      [&]() { return ATHost(7, q); }, [&]() { return ATHost(9, Ctx); },
+      [&]() { return ATShared(1919, q); },
+      [&]() { return ATShared(11, dev, Ctx); },
+      [&]() { return ATAnnotated(15, q, alloc::device); },
+      [&]() { return ATAnnotated(17, dev, Ctx, alloc::host); }});
+
+  // aligned_alloc<T>(17, N, q, ...);
+  // aligned_alloc_device<int>(17, N, q);
+  // aligned_alloc<int>(17, N, q, alloc::host);
+}
+
+int main() {
+  sycl::queue q;
+  testAlign<char>(q, 4);
+  testAlign<int>(q, 128);
+  testAlign<std::complex<double>>(q, 4);
+  return 0;
+}