#include "tensorflow/compiler/mlir/tensorflow/ir/tf_ops.h"

namespace mlir {
namespace quant {

// Returns the size in bytes of the underlying data of `const_op`. If the
// underlying type's size cannot be determined, it assumes 4 bytes per element.
int64_t GetSizeInBytes(TF::ConstOp const_op);

}  // namespace quant
}  // namespace mlir

// CHECK: "tf._TfrtSetResource"([[x]]) <{index = 0 : i64}> {device = "/CPU:0"} : (tensor<i32>) -> ()
// CHECK: [[const_1:%.*]] = "tf.Const"() <{value = dense<1> : tensor<i32>}> {device = "/CPU:0"} : () -> tensor<i32>
// CHECK: "tf._TfrtSetResource"([[const_1]]) <{index = 1 : i64}> {device = "/CPU:0"} : (tensor<i32>) -> ()

// CHECK-LABEL: func @test_hoist_const

  %2 = stablehlo.add %1, %0 : tensor<1x4x3x3xf32>
  return %2 : tensor<1x4x3x3xf32>
}
// CHECK-SAME: (%[[ARG_0:.+]]: tensor<1x3x3x4xf32>) -> tensor<1x4x3x3xf32>
// CHECK-DAG: %[[CONST_0:.+]] = stablehlo.constant
// CHECK-DAG: %[[TRANSPOSE_0:.+]] = stablehlo.transpose %[[CONST_0]], dims = [0, 2, 3, 1] : (tensor<1x4x3x3xf32>) -> tensor<1x3x3x4xf32>

// Check that the shape of the add is changed to reflect the deferred transpose.

    Value operand = op.getOperand();
    auto const_op =
        dyn_cast_or_null<mlir::stablehlo::ConstantOp>(operand.getDefiningOp());
    if (!const_op) return failure();

    // Only support float tensors.
    auto tensor_type = mlir::dyn_cast_or_null<TensorType>(const_op.getType());
    if (!tensor_type || !tensor_type.getElementType().isF32()) {
      return failure();
    }

      llvm::cl::init(false)};
};

void CopyConstantIntoFunc(int argument_index, Operation* const_op,
                          func::FuncOp func) {
  assert(
      (llvm::isa<TFL::ConstOp, TFL::QConstOp, arith::ConstantOp>(const_op)) &&
      "Expect QConst or Const op.");
  OpBuilder builder(func.getBody());
  auto cloned_const_op = const_op->clone();
  cloned_const_op->setLoc(func.getBody().getLoc());
  builder.insert(cloned_const_op);

    }
    if (!quant_type) {
      const_op->emitError("Failed to get quantized type");
      return failure();
    }

    // TODO(b/172517537): duplicate the constant when the bias is shared.
    Type expressed_type = const_op->getResult(0).getType();
    Type cast_type = quant_type.castFromExpressedType(expressed_type);
    rewriter.setInsertionPointAfter(const_op);

        "saved to checkpoint.");
    return "";
  }

  auto assigned_value_operand = assign_var_op.getOperand(1);
  auto const_op =
      llvm::dyn_cast<mlir::TF::ConstOp>(assigned_value_operand.getDefiningOp());
  if (!const_op) {
    assign_var_op->emitRemark(
        "Operand idx 1 is not a tf.ConstOp. The initializing tensor is not "
        "saved to checkpoint.");
    return "";
  }

	MOVD	p256const0<>(SB), const0
	MOVD	p256const1<>(SB), const1
	EOR	t2<<1, hlp1

	// Negate y2in based on sign
	LDP	2*16(b_ptr), (y0, y1)
	LDP	3*16(b_ptr), (y2, y3)
	MOVD	$-1, acc0

	SUBS	y0, acc0, acc0
	SBCS	y1, const0, acc1
	SBCS	y2, ZR, acc2
	SBCS	y3, const1, acc3
	SBC	$0, ZR, t0

	ADDS	$-1, acc0, acc4
	ADCS	const0, acc1, acc5
	ADCS	$0, acc2, acc6

        "//tensorflow/core:testlib",
        "@com_google_absl//absl/synchronization",
        "@eigen_archive//:eigen3",
    ],
)

cc_library_with_android_deps(
    name = "const_op",
    srcs = ["ops/const_op.cc"],
    hdrs = ["ops/const_op.h"],
    android_deps = [
        "//tensorflow/core:portable_tensorflow_lib",
    ],
    common_deps = [
        ":ops",
        ":scope",
    ],
    deps = [

  const LogicalResult parse_result =
      parseSourceString(const_op_str, &block, ParserConfig(&ctx));
  EXPECT_TRUE(succeeded(parse_result));

  auto const_op = dyn_cast_or_null<TF::ConstOp>(block.front());
  EXPECT_TRUE(const_op);

  return const_op;
}

TEST_F(GetSizeInBytesTest, Int32ScalarConstOpSizeInBytes) {
  constexpr absl::string_view kConstOpExpr =