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
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
|
#include <torch/csrc/lazy/ts_backend/ts_node_lowering.h>
#include <ATen/Functions.h>
#include <torch/csrc/jit/frontend/sugared_value.h>
#include <torch/csrc/jit/jit_log.h>
#include <torch/csrc/lazy/backend/backend_interface.h>
#include <torch/csrc/lazy/core/helpers.h>
#include <torch/csrc/lazy/core/internal_ops/ltc_ops.h>
#include <torch/csrc/lazy/core/ir_builder.h>
#include <torch/csrc/lazy/core/lazy_graph_executor.h>
#include <torch/csrc/lazy/core/ops/utils.h>
#include <torch/csrc/lazy/core/permutation_util.h>
#include <torch/csrc/lazy/ts_backend/ir_builder.h>
#include <torch/csrc/lazy/ts_backend/ts_lowering_context.h>
namespace torch::lazy {
static TSOpVector LowerBuiltin(
const torch::lazy::Node* node,
const std::shared_ptr<torch::jit::GraphFunction>& function,
const std::vector<torch::jit::NamedValue>& arguments,
const std::vector<torch::jit::NamedValue>& kwarguments = {}) {
return LowerTSBuiltin(function, node->op().op, arguments, kwarguments);
}
static TSOpVector LowerBuiltin(
c10::Symbol sym,
const std::shared_ptr<torch::jit::GraphFunction>& function,
const std::vector<torch::jit::NamedValue>& arguments,
const std::vector<torch::jit::NamedValue>& kwarguments = {}) {
return LowerTSBuiltin(function, sym, arguments, kwarguments);
}
TSOpVector LowerTSBuiltin(
const std::shared_ptr<torch::jit::GraphFunction>& function,
c10::Symbol sym,
const std::vector<torch::jit::NamedValue>& arguments,
const std::vector<torch::jit::NamedValue>& kwarguments) {
auto builtin =
std::make_shared<torch::jit::BuiltinFunction>(sym, std::nullopt);
auto magic_method = std::make_shared<torch::jit::MagicMethod>("", builtin);
auto ret = magic_method->call({}, *function, arguments, kwarguments, 0);
auto& sv = dynamic_cast<torch::jit::SimpleValue&>(*ret);
if (sv.getValue()->type()->kind() == c10::TypeKind::TupleType) {
const auto tuple_call_result = sv.asTuple({}, *function);
TSOpVector tuple_result;
for (const auto& tuple_component : tuple_call_result) {
auto tuple_component_sv =
dynamic_cast<torch::jit::SimpleValue*>(tuple_component.get());
tuple_result.push_back(tuple_component_sv->getValue());
}
return tuple_result;
}
return {sv.getValue()};
}
static torch::jit::Value* GenerateClone(
torch::jit::Value* val,
const std::shared_ptr<torch::jit::GraphFunction>& function) {
std::vector<torch::jit::NamedValue> clone_arguments;
clone_arguments.emplace_back(val);
TSOpVector cloned = LowerBuiltin(at::aten::clone, function, clone_arguments);
TORCH_CHECK_EQ(cloned.size(), 1);
return cloned.front();
}
// Node Lowerings
// Default node lowering
TSOpVector TsNode::Lower(
// NOLINTNEXTLINE(performance-unnecessary-value-param)
std::shared_ptr<torch::jit::GraphFunction> function,
TSLoweringContext* loctx) const {
std::vector<torch::jit::NamedValue> arguments;
for (const torch::lazy::Output& output : operands()) {
arguments.emplace_back(loctx->GetOutputOp(output));
}
return LowerBuiltin(this, function, arguments);
}
// Non-native ops
torch::lazy::TSOpVector Cast::Lower(
std::shared_ptr<torch::jit::GraphFunction> function,
torch::lazy::TSLoweringContext* loctx) const {
std::vector<torch::jit::NamedValue> arguments;
arguments.emplace_back(loctx->GetOutputOp(operand(0)));
arguments.emplace_back(dtype);
return LowerBuiltin(at::aten::to, function, arguments);
}
torch::lazy::TSOpVector DeviceData::Lower(
std::shared_ptr<torch::jit::GraphFunction> function,
torch::lazy::TSLoweringContext* loctx) const {
auto infoptr = data_->info();
auto deviceDataInfoPtr =
(torch::lazy::LazyGraphExecutor::DeviceDataInfo*)infoptr;
if (GRAPH_DUMP_ENABLED) {
LOG(ERROR) << "Lowering device data node, tensor id "
<< deviceDataInfoPtr->tensor_id << '\n';
}
return {loctx->GetParameter(data_)};
}
torch::lazy::TSOpVector Expand::Lower(
std::shared_ptr<torch::jit::GraphFunction> function,
torch::lazy::TSLoweringContext* loctx) const {
std::vector<torch::jit::NamedValue> arguments;
arguments.emplace_back(loctx->GetOutputOp(operand(0)));
arguments.emplace_back(size);
auto expand_out = LowerBuiltin(this, function, arguments);
if (is_scalar_expand) {
// The aten::expand operations sets all strides to 0 when the original is
// of rank 0. This leads to false positives when checking for internal
// memory overlap, because at::has_internal_overlap returns
// MemOverlap::YES when a stride is set to 0.
TORCH_CHECK_EQ(expand_out.size(), 1);
return {GenerateClone(expand_out.front(), function)};
}
return expand_out;
}
torch::lazy::TSOpVector Scalar::Lower(
std::shared_ptr<torch::jit::GraphFunction> function,
torch::lazy::TSLoweringContext* loctx) const {
auto options =
at::TensorOptions()
.device(torch::lazy::getBackend()->EagerFallbackDeviceType())
.dtype(shape().scalar_type());
return {loctx->graph()->insertConstant(at::scalar_tensor(value, options))};
}
} // namespace torch::lazy
|
