diff --git a/transactional_emulator/src/main.rs b/transactional_emulator/src/main.rs index 6303dafd..3f357ab2 100644 --- a/transactional_emulator/src/main.rs +++ b/transactional_emulator/src/main.rs @@ -249,7 +249,11 @@ impl Accelerator { assert!(element_bits.is_power_of_two()); let len_in_bits_per_load = element_bits as u32 * load_dim; - assert!(len_in_bits_per_load.is_multiple_of(8 * 64)); + // A load must be a whole number of bytes. (Was a multiple of 8*64 + // bits, i.e. a full 64-byte HBM burst — relaxed so sub-64 MLEN, which + // packs <64 bytes per row, is supported; the element-read loop below + // reads aligned 64-byte words and extracts the real bytes.) + assert!(len_in_bits_per_load.is_multiple_of(8)); let len_in_bytes_per_load = len_in_bits_per_load / 8; // Calculate scale bytes per load iteration (for Mx types) @@ -297,16 +301,34 @@ impl Accelerator { as usize + block_idx as usize * scale_len_in_bytes_per_load as usize; - // Element chunks: - for i in 0..(len_in_bytes_per_load as usize + 63) / 64 { - let chunk_offset = byte_offset + i * 64; - let chunk_size = std::cmp::min(64, total_bytes - chunk_offset); - let addr = element_addr + (i * 64) as u64; - assert!(addr.is_multiple_of(64)); - futures.push(Box::pin(async move { - let data = hbm_clone.read(addr).await; - ChunkType::Element(chunk_offset, data, chunk_size) - })); + // Element chunks: read the 64-byte-aligned HBM word(s) that + // cover this load's byte range [element_addr, element_addr+len) + // and copy only the real bytes into the packed `bytes` buffer. + // Handles element_addr that is NOT 64-aligned (sub-64 MLEN, where + // HBM rows are packed tighter than one 64-byte burst) and loads + // that straddle a 64-byte boundary. At MLEN>=64 element_addr is + // 64-aligned and len is a multiple of 64, so this reduces to the + // previous one-aligned-read-per-64-bytes behaviour. + let load_len = len_in_bytes_per_load as u64; + let load_end = element_addr + load_len; // exclusive HBM byte + let mut blk = (element_addr / 64) * 64; + while blk < load_end { + let copy_start = std::cmp::max(blk, element_addr); + let copy_end = std::cmp::min(blk + 64, load_end); + let within = (copy_start - blk) as usize; + let dst = byte_offset + (copy_start - element_addr) as usize; + let mut n = (copy_end - copy_start) as usize; + n = std::cmp::min(n, total_bytes.saturating_sub(dst)); + let addr = blk; + if n > 0 { + futures.push(Box::pin(async move { + let data = hbm_clone.read(addr).await; + let mut sel = [0u8; 64]; + sel[..n].copy_from_slice(&data[within..within + n]); + ChunkType::Element(dst, sel, n) + })); + } + blk += 64; } // Scale chunks (if Mx type) @@ -456,7 +478,9 @@ impl Accelerator { assert!(element_bits.is_power_of_two()); let len_in_bits_per_store = element_bits as u32 * store_dim; - assert!(len_in_bits_per_store.is_multiple_of(8 * 64)); + // Whole-bytes only (was a full 64-byte HBM burst); sub-64 MLEN stores + // are read-modify-written into aligned 64-byte words below. + assert!(len_in_bits_per_store.is_multiple_of(8)); let len_in_bytes_per_store = len_in_bits_per_store / 8; // Calculate scale bytes per store iteration (for Mx types) @@ -530,20 +554,28 @@ impl Accelerator { let element_addr = index + (store_iter * stride) as u64; let scale_addr = scale_index + (store_iter as f32 * stride_scale) as u64; - // Write element bytes to HBM (64-byte aligned chunks) - for i in 0..(len_in_bytes_per_store as usize + 63) / 64 { - let chunk_offset = i * 64; - let chunk_size = std::cmp::min(64, len_in_bytes_per_store as usize - chunk_offset); - let addr = element_addr + (i * 64) as u64; - assert!(addr.is_multiple_of(64)); - - let mut chunk = [0u8; 64]; - if chunk_offset < element_bytes.len() { - let copy_len = std::cmp::min(chunk_size, element_bytes.len() - chunk_offset); - chunk[..copy_len] - .copy_from_slice(&element_bytes[chunk_offset..chunk_offset + copy_len]); + // Write element bytes to HBM. Read-modify-write the 64-byte-aligned + // word(s) covering [element_addr, element_addr+len) so sub-64 MLEN + // stores (which pack <64 bytes per row and may be unaligned or share a + // 64-byte word with the next row) do not clobber neighbouring bytes. + // At MLEN>=64 each store fills whole 64-byte words, so the RMW + // overwrites all 64 bytes — identical to the previous plain write. + let store_len = len_in_bytes_per_store as u64; + let store_end = element_addr + store_len; + let mut blk = (element_addr / 64) * 64; + while blk < store_end { + let copy_start = std::cmp::max(blk, element_addr); + let copy_end = std::cmp::min(blk + 64, store_end); + let within = (copy_start - blk) as usize; + let src = (copy_start - element_addr) as usize; + let n = (copy_end - copy_start) as usize; + let mut chunk = hbm_clone.read(blk).await; + if src < element_bytes.len() { + let cn = std::cmp::min(n, element_bytes.len() - src); + chunk[within..within + cn].copy_from_slice(&element_bytes[src..src + cn]); } - hbm_clone.write(addr, chunk).await; + hbm_clone.write(blk, chunk).await; + blk += 64; } // Write scale bytes to HBM (if Mx type) diff --git a/transactional_emulator/testbench/aten/configurable.py b/transactional_emulator/testbench/aten/configurable.py index c25d69d1..dba52e59 100644 --- a/transactional_emulator/testbench/aten/configurable.py +++ b/transactional_emulator/testbench/aten/configurable.py @@ -98,13 +98,10 @@ def from_args( mlen = int(args.mlen) vlen = int(args.vlen if args.vlen is not None else mlen) blen = int(args.blen) - # HBM is modelled as 64-byte bursts; an 8-bit-element MLEN row must fill a - # whole number of bursts -> MLEN must be a multiple of 64 (see setup_hw). - if mlen % 64 != 0: - raise ValueError( - f"MLEN ({mlen}) must be a multiple of 64 (emulator models HBM as " - f"64-byte bursts). Use MLEN in {{64, 128, 192, 256, ...}}." - ) + # Sub-64 MLEN is supported: the emulator packs HBM rows tightly and reads + # aligned 64-byte words with byte-granular extraction (transfer_mx_from_hbm). + # The only HBM requirement is a byte-aligned element row, which 8-bit MXFP + # always satisfies. (mlen % blen alignment is checked in setup_hw.) hlen = int(args.hlen if args.hlen is not None else default_hlen or base["HLEN"]) broadcast_amount = int( args.broadcast_amount @@ -215,17 +212,10 @@ def setup_hw(args: argparse.Namespace, build_dir: Path) -> HardwareConfig: raise ValueError(f"MLEN ({mlen}) must be divisible by BLEN ({blen})") if vlen != mlen: raise ValueError(f"VLEN ({vlen}) must equal MLEN ({mlen}) for ATen tests") - # HBM is modelled as fixed 64-byte (512-bit) bursts. With 8-bit MXFP - # elements, one MLEN-wide row must fill a whole number of bursts, i.e. - # 8*MLEN must be a multiple of 512 -> MLEN must be a multiple of 64. - # Sub-64 MLEN would read <64B per row and trips a Rust panic deep in the - # emulator (transfer_mx_from_hbm); fail early here with a clear message. - if mlen % 64 != 0: - raise ValueError( - f"MLEN ({mlen}) must be a multiple of 64: the emulator models HBM as " - f"64-byte bursts and an 8-bit-element row narrower than 64 bytes cannot " - f"fill one burst. Use MLEN in {{64, 128, 192, 256, ...}}." - ) + # Sub-64 MLEN is supported: the Python HBM writer packs rows tightly and the + # emulator reads aligned 64-byte words with byte-granular extraction + # (transfer_mx_from_hbm / transfer_mx_to_hbm). 8-bit MXFP rows are always + # byte-aligned, so no MLEN%64 constraint is needed. base = read_behavior_config() hlen = args.hlen if args.hlen is not None else base["HLEN"] diff --git a/transactional_emulator/testbench/sim_env_utils.py b/transactional_emulator/testbench/sim_env_utils.py index 174c5c43..49aa9c8a 100644 --- a/transactional_emulator/testbench/sim_env_utils.py +++ b/transactional_emulator/testbench/sim_env_utils.py @@ -88,7 +88,6 @@ def _map_mx_byte_aligned( if blocks_array.ndim == 1: blocks_array = blocks_array.reshape(-1, 1) bias_array = np.asarray(bias).reshape(-1) - total_bytes_written = 0 with open(output_file, mode) as f: for row_idx in range(logical_rows): @@ -102,7 +101,6 @@ def _map_mx_byte_aligned( raise ValueError(f"Packed element row ({len(row_bytes)} B) > HBM row ({hbm_row_bytes} B)") row_bytes += b"\x00" * (hbm_row_bytes - len(row_bytes)) f.write(row_bytes) - total_bytes_written += len(row_bytes) for row_idx in range(logical_rows): row_start = row_idx * blocks_per_source_row @@ -115,11 +113,18 @@ def _map_mx_byte_aligned( raise ValueError(f"Packed scale row ({len(row_bytes)} B) > scale row ({scale_row_bytes} B)") row_bytes += b"\x00" * (scale_row_bytes - len(row_bytes)) f.write(row_bytes) - total_bytes_written += len(row_bytes) - remainder = total_bytes_written % 64 - if remainder != 0: - f.write(b"\x00" * (64 - remainder)) + # NOTE: no per-tensor pad to a 64-byte boundary. Each tensor is written at + # its compiler-assigned hbm_addr (forward-padded above) and packed at its + # exact element+scale byte length, matching the compiler's tight HBM base + # allocation (advances by hbm_size, not rounded to 64). A per-tensor tail + # pad to 64 would shift the NEXT tensor past its compiler base whenever a + # tensor's size is not a 64-multiple (only happens at sub-64 MLEN, e.g. + # X(16,16) = 288 B -> padded 320 -> next tensor read at the wrong addr -> + # zero weights). The emulator's MemoryBacked capacity (hbm-size, a 64-byte + # multiple sized to ~2x the preload) zero-covers the final aligned-block + # read past the file end. At MLEN>=64 tensor sizes are already 64-multiples, + # so dropping this pad is a no-op. def map_mx_data_to_hbm_for_behave_sim( @@ -157,9 +162,16 @@ def map_mx_data_to_hbm_for_behave_sim( if source_row_elements > logical_row_elements: raise ValueError(f"source_row_elements ({source_row_elements}) > logical_row_elements ({logical_row_elements})") - physical_hbm_row_bytes = (hbm_row_width + 7) // 8 + # Pack each element row tightly at its real byte width. The compiler + # addresses HBM rows at the logical element stride (e.g. MLEN bytes for 8-bit + # elements) and places scales at the tight element-region offset, so the + # writer must match that stride rather than padding rows up to a 64-byte HBM + # burst. The emulator reads aligned 64-byte words and extracts the real bytes + # (transfer_mx_from_hbm), so a tight layout is read correctly. At MLEN>=64 the + # element row is already >=64 bytes, so this equals the previous value + # (max with the 64-byte burst floor was a no-op there). element_row_bits = logical_row_elements * element_width - hbm_row_bytes = max(physical_hbm_row_bytes, (element_row_bits + 7) // 8) + hbm_row_bytes = (element_row_bits + 7) // 8 blocks_per_source_row = (source_row_elements + block_width - 1) // block_width blocks_per_logical_row = (logical_row_elements + block_width - 1) // block_width inferred_source_rows = (len(blocks) + blocks_per_source_row - 1) // blocks_per_source_row