Merge pull request #21593 from atrick/fix-align-mask

Force manual allocation (via Unsafe*Pointer) to use >= 16 alignment.

include/swift/AST/Builtins.def

@@ -423,9 +423,29 @@ BUILTIN_MISC_OPERATION(IsSameMetatype, "is_same_metatype", "n", Special)
 BUILTIN_MISC_OPERATION(Alignof, "alignof", "n", Special)
 
 /// AllocRaw has type (Int, Int) -> Builtin.RawPointer
+///
+/// Parameters: object size, object alignment.
+///
+/// This alignment is not a mask; the compiler decrements by one to provide
+/// a mask to the runtime.
+///
+/// If alignment == 0, then the runtime will use "aligned" allocation,
+/// and the memory will be aligned to _swift_MinAllocationAlignment.
 BUILTIN_MISC_OPERATION(AllocRaw, "allocRaw", "", Special)
 
 /// DeallocRaw has type (Builtin.RawPointer, Int, Int) -> ()
+///
+/// Parameters: object address, object size, object alignment.
+///
+/// This alignment is not a mask; the compiler decrements by one to provide
+/// a mask to the runtime.
+///
+/// If alignment == 0, then the runtime will use the "aligned" deallocation
+/// path, which assumes that "aligned" allocation was used.
+///
+/// Note that the alignment value provided to `deallocRaw` must be identical to
+/// the alignment value provided to `allocRaw` when the memory at this address
+/// was allocated.
 BUILTIN_MISC_OPERATION(DeallocRaw, "deallocRaw", "", Special)
 
 /// Fence has type () -> ().

stdlib/public/SwiftShims/RuntimeShims.h

@@ -56,6 +56,26 @@ int _swift_stdlib_putc_stderr(int C);
 SWIFT_RUNTIME_STDLIB_API
 __swift_size_t _swift_stdlib_getHardwareConcurrency();
 
+/// Manually allocated memory is at least 16-byte aligned in Swift.
+///
+/// When swift_slowAlloc is called with "default" alignment (alignMask ==
+/// ~(size_t(0))), it will execute the "aligned allocation path" (AlignedAlloc)
+/// using this value for the alignment.
+///
+/// This is done so users do not need to specify the allocation alignment when
+/// manually deallocating memory via Unsafe[Raw][Buffer]Pointer. Any
+/// user-specified alignment less than or equal to _swift_MinAllocationAlignment
+/// results in a runtime request for "default" alignment. This guarantees that
+/// manual allocation always uses an "aligned" runtime allocation. If an
+/// allocation is "aligned" then it must be freed using an "aligned"
+/// deallocation. The converse must also hold. Since manual Unsafe*Pointer
+/// deallocation is always "aligned", the user never needs to specify alignment
+/// during deallocation.
+///
+/// This value is inlined (and constant propagated) in user code. On Windows,
+/// the Swift runtime and user binaries need to agree on this value.
+#define _swift_MinAllocationAlignment (__swift_size_t) 16
+
 #ifdef __cplusplus
 }} // extern "C", namespace swift
 #endif

stdlib/public/core/Builtin.swift

@@ -242,8 +242,6 @@ public func _unsafeUncheckedDowncast<T : AnyObject>(_ x: AnyObject, to type: T.T
   return Builtin.castReference(x)
 }
 
-import SwiftShims
-
 @inlinable
 @inline(__always)
 public func _getUnsafePointerToStoredProperties(_ x: AnyObject)
@@ -255,6 +253,18 @@ public func _getUnsafePointerToStoredProperties(_ x: AnyObject)
     storedPropertyOffset
 }
 
+/// Get the minimum alignment for manually allocated memory.
+///
+/// Memory allocated via UnsafeMutable[Raw][Buffer]Pointer must never pass
+/// an alignment less than this value to Builtin.allocRaw. This
+/// ensures that the memory can be deallocated without specifying the
+/// alignment.
+@inlinable
+@inline(__always)
+internal func _minAllocationAlignment() -> Int {
+  return _swift_MinAllocationAlignment
+}
+
 //===----------------------------------------------------------------------===//
 // Branch hints
 //===----------------------------------------------------------------------===//

stdlib/public/core/UnsafeBufferPointer.swift.gyb

