Specialize Optional[T] to T (or subtype for Tensor) or None when executing graph

test/cpp/jit/test.cpp

@@ -80,6 +80,7 @@ namespace jit {
 
 #define TH_FORALL_TESTS_CUDA(_) \
   _(ArgumentSpec)               \
+  _(CompleteArgumentSpec)       \
   _(Fusion)                     \
   _(GraphExecutor)              \
   _(ModuleConversion)           \

test/cpp/jit/test_argument_spec.h

@@ -1,5 +1,6 @@
 #pragma once
 
+#include <torch/jit.h>
 #include "test/cpp/jit/test_utils.h"
 #include "torch/csrc/jit/argument_spec.h"
 
@@ -24,14 +25,21 @@ bool isEqual(const CompleteArgumentInfo& ti, const autograd::Variable& v) {
       isEqual(ti.strides(), v.strides());
 }
 
+bool isEqual(const ArgumentInfo& ti, const autograd::Variable& v) {
+  if (!ti.defined())
+    return ti.defined() == v.defined();
+  return ti.device() == device(v) && ti.requires_grad() == v.requires_grad() &&
+      ti.type() == v.scalar_type() && ti.dim() == v.dim();
+}
+
 autograd::Variable var(at::DeprecatedTypeProperties& t, at::IntArrayRef sizes, bool requires_grad) {
   return autograd::make_variable(at::rand(sizes, t.options()), requires_grad);
 }
 autograd::Variable undef() {
   return autograd::Variable();
 }
 
-void testArgumentSpec() {
+void testCompleteArgumentSpec() {
   auto& CF = at::CPU(at::kFloat);
   auto& CD = at::CPU(at::kDouble);
   auto& GF = at::CUDA(at::kFloat);
@@ -70,7 +78,7 @@ void testArgumentSpec() {
   ASSERT_TRUE(no_grad != a);
 
   std::unordered_set<CompleteArgumentSpec> spec;
-  spec.insert(std::move(a));
+  spec.insert(a); // we use a below, so no move
   ASSERT_TRUE(spec.count(b) > 0);
   ASSERT_EQ(spec.count(no_grad), 0);
   spec.insert(std::move(no_grad));
@@ -86,6 +94,69 @@ void testArgumentSpec() {
   ASSERT_EQ(with_const.at(2).sizes().size(), 2);
 }
 
+void testArgumentSpec() {
+  auto& CF = at::CPU(at::kFloat);
+  auto& CD = at::CPU(at::kDouble);
+  auto& GF = at::CUDA(at::kFloat);
+  auto& GD = at::CUDA(at::kDouble);
+
+  auto graph = jit::compile(R"JIT(
+   def fn(a, b, c, d, e):
+      return a, b, c, d, e
+   )JIT")
+                   ->get_function("fn")
+                   .graph();
+
+  ArgumentSpecCreator arg_spec_creator(*graph);
+
+  auto list = createStack({var(CF, {1}, true),
+                           var(CD, {1, 2}, false),
+                           var(GF, {}, true),
+                           var(GD, {4, 5, 6}, false),
+                           undef()});
+
+  // make sure we have some non-standard strides
+  list[1].toTensor().transpose_(0, 1);
+
+  // same list but different backing values
+  auto list2 = createStack({var(CF, {1}, true),
+                            var(CD, {1, 2}, false),
+                            var(GF, {}, true),
+                            var(GD, {4, 5, 6}, false),
+                            undef()});
+  list2[1].toTensor().transpose_(0, 1);
+
+
+  ArgumentSpec a = arg_spec_creator.create(true, list);
+  ArgumentSpec b = arg_spec_creator.create(true, list);
+  ASSERT_EQ(a.hashCode(), b.hashCode());
+
+  ASSERT_EQ(a, b);
+  ArgumentSpec d = arg_spec_creator.create(true, list2);
+  ASSERT_EQ(d, a);
+  ASSERT_EQ(d.hashCode(), a.hashCode());
+
+  for (size_t i = 0; i < list.size(); ++i) {
+    ASSERT_TRUE(isEqual(a.tensorAt(i), list[i].toTensor()));
+  }
+  ArgumentSpec no_grad = arg_spec_creator.create(/*with_grad=*/false, list);
+  ASSERT_TRUE(no_grad != a);
+
+  std::unordered_set<ArgumentSpec> spec;
+  spec.insert(a); // we still need a for the test below
+  ASSERT_TRUE(spec.count(b) > 0);
+  ASSERT_EQ(spec.count(no_grad), 0);
+  spec.insert(std::move(no_grad));
+  ASSERT_EQ(spec.count(arg_spec_creator.create(true, list)), 1);
+
+  list2[1].toTensor().transpose_(0, 1);
+  ArgumentSpec c = arg_spec_creator.create(
+      true, list2); // same as list, except for one stride, used to be
+                    // different, now the same
+  ASSERT_TRUE(c == a);
+  ASSERT_EQ(spec.count(c), 1);
+}
+
 } // namespace test
 } // namespace jit
 } // namespace torch

test/cpp/jit/test_autodiff.h

@@ -218,7 +218,7 @@ void testDifferentiateWithRequiresGrad() {
       at::empty_strided(2, 2, at::CPU(at::kFloat).options()), false);
 
   ArgumentSpecCreator asc(*graph);
-  asc.setInputTypes(*graph, asc.create(true, {a_var, b_var}));
+  asc.specializeTypes(*graph, asc.create(true, {a_var, b_var}));
 
   PropagateInputShapes(graph);
   PropagateRequiresGrad(graph);

test/cpp/jit/test_misc.h

@@ -835,7 +835,7 @@ void testProfiler() {
   ArgumentSpecCreator arg_spec_creator(opt_graph);
   ArgumentSpec spec =
       arg_spec_creator.create(autograd::GradMode::is_enabled(), stack);
-  arg_spec_creator.setInputTypes(opt_graph, spec);
+  arg_spec_creator.specializeTypes(opt_graph, spec);
   auto pr = ProfilingRecord::instrumentGraph(g);
   Code cd(pr->profiled_graph_);
   InterpreterState is{cd};

test/test_jit.py

@@ -5557,6 +5557,89 @@ 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 = torch.jit.last_executed_optimized_graph()
+        first_input = next(g.inputs())
+        # check if input is disconnected
+        self.assertEqual(first_input.type().kind(), 'OptionalType')
+        self.assertEqual(first_input.uses(), [])
+        t = torch.ones(1)
+        res = fn(t, 1)
+        self.assertEqual(res, 0)
+        g = torch.jit.last_executed_optimized_graph()
+        self.assertEqual(next(g.inputs()).type().kind(), 'DimensionedTensorType')
+
+        @torch.jit.script
+        def fn(x, y, b):
+            # type: (Optional[Tensor], Tensor, bool) -> Tensor
+            if b:
+                res = y
+            else:
+                res = torch.jit._unwrap_optional(x)
+            return res
+
+        t2 = torch.zeros(1)
+        res = fn(t, t2, True)
+        self.assertEqual(res, t2)
+        with self.assertRaisesRegex(RuntimeError, "Unwrapping null optional"):
+            res = fn(None, t2, False)
+        res = fn(None, t2, True)
+        g = torch.jit.last_executed_optimized_graph()
+        self.assertEqual(next(g.outputs()).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 = torch.jit.last_executed_optimized_graph()
+        first_input = next(g.inputs())
+        # check if input is disconnected
+        self.assertEqual(first_input.type().kind(), 'OptionalType')
+        self.assertEqual(first_input.uses(), [])
+        l = [2, 3]
+        res = fn(l, 1)
+        self.assertEqual(res, 5)
+        g = torch.jit.last_executed_optimized_graph()
+        self.assertEqual(next(g.inputs()).type().kind(), 'ListType')
+
+        @torch.jit.script
+        def fn(x, y, b):
+            # type: (Optional[List[int]], List[int], bool) -> List[int]
+            if b:
+                l = torch.jit._unwrap_optional(x)
+            else:
+                l = y
+            return l
+
+        l2 = [0, 1]
+        res = fn(l, l2, True)
+        self.assertEqual(res, l)
+        with self.assertRaisesRegex(RuntimeError, "Unwrapping null optional"):
+            res = fn(None, l2, True)
+        res = fn(None, l2, False)
+        g = torch.jit.last_executed_optimized_graph()
+        self.assertEqual(next(g.outputs()).type().str(), "int[]")
+
     def test_while_write_outer_then_read(self):
         def func(a, b):
             while bool(a < 10):

torch/csrc/jit/argument_spec.cpp

@@ -10,11 +10,12 @@ 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_OPTIONAL ||
+              i == SPECIALIZE_OPTIONAL_TENSOR;
         });
     if (!any_spec) {
       instructions_[pos] = SKIP;
@@ -31,6 +32,15 @@ void ArgumentSpecCreator::scan(
   if (typ->isSubtypeOf(TensorType::get())) {
     num_tensors_++;
     instructions_.emplace_back(SPECIALIZE_TENSOR);
+  } else if (typ->isSubtypeOf(OptionalType::ofTensor())) {
+    num_tensors_++;
+    num_optionals_++;
+    instructions_.emplace_back(SPECIALIZE_OPTIONAL_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_optionals_++;
+    instructions_.emplace_back(SPECIALIZE_OPTIONAL);
   } else if (auto tup = typ->cast<TupleType>()) {
     size_t pos = instructions_.size();
     instructions_.emplace_back(ENTER_TUPLE);
@@ -106,24 +116,42 @@ void ArgumentSpecCreator::dump() const {
       case SPECIALIZE_TENSOR:
         std::cout << "SpecializeTensor ";
         break;
+      case SPECIALIZE_OPTIONAL_TENSOR:
+        std::cout << "SpecializeOptionalTensor ";
+        break;
+      case SPECIALIZE_OPTIONAL:
+        std::cout << "SpecializeOptional ";
+        break;
     }
   }
   std::cout << "\n";
 }
 
 ArgumentSpec ArgumentSpecCreator::create(bool with_grad, const Stack& input)
     const {
-  ArgumentSpec spec(num_tensors_);
+  ArgumentSpec spec(num_tensors_, num_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();
   size_t stack_top = 0; // offset to the top of the stack
   for (Inst inst : instructions_) {
     switch (inst) {
+      case SPECIALIZE_OPTIONAL_TENSOR: {
+        // consume a tensor optional and add to the argspec
+        auto& arg = *stack[stack_top]++;
+        spec.addOptional(arg);
+        if (!arg.isNone()) {
+          spec.addTensor(arg, with_grad);
+        }
+      } break;
       case SPECIALIZE_TENSOR:
         // consume a tensor and add to the argspec
         spec.addTensor(*stack[stack_top]++, with_grad);
         break;
+      case SPECIALIZE_OPTIONAL:
+        // consume a non-tensor optional and add to the argspec
+        spec.addOptional(*stack[stack_top]++);
+        break;
       case ENTER_TUPLE: {
         // consume tuple
         const IValue* iv = stack[stack_top]++;
@@ -159,7 +187,7 @@ ArgumentSpec ArgumentSpecCreator::create(bool with_grad, const Stack& input)
 
 // For every input of a given graph, returns a most detailed type that can be
 // inferred for it based on this ArgumentSpec.
-std::vector<TypePtr> ArgumentSpecCreator::getSpecializedTypes(
+void ArgumentSpecCreator::specializeTypes(
     Graph& graph,
     const ArgumentSpec& spec) const {
   auto input_types =
@@ -169,21 +197,47 @@ std::vector<TypePtr> ArgumentSpecCreator::getSpecializedTypes(
   std::vector<const TypePtr*> input_stack = {input_types.data()};
   std::vector<std::function<TypePtr()>> aggregate_creators;
 
-  size_t arg_spec_offset = 0; // number of specialized tensors seen so far
+  size_t tensor_arg_spec_offset =
+      0; // number of specialized tensors seen so far
+  size_t optional_arg_spec_offset =
+      0; // number of specialized optionals seen so far
 
+  auto dim_tensor_type_from_arg = [](const ArgumentInfo& arg) {
+    return DimensionedTensorType::create(
+        arg.type(),
+        ConvertIntToCPUOrCUDA(arg.device()),
+        arg.dim(),
+        arg.requires_grad());
+  };
   for (Inst inst : instructions_) {
     switch (inst) {
+      case SPECIALIZE_OPTIONAL_TENSOR: {
+        auto& input_type = *input_stack.back()++;
+        auto is_present = spec.isPresent(optional_arg_spec_offset++);
+        if (!is_present) {
+          result_stack.back().emplace_back(input_type);
+          break;
+        }
+        auto& arg = spec.tensorAt(tensor_arg_spec_offset++);
+        AT_ASSERT(arg.defined());
+        result_stack.back().emplace_back(dim_tensor_type_from_arg(arg));
+      } break;
       case SPECIALIZE_TENSOR: {
         input_stack.back()++;
-        auto& arg = spec.at(arg_spec_offset++);
+        auto& arg = spec.tensorAt(tensor_arg_spec_offset++);
         if (!arg.defined()) {
           result_stack.back().emplace_back(AutogradZeroTensorType::get());
         } else {
-          result_stack.back().emplace_back(DimensionedTensorType::create(
-              arg.type(),
-              ConvertIntToCPUOrCUDA(arg.device()),
-              arg.dim(),
-              arg.requires_grad()));
+          result_stack.back().emplace_back(dim_tensor_type_from_arg(arg));
+        }
+      } break;
+      case SPECIALIZE_OPTIONAL: {
+        auto is_present = spec.isPresent(optional_arg_spec_offset++);
+        auto ot = (*input_stack.back()++)->expect<OptionalType>();
+        if (!is_present) {
+          result_stack.back().emplace_back(ot);
+        } else {
+          result_stack.back().emplace_back(ot->getElementType());
         }
       } break;
       case ENTER_TUPLE: {
@@ -213,15 +267,24 @@ std::vector<TypePtr> ArgumentSpecCreator::getSpecializedTypes(
     }
   }
   AT_ASSERT(result_stack.size() == 1);
-  return result_stack.back();
-}
-
-void ArgumentSpecCreator::setInputTypes(Graph& g, const ArgumentSpec& spec)
-    const {
-  auto input_types = getSpecializedTypes(g, spec);
-  auto inputs = g.inputs();
+  // FIXME: by doing this only on the inputs, we only capture graph inputs and
+  // not
+  //        optionals in tuples or objects. For that to work, we would have
+  //        to investigate the uses of the inputs in detail to change the
+  //        accesses/ unwrapping
+  auto inputs = graph.inputs();
   for (size_t i = 0; i < inputs.size(); ++i) {
-    inputs[i]->setType(input_types[i]);
+    auto t = result_stack.back()[i];
+    if (auto ot = t->cast<OptionalType>()) {
+      // if an optional input hasn't been specialized above, it is None
+      // so we disconnect the input here and replace its uses with
+      // a constant
+      WithInsertPoint guard(*graph.nodes().begin());
+      auto c = graph.insertConstant({}, ot);
+      inputs[i]->replaceAllUsesWith(c);
+    } else {
+      inputs[i]->setType(t);
+    }
   }
 }