mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-06 12:20:52 +01:00
018e9826a2
Summary: pretty simple. if planner exists, which implies that planning is enabled, create a manager for each frame. the associated serial executor will use the withMemoryPlannner fn to ensure the deallocation is done after execution completes. Test Plan: CI Differential Revision: D73635809 Pull Request resolved: https://github.com/pytorch/pytorch/pull/157053 Approved by: https://github.com/henryoier, https://github.com/georgiaphillips
99 lines
2.7 KiB
C++
99 lines
2.7 KiB
C++
#include <gtest/gtest.h>
|
|
#include <torch/nativert/executor/ExecutionFrame.h>
|
|
|
|
namespace torch::nativert {
|
|
|
|
TEST(ExecutionFrameTest, CreateFrame) {
|
|
auto graph = stringToGraph(R"(
|
|
graph(%x, %y):
|
|
%a = foo(a=%x, b=%y)
|
|
%b = foo1(a=%x, b=%y)
|
|
%c = foo2(c=%a, d=%b)
|
|
return(%c)
|
|
)");
|
|
|
|
auto frame = ExecutionFrame(*graph);
|
|
|
|
for (auto* v : graph->values()) {
|
|
frame.setIValue(v->id(), c10::IValue(at::tensor({v->id()}, at::kInt)));
|
|
auto& frame_v = frame.getIValue(v->id());
|
|
EXPECT_EQ(frame_v.tagKind(), "Tensor");
|
|
}
|
|
|
|
auto outputs = frame.tryMoveUserOutputs();
|
|
|
|
EXPECT_EQ(outputs.size(), 1);
|
|
EXPECT_EQ(outputs[0].tagKind(), "Tensor");
|
|
EXPECT_EQ(outputs[0].toTensor().item().toInt(), graph->getValue("c")->id());
|
|
}
|
|
|
|
TEST(ExecutionFrameTest, TestSetBorrowedValue) {
|
|
auto graph = stringToGraph(R"(
|
|
graph(%x, %y):
|
|
%a = foo(a=%x, b=%y)
|
|
%b = foo1(a=%x, b=%y)
|
|
%c = foo2(c=%a, d=%b)
|
|
return(%c)
|
|
)");
|
|
|
|
auto x = c10::IValue(at::tensor({1}, at::kInt));
|
|
auto y = c10::IValue(at::tensor({2}, at::kInt));
|
|
|
|
{
|
|
auto frame = ExecutionFrame(*graph);
|
|
|
|
frame.setBorrowedIValue(
|
|
graph->getValue("x")->id(),
|
|
c10::MaybeOwnedTraits<c10::IValue>::createBorrow(x));
|
|
frame.setBorrowedIValue(
|
|
graph->getValue("y")->id(),
|
|
c10::MaybeOwnedTraits<c10::IValue>::createBorrow(y));
|
|
|
|
[[maybe_unused]] auto& w = frame.getIValue(graph->getValue("x")->id());
|
|
[[maybe_unused]] auto& z = frame.getIValue(graph->getValue("y")->id());
|
|
|
|
EXPECT_EQ(x.use_count(), 1);
|
|
EXPECT_EQ(y.use_count(), 1);
|
|
|
|
EXPECT_TRUE(c10::MaybeOwnedTraits<c10::IValue>{}.debugBorrowIsValid(
|
|
frame.getIValue(graph->getValue("x")->id())));
|
|
EXPECT_TRUE(c10::MaybeOwnedTraits<c10::IValue>{}.debugBorrowIsValid(
|
|
frame.getIValue(graph->getValue("y")->id())));
|
|
}
|
|
|
|
EXPECT_EQ(x.use_count(), 1);
|
|
EXPECT_EQ(y.use_count(), 1);
|
|
}
|
|
|
|
TEST(ExecutionFrameTest, TestPersistentValue) {
|
|
auto graph = stringToGraph(R"(
|
|
graph(%x, %y, %my_weight):
|
|
%a = foo(a=%x, b=%y)
|
|
%b = foo1(a=%x, b=%y)
|
|
%c = foo2(c=%a, d=%b)
|
|
return(%c)
|
|
)");
|
|
|
|
Weights weights(graph.get());
|
|
weights.setValue("my_weight", at::tensor({1}, at::kInt));
|
|
|
|
auto new_sig = graph->signature();
|
|
// NOLINTNEXTLINE(cppcoreguidelines-pro-type-const-cast)
|
|
const_cast<std::vector<std::pair<std::string, std::string>>&>(
|
|
new_sig.inputsToWeights())
|
|
.emplace_back("my_weight", "my_weight");
|
|
graph->setSignature(new_sig);
|
|
|
|
auto frame = ExecutionFrame(*graph, weights);
|
|
|
|
EXPECT_EQ(frame.weightVersion(), 0);
|
|
auto wid = graph->getValue("my_weight")->id();
|
|
|
|
EXPECT_NO_THROW(frame.getTensor(wid));
|
|
// can't release persistent value
|
|
frame.releaseValueIfNeeded(wid);
|
|
EXPECT_FALSE(frame.getIValue(wid).isNone());
|
|
}
|
|
|
|
} // namespace torch::nativert
|