tensorflow/compiler/tf2xla/kernels/BUILD

load("//tensorflow/core/platform:rules_cc.bzl", "cc_library")
load("//tensorflow:tensorflow.bzl", "tf_kernel_library")
load(
    "//tensorflow/core/platform:build_config.bzl",
    "tf_additional_all_protos",
    "tf_proto_library",
)
load(
    "//tensorflow:tensorflow.bzl",
    "if_cuda_or_rocm",
    "tf_cuda_library",
)

package(
    default_visibility = [
        # copybara:uncomment "//learning/infra/mira:__subpackages__",
        "//third_party/cloud_tpu/inference_converter:__subpackages__",
        "//tensorflow/compiler/tf2xla:internal",
    ],
    licenses = ["notice"],
)

tf_kernel_library(
    name = "xla_ops",
    srcs = [
        "aggregate_ops.cc",
        "all_reduce_op.cc",
        "approx_topk_op.cc",
        "arg_op.cc",
        "batch_matmul_op.cc",
        "batch_norm_op.cc",
        "batchtospace_op.cc",
        "bcast_ops.cc",
        "beta_op.cc",
        "bias_ops.cc",
        "binary_ops.cc",
        "bincount_op.cc",
        "broadcast_to_op.cc",
        "bucketize_op.cc",
        "cast_op.cc",
        "categorical_op.cc",
        "cholesky_op.cc",
        "clip_by_value_op.cc",
        "concat_op.cc",
        "const_op.cc",
        "conv_ops.cc",
        "cross_op.cc",
        "cwise_ops.cc",
        "cwise_ops.h",
        "data_format_ops.cc",
        "depthtospace_op.cc",
        "dequantize_op.cc",
        "device_index_op.cc",
        "diag_op.cc",
        "dynamic_partition_op.cc",
        "dynamic_slice_ops.cc",
        "dynamic_stitch_op.cc",
        "einsum_op.cc",
        "elu_op.cc",
        "elu_op.h",
        "empty_op.cc",
        "ensure_shape_op.cc",
        "extract_image_patches_op.cc",
        "fake_param_op.cc",
        "fake_quantize_ops.cc",
        "fft_ops.cc",
        "fill_op.cc",
        "function_ops.cc",
        "gather_op.cc",
        "gather_op_helpers.h",
        "gather_scatter_ops.cc",
        "identity_op.cc",
        "image_ops.cc",
        "image_resize_ops.cc",
        "in_topk_op.cc",
        "index_ops.cc",
        "l2loss_op.cc",
        "listdiff_op.cc",
        "lower_upper_bound_ops.cc",
        "lrn_ops.cc",
        "matmul_op.cc",
        "matrix_band_part_op.cc",
        "matrix_diag_ops.cc",
        "matrix_inverse_op.cc",
        "matrix_solve_op.cc",
        "matrix_triangular_solve_op.cc",
        "mirror_pad_op.cc",
        "next_after_op.cc",
        "no_op.cc",
        "one_hot_op.cc",
        "pack_op.cc",
        "pad_op.cc",
        "pooling_ops.cc",
        "qr_op.cc",
        "quantize_and_dequantize_op.cc",
        "random_ops.cc",
        "random_ops_util.cc",
        "random_ops_util.h",
        "reduce_window_op.cc",
        "reduction_ops.cc",
        "reduction_ops.h",
        "reduction_ops_common.cc",
        "relu_op.cc",
        "relu_op.h",
        "replica_id_op.cc",
        "reshape_op.cc",
        "retval_op.cc",
        "reverse_op.cc",
        "reverse_sequence_op.cc",
        "roll_op.cc",
        "scan_ops.cc",
        "scatter_nd_op.cc",
        "segment_reduction_ops.cc",
        "select_op.cc",
        "sendrecv_ops.cc",
        "sequence_ops.cc",
        "shape_op.cc",
        "shape_util.cc",
        "sharding_op.cc",
        "sharding_util_ops.cc",
        "slice_op.cc",
        "softmax_op.cc",
        "sort_ops.cc",
        "spacetobatch_op.cc",
        "spacetodepth_op.cc",
        "sparse_to_dense_op.cc",
        "split_op.cc",
        "spmd_manual_sharding_ops.cc",
        "stack_ops.cc",
        "stateful_random_ops.cc",
        "stateless_random_ops.cc",
        "stateless_random_ops_v2.cc",
        "strided_slice_op.cc",
        "tensor_array_ops.cc",
        "tensor_list_ops.cc",
        "tile_ops.cc",
        "to_bool_op.cc",
        "topk_op.cc",
        "training_ops.cc",
        "transpose_op.cc",
        "tridiagonal_ops.cc",
        "unary_ops.cc",
        "unary_ops_composition.cc",
        "unique_op.cc",
        "unpack_op.cc",
        "variable_ops.cc",
        "where_op.cc",
        "xla_broadcast_helper_op.cc",
        "xla_call_module_op.cc",
        "xla_conv_op.cc",
        "xla_custom_call_op.cc",
        "xla_dequantize_op.cc",
        "xla_dot_op.cc",
        "xla_optimization_barrier_op.cc",
        "xla_pad_op.cc",
        "xla_reduce_op.cc",
        "xla_reduce_precision_op.cc",
        "xla_select_and_scatter_op.cc",
        "xla_self_adjoint_eig_op.cc",
        "xla_svd_op.cc",
    ],
    hdrs = [
        "image_resize_ops.h",
        "index_ops.h",
        "shape_util.h",
    ],
    tags = ["optonly"],
    deps = [
        ":case_op",
        ":conv_op_helpers",
        ":if_op",
        ":tensor_list_utils",
        ":while_op",
        "@com_google_absl//absl/algorithm:container",
        "@com_google_absl//absl/container:flat_hash_map",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/strings:str_format",
        "@com_google_absl//absl/types:optional",
        "@com_google_absl//absl/types:span",
        "@llvm-project//mlir:IR",
        "@llvm-project//mlir:Parser",
        "@llvm-project//mlir:Pass",
        "@llvm-project//mlir:Support",
        "//tensorflow/compiler/jit:xla_activity_listener",
        "//tensorflow/compiler/jit:xla_activity_proto_cc",
        "//tensorflow/compiler/xla/mlir_hlo",
        "//tensorflow/compiler/mlir/tensorflow:error_util",
        "//tensorflow/compiler/mlir/tensorflow:tensorflow_passes",
        "//tensorflow/compiler/mlir/tensorflow:tensorflow_types",
        "//tensorflow/compiler/mlir/xla:hlo_utils",
        "//tensorflow/compiler/tf2xla:common",
        "//tensorflow/compiler/tf2xla:mlir_xla_op_kernel",
        "//tensorflow/compiler/tf2xla:xla_compilation_device",
        "//tensorflow/compiler/tf2xla:xla_compiler",
        "//tensorflow/compiler/tf2xla:xla_context",
        "//tensorflow/compiler/tf2xla:xla_helpers",
        "//tensorflow/compiler/tf2xla:xla_op_registry",
        "//tensorflow/compiler/tf2xla:xla_resource",
        "//tensorflow/compiler/tf2xla/lib:broadcast",
        "//tensorflow/compiler/tf2xla/lib:data_format",
        "//tensorflow/compiler/tf2xla/lib:random",
        "//tensorflow/compiler/tf2xla/lib:scatter",
        "//tensorflow/compiler/tf2xla/lib:util",
        "//tensorflow/compiler/tf2xla/ops:xla_ops",
        "//tensorflow/compiler/xla:array4d",
        "//tensorflow/compiler/xla:comparison_util",
        "//tensorflow/compiler/xla:literal",
        "//tensorflow/compiler/xla:literal_util",
        "//tensorflow/compiler/xla:shape_util",
        "//tensorflow/compiler/xla:sharding_op_util",
        "//tensorflow/compiler/xla:status_macros",
        "//tensorflow/compiler/xla:util",
        "//tensorflow/compiler/xla:xla_data_proto_cc",
        "//tensorflow/compiler/xla/client:client_library",
        "//tensorflow/compiler/xla/client:value_inference",
        "//tensorflow/compiler/xla/client:xla_builder",
        "//tensorflow/compiler/xla/client:xla_computation",
        "//tensorflow/compiler/xla/client/lib:approx_topk",
        "//tensorflow/compiler/xla/client/lib:arithmetic",
        "//tensorflow/compiler/xla/client/lib:comparators",
        "//tensorflow/compiler/xla/client/lib:constants",
        "//tensorflow/compiler/xla/client/lib:dynamic_shaped_ops",
        "//tensorflow/compiler/xla/client/lib:loops",
        "//tensorflow/compiler/xla/client/lib:math",
        "//tensorflow/compiler/xla/client/lib:matrix",
        "//tensorflow/compiler/xla/client/lib:pooling",
        "//tensorflow/compiler/xla/client/lib:prng",
        "//tensorflow/compiler/xla/client/lib:qr",
        "//tensorflow/compiler/xla/client/lib:quantize",
        "//tensorflow/compiler/xla/client/lib:self_adjoint_eig",
        "//tensorflow/compiler/xla/client/lib:slicing",
        "//tensorflow/compiler/xla/client/lib:sorting",
        "//tensorflow/compiler/xla/client/lib:svd",
        "//tensorflow/compiler/xla/client/lib:tridiagonal",
        "//tensorflow/compiler/xla/pjrt:mlir_to_hlo",
        "//tensorflow/compiler/xla/service:hlo_proto_cc",
        "//tensorflow/core:framework",
        "//tensorflow/core:lib",
        "//tensorflow/core:protos_all_cc",
        "//tensorflow/core/kernels:resource_variable_util",
        "//tensorflow/core/kernels:scatter_nd_util",
        "//tensorflow/core/kernels:stateful_random_ops_header",
        "//tensorflow/core/kernels:stateless_random_ops_v2_header",
        "//tensorflow/core/tpu:tpu_defs",
        "//tensorflow/core/util:overflow",
        "//tensorflow/stream_executor/lib",
    ] + if_cuda_or_rocm(
        if_false = [],
        if_true = [":light_outside_compilation"],
    ),
)

