mirror of
https://github.com/pytorch/pytorch.git
synced 2025-10-27 17:54:55 +08:00
[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:
committed by
Facebook Github Bot
parent
02f055ffd9
commit
25d33a2ee8
@ -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() {
|
||||
|
||||
Reference in New Issue
Block a user