mirror of
https://github.com/zebrajr/pytorch.git
synced 2025-12-07 12:21:27 +01:00
3be2f7c8e6
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/20602 ghimport-source-id: fa3225bb5d70729d6a1bcf88295d031707d986a1 Differential Revision: D15377635 Pulled By: ZolotukhinM fbshipit-source-id: ebd385e7b9436429d0ad76ed3d932925a29f6456
489 lines
10 KiB
C++
489 lines
10 KiB
C++
#pragma once
|
|
|
|
#include "test/cpp/jit/test_base.h"
|
|
#include "torch/csrc/jit/subgraph_matcher.h"
|
|
|
|
namespace torch {
|
|
namespace jit {
|
|
namespace {
|
|
|
|
void testTrivial1() {
|
|
Graph graph, pattern;
|
|
script::parseIR(
|
|
R"IR(
|
|
graph(%0):
|
|
%a = a::aaa(%0)
|
|
return (%a))IR",
|
|
&graph);
|
|
script::parseIR(
|
|
R"IR(
|
|
graph(%0):
|
|
%x = a::aaa(%0)
|
|
return (%x))IR",
|
|
&pattern);
|
|
AT_ASSERT(!findPatternMatches(pattern, graph).empty());
|
|
}
|
|
|
|
void testTrivial2() {
|
|
Graph graph;
|
|
Var i0 = Var::asNewInput(graph);
|
|
auto o0 = i0.tanh();
|
|
graph.registerOutput(o0.value());
|
|
|
|
Graph pattern;
|
|
Var mi0 = Var::asNewInput(pattern);
|
|
auto mo0 = mi0.tanh();
|
|
pattern.registerOutput(mo0.value());
|
|
|
|
auto matches = findPatternMatches(pattern, graph);
|
|
AT_ASSERT(matches.size() == 1);
|
|
for (const Match& m : matches) {
|
|
AT_ASSERT(m.values_map.at(mi0.value()) == i0.value());
|
|
AT_ASSERT(m.values_map.at(mo0.value()) == o0.value());
|
|
AT_ASSERT(m.nodes_map.at(mo0.value()->node()) == o0.value()->node());
|
|
}
|
|
}
|
|
|
|
void testTrivial3() {
|
|
Graph graph, pattern;
|
|
script::parseIR(
|
|
R"IR(
|
|
graph(%0):
|
|
%a = a::a(%0)
|
|
%b = a::b(%0)
|
|
%c = a::c(%a, %b)
|
|
return (%c))IR",
|
|
&graph);
|
|
script::parseIR(
|
|
R"IR(
|
|
graph(%a, %b):
|
|
%c = a::c(%a, %b)
|
|
return (%c))IR",
|
|
&pattern);
|
|
AT_ASSERT(!findPatternMatches(pattern, graph).empty());
|
|
}
|
|
|
|
void testTrivial4() {
|
|
Graph graph;
|
|
Var i0 = Var::asNewInput(graph);
|
|
Var i1 = Var::asNewInput(graph);
|
|
auto o0 = i0 * i1;
|
|
graph.registerOutput(o0.value());
|
|
|
|
Graph pattern;
|
|
auto m0 = Var::asNewInput(pattern);
|
|
auto m1 = Var::asNewInput(pattern);
|
|
auto m2 = m0 * m1;
|
|
pattern.registerOutput(m2.value());
|
|
|
|
auto matches = findPatternMatches(pattern, graph);
|
|
AT_ASSERT(matches.size() == 1);
|
|
for (const Match& m : matches) {
|
|
AT_ASSERT(m.values_map.at(m0.value()) == i0.value());
|
|
AT_ASSERT(m.values_map.at(m1.value()) == i1.value());
|
|
AT_ASSERT(m.values_map.at(m2.value()) == o0.value());
|
|
AT_ASSERT(m.nodes_map.at(m2.value()->node()) == o0.value()->node());
|
|
}
|
|
}
|
|
|
|
void testLinear1() {
|
|
Graph graph, pattern;
|
|
script::parseIR(
|
|
R"IR(
|
|
graph(%0):
|
|
%a = a::aaa(%0)
|
|
%b = b::bbb(%a)
|
|
%c = c::ccc(%b)
|
|
%d = d::ddd(%c)
|
|
%a = a::aaa(%0)
|
|
return (%d))IR",
|
|
&graph);
|
|
script::parseIR(
|
|
R"IR(
|
|
graph(%0):
|
|
%x = b::bbb(%0)
|
|
%y = c::ccc(%x)
|
|
return (%y))IR",
|
|
&pattern);
|
|
AT_ASSERT(!findPatternMatches(pattern, graph).empty());
|
|
}
|
|
|
|
void testLinear2() {
|
|
Graph graph;
|
|
Var i0 = Var::asNewInput(graph);
|
|
auto o0 = i0.tanh();
|
|
auto o1 = o0.tanh();
|
|
graph.registerOutput(o1.value());
|
|
|
|
Graph pattern;
|
|
auto m0 = Var::asNewInput(pattern);
|
|
auto m1 = m0.tanh();
|
|
auto m2 = m1.tanh();
|
|
pattern.registerOutput(m2.value());
|
|
|
|
auto matches = findPatternMatches(pattern, graph);
|
|
AT_ASSERT(matches.size() == 1);
|
|
for (const Match& m : matches) {
|
|
AT_ASSERT(m.values_map.at(m0.value()) == i0.value());
|
|
AT_ASSERT(m.values_map.at(m1.value()) == o0.value());
|
|
AT_ASSERT(m.values_map.at(m2.value()) == o1.value());
|
|
AT_ASSERT(m.nodes_map.at(m1.value()->node()) == o0.value()->node());
|
|
AT_ASSERT(m.nodes_map.at(m2.value()->node()) == o1.value()->node());
|
|
}
|
|
}
|
|
|
|
/**
|
|
* Test diamond pattern:
|
|
*
|
|
* ooo
|
|
* |
|
|
* aaa
|
|
* / \
|
|
* bbb ccc
|
|
* \ /
|
|
* ddd
|
|
* |
|
|
* eee
|
|
*/
|
|
void testDiamond1() {
|
|
Graph graph, pattern1, pattern2;
|
|
script::parseIR(
|
|
R"IR(
|
|
graph(%0):
|
|
%o = o::ooo(%0)
|
|
%a = a::aaa(%o)
|
|
%b = b::bbb(%a)
|
|
%c = c::ccc(%a)
|
|
%d = d::ddd(%b, %c)
|
|
%e = e::eee(%d)
|
|
return (%e))IR",
|
|
&graph);
|
|
|
|
script::parseIR(
|
|
R"IR(
|
|
graph(%0):
|
|
%a = a::aaa(%0)
|
|
%b = b::bbb(%a)
|
|
%c = c::ccc(%a)
|
|
%d = d::ddd(%b, %c)
|
|
return (%d))IR",
|
|
&pattern1);
|
|
AT_ASSERT(!findPatternMatches(pattern1, graph).empty());
|
|
|
|
// Check that order of nodes inside the diamond does not affect the result
|
|
script::parseIR(
|
|
R"IR(
|
|
graph(%0):
|
|
%a = a::aaa(%0)
|
|
%c = c::ccc(%a)
|
|
%b = b::bbb(%a)
|
|
%d = d::ddd(%b, %c)
|
|
return (%d))IR",
|
|
&pattern2);
|
|
AT_ASSERT(!findPatternMatches(pattern2, graph).empty());
|
|
}
|
|
|
|
/**
|
|
* Test diamond pattern:
|
|
*
|
|
* i0
|
|
* |
|
|
* chunk
|
|
* / \
|
|
* os[0] os[1]
|
|
* \ /
|
|
* *
|
|
* |
|
|
* o1
|
|
*/
|
|
void testDiamond2() {
|
|
Graph graph;
|
|
Var i0 = Var::asNewInput(graph);
|
|
auto os = i0.chunk(2, 0);
|
|
auto o1 = os[0] * os[1];
|
|
|
|
Graph pattern;
|
|
auto mi0 = Var::asNewInput(pattern);
|
|
auto ms = mi0.chunk(2, 0);
|
|
auto mo = ms[0] * ms[1];
|
|
pattern.registerOutput(mo.value());
|
|
|
|
auto matches = findPatternMatches(pattern, graph);
|
|
AT_ASSERT(matches.size() == 1);
|
|
for (const Match& m : matches) {
|
|
AT_ASSERT(m.values_map.at(mi0.value()) == i0.value());
|
|
AT_ASSERT(m.values_map.at(ms[0].value()) == os[0].value());
|
|
AT_ASSERT(m.values_map.at(ms[1].value()) == os[1].value());
|
|
AT_ASSERT(m.values_map.at(mo.value()) == o1.value());
|
|
AT_ASSERT(m.nodes_map.at(mo.value()->node()) == o1.value()->node());
|
|
}
|
|
}
|
|
|
|
void testXPattern() {
|
|
Graph graph, pattern;
|
|
script::parseIR(
|
|
R"IR(
|
|
graph(%0, %1):
|
|
%b = b::bbb(%0)
|
|
%c = c::ccc(%1)
|
|
%x = x::xxx(%b, %c)
|
|
%e = e::eee(%x)
|
|
%f = f::fff(%x)
|
|
%g = g::ggg(%e, %f)
|
|
return (%g))IR",
|
|
&graph);
|
|
script::parseIR(
|
|
R"IR(
|
|
graph(%0, %1):
|
|
%b = b::bbb(%0)
|
|
%c = c::ccc(%1)
|
|
%x = x::xxx(%b, %c)
|
|
%e = e::eee(%x)
|
|
%f = f::fff(%x)
|
|
%g = g::ggg(%e, %f)
|
|
return (%g))IR",
|
|
&pattern);
|
|
AT_ASSERT(!findPatternMatches(pattern, graph).empty());
|
|
}
|
|
|
|
void testMultipleMatches() {
|
|
Graph graph, pattern;
|
|
script::parseIR(
|
|
R"IR(
|
|
graph(%t0):
|
|
%t1 = a::aaa(%t0)
|
|
%t2 = a::aaa(%t1)
|
|
%t3 = a::aaa(%t2)
|
|
%t4 = a::aaa(%t3)
|
|
return (%t4))IR",
|
|
&graph);
|
|
script::parseIR(
|
|
R"IR(
|
|
graph(%t0):
|
|
%t1 = a::aaa(%t0)
|
|
return (%t1))IR",
|
|
&pattern);
|
|
auto matches = findPatternMatches(pattern, graph);
|
|
AT_ASSERT(matches.size() == 4);
|
|
}
|
|
|
|
void testOverlappingMatches() {
|
|
Graph graph, pattern;
|
|
script::parseIR(
|
|
R"IR(
|
|
graph(%t0):
|
|
%t1 = a::aaa(%t0)
|
|
%t2 = a::aaa(%t1)
|
|
%t3 = a::aaa(%t2)
|
|
%t4 = a::aaa(%t3)
|
|
return (%t4))IR",
|
|
&graph);
|
|
script::parseIR(
|
|
R"IR(
|
|
graph(%t0):
|
|
%t1 = a::aaa(%t0)
|
|
%t2 = a::aaa(%t1)
|
|
return (%t2))IR",
|
|
&pattern);
|
|
auto matches = findPatternMatches(pattern, graph);
|
|
AT_ASSERT(matches.size() == 3);
|
|
}
|
|
|
|
void testMatchInBasicBlocks1() {
|
|
Graph graph;
|
|
script::parseIR(
|
|
R"IR(
|
|
graph(%a, %b, %c):
|
|
%d = aten::mul(%a, %b)
|
|
%x = prim::If(%c)
|
|
block0():
|
|
%x1 = aten::mul(%a, %d)
|
|
-> (%x1)
|
|
block1():
|
|
%x2 = aten::mul(%b, %d)
|
|
-> (%x2)
|
|
return (%x))IR",
|
|
&graph);
|
|
|
|
// Ensure the matches don't cross basic block boundaries
|
|
Graph pattern0;
|
|
script::parseIR(
|
|
R"IR(
|
|
graph(%x, %y):
|
|
%z = aten::mul(%x, %y)
|
|
return (%z))IR",
|
|
&pattern0);
|
|
AT_ASSERT(findPatternMatches(pattern0, graph).size() == 3);
|
|
|
|
Graph pattern1;
|
|
script::parseIR(
|
|
R"IR(
|
|
graph(%x, %y):
|
|
%z1 = aten::mul(%x, %y)
|
|
%z2 = aten::mul(%y, %z1)
|
|
return (%z2))IR",
|
|
&pattern1);
|
|
AT_ASSERT(findPatternMatches(pattern1, graph).size() == 0);
|
|
}
|
|
|
|
void testMatchInBasicBlocks2() {
|
|
Graph graph;
|
|
script::parseIR(
|
|
R"IR(
|
|
graph(%a, %b):
|
|
%x = my::mul(%a, %b)
|
|
%y = my::node_with_subblock()
|
|
block0():
|
|
%z = my::mul(%b, %x)
|
|
-> (%z)
|
|
return (%y))IR",
|
|
&graph);
|
|
|
|
// Check that we can match both mul ops
|
|
Graph pattern0;
|
|
script::parseIR(
|
|
R"IR(
|
|
graph(%x, %y):
|
|
%z = my::mul(%x, %y)
|
|
return (%z))IR",
|
|
&pattern0);
|
|
AT_ASSERT(findPatternMatches(pattern0, graph).size() == 2);
|
|
|
|
// Ensure the matches don't cross basic block boundaries
|
|
Graph pattern1;
|
|
script::parseIR(
|
|
R"IR(
|
|
graph(%x, %y):
|
|
%u = my::mul(%x, %y)
|
|
%v = my::mul(%y, %u)
|
|
return (%v))IR",
|
|
&pattern1);
|
|
AT_ASSERT(findPatternMatches(pattern1, graph).size() == 0);
|
|
}
|
|
|
|
void testAttributes() {
|
|
Graph graph;
|
|
script::parseIR(
|
|
R"IR(
|
|
graph(%0):
|
|
%a = a::a[isattr=[1,2]](%0)
|
|
%b = a::b[intattr=10, floatattr=3.14](%0)
|
|
%c = a::c[myattr="qqq"](%a, %b)
|
|
return (%c))IR",
|
|
&graph);
|
|
|
|
{
|
|
Graph pattern;
|
|
script::parseIR(
|
|
R"IR(
|
|
graph(%a, %b):
|
|
%c = a::c[myattr="qqq"](%a, %b)
|
|
return (%c))IR",
|
|
&pattern);
|
|
AT_ASSERT(!findPatternMatches(pattern, graph).empty());
|
|
}
|
|
{
|
|
Graph pattern;
|
|
script::parseIR(
|
|
R"IR(
|
|
graph(%a, %b):
|
|
%c = a::c[myattr="zzz"](%a, %b)
|
|
return (%c))IR",
|
|
&pattern);
|
|
AT_ASSERT(findPatternMatches(pattern, graph).empty());
|
|
}
|
|
{
|
|
Graph pattern;
|
|
script::parseIR(
|
|
R"IR(
|
|
graph(%0):
|
|
%b = a::b[extraattr=10](%0)
|
|
return (%b))IR",
|
|
&pattern);
|
|
AT_ASSERT(findPatternMatches(pattern, graph).empty());
|
|
}
|
|
{
|
|
Graph pattern;
|
|
script::parseIR(
|
|
R"IR(
|
|
graph(%0):
|
|
%b = a::b[intattr=10, floatattr=3.14](%0)
|
|
return (%b))IR",
|
|
&pattern);
|
|
AT_ASSERT(!findPatternMatches(pattern, graph).empty());
|
|
}
|
|
{
|
|
Graph pattern;
|
|
script::parseIR(
|
|
R"IR(
|
|
graph(%0):
|
|
%b = a::b[intattr=10, floatattr=3.14, strattr="rrr"](%0)
|
|
return (%b))IR",
|
|
&pattern);
|
|
AT_ASSERT(findPatternMatches(pattern, graph).empty());
|
|
}
|
|
{
|
|
Graph pattern;
|
|
script::parseIR(
|
|
R"IR(
|
|
graph(%0):
|
|
%a = a::a[isattr=[1,2]](%0)
|
|
return (%a))IR",
|
|
&pattern);
|
|
// Lists are not supported yet, thus we shouldn't match for now.
|
|
AT_ASSERT(findPatternMatches(pattern, graph).empty());
|
|
}
|
|
}
|
|
|
|
void testBadPattern() {
|
|
Graph graph, pattern1, pattern2;
|
|
script::parseIR(
|
|
R"IR(
|
|
graph(%0):
|
|
%a = a::aaa(%0)
|
|
return (%a))IR",
|
|
&graph);
|
|
|
|
script::parseIR(
|
|
R"IR(
|
|
graph(%x):
|
|
%y = my::node_with_subblock()
|
|
block0():
|
|
%z = my::op(%x)
|
|
-> (%z)
|
|
return (%y))IR",
|
|
&pattern1);
|
|
ASSERT_ANY_THROW(findPatternMatches(pattern1, graph));
|
|
|
|
script::parseIR(
|
|
R"IR(
|
|
graph(%x):
|
|
%y = my::op1(%x)
|
|
%z = my::op2(%x)
|
|
return (%y, %z))IR",
|
|
&pattern2);
|
|
ASSERT_ANY_THROW(findPatternMatches(pattern2, graph));
|
|
}
|
|
|
|
void testSubgraphMatching() {
|
|
testTrivial1();
|
|
testTrivial2();
|
|
testTrivial3();
|
|
testTrivial4();
|
|
testLinear1();
|
|
testLinear2();
|
|
testDiamond1();
|
|
testDiamond2();
|
|
testXPattern();
|
|
testMultipleMatches();
|
|
testOverlappingMatches();
|
|
testMatchInBasicBlocks1();
|
|
testMatchInBasicBlocks2();
|
|
testAttributes();
|
|
testBadPattern();
|
|
}
|
|
|
|
} // namespace
|
|
} // namespace jit
|
|
} // namespace torch
|