tf_cuda_library(
    name = "light_outside_compilation",
    srcs = ["light_outside_compilation.cc"],
    hdrs = ["light_outside_compilation.h"],
    visibility = ["//visibility:public"],
    deps = [
        ":callback_proto_cc",
        "//tensorflow/compiler/tf2xla:common",
        "//tensorflow/compiler/tf2xla:xla_compiler",
        "//tensorflow/compiler/tf2xla:xla_helpers",
        "//tensorflow/compiler/tf2xla:xla_op_registry",
        "//tensorflow/compiler/tf2xla/lib:util",
        "//tensorflow/compiler/xla:shape_util",
        "//tensorflow/compiler/xla:xla_data_proto_cc",
        "//tensorflow/compiler/xla/service:custom_call_status",
        "//tensorflow/compiler/xla/service:custom_call_target_registry",
        "//tensorflow/core:framework",
        "//tensorflow/core:gpu_runtime",
        "//tensorflow/core:protos_all_cc",
        "//tensorflow/core/platform:status",
        "//tensorflow/stream_executor/gpu:gpu_executor_header",
        "//tensorflow/stream_executor/gpu:gpu_stream_header",
        "//tensorflow/stream_executor/gpu:gpu_types_header",
        "@com_google_absl//absl/strings",
        "@com_google_absl//absl/types:span",
    ],
    alwayslink = 1,
)

