[XLA] Replace convert(constant) with constant

PiperOrigin-RevId: 652328778

Author: tensorflower-gardener

third_party/xla/xla/service/algebraic_simplifier.cc

@@ -5200,6 +5200,19 @@ absl::Status AlgebraicSimplifierVisitor::HandleConvert(
                               convert->mutable_operand(0)->mutable_operand(0));
   }
 
+  // Try to replace convert(constant) with a constant of the right type to begin
+  // with. Disallow moving int4 since it is not supported for many ops
+  HloInstruction* constant;
+  if (Match(convert, m::Convert(m::Constant(&constant))) &&
+      !primitive_util::IsSubByteNonPredType(src_type) &&
+      !primitive_util::IsSubByteNonPredType(dest_type)) {
+    TF_ASSIGN_OR_RETURN(Literal dest_literal,
+                        constant->literal().Convert(dest_type));
+    VLOG(10) << "Replacing convert(constant) with constant";
+    return ReplaceWithNewInstruction(
+        convert, HloInstruction::CreateConstant(std::move(dest_literal)));
+  }
+
   return TryRemoveUpcastAndDowncastSurroundingBinaryOp(convert);
 }
 

third_party/xla/xla/service/algebraic_simplifier_test.cc

@@ -11676,5 +11676,43 @@ ENTRY main.1 {
             HloOpcode::kParameter);
 }
 
+TEST_F(AlgebraicSimplifierTest, RemoveConvertConstant) {
+  const std::string hlo_string = R"(
+    HloModule module
+
+    add {
+      p0 = f32[] parameter(0)
+      p1 = f32[] parameter(1)
+      ROOT r = f32[] add(p0, p1)
+    }
+
+    ENTRY test {
+        a = f32[32,64] parameter(0)
+        b = s32[] constant(0)
+        c = f32[] convert(b)
+        ROOT reduce = f32[32] reduce(a, c),
+                      dimensions={1}, to_apply=add
+      }
+    )";
+  TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(hlo_string));
+  EXPECT_TRUE(AlgebraicSimplifier(default_options_).Run(m.get()).value());
+  HloInstruction* root = m->entry_computation()->root_instruction();
+  EXPECT_THAT(root, GmockMatch(m::Reduce(m::Parameter(0),
+                                         m::Constant().WithShape(F32, {}))));
+}
+
+TEST_F(AlgebraicSimplifierTest, KeepInt4ConvertConstant) {
+  const std::string hlo_string = R"(
+    HloModule module
+
+    ENTRY test {
+        a = s4[] constant(0)
+        ROOT b = s8[] convert(a)
+      }
+    )";
+  TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(hlo_string));
+  ASSERT_FALSE(AlgebraicSimplifier(default_options_).Run(m.get()).value());
+}
+
 }  // namespace
 }  // namespace xla