Add XPlane to OpMetricsDb converter.
PiperOrigin-RevId: 286225031 Change-Id: I866ae69a24a47517aba285c40dbd9f2ecd615026
This commit is contained in:
parent
b35e03a73b
commit
05df1eb863
@ -8,6 +8,7 @@ cc_library(
|
||||
srcs = ["host_threads_xplane_to_tf_metrics_db.cc"],
|
||||
hdrs = ["host_threads_xplane_to_tf_metrics_db.h"],
|
||||
deps = [
|
||||
":op_metrics_db_combiner",
|
||||
":op_stack",
|
||||
"//tensorflow/core/platform:types",
|
||||
"//tensorflow/core/profiler/protobuf:op_metrics_proto_cc",
|
||||
|
@ -20,6 +20,8 @@ limitations under the License.
|
||||
#include "absl/algorithm/container.h"
|
||||
#include "absl/container/flat_hash_map.h"
|
||||
#include "tensorflow/core/profiler/convert/op_stack.h"
|
||||
#include "tensorflow/core/profiler/protobuf/op_metrics.pb.h"
|
||||
#include "tensorflow/core/profiler/protobuf/xplane.pb.h"
|
||||
#include "tensorflow/core/profiler/utils/op_utils.h"
|
||||
#include "tensorflow/core/profiler/utils/timespan.h"
|
||||
#include "tensorflow/core/profiler/utils/xplane_visitor.h"
|
||||
@ -137,6 +139,23 @@ void CollectTfActivities(const XLineVisitor& line,
|
||||
|
||||
} // namespace
|
||||
|
||||
absl::flat_hash_map<int64, TfOp> CollectTfOpsFromHostThreadsXPlane(
|
||||
const XPlane& host_trace) {
|
||||
absl::flat_hash_map<int64, TfOp> tf_ops;
|
||||
for (const auto& id_metadata : host_trace.event_metadata()) {
|
||||
const XEventMetadata& metadata = id_metadata.second;
|
||||
// On the host, we have added some user-specified TraceMe's in addition to
|
||||
// the TraceMe's added to every TensorFlow op by the system. These
|
||||
// user-inserted TraceMe's have "unknown" type. We don't count them in
|
||||
// Tf-stats.
|
||||
TfOp tf_op = ParseTfOpFullname(metadata.name());
|
||||
if (!IsUnknownOp(tf_op.type)) {
|
||||
tf_ops.try_emplace(metadata.id(), tf_op);
|
||||
}
|
||||
}
|
||||
return tf_ops;
|
||||
}
|
||||
|
||||
TfMetricsDbData ConvertHostThreadsXLineToTfMetricsDbData(
|
||||
const XLineVisitor& line, const absl::flat_hash_map<int64, TfOp>& tf_ops) {
|
||||
TfMetricsDbData tf_metrics_db_data;
|
||||
@ -148,5 +167,24 @@ TfMetricsDbData ConvertHostThreadsXLineToTfMetricsDbData(
|
||||
return tf_metrics_db_data;
|
||||
}
|
||||
|
||||
void ConsumeTfMetricsDbData(TfMetricsDbData src, OpMetricsDbCombiner* dst) {
|
||||
AddIdleOp(&src.tf_metrics_db);
|
||||
dst->Combine(src.tf_metrics_db);
|
||||
src.tf_metrics_db.Clear();
|
||||
}
|
||||
|
||||
OpMetricsDb ConvertHostThreadsXPlaneToTfMetricsDb(const XPlane& host_trace) {
|
||||
absl::flat_hash_map<int64, TfOp> tf_ops =
|
||||
CollectTfOpsFromHostThreadsXPlane(host_trace);
|
||||
OpMetricsDb result;
|
||||
OpMetricsDbCombiner combiner(&result);
|
||||
XPlaneVisitor plane(&host_trace);
|
||||
plane.ForEachLine([&tf_ops, &combiner](const XLineVisitor& line) {
|
||||
ConsumeTfMetricsDbData(
|
||||
ConvertHostThreadsXLineToTfMetricsDbData(line, tf_ops), &combiner);
|
||||
});
|
||||
return result;
|
||||
}
|
||||
|
||||
} // namespace profiler
|
||||
} // namespace tensorflow
|
||||
|
@ -18,6 +18,7 @@ limitations under the License.
|
||||
|
||||
#include "absl/container/flat_hash_map.h"
|
||||
#include "tensorflow/core/platform/types.h"
|
||||
#include "tensorflow/core/profiler/convert/op_metrics_db_combiner.h"
|
||||
#include "tensorflow/core/profiler/protobuf/op_metrics.pb.h"
|
||||
#include "tensorflow/core/profiler/protobuf/xplane.pb.h"
|
||||
#include "tensorflow/core/profiler/utils/event_span.h"
|
||||
@ -40,9 +41,16 @@ struct TfMetricsDbData {
|
||||
HostOpMetricsDbBuilder tf_metrics_db_builder{&tf_metrics_db};
|
||||
};
|
||||
|
||||
absl::flat_hash_map<int64, TfOp> CollectTfOpsFromHostThreadsXPlane(
|
||||
const XPlane& host_trace);
|
||||
|
||||
TfMetricsDbData ConvertHostThreadsXLineToTfMetricsDbData(
|
||||
const XLineVisitor& line, const absl::flat_hash_map<int64, TfOp>& tf_ops);
|
||||
|
||||
void ConsumeTfMetricsDbData(TfMetricsDbData src, OpMetricsDbCombiner* dst);
|
||||
|
||||
OpMetricsDb ConvertHostThreadsXPlaneToTfMetricsDb(const XPlane& host_trace);
|
||||
|
||||
} // namespace profiler
|
||||
} // namespace tensorflow
|
||||
|
||||
|
Loading…
Reference in New Issue
Block a user