# A separate cc_library for resampler_ops is needed because resampler is in
# contrib/, and thus the declaration of resampler cannot be pulled into the deps
# of xla_ops. Therefore, resampler_ops is its own cc_library target, and its
# corresponding tf_kernel_library is defined in contrib/resampler/BUILD.
cc_library(
    name = "resampler_ops",
    srcs = ["resampler_ops.cc"],
    hdrs = ["resampler_ops.h"],
    visibility = ["//visibility:public"],
    deps = [
        "//tensorflow/compiler/tf2xla:common",
        "//tensorflow/compiler/tf2xla:xla_compiler",
        "//tensorflow/compiler/xla:array4d",
        "//tensorflow/compiler/xla:literal",
        "//tensorflow/compiler/xla:shape_util",
        "//tensorflow/compiler/xla:xla_data_proto_cc",
        "//tensorflow/compiler/xla/client:xla_builder",
        "//tensorflow/compiler/xla/client/lib:arithmetic",
        "//tensorflow/compiler/xla/client/lib:constants",
        "//tensorflow/core:framework",
        "//tensorflow/core:lib",
        "//tensorflow/core:protos_all_cc",
    ],
    alwayslink = 1,
)

# See note on resampler_ops target for why this is also a seprate op.
cc_library(
    name = "resampler_addon_ops",
    srcs = ["resampler_addon_ops.cc"],
    visibility = ["//visibility:public"],
    deps = [
        ":resampler_ops",
        "//tensorflow/compiler/tf2xla:common",
        "//tensorflow/compiler/tf2xla:xla_compiler",
        "//tensorflow/compiler/tf2xla:xla_helpers",
        "//tensorflow/compiler/tf2xla:xla_op_registry",
    ],
    alwayslink = 1,
)

