xref: /aosp_15_r20/external/tensorflow/tensorflow/compiler/xla/service/async_op_canonicalizer.h (revision b6fb3261f9314811a0f4371741dbb8839866f948)
1 /* Copyright 2022 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_XLA_SERVICE_ASYNC_OP_CANONICALIZER_H_
17 #define TENSORFLOW_COMPILER_XLA_SERVICE_ASYNC_OP_CANONICALIZER_H_
18 
19 #include "tensorflow/compiler/xla/service/hlo_module.h"
20 #include "tensorflow/compiler/xla/service/hlo_pass_interface.h"
21 
22 namespace xla {
23 
24 // This pass looks at all of the async operations in the module and assigns the
25 // async operations that participate in the same async action a unique async
26 // group id. Async operations in the same group id typically consist of one
27 // async-start operation, one async-done operation, and zero or more
28 // async-update operations. Then, this pass ensures all of the async operations
29 // with the same group id wrap the same computation such that each async
30 // computation is associated with all of the async operations that have the same
31 // group id.
32 class AsyncOpCanonicalizer : public HloModulePass {
33  public:
34   ~AsyncOpCanonicalizer() override = default;
name()35   absl::string_view name() const override { return "async-op-canonicalizer"; }
36   using HloPassInterface::Run;
37   StatusOr<bool> Run(
38       HloModule* module,
39       const absl::flat_hash_set<absl::string_view>& execution_threads) override;
40 };
41 
42 }  // namespace xla
43 
44 #endif  // TENSORFLOW_COMPILER_XLA_SERVICE_ASYNC_OP_CANONICALIZER_H_
45