namespace tac {
namespace {

// This pass is used to fold tfl.const ops to each subgraph (func::FuncOp):
// See the example below:
//
// In main:
// %0 = tfl.const...
// %1 = tfl.const...
// %2 = call func_1(..., %0,...)
// %3 = call func_2(..., %0, ..., %1...)
// ...
//
// Then those consts will be copied into each function and replace their usage.
// func_1:
//   %0 = tfl.const...
// func_2:
//   %0 = tfl.const...

// CHECK-DAG: %[[CONST_2:.+]] = stablehlo.constant dense<3> : tensor<1xi32>
// CHECK-DAG: %[[CONST_3:.+]] = stablehlo.constant dense<4> : tensor<1xi32>
// CHECK-DAG: %[[CONST_4:.+]] = stablehlo.constant dense<2> : tensor<1xi32>

// Populates a TPUCompileMetadataProto from attributes of a
// `tf_device::ClusterFuncOp`. If any necessary attributes are missing from the
// op, a failure will be returned.
// TODO(lyandy): Support session handle and guaranteed consts.
LogicalResult SetMetadataProtoFromClusterFuncOp(
    tf_device::ClusterFuncOp op, int num_replicas, int num_cores_per_replica,
    std::optional<xla::DeviceAssignmentProto>&& xla_device_assignment,

class QuantizeConstWeights : public OpRewritePattern<TF::ConstOp> {
 public:
  explicit QuantizeConstWeights(
      MLIRContext* context,
      const tensorflow::quantization::QuantizationOptions& quantization_options)
      : OpRewritePattern<TF::ConstOp>(context),
        quant_options_(quantization_options) {}

  LogicalResult matchAndRewrite(TF::ConstOp op,

	VREV64	V14.B16, V14.B16
	VREV64	V15.B16, V15.B16
	VREV64	V16.B16, V16.B16
	VREV64	V17.B16, V17.B16
	VREV64	V18.B16, V18.B16
	VREV64	V19.B16, V19.B16

	MOVD	R3, R4
	// load first 4 round consts in V20-V23
	VLD1.P	64(R4), [V20.D2, V21.D2, V22.D2, V23.D2]

	SHA512ROUND(V0, V1, V2, V3, V4, V20, V24, V12, V13, V19, V16, V17)
	SHA512ROUND(V3, V0, V4, V2, V1, V21, V25, V13, V14, V12, V17, V18)

	flagVariants := make(map[string]bool)

	for _, lib := range ctxt.Library {

		consts := d.ldr.Lookup(dwarf.ConstInfoPrefix+lib.Pkg, 0)
		for _, unit := range lib.Units {
			// We drop the constants into the first CU.
			if consts != 0 {
				unit.Consts = sym.LoaderSym(consts)
				d.importInfoSymbol(consts)
				consts = 0
			}
			ctxt.compUnits = append(ctxt.compUnits, unit)

DATA consts<>+0x48(SB)/8, $0x0102030005060704
DATA consts<>+0x50(SB)/8, $0x6170786561707865
DATA consts<>+0x58(SB)/8, $0x6170786561707865
DATA consts<>+0x60(SB)/8, $0x3320646e3320646e
DATA consts<>+0x68(SB)/8, $0x3320646e3320646e
DATA consts<>+0x70(SB)/8, $0x79622d3279622d32
DATA consts<>+0x78(SB)/8, $0x79622d3279622d32
DATA consts<>+0x80(SB)/8, $0x6b2065746b206574
DATA consts<>+0x88(SB)/8, $0x6b2065746b206574

// Creates a pass that creates a RestoreV2 op in the initializer function with
// type "restore_op" that initializes variables from the checkpoint. It finds
// tf.AssignVariableOp(tf.VarHandleOp, tf.Const) patterns in the initializer
// function and replaces tf.Consts with the results of RestoreV2.
std::unique_ptr<OperationPass<ModuleOp>> CreateInsertRestoreOpPass();

// Creates a pass that creates a new function that wraps the newly created

StatusOrOptional<SliceInputs> GetSliceInputs(Node* slice) {
  const int kSliceInputIndex = 0;
  const int kSliceBeginIndex = 1;
  const int kSliceSizeIndex = 2;

  const Edge* slice_input_edge;
  TF_RETURN_IF_ERROR(slice->input_edge(kSliceInputIndex, &slice_input_edge));
  const Edge* slice_size_edge;
  TF_RETURN_IF_ERROR(slice->input_edge(kSliceSizeIndex, &slice_size_edge));
  const Edge* slice_begin_edge;

(MUL(Q|L)const [-3] x) => (NEG(Q|L) (LEA(Q|L)2 <v.Type> x x))
(MUL(Q|L)const [-1] x) => (NEG(Q|L) x)
(MUL(Q|L)const [ 0] _) => (MOV(Q|L)const [0])
(MUL(Q|L)const [ 1] x) => x
(MUL(Q|L)const [ 3] x) => (LEA(Q|L)2 x x)
(MUL(Q|L)const [ 5] x) => (LEA(Q|L)4 x x)
(MUL(Q|L)const [ 7] x) => (LEA(Q|L)2 x (LEA(Q|L)2 <v.Type> x x))
(MUL(Q|L)const [ 9] x) => (LEA(Q|L)8 x x)