mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 00:21:07 +01:00
9cbeb0faed
Summary: With https://github.com/pytorch/pytorch/pull/35562, we are running peephole optimization on inlining to reduce the number of nodes that are copied. The tracer encodes the sizes in the graph like: ``` graph(%0 : Double(7)): %1 : Function = prim::Constant[name="tensor_size"]() %2 : Tensor = prim::CallFunction(%1, %0) return (%2) ``` however people would like to reuse the graph with different shapes so running size invalidations would invalidate that. long term it might be better for the tracer to not include shape information but there are downstream users of that. Separates out FuseAddMM from peephole so that now there is a single `disable_size_optimizations` parameter, and onnx explicitly invokes fuseaddmm. Pull Request resolved: https://github.com/pytorch/pytorch/pull/36404 Differential Revision: D20968974 Pulled By: eellison fbshipit-source-id: 56f8f1699e3b0adeeccdfd5a67bb975fd41a2913
118 lines
2.7 KiB
C++
118 lines
2.7 KiB
C++
#include <test/cpp/jit/test_base.h>
|
|
#include <test/cpp/jit/test_utils.h>
|
|
|
|
#include <torch/csrc/jit/ir/ir.h>
|
|
#include <torch/csrc/jit/ir/irparser.h>
|
|
#include <torch/csrc/jit/passes/peephole.h>
|
|
|
|
namespace torch {
|
|
namespace jit {
|
|
|
|
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);
|
|
}
|
|
|
|
// tests addmm fusion
|
|
{
|
|
auto graph = std::make_shared<Graph>();
|
|
parseIR(
|
|
R"IR(
|
|
graph(
|
|
%0 : Float(2, 3, 4),
|
|
%1 : Float(2, 3, 4),
|
|
%2 : Float(1, 1, 1)):
|
|
%3 : int = prim::Constant[value=1]()
|
|
%4 : Tensor = aten::mm(%0, %1)
|
|
%5 : Tensor = aten::add(%4, %2, %3)
|
|
%6 : Tensor = aten::add(%5, %2, %3)
|
|
return (%6)
|
|
)IR",
|
|
graph.get());
|
|
FuseAddMM(graph);
|
|
testing::FileCheck().check("addmm")->run(*graph);
|
|
}
|
|
}
|
|
} // namespace jit
|
|
} // namespace torch
|