@@ -63,6 +63,19 @@ class ProfilerETDumpTest : public ::testing::Test {
6363 free (buf);
6464 }
6565
66+ // Triggers ET_EXPECT_DEATH if log_intermediate_output_delegate has no
67+ // data sink
68+ void expect_log_intermediate_delegate_death (
69+ ETDumpGen* gen,
70+ TensorFactory<ScalarType::Float>& tf) {
71+ ET_EXPECT_DEATH (
72+ gen->log_intermediate_output_delegate (
73+ " test_event_tensor" ,
74+ static_cast <torch::executor::DebugHandle>(-1 ),
75+ tf.ones ({3 , 2 })),
76+ " Must set data sink before writing tensor-like data" );
77+ }
78+
6679 ETDumpGen* etdump_gen[2 ];
6780 uint8_t * buf = nullptr ;
6881 std::unique_ptr<TempFile> temp_file;
@@ -188,11 +201,16 @@ TEST_F(ProfilerETDumpTest, AllocationEvents) {
188201}
189202
190203TEST_F (ProfilerETDumpTest, DebugEvent) {
204+ const size_t debug_buf_size = 2048 ;
205+ const size_t etdump_buf_size = 512 * 1024 ;
206+ ASSERT_NE (this ->buf , nullptr );
207+ Span<uint8_t > span_buf = Span<uint8_t >(this ->buf , etdump_buf_size);
208+
191209 for (size_t i = 0 ; i < 2 ; i++) {
192210 for (size_t j = 0 ; j < 3 ; j++) {
193211 etdump_gen[i]->create_event_block (" test_block" );
194212
195- void * ptr = malloc (2048 );
213+ void * ptr = malloc (debug_buf_size );
196214
197215 EValue evalue_int ((int64_t )5 );
198216 etdump_gen[i]->log_evalue (evalue_int);
@@ -208,24 +226,37 @@ TEST_F(ProfilerETDumpTest, DebugEvent) {
208226 TensorFactory<ScalarType::Float> tf;
209227 EValue evalue_tensor (tf.ones ({3 , 2 }));
210228
211- // using span to record debug data
212- Span<uint8_t > buffer ((uint8_t *)ptr, 2048 );
213- auto buffer_data_sink = BufferDataSink::create (ptr, 2048 );
229+ // Create span to record debug data
230+ Span<uint8_t > buffer ((uint8_t *)ptr, debug_buf_size );
231+ auto buffer_data_sink = BufferDataSink::create (ptr, debug_buf_size );
214232 auto file_data_sink = FileDataSink::create (dump_file_path.c_str ());
215233
216234 if (j == 0 ) {
217235 ET_EXPECT_DEATH (
218236 etdump_gen[i]->log_evalue (evalue_tensor),
219237 " Must set data sink before writing tensor-like data" );
238+
239+ // Set debug buffer with span
220240 etdump_gen[i]->set_debug_buffer (buffer);
221- }
222- // using buffer data sink to record debug data
223- else if (j == 1 ) {
224- etdump_gen[i]->set_data_sink (&buffer_data_sink.get ());
225- }
226- // using file data sink to record debug data
227- else {
228- etdump_gen[i]->set_data_sink (&file_data_sink.get ());
241+ } else {
242+ // Reset ETDumpGen to trigger ET_EXPECT_DEATH before setting data sink
243+ delete etdump_gen[i];
244+
245+ // Recreate ETDumpGen; set span buffer only for etdump_gen[1]
246+ etdump_gen[i] = (i == 0 ) ? new ETDumpGen () : new ETDumpGen (span_buf);
247+ etdump_gen[i]->create_event_block (" test_block" );
248+
249+ ET_EXPECT_DEATH (
250+ etdump_gen[i]->log_evalue (evalue_tensor),
251+ " Must set data sink before writing tensor-like data" );
252+
253+ if (j == 1 ) {
254+ // Set buffer data sink
255+ etdump_gen[i]->set_data_sink (&buffer_data_sink.get ());
256+ } else {
257+ // Set file data sink
258+ etdump_gen[i]->set_data_sink (&file_data_sink.get ());
259+ }
229260 }
230261
231262 etdump_gen[i]->log_evalue (evalue_tensor);
@@ -500,35 +531,43 @@ TEST_F(ProfilerETDumpTest, VerifyData) {
500531}
501532
502533TEST_F (ProfilerETDumpTest, LogDelegateIntermediateOutput) {
534+ const size_t debug_buf_size = 2048 ;
535+ const size_t etdump_buf_size = 512 * 1024 ;
536+ ASSERT_NE (this ->buf , nullptr );
537+ Span<uint8_t > span_buf = Span<uint8_t >(this ->buf , etdump_buf_size);
538+
503539 for (size_t i = 0 ; i < 2 ; i++) {
504540 for (size_t j = 0 ; j < 3 ; j++) {
505- void * ptr = malloc (2048 );
506- Span<uint8_t > buffer ((uint8_t *)ptr, 2048 );
541+ void * ptr = malloc (debug_buf_size );
542+ Span<uint8_t > buffer ((uint8_t *)ptr, debug_buf_size );
507543
508- auto buffer_data_sink = BufferDataSink::create (ptr, 2048 );
544+ auto buffer_data_sink = BufferDataSink::create (ptr, debug_buf_size );
509545 auto file_data_sink = FileDataSink::create (dump_file_path.c_str ());
510546
511- etdump_gen[i]->create_event_block (" test_block" );
512547 TensorFactory<ScalarType::Float> tf;
513548
514- // using span to record debug data
515549 if (j == 0 ) {
516- // TODO(gasoonjia): add similar ET_EXPECT_DEATH on BufferDataSink branch
517- ET_EXPECT_DEATH (
518- etdump_gen[i]->log_intermediate_output_delegate (
519- " test_event_tensor" ,
520- static_cast <torch::executor::DebugHandle>(-1 ),
521- tf.ones ({3 , 2 })),
522- " Must set data sink before writing tensor-like data" );
550+ etdump_gen[i]->create_event_block (" test_block" );
551+ expect_log_intermediate_delegate_death (etdump_gen[i], tf);
552+
553+ // Set debug buffer with span
523554 etdump_gen[i]->set_debug_buffer (buffer);
524- }
525- // using buffer data sink to record debug data
526- else if (j == 1 ) {
527- etdump_gen[i]->set_data_sink (&buffer_data_sink.get ());
528- }
529- // using file data sink to record debug data
530- else {
531- etdump_gen[i]->set_data_sink (&file_data_sink.get ());
555+ } else {
556+ // Reset ETDumpGen to trigger ET_EXPECT_DEATH before setting data sink
557+ delete etdump_gen[i];
558+
559+ // Recreate ETDumpGen; set span buffer only for etdump_gen[1]
560+ etdump_gen[i] = (i == 0 ) ? new ETDumpGen () : new ETDumpGen (span_buf);
561+ etdump_gen[i]->create_event_block (" test_block" );
562+ expect_log_intermediate_delegate_death (etdump_gen[i], tf);
563+
564+ if (j == 1 ) {
565+ // Set buffer data sink
566+ etdump_gen[i]->set_data_sink (&buffer_data_sink.get ());
567+ } else {
568+ // Set file data sink
569+ etdump_gen[i]->set_data_sink (&file_data_sink.get ());
570+ }
532571 }
533572
534573 // Only a valid name or delegate debug index should be passed in. If valid
0 commit comments