allow metadata tf_op override level0's tf_op.

PiperOrigin-RevId: 347040091
Change-Id: I13bed86c9c5b5acd7d95a3c4db96a279755f6a76
This commit is contained in:
A. Unique TensorFlower 2020-12-11 11:31:02 -08:00 committed by TensorFlower Gardener
parent cae04d6d4d
commit bea42a1b32
2 changed files with 20 additions and 16 deletions

View File

@ -224,27 +224,25 @@ struct PerDeviceCollector {
std::vector<Annotation> annotation_stack =
ParseAnnotationStack(event.annotation);
// If multiple metadata have the same key name, show the values from the top
// of the stack (innermost annotation). Concatenate the values from
// "hlo_op".
absl::flat_hash_set<absl::string_view> key_set;
std::vector<absl::string_view> hlo_op_names;
for (auto annotation = annotation_stack.rbegin();
annotation != annotation_stack.rend(); ++annotation) {
for (const Annotation::Metadata& metadata : annotation->metadata) {
if (metadata.key == "tf_op") {
continue; // ignored, obtained from HLO proto via DebugInfoMap
} else if (key_set.insert(metadata.key).second) {
xevent.ParseAndAddStatValue(
*plane->GetOrCreateStatMetadata(metadata.key), metadata.value);
}
}
}
if (!annotation_stack.empty()) {
xevent.AddStatValue(
*plane->GetOrCreateStatMetadata(GetStatTypeStr(StatType::kTfOp)),
*plane->GetOrCreateStatMetadata(annotation_stack.begin()->name));
}
// If multiple metadata have the same key name, show the values from the top
// of the stack (innermost annotation). Concatenate the values from
// "hlo_op".
absl::flat_hash_set<absl::string_view> key_set;
for (auto annotation = annotation_stack.rbegin();
annotation != annotation_stack.rend(); ++annotation) {
for (const Annotation::Metadata& metadata : annotation->metadata) {
if (key_set.insert(metadata.key).second) {
xevent.ParseAndAddStatValue(
*plane->GetOrCreateStatMetadata(metadata.key), metadata.value);
}
}
}
}
absl::optional<int> GetDeviceAttribute(CUdevice device,

View File

@ -133,6 +133,12 @@ TF_ATTRIBUTE_ALWAYS_INLINE inline std::string TraceMeOp(
absl::string_view op_name, absl::string_view op_type) {
return absl::StrCat(op_name, ":", op_type);
}
TF_ATTRIBUTE_ALWAYS_INLINE inline std::string TraceMeOp(const char* op_name,
const char* op_type) {
return absl::StrCat(op_name, ":", op_type);
}
TF_ATTRIBUTE_ALWAYS_INLINE inline std::string TraceMeOp(
std::string&& op_name, absl::string_view op_type) {
absl::StrAppend(&op_name, ":", op_type);