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_XLA_COMPARISON_UTIL_H_ 17 #define TENSORFLOW_COMPILER_XLA_COMPARISON_UTIL_H_ 18 19 #include "absl/base/macros.h" 20 #include "absl/types/optional.h" 21 #include "tensorflow/compiler/xla/primitive_util.h" 22 #include "tensorflow/compiler/xla/statusor.h" 23 #include "tensorflow/compiler/xla/types.h" 24 25 namespace xla { 26 27 class Comparison { 28 public: 29 // Represents type of comparison 30 enum class Type : uint8 { 31 kFloat, 32 kFloatTotalOrder, 33 kSigned, 34 kUnsigned, 35 }; 36 // 37 // Represents different comparison operations. 38 enum class Direction : uint8 { 39 kEq, 40 kNe, 41 kGe, 42 kGt, 43 kLe, 44 kLt, 45 }; 46 47 Comparison() = delete; Comparison(Direction dir,Type type)48 explicit Comparison(Direction dir, Type type) : dir_(dir), type_(type) {} 49 explicit Comparison(Direction dir, PrimitiveType type); 50 GetDirection()51 Direction GetDirection() const { return dir_; } GetType()52 Type GetType() const { return type_; } 53 IsEq()54 inline bool IsEq() const { return dir_ == Direction::kEq; } IsNe()55 inline bool IsNe() const { return dir_ == Direction::kNe; } IsGe()56 inline bool IsGe() const { return dir_ == Direction::kGe; } IsGt()57 inline bool IsGt() const { return dir_ == Direction::kGt; } IsLt()58 inline bool IsLt() const { return dir_ == Direction::kLt; } IsFloat()59 inline bool IsFloat() const { return type_ == Type::kFloat; } IsFloatTotalOrder()60 inline bool IsFloatTotalOrder() const { 61 return type_ == Type::kFloatTotalOrder; 62 } IsSigned()63 inline bool IsSigned() const { return type_ == Type::kSigned; } IsUnsigned()64 inline bool IsUnsigned() const { return type_ == Type::kUnsigned; } 65 66 // Returns true for comparisons, for which (a dir a) is always true. 67 bool IsReflexive() const; 68 69 // Returns true for comparisons, for which (a dir a) is always false. 70 bool IsAntireflexive() const; 71 72 // Gets the converse of the given comparison direction (e.g. >= turns to <=). 73 // Useful when commuting operands to get constants into 74 // immediate-accepting positions in the ISA. 75 Comparison Converse() const; 76 77 // Gets the inverse of the given comparison if it exists (e.g. >= turns to <). 78 // Returns optional value because not all inversions may be supported. 79 absl::optional<Comparison> Inverse() const; 80 81 std::string ToString(std::string prefix1 = ".", 82 std::string prefix2 = ".") const; 83 84 template <typename T, typename Comparator = bool (*)(const T, const T)> GetComparator()85 Comparator GetComparator() const { 86 switch (GetDirection()) { 87 case Direction::kEq: 88 return +[](const T a, const T b) { return a == b; }; 89 case Direction::kNe: 90 return +[](const T a, const T b) { return a != b; }; 91 case Direction::kGe: 92 return +[](const T a, const T b) { return a >= b; }; 93 case Direction::kGt: 94 return +[](const T a, const T b) { return a > b; }; 95 case Direction::kLe: 96 return +[](const T a, const T b) { return a <= b; }; 97 case Direction::kLt: 98 return +[](const T a, const T b) { return a < b; }; 99 } 100 } 101 102 template <typename T> Compare(const T a,const T b)103 bool Compare(const T a, const T b) const { 104 return GetComparator<T>()(a, b); 105 } 106 static Type DefaultComparisonType(PrimitiveType t); 107 108 private: 109 static Direction Converse(Direction dir); 110 static Direction Inverse(Direction dir); 111 112 const Direction dir_; 113 Type type_; 114 }; 115 116 inline std::ostream& operator<<(std::ostream& os, const Comparison& cmp) { 117 return os << cmp.ToString(); 118 } 119 string ComparisonDirectionToString(Comparison::Direction direction); 120 std::string ComparisonTypeToString(Comparison::Type type); 121 122 StatusOr<Comparison::Direction> StringToComparisonDirection( 123 absl::string_view direction_name); 124 125 StatusOr<Comparison::Type> StringToComparisonType( 126 absl::string_view compare_type_name); 127 128 using ComparisonDirection = Comparison::Direction; 129 130 } // namespace xla 131 #endif // TENSORFLOW_COMPILER_XLA_COMPARISON_UTIL_H_ 132