allow metadata tf_op override level0's tf_op.
PiperOrigin-RevId: 347040091 Change-Id: I13bed86c9c5b5acd7d95a3c4db96a279755f6a76
This commit is contained in:
parent
cae04d6d4d
commit
bea42a1b32
@ -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,
|
||||
|
@ -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);
|
||||
|
Loading…
Reference in New Issue
Block a user