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 file defines attributes for DTensor. 17 18 #ifndef TENSORFLOW_DTENSOR_MLIR_DTENSOR_DIALECT_IR_DTENSOR_ATTRIBUTES_H_ 19 #define TENSORFLOW_DTENSOR_MLIR_DTENSOR_DIALECT_IR_DTENSOR_ATTRIBUTES_H_ 20 21 #include "mlir/IR/Attributes.h" // from @llvm-project 22 #include "tensorflow/dtensor/cc/tensor_layout.h" 23 #include "tensorflow/dtensor/proto/layout.pb.h" 24 25 namespace mlir { 26 namespace dtensor { 27 28 namespace detail { 29 struct LayoutAttrStorage; 30 struct MeshAttrStorage; 31 } // namespace detail 32 33 // Attribute to keep track of a mesh. 34 class MeshAttr 35 : public Attribute::AttrBase<MeshAttr, Attribute, detail::MeshAttrStorage> { 36 public: 37 using Base::Base; 38 using Mesh = tensorflow::dtensor::Mesh; 39 40 // Constructor of attribute 41 static MeshAttr get(MLIRContext* context, const Mesh& mesh); 42 43 // Returns Mesh 44 const Mesh& getValue() const; 45 }; 46 47 // Custom attribute to keep track of dtensor layouts. 48 class LayoutAttr : public Attribute::AttrBase<LayoutAttr, Attribute, 49 detail::LayoutAttrStorage> { 50 public: 51 using Base::Base; 52 using Layout = tensorflow::dtensor::Layout; 53 using Mesh = tensorflow::dtensor::Mesh; 54 55 // Create a layout attribute. 56 static LayoutAttr get(MLIRContext* context, Layout layout); 57 58 // Get layout. 59 const Layout& getValue() const; 60 }; 61 62 } // namespace dtensor 63 } // namespace mlir 64 65 #endif // TENSORFLOW_DTENSOR_MLIR_DTENSOR_DIALECT_IR_DTENSOR_ATTRIBUTES_H_ 66