diff --git a/onnxruntime/core/framework/execution_frame.cc b/onnxruntime/core/framework/execution_frame.cc index c5323c4e0a..4f4f42c091 100644 --- a/onnxruntime/core/framework/execution_frame.cc +++ b/onnxruntime/core/framework/execution_frame.cc @@ -89,7 +89,7 @@ Status ExecutionFrame::AllocateMLValueTensorSelfOwnBufferHelper(int mlvalue_inde if (len < 0) { return Status(ONNXRUNTIME, INVALID_ARGUMENT, "Tensor shape cannot contain any negative value"); } - if (!IAllocator::CalcMemSizeForArray(len, element_type->Size(), &size)) { + if (!IAllocator::CalcMemSizeForArrayWithAlignment<64>(len, element_type->Size(), &size)) { return Status(ONNXRUNTIME, FAIL, "size overflow"); } } diff --git a/onnxruntime/test/framework/execution_frame_test.cc b/onnxruntime/test/framework/execution_frame_test.cc index 8125bcb0dd..e2b741cffd 100644 --- a/onnxruntime/test/framework/execution_frame_test.cc +++ b/onnxruntime/test/framework/execution_frame_test.cc @@ -257,9 +257,9 @@ TEST(ExecutionFrameTest, MemPatternTest) { EXPECT_EQ(pattern.patterns.size(), pattern.locations.size()); EXPECT_EQ(pattern.patterns.size(), 1); auto p = pattern.GetPatterns(cpu_allocator->Info()); - EXPECT_EQ(p->PeakSize(), sizeof(float) * (4 + 6)); + EXPECT_EQ(p->PeakSize(), 2 * 64); // each allocation is 64-byte aligned EXPECT_EQ(p->GetBlock(3)->offset_, 0); - EXPECT_EQ(p->GetBlock(4)->offset_, sizeof(float) * 4); + EXPECT_EQ(p->GetBlock(4)->offset_, 64); } } // namespace test } // namespace onnxruntime