@@ -567,7 +567,22 @@ extension Unsafe${Mutable}BufferPointer {
   public static func allocate(capacity count: Int) 
     -> UnsafeMutableBufferPointer<Element> {
     let size = MemoryLayout<Element>.stride * count
-    let raw  = Builtin.allocRaw(size._builtinWordValue, Builtin.alignof(Element.self))
+    // For any alignment <= _minAllocationAlignment, force alignment = 0.
+    // This forces the runtime's "aligned" allocation path so that
+    // deallocation does not require the original alignment.
+    //
+    // The runtime guarantees:
+    //
+    // align == 0 || align > _minAllocationAlignment:
+    //   Runtime uses "aligned allocation".
+    //
+    // 0 < align <= _minAllocationAlignment:
+    //   Runtime may use either malloc or "aligned allocation".
+    var align = Builtin.alignof(Element.self)
+    if Int(align) <= _minAllocationAlignment() {
+      align = (0)._builtinWordValue
+    }
+    let raw  = Builtin.allocRaw(size._builtinWordValue, align)
     Builtin.bindMemory(raw, count._builtinWordValue, Element.self)
     return UnsafeMutableBufferPointer(
       start: UnsafeMutablePointer(raw), count: count)

stdlib/public/core/UnsafePointer.swift

@@ -225,7 +225,11 @@ public struct UnsafePointer<Pointee>: _Pointer {
   /// block. The memory must not be initialized or `Pointee` must be a trivial type.
   @inlinable
   public func deallocate() {
-    Builtin.deallocRaw(_rawValue, (-1)._builtinWordValue, (-1)._builtinWordValue)
+    // Passing zero alignment to the runtime forces "aligned
+    // deallocation". Since allocation via `UnsafeMutable[Raw][Buffer]Pointer`
+    // always uses the "aligned allocation" path, this ensures that the
+    // runtime's allocation and deallocation paths are compatible.
+    Builtin.deallocRaw(_rawValue, (-1)._builtinWordValue, (0)._builtinWordValue)
   }
 
   /// Accesses the instance referenced by this pointer.
@@ -568,8 +572,22 @@ public struct UnsafeMutablePointer<Pointee>: _Pointer {
   public static func allocate(capacity count: Int)
     -> UnsafeMutablePointer<Pointee> {
     let size = MemoryLayout<Pointee>.stride * count
-    let rawPtr =
-      Builtin.allocRaw(size._builtinWordValue, Builtin.alignof(Pointee.self))
+    // For any alignment <= _minAllocationAlignment, force alignment = 0.
+    // This forces the runtime's "aligned" allocation path so that
+    // deallocation does not require the original alignment.
+    //
+    // The runtime guarantees:
+    //
+    // align == 0 || align > _minAllocationAlignment:
+    //   Runtime uses "aligned allocation".
+    //
+    // 0 < align <= _minAllocationAlignment:
+    //   Runtime may use either malloc or "aligned allocation".
+    var align = Builtin.alignof(Pointee.self)
+    if Int(align) <= _minAllocationAlignment() {
+      align = (0)._builtinWordValue
+    }
+    let rawPtr = Builtin.allocRaw(size._builtinWordValue, align)
     Builtin.bindMemory(rawPtr, count._builtinWordValue, Pointee.self)
     return UnsafeMutablePointer(rawPtr)
   }
@@ -580,8 +598,11 @@ public struct UnsafeMutablePointer<Pointee>: _Pointer {
   /// block. The memory must not be initialized or `Pointee` must be a trivial type.
   @inlinable
   public func deallocate() {
-    Builtin.deallocRaw(_rawValue, (-1)._builtinWordValue,
-                       Builtin.alignof(Pointee.self))
+    // Passing zero alignment to the runtime forces "aligned
+    // deallocation". Since allocation via `UnsafeMutable[Raw][Buffer]Pointer`
+    // always uses the "aligned allocation" path, this ensures that the
+    // runtime's allocation and deallocation paths are compatible.
+    Builtin.deallocRaw(_rawValue, (-1)._builtinWordValue, (0)._builtinWordValue)
   }
 
   /// Accesses the instance referenced by this pointer.

stdlib/public/core/UnsafeRawPointer.swift

@@ -241,7 +241,11 @@ public struct UnsafeRawPointer: _Pointer {
   /// trivial type.
   @inlinable
   public func deallocate() {
-    Builtin.deallocRaw(_rawValue, (-1)._builtinWordValue, (-1)._builtinWordValue)
+    // Passing zero alignment to the runtime forces "aligned
+    // deallocation". Since allocation via `UnsafeMutable[Raw][Buffer]Pointer`
+    // always uses the "aligned allocation" path, this ensures that the
+    // runtime's allocation and deallocation paths are compatible.
+    Builtin.deallocRaw(_rawValue, (-1)._builtinWordValue, (0)._builtinWordValue)
   }
 
   /// Binds the memory to the specified type and returns a typed pointer to the
@@ -577,6 +581,21 @@ public struct UnsafeMutableRawPointer: _Pointer {
   public static func allocate(
     byteCount: Int, alignment: Int
   ) -> UnsafeMutableRawPointer {
+    // For any alignment <= _minAllocationAlignment, force alignment = 0.
+    // This forces the runtime's "aligned" allocation path so that
+    // deallocation does not require the original alignment.
+    //
+    // The runtime guarantees:
+    //
+    // align == 0 || align > _minAllocationAlignment:
+    //   Runtime uses "aligned allocation".
+    //
+    // 0 < align <= _minAllocationAlignment:
+    //   Runtime may use either malloc or "aligned allocation".
+    var alignment = alignment
+    if alignment <= _minAllocationAlignment() {
+      alignment = 0
+    }
     return UnsafeMutableRawPointer(Builtin.allocRaw(
         byteCount._builtinWordValue, alignment._builtinWordValue))
   }
@@ -587,7 +606,11 @@ public struct UnsafeMutableRawPointer: _Pointer {
   /// trivial type.
   @inlinable
   public func deallocate() {
-    Builtin.deallocRaw(_rawValue, (-1)._builtinWordValue, (-1)._builtinWordValue)
+    // Passing zero alignment to the runtime forces "aligned
+    // deallocation". Since allocation via `UnsafeMutable[Raw][Buffer]Pointer`
+    // always uses the "aligned allocation" path, this ensures that the
+    // runtime's allocation and deallocation paths are compatible.
+    Builtin.deallocRaw(_rawValue, (-1)._builtinWordValue, (0)._builtinWordValue)
   }
 
   /// Binds the memory to the specified type and returns a typed pointer to the

stdlib/public/runtime/Heap.cpp

@@ -18,6 +18,7 @@
 #include "swift/Runtime/Heap.h"
 #include "Private.h"
 #include "swift/Runtime/Debug.h"
+#include "../SwiftShims/RuntimeShims.h"
 #include <algorithm>
 #include <stdlib.h>
 
@@ -48,19 +49,61 @@ using namespace swift;
 
 #endif
 
+// This assert ensures that manually allocated memory always uses the
+// AlignedAlloc path. The stdlib will use "default" alignment for any user
+// requested alignment less than or equal to _swift_MinAllocationAlignment. The
+// runtime must ensure that any alignment > _swift_MinAllocationAlignment also
+// uses the "aligned" deallocation path.
+static_assert(_swift_MinAllocationAlignment > MALLOC_ALIGN_MASK,
+              "Swift's default alignment must exceed platform malloc mask.");
 
-
+// When alignMask == ~(size_t(0)), allocation uses the "default"
+// _swift_MinAllocationAlignment. This is different than calling swift_slowAlloc
+// with `alignMask == _swift_MinAllocationAlignment - 1` because it forces
+// the use of AlignedAlloc. This allows manually allocated to memory to always
+// be deallocated with AlignedFree without knowledge of its original allocation
+// alignment.
+//
+// For alignMask > (_minAllocationAlignment-1)
+// i.e. alignment == 0 || alignment > _minAllocationAlignment:
+//   The runtime must use AlignedAlloc, and the standard library must
+//   deallocate using an alignment that meets the same condition.
+//
+// For alignMask <= (_minAllocationAlignment-1)
+// i.e. 0 < alignment <= _minAllocationAlignment:
+//   The runtime may use either malloc or AlignedAlloc, and the standard library
+//   must deallocate using an identical alignment.
 void *swift::swift_slowAlloc(size_t size, size_t alignMask) {
   void *p;
+  // This check also forces "default" alignment to use AlignedAlloc.
   if (alignMask <= MALLOC_ALIGN_MASK) {
     p = malloc(size);
   } else {
-    p = AlignedAlloc(size, alignMask + 1);
+    size_t alignment = (alignMask == ~(size_t(0)))
+                           ? _swift_MinAllocationAlignment
+                           : alignMask + 1;
+    p = AlignedAlloc(size, alignment);
   }
   if (!p) swift::crash("Could not allocate memory.");
   return p;
 }
 
+// Unknown alignment is specified by passing alignMask == ~(size_t(0)), forcing
+// the AlignedFree deallocation path for unknown alignment. The memory
+// deallocated with unknown alignment must have been allocated with either
+// "default" alignment, or alignment > _swift_MinAllocationAlignment, to
+// guarantee that it was allocated with AlignedAlloc.
+//
+// The standard library assumes the following behavior:
+//
+// For alignMask > (_minAllocationAlignment-1)
+// i.e. alignment == 0 || alignment > _minAllocationAlignment:
+//   The runtime must use AlignedFree.
+//
+// For alignMask <= (_minAllocationAlignment-1)
+// i.e. 0 < alignment <= _minAllocationAlignment:
+//   The runtime may use either `free` or AlignedFree as long as it is
+//   consistent with allocation with the same alignment.
 void swift::swift_slowDealloc(void *ptr, size_t bytes, size_t alignMask) {
   if (alignMask <= MALLOC_ALIGN_MASK) {
     free(ptr);