From 287dc5ea54c5415e44e1a65455839a9d370d0ca3 Mon Sep 17 00:00:00 2001 From: Nikolay Korovaiko Date: Thu, 9 Sep 2021 22:50:11 +0000 Subject: [PATCH 1/6] build a graph --- lazy_tensor_core/example_size.py | 30 ++++++++++ lazy_tensor_core/jie.py | 36 +++++++++++ lazy_tensor_core/jie_empty.py | 34 +++++++++++ .../csrc/init_python_bindings.cpp | 20 +++++++ .../lazy_tensor_core/csrc/tensor.cpp | 46 ++++++++++++++ .../csrc/ts_backend/ops/add.cpp | 16 +++++ .../csrc/ts_backend/ops/add.h | 14 +++++ .../csrc/ts_backend/ts_node_lowering.cpp | 9 ++- lazy_tensor_core/lazy_tensors/shape.h | 8 ++- lazy_tensor_core/test/dynamic_lazy_tensor.py | 55 +++++++++++++++++ lazy_tensor_core/test/dynamic_lazy_tensor2.py | 60 +++++++++++++++++++ test/test_jit_fuser_te.py | 1 + 12 files changed, 326 insertions(+), 3 deletions(-) create mode 100644 lazy_tensor_core/example_size.py create mode 100644 lazy_tensor_core/jie.py create mode 100644 lazy_tensor_core/jie_empty.py create mode 100644 lazy_tensor_core/test/dynamic_lazy_tensor.py create mode 100644 lazy_tensor_core/test/dynamic_lazy_tensor2.py diff --git a/lazy_tensor_core/example_size.py b/lazy_tensor_core/example_size.py new file mode 100644 index 0000000000000..3d3de837c83a0 --- /dev/null +++ b/lazy_tensor_core/example_size.py @@ -0,0 +1,30 @@ +import torch +import lazy_tensor_core +import lazy_tensor_core.debug.metrics as metrics +import lazy_tensor_core.core.lazy_model as ltm + +lazy_tensor_core._LAZYC._ltc_init_ts_backend() + +import os +path = os.path.abspath(lazy_tensor_core.__file__) +print(path) +torch.manual_seed(42) + +device = 'lazy' +dtype = torch.float32 + + +x = torch.tensor([[0.1, 1.2], [2.2, 3.1], [4.9, 5.2]], device=device) +y = torch.tensor([[0.1, 1.2], [2.2, 3.1], [4.9, 5.2]], device=device) +z = x.normal_() +w = z + y +#z = x + y +#w = x + z.size()[1] +ltm.mark_step() +#print(w) + + + +#print(metrics.metrics_report()) + + diff --git a/lazy_tensor_core/jie.py b/lazy_tensor_core/jie.py new file mode 100644 index 0000000000000..a63292d8d1def --- /dev/null +++ b/lazy_tensor_core/jie.py @@ -0,0 +1,36 @@ +import torch +import lazy_tensor_core +import lazy_tensor_core.debug.metrics as metrics +import lazy_tensor_core.core.lazy_model as ltm + + + +lazy_tensor_core._LAZYC._ltc_init_ts_backend() + +# from caffe2.python import workspace +# workspace.GlobalInit(['caffe2', '--caffe2_log_level=-1000']) + +lazy_tensor_core._LAZYC._ltc_set_dynamic_shapes_mode() +device = 'lazy' +dtype = torch.float32 + +x1 = torch.randn(2, 4, device=device, dtype=dtype) +ltm.mark_step() +x1 = torch.randn(2, 6, device=device, dtype=dtype) +ltm.mark_step() +# x2 = torch.randn(2, 8, 8, device=device, dtype=dtype) + +# def f(x): +# o = x + 1.0 +# o = torch.nn.functional.gelu(o) +# return o + +# for i in range(2): +# o = f(x1) +# ltm.mark_step() + +# o = f(x2) +# ltm.mark_step() + +print(metrics.metrics_report()) + diff --git a/lazy_tensor_core/jie_empty.py b/lazy_tensor_core/jie_empty.py new file mode 100644 index 0000000000000..8cd9e1d6d24a8 --- /dev/null +++ b/lazy_tensor_core/jie_empty.py @@ -0,0 +1,34 @@ +import torch +import lazy_tensor_core +import lazy_tensor_core.debug.metrics as metrics +import lazy_tensor_core.core.lazy_model as ltm + + + +lazy_tensor_core._LAZYC._ltc_init_ts_backend() + +# from caffe2.python import workspace +# workspace.GlobalInit(['caffe2', '--caffe2_log_level=-1000']) + +lazy_tensor_core._LAZYC._ltc_set_dynamic_shapes_mode() +device = 'lazy' +dtype = torch.float32 + +x1 = torch.empty(2, 4, device=device, dtype=dtype) +ltm.mark_step() +# x2 = torch.randn(2, 8, 8, device=device, dtype=dtype) + +# def f(x): +# o = x + 1.0 +# o = torch.nn.functional.gelu(o) +# return o + +# for i in range(2): +# o = f(x1) +# ltm.mark_step() + +# o = f(x2) +# ltm.mark_step() + +print(metrics.metrics_report()) + diff --git a/lazy_tensor_core/lazy_tensor_core/csrc/init_python_bindings.cpp b/lazy_tensor_core/lazy_tensor_core/csrc/init_python_bindings.cpp index 70d297bebf602..b272528d9fb8c 100644 --- a/lazy_tensor_core/lazy_tensor_core/csrc/init_python_bindings.cpp +++ b/lazy_tensor_core/lazy_tensor_core/csrc/init_python_bindings.cpp @@ -32,6 +32,7 @@ #include "torch/csrc/autograd/variable.h" #include "torch/csrc/jit/python/pybind.h" #include "torch/csrc/utils/cuda_lazy_init.h" +#include "lazy_tensor_core/csrc/ts_backend/ops/add.h" namespace torch_lazy_tensors { namespace { @@ -413,6 +414,25 @@ void InitLtcModuleBindings(py::module m) { }; return GetTensorsDump(tensors, coverter); }); + m.def("_dynamic_size", + [](at::Tensor& self) { + + std::cerr << "in _dynamic_size\n"; + LazyTensor self_lazy_tensor = bridge::GetLtcTensor(self); + return bridge::AtenFromLtcTensor( + self_lazy_tensor.CreateFrom(ir::MakeNode( + self_lazy_tensor.GetIrValue()))); + }); + m.def("_dynamic_expand", + [](at::Tensor& self, at::Tensor& other) { + std::cerr << "in _expand\n"; + LazyTensor self_lazy_tensor = bridge::GetLtcTensor(self); + LazyTensor other_lazy_tensor = + bridge::GetOrCreateLtcTensor(other, self_lazy_tensor.GetDevice()); + return bridge::AtenFromLtcTensor( + self_lazy_tensor.CreateFrom(ir::MakeNode( + self_lazy_tensor.GetIrValue(), other_lazy_tensor.GetIrValue()))); + }); m.def("_get_ltc_tensors_text", [](const std::vector& tensors) -> std::string { auto coverter = [](lazy_tensors::Span nodes) { diff --git a/lazy_tensor_core/lazy_tensor_core/csrc/tensor.cpp b/lazy_tensor_core/lazy_tensor_core/csrc/tensor.cpp index 7b5ab92116b07..0c887cc0d2628 100644 --- a/lazy_tensor_core/lazy_tensor_core/csrc/tensor.cpp +++ b/lazy_tensor_core/lazy_tensor_core/csrc/tensor.cpp @@ -38,6 +38,28 @@ #include "lazy_tensors/str_join.h" #include "torch/csrc/autograd/variable.h" + +#include "torch/csrc/jit/runtime/custom_operator.h" + +int get_current_level() { + static const auto PRINT_VLOG = std::getenv("PRINT_VLOG"); + if (PRINT_VLOG) { + return std::atoi(PRINT_VLOG); + } + return 3; +} + +std::ostream& get_ostream(int level) { + std::cerr << "in get_ostream\n"; + static std::stringstream dummy{}; + + static const auto cur_level = get_current_level(); + if (true || level >= cur_level) { + return std::cerr; + } + return dummy; +} + namespace torch_lazy_tensors { namespace { @@ -1618,6 +1640,11 @@ std::shared_ptr LazyTensor::SyncTensorsGraphInternal( &coll.indices); PostOrderData po_data = RunPostOrder(*tensors, coll.indices); + + for (auto n: po_data.post_order) { + LTC_VLOG(5) << "node = " << *n << std::endl; + } + coll.hash = lazy_tensors::util::HashCombine( coll.hash, lazy_tensors::util::Hash(po_data.parameter_sequence)); LTC_VLOG(4) << "Parameter sequence graph hash " @@ -1660,3 +1687,22 @@ lazy_tensors::uint64 LazyTensor::GetRunningSeed(const Device& device) { } } // namespace torch_lazy_tensors + + +// void DynamicSize(torch::jit::Stack* stack) { +// at::Tensor t = torch::jit::pop(stack).toTensor(); +// torch::jit::push(stack, t.sizes()); +// } + +// const torch::jit::RegisterOperators DynamicSizeOp({ +// torch::jit::Operator( +// "aten::dynamic_size(Tensor a) -> int[]", +// [](const torch::jit::Node*) -> torch::jit::Operation { +// return [](torch::jit::Stack* stack) { +// at::Tensor t = torch::jit::pop(stack).toTensor(); +// torch::jit::push(stack, t.sizes()); +// }; +// }, +// //DynamicSize, +// c10::AliasAnalysisKind::FROM_SCHEMA), +// }); \ No newline at end of file diff --git a/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/ops/add.cpp b/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/ops/add.cpp index e6d7383b704a0..75ced73deb1a0 100644 --- a/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/ops/add.cpp +++ b/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/ops/add.cpp @@ -16,6 +16,22 @@ NodePtr Add::Clone(OpList operands) const { return MakeNode(operands.at(0), operands.at(1)); } +DynamicExpand::DynamicExpand(const Value& lhs, const Value& rhs) + : Node(ir::OpKind(at::aten::expand), {lhs, rhs}, lhs.shape(), + /*num_outputs=*/1, /*hash_seed=*/0x5a2d296e9) {} + +NodePtr DynamicExpand::Clone(OpList operands) const { + return MakeNode(operands.at(0), operands.at(1)); +} + +DynamicSize::DynamicSize(const Value& lhs) + : Node(ir::OpKind(at::aten::size), {lhs}, lazy_tensors::Shape{}, + /*num_outputs=*/1, /*hash_seed=*/0x5a2d296e9) {} + +NodePtr DynamicSize::Clone(OpList operands) const { + return MakeNode(operands.at(0)); +} + } // namespace ops } // namespace ir } // namespace torch_lazy_tensors diff --git a/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/ops/add.h b/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/ops/add.h index a3f6020c81c17..f420823211a15 100644 --- a/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/ops/add.h +++ b/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/ops/add.h @@ -14,6 +14,20 @@ class Add : public Node { NodePtr Clone(OpList operands) const override; }; +class DynamicSize : public Node { + public: + DynamicSize(const Value& lhs); + + NodePtr Clone(OpList operands) const override; +}; + +class DynamicExpand : public Node { + public: + DynamicExpand(const Value& lhs, const Value& sz); + + NodePtr Clone(OpList operands) const override; +}; + } // namespace ops } // namespace ir } // namespace torch_lazy_tensors diff --git a/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/ts_node_lowering.cpp b/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/ts_node_lowering.cpp index 6a93db3c96d99..341841e3d467b 100644 --- a/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/ts_node_lowering.cpp +++ b/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/ts_node_lowering.cpp @@ -306,6 +306,12 @@ class TSNodeLowering : public NodeLowering { ir::NodeCast(node, *ir::ops::ltc_device_data); return {loctx()->GetParameter(device_data_node->data())}; } + + if (node->op().op == at::aten::size) { + auto size_val = loctx()->graph()->insert(at::aten::size, {loctx()->GetOutputOp(node->operands().at(0))}); + return {size_val}; + } + std::vector arguments; for (const ir::Output& output : node->operands()) { arguments.emplace_back(loctx()->GetOutputOp(output)); @@ -617,7 +623,8 @@ class TSNodeLowering : public NodeLowering { TSOpVector LowerExpand(const ir::ops::Expand* node) { std::vector arguments; arguments.emplace_back(loctx()->GetOutputOp(node->operand(0))); - arguments.emplace_back(node->size()); + arguments.emplace_back(loctx()->GetOutputOp(node->operand(1))); + //arguments.emplace_back(node->size()); auto expand_out = LowerBuiltin(node, arguments); if (node->is_scalar_expand()) { // The aten::expand operations sets all strides to 0 when the original is diff --git a/lazy_tensor_core/lazy_tensors/shape.h b/lazy_tensor_core/lazy_tensors/shape.h index 0625d24c930e3..acd006151a8ef 100644 --- a/lazy_tensor_core/lazy_tensors/shape.h +++ b/lazy_tensor_core/lazy_tensors/shape.h @@ -74,7 +74,9 @@ class Shape { int dimensions_size() const { return dimensions_.size(); } int64 dimensions(int index) const { if (dynamic_mode_.load()) { - throw std::runtime_error("Exact shape not known"); + LTC_LOG(INFO) << "@@@ 1 dimensions shape is " << this->ToString() << std::endl; + //std::cerr << "@@@ 1 dimensions shape is " << this->ToString() << std::endl; + //throw std::runtime_error("Exact shape not known"); } LTC_CHECK_LT(index, dimensions_.size()); return dimensions_[index]; @@ -87,7 +89,9 @@ class Shape { lazy_tensors::Span dimensions() const { if (dynamic_mode_.load()) { - throw std::runtime_error("Exact shape not known"); + LTC_LOG(INFO) << "@@@ 2 dimensions shape is " << this->ToString() << std::endl; + //std::cerr << "@@@ 2 dimensions shape is " << this->ToString() << std::endl; + //throw std::runtime_error("Exact shape not known"); } return MakeSpan(dimensions_); } diff --git a/lazy_tensor_core/test/dynamic_lazy_tensor.py b/lazy_tensor_core/test/dynamic_lazy_tensor.py new file mode 100644 index 0000000000000..bf1b1eb2c17dc --- /dev/null +++ b/lazy_tensor_core/test/dynamic_lazy_tensor.py @@ -0,0 +1,55 @@ +import torch +import lazy_tensor_core +import lazy_tensor_core.debug.metrics as metrics +import lazy_tensor_core.core.lazy_model as ltm + + + +lazy_tensor_core._LAZYC._ltc_init_ts_backend() + +# class DynamicLazyTensor(torch.Tensor): + +# def __init__(*args, **kwargs): +# super().__init__(*args, **kwargs) + +# def __torch_function__(self, func, types, args=(), kwargs=None): +# return lazy_tensor_core._LAZYC._dynamic_size(self) + + +# a = DynamicLazyTensor([2.0, 4.0]) +# a.size() + + +class DynamicLazyTensor2(torch.Tensor): + + def __init__(self, t): + super().__init__() + self._t = t.to(device='lazy') + print(id(self)) + + def expand(self, sz): + return lazy_tensor_core._LAZYC._dynamic_expand(self._t, sz) + + def __torch_function__(self, func, types, args=(), kwargs=None): + if func.__name__ == 'size': + return lazy_tensor_core._LAZYC._dynamic_size(self._t) + elif func.__name__ == 'expand': + print(type(args[0]).__name__) + print(type(args[1]).__name__) + return lazy_tensor_core._LAZYC._dynamic_expand(self._t, args[0]) + else: + print("Not Implemented!") + + +b = DynamicLazyTensor2(torch.rand(3, 4, 5)) +c = b.size() +d = b.expand(c) +#d = b.expand((3, 4, 5)) + +# _get_ltc_tensors_text _get_ltc_tensors_backend +print(lazy_tensor_core._LAZYC._get_ltc_tensors_text([d])) +print(lazy_tensor_core._LAZYC._get_ltc_tensors_backend([d])) + +# Questions +# How do add dynamic_ and override them in LTC +# why can't I use `expand` \ No newline at end of file diff --git a/lazy_tensor_core/test/dynamic_lazy_tensor2.py b/lazy_tensor_core/test/dynamic_lazy_tensor2.py new file mode 100644 index 0000000000000..d70936038083f --- /dev/null +++ b/lazy_tensor_core/test/dynamic_lazy_tensor2.py @@ -0,0 +1,60 @@ +import torch +import lazy_tensor_core +import lazy_tensor_core.debug.metrics as metrics +import lazy_tensor_core.core.lazy_model as ltm + + + +lazy_tensor_core._LAZYC._ltc_init_ts_backend() + +lazy_tensor_core._LAZYC._ltc_set_dynamic_shapes_mode() + +class DynamicLazyTensor2: + + def __init__(self, t): + super().__init__() + self._t = t + + def __getattr__(self, name): + meth = getattr(lazy_tensor_core._LAZYC, f"_dynamic_{name}") + + def wrapper(*args, **kwargs): + args = (self,) + args + # TODO: we need to unwrap arguments for kwargs as well + args = [x._t if isinstance(x, DynamicLazyTensor2) else x for x in args] + t = meth(*args, **kwargs) + # TODO: this may return a tuple, dict, etc. + return DynamicLazyTensor2(t) + + return wrapper + + +b = DynamicLazyTensor2(torch.rand(3, 4, 5).to(device='lazy')) +c = b.size() +d = b.expand(c) +#d = b.expand((3, 4, 5)) +print(lazy_tensor_core._LAZYC._get_ltc_tensors_text([d._t])) +print(lazy_tensor_core._LAZYC._get_ltc_tensors_backend([d._t])) +ltm.mark_step() + +b = DynamicLazyTensor2(torch.rand(1, 2, 3).to(device='lazy')) +c = b.size() +d = b.expand(c) +ltm.mark_step() +print("DONE!") + +# print(lazy_tensor_core._LAZYC._get_ltc_tensors_text([d._t])) +# print(lazy_tensor_core._LAZYC._get_ltc_tensors_backend([d._t])) + + +# _get_ltc_tensors_text _get_ltc_tensors_backend + +# Questions +# How do add dynamic_ and override them in LTC +# why can't I use `expand` + + # def expand(self, sz): + # return lazy_tensor_core._LAZYC._dynamic_expand(self._t, sz) + + # supported_methods = {'expand': } + \ No newline at end of file diff --git a/test/test_jit_fuser_te.py b/test/test_jit_fuser_te.py index 66194850b4e61..5f7e016ef314b 100644 --- a/test/test_jit_fuser_te.py +++ b/test/test_jit_fuser_te.py @@ -1260,6 +1260,7 @@ def apply(fn): F.hardtanh, F.hardsigmoid, F.hardswish, + F.tanhshrink, torch.sqrt, torch.rsqrt, F.gelu, From 82b59879e806f5661f41cfe6085a7423d2429708 Mon Sep 17 00:00:00 2001 From: Nikolay Korovaiko Date: Tue, 14 Sep 2021 06:00:56 +0000 Subject: [PATCH 2/6] pass simple tests --- .../lazy_tensor_core/csrc/tensor.cpp | 41 +++++++++++++------ .../csrc/ts_backend/ts_node_lowering.cpp | 23 ++++++++++- 2 files changed, 50 insertions(+), 14 deletions(-) diff --git a/lazy_tensor_core/lazy_tensor_core/csrc/tensor.cpp b/lazy_tensor_core/lazy_tensor_core/csrc/tensor.cpp index 0c887cc0d2628..eb1e2b4024834 100644 --- a/lazy_tensor_core/lazy_tensor_core/csrc/tensor.cpp +++ b/lazy_tensor_core/lazy_tensor_core/csrc/tensor.cpp @@ -39,6 +39,7 @@ #include "torch/csrc/autograd/variable.h" +#include "torch/torch.h" #include "torch/csrc/jit/runtime/custom_operator.h" int get_current_level() { @@ -1694,15 +1695,31 @@ lazy_tensors::uint64 LazyTensor::GetRunningSeed(const Device& device) { // torch::jit::push(stack, t.sizes()); // } -// const torch::jit::RegisterOperators DynamicSizeOp({ -// torch::jit::Operator( -// "aten::dynamic_size(Tensor a) -> int[]", -// [](const torch::jit::Node*) -> torch::jit::Operation { -// return [](torch::jit::Stack* stack) { -// at::Tensor t = torch::jit::pop(stack).toTensor(); -// torch::jit::push(stack, t.sizes()); -// }; -// }, -// //DynamicSize, -// c10::AliasAnalysisKind::FROM_SCHEMA), -// }); \ No newline at end of file +const torch::jit::RegisterOperators DynamicSizeOp({ + torch::jit::Operator( + "aten::dynamic_size(Tensor a) -> Tensor", + [](const torch::jit::Node*) -> torch::jit::Operation { + return [](torch::jit::Stack* stack) { + auto t = torch::jit::pop(stack).toTensor(); + auto sz_ten = torch::tensor(t.sizes(), c10::TensorOptions(c10::kLong)); + std::cerr << sz_ten; + torch::jit::push(stack, sz_ten); + }; + }, + c10::AliasAnalysisKind::FROM_SCHEMA), + torch::jit::Operator( + "prim::tensor_to_list(Tensor a) -> int[]", + [](const torch::jit::Node*) -> torch::jit::Operation { + return [](torch::jit::Stack* stack) { + auto t = torch::jit::pop(stack).toTensor(); + auto n = t.numel(); + std::vector r; + auto t_data = t.data(); + for (auto i : c10::irange(n)) { + r.push_back(t_data[i]); + } + torch::jit::push(stack, r); + }; + }, + c10::AliasAnalysisKind::FROM_SCHEMA), +}); \ No newline at end of file diff --git a/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/ts_node_lowering.cpp b/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/ts_node_lowering.cpp index 341841e3d467b..432348eaffd29 100644 --- a/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/ts_node_lowering.cpp +++ b/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/ts_node_lowering.cpp @@ -4,6 +4,7 @@ #include #include +#include "ATen/core/interned_strings.h" #include "lazy_tensor_core/csrc/compiler/node_lowering.h" #include "lazy_tensor_core/csrc/data_ops.h" #include "lazy_tensor_core/csrc/helpers.h" @@ -40,6 +41,7 @@ #include "lazy_tensor_core/csrc/ts_backend/ts_computation_client.h" #include "lazy_tensor_core/csrc/ts_backend/ts_lowering_context.h" #include "lazy_tensors/permutation_util.h" +#include "torch/csrc/jit/ir/constants.h" namespace torch_lazy_tensors { namespace compiler { @@ -308,7 +310,8 @@ class TSNodeLowering : public NodeLowering { } if (node->op().op == at::aten::size) { - auto size_val = loctx()->graph()->insert(at::aten::size, {loctx()->GetOutputOp(node->operands().at(0))}); + //auto size_val = loctx()->graph()->insert(at::aten::size, {loctx()->GetOutputOp(node->operands().at(0))}); + auto size_val = loctx()->graph()->insert(c10::Symbol::aten("dynamic_size"), {loctx()->GetOutputOp(node->operands().at(0))}); return {size_val}; } @@ -622,8 +625,24 @@ class TSNodeLowering : public NodeLowering { TSOpVector LowerExpand(const ir::ops::Expand* node) { std::vector arguments; + auto sz_val = loctx()->GetOutputOp(node->operand(1)); + if (sz_val->type()->cast()) { + // auto cnst_zero = loctx()->graph()->insertConstant(0); + // auto cnst_one = loctx()->graph()->insertConstant(1); + // std::vector values; + // values.push_back(sz_val); + // values.push_back(cnst_zero); + // values.push_back(cnst_one); + // auto n = loctx()->graph()->create(c10::Symbol::prim("tolist"), values, 1); + // loctx()->graph()->insertNode(n); + // sz_val = n->output(); + // sz_val->setType(c10::ListType::ofInts()); + sz_val = loctx()->graph()->insert(c10::Symbol::prim("tensor_to_list"), {sz_val}); + //loctx()->graph()->insertNode(c10::Symbol::prim("tolist"), {sz_val}); + + } arguments.emplace_back(loctx()->GetOutputOp(node->operand(0))); - arguments.emplace_back(loctx()->GetOutputOp(node->operand(1))); + arguments.emplace_back(sz_val); //arguments.emplace_back(node->size()); auto expand_out = LowerBuiltin(node, arguments); if (node->is_scalar_expand()) { From 13fa2c88147bd8b65c44dcf271be6ded83d45adc Mon Sep 17 00:00:00 2001 From: Nikolay Korovaiko Date: Tue, 21 Sep 2021 18:57:18 +0000 Subject: [PATCH 3/6] view --- .../csrc/init_python_bindings.cpp | 48 +++++++ .../lazy_tensor_core/csrc/tensor.cpp | 51 ++++++- .../csrc/ts_backend/ops/add.cpp | 45 +++++- .../csrc/ts_backend/ops/add.h | 36 +++++ .../csrc/ts_backend/ts_node_lowering.cpp | 55 +++++++- lazy_tensor_core/test/dynamic_lazy_tensor2.py | 133 ++++++++++++++++-- 6 files changed, 347 insertions(+), 21 deletions(-) diff --git a/lazy_tensor_core/lazy_tensor_core/csrc/init_python_bindings.cpp b/lazy_tensor_core/lazy_tensor_core/csrc/init_python_bindings.cpp index b272528d9fb8c..c2f47fba11b7e 100644 --- a/lazy_tensor_core/lazy_tensor_core/csrc/init_python_bindings.cpp +++ b/lazy_tensor_core/lazy_tensor_core/csrc/init_python_bindings.cpp @@ -7,6 +7,7 @@ #include #include +#include "ATen/core/functional.h" #include "lazy_tensor_core/csrc/aten_ltc_bridge.h" #include "lazy_tensor_core/csrc/compiler/backend_impl_interface.h" #include "lazy_tensor_core/csrc/device.h" @@ -33,6 +34,7 @@ #include "torch/csrc/jit/python/pybind.h" #include "torch/csrc/utils/cuda_lazy_init.h" #include "lazy_tensor_core/csrc/ts_backend/ops/add.h" +#include "torch/torch.h" namespace torch_lazy_tensors { namespace { @@ -423,6 +425,17 @@ void InitLtcModuleBindings(py::module m) { self_lazy_tensor.CreateFrom(ir::MakeNode( self_lazy_tensor.GetIrValue()))); }); + m.def("_add_dim", + [](at::Tensor& self, at::Tensor& other) { + + std::cerr << "in _add_dim\n"; + LazyTensor self_lazy_tensor = bridge::GetLtcTensor(self); + LazyTensor other_lazy_tensor = + bridge::GetOrCreateLtcTensor(other, self_lazy_tensor.GetDevice()); + return bridge::AtenFromLtcTensor( + self_lazy_tensor.CreateFrom(ir::MakeNode( + self_lazy_tensor.GetIrValue(), other_lazy_tensor.GetIrValue()))); + }); m.def("_dynamic_expand", [](at::Tensor& self, at::Tensor& other) { std::cerr << "in _expand\n"; @@ -433,6 +446,41 @@ void InitLtcModuleBindings(py::module m) { self_lazy_tensor.CreateFrom(ir::MakeNode( self_lazy_tensor.GetIrValue(), other_lazy_tensor.GetIrValue()))); }); + m.def("_dynamic_view", + [](std::vector& self_and_dims) { + std::cerr << "in _dynamic_view\n"; + auto self_lazy_tensor = bridge::GetLtcTensor(self_and_dims[0]); + auto ir_values = c10::fmap(self_and_dims, [&self_lazy_tensor](const at::Tensor& t) { + return bridge::GetOrCreateLtcTensor(t, self_lazy_tensor.GetDevice()).GetIrValue(); + }); + return bridge::AtenFromLtcTensor( + self_lazy_tensor.CreateFrom(ir::MakeNode(ir_values))); + }); + m.def("_dynamic_linear", + //TODO: figure out how to do optional bias + [](at::Tensor& self, at::Tensor& weight, at::Tensor& bias) { + std::cerr << "in _linear\n"; + LazyTensor self_lazy_tensor = bridge::GetLtcTensor(self); + LazyTensor weight_lazy_tensor = + bridge::GetOrCreateLtcTensor(weight, self_lazy_tensor.GetDevice()); + LazyTensor bias_lazy_tensor = + bridge::GetOrCreateLtcTensor(bias, self_lazy_tensor.GetDevice()); + return bridge::AtenFromLtcTensor( + self_lazy_tensor.CreateFrom(ir::MakeNode( + self_lazy_tensor.GetIrValue(), weight_lazy_tensor.GetIrValue(), bias_lazy_tensor.GetIrValue()))); + }); + m.def("_dynamic_getitem", + [](at::Tensor& self, int index) { + std::cerr << "in _expand\n"; + LazyTensor self_lazy_tensor = bridge::GetLtcTensor(self); + auto at_index_ten = torch::tensor({index}, c10::TensorOptions(c10::kLong)); + LazyTensor other_lazy_tensor = + bridge::GetOrCreateLtcTensor(at_index_ten, self_lazy_tensor.GetDevice()); + return bridge::AtenFromLtcTensor( + self_lazy_tensor.CreateFrom(ir::MakeNode( + self_lazy_tensor.GetIrValue(), other_lazy_tensor.GetIrValue()))); + }); + // IrValueFromScalar m.def("_get_ltc_tensors_text", [](const std::vector& tensors) -> std::string { auto coverter = [](lazy_tensors::Span nodes) { diff --git a/lazy_tensor_core/lazy_tensor_core/csrc/tensor.cpp b/lazy_tensor_core/lazy_tensor_core/csrc/tensor.cpp index eb1e2b4024834..06ad5730218c9 100644 --- a/lazy_tensor_core/lazy_tensor_core/csrc/tensor.cpp +++ b/lazy_tensor_core/lazy_tensor_core/csrc/tensor.cpp @@ -51,11 +51,10 @@ int get_current_level() { } std::ostream& get_ostream(int level) { - std::cerr << "in get_ostream\n"; static std::stringstream dummy{}; static const auto cur_level = get_current_level(); - if (true || level >= cur_level) { + if (level >= cur_level) { return std::cerr; } return dummy; @@ -1707,6 +1706,39 @@ const torch::jit::RegisterOperators DynamicSizeOp({ }; }, c10::AliasAnalysisKind::FROM_SCHEMA), + torch::jit::Operator( + "prim::list_to_tensor(int[] a) -> Tensor", + [](const torch::jit::Node*) -> torch::jit::Operation { + return [](torch::jit::Stack* stack) { + auto sz_vec = torch::jit::pop(stack).toIntVector(); + auto sz_ten = torch::tensor(sz_vec, c10::TensorOptions(c10::kLong)); + std::cerr << sz_ten; + torch::jit::push(stack, sz_ten); + }; + }, + c10::AliasAnalysisKind::FROM_SCHEMA), + torch::jit::Operator( + "prim::dim_to_tensor(int a) -> Tensor", + [](const torch::jit::Node*) -> torch::jit::Operation { + return [](torch::jit::Stack* stack) { + auto dim = torch::jit::pop(stack).toInt(); + auto sz_ten = torch::tensor({dim}, c10::TensorOptions(c10::kLong)); + std::cerr << sz_ten; + torch::jit::push(stack, sz_ten); + }; + }, + c10::AliasAnalysisKind::FROM_SCHEMA), + torch::jit::Operator( + "prim::dim_to_tensor(...) -> int[]", + [](const torch::jit::Node*) -> torch::jit::Operation { + return [](torch::jit::Stack* stack) { + auto dim = torch::jit::pop(stack).toInt(); + auto sz_ten = torch::tensor({dim}, c10::TensorOptions(c10::kLong)); + std::cerr << sz_ten; + torch::jit::push(stack, sz_ten); + }; + }, + c10::AliasAnalysisKind::FROM_SCHEMA), torch::jit::Operator( "prim::tensor_to_list(Tensor a) -> int[]", [](const torch::jit::Node*) -> torch::jit::Operation { @@ -1722,4 +1754,19 @@ const torch::jit::RegisterOperators DynamicSizeOp({ }; }, c10::AliasAnalysisKind::FROM_SCHEMA), + torch::jit::Operator( + "prim::dim_tensors_to_list(...) -> int[]", + [](const torch::jit::Node* n) -> torch::jit::Operation { + auto num_inputs = n->inputs().size(); + return [num_inputs](torch::jit::Stack* stack) { + std::vector dims; + auto ivals = torch::jit::last(stack, num_inputs); + for (auto iv : ivals) { + dims.push_back(iv.toTensor().item()); + } + torch::jit::drop(stack, num_inputs); + torch::jit::push(stack, dims); + }; + }, + c10::AliasAnalysisKind::FROM_SCHEMA), }); \ No newline at end of file diff --git a/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/ops/add.cpp b/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/ops/add.cpp index 75ced73deb1a0..9cc786a799b9b 100644 --- a/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/ops/add.cpp +++ b/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/ops/add.cpp @@ -16,8 +16,32 @@ NodePtr Add::Clone(OpList operands) const { return MakeNode(operands.at(0), operands.at(1)); } +DynamicView::DynamicView(lazy_tensors::Span values) + : Node(ir::OpKind(c10::Symbol::prim("_dynamic_view")), values, + /*num_outputs=*/1, /*hash_seed=*/0x5a2d296e9) {} + +NodePtr DynamicView::Clone(OpList operands) const { + return MakeNode(operands); +} + +AddDim::AddDim(const Value& lhs, const Value& rhs) + : Node(ir::OpKind(c10::Symbol::prim("_add_dim")), {lhs, rhs}, lazy_tensors::Shape{}, + /*num_outputs=*/1, /*hash_seed=*/0x5a2d296e9) {} + +NodePtr AddDim::Clone(OpList operands) const { + return MakeNode(operands.at(0), operands.at(1)); +} + +MulDim::MulDim(const Value& lhs, const Value& rhs) + : Node(ir::OpKind(c10::Symbol::prim("_mul_dim")), {lhs, rhs}, lazy_tensors::Shape{}, + /*num_outputs=*/1, /*hash_seed=*/0x5a2d296e9) {} + +NodePtr MulDim::Clone(OpList operands) const { + return MakeNode(operands.at(0), operands.at(1)); +} + DynamicExpand::DynamicExpand(const Value& lhs, const Value& rhs) - : Node(ir::OpKind(at::aten::expand), {lhs, rhs}, lhs.shape(), + : Node(ir::OpKind(c10::Symbol::prim("_dynamic_expand")), {lhs, rhs}, lhs.shape(), /*num_outputs=*/1, /*hash_seed=*/0x5a2d296e9) {} NodePtr DynamicExpand::Clone(OpList operands) const { @@ -25,13 +49,30 @@ NodePtr DynamicExpand::Clone(OpList operands) const { } DynamicSize::DynamicSize(const Value& lhs) - : Node(ir::OpKind(at::aten::size), {lhs}, lazy_tensors::Shape{}, + : Node(ir::OpKind(c10::Symbol::prim("_dynamic_size")), {lhs}, lazy_tensors::Shape{}, /*num_outputs=*/1, /*hash_seed=*/0x5a2d296e9) {} NodePtr DynamicSize::Clone(OpList operands) const { return MakeNode(operands.at(0)); } +// TODO: figure out how to do optional in LTC IR +DynamicLinear::DynamicLinear(const Value& input, const Value& weight, const Value& bias) + : Node(ir::OpKind(c10::Symbol::prim("_dynamic_linear")), {input, weight, bias}, input.shape(), + /*num_outputs=*/1, /*hash_seed=*/0x5a2d296e9) {} + +NodePtr DynamicLinear::Clone(OpList operands) const { + return MakeNode(operands.at(0), operands.at(1), operands.at(2)); +} + +DynamicGetItem::DynamicGetItem(const Value& lhs, const Value& rhs) + : Node(ir::OpKind(c10::Symbol::prim("_dynamic_getitem")), {lhs, rhs}, lazy_tensors::Shape{}, + /*num_outputs=*/1, /*hash_seed=*/0x5a2d296e9) {} + +NodePtr DynamicGetItem::Clone(OpList operands) const { + return MakeNode(operands.at(0), operands.at(1)); +} + } // namespace ops } // namespace ir } // namespace torch_lazy_tensors diff --git a/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/ops/add.h b/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/ops/add.h index f420823211a15..7cfae1952d3a7 100644 --- a/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/ops/add.h +++ b/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/ops/add.h @@ -14,6 +14,20 @@ class Add : public Node { NodePtr Clone(OpList operands) const override; }; +class AddDim : public Node { + public: + AddDim(const Value& lhs, const Value& rhs); + + NodePtr Clone(OpList operands) const override; +}; + +class MulDim : public Node { + public: + MulDim(const Value& lhs, const Value& rhs); + + NodePtr Clone(OpList operands) const override; +}; + class DynamicSize : public Node { public: DynamicSize(const Value& lhs); @@ -28,6 +42,28 @@ class DynamicExpand : public Node { NodePtr Clone(OpList operands) const override; }; +class DynamicLinear : public Node { + public: + DynamicLinear(const Value& input, const Value& weight, const Value& bias); + + NodePtr Clone(OpList operands) const override; +}; + +class DynamicGetItem : public Node { + public: + DynamicGetItem(const Value& lhs, const Value& rhs); + + NodePtr Clone(OpList operands) const override; +}; + +class DynamicView : public Node { + public: + DynamicView(lazy_tensors::Span values); + + NodePtr Clone(OpList operands) const override; + +}; + } // namespace ops } // namespace ir } // namespace torch_lazy_tensors diff --git a/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/ts_node_lowering.cpp b/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/ts_node_lowering.cpp index 432348eaffd29..cff41fa73cce3 100644 --- a/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/ts_node_lowering.cpp +++ b/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/ts_node_lowering.cpp @@ -309,12 +309,61 @@ class TSNodeLowering : public NodeLowering { return {loctx()->GetParameter(device_data_node->data())}; } - if (node->op().op == at::aten::size) { - //auto size_val = loctx()->graph()->insert(at::aten::size, {loctx()->GetOutputOp(node->operands().at(0))}); - auto size_val = loctx()->graph()->insert(c10::Symbol::aten("dynamic_size"), {loctx()->GetOutputOp(node->operands().at(0))}); + if (node->op().op == c10::Symbol::prim("_dynamic_size")) { + // We are replacing _dynamic_expand with at::aten::expand + auto size_val = loctx()->graph()->insert(at::aten::size, {loctx()->GetOutputOp(node->operands().at(0))}); + size_val = loctx()->graph()->insert(c10::Symbol::prim("list_to_tensor"), {size_val}); return {size_val}; } + if (node->op().op == c10::Symbol::prim("_dynamic_expand")) { + auto sz_val = loctx()->GetOutputOp(node->operand(1)); + sz_val = loctx()->graph()->insert(c10::Symbol::prim("tensor_to_list"), {sz_val}); + // We are replacing _dynamic_expand with at::aten::expand + auto expand = loctx()->graph()->insert(at::aten::expand, {loctx()->GetOutputOp(node->operand(0)), sz_val}); + // TODO: do we need to treat scalar expands differently? + return {expand}; + } + + if (node->op().op == c10::Symbol::prim("_dynamic_linear")) { + auto input = loctx()->GetOutputOp(node->operand(0)); + auto weight = loctx()->GetOutputOp(node->operand(1)); + auto bias = loctx()->GetOutputOp(node->operand(2)); + auto linear = loctx()->graph()->insert(at::aten::linear, {input, weight, bias}); + return {linear}; + } + + if (node->op().op == c10::Symbol::prim("_dynamic_getitem")) { + auto tensor_list = loctx()->GetOutputOp(node->operand(0)); + auto list_val = loctx()->graph()->insert(c10::Symbol::prim("tensor_to_list"), {tensor_list}); + auto tensor_index = loctx()->GetOutputOp(node->operand(1)); + auto index_val = loctx()->graph()->insert(at::aten::item, {tensor_index}); + auto dim_val = loctx()->graph()->insert(at::aten::__getitem__, {list_val, index_val}); + auto dim_tensor = loctx()->graph()->insert(c10::Symbol::prim("dim_to_tensor"), {dim_val}); + return {dim_tensor}; + } + + if (node->op().op == c10::Symbol::prim("_dynamic_view")) { + auto self = loctx()->GetOutputOp(node->operand(0)); + std::vector tensor_dims; + for (size_t i = 1; i < node->operands().size(); i++) { + tensor_dims.push_back(loctx()->GetOutputOp(node->operand(i))); + } + auto list_val = loctx()->graph()->insert(c10::Symbol::prim("dim_tensors_to_list"), tensor_dims); + auto view_val = loctx()->graph()->insert(at::aten::view, {self, list_val}); + return {view_val}; + } + + if (node->op().op == c10::Symbol::prim("_add_dim")) { + auto tensor_dim1 = loctx()->GetOutputOp(node->operand(0)); + auto dim1 = loctx()->graph()->insert(at::aten::item, {tensor_dim1}); + auto tensor_dim2 = loctx()->GetOutputOp(node->operand(1)); + auto dim2 = loctx()->graph()->insert(at::aten::item, {tensor_dim2}); + auto add = loctx()->graph()->insert(at::aten::add, {dim1, dim2}); + auto dim_tensor = loctx()->graph()->insert(c10::Symbol::prim("dim_to_tensor"), {add}); + return {dim_tensor}; + } + std::vector arguments; for (const ir::Output& output : node->operands()) { arguments.emplace_back(loctx()->GetOutputOp(output)); diff --git a/lazy_tensor_core/test/dynamic_lazy_tensor2.py b/lazy_tensor_core/test/dynamic_lazy_tensor2.py index d70936038083f..530a9e8d44286 100644 --- a/lazy_tensor_core/test/dynamic_lazy_tensor2.py +++ b/lazy_tensor_core/test/dynamic_lazy_tensor2.py @@ -2,46 +2,142 @@ import lazy_tensor_core import lazy_tensor_core.debug.metrics as metrics import lazy_tensor_core.core.lazy_model as ltm - - +from collections.abc import Iterable +from enum import Enum lazy_tensor_core._LAZYC._ltc_init_ts_backend() lazy_tensor_core._LAZYC._ltc_set_dynamic_shapes_mode() + +class Type(Enum): + TENSOR = 1 + SHAPE = 2 + DIMENSION = 3 + + +def promote_type(type1, type2): + pass + + +# DynamicLazyTensor2 has two fields: a cpp tensor object and its type +# In python we often need to know which type of a tensor we are dealing with: Tensor, List or Dim, so +# we can dispatch it correctly. +# We need to do it because the current type system in LTC only knows about tensors. class DynamicLazyTensor2: def __init__(self, t): super().__init__() self._t = t + self._type = Type.TENSOR + + def __str__(self): + return str(self._t.cpu()) + + # self should have _type == Tensor + # and `index` needs to be either _type == DIMENSION or python int + # which we will wrap in a lazy tensor object. + def __getitem__(self, index): + print("running __getitem__") + if isinstance(index, Iterable): + assert(not self._t.shape) + else: + # TODO: index should also be probably symbolic + t = lazy_tensor_core._LAZYC._dynamic_getitem(self._t, index) + wt = DynamicLazyTensor2(t) + wt._type = Type.DIMENSION + return wt - def __getattr__(self, name): + # N.B. For arithmetic operations, we essentialy have two cases add.tensor and add.dim + # in c++ land we can't easily tell which one we need to do (TODO: check if we could use shape to tell tensor vs scalar case) + # so we will have to use two different opcodes: add and _add_dim to tell the difference and lower arithmetic ops correctly. + def add(self, other): + assert(self._type == other._type) + if self._type == Type.TENSOR: + t = self._t.add(other._t) + return DynamicLazyTensor2(t) + elif self._type == Type.DIMENSION: + t = lazy_tensor_core._LAZYC._add_dim(self._t, other._t) + return DynamicLazyTensor2(t) + + + # again we keep track if `dims` are actually tensors that represent individual dimensions. + def view(*args): + self = args[0] + dims = args[1:] + # TODO we should be wrapping any ints into tensors + assert(all([x._type == Type.DIMENSION for x in dims])) + stripped_wrappers = [x._t for x in args] + t = lazy_tensor_core._LAZYC._dynamic_view(stripped_wrappers) + return DynamicLazyTensor2(t) + + # as much of dispatching logic above should generalized and folded into __getattr__ + def __getattr__(self, name): meth = getattr(lazy_tensor_core._LAZYC, f"_dynamic_{name}") - + def wrapper(*args, **kwargs): args = (self,) + args # TODO: we need to unwrap arguments for kwargs as well args = [x._t if isinstance(x, DynamicLazyTensor2) else x for x in args] t = meth(*args, **kwargs) # TODO: this may return a tuple, dict, etc. - return DynamicLazyTensor2(t) + wt = DynamicLazyTensor2(t) + if name == 'size': + wt._type = Type.SHAPE + return wt return wrapper -b = DynamicLazyTensor2(torch.rand(3, 4, 5).to(device='lazy')) +b = DynamicLazyTensor2((torch.ones(1, 10) * 777).to(device='lazy')) +w = DynamicLazyTensor2((torch.ones(10, 10) * 777).to(device='lazy')) +#dim1 = DynamicLazyTensor2((torch.tensor([10], dtype=torch.int64)).to(device='lazy')) +#dim2 = DynamicLazyTensor2((torch.tensor([1], dtype=torch.int64)).to(device='lazy')) +#dim1._type = Type.DIMENSION +#dim2._type = Type.DIMENSION c = b.size() -d = b.expand(c) -#d = b.expand((3, 4, 5)) -print(lazy_tensor_core._LAZYC._get_ltc_tensors_text([d._t])) -print(lazy_tensor_core._LAZYC._get_ltc_tensors_backend([d._t])) +dim = c[1] +dim3 = dim.add(c[0]) +e = b.view(c[0], c[1]) +print(dim._type) +print(lazy_tensor_core._LAZYC._get_ltc_tensors_text([e._t])) +print(lazy_tensor_core._LAZYC._get_ltc_tensors_backend([e._t])) ltm.mark_step() -b = DynamicLazyTensor2(torch.rand(1, 2, 3).to(device='lazy')) +b = DynamicLazyTensor2((torch.ones(1, 10) * 777).to(device='lazy')) +w = DynamicLazyTensor2((torch.ones(10, 10) * 777).to(device='lazy')) +#dim1 = DynamicLazyTensor2((torch.tensor([10], dtype=torch.int64)).to(device='lazy')) +#dim2 = DynamicLazyTensor2((torch.tensor([1], dtype=torch.int64)).to(device='lazy')) +#dim1._type = Type.DIMENSION +#dim2._type = Type.DIMENSION c = b.size() -d = b.expand(c) +dim = c[1] +dim3 = dim.add(c[0]) +e = b.view(c[0], c[1]) +print(dim._type) +print(lazy_tensor_core._LAZYC._get_ltc_tensors_text([e._t])) +print(lazy_tensor_core._LAZYC._get_ltc_tensors_backend([e._t])) ltm.mark_step() -print("DONE!") + +# b = DynamicLazyTensor2((torch.ones(1, 10) * 777).to(device='lazy')) +# w = DynamicLazyTensor2((torch.ones(10, 10) * 777).to(device='lazy')) +# c = b.size() +# print(c._type) +# d = b.expand(c) +# e = d.linear(w, b) # TODO: figure out how to override torch.nn.linear +# print(lazy_tensor_core._LAZYC._get_ltc_tensors_text([e._t])) +# print(lazy_tensor_core._LAZYC._get_ltc_tensors_backend([e._t])) +# ltm.mark_step() + +# b = DynamicLazyTensor2((torch.ones(1, 10) * 777).to(device='lazy')) +# w = DynamicLazyTensor2((torch.ones(10, 10) * 777).to(device='lazy')) +# c = b.size() +# print(c._type) +# d = b.expand(c) +# e = d.linear(w, b) # TODO: figure out how to override torch.nn.linear +# print(lazy_tensor_core._LAZYC._get_ltc_tensors_text([e._t])) +# print(lazy_tensor_core._LAZYC._get_ltc_tensors_backend([e._t])) +# ltm.mark_step() # print(lazy_tensor_core._LAZYC._get_ltc_tensors_text([d._t])) # print(lazy_tensor_core._LAZYC._get_ltc_tensors_backend([d._t])) @@ -57,4 +153,13 @@ def wrapper(*args, **kwargs): # return lazy_tensor_core._LAZYC._dynamic_expand(self._t, sz) # supported_methods = {'expand': } - \ No newline at end of file + + + +# graph(%p0 : Tensor): +# %1 : int[] = aten::size(%p0) +# %2 : Tensor = prim::list_to_tensor(%1) +# %3 : int[] = prim::tensor_to_list(%2) +# %4 : bool = prim::Constant[value=0]() +# %5 : Tensor = aten::expand(%p0, %3, %4) +# return (%5) \ No newline at end of file From 5651a3c44e9e25282cc549ec8ca5ebb3e72ae967 Mon Sep 17 00:00:00 2001 From: Nikolay Korovaiko Date: Tue, 21 Sep 2021 22:56:51 +0000 Subject: [PATCH 4/6] remove prints --- .../lazy_tensor_core/csrc/init_python_bindings.cpp | 10 +--------- 1 file changed, 1 insertion(+), 9 deletions(-) diff --git a/lazy_tensor_core/lazy_tensor_core/csrc/init_python_bindings.cpp b/lazy_tensor_core/lazy_tensor_core/csrc/init_python_bindings.cpp index c2f47fba11b7e..fddfe3fac96e7 100644 --- a/lazy_tensor_core/lazy_tensor_core/csrc/init_python_bindings.cpp +++ b/lazy_tensor_core/lazy_tensor_core/csrc/init_python_bindings.cpp @@ -418,8 +418,6 @@ void InitLtcModuleBindings(py::module m) { }); m.def("_dynamic_size", [](at::Tensor& self) { - - std::cerr << "in _dynamic_size\n"; LazyTensor self_lazy_tensor = bridge::GetLtcTensor(self); return bridge::AtenFromLtcTensor( self_lazy_tensor.CreateFrom(ir::MakeNode( @@ -427,8 +425,6 @@ void InitLtcModuleBindings(py::module m) { }); m.def("_add_dim", [](at::Tensor& self, at::Tensor& other) { - - std::cerr << "in _add_dim\n"; LazyTensor self_lazy_tensor = bridge::GetLtcTensor(self); LazyTensor other_lazy_tensor = bridge::GetOrCreateLtcTensor(other, self_lazy_tensor.GetDevice()); @@ -438,7 +434,6 @@ void InitLtcModuleBindings(py::module m) { }); m.def("_dynamic_expand", [](at::Tensor& self, at::Tensor& other) { - std::cerr << "in _expand\n"; LazyTensor self_lazy_tensor = bridge::GetLtcTensor(self); LazyTensor other_lazy_tensor = bridge::GetOrCreateLtcTensor(other, self_lazy_tensor.GetDevice()); @@ -448,7 +443,6 @@ void InitLtcModuleBindings(py::module m) { }); m.def("_dynamic_view", [](std::vector& self_and_dims) { - std::cerr << "in _dynamic_view\n"; auto self_lazy_tensor = bridge::GetLtcTensor(self_and_dims[0]); auto ir_values = c10::fmap(self_and_dims, [&self_lazy_tensor](const at::Tensor& t) { return bridge::GetOrCreateLtcTensor(t, self_lazy_tensor.GetDevice()).GetIrValue(); @@ -457,9 +451,8 @@ void InitLtcModuleBindings(py::module m) { self_lazy_tensor.CreateFrom(ir::MakeNode(ir_values))); }); m.def("_dynamic_linear", - //TODO: figure out how to do optional bias + //TODO: figure out how to do optional bias [](at::Tensor& self, at::Tensor& weight, at::Tensor& bias) { - std::cerr << "in _linear\n"; LazyTensor self_lazy_tensor = bridge::GetLtcTensor(self); LazyTensor weight_lazy_tensor = bridge::GetOrCreateLtcTensor(weight, self_lazy_tensor.GetDevice()); @@ -471,7 +464,6 @@ void InitLtcModuleBindings(py::module m) { }); m.def("_dynamic_getitem", [](at::Tensor& self, int index) { - std::cerr << "in _expand\n"; LazyTensor self_lazy_tensor = bridge::GetLtcTensor(self); auto at_index_ten = torch::tensor({index}, c10::TensorOptions(c10::kLong)); LazyTensor other_lazy_tensor = From 73955237c6eea99ec0998fb3e215c19851ae1750 Mon Sep 17 00:00:00 2001 From: Nikolay Korovaiko Date: Wed, 22 Sep 2021 16:47:13 +0000 Subject: [PATCH 5/6] remove dead code --- lazy_tensor_core/example_size.py | 30 ---------- lazy_tensor_core/jie.py | 36 ------------ lazy_tensor_core/jie_empty.py | 34 ------------ .../csrc/ts_backend/ts_node_lowering.cpp | 19 +------ lazy_tensor_core/lazy_tensors/shape.h | 8 +-- lazy_tensor_core/test/dynamic_lazy_tensor.py | 55 ------------------- test/test_jit_fuser_te.py | 1 - 7 files changed, 3 insertions(+), 180 deletions(-) delete mode 100644 lazy_tensor_core/example_size.py delete mode 100644 lazy_tensor_core/jie.py delete mode 100644 lazy_tensor_core/jie_empty.py delete mode 100644 lazy_tensor_core/test/dynamic_lazy_tensor.py diff --git a/lazy_tensor_core/example_size.py b/lazy_tensor_core/example_size.py deleted file mode 100644 index 3d3de837c83a0..0000000000000 --- a/lazy_tensor_core/example_size.py +++ /dev/null @@ -1,30 +0,0 @@ -import torch -import lazy_tensor_core -import lazy_tensor_core.debug.metrics as metrics -import lazy_tensor_core.core.lazy_model as ltm - -lazy_tensor_core._LAZYC._ltc_init_ts_backend() - -import os -path = os.path.abspath(lazy_tensor_core.__file__) -print(path) -torch.manual_seed(42) - -device = 'lazy' -dtype = torch.float32 - - -x = torch.tensor([[0.1, 1.2], [2.2, 3.1], [4.9, 5.2]], device=device) -y = torch.tensor([[0.1, 1.2], [2.2, 3.1], [4.9, 5.2]], device=device) -z = x.normal_() -w = z + y -#z = x + y -#w = x + z.size()[1] -ltm.mark_step() -#print(w) - - - -#print(metrics.metrics_report()) - - diff --git a/lazy_tensor_core/jie.py b/lazy_tensor_core/jie.py deleted file mode 100644 index a63292d8d1def..0000000000000 --- a/lazy_tensor_core/jie.py +++ /dev/null @@ -1,36 +0,0 @@ -import torch -import lazy_tensor_core -import lazy_tensor_core.debug.metrics as metrics -import lazy_tensor_core.core.lazy_model as ltm - - - -lazy_tensor_core._LAZYC._ltc_init_ts_backend() - -# from caffe2.python import workspace -# workspace.GlobalInit(['caffe2', '--caffe2_log_level=-1000']) - -lazy_tensor_core._LAZYC._ltc_set_dynamic_shapes_mode() -device = 'lazy' -dtype = torch.float32 - -x1 = torch.randn(2, 4, device=device, dtype=dtype) -ltm.mark_step() -x1 = torch.randn(2, 6, device=device, dtype=dtype) -ltm.mark_step() -# x2 = torch.randn(2, 8, 8, device=device, dtype=dtype) - -# def f(x): -# o = x + 1.0 -# o = torch.nn.functional.gelu(o) -# return o - -# for i in range(2): -# o = f(x1) -# ltm.mark_step() - -# o = f(x2) -# ltm.mark_step() - -print(metrics.metrics_report()) - diff --git a/lazy_tensor_core/jie_empty.py b/lazy_tensor_core/jie_empty.py deleted file mode 100644 index 8cd9e1d6d24a8..0000000000000 --- a/lazy_tensor_core/jie_empty.py +++ /dev/null @@ -1,34 +0,0 @@ -import torch -import lazy_tensor_core -import lazy_tensor_core.debug.metrics as metrics -import lazy_tensor_core.core.lazy_model as ltm - - - -lazy_tensor_core._LAZYC._ltc_init_ts_backend() - -# from caffe2.python import workspace -# workspace.GlobalInit(['caffe2', '--caffe2_log_level=-1000']) - -lazy_tensor_core._LAZYC._ltc_set_dynamic_shapes_mode() -device = 'lazy' -dtype = torch.float32 - -x1 = torch.empty(2, 4, device=device, dtype=dtype) -ltm.mark_step() -# x2 = torch.randn(2, 8, 8, device=device, dtype=dtype) - -# def f(x): -# o = x + 1.0 -# o = torch.nn.functional.gelu(o) -# return o - -# for i in range(2): -# o = f(x1) -# ltm.mark_step() - -# o = f(x2) -# ltm.mark_step() - -print(metrics.metrics_report()) - diff --git a/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/ts_node_lowering.cpp b/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/ts_node_lowering.cpp index cff41fa73cce3..9e3fb65ee10ef 100644 --- a/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/ts_node_lowering.cpp +++ b/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/ts_node_lowering.cpp @@ -674,25 +674,8 @@ class TSNodeLowering : public NodeLowering { TSOpVector LowerExpand(const ir::ops::Expand* node) { std::vector arguments; - auto sz_val = loctx()->GetOutputOp(node->operand(1)); - if (sz_val->type()->cast()) { - // auto cnst_zero = loctx()->graph()->insertConstant(0); - // auto cnst_one = loctx()->graph()->insertConstant(1); - // std::vector values; - // values.push_back(sz_val); - // values.push_back(cnst_zero); - // values.push_back(cnst_one); - // auto n = loctx()->graph()->create(c10::Symbol::prim("tolist"), values, 1); - // loctx()->graph()->insertNode(n); - // sz_val = n->output(); - // sz_val->setType(c10::ListType::ofInts()); - sz_val = loctx()->graph()->insert(c10::Symbol::prim("tensor_to_list"), {sz_val}); - //loctx()->graph()->insertNode(c10::Symbol::prim("tolist"), {sz_val}); - - } arguments.emplace_back(loctx()->GetOutputOp(node->operand(0))); - arguments.emplace_back(sz_val); - //arguments.emplace_back(node->size()); + arguments.emplace_back(node->size()); auto expand_out = LowerBuiltin(node, arguments); if (node->is_scalar_expand()) { // The aten::expand operations sets all strides to 0 when the original is diff --git a/lazy_tensor_core/lazy_tensors/shape.h b/lazy_tensor_core/lazy_tensors/shape.h index acd006151a8ef..0625d24c930e3 100644 --- a/lazy_tensor_core/lazy_tensors/shape.h +++ b/lazy_tensor_core/lazy_tensors/shape.h @@ -74,9 +74,7 @@ class Shape { int dimensions_size() const { return dimensions_.size(); } int64 dimensions(int index) const { if (dynamic_mode_.load()) { - LTC_LOG(INFO) << "@@@ 1 dimensions shape is " << this->ToString() << std::endl; - //std::cerr << "@@@ 1 dimensions shape is " << this->ToString() << std::endl; - //throw std::runtime_error("Exact shape not known"); + throw std::runtime_error("Exact shape not known"); } LTC_CHECK_LT(index, dimensions_.size()); return dimensions_[index]; @@ -89,9 +87,7 @@ class Shape { lazy_tensors::Span dimensions() const { if (dynamic_mode_.load()) { - LTC_LOG(INFO) << "@@@ 2 dimensions shape is " << this->ToString() << std::endl; - //std::cerr << "@@@ 2 dimensions shape is " << this->ToString() << std::endl; - //throw std::runtime_error("Exact shape not known"); + throw std::runtime_error("Exact shape not known"); } return MakeSpan(dimensions_); } diff --git a/lazy_tensor_core/test/dynamic_lazy_tensor.py b/lazy_tensor_core/test/dynamic_lazy_tensor.py deleted file mode 100644 index bf1b1eb2c17dc..0000000000000 --- a/lazy_tensor_core/test/dynamic_lazy_tensor.py +++ /dev/null @@ -1,55 +0,0 @@ -import torch -import lazy_tensor_core -import lazy_tensor_core.debug.metrics as metrics -import lazy_tensor_core.core.lazy_model as ltm - - - -lazy_tensor_core._LAZYC._ltc_init_ts_backend() - -# class DynamicLazyTensor(torch.Tensor): - -# def __init__(*args, **kwargs): -# super().__init__(*args, **kwargs) - -# def __torch_function__(self, func, types, args=(), kwargs=None): -# return lazy_tensor_core._LAZYC._dynamic_size(self) - - -# a = DynamicLazyTensor([2.0, 4.0]) -# a.size() - - -class DynamicLazyTensor2(torch.Tensor): - - def __init__(self, t): - super().__init__() - self._t = t.to(device='lazy') - print(id(self)) - - def expand(self, sz): - return lazy_tensor_core._LAZYC._dynamic_expand(self._t, sz) - - def __torch_function__(self, func, types, args=(), kwargs=None): - if func.__name__ == 'size': - return lazy_tensor_core._LAZYC._dynamic_size(self._t) - elif func.__name__ == 'expand': - print(type(args[0]).__name__) - print(type(args[1]).__name__) - return lazy_tensor_core._LAZYC._dynamic_expand(self._t, args[0]) - else: - print("Not Implemented!") - - -b = DynamicLazyTensor2(torch.rand(3, 4, 5)) -c = b.size() -d = b.expand(c) -#d = b.expand((3, 4, 5)) - -# _get_ltc_tensors_text _get_ltc_tensors_backend -print(lazy_tensor_core._LAZYC._get_ltc_tensors_text([d])) -print(lazy_tensor_core._LAZYC._get_ltc_tensors_backend([d])) - -# Questions -# How do add dynamic_ and override them in LTC -# why can't I use `expand` \ No newline at end of file diff --git a/test/test_jit_fuser_te.py b/test/test_jit_fuser_te.py index 5f7e016ef314b..66194850b4e61 100644 --- a/test/test_jit_fuser_te.py +++ b/test/test_jit_fuser_te.py @@ -1260,7 +1260,6 @@ def apply(fn): F.hardtanh, F.hardsigmoid, F.hardswish, - F.tanhshrink, torch.sqrt, torch.rsqrt, F.gelu, From 1678573b57660b0c3e48772b3c6bffda3f37a4ad Mon Sep 17 00:00:00 2001 From: Nikolay Korovaiko Date: Wed, 29 Sep 2021 17:08:30 +0000 Subject: [PATCH 6/6] using IR nodes --- .../csrc/init_python_bindings.cpp | 13 +++ .../lazy_tensor_core/csrc/tensor_impl.cpp | 7 ++ .../csrc/ts_backend/ops/add.cpp | 16 ++++ .../csrc/ts_backend/ops/add.h | 14 +++ .../csrc/ts_backend/ts_node_lowering.cpp | 11 +++ lazy_tensor_core/test/dynamic_lazy_tensor2.py | 54 ++++++------ lazy_tensor_core/test/dynamic_lazy_tensor3.py | 87 +++++++++++++++++++ 7 files changed, 176 insertions(+), 26 deletions(-) create mode 100644 lazy_tensor_core/test/dynamic_lazy_tensor3.py diff --git a/lazy_tensor_core/lazy_tensor_core/csrc/init_python_bindings.cpp b/lazy_tensor_core/lazy_tensor_core/csrc/init_python_bindings.cpp index fddfe3fac96e7..0dafc99c9033f 100644 --- a/lazy_tensor_core/lazy_tensor_core/csrc/init_python_bindings.cpp +++ b/lazy_tensor_core/lazy_tensor_core/csrc/init_python_bindings.cpp @@ -423,6 +423,18 @@ void InitLtcModuleBindings(py::module m) { self_lazy_tensor.CreateFrom(ir::MakeNode( self_lazy_tensor.GetIrValue()))); }); + m.def("_dynamic_size2", + [](at::Tensor& self) { + LazyTensor self_lazy_tensor = bridge::GetLtcTensor(self); + return ir::MakeNode(self_lazy_tensor.GetIrValue()); + }); + m.def("_dynamic_expand2", + [](at::Tensor& self, std::shared_ptr val) { + LazyTensor self_lazy_tensor = bridge::GetLtcTensor(self); + return bridge::AtenFromLtcTensor( + self_lazy_tensor.CreateFrom(ir::MakeNode( + self_lazy_tensor.GetIrValue(),val))); + }); m.def("_add_dim", [](at::Tensor& self, at::Tensor& other) { LazyTensor self_lazy_tensor = bridge::GetLtcTensor(self); @@ -551,6 +563,7 @@ void InitLtcModuleBindings(py::module m) { }); py::class_>(m, "IrValue"); + py::class_>(m, "IrNode"); m.def("_ltc_create_token", [](const std::string& device) { return CreateToken(device); }); m.def("_ltc_all_reduce_inplace", [](const std::string& reduce_type, diff --git a/lazy_tensor_core/lazy_tensor_core/csrc/tensor_impl.cpp b/lazy_tensor_core/lazy_tensor_core/csrc/tensor_impl.cpp index ad2afc63cc1c0..ce889c9a4483f 100644 --- a/lazy_tensor_core/lazy_tensor_core/csrc/tensor_impl.cpp +++ b/lazy_tensor_core/lazy_tensor_core/csrc/tensor_impl.cpp @@ -107,6 +107,13 @@ void LTCTensorImpl::shallow_copy_from( } at::IntArrayRef LTCTensorImpl::sizes() const { + // get data directly from a tensor if it was materialized + // this would be used if the next op is a fallback + // and this tensor is an input to the op + auto opt_ten = tensor_.CurrentTensorData(); + if (opt_ten) { + return opt_ten->sizes(); + } const_cast(this)->SetupSizeProperties(); return c10::TensorImpl::sizes(); } diff --git a/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/ops/add.cpp b/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/ops/add.cpp index 9cc786a799b9b..a7cdf150915a4 100644 --- a/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/ops/add.cpp +++ b/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/ops/add.cpp @@ -48,6 +48,14 @@ NodePtr DynamicExpand::Clone(OpList operands) const { return MakeNode(operands.at(0), operands.at(1)); } +DynamicExpand2::DynamicExpand2(const Value& lhs, const Value& rhs) + : Node(ir::OpKind(c10::Symbol::prim("_dynamic_expand2")), {lhs, rhs}, lhs.shape(), + /*num_outputs=*/1, /*hash_seed=*/0x5a2d296e9) {} + +NodePtr DynamicExpand2::Clone(OpList operands) const { + return MakeNode(operands.at(0), operands.at(1)); +} + DynamicSize::DynamicSize(const Value& lhs) : Node(ir::OpKind(c10::Symbol::prim("_dynamic_size")), {lhs}, lazy_tensors::Shape{}, /*num_outputs=*/1, /*hash_seed=*/0x5a2d296e9) {} @@ -56,6 +64,14 @@ NodePtr DynamicSize::Clone(OpList operands) const { return MakeNode(operands.at(0)); } +DynamicSize2::DynamicSize2(const Value& lhs) + : Node(ir::OpKind(c10::Symbol::prim("_dynamic_size2")), {lhs}, lazy_tensors::Shape{}, + /*num_outputs=*/1, /*hash_seed=*/0x5a2d296e9) {} + +NodePtr DynamicSize2::Clone(OpList operands) const { + return MakeNode(operands.at(0)); +} + // TODO: figure out how to do optional in LTC IR DynamicLinear::DynamicLinear(const Value& input, const Value& weight, const Value& bias) : Node(ir::OpKind(c10::Symbol::prim("_dynamic_linear")), {input, weight, bias}, input.shape(), diff --git a/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/ops/add.h b/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/ops/add.h index 7cfae1952d3a7..5bcedbb00769b 100644 --- a/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/ops/add.h +++ b/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/ops/add.h @@ -35,6 +35,13 @@ class DynamicSize : public Node { NodePtr Clone(OpList operands) const override; }; +class DynamicSize2 : public Node { + public: + DynamicSize2(const Value& lhs); + + NodePtr Clone(OpList operands) const override; +}; + class DynamicExpand : public Node { public: DynamicExpand(const Value& lhs, const Value& sz); @@ -42,6 +49,13 @@ class DynamicExpand : public Node { NodePtr Clone(OpList operands) const override; }; +class DynamicExpand2 : public Node { + public: + DynamicExpand2(const Value& lhs, const Value& sz); + + NodePtr Clone(OpList operands) const override; +}; + class DynamicLinear : public Node { public: DynamicLinear(const Value& input, const Value& weight, const Value& bias); diff --git a/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/ts_node_lowering.cpp b/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/ts_node_lowering.cpp index 9e3fb65ee10ef..6721e45d08e68 100644 --- a/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/ts_node_lowering.cpp +++ b/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/ts_node_lowering.cpp @@ -316,6 +316,11 @@ class TSNodeLowering : public NodeLowering { return {size_val}; } + if (node->op().op == c10::Symbol::prim("_dynamic_size2")) { + auto size_val = loctx()->graph()->insert(at::aten::size, {loctx()->GetOutputOp(node->operands().at(0))}); + return {size_val}; + } + if (node->op().op == c10::Symbol::prim("_dynamic_expand")) { auto sz_val = loctx()->GetOutputOp(node->operand(1)); sz_val = loctx()->graph()->insert(c10::Symbol::prim("tensor_to_list"), {sz_val}); @@ -325,6 +330,12 @@ class TSNodeLowering : public NodeLowering { return {expand}; } + if (node->op().op == c10::Symbol::prim("_dynamic_expand2")) { + auto sz_val = loctx()->GetOutputOp(node->operand(1)); + auto expand = loctx()->graph()->insert(at::aten::expand, {loctx()->GetOutputOp(node->operand(0)), sz_val}); + return {expand}; + } + if (node->op().op == c10::Symbol::prim("_dynamic_linear")) { auto input = loctx()->GetOutputOp(node->operand(0)); auto weight = loctx()->GetOutputOp(node->operand(1)); diff --git a/lazy_tensor_core/test/dynamic_lazy_tensor2.py b/lazy_tensor_core/test/dynamic_lazy_tensor2.py index 530a9e8d44286..88f785e92d3f6 100644 --- a/lazy_tensor_core/test/dynamic_lazy_tensor2.py +++ b/lazy_tensor_core/test/dynamic_lazy_tensor2.py @@ -58,7 +58,9 @@ def add(self, other): return DynamicLazyTensor2(t) elif self._type == Type.DIMENSION: t = lazy_tensor_core._LAZYC._add_dim(self._t, other._t) - return DynamicLazyTensor2(t) + wt = DynamicLazyTensor2(t) + wt._type = Type.DIMENSION + return wt # again we keep track if `dims` are actually tensors that represent individual dimensions. @@ -89,35 +91,35 @@ def wrapper(*args, **kwargs): return wrapper -b = DynamicLazyTensor2((torch.ones(1, 10) * 777).to(device='lazy')) -w = DynamicLazyTensor2((torch.ones(10, 10) * 777).to(device='lazy')) -#dim1 = DynamicLazyTensor2((torch.tensor([10], dtype=torch.int64)).to(device='lazy')) -#dim2 = DynamicLazyTensor2((torch.tensor([1], dtype=torch.int64)).to(device='lazy')) -#dim1._type = Type.DIMENSION -#dim2._type = Type.DIMENSION -c = b.size() +b = DynamicLazyTensor2((torch.ones(1, 10, device='cuda') * 777).to(device='lazy')) +w = DynamicLazyTensor2((torch.ones(10, 10, device='cuda') * 777).to(device='lazy')) +dummy5 = DynamicLazyTensor2((torch.ones(5, 1, device='cuda') * 777).to(device='lazy')) +c = w.size() dim = c[1] -dim3 = dim.add(c[0]) -e = b.view(c[0], c[1]) +dim20 = dim.add(c[0]) +w2 = w.view(dummy5.size()[0], dim20) +e = w.view(c[1], c[0]) +f = b.linear(w, b) print(dim._type) -print(lazy_tensor_core._LAZYC._get_ltc_tensors_text([e._t])) -print(lazy_tensor_core._LAZYC._get_ltc_tensors_backend([e._t])) +print(lazy_tensor_core._LAZYC._get_ltc_tensors_text([f._t])) +print(lazy_tensor_core._LAZYC._get_ltc_tensors_backend([f._t])) ltm.mark_step() -b = DynamicLazyTensor2((torch.ones(1, 10) * 777).to(device='lazy')) -w = DynamicLazyTensor2((torch.ones(10, 10) * 777).to(device='lazy')) -#dim1 = DynamicLazyTensor2((torch.tensor([10], dtype=torch.int64)).to(device='lazy')) -#dim2 = DynamicLazyTensor2((torch.tensor([1], dtype=torch.int64)).to(device='lazy')) -#dim1._type = Type.DIMENSION -#dim2._type = Type.DIMENSION -c = b.size() -dim = c[1] -dim3 = dim.add(c[0]) -e = b.view(c[0], c[1]) -print(dim._type) -print(lazy_tensor_core._LAZYC._get_ltc_tensors_text([e._t])) -print(lazy_tensor_core._LAZYC._get_ltc_tensors_backend([e._t])) -ltm.mark_step() +# b = DynamicLazyTensor2((torch.ones(1, 10, device='cuda') * 777).to(device='lazy')) +# w = DynamicLazyTensor2((torch.ones(10, 10, device='cuda') * 777).to(device='lazy')) +# dummy5 = DynamicLazyTensor2((torch.ones(5, 1, device='cuda') * 777).to(device='lazy')) +# c = w.size() +# dim = c[1] +# print("adding") +# dim20 = dim.add(c[0]) +# print("end adding") +# w2 = w.view(dummy5.size()[0], dim20) +# e = w.view(c[1], c[0]) +# f = b.linear(w, b) +# print(dim._type) +# print(lazy_tensor_core._LAZYC._get_ltc_tensors_text([f._t])) +# print(lazy_tensor_core._LAZYC._get_ltc_tensors_backend([f._t])) +# ltm.mark_step() # b = DynamicLazyTensor2((torch.ones(1, 10) * 777).to(device='lazy')) # w = DynamicLazyTensor2((torch.ones(10, 10) * 777).to(device='lazy')) diff --git a/lazy_tensor_core/test/dynamic_lazy_tensor3.py b/lazy_tensor_core/test/dynamic_lazy_tensor3.py new file mode 100644 index 0000000000000..2e0f111af23f4 --- /dev/null +++ b/lazy_tensor_core/test/dynamic_lazy_tensor3.py @@ -0,0 +1,87 @@ +import torch +import lazy_tensor_core +import lazy_tensor_core.debug.metrics as metrics +import lazy_tensor_core.core.lazy_model as ltm +from collections.abc import Iterable +from enum import Enum + +lazy_tensor_core._LAZYC._ltc_init_ts_backend() +lazy_tensor_core._LAZYC._ltc_set_dynamic_shapes_mode() + + +# class LazyConv2d(torch.nn.Module): +# def __init__(self, in_chan, out_chan, kernel_size, stride, padding, dilation, groups = 1, bias = True, padding_mode='zeros'): +# super(Sub, self).__init__() +# self.weight = nn.Parameter(torch.randn(2)) + + +class Type(Enum): + TENSOR = 1 + SHAPE = 2 + DIMENSION = 3 + + +# DynamicLazyTensor2 has two fields: a cpp tensor object and its type +# In python we often need to know which type of a tensor we are dealing with: Tensor, List or Dim, so +# we can dispatch it correctly. +# We need to do it because the current type system in LTC only knows about tensors. +class DynamicLazyTensor3: + + def __init__(self, t = None, n = None): + super().__init__() + self._t = t + self._n = n + self._type = Type.TENSOR + + + def size(self): + node = lazy_tensor_core._LAZYC._dynamic_size2(self._t) + wt = DynamicLazyTensor3(n = node) + return wt + + def expand(self, size): + t = lazy_tensor_core._LAZYC._dynamic_expand2(self._t, size._n) + wt = DynamicLazyTensor3(t = t) + return wt + + + def abs(self): + t = self._t.abs() + wt = DynamicLazyTensor3(t = t) + return wt + + def __getattr__(self, name): + + if name in ('add', 'sub', 'div', 'mul', 'backward', 'abs'): + def wrapper(*args, **kwargs): + m = getattr(self._t, name) + args = [x._t if isinstance(x, DynamicLazyTensor3) else x for x in args] + # TODO: kwargs + print(len(args)) + t = m(*args, **kwargs) + wt = DynamicLazyTensor3(t = t) + return wt + else: + raise RuntimeError("NYI") + + return wrapper + + +# __torch_dispatch__ +# extend from Tensor +# .size -> MetaTensor in autograd engine (Tensor.size()) +# autograd might need to work differently to work in eager and lazy tensor +# repeat an excercise with var.dim() + + +a = DynamicLazyTensor3((torch.ones(1, 10, device='cpu')).to(device='lazy')) +b = DynamicLazyTensor3((torch.ones(10, 10, device='cpu')).to(device='lazy')) +d = DynamicLazyTensor3((torch.ones(10, 10, device='cpu')).to(device='lazy')) +c = a.abs() +c.backward() +# b = a.size() +# c = a.expand(b) + +print(lazy_tensor_core._LAZYC._get_ltc_tensors_text([c._t])) +#print(lazy_tensor_core._LAZYC._get_ltc_tensors_backend([c._t])) +#ltm.mark_step()