Asuhan/dynamic shapes experiment

lazy_tensor_core/lazy_tensor_core/csrc/init_python_bindings.cpp

@@ -7,6 +7,7 @@
 #include <thread>
 #include <vector>
 
+#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<ir::ops::DynamicSize>(
+                  self_lazy_tensor.GetIrValue())));
+        });
+  m.def("_dynamic_size2",
+        [](at::Tensor& self) {
+          LazyTensor self_lazy_tensor = bridge::GetLtcTensor(self);
+          return ir::MakeNode<ir::ops::DynamicSize>(self_lazy_tensor.GetIrValue());
+        });
+  m.def("_dynamic_expand2",
+        [](at::Tensor& self, std::shared_ptr<ir::Node> val) {
+          LazyTensor self_lazy_tensor = bridge::GetLtcTensor(self);
+          return bridge::AtenFromLtcTensor(
+              self_lazy_tensor.CreateFrom(ir::MakeNode<ir::ops::DynamicExpand>(
+                  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<ir::ops::AddDim>(
+                  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<ir::ops::DynamicExpand>(
+                  self_lazy_tensor.GetIrValue(), other_lazy_tensor.GetIrValue())));
+        });
+  m.def("_dynamic_view",
+        [](std::vector<at::Tensor>& 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::ops::DynamicView>(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<ir::ops::DynamicLinear>(
+                  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<ir::ops::DynamicGetItem>(
+                  self_lazy_tensor.GetIrValue(), other_lazy_tensor.GetIrValue())));
+        });
+        // IrValueFromScalar
   m.def("_get_ltc_tensors_text",
         [](const std::vector<at::Tensor>& tensors) -> std::string {
           auto coverter = [](lazy_tensors::Span<const ir::Node* const> nodes) {
@@ -491,6 +563,7 @@ void InitLtcModuleBindings(py::module m) {
   });
 
   py::class_<ir::Value, std::shared_ptr<ir::Value>>(m, "IrValue");
+  py::class_<ir::Node, std::shared_ptr<ir::Node>>(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,

lazy_tensor_core/lazy_tensor_core/csrc/tensor.cpp

@@ -9,7 +9,9 @@
 #include <mutex>
 #include <set>
 #include <stdexcept>
+#include <unordered_map>
 #include <unordered_set>
+#include <vector>
 
 #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<int64_t> generate_sizes_node(at::Tensor& t) {
+  static std::mutex mut;
+  static std::vector<ir::Value> handle_to_ir_value;
+  std::lock_guard<std::mutex> lock(mut);
+  int64_t index = handle_to_ir_value.size();
+  auto lt = bridge::GetLtcTensor(t);
+  auto node = ir::MakeNode<ir::ops::DynamicSize>(lt.GetIrValue());
+  handle_to_ir_value.push_back(node);
+  return index;
+}
+
 namespace {
 
 struct TlsData {
@@ -1344,8 +1386,10 @@ std::shared_ptr<LazyTensor::Async> 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::Async> 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<int64_t> r;
+            auto t_data = t.data<int64_t>();
+            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<int64_t> dims;
+            auto ivals = torch::jit::last(stack, num_inputs);
+            for (auto iv : ivals) {
+              dims.push_back(iv.toTensor().item<int64_t>());
+            }
+            torch::jit::drop(stack, num_inputs);
+            torch::jit::push(stack, dims);
+          };
+        },
+        c10::AliasAnalysisKind::FROM_SCHEMA),
+});

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<LTCTensorImpl*>(this)->SetupSizeProperties();
   return c10::TensorImpl::sizes();
 }

lazy_tensor_core/lazy_tensor_core/csrc/ts_backend/EagerFallback.cpp

@@ -4,17 +4,39 @@
 #include <ATen/core/boxing/KernelFunction.h>
 #include <ATen/native/CPUFallback.h>
 #include <torch/library.h>
+#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<at::Tensor> _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<std::string> 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<at::Tensor> 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()) {

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,