[SYCL] Fix handling of unsupported alignment by aligned_alloc_xxx<T> (#12569)

Implementation is supposed to return `nullptr` when requested alignment is not supported. Since our runtime performs all allocations through Unified Runtime that means for us that any alignment which is not a power of two is unsupported.

Note that the resulting alignment may not be the same as requested one (per the SYCL 2020 specification) and therefore we can't just rely on return value of underlying non-templated version of the alloc function and have to perform the check explicitly.

There is an issue with some backends not properly returning error on an unsupported alignment, it wis reported in #11642

---------

Signed-off-by: Hu, Peisen <[email protected]>

Author: HPS-1

sycl/include/sycl/usm.hpp

@@ -154,6 +154,13 @@ __SYCL_EXPORT void *aligned_alloc(
     const property_list &propList,
     const detail::code_location &CodeLoc = detail::code_location::current());
 
+///
+// Helper function used to determine if the Alignment argument is a power of 2
+///
+inline size_t is_not_power_of_two(size_t Alignment) {
+  return (Alignment & (Alignment - 1));
+}
+
 ///
 // Template forms
 ///
@@ -179,6 +186,9 @@ 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()) {
+  if (is_not_power_of_two(Alignment)) {
+    return nullptr;
+  }
   return static_cast<T *>(aligned_alloc_device(max(Alignment, alignof(T)),
                                                Count * sizeof(T), Dev, Ctxt,
                                                PropList, CodeLoc));
@@ -189,6 +199,9 @@ 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()) {
+  if (is_not_power_of_two(Alignment)) {
+    return nullptr;
+  }
   return aligned_alloc_device<T>(Alignment, Count, Q.get_device(),
                                  Q.get_context(), PropList, CodeLoc);
 }
@@ -230,6 +243,9 @@ 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()) {
+  if (is_not_power_of_two(Alignment)) {
+    return nullptr;
+  }
   return static_cast<T *>(aligned_alloc_host(std ::max(Alignment, alignof(T)),
                                              Count * sizeof(T), Ctxt, PropList,
                                              CodeLoc));
@@ -240,6 +256,9 @@ 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()) {
+  if (is_not_power_of_two(Alignment)) {
+    return nullptr;
+  }
   return aligned_alloc_host<T>(Alignment, Count, Q.get_context(), PropList,
                                CodeLoc);
 }
@@ -249,6 +268,9 @@ 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()) {
+  if (is_not_power_of_two(Alignment)) {
+    return nullptr;
+  }
   return static_cast<T *>(aligned_alloc_shared(max(Alignment, alignof(T)),
                                                Count * sizeof(T), Dev, Ctxt,
                                                PropList, CodeLoc));
@@ -259,6 +281,9 @@ 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()) {
+  if (is_not_power_of_two(Alignment)) {
+    return nullptr;
+  }
   return aligned_alloc_shared<T>(Alignment, Count, Q.get_device(),
                                  Q.get_context(), PropList, CodeLoc);
 }
@@ -286,6 +311,9 @@ 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()) {
+  if (is_not_power_of_two(Alignment)) {
+    return nullptr;
+  }
   return static_cast<T *>(aligned_alloc(max(Alignment, alignof(T)),
                                         Count * sizeof(T), Dev, Ctxt, Kind,
                                         PropList, CodeLoc));
@@ -296,6 +324,9 @@ 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()) {
+  if (is_not_power_of_two(Alignment)) {
+    return nullptr;
+  }
   return aligned_alloc<T>(Alignment, Count, Q.get_device(), Q.get_context(),
                           Kind, PropList, CodeLoc);
 }

sycl/test-e2e/USM/align.cpp

@@ -0,0 +1,102 @@
+// 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 non-templated aligned_alloc_xxx functions to return nullptr,
+// e.g. when the alignment argument is not a power of 2, while they fail to do
+// so when run on gpu. This maybe because the gpu runtime has different
+// behavior. Therefore, GPU is unsupported until issue #12638 gets resolved.
+
+#include <sycl/sycl.hpp>
+
+#include <complex>
+#include <numeric>
+
+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, int Case) {
+    decltype(AllocFn()) Ptr = AllocFn();
+    if (Ptr != nullptr) {
+      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); }});
+}
+
+int main() {
+  sycl::queue q;
+  testAlign<char>(q, 4);
+  testAlign<int>(q, 128);
+  testAlign<std::complex<double>>(q, 4);
+  return 0;
+}