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
55
|
#pragma once
#include <ATen/core/jit_type.h>
#include <torch/csrc/jit/ir/ir.h>
namespace torch {
namespace jit {
// Simple data structure for containing a type T in nested control blocks
// Should only be used after initial compilation where type checking and
// loads and stores are emitted
template <typename T>
struct MiniEnvironment {
MiniEnvironment(Block* b, std::shared_ptr<MiniEnvironment> next = nullptr)
: next(std::move(next)) {}
std::shared_ptr<MiniEnvironment<T>> next;
T findInThisFrame(const std::string& name) {
auto it = table.find(name);
if (it != table.end()) {
return it->second;
}
return nullptr;
}
T findInAnyFrame(const std::string& name) {
for (auto runner = this; runner; runner = runner->next.get()) {
if (auto r = runner->findInThisFrame(name)) {
return r;
}
}
return nullptr;
}
void setVar(const std::string& name, T value) {
table[name] = value;
}
std::vector<std::string> definedVariables() {
std::vector<std::string> result;
for (auto& kv : table) {
result.push_back(kv.first);
}
std::sort(result.begin(), result.end());
return result;
}
private:
std::unordered_map<std::string, T> table;
};
} // namespace jit
} // namespace torch
|
