From a3f1c2a668684fd48d3fc6574d8fc0472655b2a9 Mon Sep 17 00:00:00 2001 From: Abdurrahman Akkas Date: Tue, 21 Apr 2020 18:59:04 -0700 Subject: [PATCH] Add an interface to add separator before appending numbers in OpOrArgNameMapper. PiperOrigin-RevId: 307725706 Change-Id: If75e1eb68eb37ece533de7cceccf1cc7da290be8 --- tensorflow/compiler/mlir/op_or_arg_name_mapper.cc | 4 +++- tensorflow/compiler/mlir/op_or_arg_name_mapper.h | 3 +++ 2 files changed, 6 insertions(+), 1 deletion(-) diff --git a/tensorflow/compiler/mlir/op_or_arg_name_mapper.cc b/tensorflow/compiler/mlir/op_or_arg_name_mapper.cc index 272fab9cd1c..bce0ed4a33d 100644 --- a/tensorflow/compiler/mlir/op_or_arg_name_mapper.cc +++ b/tensorflow/compiler/mlir/op_or_arg_name_mapper.cc @@ -55,8 +55,10 @@ llvm::StringRef OpOrArgNameMapper::GetUniqueName(llvm::StringRef prefix) { // to be unique. auto& val = prefix_it.first->second; llvm::SmallString<64> probe_name(prefix); + probe_name.append(GetSuffixSeparator()); + const int probe_prefix_size = probe_name.size(); while (true) { - probe_name.resize(prefix.size()); + probe_name.resize(probe_prefix_size); // TODO(jpienaar): Subtract one so that the initial suffix is 0 instead // of 1. // TODO(jpienaar): Switch to radix 36 and update tests. diff --git a/tensorflow/compiler/mlir/op_or_arg_name_mapper.h b/tensorflow/compiler/mlir/op_or_arg_name_mapper.h index 108496e2283..6a52d13fbc0 100644 --- a/tensorflow/compiler/mlir/op_or_arg_name_mapper.h +++ b/tensorflow/compiler/mlir/op_or_arg_name_mapper.h @@ -64,6 +64,9 @@ class OpOrArgNameMapper { return op_or_val_to_name_; } + // Returns the separator used before uniqueing suffix. + virtual llvm::StringRef GetSuffixSeparator() { return ""; } + private: // Returns name from the location of the operation or value. virtual std::string GetName(OpOrVal op_or_val) = 0;