Fix alignment crashes in AVX512 builds (#19121)

* Fix issue #15588 by simplifying the code

The allocator.h code tried to be clever and use 32 byte alignment for SSE/AVX2/etc use,
and 64 byte alignment for AVX512.

Unfortunately, the #ifdef in use (from EIGEN) is not useful; the bazel BUILD files do
not propagate the tf_copts() compiler flags when the allocator.cc/allocator.h files get
compiled, to EIGEN does not see the actual AVX512 using compiler flags...

Rather than changing compiler flag propagation throughout a whole bunch of code,
there's an opportunity to just simplify the code and always use 64 byte alignment.
Yes it wastes a bit of space, but on the other hand now these allocations are
cache line aligned which isn't a bad thing... and an ifdef can be dropped

Signed-off-by: Arjan van de Ven <arjan@linux.intel.com>

* Set EIGEN_MAX_ALIGN_BYTES=64

This patch sets a 64 byte upper bound on the alignment of memory allocated by
eigen.  This is necessary to prevent crashes during the execution of the unit
tests when they are compiled with AVX512 support.

Signed-off-by: Mark Ryan <mark.d.ryan@intel.com>

* Update the tensorflow/compiler/aot tests for 64 byte alignment

Modifications to the tensorflow/core/framework/allocator.h to always
use 64 byte alignment causes failures in the tensorflow/compiler/aot
unit tests.  This patch updates these tests so that they pass with
64 byte aligned allocated memory.

Signed-off-by: Mark Ryan <mark.d.ryan@intel.com>

* Update Tensor.Slice_Basic for 64 byte alignment

The test case

//tensorflow/core:framework_tensor_test:Tensor.Slice_Basic

fails with EIGEN_MAX_ALIGN_BYTES set to 64.  The reason is that the
slices it takes of the sample tensor are 32 byte and not 64 byte
aligned.  This commit increases one of the dimensions of the original
tensor to ensure that the slices taken by the test cases are indeed 64
byte aligned.

Signed-off-by: Mark Ryan <mark.d.ryan@intel.com>

* Update ScopedAllocatorConcatOpTest.Reshape for 64 byte alignment

The ScopedAllocatorConcatOpTest.Reshape test requires that the elements
of the field_shapes parameter of ExecOp are multiples of
Allocator::kAllocatorAlignment in size.  If they are not, the backing
tensor allocated by PrepOp will have too many elements and reshaping
will fail.  This commit modifies the test case, making the elements
64 bytes in size, the new value for Allocator::kAllocatorAlignment.

Signed-off-by: Mark Ryan <mark.d.ryan@intel.com>

1 files changed, 1 insertions, 0 deletions

diff --git a/third_party/eigen.BUILD b/third_party/eigen.BUILD
index 07bb6645eb..e54c1a4501 100644
--- a/third_party/eigen.BUILD
+++ b/third_party/eigen.BUILD
@@ -64,6 +64,7 @@ cc_library(
         # This define (mostly) guarantees we don't link any problematic
         # code. We use it, but we do not rely on it, as evidenced above.
         "EIGEN_MPL2_ONLY",
+        "EIGEN_MAX_ALIGN_BYTES=64",
     ],
     includes = ["."],
     visibility = ["//visibility:public"],