[JIT] Optimize FunctionSchema::checkArg for the Tensor case.

aten/src/ATen/core/ivalue.h

@@ -685,6 +685,10 @@ struct CAFFE2_API IValue final {
     return "InvalidTag(" + c10::guts::to_string(static_cast<int>(tag)) + ")";
   }
 
+  uint32_t tagAsInt() const {
+    return static_cast<uint32_t>(tag);
+  }
+
   // generic v.to<at::Tensor>() implementations
   // that can be used in special functions like pop/push
   // that use template meta-programming.

torch/csrc/jit/api/function_impl.cpp

@@ -45,10 +45,33 @@ c10::intrusive_ptr<c10::ivalue::Future> GraphFunction::runAsync(
   return get_executor().runAsync(stack, std::move(taskLauncher));
 }
 
+size_t GraphFunction::computeInputTypesHash(
+    const std::vector<IValue>& stack) const {
+  // Use an algorithm similar to boost::hash_combine to compute the vector hash
+  size_t r = 0;
+  const size_t magic_number = 0x9e3779b9;
+  for (const IValue& iv : stack) {
+    r ^= std::hash<uint32_t>{}(iv.tagAsInt()) + magic_number + (r << 6) +
+        (r >> 2);
+  }
+  return r;
+}
+
 IValue GraphFunction::operator()(
     std::vector<IValue> stack,
     const Kwargs& kwargs) {
-  getSchema().checkAndNormalizeInputs(stack, kwargs);
+  bool need_schema_check = true;
+  if (!kwargs.size()) { // Fast path
+    size_t input_types_hash = computeInputTypesHash(stack);
+    if (!schema_checks_cache_.count(input_types_hash)) {
+      getSchema().checkAndNormalizeInputs(stack, kwargs);
+      schema_checks_cache_.insert(input_types_hash);
+    }
+    need_schema_check = false;
+  }
+  if (need_schema_check) {
+    getSchema().checkAndNormalizeInputs(stack, kwargs);
+  }
   run(stack);
   return stack.front();
 }

torch/csrc/jit/api/function_impl.h

@@ -141,6 +141,14 @@ struct TORCH_API GraphFunction : public Function {
   // mutable because getSchema caches the default schema if one is requested
   // before a call to setSchema
   mutable std::unique_ptr<FunctionSchema> schema_;
+
+  // If we're seeing inputs of the same types over and over again (a frequent
+  // use case when a model is run in inference), and all of the inputs are
+  // positional, we can shortcut expensive schema checks. We do this by
+  // computing a hash of types of the inputs and if that combination of types
+  // has been seen before, we skip the schema check.
+  mutable std::unordered_set<size_t> schema_checks_cache_;
+  size_t computeInputTypesHash(const std::vector<IValue>& stack) const;
 };
 } // namespace jit
 } // namespace torch