Access JIT decomposition utils through an interface

Author: soulitzer

build_variables.bzl

@@ -166,6 +166,7 @@ core_trainer_sources = [
     "torch/csrc/autograd/saved_variable.cpp",
     "torch/csrc/autograd/variable.cpp",
     "torch/csrc/autograd/utils/warnings.cpp",
+    "torch/csrc/jit_decomp_interface.cpp",
     "torch/csrc/jit/frontend/name_mangler.cpp",
     "torch/csrc/jit/ir/type_hashing.cpp",
     "torch/csrc/jit/serialization/pickler.cpp",

torch/csrc/autograd/VariableTypeUtilsDependOnOps.h

@@ -1,6 +1,8 @@
 #pragma once
 
-#include <torch/csrc/jit/runtime/decomposition_registry.h>
+#include <ATen/core/boxing/KernelFunction.h>
+#include <ATen/core/dispatch/Dispatcher.h>
+#include <torch/csrc/jit_decomp_interface.h>
 
 // This is the set of helpers in VariableTypeUtils have a dependency on
 // native_functions.yaml meaning the file will need to be re-compiled every time
@@ -12,14 +14,28 @@ namespace torch {
 namespace autograd {
 namespace impl {
 
+class MyFunctor final : public c10::OperatorKernel {
+ public:
+  MyFunctor(JitDecompInterface* fns) : fns_(fns){};
+
+  void operator()(
+      const c10::OperatorHandle& op,
+      c10::DispatchKeySet ks,
+      torch::jit::Stack* stack) {
+    fns_->run_jit_decomposition_(op, stack);
+  }
+  JitDecompInterface* fns_;
+};
+
 // Depends on torch/csrc/jit/ir/ir.h -> aten/src/ATen/core/interned_strings.h
 template <class Return, class... Args>
 Return run_jit_decomposition_with_args_for_jvp(
     c10::string_view name,
     const c10::OperatorHandle& opHandle,
     c10::DispatchKeySet dispatchKeySet,
     Args&&... args) {
-  bool has_decomp = jit::has_jit_decomposition(opHandle.schema());
+  JitDecompInterface* fns = getJitDecomp();
+  bool has_decomp = fns->has_jit_decomposition_(opHandle.schema());
 
   TORCH_CHECK_NOT_IMPLEMENTED(
       has_decomp,
@@ -33,7 +49,8 @@ Return run_jit_decomposition_with_args_for_jvp(
       "PYTORCH_JIT=0 is set, some operators may no longer be used with forward AD.");
 
   return c10::KernelFunction::makeFromBoxedKernel(
-             c10::BoxedKernel::makeFromFunction<&jit::run_jit_decomposition>())
+             c10::BoxedKernel::makeFromFunctor(
+                 std::make_unique<MyFunctor>(fns)))
       .call<Return, Args...>(
           opHandle, dispatchKeySet, std::forward<Args>(args)...);
 }

torch/csrc/jit/runtime/decomposition_registry.cpp

@@ -13,6 +13,7 @@
 #include <torch/csrc/jit/passes/inliner.h>
 #include <torch/csrc/jit/passes/peephole.h>
 #include <torch/csrc/jit/runtime/graph_executor.h>
+#include <torch/csrc/jit_decomp_interface.h>
 #include <memory>
 #include <unordered_map>
 
@@ -160,6 +161,26 @@ void RegisterDecomposition(
   schema_to_decomposition[&schema] = g;
 }
 
+struct JitDecomp final : torch::autograd::impl::JitDecompInterface {
+  bool has_jit_decomposition_(const c10::FunctionSchema& schema) const override;
+  void run_jit_decomposition_(
+      const c10::OperatorHandle& op,
+      torch::jit::Stack* stack) const override;
+};
+
+JitDecomp jitDecomp;
+torch::autograd::impl::JitDecompRegisterer registerJitDecomp(&jitDecomp);
+
+void JitDecomp::run_jit_decomposition_(
+    const c10::OperatorHandle& op,
+    torch::jit::Stack* stack) const {
+  run_jit_decomposition(op, stack);
+}
+
+bool JitDecomp::has_jit_decomposition_(const FunctionSchema& schema) const {
+  return has_jit_decomposition(schema);
+}
+
 void run_jit_decomposition(
     const c10::OperatorHandle& op,
     torch::jit::Stack* stack) {

torch/csrc/jit_decomp_interface.cpp

@@ -0,0 +1,23 @@
+#include <torch/csrc/jit_decomp_interface.h>
+
+namespace torch {
+namespace autograd {
+namespace impl {
+
+namespace {
+JitDecompInterface* fns = nullptr;
+}
+
+void setJitDecompInterface(JitDecompInterface* f) {
+  fns = f;
+}
+JitDecompInterface* getJitDecomp() {
+  TORCH_CHECK(
+      fns,
+      "Support for JIT decompositions has not been loaded; have you linked against TBD?")
+  return fns;
+}
+
+} // namespace impl
+} // namespace autograd
+} // namespace torch

torch/csrc/jit_decomp_interface.h

@@ -0,0 +1,31 @@
+#pragma once
+
+#include <ATen/core/Tensor.h>
+#include <ATen/core/function_schema.h>
+#include <c10/macros/Export.h>
+
+namespace torch {
+namespace autograd {
+namespace impl {
+
+struct TORCH_API JitDecompInterface {
+  virtual ~JitDecompInterface() = default;
+  virtual bool has_jit_decomposition_(
+      const c10::FunctionSchema& schema) const = 0;
+  virtual void run_jit_decomposition_(
+      const c10::OperatorHandle& op,
+      jit::Stack* stack) const = 0;
+};
+
+TORCH_API void setJitDecompInterface(JitDecompInterface* fns);
+TORCH_API JitDecompInterface* getJitDecomp();
+
+struct TORCH_API JitDecompRegisterer {
+  explicit JitDecompRegisterer(JitDecompInterface* fns) {
+    setJitDecompInterface(fns);
+  }
+};
+
+} // namespace impl
+} // namespace autograd
+} // namespace torch