1 /* Copyright 2021 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_HLO_OPS_COMMON_H
17 #define MLIR_HLO_DIALECT_MHLO_IR_HLO_OPS_COMMON_H
18
19 // This file defines functionality shared between chlo/mhlo/lhlo dialects.
20
21 #include <algorithm>
22
23 #include "llvm/ADT/SmallSet.h"
24 #include "mlir/IR/BuiltinAttributes.h"
25 #include "mlir/IR/OpDefinition.h"
26 #include "mlir/IR/OpImplementation.h"
27 #include "mlir/IR/Operation.h"
28
29 namespace mlir {
30 namespace hlo {
31
32 // TODO(b/236017415): remove when mhlo uses prefix accessor.
33 namespace accessor_dispatch {
34 template <typename OpT>
35 auto getReplicaGroups(OpT op, int)
36 -> decltype(op.getReplicaGroups(), DenseIntElementsAttr{}) {
37 return op.getReplicaGroups();
38 }
39 template <typename OpT>
40 auto getReplicaGroups(OpT op, char)
41 -> decltype(op.replica_groups(), DenseIntElementsAttr{}) {
42 return op.replica_groups();
43 }
44 } // namespace accessor_dispatch
45
46 // Verifies replica groups attached to collective communication operations.
47 // If the attribute is not empty, it must be a rank 2 tensor, and each replica
48 // should appear exactly once. If `is_uniform_sized` is true, then we also check
49 // that each group is of the same size. If the operation has
50 // `use_global_device_ids` set, then replica group cannot be empty.
51 template <typename OpT>
verifyReplicaGroups(OpT op,bool isUniformSized)52 LogicalResult verifyReplicaGroups(OpT op, bool isUniformSized) {
53 DenseIntElementsAttr attr = accessor_dispatch::getReplicaGroups(op, 0);
54 auto replicaGroupType = attr.getType().dyn_cast<RankedTensorType>();
55 if (!replicaGroupType || replicaGroupType.getRank() != 2 ||
56 !replicaGroupType.getElementType().isInteger(/*width=*/64))
57 return op.emitOpError(
58 "replica groups should be a rank 2 tensor of 64 bit integers");
59
60 if (replicaGroupType.getShape().equals(ArrayRef<int64_t>{0, 0})) {
61 // verifyReplicaGroups() is used by MHLO and LMHLO, note that MHLO does not
62 // have attr 'use_global_device_ids' actually.
63 if (op->hasAttr("use_global_device_ids") &&
64 op->getAttr("use_global_device_ids")
65 .template cast<BoolAttr>()
66 .getValue()) {
67 return op.emitOpError(
68 "if `use_global_device_ids` is set, the replica groups cannot be "
69 "empty");
70 }
71 return success();
72 }
73
74 int64_t maxReplicaIdSeen = 0;
75 llvm::SmallSet<int64_t, 8> replicaSeen;
76 for (int64_t id : attr.getValues<int64_t>()) {
77 // Replica groups are stored in a 2D tensor. If the op supports non-uniform
78 // groups, null replica IDs are stored as -1.
79 if (id == -1) {
80 if (isUniformSized) {
81 return op.emitOpError("Invalid replica id -1");
82 }
83 continue;
84 }
85
86 if (!replicaSeen.insert(id).second) {
87 return op.emitOpError("replica id #") << id << " seen more than once";
88 }
89 maxReplicaIdSeen = std::max(maxReplicaIdSeen, id);
90 }
91
92 for (int64_t id = 0; id <= maxReplicaIdSeen; id++) {
93 if (!replicaSeen.contains(id)) {
94 return op.emitOpError("replica id #")
95 << id << " not seen in replica groups";
96 }
97 }
98 return success();
99 }
100
101 // Verifies the source target pairs attached to collective permute.
102 LogicalResult verifyCollectivePermuteSourceTargetPairs(
103 Operation* op, DenseIntElementsAttr attr);
104
105 LogicalResult verifyReduceScatter(Operation* op, TypeRange operandTypes,
106 TypeRange resultTypes,
107 uint64_t scatterDimension);
108
109 // Custom formatting for convolution window attributes.
110 void printWindowAttributes(OpAsmPrinter& p, Operation* op,
111 llvm::Optional<DenseIntElementsAttr> windowStrides,
112 llvm::Optional<DenseIntElementsAttr> padding,
113 llvm::Optional<DenseIntElementsAttr> lhsDilation,
114 llvm::Optional<DenseIntElementsAttr> rhsDilation,
115 llvm::Optional<DenseElementsAttr> windowReversal);
116
117 ParseResult parseWindowAttributes(OpAsmParser& parser,
118 DenseIntElementsAttr& windowStrides,
119 DenseIntElementsAttr& padding,
120 DenseIntElementsAttr& lhsDilation,
121 DenseIntElementsAttr& rhsDilation,
122 DenseElementsAttr& windowReversal);
123
124 } // namespace hlo
125 } // namespace mlir
126
127 #endif // MLIR_HLO_DIALECT_MHLO_IR_HLO_OPS_COMMON_H
128