[JIT] Use Type Level Granularity in Alias Analysis Wildcards (#32251)

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/32251

Previously wildcard sets were associated by TypeKind, meaning all Lists were in one alias set, all Classes were in one alias set, etc. We can improve analysis by bucketing wildcard sets by TypePtr instead. Any two mutable types which can unify should be in the same wildcard set bucket.

This also allows us do much simpler `mayContainAlias` analysis, and also improves `analyzeConservative` analysis because now we can recurse through all contained memory locations and mark writes, instead of just recursing only level deep in contained elements.

Test Plan: Imported from OSS

Differential Revision: D19563263

Pulled By: eellison

fbshipit-source-id: 371a37d1a8596abc6c53f41c09840b6c140ea362
This commit is contained in:
Elias Ellison
2020-01-28 18:05:52 -08:00
committed by Facebook Github Bot
parent 02f055ffd9
commit 25d33a2ee8
9 changed files with 282 additions and 88 deletions

View File

@ -541,9 +541,10 @@ void testWriteTracking() {
void testContainerAliasing() {
{
auto graph = std::make_shared<Graph>();
std::unordered_map<std::string, Value*> vmap;
script::parseIR(
R"IR(
graph():
graph(%inp: Tensor[]):
%x : str = prim::Constant[value="a"]()
%y : Tensor = prim::Constant()
%a : (Tensor) = prim::TupleConstruct(%y)
@ -551,19 +552,22 @@ void testContainerAliasing() {
%c : Tensor[] = prim::ListConstruct(%y)
return (%a, %b, %c)
)IR",
&*graph);
&*graph,
vmap);
auto node_iter = graph->block()->nodes().begin();
auto str_node = node_iter++; // string
Node* ten_node = *node_iter++;
auto str_output = vmap["x"];
auto ten_output = vmap["y"];
AliasDb aliasDb(graph);
AT_ASSERT(graph->outputs().size() == 3);
for (auto out : graph->outputs()) {
AT_ASSERT(aliasDb.mayContainAlias(ten_node->output(), out));
AT_ASSERT(aliasDb.mayContainAlias(ten_output, out));
}
AT_ASSERT(aliasDb.mayContainAlias({ten_node->output()}, graph->outputs()));
AT_ASSERT(!aliasDb.mayContainAlias(str_node->output(), graph->outputs()));
AT_ASSERT(aliasDb.mayContainAlias(ten_output, graph->inputs()));
AT_ASSERT(aliasDb.mayContainAlias({ten_output}, graph->outputs()));
AT_ASSERT(!aliasDb.mayContainAlias(str_output, graph->outputs()));
}
{
@ -940,6 +944,90 @@ void testWildcards() {
wildcardWrite, std::unordered_set<const Value*>{a}));
ASSERT_TRUE(aliasDb.hasWriters(wildcard->node()));
}
// test that wildcards are correctly divided by type
{
auto graph = std::make_shared<Graph>();
std::unordered_map<std::string, Value*> vmap;
script::parseIR(
R"IR(
graph(%ten_list : Tensor[], %int_list : int[], %opt_ten_list : Tensor[]?):
%ten : Tensor = prim::Constant()
%4 : Tensor[] = aten::append(%ten_list, %ten)
%ten_ten_list : Tensor[][] = prim::Constant()
%int_int_list : int[][] = prim::Constant()
return ()
)IR",
&*graph,
vmap);
AliasDb aliasDb(graph);
auto opt_ten_list = vmap["opt_ten_list"];
auto ten_list = vmap["ten_list"];
auto int_list = vmap["int_list"];
AT_ASSERT(!aliasDb.hasWriters(int_list));
AT_ASSERT(aliasDb.hasWriters(opt_ten_list));
AT_ASSERT(aliasDb.hasWriters(ten_list));
AT_ASSERT(aliasDb.mayContainAlias(ten_list, opt_ten_list));
AT_ASSERT(aliasDb.mayAlias(ten_list, opt_ten_list));
auto list_of_tensor_lists = vmap["ten_ten_list"];
AT_ASSERT(aliasDb.mayContainAlias(ten_list, list_of_tensor_lists));
AT_ASSERT(aliasDb.mayContainAlias(ten_list, vmap["ten"]));
}
// test invariant container aliasing
// the containers of different type cannot alias each other,
// however they may contain elements which alias each other
{
auto graph = std::make_shared<Graph>();
std::unordered_map<std::string, Value*> vmap;
script::parseIR(
R"IR(
graph(%ten_list : Tensor[], %ten_opt_list : Tensor?[]):
%ten : Tensor = prim::Constant()
%4 : Tensor[] = aten::append(%ten_list, %ten)
return ()
)IR",
&*graph,
vmap);
AliasDb aliasDb(graph);
auto ten_opt_list = vmap["ten_opt_list"];
auto ten_list = vmap["ten_list"];
AT_ASSERT(!aliasDb.hasWriters(ten_opt_list));
AT_ASSERT(aliasDb.hasWriters(ten_list));
AT_ASSERT(aliasDb.mayContainAlias(ten_list, ten_opt_list));
AT_ASSERT(!aliasDb.mayAlias(ten_list, ten_opt_list));
}
{
auto graph = std::make_shared<Graph>();
std::unordered_map<std::string, Value*> vmap;
script::parseIR(
R"IR(
graph(%float_3D : Float(*, *, *), %float_2D : Float(*, *)):
return ()
)IR",
&*graph,
vmap);
AliasDb aliasDb(graph);
AT_ASSERT(aliasDb.mayAlias(vmap["float_3D"], vmap["float_2D"]));
}
{
auto graph = std::make_shared<Graph>();
std::unordered_map<std::string, Value*> vmap;
script::parseIR(
R"IR(
graph(%float_3D_list : Float(*, *, *)[], %float_2D_list : Float(*, *)[], %ten: Tensor):
return ()
)IR",
&*graph,
vmap);
AliasDb aliasDb(graph);
AT_ASSERT(aliasDb.mayAlias(vmap["float_3D_list"], vmap["float_2D_list"]));
AT_ASSERT(aliasDb.mayContainAlias(vmap["float_3D_list"], vmap["ten"]));
AT_ASSERT(aliasDb.mayContainAlias(vmap["float_2D_list"], vmap["ten"]));
}
}
void testMemoryDAG() {