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 TENSORFLOW_COMPILER_MLIR_TOOLS_KERNEL_GEN_TRANSFORMS_PASSES_H_ 17 #define TENSORFLOW_COMPILER_MLIR_TOOLS_KERNEL_GEN_TRANSFORMS_PASSES_H_ 18 19 #include <memory> 20 21 #include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project 22 #include "mlir/Dialect/GPU/IR/GPUDialect.h" // from @llvm-project 23 #include "mlir/Dialect/LLVMIR/LLVMDialect.h" // from @llvm-project 24 #include "mlir/IR/BuiltinOps.h" // from @llvm-project 25 #include "mlir/Pass/Pass.h" // from @llvm-project 26 27 namespace mlir { 28 namespace kernel_gen { 29 namespace tf_framework { 30 31 // Pass to replace some of the Standard ops with TF Framework ops. 32 // * adds tf_framework::OpKernelContextType argument to the function 33 // * std.alloc becomes tf_framework.alloc_raw 34 // * std.dealloc becomes tf_framework.dealloc_raw 35 // * std.assert becomes tf_framework.assert 36 std::unique_ptr<OperationPass<ModuleOp>> CreateEmbedTFFrameworkPass(); 37 38 // Pass to convert tf_framework.assert operations to calls to 39 // tf_framework.report_error and create the required control flow to abort the 40 // function on failed execution. 41 std::unique_ptr<OperationPass<ModuleOp>> CreateRewriteTFFrameworkAssert(); 42 43 } // namespace tf_framework 44 45 namespace transforms { 46 47 // Pass to find and annotate candidates for buffer reuse. 48 std::unique_ptr<OperationPass<func::FuncOp>> CreateBufferReusePass(); 49 50 // Pass to rewrite all TF operations to JIT invocations through the TF 51 // framework. 52 std::unique_ptr<OperationPass<func::FuncOp>> CreateTFToJITInvocationPass( 53 llvm::ArrayRef<int64_t> tile_sizes = {}, 54 llvm::ArrayRef<int64_t> unroll_factors = {}, int64_t max_supported_rank = 5, 55 bool enable_ftz = false, bool index_64bit = false, bool cpu_codegen = false, 56 bool jit_i64_indexed_for_large_tensors = false); 57 58 // Pass for applying LLVM legalization patterns. 59 std::unique_ptr<OperationPass<ModuleOp>> CreateTFKernelToLLVMPass( 60 mlir::StringRef blob_annotation = {}); 61 62 // Pass to tranform shape computations in shape dialect to standard and scf 63 // using memref descriptors. 64 std::unique_ptr<OperationPass<ModuleOp>> CreateShapeToDescriptorsPass(); 65 66 // Pass to convert scf::ParallelOp to scf::ForOp. 67 std::unique_ptr<OperationPass<func::FuncOp>> CreateParallelLoopsToSequential(); 68 69 // Pass to annotate GPU Module with its PTX. 70 std::unique_ptr<OperationPass<gpu::GPUModuleOp>> CreateGpuKernelToBlobPass( 71 mlir::StringRef blob_annotation = {}, 72 ArrayRef<std::string> architectures = {}, bool print_ptx = false, 73 bool print_llvmir = false, bool enable_ftz = false); 74 75 // Pass to propagate tensorflow runtime ABI knowledge across kernel boundaries. 76 std::unique_ptr<OperationPass<func::FuncOp>> 77 CreatePropagateTfAbiKnowledgeToKernels(); 78 79 // Pass to propagate shape equalities across kernel boundaries. 80 std::unique_ptr<OperationPass<func::FuncOp>> 81 CreatePropagateShapeKnowledgeToKernels(); 82 83 // Pass to print content of memrefs. 84 std::unique_ptr<OperationPass<ModuleOp>> CreateEmbedMemRefPrintsPass(); 85 86 /// Greedily maps loops to GPU hardware dimensions. 87 std::unique_ptr<mlir::OperationPass<func::FuncOp>> CreateMapParallelLoopsPass(); 88 89 /// We need to direct fusion to the inner loops. This cannot be done with 90 /// a passmanager alone ATM, as nested pass managers require operations to 91 /// be closed from above. 92 std::unique_ptr<mlir::OperationPass<func::FuncOp>> 93 CreateFuseInnerParallelLoopsPass(); 94 95 // Pass to remove copies which are consumed by a GenericOp. 96 std::unique_ptr<OperationPass<func::FuncOp>> CreateCopyCleanupPass(); 97 98 std::unique_ptr<OperationPass<ModuleOp>> CreateKernelgenFinalBufferizePass(); 99 100 } // namespace transforms 101 102 #define GEN_PASS_REGISTRATION 103 #include "tensorflow/compiler/mlir/tools/kernel_gen/transforms/kernel_gen_passes.h.inc" 104 105 } // namespace kernel_gen 106 } // namespace mlir 107 108 #endif // TENSORFLOW_COMPILER_MLIR_TOOLS_KERNEL_GEN_TRANSFORMS_PASSES_H_ 109