/*
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 <memory>
#include "oneflow/core/common/cpp_attribute.h"
#include "oneflow/core/common/maybe.h"
#include "oneflow/core/common/cpp_attribute.h"
#include "oneflow/core/common/container_util.h"
#include "oneflow/core/framework/consistency_check.h"
#include "oneflow/core/framework/user_op_conf.h"
#include "oneflow/core/framework/op_expr.h"
#include "oneflow/core/framework/op_builder.h"
#include "oneflow/core/framework/multi_client_session_context.h"
#include "oneflow/core/framework/op_interpreter.h"
#include "oneflow/core/framework/op_interpreter/op_interpreter_util.h"
#include "oneflow/core/framework/instructions_builder.h"
#include "oneflow/core/framework/scope_util.h"
#include "oneflow/core/framework/session_util.h"
#include "oneflow/core/framework/symbol_storage_util.h"
#include "oneflow/core/framework/tensor.h"
#include "oneflow/core/framework/tensor_name_scope.h"
#include "oneflow/core/framework/tensor_tuple.h"
#include "oneflow/core/framework/user_op_registry.h"
#include "oneflow/core/framework/nd_sbp.h"
#include "oneflow/core/job/lazy_mode.h"
#include "oneflow/core/operator/operator.h"
#include "oneflow/core/job/sbp_parallel.h"
#include "oneflow/core/job/job_build_and_infer_ctx_mgr.h"
#include "oneflow/core/vm/vm_util.h"
#include "oneflow/core/functional/functional.h"
namespace oneflow {

namespace one {

namespace {

Maybe<Tensor> BuildTensor(const OpAttribute& op_attribute, const std::string& bn_in_op,
                          const std::shared_ptr<ParallelDesc>& parallel_desc, const bool is_lazy,
                          const bool is_local) {
  CHECK_OR_RETURN(op_attribute.has_logical_blob_desc_signature());
  const auto& blob_desc_sign_map = op_attribute.logical_blob_desc_signature().bn_in_op2blob_desc();
  auto blob_desc_it = blob_desc_sign_map.find(bn_in_op);
  CHECK_OR_RETURN(blob_desc_it != blob_desc_sign_map.end())
      << "blob_desc of " << bn_in_op << " not found in op " << op_attribute.op_conf().name();

  auto shape = std::make_shared<Shape>(blob_desc_it->second.shape());
  auto stride = std::make_shared<Stride>(shape);
  auto dtype = blob_desc_it->second.data_type();
  if (is_local) {
    const auto& device = JUST(Device::MakeDeviceByParallelDesc(*parallel_desc));
    const auto& tensor =
        JUST(MirroredTensor::MakeTensor(shape, stride, dtype, device, is_lazy,
                                        /* requires_grad= */ false, /* is_leaf= */ true));
    return static_cast<std::shared_ptr<Tensor>>(tensor);
  } else {
    const auto& nd_sbp_sign_map = op_attribute.nd_sbp_signature().bn_in_op2nd_sbp();
    auto nd_sbp_it = nd_sbp_sign_map.find(bn_in_op);
    CHECK_OR_RETURN(nd_sbp_it != nd_sbp_sign_map.end())
        << "nd_sbp of " << bn_in_op << " not found in op " << op_attribute.op_conf().name();
    NdSbp nd_sbp(nd_sbp_it->second);
    const auto& tensor = JUST(ConsistentTensor::MakeTensor(
        shape, dtype, SymbolOf(nd_sbp), SymbolOf(*parallel_desc), is_lazy,
        /*requires_grad=*/false, /*is_leaf=*/true));
    return static_cast<std::shared_ptr<Tensor>>(tensor);
  }
}

Maybe<void> CheckTensorMatchAttr(const std::shared_ptr<Tensor>& tensor,
                                 const OpAttribute& op_attribute, const std::string& bn_in_op,
                                 const std::shared_ptr<ParallelDesc>& parallel_desc,
                                 const bool is_local) {
  CHECK_EQ_OR_RETURN(tensor->is_local(), is_local);

  CHECK_OR_RETURN(op_attribute.has_logical_blob_desc_signature());
  const auto& blob_desc_sign_map = op_attribute.logical_blob_desc_signature().bn_in_op2blob_desc();
  auto blob_desc_it = blob_desc_sign_map.find(bn_in_op);
  CHECK_OR_RETURN(blob_desc_it != blob_desc_sign_map.end())
      << "blob_desc of " << bn_in_op << " not found in op " << op_attribute.op_conf().name();

  auto shape = std::make_shared<Shape>(blob_desc_it->second.shape());
  auto dtype = blob_desc_it->second.data_type();
  CHECK_EQ_OR_RETURN(*tensor->shape(), *shape);
  CHECK_EQ_OR_RETURN(tensor->dtype()->data_type(), dtype);

  if (is_local) {
    const auto& device = JUST(Device::MakeDeviceByParallelDesc(*parallel_desc));
    CHECK_OR_RETURN(JUST(tensor->device()) == device);
  } else {
    const auto& nd_sbp_sign_map = op_attribute.nd_sbp_signature().bn_in_op2nd_sbp();
    auto nd_sbp_it = nd_sbp_sign_map.find(bn_in_op);
    CHECK_OR_RETURN(nd_sbp_it != nd_sbp_sign_map.end())
        << "nd_sbp of " << bn_in_op << " not found in op " << op_attribute.op_conf().name();
    NdSbp nd_sbp(nd_sbp_it->second);
    CHECK_OR_RETURN(JUST(tensor->nd_sbp()) == SymbolOf(nd_sbp))
        << "The input sbp is not valid for an inplace operation, please try to use non-inplace. "
        << NdSbpToString(JUST(tensor->nd_sbp())) << " vs " << NdSbpToString(nd_sbp);
    CHECK_OR_RETURN(JUST(tensor->parallel_desc()) == SymbolOf(*parallel_desc));
  }
  return Maybe<void>::Ok();
}

Maybe<const std::string&> GetDeviceTagOfTensor(const std::shared_ptr<Tensor>& tensor) {
  if (tensor->is_consistent()) { return JUST(tensor->parallel_desc())->device_tag(); }
  return JUST(tensor->device())->type();
}

bool GetIsDynamicOfTensor(const std::shared_ptr<Tensor>& tensor) {
  if (tensor->is_consistent()) {
    return false;
  } else {
    return true;
  }
}

Maybe<void> GenNdSbpByTensor(NdSbp* nd_sbp, const std::shared_ptr<Tensor>& tensor) {
  nd_sbp->clear_sbp_parallel();
  if (tensor->is_local()) {
    // NOTE(chengcheng):
    //   OneFlow Lazy is always consistent. LocalTensor is a special case of ConsistentTensor which
    //   placement is only this rank, and SbpParallel is Broadcast.
    nd_sbp->add_sbp_parallel()->mutable_broadcast_parallel();
  } else {
    *nd_sbp = *JUST(tensor->nd_sbp());
  }
  return Maybe<void>::Ok();
}

Maybe<void> GenVariableOpConfNdSbpStringByTensor(VariableOpConf* var_conf,
                                                 const std::shared_ptr<Tensor>& tensor) {
  var_conf->clear_nd_sbp();
  if (tensor->is_local()) {
    SbpParallel broadcast;
    broadcast.mutable_broadcast_parallel();
    var_conf->add_nd_sbp(SbpParallelToString(broadcast));
  } else {
    const NdSbp& nd_sbp = *JUST(tensor->nd_sbp());
    for (const auto& sbp_parallel : nd_sbp.sbp_parallel()) {
      var_conf->add_nd_sbp(SbpParallelToString(sbp_parallel));
    }
  }
  return Maybe<void>::Ok();
}