cc_library(
    name = "conv_op_helpers",
    srcs = ["conv_op_helpers.cc"],
    hdrs = ["conv_op_helpers.h"],
    deps = [
        "//tensorflow/compiler/tf2xla:common",
        "//tensorflow/compiler/tf2xla:xla_compiler",
        "//tensorflow/compiler/tf2xla:xla_helpers",
        "//tensorflow/compiler/tf2xla:xla_op_registry",
        "//tensorflow/compiler/xla:literal_util",
        "//tensorflow/compiler/xla:statusor",
        "//tensorflow/compiler/xla:util",
        "//tensorflow/compiler/xla/client:xla_builder",
        "//tensorflow/compiler/xla/client/lib:arithmetic",
        "//tensorflow/compiler/xla/client/lib:constants",
        "//tensorflow/core:framework",
        "//tensorflow/core/framework:bounds_check",
        "//tensorflow/core/kernels:conv_grad_shape_utils",
        "@com_google_absl//absl/types:span",
    ],
)

cc_library(
    name = "tensor_list_utils",
    srcs = ["tensor_list_utils.cc"],
    hdrs = ["tensor_list_utils.h"],
    deps = [
        "//tensorflow/compiler/tf2xla:common",
        "//tensorflow/compiler/tf2xla:xla_compiler",
        "//tensorflow/compiler/xla:literal_util",
        "//tensorflow/compiler/xla:shape_util",
        "//tensorflow/compiler/xla:status_macros",
        "//tensorflow/compiler/xla:statusor",
        "//tensorflow/compiler/xla:util",
        "//tensorflow/compiler/xla:xla_data_proto_cc",
        "//tensorflow/compiler/xla/client:xla_builder",
        "//tensorflow/core:framework",
        "//tensorflow/core:lib",
    ],
)

cc_library(
    name = "if_while_utils",
    srcs = ["if_while_utils.cc"],
    hdrs = ["if_while_utils.h"],
    deps = [
        "//tensorflow/compiler/tf2xla:common",
        "//tensorflow/compiler/tf2xla:xla_compiler",
        "//tensorflow/compiler/tf2xla/ops:xla_ops",
        "//tensorflow/compiler/xla:literal",
        "//tensorflow/core:lib",
    ],
)

