[llvm] [offload] add support for aligned allocations (PR #203353)

[Piotr Balcer via llvm-commits]

================
@@ -324,10 +328,31 @@ struct AMDGPUMemoryPoolTy {
   /// Get the allocation granularity of the pool.
   size_t getGranule() const { return Granule; }
 
+  /// Get the allocation alignment of the pool.
+  size_t getAlignment() const { return AllocationAlignment; }
+
   /// Allocate memory on the memory pool.
-  Error allocate(size_t Size, void **PtrStorage) {
+  Error allocate(size_t Size, void **PtrStorage, size_t Alignment) {
+    // A non-zero value passed as the Alignment indicates that the user expects
+    // the allocation to have a specific alignment. However, the HSA API does
+    // not allow users to define alignment. Therefore, the passed alignment is
+    // compared with the alignment of the memory allocated using the given pool.
+    // If the default alignment is greater than or equal to the alignment
+    // requested by the user, it would still meet the user's requirements.
+    if (Alignment > 0 && Alignment >= AllocationAlignment) {
+      return Plugin::error(
+          ErrorCode::UNSUPPORTED,
+          "too large alignment size for the device allocation");
+    }
+
     hsa_status_t Status =
         hsa_amd_memory_pool_allocate(MemoryPool, Size, 0, PtrStorage);
+
+    if (Alignment > 0 && !isAddrAligned(Align(Alignment), *PtrStorage)) {
+      return Plugin::error(ErrorCode::UNSUPPORTED,
+                           "unsupported alignment size");
----------------
pbalcer wrote:

If we enter this error scenario, `PtrStorage` will leak.

https://github.com/llvm/llvm-project/pull/203353