Allow autograd to work even when the shape of values cannot be determined (pytorch#8641)

This commit implements the solution proposed in pytorch#8410
to workaround the need to create zero tensors with the same shape as inputs.
It introduces the concept of a LinearBlock which marks places in the code
where we know if all the inputs to the node are zero, then the outputs
to the node are also zero. Autodiff introduces LinearBlocks around
backwards functions, which have this property. specializeUndef then
propagates Undef nodes using this information.

Notes:
* Since we do not always specialize, we have a pass LowerLinearBlocks
that replaces the block with an if statement that dynamically guards
the Undef case.
* We introduce AutogradAdd which is addition that still works when
its inputs might be undefined. In cases where we specialize this will
get removed in favor of a normal add, but there are cases where
gradient graphs do not specialize (e.g. when they are not differentiable,
but a derivative is required) so it is important for this op to be executable.

Author: zdevito

torch/csrc/jit/interned_strings.h

@@ -9,13 +9,27 @@
 
 namespace torch { namespace jit {
 
-#define FORALL_NS_SYMBOLS(_) \
-_(namespaces, prim) \
-_(namespaces, aten) \
-_(namespaces, onnx) \
-_(namespaces, attr) \
-_(namespaces, scope) \
-_(namespaces, namespaces) \
+// Every symbol is classified in a namespace, specifying what kind of symbol it
+// is.  Unsigned char to ensure widening to unique_t (also an unsigned type)
+enum class SymbolNamespace : unsigned char {
+  onnx  = 'o',
+  prim  = 'p',
+  aten  = 't',
+  // NB: ONNX and ATen attributes all live in a unified namespace, as
+  // their interpretation depends on the operator name (which is namespaced)
+  attr  = 'a',
+  // TODO: eliminate me
+  scope = 's'
+};
+
+// Primitive symbols are synthetic operators that occur only in the IR
+// and don't have corresponding implementations in ATen.
+//
+// TODO: We need documentation for all of these symbols.
+//
+// TODO: Consider moving the synthetic onnx operators to their own
+// namespace.
+#define FORALL_PRIM_SYMBOLS(_) \
 _(prim, Assign) \
 _(prim, Constant) \
 _(prim, CppOp) \
@@ -48,9 +62,32 @@ _(prim, NumToTensor) \
 _(prim, TensorToNum) \
 _(prim, AutogradAdd) \
 _(prim, GradOf) \
-_(prim, AnyDefined) \
+_(prim, AnyDefined)
+/* end */
+
+// Workaround for some not-yet-defined ATen symbols, see
+//  - __not__: https://github.com/pytorch/pytorch/issues/5495
+//  - ones, zeros: https://github.com/pytorch/pytorch/issues/5496
+
+#define FORALL_ATEN_EXTRA_SYMBOLS(_) \
 _(aten, __not__) \
+/* end */
+
+#define FORALL_ATEN_SYMBOLS(_) \
 FORALL_ATEN_BASE_SYMBOLS(_) \
+FORALL_ATEN_EXTRA_SYMBOLS(_)
+
+// These symbols correspond to ONNX operators.  Their semantics
+// are defined in https://github.com/onnx/onnx/blob/master/docs/Operators.md
+// The particular version we are targeting is specified by '_onnx_opset_version'
+// in torch.onnx.symbolic
+//
+// In general, most ONNX operators won't get an entry here, because they
+// are handled from the Python end.  However, you may occasionally need
+// to intern an ONNX symbol here so that you can conveniently write an
+// optimization on ONNX operations.
+
+#define FORALL_ONNX_SYMBOLS(_) \
 _(onnx, Add) \
 _(onnx, Concat) \
 _(onnx, Constant) \
@@ -72,8 +109,15 @@ _(onnx, Transpose) \
 _(onnx, Unsqueeze) \
 _(onnx, Loop) \
 _(onnx, If) \
-_(onnx, Reshape) \
-FORALL_ATTR_BASE_SYMBOLS(_) \
+_(onnx, Reshape)
+/* end */
+
+// These symbols are attribute keys.  They are shared between both ONNX and ATen
+// operators (you disambiguate their meaning by looking at the operator itself).
+// In general, you only need to define attribute keys that are used by
+// onnx or prim; ATen attributes are automatically generated in FORALL_ATTR_BASE_SYMBOLS.
+
+#define FORALL_ATTR_EXTRA_SYMBOLS(_) \
 _(attr, Subgraph) \
 _(attr, axes) \
 _(attr, axis) \
@@ -89,33 +133,40 @@ _(attr, starts) \
 _(attr, transA) \
 _(attr, transB) \
 _(attr, name)
+/* end */
 
-// 'prim' symbols are synthetic operators that occur only in the IR
-// and don't have corresponding implementations in ATen.
-
-// 'onnx' symbols correspond to ONNX operators.  Their semantics
-// are defined in https://github.com/onnx/onnx/blob/master/docs/Operators.md
-// The particular version we are targeting is specified by '_onnx_opset_version'
-// in torch.onnx.symbolic
-//
-// In general, most ONNX operators won't get an entry here, because they
-// are handled from the Python end.  However, you may occasionally need
-// to intern an ONNX symbol here so that you can conveniently write an
-// optimization on ONNX operations.
+#define FORALL_ATTR_SYMBOLS(_) \
+FORALL_ATTR_BASE_SYMBOLS(_) \
+FORALL_ATTR_EXTRA_SYMBOLS(_)
 
-// 'attr' symbols are attribute keys.  They are shared between both ONNX and ATen
-// operators (you disambiguate their meaning by looking at the operator itself).
-// In general, you only need to define attribute keys that are used by
-// onnx or prim; ATen attributes are automatically generated in FORALL_ATTR_BASE_SYMBOLS.
+#define FORALL_BUILTIN_SYMBOLS(_) \
+  FORALL_ONNX_SYMBOLS(_) \
+  FORALL_ATEN_SYMBOLS(_) \
+  FORALL_ATTR_SYMBOLS(_) \
+  FORALL_PRIM_SYMBOLS(_) \
+  /* end */
 
 // Note [Symbol allocation]
 // ~~~~~~~~~~~~~~~~~~~~~~~~
 //
-//  1. Symbol namespace is split up into namespaces.
+//  1. Symbol namespace is split up into namespaces.  The hex structure
+//  of our symbols is TTUUUUUU, where TT is the tag byte and U are the unique
+//  bytes.
 //
-//  2. The intended access pattern for built-in symbols is onnx::MatMul
+//  2. We only maintain a single counter for the unique bytes, which means that
+//  we take 256 more space than we would have if we maintained multiple
+//  counters.
+//
+//  3. The first unique_start symbols are reserved for "built-in" symbols.
+//  These symbols are allocated at compile time and get put into the intern
+//  table at process startup time.  Since it's pretty easy to maintain a
+//  distinct counter for every built-in namespace, we let the unique bytes of
+//  built-in symbols to overlap (this is why unique_start is a max)
+//
+//  4. The intended access pattern for built-in symbols is onnx::MatMul
 //  in the torch::jit namespace (this is a Symbol).
 //
+// The code here is not very economical but it gets the job done.
 
 
 // Built-in constant definition strategy:
@@ -129,15 +180,19 @@ _(attr, name)
 
 typedef uint32_t unique_t;
 
+constexpr size_t unique_tag_bits = 8;
+constexpr size_t unique_bits = sizeof(unique_t) * 8 - unique_tag_bits;
+constexpr unique_t unique_mask = (1ULL << unique_bits) - 1;
+
 static const std::string domain_prefix = "org.pytorch.";
 
 // A Symbol is like an interned string, but with a little extra
 // structure; it is namespaced via SymbolNamespace and the resulting
 // intern pointers support efficient namespace testing.
 struct Symbol {
   explicit constexpr Symbol() : value(0) {};
-  explicit constexpr Symbol(unique_t uniq)
-  : value(uniq) {}
+  explicit constexpr Symbol(SymbolNamespace ns, uint32_t uniq)
+  : value((static_cast<uint32_t>(ns) << unique_bits) | (uniq & unique_mask)) {};
 
   // Get a Symbol for a qualified string like "attr::bar"
   static Symbol fromQualString(const std::string & s);
@@ -150,24 +205,26 @@ struct Symbol {
   // argument "foo", and then attempt to intern it.  DO NOT USE THIS
   // with a string literal; attr::foo should be available in that case
   // (and if it's not, you should add it to the built-ins list above.)
-  static Symbol attr(const std::string & s);
-  static Symbol aten(const std::string & s);
-  static Symbol onnx(const std::string & s);
-  static Symbol prim(const std::string & s);
+  static Symbol attr(const std::string & s) { return Symbol(SymbolNamespace::attr, s); };
+  static Symbol aten(const std::string & s) { return Symbol(SymbolNamespace::aten, s); };
+  static Symbol onnx(const std::string & s) { return Symbol(SymbolNamespace::onnx, s); };
+  static Symbol prim(const std::string & s) { return Symbol(SymbolNamespace::prim, s); };
   // TODO: eliminate me
-  static Symbol scope(const std::string & s);
+  static Symbol scope(const std::string & s) { return Symbol(SymbolNamespace::scope, s); };
 
-  bool is_attr() const;
-  bool is_aten() const;
-  bool is_prim() const;
-  bool is_onnx() const;
+  constexpr bool is_attr() const { return ns() == SymbolNamespace::attr; };
+  constexpr bool is_aten() const { return ns() == SymbolNamespace::aten; };
+  constexpr bool is_prim() const { return ns() == SymbolNamespace::prim; };
+  constexpr bool is_onnx() const { return ns() == SymbolNamespace::onnx; };
 
   // So we can switch on this
   constexpr operator unique_t() const {
     return value;
   }
 
-  Symbol ns() const;
+  constexpr SymbolNamespace ns() const {
+    return static_cast<SymbolNamespace>(value >> unique_bits);
+  }
 
   // Give a string corresponding to the unqualified version of this name, e.g.,
   // "mm". Use this in a context where the intended namespace of the string is
@@ -189,41 +246,33 @@ struct Symbol {
   std::string domainString() const;
 
 private:
-  explicit Symbol(Symbol ns, const std::string & s);
+  explicit Symbol(SymbolNamespace ns, const std::string & s);
   unique_t value;
 };
 
 static inline bool operator==(Symbol lhs, Symbol rhs) {
   return static_cast<unique_t>(lhs) == static_cast<unique_t>(rhs);
 }
 
-enum class _keys : unique_t {
-    #define DEFINE_KEY(ns, s) ns##_##s,
-    FORALL_NS_SYMBOLS(DEFINE_KEY)
-    #undef DEFINE_KEY
-    num_symbols
-};
-
-#define DEFINE_SYMBOL(s) \
-  constexpr Symbol s(static_cast<unique_t>(_keys::s));
+#define DEFINE_KEY(ns, s) s,
+#define DEFINE_SYMBOL(ns, s) constexpr Symbol s(SymbolNamespace::ns, static_cast<unique_t>(_keys::s));
+#define DEFINE_BUILTINS(ns, forall_symbols) \
+  namespace ns { \
+    enum class _keys : unique_t { \
+      forall_symbols(DEFINE_KEY) \
+      num_symbols \
+    }; \
+    forall_symbols(DEFINE_SYMBOL) \
+  }
 
-#undef DEFINE_SYMBOL
+DEFINE_BUILTINS(onnx, FORALL_ONNX_SYMBOLS)
+DEFINE_BUILTINS(aten, FORALL_ATEN_SYMBOLS)
+DEFINE_BUILTINS(attr, FORALL_ATTR_SYMBOLS)
+DEFINE_BUILTINS(prim, FORALL_PRIM_SYMBOLS)
 
-#define DEFINE_SYMBOL(ns, s) \
-  namespace ns { constexpr Symbol s(static_cast<unique_t>(_keys::ns##_##s)); }
-FORALL_NS_SYMBOLS(DEFINE_SYMBOL)
+#undef DEFINE_KEY
 #undef DEFINE_SYMBOL
 
-inline Symbol Symbol::attr(const std::string & s) { return Symbol::fromQualString("attr::" + s); }
-inline Symbol Symbol::aten(const std::string & s)  { return Symbol::fromQualString("aten::" + s); }
-inline Symbol Symbol::onnx(const std::string & s)  { return Symbol::fromQualString("onnx::" + s); }
-inline Symbol Symbol::prim(const std::string & s)  { return Symbol::fromQualString("prim::" + s); }
-inline Symbol Symbol::scope(const std::string & s) { return Symbol::fromQualString("scope::" + s); }
-inline bool Symbol::is_attr() const { return ns() == namespaces::attr; }
-inline bool Symbol::is_aten() const { return ns() == namespaces::aten; }
-inline bool Symbol::is_prim() const { return ns() == namespaces::prim; }
-inline bool Symbol::is_onnx() const { return ns() == namespaces::onnx; }
-
 }} // namespace torch::jit
 
 // make symbol behave like an integer in hash tables