op_proto_maker.h 2.4 KB
Newer Older
1
/* Copyright (c) 2016 PaddlePaddle Authors. All Rights Reserved.
2 3 4 5 6 7 8 9 10 11 12 13 14 15
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. */

#pragma once

16
#include <string>
R
reyoung 已提交
17
#include "glog/logging.h"
Y
Yi Wang 已提交
18 19
#include "paddle/fluid/framework/attribute.h"
#include "paddle/fluid/framework/framework.pb.h"
20 21 22 23 24 25
namespace paddle {
namespace framework {

// this class not only make proto but also init attribute checkers.
class OpProtoAndCheckerMaker {
 public:
Y
Yu Yang 已提交
26
  virtual void Make() = 0;
27

28
  virtual ~OpProtoAndCheckerMaker() {
R
reyoung 已提交
29
    CHECK(validated_) << "should call Validate after build";
30 31
  }

R
reyoung 已提交
32
  void SetProto(proto::OpProto *proto) { proto_ = proto; }
Y
Yu Yang 已提交
33

R
reyoung 已提交
34
  void SetChecker(OpAttrChecker *attr_checker) { op_checker_ = attr_checker; }
Y
Yu Yang 已提交
35

36 37 38 39
  void Validate();

 protected:
  struct VariableBuilder {
R
reyoung 已提交
40
    proto::OpProto::Var *var_;
41

R
reyoung 已提交
42
    VariableBuilder &AsDuplicable() {
43 44 45 46
      var_->set_duplicable(true);
      return *this;
    }

R
reyoung 已提交
47
    VariableBuilder &AsIntermediate() {
48 49 50
      var_->set_intermediate(true);
      return *this;
    }
Y
Yang Yang 已提交
51

R
reyoung 已提交
52
    VariableBuilder &AsDispensable() {
Y
Yang Yang 已提交
53 54 55
      var_->set_dispensable(true);
      return *this;
    }
56 57
  };

R
reyoung 已提交
58
  VariableBuilder AddInput(const std::string &name, const std::string &comment);
59

R
reyoung 已提交
60 61
  VariableBuilder AddOutput(const std::string &name,
                            const std::string &comment);
62 63

  template <typename T>
R
reyoung 已提交
64 65
  TypedAttrChecker<T> &AddAttr(const std::string &name,
                               const std::string &comment,
66
                               bool generated = false) {
R
reyoung 已提交
67
    auto *attr = proto_->add_attrs();
68 69 70 71 72 73 74
    attr->set_name(name);
    attr->set_comment(comment);
    attr->set_generated(generated);
    attr->set_type(AttrTypeID<T>());
    return op_checker_->AddAttrChecker<T>(name);
  }

R
reyoung 已提交
75
  void AddComment(const std::string &comment) { proto_->set_comment(comment); }
76 77 78 79

 private:
  void CheckNoDuplicatedInOutAttrs();

R
reyoung 已提交
80 81
  proto::OpProto *proto_;
  OpAttrChecker *op_checker_;
82 83 84 85
  bool validated_{false};
};
}  // namespace framework
}  // namespace paddle