| 31 | using SortSimplifierTest = HloTestBase; |
| 32 | |
| 33 | TEST_F(SortSimplifierTest, RemoveUnusedSortOperandArrayResult) { |
| 34 | const char* hlo_string = R"( |
| 35 | HloModule permutation_sort |
| 36 | |
| 37 | compare { |
| 38 | p.0.lhs = f32[] parameter(0) |
| 39 | p.0.rhs = f32[] parameter(1) |
| 40 | p.1.lhs = s32[] parameter(2) |
| 41 | p.1.rhs = s32[] parameter(3) |
| 42 | ROOT lt = pred[] compare(p.0.lhs, p.0.rhs), direction=LT |
| 43 | } |
| 44 | |
| 45 | ENTRY sort_computation { |
| 46 | keys = f32[64,8732]{1,0} parameter(0) |
| 47 | values = s32[64,8732]{1,0} parameter(1) |
| 48 | sort = (f32[64,8732]{1,0}, s32[64,8732]{1,0}) sort(keys, values), |
| 49 | dimensions={1}, to_apply=compare |
| 50 | ROOT gte = f32[64,8732]{1,0} get-tuple-element(sort), index=0 |
| 51 | })"; |
| 52 | TF_ASSERT_OK_AND_ASSIGN(auto module, |
| 53 | ParseAndReturnVerifiedModule(hlo_string)); |
| 54 | |
| 55 | SortSimplifier simplifier; |
| 56 | uint64 num_executions = 0; |
| 57 | do { |
| 58 | num_executions++; |
| 59 | } while (simplifier.Run(module.get()).ValueOrDie()); |
| 60 | EXPECT_EQ(num_executions, 2); |
| 61 | auto root = module->entry_computation()->root_instruction(); |
| 62 | EXPECT_THAT(root, GmockMatch(m::Sort(m::Parameter(0)))); |
| 63 | } |
| 64 | |
| 65 | TEST_F(SortSimplifierTest, RemoveUnusedSortOperandTuple) { |
| 66 | const char* hlo_string = R"( |
nothing calls this directly
no test coverage detected