Improve encapsulation in StepDetails
PiperOrigin-RevId: 387004520
Change-Id: Ibe1c2b221818a1c430a96bd21f5cd8f0433f9bfd
diff --git a/tensorflow/core/profiler/utils/event_span.cc b/tensorflow/core/profiler/utils/event_span.cc
index 7bf1794..0176b5c 100644
--- a/tensorflow/core/profiler/utils/event_span.cc
+++ b/tensorflow/core/profiler/utils/event_span.cc
@@ -119,14 +119,6 @@
}
};
-void CombineStepDetails(const StepDetails& src, StepDetails* dst) {
- dst->AppendMarkers(src.Markers());
- dst->AppendEvents(src.Events());
- dst->AppendCollectives(src.Collectives());
- dst->AggregateDeviceMemoryTransfers(src.DeviceMemoryTransfers());
- if (dst->StepName().empty()) dst->SetStepName(src.StepName());
-}
-
constexpr int kNumGenericEventTypes = GenericEventType::kLastGenericEventType -
GenericEventType::kFirstGenericEventType +
1;
@@ -235,7 +227,7 @@
int64_t step_id = step_details.first;
const StepDetails& src_details = step_details.second;
StepDetails* dst_details = &(*dst)[step_id];
- CombineStepDetails(src_details, dst_details);
+ dst_details->Combine(src_details);
}
}
@@ -262,15 +254,8 @@
for (const auto& step_events : overlapped_step_events) {
const auto& step_id = step_events.first;
const auto& step_details = step_events.second;
- *non_overlapped_step_events[step_id].MutableMarkers() =
- step_details.Markers();
- *non_overlapped_step_events[step_id].MutableEvents() =
- ToNonOverlappedEvents(step_details.Events());
- *non_overlapped_step_events[step_id].MutableCollectives() =
- step_details.Collectives();
- *non_overlapped_step_events[step_id].MutableDeviceMemoryTransfers() =
- step_details.DeviceMemoryTransfers();
- non_overlapped_step_events[step_id].SetStepName(step_details.StepName());
+ non_overlapped_step_events.try_emplace(step_id,
+ step_details.ToNonOverlapped());
}
return non_overlapped_step_events;
}
@@ -279,21 +264,6 @@
void StepDetails::AddEvent(const EventTypeSpan& e) { events_.push_back(e); }
-void StepDetails::AppendMarkers(const std::vector<StepMarker>& other_markers) {
- markers_.insert(markers_.end(), other_markers.begin(), other_markers.end());
-}
-
-void StepDetails::AppendEvents(const std::vector<EventTypeSpan>& other_events) {
- events_.insert(events_.end(), other_events.begin(), other_events.end());
-}
-
-void StepDetails::AppendCollectives(
- const absl::flat_hash_map<uint32, AllReduceDbResult>& collectives) {
- for (const auto& it : collectives) {
- collectives_[it.first] = it.second;
- }
-}
-
void StepDetails::AggregateDeviceMemoryTransfers(
const std::vector<DeviceMemoryTransfer> device_memory_transfers) {
if (device_memory_transfers.size() != device_memory_transfers_.size()) {
@@ -367,6 +337,25 @@
return max_device_step_time;
}
+StepDetails StepDetails::ToNonOverlapped() const {
+ StepDetails non_overlapped_step_details;
+ non_overlapped_step_details.markers_ = markers_;
+ non_overlapped_step_details.events_ = ToNonOverlappedEvents(events_);
+ non_overlapped_step_details.collectives_ = collectives_;
+ non_overlapped_step_details.device_memory_transfers_ =
+ device_memory_transfers_;
+ non_overlapped_step_details.step_name_ = step_name_;
+ return non_overlapped_step_details;
+}
+
+void StepDetails::Combine(const StepDetails& other) {
+ markers_.insert(markers_.end(), other.markers_.begin(), other.markers_.end());
+ events_.insert(events_.end(), other.events_.begin(), other.events_.end());
+ collectives_.insert(other.collectives_.begin(), other.collectives_.end());
+ AggregateDeviceMemoryTransfers(other.device_memory_transfers_);
+ if (step_name_.empty()) step_name_ = other.step_name_;
+}
+
std::string StepDetails::DebugString() const {
std::string result = "([";
for (int i = 0, end = markers_.size(); i < end; i++) {
diff --git a/tensorflow/core/profiler/utils/event_span.h b/tensorflow/core/profiler/utils/event_span.h
index dd94b03..ddc44db 100644
--- a/tensorflow/core/profiler/utils/event_span.h
+++ b/tensorflow/core/profiler/utils/event_span.h
@@ -160,14 +160,6 @@
}
// Returns the step time.
Timespan StepTime() const;
- std::vector<StepMarker>* MutableMarkers() { return &markers_; }
- std::vector<EventTypeSpan>* MutableEvents() { return &events_; }
- absl::flat_hash_map<uint32, AllReduceDbResult>* MutableCollectives() {
- return &collectives_;
- }
- std::vector<DeviceMemoryTransfer>* MutableDeviceMemoryTransfers() {
- return &device_memory_transfers_;
- }
// Adds a step-marker to this step.
void AddMarker(const StepMarker& m);
// Adds an EventTypeSpan to this step.
@@ -179,28 +171,30 @@
// allowed.
void AddDeviceMemoryTransferEvent(EventType event_type,
const Timespan& time_span, uint64 bytes);
- // Appends the step-markers from another step to this step.
- void AppendMarkers(const std::vector<StepMarker>& other_markers);
- // Appends the events from another step to this step.
- void AppendEvents(const std::vector<EventTypeSpan>& other_events);
- // Appends the collectives from another step to this step.
- void AppendCollectives(
- const absl::flat_hash_map<uint32, AllReduceDbResult>& collectives);
- // Accumulates the device memory transfers from another step to this step.
- void AggregateDeviceMemoryTransfers(
- const std::vector<DeviceMemoryTransfer> device_memory_transfers);
// Returns the step name.
std::string StepName() const { return step_name_; }
// Sets the name of this step.
void SetStepName(std::string step_name) { step_name_ = step_name; }
+
+ // Converts from overlapped events to non-overlapped events.
+ StepDetails ToNonOverlapped() const;
+
+ // Combines other.
+ void Combine(const StepDetails& other);
+
// Equality test.
bool operator==(const StepDetails& other) const;
// Inequality test.
bool operator!=(const StepDetails& other) const { return !(*this == other); }
+
// Returns a string that prints the content of this object.
std::string DebugString() const;
private:
+ // Accumulates the device memory transfers from another step to this step.
+ void AggregateDeviceMemoryTransfers(
+ const std::vector<DeviceMemoryTransfer> device_memory_transfers);
+
// All step-markers found for marking this step in the traces. There could be
// multiple step-markers for a single step for different reasons. One such
// reason is that there may be one step-marker for the same step on each core;
