chore: import upstream snapshot with attribution

tensorflow/compiler/mlir/tfr/ir/tfr_ops.h

@@ -0,0 +1,66 @@
+/* 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_TFR_IR_TFR_OPS_H_
+#define TENSORFLOW_COMPILER_MLIR_TFR_IR_TFR_OPS_H_
+
+#include "llvm/ADT/StringSet.h"
+#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
+#include "mlir/Dialect/Shape/IR/Shape.h"  // from @llvm-project
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Dialect.h"  // from @llvm-project
+#include "mlir/IR/DialectImplementation.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+#include "mlir/Interfaces/CallInterfaces.h"  // from @llvm-project
+#include "mlir/Interfaces/ControlFlowInterfaces.h"  // from @llvm-project
+#include "mlir/Interfaces/FunctionInterfaces.h"  // from @llvm-project
+#include "mlir/Interfaces/InferTypeOpInterface.h"  // from @llvm-project
+#include "mlir/Interfaces/SideEffectInterfaces.h"  // from @llvm-project
+#include "mlir/Support/LLVM.h"  // from @llvm-project
+
+namespace mlir {
+namespace TFR {
+
+constexpr char kAttrArgumentNameAttr[] = "tfr.name";
+constexpr char kAttrArgumentDefaultAttr[] = "tfr.default";
+constexpr char kAttrArgumentTypeAttr[] = "tfr.type";
+
+class TFRDialect : public Dialect {
+ public:
+  explicit TFRDialect(MLIRContext *context);
+
+  static StringRef getDialectNamespace() { return "tfr"; }
+
+  Operation *materializeConstant(OpBuilder &builder, Attribute value, Type type,
+                                 Location loc) override;
+
+  // Parse a type registered to this dialect.
+  Type parseType(DialectAsmParser &parser) const override;
+
+  // Prints a type registered to this dialect.
+  void printType(Type ty, DialectAsmPrinter &os) const override;
+};
+
+}  // namespace TFR
+}  // namespace mlir
+
+#define GET_OP_CLASSES
+#include "tensorflow/compiler/mlir/tfr/ir/tfr_ops.h.inc"
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFR_IR_TFR_OPS_H_

tensorflow/compiler/mlir/tfr/ir/tfr_ops.td

@@ -0,0 +1,615 @@
+/* Copyright 2019 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 is the operation definition file for TFR
+
+#ifndef DIALECT_TFR_OPS_
+#define DIALECT_TFR_OPS_
+
+include "mlir/Dialect/Shape/IR/ShapeBase.td"
+include "mlir/IR/OpBase.td"
+include "mlir/Interfaces/FunctionInterfaces.td"
+include "mlir/Dialect/Quant/IR/QuantBase.td"
+include "mlir/IR/SymbolInterfaces.td"
+include "mlir/Interfaces/CallInterfaces.td"
+include "mlir/Interfaces/ControlFlowInterfaces.td"
+include "mlir/Interfaces/InferTypeOpInterface.td"
+include "mlir/Interfaces/SideEffectInterfaces.td"
+include "tensorflow/compiler/mlir/tensorflow/ir/tf_op_base.td"
+
+//===----------------------------------------------------------------------===//
+// Dialect
+//===----------------------------------------------------------------------===//
+
+def TFR_Dialect : Dialect {
+  let name = "tfr";
+
+  let description = [{
+    The TensorFlow Composition dialect.
+  }];
+
+  let cppNamespace = "::mlir::TFR";
+}
+
+//===----------------------------------------------------------------------===//
+// Type classes
+//===----------------------------------------------------------------------===//
+
+// tensor argument types
+class TFR_Type<string name> : DialectType<TFR_Dialect,
+    CPred<"llvm::isa<mlir::TFR::" # name # "Type>($_self)">,
+    "TFR " # name #" type">,
+    BuildableType<"$_builder.getType<mlir::TFR::" # name # "Type>()">;
+def TFR_TensorType : TFR_Type<"TFRTensor">;
+def TFR_TensorListType : TFR_Type<"TFRTensorList">;
+def TFR_AllTensorTypes : Type<Or<[
+    TFR_TensorType.predicate,
+    TFR_TensorListType.predicate]>, "all tensor related types">;
+
+// attribute argument types
+def TFR_AttrType : TFR_Type<"TFRAttr">;
+def TFR_AttrScalarType: TypeAlias<TF_ElementType, "scalar attribute">;
+def TFR_AttrVectorType : VectorOfNonZeroRankOf<[TF_ElementType, TFR_AttrType]>;
+def TFR_AllAttrTypes : Type<Or<[
+    TFR_AttrType.predicate,
+    Index.predicate,
+    TFR_AttrScalarType.predicate,
+    TFR_AttrVectorType.predicate]>, "all attribute related types">;
+
+// all allowed arguments types
+def TFR_allowedArgType : Type<Or<[
+    TFR_AllTensorTypes.predicate,
+    TFR_AllAttrTypes.predicate]>, "allowed tfr.call operand types">;
+
+def TFR_allowedConstValues : Attr<Or<[
+    FlatSymbolRefAttr.predicate,
+    TypeAttr.predicate,
+    StrAttr.predicate,
+    ArrayAttr.predicate]>, "allowed tfr.constant value"> {
+  let storageType = "Attribute";
+  let returnType = "Attribute";
+  let convertFromStorage = "$_self";
+  let constBuilderCall = "$0";
+}
+
+// all allowed result types
+def TFR_allowedResultType : TypeAlias<TFR_AllTensorTypes,
+  "allowed tfr.call result types">;
+
+// standard tensor type and tfr.tensor types can be casted to each other.
+def TFR_singleTensorType : Type<Or<[
+    TFR_TensorType.predicate,
+    TF_Tensor.predicate,
+    TensorOf<[quant_QuantizedType]>.predicate]>, "single tensor or tfr.tensor type">;
+
+// all allowed build list input types
+def TFR_allowedBuiltListType : Type<Or<[
+    TFR_TensorType.predicate,
+    TF_ElementType.predicate,
+    TFR_AttrType.predicate]>, "single tfr.tensor or tensor element type">;
+
+// all allowed build list result types
+def TFR_allowedListResultType : Type<Or<[
+    TFR_TensorListType.predicate,
+    TFR_AttrType.predicate]>, "tfr.tensor_list or tfr.attr type">;
+
+//===----------------------------------------------------------------------===//
+// Op classes
+//===----------------------------------------------------------------------===//
+
+class TFR_Op<string mnemonic, list<Trait> traits> :
+    Op<TFR_Dialect, mnemonic, traits>;
+
+def TFR_CallOp : TFR_Op<"call", [CallOpInterface]> {
+  let description = [{
+    The `call` operation represents a direct call to a function that is within
+    the same symbol scope as the callee. The operands and result types of the
+    call must match the specified function type. The callee is encoded as a
+    symbol reference attribute named "callee".
+
+    Example:
+
+    ```mlir
+    %2 = tfr.call @my_add(%0, %1) : (tfr.tensor, f32) -> tfr.tensor_list
+    ```
+
+    Note that the operands of the `call` operation can only be with tfr.tensor,
+    tfr.tensor_list, tfr.attr and mlir float and integer types. The results of
+    the `call` operation can only be with tfr.tensor and tfr.tensor_list types.
+  }];
+
+  let arguments = (ins
+    FlatSymbolRefAttr:$callee,
+    Variadic<TFR_allowedArgType>:$args,
+    OptionalAttr<DictArrayAttr>:$arg_attrs,
+    OptionalAttr<DictArrayAttr>:$res_attrs);
+
+  let results = (outs
+    Variadic<TFR_allowedResultType>:$outs);
+
+  let extraClassDeclaration = [{
+    // Get the argument operands to the called function.
+    operand_range getArgOperands() { return getArgs(); }
+    MutableOperandRange getArgOperandsMutable() {
+      return getArgsMutable();
+    }
+
+    // Return the callee of this operation.
+    CallInterfaceCallable getCallableForCallee() { return getCalleeAttr(); }
+    // Sets the callee from the callable
+    void setCalleeFromCallable(CallInterfaceCallable callee);
+  }];
+
+  let assemblyFormat = [{
+    $callee `(` $args `)` attr-dict `:` functional-type($args, results)
+  }];
+}
+
+def TFR_CastOp : TFR_Op<"cast", [Pure]> {
+  let description = [{
+    The `cast` operation converts the operand with built-in tensor type to
+    tfr.tensor type, or vice versa.
+
+    Example:
+
+    ```mlir
+    %1 = tfr.cast(%0) : tensor<f32> -> !tfr.tensor
+    %3 = tfr.cast(%1) : !tfr.tensor -> tensor<f32>
+    ```
+  }];
+
+  let arguments = (ins TFR_singleTensorType:$arg);
+
+  let results = (outs TFR_singleTensorType:$out);
+
+  let extraClassDeclaration = [{
+    // Return element type of the input tensor type. Only available when the
+    // input is a MLIR built-in tensor type.
+    Attribute getInputElementType() {
+      if (auto ty = llvm::dyn_cast<TensorType>(getArg().getType())) {
+        return TypeAttr::get(ty.getElementType());
+      }
+      return {};
+    }
+  }];
+
+  let hasCanonicalizer = 1;
+}
+
+def TFR_GetShapeOp : TFR_Op<"get_shape", [Pure]> {
+  let description = [{
+    The `get_shape` operation gets the shape of a tfr.tensor and returns
+    !shape.shape type.
+
+    Example:
+
+    ```mlir
+    %1 = "tfr.get_shape"(%0) : !tfr.tensor -> !shape.shape
+    %1 = tfr.get_shape %0 -> !shape.shape
+    ```
+  }];
+
+  let arguments = (ins TFR_TensorType:$arg);
+
+  let results = (outs Shape_ShapeType:$out);
+
+  let assemblyFormat = "$arg attr-dict `->` type($out)";
+
+  let hasCanonicalizer = 1;
+}
+
+def TFR_GetElementTypeOp : TFR_Op<"get_element_type", [Pure]> {
+  let description = [{
+    The `get_element_type` operation gets the element type of a tfr.tensor and
+    returns !tfr.attr.
+
+    Example:
+
+    ```mlir
+    %1 = "tfr.get_element_type"(%0) : !tfr.tensor -> !tfr.attr
+    %1 = tfr.get_element_type %0 -> !tfr.attr
+    ```
+  }];
+
+  let arguments = (ins TFR_TensorType:$arg);
+
+  let results = (outs TFR_AttrType:$out);
+
+  let assemblyFormat = "$arg attr-dict `->` type($out)";
+}
+
+def TFR_EqualOp : TFR_Op<"equal", [Pure, SameTypeOperands]> {
+  let description = [{
+    The `equal` operation compares the values of the tfr.attr type arguments.
+    The operation returns an i1 boolean indicating if the two values are the
+    same.
+    Example:
+
+    ```mlir
+    %x = tfr.equal %lhs, %rhs -> i1
+    %x = "tfr.equal"(%lhs, %rhs) : (!tfr.attr, !tfr.attr) -> i1
+    ```
+  }];
+
+  let arguments = (ins
+      TFR_AttrType:$lhs,
+      TFR_AttrType:$rhs
+  );
+  let results = (outs BoolLike:$result);
+
+  let hasFolder = 1;
+
+  let assemblyFormat = "$lhs `,` $rhs attr-dict `->` type($result)";
+}
+
+def TFR_ConstOp : TFR_Op<"constant", [ConstantLike, Pure]> {
+  let description = [{
+    The `attr` operation stores TF op's attribute, which doesn't support
+    arithmetic operations.
+
+    Example:
+
+    ```mlir
+    %1 = "tfr.constant"() { value: i32 } : () -> !tfr.attr
+    %2 = "tfr.constant"() { value: [i32, f32] } : () -> !tfr.attr
+    %3 = tfr.constant [i32, f32] -> !tfr.attr
+    %4 = tfr.constant f32 -> !tfr.attr
+    ```
+  }];
+
+  let arguments = (ins TFR_allowedConstValues:$value);
+
+  let results = (outs TFR_AttrType:$out);
+
+  let hasFolder = 1;
+
+  let builders = [
+    OpBuilder<(ins "Attribute":$value),
+    [{
+      auto* ctx = value.getContext();
+      $_state.addAttribute("value", value);
+      $_state.addTypes(TFRAttrType::get(ctx));
+    }]>
+  ];
+
+  let assemblyFormat = [{
+    $value attr-dict `->` type($out)
+  }];
+}
+
+def TFR_ConstantTensorOp : TFR_Op<"constant_tensor", [Pure]> {
+  let description = [{
+    The `constant_tensor` operation converts the operand with non-built-in
+    tensor type to built-in tensor type or tfr.tensor type. If it is built-in
+    tensor type, the shape shouldn't be changed during the conversion.
+
+    Example:
+
+    ```mlir
+    %1 = tfr.constant_tensor(%0) : f32 -> tensor<f32>
+    %3 = tfr.constant_tensor(%2) : vector<1xf32> -> tensor<1xf32>
+    ```
+  }];
+
+  let arguments = (ins TFR_AllAttrTypes:$arg);
+
+  let results = (outs TFR_singleTensorType:$out);
+
+  let hasCanonicalizer = 1;
+
+  let hasVerifier = 1;
+}
+
+def TFR_GetElementOp : TFR_Op<"get_element", [Pure]> {
+  let description = [{
+    The `get_element` operation extracts one tfr.tensor element from a
+    tfr.tensor_list.
+
+    Example:
+
+    ```mlir
+    %2 = tfr.get_element %1[%0] : (tfr.tensor, index) -> tfr.tensor
+    ```
+  }];
+
+  let arguments = (ins
+    TFR_TensorListType:$tensor_list,
+    Index:$index);
+
+  let results = (outs TFR_TensorType:$out);
+
+  let hasCanonicalizer = 1;
+
+  let assemblyFormat = [{
+    $tensor_list `[` $index `]` attr-dict `:`
+      `(` type($tensor_list) `,` type($index) `)` `->` type($out)
+  }];
+}
+
+def TFR_BuildListOp : TFR_Op<"build_list", [Pure]> {
+  let description = [{
+   The `build_list` operation builds a tensor list from a list of tensors, or
+   an tfr.attr from a list of scalars.
+
+    Example:
+
+    ```mlir
+    %3 = tfr.build_list(%2, %1, %0) :
+      (tfr.tensor, tfr.tensor, tfr.tensor) -> tfr.tensor_list
+    %3 = tfr.build_list(%2, %1, %0) : (i32, i32, i32) -> tfr.attr
+    ```
+  }];
+
+  let arguments = (ins Variadic<TFR_allowedBuiltListType>:$tensors);
+
+  let results = (outs TFR_allowedListResultType:$out);
+
+  let hasCanonicalizer = 1;
+}
+
+def TFR_GetLengthOp : TFR_Op<"get_length", [Pure]> {
+  let description = [{
+    The `get_length` operation returns the number of tensors for a
+    tfr.tensor_list.
+
+    Example:
+
+    ```mlir
+    %2 = tfr.get_length(%1) : tfr.tensor -> index
+    %2 = tfr.get_length %1 -> index
+    ```
+  }];
+
+  let arguments = (ins TFR_TensorListType:$tensor_list);
+
+  let results = (outs Index:$out);
+
+  let hasCanonicalizer = 1;
+
+  let assemblyFormat = [{
+    $tensor_list attr-dict `->` type($out)
+  }];
+}
+
+//===----------------------------------------------------------------------===//
+// Function related classes
+//===----------------------------------------------------------------------===//
+
+def TFR_TFRFuncOp : TFR_Op<"func", [HasParent<"ModuleOp">,
+                                    FunctionOpInterface,
+                                    IsolatedFromAbove, Symbol]> {
+  let summary = "TFR Function defines a composition of other ops";
+
+  let description = [{
+    Defines a function that can be used to decompose an TF function call to
+    the invocation of a set of other TF ops.
+
+    Syntax:
+
+    ```
+    op ::= `tfr.func` visibility? symbol-ref-id `(` argument-list `)` (`->`
+    function-result-list)? function-attributes? region
+    ```
+
+    Example:
+
+    ```mlir
+    tfr.func @foo(%arg0: !tfr.tensor, %arg1: !tfr.tensor_list<T>,
+                  %arg2: int {tfr.name="T", tfr.default=1})
+        attributes {qux: "quux"} {
+      tfr.return
+    }
+    ```
+
+    Note the arguments are ordered by the following rule:
+      tfr.tensor > tfr.tensor_list > tfr.attr/i32/...,
+    and only one trfr.tensor_list argument is allowed.
+  }];
+
+  let arguments = (ins
+    TypeAttrOf<FunctionType>:$function_type,
+    StrAttr:$sym_name,
+    OptionalAttr<DictArrayAttr>:$arg_attrs,
+    OptionalAttr<DictArrayAttr>:$res_attrs
+  );
+
+  let results = (outs);
+
+  // When the regions is empty, the tfr.func is an external function and used
+  // to model the element type constraints of the tf op. Otherwise, there is one
+  // region containing the composition.
+  let regions = (region VariadicRegion<AnyRegion>:$body);
+
+  let extraClassDeclaration = [{
+    LogicalResult verifyType() { return success(); }
+
+    mlir::Region *getCallableRegion();
+
+    /// Returns the argument types of this function.
+    ArrayRef<Type> getArgumentTypes() { return getFunctionType().getInputs(); }
+
+    /// Returns the result types of this function.
+    ArrayRef<Type> getResultTypes() { return getFunctionType().getResults(); }
+
+    // Get the names of all defined attributes, including both derived and
+    // non-derived ones.
+    llvm::StringSet<> getDefinedAttributeNames() {
+      llvm::StringSet<> all_attrs;
+      for (auto& attr : (*this)->getAttrs()) {
+        all_attrs.insert(attr.getName().strref());
+      }
+      for (const auto& operand : llvm::enumerate(getArgumentTypes())) {
+        if (auto attr_name = getArgAttrOfType<StringAttr>(
+            operand.index(), kAttrArgumentNameAttr)) {
+          all_attrs.insert(attr_name.getValue());
+        }
+      }
+      return all_attrs;
+    }
+  }];
+
+  let hasVerifier = 1;
+  let hasCustomAssemblyFormat = 1;
+}
+
+def TFR_TFRReturnOp : TFR_Op<"return", [HasParent<"TFRFuncOp">, Pure,
+                                        ReturnLike, Terminator]> {
+  let description = [{
+    A terminator operation for regions that appear in the body of  `tfr.func`
+    functions. The operands to the `tfr.return` are the result values returned
+    by an invocation of the `tfr.func`.
+
+    Note that only the tfr.tensor and tfr.tensor_list can be returned.
+  }];
+
+  let arguments = (ins Variadic<TFR_allowedResultType>:$arguments);
+
+  let assemblyFormat = "attr-dict ($arguments^ `:` type($arguments))?";
+}
+
+//===----------------------------------------------------------------------===//
+// Quantization related operations
+//===----------------------------------------------------------------------===//
+
+def TFR_TFRQuantActRangeOp : TFR_Op<"quant_act_range", [Pure]> {
+  let description = [{
+   The `quant_act_range` returns the a pair of integers to indicate the fixed
+   range for the fused activation `act` with the quantization defined by the
+   `scale` and `zero point`. Currently, the allowed activations are
+   `NONE`, `RELU`, `RELU6` and `RELU_N1_TO_1`.
+
+    Example:
+
+    ```mlir
+    %3, %4 = tfr.quant_act_range(%2, %1, %0) :
+        (tfr.attr, float, i64) -> (tfr.tensor, tfr.tensor)
+    ```
+  }];
+
+  let arguments = (ins
+      TFR_AttrType:$act,
+      F32:$scale,
+      I64:$zp);
+
+  let results = (outs TFR_TensorType:$min, TFR_TensorType:$max);
+
+  let assemblyFormat = [{
+    `(` $act `,` $scale `,` $zp `)` attr-dict `:` functional-type(operands, results)
+  }];
+}
+
+def TFR_TFRQuantRescaleOp : TFR_Op<"quant_rescale", [Pure]> {
+  let description = [{
+   The `quant_rescale` rescales the elements of the integer tensor by the
+   floating-point rescale factor. This op needs to be legalized to the preferred
+   operations of the backends.
+
+    Example:
+
+    ```mlir
+    %3 = tfr.quant_rescale(%2, %1, %0) :
+        (tfr.tensor, tfr.tensor, i64) -> (tfr.tensor)
+    ```
+  }];
+
+  let arguments = (ins
+      TFR_TensorType:$input,
+      TFR_TensorType:$scale,
+      I64:$zp);
+
+  let results = (outs TFR_TensorType:$output);
+
+  let assemblyFormat = [{
+    `(` $input `,` $scale `,` $zp `)` attr-dict `:` functional-type(operands, results)
+  }];
+
+  let hasCanonicalizer = 1;
+}
+
+def TFR_TFRQuantRawDataOp : TFR_Op<"quant_raw_data", [
+    Pure,
+    SameOperandsAndResultType]> {
+  let description = [{
+   The `quant_raw_data` removes the quantization parameter from the intput
+   tensor(s).
+
+    Example:
+
+    ```mlir
+    %3 = tfr.quant_raw_data(%0) : (tfr.tensor) -> (tfr.tensor)
+    ```
+  }];
+
+  let arguments = (ins TFR_AllTensorTypes:$input);
+
+  let results = (outs TFR_AllTensorTypes:$output);
+
+  let assemblyFormat = [{
+    `(` $input `)` attr-dict `:` functional-type($input, results)
+  }];
+
+  let hasCanonicalizer = 1;
+}
+
+def TFR_TFRQuantQParamsOp : TFR_Op<"quant_qparam", [Pure]> {
+  let description = [{
+   The `quant_qparam` returns the quantization parameter of the input
+   tensors.
+
+    Example:
+
+    ```mlir
+    %3 = tfr.quant_qparam(%0) : (tfr.tensor) -> (float, tfr.tensor)
+    ```
+  }];
+
+  let arguments = (ins TFR_TensorType:$input);
+
+  let results = (outs TFR_TensorType:$scale, TFR_TensorType:$zp);
+
+  let assemblyFormat = [{
+    `(` $input `)` attr-dict `:` functional-type($input, results)
+  }];
+
+  let hasCanonicalizer = 1;
+}
+
+
+def TFR_TFRQuantScaleFactorOp : TFR_Op<"quant_scale_factor", [Pure]> {
+  let description = [{
+   The `quant_scale_factor` computes the effective scale factor according to the
+   output scale and input scales.
+
+    Example:
+
+    ```mlir
+    %3 = tfr.quant_scale_factor(%0) : (f32, tfr.tensor_list) -> (tfr.tensor)
+    ```
+  }];
+
+  let arguments = (ins
+      F32:$out_scale,
+      TFR_TensorListType:$in_scales);
+
+  let results = (outs TFR_TensorType:$scale_factor);
+
+  let assemblyFormat = [{
+    `(` $out_scale `,` $in_scales `)` attr-dict `:` functional-type(operands, results)
+  }];
+
+  let hasCanonicalizer = 1;
+}
+
+#endif // DIALECT_TFR_OPS_

tensorflow/compiler/mlir/tfr/ir/tfr_types.h

@@ -0,0 +1,124 @@
+/* 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.
+==============================================================================*/
+
+#ifndef TENSORFLOW_COMPILER_MLIR_TFR_IR_TFR_TYPES_H_
+#define TENSORFLOW_COMPILER_MLIR_TFR_IR_TFR_TYPES_H_
+
+#include <memory>
+#include <string>
+
+#include "mlir/IR/Attributes.h"  // from @llvm-project
+#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
+#include "mlir/IR/Diagnostics.h"  // from @llvm-project
+#include "mlir/IR/Location.h"  // from @llvm-project
+#include "mlir/IR/MLIRContext.h"  // from @llvm-project
+#include "mlir/IR/Operation.h"  // from @llvm-project
+#include "mlir/IR/TypeSupport.h"  // from @llvm-project
+#include "mlir/IR/Types.h"  // from @llvm-project
+
+namespace mlir {
+namespace TFR {
+
+class TFRType : public Type {
+ public:
+  using Type::Type;
+
+  static bool classof(Type type);
+};
+
+namespace detail {
+
+struct TFRTypeStorage final
+    : public TypeStorage,
+      public llvm::TrailingObjects<TFRTypeStorage, StringAttr> {
+  using KeyTy = ArrayRef<StringAttr>;
+
+  explicit TFRTypeStorage(unsigned num_attrs) : num_attrs(num_attrs) {}
+
+  static TFRTypeStorage* construct(TypeStorageAllocator& allocator, KeyTy key) {
+    // Allocate a new storage instance.
+    auto byteSize = TFRTypeStorage::totalSizeToAlloc<StringAttr>(key.size());
+    auto rawMem = allocator.allocate(byteSize, alignof(TFRTypeStorage));
+    auto result = ::new (rawMem) TFRTypeStorage(key.size());
+
+    // Copy in the string attributes into the trailing storage.
+    std::uninitialized_copy(key.begin(), key.end(),
+                            result->getTrailingObjects());
+    return result;
+  }
+
+  bool operator==(const KeyTy& attrs) const { return attrs == GetAttrs(); }
+
+  KeyTy GetAttrs() const { return {getTrailingObjects(), num_attrs}; }
+
+  unsigned num_attrs;
+};
+
+template <typename Derived>
+class TFRTypeImpl : public Type::TypeBase<Derived, TFRType, TFRTypeStorage> {
+ public:
+  using Base = Type::TypeBase<Derived, TFRType, TFRTypeStorage>;
+  using TFRBase = TFRTypeImpl<Derived>;
+  using Base::Base;
+
+  static Derived get(ArrayRef<StringAttr> attrs, MLIRContext* context) {
+    return Base::get(context, attrs);
+  }
+
+  static Derived getChecked(ArrayRef<StringAttr> attrs, Location loc) {
+    return Base::getChecked(loc, loc.getContext(), attrs);
+  }
+  static Derived getChecked(function_ref<InFlightDiagnostic()> emitError,
+                            MLIRContext* context, ArrayRef<StringAttr> attrs) {
+    return Base::getChecked(emitError, context, attrs);
+  }
+
+  static Derived get(MLIRContext* context) { return get({}, context); }
+
+  // TODO(fengliuai): fix the implementation
+  static LogicalResult verify(function_ref<InFlightDiagnostic()> emitError,
+                              ArrayRef<StringAttr> attrs) {
+    return success();
+  }
+
+  ArrayRef<StringAttr> getAttrKeys() { return Base::getImpl()->GetAttrs(); }
+};
+}  // namespace detail
+
+class TFRTensorType : public detail::TFRTypeImpl<TFRTensorType> {
+ public:
+  using TFRBase::TFRBase;
+  static constexpr StringLiteral name = "tfr.tensor";
+  static std::string getTypeName() { return "TFRTensorType"; }
+};
+
+class TFRTensorListType : public detail::TFRTypeImpl<TFRTensorListType> {
+ public:
+  using TFRBase::TFRBase;
+  static constexpr StringLiteral name = "tfr.tensor_list";
+  static std::string getTypeName() { return "TFRTensorListType"; }
+};
+
+class TFRAttrType : public Type::TypeBase<TFRAttrType, TFRType, TypeStorage> {
+ public:
+  using Base::Base;
+  static constexpr StringLiteral name = "tfr.attr";
+  static std::string getTypeName() { return "TFRAttrType"; }
+};
+
+}  // namespace TFR
+}  // namespace mlir
+
+#endif  // TENSORFLOW_COMPILER_MLIR_TFR_IR_TFR_TYPES_H_