diff --git a/xla/service/hlo_verifier.cc b/xla/service/hlo_verifier.cc index 22592935ec498..3bbb4c82b59bf 100644 --- a/xla/service/hlo_verifier.cc +++ b/xla/service/hlo_verifier.cc @@ -1949,26 +1949,6 @@ absl::Status ShapeVerifier::CheckShape( } return ShapesSame(instruction->shape(), inferred_shape, equal); } - case HloOpcode::kCopy: { - // Disallow host offloading copies which change FpPrecision. - if (opts_.IsLayoutSensitive()) { - if (instruction->shape().has_layout() && - inferred_shape.has_layout()) { - int64_t instruction_memory_space = - instruction->shape().layout().memory_space(); - int64_t operand_memory_space = - inferred_shape.layout().memory_space(); - if (instruction_memory_space != operand_memory_space && - (instruction_memory_space == Layout::kHostMemorySpace || - operand_memory_space == Layout::kHostMemorySpace)) { - // Is a host->device copy for a device->host copy. - return Shape::Equal().IgnoreMemorySpaceInLayout()( - instruction->shape(), inferred_shape); - } - } - } - [[fallthrough]]; - } // We allow arbitrary layout and f32->bf16 transformations on all other // instructions, although this may be made more strict pending discussion diff --git a/xla/service/hlo_verifier_test.cc b/xla/service/hlo_verifier_test.cc index 07c47980640b7..77712517ed633 100644 --- a/xla/service/hlo_verifier_test.cc +++ b/xla/service/hlo_verifier_test.cc @@ -83,15 +83,6 @@ class HloVerifierTestLayoutSensitive : public HloTestBase { LayoutAssignment::InstructionCanChangeLayout) {} }; -class HloVerifierTestLayoutSensitiveAndAllowMixedPrecision - : public HloTestBase { - public: - HloVerifierTestLayoutSensitiveAndAllowMixedPrecision() - : HloTestBase(/*verifier_layout_sensitive=*/true, - /*allow_mixed_precision_in_hlo_verifier=*/true, - LayoutAssignment::InstructionCanChangeLayout) {} -}; - class HloVerifierTestLayoutFusion : public HloTestBase { public: HloVerifierTestLayoutFusion() @@ -3339,49 +3330,6 @@ TEST_F(HloVerifierTestLayoutSensitive, "memory space from device to host")); } -TEST_F(HloVerifierTestLayoutSensitiveAndAllowMixedPrecision, - HostOffloadingCopyCannotChangeType) { - const char* const hlo_string = R"( -HloModule m - -ENTRY main { - param = f32[1024,1024]{1,0:T(8,128)S(5)} parameter(0) - copy = bf16[1024,1024]{1,0:T(8,128)} copy(param) - ROOT dot = f32[1024,1024]{1,0:T(8,128)} dot(copy, copy), lhs_contracting_dims={1}, rhs_contracting_dims={0} -} -)"; - TF_ASSERT_OK_AND_ASSIGN(auto module, - ParseAndReturnUnverifiedModule(hlo_string)); - - auto status = verifier().Run(module.get()).status(); - ASSERT_FALSE(status.ok()); - EXPECT_THAT(status.message(), - HasSubstr("Expected instruction to have shape equal to " - "f32[1024,1024]{1,0:T(8,128)S(5)}, actual shape is " - "bf16[1024,1024]{1,0:T(8,128)}")); -} - -TEST_F(HloVerifierTestLayoutSensitiveAndAllowMixedPrecision, - HostOffloadingCopyCannotChangeLayout) { - const char* const hlo_string = R"( -HloModule m - -ENTRY main { - param = f32[1024,1024]{1,0:T(8,128)S(5)} parameter(0) - ROOT copy = f32[1024,1024]{0,1:T(8,128)} copy(param) -} -)"; - TF_ASSERT_OK_AND_ASSIGN(auto module, - ParseAndReturnUnverifiedModule(hlo_string)); - - auto status = verifier().Run(module.get()).status(); - ASSERT_FALSE(status.ok()); - EXPECT_THAT(status.message(), - HasSubstr("Expected instruction to have shape equal to " - "f32[1024,1024]{1,0:T(8,128)S(5)}, actual shape is " - "f32[1024,1024]{0,1:T(8,128)}")); -} - TEST_F(HloVerifierTestLayoutSensitive, MismatchedMinorToMajorSizeAndDimensionSize) { const char* const hlo_string = R"(