98 lines
3.3 KiB
C
98 lines
3.3 KiB
C
|
/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
|
||
|
|
||
|
Licensed under the Apache License, Version 2.0 (the "License");
|
||
|
you may not use this file except in compliance with the License.
|
||
|
You may obtain a copy of the License at
|
||
|
|
||
|
http://www.apache.org/licenses/LICENSE-2.0
|
||
|
|
||
|
Unless required by applicable law or agreed to in writing, software
|
||
|
distributed under the License is distributed on an "AS IS" BASIS,
|
||
|
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||
|
See the License for the specific language governing permissions and
|
||
|
limitations under the License.
|
||
|
==============================================================================*/
|
||
|
|
||
|
// This file defines the attributes used in the TensorFlow dialect.
|
||
|
|
||
|
#ifndef TENSORFLOW_COMPILER_MLIR_TENSORFLOW_IR_TF_ATTRIBUTES_H_
|
||
|
#define TENSORFLOW_COMPILER_MLIR_TENSORFLOW_IR_TF_ATTRIBUTES_H_
|
||
|
|
||
|
#include "llvm/ADT/StringRef.h"
|
||
|
#include "mlir/IR/BuiltinAttributes.h" // from @llvm-project
|
||
|
#include "mlir/IR/BuiltinTypes.h" // from @llvm-project
|
||
|
#include "mlir/IR/MLIRContext.h" // from @llvm-project
|
||
|
|
||
|
namespace mlir {
|
||
|
namespace TF {
|
||
|
|
||
|
namespace detail {
|
||
|
|
||
|
struct ShapeAttrStorage;
|
||
|
struct FuncAttrStorage;
|
||
|
|
||
|
} // namespace detail
|
||
|
|
||
|
class ShapeAttr : public Attribute::AttrBase<ShapeAttr, Attribute,
|
||
|
detail::ShapeAttrStorage>
|
||
|
{
|
||
|
public:
|
||
|
using Base::Base;
|
||
|
|
||
|
// Get or create a shape attribute. If shape is llvm::None, then it is
|
||
|
// unranked. Otherwise it is ranked. And for ranked shapes, the value of the
|
||
|
// dimension size must be >= -1. The value of -1 means the dimension is
|
||
|
// dynamic. Otherwise, the dimension is static.
|
||
|
static ShapeAttr get(mlir::MLIRContext* context,
|
||
|
llvm::Optional<ArrayRef<int64_t> > shape);
|
||
|
|
||
|
// Get or create a shape attribute from a ShapedType type.
|
||
|
static ShapeAttr get(mlir::MLIRContext* context, ShapedType shaped_type);
|
||
|
|
||
|
llvm::Optional<ArrayRef<int64_t> > getValue() const;
|
||
|
|
||
|
bool hasRank() const;
|
||
|
|
||
|
// If this is ranked, return the rank. Otherwise, abort.
|
||
|
int64_t getRank() const;
|
||
|
|
||
|
// If this is ranked, return the shape. Otherwise, abort.
|
||
|
ArrayRef<int64_t> getShape() const;
|
||
|
|
||
|
// If this is unranked type or any dimension has unknown size (<0), it doesn't
|
||
|
// have static shape. If all dimensions have known size (>= 0), it has static
|
||
|
// shape.
|
||
|
bool hasStaticShape() const;
|
||
|
};
|
||
|
|
||
|
// Custom attribute to model AttrValue.value.func (NameAttrList type attribute).
|
||
|
// This attribute holds a SymbolRefAttr, for the NameAttrList.name string and a
|
||
|
// DictionaryAttr for the NameAttrList.attr map<string, AttrValue>. It is
|
||
|
// currently printed and parsed for the following format:
|
||
|
//
|
||
|
// #tf.func<@symbol, {attr = "value"}>
|
||
|
//
|
||
|
// where the first element is the SymbolRefAttr and the second element is the
|
||
|
// DictionaryAttr.
|
||
|
class FuncAttr
|
||
|
: public Attribute::AttrBase<FuncAttr, Attribute, detail::FuncAttrStorage>
|
||
|
{
|
||
|
public:
|
||
|
using Base::Base;
|
||
|
|
||
|
static FuncAttr get(mlir::MLIRContext* context, llvm::StringRef name,
|
||
|
DictionaryAttr attr);
|
||
|
|
||
|
static FuncAttr get(mlir::MLIRContext* context, SymbolRefAttr symbol,
|
||
|
DictionaryAttr attr);
|
||
|
|
||
|
SymbolRefAttr GetName() const;
|
||
|
|
||
|
DictionaryAttr GetAttrs() const;
|
||
|
};
|
||
|
|
||
|
} // namespace TF
|
||
|
} // namespace mlir
|
||
|
|
||
|
#endif // TENSORFLOW_COMPILER_MLIR_TENSORFLOW_IR_TF_ATTRIBUTES_H_
|