Maybe<const ParallelDesc> GetParallelDescOfTensor(const std::shared_ptr<Tensor>& tensor) {
  if (tensor->is_local()) {
    const auto& device = JUST(tensor->device());
    const auto& placement = JUST(Placement4Device(device));
    return placement.shared_from_symbol();
  } else {
    return JUST(tensor->parallel_desc()).shared_from_symbol();
  }
}

Maybe<Scope> NewScopeWithParallelConfAndCurScope(const ParallelConf& parallel_conf) {
  std::shared_ptr<Scope> new_scope;
  const auto& old_scope = JUST(GetCurrentScope());
  JUST(PhysicalRun([&](InstructionsBuilder* builder) -> Maybe<void> {
    new_scope = JUST(builder->BuildScopeWithNewParallelConf(old_scope, parallel_conf));
    return Maybe<void>::Ok();
  }));
  // NOTE(chengcheng): need sync vm for get scope right now
  JUST(vm::CurrentRankSync());
  CHECK_OR_RETURN(new_scope);
  return new_scope;
}

Maybe<Scope> NewScopeWithParallelDescByTensor(const std::shared_ptr<Tensor>& tensor) {
  return NewScopeWithParallelConfAndCurScope(
      JUST(GetParallelDescOfTensor(tensor))->parallel_conf());
}

int32_t GetGradAccStep(const JobConfigProto& job_conf) {
  if (job_conf.has_train_conf() && job_conf.has_num_gradient_accumulation_steps()
      && job_conf.num_gradient_accumulation_steps() > 1) {
    return job_conf.num_gradient_accumulation_steps();
  } else {
    return 1;
  }
}

Maybe<Tensor> GradAccTryInsertUnpackAfterInput(
    const OperatorConf& input_conf, const std::shared_ptr<ParallelDesc>& blob_parallel_desc,
    const std::shared_ptr<Tensor>& input_tensor) {
  auto infer_ctx = JUST(GetCurInferCtx());
  int64_t grad_acc_step = GetGradAccStep(infer_ctx->job().job_conf());
  if (grad_acc_step > 1) {
    // NOTE(chengcheng):
    //   We assume that the input data is one mini-batch which containing multi micro-batches.
    //   So we need unpack input data for each micro-batch.
    VLOG(2)
        << " Current OneFlow nn.Graph grad acc semantics is different from Torch. \n"
        << " Once call nn.Graph in OneFlow, it indicates a mini-batch. When grad acc steps > 1, \n"
        << " the input tensor of nn.Graph will be unpacked by 0th dim into multiple micro-batches "
        << " and exec them in order.\n";

    user_op::UserOpConfWrapperBuilder unpack_builder("System-GradientAccumulation-InputUnpack-"
                                                     + input_conf.name() + "-" + NewUniqueId());
    const std::string input_tensor_lbn = GenLogicalBlobName(input_conf.name(), "out");
    const auto unpack_op = unpack_builder.OpTypeName("unpack")
                               .Input("in", input_tensor_lbn)
                               .Output("out")
                               .Attr<int32_t>("unpack_num", grad_acc_step)
                               .ScopeSymbolId(input_conf.scope_symbol_id())
                               .DeviceTag(input_conf.device_tag())
                               .Build();

    OpAttribute unpack_op_attr = *JUST(infer_ctx->AddAndInferConsistentOp(unpack_op.op_conf()));
    VLOG(2) << "Lazy nn.Graph name " << infer_ctx->job().job_conf().job_name() << " add op: \n"
            << unpack_op.op_conf().DebugString() << std::endl;
    VLOG(3) << "Lazy nn.Graph name " << infer_ctx->job().job_conf().job_name()
            << " infer and and op attr : \n"
            << unpack_op_attr.DebugString() << std::endl;

    const std::string unpack_lbn = unpack_op.output("out", 0);
    auto unpack_input =
        JUST(BuildTensor(unpack_op_attr, "out_0", blob_parallel_desc, /* is_lazy= */ true,
                         /* is_local= */ input_tensor->is_local()));
    TensorNameScope::Global()->Record(unpack_input, unpack_lbn);
    return unpack_input;
  } else {
    return input_tensor;
  }
}

Maybe<Tensor> GradAccTryInsertRepeatAfterVar(
    const OperatorConf& var_conf, const std::shared_ptr<ParallelDesc>& blob_parallel_desc,
    const std::shared_ptr<Tensor>& var_tensor) {
  auto infer_ctx = JUST(GetCurInferCtx());
  int64_t grad_acc_step = GetGradAccStep(infer_ctx->job().job_conf());
  if (grad_acc_step > 1) {
    // NOTE(chengcheng):
    //   We assume that the nn.Graph once call is one mini-batch which containing multi
    //   micro-batches. So we just repeat variable tensor for each micro-batch.
    VLOG(2)
        << " Current OneFlow nn.Graph grad acc semantics is different from Torch. \n"
        << " Once call nn.Graph in OneFlow, it indicates a mini-batch. When grad acc steps > 1, \n"
        << " the var tensor of nn.Graph will be repeated exec for multiple micro-batches. \n";

    const std::string var_tensor_lbn = GenLogicalBlobName(var_conf.name(), "out");
    user_op::UserOpConfWrapperBuilder repeat_builder("System-GradientAccumulation-VariableRepeat-"
                                                     + var_conf.name() + "-" + NewUniqueId());
    const auto repeat_op = repeat_builder.OpTypeName("repeat")
                               .Input("in", var_tensor_lbn)
                               .Output("out")
                               .Attr<int32_t>("repeat_num", grad_acc_step)
                               .ScopeSymbolId(var_conf.scope_symbol_id())
                               .DeviceTag(var_conf.device_tag())
                               .Build();

    OpAttribute repeat_op_attr = *JUST(infer_ctx->AddAndInferConsistentOp(repeat_op.op_conf()));
    VLOG(2) << "Lazy nn.Graph name " << infer_ctx->job().job_conf().job_name() << " add op: \n"
            << repeat_op.op_conf().DebugString() << std::endl;
    VLOG(3) << "Lazy nn.Graph name " << infer_ctx->job().job_conf().job_name()
            << " infer and and op attr : \n"
            << repeat_op_attr.DebugString() << std::endl;

    const std::string repeat_lbn = repeat_op.output("out", 0);
    auto repeat_var =
        JUST(BuildTensor(repeat_op_attr, "out_0", blob_parallel_desc, /* is_lazy= */ true,
                         /* is_local= */ var_tensor->is_local()));
    TensorNameScope::Global()->Record(repeat_var, repeat_lbn);
    return repeat_var;
  } else {
    return var_tensor;
  }
}

