| 30 | class DynamicIndexSplitterTest : public HloTestBase {}; |
| 31 | |
| 32 | TEST_F(DynamicIndexSplitterTest, DynamicSlice) { |
| 33 | const char* const kDynamicSlice = R"( |
| 34 | HloModule DynamicSlice_module |
| 35 | |
| 36 | ENTRY entry (operand: s32[4,5,6], indices: s32[3]) -> s32[1,1,1] { |
| 37 | operand = s32[4,5,6] parameter(0) |
| 38 | indices = s32[3] parameter(1) |
| 39 | ROOT dynamic-slice = s32[1,1,1] dynamic-slice(operand, indices), dynamic_slice_sizes={1,1,1} |
| 40 | } |
| 41 | )"; |
| 42 | |
| 43 | HloModuleConfig config; |
| 44 | DebugOptions debug_options = config.debug_options(); |
| 45 | debug_options.set_xla_allow_scalar_index_dynamic_ops(true); |
| 46 | config.set_debug_options(debug_options); |
| 47 | |
| 48 | TF_ASSERT_OK_AND_ASSIGN(auto module, |
| 49 | ParseAndReturnVerifiedModule(kDynamicSlice, config)); |
| 50 | TF_ASSERT_OK_AND_ASSIGN(bool changed, |
| 51 | DynamicIndexSplitter().Run(module.get())); |
| 52 | EXPECT_TRUE(changed); |
| 53 | ASSERT_THAT(module->entry_computation()->root_instruction(), |
| 54 | op::DynamicSlice(op::Parameter(0), |
| 55 | op::Reshape(op::Slice(op::Parameter(1))), |
| 56 | op::Reshape(op::Slice(op::Parameter(1))), |
| 57 | op::Reshape(op::Slice(op::Parameter(1))))); |
| 58 | |
| 59 | for (int i = 0; i < 3; ++i) { |
| 60 | const HloInstruction* slice = module->entry_computation() |
| 61 | ->root_instruction() |
| 62 | ->operand(i + 1) |
| 63 | ->operand(0); |
| 64 | EXPECT_EQ(slice->slice_starts(0), i); |
| 65 | EXPECT_EQ(slice->slice_limits(0), i + 1); |
| 66 | } |
| 67 | } |
| 68 | |
| 69 | TEST_F(DynamicIndexSplitterTest, DynamicUpdateSlice) { |
| 70 | const char* const kDynamicUpdateSlice = R"( |
nothing calls this directly
no test coverage detected