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 // This transformation pass converts unified compilation and replication
17 // attributes into legacy attributes. For example,  _replication_info=X
18 // and _xla_compile_device_type=TPU should be replaced with _tpu_replicate=X.
19 // This ensures the unified attributes not get exposed outside of the MLIR
20 // bridge with V1 pipeline in some cases.
21 
22 #include "llvm/ADT/StringRef.h"
23 #include "llvm/Support/Debug.h"
24 #include "mlir/IR/Builders.h"  // from @llvm-project
25 #include "mlir/IR/Operation.h"  // from @llvm-project
26 #include "mlir/Pass/Pass.h"  // from @llvm-project
27 #include "mlir/Support/LLVM.h"  // from @llvm-project
28 #include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"
29 #include "tensorflow/compiler/mlir/tensorflow/transforms/passes_detail.h"
30 #include "tensorflow/compiler/mlir/tensorflow/utils/attribute_utils.h"
31 
32 namespace mlir {
33 namespace TFTPU {
34 
35 namespace {
36 
37 struct ConvertToLegacyCompileAndReplicateAttributesPass
38     : public TF::ConvertToLegacyCompileAndReplicateAttributesPassBase<
39           ConvertToLegacyCompileAndReplicateAttributesPass> {
40   void runOnOperation() override;
41 };
42 
ConvertToLegacyAttributes(func::FuncOp func_op)43 LogicalResult ConvertToLegacyAttributes(func::FuncOp func_op) {
44   auto result = func_op->walk([&](mlir::Operation* op) {
45     if (failed(TF::HasValidCompilationAndReplicationAttributes(*op)))
46       return WalkResult::interrupt();
47     if (op->hasAttr(TF::kReplicationInfoAttr)) {
48       op->setAttr(TF::kTpuReplicateAttr, op->getAttr(TF::kReplicationInfoAttr));
49       op->removeAttr(TF::kReplicationInfoAttr);
50       op->removeAttr(TF::kCompileDeviceTypeAttr);
51     }
52     return mlir::WalkResult::advance();
53   });
54   return failure(result.wasInterrupted());
55 }
56 
runOnOperation()57 void ConvertToLegacyCompileAndReplicateAttributesPass::runOnOperation() {
58   func::FuncOp func_op = getOperation();
59   if (failed(ConvertToLegacyAttributes(func_op))) return signalPassFailure();
60 }
61 
62 }  // namespace
63 
64 std::unique_ptr<OperationPass<func::FuncOp>>
CreateConvertToLegacyCompileAndReplicateAttributesPass()65 CreateConvertToLegacyCompileAndReplicateAttributesPass() {
66   return std::make_unique<ConvertToLegacyCompileAndReplicateAttributesPass>();
67 }
68 
69 }  // namespace TFTPU
70 }  // namespace mlir
71