mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
6384c2d81b
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/35115 This commit runs the newly added tools/clang_format.py on the JIT codebase and includes all of the formatting changes thus produced. Testing: Ran the script, CI. Test Plan: Imported from OSS Reviewed By: eellison Differential Revision: D20568523 Pulled By: SplitInfinity fbshipit-source-id: e09bdb982ccf090eecfb7c7b461b8d0681eef82b
98 lines
2.4 KiB
C++
98 lines
2.4 KiB
C++
#include <test/cpp/jit/test_base.h>
|
|
#include <test/cpp/jit/test_utils.h>
|
|
#include <torch/csrc/jit/ir/subgraph_matcher.h>
|
|
#include <torch/csrc/jit/passes/subgraph_rewrite.h>
|
|
#include <torch/csrc/jit/testing/file_check.h>
|
|
|
|
namespace torch {
|
|
namespace jit {
|
|
using namespace testing;
|
|
|
|
void testFilterMatch() {
|
|
auto graph = std::make_shared<Graph>();
|
|
|
|
parseIR(
|
|
R"IR(
|
|
graph(%0):
|
|
%a = a::aaa(%0)
|
|
%b = prim::Constant[value=1]()
|
|
%c = c::ccc(%a, %b)
|
|
return (%c))IR",
|
|
graph.get());
|
|
|
|
std::string pattern = R"IR(
|
|
graph(%a, %b):
|
|
%c = c::ccc(%a, %b)
|
|
return (%c))IR";
|
|
Graph pattern_graph;
|
|
std::unordered_map<std::string, Value*> vmap;
|
|
|
|
parseIR(pattern, &pattern_graph, vmap);
|
|
|
|
auto filter = [](const Match& match,
|
|
const std::unordered_map<std::string, Value*>& vmap) {
|
|
const auto& match_vmap = match.values_map;
|
|
auto b_node = match_vmap.at(vmap.at("b"))->node();
|
|
return b_node->kind() == prim::Constant;
|
|
};
|
|
|
|
std::string replacement = R"IR(
|
|
graph(%a, %b):
|
|
%d = d::ddd(%a, %b)
|
|
return (%d))IR";
|
|
|
|
SubgraphRewriter rewriter;
|
|
rewriter.RegisterRewritePattern(pattern, replacement);
|
|
rewriter.runOnGraph(graph, filter);
|
|
|
|
FileCheck().check("d::ddd")->check_not("c::ccc")->run(*graph);
|
|
}
|
|
|
|
void testFilterNoMatch() {
|
|
auto graph = std::make_shared<Graph>();
|
|
parseIR(
|
|
R"IR(
|
|
graph(%0):
|
|
%a = a::aaa(%0)
|
|
%b = prim::Constant[value=1]()
|
|
%c = c::ccc(%a, %b)
|
|
return (%c))IR",
|
|
graph.get());
|
|
|
|
std::string pattern = R"IR(
|
|
graph(%a, %b):
|
|
%c = c::ccc(%a, %b)
|
|
return (%c))IR";
|
|
Graph pattern_graph;
|
|
std::unordered_map<std::string, Value*> vmap;
|
|
|
|
parseIR(pattern, &pattern_graph, vmap);
|
|
|
|
auto filter = [](const Match& match,
|
|
const std::unordered_map<std::string, Value*>& vmap) {
|
|
const auto& match_vmap = match.values_map;
|
|
auto b_node = match_vmap.at(vmap.at("b"))->node();
|
|
// b_node is not Constant, so this won't match and we'll skip the rewrite
|
|
return b_node->kind() == prim::Assign;
|
|
};
|
|
|
|
std::string replacement = R"IR(
|
|
graph(%a, %b):
|
|
%d = d::ddd(%a, %b)
|
|
return (%d))IR";
|
|
|
|
SubgraphRewriter rewriter;
|
|
rewriter.RegisterRewritePattern(pattern, replacement);
|
|
rewriter.runOnGraph(graph, filter);
|
|
|
|
FileCheck().check("c::ccc")->check_not("d::ddd")->run(*graph);
|
|
}
|
|
|
|
void testSubgraphRewriter() {
|
|
testFilterMatch();
|
|
testFilterNoMatch();
|
|
}
|
|
|
|
} // namespace jit
|
|
} // namespace torch
|