1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
|
#pragma once
#include <ATen/core/Tensor.h>
#include <ATen/core/function_schema.h>
#include <c10/macros/Export.h>
// NOTE: [Jit Decomposition Interface]
//
// For some context of why we need this at all, see NOTE: [forward-mode AD
// decompositions mechanism]
//
// Introducing that mechanism from the NOTE is problematic because:
// - it relies on TorchScript, so now VariableTypeX.cpp depends on TorchScript.
// - there exist internal builds like lite_trainer, which depend on VariableType
// but do not depend on TorchScript.
//
// For internal builds like lite_trainer builds to pass, and for OSS builds that
// do depend on TorchScript to still support the forward AD decomp mechanism, we
// implement a PImpl pattern to avoid a static dependency in favor of a dynamic
// one
// - during static initialization time, if the library is built with TorchScript
// setJitDecompImpl is called in decomposition_registry.cpp setting a global
// ptr to the impl
// - when the program is run,if getJitDecompImpl returns a non null ptr, we can
// carry on normally, otherwise we gracefully error out
//
// For extra context, see VariableHooksInterface.h, where a similar technique
// is used
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 setJitDecompImpl(JitDecompInterface* impl);
TORCH_API JitDecompInterface* getJitDecompImpl();
struct TORCH_API JitDecompRegisterer {
explicit JitDecompRegisterer(JitDecompInterface* impl) {
setJitDecompImpl(impl);
}
};
} // namespace impl
} // namespace autograd
} // namespace torch
|