Maybe<Tensor> GradAccTryInsertPackBeforeOutput(const std::shared_ptr<Scope>& scope,
                                               const std::string& output_in_lbn,
                                               const std::string& output_op_name,
                                               const std::shared_ptr<Tensor>& output_tensor) {
  auto infer_ctx = JUST(GetCurInferCtx());
  int64_t grad_acc_step = GetGradAccStep(infer_ctx->job().job_conf());
  if (grad_acc_step > 1) {
    // NOTE(chengcheng):
    //   We assume that the nn.Graph once call is one mini-batch which containing multi
    //   micro-batches. So we need pack output tensor for each micro-batch to one micro-batch.
    VLOG(2)
        << " Current OneFlow nn.Graph grad acc semantics is different from Torch. \n"
        << " Once call nn.Graph in OneFlow, it indicates a mini-batch. When grad acc steps > 1, \n"
        << " the output tensor of nn.Graph will be packed to a big tensor by 0th dim, after exec \n"
        << " for multiple micro-batches. \n";

    user_op::UserOpConfWrapperBuilder pack_builder("System-GradientAccumulation-OutputPack-"
                                                   + output_op_name);
    const auto output_pack_op = pack_builder.OpTypeName("pack")
                                    .Input("in", output_in_lbn)
                                    .Output("out")
                                    .Attr<int32_t>("pack_num", grad_acc_step)
                                    .ScopeSymbolId(JUST(scope->symbol_id()))
                                    .DeviceTag(JUST(GetDeviceTagOfTensor(output_tensor)))
                                    .Build();

    int64_t parallel_desc_sym_id = JUST(scope->GetParallelDescSymbolId(output_pack_op.op_conf()));
    auto blob_parallel_desc = JUST(GetSymbol<ParallelDesc>(parallel_desc_sym_id));

    OpAttribute pack_op_attr = *JUST(infer_ctx->AddAndInferConsistentOp(output_pack_op.op_conf()));
    VLOG(2) << "Lazy nn.Graph name " << infer_ctx->job().job_conf().job_name() << " add op: \n"
            << output_pack_op.op_conf().DebugString() << std::endl;
    VLOG(3) << "Lazy nn.Graph name " << infer_ctx->job().job_conf().job_name()
            << " infer and and op attr : \n"
            << pack_op_attr.DebugString() << std::endl;

    const std::string pack_lbn = output_pack_op.output("out", 0);
    auto packed_output =
        JUST(BuildTensor(pack_op_attr, "out_0", blob_parallel_desc, /* is_lazy= */ true,
                         /* is_local= */ output_tensor->is_local()));
    TensorNameScope::Global()->Record(packed_output, pack_lbn);
    return packed_output;
  } else {
    return output_tensor;
  }
}

Maybe<void> GradAccTryInsertRepeatTickBeforeSource(
    const std::shared_ptr<OperatorConf>& source_op_conf) {
  auto infer_ctx = JUST(GetCurInferCtx());
  int64_t grad_acc_step = GetGradAccStep(infer_ctx->job().job_conf());
  if (grad_acc_step > 1) {
    // NOTE(chengcheng):
    //   We assume that the nn.Graph once call is one mini-batch which containing multi
    //   micro-batches. So we need repeat source op for each micro-batch in one micro-batch.
    VLOG(2)
        << " Current OneFlow nn.Graph grad acc semantics is different from Torch. \n"
        << " Once call nn.Graph in OneFlow, it indicates a mini-batch. When grad acc steps > 1, \n"
        << " the source op of nn.Graph will be repeated exec n-times for multiple micro-batches.\n";

    // Insert Tick
    OperatorConf tick_conf{};
    tick_conf.set_name("System-GradientAccumulation-RepeatTick-DeviceTick-"
                       + source_op_conf->name());
    tick_conf.set_device_tag(source_op_conf->device_tag());
    tick_conf.mutable_device_tick_conf()->set_out("out");
    tick_conf.set_scope_symbol_id(source_op_conf->scope_symbol_id());
    auto tick_lbn = GenLogicalBlobName(tick_conf.name(), tick_conf.device_tick_conf().out());
    OpAttribute tick_op_attr = *JUST(infer_ctx->AddAndInferConsistentOp(tick_conf));
    VLOG(2) << "Lazy nn.Graph name " << infer_ctx->job().job_conf().job_name() << " add op: \n"
            << tick_conf.DebugString() << std::endl;
    VLOG(3) << "Lazy nn.Graph name " << infer_ctx->job().job_conf().job_name()
            << " infer and and op attr : \n"
            << tick_op_attr.DebugString() << std::endl;

    user_op::UserOpConfWrapperBuilder repeat_builder(
        "System-GradientAccumulation-RepeatTick-Repeat-" + source_op_conf->name());
    const auto repeat_op = repeat_builder.OpTypeName("repeat")
                               .Input("in", tick_lbn)
                               .Output("out")
                               .Attr<int32_t>("repeat_num", grad_acc_step)
                               .ScopeSymbolId(source_op_conf->scope_symbol_id())
                               .DeviceTag(source_op_conf->device_tag())
                               .Build();

    OpAttribute repeat_op_attr = *JUST(infer_ctx->AddAndInferConsistentOp(repeat_op.op_conf()));
    VLOG(2) << "Lazy nn.Graph name " << infer_ctx->job().job_conf().job_name() << " add op: \n"
            << repeat_op.op_conf().DebugString() << std::endl;
    VLOG(3) << "Lazy nn.Graph name " << infer_ctx->job().job_conf().job_name()
            << " infer and and op attr : \n"
            << repeat_op_attr.DebugString() << std::endl;

    const std::string repeat_tick_lbn = repeat_op.output("out", 0);
    (*source_op_conf->mutable_user_conf()->mutable_input())[user_op::kUserSourceOpTickInputArgName]
        .add_s(repeat_op.output("out", 0));
  }
  return Maybe<void>::Ok();
}

Maybe<std::string> GradAccTryInsertRepeatAfterFreeVar(const OperatorConf& var_conf) {
  const std::string var_tensor_lbn = GenLogicalBlobName(var_conf.name(), "out");
  auto infer_ctx = JUST(GetCurInferCtx());
  int64_t grad_acc_step = GetGradAccStep(infer_ctx->job().job_conf());
  if (grad_acc_step > 1) {
    // NOTE(chengcheng):
    //   We assume that the nn.Graph once call is one mini-batch which containing multi
    //   micro-batches. So we just repeat variable tensor for each micro-batch.
    VLOG(2)
        << " Current OneFlow nn.Graph grad acc semantics is different from Torch. \n"
        << " Once call nn.Graph in OneFlow, it indicates a mini-batch. When grad acc steps > 1, \n"
        << " the free var tensor of nn.Graph will be repeated exec for multiple micro-batches. \n";

    user_op::UserOpConfWrapperBuilder repeat_builder("System-GradientAccumulation-VariableRepeat-"
                                                     + var_conf.name() + "-" + NewUniqueId());
    const auto repeat_op = repeat_builder.OpTypeName("repeat")
                               .Input("in", var_tensor_lbn)
                               .Output("out")
                               .Attr<int32_t>("repeat_num", grad_acc_step)
                               .ScopeSymbolId(var_conf.scope_symbol_id())
                               .DeviceTag(var_conf.device_tag())
                               .Build();

    OpAttribute repeat_op_attr = *JUST(infer_ctx->AddAndInferConsistentOp(repeat_op.op_conf()));
    VLOG(2) << "Lazy nn.Graph name " << infer_ctx->job().job_conf().job_name() << " add op: \n"
            << repeat_op.op_conf().DebugString() << std::endl;
    VLOG(3) << "Lazy nn.Graph name " << infer_ctx->job().job_conf().job_name()
            << " infer and and op attr : \n"
            << repeat_op_attr.DebugString() << std::endl;

    const std::string repeat_lbn = repeat_op.output("out", 0);
    return repeat_lbn;
  } else {
    return var_tensor_lbn;
  }
}

