xref: /aosp_15_r20/external/tensorflow/tensorflow/compiler/mlir/tools/kernel_gen/transforms/passes.h (revision b6fb3261f9314811a0f4371741dbb8839866f948)
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