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_PASSES_H_ 17 #define TENSORFLOW_COMPILER_MLIR_HLO_INCLUDE_MLIR_HLO_DIALECT_MHLO_TRANSFORMS_PASSES_H_ 18 19 #include <memory> 20 21 #include "llvm/ADT/ArrayRef.h" 22 23 namespace mlir { 24 25 class FuncOp; 26 class FunctionPass; 27 class ModuleOp; 28 class Operation; 29 template <typename T> 30 class OperationPass; 31 class Pass; 32 33 // Transforms unranked HLO operations to ranked ones where possible. 34 std::unique_ptr<FunctionPass> createTransformUnrankedHloPass(); 35 36 namespace mhlo { 37 38 /// Lowers HLO control flow ops to the Standard dialect. 39 std::unique_ptr<OperationPass<FuncOp>> createLegalizeControlFlowPass(); 40 41 /// Lowers MHLO control flow ops to the SCF dialect. 42 std::unique_ptr<OperationPass<FuncOp>> createControlFlowToScfPass(); 43 44 /// Lowers from HLO dialect to Standard dialect. 45 std::unique_ptr<OperationPass<FuncOp>> createLegalizeToStdPass(); 46 47 /// Lowers from the CHLO dialect to the HLO dialect. 48 std::unique_ptr<FunctionPass> createChloLegalizeToHloPass( 49 bool broadcast_only = false); 50 51 /// Lowers from HLO dialect to LHLO dialect allocating/deallocating temporary 52 /// buffers if necessary. 53 std::unique_ptr<OperationPass<ModuleOp>> createLegalizeToLhloPass(); 54 55 // Lowers from HLO dialect to Linalg dialect. 56 std::unique_ptr<OperationPass<FuncOp>> createLegalizeHloToLinalgPass(); 57 58 // Sinks constants implicitly captured in control flow regions. This is 59 // necessary to export to XLA. 60 std::unique_ptr<OperationPass<FuncOp>> createSinkConstantsToControlFlowPass(); 61 62 // fuse mhlo ops to kLoop/kInput fusion patterns 63 std::unique_ptr<OperationPass<FuncOp>> createMhloFusionPass(); 64 65 /// Lowers trigonometric operations from the standard dialect to approximations 66 /// that do not use intrinsics. 67 std::unique_ptr<OperationPass<FuncOp>> 68 createLegalizeTrigonometricToApproximationPass(); 69 70 std::unique_ptr<FunctionPass> createOptimizeMhloPass(); 71 std::unique_ptr<FunctionPass> createLowerComplexPass(); 72 std::unique_ptr<::mlir::Pass> createLegalizeGeneralDotPass(); 73 std::unique_ptr<FunctionPass> createLegalizeGatherToTorchIndexSelectPass(); 74 75 } // namespace mhlo 76 77 namespace lmhlo { 78 79 // Lowers from LHLO dialect to Affine dialect. 80 std::unique_ptr<OperationPass<FuncOp>> createLhloLegalizeToAffinePass(); 81 82 // Lowers from LHLO dialect to Linalg dialect. 83 std::unique_ptr<OperationPass<FuncOp>> createLegalizeLhloToLinalgPass(); 84 85 // Lowers from LHLO dialect to GPU dialect. 86 std::unique_ptr<FunctionPass> createLegalizeToGpuPass(); 87 88 // Fuses linalg ops obtained after LHLO lowering. To enable fusion, 89 // operations are first tiled. 90 // 91 // When 'use_parallel_loops' is set, the tiling will use scf.parallel 92 // operations. Otherwise, scf.for operations are used. 93 // 94 // 'tile_sizes' provides the tile sizes to use for tiling. If the linalg 95 // operation has more dimensions than tile sizes provided, 1 is used as 96 // default. 97 std::unique_ptr<FunctionPass> createLhloFuseLinalgPass( 98 bool use_parallel_loops = false, llvm::ArrayRef<unsigned> tile_sizes = {}); 99 100 // Lowers from LHLO dialect to parallel loops. 101 std::unique_ptr<OperationPass<FuncOp>> createLegalizeLhloToParallelLoopsPass(); 102 103 } // namespace lmhlo 104 105 } // namespace mlir 106 107 #endif // TENSORFLOW_COMPILER_MLIR_HLO_INCLUDE_MLIR_HLO_DIALECT_MHLO_TRANSFORMS_PASSES_H_ 108