Disable bfloat16 propagation for dynamic slice op with operand in host memory

fhoushmand · tensorflower-gardener · commit 0a2dff13ae9d · 2025-06-02T16:28:31.000-07:00
PiperOrigin-RevId: 766384513
diff --git a/third_party/xla/xla/hlo/transforms/bfloat16_propagation.cc b/third_party/xla/xla/hlo/transforms/bfloat16_propagation.cc
@@ -34,6 +34,7 @@ limitations under the License.
 #include "xla/hlo/ir/hlo_opcode.h"
 #include "xla/hlo/transforms/simplifiers/hlo_dce.h"
 #include "xla/hlo/transforms/simplifiers/tuple_simplifier.h"
+#include "xla/layout.h"
 #include "xla/literal.h"
 #include "xla/map_util.h"
 #include "xla/service/float_support.h"
@@ -568,6 +569,12 @@ bool BFloat16Propagation::InstructionIsCandidateForBF16Output(
       }
     }
   }
+  if (hlo->opcode() == HloOpcode::kDynamicSlice &&
+      hlo->operand(0)->shape().has_layout() &&
+      hlo->operand(0)->shape().layout().memory_space() ==
+          Layout::kHostMemorySpace) {
+    return false;
+  }
   return true;
 }
 
diff --git a/third_party/xla/xla/hlo/transforms/bfloat16_propagation_test.cc b/third_party/xla/xla/hlo/transforms/bfloat16_propagation_test.cc
@@ -1252,6 +1252,29 @@ ENTRY main {
   EXPECT_FALSE(OutputsBF16(dus));
 }
 
+TEST_F(BFloat16PropagationTest, DynamicSliceWithHostMemory) {
+  // In the case of dynamic-slice from host memory, we should not propagate
+  // bf16.
+  const std::string module_str = R"(
+  HloModule Module
+
+  ENTRY main {
+    param = f32[128,128]{1,0:S(5)} parameter(0)
+    constant.3 = s32[] constant(0)
+    dynamic-slice = f32[128,8] dynamic-slice(param, constant.3, constant.3), dynamic_slice_sizes={128,8}
+    ROOT dot = f32[128,128] dot(dynamic-slice, dynamic-slice), lhs_contracting_dims={1}, rhs_contracting_dims={1}
+  }
+  )";
+
+  TF_ASSERT_OK_AND_ASSIGN(std::unique_ptr<HloModule> module,
+                          ParseAndReturnVerifiedModule(module_str));
+  EXPECT_FALSE(PropagatePrecision(module.get()));
+
+  HloInstruction* dus =
+      module->entry_computation()->GetInstructionWithName("dynamic-slice");
+  EXPECT_FALSE(OutputsBF16(dus));
+}
+
 // This test demonstrates the need for invoking the ResolveAliasingBuffer
 // multiple times via a fixed-point algorithm. The key was the aliasing of the
 // two output buffers of the conditional, at subshape 0 (first element). This

Original file line number	Diff line number	Diff line change
`@@ -34,6 +34,7 @@ limitations under the License.`
`34`	`34`	`#include "xla/hlo/ir/hlo_opcode.h"`
`35`	`35`	`#include "xla/hlo/transforms/simplifiers/hlo_dce.h"`
`36`	`36`	`#include "xla/hlo/transforms/simplifiers/tuple_simplifier.h"`
	`37`	`+#include "xla/layout.h"`
`37`	`38`	`#include "xla/literal.h"`
`38`	`39`	`#include "xla/map_util.h"`
`39`	`40`	`#include "xla/service/float_support.h"`
`@@ -568,6 +569,12 @@ bool BFloat16Propagation::InstructionIsCandidateForBF16Output(`
`568`	`569`	`}`
`569`	`570`	`}`
`570`	`571`	`}`
	`572`	`+ if (hlo->opcode() == HloOpcode::kDynamicSlice &&`
	`573`	`+ hlo->operand(0)->shape().has_layout() &&`
	`574`	`+ hlo->operand(0)->shape().layout().memory_space() ==`
	`575`	`+ Layout::kHostMemorySpace) {`
	`576`	`+ return false;`
	`577`	`+ }`
`571`	`578`	`return true;`
`572`	`579`	`}`
`573`	`580`