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 // This file implements logic for lowering MHLO general dot to a regular dot.
17
18 #include "llvm/ADT/STLExtras.h"
19 #include "llvm/ADT/StringSwitch.h"
20 #include "mlir-hlo/Dialect/mhlo/IR/hlo_ops.h"
21 #include "mlir-hlo/Dialect/mhlo/transforms/PassDetail.h"
22 #include "mlir-hlo/Dialect/mhlo/transforms/passes.h"
23 #include "mlir-hlo/Dialect/mhlo/transforms/rewriters.h"
24 #include "mlir/Dialect/StandardOps/IR/Ops.h"
25 #include "mlir/IR/Attributes.h"
26 #include "mlir/IR/BuiltinOps.h"
27 #include "mlir/IR/BuiltinTypes.h"
28 #include "mlir/IR/Location.h"
29 #include "mlir/IR/Operation.h"
30 #include "mlir/IR/TypeUtilities.h"
31 #include "mlir/Pass/Pass.h"
32 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
33
34 namespace mlir {
35 namespace mhlo {
36 namespace {
37
TransposeReshape(Value arg,Location loc,llvm::ArrayRef<int64_t> left_dims,llvm::ArrayRef<int64_t> right_dims,llvm::ArrayRef<int64_t> arg_shape,PatternRewriter * rewriter)38 Value TransposeReshape(Value arg, Location loc,
39 llvm::ArrayRef<int64_t> left_dims,
40 llvm::ArrayRef<int64_t> right_dims,
41 llvm::ArrayRef<int64_t> arg_shape,
42 PatternRewriter *rewriter) {
43 auto element_type = getElementTypeOrSelf(arg.getType());
44
45 int64_t left_size = 1;
46 for (auto dim : left_dims) {
47 left_size *= arg_shape[dim];
48 }
49
50 int64_t right_size = 1;
51 for (auto dim : right_dims) {
52 right_size *= arg_shape[dim];
53 }
54
55 // Generate the transpose permutation attribute.
56 llvm::SmallVector<int64_t, 5> transpose_permutation(left_dims.begin(),
57 left_dims.end());
58 transpose_permutation.append(right_dims.begin(), right_dims.end());
59
60 TensorType transpose_permutation_type = RankedTensorType::get(
61 {static_cast<int64_t>(transpose_permutation.size())},
62 rewriter->getIntegerType(64));
63
64 auto transpose_permutation_attr =
65 DenseIntElementsAttr::get(transpose_permutation_type,
66 llvm::makeArrayRef(transpose_permutation))
67 .cast<DenseIntElementsAttr>();
68
69 // Compute the resulting shape.
70 llvm::SmallVector<int64_t, 5> transposed_shape;
71 for (auto val : transpose_permutation) {
72 transposed_shape.push_back(arg_shape[val]);
73 }
74 auto transpose_type = RankedTensorType::get(transposed_shape, element_type);
75 auto transpose_result = rewriter->create<TransposeOp>(
76 loc, transpose_type, arg, transpose_permutation_attr);
77
78 // Return the final result.
79 auto reshaped_type =
80 RankedTensorType::get({left_size, right_size}, element_type);
81 return rewriter->create<ReshapeOp>(loc, reshaped_type, transpose_result);
82 }
83
ProcessDotArg(Value arg,Location loc,ElementsAttr contract_dims_attr,bool outer_dims_first,PatternRewriter * rewriter)84 Value ProcessDotArg(Value arg, Location loc, ElementsAttr contract_dims_attr,
85 bool outer_dims_first, PatternRewriter *rewriter) {
86 auto shape = arg.getType().cast<ShapedType>().getShape();
87
88 llvm::SmallVector<bool, 5> is_outer_dim;
89 is_outer_dim.resize(shape.size(), true);
90
91 // Compute the contract dimension ordering.
92 llvm::SmallVector<int64_t, 5> contract_dims;
93 for (auto dim : contract_dims_attr.getValues<int64_t>()) {
94 contract_dims.push_back(dim);
95 is_outer_dim[dim] = false;
96 }
97
98 // Compute the outer dimension orderings.
99 llvm::SmallVector<int64_t, 5> outer_dims;
100 for (auto it : llvm::enumerate(is_outer_dim)) {
101 if (it.value()) {
102 outer_dims.push_back(it.index());
103 }
104 }
105
106 if (outer_dims_first) {
107 return TransposeReshape(arg, loc, outer_dims, contract_dims, shape,
108 rewriter);
109 }
110
111 return TransposeReshape(arg, loc, contract_dims, outer_dims, shape, rewriter);
112 }
113
114 struct GeneralDotConvert : public OpRewritePattern<DotGeneralOp> {
115 // Attempts to lower a General Dot operator to a standard Dot operator.
116 // General dots include batching dimensions and can have collapsing
117 // dimensions along any axis. Inserting correctly arrange transpose and
118 // reshape operators organizes the tensors and allows the General Dot to be
119 // replaced with the standard Dot operator.
120 //
121 // Note: This requires an empty list of batch dimensions.
122
GeneralDotConvertmlir::mhlo::__anonf50d5d480111::GeneralDotConvert123 explicit GeneralDotConvert(MLIRContext *context)
124 : OpRewritePattern(context) {}
125
matchAndRewritemlir::mhlo::__anonf50d5d480111::GeneralDotConvert126 LogicalResult matchAndRewrite(DotGeneralOp op,
127 PatternRewriter &rewriter) const override {
128 auto dot_element_type = getElementTypeOrSelf(op);
129
130 auto dot_numbers = op.dot_dimension_numbers();
131 if (dot_numbers.lhs_batching_dimensions().getNumElements() != 0 ||
132 dot_numbers.rhs_batching_dimensions().getNumElements() != 0) {
133 return failure();
134 }
135
136 auto lhs = ProcessDotArg(op.lhs(), op.getLoc(),
137 dot_numbers.lhs_contracting_dimensions(),
138 /*outer_dims_first=*/true, &rewriter);
139
140 auto rhs = ProcessDotArg(op.rhs(), op.getLoc(),
141 dot_numbers.rhs_contracting_dimensions(),
142 /*outer_dims_first=*/false, &rewriter);
143
144 // Accept only static shaped types.
145 auto lhs_shape_type = lhs.getType().dyn_cast_or_null<ShapedType>();
146 auto rhs_shape_type = rhs.getType().dyn_cast_or_null<ShapedType>();
147 if (!lhs_shape_type || !rhs_shape_type) return failure();
148 if (!lhs_shape_type.hasStaticShape() || !rhs_shape_type.hasStaticShape())
149 return failure();
150
151 // Dot resulting shape.
152 auto lhs_shape = lhs_shape_type.getShape();
153 auto rhs_shape = rhs_shape_type.getShape();
154 auto new_dot_type =
155 RankedTensorType::get({lhs_shape[0], rhs_shape[1]}, dot_element_type);
156
157 ArrayAttr precision_config;
158 if (op.precision_config()) precision_config = *op.precision_config();
159 auto new_dot_op = rewriter.create<DotOp>(op.getLoc(), new_dot_type, lhs,
160 rhs, precision_config);
161
162 rewriter.replaceOpWithNewOp<ReshapeOp>(op, op.getType(), new_dot_op);
163 return success();
164 }
165 };
166
167 struct LegalizeGeneralDotPass
168 : public LegalizeGeneralDotPassBase<LegalizeGeneralDotPass> {
169 /// Lower all general dots that can be represented as a non-batched matmul.
runOnFunctionmlir::mhlo::__anonf50d5d480111::LegalizeGeneralDotPass170 void runOnFunction() override {
171 OwningRewritePatternList patterns(&getContext());
172 PopulateGeneralDotOpLoweringPatterns(&patterns, &getContext());
173 (void)applyPatternsAndFoldGreedily(getFunction(), std::move(patterns));
174 }
175 };
176
177 } // namespace
178 } // namespace mhlo
179 } // namespace mlir
180
PopulateGeneralDotOpLoweringPatterns(OwningRewritePatternList * patterns,MLIRContext * ctx)181 void mlir::mhlo::PopulateGeneralDotOpLoweringPatterns(
182 OwningRewritePatternList *patterns, MLIRContext *ctx) {
183 patterns->insert<GeneralDotConvert>(ctx);
184 }
185
createLegalizeGeneralDotPass()186 std::unique_ptr<::mlir::Pass> mlir::mhlo::createLegalizeGeneralDotPass() {
187 return std::make_unique<LegalizeGeneralDotPass>();
188 }
189