// StableHLO+TFL --> VHLO+TFL Ops
//===----------------------------------------------------------------------===//

// Wrap op result uses in an unrealized cast to create a cast to buffer
// any type changes to result, and apply type converter to result:
//   result = op(V0)
//   V1     = op2(result)
//   ==>
//   result = op(V0)
//   V1     = unrealized_cast(result)

  StringRef caller_name =
      mlir::cast<NameLoc>(callsite_loc.getCaller()).getName().strref();
  return caller_name.starts_with(kQuantizationUnitPrefix) &&
         caller_name.ends_with(kQuantizationUnitSuffix);
}

std::optional<QuantizationUnitLoc::QuantizationUnit>
FindQuantizationUnitFromLoc(Location loc) {
  if (isa<QuantizationUnitLoc>(loc)) {
    Location caller = mlir::cast<CallSiteLoc>(loc).getCaller();

  // etc.
  constexpr int64_t kLookupTableFindCostScale = 8;
  constexpr int64_t kLookupTableFindStringKeyCostScale = 16;

  auto value_type = mlir::cast<mlir::TensorType>(op.getValues().getType());
  auto key_type = mlir::cast<mlir::TensorType>(op.getKeys().getType());

  int64_t output_size = InferTensorSize(context, value_type);

  int64_t cost = kLookupTableFindCostScale * output_size;

def LegalizeTopKV2 : Pat<(TF_TopKV2Op $input, $k, $ignored_sorted),
                         (TFL_TopKV2Op $input, $k)>;

def ReductionDimensionIsLastDim : Constraint<CPred<"($0.cast<IntegerAttr>().getInt() == "
  "$1.getType().cast<ShapedType>().getRank() - 1 || $0.cast<IntegerAttr>().getInt() == -1)">>;

// Legalizes TF_ApproxTopKOp to TFL_TopKV2Op with the following constraints:
//    1. It computes max k

    return nullptr;
  }
  DenseElementsAttr quantDenseAttr =
      convertDenseFPElementsAttr(mlir::cast<DenseFPElementsAttr>(realDenseAttr),
                                 quantizedElementType, converter);
  if (!quantDenseAttr) {
    return nullptr;
  }

  // Cast from an expressed-type-based type to storage-type-based type,
  // preserving the sparse shape (i.e. tensor<4xf32> -> tensor<4xi8>).

    if (!mlir::cast<ShapedType>(conv_op.getRhs().getType()).hasStaticShape()) {
      return false;
    }
    if (!mlir::cast<ShapedType>(conv_op.getLhs().getType()).hasStaticShape() &&
        !mlir::cast<ShapedType>(conv_op.getType()).hasStaticShape()) {
      auto dnums = conv_op.getDimensionNumbers();
      auto lhs_type = mlir::cast<ShapedType>(conv_op.getLhs().getType());

// TFR function call defines the signatures of the TF ops.
//
namespace mlir {
namespace TFR {

namespace {

// This pattern is to rewrite the "tfr.call" op and the "tfr.cast" ops on the
// operands by a TF op with "tfr.cast" ops on the results. The result type of
// the new TF op is an unranked tensor with element type derived.
class RewriteTFRCallOp : public OpRewritePattern<CallOp> {

      /*effective_hidden_scale_intermediate=*/mlir::TypeAttr());

  // Cast the static shaped lstm result to FuncOp's signature -
  // Ranked but unknown 2nd dimension to support stacking these.
  SmallVector<int64_t, 2> func_output_shape = {1, tensorflow::kTFDynamicSize};
  auto func_result_type = tensorflow::GetTypeFromTFTensorShape(
      func_output_shape,
      mlir::cast<RankedTensorType>(input_.getType()).getElementType());

LogicalResult VerifyTensorTypesCompatible(Type t1, Type t2) {
  if (!mlir::isa<TensorType>(t1) || !mlir::isa<TensorType>(t2)) {
    return failure();
  }
  return verifyCompatibleShape(mlir::cast<TensorType>(t1),
                               mlir::cast<TensorType>(t2));
}

LogicalResult GlobalTensorOp::verify() {
  GlobalTensorOp global_tensor = *this;
  if (global_tensor.getValue()) {
    if (failed(VerifyTensorTypesCompatible(

    // This is the argument used to refer to the pass in
    // the textual format (on the commandline for example).
    return "quant-replace-cast-hacks-with-tf-xla-ops";
  }
  StringRef getDescription() const final {
    // This is a brief description of the pass.
    return "Replace mixed-type Conv and Matmul cast hacks with TF XLA ops.";
  }

  void runOnOperation() override;
};

// Generates params for the XLA Convolution op.