Merge pull request #60025 from glessard/rdar96841899

[stdlib] fix unaligned loads of large SIMD vectors on x86_64

Author: glessard

stdlib/public/core/StringUTF16View.swift

@@ -597,11 +597,11 @@ extension String.UTF16View {
 
     while readPtr + MemoryLayout<U>.stride < endPtr {
       //Find the number of continuations (0b10xxxxxx)
-      let sValue = Builtin.loadRaw(readPtr._rawValue) as S
+      let sValue = readPtr.loadUnaligned(as: S.self)
       let continuations = S.zero.replacing(with: S.one, where: sValue .< -65 + 1)
 
       //Find the number of 4 byte code points (0b11110xxx)
-      let uValue = Builtin.loadRaw(readPtr._rawValue) as U
+      let uValue = readPtr.loadUnaligned(as: U.self)
       let fourBytes = S.zero.replacing(
         with: S.one,
         where: unsafeBitCast(

stdlib/public/core/UnsafeRawPointer.swift

@@ -244,7 +244,7 @@ public struct UnsafeRawPointer: _Pointer {
   /// - Parameter other: The typed pointer to convert.
   @_transparent
   public init<T>(@_nonEphemeral _ other: UnsafeMutablePointer<T>) {
-   _rawValue = other._rawValue
+    _rawValue = other._rawValue
   }
 
   /// Creates a new raw pointer from the given typed pointer.
@@ -257,8 +257,8 @@ public struct UnsafeRawPointer: _Pointer {
   ///   result is `nil`.
   @_transparent
   public init?<T>(@_nonEphemeral _ other: UnsafeMutablePointer<T>?) {
-   guard let unwrapped = other else { return nil }
-   _rawValue = unwrapped._rawValue
+    guard let unwrapped = other else { return nil }
+    _rawValue = unwrapped._rawValue
   }
 
   /// Deallocates the previously allocated memory block referenced by this pointer.
@@ -429,7 +429,7 @@ public struct UnsafeRawPointer: _Pointer {
                               MemoryLayout<T>.alignment._builtinWordValue)
     return Builtin.loadRaw(alignedPointer)
 #else
-  return Builtin.loadRaw(rawPointer)
+    return Builtin.loadRaw(rawPointer)
 #endif
   }
 
@@ -455,13 +455,25 @@ public struct UnsafeRawPointer: _Pointer {
   /// - Returns: A new instance of type `T`, read from the raw bytes at
   ///   `offset`. The returned instance isn't associated
   ///   with the value in the range of memory referenced by this pointer.
+  @inlinable
   @_alwaysEmitIntoClient
   public func loadUnaligned<T>(
     fromByteOffset offset: Int = 0,
     as type: T.Type
   ) -> T {
     _debugPrecondition(_isPOD(T.self))
-    return Builtin.loadRaw((self + offset)._rawValue)
+    return withUnsafeTemporaryAllocation(of: T.self, capacity: 1) {
+      let temporary = $0.baseAddress._unsafelyUnwrappedUnchecked
+      Builtin.int_memcpy_RawPointer_RawPointer_Int64(
+        temporary._rawValue,
+        (self + offset)._rawValue,
+        UInt64(MemoryLayout<T>.size)._value,
+        /*volatile:*/ false._value
+      )
+      return temporary.pointee
+    }
+    //FIXME: reimplement with `loadRaw` when supported in SIL (rdar://96956089)
+    // e.g. Builtin.loadRaw((self + offset)._rawValue)
   }
 }
 
@@ -1180,13 +1192,25 @@ public struct UnsafeMutableRawPointer: _Pointer {
   /// - Returns: A new instance of type `T`, read from the raw bytes at
   ///   `offset`. The returned instance isn't associated
   ///   with the value in the range of memory referenced by this pointer.
+  @inlinable
   @_alwaysEmitIntoClient
   public func loadUnaligned<T>(
     fromByteOffset offset: Int = 0,
     as type: T.Type
   ) -> T {
     _debugPrecondition(_isPOD(T.self))
-    return Builtin.loadRaw((self + offset)._rawValue)
+    return withUnsafeTemporaryAllocation(of: T.self, capacity: 1) {
+      let temporary = $0.baseAddress._unsafelyUnwrappedUnchecked
+      Builtin.int_memcpy_RawPointer_RawPointer_Int64(
+        temporary._rawValue,
+        (self + offset)._rawValue,
+        UInt64(MemoryLayout<T>.size)._value,
+        /*volatile:*/ false._value
+      )
+      return temporary.pointee
+    }
+    //FIXME: reimplement with `loadRaw` when supported in SIL (rdar://96956089)
+    // e.g. Builtin.loadRaw((self + offset)._rawValue)
   }
 
   /// Stores the given value's bytes into raw memory at the specified offset.

test/stdlib/UnsafeRawPointer.swift

@@ -114,6 +114,35 @@ UnsafeMutableRawPointerExtraTestSuite.test("load.unaligned")
   expectEqual(result, 0xffff_0000)
 }
 
+UnsafeMutableRawPointerExtraTestSuite.test("load.unaligned.SIMD")
+.skip(.custom({
+  if #available(SwiftStdlib 5.7, *) { return false }
+  return true
+}, reason: "Requires Swift 5.7's stdlib"))
+.code {
+  guard #available(SwiftStdlib 5.7, *) else { return }
+  var bytes: [UInt8] = Array(0..<64)
+  var offset = 3
+  let vector16 = bytes.withUnsafeBytes { buffer -> SIMD16<UInt8> in
+    let aligned = buffer.baseAddress!.alignedUp(for: SIMD16<UInt8>.self)
+    offset += buffer.baseAddress!.distance(to: aligned)
+    return buffer.loadUnaligned(fromByteOffset: offset, as: SIMD16<UInt8>.self)
+  }
+  expectEqual(Int(vector16[0]), offset)
+  var lastIndex = vector16.indices.last!
+  expectEqual(Int(vector16[lastIndex]), offset+lastIndex)
+
+  offset = 11
+  let vector32 = bytes.withUnsafeMutableBytes { buffer -> SIMD32<UInt8> in
+    let aligned = buffer.baseAddress!.alignedUp(for: SIMD32<UInt8>.self)
+    offset += buffer.baseAddress!.distance(to: aligned)
+    return buffer.loadUnaligned(fromByteOffset: offset, as: SIMD32<UInt8>.self)
+  }
+  expectEqual(Int(vector32[0]), offset)
+  lastIndex = vector32.indices.last!
+  expectEqual(Int(vector32[lastIndex]), offset+lastIndex)
+}
+
 UnsafeMutableRawPointerExtraTestSuite.test("load.invalid")
 .skip(.custom({ !_isDebugAssertConfiguration() },
               reason: "This tests a debug precondition.."))