Class UserOpExprInferContext

oneflow/core/framework/op_expr.cpp:169–401

namespace {

class UserOpExprInferContext : public user_op::InferContext {
 public:
  UserOpExprInferContext(const UserOpExpr* user_op_expr, const AttrMap& attrs,
                         const std::string& device_tag,
                         const std::function<const TensorMeta*(int32_t)>& TensorMeta4InputIndex,
                         const std::function<TensorMeta*(int32_t)>& TensorMeta4OutputIndex)
      : user_op_expr_(user_op_expr),
        composed_attrs_(attrs, user_op_expr->base_attrs()),
        tensor_meta4input_index_(TensorMeta4InputIndex),
        tensor_meta4output_index_(TensorMeta4OutputIndex) {
    loc_ = DispatchFrame::get_str();
  }
  virtual ~UserOpExprInferContext() override = default;

  const std::vector<std::pair<std::string, int32_t>>& inputs() const override {
    return user_op_expr_->indexed_input_pairs();
  }

  const std::vector<std::pair<std::string, int32_t>>& outputs() const override {
    return user_op_expr_->indexed_output_pairs();
  }

  const user_op::TensorDesc& InputTensorDesc(const std::string& arg_name,
                                             int32_t index) const override {
    return *TensorDesc4ArgNameAndIndex(arg_name, index);
  }
  const user_op::TensorDesc& OutputTensorDesc(const std::string& arg_name,
                                              int32_t index) const override {
    return *TensorDesc4ArgNameAndIndex(arg_name, index);
  }
  user_op::TensorDesc* MutOutputTensorDesc(const std::string& name, int32_t index) override {
    return MutTensorDesc4ArgNameAndIndex(name, index);
  }

  const user_op::TensorDesc* TensorDesc4ArgNameAndIndex(const std::string& name,
                                                        int32_t index) const {
    {
      const auto& arg_tuple = *user_op_expr_->output_arg_tuple();
      int32_t tuple_index = arg_tuple.TensorTupleIndex4ArgNameAndIndex(name, index);
      if (tuple_index >= 0) { return tensor_meta4output_index_(tuple_index); }
    }
    {
      const auto& arg_tuple = *user_op_expr_->input_arg_tuple();
      int32_t tuple_index = arg_tuple.TensorTupleIndex4ArgNameAndIndex(name, index);
      if (tuple_index >= 0) { return tensor_meta4input_index_(tuple_index); }
    }
    return nullptr;
  }

  user_op::TensorDesc* MutTensorDesc4ArgNameAndIndex(const std::string& name, int32_t index) {
    {
      const auto& arg_tuple = *user_op_expr_->output_arg_tuple();
      int32_t tuple_index = arg_tuple.TensorTupleIndex4ArgNameAndIndex(name, index);
      if (tuple_index >= 0) {
        TensorMeta* tensor_meta_ptr = tensor_meta4output_index_(tuple_index);
        CHECK_NOTNULL(dynamic_cast<MutTensorMeta*>(tensor_meta_ptr));
        return tensor_meta_ptr;
      }