op_variant.cc 3.0 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19
// Copyright (c) 2019 PaddlePaddle 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 "paddle/fluid/operators/controlflow/op_variant.h"

namespace paddle {
namespace operators {

20
struct InputsVisitor {
21 22 23 24 25 26
  template <typename OpType>
  const framework::VariableNameMap *operator()(const OpType *op) const {
    return &(op->Inputs());
  }
};

27
struct OutputsVisitor {
28 29 30 31 32 33
  template <typename OpType>
  const framework::VariableNameMap *operator()(const OpType *op) const {
    return &(op->Outputs());
  }
};

34
struct AttributeMapVisitor {
35 36 37 38 39 40 41 42 43 44
  const framework::AttributeMap *operator()(const framework::OpDesc *op) const {
    return &(op->GetAttrMap());
  }

  const framework::AttributeMap *operator()(
      const framework::OperatorBase *op) const {
    return &(op->Attrs());
  }
};

45
struct RawPointerVisitor {
46 47 48 49 50 51 52
  template <typename OpType>
  const void *operator()(const OpType *op) const {
    return op;
  }
};

const framework::VariableNameMap &OpVariant::Inputs() const {
R
Ruibiao Chen 已提交
53
  return *paddle::visit(InputsVisitor(), op_);
54 55 56
}

const framework::VariableNameMap &OpVariant::Outputs() const {
R
Ruibiao Chen 已提交
57
  return *paddle::visit(OutputsVisitor(), op_);
58 59 60
}

const framework::AttributeMap &OpVariant::Attrs() const {
R
Ruibiao Chen 已提交
61
  return *paddle::visit(AttributeMapVisitor(), op_);
62 63 64
}

const void *OpVariant::RawPointer() const {
R
Ruibiao Chen 已提交
65
  return paddle::visit(RawPointerVisitor(), op_);
66 67
}

68 69 70 71 72 73 74
void AppendOpVariantByOpName(const std::vector<framework::OpDesc *> &op_descs,
                             const std::string &candidate_op_name,
                             std::vector<OpVariant> *result_ops) {
  PADDLE_ENFORCE_NOT_NULL(
      result_ops,
      platform::errors::Unavailable("result_ops should not be a null_ptr."));
  for (auto *op_desc : op_descs) {
75 76 77
    PADDLE_ENFORCE_NOT_NULL(
        op_desc,
        platform::errors::Unavailable("op_desc should not be a null_ptr."));
78 79 80 81 82 83
    if (op_desc->Type() == candidate_op_name) {
      result_ops->emplace_back(op_desc);
    }
  }
}

84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100
void AppendOpVariantByOpName(
    const std::vector<framework::OpDesc *> &op_descs,
    const std::string &candidate_op_name,
    std::unordered_set<OpVariant, OpVariant::Hasher> *result_ops) {
  PADDLE_ENFORCE_NOT_NULL(
      result_ops,
      platform::errors::Unavailable("result_ops should not be a null_ptr."));
  for (auto *op_desc : op_descs) {
    PADDLE_ENFORCE_NOT_NULL(
        op_desc,
        platform::errors::Unavailable("op_desc should not be a null_ptr."));
    if (op_desc->Type() == candidate_op_name) {
      result_ops->emplace(op_desc);
    }
  }
}

101 102
}  // namespace operators
}  // namespace paddle