1 /* Copyright 2018 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_JIT_INCREASE_DYNAMISM_FOR_AUTO_JIT_PASS_H_
17 #define TENSORFLOW_COMPILER_JIT_INCREASE_DYNAMISM_FOR_AUTO_JIT_PASS_H_
18 
19 #include "tensorflow/core/common_runtime/optimization_registry.h"
20 
21 namespace tensorflow {
22 
23 // Increases the amount of "dynamism" representable by XLA clusters by rewriting
24 // the TensorFlow graph.  This pass does the following rewrites:
25 //
26 // Slice
27 // -----
28 //
29 //   Slice(op, begin, size <must be constant>) =>
30 //     Slice(op, begin, actual_size(op.shape(), size, begin));
31 //       _XlaCompileTimeConstantInputs={2}
32 //
33 // where
34 //
35 //   actual_size(op_shape, size, begin)[i] =
36 //     size[i] == -1 ? (op_shape[i] - size[i])
37 //                   : size[i]
38 //
39 // This pass, combined with jit/partially_decluster_pass, reduces the number of
40 // unnecessary cluster recompilations in some common cases.  After the rewrite
41 // shown above jit/partially_decluster_pass extracts the actual_size(...)
42 // computation to outside the XLA cluster, causing the cluster to be versioned
43 // only on the actual size of the XlaDynamicSlice.  This avoids recompilation
44 // due to superficial changes that don't affect tensor shapes.
45 //
46 // Future Work TODO(b/111210515)
47 // -----------------------------
48 //
49 // In the future we will also translate StridedSlice and Pad a similar way.
50 class IncreaseDynamismForAutoJitPass : public GraphOptimizationPass {
51  public:
52   Status Run(const GraphOptimizationPassOptions& options) override;
53 };
54 
55 }  // namespace tensorflow
56 
57 #endif  // TENSORFLOW_COMPILER_JIT_INCREASE_DYNAMISM_FOR_AUTO_JIT_PASS_H_
58