wehub-resource-sync
103 lines
2.7 KiB
C++
103 lines
2.7 KiB
C++
// Copyright (c) 2025 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 <chrono>
|
|
#include <map>
|
|
#include <memory>
|
|
#include <set>
|
|
#include "paddle/common/macros.h"
|
|
#include "paddle/phi/api/include/tensor.h"
|
|
#include "paddle/phi/core/dense_tensor.h"
|
|
#include "paddle/utils/optional.h"
|
|
|
|
namespace egr {
|
|
|
|
class ActivationOffloaderWithPlace;
|
|
|
|
class PADDLE_API ReloadFunctor {
|
|
public:
|
|
explicit ReloadFunctor(std::weak_ptr<phi::DenseTensor> tensor,
|
|
ActivationOffloaderWithPlace *offloader);
|
|
|
|
void Reload();
|
|
|
|
private:
|
|
std::weak_ptr<phi::DenseTensor> tensor_;
|
|
ActivationOffloaderWithPlace *offloader_;
|
|
};
|
|
|
|
class ActivationOffloaderWithPlace {
|
|
public:
|
|
explicit ActivationOffloaderWithPlace(phi::GPUPlace place);
|
|
|
|
void SetSkipTensors(const std::vector<paddle::Tensor> &tensors);
|
|
|
|
paddle::optional<ReloadFunctor> Add(const paddle::Tensor &activation);
|
|
|
|
size_t Offload(size_t size);
|
|
|
|
void Remove(const std::weak_ptr<phi::DenseTensor> &tensor);
|
|
|
|
phi::GPUPlace Place() const { return place_; }
|
|
|
|
size_t CachedSize() const;
|
|
|
|
private:
|
|
void Shrink();
|
|
|
|
DISABLE_COPY_AND_ASSIGN(ActivationOffloaderWithPlace);
|
|
|
|
private:
|
|
using WeakTensorSet =
|
|
std::set<std::weak_ptr<phi::DenseTensor>,
|
|
std::owner_less<std::weak_ptr<phi::DenseTensor>>>;
|
|
using WeakTensorMap =
|
|
std::map<std::weak_ptr<phi::DenseTensor>,
|
|
size_t,
|
|
std::owner_less<std::weak_ptr<phi::DenseTensor>>>;
|
|
phi::GPUPlace place_;
|
|
WeakTensorMap activations_;
|
|
WeakTensorSet skip_tensors_;
|
|
};
|
|
|
|
class ActivationOffloader {
|
|
private:
|
|
ActivationOffloader() = default;
|
|
|
|
public:
|
|
void SetSkipTensors(const std::vector<paddle::Tensor> &tensors);
|
|
|
|
PADDLE_API paddle::optional<ReloadFunctor> Add(
|
|
const paddle::Tensor &activation);
|
|
|
|
size_t Offload(phi::Place place, size_t size);
|
|
|
|
size_t CachedSize() const;
|
|
|
|
PADDLE_API static ActivationOffloader *Instance();
|
|
|
|
private:
|
|
ActivationOffloaderWithPlace *GetOrCreateOffloader(phi::Place place);
|
|
|
|
DISABLE_COPY_AND_ASSIGN(ActivationOffloader);
|
|
|
|
private:
|
|
std::map<phi::Place, std::unique_ptr<ActivationOffloaderWithPlace>>
|
|
offloaders_;
|
|
};
|
|
|
|
} // namespace egr
|