[quant][graph] Add useQuantizable function

supriyar · supriyar · commit 53424ae47f49 · 2020-04-15T16:41:34.000-07:00
Summary: Enables to selectively insert observers at the inputs of aten/call functionc

Test Plan:
test_quantize_script.py
Reviewers:

Subscribers:

Tasks:

Tags:

[ghstack-poisoned]
diff --git a/torch/csrc/jit/passes/quantization.cpp b/torch/csrc/jit/passes/quantization.cpp
@@ -61,6 +61,7 @@ std::vector<std::string> _quantizable_aten_funcs = {
     "add_",
     "add",
     "cat",
+    "lstm",
 };
 
 // These are the prim::CallFunctions that doesn't require observation and
@@ -121,6 +122,20 @@ std::vector<std::string> _single_input_general_aten_funcs = {
     "relu",
 };
 
+struct FuncArg {
+  std::string func_name;
+  int arg_index;
+};
+
+using AtenFuncArgs = std::vector<FuncArg>;
+using CallFuncArgs = std::vector<FuncArg>;
+
+// Special checks for ops that do not require observers for all input tensors.
+// For each operator in this list observers are inserted for the input based
+// on the index specified.
+AtenFuncArgs _observe_inputs_aten_func = {};
+CallFuncArgs _observe_inputs_call_func = {};
+
 void fillQConfigMap(
     const Module& module,
     const QConfigDict& qconfig_dict,
@@ -686,33 +701,28 @@ graph(%self, %a, %b, %inplace):
 };
 
 // Check if `use` is an aten function of name `func_name` and if value
-// `v` is the nth argument of the function
-bool isAtenFuncNthArg(
-    Value* v,
-    Node* use,
+// `v` is the nth argument (if provided) of the function.
+bool matchAtenFuncToUse(
+    const Use& use,
     const std::string& func_name,
-    int n) {
-  return use->kind() == Symbol::aten(func_name) && v == use->inputs().at(n);
+    c10::optional<int> n) {
+  Node* node = use.user;
+  return node->kind() == Symbol::aten(func_name) &&
+      (!n.has_value() || n.value() == use.offset);
 }
 
 // Check if `use` is a CallFunction of name `func_name` and if value
-// `v` is the nth argument of the function
-bool isCallFunctionNthArg(
-    Value* v,
-    Node* use,
+// `v` is the nth argument (if provided) of the function
+bool matchCallFuncToUse(
+    const Use& use,
     const std::string& func_name,
-    int n) {
-  return use->kind() == prim::CallFunction &&
-      getFuncName(use->inputs()[0]) == func_name && v == use->inputs().at(n);
+    c10::optional<int> n) {
+  Node* node = use.user;
+  return node->kind() == prim::CallFunction &&
+      getFuncName(node->inputs()[0]) == func_name &&
+      (!n.has_value() || n.value() == use.offset);
 }
 
-struct FuncArg {
-  std::string func_name;
-  int arg_index;
-};
-using AtenFuncArgs = std::vector<FuncArg>;
-using CallFuncArgs = std::vector<FuncArg>;
-
 // Check any use of `v` matches the aten function call
 // or CallFunction patterns
 bool matchArgPattern(
@@ -721,14 +731,13 @@ bool matchArgPattern(
     const CallFuncArgs& call_func_args) {
   for (const Use& u : v->uses()) {
     for (const auto& func_arg : aten_func_args) {
-      if (isAtenFuncNthArg(v, u.user, func_arg.func_name, func_arg.arg_index)) {
+      if (matchAtenFuncToUse(u, func_arg.func_name, func_arg.arg_index)) {
         return true;
       }
     }
 
     for (const auto& func_arg : call_func_args) {
-      if (isCallFunctionNthArg(
-              v, u.user, func_arg.func_name, func_arg.arg_index)) {
+      if (matchCallFuncToUse(u, func_arg.func_name, func_arg.arg_index)) {
         return true;
       }
     }
@@ -997,10 +1006,24 @@ void InsertObserversHelper::preprocess(
   }
 }
 
-// Returns true if the value is the weight to LSTM operator.
-bool isDynamicLSTMWeight(Value* v, Use use, bool is_dynamic) {
-  return is_dynamic && use.user->kind() == Symbol::aten("lstm") &&
-      (use.offset == 2);
+bool useQuantizable(const Use& use, bool is_dynamic) {
+  for (const auto& func_input : _observe_inputs_aten_func) {
+    if (matchAtenFuncToUse(use, func_input.func_name, c10::nullopt)) {
+      return use.offset == func_input.arg_index;
+    }
+  }
+
+  for (const auto& func_input : _observe_inputs_call_func) {
+    if (matchCallFuncToUse(use, func_input.func_name, c10::nullopt)) {
+      return use.offset == func_input.arg_index;
+    }
+  }
+  // Dynamic quantized ops that require special handling for inputs.
+  if (is_dynamic && matchAtenFuncToUse(use, "lstm", c10::nullopt)) {
+    return use.offset == 2;
+  }
+
+  return nodeQuantizable(use.user);
 }
 
 // TODO: remove this as a class method
@@ -1018,9 +1041,9 @@ bool InsertObserversHelper::valueNeedsToBeQuantized(Value* v) {
       return true;
     }
   }
-  // Check whether user is quantizable
+  // Check whether node input value is quantizable
   for (const auto& use : v->uses()) {
-    if (nodeQuantizable(use.user) || isDynamicLSTMWeight(v, use, is_dynamic)) {
+    if (useQuantizable(use, is_dynamic)) {
       return true;
     }
   }
