[SYCL][NativeCPU] Fix alignment of global and local memory. (#19076)

For parameters, we were inferring a suitable alignment from the
requested size, but for global and local memory we were not doing the
same.

Ideally we would keep track of what alignment we need, but UR does not
expose this information to us so assume the worst.

Author: hvdijk

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

@@ -81,6 +81,20 @@ inline void *aligned_malloc(size_t alignment, size_t size) {
   return ptr;
 }
 
+// In many cases we require aligned memory without being told what the alignment
+// requirement is. This helper function returns maximally aligned memory based
+// on the size.
+inline void *aligned_malloc(size_t size) {
+  constexpr size_t max_alignment = 16 * sizeof(double);
+  size_t alignment = max_alignment;
+  while (alignment > size) {
+    alignment >>= 1;
+  }
+  // aligned_malloc requires size to be a multiple of alignment; round up.
+  size = (size + alignment - 1) & ~(alignment - 1);
+  return aligned_malloc(alignment, size);
+}
+
 inline void aligned_free(void *ptr) {
 #ifdef _MSC_VER
   _aligned_free(ptr);

unified-runtime/source/adapters/native_cpu/kernel.hpp

@@ -42,7 +42,7 @@ struct ur_kernel_handle_t_ : RefCounted {
 
   ~ur_kernel_handle_t_() {
     removeArgReferences();
-    free(_localMemPool);
+    native_cpu::aligned_free(_localMemPool);
   }
 
   ur_kernel_handle_t_(ur_program_handle_t hProgram, const char *name,
@@ -59,7 +59,6 @@ struct ur_kernel_handle_t_ : RefCounted {
     args_index_t Indices;
     std::vector<size_t> ParamSizes;
     std::vector<bool> OwnsMem;
-    static constexpr size_t MaxAlign = 16 * sizeof(double);
 
     arguments() = default;
 
@@ -109,11 +108,7 @@ struct ur_kernel_handle_t_ : RefCounted {
         }
       }
       if (NeedAlloc) {
-        size_t Align = MaxAlign;
-        while (Align > Size) {
-          Align >>= 1;
-        }
-        Indices[Index] = native_cpu::aligned_malloc(Align, Size);
+        Indices[Index] = native_cpu::aligned_malloc(Size);
         ParamSizes[Index] = Size;
         OwnsMem[Index] = true;
       }
@@ -158,8 +153,8 @@ struct ur_kernel_handle_t_ : RefCounted {
     if (reqSize == 0 || reqSize == _localMemPoolSize) {
       return;
     }
-    // realloc handles nullptr case
-    _localMemPool = (char *)realloc(_localMemPool, reqSize);
+    native_cpu::aligned_free(_localMemPool);
+    _localMemPool = static_cast<char *>(native_cpu::aligned_malloc(reqSize));
     _localMemPoolSize = reqSize;
   }
 

unified-runtime/source/adapters/native_cpu/memory.hpp

@@ -19,12 +19,11 @@
 
 struct ur_mem_handle_t_ : ur_object {
   ur_mem_handle_t_(size_t Size, bool _IsImage)
-      : _mem{static_cast<char *>(malloc(Size))}, _ownsMem{true},
-        IsImage{_IsImage} {}
+      : _mem{static_cast<char *>(native_cpu::aligned_malloc(Size))},
+        _ownsMem{true}, IsImage{_IsImage} {}
 
   ur_mem_handle_t_(void *HostPtr, size_t Size, bool _IsImage)
-      : _mem{static_cast<char *>(malloc(Size))}, _ownsMem{true},
-        IsImage{_IsImage} {
+      : ur_mem_handle_t_(Size, _IsImage) {
     memcpy(_mem, HostPtr, Size);
   }
 
@@ -34,7 +33,7 @@ struct ur_mem_handle_t_ : ur_object {
 
   ~ur_mem_handle_t_() {
     if (_ownsMem) {
-      free(_mem);
+      native_cpu::aligned_free(_mem);
     }
   }