Fix decoding multiple frames

Signed-off-by: Robert Winkler <[email protected]>
google · Jan 15, 2025 · 50f6140 · 50f6140
1 parent ddde31a
commit 50f6140
Show file tree

Hide file tree

Showing 11 changed files with 399 additions and 198 deletions.
diff --git a/xls/modules/zstd/command_constructor.x b/xls/modules/zstd/command_constructor.x
@@ -27,7 +27,7 @@ type CopyOrMatchContent = common::CopyOrMatchContent;
 type CopyOrMatchLength = common::CopyOrMatchLength;
 
 type SequenceExecutorPacket = common::SequenceExecutorPacket<common::SYMBOL_WIDTH>;
-type LiteralsBufferCtrl = common::LiteralsBufferCtrl; 
+type LiteralsBufferCtrl = common::LiteralsBufferCtrl;
 type CommandConstructorData = common::CommandConstructorData;
 type ExtendedBlockDataPacket = common::ExtendedBlockDataPacket;
 type BlockSyncData = common::BlockSyncData;
@@ -42,7 +42,7 @@ struct State {
     status: Status,
     received_literals: CopyOrMatchLength,
     literals_to_receive: CopyOrMatchLength,
-    sync: BlockSyncData,
+    command: CommandConstructorData,
 }
 
 pub proc CommandConstructor {
@@ -64,24 +64,29 @@ pub proc CommandConstructor {
         let tok0 = join();
 
         let recv_command = state.status == Status::RECV_COMMAND;
-        let (tok1_0, command) =
-            recv_if(tok0, sequence_decoder_r, recv_command, zero!<CommandConstructorData>());
+        let (tok1_0, command) = recv_if(tok0, sequence_decoder_r, recv_command, state.command);
+        if recv_command {
+            trace_fmt!("[CommandConstructor] Received command: {:#x}", command);
+        } else {};
 
         let recv_literals = state.status == Status::RECV_LITERALS;
-        let (tok1_1, literals) =
-            recv_if(tok0, literals_buffer_resp_r, recv_literals, zero!<SequenceExecutorPacket>());
+        let (tok1_1, literals) = recv_if(tok0, literals_buffer_resp_r, recv_literals, zero!<SequenceExecutorPacket>());
+        if recv_literals {
+            trace_fmt!("[CommandConstructor] Received literals: {:#x}", literals);
+        } else {};
 
         let tok1 = join(tok1_0, tok1_1);
 
         let (new_state, do_send_command, do_send_literals_req) = match (state.status) {
             Status::RECV_COMMAND => {
-                if command.data.msg_type == SequenceExecutorMessageType::LITERAL {
+                if (command.data.msg_type == SequenceExecutorMessageType::LITERAL) &&
+                   (command.data.length != CopyOrMatchLength:0) {
                     (
                         State {
                             status: Status::RECV_LITERALS,
                             received_literals: CopyOrMatchLength:0,
                             literals_to_receive: command.data.length,
-                            sync: command.sync,
+                            command: command,
                         }, false, true,
                     )
                 } else {
@@ -108,7 +113,7 @@ pub proc CommandConstructor {
         let resp = match(state.status) {
             // sent only if the original message was of type SEQUENCE
             Status::RECV_COMMAND => ExtendedBlockDataPacket {
-                msg_type: SequenceExecutorMessageType::SEQUENCE,
+                msg_type: command.data.msg_type,
                 packet: BlockDataPacket {
                     last: command.data.last,
                     last_block: command.sync.last_block,
@@ -120,16 +125,19 @@ pub proc CommandConstructor {
             Status::RECV_LITERALS => ExtendedBlockDataPacket {
                 msg_type: SequenceExecutorMessageType::LITERAL,
                 packet: BlockDataPacket {
-                    last: literals.last,
-                    last_block: state.sync.last_block,
-                    id: state.sync.id,
+                    last: literals.last & command.data.last,
+                    last_block: command.sync.last_block,
+                    id: command.sync.id,
                     data: literals.content,
                     length: literals.length as u32, // FIXME: remove cast after unifying types of 'length' fields
                 }
             },
             _ => fail!("resp_match_unreachable", zero!<ExtendedBlockDataPacket>())
         };
         send_if(tok1, command_aggregator_s, do_send_command, resp);
+        if do_send_command {
+            trace_fmt!("[CommandConstructor] Sending command: {:#x}", resp);
+        } else {};
 
         new_state
     }

diff --git a/xls/modules/zstd/comp_block_dec.x b/xls/modules/zstd/comp_block_dec.x
diff --git a/xls/modules/zstd/comp_frame.x b/xls/modules/zstd/comp_frame.x
@@ -5,17 +5,17 @@ pub struct DataArray<BITS_PER_WORD: u32, LENGTH: u32>{
 }
 pub const FRAMES:DataArray<
   u32:64,
-  u32:5
->[1] = [DataArray<64, 5>{
-  length: u32:33,
-  array_length: u32:5,
-  data: uN[64][5]:[uN[64]:0x001a3384fd2fb528, uN[64]:0xc1d3500000850000, uN[64]:0xdcf0529b98db8a06, uN[64]:0x308fa3120a430001, uN[64]:0x50]
+  u32:7
+>[1] = [DataArray<64, 7>{
+  length: u32:51,
+  array_length: u32:7,
+  data: uN[64][7]:[uN[64]:0x00504784fd2fb528, uN[64]:0xcf95700001150000, uN[64]:0xe17d50b989ac93c4, uN[64]:0x0daf000895a6e608, uN[64]:0xb96010b86f7602a4, uN[64]:0x05b0e051238666e8, uN[64]:0x8470e3]
 }];
 pub const DECOMPRESSED_FRAMES:DataArray<
   u32:64,
-  u32:4
->[1] = [DataArray<64, 4>{
-  length: u32:26,
-  array_length: u32:4,
-  data: uN[64][4]:[uN[64]:0x529b98db8a06c1d3, uN[64]:0x529b98db8a06dcf0, uN[64]:0x529b98db8a06dcf0, uN[64]:0xdcf0]
+  u32:10
+>[1] = [DataArray<64, 10>{
+  length: u32:80,
+  array_length: u32:10,
+  data: uN[64][10]:[uN[64]:0xc4c4cf95cf95cf95, uN[64]:0x93c4c4c4c4c4c4c4, uN[64]:0xacc493c493c493c4, uN[64]:0xc493c493c493c489, uN[64]:0x93c493c489acc493, uN[64]:0x08e17d50b9c493c4, uN[64]:0xc4c4c4cf9595a6e6, uN[64]:0x93c493c4c4c4c4c4, uN[64]:0xc489acc493c493c4, uN[64]:0xc493c493c493c493]
 }];
diff --git a/xls/modules/zstd/dec_mux.x b/xls/modules/zstd/dec_mux.x
@@ -130,7 +130,7 @@ pub proc DecoderMux {
                  (((state.raw_data.packet.id == (state.prev_id + u32:1)) && state.prev_last) ||
                   ((state.raw_data.packet.id == state.prev_id) && !state.prev_last))) {
                     assert!(!state.raw_data_valid_next_frame, "raw_packet_valid_in_current_and_next_frame");
-                    (true,
+                    ((state.raw_data.packet.length as CopyOrMatchLength) != CopyOrMatchLength:0,
                      SequenceExecutorPacket {
                          msg_type: state.raw_data.msg_type,
                          length: state.raw_data.packet.length as CopyOrMatchLength,
@@ -153,7 +153,7 @@ pub proc DecoderMux {
                         (((state.rle_data.packet.id == (state.prev_id + u32:1)) && state.prev_last) ||
                          ((state.rle_data.packet.id == state.prev_id) && !state.prev_last))) {
                     assert!(!state.rle_data_valid_next_frame, "rle_packet_valid_in_current_and_next_frame");
-                    (true,
+                    ((state.rle_data.packet.length as CopyOrMatchLength) != CopyOrMatchLength:0,
                      SequenceExecutorPacket {
                          msg_type: state.rle_data.msg_type,
                          length: state.rle_data.packet.length as CopyOrMatchLength,
@@ -282,9 +282,7 @@ proc DecoderMuxEmptyRawBlocksTest {
 
     let (tok, data) = recv(tok, output_r); assert_eq(data, SequenceExecutorPacket {last: bool: false, msg_type: SequenceExecutorMessageType::LITERAL, content: CopyOrMatchContent:0x11111111, length: CopyOrMatchLength:4 });
     let (tok, data) = recv(tok, output_r); assert_eq(data, SequenceExecutorPacket {last: bool: false, msg_type: SequenceExecutorMessageType::LITERAL, content: CopyOrMatchContent:0x22222222, length: CopyOrMatchLength:4 });
-    let (tok, data) = recv(tok, output_r); assert_eq(data, SequenceExecutorPacket {last: bool: false, msg_type: SequenceExecutorMessageType::LITERAL, content: CopyOrMatchContent:0x0,        length: CopyOrMatchLength:0  });
     let (tok, data) = recv(tok, output_r); assert_eq(data, SequenceExecutorPacket {last: bool: false, msg_type: SequenceExecutorMessageType::LITERAL, content: CopyOrMatchContent:0x33333333, length: CopyOrMatchLength:4 });
-    let (tok, data) = recv(tok, output_r); assert_eq(data, SequenceExecutorPacket {last: bool: true,  msg_type: SequenceExecutorMessageType::LITERAL, content: CopyOrMatchContent:0x0       , length: CopyOrMatchLength:0  });
 
     send(tok, terminator, true);
   }
@@ -320,9 +318,7 @@ proc DecoderMuxEmptyRleBlocksTest {
 
     let (tok, data) = recv(tok, output_r); assert_eq(data, SequenceExecutorPacket {last: bool: false, msg_type: SequenceExecutorMessageType::LITERAL, content: CopyOrMatchContent:0x11111111, length: CopyOrMatchLength:4 });
     let (tok, data) = recv(tok, output_r); assert_eq(data, SequenceExecutorPacket {last: bool: false, msg_type: SequenceExecutorMessageType::LITERAL, content: CopyOrMatchContent:0x22222222, length: CopyOrMatchLength:4 });
-    let (tok, data) = recv(tok, output_r); assert_eq(data, SequenceExecutorPacket {last: bool: false, msg_type: SequenceExecutorMessageType::LITERAL, content: CopyOrMatchContent:0x0,        length: CopyOrMatchLength:0  });
     let (tok, data) = recv(tok, output_r); assert_eq(data, SequenceExecutorPacket {last: bool: false, msg_type: SequenceExecutorMessageType::LITERAL, content: CopyOrMatchContent:0x33333333, length: CopyOrMatchLength:4 });
-    let (tok, data) = recv(tok, output_r); assert_eq(data, SequenceExecutorPacket {last: bool: true,  msg_type: SequenceExecutorMessageType::LITERAL, content: CopyOrMatchContent:0x0,        length: CopyOrMatchLength:0  });
 
     send(tok, terminator, true);
   }
@@ -361,7 +357,6 @@ proc DecoderMuxEmptyBlockBetweenRegularBlocksOnTheSameInputChannelTest {
     let (tok, data) = recv(tok, output_r); assert_eq(data, SequenceExecutorPacket {last: bool: false, msg_type: SequenceExecutorMessageType::LITERAL, content: CopyOrMatchContent:0x11111111, length: CopyOrMatchLength:4 });
     let (tok, data) = recv(tok, output_r); assert_eq(data, SequenceExecutorPacket {last: bool: false, msg_type: SequenceExecutorMessageType::LITERAL, content: CopyOrMatchContent:0x22222222, length: CopyOrMatchLength:4 });
     let (tok, data) = recv(tok, output_r); assert_eq(data, SequenceExecutorPacket {last: bool: false, msg_type: SequenceExecutorMessageType::LITERAL, content: CopyOrMatchContent:0x33333333, length: CopyOrMatchLength:4 });
-    let (tok, data) = recv(tok, output_r); assert_eq(data, SequenceExecutorPacket {last: bool: false, msg_type: SequenceExecutorMessageType::LITERAL, content: CopyOrMatchContent:0x0,        length: CopyOrMatchLength:0  });
     let (tok, data) = recv(tok, output_r); assert_eq(data, SequenceExecutorPacket {last: bool: false, msg_type: SequenceExecutorMessageType::LITERAL, content: CopyOrMatchContent:0xAAAAAAAA, length: CopyOrMatchLength:4 });
     let (tok, data) = recv(tok, output_r); assert_eq(data, SequenceExecutorPacket {last: bool: false, msg_type: SequenceExecutorMessageType::LITERAL, content: CopyOrMatchContent:0xBBBBBBBB, length: CopyOrMatchLength:4 });
     let (tok, data) = recv(tok, output_r); assert_eq(data, SequenceExecutorPacket {last: bool: true,  msg_type: SequenceExecutorMessageType::LITERAL, content: CopyOrMatchContent:0x00000000, length: CopyOrMatchLength:4 });
@@ -403,7 +398,6 @@ proc DecoderMuxEmptyBlockBetweenRegularBlocksOnDifferentInputChannelsTest {
     let (tok, data) = recv(tok, output_r); assert_eq(data, SequenceExecutorPacket {last: bool: false, msg_type: SequenceExecutorMessageType::LITERAL, content: CopyOrMatchContent:0x11111111, length: CopyOrMatchLength:4 });
     let (tok, data) = recv(tok, output_r); assert_eq(data, SequenceExecutorPacket {last: bool: false, msg_type: SequenceExecutorMessageType::LITERAL, content: CopyOrMatchContent:0x22222222, length: CopyOrMatchLength:4 });
     let (tok, data) = recv(tok, output_r); assert_eq(data, SequenceExecutorPacket {last: bool: false, msg_type: SequenceExecutorMessageType::LITERAL, content: CopyOrMatchContent:0x33333333, length: CopyOrMatchLength:4 });
-    let (tok, data) = recv(tok, output_r); assert_eq(data, SequenceExecutorPacket {last: bool: false, msg_type: SequenceExecutorMessageType::LITERAL, content: CopyOrMatchContent:0x0,        length: CopyOrMatchLength:0  });
     let (tok, data) = recv(tok, output_r); assert_eq(data, SequenceExecutorPacket {last: bool: false, msg_type: SequenceExecutorMessageType::LITERAL, content: CopyOrMatchContent:0xAAAAAAAA, length: CopyOrMatchLength:4 });
     let (tok, data) = recv(tok, output_r); assert_eq(data, SequenceExecutorPacket {last: bool: false, msg_type: SequenceExecutorMessageType::LITERAL, content: CopyOrMatchContent:0xBBBBBBBB, length: CopyOrMatchLength:4 });
     let (tok, data) = recv(tok, output_r); assert_eq(data, SequenceExecutorPacket {last: bool: true,  msg_type: SequenceExecutorMessageType::LITERAL, content: CopyOrMatchContent:0x00000000, length: CopyOrMatchLength:4 });
@@ -474,19 +468,9 @@ proc DecoderMuxMultipleFramesTest {
     // Frame #2
     let (tok, data) = recv(tok, output_r); assert_eq(data, SequenceExecutorPacket {last: bool: false, msg_type: SequenceExecutorMessageType::LITERAL, content: CopyOrMatchContent:0x44444444, length: CopyOrMatchLength:4 });
     let (tok, data) = recv(tok, output_r); assert_eq(data, SequenceExecutorPacket {last: bool: false, msg_type: SequenceExecutorMessageType::LITERAL, content: CopyOrMatchContent:0x0       , length: CopyOrMatchLength:0  });
-    let (tok, data) = recv(tok, output_r); assert_eq(data, SequenceExecutorPacket {last: bool: false, msg_type: SequenceExecutorMessageType::LITERAL, content: CopyOrMatchContent:0x0       , length: CopyOrMatchLength:0  });
-    let (tok, data) = recv(tok, output_r); assert_eq(data, SequenceExecutorPacket {last: bool: true,  msg_type: SequenceExecutorMessageType::LITERAL, content: CopyOrMatchContent:0x0       , length: CopyOrMatchLength:0  });
+    let (tok, data) = recv(tok, output_r); assert_eq(data, SequenceExecutorPacket {last: bool: false,  msg_type: SequenceExecutorMessageType::LITERAL, content: CopyOrMatchContent:0x0       , length: CopyOrMatchLength:0  });
     // Frame #3
     let (tok, data) = recv(tok, output_r); assert_eq(data, SequenceExecutorPacket {last: bool: false, msg_type: SequenceExecutorMessageType::LITERAL, content: CopyOrMatchContent:0x55555555, length: CopyOrMatchLength:4 });
-    let (tok, data) = recv(tok, output_r); assert_eq(data, SequenceExecutorPacket {last: bool: false, msg_type: SequenceExecutorMessageType::LITERAL, content: CopyOrMatchContent:0x0       , length: CopyOrMatchLength:0  });
-    let (tok, data) = recv(tok, output_r); assert_eq(data, SequenceExecutorPacket {last: bool: false, msg_type: SequenceExecutorMessageType::LITERAL, content: CopyOrMatchContent:0x0       , length: CopyOrMatchLength:0  });
-    let (tok, data) = recv(tok, output_r); assert_eq(data, SequenceExecutorPacket {last: bool: false, msg_type: SequenceExecutorMessageType::LITERAL, content: CopyOrMatchContent:0x0       , length: CopyOrMatchLength:0  });
-    let (tok, data) = recv(tok, output_r); assert_eq(data, SequenceExecutorPacket {last: bool: false, msg_type: SequenceExecutorMessageType::LITERAL, content: CopyOrMatchContent:0x0       , length: CopyOrMatchLength:0  });
-    let (tok, data) = recv(tok, output_r); assert_eq(data, SequenceExecutorPacket {last: bool: false, msg_type: SequenceExecutorMessageType::LITERAL, content: CopyOrMatchContent:0x0       , length: CopyOrMatchLength:0  });
-    let (tok, data) = recv(tok, output_r); assert_eq(data, SequenceExecutorPacket {last: bool: false, msg_type: SequenceExecutorMessageType::LITERAL, content: CopyOrMatchContent:0x0       , length: CopyOrMatchLength:0  });
-    let (tok, data) = recv(tok, output_r); assert_eq(data, SequenceExecutorPacket {last: bool: false, msg_type: SequenceExecutorMessageType::LITERAL, content: CopyOrMatchContent:0x0       , length: CopyOrMatchLength:0  });
-    let (tok, data) = recv(tok, output_r); assert_eq(data, SequenceExecutorPacket {last: bool: false, msg_type: SequenceExecutorMessageType::LITERAL, content: CopyOrMatchContent:0x0       , length: CopyOrMatchLength:0  });
-    let (tok, data) = recv(tok, output_r); assert_eq(data, SequenceExecutorPacket {last: bool: true,  msg_type: SequenceExecutorMessageType::LITERAL, content: CopyOrMatchContent:0x0       , length: CopyOrMatchLength:0  });
 
     send(tok, terminator, true);
   }

diff --git a/xls/modules/zstd/fse_dec.x b/xls/modules/zstd/fse_dec.x
@@ -343,13 +343,13 @@ pub proc FseDecoder<
             }
         };
         let do_send_command = (
-            (command_data.length != CopyOrMatchLength:0) &&
             (
                 state.fsm == FseDecoderFSM::SEND_COMMAND_LITERAL ||
                 state.fsm == FseDecoderFSM::SEND_COMMAND_SEQUENCE ||
                 state.fsm == FseDecoderFSM::SEND_LEFTOVER_LITERALS_REQ
             )
         );
+
         let command = CommandConstructorData {
             sync: state.ctrl.sync,
             data: command_data,

diff --git a/xls/modules/zstd/literals_buffer.x b/xls/modules/zstd/literals_buffer.x
@@ -709,7 +709,7 @@ proc LiteralsBufferReader<
             last: ctrl_last,
         };
 
-        let (read_reqs, read_start, read_len, packet, _) = parallel_rams::sequence_packet_to_read_reqs<
+        let (read_reqs, read_start, read_len, _, _) = parallel_rams::sequence_packet_to_read_reqs<
             HISTORY_BUFFER_SIZE_KB, RAM_ADDR_WIDTH, RAM_DATA_WIDTH
         >(
             state.hyp_ptr, packet, state.hb_len

diff --git a/xls/modules/zstd/parallel_rams.x b/xls/modules/zstd/parallel_rams.x
@@ -443,9 +443,20 @@ pub fn sequence_packet_to_read_reqs<
     type ReadReq = ram::ReadReq<RAM_ADDR_WIDTH, RAM_NUM_PARTITIONS>;
     type Packet = SequenceExecutorPacket<RAM_DATA_WIDTH>;
 
-    let max_len = std::umin(seq.length as u32, std::umin(RAM_NUM, hb_len as u32));
+    let max_len = std::umin(seq.length as u32, std::umin(RAM_NUM, std::umin(hb_len as u32, seq.content as u32)));
 
     let (curr_seq, next_seq, next_seq_valid) = if seq.length > max_len as CopyOrMatchLength {
+        (
+            seq,
+            Packet {
+                msg_type: SequenceExecutorMessageType::SEQUENCE,
+                length:  seq.length - max_len as CopyOrMatchLength,
+                content: seq.content,
+                last: false,
+            },
+            true,
+        )
+    } else if seq.length > seq.content as CopyOrMatchLength {
         (
             Packet {
                 msg_type: SequenceExecutorMessageType::SEQUENCE,
@@ -545,7 +556,7 @@ fn test_sequence_packet_to_read_reqs() {
         RamReadLen:8,
         Packet {
             msg_type: SequenceExecutorMessageType::SEQUENCE,
-            content: CopyOrMatchContent:18,
+            content: CopyOrMatchContent:10,
             length: CopyOrMatchLength:1,
             last: false
         }, true,