Backend Cleanroom port (#13)

Implements CheckpointTensor backend Hook up the overload with a passing test. (#15) * save * save * review * review * add back files Adding some overload (#16) * save * save * review * remove outdated file * save * can compile again * save * save * up up array * finsih the overloads * replace vec[n] with vec.at(n) make resnet pass again, adding more logs for logging simulator (#17) * commit * add more overloads * fix log * save Logging everything for resnet (#18) * save * save * save * save * use release instead of free Fix logger (special handling for small constants) (#19) * save * fix comment TreeLSTM overloads (#21) Overloads for mean and mean.dim (needed for densenet) Add overloads for U-Net Add ability to annotate log and make new log file (#25) Restore overloads needed for LSTM and GRU Implemented unrolled_gan (#23) * save * save * save Overload bitwise and for ParityACT Add various overloads for ACT model Log aliases of operator outputs Overloads and changes needed for transformer Overloads for topk functions Add overloads for deepspeech More overloads for ACT DTR implementation first pass (#37) save before change save before change add aliaspool design, and add weak/strong pointer discussion add more code rebase add allocator hook save metadata to prepare for eviction save refactor: move log to a seprate file add file raii save save comment save save save save find error, bisecting code save save address review comment address comment address comment fix segfault save save pin save save save save save refactor - remove stitch save save [ impl ] overload diag & mv refactor - remove stitch save save restore equivalentclassnode save save save 50% resnet here we go 50% resnet here we go save save save save save save More opts save Try removing pop-ing from loop save save save save mike try this mike try this save undo system optimization for now save save save fix compile error save save save save fix compile save fix save save [ restore ] overload for NLL loss2D Add an off switch for sampling, fix tense on 'shrinked' Add switch for ignoring small tensors
uwsampl · Apr 4, 2023 · a21ed14 · a21ed14
1 parent 8f2f1a0
commit a21ed14
Show file tree

Hide file tree

Showing 15 changed files with 4,557 additions and 12 deletions.
diff --git a/aten/src/ATen/CheckpointTensorImpl.cpp b/aten/src/ATen/CheckpointTensorImpl.cpp
diff --git a/aten/src/ATen/CheckpointTensorImpl.h b/aten/src/ATen/CheckpointTensorImpl.h
diff --git a/aten/src/ATen/Logger.h b/aten/src/ATen/Logger.h
@@ -0,0 +1,197 @@
+#pragma once
+
+#include <string>
+#include <iostream>
+#include <../../../third_party/json/single_include/nlohmann/json.hpp>
+
+namespace at {
+
+struct DTRLogger {
+  std::string time_prefix;
+  std::ofstream out;
+  static std::string get_time_prefix() {
+    std::time_t t = std::time(nullptr);
+    std::tm* tm = std::localtime(&t);
+    return
+      std::to_string(1900+tm->tm_year) + "-" +
+      std::to_string(1+tm->tm_mon) + "-" +
+      std::to_string(tm->tm_mday) + "-" +
+      std::to_string(tm->tm_hour) + "-" +
+      std::to_string(tm->tm_min) + "-" +
+      std::to_string(tm->tm_sec);
+  }
+  std::string get_filename(const std::string& name) {
+    return time_prefix + "-" + name + ".log";
+  }
+  DTRLogger() : time_prefix(get_time_prefix()), out(get_filename("default")) { }
+  void log(const std::string& str) {
+    out << str << std::endl;
+  }
+  static DTRLogger& logger() {
+    static DTRLogger ret;
+    return ret;
+  }
+
+};
+
+using json = nlohmann::json;
+constexpr bool log_json = true;
+const std::string INSTRUCTION = "INSTRUCTION";
+const std::string ANNOTATION = "ANNOTATION";
+const std::string RELEASE = "RELEASE";
+const std::string PIN = "PIN";
+const std::string TIME = "TIME";
+const std::string ARGS = "ARGS";
+const std::string MEMORY = "MEMORY";
+const std::string ALIAS = "ALIAS";
+const std::string NAME = "NAME";
+const std::string CONSTANT = "CONSTANT";
+
+void DTRLogConstant(const std::string& name) {
+  if (log_json) {
+    json j;
+    j[INSTRUCTION] = CONSTANT;
+    j[NAME] = name;
+    DTRLogger::logger().log(j.dump());
+  } else {
+    DTRLogger::logger().log(CONSTANT + " " + name);
+  }
+}
+
+void DTRLogMemory(const std::string& name, size_t memory) {
+  if (log_json) {
+    json j;
+    j[INSTRUCTION] = MEMORY;
+    j[NAME] = name;
+    j[MEMORY] = std::to_string(memory);
+    DTRLogger::logger().log(j.dump());
+  } else {
+    DTRLogger::logger().log(name + " " + MEMORY + ": " + std::to_string(memory));
+  }
+}
+
+void DTRLogAlias(const std::string& name, int index) {
+  if (log_json) {
+    json j;
+    j[INSTRUCTION] = ALIAS;
+    j[NAME] = name;
+    j[ALIAS] = std::to_string(index);
+    DTRLogger::logger().log(j.dump());
+  } else {
+    DTRLogger::logger().log(name + " " + ALIAS + ": " + std::to_string(index));
+  }
+}
+
+void DTRLogCopyFrom(const std::string& to, const std::string& from) {
+  if (log_json) {
+    json j;
+    j[INSTRUCTION] = "COPY_FROM";
+    j["DST"] = to;
+    j["SRC"] = from;
+    DTRLogger::logger().log(j.dump());
+  } else {
+    DTRLogger::logger().log(to + " <- " + from);
+  }
+}
+
+void DTRLogCopy(const std::string& new_name, const std::string& old_name) {
+  if (log_json) {
+    json j;
+    j[INSTRUCTION] = "COPY";
+    j["DST"] = new_name;
+    j["SRC"] = old_name;
+    DTRLogger::logger().log(j.dump());
+  } else {
+    DTRLogger::logger().log(new_name + " = " + old_name);
+  }
+}
+
+void DTRLogMutate(const std::string& name,
+                  const std::vector<std::string>& args,
+                  const std::vector<size_t>& mutate,
+                  const std::string& time) {
+  if (log_json) {
+    json j;
+    j[INSTRUCTION] = "MUTATE";
+    j[NAME] = name;
+    j[ARGS] = args;
+    j["MUTATE"] = mutate;
+    j[TIME] = time;
+    DTRLogger::logger().log(j.dump());
+  } else {
+    std::string log = name;
+    log += "(";
+    for (const auto& s : args) {
+      log += s;
+      log += ", ";
+    }
+    log += ") ";
+    log += " MUTATING: ";
+    log += "(";
+    for (const size_t i : mutate) {
+      log += std::to_string(i);
+      log += ", ";
+    }
+    log += ") ";
+    log += TIME;
+    log += ": ";
+    log += time;
+    DTRLogger::logger().log(log);
+  }
+}
+
+void DTRLogRelease(const std::string& name) {
+  if (log_json) {
+    json j;
+    j[INSTRUCTION] = RELEASE;
+    j[NAME] = name;
+    DTRLogger::logger().log(j.dump());
+  } else {
+    DTRLogger::logger().log(RELEASE + ": " + name);
+  }
+}
+
+void DTRLogPin(const std::string& name) {
+  if (log_json) {
+    json j;
+    j[INSTRUCTION] = PIN;
+    j[NAME] = name;
+    DTRLogger::logger().log(j.dump());
+  } else {
+    DTRLogger::logger().log(RELEASE + ": " + name);
+  }
+}
+
+void DTRLogCall(const std::vector<std::string>& res,
+                const std::string& name,
+                const std::vector<std::string>& args,
+                const std::string& time) {
+  if (log_json) {
+    json j;
+    j[INSTRUCTION] = "CALL";
+    j[NAME] = name;
+    j["RESULT"] = res;
+    j[ARGS] = args;
+    j[TIME] = time;
+    DTRLogger::logger().log(j.dump());
+  } else {
+    std::string arg = name + "(";
+    for (const auto& s : args) {
+      arg += s;
+      arg += ", ";
+    }
+    arg += ")";
+    std::string log = "(";
+    for (const auto& s: res) {
+      log += s;
+      log += ", ";
+    }
+    log += ") = ";
+    log += arg;
+    log += " TIME: ";
+    log += time;
+    DTRLogger::logger().log(log);
+  }
+}
+
+}
diff --git a/aten/src/ATen/native/Activation.cpp b/aten/src/ATen/native/Activation.cpp
@@ -829,4 +829,94 @@ Tensor& log_sigmoid_backward_cpu_out(const Tensor& grad_output,
 DEFINE_DISPATCH(GeluKernel);
 DEFINE_DISPATCH(GeluBackwardKernel);
 
+Tensor slice_backward(const Tensor& grad, IntArrayRef input_sizes, int64_t dim, int64_t start, int64_t end, int64_t step) {
+  auto grad_input = at::zeros(input_sizes, grad.options());
+  grad_input.slice(dim, start, end, step).copy_(grad);
+  return grad_input;
+}
+
+Tensor select_backward(const Tensor& grad, IntArrayRef input_sizes, int64_t dim, int64_t index) {
+  auto grad_input = at::zeros(input_sizes, grad.options());
+  grad_input.select(dim, index).copy_(grad);
+  return grad_input;
+}
+
+std::vector<Tensor> split(const Tensor& self, int64_t split_size, int64_t dim) {
+  TORCH_CHECK(self.dim() != 0, "split expects at least a 1-dimensional tensor");
+  TORCH_CHECK(split_size >= 0,  "split expects split_size be non-negative, but got split_size=", split_size);
+  int64_t dim_size = self.size(dim);
+  TORCH_CHECK(split_size > 0 || self.size(dim) == 0,
+           "split_size can only be 0 if dimension size is 0, "
+           "but got dimension size of ", dim_size);
+  // if split_size is 0 and dimension size is 0, there is 1 split.
+  int64_t num_splits = 1;
+  if (split_size != 0) {
+    // ensuring num_splits is at least 1 makes consistent the case where split_size > dim_size
+    // (returns a single split).  We might want to error here, but keep it for BC.
+    num_splits = std::max<int64_t>((dim_size + split_size - 1) / split_size, 1);
+  }
+  std::vector<Tensor> splits(num_splits);
+  int64_t last_split_size = split_size - (split_size * num_splits - dim_size);
+
+  for (int64_t i = 0; i < num_splits; ++i) {
+    auto length = i < num_splits - 1 ? split_size : last_split_size;
+    splits[i] = self.narrow(dim, i * split_size, length);
+  }
+  return splits;
+}
+
+std::vector<Tensor> split_with_sizes(const Tensor& self, IntArrayRef split_sizes, int64_t dim) {
+  TORCH_CHECK(self.dim() != 0, "split expects at least a 1-dimensional tensor");
+  int64_t dim_size = self.size(dim);
+  int64_t num_splits = split_sizes.size();
+  std::vector<Tensor> splits(num_splits);
+  int64_t start_idx = 0;
+  int64_t i;
+
+  for (i = 0; i < num_splits; ++i) {
+    auto length = split_sizes[i];
+    TORCH_CHECK(length >= 0,
+             "split_with_sizes expects split_sizes have only non-negative ",
+             "entries, but got split_sizes=", split_sizes);
+    splits[i] = self.narrow(dim, start_idx, length);
+    start_idx += length;
+  }
+  TORCH_CHECK(start_idx == dim_size,
+           "split_with_sizes expects split_sizes to sum exactly to ", dim_size,
+           " (input tensor's size at dimension ", dim, "), ", "but got split_sizes=", split_sizes);
+  return splits;
+}
+
+Tensor split_backward(c10::ArrayRef<Tensor> grads,
+                      int64_t split_size, int64_t dim, IntArrayRef sizes, const at::TensorOptions &options) {
+  dim = at::maybe_wrap_dim(dim, sizes.size());
+  int64_t dim_size = sizes[dim];
+  int64_t num_splits = grads.size();
+  std::vector<int64_t> split_sizes(num_splits, split_size);
+  split_sizes[num_splits - 1] = split_size - (split_size * num_splits - dim_size);
+  return at::native::split_with_sizes_backward(grads, split_sizes, dim, sizes, options);
+}
+
+Tensor split_with_sizes_backward(c10::ArrayRef<Tensor> grads,
+                                 IntArrayRef split_sizes, int64_t dim, IntArrayRef sizes, const at::TensorOptions &options) {
+  dim = at::maybe_wrap_dim(dim, sizes.size());
+
+  // it's possible some of the grads are not defined (represents tensors of all 0s).
+  // Since at::cat can't handle those, let's define them
+  std::vector<Tensor> grads_all_defined(grads.size());
+  for (size_t j = 0; j < grads.size(); ++j) {
+    if (grads[j].defined()) {
+      grads_all_defined[j] = grads[j];
+    } else {
+      auto length = split_sizes[j];
+      auto grad_size = sizes.vec();
+      grad_size[dim] = length;
+      grads_all_defined[j] = at::zeros(grad_size, options);
+    }
+  }
+
+  auto ret =  at::cat(grads_all_defined, dim);
+  return ret;
+}
+
 }}  // namespace at::native