See the License for the specific language governing permissions and
limitations under the License.
==============================================================================*/

#include "tensorflow/cc/ops/const_op.h"
#include "tensorflow/core/framework/node_def_util.h"
#include "tensorflow/core/framework/tensor_testutil.h"
#include "tensorflow/core/lib/core/status_test_util.h"
#include "tensorflow/core/platform/test.h"

// `kAssumedNumBytesPerElem` bytes.
int64_t GetSizeOfUnsupportedTypeConst(TF::ConstOp const_op) {
  return kAssumedNumBytesPerElem * const_op.getValue().getNumElements();
}

}  // namespace

int64_t GetSizeInBytes(TF::ConstOp const_op) {
  const Type dtype = const_op.getDtype();

  if (dtype.isIntOrFloat()) {
    return GetSizeOfIntOrFloatConst(const_op);
  } else if (isa<TF::StringType>(dtype)) {

                                        /*resource=*/resource_type,
                                        /*container=*/"", shared_name);

    auto read_variable_op = builder.create<TF::ReadVariableOp>(
        const_op.getLoc(), const_op.getType(), var_handle_op);

    // Replace each usage of ConstOp with the corresponding ReadVariableOp.
    const_op.getResult().replaceAllUsesWith(read_variable_op);

  %1 = "tf.ReadVariableOp"(%arg0) : (tensor<!tf_type.resource<tensor<f32>>>) -> tensor<f32>
  // CHECK-NEXT: %[[ADD:[a-z0-9]+]] = "tf.AddV2"(%[[CONST_0]], %[[CONST_0]])
  %2 = "tf.AddV2"(%1, %1) : (tensor<f32>, tensor<f32>) -> tensor<f32>
  // CHECK-NEXT: %[[CONST_1:.*]] = "tf.Const"() <{value = dense<1.050000e+03> : tensor<f32>}>
  %3 = "tf.Const"() {value = dense<1.050000e+03> : tensor<f32>} : () -> tensor<f32>

  %2 = "tf.AddV2"(%0, %1): (tensor<15x28x28x1xf32>, tensor<1xf32>) -> tensor<15x28x28x1xf32>
  func.return %2 : tensor<15x28x28x1xf32>
// CHECK: %[[CONST_0:.*]] = "tf.Const"() <{value = dense<1.000000e+00> : tensor<1xf32>}> : () -> tensor<1xf32>

// CHECK: %[[CONST_0:.*]] = "tf.Const"() <{value = dense<0.000000e+00> : tensor<10xf32>}>
// CHECK: %[[PARTITIONEDCALL_0:.*]] = "tf.PartitionedCall"(%arg0, %arg1, %[[CONST_0]])
// CHECK-SAME: f = @composite_matmul_with_bias_and_relu6_fn_1}>
// CHECK-SAME: {_tfl_quant_trait = "fully_quantizable"
// CHECK: %[[PARTITIONEDCALL_1:.*]] = "tf.PartitionedCall"(%arg0, %arg1, %[[CONST_0]])

    module.walk([&](mlir::arith::ConstantOp const_op) {
      MaybeUnfoldLargeSplatConstant(&op_builder, const_op);
    });
  }

 private:
  void MaybeUnfoldLargeSplatConstant(mlir::OpBuilder* op_builder,
                                     mlir::arith::ConstantOp const_op) const {
    auto splat_elements_attr =
        mlir::dyn_cast<SplatElementsAttr>(const_op.getValue());
    if (!splat_elements_attr) {
      return;

    // default.
    module.walk([&](mhlo::ConstantOp const_op) {
      UnfoldSplatConstant(&op_builder, const_op);
    });
  }

 private:
  void UnfoldSplatConstant(mlir::OpBuilder* op_builder,
                           mhlo::ConstantOp const_op) const {
    auto splat_elements_attr =
        mlir::dyn_cast<SplatElementsAttr>(const_op.getValue());
    if (!splat_elements_attr) {
      return;
    }

      visitUsedValuesDefinedAbove(body, [&](OpOperand *use) {
        Value constant = use->get();
        auto const_op = dyn_cast_or_null<TF::ConstOp>(constant.getDefiningOp());
        if (!const_op) return;

        // Filter constants using user provided predicate function.
        if (filter && !filter(cluster, const_op.getValue())) return;

        // We found a constant, try to insert it in the map and re-use its

  Output id1 = ops::Identity(root.WithOpName("id1"), sw.output_true);

  Output const0 = ops::Const(root.WithOpName("const0"), 1);
  Output const1 = ops::Const(root.WithOpName("const1"), 2);

  Output add = ops::Add(root.WithOpName("add"), const0, const1);

  root.graph()->AddControlEdge(id0.node(), const0.node());
  root.graph()->AddControlEdge(id1.node(), const1.node());