text: /*
  Copyright 2020 The OneFlow Authors. All rights reserved.
  
  Licensed under the Apache License, Version 2.0 (the "License");
  you may not use this file except in compliance with the License.
  You may obtain a copy of the License at
  
      http://www.apache.org/licenses/LICENSE-2.0
  
  Unless required by applicable law or agreed to in writing, software
  distributed under the License is distributed on an "AS IS" BASIS,
  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  See the License for the specific language governing permissions and
  limitations under the License.
  */
  
  #include "oneflow/core/framework/op_expr_grad_function.h"
  
  #include "oneflow/core/framework/op_builder.h"
  #include "oneflow/core/framework/op_expr.h"
  #include "oneflow/core/framework/op_interpreter/op_interpreter_util.h"
  #include "oneflow/core/job_rewriter/autograd.h"
  #include "oneflow/core/operator/op_conf.pb.h"
  #include "oneflow/core/operator/operator.h"
  #include "oneflow/core/register/logical_blob_id.pb.h"
  
  namespace oneflow {
  namespace one {
  
  struct DefaultOpExprInterpState : public OpExprInterpState {
    std::vector<bool> requires_grads;
  };
  
  class DefaultOpExprGradFunction : public OpExprGradFunction<DefaultOpExprInterpState> {
   public:
    // The snapshot indicates the indices of the required forward inputs and outputs
    // by each backward operator.
    struct Snapshot {
      std::vector<int> input_indices;
      std::vector<int> output_indices;
      int count = 0;
    };
  
    Maybe<void> Init(const OpExpr& op) override;
  
    Maybe<void> Capture(DefaultOpExprInterpState* ctx, const TensorTuple& inputs,
                        const TensorTuple& outputs, const AttrMap& attrs) const override;
  
    Maybe<void> Apply(const DefaultOpExprInterpState* ctx, const TensorTuple& out_grads,
                      TensorTuple* in_grads) const override;
  
   private:
    Maybe<void> GenerateOpGradConf(const Operator& op, std::vector<OperatorConf>* bw_op_confs);
  
    Maybe<void> UpdateRequiresBackward(DefaultOpExprInterpState* ctx,
                                       const TensorTuple& inputs) const;
  
   private:
    struct BackwardEntry {
      std::shared_ptr<OpExpr> backward_op;
      // Each output of the backward op maybe related to multiple input gradients.
      std::vector<std::vector<int>> in_grad_indices_for_each_output;
      std::vector<int> out_grad_indices;
  
      // Snapshot information for each backward op.
      Snapshot snapshot;
    };
    std::vector<BackwardEntry> backward_entries_;
  
    int input_size_;
  
    // The input gradient logical blob id for each forward input blob name which
    // needs backward.
    HashMap<std::string, LogicalBlobId> ibn_to_grad_lbi_map_;
    // The output gradient logical blob id for each forward output blob name.
    HashMap<std::string, LogicalBlobId> obn_to_grad_lbi_map_;
  };
  
  namespace {
  
  Maybe<void> ProcessInput(const std::string& bn, const std::string& lbn,
                           const HashMap<std::string, int>& input_lbn_to_index_map,
                           const HashMap<std::string, int>& output_lbn_to_index_map,
                           const HashMap<std::string, std::string>& out_grad_lbn_to_obn_map,
                           const HashMap<std::string, int>& obn_to_index_map,
                           DefaultOpExprGradFunction::Snapshot* snapshot,
                           std::vector<int>* out_grad_indices,
                           std::vector<std::vector<std::string>>* typed_ibns) {
    CHECK_EQ_OR_RETURN(typed_ibns->size(), 3);
    if (input_lbn_to_index_map.count(lbn)) {
      snapshot->input_indices.emplace_back(input_lbn_to_index_map.at(lbn));
      snapshot->count += 1;
      (*typed_ibns)[0].emplace_back(bn);
    } else if (output_lbn_to_index_map.count(lbn)) {
      snapshot->output_indices.emplace_back(output_lbn_to_index_map.at(lbn));
      snapshot->count += 1;
      (*typed_ibns)[1].emplace_back(bn);
    } else {
      const auto& it = out_grad_lbn_to_obn_map.find(lbn);
      // Otherwise this input should be output gradient.
      CHECK_OR_RETURN(it != out_grad_lbn_to_obn_map.end());
      // The output gradient index is equal to the forward output index.
      out_grad_indices->emplace_back(obn_to_index_map.at(it->second));
      (*typed_ibns)[2].emplace_back(bn);
    }
    return Maybe<void>::Ok();
  }
  
  Maybe<void> ProcessOutput(
      const std::string& bn, const std::string& lbn,
      const HashMap<std::string, std::vector<std::string>>& in_grad_lbn_to_ibns_map,
      const HashMap<std::string, int>& ibn_to_index_map,
      std::vector<std::vector<int>>* in_grad_indices_for_each_output,
      std::vector<std::string>* indexed_obns, HashSet<int>* reached_in_grad_indices_for_each_output) {
    indexed_obns->emplace_back(bn);
    const auto& it = in_grad_lbn_to_ibns_map.find(lbn);
    CHECK_OR_RETURN(it != in_grad_lbn_to_ibns_map.end());
    // The input gradient index is equal to the foward input index.
    in_grad_indices_for_each_output->emplace_back();
    for (const std::string& ibn : it->second) {
      int idx = ibn_to_index_map.at(ibn);
      in_grad_indices_for_each_output->back().emplace_back(idx);
      reached_in_grad_indices_for_each_output->emplace(idx);
    }
    return Maybe<void>::Ok();
  }
  
  }  // namespace
  
  Maybe<void> DefaultOpExprGradFunction::GenerateOpGradConf(const Operator& op,
                                                            std::vector<OperatorConf>* bw_op_confs) {
    auto DiffLbi4BnInOp = [&](const std::string& bn) -> LogicalBlobId* {
      const auto& input_bns = op.input_bns();
      const auto& output_bns = op.output_bns();
      if (std::find(input_bns.begin(), input_bns.end(), bn) != input_bns.end()) {
        if (!op.InputBlobModifier4Ibn(bn).requires_grad()) { return nullptr; }
        return &ibn_to_grad_lbi_map_[bn];
      } else if (std::find(output_bns.begin(), output_bns.end(), bn) != output_bns.end()) {
        if (!op.OutputBlobModifier4Obn(bn).requires_grad()) { return nullptr; }
        auto it = obn_to_grad_lbi_map_.find(bn);
        if (it == obn_to_grad_lbi_map_.end()) {
          LogicalBlobId lbi;
          lbi.set_op_name(GradientOpName(op.op_name()));
          lbi.set_blob_name(bn);
          it = obn_to_grad_lbi_map_.emplace(bn, lbi).first;
        }
        return &(it->second);
      } else {
        LOG(FATAL) << "Blob name (" << op.op_name() << "/" << bn << ") is missing.";
      }
      return nullptr;
    };
    const auto& dummy_blob_desc = BlobDesc(Shape(), DataType::kInvalidDataType);
    auto LogicalBlobDesc4BnInOp = [&](const std::string& bn) -> const BlobDesc& {
      return dummy_blob_desc;
    };
    JUST(GenerateBackwardOpConfIf(op, bw_op_confs, DiffLbi4BnInOp, LogicalBlobDesc4BnInOp));
    return Maybe<void>::Ok();
  }
  
  Maybe<void> DefaultOpExprGradFunction::Init(const OpExpr& op) {
    if (op.input_size() == 0) { return Maybe<void>::Ok(); }
    const auto* fw_op_expr = dynamic_cast<const BuiltinOpExpr*>(&op);
    CHECK_NOTNULL_OR_RETURN(fw_op_expr);
  
    input_size_ = fw_op_expr->indexed_ibns().size();
  
    OperatorConf fw_op_conf;
    fw_op_expr->BuildOpConf(&fw_op_conf, /*attrs=*/{});
  
    // Generate backward operator conf for each input. The `LogicalBlobId` for
    // backward output gradient is dummy due to inaccessibility.
    std::vector<OperatorConf> bw_op_confs;
    std::shared_ptr<Operator> op_adapter = ConstructOp(fw_op_conf, DeviceType::kCPU);
    JUST(GenerateOpGradConf(*op_adapter, &bw_op_confs));
  
    backward_entries_.resize(bw_op_confs.size());
  
    HashMap<std::string, int> input_lbn_to_index_map;
    HashMap<std::string, int> output_lbn_to_index_map;
    HashMap<std::string, int> ibn_to_index_map;
    HashMap<std::string, int> obn_to_index_map;
    for (int i = 0; i < fw_op_expr->indexed_ibns().size(); ++i) {
      const auto& bn = fw_op_expr->indexed_ibns().at(i);
      ibn_to_index_map.emplace(bn, i);
      std::string lbn = GenLogicalBlobName(op_adapter->BnInOp2Lbi(bn));
      input_lbn_to_index_map.emplace(lbn, i);
    }
    for (int i = 0; i < fw_op_expr->indexed_obns().size(); ++i) {
      const auto& bn = fw_op_expr->indexed_obns().at(i);
      obn_to_index_map.emplace(bn, i);
      std::string lbn = GenLogicalBlobName(op_adapter->BnInOp2Lbi(bn));
      output_lbn_to_index_map.emplace(lbn, i);
    }
  
    HashMap<std::string, std::string> out_grad_lbn_to_obn_map;
    HashMap<std::string, std::vector<std::string>> in_grad_lbn_to_ibns_map;
    for (const auto& it : obn_to_grad_lbi_map_) {
      std::string lbn = GenLogicalBlobName(it.second);
      out_grad_lbn_to_obn_map[lbn] = it.first;
    }
    for (const auto& it : ibn_to_grad_lbi_map_) {
      std::string lbn = GenLogicalBlobName(it.second);
      in_grad_lbn_to_ibns_map[lbn].emplace_back(it.first);
    }
    HashSet<int> reached_in_grad_indices_for_each_output;
    for (int i = 0; i < bw_op_confs.size(); ++i) {
      const auto& op_conf = bw_op_confs.at(i);
      VLOG(3) << op_conf.DebugString() << std::endl;
      std::shared_ptr<Operator> bw_op_adapter = ConstructOp(op_conf, DeviceType::kCPU);
      std::vector<std::string> bw_indexed_ibns, bw_indexed_obns;
  
      BackwardEntry* entry = &(backward_entries_[i]);
      // Input blob names in the backward op will be divided into 3 types which
      // means the input comes from either forward inputs, or forward outputs, or
      // backward output gradients.
      std::vector<std::vector<std::string>> typed_ibns(3);
      entry->snapshot.count = 0;
      for (const auto& bn : bw_op_adapter->input_bns()) {
        std::string lbn = GenLogicalBlobName(bw_op_adapter->BnInOp2Lbi(bn));
        JUST(ProcessInput(bn, lbn, input_lbn_to_index_map, output_lbn_to_index_map,
                          out_grad_lbn_to_obn_map, obn_to_index_map, &entry->snapshot,
                          &entry->out_grad_indices, &typed_ibns));
      }
      for (const auto& ibns : typed_ibns) {
        for (const auto& v : ibns) { bw_indexed_ibns.emplace_back(v); }
      }
      for (const auto& bn : bw_op_adapter->output_bns()) {
        std::string lbn = GenLogicalBlobName(bw_op_adapter->BnInOp2Lbi(bn));
        JUST(ProcessOutput(bn, lbn, in_grad_lbn_to_ibns_map, ibn_to_index_map,
                           &entry->in_grad_indices_for_each_output, &bw_indexed_obns,
                           &reached_in_grad_indices_for_each_output));
      }
  
      // Currently only user op is considered.
      if (op_conf.has_user_conf()) {
        UserOpConf user_conf(op_conf.user_conf());
        entry->backward_op = JUST(
            UserOpExpr::New(op_conf.name(), std::move(user_conf), bw_indexed_ibns, bw_indexed_obns));
      } else {
        // TODO()
        UNIMPLEMENTED();
      }
    }
  
    // Create identity ops for the inplaced outputs.
    for (const auto& it : ibn_to_index_map) {
      const int& in_grad_idx = it.second;
      if (reached_in_grad_indices_for_each_output.count(in_grad_idx)) { continue; }
      const auto& grad_it = ibn_to_grad_lbi_map_.find(/*in_grad_bn=*/it.first);
      // Skip the input that does not require grad.
      if (grad_it == ibn_to_grad_lbi_map_.end()) { continue; }
      std::string lbn = GenLogicalBlobName(grad_it->second);
      std::vector<std::vector<std::string>> typed_ibns(3);
      backward_entries_.emplace_back();
      BackwardEntry* entry = &(backward_entries_.back());
      entry->snapshot.count = 0;
      JUST(ProcessInput(it.first, lbn, input_lbn_to_index_map, output_lbn_to_index_map,
                        out_grad_lbn_to_obn_map, obn_to_index_map, &entry->snapshot,
                        &entry->out_grad_indices, &typed_ibns));
      entry->in_grad_indices_for_each_output.emplace_back(std::vector<int>{in_grad_idx});
      entry->backward_op = JUST(OpBuilder("identity").Input("in").Output("out").Build());
    }
    return Maybe<void>::Ok();
  }
  
  Maybe<void> DefaultOpExprGradFunction::UpdateRequiresBackward(DefaultOpExprInterpState* ctx,
                                                                const TensorTuple& inputs) const {
    ctx->requires_grads.resize(backward_entries_.size());
    for (int i = 0; i < backward_entries_.size(); ++i) {
      const auto& entry = backward_entries_.at(i);
      const auto& in_grad_indices = entry.in_grad_indices_for_each_output;
      ctx->requires_grads.at(i) =  // NOLINT
          std::any_of(in_grad_indices.begin(), in_grad_indices.end(),
                      [&](const std::vector<int>& indices) {
                        for (const int& j : indices) {
                          if (inputs.at(j)->requires_grad()) { return true; }
                        }
                        return false;
                      });
    }
    return Maybe<void>::Ok();
  }
  
  Maybe<void> DefaultOpExprGradFunction::Capture(DefaultOpExprInterpState* ctx,
                                                 const TensorTuple& inputs,
                                                 const TensorTuple& outputs,
                                                 const AttrMap& attrs) const {
    CHECK_OR_RETURN(attrs.empty())
        << "The default op expr gradient func does not support dynamic attributes.";
    JUST(UpdateRequiresBackward(ctx, inputs));
    for (int i = 0; i < backward_entries_.size(); ++i) {
      if (!ctx->requires_grads.at(i)) { continue; }
      const auto& entry = backward_entries_.at(i);
      for (const int& idx : entry.snapshot.input_indices) {
        ctx->SaveTensorForBackward(inputs.at(idx));
      }
      for (const int& idx : entry.snapshot.output_indices) {
        ctx->SaveTensorForBackward(outputs.at(idx));
      }
    }
    return Maybe<void>::Ok();
  }
  
  Maybe<void> DefaultOpExprGradFunction::Apply(const DefaultOpExprInterpState* ctx,
                                               const TensorTuple& out_grads,
                                               TensorTuple* in_grads) const {
    CHECK_EQ_OR_RETURN(in_grads->size(), input_size_);
    const auto& saved_tensors = ctx->SavedTensors();
    int offset = 0;
    for (int i = 0; i < backward_entries_.size(); ++i) {
      if (!ctx->requires_grads.at(i)) { continue; }
      const auto& entry = backward_entries_.at(i);
      TensorTuple inputs;
      for (int j = 0; j < entry.snapshot.count; ++j) {
        inputs.emplace_back(saved_tensors.at(offset + j));
      }
      for (const int& idx : entry.out_grad_indices) { inputs.emplace_back(out_grads.at(idx)); }
      TensorTuple outputs(entry.in_grad_indices_for_each_output.size());
      const auto& interpreter = JUST(OpInterpUtil::GetInterpreter());
      JUST(interpreter->Apply(*(entry.backward_op), inputs, &outputs));
  
      for (int j = 0; j < entry.in_grad_indices_for_each_output.size(); ++j) {
        for (const int& idx : entry.in_grad_indices_for_each_output.at(j)) {
          in_grads->at(idx) = outputs.at(j);
        }
      }
      offset += entry.snapshot.count;
    }
    return Maybe<void>::Ok();
  }
  
  REGISTER_OP_EXPR_GRAD_FUNCTION("default", DefaultOpExprGradFunction);
  
  }  // namespace one
  }  // namespace oneflow
  