Caffe2 - C++ API
A deep learning, cross platform ML framework
compiler.h
1 #pragma once
2 #include <functional>
3 #include <memory>
4 #include <string>
5 
6 #include <torch/csrc/jit/ir.h>
7 #include <torch/csrc/jit/script/error_report.h>
8 #include <torch/csrc/jit/script/module.h>
9 #include <torch/csrc/jit/script/sugared_value.h>
10 #include <torch/csrc/jit/script/tree_views.h>
11 
12 namespace torch {
13 namespace jit {
14 namespace script {
15 
16 using Resolver = std::function<std::shared_ptr<
17  SugaredValue>(const std::string& name, Method& m, const SourceRange& loc)>;
18 
19 inline std::shared_ptr<SugaredValue> nativeResolver(
20  const std::string& name,
21  Method& m,
22  const SourceRange& loc) {
23  if (name == "torch") {
24  return std::make_shared<BuiltinModule>("aten");
25  }
26  return nullptr;
27 }
28 
29 // Represents the `self` argument to a method. This wrapper class is necessary
30 // because sometimes `self` sometimes is first class and sometimes not.
31 //
32 // `self` is first class when it refers to a ClassType. It will be bound as a
33 // graph input argument.
34 // `self` is sugared when it refers to a ModuleValue.
35 class Self {
36  public:
37  explicit Self(std::shared_ptr<SugaredValue> sugared)
38  : sugared_(std::move(sugared)) {}
39  explicit Self(ClassTypePtr type) : firstClass_(std::move(type)) {}
40 
41  ClassTypePtr asFirstClass() const {
42  return firstClass_;
43  }
44  std::shared_ptr<SugaredValue> asSugared() const {
45  return sugared_;
46  }
47 
48  private:
49  // Used when `self` is not first-class and so we don't represent it in the
50  // graph. This is only ModuleValue.
51  std::shared_ptr<SugaredValue> sugared_ = nullptr;
52  // Used when `self` is a first-class type
53  ClassTypePtr firstClass_ = nullptr;
54 };
55 
56 TORCH_API void defineMethodsInModule(
57  const std::shared_ptr<Module>& m,
58  const std::vector<Def>& definitions,
59  const std::vector<Resolver>& resolvers, /* determines how we handle free
60  variables in each definition*/
61  // if non-null, the first argument to each def, is bound to this value
62  const c10::optional<Self>& self);
63 
64 // same as above but parse the definitions from source
65 TORCH_API void defineMethodsInModule(
66  const std::shared_ptr<Module>& m,
67  const std::string& source,
68  const Resolver& resolver,
69  const c10::optional<Self>& self);
70 
71 TORCH_API void lambdaLiftFork(Node* fork_node);
72 
73 } // namespace script
74 } // namespace jit
75 } // namespace torch
Definition: jit_type.h:17