diff options
Diffstat (limited to 'tensorflow/compiler/xla/service/hlo_domain_map.h')
| -rw-r--r-- | tensorflow/compiler/xla/service/hlo_domain_map.h | 4 |
1 files changed, 4 insertions, 0 deletions
diff --git a/tensorflow/compiler/xla/service/hlo_domain_map.h b/tensorflow/compiler/xla/service/hlo_domain_map.h index e62ef763fb..1ca7159725 100644 --- a/tensorflow/compiler/xla/service/hlo_domain_map.h +++ b/tensorflow/compiler/xla/service/hlo_domain_map.h @@ -65,6 +65,10 @@ class HloDomainMap { // currently processing. bool IsDomainInstruction(HloInstruction* instruction) const; + // Retrieves the domain identifier of the instruction, or -1 in case + // instruction is not found within any domain. + int64 GetDomainId(HloInstruction* instruction) const; + private: HloDomainMap(string domain_kind) : domain_kind_(std::move(domain_kind)) {} |