diff --git a/onnxruntime/core/providers/cpu/cpu_execution_provider.cc b/onnxruntime/core/providers/cpu/cpu_execution_provider.cc index cbdf79caf3..813fdc54ec 100644 --- a/onnxruntime/core/providers/cpu/cpu_execution_provider.cc +++ b/onnxruntime/core/providers/cpu/cpu_execution_provider.cc @@ -2,6 +2,7 @@ // Licensed under the MIT License. #include "core/providers/cpu/cpu_execution_provider.h" +#include #include "core/framework/op_kernel.h" #include "core/framework/kernel_registry.h" #include "core/mlas/inc/mlas.h" @@ -29,7 +30,7 @@ CPUExecutionProvider::CPUExecutionProvider(const CPUExecutionProviderInfo& info) std::vector CPUExecutionProvider::CreatePreferredAllocators() { bool create_arena = info_.create_arena; -#if defined(USE_JEMALLOC) || defined(USE_MIMALLOC) +#if defined(USE_JEMALLOC) || defined(USE_MIMALLOC) || defined(ABSL_HAVE_ADDRESS_SANITIZER) // JEMalloc/mimalloc already have memory pool, so just use device allocator. create_arena = false; #elif !(defined(__amd64__) || defined(_M_AMD64) || defined(__aarch64__) || defined(_M_ARM64)) diff --git a/onnxruntime/test/framework/allocator_test.cc b/onnxruntime/test/framework/allocator_test.cc index 2c1cd48d3d..8961058628 100644 --- a/onnxruntime/test/framework/allocator_test.cc +++ b/onnxruntime/test/framework/allocator_test.cc @@ -1,5 +1,6 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#include #include "core/framework/allocator.h" @@ -15,7 +16,7 @@ TEST(AllocatorTest, CPUAllocatorTest) { EXPECT_EQ(cpu_arena->Info().id, 0); // arena is disabled for CPUExecutionProvider on x86 and JEMalloc -#if (defined(__amd64__) || defined(_M_AMD64) || defined(__aarch64__) || defined(_M_ARM64)) && !defined(USE_JEMALLOC) && !defined(USE_MIMALLOC) +#if (defined(__amd64__) || defined(_M_AMD64) || defined(__aarch64__) || defined(_M_ARM64)) && !defined(USE_JEMALLOC) && !defined(USE_MIMALLOC) && !defined(ABSL_HAVE_ADDRESS_SANITIZER) EXPECT_EQ(cpu_arena->Info().alloc_type, OrtAllocatorType::OrtArenaAllocator); #else EXPECT_EQ(cpu_arena->Info().alloc_type, OrtAllocatorType::OrtDeviceAllocator); diff --git a/onnxruntime/test/framework/session_state_test.cc b/onnxruntime/test/framework/session_state_test.cc index 0c2d8bcb2e..ed698ab920 100644 --- a/onnxruntime/test/framework/session_state_test.cc +++ b/onnxruntime/test/framework/session_state_test.cc @@ -2,6 +2,7 @@ // Licensed under the MIT License. #include +#include #include "asserts.h" #include "core/framework/execution_providers.h" @@ -215,7 +216,7 @@ TEST_P(SessionStateTestP, TestInitializerProcessing) { // if the relevant session option config flag is set // For this test we need to enable the arena-based allocator which is not supported on x86 builds, so // enable this test only on x64 builds -#if (defined(__amd64__) || defined(_M_AMD64) || defined(__aarch64__) || defined(_M_ARM64)) && !defined(USE_MIMALLOC) +#if (defined(__amd64__) || defined(_M_AMD64) || defined(__aarch64__) || defined(_M_ARM64)) && !defined(USE_MIMALLOC) && !defined(ABSL_HAVE_ADDRESS_SANITIZER) TEST(SessionStateTest, TestInitializerMemoryAllocatedUsingNonArenaMemory) { AllocatorPtr cpu_allocator = std::make_shared(); // Part 1: Feature turned ON (i.e.) allocate from non-arena memory diff --git a/onnxruntime/test/framework/tensor_test.cc b/onnxruntime/test/framework/tensor_test.cc index 38e3f184eb..9202543b75 100644 --- a/onnxruntime/test/framework/tensor_test.cc +++ b/onnxruntime/test/framework/tensor_test.cc @@ -6,7 +6,7 @@ #include "gmock/gmock.h" #include "gtest/gtest.h" - +#include #include namespace onnxruntime { @@ -138,7 +138,7 @@ TEST(TensorTest, EmptyTensorTest) { EXPECT_EQ(location.id, 0); // arena is disabled for CPUExecutionProvider on x86 and JEMalloc -#if (defined(__amd64__) || defined(_M_AMD64) || defined(__aarch64__) || defined(_M_ARM64)) && !defined(USE_JEMALLOC) && !defined(USE_MIMALLOC) +#if (defined(__amd64__) || defined(_M_AMD64) || defined(__aarch64__) || defined(_M_ARM64)) && !defined(USE_JEMALLOC) && !defined(USE_MIMALLOC) && !defined(ABSL_HAVE_ADDRESS_SANITIZER) EXPECT_EQ(location.alloc_type, OrtAllocatorType::OrtArenaAllocator); #else EXPECT_EQ(location.alloc_type, OrtAllocatorType::OrtDeviceAllocator); diff --git a/onnxruntime/test/quantization/quantization_test.cc b/onnxruntime/test/quantization/quantization_test.cc index bdfac77b33..773f56de53 100644 --- a/onnxruntime/test/quantization/quantization_test.cc +++ b/onnxruntime/test/quantization/quantization_test.cc @@ -99,24 +99,22 @@ void EnsureQuantizedTensorParam(const float scale, const T zero_point) { // First, create the scale tensor: auto alloc = TestCPUExecutionProvider()->CreatePreferredAllocators()[0]; - auto num_bytes = shape.Size() * sizeof(float); - void* data = alloc->Alloc(num_bytes); - float* float_data = static_cast(data); + IAllocatorUniquePtr buffer = IAllocator::MakeUniquePtr(alloc, shape.Size()); + float* float_data = buffer.get(); float_data[0] = scale; Tensor scale_tensor(DataTypeImpl::GetType(), shape, - data, + float_data, alloc->Info(), /*offset=*/0); // Next, create the zero_point tensor: - auto T_num_bytes = shape.Size() * sizeof(T); - void* T_data = alloc->Alloc(T_num_bytes); - T* typed_data = static_cast(T_data); + IAllocatorUniquePtr buffer2 = IAllocator::MakeUniquePtr(alloc, shape.Size()); + T* typed_data = buffer2.get(); typed_data[0] = zero_point; Tensor zero_point_tensor(DataTypeImpl::GetType(), shape, - T_data, + typed_data, alloc->Info(), /*offset=*/0);