Maybe<void> AddFreeEagerTensorToVariableOp(const std::shared_ptr<Tensor>& input_tensor) {
  if (!input_tensor->is_contiguous()) {
    LazyMode::Guard lazy_mode_disabled_guard(false);
    JUST(functional::InplaceToContiguous(input_tensor));
    JUST(vm::CurrentRankSync());
  }

  CHECK_OR_RETURN(input_tensor->is_eager());
  const std::string& empty_lbn = TensorNameScope::Global()->Lookup(input_tensor);
  CHECK_OR_RETURN(empty_lbn.empty());
  std::shared_ptr<Scope> scope = JUST(NewScopeWithParallelDescByTensor(input_tensor));
  OperatorConf op_conf;
  op_conf.set_scope_symbol_id(JUST(scope->symbol_id()));
  op_conf.set_device_tag(JUST(GetDeviceTagOfTensor(input_tensor)));
  VariableOpConf* var_conf = op_conf.mutable_variable_conf();
  var_conf->set_out("out");
  input_tensor->shape()->ToProto(var_conf->mutable_shape());
  var_conf->set_data_type(input_tensor->dtype()->data_type());
  // NOTE(chengcheng): VariableOpConf initializer_conf is useless because variable is inited
  //   by EagerTensor.
  var_conf->mutable_initializer()->mutable_empty_conf();
  JUST(GenVariableOpConfNdSbpStringByTensor(var_conf, input_tensor));
  // NOTE(chengcheng): Free EagerTensor not trainable
  var_conf->set_trainable(false);

  auto infer_ctx = JUST(GetCurInferCtx());
  // NOTE(chengcheng): MUST reset unique op name before InferCtx::AddOp, FreeEagerTensor has no
  //  name so just new a unique name for it.
  const std::string new_op_name = *JUST(infer_ctx->NewUniqueOpNameByFunctionalOpConf(op_conf));
  op_conf.set_name(new_op_name);

  VLOG(2) << "Lazy nn.Graph name " << infer_ctx->job().job_conf().job_name() << " try to add op: \n"
          << op_conf.DebugString() << std::endl;
  OpAttribute op_attr = *JUST(infer_ctx->AddAndInferConsistentOp(op_conf));
  VLOG(2) << "Lazy nn.Graph name " << infer_ctx->job().job_conf().job_name() << " add op : \n"
          << op_conf.name() << " for FreeEagerTensor.\n";
  VLOG(3) << "Lazy nn.Graph name " << infer_ctx->job().job_conf().job_name()
          << " infer and and op attr : \n"
          << op_attr.DebugString() << " for FreeEagerTensor.\n";

  // NOTE(chengcheng): MUST store this tensor to MultiClientSessionContext for graph runtime bind.
  const std::string graph_name = *JUST(JUST(GlobalJobBuildAndInferCtxMgr())->GetCurrentJobName());
  const std::string lbn = GenLogicalBlobName(new_op_name, "out");
  Singleton<MultiClientSessionContext>::Get()->StoreFreeEagerTensorWithNameByGraphName(
      graph_name, input_tensor, new_op_name);
  // NOTE(chengcheng): MUST record this eager_tensor name as new variable output lbn.
  // NOTE(chengcheng): in GradAcc FreeEagerTensor need insert repeat op, but there is no need to
  //  create a new tensor for repeat op out. We just set repeat lbn as this free eager tensor's lbn.
  TensorNameScope::Global()->Record(input_tensor,
                                    *JUST(GradAccTryInsertRepeatAfterFreeVar(op_conf)));
  return Maybe<void>::Ok();
}

}  // namespace

Maybe<void> LazyInterpreter::ApplyImpl(const FeedInputOpExpr& op_expr, const TensorTuple& inputs,
                                       TensorTuple* outputs, const OpExprInterpContext& ctx) const {
  // NOTE(chengcheng): inputs[0] is the EagerTensor
  CHECK_EQ_OR_RETURN(inputs.size(), 1);
  CHECK_EQ_OR_RETURN(op_expr.input_size(), 1);
  const std::shared_ptr<Tensor>& input_tensor = inputs.at(0);
  CHECK_OR_RETURN(input_tensor->is_eager());

  std::shared_ptr<Scope> scope = JUST(NewScopeWithParallelDescByTensor(input_tensor));

  OperatorConf op_conf;
  op_conf.set_name(op_expr.op_name());  // construct by python nn.Graph
  op_conf.set_scope_symbol_id(JUST(scope->symbol_id()));
  op_conf.set_device_tag(JUST(GetDeviceTagOfTensor(input_tensor)));
  // NOTE(chengcheng):
  //   We contruct InputOpConf instead of FeedInputOpConf because FeedInputOpExpr JUST for getting
  //   input EagerTensor.
  InputOpConf* input_conf = op_conf.mutable_input_conf();
  input_conf->set_out("out");
  InterfaceBlobConf* blob_conf = input_conf->mutable_blob_conf();

  input_tensor->shape()->ToProto(blob_conf->mutable_shape());
  blob_conf->set_data_type(input_tensor->dtype()->data_type());
  // NOTE(chengcheng): is_dynamic true has conflict in consistent lazy job even if world size 1.
  //     this flag will be removed in the future.
  // blob_conf->set_is_dynamic(GetIsDynamicOfTensor(input_tensor));
  blob_conf->set_is_dynamic(false);
  JUST(GenNdSbpByTensor(blob_conf->mutable_nd_sbp(), input_tensor));

  auto infer_ctx = JUST(GetCurInferCtx());
  VLOG(2) << "Lazy nn.Graph name " << infer_ctx->job().job_conf().job_name()
          << " try to add op: \n: " << op_conf.DebugString() << std::endl;
  OpAttribute op_attr = *JUST(infer_ctx->AddAndInferConsistentOp(op_conf));
  VLOG(2) << "Lazy nn.Graph name " << infer_ctx->job().job_conf().job_name() << " add op : \n"
          << op_conf.name() << std::endl;
  VLOG(3) << "Lazy nn.Graph name " << infer_ctx->job().job_conf().job_name()
          << " infer and and op attr : \n"
          << op_attr.DebugString() << std::endl;

  int64_t parallel_desc_sym_id = JUST(scope->GetParallelDescSymbolId(op_conf));
  auto blob_parallel_desc = JUST(GetSymbol<ParallelDesc>(parallel_desc_sym_id));

  // Check outputs num and setup output tensor properties.
  CHECK_EQ_OR_RETURN(outputs->size(), 1);
  CHECK_EQ_OR_RETURN(op_expr.output_size(), 1);
  CHECK_OR_RETURN(!(*outputs)[0]);
  const std::string obn = "out";  // NOTE(chengcheng): obn is NOT op_expr.indexed_obns
  auto origin_input = JUST(BuildTensor(op_attr, obn, blob_parallel_desc, /* is_lazy= */ true,
                                       /* is_local= */ input_tensor->is_local()));
  TensorNameScope::Global()->Record(origin_input, GenLogicalBlobName(op_conf.name(), obn));

  // NOTE(chengcheng): Do GradAcc pass when add input op.
  (*outputs)[0] = JUST(GradAccTryInsertUnpackAfterInput(op_conf, blob_parallel_desc, origin_input));
  return Maybe<void>::Ok();
}

