Home
last modified time | relevance | path

Searched refs:consumer_id (Results 1 – 6 of 6) sorted by relevance

/external/tensorflow/tensorflow/compiler/xla/python/
Doutfeed_receiver.cc105 OutfeedData(PjRtDevice* device, uint32_t consumer_id, Shape shape) in OutfeedData() argument
107 consumer_id_(consumer_id), in OutfeedData()
113 uint32_t consumer_id() const { return consumer_id_; } in consumer_id() function in xla::OutfeedData
168 uint32_t consumer_id,
295 uint32_t consumer_id = header_data[1]; in DeviceListenerThreadLoop() local
299 auto registered_shape = shape_registry_.find(consumer_id); in DeviceListenerThreadLoop()
303 << "] Cannot find registered shape for consumer ID " << consumer_id in DeviceListenerThreadLoop()
309 auto received = absl::make_unique<OutfeedData>(device, consumer_id, shape); in DeviceListenerThreadLoop()
311 if (consumer_id == kOutfeedCidShutdown) { in DeviceListenerThreadLoop()
368 if (received->consumer_id() == kOutfeedCidShutdown) { in CallbackThreadLoop()
[all …]
Doutfeed_receiver_test.cc59 uint32_t consumer_id; member
63 void Receive(uint32_t consumer_id, std::shared_ptr<Literal> data) { in Receive() argument
65 received_.push_back(Data{consumer_id, data}); in Receive()
114 [&receiver](PjRtDevice* device, uint32_t consumer_id, in TEST()
116 receiver->Receive(consumer_id, data); in TEST()
136 EXPECT_EQ(consumer_id0, received[0].consumer_id); in TEST()
147 [&receiver](PjRtDevice* device, uint32_t consumer_id, in TEST()
149 receiver->Receive(consumer_id, data); in TEST()
179 EXPECT_EQ(consumer_id0, received[0].consumer_id); in TEST()
181 EXPECT_EQ(consumer_id1, received[1].consumer_id); in TEST()
[all …]
Doutfeed_receiver_py.cc51 [this](PjRtDevice* device, uint32_t consumer_id, in OutfeedReceiverForPython()
53 this->Callback(device, consumer_id, std::move(literal)); in OutfeedReceiverForPython()
84 uint32_t consumer_id, std::vector<XlaOp> arrays) { in AddOutfeed() argument
85 return outfeed_receiver_->AddOutfeedToBuilder(builder, token, consumer_id, in AddOutfeed()
89 void Callback(PjRtDevice* device, uint32_t consumer_id, in Callback() argument
109 callback_python_(WrapWithClient<PjRtDevice>(*it, device), consumer_id, in Callback()
Doutfeed_receiver.h66 uint32_t consumer_id,
/external/XNNPACK/src/
Dsubgraph.c550 const uint32_t consumer_id = value->first_consumer; in xnn_subgraph_optimize() local
551 if (consumer_id == XNN_INVALID_NODE_ID) { in xnn_subgraph_optimize()
554 assert(consumer_id < subgraph->num_nodes); in xnn_subgraph_optimize()
558 struct xnn_node* consumer = &subgraph->nodes[consumer_id]; in xnn_subgraph_optimize()
575 … xnn_log_info("fuse Clamp Node #%"PRIu32" into upstream Node #%"PRIu32, consumer_id, producer_id); in xnn_subgraph_optimize()
609 consumer_id, producer_id); in xnn_subgraph_optimize()
623 subgraph->values[fused_input_id].first_consumer = consumer_id; in xnn_subgraph_optimize()
633 consumer_id, producer_id); in xnn_subgraph_optimize()
651 subgraph->values[fused_input_id].first_consumer = consumer_id; in xnn_subgraph_optimize()
/external/tensorflow/tensorflow/core/profiler/utils/
Dgroup_events.cc335 absl::optional<uint64> consumer_id; in EventNode() local
350 consumer_id = stat.IntOrUintValue(); in EventNode()
370 if (!consumer_type.has_value() || !consumer_id.has_value()) { in EventNode()
373 consumer_id = consumer_context->id; in EventNode()
381 if (consumer_type.has_value() && consumer_id.has_value()) { in EventNode()
382 consumer_context_ = {*consumer_type, *consumer_id}; in EventNode()