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<

/// TODO(b/130756570): Support OpBase constraints in PatternRewrites.
def SingleResultAndOperandHaveSameElementType : Constraint<
  CPred<"getElementTypeOrSelf($0) == getElementTypeOrSelf($1)">>;

def SingleResultAndOperandHaveSameType : Constraint<
  CPred<"$0.getType() == $1.getType()">>;

def IsRank2Tensor : Type<HasAnyRankOfPred<[2]>, "Rank 2 tensor">;

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

//===----------------------------------------------------------------------===//

// The base class of a quantized type.
class QuantizedType<string n, list<int> params, bit signed>
  : Type<And<[CPred<"$_self.isa<mlir::quant::QuantizedType>()">,
              CPred<"$_self.cast<mlir::quant::QuantizedType>()" #
                    ".getStorageTypeIntegralWidth() == " # !head(params)>]>,
    "Q" # !if (signed, "I", "UI") # !head(params) # " type"> {

                    (variadic $input),
                    ConstantStrAttr<StrAttr, "aten.hardswish.default">, $_, $_, $_),
                    (TFL_HardSwishOp $input)>;

def IsNchwLayoutOp: Constraint<CPred<
  "$0.get(\"is_nchw_op\") && $0.get(\"is_nchw_op\").dyn_cast<BoolAttr>() "
      "== mlir::BoolAttr::get($_builder.getContext(), true)">>;

def IsNhwcLayoutOp: Constraint<Neg<IsNchwLayoutOp.predicate>>;