| 97 | } |
| 98 | |
| 99 | Status PropagateFromArgOp( |
| 100 | const Node& n, const absl::optional<std::string>& function_name, |
| 101 | const absl::flat_hash_set<int>& resource_arg_indices, |
| 102 | absl::flat_hash_map<const Edge*, ResourceUsageAnalysis::NodeInfo>* |
| 103 | user_to_source) { |
| 104 | TF_RET_CHECK(n.type_string() == kArgOp); |
| 105 | |
| 106 | int index; |
| 107 | TF_RETURN_IF_ERROR(GetNodeAttr(n.attrs(), "index", &index)); |
| 108 | if (!resource_arg_indices.contains(index)) return Status::OK(); |
| 109 | |
| 110 | TF_RET_CHECK(function_name.has_value()) |
| 111 | << "ResourceUsageAnalysis does not support analyzing _Arg nodes " |
| 112 | "carrying Stack/TensorArray resource in given graph unless they " |
| 113 | "are in function calls."; |
| 114 | |
| 115 | const ResourceUsageAnalysis::NodeInfo src_node_info(function_name, n.name(), |
| 116 | n.type_string()); |
| 117 | |
| 118 | for (const Edge* o : n.out_edges()) { |
| 119 | if (o->IsControlEdge()) continue; |
| 120 | if (o->dst()->input_type(o->dst_input()) != DataType::DT_RESOURCE) { |
| 121 | continue; |
| 122 | } |
| 123 | (*user_to_source)[o] = src_node_info; |
| 124 | } |
| 125 | |
| 126 | return Status::OK(); |
| 127 | } |
| 128 | |
| 129 | Status UpdateResourceUsageFromFunctionBodyAnalysis( |
| 130 | const Node& call_node, |
no test coverage detected