1 /** 2 * Copyright 2020 Huawei Technologies Co., Ltd 3 * 4 * Licensed under the Apache License, Version 2.0 (the "License"); 5 * you may not use this file except in compliance with the License. 6 * You may obtain a copy of the License at 7 * 8 * http://www.apache.org/licenses/LICENSE-2.0 9 * 10 * Unless required by applicable law or agreed to in writing, software 11 * distributed under the License is distributed on an "AS IS" BASIS, 12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 * See the License for the specific language governing permissions and 14 * limitations under the License. 15 */ 16 17 #ifndef MINDSPORE_CCSRC_MINDDATA_DATASET_ENGINE_IR_DATASETOPS_REPEAT_NODE_H_ 18 #define MINDSPORE_CCSRC_MINDDATA_DATASET_ENGINE_IR_DATASETOPS_REPEAT_NODE_H_ 19 20 #include <map> 21 #include <memory> 22 #include <set> 23 #include <string> 24 #include <vector> 25 26 #include "minddata/dataset/engine/ir/datasetops/dataset_node.h" 27 28 namespace mindspore { 29 namespace dataset { 30 class RepeatOp; 31 32 class RepeatNode : public DatasetNode { 33 // Allow GeneratorNode to access internal members 34 friend class GeneratorNode; 35 36 public: 37 /// \brief Constructor RepeatNode()38 RepeatNode() : op_(nullptr), repeat_count_(-1) {} 39 40 /// \brief Constructor 41 RepeatNode(std::shared_ptr<DatasetNode> child, int32_t count); 42 43 /// \brief Destructor 44 ~RepeatNode() override = default; 45 46 /// \brief Node name getter 47 /// \return Name of the current node Name()48 std::string Name() const override { return kRepeatNode; } 49 50 /// \brief Print the description 51 /// \param out - The output stream to write output to 52 void Print(std::ostream &out) const override; 53 54 /// \brief Copy the node to a new object 55 /// \return A shared pointer to the new copy 56 std::shared_ptr<DatasetNode> Copy() override; 57 58 /// \brief a base class override function to create the required runtime dataset op objects for this class 59 /// \param node_ops - A vector containing shared pointer to the Dataset Ops that this object will create 60 /// \return Status Status::OK() if build successfully 61 Status Build(std::vector<std::shared_ptr<DatasetOp>> *const node_ops) override; 62 63 /// \brief Parameters validation 64 /// \return Status Status::OK() if all the parameters are valid 65 Status ValidateParams() override; 66 67 /// \brief Getter 68 /// \return Number of cycles to repeat the execution Count()69 int32_t Count() const { return repeat_count_; } 70 71 /// \brief Base-class override for GetDatasetSize 72 /// \param[in] size_getter Shared pointer to DatasetSizeGetter 73 /// \param[in] estimate This is only supported by some of the ops and it's used to speed up the process of getting 74 /// dataset size at the expense of accuracy. 75 /// \param[out] dataset_size the size of the dataset 76 /// \return Status of the function 77 Status GetDatasetSize(const std::shared_ptr<DatasetSizeGetter> &size_getter, bool estimate, 78 int64_t *dataset_size) override; 79 80 /// \brief Base-class override for accepting IRNodePass visitor 81 /// \param[in] p The node to visit 82 /// \param[out] modified Indicator if the node was modified 83 /// \return Status of the node visit 84 Status Accept(IRNodePass *const p, bool *const modified) override; 85 86 /// \brief Base-class override for accepting IRNodePass visitor 87 /// \param[in] p The node to visit 88 /// \param[out] modified Indicator if the node was modified 89 /// \return Status of the node visit 90 Status AcceptAfter(IRNodePass *const p, bool *const modified) override; 91 92 /// \brief Record the Repeat/EpochCtrl node that is the closest ancestor of this node 93 /// \param[in] the ancestor node 94 /// \return Status of the function AddResetAncestor(const std::shared_ptr<RepeatNode> & src)95 Status AddResetAncestor(const std::shared_ptr<RepeatNode> &src) { 96 /* 97 * This check is to ensure we don't overwrite an existing value of its ancestor. 98 * It is okay to assign to the same value more than once in RepeatNode (but not in GeneratorNode). 99 * Consider the following scenario 100 * EpochCtrl(-1) 101 * | 102 * Repeat 103 * | 104 * Concat 105 * / \ 106 * GenData1 GenData2 107 * 108 * We will record the ancestor relationship of (Repeat, EpochCtrl) twice, one at Visit(GenData1), the other at 109 * Vist(GenData2). 110 */ 111 std::shared_ptr<RepeatNode> tmp_repeat_node = reset_ancestor_.lock(); 112 CHECK_FAIL_RETURN_UNEXPECTED(tmp_repeat_node == nullptr || tmp_repeat_node == src, 113 "Internal error: Overwriting an existing value"); 114 reset_ancestor_ = src; 115 return Status::OK(); 116 } 117 118 /// \brief Getter functions RepeatCount()119 int32_t RepeatCount() const { return repeat_count_; } 120 121 /// \brief Get the arguments of node 122 /// \param[out] out_json JSON string of all attributes 123 /// \return Status of the function 124 Status to_json(nlohmann::json *out_json) override; 125 126 /// \brief Function for read dataset operation from json 127 /// \param[in] json_obj The JSON object to be deserialized 128 /// \param[in] ds dataset node constructed 129 /// \param[out] result Deserialized dataset after the operation 130 /// \return Status The status code returned 131 static Status from_json(nlohmann::json json_obj, std::shared_ptr<DatasetNode> ds, 132 std::shared_ptr<DatasetNode> *result); 133 134 protected: 135 std::shared_ptr<RepeatOp> op_; // keep its corresponding run-time op of EpochCtrlNode and RepeatNode 136 std::weak_ptr<RepeatNode> reset_ancestor_; // updated its immediate Repeat/EpochCtrl ancestor in GeneratorNodePass 137 int32_t repeat_count_; 138 }; 139 } // namespace dataset 140 } // namespace mindspore 141 #endif // MINDSPORE_CCSRC_MINDDATA_DATASET_ENGINE_IR_DATASETOPS_REPEAT_NODE_H_ 142