Files
pytorch/test/cpp/jit/test_constant_pooling.cpp
Michael Suo 6d21d5f0b3 gtest-ify JIT tests, through the letter c (#45249)
Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/45249

Reland of https://github.com/pytorch/pytorch/pull/45055 and
https://github.com/pytorch/pytorch/pull/45020

See https://github.com/pytorch/pytorch/pull/45018 for context.

Test Plan: Imported from OSS

Reviewed By: jamesr66a

Differential Revision: D23892645

Pulled By: suo

fbshipit-source-id: e7fe58d5e1a5a0c44f4e2aec9694145afabde0fd
2020-09-24 00:21:20 -07:00

93 lines
2.5 KiB
C++

#include <gtest/gtest.h>
#include <torch/csrc/jit/ir/ir.h>
#include <torch/csrc/jit/ir/irparser.h>
#include <torch/csrc/jit/passes/constant_pooling.h>
#include <torch/csrc/jit/passes/constant_propagation.h>
#include <torch/csrc/jit/testing/file_check.h>
#include <sstream>
#include <string>
namespace torch {
namespace jit {
TEST(ConstantPoolingTest, Int) {
auto graph = std::make_shared<Graph>();
parseIR(
R"IR(
graph():
%8 : int = prim::Constant[value=1]()
%10 : int = prim::Constant[value=1]()
return (%8, %10)
)IR",
&*graph);
ConstantPooling(graph);
testing::FileCheck()
.check_count("prim::Constant", 1, /*exactly*/ true)
->run(*graph);
}
TEST(ConstantPoolingTest, PoolingAcrossBlocks) {
auto graph = std::make_shared<Graph>();
parseIR(
R"IR(
graph(%cond : Tensor):
%a : str = prim::Constant[value="bcd"]()
%3 : bool = aten::Bool(%cond)
%b : str = prim::If(%3)
block0():
%b.1 : str = prim::Constant[value="abc"]()
-> (%b.1)
block1():
%b.2 : str = prim::Constant[value="abc"]()
-> (%b.2)
%7 : (str, str) = prim::TupleConstruct(%a, %b)
return (%7)
)IR",
&*graph);
ConstantPooling(graph);
testing::FileCheck()
.check_count("prim::Constant[value=\"abc\"]", 1, /*exactly*/ true)
->check_count("prim::Constant[value=\"bcd\"]", 1, /*exactly*/ true)
->run(*graph);
}
TEST(ConstantPoolingTest, PoolingDifferentDevices) {
auto graph = std::make_shared<Graph>();
parseIR(
R"IR(
graph():
%2 : int = prim::Constant[value=2]()
%1 : int = prim::Constant[value=1]()
%5 : int? = prim::Constant()
%7 : Device? = prim::Constant()
%15: bool = prim::Constant[value=0]()
%10 : int = prim::Constant[value=6]()
%3 : int[] = prim::ListConstruct(%1, %2)
%x : Tensor = aten::tensor(%3, %5, %7, %15)
%y : Tensor = aten::tensor(%3, %10, %7, %15)
%9 : int[] = prim::ListConstruct(%1, %2)
%z : Tensor = aten::tensor(%9, %10, %7, %15)
prim::Print(%x, %y, %z)
return (%1)
)IR",
&*graph);
// three tensors created - two different devices among the three
// don't have good support for parsing tensor constants
ConstantPropagation(graph);
ConstantPooling(graph);
testing::FileCheck()
.check_count(
"Float(2:1, requires_grad=0, device=cpu) = prim::Constant",
1,
/*exactly*/ true)
->check_count(
"Long(2:1, requires_grad=0, device=cpu) = prim::Constant",
1,
/*exactly*/ true)
->run(*graph);
}
} // namespace jit
} // namespace torch