Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

JIT: Elmiminate SumToSize by using Optional Lists #18697

Open
wants to merge 26 commits into
base: master
from
Open
Changes from 14 commits
Commits
Show all changes
26 commits
Select commit Hold shift + click to select a range
4ecb33a
Specialize Optional (Tensor) to None when executing graph
t-vi Mar 24, 2019
431edca
don't unpack NoneType in _unwrap_optional
t-vi Mar 24, 2019
f557aeb
fix worst stuff
t-vi Mar 24, 2019
c032d5a
Merge branch 'master' into optional_None
t-vi Mar 25, 2019
60a3c93
Merge branch 'master' into optional_None
t-vi Apr 1, 2019
131739f
[WIP] JIT: Elmiminate SumToSize by using Optional Lists
t-vi Apr 1, 2019
bf3de65
no dump()
t-vi Apr 1, 2019
3b4acf9
Fixes
t-vi Apr 2, 2019
e18cf4f
incorporate some feedback from eellison, thanks
t-vi Apr 2, 2019
6650d55
Merge branch 'master' into optional_None
t-vi Apr 3, 2019
9cd8ac9
clean up specialize_autogradzero
t-vi Apr 3, 2019
c43e367
add test for optional list
t-vi Apr 4, 2019
00edd17
Merge branch 'optional_None' into optional_list_sts
t-vi Apr 4, 2019
f2bc16a
change size_if_not_equal to take sizes
t-vi Apr 4, 2019
c21b407
this==other is much more interesting then this==this and works better…
t-vi Apr 5, 2019
9a808fa
Merge branch 'master' into optional_list_sts
t-vi Apr 5, 2019
83b5a74
fixes
t-vi Apr 5, 2019
fb45987
Merge branch 'master' into optional_list_sts
t-vi Apr 6, 2019
7f3f1c6
Merge branch 'master' into optional_list_sts
t-vi May 7, 2019
b518e61
Merge remote-tracking branch 'origin/master' into HEAD
May 9, 2019
2042e22
clang format my changes
t-vi May 9, 2019
b7000b0
update TestFuser
t-vi May 9, 2019
5967ad0
Merge branch 'master' into optional_list_sts
t-vi May 9, 2019
2e6e165
add test
t-vi May 9, 2019
fbb2a7f
only one forward multiple backwards to sharpen test
t-vi May 10, 2019
13a6932
formatting typo
t-vi May 10, 2019
File filter...
Filter file types
Jump to…
Jump to file or symbol
Failed to load files and symbols.
+260 −69
Diff settings

Always

Just for now

