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..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 @@ -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" @@ -32,6 +33,8 @@ #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" +#include "torch/torch.h" namespace torch_lazy_tensors { namespace { @@ -413,6 +416,75 @@ void InitLtcModuleBindings(py::module m) { }; return GetTensorsDump(tensors, coverter); }); + m.def("_dynamic_size", + [](at::Tensor& self) { + LazyTensor self_lazy_tensor = bridge::GetLtcTensor(self); + return bridge::AtenFromLtcTensor( + 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); + 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) { + 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_view", + [](std::vector& self_and_dims) { + 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) { + 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) { + 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) { @@ -491,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.cpp b/lazy_tensor_core/lazy_tensor_core/csrc/tensor.cpp index 7b5ab92116b07..06ad5730218c9 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/torch.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) { + static std::stringstream dummy{}; + + static const auto cur_level = get_current_level(); + if (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,86 @@ 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) -> 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::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 { + 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), + 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/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 e6d7383b704a0..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 @@ -16,6 +16,79 @@ 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(c10::Symbol::prim("_dynamic_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)); +} + +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) {} + +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(), + /*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 a3f6020c81c17..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 @@ -14,6 +14,70 @@ 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); + + 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); + + 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); + + 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 6a93db3c96d99..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 @@ -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 { @@ -306,6 +308,73 @@ class TSNodeLowering : public NodeLowering { ir::NodeCast(node, *ir::ops::ltc_device_data); return {loctx()->GetParameter(device_data_node->data())}; } + + 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_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}); + // 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_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)); + 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 new file mode 100644 index 0000000000000..88f785e92d3f6 --- /dev/null +++ b/lazy_tensor_core/test/dynamic_lazy_tensor2.py @@ -0,0 +1,167 @@ +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 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 + + # 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) + wt = DynamicLazyTensor2(t) + wt._type = Type.DIMENSION + return wt + + + # 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. + wt = DynamicLazyTensor2(t) + if name == 'size': + wt._type = Type.SHAPE + return wt + + return wrapper + + +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] +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([f._t])) +print(lazy_tensor_core._LAZYC._get_ltc_tensors_backend([f._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')) +# 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])) + + +# _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': } + + + +# 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 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()