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

def HasSupportedElementType :
  Constraint<CPred<"HasSupportedElementType($0.getDefiningOp())">>;

def IsSupportedElementType :
  Constraint<CPred<"IsSupportedElementType($0.getType().cast<ShapedType>())">>;

def LegalizeVarHandle : Pat<
  (TF_VarHandleOp:$result $container, $shared_name),

class HasEqualElementSize<list<int> shape_1, list<int> shape_2> : Constraint<
  CPred<"quant::HasEqualElementSize($0, $1,"
  "llvm::ArrayRef<int>({" # !interleave(shape_1, ", ") # "}),"
  "llvm::ArrayRef<int>({" # !interleave(shape_2, ", ") # "}))">,
  "Checks if the given dimensions contain the same number of elements.">;

def ReshapableTo1DTensor : Constraint<
  CPred<"quant::ReshapableTo1DTensor($0.getType().cast<ShapedType>())">,

def IsNotInLiftedFunc :
      Constraint<CPred<"!IsInLiftedFunc($0.getDefiningOp())">>;

// Checks if the value is not inside a StableHLO op with region.
def IsNotInStableHloOpRegion :
      Constraint<CPred<"!IsInStableHloOpRegion($0.getDefiningOp())">>;

// Checks if the given einsum op is supported for XlaDotV2 quantization.
def IsEinsumSupportedByXlaDotV2 :
      Constraint<CPred<"IsEinsumSupportedByXlaDotV2($0)">>;

class IsBoolAttrEqual<string true_or_false> : Constraint<CPred<
  "$0.getValue() == "#true_or_false#"">>;

// Receives a composite DictionaryAttr as an argument and checks if one of the
// its attributes (with the name `attr_name`) is of type `attribute` and has
// the value `val`.
class IsCompositeAttribute<string attr_name, Attr attribute, string val>:
  Constraint<CPred<

limitations under the License.
==============================================================================*/
include "stablehlo/dialect/StablehloOps.td"

class IsStringAttrOf<string value> : Constraint<
  CPred<"::llvm::isa_and_nonnull<StringAttr>($_self) && $_self.cast<StringAttr>().getValue() == \"" # value # "\"">,
  "Is a string attribute whose value is \"" # value # "\""
>;

include "mlir/IR/CommonTypeConstraints.td"

def CloneF32ElementsAttrWithOnes
  : NativeCodeCall<"DenseElementsAttr::get($0.getType().cast<ShapedType>(), (float)1.0)">;

def NotConstant : Constraint<
    CPred<"$0.isa<BlockArgument>() || !llvm::isa<stablehlo::ConstantOp>($0.getDefiningOp())">,
    "Is not a constant.">;

def : Pat<(StableHLO_DivOp $l,
            (StableHLO_ConstantOp:$divisor FloatElementsAttr<32>:$cst)),

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 :
  Pattern<
    (TFR_TFRQuantActRangeOp
     (TFR_ConstOp HasStringAttr<"NONE">:$act), $scale, $zp),

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

def NotFromDequant : Constraint<CPred<"NotFromDequant($0)">>;

// Fuses TFL_FullyConnectedOp and TFL_TransposeOp Rhs to TFL_BatchMatMulOp when
// it's used by TFL_BatchMatMulOp and "transpose_lhs" is true.
def FuseTransposeFCLhsToBatchMatmul : Pat<

// 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">;

// Squash tfl.dequantize and tfl.quantize pairs.
// TODO(fengliuai): Compare the scale of input and output. This can also be

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

// Check to see if the tensor dimensions can be Squeezed by eliminating 1s'
def CanSqueezeTensor : Constraint<CPred<
  "GetShape($0).getNumElements() > GetSqueezedShape($0).getNumElements()">>;


// Pattern to convert TFL_TransposeOp with rank>6 to rank<=6 if there are