1 /* Copyright 2020 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 MLIR_HLO_DIALECT_MHLO_IR_CHLO_OPS_H
17 #define MLIR_HLO_DIALECT_MHLO_IR_CHLO_OPS_H
18
19 #include "llvm/ADT/StringRef.h"
20 #include "mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
21 #include "mlir-hlo/utils/hlo_utils.h"
22 #include "mlir/Dialect/Complex/IR/Complex.h"
23 #include "mlir/IR/BuiltinTypes.h"
24 #include "mlir/IR/Dialect.h"
25 #include "mlir/IR/DialectImplementation.h"
26 #include "mlir/IR/MLIRContext.h"
27 #include "mlir/IR/OpDefinition.h"
28 #include "mlir/IR/Operation.h"
29 #include "mlir/IR/TypeUtilities.h"
30 #include "mlir/IR/Types.h"
31 #include "mlir/Interfaces/ControlFlowInterfaces.h"
32 #include "mlir/Interfaces/InferTypeOpInterface.h"
33 #include "mlir/Interfaces/SideEffectInterfaces.h"
34
35 namespace mlir {
36 namespace chlo {
37
38 class ChloDialect : public Dialect {
39 void initialize();
40
41 public:
ChloDialect(MLIRContext * context)42 explicit ChloDialect(MLIRContext* context)
43 : Dialect(getDialectNamespace(), context, TypeID::get<ChloDialect>()) {
44 initialize();
45 }
46 Operation* materializeConstant(OpBuilder& builder, Attribute value, Type type,
47 Location loc) override;
getDialectNamespace()48 static StringRef getDialectNamespace() { return "chlo"; }
49 };
50
51 } // namespace chlo
52 } // namespace mlir
53
54 namespace mlir {
55 namespace chlo {
56 namespace OpTrait {
57
58 template <typename ConcreteType>
59 class Broadcasting
60 : public mlir::OpTrait::TraitBase<ConcreteType, Broadcasting> {};
61
62 } // namespace OpTrait
63 } // namespace chlo
64 } // namespace mlir
65
66 #define GET_OP_CLASSES
67 #include "mlir-hlo/Dialect/mhlo/IR/chlo_ops.h.inc"
68
69 namespace mlir {
70 namespace chlo {
71
72 template <typename T>
getConstantLike(OpBuilder & b,Location loc,T constant,Value val)73 static Value getConstantLike(OpBuilder& b, Location loc, T constant,
74 Value val) {
75 Type ty = getElementTypeOrSelf(val.getType());
76 auto getAttr = [&]() -> Attribute {
77 if (ty.isa<IntegerType>()) return b.getIntegerAttr(ty, constant);
78 if (ty.isa<FloatType>()) return b.getFloatAttr(ty, constant);
79 if (auto complexTy = ty.dyn_cast<ComplexType>())
80 return complex::NumberAttr::get(complexTy, constant, 0);
81 llvm_unreachable("unhandled element type");
82 };
83 return b.create<ConstantLikeOp>(loc, getAttr(), val);
84 }
85
86 Value getConstantLike(OpBuilder& b, Location loc, const APFloat& constant,
87 Value val);
88
89 Value getConstantLikeMaxFiniteValue(OpBuilder& b, Location loc, Value val);
90
91 Value getConstantLikeInfValue(OpBuilder& b, Location loc, Value val,
92 bool negative);
93
94 Value getConstantLikeSmallestFiniteValue(OpBuilder& b, Location loc, Value val);
95
96 } // namespace chlo
97 } // namespace mlir
98
99 #endif // MLIR_HLO_DIALECT_MHLO_IR_CHLO_OPS_H
100