diff --git a/tensorflow/core/kernels/image/extract_volume_patches_op.cc b/tensorflow/core/kernels/image/extract_volume_patches_op.cc index 52774b698afe20..8b4c17f5cb49d0 100644 --- a/tensorflow/core/kernels/image/extract_volume_patches_op.cc +++ b/tensorflow/core/kernels/image/extract_volume_patches_op.cc @@ -130,9 +130,9 @@ class ExtractVolumePatchesOp : public UnaryOp { MultiplyWithoutOverflow(ksize_rows, MultiplyWithoutOverflow(ksize_cols, depth))); OP_REQUIRES(context, patch_size >= 0, - errors::InvalidArgument( + absl::InvalidArgumentError(absl::StrCat( "Output size would overflow: ", ksize_planes, " x ", - ksize_rows, " x ", ksize_cols, " x ", depth)); + ksize_rows, " x ", ksize_cols, " x ", depth))); const std::vector out_sizes = {batch, out_planes, out_rows, out_cols, patch_size}; diff --git a/tensorflow/lite/experimental/acceleration/compatibility/BUILD b/tensorflow/lite/experimental/acceleration/compatibility/BUILD index 3d3f23a017b6fd..db4dc4f6a979ef 100644 --- a/tensorflow/lite/experimental/acceleration/compatibility/BUILD +++ b/tensorflow/lite/experimental/acceleration/compatibility/BUILD @@ -100,9 +100,9 @@ genrule( ], outs = ["devicedb-sample.bin"], cmd = """ - $(location :json_to_fb) \ - --fbs=$(location :database.fbs) \ - --json_input=$(location :devicedb-sample.json) \ + $(execpath :json_to_fb) \ + --fbs=$(execpath :database.fbs) \ + --json_input=$(execpath :devicedb-sample.json) \ --fb_output=$(@) """, tools = [":json_to_fb"], @@ -125,10 +125,10 @@ genrule( # convert_file_to_c_source for some reason doesn't define the global with # 'extern', which is needed for global const variables in C++. cmd = """ - $(location :convert_binary_to_cc_source) \ - --input_binary_file $(location :devicedb-sample.bin) \ - --output_header_file $(location :devicedb-sample.h) \ - --output_source_file $(location :devicedb-sample.cc) \ + $(execpath :convert_binary_to_cc_source) \ + --input_binary_file $(execpath :devicedb-sample.bin) \ + --output_header_file $(execpath :devicedb-sample.h) \ + --output_source_file $(execpath :devicedb-sample.cc) \ --array_variable_name g_tflite_acceleration_devicedb_sample_binary """, tools = [":convert_binary_to_cc_source"], @@ -162,6 +162,7 @@ cc_test( exports_files([ "database.fbs", "gpu_compatibility.bin", + "mldrift_compatibility.bin", ]) genrule( @@ -174,15 +175,34 @@ genrule( # convert_file_to_c_source for some reason doesn't define the global with # 'extern', which is needed for global const variables in C++. cmd = """ - $(location :convert_binary_to_cc_source) \ - --input_binary_file $(location :gpu_compatibility.bin) \ - --output_header_file $(location :gpu_compatibility_binary.h) \ - --output_source_file $(location :gpu_compatibility_binary.cc) \ + $(execpath :convert_binary_to_cc_source) \ + --input_binary_file $(execpath :gpu_compatibility.bin) \ + --output_header_file $(execpath :gpu_compatibility_binary.h) \ + --output_source_file $(execpath :gpu_compatibility_binary.cc) \ --array_variable_name g_tflite_acceleration_gpu_compatibility_binary """, tools = [":convert_binary_to_cc_source"], ) +genrule( + name = "mldrift_compatibility_binary", + srcs = ["mldrift_compatibility.bin"], + outs = [ + "mldrift_compatibility_binary.h", + "mldrift_compatibility_binary.cc", + ], + # The array variable name must differ from the gpu_compatibility one above: + # both allowlists may end up linked into the same binary. + cmd = """ + $(execpath :convert_binary_to_cc_source) \ + --input_binary_file $(execpath :mldrift_compatibility.bin) \ + --output_header_file $(execpath :mldrift_compatibility_binary.h) \ + --output_source_file $(execpath :mldrift_compatibility_binary.cc) \ + --array_variable_name g_tflite_acceleration_mldrift_compatibility_binary + """, + tools = [":convert_binary_to_cc_source"], +) + cc_library( name = "android_info", srcs = ["android_info.cc"], @@ -198,6 +218,16 @@ cc_library( hdrs = [":gpu_compatibility_binary.h"], ) +# Deliberately a target of its own rather than an addition to the library above: +# a client that only cares about ML Drift should not link the (much larger) +# TFLite GPU allowlist, and vice versa. Both binaries are consumed through the +# generic GPUCompatibilityList::Create(flatbuffer, length) overload. +cc_library( + name = "mldrift_compatibility_binary_embed", + srcs = [":mldrift_compatibility_binary.cc"], + hdrs = [":mldrift_compatibility_binary.h"], +) + cc_library( name = "gpu_compatibility", hdrs = ["gpu_compatibility.h"], diff --git a/tensorflow/lite/experimental/acceleration/compatibility/mldrift_compatibility.bin b/tensorflow/lite/experimental/acceleration/compatibility/mldrift_compatibility.bin new file mode 100755 index 00000000000000..2d8ae7f3fc66b8 Binary files /dev/null and b/tensorflow/lite/experimental/acceleration/compatibility/mldrift_compatibility.bin differ diff --git a/third_party/xla/xla/service/BUILD b/third_party/xla/xla/service/BUILD index 9fd8b08de32c83..7ee9978cea8592 100644 --- a/third_party/xla/xla/service/BUILD +++ b/third_party/xla/xla/service/BUILD @@ -4780,8 +4780,7 @@ cc_library( "//xla/hlo/pass:hlo_pass", "//xla/hlo/transforms/simplifiers:hlo_dce", "//xla/hlo/transforms/simplifiers:tuple_simplifier", - "//xla/tsl/platform:errors", - "//xla/tsl/platform:statusor", + "//xla/tsl/platform:logging", "@com_google_absl//absl/algorithm:container", "@com_google_absl//absl/container:flat_hash_map", "@com_google_absl//absl/container:flat_hash_set", @@ -4810,6 +4809,7 @@ xla_cc_test( "//xla/tests:xla_internal_test_main", "@com_google_absl//absl/log", "@com_google_absl//absl/status:status_matchers", + "@com_google_absl//absl/status:statusor", "@com_google_googletest//:gtest", ], ) diff --git a/third_party/xla/xla/service/while_loop_invariant_code_motion.cc b/third_party/xla/xla/service/while_loop_invariant_code_motion.cc index 208cd413ef499c..efeb41a17ff8fb 100644 --- a/third_party/xla/xla/service/while_loop_invariant_code_motion.cc +++ b/third_party/xla/xla/service/while_loop_invariant_code_motion.cc @@ -27,6 +27,7 @@ limitations under the License. #include "absl/log/check.h" #include "absl/log/log.h" #include "absl/status/status_macros.h" +#include "absl/status/statusor.h" #include "absl/strings/string_view.h" #include "xla/hlo/analysis/while_loop_analysis.h" #include "xla/hlo/ir/hlo_computation.h" @@ -41,8 +42,7 @@ limitations under the License. #include "xla/service/while_util.h" #include "xla/shape.h" #include "xla/shape_util.h" -#include "xla/tsl/platform/errors.h" -#include "xla/tsl/platform/statusor.h" +#include "xla/tsl/platform/logging.h" #include "xla/util.h" namespace xla { @@ -100,7 +100,11 @@ WhileLoopInvariantCodeMotion::TryHoistingInvariantInstructionsFromWhileBody( return false; } - std::string while_instr_name = while_instr->ToString(print_no_metadata); + // Only used by VLOG, and the string carries the whole loop state shape. + std::string while_instr_name; + if (VLOG_IS_ON(1)) { + while_instr_name = while_instr->ToString(print_no_metadata); + } VLOG(2) << "Trying to hoist from " << while_instr_name; auto maybe_upper_bound = ComputeWhileLoopTripCountUpperBound(while_instr); @@ -146,7 +150,8 @@ WhileLoopInvariantCodeMotion::TryHoistingInvariantInstructionsFromWhileBody( return false; } - for (auto* instruction : while_body->MakeInstructionPostOrder()) { + // This scan bails on any match, so it does not need the post order. + for (const auto* instruction : while_body->instructions()) { // LICM in the presence of domain instructions is complex, bail. if (instruction->opcode() == HloOpcode::kDomain || instruction->IsCustomCall("SPMDFullToShardShape") || @@ -167,6 +172,12 @@ WhileLoopInvariantCodeMotion::TryHoistingInvariantInstructionsFromWhileBody( std::vector instructions_to_replace; std::vector replacement_instructions; + const auto is_invariant = [&](HloInstruction* op) { + return op->opcode() == HloOpcode::kConstant || + hoisted_instructions.contains(op) || + unhoisted_invariant_instructions.contains(op); + }; + for (auto* instruction : while_body->MakeInstructionPostOrder()) { allowance->DeductCost(1); if (!allowance->ContinueAnalysis()) { @@ -185,6 +196,13 @@ WhileLoopInvariantCodeMotion::TryHoistingInvariantInstructionsFromWhileBody( instruction->opcode() != HloOpcode::kReshape) { continue; } + + // Cheap check first. Both checks are pure filters, so their order does not + // change which instructions are hoisted. + if (!absl::c_all_of(instruction->operands(), is_invariant)) { + continue; + } + // Constants don't inflate, so size inflation check doesn't make sense for // constants. if (hoist_size_inflation_ratio_ && @@ -220,16 +238,6 @@ WhileLoopInvariantCodeMotion::TryHoistingInvariantInstructionsFromWhileBody( } } - auto is_invariant = [&](HloInstruction* op) { - return hoisted_instructions.find(op) != hoisted_instructions.end() || - unhoisted_invariant_instructions.contains(op) || - op->opcode() == HloOpcode::kConstant; - }; - - if (!absl::c_all_of(instruction->operands(), is_invariant)) { - continue; - } - if (NotWorthHoistingIndividually(*instruction)) { VLOG(2) << "Adding " << instruction->ToString(print_no_metadata) << " to unhoisted invariant set."; diff --git a/third_party/xla/xla/service/while_loop_invariant_code_motion_test.cc b/third_party/xla/xla/service/while_loop_invariant_code_motion_test.cc index 50ddfab37252d3..6772be7175bc99 100644 --- a/third_party/xla/xla/service/while_loop_invariant_code_motion_test.cc +++ b/third_party/xla/xla/service/while_loop_invariant_code_motion_test.cc @@ -17,10 +17,12 @@ limitations under the License. #include #include +#include #include #include "absl/log/log.h" #include "absl/status/status_matchers.h" +#include "absl/status/statusor.h" #include "xla/hlo/ir/hlo_computation.h" #include "xla/hlo/ir/hlo_instruction.h" #include "xla/hlo/ir/hlo_module.h" @@ -501,7 +503,8 @@ ENTRY entry { )"; TEST_F(WhileLoopInvariantCodeMotionTest, HoistsConstantWhenAsked) { - auto m = ParseAndReturnVerifiedModule(kConstantHoistingTestCase).value(); + ASSERT_OK_AND_ASSIGN(std::unique_ptr m, + ParseAndReturnVerifiedModule(kConstantHoistingTestCase)); ASSERT_OK_AND_ASSIGN( bool simplified_loop, @@ -533,7 +536,8 @@ TEST_F(WhileLoopInvariantCodeMotionTest, HoistsConstantWhenAsked) { } TEST_F(WhileLoopInvariantCodeMotionTest, DoesNotHoistConstantByDefault) { - auto m = ParseAndReturnVerifiedModule(kConstantHoistingTestCase).value(); + ASSERT_OK_AND_ASSIGN(std::unique_ptr m, + ParseAndReturnVerifiedModule(kConstantHoistingTestCase)); ASSERT_OK_AND_ASSIGN(bool simplified_loop, WhileLoopInvariantCodeMotion{}.Run(m.get())); @@ -605,7 +609,8 @@ ENTRY entry { )"; TEST_F(WhileLoopInvariantCodeMotionTest, HoistsInflatingByDefault) { - auto m = ParseAndReturnVerifiedModule(kInflatingTestCase).value(); + ASSERT_OK_AND_ASSIGN(std::unique_ptr m, + ParseAndReturnVerifiedModule(kInflatingTestCase)); ASSERT_OK_AND_ASSIGN( bool simplified_loop, @@ -617,16 +622,259 @@ TEST_F(WhileLoopInvariantCodeMotionTest, HoistsInflatingByDefault) { EXPECT_THAT(while_body->instructions(), Not(Contains(op::Iota()))); } +// The previous call misbound its positional arguments (1.0 landed on +// hoist_other), so the ratio was never set and the check never ran. TEST_F(WhileLoopInvariantCodeMotionTest, NoHoistInflating) { - auto m = ParseAndReturnVerifiedModule(kInflatingTestCase).value(); + ASSERT_OK_AND_ASSIGN(std::unique_ptr m, + ParseAndReturnVerifiedModule(kInflatingTestCase)); + + ASSERT_OK_AND_ASSIGN( + bool simplified_loop, + WhileLoopInvariantCodeMotion(/*hoist_constants=*/true, + /*hoist_reshapes=*/false, + /*hoist_other=*/true, + /*hoist_size_inflation_ratio=*/1.0) + .Run(m.get())); + // The constant is hoisted; the iota inflates and stays in the loop. + EXPECT_TRUE(simplified_loop); + + HloComputation* while_body = m->GetComputationWithName("wide.body"); + ASSERT_NE(while_body, nullptr); + EXPECT_THAT(while_body->instructions(), Contains(op::Iota())); + EXPECT_THAT(while_body->instructions(), Not(Contains(op::Constant()))); +} + +const char* const kInflatingChainTestCase = R"( +HloModule ModuleWithWhile + +body { + p_body = (f32[8], f32[8], f32[8,1024]) parameter(0) + invariant = f32[8] get-tuple-element(p_body), index=0 + negate = f32[8] negate(invariant) + broadcast = f32[8,1024] broadcast(invariant), dimensions={0} + negate_broadcast = f32[8,1024] negate(broadcast) + ROOT root = (f32[8], f32[8], f32[8,1024]) tuple(invariant, negate, negate_broadcast) +} + +condition { + p_cond = (f32[8], f32[8], f32[8,1024]) parameter(0) + ROOT result = pred[] constant(true) +} + +ENTRY entry { + param = (f32[8], f32[8], f32[8,1024]) parameter(0) + ROOT while = (f32[8], f32[8], f32[8,1024]) while(param), condition=condition, body=body +} +)"; + +TEST_F(WhileLoopInvariantCodeMotionTest, HoistsInflatingChainWithoutRatio) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr m, + ParseAndReturnVerifiedModule(kInflatingChainTestCase)); + + ASSERT_OK_AND_ASSIGN(bool simplified_loop, + WhileLoopInvariantCodeMotion{}.Run(m.get())); + EXPECT_TRUE(simplified_loop); + + HloInstruction* transformed_while; + FindOnlyWhileInstruction(m->entry_computation(), &transformed_while); + EXPECT_THAT(transformed_while->while_body()->instructions(), + Each(Not(AnyOf(op::Negate(), op::Broadcast())))); + EXPECT_THAT(m->entry_computation()->instructions(), + Contains(op::Negate(op::Broadcast()))); +} + +TEST_F(WhileLoopInvariantCodeMotionTest, + SizeInflationRatioRejectsOnlyInflatingInstructions) { + ASSERT_OK_AND_ASSIGN(std::unique_ptr m, + ParseAndReturnVerifiedModule(kInflatingChainTestCase)); ASSERT_OK_AND_ASSIGN( bool simplified_loop, WhileLoopInvariantCodeMotion(/*hoist_constants=*/false, - /*hoist_non_constants=*/true, + /*hoist_reshapes=*/false, + /*hoist_other=*/true, /*hoist_size_inflation_ratio=*/1.0) .Run(m.get())); - EXPECT_FALSE(simplified_loop); + EXPECT_TRUE(simplified_loop); + + // negate(invariant) does not inflate and is hoisted. The broadcast inflates + // and stays in the loop, so its negate has an unhoisted operand and stays + // too. + HloInstruction* transformed_while; + FindOnlyWhileInstruction(m->entry_computation(), &transformed_while); + EXPECT_THAT( + m->entry_computation()->instructions(), + AllOf(Contains(op::Negate(op::GetTupleElement(op::Parameter(0), 0))), + Not(Contains(op::Broadcast())))); + EXPECT_THAT(transformed_while->while_body()->instructions(), + AllOf(Contains(op::Negate(op::Broadcast())), + Each(Not(op::Negate(op::GetTupleElement()))))); +} + +TEST_F(WhileLoopInvariantCodeMotionTest, + SizeInflationRatioNeverHoistsLoopVaryingOperands) { + // Both converts shrink their input, so the ratio permits them; only the one + // with an invariant operand may be hoisted. + const char* const kHloModule = R"( +HloModule ModuleWithWhile + +body { + p_body = (f32[8,128], f32[8,128], bf16[8,128], bf16[8,128]) parameter(0) + invariant = f32[8,128] get-tuple-element(p_body), index=0 + varying = f32[8,128] get-tuple-element(p_body), index=1 + convert_invariant = bf16[8,128] convert(invariant) + convert_varying = bf16[8,128] convert(varying) + negate = f32[8,128] negate(varying) + ROOT root = (f32[8,128], f32[8,128], bf16[8,128], bf16[8,128]) tuple(invariant, negate, convert_invariant, convert_varying) +} + +condition { + p_cond = (f32[8,128], f32[8,128], bf16[8,128], bf16[8,128]) parameter(0) + ROOT result = pred[] constant(true) +} + +ENTRY entry { + param = (f32[8,128], f32[8,128], bf16[8,128], bf16[8,128]) parameter(0) + ROOT while = (f32[8,128], f32[8,128], bf16[8,128], bf16[8,128]) while(param), condition=condition, body=body +} +)"; + ASSERT_OK_AND_ASSIGN(std::unique_ptr m, + ParseAndReturnVerifiedModule(kHloModule)); + + ASSERT_OK_AND_ASSIGN( + bool simplified_loop, + WhileLoopInvariantCodeMotion(/*hoist_constants=*/false, + /*hoist_reshapes=*/false, + /*hoist_other=*/true, + /*hoist_size_inflation_ratio=*/1.0) + .Run(m.get())); + EXPECT_TRUE(simplified_loop); + + HloInstruction* transformed_while; + FindOnlyWhileInstruction(m->entry_computation(), &transformed_while); + EXPECT_THAT(m->entry_computation()->instructions(), + Contains(op::Convert(op::GetTupleElement(op::Parameter(0), 0)))); + EXPECT_THAT(transformed_while->while_body()->instructions(), + Contains(op::Convert(op::GetTupleElement(op::Parameter(0), 1)))); + EXPECT_THAT(transformed_while->while_body()->instructions(), + Each(Not(op::Convert(op::GetTupleElement(op::Parameter(0), 0))))); +} + +TEST_F(WhileLoopInvariantCodeMotionTest, + SizeInflationCountsAllLeavesOfTupleOperands) { + // get-tuple-element(custom-call), index=1 produces 2048 bytes out of a 2080 + // byte tuple operand: hoistable at ratio 1.0 only if the operand tuple's + // leaves are summed; a tuple counted as size 0 rejects it. + const char* const kHloModule = R"( +HloModule ModuleWithWhile + +body { + p_body = (f32[8], f32[8,64], f32[8,64]) parameter(0) + invariant_a = f32[8] get-tuple-element(p_body), index=0 + invariant_b = f32[8,64] get-tuple-element(p_body), index=1 + custom_call = (f32[8], f32[8,64]) custom-call(invariant_a, invariant_b), custom_call_target="Foo" + element = f32[8,64] get-tuple-element(custom_call), index=1 + negate = f32[8,64] negate(element) + ROOT root = (f32[8], f32[8,64], f32[8,64]) tuple(invariant_a, invariant_b, negate) +} + +condition { + p_cond = (f32[8], f32[8,64], f32[8,64]) parameter(0) + ROOT result = pred[] constant(true) +} + +ENTRY entry { + param = (f32[8], f32[8,64], f32[8,64]) parameter(0) + ROOT while = (f32[8], f32[8,64], f32[8,64]) while(param), condition=condition, body=body +} +)"; + ASSERT_OK_AND_ASSIGN(std::unique_ptr m, + ParseAndReturnVerifiedModule(kHloModule)); + + ASSERT_OK_AND_ASSIGN( + bool simplified_loop, + WhileLoopInvariantCodeMotion(/*hoist_constants=*/false, + /*hoist_reshapes=*/false, + /*hoist_other=*/true, + /*hoist_size_inflation_ratio=*/1.0) + .Run(m.get())); + EXPECT_TRUE(simplified_loop); + + HloInstruction* transformed_while; + FindOnlyWhileInstruction(m->entry_computation(), &transformed_while); + EXPECT_THAT(m->entry_computation()->instructions(), + Contains(op::Negate(op::GetTupleElement(op::CustomCall(), 1)))); + EXPECT_THAT(transformed_while->while_body()->instructions(), + Each(Not(AnyOf(op::CustomCall(), op::Negate())))); +} + +// Runs the pass with hoist_size_inflation_ratio = 1.0 and counts the +// shape_size_function calls, which is how often the size inflation test +// visits an array leaf. +class ShapeSizeCallCountTest : public WhileLoopInvariantCodeMotionTest { + protected: + static constexpr int kStateWidth = 64; + // Loose linear bound on shape_size_function calls: the fixed pass makes + // 2 * kStateWidth + 2 calls in the test below; a walk of the tuple shape per + // user makes about kStateWidth * kStateWidth. + static constexpr int kMaxLinearCalls = 4 * kStateWidth + 16; + + absl::StatusOr RunWithCountingShapeSizeFunction(HloModule* module) { + return WhileLoopInvariantCodeMotion( + /*hoist_constants=*/false, /*hoist_reshapes=*/false, + /*hoist_other=*/true, /*hoist_size_inflation_ratio=*/1.0, + /*shape_size_function=*/ + [this](const Shape& shape) { + ++shape_size_calls_; + return ShapeUtil::ByteSizeOfElements(shape); + }) + .Run(module); + } + + int64_t shape_size_calls_ = 0; +}; + +TEST_F(ShapeSizeCallCountTest, LinearInLoopStateWidth) { + auto m = CreateNewVerifiedModule(); + Shape array_f32 = ShapeUtil::MakeShape(F32, {8}); + Shape while_shape = + ShapeUtil::MakeTupleShape(std::vector(kStateWidth, array_f32)); + + // All state elements are loop invariant except the last, which receives + // negate(element 0): the single hoisting candidate. + HloComputation* while_body = [&]() { + HloComputation::Builder builder(TestName() + ".while_body"); + HloInstruction* param = builder.AddInstruction( + HloInstruction::CreateParameter(0, while_shape, "param")); + std::vector root_operands; + for (int i = 0; i + 1 < kStateWidth; ++i) { + root_operands.push_back(builder.AddInstruction( + HloInstruction::CreateGetTupleElement(array_f32, param, i))); + } + root_operands.push_back(builder.AddInstruction(HloInstruction::CreateUnary( + array_f32, HloOpcode::kNegate, root_operands.front()))); + builder.AddInstruction(HloInstruction::CreateTuple(root_operands)); + return m->AddEmbeddedComputation(builder.Build()); + }(); + + HloComputation::Builder builder(TestName()); + HloInstruction* init_value = builder.AddInstruction( + HloInstruction::CreateParameter(0, while_shape, "init_value")); + builder.AddInstruction(HloInstruction::CreateWhile( + while_shape, MakeAlwaysTrueComputation(while_shape, m.get()), while_body, + init_value)); + HloComputation* entry_computation = m->AddEntryComputation(builder.Build()); + + ASSERT_OK_AND_ASSIGN(bool simplified_loop, + RunWithCountingShapeSizeFunction(m.get())); + EXPECT_TRUE(simplified_loop); + EXPECT_THAT(entry_computation->instructions(), Contains(op::Negate())); + + // Every get-tuple-element of the state has the parameter as its operand and + // cannot be hoisted. Walking the kStateWidth element parameter shape for + // each of them would cost kStateWidth * kStateWidth calls. + EXPECT_GT(shape_size_calls_, 0); + EXPECT_LE(shape_size_calls_, kMaxLinearCalls); } TEST_F(WhileLoopInvariantCodeMotionTest, DoesNotHoistSPMDFullToShardShape) {