PR #17579: Algebraic simplifier: mark iota non-negative.

Imported from GitHub PR #17579 Copybara import of the project: -- 02c09a8 by Ilia Sergachev <isergachev@nvidia.com>: Algebraic simplifier: mark iota non-negative. -- 4735edc by Ilia Sergachev <isergachev@nvidia.com>: Fix unrelated clang-format issues to make CI happy Merging this change closes #17579 COPYBARA_INTEGRATE_REVIEW=#17579 from openxla:iota_non_neg 4735edc PiperOrigin-RevId: 678640567
openxla · Sep 25, 2024 · 464d309 · 464d309
1 parent d31a740
commit 464d309
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 1 deletion.
diff --git a/xla/service/algebraic_simplifier.cc b/xla/service/algebraic_simplifier.cc
@@ -530,7 +530,8 @@ bool AlgebraicSimplifierVisitor::IsNonNegative(
       return hlo->operand(0) == hlo->operand(1);
     }
     case HloOpcode::kAbs:
-    case HloOpcode::kExp: {
+    case HloOpcode::kExp:
+    case HloOpcode::kIota: {
       return true;
     }
     case HloOpcode::kBroadcast: {

diff --git a/xla/service/algebraic_simplifier_test.cc b/xla/service/algebraic_simplifier_test.cc
@@ -10561,6 +10561,18 @@ TEST_F(AlgebraicSimplifierTest, AbsEliminationSelMaxBcast) {
                                      m::Broadcast(m::ConstantScalar())))));
 }
 
+TEST_F(AlgebraicSimplifierTest, AbsEliminationIota) {
+  TF_ASSERT_OK_AND_ASSIGN(auto m, ParseAndReturnVerifiedModule(R"(
+    e {
+      i = s32[3,2] iota(), iota_dimension=0
+      ROOT a = s32[3,2] abs(i)
+    }
+  )"));
+  ASSERT_TRUE(AlgebraicSimplifier(default_options_).Run(m.get()).value());
+  EXPECT_THAT(m->entry_computation()->root_instruction(),
+              GmockMatch(m::Iota()));
+}
+
 TEST_F(AlgebraicSimplifierTest, SimplifyRedundantBitcastConvert) {
   const char* kModuleStr = R"(
     HloModule m