Yu Yang
6 years ago
parent
02631965c8
commit
ea81f8eed2
@ -1,69 +0,0 @@
|
||||
// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#include "paddle/fluid/memory/allocation/naive_managed_allocator.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace memory {
|
||||
namespace allocation {
|
||||
|
||||
NaiveManagedAllocator::NaiveManagedAllocator(
|
||||
std::unique_ptr<Allocator> &&allocator) {
|
||||
auto *underlying_allocator =
|
||||
dynamic_cast<UnmanagedAllocator *>(allocator.get());
|
||||
PADDLE_ENFORCE_NOT_NULL(underlying_allocator);
|
||||
allocator.release();
|
||||
Init(std::unique_ptr<UnmanagedAllocator>(underlying_allocator));
|
||||
}
|
||||
|
||||
NaiveManagedAllocator::NaiveManagedAllocator(
|
||||
std::unique_ptr<UnmanagedAllocator> &&allocator) {
|
||||
Init(std::move(allocator));
|
||||
}
|
||||
void NaiveManagedAllocator::Init(
|
||||
std::unique_ptr<UnmanagedAllocator> &&allocator) {
|
||||
underlying_allocator_ = std::move(allocator);
|
||||
}
|
||||
bool NaiveManagedAllocator::IsAllocThreadSafe() const {
|
||||
return underlying_allocator_->IsAllocThreadSafe();
|
||||
}
|
||||
std::unique_ptr<Allocation> NaiveManagedAllocator::Allocate(size_t size,
|
||||
Attr attr) {
|
||||
std::unique_ptr<Allocation> allocation =
|
||||
underlying_allocator_->Allocate(size, attr);
|
||||
return std::unique_ptr<Allocation>(
|
||||
new NaiveManagedAllocation(std::move(allocation), shared_from_this()));
|
||||
}
|
||||
std::shared_ptr<Allocation> NaiveManagedAllocator::AllocateShared(size_t size,
|
||||
Attr attr) {
|
||||
std::unique_ptr<Allocation> allocation =
|
||||
underlying_allocator_->Allocate(size, attr);
|
||||
return std::shared_ptr<Allocation>(
|
||||
new NaiveManagedAllocation(std::move(allocation), shared_from_this()));
|
||||
}
|
||||
|
||||
NaiveManagedAllocation::~NaiveManagedAllocation() {
|
||||
auto allocator = allocator_.lock();
|
||||
if (UNLIKELY(allocator == nullptr)) {
|
||||
// the allocator is destructed before allocations.
|
||||
// do nothing.
|
||||
return;
|
||||
}
|
||||
// invoke Free
|
||||
allocator->UnderlyingAllocator().FreeUniquePtr(
|
||||
std::move(underlying_allocation_));
|
||||
}
|
||||
} // namespace allocation
|
||||
} // namespace memory
|
||||
} // namespace paddle
|
||||
@ -1,76 +0,0 @@
|
||||
// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#pragma once
|
||||
#include <memory>
|
||||
#include "paddle/fluid/memory/allocation/allocator.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace memory {
|
||||
namespace allocation {
|
||||
|
||||
// An allocator to wrap an UnmanagedAllocator and make the allocation managed
|
||||
// by C++ smart ptr.
|
||||
//
|
||||
// NOTE: if the NaiveManagedAllocator is destroyed before
|
||||
// NaiveManagedAllocations, the allocation will never be released.
|
||||
class NaiveManagedAllocator;
|
||||
class NaiveManagedAllocation : public Allocation {
|
||||
public:
|
||||
NaiveManagedAllocation(std::unique_ptr<Allocation>&& underlying_allocation,
|
||||
std::shared_ptr<NaiveManagedAllocator> allocator)
|
||||
: Allocation(underlying_allocation->ptr(), underlying_allocation->size(),
|
||||
underlying_allocation->place()),
|
||||
underlying_allocation_(std::move(underlying_allocation)),
|
||||
allocator_(allocator) {}
|
||||
|
||||
~NaiveManagedAllocation() final;
|
||||
|
||||
private:
|
||||
std::unique_ptr<Allocation> underlying_allocation_;
|
||||
std::weak_ptr<NaiveManagedAllocator> allocator_;
|
||||
};
|
||||
|
||||
class NaiveManagedAllocator
|
||||
: public ManagedAllocator,
|
||||
public std::enable_shared_from_this<NaiveManagedAllocator> {
|
||||
public:
|
||||
template <typename... ARGS>
|
||||
static std::shared_ptr<ManagedAllocator> Create(ARGS... args) {
|
||||
return std::static_pointer_cast<ManagedAllocator>(
|
||||
std::shared_ptr<NaiveManagedAllocator>(
|
||||
new NaiveManagedAllocator(std::move(args)...)));
|
||||
}
|
||||
|
||||
inline UnmanagedAllocator& UnderlyingAllocator() {
|
||||
return *underlying_allocator_;
|
||||
}
|
||||
|
||||
bool IsAllocThreadSafe() const override;
|
||||
std::unique_ptr<Allocation> Allocate(size_t size,
|
||||
Attr attr = kDefault) override;
|
||||
std::shared_ptr<Allocation> AllocateShared(size_t size,
|
||||
Attr attr = kDefault) override;
|
||||
|
||||
private:
|
||||
explicit NaiveManagedAllocator(std::unique_ptr<Allocator>&& allocator);
|
||||
explicit NaiveManagedAllocator(
|
||||
std::unique_ptr<UnmanagedAllocator>&& allocator);
|
||||
void Init(std::unique_ptr<UnmanagedAllocator>&& allocator);
|
||||
|
||||
std::unique_ptr<UnmanagedAllocator> underlying_allocator_;
|
||||
};
|
||||
} // namespace allocation
|
||||
} // namespace memory
|
||||
} // namespace paddle
|
||||
@ -1,82 +0,0 @@
|
||||
// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#include "paddle/fluid/memory/allocation/naive_managed_allocator.h"
|
||||
#include <atomic> // NOLINT
|
||||
#include <random>
|
||||
#include <thread> // NOLINT
|
||||
#include <vector>
|
||||
#include "gtest/gtest.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace memory {
|
||||
namespace allocation {
|
||||
|
||||
class StubAllocator : public UnmanagedAllocator {
|
||||
public:
|
||||
std::unique_ptr<Allocation> Allocate(size_t size,
|
||||
Attr attr = kDefault) override {
|
||||
counter_.fetch_add(1);
|
||||
return std::unique_ptr<Allocation>(
|
||||
new Allocation(nullptr, size, platform::CPUPlace()));
|
||||
}
|
||||
void FreeUniquePtr(std::unique_ptr<Allocation> allocation) override {
|
||||
counter_.fetch_sub(1);
|
||||
}
|
||||
bool IsAllocThreadSafe() const override { return true; }
|
||||
|
||||
std::atomic<int> counter_{0};
|
||||
};
|
||||
|
||||
TEST(NaiveManagedAllocator, main) {
|
||||
auto allocator = NaiveManagedAllocator::Create(
|
||||
std::unique_ptr<Allocator>(new StubAllocator()));
|
||||
|
||||
auto th_main = [=] {
|
||||
std::random_device dev;
|
||||
std::default_random_engine engine(dev());
|
||||
std::uniform_int_distribution<int> dist(0, 1);
|
||||
|
||||
std::vector<std::shared_ptr<Allocation>> allocations;
|
||||
|
||||
for (int j = 0; j < 1024; ++j) {
|
||||
bool to_insert = static_cast<bool>(dist(engine));
|
||||
if (to_insert) {
|
||||
allocations.emplace_back(allocator->AllocateShared(10));
|
||||
} else {
|
||||
if (!allocations.empty()) {
|
||||
allocations.pop_back();
|
||||
}
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
{
|
||||
std::vector<std::thread> threads;
|
||||
for (size_t i = 0; i < 1024; ++i) {
|
||||
threads.emplace_back(th_main);
|
||||
}
|
||||
for (auto& th : threads) {
|
||||
th.join();
|
||||
}
|
||||
}
|
||||
ASSERT_EQ(reinterpret_cast<StubAllocator&>(
|
||||
std::dynamic_pointer_cast<NaiveManagedAllocator>(allocator)
|
||||
->UnderlyingAllocator())
|
||||
.counter_,
|
||||
0);
|
||||
}
|
||||
} // namespace allocation
|
||||
} // namespace memory
|
||||
} // namespace paddle
|
||||
@ -0,0 +1,35 @@
|
||||
// Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
#pragma once
|
||||
|
||||
#include "paddle/fluid/memory/allocation/allocator.h"
|
||||
|
||||
namespace paddle {
|
||||
namespace memory {
|
||||
namespace allocation {
|
||||
|
||||
class UnderlyingManualAllocation : public MannualFreeAllocation {
|
||||
public:
|
||||
UnderlyingManualAllocation(MannualFreeAllocator* allocator,
|
||||
std::unique_ptr<Allocation> allocation)
|
||||
: MannualFreeAllocation(allocator, allocation->ptr(), allocation->size(),
|
||||
allocation->place()),
|
||||
allocation_(std::move(allocation)) {}
|
||||
std::unique_ptr<Allocation> allocation_;
|
||||
};
|
||||
|
||||
} // namespace allocation
|
||||
} // namespace memory
|
||||
} // namespace paddle
|
||||
Some files were not shown because too many files have changed in this diff Show More
Loading…
Reference in new issue