Maybe<void> LazyInterpreter::ApplyImpl(const FeedVariableOpExpr& op_expr, const TensorTuple& inputs,
                                       TensorTuple* outputs, const OpExprInterpContext& ctx) const {
  // NOTE(chengcheng): inputs[0] is the EagerTensor
  CHECK_EQ_OR_RETURN(inputs.size(), 1);
  CHECK_EQ_OR_RETURN(op_expr.input_size(), 1);
  const std::shared_ptr<Tensor>& input_tensor = inputs.at(0);
  CHECK_OR_RETURN(input_tensor->is_eager());

  std::shared_ptr<Scope> scope = JUST(NewScopeWithParallelDescByTensor(input_tensor));

  OperatorConf op_conf;
  op_conf.set_name(op_expr.op_name());  // construct by python nn.Graph
  op_conf.set_scope_symbol_id(JUST(scope->symbol_id()));
  op_conf.set_device_tag(JUST(GetDeviceTagOfTensor(input_tensor)));
  // NOTE(chengcheng):
  //   We contruct VariableOpConf instead of FeedVariableOpConf because FeedVariableOpExpr JUST
  //   for getting input EagerTensor.
  VariableOpConf* var_conf = op_conf.mutable_variable_conf();
  var_conf->set_out("out");
  input_tensor->shape()->ToProto(var_conf->mutable_shape());
  var_conf->set_data_type(input_tensor->dtype()->data_type());
  // NOTE(chengcheng): VariableOpConf initializer_conf is useless because variable is inited
  //   by EagerTensor.
  var_conf->mutable_initializer()->mutable_empty_conf();
  JUST(GenVariableOpConfNdSbpStringByTensor(var_conf, input_tensor));
  if (!input_tensor->requires_grad()) { var_conf->set_trainable(false); }
  if (input_tensor->requires_grad()) {
    double l2 = JUST(ctx.attrs.GetAttr<double>("l2"));
    if (unlikely(l2 != 0.0)) { var_conf->mutable_regularizer()->mutable_l1_l2_conf()->set_l2(l2); }
  }

  auto infer_ctx = JUST(GetCurInferCtx());
  VLOG(2) << "Lazy nn.Graph name " << infer_ctx->job().job_conf().job_name()
          << " try to add op: \n: " << op_conf.DebugString() << std::endl;
  OpAttribute op_attr = *JUST(infer_ctx->AddAndInferConsistentOp(op_conf));
  VLOG(2) << "Lazy nn.Graph name " << infer_ctx->job().job_conf().job_name() << " add op : \n"
          << op_conf.name() << std::endl;
  VLOG(3) << "Lazy nn.Graph name " << infer_ctx->job().job_conf().job_name()
          << " infer and and op attr : \n"
          << op_attr.DebugString() << std::endl;

  int64_t parallel_desc_sym_id = JUST(scope->GetParallelDescSymbolId(op_conf));
  auto blob_parallel_desc = JUST(GetSymbol<ParallelDesc>(parallel_desc_sym_id));

  // Check outputs num and setup output tensor properties.
  CHECK_EQ_OR_RETURN(outputs->size(), 1);
  CHECK_EQ_OR_RETURN(op_expr.output_size(), 1);
  CHECK_OR_RETURN(!(*outputs)[0]);

  const std::string obn = "out";  // NOTE(chengcheng): obn is NOT op_expr.indexed_obns
  auto origin_var = JUST(BuildTensor(op_attr, obn, blob_parallel_desc, /* is_lazy= */ true,
                                     /* is_local */ input_tensor->is_local()));

  // NOTE(chengcheng): Record variable op output LazyTenosr
  TensorNameScope::Global()->Record(origin_var, GenLogicalBlobName(op_conf.name(), obn));
  // NOTE(chengcheng): Record EagerTensor as variable tensor name
  TensorNameScope::Global()->Record(input_tensor, GenLogicalBlobName(op_conf.name(), obn));

  (*outputs)[0] = JUST(GradAccTryInsertRepeatAfterVar(op_conf, blob_parallel_desc, origin_var));
  return Maybe<void>::Ok();
}

Maybe<void> LazyInterpreter::ApplyImpl(const FetchOutputOpExpr& op_expr, const TensorTuple& inputs,
                                       TensorTuple* outputs, const OpExprInterpContext& ctx) const {
  // NOTE(chengcheng): inputs[0] is the LazyTensor
  CHECK_EQ_OR_RETURN(inputs.size(), 1);
  CHECK_EQ_OR_RETURN(op_expr.input_size(), 1);
  const std::shared_ptr<Tensor>& input_tensor = inputs.at(0);
  std::string input_lbn = TensorNameScope::Global()->Lookup(input_tensor);
  // Lazy tensor must has lbn.
  // Eager tensor may has lbn if it has already been treated as an output of a variable op
  // or an output of an inplace op.
  if (input_lbn.empty()) {
    CHECK_OR_RETURN(input_tensor->is_eager());
    // This output tensor is a new free eager tensor, so treat it as a new variable op output.
    JUST(AddFreeEagerTensorToVariableOp(input_tensor));
    input_lbn = TensorNameScope::Global()->Lookup(input_tensor);
  }
  CHECK_OR_RETURN(!input_lbn.empty());  // lbn must exist.

  std::shared_ptr<Scope> scope = JUST(NewScopeWithParallelDescByTensor(input_tensor));

  std::shared_ptr<Tensor> output_tensor =
      JUST(GradAccTryInsertPackBeforeOutput(scope, input_lbn, op_expr.op_name(), input_tensor));

  const std::string output_lbn = TensorNameScope::Global()->Lookup(output_tensor);

  OperatorConf op_conf;
  op_conf.set_name(op_expr.op_name());  // construct by python nn.Graph
  op_conf.set_scope_symbol_id(JUST(scope->symbol_id()));
  op_conf.set_device_tag(JUST(GetDeviceTagOfTensor(input_tensor)));
  // NOTE(chengcheng):
  //   We contruct OutputOpConf instead of FetchOutputOpConf because FetchOutputOpExpr JUST
  //   for get nn.Graph output LazyTensor.
  OutputOpConf* output_conf = op_conf.mutable_output_conf();
  output_conf->set_in(output_lbn);
  output_conf->set_out("out");
  InterfaceBlobConf* blob_conf = output_conf->mutable_blob_conf();
  output_tensor->shape()->ToProto(blob_conf->mutable_shape());
  blob_conf->set_data_type(output_tensor->dtype()->data_type());
  // NOTE(chengcheng): is_dynamic true has conflict in consistent lazy job even if world size 1.
  //     this flag will be removed in the future.
  // blob_conf->set_is_dynamic(GetIsDynamicOfTensor(output_tensor));
  blob_conf->set_is_dynamic(false);
  JUST(GenNdSbpByTensor(blob_conf->mutable_nd_sbp(), output_tensor));

  auto infer_ctx = JUST(GetCurInferCtx());
  VLOG(2) << "Lazy nn.Graph name " << infer_ctx->job().job_conf().job_name() << " try to add op: \n"
          << op_conf.DebugString() << std::endl;
  OpAttribute op_attr = *JUST(infer_ctx->AddAndInferConsistentOp(op_conf));
  VLOG(2) << "Lazy nn.Graph name " << infer_ctx->job().job_conf().job_name() << " add op : \n"
          << op_conf.name() << std::endl;
  VLOG(3) << "Lazy nn.Graph name " << infer_ctx->job().job_conf().job_name()
          << " infer and and op attr : \n"
          << op_attr.DebugString() << std::endl;

  int64_t parallel_desc_sym_id = JUST(scope->GetParallelDescSymbolId(op_conf));
  auto blob_parallel_desc = JUST(GetSymbol<ParallelDesc>(parallel_desc_sym_id));

  // Check outputs num and setup output tensor properties.
  CHECK_EQ_OR_RETURN(outputs->size(), 1);
  CHECK_EQ_OR_RETURN(op_expr.output_size(), 1);
  CHECK_OR_RETURN(!(*outputs)[0]);
  const std::string obn = "out";  // NOTE(chengcheng): obn is NOT op_expr.indexed_obns
  (*outputs)[0] = JUST(BuildTensor(op_attr, obn, blob_parallel_desc, /* is_lazy= */ false,
                                   /* is_local= */ output_tensor->is_local()));
  return Maybe<void>::Ok();
}

