From fc76ae0cb04e6a8c5a21bc5bdd4b488a4e10a6b3 Mon Sep 17 00:00:00 2001
From: Victor Stone <victorstone@google.com>
Date: Thu, 12 Sep 2024 16:33:19 -0700
Subject: [PATCH] Remove host offloading copy checks from HloVerifier. Due to
 float propagation and float normalization being separate passes, host
 offloading copies can be multi-precision (at least temporarily). Float
 normalization will fix that and insert a convert.

Reverts 87494d1fd1936f5552acd0455693775742c49631

PiperOrigin-RevId: 674055866
---
 xla/service/hlo_verifier.cc      | 20 ------------
 xla/service/hlo_verifier_test.cc | 52 --------------------------------
 2 files changed, 72 deletions(-)

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"(