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 | #include "tensorflow/dtensor/mlir/op_utils.h" |
17 | |
18 | #include <string> |
19 | |
20 | #include "llvm/Support/raw_ostream.h" |
21 | |
22 | namespace tensorflow { |
23 | namespace dtensor { |
24 | |
25 | // OpHash prints the `op` into a string and performs hash value on the output |
26 | // string. |
27 | // |
28 | // The `print` includes the full representation of the `op`, e.g., target |
29 | // register, layout, shape, etc. This should be sufficient to uniquely |
30 | // identify the operation in most cases. This does not capture function scope |
31 | // (identical op in 2 separate functions). |
32 | uint64_t OpHash(mlir::Operation* op) { |
33 | std::string output; |
34 | llvm::raw_string_ostream output_stream(output); |
35 | mlir::OpPrintingFlags flags; |
36 | flags.elideLargeElementsAttrs(1024); |
37 | op->print(output_stream, flags); |
38 | return llvm::hash_value(output); |
39 | } |
40 | |
41 | // Returns FuncOp if `op` is a callable. |
42 | absl::optional<mlir::func::FuncOp> MaybeFindFunction(mlir::Operation* op) { |
43 | auto call_op = llvm::dyn_cast<mlir::CallOpInterface>(op); |
44 | if (!call_op) return absl::nullopt; |
45 | |
46 | mlir::CallInterfaceCallable callable = call_op.getCallableForCallee(); |
47 | mlir::SymbolRefAttr sym = callable.dyn_cast<mlir::SymbolRefAttr>(); |
48 | if (!sym) return absl::nullopt; |
49 | |
50 | mlir::func::FuncOp func = llvm::dyn_cast<mlir::func::FuncOp>( |
51 | mlir::SymbolTable::lookupNearestSymbolFrom(op, sym)); |
52 | if (!func) return absl::nullopt; |
53 | |
54 | return func; |
55 | } |
56 | |
57 | void RemoveDTensorLayoutOp(mlir::TF::DTensorLayout layout) { |
58 | layout.output().replaceAllUsesWith(layout.input()); |
59 | layout.erase(); |
60 | } |
61 | |
62 | } // namespace dtensor |
63 | } // namespace tensorflow |
64 | |