mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
dc6b5b2a52
Summary: Simplify or eliminate boolean and/or expressions, optimize unwrapping a value that cannot be None, and optimize using `is` with a None and a non-None value Since peephole optimize is now introducing constants, i added another constant propagation pass after running it. Previously i had a PR that did this & optimized shape ops - i will add the shape optimizations in a separate PR. Pull Request resolved: https://github.com/pytorch/pytorch/pull/18259 Differential Revision: D14602749 Pulled By: eellison fbshipit-source-id: 1c3f5a67067d8dfdf55d7b78dcb616472ea8a267
105 lines
2.3 KiB
C++
105 lines
2.3 KiB
C++
#pragma once
|
|
|
|
#include <test/cpp/jit/test_base.h>
|
|
#include <test/cpp/jit/test_utils.h>
|
|
|
|
#include <torch/csrc/jit/irparser.h>
|
|
#include <torch/csrc/jit/passes/peephole.h>
|
|
|
|
namespace torch {
|
|
namespace jit {
|
|
|
|
using namespace script;
|
|
using namespace testing;
|
|
|
|
namespace test {
|
|
|
|
void testPeepholeOptimize() {
|
|
// test is / is not none optimization
|
|
{
|
|
auto graph = std::make_shared<Graph>();
|
|
parseIR(
|
|
R"IR(
|
|
graph(%0 : int):
|
|
%1 : None = prim::Constant()
|
|
%2 : bool = aten::__is__(%0, %1)
|
|
%3 : bool = aten::__isnot__(%0, %1)
|
|
return (%2, %3)
|
|
)IR",
|
|
graph.get());
|
|
PeepholeOptimize(graph);
|
|
testing::FileCheck()
|
|
.check_not("aten::__is__")
|
|
->check_not("aten::__isnot__")
|
|
->run(*graph);
|
|
}
|
|
{
|
|
auto graph = std::make_shared<Graph>();
|
|
parseIR(
|
|
R"IR(
|
|
graph(%0: int?):
|
|
%1 : None = prim::Constant()
|
|
%2 : bool = aten::__is__(%0, %1)
|
|
%3 : bool = aten::__isnot__(%0, %1)
|
|
return (%2, %3)
|
|
)IR",
|
|
graph.get());
|
|
PeepholeOptimize(graph);
|
|
testing::FileCheck()
|
|
.check("aten::__is__")
|
|
->check("aten::__isnot__")
|
|
->run(*graph);
|
|
}
|
|
|
|
{
|
|
auto graph = std::make_shared<Graph>();
|
|
parseIR(
|
|
R"IR(
|
|
graph(%0: int?):
|
|
%1 : Tensor = prim::AutogradZero()
|
|
%2 : None = prim::Constant()
|
|
%4 : bool = aten::__is__(%0, %1)
|
|
%5 : bool = aten::__isnot__(%1, %2)
|
|
return (%4, %5)
|
|
)IR",
|
|
graph.get());
|
|
PeepholeOptimize(graph);
|
|
testing::FileCheck()
|
|
.check("aten::__is__")
|
|
->check_not("aten::__isnot__")
|
|
->run(*graph);
|
|
}
|
|
|
|
// test unwrap optional
|
|
{
|
|
auto graph = std::make_shared<Graph>();
|
|
parseIR(
|
|
R"IR(
|
|
graph():
|
|
%1 : Float(*, *, *) = prim::Constant()
|
|
%2 : bool = aten::_unwrap_optional(%1)
|
|
%3 : bool = prim::unchecked_unwrap_optional(%1)
|
|
return (%2, %3)
|
|
)IR",
|
|
graph.get());
|
|
PeepholeOptimize(graph);
|
|
testing::FileCheck().check_not("unwrap")->run(*graph);
|
|
}
|
|
{
|
|
auto graph = std::make_shared<Graph>();
|
|
parseIR(
|
|
R"IR(
|
|
graph(%1 : Float(*, *, *)?):
|
|
%2 : bool = aten::_unwrap_optional(%1)
|
|
%3 : bool = prim::unchecked_unwrap_optional(%1)
|
|
return (%2, %3)
|
|
)IR",
|
|
graph.get());
|
|
PeepholeOptimize(graph);
|
|
testing::FileCheck().check_count("unwrap", 2)->run(*graph);
|
|
}
|
|
}
|
|
} // namespace test
|
|
} // namespace jit
|
|
} // namespace torch
|