Handle ragged dot in precision config methods.

PiperOrigin-RevId: 696699435
openxla · Nov 15, 2024 · fb0fa1f · fb0fa1f
1 parent a7aff7f
commit fb0fa1f
Show file tree

Hide file tree

Showing 2 changed files with 24 additions and 0 deletions.
diff --git a/xla/hlo/ir/hlo_instruction.cc b/xla/hlo/ir/hlo_instruction.cc
@@ -5239,6 +5239,9 @@ const PrecisionConfig& HloInstruction::precision_config() const {
   if (auto* dot = DynCast<HloDotInstruction>(this)) {
     return dot->precision_config();
   }
+  if (auto* ragged_dot = DynCast<HloRaggedDotInstruction>(this)) {
+    return ragged_dot->precision_config();
+  }
 
   if (auto* custom_call = DynCast<HloCustomCallInstruction>(this)) {
     return custom_call->precision_config();
@@ -5253,6 +5256,9 @@ PrecisionConfig* HloInstruction::mutable_precision_config() {
   if (auto* dot = DynCast<HloDotInstruction>(this)) {
     return dot->mutable_precision_config();
   }
+  if (auto* ragged_dot = DynCast<HloRaggedDotInstruction>(this)) {
+    return ragged_dot->mutable_precision_config();
+  }
   if (auto* custom_call = DynCast<HloCustomCallInstruction>(this)) {
     return custom_call->mutable_precision_config();
   }

diff --git a/xla/service/hlo_instruction_test.cc b/xla/service/hlo_instruction_test.cc
@@ -3181,5 +3181,23 @@ TEST_F(HloInstructionTest, UnfuseInstructionWithConstantOperand) {
               GmockMatch(m::Add(m::Parameter(0), m::Broadcast(m::Constant()))));
 }
 
+TEST_F(HloInstructionTest, RaggedDotHasPrecisionConfig) {
+  constexpr char kHloString[] = R"(
+  HloModule module
+  ENTRY entry_computation {
+    a = f32[11,5] parameter(0)
+    b = f32[3,5,7] parameter(1)
+    c = u32[3] parameter(2)
+    ROOT dot = f32[11,7] ragged-dot(a, b, c), lhs_contracting_dims={1}, rhs_contracting_dims={1}, lhs_ragged_dims={0}, rhs_group_dims={0}
+  })";
+  TF_ASSERT_OK_AND_ASSIGN(auto module,
+                          ParseAndReturnVerifiedModule(kHloString));
+  auto* ragged_dot = module->entry_computation()->root_instruction();
+
+  EXPECT_THAT(ragged_dot->precision_config().operand_precision(),
+              ::testing::ElementsAre(PrecisionConfig::DEFAULT,
+                                     PrecisionConfig::DEFAULT));
+}
+
 }  // namespace
 }  // namespace xla