mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-20 21:14:14 +08:00
8892b782a8
Summary: att Test Plan: ci Rollback Plan: Differential Revidsion: D76167093 Pull Request resolved: https://github.com/pytorch/pytorch/pull/155374 Approved by: https://github.com/zhxchen17
33 lines
843 B
C++
33 lines
843 B
C++
#pragma once
|
|
|
|
#include <c10/util/FbcodeMaps.h>
|
|
|
|
#include <torch/nativert/graph/Graph.h>
|
|
|
|
namespace torch::nativert {
|
|
|
|
// ExecutionPlan is the result produced by ExecutionPlanner
|
|
// ATM, it only contains value deallocation plan.
|
|
struct ExecutionPlan {
|
|
// i-th entry in this list are the Values can be freed *after* execution i-th
|
|
// node
|
|
std::vector<std::vector<ValueId>> valuesToFree;
|
|
};
|
|
|
|
class ExecutionPlanner {
|
|
public:
|
|
explicit ExecutionPlanner(const Graph& graph) : graph_(graph) {}
|
|
|
|
std::unique_ptr<ExecutionPlan> createPlan();
|
|
// get list of values we can't free
|
|
static c10::FastSet<ValueId> staticValues(const Graph& graph);
|
|
|
|
private:
|
|
void generateDeallocationPlan(ExecutionPlan& plan);
|
|
const Graph& graph_;
|
|
};
|
|
|
|
std::ostream& operator<<(std::ostream& out, const ExecutionPlan& plan);
|
|
|
|
} // namespace torch::nativert
|