[go: nahoru, domu]
Skip to content

Commit

Permalink
Stop using xla/statusor.h now that it just contains an alias for absl…
Browse files Browse the repository at this point in the history 
…::Status.

In some situations, this meant also changing unrelated files to directly include tsl/platform/statusor.h to get the definitions for TF_ASSIGN_OR_RETURN, etc., where they were getting transitively included for free.

PiperOrigin-RevId: 644301234
  • Loading branch information
@klucke @tensorflower-gardener
klucke authored and tensorflower-gardener committed Jun 19, 2024
1 parent 8c71440 commit d724a03
Show file tree
Hide file tree
Showing 15 changed files with 54 additions and 37 deletions.
2 changes: 1 addition & 1 deletion tensorflow/core/kernels/scatter_nd_op_test.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -363,7 +363,7 @@ TEST_F(ScatterNdOpConstructionTest, Error_BadIndicesPolicyInvalid) {
.Input(FakeInput(DT_INT32))
.Attr("bad_indices_policy", "AN_UNRECOGNIZED_POLICY")
.Finalize(node_def()));
EXPECT_NE(InitOp(), OkStatus());
EXPECT_NE(InitOp(), absl::OkStatus());
}

class ScatterNdUpdateBM : public ScatterNdUpdateOpTest {
Expand Down
33 changes: 7 additions & 26 deletions third_party/xla/xla/service/gpu/BUILD
View file
Original file line number Diff line number Diff line change
Expand Up @@ -327,7 +327,6 @@ cc_library(
"//xla:literal",
"//xla:shape_util",
"//xla:status_macros",
"//xla:statusor",
"//xla:util",
"//xla:xla_data_proto_cc",
"//xla/ffi:attribute_map",
Expand Down Expand Up @@ -393,6 +392,7 @@ cc_library(
"@com_google_absl//absl/log",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:str_format",
"@com_google_absl//absl/types:span",
Expand Down Expand Up @@ -507,7 +507,6 @@ cc_library(
"//xla:literal",
"//xla:shape_util",
"//xla:status_macros",
"//xla:statusor",
"//xla:util",
"//xla:xla_data_proto_cc",
"//xla/hlo/ir:hlo",
Expand Down Expand Up @@ -809,7 +808,6 @@ cc_library(
"//xla:shape_util",
"//xla:status_macros",
"//xla/tools:hlo_decomposer_lib",
"//xla:statusor",
"//xla:util",
"//xla:xla_data_proto_cc",
"//xla:xla_proto_cc",
Expand Down Expand Up @@ -1209,7 +1207,6 @@ cc_library(
"//xla:shape_util",
"//xla:status",
"//xla:status_macros",
"//xla:statusor",
"//xla:types",
"//xla:util",
"//xla:xla_data_proto_cc",
Expand Down Expand Up @@ -1382,7 +1379,6 @@ xla_cc_test(
deps = [
":gemm_fusion",
":triton_fusion_analysis",
"//xla:statusor",
"//xla/hlo/ir:hlo",
"//xla/stream_executor:device_description",
"//xla/tests:hlo_test_base",
Expand Down Expand Up @@ -1501,7 +1497,6 @@ cc_library(
"//xla:autotuning_proto_cc",
"//xla:literal_util",
"//xla:shape_util",
"//xla:statusor",
"//xla:util",
"//xla:xla_data_proto_cc",
"//xla/hlo/ir:hlo",
Expand All @@ -1511,6 +1506,7 @@ cc_library(
"@com_google_absl//absl/container:flat_hash_set",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@com_google_absl//absl/strings:cord",
"@com_google_absl//absl/types:span",
Expand Down Expand Up @@ -1618,7 +1614,6 @@ cc_library(
"//xla:util",
"//xla:autotuning_proto_cc",
"//xla:shape_util",
"//xla:statusor",
"@local_tsl//tsl/platform:errors",
"@local_tsl//tsl/platform:logging",
"@local_tsl//tsl/platform:statusor",
Expand Down Expand Up @@ -1650,7 +1645,6 @@ cc_library(
"//xla:autotuning_proto_cc",
"//xla:shape_util",
"//xla:status_macros",
"//xla:statusor",
"//xla:types",
"//xla:util",
"//xla:xla_proto_cc",
Expand Down Expand Up @@ -1696,7 +1690,6 @@ cc_library(
"//xla/stream_executor/gpu:redzone_allocator",
"//xla:executable_run_options",
"//xla:shape_util",
"//xla:statusor",
"//xla:util",
"//xla:xla_proto_cc",
"@local_tsl//tsl/platform:errors",
Expand Down Expand Up @@ -1769,7 +1762,6 @@ cc_library(
"//xla:autotuning_proto_cc",
"//xla:shape_util",
"//xla:status_macros",
"//xla:statusor",
"//xla:types",
"//xla:util",
"//xla:xla_data_proto_cc",
Expand Down Expand Up @@ -1999,7 +1991,6 @@ cc_library(
"//xla:debug_options_flags",
"//xla:literal_util",
"//xla:shape_util",
"//xla:statusor",
"//xla:util",
"//xla:xla_data_proto_cc",
"//xla/hlo/ir:hlo",
Expand Down Expand Up @@ -2096,7 +2087,6 @@ cc_library(
":stream_executor_util",
"//xla:shape_util",
"//xla:status_macros",
"//xla:statusor",
"//xla:types",
"//xla:util",
"//xla:xla_data_proto_cc",
Expand All @@ -2105,6 +2095,7 @@ cc_library(
"//xla/stream_executor:dnn",
"//xla/stream_executor:lazy_op_runner",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings",
"@local_tsl//tsl/platform:statusor",
] + if_cuda_is_configured([
Expand Down Expand Up @@ -2279,7 +2270,6 @@ cc_library(
]),
deps = [
"//xla:comparison_util",
"//xla:statusor",
"//xla:types",
"//xla:util",
"//xla:xla_data_proto_cc",
Expand Down Expand Up @@ -2317,7 +2307,6 @@ cc_library(
"//xla:literal",
"//xla:literal_util",
"//xla:shape_util",
"//xla:statusor",
"//xla:util",
"//xla:xla_data_proto_cc",
"//xla/hlo/ir:hlo",
Expand Down Expand Up @@ -2830,7 +2819,6 @@ cc_library(
":cudnn_support_utils",
":stream_executor_util",
"//xla:shape_util",
"//xla:statusor",
"//xla:util",
"//xla/client:xla_builder",
"//xla/client:xla_computation",
Expand Down Expand Up @@ -2859,7 +2847,6 @@ xla_cc_test(
":backend_configs_cc",
":cublas_cudnn",
":cudnn_vectorize_convolutions",
"//xla:statusor",
"//xla:util",
"//xla/service:call_inliner",
"//xla/service:hlo_parser",
Expand Down Expand Up @@ -3503,7 +3490,7 @@ cc_library(
"@llvm-project//mlir:Support",
"//xla:autotune_results_proto_cc",
"//xla:status_macros",
"//xla:statusor",
"@com_google_absl//absl/status:statusor",
"//xla:types",
"//xla:util",
"//xla:xla_data_proto_cc",
Expand Down Expand Up @@ -3820,7 +3807,6 @@ cc_library(
"@llvm-project//llvm:Support",
"//xla:autotune_results_proto_cc",
"//xla:status_macros",
"//xla:statusor",
"//xla:types",
"//xla:util",
"//xla:xla_proto_cc",
Expand Down Expand Up @@ -4336,7 +4322,6 @@ xla_cc_test(
],
deps = [
":gpu_p2p_pipeliner",
"//xla:statusor",
"//xla:util",
"//xla/hlo/ir:hlo",
"//xla/service:hlo_module_config",
Expand All @@ -4346,6 +4331,7 @@ xla_cc_test(
"//xla/tests:hlo_test_base",
"//xla/tests:xla_internal_test_main",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/status:statusor",
"@com_google_googletest//:gtest",
],
)
Expand Down Expand Up @@ -4512,13 +4498,13 @@ cc_library(
":ir_emission_utils",
":reduction_utils",
"//xla:shape_util",
"//xla:statusor",
"//xla/hlo/ir:hlo",
"//xla/stream_executor:device_description",
"@com_google_absl//absl/algorithm:container",
"@com_google_absl//absl/container:inlined_vector",
"@com_google_absl//absl/log",
"@com_google_absl//absl/log:check",
"@com_google_absl//absl/status:statusor",
"@com_google_absl//absl/strings:string_view",
"@com_google_absl//absl/types:span",
"@llvm-project//llvm:Support",
Expand Down Expand Up @@ -4559,7 +4545,6 @@ cc_library(
"@eigen_archive//:eigen3",
"//xla:shape_util",
"//xla:status_macros",
"//xla:statusor",
"//xla:util",
"//xla/service:hlo_module_config",
"//xla/stream_executor",
Expand Down Expand Up @@ -4779,7 +4764,6 @@ cc_library(
":backend_configs_cc",
":cublas_cudnn",
"//xla:shape_util",
"//xla:statusor",
"//xla:types",
"//xla:util",
"//xla:window_util",
Expand Down Expand Up @@ -4840,7 +4824,6 @@ cc_library(
"//xla:permutation_util",
"//xla:shape_util",
"//xla:status_macros",
"//xla:statusor",
"//xla:types",
"//xla:util",
"//xla:xla_data_proto_cc",
Expand Down Expand Up @@ -5402,14 +5385,14 @@ cc_library(
srcs = if_gpu_is_configured(["make_batch_pointers.cc"]),
hdrs = if_gpu_is_configured(["make_batch_pointers.h"]),
deps = [
"//xla:statusor",
"//xla:types",
"//xla:util",
"//xla/stream_executor",
"//xla/stream_executor:device_memory",
"//xla/stream_executor:typed_kernel_factory",
"//xla/stream_executor/gpu:gpu_stream_header",
"@com_google_absl//absl/status",
"@com_google_absl//absl/status:statusor",
"@local_tsl//tsl/platform:errors",
"@local_tsl//tsl/platform:statusor",
] + if_cuda_is_configured([
Expand Down Expand Up @@ -5596,7 +5579,6 @@ cc_library(
hdrs = ["topk_splitter.h"],
deps = [
"//xla:shape_util",
"//xla:statusor",
"//xla:xla_data_proto_cc",
"//xla/hlo/ir:hlo",
"//xla/service:hlo_creation_utils",
Expand Down Expand Up @@ -6041,7 +6023,6 @@ cc_library(
deps = [
":backend_configs_cc",
":gpu_fusible",
"//xla:statusor",
"//xla:util",
"//xla:xla_data_proto_cc",
"//xla/hlo/ir:hlo",
Expand Down
2 changes: 1 addition & 1 deletion third_party/xla/xla/service/gpu/cudnn_fused_mha_rewriter.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ limitations under the License.
#include "absl/log/check.h"
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_format.h"
#include "absl/strings/str_join.h"
#include "absl/strings/string_view.h"
Expand All @@ -46,7 +47,6 @@ limitations under the License.
#include "xla/shape.h"
#include "xla/shape_util.h"
#include "xla/status_macros.h"
#include "xla/statusor.h"
#include "xla/stream_executor/device_description.h"
#include "xla/stream_executor/dnn.h"
#include "xla/types.h"
Expand Down
2 changes: 1 addition & 1 deletion third_party/xla/xla/service/gpu/cudnn_vectorize_convolutions.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ limitations under the License.
#include "absl/container/flat_hash_set.h"
#include "absl/container/inlined_vector.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "xla/client/xla_builder.h"
Expand All @@ -42,7 +43,6 @@ limitations under the License.
#include "xla/service/hlo_module_config.h"
#include "xla/shape.h"
#include "xla/shape_util.h"
#include "xla/statusor.h"
#include "xla/stream_executor/device_description.h"
#include "xla/stream_executor/dnn.h"
#include "xla/util.h"
Expand Down
1 change: 0 additions & 1 deletion third_party/xla/xla/service/gpu/cudnn_vectorize_convolutions_test.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -29,7 +29,6 @@ limitations under the License.
#include "xla/service/hlo_parser.h"
#include "xla/service/pattern_matcher.h"
#include "xla/service/pattern_matcher_gmock.h"
#include "xla/statusor.h"
#include "xla/stream_executor/device_description.h"
#include "xla/stream_executor/dnn.h"
#include "xla/tests/hlo_test_base.h"
Expand Down
36 changes: 36 additions & 0 deletions third_party/xla/xla/service/gpu/fusions/loop_mlir_test.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -181,6 +181,42 @@ TEST_F(MlirLoopFusionTest, ThreadId_Broadcast) {
)"));
}

TEST_F(MlirLoopFusionTest, Constant_Broadcast) {
auto kHloString = R"(
HloModule module
bcast {
zero = bf16[] constant(0)
ROOT broadcast = bf16[2,16,48]{2,1,0} broadcast(zero), dimensions={}
}
ENTRY entry {
ROOT %fusion = bf16[2,16,48]{2,1,0} fusion(), kind=kLoop, calls=bcast
}
)";
TF_ASSERT_OK(EmitAndCheckIR(kHloString, R"(
// CHECK: #[[MAP0:.*]] = affine_map<(d0, d1) -> (d1 * 1024 + d0)>
// CHECK: #[[MAP1:.*]] = affine_map<(d0, d1) -> (((d1 * 4 + d0 floordiv 256) floordiv 3) mod 2)>
// CHECK: #[[MAP2:.*]] = affine_map<(d0, d1) -> (((d1 * 64 + d0 floordiv 16) floordiv 3) mod 16)>
// CHECK: #[[MAP3:.*]] = affine_map<(d0, d1) -> ((d1 * 1024 + d0) mod 48)>
// CHECK: func.func @fused_computation(%[[ARG0:.*]]: tensor<2x16x48xbf16>
// CHECK: %[[UPPER_BOUND:.*]] = arith.constant 1535 : index
// CHECK: %[[THREAD_ID:.*]] = gpu.thread_id
// CHECK: %[[BLOCK_ID:.*]] = gpu.block_id
// CHECK: %[[LINEAR:.*]] = xla_gpu.apply_indexing #[[MAP0]]
// CHECL: %[[IN_BOUNDS:.*]] = arith.cmpi sle, %[[LINEAR]], %[[UPPER_BOUND]] : index
// scf.if %[[IN_BOUNDS]]
// CHECK: %[[I0:.*]] = xla_gpu.apply_indexing #[[MAP1]]
// CHECK: %[[I1:.*]] = xla_gpu.apply_indexing #[[MAP2]]
// CHECK: %[[I2:.*]] = xla_gpu.apply_indexing #[[MAP3]]
// CHECK: %[[BCAST:.*]] = xla_gpu.pure_call @bcast_broadcast
// CHECK: %[[INSERTED:.*]] = tensor.insert %[[BCAST]] into %[[ARG0]][%[[I0]], %[[I1]], %[[I2]]]
// CHECK: func.func private @bcast_broadcast
// CHECK: arith.constant 0.000000e+00
)"));
EXPECT_TRUE(RunAndCompareNoHloPasses(kHloString, ErrorSpec{0}));
}

TEST_F(MlirLoopFusionTest, NoCodeDuplication) {
// This test HLO is copied from
// xla/service/fusion_node_indexing_evaluation_test.cc.
Expand Down
2 changes: 1 addition & 1 deletion third_party/xla/xla/service/gpu/gemm_algorithm_picker.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,6 +26,7 @@ limitations under the License.

#include "absl/container/flat_hash_set.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/str_cat.h"
#include "absl/strings/string_view.h"
#include "absl/synchronization/mutex.h"
Expand All @@ -43,7 +44,6 @@ limitations under the License.
#include "xla/service/gpu/variant_visitor.h"
#include "xla/shape.h"
#include "xla/shape_util.h"
#include "xla/statusor.h"
#include "xla/stream_executor/blas.h"
#include "xla/stream_executor/device_description.h"
#include "xla/stream_executor/device_memory.h"
Expand Down
2 changes: 1 addition & 1 deletion third_party/xla/xla/service/gpu/gemm_rewriter.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@ limitations under the License.
#include "absl/container/flat_hash_set.h"
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/string_view.h"
#include "absl/types/span.h"
#include "xla/hlo/evaluator/hlo_evaluator.h"
Expand All @@ -57,7 +58,6 @@ limitations under the License.
#include "xla/shape_util.h"
#include "xla/status.h"
#include "xla/status_macros.h"
#include "xla/statusor.h"
#include "xla/stream_executor/blas.h"
#include "xla/stream_executor/device_description.h"
#include "xla/stream_executor/gpu/gpu_blas_lt.h"
Expand Down
1 change: 1 addition & 0 deletions third_party/xla/xla/service/gpu/gpu_hlo_schedule.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -30,6 +30,7 @@ limitations under the License.
#include "absl/log/check.h"
#include "absl/log/log.h"
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/match.h"
#include "absl/strings/numbers.h"
#include "absl/strings/str_format.h"
Expand Down
Loading

0 comments on commit d724a03

Please sign in to comment.