Maybe<void> LazyInterpreter::ApplyImpl(const ImageDecoderRandomCropResizeOpExpr& op_expr,
                                       const TensorTuple& inputs, TensorTuple* outputs,
                                       const OpExprInterpContext& ctx) const {
  CHECK_EQ_OR_RETURN(inputs.size(), 1);
  CHECK_EQ_OR_RETURN(op_expr.input_size(), 1);
  const std::shared_ptr<Tensor>& input_tensor = inputs.at(0);
  const std::string& input_lbn = TensorNameScope::Global()->Lookup(input_tensor);
  CHECK_OR_RETURN(!input_lbn.empty());  // lbn must exist.

  auto op_conf = JUST(OpInterpUtil::GenBuiltinOpConf(op_expr, ctx.attrs));
  std::string device_tag;
  if (IsCpuOnly(*op_conf)) {
    device_tag = "cpu";
  } else {
    device_tag = "cuda";
  }

  ParallelConf parallel_conf = JUST(GetParallelDescOfTensor(input_tensor))->parallel_conf();
  parallel_conf.set_device_tag(device_tag);  // NOTE(chengcheng): only support gpu decode.
  const auto& scope = JUST(NewScopeWithParallelConfAndCurScope(parallel_conf));

  op_conf->set_scope_symbol_id(JUST(scope->symbol_id()));
  op_conf->set_device_tag(device_tag);

  // NOTE(chengcheng): replace right input_lbn and obn
  ReplaceInputLbnInOpCustomizedConf(op_conf.get(), /* ibn */ "in", input_lbn);
  op_conf->mutable_image_decoder_random_crop_resize_conf()->set_out("out");

  auto infer_ctx = JUST(GetCurInferCtx());
  // NOTE(chengcheng): MUST reset unique op name before InferCtx::AddOp
  const std::string new_op_name = *JUST(infer_ctx->NewUniqueOpNameByFunctionalOpConf(*op_conf));
  op_conf->set_name(new_op_name);
  VLOG(2) << "Lazy nn.Graph name " << infer_ctx->job().job_conf().job_name() << " try to add op: \n"
          << op_conf->DebugString() << std::endl;
  OpAttribute op_attr = *JUST(infer_ctx->AddAndInferConsistentOp(*op_conf));
  VLOG(2) << "Lazy nn.Graph name " << infer_ctx->job().job_conf().job_name() << " add op : \n"
          << op_conf->name() << std::endl;
  VLOG(3) << "Lazy nn.Graph name " << infer_ctx->job().job_conf().job_name()
          << " infer and and op attr : \n"
          << op_attr.DebugString() << std::endl;

  int64_t parallel_desc_sym_id = JUST(scope->GetParallelDescSymbolId(*op_conf));
  auto blob_parallel_desc = JUST(GetSymbol<ParallelDesc>(parallel_desc_sym_id));

  // Check outputs num and setup output tensor properties.
  CHECK_EQ_OR_RETURN(outputs->size(), 1);
  CHECK_EQ_OR_RETURN(op_expr.output_size(), 1);
  CHECK_OR_RETURN(!(*outputs)[0]);
  const std::string obn = "out";  // NOTE(chengcheng): obn is NOT op_expr.indexed_obns
  (*outputs)[0] = JUST(BuildTensor(op_attr, obn, blob_parallel_desc, /* is_lazy= */ true,
                                   /* is_local= */ input_tensor->is_local()));
  TensorNameScope::Global()->Record((*outputs)[0], GenLogicalBlobName(new_op_name, obn));
  return Maybe<void>::Ok();
}

namespace {

Maybe<void> LazyInterpreterApplyImplForSourceUserOpExpr(const UserOpExpr& op_expr,
                                                        TensorTuple* outputs,
                                                        const OpExprInterpContext& ctx) {
  NonRecursiveMetaInfoConsistencyCheckScope non_scope;
  bool is_local;
  std::shared_ptr<const ParallelDesc> parallel_desc;
  if (ctx.parallel_desc.has_value()) {
    // NOTE(chengcheng): consistent
    CHECK_OR_RETURN(!ctx.device.has_value());
    const auto& parallel_desc_sym = JUST(ctx.parallel_desc);
    parallel_desc = parallel_desc_sym.shared_from_symbol();
    JUST(MetaInfoConsistencyCheck(parallel_desc_sym, ctx.nd_sbp, 1, /* force_check */ false));
    is_local = false;
  } else {
    // NOTE(chengcheng): local
    CHECK_OR_RETURN(!ctx.nd_sbp.has_value());
    if (ctx.device.has_value()) {
      const auto& device = JUST(ctx.device);
      const auto& placement = JUST(Placement4Device(device));
      parallel_desc = placement.shared_from_symbol();
    } else {
      // NOTE(chengcheng): if functor NOT set device, using cpu device default.
      const auto& device = JUST(Device::New("cpu"));
      const auto& placement = JUST(Placement4Device(device));
      parallel_desc = placement.shared_from_symbol();
    }
    is_local = true;
  }
  const auto& parallel_conf = parallel_desc->parallel_conf();
  const auto& scope = JUST(NewScopeWithParallelConfAndCurScope(parallel_conf));
  auto op_conf = JUST(OpInterpUtil::GenBuiltinOpConf(op_expr, ctx.attrs));
  op_conf->set_scope_symbol_id(JUST(scope->symbol_id()));
  op_conf->set_device_tag(parallel_conf.device_tag());

  auto infer_ctx = JUST(GetCurInferCtx());
  // NOTE(chengcheng): MUST reset unique op name before InferCtx::AddOp
  const std::string new_op_name = *JUST(infer_ctx->NewUniqueOpNameByFunctionalOpConf(*op_conf));

  // NOTE(chengcheng): for UserOp, NOT only reset op_name, but also the output values.
  op_conf->set_name(new_op_name);
  for (auto& pair : *(op_conf->mutable_user_conf()->mutable_output())) {
    auto& list_s = pair.second;
    for (int i = 0; i < list_s.s_size(); ++i) {
      std::string old_lbn = list_s.s(i);
      LogicalBlobId old_lbi = GenLogicalBlobId(old_lbn);
      // NOTE(chengcheng): MUST change the old_lbn to new op name.
      std::string new_lbn = GenLogicalBlobName(new_op_name, old_lbi.blob_name());
      list_s.set_s(i, new_lbn);
    }
  }

  JUST(GradAccTryInsertRepeatTickBeforeSource(op_conf));

  VLOG(2) << "Lazy nn.Graph name " << infer_ctx->job().job_conf().job_name() << " try to add op: \n"
          << op_conf->DebugString() << std::endl;
  OpAttribute op_attr = *JUST(infer_ctx->AddAndInferConsistentOp(*op_conf));
  VLOG(2) << "Lazy nn.Graph name " << infer_ctx->job().job_conf().job_name() << " add op : \n"
          << op_conf->name() << std::endl;
  VLOG(3) << "Lazy nn.Graph name " << infer_ctx->job().job_conf().job_name()
          << " infer and and op attr : \n"
          << op_attr.DebugString() << std::endl;

  int64_t parallel_desc_sym_id = JUST(scope->GetParallelDescSymbolId(*op_conf));
  auto blob_parallel_desc = JUST(GetSymbol<ParallelDesc>(parallel_desc_sym_id));

  // Check outputs num and setup output tensor properties.
  CHECK_EQ_OR_RETURN(outputs->size(), op_expr.output_size());
  for (int i = 0; i < op_expr.output_size(); ++i) {
    CHECK_OR_RETURN(!(*outputs)[i]);
    const std::string& obn = op_expr.indexed_obns().at(i);
    (*outputs)[i] =
        JUST(BuildTensor(op_attr, obn, blob_parallel_desc, /* is_lazy= */ true, is_local));
    TensorNameScope::Global()->Record((*outputs)[i], GenLogicalBlobName(new_op_name, obn));
  }
  return Maybe<void>::Ok();
}

Maybe<void> LazyInterpreterApplyImplForCopyUserOpExpr(const UserOpExpr& op_expr,
                                                      const TensorTuple& inputs,
                                                      TensorTuple* outputs,
                                                      const OpExprInterpContext& ctx) {
  CHECK_OR_RETURN(op_expr.op_type_name() == "copy");
  CHECK_EQ_OR_RETURN(inputs.size(), 1);
  CHECK_EQ_OR_RETURN(op_expr.input_size(), 1);
  const std::shared_ptr<Tensor>& input_tensor = inputs.at(0);
  std::string input_lbn = TensorNameScope::Global()->Lookup(input_tensor);
  if (input_lbn.empty()) {
    JUST(AddFreeEagerTensorToVariableOp(input_tensor));
    input_lbn = TensorNameScope::Global()->Lookup(input_tensor);
  }
  CHECK_OR_RETURN(!input_lbn.empty());  // lbn must exist.
  std::string device_type = JUST(ctx.attrs.GetAttr<std::string>("device_type"));
  int64_t device_id = JUST(ctx.attrs.GetAttr<int64_t>("device_id"));

  CHECK_EQ_OR_RETURN(outputs->size(), 1);
  CHECK_EQ_OR_RETURN(op_expr.output_size(), 1);
  if (input_tensor->is_local()) {
    (*outputs)[0] = JUST(MirroredTensor::MakeTensor(
        input_tensor->shape(), JUST(input_tensor->stride()), input_tensor->dtype()->data_type(),
        JUST(Device::New(device_type, device_id)),
        /* is_lazy= */ true,
        /*requires_grad=*/false, /*is_leaf=*/true));
  } else {
    ParallelConf parallel_conf = JUST(input_tensor->parallel_desc())->parallel_conf();
    parallel_conf.set_device_tag(device_type);
    ParallelDesc parallel_desc(parallel_conf);
    (*outputs)[0] =
        JUST(ConsistentTensor::MakeTensor(input_tensor->shape(), input_tensor->dtype()->data_type(),
                                          JUST(input_tensor->nd_sbp()), SymbolOf(parallel_desc),
                                          /* is_lazy= */ true,
                                          /*requires_grad=*/false, /*is_leaf=*/true));
  }
  // NOTE(chengcheng): output tensor lbn is SAME with input tensor.
  TensorNameScope::Global()->Record(outputs->at(0), input_lbn);
  return Maybe<void>::Ok();
}

}  // namespace

