DeviceIndex op: given a list of device names, this operation returns the index of the device this op runs. In the case of XLA, we are not executing on any device, we return the length of the list. PiperOrigin-RevId: 317740778 Change-Id: I0679aa0adc5508b83502eee0d2044584577ed5b4
343 lines
11 KiB
Python
343 lines
11 KiB
Python
load("//tensorflow:tensorflow.bzl", "tf_kernel_library")
|
|
|
|
package(
|
|
default_visibility = ["//tensorflow/compiler/tf2xla:internal"],
|
|
licenses = ["notice"], # Apache 2.0
|
|
)
|
|
|
|
tf_kernel_library(
|
|
name = "xla_ops",
|
|
srcs = [
|
|
"aggregate_ops.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",
|
|
"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_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.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",
|
|
"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",
|
|
"strided_slice_op.cc",
|
|
"tensor_array_ops.cc",
|
|
"tensor_list_ops.cc",
|
|
"tile_ops.cc",
|
|
"topk_op.cc",
|
|
"training_ops.cc",
|
|
"transpose_op.cc",
|
|
"tridiagonal_ops.cc",
|
|
"unary_ops.cc",
|
|
"unary_ops_composition.cc",
|
|
"unpack_op.cc",
|
|
"variable_ops.cc",
|
|
"xla_broadcast_helper_op.cc",
|
|
"xla_conv_op.cc",
|
|
"xla_dequantize_op.cc",
|
|
"xla_dot_op.cc",
|
|
"xla_pad_op.cc",
|
|
"xla_reduce_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",
|
|
"//tensorflow/compiler/jit:xla_activity_listener",
|
|
"//tensorflow/compiler/jit:xla_activity_proto_cc",
|
|
"//tensorflow/compiler/tf2xla:common",
|
|
"//tensorflow/compiler/tf2xla:xla_compiler",
|
|
"//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:status_macros",
|
|
"//tensorflow/compiler/xla:util",
|
|
"//tensorflow/compiler/xla:xla_data_proto_cc",
|
|
"//tensorflow/compiler/xla/client:client_library",
|
|
"//tensorflow/compiler/xla/client:xla_builder",
|
|
"//tensorflow/compiler/xla/client:xla_computation",
|
|
"//tensorflow/compiler/xla/client/lib:arithmetic",
|
|
"//tensorflow/compiler/xla/client/lib:comparators",
|
|
"//tensorflow/compiler/xla/client/lib:constants",
|
|
"//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/core:framework",
|
|
"//tensorflow/core:lib",
|
|
"//tensorflow/core:protos_all_cc",
|
|
"//tensorflow/core/kernels:stateful_random_ops_header",
|
|
"@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",
|
|
],
|
|
)
|
|
|
|
# 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"],
|
|
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,
|
|
)
|
|
|
|
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/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:xla_compiler",
|
|
"//tensorflow/compiler/tf2xla/ops:xla_ops",
|
|
"//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:xla_compiler",
|
|
"//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/ops:xla_ops",
|
|
"//tensorflow/compiler/xla:literal",
|
|
"//tensorflow/compiler/xla/client:xla_builder",
|
|
"//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/ops:xla_ops",
|
|
"//tensorflow/compiler/xla:literal",
|
|
"//tensorflow/compiler/xla/client:xla_builder",
|
|
"//tensorflow/compiler/xla/client/lib:constants",
|
|
"//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:xla_compiler",
|
|
"//tensorflow/core:array_ops_op_lib",
|
|
"//tensorflow/core:framework",
|
|
"//tensorflow/core:lib",
|
|
"//tensorflow/core:logging_ops_op_lib",
|
|
],
|
|
alwayslink = 1,
|
|
)
|