tf_kernel_library(
    name = "while_op",
    srcs = ["while_op.cc"],
    hdrs = ["while_op.h"],
    deps = [
        ":if_while_utils",
        ":tensor_list_utils",
        "//tensorflow/compiler/tf2xla:common",
        "//tensorflow/compiler/tf2xla:side_effect_util",
        "//tensorflow/compiler/tf2xla:tf2xla_util",
        "//tensorflow/compiler/tf2xla:xla_compiler",
        "//tensorflow/compiler/tf2xla:xla_helpers",
        "//tensorflow/compiler/tf2xla:xla_op_registry",
        "//tensorflow/compiler/tf2xla/ops:xla_ops",
        "//tensorflow/compiler/xla:literal",
        "//tensorflow/compiler/xla:status_macros",
        "//tensorflow/compiler/xla/client:xla_builder",
        "//tensorflow/compiler/xla/client:xla_computation",
        "//tensorflow/core:framework",
        "//tensorflow/core:protos_all_cc",
        "@com_google_absl//absl/strings",
    ],
)

tf_kernel_library(
    name = "if_op",
    srcs = ["if_op.cc"],
    hdrs = ["if_op.h"],
    deps = [
        ":if_while_utils",
        "//tensorflow/compiler/tf2xla:common",
        "//tensorflow/compiler/tf2xla:side_effect_util",
        "//tensorflow/compiler/tf2xla:xla_compiler",
        "//tensorflow/compiler/tf2xla:xla_context",
        "//tensorflow/compiler/tf2xla:xla_op_registry",
        "//tensorflow/compiler/tf2xla/ops:xla_ops",
        "//tensorflow/compiler/xla:literal",
        "//tensorflow/compiler/xla/client:xla_builder",
        "//tensorflow/compiler/xla/client/lib:dynamic_shaped_ops",
        "//tensorflow/core:framework",
        "//tensorflow/core:lib",
        "//tensorflow/core:protos_all_cc",
    ],
)

tf_kernel_library(
    name = "case_op",
    srcs = ["case_op.cc"],
    hdrs = ["case_op.h"],
    deps = [
        ":if_while_utils",
        "//tensorflow/compiler/tf2xla:common",
        "//tensorflow/compiler/tf2xla:side_effect_util",
        "//tensorflow/compiler/tf2xla:xla_compiler",
        "//tensorflow/compiler/tf2xla:xla_context",
        "//tensorflow/compiler/tf2xla:xla_op_registry",
        "//tensorflow/compiler/tf2xla/ops:xla_ops",
        "//tensorflow/compiler/xla:literal",
        "//tensorflow/compiler/xla/client:xla_builder",
        "//tensorflow/compiler/xla/client/lib:constants",
        "//tensorflow/compiler/xla/client/lib:dynamic_shaped_ops",
        "//tensorflow/core:framework",
        "//tensorflow/core:lib",
        "//tensorflow/core:protos_all_cc",
    ],
)

# Kernels that have a dummy (no-op) implementation.
tf_kernel_library(
    name = "xla_dummy_ops",
    srcs = [
        "assert_op.cc",
        "check_numerics_op.cc",
    ],
    deps = [
        "//tensorflow/compiler/tf2xla:mlir_xla_op_kernel",
        "//tensorflow/compiler/tf2xla:xla_compiler",
        "//tensorflow/compiler/tf2xla:xla_op_registry",
        "//tensorflow/core:array_ops_op_lib",
        "//tensorflow/core:framework",
        "//tensorflow/core:lib",
        "//tensorflow/core:logging_ops_op_lib",
    ],
    alwayslink = 1,
)

tf_proto_library(
    name = "callback_proto",
    srcs = ["callback.proto"],
    cc_api_version = 2,
    protodeps = tf_additional_all_protos(),
)