frozenleaves/pytorch
1
0
Fork 0
mirror of https://github.com/pytorch/pytorch.git synced 2025-10-20 12:54:11 +08:00
Code Issues Packages Projects Releases Wiki Activity
Files
main
pytorch/test/cpp/jit/test_schema_info.cpp
goldenxuett 8d5951e7e8 [JIT] Add is_aliasing method to FunctionSchema (#82255) 
- Add is_aliasing method in function schema to be able to indicate if an argument has an alias_set attached to it. This is utilized in the integration with autograd (see next PR)
- Tested in test_schema_info
Pull Request resolved: https://github.com/pytorch/pytorch/pull/82255
Approved by: https://github.com/davidberard98
2022-07-27 20:19:21 +00:00

395 lines
17 KiB
C++
Raw Permalink Blame History

#include <gtest/gtest.h>
#include <torch/csrc/autograd/generated/variable_factories.h>
#include <torch/csrc/utils/schema_info.h>
namespace torch {
namespace utils {
using c10::SchemaArgType;
TEST(FunctionSchemaIsAliasingTest, Basic) {
c10::FunctionSchema schema = torch::jit::parseSchema(
"aten::test.Tensor(Tensor(a) self, Tensor(b!) other, Tensor more_other) -> (Tensor(a), Tensor(b!))");
ASSERT_TRUE(schema.is_aliasing({SchemaArgType::output, 0}));
ASSERT_TRUE(schema.is_aliasing({SchemaArgType::output, 1}));
ASSERT_TRUE(schema.is_aliasing({SchemaArgType::input, 0}));
ASSERT_TRUE(schema.is_aliasing({SchemaArgType::input, 1}));
ASSERT_FALSE(schema.is_aliasing({SchemaArgType::input, 2}));
}
TEST(FunctionSchemaIsAliasingTest, InvalidArgument) {
c10::FunctionSchema schema = torch::jit::parseSchema(
"aten::sub_.Tensor(Tensor(a!) self, Tensor other, *, Scalar alpha=1) -> (Tensor(a!))");
ASSERT_THROW(schema.is_aliasing({SchemaArgType::input, 4}), c10::Error);
ASSERT_THROW(schema.is_aliasing({SchemaArgType::output, 4}), c10::Error);
}
TEST(FunctionSchemaIsMutableTest, Basic) {
c10::FunctionSchema schema = torch::jit::parseSchema(
"aten::sub_.Tensor(Tensor(a!) self, Tensor other, *, Scalar alpha=1) -> (Tensor(a!))");
ASSERT_TRUE(schema.is_mutable({SchemaArgType::output, 0}));
ASSERT_TRUE(schema.is_mutable({SchemaArgType::input, 0}));
ASSERT_TRUE(schema.is_mutable("self"));
ASSERT_FALSE(schema.is_mutable({SchemaArgType::input, 1}));
ASSERT_FALSE(schema.is_mutable("other"));
ASSERT_FALSE(schema.is_mutable({SchemaArgType::input, 2}));
ASSERT_FALSE(schema.is_mutable("alpha"));
}
TEST(FunctionSchemaIsMutableTest, InvalidArgument) {
c10::FunctionSchema schema = torch::jit::parseSchema(
"aten::sub_.Tensor(Tensor(a!) self, Tensor other, *, Scalar alpha=1) -> (Tensor(a!))");
ASSERT_THROW(schema.is_mutable({SchemaArgType::input, 4}), c10::Error);
ASSERT_THROW(schema.is_mutable({SchemaArgType::output, 4}), c10::Error);
ASSERT_THROW(schema.is_mutable("named_argument"), c10::Error);
}
TEST(SchemaInfoIsMutableTest, Basic) {
SchemaInfo schema(
"aten::sub_.Tensor(Tensor(a!) self, Tensor other, *, Scalar alpha=1) -> (Tensor(a!))");
ASSERT_TRUE(schema.is_mutable({SchemaArgType::input, 0}));
ASSERT_TRUE(schema.is_mutable("self"));
ASSERT_FALSE(schema.is_mutable({SchemaArgType::input, 1}));
ASSERT_FALSE(schema.is_mutable("other"));
ASSERT_FALSE(schema.is_mutable({SchemaArgType::input, 2}));
ASSERT_FALSE(schema.is_mutable("alpha"));
}
TEST(SchemaInfoIsMutableTest, InvalidArgument) {
SchemaInfo schema(
"aten::sub_.Tensor(Tensor(a!) self, Tensor other, *, Scalar alpha=1) -> (Tensor(a!))");
ASSERT_THROW(schema.is_mutable({SchemaArgType::input, 4}), c10::Error);
ASSERT_THROW(schema.is_mutable("named_argument"), c10::Error);
}
TEST(SchemaInfoIsMutableTest, AliasingInputs) {
SchemaInfo schema(
"aten::test.Tensor(Tensor(a!) self, Tensor(b) other, *, Scalar alpha=1) -> (Tensor(a!), Tensor(b))");
ASSERT_TRUE(schema.is_mutable({SchemaArgType::input, 0}));
ASSERT_TRUE(schema.is_mutable({SchemaArgType::output, 0}));
ASSERT_TRUE(schema.is_mutable("self"));
ASSERT_FALSE(schema.is_mutable({SchemaArgType::input, 1}));
ASSERT_FALSE(schema.is_mutable({SchemaArgType::output, 1}));
ASSERT_FALSE(schema.is_mutable("other"));
at::Tensor input = at::randn({3, 3});
schema.addArgumentValue("self", input);
schema.addArgumentValue("other", input);
ASSERT_TRUE(schema.is_mutable({SchemaArgType::input, 1}));
ASSERT_TRUE(schema.is_mutable({SchemaArgType::output, 1}));
ASSERT_TRUE(schema.is_mutable("other"));
}
TEST(SchemaInfoIsMutableTest, InstanceNorm) {
SchemaInfo schema_info(
"aten::instance_norm(Tensor input, Tensor? weight, Tensor? bias, Tensor? running_mean, Tensor? running_var, bool use_input_stats, float momentum, float eps, bool cudnn_enabled) -> Tensor");
ASSERT_TRUE(schema_info.is_mutable("running_mean"));
ASSERT_TRUE(schema_info.is_mutable("running_var"));
schema_info.addArgumentValue("use_input_stats", false);
ASSERT_FALSE(schema_info.is_mutable("running_mean"));
ASSERT_FALSE(schema_info.is_mutable("running_var"));
}
TEST(SchemaInfoIsMutableTest, BatchNorm) {
SchemaInfo schema_info(
"aten::batch_norm(Tensor input, Tensor? weight, Tensor? bias, Tensor? running_mean, Tensor? running_var, bool training, float momentum, float eps, bool cudnn_enabled) -> Tensor");
ASSERT_TRUE(schema_info.is_mutable("running_mean"));
ASSERT_TRUE(schema_info.is_mutable("running_var"));
schema_info.addArgumentValue("training", false);
ASSERT_FALSE(schema_info.is_mutable("running_mean"));
ASSERT_FALSE(schema_info.is_mutable("running_var"));
}
TEST(SchemaInfoIsNonDeterministicTest, Basic) {
SchemaInfo deterministic_schema_info(
"aten::sub_.Tensor(Tensor(a!) self, Tensor other, *, Scalar alpha=1) -> (Tensor(a!))");
SchemaInfo nondeterministic_schema_info(
"aten::bernoulli(Tensor self, *, Generator? generator) -> Tensor");
ASSERT_FALSE(deterministic_schema_info.is_nondeterministic());
ASSERT_TRUE(nondeterministic_schema_info.is_nondeterministic());
}
TEST(SchemaInfoIsNonDeterministicTest, Dropout) {
SchemaInfo droupout_schema_info(
"aten::dropout(Tensor input, float p, bool train) -> Tensor");
ASSERT_TRUE(droupout_schema_info.is_nondeterministic());
droupout_schema_info.addArgumentValue("train", false);
ASSERT_FALSE(droupout_schema_info.is_nondeterministic());
}
TEST(FunctionSchemaMayAliasTest, Basic) {
c10::FunctionSchema schema = torch::jit::parseSchema(
"aten::sub_.Tensor(Tensor(a!) self, Tensor other, *, Scalar alpha=1) -> (Tensor(a!))");
ASSERT_TRUE(
schema.may_alias({SchemaArgType::input, 0}, {SchemaArgType::output, 0}));
ASSERT_FALSE(
schema.may_alias({SchemaArgType::input, 1}, {SchemaArgType::output, 0}));
ASSERT_FALSE(
schema.may_alias({SchemaArgType::input, 1}, {SchemaArgType::input, 0}));
}
TEST(FunctionSchemaMayAliasTest, InvalidArgument) {
c10::FunctionSchema schema = torch::jit::parseSchema(
"aten::sub_.Tensor(Tensor(a!) self, Tensor other, *, Scalar alpha=1) -> (Tensor(a!))");
ASSERT_THROW(
schema.may_alias({SchemaArgType::input, 15}, {SchemaArgType::output, 0}),
c10::Error);
ASSERT_THROW(
schema.may_alias({SchemaArgType::input, 0}, {SchemaArgType::output, 15}),
c10::Error);
}
TEST(FunctionSchemaMayAliasTest, Wildcard) {
c10::FunctionSchema schema = torch::jit::parseSchema(
"aten::test.Tensor(Tensor(*) self) -> (Tensor(*), Tensor)");
ASSERT_TRUE(
schema.may_alias({SchemaArgType::output, 0}, {SchemaArgType::input, 0}));
ASSERT_FALSE(
schema.may_alias({SchemaArgType::output, 1}, {SchemaArgType::input, 0}));
}
TEST(SchemaInfoMayAliasTest, AliasingInputs) {
SchemaInfo schema(
"aten::sub.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor");
ASSERT_FALSE(
schema.may_alias({SchemaArgType::input, 0}, {SchemaArgType::input, 1}));
at::Tensor input = at::randn({3, 3});
schema.addArgumentValue("self", input);
schema.addArgumentValue("other", input);
ASSERT_TRUE(
schema.may_alias({SchemaArgType::input, 0}, {SchemaArgType::input, 1}));
}
TEST(SchemaInfoMayAliasTest, AliasingOutputs) {
SchemaInfo schema(
"aten::aminmax.out(Tensor self, *, int? dim=None, bool keepdim=False, Tensor(a!) min, Tensor(b!) max) -> (Tensor(a!) min, Tensor(b!) max)");
ASSERT_FALSE(
schema.may_alias({SchemaArgType::output, 0}, {SchemaArgType::output, 1}));
at::Tensor input = at::randn({3, 3});
schema.addArgumentValue("min", input);
schema.addArgumentValue("max", input);
ASSERT_TRUE(
schema.may_alias({SchemaArgType::output, 0}, {SchemaArgType::output, 1}));
}
TEST(SchemaInfoMayAliasTest, AliasingInputOutput) {
SchemaInfo schema(
"aten::sub_.Tensor(Tensor(a!) self, Tensor other, *, Scalar alpha=1) -> (Tensor(a!))");
ASSERT_TRUE(
schema.may_alias({SchemaArgType::input, 0}, {SchemaArgType::output, 0}));
ASSERT_FALSE(
schema.may_alias({SchemaArgType::input, 1}, {SchemaArgType::output, 0}));
at::Tensor input = at::randn({3, 3});
schema.addArgumentValue("self", input);
schema.addArgumentValue("other", input);
ASSERT_TRUE(
schema.may_alias({SchemaArgType::input, 0}, {SchemaArgType::output, 0}));
ASSERT_TRUE(
schema.may_alias({SchemaArgType::input, 1}, {SchemaArgType::output, 0}));
}
TEST(SchemaInfoMayAliasTest, MultipleWildcardInputs) {
SchemaInfo schema(
"aten::test.Tensor(Tensor(a) a, Tensor(*) b, Tensor(*) c) -> (Tensor(a), Tensor(*))");
ASSERT_TRUE(
schema.may_alias({SchemaArgType::input, 0}, {SchemaArgType::output, 0}));
ASSERT_TRUE(
schema.may_alias({SchemaArgType::input, 1}, {SchemaArgType::output, 1}));
ASSERT_TRUE(
schema.may_alias({SchemaArgType::input, 2}, {SchemaArgType::output, 1}));
ASSERT_FALSE(
schema.may_alias({SchemaArgType::input, 0}, {SchemaArgType::input, 1}));
ASSERT_FALSE(
schema.may_alias({SchemaArgType::input, 0}, {SchemaArgType::input, 2}));
ASSERT_FALSE(
schema.may_alias({SchemaArgType::input, 0}, {SchemaArgType::output, 1}));
ASSERT_FALSE(
schema.may_alias({SchemaArgType::input, 1}, {SchemaArgType::output, 0}));
at::Tensor input = at::randn({3, 3});
schema.addArgumentValue("a", input);
schema.addArgumentValue("b", input);
ASSERT_TRUE(
schema.may_alias({SchemaArgType::input, 0}, {SchemaArgType::output, 0}));
ASSERT_TRUE(
schema.may_alias({SchemaArgType::input, 1}, {SchemaArgType::output, 1}));
ASSERT_TRUE(
schema.may_alias({SchemaArgType::input, 2}, {SchemaArgType::output, 1}));
ASSERT_TRUE(
schema.may_alias({SchemaArgType::input, 0}, {SchemaArgType::input, 1}));
ASSERT_TRUE(
schema.may_alias({SchemaArgType::input, 0}, {SchemaArgType::input, 2}));
ASSERT_TRUE(
schema.may_alias({SchemaArgType::input, 0}, {SchemaArgType::output, 1}));
ASSERT_TRUE(
schema.may_alias({SchemaArgType::input, 1}, {SchemaArgType::output, 0}));
}
TEST(SchemaInfoMayAliasTest, MultipleNonWildcardInputs) {
SchemaInfo schema(
"aten::test.Tensor(Tensor(a) a, Tensor(a) b, Tensor(*) c, Tensor(b) d) -> (Tensor(a), Tensor(*))");
ASSERT_TRUE(
schema.may_alias({SchemaArgType::input, 0}, {SchemaArgType::input, 1}));
ASSERT_TRUE(
schema.may_alias({SchemaArgType::input, 0}, {SchemaArgType::input, 2}));
ASSERT_TRUE(
schema.may_alias({SchemaArgType::input, 2}, {SchemaArgType::input, 1}));
ASSERT_TRUE(
schema.may_alias({SchemaArgType::input, 2}, {SchemaArgType::output, 0}));
}
TEST(SchemaInfoMayAliasTest, MultipleNonWildcardOutputs) {
SchemaInfo schema(
"aten::test.Tensor(Tensor(a) a, Tensor(*) b) -> (Tensor(a), Tensor(a))");
ASSERT_TRUE(
schema.may_alias({SchemaArgType::input, 0}, {SchemaArgType::input, 1}));
ASSERT_TRUE(
schema.may_alias({SchemaArgType::output, 0}, {SchemaArgType::output, 1}));
ASSERT_TRUE(
schema.may_alias({SchemaArgType::output, 0}, {SchemaArgType::input, 1}));
}
TEST(SchemaInfoMayAliasTest, MismatchingTypes) {
SchemaInfo schema("aten::test.Tensor(Tensor(a) a) -> int(a)");
ASSERT_FALSE(
schema.may_alias({SchemaArgType::input, 0}, {SchemaArgType::output, 0}));
}
TEST(FunctionSchemaMayContainAliasTest, Basic) {
c10::FunctionSchema schema = torch::jit::parseSchema(
"aten::sub_.Tensor(Tensor(a!) self, Tensor other, *, Scalar alpha=1) -> (Tensor(a!))");
ASSERT_TRUE(schema.may_contain_alias(
{SchemaArgType::input, 0}, {SchemaArgType::output, 0}));
ASSERT_FALSE(schema.may_contain_alias(
{SchemaArgType::input, 1}, {SchemaArgType::output, 0}));
ASSERT_FALSE(schema.may_contain_alias(
{SchemaArgType::input, 1}, {SchemaArgType::input, 0}));
}
TEST(FunctionSchemaMayContainAliasTest, Wildcard) {
c10::FunctionSchema schema = torch::jit::parseSchema(
"aten::test.Tensor(Tensor(*) self) -> (Tensor[], Tensor)");
ASSERT_FALSE(
schema.may_alias({SchemaArgType::output, 0}, {SchemaArgType::input, 0}));
ASSERT_TRUE(schema.may_contain_alias(
{SchemaArgType::output, 0}, {SchemaArgType::input, 0}));
ASSERT_TRUE(schema.may_contain_alias(
{SchemaArgType::output, 0}, {SchemaArgType::input, 0}, false));
ASSERT_FALSE(schema.may_contain_alias(
{SchemaArgType::input, 0}, {SchemaArgType::output, 0}, false));
ASSERT_FALSE(
schema.may_alias({SchemaArgType::output, 1}, {SchemaArgType::input, 0}));
}
TEST(FunctionSchemaMayContainAliasTest, InputAndOutputContainers) {
c10::FunctionSchema schema =
torch::jit::parseSchema("aten::test.Tensor(Tensor[] self) -> Tensor[]");
ASSERT_FALSE(
schema.may_alias({SchemaArgType::output, 0}, {SchemaArgType::input, 0}));
ASSERT_TRUE(schema.may_contain_alias(
{SchemaArgType::output, 0}, {SchemaArgType::input, 0}));
ASSERT_TRUE(schema.may_contain_alias(
{SchemaArgType::output, 0}, {SchemaArgType::input, 0}, false));
ASSERT_TRUE(schema.may_contain_alias(
{SchemaArgType::input, 0}, {SchemaArgType::output, 0}, false));
}
TEST(SchemaInfoMayContainAliasTest, ContainAliasInputsEqual) {
SchemaInfo schema(
"aten::sub.Tensor(Tensor self, Tensor other, *, Scalar alpha=1) -> Tensor");
ASSERT_FALSE(schema.may_contain_alias(
{SchemaArgType::input, 0}, {SchemaArgType::input, 1}));
at::Tensor input = at::randn({3, 3});
schema.addArgumentValue("self", input);
schema.addArgumentValue("other", input);
ASSERT_TRUE(schema.may_contain_alias(
{SchemaArgType::input, 0}, {SchemaArgType::input, 1}));
ASSERT_TRUE(schema.may_contain_alias(
{SchemaArgType::input, 0}, {SchemaArgType::input, 1}, false));
ASSERT_TRUE(schema.may_contain_alias(
{SchemaArgType::input, 1}, {SchemaArgType::input, 0}, false));
}
TEST(SchemaInfoMayContainAliasTest, ContainAliasInputsContained) {
SchemaInfo schema(
"aten::test.Tensor(Tensor[] self, Tensor other, *, Scalar alpha=1) -> Tensor");
ASSERT_FALSE(schema.may_contain_alias(
{SchemaArgType::input, 0}, {SchemaArgType::input, 1}));
at::Tensor input = at::randn({3, 3});
schema.addArgumentValue("self", c10::List<at::Tensor>({input}));
schema.addArgumentValue("other", input);
ASSERT_TRUE(schema.may_contain_alias(
{SchemaArgType::input, 0}, {SchemaArgType::input, 1}));
ASSERT_TRUE(schema.may_contain_alias(
{SchemaArgType::input, 0}, {SchemaArgType::input, 1}, false));
ASSERT_FALSE(schema.may_contain_alias(
{SchemaArgType::input, 1}, {SchemaArgType::input, 0}, false));
}
TEST(SchemaInfoMayContainAliasTest, ContainAliasOutputs) {
SchemaInfo schema(
"aten::aminmax.out(Tensor self, *, int? dim=None, bool keepdim=False, Tensor(a!) min, Tensor(b!) max) -> (Tensor(a!) min, Tensor(b!) max)");
ASSERT_FALSE(schema.may_contain_alias(
{SchemaArgType::output, 0}, {SchemaArgType::output, 1}));
at::Tensor input = at::randn({3, 3});
schema.addArgumentValue("min", input);
schema.addArgumentValue("max", input);
ASSERT_TRUE(schema.may_contain_alias(
{SchemaArgType::output, 0}, {SchemaArgType::output, 1}));
}
TEST(SchemaInfoMayContainAliasTest, ContainAliasInputOutput) {
SchemaInfo schema(
"aten::test.tensor(Tensor(a) self, Tensor[] other) -> Tensor(a)");
ASSERT_FALSE(schema.may_contain_alias(
{SchemaArgType::output, 0}, {SchemaArgType::input, 1}));
at::Tensor input = at::randn({3, 3});
schema.addArgumentValue("other", c10::List<at::Tensor>({input}));
schema.addArgumentValue("self", input);
ASSERT_TRUE(schema.may_contain_alias(
{SchemaArgType::output, 0}, {SchemaArgType::input, 1}));
ASSERT_FALSE(schema.may_contain_alias(
{SchemaArgType::output, 0}, {SchemaArgType::input, 1}, false));
ASSERT_TRUE(schema.may_contain_alias(
{SchemaArgType::input, 1}, {SchemaArgType::output, 0}, false));
}
TEST(SchemaInfoMayContainAliasTest, InputAndOutputContainers) {
SchemaInfo schema(
"aten::test.tensor(Tensor self, Tensor[] other) -> Tensor[]");
ASSERT_TRUE(schema.may_contain_alias(
{SchemaArgType::output, 0}, {SchemaArgType::input, 1}));
ASSERT_FALSE(schema.may_contain_alias(
{SchemaArgType::output, 0}, {SchemaArgType::input, 0}));
ASSERT_FALSE(schema.may_contain_alias(
{SchemaArgType::input, 0}, {SchemaArgType::input, 1}));
at::Tensor input = at::randn({3, 3});
schema.addArgumentValue("other", c10::List<at::Tensor>({input}));
schema.addArgumentValue("self", input);
ASSERT_TRUE(schema.may_contain_alias(
{SchemaArgType::output, 0}, {SchemaArgType::input, 1}));
ASSERT_TRUE(schema.may_contain_alias(
{SchemaArgType::output, 0}, {SchemaArgType::input, 0}));
ASSERT_TRUE(schema.may_contain_alias(
{SchemaArgType::input, 0}, {SchemaArgType::input, 1}));
}
TEST(SchemaInfoMayContainAliasTest, Wildcard) {
SchemaInfo schema(
"aten::test.tensor(Tensor a, Tensor[] b, Tensor(*) c) -> Tensor[]");
ASSERT_FALSE(schema.may_contain_alias(
{SchemaArgType::input, 0}, {SchemaArgType::input, 2}));
ASSERT_FALSE(schema.may_contain_alias(
{SchemaArgType::input, 0}, {SchemaArgType::input, 1}));
ASSERT_TRUE(schema.may_contain_alias(
{SchemaArgType::input, 2}, {SchemaArgType::input, 1}));
at::Tensor input = at::randn({3, 3});
schema.addArgumentValue("b", c10::List<at::Tensor>({input}));
schema.addArgumentValue("a", input);
ASSERT_TRUE(schema.may_contain_alias(
{SchemaArgType::input, 0}, {SchemaArgType::input, 2}));
ASSERT_TRUE(schema.may_contain_alias(
{SchemaArgType::input, 0}, {SchemaArgType::input, 1}));
ASSERT_TRUE(schema.may_contain_alias(
{SchemaArgType::input, 2}, {SchemaArgType::input, 1}));
}
} // namespace utils
} // namespace torch
Reference in New Issue View Git Blame Copy Permalink