include "mlir/IR/PatternBase.td"
include "tensorflow/compiler/mlir/quantization/common/attrs_and_constraints.td"

// Combines the two variadic arguments ($in_tensors and $captured_tensors).
def GetBatchFunctionOpArgOperands:
    NativeCodeCall<"cast<TF::BatchFunctionOp>($0[0].getDefiningOp()).getArgOperands()">;

// Replaces `TF_BatchFunctionOp` into `TF_PartitionedCallOp` that calls the
// same $f. This may be required, for example, when inlining is desired,

include "mlir/IR/OpBase.td"
include "mlir/IR/PatternBase.td"
include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td"

// TODO(jpienaar): Move this somewhere general.
class GetI64ScalarElementsAttr<int value> :
  NativeCodeCall<"DenseElementsAttr::get<int64_t>(RankedTensorType::get({}, $_builder.getIntegerType(64)), " # value # ")">;

def FuseMapAndBatch : Pat<
  (TF_BatchDatasetV2Op

                $axis),
  [(F32ElementsVal $output)]>;

def Get1DShapeValue: NativeCodeCall<"Get1DShapeValue($_builder, $0)">;

class GetIthValue<int index> : NativeCodeCall<"$0[" # index # "]">;

def GetEmbeddingLookupShape: NativeCodeCall<"GetEmbeddingLookupShape($0, $1)">;

// Replace
//   OneHot(index, depth, on=1.0f, off=0.0f, axis=-1) * filter
// With

include "mlir/IR/PatternBase.td"
include "mlir/Dialect/Arith/IR/ArithOps.td"
include "mlir/Dialect/Func/IR/FuncOps.td"
include "tensorflow/compiler/mlir/tfr/ir/tfr_ops.td"

class Quantize<string value> : NativeCodeCall<"TFR::Quantize(" # value # ", $0, $1, $_builder)">;

class HasStringAttr<string value> : AttrConstraint<
    CPred<"$_self.cast<StringAttr>().getValue() == \"" # value # "\"">>;

def QuantActRangeNonePattern :

include "tensorflow/compiler/mlir/lite/ir/tfl_ops.td"

// Quantize attribute $0 by using quantization parameter from %1.
def QuantizeByQuantizedType : NativeCodeCall<"quant::Quantize($0, $1.getValue())">;
def F32ElementsAttr : ElementsAttrBase<
  CPred<"$_self.cast<ElementsAttr>().getShapedType().getElementType().isF32()">, "float constant tensor">;

include "mlir/Dialect/Func/IR/FuncOps.td"
include "mlir/Dialect/Arith/IR/ArithOps.td"
include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td"

// Marks the op as no fallback.
def MarkNoFallback : NativeCodeCall<"SetNoFallbackAttr($_builder, $0)">;

def NoFallbackAttrNotSet : Constraint<CPred<
  "!$0.getDefiningOp()->hasAttr(kNoFallbackAttr)">>;

class FloatValueEquals<string val> : Constraint<CPred<

include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.td"

def IsDataFormatNHWC : ConstantAttr<TF_ConvnetDataFormatAttr, "\"NHWC\"">;

// Get the last dimension size as a 1-d single element attr.
def GetLastDimSizeAsI32 : NativeCodeCall<
  "DenseElementsAttr::get(RankedTensorType::get({1}, $_builder.getIntegerType(32)), "
  "static_cast<int32_t>($0.getType().cast<RankedTensorType>().getDimSize(  "
  "  $0.getType().cast<RankedTensorType>().getRank() - 1)))">;