@@ -76,6 +76,7 @@ namespace c10 {
_(prim, max) \
_(prim, rangelist) \
_(aten, _grad_sum_to_size) \
_(aten, _size_if_not_equal) \
_(aten, _ncf_unsqueeze) \
_(aten, warn) \
_(aten, floordiv) \
@@ -4655,6 +4655,73 @@ def named_var_and(x, y):
if y is not None and x_none:
print(x + y) # noqa: T484

def test_optional_tensor(self):
@torch.jit.script
def fn(x, y):
# type: (Optional[Tensor], int) -> int
if x is None:
return y
else:
return 0

res = fn(None, 1)
self.assertEqual(res, 1)
g = fn.graph_for(None, 1)
self.assertEqual(next(g.inputs()).type().kind(), 'NoneType')
t = torch.ones(1)
res = fn(t, 1)
self.assertEqual(res, 0)
g = fn.graph_for(t, 1)
self.assertEqual(next(g.inputs()).type().kind(), 'DimensionedTensorType')

@torch.jit.script
def fn(x, y):
# type: (Optional[Tensor], float) -> Tensor
y = torch.jit._unwrap_optional(x) * y
return y

res = fn(t, 2.0)
self.assertEqual(res, t * 2.0)
with self.assertRaisesRegex(RuntimeError, "Unwrapping null optional"):
res = fn(None, 2.0)
g = fn.graph_for(None, 2.0)
self.assertEqual(next(g.nodes()).output().type().str(), "Tensor")

def test_optional_list(self):
@torch.jit.script
def fn(x, y):
# type: (Optional[List[int]], int) -> int
if x is None:
return y
else:
res = 0
for d in x:
res += d
return res

res = fn(None, 1)
self.assertEqual(res, 1)
g = fn.graph_for(None, 1)
self.assertEqual(next(g.inputs()).type().kind(), 'NoneType')
l = [2, 3]
res = fn(l, 1)
self.assertEqual(res, 5)
g = fn.graph_for(l, 1)
self.assertEqual(next(g.inputs()).type().kind(), 'ListType')

@torch.jit.script
def fn(x, y):
# type: (Optional[List[int]], int) -> List[int]
l = torch.jit._unwrap_optional(x)
return l[:y]

res = fn(l, 1)
self.assertEqual(res, l[:1])
with self.assertRaisesRegex(RuntimeError, "Unwrapping null optional"):
res = fn(None, 1)
g = fn.graph_for(None, 1)
self.assertEqual(list(g.nodes())[-1].output().type().str(), "int[]")

def test_while_write_outer_then_read(self):
def func(a, b):
while bool(a < 10):
@@ -10,11 +10,11 @@ void ArgumentSpecCreator::scan(
const WrittenSlots& written_slots) {
auto finishAggregate = [&](size_t pos) {
// it is possible after all the work we did to scan this aggregate,
// we found no tensors to specialize. In this case, just generate
// a skip for the whole aggregate.
// we found no tensors or optionals to specialize. In this case, just
// generate a skip for the whole aggregate.
bool any_spec = std::any_of(
instructions_.begin() + pos, instructions_.end(), [](Inst i) {
return i == SPECIALIZE_TENSOR;
return i == SPECIALIZE_TENSOR || i == SPECIALIZE_NONTENSOR_OPTIONAL;
});
if (!any_spec) {
instructions_[pos] = SKIP;
@@ -28,9 +28,14 @@ void ArgumentSpecCreator::scan(
if (depth >= DEPTH_LIMIT) {
instructions_.emplace_back(SKIP);
}
if (typ->isSubtypeOf(TensorType::get())) {
num_tensors_++;
if (typ->isSubtypeOf(TensorType::get()) || typ->isSubtypeOf(OptionalType::ofTensor())) {
num_tensors_or_optionals_++;
instructions_.emplace_back(SPECIALIZE_TENSOR);
} else if (typ->kind() == TypeKind::OptionalType) {
// note that Optional[Tuple] or Optional[Class] will just register
// as optional (previously they didn't at all, so it's not a regression).
num_tensors_or_optionals_++;
instructions_.emplace_back(SPECIALIZE_NONTENSOR_OPTIONAL);
} else if (auto tup = typ->cast<TupleType>()) {
size_t pos = instructions_.size();
instructions_.emplace_back(ENTER_TUPLE);
@@ -106,14 +111,17 @@ void ArgumentSpecCreator::dump() const {
case SPECIALIZE_TENSOR:
std::cout << "SpecializeTensor ";
break;
case SPECIALIZE_NONTENSOR_OPTIONAL:
std::cout << "SpecializeNonTensorOptional ";
break;
}
}
std::cout << "\n";
}

ArgumentSpec ArgumentSpecCreator::create(bool with_grad, const Stack& input)
const {
ArgumentSpec spec(num_tensors_);
ArgumentSpec spec(num_tensors_or_optionals_);
const IValue* stack[DEPTH_LIMIT]; // The stack of IValue lists
// The stack gets initialized with the input list
stack[0] = last(input, num_inputs_).begin();
@@ -124,6 +132,10 @@ ArgumentSpec ArgumentSpecCreator::create(bool with_grad, const Stack& input)
// consume a tensor and add to the argspec
spec.addTensor(*stack[stack_top]++, with_grad);
break;
case SPECIALIZE_NONTENSOR_OPTIONAL:
// consume a tensor and add to the argspec
spec.addNontensorOptional(*stack[stack_top]++);
break;
case ENTER_TUPLE: {
// consume tuple
const IValue* iv = stack[stack_top]++;
@@ -176,7 +188,9 @@ std::vector<TypePtr> ArgumentSpecCreator::getSpecializedTypes(
case SPECIALIZE_TENSOR: {
input_stack.back()++;
auto& arg = spec.at(arg_spec_offset++);
if (!arg.defined()) {
if (arg.isNone()) {
result_stack.back().emplace_back(NoneType::get());
} else if (!arg.defined()) {
result_stack.back().emplace_back(AutogradZeroTensorType::get());
} else {
result_stack.back().emplace_back(DimensionedTensorType::create(
@@ -186,6 +200,15 @@ std::vector<TypePtr> ArgumentSpecCreator::getSpecializedTypes(
arg.requires_grad()));
}
} break;
case SPECIALIZE_NONTENSOR_OPTIONAL: {
auto ot = (*input_stack.back()++)->expect<OptionalType>();
auto& arg = spec.at(arg_spec_offset++);
if (arg.isNone()) {
result_stack.back().emplace_back(NoneType::get());
} else {
result_stack.back().emplace_back(ot->getElementType());
}
} break;
case ENTER_TUPLE: {
auto tup = (*input_stack.back()++)->expect<TupleType>();
input_stack.emplace_back(tup->elements().data());
@@ -25,6 +25,9 @@ struct ArgumentInfo {
bool defined() const {
return defined_;
}
bool isNone() const {
return is_none_;
}
int device() const {
return device_;
}
@@ -47,9 +50,15 @@ struct ArgumentInfo {
}

private:
unsigned is_nontensor_optional_ : 1;
// =0 if it is an Optional[Tensor] or Tensor. =1 other optional type.
// if this is set, values below is_none_ must be 0
// Non-optional other types are not considered
unsigned is_none_ : 1;
// is_none_ =1 if an Optional[T] is None, remaining bits undefined
unsigned defined_ : 1;
unsigned requires_grad_ : 1;
unsigned : 5;
unsigned : 4;
unsigned dim_ : 8;
int device_ : 8; // NOTE: this needs to be signed because we use -1 to
// represent CPU
@@ -69,25 +78,40 @@ struct ArgumentSpec {
args.reserve(num_flat_inputs);
}

void addNontensorOptional(const IValue& input) {
args.emplace_back();
auto& arg = args.back();
// Initialize all fields to 0. This is convenient, because e.g.
// requires_grad() can be checked even on tensors AND will make
// padding bits all 0s.
std::memset(&arg, 0, sizeof(ArgumentInfo));
arg.is_nontensor_optional_ = true;
arg.is_none_ = input.isNone();
}

void addTensor(const IValue& input, bool with_grad) {
AT_ASSERT(input.isTensor());
args.emplace_back();
auto& arg = args.back();
// Initialize all fields to 0. This is convenient, because e.g.
// requires_grad() can be checked even on tensors AND will make
// padding bits all 0s.
std::memset(&arg, 0, sizeof(ArgumentInfo));

// [argspec refcounting] reinterpret the IValue to avoid having to refcount
// the Tensor microbenchmarks
// https://github.com/zdevito/pytorch/commit/21e7200a0a0fc456bea2f10e95b1781f83933d10
// show overhead in extra refcounting along this path
const at::Tensor* t = reinterpret_cast<const at::Tensor*>(&input);
if ((arg.defined_ = t->defined())) {
arg.requires_grad_ = with_grad && autograd::Variable(*t).requires_grad();
arg.dim_ = t->dim();
arg.device_ = t->is_cuda() ? t->get_device() : -1;
arg.type_ = static_cast<unsigned>(t->scalar_type());
if (input.isNone()) {
arg.is_none_ = 1;
} else {
AT_ASSERT(input.isTensor());
// [argspec refcounting] reinterpret the IValue to avoid having to refcount
// the Tensor microbenchmarks
// https://github.com/zdevito/pytorch/commit/21e7200a0a0fc456bea2f10e95b1781f83933d10
// show overhead in extra refcounting along this path
const at::Tensor* t = reinterpret_cast<const at::Tensor*>(&input);
if ((arg.defined_ = t->defined())) {
arg.requires_grad_ = with_grad && autograd::Variable(*t).requires_grad();
arg.dim_ = t->dim();
arg.device_ = t->is_cuda() ? t->get_device() : -1;
arg.type_ = static_cast<unsigned>(t->scalar_type());
}
}
combineHash(arg);
}
@@ -144,8 +168,11 @@ struct ArgumentSpecCreator {
ENTER_OBJECT, // same as ENTER_TUPLE, but the input is a class
LEAVE, // pop the top-most list from the stack
SKIP, // consume an element from the top-most list, and discard
SPECIALIZE_TENSOR, // consume a tensor for the top-most list, and
// add it to the ArgSpec key being created
SPECIALIZE_TENSOR, // consume a tensor or optional tensor for the top-most
// list, and add it to the ArgSpec key being created
SPECIALIZE_NONTENSOR_OPTIONAL,
// consume a nontensor optional from the top-most list,
// and add it to the ArgSpec key being created
};
ArgumentSpecCreator(Graph& graph);
ArgumentSpec create(bool with_grad, const Stack& stack) const;
@@ -163,7 +190,7 @@ struct ArgumentSpecCreator {
size_t depth,
const WrittenSlots& written_slots);
size_t num_inputs_;
size_t num_tensors_ = 0;
size_t num_tensors_or_optionals_ = 0;
std::vector<Inst> instructions_;
};

@@ -176,11 +203,12 @@ struct ArgumentSpecCreator {
// API users should use ArgumentInfo
struct CompleteArgumentInfoPOD {
// total size is 64-bit
unsigned is_tensor : 8; // all other fields are invalid if this is false
unsigned is_tensor : 8; // all other fields except is_none are invalid if this is false
unsigned type : 8; // scalar type
unsigned is_none : 1; // for optionals
unsigned defined : 1;
unsigned requires_grad : 1;
signed device : 14;
signed device : 13;
uint32_t total_dims; // all TensorInfoPODs are in CompleteArgumentSpec's
// tensor_info() array. total_dims is the total number of
// dimensions seen so far in all previous members of
Oops, something went wrong.
ProTip! Use n and p to navigate between commits in a pull request.
You can’t perform that action at this time.