Maybe<void> LazyInterpreter::ApplyImpl(const UserOpExpr& op_expr, const TensorTuple& inputs,
                                       TensorTuple* outputs, const OpExprInterpContext& ctx) const {
  CHECK_EQ_OR_RETURN(inputs.size(), op_expr.input_size());

  // NOTE(chengcheng): Handle special UserOp such as:
  //     1. [Source UserOp] : OFRecordReader, CoinFlip
  //     2. [Change Placement/ParallelDesc UserOp] : to(copy)/to_global/parallel_cast
  //     3. [Multi-Inputs & Different ParallelDesc for each input UserOp] : like there are 2 inputs,
  //             one from CPU and the other from GPU.
  //     ..., etc.
  //
  //     Need add if for each special UserOp for infer:
  //     1. op_conf: device_tag,
  //     2. output tensor: is_local,
  //     3. op_parallel_conf for build new scope with parallel_desc
  //     4. output blob (different with tensor) -> parallel_conf
  //     5. need add to JobBuildAndInferCtx (like copy will NOT need)
  if (inputs.size() == 0) {
    // NOTE(chengcheng): handle for source UserOp like OFRecordReader, CoinFlip
    return LazyInterpreterApplyImplForSourceUserOpExpr(op_expr, outputs, ctx);
  }
  if (op_expr.op_type_name() == "copy") {
    // NOTE(chengcheng): handle for copy UserOp which will NOT add op to job.
    return LazyInterpreterApplyImplForCopyUserOpExpr(op_expr, inputs, outputs, ctx);
  }

  // NOTE(chengcheng):
  //   Normal UserOp inputs size >= 1 for infer parallel_desc.
  CHECK_GE_OR_RETURN(inputs.size(), 1);
  auto op_conf = JUST(OpInterpUtil::GenBuiltinOpConf(op_expr, ctx.attrs));
  std::shared_ptr<Scope> scope = JUST(NewScopeWithParallelDescByTensor(JUST(VectorAt(inputs, 0))));
  op_conf->set_scope_symbol_id(JUST(scope->symbol_id()));
  const std::string device_tag = JUST(GetDeviceTagOfTensor(JUST(VectorAt(inputs, 0))));
  const bool is_local = inputs.at(0)->is_local();
  const std::shared_ptr<const ParallelDesc> parallel_desc =
      JUST(GetParallelDescOfTensor(inputs.at(0)));

  op_conf->set_device_tag(device_tag);
  auto infer_ctx = JUST(GetCurInferCtx());
  // NOTE(chengcheng): MUST reset unique op name before InferCtx::AddOp
  const std::string new_op_name = *JUST(infer_ctx->NewUniqueOpNameByFunctionalOpConf(*op_conf));
  const std::string graph_name = infer_ctx->job().job_conf().job_name();

  for (int i = 0; i < inputs.size(); ++i) {
    const auto& input_tensor = inputs.at(i);
    CHECK_EQ_OR_RETURN(is_local, input_tensor->is_local());
    if (is_local) {
      CHECK_OR_RETURN(device_tag == JUST(GetDeviceTagOfTensor(input_tensor)))
          << Error::RuntimeError() << "Lazy nn.Graph name: " << graph_name
          << " encountered ERROR in module/op_name: " << new_op_name
          << ". Expected all tensors to be on the same device, but found at least two devices, "
          << JUST(JUST(VectorAt(inputs, 0))->device())->ToString() << " (positional 0) and "
          << JUST(JUST(VectorAt(inputs, i))->device())->ToString() << " (positional " << i
          << ")! Please use tensor.to() to synchronize all the input with the same device.";
    } else {
      // TODO: Print out all the placement
      CHECK_OR_RETURN(parallel_desc->Equals(*JUST(GetParallelDescOfTensor(input_tensor))))
          << Error::RuntimeError() << "Lazy nn.Graph name: " << graph_name
          << " encountered ERROR in module/op_name: " << new_op_name
          << ". Expected all tensors to be on the same placement, but found at least two "
             "placements, "
          << *JUST(PlacementToString(JUST(JUST(VectorAt(inputs, 0))->parallel_desc())))
          << " (positional 0) and "
          << *JUST(PlacementToString(JUST(JUST(VectorAt(inputs, i))->parallel_desc())))
          << " (positional " << i
          << ")! Please use tensor.to_global() to synchronize all the input with the same "
             "placement.";
    }
    const std::string& ibn = op_expr.indexed_ibns().at(i);
    std::string lbn = TensorNameScope::Global()->Lookup(input_tensor);
    if (lbn.empty()) {
      JUST(AddFreeEagerTensorToVariableOp(input_tensor));
      lbn = TensorNameScope::Global()->Lookup(input_tensor);
    }
    CHECK_OR_RETURN(!lbn.empty());  // NOTE(chengcheng): lbn must not empty now.
    ReplaceInputLbnInOpCustomizedConf(op_conf.get(), ibn, lbn);
  }

  // NOTE(chengcheng): for UserOp, NOT only reset op_name, but also the output values.
  op_conf->set_name(new_op_name);
  for (auto& pair : *(op_conf->mutable_user_conf()->mutable_output())) {
    auto& list_s = pair.second;
    for (int i = 0; i < list_s.s_size(); ++i) {
      std::string old_lbn = list_s.s(i);
      LogicalBlobId old_lbi = GenLogicalBlobId(old_lbn);
      // NOTE(chengcheng): MUST change the old_lbn to new op name.
      std::string new_lbn = GenLogicalBlobName(new_op_name, old_lbi.blob_name());
      list_s.set_s(i, new_lbn);
    }
  }

  // Check outputs num and setup output tensor properties.
  CHECK_EQ_OR_RETURN(outputs->size(), op_expr.output_size());

  // Disable boxing if the computation is inplace.
  for (int i = 0; i < op_expr.output_size(); ++i) {
    const auto& output = outputs->at(i);
    if (output) {
      const std::string& lbn = TensorNameScope::Global()->Lookup(output);
      CHECK_OR_RETURN(!lbn.empty()) << "The output which index is " << i
                                    << " has no tensor name, please check whether the inplaced "
                                       "output is also an input of the operation "
                                    << new_op_name;
      JUST(infer_ctx->DisableBoxing(lbn));
    }
  }
  VLOG(2) << "Lazy nn.Graph name " << graph_name << " try to add op: \n"
          << op_conf->DebugString() << std::endl;
  OpAttribute op_attr = *JUST(infer_ctx->AddAndInferConsistentOp(*op_conf));
  VLOG(2) << "Lazy nn.Graph name " << graph_name << " add op : \n" << op_conf->name() << std::endl;
  VLOG(3) << "Lazy nn.Graph name " << graph_name << " infer and and op attr : \n"
          << op_attr.DebugString() << std::endl;

  int64_t parallel_desc_sym_id = JUST(scope->GetParallelDescSymbolId(*op_conf));
  auto blob_parallel_desc = JUST(GetSymbol<ParallelDesc>(parallel_desc_sym_id));
  for (int i = 0; i < op_expr.output_size(); ++i) {
    const std::string& obn = op_expr.indexed_obns().at(i);
    if (!(*outputs)[i]) {
      (*outputs)[i] =
          JUST(BuildTensor(op_attr, obn, blob_parallel_desc, /* is_lazy= */ true, is_local));
    } else {
      VLOG(2) << "Lazy nn.Graph name " << graph_name << " op name " << new_op_name
              << " run with inplace.";
      const std::shared_ptr<Tensor>& inplace_out = (*outputs)[i];
      JUST(CheckTensorMatchAttr(inplace_out, op_attr, obn, blob_parallel_desc, is_local));
    }
    TensorNameScope::Global()->Record((*outputs)[i], GenLogicalBlobName(new_op_name, obn));
  }
  return Maybe<void>::Ok();
}

