1 /* Copyright 2021 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_TFRT_JIT_TF_JITRT_PIPELINE_H_
17 #define TENSORFLOW_COMPILER_MLIR_TFRT_JIT_TF_JITRT_PIPELINE_H_
18
19 #include "mlir/Pass/PassManager.h"
20 #include "mlir/Pass/PassOptions.h"
21 #include "llvm/ADT/Hashing.h"
22
23 namespace tensorflow {
24
25 struct TfJitRtPipelineOptions
26 : public mlir::PassPipelineOptions<TfJitRtPipelineOptions> {
27 Option<bool> one_shot_bufferize{
28 *this, "one-shot-bufferize",
29 llvm::cl::desc("Enable one-shot bufferization."), llvm::cl::init(false)};
30 Option<bool> vectorize{*this, "vectorize",
31 llvm::cl::desc("Enable tiling for vectorization."),
32 llvm::cl::init(false)};
33
34 Option<bool> peel{*this, "peel", llvm::cl::desc("Enable loop peeling."),
35 llvm::cl::init(true)};
36
37 Option<bool> fuse_fill{
38 *this, "fuse-fill",
39 llvm::cl::desc("Enable fusion of `linalg.fill` into a tiled reduction."),
40 llvm::cl::init(true)};
41
42 Option<int64_t> vector_size{*this, "vector-size",
43 llvm::cl::desc("Vector size for a 1D reduction."),
44 llvm::cl::init(8)};
45
46 Option<int64_t> reduction_1d_tile_size{
47 *this, "reduction-1d-tile-size",
48 llvm::cl::desc("Tile size for a 1D reduction."), llvm::cl::init(32)};
49
50 ListOption<int64_t> reduction_2d_tile_sizes{
51 *this, "reduction-2d-tile-sizes",
52 llvm::cl::desc("Tile sizes for a 2D reduction."), llvm::cl::ZeroOrMore};
53
54 Option<bool> legalize_i1_tensors{
55 *this, "legalize-i1-tensors",
56 llvm::cl::desc("Convert i1 tensors to i8 tensors."),
57 llvm::cl::init(false)};
58
59 Option<bool> codegen_transpose{
60 *this, "codegen-transpose",
61 llvm::cl::desc(
62 "Enable the specific code generation for transpose operations."),
63 llvm::cl::init(false)};
64 };
65
66 // Make TfJitRtPipelineOptions hashable.
hash_value(const TfJitRtPipelineOptions & opts)67 inline ::llvm::hash_code hash_value(const TfJitRtPipelineOptions& opts) {
68 return ::llvm::hash_value(static_cast<bool>(opts.vectorize));
69 }
70
71 // Creates a pipeline that lowers modules from the Tensorflow dialect to
72 // the Linalg on buffers. `TfJitRtPipelineOptions` contains flags to
73 // enable/disable experimental features.
74 void CreateTfJitRtPipeline(mlir::OpPassManager& pm,
75 const TfJitRtPipelineOptions& options);
76
77 // Calls CreateTfJitRtPipeline with the default TfJitRtPipelineOptions.
78 void CreateDefaultTfJitRtPipeline(mlir::OpPassManager& pm);
79
80 // Creates a pipeline that runs on compiled module specialization. It runs the
81 // Tensorflow shape inference and canonicalization, so that specialized function
82 // always has ranked inputs and results to infer JitRt ABI requirements.
83 void CreateJitRtSpecializationPipeline(mlir::OpPassManager& pm);
84
85 } // namespace tensorflow
86
87 #endif // TENSORFLOW_COMPILER_MLIR_TFRT_JIT_TF_JITRT_PIPELINE_H_
88