IndexReduce as core function to calculate memory ii

Former-commit-id: fe13fb22c5725ea200589868fce5e7be48ff830e

oneflow/core/auto_placement/df_demo.cpp

@@ -4,8 +4,59 @@ namespace oneflow {
 
 namespace df {
 
-void DifferentialDemo() {
+namespace {
+
+void AutoPlacementMemoryDemo() {
+  Tensor var(Shape({4, 4}), [](size_t index) { return index % 2 ? 0 : 1000; });
+  Tensor row_ones(Shape({var.shape().At(0)}), 1);
+  Tensor col_ones(Shape({var.shape().At(1)}), 1);
+  Tensor epsilon(0.000000001);
+  FOR_RANGE(int, i, 0, 2000) {
+    double lr = 1;
+    if (i < 400) {
+      lr = 0.1;
+    } else if (i < 800) {
+      lr = 0.01;
+    } else if (i < 1200) {
+      lr = 0.001;
+    } else {
+      lr = 0.0001;
+    }
+
+    Tensor x = Add(Square(FixedExpectation(Update(&var, lr), 1)), epsilon);
+    const auto& x_copies = Clone(x, 4);
+    Tensor row = MatrixRowSum(x_copies.at(0));
+    Tensor col = MatrixColSum(x_copies.at(1));
+    Tensor load =
+        ElemWiseMul(x_copies.at(2), TensorProduct(row_ones, Reciprocal(col)));
+    Tensor time = ElemWiseMul(TensorProduct(row, col_ones), load);
+    Tensor ii = Max(time);
+    Backward(Add(ii, AvgAbsDeviation(MatrixColMax(x_copies.at(3)))));
+
+    std::cout << "x: ";
+    for (double i : x.buffer().data()) { std::cout << i << " "; }
+    std::cout << std::endl;
+    std::cout << "row: ";
+    for (double i : row.buffer().data()) { std::cout << i << " "; }
+    std::cout << std::endl;
+    std::cout << "col: ";
+    for (double i : col.buffer().data()) { std::cout << i << " "; }
+    std::cout << std::endl;
+    std::cout << "time: ";
+    for (double i : time.buffer().data()) { std::cout << i << " "; }
+    std::cout << std::endl << std::endl;
+
+    //    Backward(Variance(MatrixColMax(Update(&var, lr))));
+    //    std::cout << "var: ";
+    //    for (double i : var.buffer().data()) { std::cout << i << " "; }
+    //    std::cout << std::endl;
+  }
+}
+
+void AutoPlacementComputationDemo() {
   Tensor var(Shape({4, 4}), [](size_t index) { return index % 2 ? 0 : 1000; });
+  Tensor row_ones(Shape({var.shape().At(0)}), 1);
+  Tensor col_ones(Shape({var.shape().At(1)}), 1);
   Tensor epsilon(0.000000001);
   FOR_RANGE(int, i, 0, 2000) {
     double lr = 1;
@@ -23,8 +74,9 @@ void DifferentialDemo() {
     const auto& x_copies = Clone(x, 4);
     Tensor row = MatrixRowSum(x_copies.at(0));
     Tensor col = MatrixColSum(x_copies.at(1));
-    Tensor table =
-        ElemWiseMul(TensorProduct(row, Reciprocal(col)), x_copies.at(2));
+    Tensor load =
+        ElemWiseMul(x_copies.at(2), TensorProduct(row_ones, Reciprocal(col)));
+    Tensor table = ElemWiseMul(TensorProduct(row, col_ones), load);
     Tensor ii = Max(table);
     Backward(Add(ii, AvgAbsDeviation(MatrixColMax(x_copies.at(3)))));
 
@@ -48,12 +100,18 @@ void DifferentialDemo() {
   }
 }
 
+void DifferentialDemo() {
+  //  AutoPlacementComputationDemo();
+  AutoPlacementMemoryDemo();
+}
+
+}  // namespace
+
 }  // namespace df
 
 }  // namespace oneflow
 
 int main(int argc, char** argv) {
-  using namespace oneflow;
-  df::DifferentialDemo();
+  oneflow::df::DifferentialDemo();
   return 0;
 }

oneflow/core/auto_placement/df_func.cpp

@@ -4,6 +4,26 @@ namespace oneflow {
 
 namespace df {
 
+Tensor IndexReduce(const Tensor& input,
+                   const std::vector<std::vector<int64_t>>& reduce_indexes) {
+  int64_t size = reduce_indexes.size();
+  std::shared_ptr<Buffer> out(new Buffer(Shape({size}), 0));
+  FOR_RANGE(int, i, 0, out->Size()) {
+    for (int64_t index : reduce_indexes.at(i)) {
+      out->At(i) += input.At(index);
+    }
+  }
+  return Tensor(out, [=](const Buffer& out_diff) {
+    Buffer input_diff(input.shape(), 0);
+    FOR_RANGE(int, i, 0, out_diff.Size()) {
+      for (int64_t index : reduce_indexes.at(i)) {
+        input_diff.At(index) += out_diff.At(i);
+      }
+    }
+    input.HandleDiff(input_diff);
+  });
+}
+
 Tensor Update(Tensor* var, double lr) {
   auto buffer = var->mut_buffer_ptr();
   return Tensor(*var, [=](const Buffer& diff) {
@@ -68,14 +88,31 @@ Tensor Tee(const Tensor& input, Tensor* out) {
 Tensor Exp(const Tensor& input) {
   std::shared_ptr<Buffer> out(new Buffer(input.buffer()));
   FOR_RANGE(int, i, 0, out->Size()) {
-    double& x = out->mut_data()->at(i);
-    x = exp(x);
+    double& x = out->At(i);
+    x = std::exp(x);
   }
   return Tensor(out, [=](const Buffer& out_diff) {
     Buffer input_diff(out_diff);
     FOR_RANGE(int, i, 0, input_diff.Size()) {
-      double& diff = input_diff.mut_data()->at(i);
-      diff *= out->data().at(i);
+      double& diff = input_diff.At(i);
+      diff *= out->At(i);
+    }
+    input.HandleDiff(input_diff);
+  });
+}
+
+Tensor Tanh(const Tensor& input) {
+  std::shared_ptr<Buffer> out(new Buffer(input.buffer()));
+  FOR_RANGE(int, i, 0, out->Size()) {
+    double& x = out->At(i);
+    x = std::tanh(x);
+  }
+  return Tensor(out, [=](const Buffer& out_diff) {
+    Buffer input_diff(out_diff);
+    FOR_RANGE(int, i, 0, input_diff.Size()) {
+      double& diff = input_diff.At(i);
+      double o = out->At(i);
+      diff *= 1 - o * o;
     }
     input.HandleDiff(input_diff);
   });

oneflow/core/auto_placement/df_func.h

@@ -7,14 +7,21 @@ namespace oneflow {
 
 namespace df {
 
+Tensor IndexReduce(const Tensor& input,
+                   const std::vector<std::vector<int64_t>>& reduce_indexes);
+
 Tensor Update(Tensor* var, double lr);
 
 std::vector<Tensor> Clone(const Tensor& input, size_t n);
 
 Tensor Minus(const Tensor& input);
+
 Tensor Abs(const Tensor& input);
+
 Tensor Exp(const Tensor& input);
 
+Tensor Tanh(const Tensor& input);
+
 Tensor Tee(const Tensor& input, Tensor* out);
 
 Tensor Add(const Tensor& a, const Tensor& b);