1 /* Copyright 2019 The TensorFlow Authors. All Rights Reserved. 2 3 Licensed under the Apache License, Version 2.0 (the "License"); 4 you may not use this file except in compliance with the License. 5 You may obtain a copy of the License at 6 7 http://www.apache.org/licenses/LICENSE-2.0 8 9 Unless required by applicable law or agreed to in writing, software 10 distributed under the License is distributed on an "AS IS" BASIS, 11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 See the License for the specific language governing permissions and 13 limitations under the License. 14 ==============================================================================*/ 15 16 #ifndef TENSORFLOW_COMPILER_MLIR_op_or_val_NAME_MAPPER_H_ 17 #define TENSORFLOW_COMPILER_MLIR_op_or_val_NAME_MAPPER_H_ 18 19 #include <string> 20 21 #include "absl/strings/string_view.h" 22 #include "llvm/ADT/DenseMap.h" 23 #include "llvm/ADT/PointerUnion.h" 24 #include "llvm/ADT/StringMap.h" 25 #include "llvm/ADT/StringRef.h" 26 #include "mlir/IR/Operation.h" // from @llvm-project 27 #include "mlir/IR/Value.h" // from @llvm-project 28 29 namespace tensorflow { 30 31 // PointerUnion for operation and value. 32 // TODO(jpienaar): Rename the files. 33 using OpOrVal = llvm::PointerUnion<mlir::Operation*, mlir::Value>; 34 35 // Mapper from operation or value to name. 36 class OpOrArgNameMapper { 37 public: 38 // Returns unique name for the given prefix. 39 llvm::StringRef GetUniqueName(llvm::StringRef prefix); 40 41 // Returns unique name for the operation or value. 42 llvm::StringRef GetUniqueName(OpOrVal op_or_val); 43 44 // Returns unique name as a string_view for the operation or value. 45 absl::string_view GetUniqueNameView(OpOrVal op_or_val); 46 47 // Initializes operation or value to map to name. Returns number of 48 // operations or value already named 'name' which should be 0 else 49 // GetUniqueName could return the same names for different operations or 50 // values. 51 // Note: Its up to the caller to decide the behavior when assigning two 52 // operations or values to the same name. 53 int InitOpName(OpOrVal op_or_val, llvm::StringRef name); 54 55 virtual ~OpOrArgNameMapper(); 56 57 protected: 58 // Returns true if the name is unique. A derived class can override it if the 59 // class maintains uniqueness in a different scope. 60 virtual bool IsUnique(llvm::StringRef name); 61 62 // Returns a constant view of the underlying map. GetMap()63 const llvm::DenseMap<OpOrVal, absl::string_view>& GetMap() const { 64 return op_or_val_to_name_; 65 } 66 67 // Returns the separator used before uniqueing suffix. GetSuffixSeparator()68 virtual llvm::StringRef GetSuffixSeparator() { return ""; } 69 70 private: 71 // Returns name from the location of the operation or value. 72 virtual std::string GetName(OpOrVal op_or_val) = 0; 73 74 // Maps string name to count. This map is used to help keep track of unique 75 // names for operations or values. 76 llvm::StringMap<int64_t> name_to_count_; 77 // Maps operation or values to name. Value in map is a view of the string 78 // name in `name_to_count_`. Names in `name_to_count_` are never removed. 79 llvm::DenseMap<OpOrVal, absl::string_view> op_or_val_to_name_; 80 }; 81 82 // OpOrArgNameMapper that returns, for operations or values not initialized 83 // to a specific name, a name based on the location of the operation or 84 // value. 85 class OpOrArgLocNameMapper : public OpOrArgNameMapper { 86 protected: 87 std::string GetName(OpOrVal op_or_val) override; 88 }; 89 90 // OpOrArgNameMapper that returns, for operations or values not initialized 91 // to a specific name, a short name. 92 class OpOrArgStripNameMapper : public OpOrArgNameMapper { 93 private: 94 std::string GetName(OpOrVal op_or_val) override; 95 96 // Number of ops mapped. 97 int count_ = 0; 98 }; 99 100 } // namespace tensorflow 101 102 #endif // TENSORFLOW_COMPILER_MLIR_op_or_val_NAME_MAPPER_H_ 103