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
56
57
58
59
60
61
62
63
|
#include <torch/csrc/jit/tensorexpr/codegen.h>
#include <sstream>
namespace torch {
namespace jit {
namespace tensorexpr {
RegisterCodeGenList::StmtFactoryMethod RegisterCodeGenList::
FindStmtFactoryMethod(const std::string& name) {
auto iter = stmt_factory_methods_.find(name);
if (iter == stmt_factory_methods_.end()) {
std::ostringstream oss;
oss << "Invalid stmt codegen name: " << name << ". ";
oss << "Existing codegen names: [";
int index = 0;
for (const auto& entry : stmt_factory_methods_) {
if (index != 0) {
oss << ", ";
}
oss << entry.first;
index++;
}
oss << "]";
throw std::runtime_error(oss.str());
}
return iter->second;
}
void RegisterCodeGenList::AddStmtFactoryMethod(
const std::string& name,
const StmtFactoryMethod& stmt_factory_method) {
auto insert_ret =
stmt_factory_methods_.insert(std::make_pair(name, stmt_factory_method));
if (!insert_ret.second) {
throw std::runtime_error("Duplicated CodeGen names: " + name);
}
}
std::unique_ptr<CodeGen> CreateCodeGen(
const std::string& name,
Stmt* stmt,
const std::vector<CodeGen::BufferArg>& params,
at::Device device) {
RegisterCodeGenList::StmtFactoryMethod method =
RegisterCodeGenList::GetInstance().FindStmtFactoryMethod(name);
return method(stmt, params, device);
}
const Expr* GenericIntrinsicsExpander::mutate(const Intrinsics* v) {
if (v->op_type() == kSigmoid) {
auto x = v->param(0)->accept_mutator(this);
auto one = ExprHandle(getImmediateByType(v->dtype(), 1.0));
auto zero = ExprHandle(getImmediateByType(v->dtype(), 0.0));
ExprHandle y = one / (one + exp(zero - ExprHandle(x)));
return y.node();
}
return IRMutator::mutate(v);
}
} // namespace tensorexpr
} // namespace jit
} // namespace torch
|
