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..bf3404ec2e04b 100644 --- a/lazy_tensor_core/lazy_tensor_core/csrc/tensor.cpp +++ b/lazy_tensor_core/lazy_tensor_core/csrc/tensor.cpp @@ -9,7 +9,9 @@ #include #include #include +#include #include +#include #include "lazy_tensor_core/csrc/debug_util.h" #include "lazy_tensor_core/csrc/helpers.h" @@ -38,7 +40,47 @@ #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; +} + +int lt_vvveak_symbol() { + return 2; +} + + + namespace torch_lazy_tensors { + +std::vector generate_sizes_node(at::Tensor& t) { + static std::mutex mut; + static std::vector handle_to_ir_value; + std::lock_guard lock(mut); + int64_t index = handle_to_ir_value.size(); + auto lt = bridge::GetLtcTensor(t); + auto node = ir::MakeNode(lt.GetIrValue()); + handle_to_ir_value.push_back(node); + return index; +} + namespace { struct TlsData { @@ -1344,8 +1386,10 @@ std::shared_ptr LazyTensor::ScheduleSyncTensorsGraph( for (size_t i = 0; i < results.size(); ++i) { if (async->tensors_data[i] != nullptr) { + std::cerr << "running assing\n"; async->tensors_data[i]->Assign(*results[i]); } else { + std::cerr << "running move\n"; async->tensors_data[i] = std::move(results[i]); } } @@ -1618,6 +1662,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 +1709,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..b332558d0edf1 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,14 @@ 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 ten_handle = tensor_.CurrentDataHandle(); + std::cerr << "in sizes\n"; + if (ten_handle) { + return ten_handle->shape().dimensions(); + } const_cast(this)->SetupSizeProperties(); return c10::TensorImpl::sizes(); } diff --git a/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/EagerFallback.cpp b/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/EagerFallback.cpp index 6cd89b272aecb..364ca94413064 100644 --- a/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/EagerFallback.cpp +++ b/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/EagerFallback.cpp @@ -4,17 +4,39 @@ #include #include #include +#include "ATen/core/functional.h" +#include "c10/core/DeviceType.h" +#include "lazy_tensor_core/csrc/aten_ltc_bridge.h" +#include "lazy_tensor_core/csrc/tensor.h" namespace torch_lazy_tensors { namespace { std::vector _to_eager(at::TensorList tensors, c10::DeviceType device_type) { + + switch (device_type) { case at::kCPU: { return at::_to_cpu(tensors); } default: { + // TODO we need to collect the right devices + auto lazy_tensors = torch_lazy_tensors::bridge::GetLtcTensors(tensors); + std::vector devices {std::string("cuda")}; + LazyTensor::SyncTensorsGraph(&lazy_tensors, devices, true, false); + if (lazy_tensors.size()) { + auto t = lazy_tensors[0].CurrentTensorData(); + auto dp = lazy_tensors[0].CurrentDataHandle(); + if (dp) { + std::cerr << "shape = " << c10::Join(",", dp->shape().dimensions()) << "dp = " << dp.get() << std::endl; + } + else { + std::cerr << "no tensor dp\n"; + } + + } + std::cerr << "we are in default " << c10::Device(device_type) << "\n"; std::vector eager_tensors; for (const auto& t : tensors) { c10::TensorOptions options = t.options().device(device_type); @@ -83,6 +105,7 @@ void eager_fallback(const c10::OperatorHandle& op, torch::jit::Stack* stack, // Step 1: Convert all non-eager tensor inputs into eager tensors and put them // on the stack at the correct indices. + std::cerr << "going through arguments\n"; for (int64_t idx = 0; idx < arguments.size(); ++idx) { const auto& ivalue = arguments[idx]; if (ivalue.isTensor()) { diff --git a/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/aten_eager_fallback.cpp b/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/aten_eager_fallback.cpp index 8704763e4e2de..0cb9059ff82da 100644 --- a/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/aten_eager_fallback.cpp +++ b/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/aten_eager_fallback.cpp @@ -19,6 +19,11 @@ void ltc_eager_fallback(const c10::OperatorHandle& op, LTC_FN_TRACK(3); const auto name = c10::toString(op.operator_name()); + auto static const printFallback = std::getenv("PRINT_FALLBACK"); + if (printFallback) { + std::cerr << "running fallback for op " << name << std::endl; + } + // Manually applying the LTC_COUNTER macro. // We need to do it ourselves and explicitly keep a mapping of counters // because this boxed fallback kernel is used by multiple operators, 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_computation_client.cpp b/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/ts_computation_client.cpp index bdffd40f7587a..28d7ed6117376 100644 --- a/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/ts_computation_client.cpp +++ b/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/ts_computation_client.cpp @@ -67,8 +67,13 @@ std::vector TSComputationClient::ExecuteComputation( client::ShapeData shape( element_type, std::vector(result_sizes.begin(), result_sizes.end())); - results.push_back( - std::make_shared(result, shape, device)); + + auto tmp = std::make_shared(result, shape, device); + std::cerr << "data ptr = " << tmp.get() << " shape = " << c10::Join(",", shape.dimensions()) << std::endl; + results.push_back(tmp); + + // results.push_back( + // std::make_shared(result, shape, device)); } return results; } diff --git a/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/ts_computation_client.h b/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/ts_computation_client.h index ae126a29dd82a..616a96c020181 100644 --- a/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/ts_computation_client.h +++ b/lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/ts_computation_client.h @@ -21,7 +21,10 @@ class TSComputationClient : public ComputationClient { } void Assign(const Data& data) override { + // TODO: this should really be an assignment constructor data_ = static_cast(data).data_; + shape_ = static_cast(data).shape_; + device_ = static_cast(data).device_; } bool HasValue() const override { return data_.defined(); } 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/lazy_tensors/shape.cc b/lazy_tensor_core/lazy_tensors/shape.cc index f04fd6de56957..a38ceb4c14722 100644 --- a/lazy_tensor_core/lazy_tensors/shape.cc +++ b/lazy_tensor_core/lazy_tensors/shape.cc @@ -2,6 +2,14 @@ namespace lazy_tensors { + bool thread_local shape_check = true; + void set_dynamic_mode_shape_check(bool v) { + shape_check = v; + } + bool dynamic_mode_shape_check() { + return shape_check; + } + void Shape::DeleteDimension(int64 dim_to_delete) { LTC_CHECK(IsArray()); LTC_CHECK_GE(dim_to_delete, 0); diff --git a/lazy_tensor_core/lazy_tensors/shape.h b/lazy_tensor_core/lazy_tensors/shape.h index 0625d24c930e3..de5b1117190e9 100644 --- a/lazy_tensor_core/lazy_tensors/shape.h +++ b/lazy_tensor_core/lazy_tensors/shape.h @@ -15,6 +15,9 @@ namespace lazy_tensors { + void set_dynamic_mode_shape_check(bool v); + bool dynamic_mode_shape_check(); + class Shape { public: Shape() : element_type_(PrimitiveType::INVALID) {} @@ -73,7 +76,7 @@ class Shape { // Methods for accessing the dimensions array. int dimensions_size() const { return dimensions_.size(); } int64 dimensions(int index) const { - if (dynamic_mode_.load()) { + if (dynamic_mode_shape_check() && dynamic_mode_.load()) { throw std::runtime_error("Exact shape not known"); } LTC_CHECK_LT(index, dimensions_.size()); @@ -86,7 +89,7 @@ class Shape { } lazy_tensors::Span dimensions() const { - if (dynamic_mode_.load()) { + if (dynamic_mode_shape_check() && dynamic_mode_.load()) { throw std::runtime_error("Exact shape not known"); } return MakeSpan(dimensions_); diff --git a/lazy_tensor_core/test/cpp/test_aten_ltc_ts_tensor.cpp b/lazy_tensor_core/test/cpp/test_aten_ltc_ts_tensor.cpp index 9e0d5b59ecf7c..695a5fd3a6473 100644 --- a/lazy_tensor_core/test/cpp/test_aten_ltc_ts_tensor.cpp +++ b/lazy_tensor_core/test/cpp/test_aten_ltc_ts_tensor.cpp @@ -121,6 +121,17 @@ TEST_F(AtenLtcTsTensorTest, TestCastInt) { }); } +// TEST_F(AtenLtcTsTensorTest, TestFallbackForNotRegistered) { +// torch::Tensor a = +// torch::rand({2, 2}, torch::TensorOptions(torch::kFloat)) * 100.0; +// const Device* device = GetDefaultDevice(); +// torch::Tensor xla_a = CopyToDevice(a, device); +// auto xla_b = xla_a.relu(); +// auto xla_c = xla_b.hardshrink(); +// AllEqual(xla_c, a.relu().hardshrink()); +// } + + TEST_F(AtenLtcTsTensorTest, TestCastLong) { torch::Tensor a = torch::rand({2, 2}, torch::TensorOptions(torch::kFloat)) * 100.0; 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..00df0298f0215 --- /dev/null +++ b/lazy_tensor_core/test/dynamic_lazy_tensor3.py @@ -0,0 +1,458 @@ +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 + + +class AutogradAbs(torch.autograd.Function): + @staticmethod + def forward(ctx, input): + mask = input.gt(0.0) + #ctx.save_for_backward(mask) + return input.mul(mask) + + @staticmethod + def backward(ctx, grad_output): + (mask,) = ctx.saved_tensors + return grad_output.mul(mask) + +# 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() + t = AutogradAbs.apply(self._t) + 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() + + +# C++ structures for + +""" +enum ValuType { + Tensor, + IntList, + Int +}; + +class Node { + // ... + ValuType type_; +} +""" + + +# Dimensions can come from 3 sources: +# * we wrap python ints that are used as arguments to `__getitem__`, or `sum.axis` +# * __getitem__ on LazyDynamicSize +# * arithmetic ops on `LazyDynamicDimension` e.g. a + b +class LazyDynamicDimension: + + def __init__(self, n): + super().__init__() + self._n = n # IR Value + self._cached = None + + # this makes LazyDynamicDimension a parameter to a computation graph + # Do we need this? + @staticmethod + def fromInt(num): + self._n = lazy_tensor_core._LAZYC._dim_from_int(self, num) + """ + return ir::MakeNode(self_lazy_tensor.GetIrValue(), ValuType::Int); + """ + + def __str__(self): + if not self._cached: + self._cached = lazy_tensor_core._LAZYC._execute_graph(self._n) + + return self._cached + #Option 1 + # materialize the underlying root (_t) + # convert LTC IR to TS + # execute TS IR to get the size + # pros: can be reused by both XLA and Lazy + # cons: somewhat second execution engine + + # Option 2 + # ExecuteComputation's needs to deal with different result types + # We also probably need to squeeze int, IntLists into Data's tensor? + + def __int_(self): + pass + + def add(self, other): + if isinstance(other, int): + # what do we do? + else: + self._n = lazy_tensor_core._LAZYC._add_dim(self, other) + """ + auto dim_self = py::cast<>self.attr("n").cast(); + auto dim_other = py::cast<>self.attr("n").cast(); + // We should be able to reuse the same opcodes + return ir::MakeNode(self_lazy_tensor.GetIrValue(), ValuType::Int); + """ + + def sub(self, other): + # same as add + pass + +# this code below can hopefully be autogenerated +# we do need to parse expressions to figure out +# the tensors that require sizes() +""" +- name: sum(Tensor self, *, ScalarType? dtype=None) -> Tensor + self: grad.expand(self.sizes()) + result: auto_linear + +""" +class AutogradSum(torch.autograd.Function): + @staticmethod + def forward(ctx, self): + ctx.save_for_backward(self) + return self.sum() + + @staticmethod + def backward(ctx, grad_output): + (self,) = ctx.saved_tensors() + return grad_output.expand(self.sizes()) + +class AutogradExpand(torch.autograd.Function): + @staticmethod + def forward(ctx, self, *sizes): + ctx.save_for_backward(self) + dyn_size = LazyDynamicSize.fromInts(sizes) + return lazy_tensor_core._LAZYC._dynamic_expand(self, sizes) + """ + LazyTensor self_lazy_tensor = bridge::GetLtcTensor(self); + auto sizes = py::cast<>self.attr("n").cast(); + // We should be able to reuse the same opcodes + return ir::MakeNode(self_lazy_tensor.GetIrValue(), sizes); + """ + + @staticmethod + def backward(ctx, grad_output): + (self,) = ctx.saved_tensors() + return grad_output.sum_to(grad_output, self.sizes()) + +class LazyDynamicSize: + def __init__(self, n): + super().__init__() + self._n = n + + + def add(self, other): + pass + + @staticmethod + def fromTensor(t): + # this creates an IR Node of type IntList + """ + LazyTensor self_lazy_tensor = bridge::GetLtcTensor(self); + return ir::MakeNode(self_lazy_tensor.GetIrValue(), ValuType::IntList); + """ + n = lazy_tensor_core._LAZYC._dynamic_size(self._t) + return LazyDynamicSize(n) + + @staticmethod + def fromInts(*args): + args = [a if isinstance(a, LazyDynamicDimension) else LazyDynamicDimension.fromInt(a) for a in args] + n = lazy_tensor_core._LAZYC._dynamic_size_from_ints(args) + return LazyDynamicSize(n) + + def __getitem__(self, index): + w_index = index + if isinstance(index, int): + w_index = LazyDynamicDimension.fromInt(index) + n = lazy_tensor_core._LAZYC._get_item(self, index) + return LazyDynamicDimension(self, n) + + + def __setitem__(self, index): + # similar to __getitem__ ? + + +class DynamicLazyTensor4(torch.Tensor): + """ + Non-compound ops + """ + + # Category 1: non-compound ops that don't take size or dim and don't use sizes() in backward + # e.g. add, relu, abs, mul, div, etc + # No work necessary as long we can rely on `torch.autograd.register_py_tensor_class_for_device("lazy", DynamicLazyTensor4)` + + # Category 2: non-compound ops that return sizes() or dim can be implemented with thin wrappers around IR + # We would need to handwrite these ops + def sizes(): + return LazyDynamicSize.fromTensor(self) + + def size(index): + pass + + # Category 3: non-compound ops that use sizes in backward + # We would codegen both forward and backward + + # this only implements the `sum` overload, + # we will need to implement `sum.dim_IntList` + # overload here as well + def sum(self): + return AutogradSum.apply(*args, **kwargs) + + + # Category 4: non-compound or compound ops that take sizes or individual dimensions (e.g. expand) (CODEGENABLE) + def expand(self, *sizes): + return AutogradExpand.apply(*args, sizes) + + + """ + Compound Ops + """ + + + # Category 5: compound ops that use sizes in forward and don't have backward formulas in derivatives + # We would need to handwrite + def dropout(...) + + + # Category 6: compound ops that use sizes in forward and have backward formulas in derivatives + # We would handwrite forward and codegen the backward + def addmm(self, mat1, mat2, beta, alpha): + pass + + +# register the python class +torch.autograd.register_py_tensor_class_for_device("lazy", DynamicLazyTensor4) + + + + + + + +""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""" +""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""" + +# C++ structures for + +""" +enum ValueType { + Tensor, + IntList, + Int +}; + +class Node { + // ... + ValuType type_; +} +""" + +class Type(Enum): + TENSOR = 1 + SHAPE = 2 + DIMENSION = 3 + +# we will use DynamicLazyTensor4 to back up DynamicSize and DynamicDimension +class DynamicLazyTensor4(torch.Tensor, *args, **kwargs): + super().__init__(*args, **kwargs) + + # Option 1: store type explicitly in Python e.g. + + self._type = Type.TENSOR + + # Option 2: get type from the underlying Tensor if ValueType is typed + def get_value_type(self): + """ + LazyTensor self_lazy_tensor = bridge::GetLtcTensor(self); + return self_lazy_tensor.GetIrValue().getNode().getValueType(); + return ir::MakeNode(self_lazy_tensor.GetIrValue()); + """ + + + """ + Non-compound ops + """ + + # Category 0a: all ops that can take either TENSOR, SHAPE, DIMENSION need to be either overriden if we store types explicitly, + # so we can type them in Python as TENSOR, SHAPE or DIMENSION. We can avoid this with `get_value_type` + def add(self, other): + if self._type == Type.TENSOR or other._type == Type.TENSOR: + """ + 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(), ValueType::Tensor))); + """ + t = lazy_tensor_core._LAZYC._dynamic_add(self, other) + t._type = Type.TENSOR + else self._type == Type.SHAPE or other._type == Type.SHAPE + # no other case + assert(self._type == Type.SHAPE and other._type == Type.SHAPE) + # // ValueType::IntList + t = lazy_tensor_core._LAZYC._dynamic_add(self, other) + t._type = Type.SHAPE + else: + pass + # DIMENSION case : ValueType::Int + + + # Category 0b: all ops that can only take TENSORs need to throw on SHAPE and DIMENSION + # this could be quite a few ops + def relu(self): + if self._type != Type.TENSOR: + raise RuntimeError() + else: + super().relu() + + def __getitem__(self, index): + if self._type != Type.SHAPE + raise RuntimeError() + else: + t = lazy_tensor_core._LAZYC._get_item(self, index) + t._type = Type.DIMENSION + def __int__(self): + pass + + def __str__(self): + pass + + # Category 1: non-compound ops that don't take size or dim and don't use sizes() in backward + # e.g. add, relu, abs, mul, div, etc + # See caveats in Category 0 + + # Category 2: non-compound ops that return sizes() or dim can be implemented with thin wrappers around IR + # We would need to handwrite these ops + def sizes(): + # this creates an IR Node of type IntList + """ + LazyTensor self_lazy_tensor = bridge::GetLtcTensor(self); + auto size = ir::MakeNode(self_lazy_tensor.GetIrValue(), ValuType::IntList); + return bridge::AtenFromLtcTensor( + self_lazy_tensor.CreateFrom(size)); + """ + t = lazy_tensor_core._LAZYC._dynamic_size(self._t) + t._type = Type.SHAPE + return t + + def size(index): + pass + + # Category 3: non-compound ops that use sizes in backward + # We would codegen both forward and backward + + # this only implements the `sum` overload, + # we will need to implement `sum.dim_IntList` + # overload here as well + def sum(self): + return AutogradSum.apply(*args, **kwargs) + + + # Category 4: non-compound or compound ops that take sizes or individual dimensions (e.g. expand) (CODEGENABLE) + def expand(self, *sizes): + return AutogradExpand.apply(*args, sizes) + + + """ + Compound Ops + """ + + + # Category 5: compound ops that use sizes in forward and don't have backward formulas in derivatives + # We would need to handwrite + def dropout(...) + + + # Category 6: compound ops that use sizes in forward and have backward formulas in derivatives + # We would handwrite forward and codegen the backward + def addmm(self, mat1, mat2, beta, alpha): + pass + + +# register the python class +torch.autograd.register_py_tensor_class_for_device("lazy", DynamicLazyTensor4) + + + + + + + + + +# copy the registration code for device \ No newline at end of file diff --git a/lazy_tensor_core/test/dynamic_lazy_tensor4.py b/lazy_tensor_core/test/dynamic_lazy_tensor4.py new file mode 100644 index 0000000000000..ff16d4c61564b --- /dev/null +++ b/lazy_tensor_core/test/dynamic_lazy_tensor4.py @@ -0,0 +1,100 @@ +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 + + +class AutogradAbs(torch.autograd.Function): + @staticmethod + def forward(ctx, input): + mask = input.gt(0.0) + #ctx.save_for_backward(mask) + return input.mul(mask) + + @staticmethod + def backward(ctx, grad_output): + (mask,) = ctx.saved_tensors + return grad_output.mul(mask) + +# 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() + t = AutogradAbs.apply(self._t) + 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() diff --git a/torch/_C/_autograd.pyi b/torch/_C/_autograd.pyi index 07e55f1e7aa47..4fe6e59923244 100644 --- a/torch/_C/_autograd.pyi +++ b/torch/_C/_autograd.pyi @@ -93,3 +93,4 @@ def _reset_default_hooks() -> None: ... def _enable_profiler_legacy(config: ProfilerConfig) -> None: ... def _disable_profiler_legacy() -> List[List[ProfilerEvent]]: ... +def _register_py_class_for_device(device: str, cls: Any) -> None: ... diff --git a/torch/autograd/__init__.py b/torch/autograd/__init__.py index ece09fb9adfd6..b801f31271d04 100644 --- a/torch/autograd/__init__.py +++ b/torch/autograd/__init__.py @@ -271,3 +271,8 @@ def variable(*args, **kwargs): _prepare_profiler, _enable_profiler, _disable_profiler) from . import profiler + +def register_py_tensor_class_for_device(device, cls): + if not isinstance(cls, type): + raise RuntimeError("cls isn't a typeinfo object") + torch._C._autograd._register_py_class_for_device(device, cls) diff --git a/torch/autograd/variable.py b/torch/autograd/variable.py index b307a9ac4b6c7..02305bd5f5888 100644 --- a/torch/autograd/variable.py +++ b/torch/autograd/variable.py @@ -6,11 +6,9 @@ class VariableMeta(type): def __instancecheck__(cls, other): return isinstance(other, torch.Tensor) - # mypy doesn't understand torch._six.with_metaclass class Variable(with_metaclass(VariableMeta, torch._C._LegacyVariableBase)): # type: ignore[misc] pass - from torch._C import _ImperativeEngine as ImperativeEngine Variable._execution_engine = ImperativeEngine() diff --git a/torch/csrc/autograd/engine.cpp b/torch/csrc/autograd/engine.cpp index 252a74b4c07c7..d3a3ccb82a014 100644 --- a/torch/csrc/autograd/engine.cpp +++ b/torch/csrc/autograd/engine.cpp @@ -37,7 +37,14 @@ #include #include #include +#include "c10/core/DeviceType.h" #include +#include +#include + +namespace torch_lazy_tensors { +extern __attribute__((weak)) at::Tensor generate_size_check_for(std::vector sz, at::Tensor& grad); +} namespace torch { namespace autograd { @@ -668,17 +675,34 @@ void validate_outputs( // AT_ERROR(format_error(ss.str())); continue; } - if (!grad.sizes().equals(metadata.shape())) { - if (!at::is_expandable_to(metadata.shape(), grad.sizes())) { - std::stringstream ss; - ss << "invalid gradient at index " << i << " - got "; - ss << grad.sizes() << " but expected shape compatible with "; - ss << metadata.shape(); - AT_ERROR(format_error(ss.str())); + + // TODO: Consider if metadata can keep a pointer to a tensor + // TODO: we could consider adding an op `is_expandable_or_throw` but + // we still need to check the device to treat metadata.shape() differently for metatensors + if (metadata.device().type() == at::kLazy) { + // TODO: we can remove weak symbol linking when we move these helpers to master + if (torch_lazy_tensors::generate_size_check_for == nullptr) { + std::cerr << "torch_lazy_tensors::generate_size_check_for wasn't defined!\n"; + TORCH_CHECK(false, "torch_lazy_tensors::generate_size_check_for"); + } + std::vector isz (metadata.shape().begin(), metadata.shape().end()); + grad = torch_lazy_tensors::generate_size_check_for(isz, grad); + } else { + if (!grad.sizes().equals(metadata.shape())) { + if (!at::is_expandable_to(metadata.shape(), grad.sizes())) { + std::stringstream ss; + ss << "invalid gradient at index " << i << " - got "; + ss << grad.sizes() << " but expected shape compatible with "; + ss << metadata.shape(); + AT_ERROR(format_error(ss.str())); + } + grad = at::sum_to(std::move(grad), metadata.shape()); } - grad = at::sum_to(std::move(grad), metadata.shape()); } + + + bool input_is_complex = isComplexType(c10::typeMetaToScalarType(metadata.options().dtype())); bool grad_is_complex = isComplexType(grad.scalar_type()); diff --git a/torch/csrc/autograd/function.cpp b/torch/csrc/autograd/function.cpp index 83509c9dae106..4aa4be6c60dbc 100644 --- a/torch/csrc/autograd/function.cpp +++ b/torch/csrc/autograd/function.cpp @@ -13,6 +13,14 @@ #include #include #include +#include "c10/core/DeviceType.h" + +extern __attribute__((weak)) int lt_vvveak_symbol(); + +namespace torch_lazy_tensors { +extern __attribute__((weak)) std::vector generate_sizes_node(at::Tensor& t); +} + namespace torch { namespace autograd { @@ -39,6 +47,32 @@ auto Node::name() const -> std::string { return c10::demangle(typeid(*this).name()); } +uint32_t Node::add_input_metadata(const at::Tensor& t) noexcept { + // NOLINTNEXTLINE(cppcoreguidelines-init-variables) + uint32_t input_nr = input_metadata_.size(); + if (t.device().type() == c10::kLazy) { + + if (lt_vvveak_symbol) { + std::cerr << "lt_weak_symbol_used" << lt_vvveak_symbol() << std::endl; + } else { + std::cerr << "no lt_weak_symbol\n"; + } + + auto nt = const_cast(t); + if (torch_lazy_tensors::generate_sizes_node) { + auto sz = torch_lazy_tensors::generate_sizes_node(nt); + InputMetadata im(t.options(), sz, t.device()); + input_metadata_.push_back(std::move(im)); + } else { + std::cerr << "torch_lazy_tensors::generate_sizes_node wasn't defined!\n"; + TORCH_CHECK(false, "torch_lazy_tensors::generate_sizes_node"); + } + } else { + input_metadata_.emplace_back(t); + } + return input_nr; +} + AnomalyMetadata* Node::metadata() noexcept { if (!anomaly_metadata_) { anomaly_metadata_ = Engine::get_default_engine().make_anomaly_metadata(); diff --git a/torch/csrc/autograd/init.cpp b/torch/csrc/autograd/init.cpp index 0778cbc6012a7..d2b2b4c084368 100644 --- a/torch/csrc/autograd/init.cpp +++ b/torch/csrc/autograd/init.cpp @@ -14,6 +14,7 @@ #include #include #include +#include #include #include @@ -274,6 +275,11 @@ PyObject* THPAutograd_initExtension(PyObject* _unused, PyObject *unused) { py::class_(_C_m, "_InferenceMode") .def(py::init()); + _C_m.def("_register_py_class_for_device", [](const std::string& device, py::object python_type_class) { + auto cls = python_type_class.ptr(); + registerPythonTensorClass(device, cls); + }); + py::class_(_C_m, "_DisableTorchDispatch") .def(py::init<>()); diff --git a/torch/csrc/autograd/python_variable.cpp b/torch/csrc/autograd/python_variable.cpp index f0cd80d0d2883..c41ecfb76bde3 100644 --- a/torch/csrc/autograd/python_variable.cpp +++ b/torch/csrc/autograd/python_variable.cpp @@ -27,9 +27,11 @@ #include #include #include +#include #include #include + #include #include @@ -44,6 +46,8 @@ #include #include + + using namespace at; using namespace torch; using namespace torch::autograd; @@ -152,6 +156,23 @@ static PyObject* THPVariable_NewWithVar( } return obj; } +// NOLINTNEXTLINE +static PyObject* device_to_py_class_ [static_cast(c10::DeviceType::COMPILE_TIME_MAX_DEVICE_TYPES)]; + +void registerPythonTensorClass(const std::string& device, PyObject* python_tensor_class) { + c10::Device dev(device); + + TORCH_CHECK(dev.type() == kXLA, "Only the python class for XLA can be overriden"); + if (device_to_py_class_[static_cast(dev.type())] != nullptr) { + TORCH_WARN("Overriding a previously registered python class for ", dev.str()); + } + + device_to_py_class_[static_cast(dev.type())] = python_tensor_class; +} + +static PyObject* getPythonTensorClass(c10::Device d) { + return device_to_py_class_[static_cast(d.type())]; +} // TODO: Make this take Variable by const reference PyObject * THPVariable_Wrap(Variable var) @@ -198,6 +219,17 @@ PyObject * THPVariable_Wrap(Variable var) status = c10::impl::PyInterpreterStatus::MAYBE_UNINITIALIZED; } } + + if (C10_LIKELY(var.device().type() != c10::kXLA)) { + return THPVariable_NewWithVar( + (PyTypeObject*)THPVariableClass, std::move(var), status); + } + + if (auto clazz = getPythonTensorClass(var.device())) { + return THPVariable_NewWithVar( + (PyTypeObject*)clazz, std::move(var), status); + } + return THPVariable_NewWithVar( (PyTypeObject*)THPVariableClass, std::move(var), status); } diff --git a/torch/csrc/autograd/python_variable.h b/torch/csrc/autograd/python_variable.h index faf882d7edc2c..4db6c910d29cc 100644 --- a/torch/csrc/autograd/python_variable.h +++ b/torch/csrc/autograd/python_variable.h @@ -18,6 +18,8 @@ struct THPVariable { PyObject* backward_hooks = nullptr; }; +TORCH_API void registerPythonTensorClass(const std::string& device, PyObject* python_tensor_class); + THP_API PyObject *THPVariableClass; THP_API PyObject *ParameterClass; diff --git a/torch/csrc/jit/runtime/profiling_graph_executor_impl.cpp b/torch/csrc/jit/runtime/profiling_graph_executor_impl.cpp index 747fb943c210a..495ddc7dbf686 100644 --- a/torch/csrc/jit/runtime/profiling_graph_executor_impl.cpp +++ b/torch/csrc/jit/runtime/profiling_graph_executor_impl.cpp @@ -56,6 +56,8 @@ C10_DEFINE_int64( kDefaultBailoutDepth, "Number of re-specializations"); +__attribute__((weak)) int lt_vvveak_symbol(); + namespace torch { namespace jit { @@ -695,6 +697,13 @@ const ExecutionPlan& ProfilingGraphExecutorImpl::getPlanFor( size_t remaining_bailout_depth) { std::lock_guard lock(compile_mutex); + std::cerr << "running profiling executor!\n"; + if (lt_vvveak_symbol) { + std::cerr << "lt_weak_symbol_used" << lt_vvveak_symbol() << std::endl; + } else { + std::cerr << "no lt_weak_symbol\n"; + } + // IMPORTANT: This is a hot path of calling a torchscript function. Try not to // add any code above this. if (optimized_plan_) {