Add flops/bytes of non XLA GPU models to tensorflow stats.
PiperOrigin-RevId: 294749986 Change-Id: I1142617e7cafc259be8ea7ffdc1140918d4d5326
This commit is contained in:
parent
24fa7c6822
commit
356e90472b
@ -16,6 +16,7 @@ cc_library(
|
||||
"//tensorflow/core:lib_internal",
|
||||
"//tensorflow/core/profiler/protobuf:op_metrics_proto_cc",
|
||||
"//tensorflow/core/profiler/protobuf:xplane_proto_cc",
|
||||
"//tensorflow/core/profiler/utils:cost_utils",
|
||||
"//tensorflow/core/profiler/utils:event_span",
|
||||
"//tensorflow/core/profiler/utils:op_utils",
|
||||
"//tensorflow/core/profiler/utils:tf_op_utils",
|
||||
|
@ -24,7 +24,9 @@ limitations under the License.
|
||||
#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/cost_utils.h"
|
||||
#include "tensorflow/core/profiler/utils/op_utils.h"
|
||||
#include "tensorflow/core/profiler/utils/tf_op_utils.h"
|
||||
#include "tensorflow/core/profiler/utils/timespan.h"
|
||||
#include "tensorflow/core/profiler/utils/trace_utils.h"
|
||||
|
||||
@ -198,6 +200,7 @@ OpMetricsDb ConvertDeviceTraceXPlaneToOpMetricsDb(
|
||||
int64 first_op_offset_ps = kint64max;
|
||||
int64 last_op_offset_ps = 0;
|
||||
|
||||
TfOpRoofLineCostEstimator op_level_cost_estimator;
|
||||
XPlaneVisitor plane = CreateTfXPlaneVisitor(&device_trace);
|
||||
plane.ForEachLine([&](const XLineVisitor& line) {
|
||||
if (IsDerivedThreadId(line.Id())) return;
|
||||
@ -210,11 +213,14 @@ OpMetricsDb ConvertDeviceTraceXPlaneToOpMetricsDb(
|
||||
absl::string_view tf_op_fullname = stat->str_value();
|
||||
if (tf_op_fullname.empty()) return;
|
||||
TfOp tf_op = ParseTfOpFullname(tf_op_fullname);
|
||||
TfOpRoofLineCostEstimator::OpRoofLineStats costs;
|
||||
if (tf_op.type != kUnknownOp) {
|
||||
costs = op_level_cost_estimator.Predict(event);
|
||||
}
|
||||
device_op_metrics_db_builder.EnterOp(
|
||||
/*program_id=*/0, tf_op.name, tf_op.type, tf_op_fullname,
|
||||
/*occurrences=*/1, event.DurationPs(),
|
||||
/*children_time_ps=*/0, /*flops=*/0,
|
||||
/*bytes_accessed=*/0);
|
||||
/*children_time_ps=*/0, costs.flops, costs.bytes_accessed);
|
||||
});
|
||||
});
|
||||
result.set_total_time_ps(last_op_offset_ps - first_op_offset_ps);
|
||||
|
Loading…
x
Reference in New Issue
Block a user