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_DTENSOR_MLIR_LAYOUT_PARSING_H_ 17 #define TENSORFLOW_DTENSOR_MLIR_LAYOUT_PARSING_H_ 18 19 #include <string> 20 21 #include "absl/container/flat_hash_map.h" 22 #include "absl/container/flat_hash_set.h" 23 #include "mlir/Dialect/Func/IR/FuncOps.h" // from @llvm-project 24 #include "mlir/IR/BuiltinOps.h" // from @llvm-project 25 #include "mlir/IR/Operation.h" // from @llvm-project 26 #include "tensorflow/core/platform/status.h" 27 #include "tensorflow/dtensor/cc/dstatus.h" 28 #include "tensorflow/dtensor/cc/tensor_layout.h" 29 #include "tensorflow/dtensor/proto/layout.pb.h" 30 #include "tensorflow/stream_executor/lib/statusor.h" 31 32 namespace tensorflow { 33 namespace dtensor { 34 35 // Extracts `_layout` attribute from `op` and assert a single layout. 36 StatusOr<absl::optional<Layout>> ExtractSingleLayoutFromOp(mlir::Operation* op); 37 38 // Extracts `_layout` attribute from `op`, and returns an error is the layout 39 // is missing. 40 StatusOr<Layout> ExtractRequiredSingleLayoutFromOp(mlir::Operation* op); 41 42 // Extracts `_layout` attribute from `op` and assert a single layout. 43 StatusOr<absl::optional<Layout>> ExtractSingleLayoutFromOp( 44 mlir::Operation* op, std::string attr_name); 45 46 // Extracts `_layout` attribute from `op`. 47 StatusOr<std::vector<absl::optional<Layout>>> ExtractLayoutFromOp( 48 mlir::Operation* op); 49 50 // Extracts `_layout` attribute from `op` and returns an error if any are 51 // missing. 52 StatusOr<std::vector<Layout>> ExtractRequiredLayoutFromOp(mlir::Operation* op); 53 54 // Extract and deserialize a tensor layout from `attr_name`. 55 StatusOr<std::vector<absl::optional<Layout>>> ExtractLayoutFromOp( 56 mlir::Operation* op, std::string attr_name); 57 58 // Extracts '_layout' attribute from `operand`. 59 StatusOr<absl::optional<Layout>> ExtractLayoutFromOperand(mlir::Value operand); 60 61 // Extracts '_layout' attribute from `operand` and returns an error if missing. 62 StatusOr<Layout> ExtractRequiredLayoutFromOperand(mlir::Value operand); 63 64 // Extracts `_layout` attribute from `op`'s operands and returns an error if 65 // any are missing. 66 StatusOr<std::vector<Layout>> ExtractRequiredLayoutFromOperands( 67 mlir::Operation* op); 68 69 // Set `_layout` attribute for op. For layouts without value, an empty string is 70 // used as place holder. 71 void SetLayoutOnOp(mlir::Operation* op, mlir::OpBuilder builder, 72 absl::Span<const absl::optional<Layout>> layouts); 73 74 void SetLayoutOnOp(mlir::Operation* op, 75 absl::Span<const absl::optional<Layout>> layouts); 76 77 void SetSingleLayoutOnOp(mlir::Operation* op, const Layout& layout); 78 79 // Extracts device mesh configuration from op's enclosing tf_device.Cluster op. 80 StatusOr<Mesh> ExtractDeviceMeshEnclosingCluster(mlir::Operation* op); 81 82 // Extracts device mesh configuration from op's `_mesh` attribute. 83 StatusOr<absl::optional<Mesh>> ExtractDeviceMeshFromOp(mlir::Operation* op); 84 85 // Extracts default layout information from function return attribute. 86 StatusOr<absl::optional<Layout>> ExtractLayoutFromFunctionReturnAttr( 87 mlir::func::ReturnOp return_op, const int return_index); 88 89 } // namespace dtensor 90 } // namespace tensorflow 91 92 #endif // TENSORFLOW_DTENSOR_MLIR_LAYOUT_PARSING_H_ 93