| 43 | class SliceSinkerTest : public HloTestBase {}; |
| 44 | |
| 45 | TEST_F(SliceSinkerTest, TernaryOperation) { |
| 46 | const char* kModuleStr = R"( |
| 47 | HloModule m |
| 48 | test { |
| 49 | p0 = pred[8,9] parameter(0) |
| 50 | p1 = f32[8,9] parameter(1) |
| 51 | p2 = f32[8,9] parameter(2) |
| 52 | s00 = pred[2,9] slice(pred[8,9] p0), slice={[0:2], [0:9]} |
| 53 | s01 = pred[6,9] slice(pred[8,9] p0), slice={[2:8], [0:9]} |
| 54 | s10 = f32[2,9] slice(f32[8,9] p1), slice={[0:2], [0:9]} |
| 55 | s11 = f32[6,9] slice(f32[8,9] p1), slice={[2:8], [0:9]} |
| 56 | s20 = f32[2,9] slice(f32[8,9] p2), slice={[0:2], [0:9]} |
| 57 | s21 = f32[6,9] slice(f32[8,9] p2), slice={[2:8], [0:9]} |
| 58 | sel0 = f32[2,9] select(pred[2,9] s00, f32[2,9] s10, f32[2,9] s20) |
| 59 | sel1 = f32[6,9] select(pred[6,9] s01, f32[6,9] s11, f32[6,9] s21) |
| 60 | ROOT tuple = (f32[2,9], f32[6,9]) tuple(sel0, sel1) |
| 61 | } |
| 62 | )"; |
| 63 | TF_ASSERT_OK_AND_ASSIGN(auto module, |
| 64 | ParseAndReturnVerifiedModule(kModuleStr)); |
| 65 | SliceSinker slice_sinker; |
| 66 | TF_ASSERT_OK_AND_ASSIGN(bool result, RunHloPass(&slice_sinker, module.get())); |
| 67 | EXPECT_TRUE(result); |
| 68 | HloInstruction* inst = module->entry_computation()->root_instruction(); |
| 69 | const HloInstruction* slice0; |
| 70 | const HloInstruction* slice1; |
| 71 | EXPECT_THAT(inst, |
| 72 | GmockMatch(m::Tuple( |
| 73 | m::Slice(&slice0, m::Select(m::Parameter(0), m::Parameter(1), |
| 74 | m::Parameter(2))), |
| 75 | m::Slice(&slice1, m::Select(m::Parameter(0), m::Parameter(1), |
| 76 | m::Parameter(2)))))); |
| 77 | EXPECT_THAT(slice0->slice_starts(), ElementsAre(0, 0)); |
| 78 | EXPECT_THAT(slice0->slice_limits(), ElementsAre(2, 9)); |
| 79 | EXPECT_THAT(slice0->slice_strides(), ElementsAre(1, 1)); |
| 80 | EXPECT_THAT(slice1->slice_starts(), ElementsAre(2, 0)); |
| 81 | EXPECT_THAT(slice1->slice_limits(), ElementsAre(8, 9)); |
| 82 | EXPECT_THAT(slice1->slice_strides(), ElementsAre(1, 1)); |
| 83 | } |
| 84 | |
| 85 | TEST_F(SliceSinkerTest, OverlappingPartialSlicesBeneficial) { |
| 86 | const char* kModuleStr = R"( |
nothing calls this directly
no test coverage detected