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_HLO_INCLUDE_MLIR_HLO_DIALECT_MHLO_TRANSFORMS_REWRITERS_H_
17 #define TENSORFLOW_COMPILER_MLIR_HLO_INCLUDE_MLIR_HLO_DIALECT_MHLO_TRANSFORMS_REWRITERS_H_
18 
19 #include <memory>
20 
21 #include "mlir/IR/MLIRContext.h"
22 #include "mlir/IR/PatternMatch.h"
23 #include "mlir/Transforms/Bufferize.h"
24 #include "mlir/Transforms/DialectConversion.h"
25 
26 namespace mlir {
27 class OwningRewritePatternList;
28 
29 // Populates a collection of rewrite patterns to realize element-wise operations
30 // on ranked tensors where possible.
31 void PopulateTransformUnrankedHloPatterns(MLIRContext *context,
32                                           OwningRewritePatternList *patterns);
33 
34 namespace mhlo {
35 
36 // Collection of rewrite patterns for lowering a general dot product.
37 void PopulateGeneralDotOpLoweringPatterns(OwningRewritePatternList *patterns,
38                                           MLIRContext *ctx);
39 
40 // Collection of rewrite patterns for lowering complex operations to equivalent
41 // float operations.
42 void PopulateComplexLoweringPatterns(MLIRContext *context,
43                                      OwningRewritePatternList *patterns);
44 
45 void PopulateOptimizeMHLOPatterns(MLIRContext *context,
46                                   OwningRewritePatternList *patterns);
47 
48 // Rewrite patterns for gather to equivalent torch index select legalization.
49 void PopulateGatherToTorchIndexSelectPatterns(
50     mlir::MLIRContext *context, OwningRewritePatternList *patterns);
51 
52 void PopulateMhloToStdPatterns(OwningRewritePatternList *patterns,
53                                MLIRContext *ctx);
54 
55 // Collection of rewrite patterns for lowering of dynamic HLOs to LHLO dialect.
56 void populateDynamicHLOToLHLOConversionPattern(
57     MLIRContext *context, BufferizeTypeConverter *converter,
58     OwningRewritePatternList *patterns, bool insert_copy = true);
59 
60 // Collection of rewrite patterns for lowering of HLO to LHLO dialect.
61 void populateHLOToLHLOConversionPattern(MLIRContext *context,
62                                         BufferizeTypeConverter *converter,
63                                         OwningRewritePatternList *patterns);
64 
65 // Collection of rewrite patterns for lowering of HLO to Linalg dialect.
66 void populateHLOToLinalgConversionPattern(MLIRContext *context,
67                                           OwningRewritePatternList *patterns);
68 
69 // Sets up legality definitions for materializing broadcasts.
70 void SetupMaterializeBroadcastsLegality(MLIRContext *context,
71                                         ConversionTarget *conversionTarget);
72 
73 // Populates a collection of rewrite patterns for materializing broadcast
74 // attributes to equivalent sequences of ops.
75 void PopulateMaterializeBroadcastsPatterns(MLIRContext *context,
76                                            OwningRewritePatternList *patterns);
77 
78 // Sets up legality definitions for element-wise operations on ranked tensors.
79 void SetupTransformUnrankedHloLegality(MLIRContext *context,
80                                        ConversionTarget *conversionTarget);
81 
82 // Populates a collection of rewrite patterns to realize element-wise operations
83 // on ranked tensors where possible.
84 void PopulateTransformUnrankedHloPatterns(MLIRContext *context,
85                                           OwningRewritePatternList *patterns);
86 
87 // Populate a collection of conversion patterns for un-fusing
88 // batch_norm_inference and batch_norm_training into constituent HLO ops.
89 // TODO(laurenzo): Implement un-fusing of batch_norm_training.
90 void PopulateUnfuseBatchNormPatterns(MLIRContext *context,
91                                      OwningRewritePatternList *patterns);
92 
93 // Populates patterns that translate the trigonometric operations from the
94 // standard dialect to approximations that do not use intrinsics.
95 void PopulateTrigonometricToApproximationPatterns(
96     MLIRContext *context, OwningRewritePatternList *patterns);
97 
98 }  // namespace mhlo
99 
100 namespace chlo {
101 
102 // Populates a collection of conversion patterns for legalizing broadcasting
103 // client-HLO to their non-broadcasting counterparts.
104 void PopulateChloBroadcastingPatterns(MLIRContext *context,
105                                       OwningRewritePatternList *patterns);
106 
107 // Populates a collection of conversion patterns for legalizing client-HLO to
108 // HLO. Includes decomposition of operations and inserting of explicit
109 // broadcasts.
110 void PopulateLegalizeChloToHloPatterns(MLIRContext *context,
111                                        OwningRewritePatternList *patterns);
112 
113 }  // namespace chlo
114 
115 }  // namespace mlir
116 
117 #endif  // TENSORFLOW_COMPILER_MLIR_HLO_INCLUDE_MLIR_HLO_DIALECT_MHLO_TRANSFORMS_REWRITERS_H_
118