Maybe<void> LazyInterpreter::ApplyImpl(const FunctionOpExpr& op_expr, const TensorTuple& inputs,
                                       TensorTuple* outputs, const OpExprInterpContext& ctx) const {
  // TODO(hjchen2)
  OF_UNIMPLEMENTED() << "The type " << op_expr.op_type_name()
                     << " has not been supported in LazyInterpreter::Apply.";
  return Maybe<void>::Ok();
}

Maybe<void> LazyInterpreter::ApplyImpl(const ConsistentToConsistentOpExpr& op_expr,
                                       const TensorTuple& inputs, TensorTuple* outputs,
                                       const OpExprInterpContext& ctx) const {
  CHECK_EQ_OR_RETURN(op_expr.input_size(), 1);
  CHECK_EQ_OR_RETURN(inputs.size(), 1);
  const auto& input_tensor = inputs[0];
  CHECK_OR_RETURN(input_tensor->is_consistent());

  CHECK_OR_RETURN(ctx.parallel_desc.has_value());
  const auto& parallel_desc_sym = JUST(ctx.parallel_desc);
  CHECK_OR_RETURN(ctx.nd_sbp.has_value());
  const auto& sbp_sym = JUST(ctx.nd_sbp);

  std::string input_lbn = TensorNameScope::Global()->Lookup(input_tensor);
  if (input_lbn.empty()) {
    JUST(AddFreeEagerTensorToVariableOp(input_tensor));
    input_lbn = TensorNameScope::Global()->Lookup(input_tensor);
    CHECK_OR_RETURN(!input_lbn.empty());
  }

  std::shared_ptr<Tensor> input_proxy;
  if (!JUST(GetParallelDescOfTensor(input_tensor))
           ->Equals(*parallel_desc_sym.shared_from_symbol())) {
    // NOTE(zwx): The input tensor's parallel_desc is not equal to that of op's,
    // create a proxy input with the parallel_desc that is the same as op's
    input_proxy =
        JUST(ConsistentTensor::MakeTensor(input_tensor->shape(), input_tensor->dtype()->data_type(),
                                          JUST(input_tensor->nd_sbp()), parallel_desc_sym,
                                          /* is_lazy= */ true,
                                          /*requires_grad=*/false, /*is_leaf=*/true));
    TensorNameScope::Global()->Record(input_proxy, input_lbn);
  }

  CHECK_EQ_OR_RETURN(op_expr.output_size(), 1);
  CHECK_EQ_OR_RETURN(outputs->size(), 1);
  CHECK_OR_RETURN(!(*outputs)[0]);

  if (!op_expr.grad_nd_sbp().has_value() && sbp_sym == JUST(input_tensor->nd_sbp())) {
    // NOTE(chengcheng):  if to_global ONLY change placement (nd_sbp and grad_nd_sbp is same),
    //    there is no need to build hierarchical_parallel_cast op.
    if (input_proxy) {
      (*outputs)[0] = input_proxy;
    } else {
      (*outputs)[0] = input_tensor;
    }
    return Maybe<void>::Ok();
  }

  // build parallel cast op expr
  std::shared_ptr<std::vector<std::string>> sbp_list_ptr = JUST(GetNdSbpStrList(sbp_sym));
  std::string grad_mode;
  std::vector<std::string> grad_sbp_str_list;
  if (op_expr.grad_nd_sbp().has_value()) {
    grad_mode = "manual";
    grad_sbp_str_list = *JUST(GetNdSbpStrList(JUST(op_expr.grad_nd_sbp())));
  } else {
    grad_mode = "identity";
  }
  std::shared_ptr<UserOpExpr> parallel_cast_op_expr =
      JUST(OpBuilder("hierarchical_parallel_cast", "trivial_op_name")
               .Input("in")
               .Output("out")
               .Attr<std::vector<std::string>>("nd_sbp", *sbp_list_ptr)
               .Attr<std::string>("grad_mode", grad_mode)
               .Attr<std::vector<std::string>>("grad_nd_sbp", grad_sbp_str_list)
               .Build());

  if (input_proxy) {
    (*outputs)[0] =
        JUST(OpInterpUtil::Dispatch<one::Tensor>(*parallel_cast_op_expr, {input_proxy}));
  } else {
    (*outputs)[0] =
        JUST(OpInterpUtil::Dispatch<one::Tensor>(*parallel_cast_op_expr, {input_tensor}));
  }

  return Maybe<void>::Ok();
}

}  // namespace one
}  // namespace oneflow
