[Serialization] Fix buffer alignment issues (#11888)

* fix buffer alignment issues

* remove unused fields

* aligned memory allocation

* windows compat

* license. fix compiler warnings

* fix compilation error

* reinterpret_cast

src/ray/common/buffer.h

@@ -19,6 +19,9 @@
 
 #include "arrow/buffer.h"
 #include "ray/common/status.h"
+#include "ray/thirdparty/aligned_alloc.h"
+
+#define BUFFER_ALIGNMENT 64
 
 namespace ray {
 
@@ -48,6 +51,7 @@ class Buffer {
 };
 
 /// Represents a byte buffer in local memory.
+/// TODO(suquark): In C++17, we can use std::aligned_alloc
 class LocalMemoryBuffer : public Buffer {
  public:
   /// Constructor.
@@ -65,10 +69,10 @@ class LocalMemoryBuffer : public Buffer {
       : has_data_copy_(copy_data) {
     if (copy_data) {
       RAY_CHECK(data != nullptr);
-      buffer_.resize(size);
-      std::copy(data, data + size, buffer_.begin());
-      data_ = buffer_.data();
-      size_ = buffer_.size();
+      buffer_ = reinterpret_cast<uint8_t *>(aligned_malloc(size, BUFFER_ALIGNMENT));
+      std::copy(data, data + size, buffer_);
+      data_ = buffer_;
+      size_ = size;
     } else {
       data_ = data;
       size_ = size;
@@ -77,9 +81,9 @@ class LocalMemoryBuffer : public Buffer {
 
   /// Construct a LocalMemoryBuffer of all zeros of the given size.
   LocalMemoryBuffer(size_t size) : has_data_copy_(true) {
-    buffer_.resize(size, 0);
-    data_ = buffer_.data();
-    size_ = buffer_.size();
+    buffer_ = reinterpret_cast<uint8_t *>(aligned_malloc(size, BUFFER_ALIGNMENT));
+    data_ = buffer_;
+    size_ = size;
   }
 
   uint8_t *Data() const override { return data_; }
@@ -90,7 +94,12 @@ class LocalMemoryBuffer : public Buffer {
 
   bool IsPlasmaBuffer() const override { return false; }
 
-  ~LocalMemoryBuffer() { size_ = 0; }
+  ~LocalMemoryBuffer() {
+    size_ = 0;
+    if (buffer_ != NULL) {
+      aligned_free(buffer_);
+    }
+  }
 
  private:
   /// Disable copy constructor and assignment, as default copy will
@@ -105,7 +114,7 @@ class LocalMemoryBuffer : public Buffer {
   /// Whether this buffer holds a copy of data.
   bool has_data_copy_;
   /// This is only valid when `should_copy` is true.
-  std::vector<uint8_t> buffer_;
+  uint8_t *buffer_ = NULL;
 };
 
 /// Represents a byte buffer for plasma object. This can be used to hold the

src/ray/protobuf/serialization.proto

@@ -21,9 +21,9 @@ package ray.serialization;
 // ## About Pickle 5 Protocol
 // Pickle5 will create two things during serialization:
 //   1. Inband data. This is the framed pickle data for most objects.
-//   2. Buffers. They are python buffers referring internal data of objects.
-//       They contain metadata of the buffer and a native pointer.
-//       Thus they provide interface for zero-copy serialization.
+//   2. Buffers. They are python buffers wrapping internal data chunks of objects.
+//       They contain metadata of the buffer and a native pointer, which is
+//       intended for zero-copy serialization.
 //
 // ## Protobuf object
 // A PythonObject protobuf object will be created for each python object.
@@ -34,19 +34,21 @@ package ray.serialization;
 //
 // ## Python object serialization memory layout
 // This section describes the memory layout in the Plasma store buffer.
-// Unfortunately, no frame info is included in protobuf data, so we have to specify
-// the length and offset of PythonObject.
 // ---------------------
-// i64 offset(PythonObject):
-//     Offset of the PythonObject relative to the start of this buffer.
+// i64 len(inband_data):
+//     Length of the inband data.
 // i64 len(PythonObject):
 //     Length of the PythonObject.
-// inband_data | pad(64)
+// inband_data
 //     Inband data, padded with 64 bytes for the alignment of buffers.
-// buffers | pad(8)
-//     Raw data of buffers, padded with 8 bytes for the alignment of PythonObject.
 // PythonObject
-//     PythonObject is stored at the end because its size will be variable.
+//     PythonObject protobuf defined as below.
+// buffers | aligned(64)
+//     Raw data of buffers, this section is 64-byte aligned.
+//     Inside the section, large buffers (>=2048 bytes) are aligned
+//     to 64 bytes for faster SIMD, small buffers (<2048 bytes) are
+//     aligned to 8 bytes.
+
 // ---------------------
 
 // The message for metadata of python buffer objects.
@@ -81,16 +83,12 @@ message PythonBuffer {
 
 // The message for pickle5 serialized python object.
 message PythonObject {
-  // The offset of the inband data section relative to the beginning of the Plasma buffer.
-  uint64 inband_data_offset = 1;
   // The size of the inband data section.
-  uint64 inband_data_size = 2;
-  // The offset of the raw buffers section relative to the beginning of the Plasma buffer.
-  uint64 raw_buffers_offset = 3;
+  uint64 inband_data_size = 1;
   // The size of the buffers section. It is not used in deserialization
   // because we already have the length and address of every buffer. However, it could
   // be useful for debugging or future adjustment, so we just keep it.
-  uint64 raw_buffers_size = 4;
+  uint64 raw_buffers_size = 2;
   // The metadata of python buffer objects.
-  repeated PythonBuffer buffer = 5;
+  repeated PythonBuffer buffer = 3;
 }

src/ray/thirdparty/aligned_alloc.c

@@ -0,0 +1,78 @@
+/*
+Adopted from https://github.com/NickStrupat/AlignedMalloc
+
+The MIT License (MIT)
+
+Copyright (c) 2015 Nick Strupat
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+*/
+
+#include "ray/thirdparty/aligned_alloc.h"
+
+#if defined(__APPLE__) || defined(__linux__)
+
+#include <stdlib.h>
+
+void *aligned_malloc(size_t size, size_t alignment) {
+	void *pointer = NULL;
+	int rv = posix_memalign(&pointer, alignment, size);
+	if (rv != 0) {
+	  pointer = NULL;
+	}
+	return pointer;
+}
+
+void aligned_free(void *pointer) {
+	free(pointer);
+}
+
+#elif defined(_WIN32)
+
+#include <malloc.h>
+
+void *aligned_malloc(size_t size, size_t alignment) {
+	return _aligned_malloc(size, alignment);
+}
+
+void aligned_free(void *pointer) {
+	_aligned_free(pointer);
+}
+
+#else
+
+// https://sites.google.com/site/ruslancray/lab/bookshelf/interview/ci/low-level/write-an-aligned-malloc-free-function
+#include <stdlib.h>
+
+void *aligned_malloc(size_t size, size_t alignment) {
+	void *p1; // original block
+	void **p2; // aligned block
+	int offset = alignment - 1 + sizeof(void *);
+	if ((p1 = (void *)malloc(size + offset)) == NULL)
+		return NULL;
+	p2 = (void **)(((size_t)(p1) + offset) & ~(alignment - 1));
+	p2[-1] = p1;
+	return p2;
+}
+
+void aligned_free(void *pointer) {
+	free(((void **)pointer)[-1]);
+}
+
+#endif

src/ray/thirdparty/aligned_alloc.h

@@ -0,0 +1,43 @@
+/*
+Adopted from https://github.com/NickStrupat/AlignedMalloc
+
+The MIT License (MIT)
+
+Copyright (c) 2015 Nick Strupat
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.
+*/
+
+#ifndef ALIGNED_ALLOC_H_INCLUDED
+#define ALIGNED_ALLOC_H_INCLUDED
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+#include <stddef.h>
+
+void * aligned_malloc(size_t size, size_t alignment);
+void aligned_free(void * pointer);
+
+#ifdef __cplusplus
+}
+#endif
+
+#endif