From 396c7ec83e7250133769e84b1a807cb7d588931b Mon Sep 17 00:00:00 2001 From: Timi Adeniran Date: Mon, 3 Nov 2025 10:29:05 +0000 Subject: [PATCH 01/34] Ignore target directory --- .gitignore | 1 + 1 file changed, 1 insertion(+) diff --git a/.gitignore b/.gitignore index b28d259f..7f3fad77 100644 --- a/.gitignore +++ b/.gitignore @@ -61,6 +61,7 @@ ckpts_* *.db *.db-journal co_design/interface/config.toml +target/ # behavioral simulator behavioral_simulator/target From 7a13643f18c7c89023fbc7c0b48fca838fc494cc Mon Sep 17 00:00:00 2001 From: Timi Adeniran Date: Mon, 3 Nov 2025 10:31:11 +0000 Subject: [PATCH 02/34] First pass: monitor from transfer_from_hbm function --- behavioral_simulator/src/main.rs | 94 ++++++++++++++++++-------------- 1 file changed, 52 insertions(+), 42 deletions(-) diff --git a/behavioral_simulator/src/main.rs b/behavioral_simulator/src/main.rs index 69720e67..bce7b7fb 100644 --- a/behavioral_simulator/src/main.rs +++ b/behavioral_simulator/src/main.rs @@ -713,6 +713,8 @@ struct Accelerator { reg_file: AcceeleratorRegFile, intsram: Vec, fpsram: Vec, + total_hbm_bytes_read: usize, + total_hbm_bytes_written: usize, } struct AcceeleratorRegFile { @@ -761,53 +763,54 @@ impl Accelerator { } else { load_dim }; - Executor::current().spawn(async move { - let element_ty = hbm_type.element_type(); - let element_bits = element_ty.size_in_bits(); - // Extract scale bits and block size if Mx type, otherwise use element_bits/1 as default - let (scale_bits, blocksize) = match hbm_type { - MxDataType::Mx { - elem: _, - scale, - block, - } => (scale.size_in_bits(), block), - _ => (element_bits, 1), // Plain type: each element is "scaled" by 1 - }; - - let element_scale_ratio = (element_bits * blocksize as u8) / scale_bits; - let stride_scale = stride as f32 / element_scale_ratio as f32; - assert!(element_bits.is_power_of_two()); - - let len_in_bits_per_load = element_bits as u32 * load_dim; - // println!("element_bits = {:?}", element_bits); - // println!("load_dim = {:?}", load_dim); - // println!("len_in_bits_per_load = {:?}", len_in_bits_per_load); - assert!(len_in_bits_per_load.is_multiple_of(8 * 64)); - let len_in_bytes_per_load = len_in_bits_per_load / 8; - - // Calculate scale bytes per load iteration (for Mx types) - let (scale_len_in_bytes_per_load, block) = if let MxDataType::Mx { + let element_ty = hbm_type.element_type(); + let element_bits = element_ty.size_in_bits(); + + // Extract scale bits and block size if Mx type, otherwise use element_bits/1 as default + let (scale_bits, blocksize) = match hbm_type { + MxDataType::Mx { elem: _, scale, block, - } = hbm_type - { - let scale_bits = scale.size_in_bits(); - assert!(scale_bits.is_power_of_two()); - let scale_len_in_bits_per_load = scale_bits as u32 * (load_dim / block); - assert!(scale_len_in_bits_per_load.is_multiple_of(8)); - (scale_len_in_bits_per_load / 8, block as usize) - } else { - (0, usize::MAX) - }; - - // Total elements/bytes for all writes: - let total_elements = (write_dim * num_writes) as usize; - let total_bytes = (len_in_bytes_per_load * write_amount * num_writes) as usize; - let total_scale_bytes = - (scale_len_in_bytes_per_load * write_amount * num_writes) as usize; + } => (scale.size_in_bits(), block), + _ => (element_bits, 1), // Plain type: each element is "scaled" by 1 + }; + let element_scale_ratio = (element_bits * blocksize as u8) / scale_bits; + let stride_scale = stride as f32 / element_scale_ratio as f32; + assert!(element_bits.is_power_of_two()); + + let len_in_bits_per_load = element_bits as u32 * load_dim; + // println!("element_bits = {:?}", element_bits); + // println!("load_dim = {:?}", load_dim); + // println!("len_in_bits_per_load = {:?}", len_in_bits_per_load); + assert!(len_in_bits_per_load.is_multiple_of(8 * 64)); + let len_in_bytes_per_load = len_in_bits_per_load / 8; + + // Calculate scale bytes per load iteration (for Mx types) + let (scale_len_in_bytes_per_load, block) = if let MxDataType::Mx { + elem: _, + scale, + block, + } = hbm_type + { + let scale_bits = scale.size_in_bits(); + assert!(scale_bits.is_power_of_two()); + let scale_len_in_bits_per_load = scale_bits as u32 * (load_dim / block); + assert!(scale_len_in_bits_per_load.is_multiple_of(8)); + (scale_len_in_bits_per_load / 8, block as usize) + } else { + (0, usize::MAX) + }; + + // Total elements/bytes for all writes: + let total_elements = (write_dim * num_writes) as usize; + let total_bytes = (len_in_bytes_per_load * write_amount * num_writes) as usize; + let total_scale_bytes = (scale_len_in_bytes_per_load * write_amount * num_writes) as usize; + self.total_hbm_bytes_read += total_bytes + total_scale_bytes; + + Executor::current().spawn(async move { let mut bytes = vec![0u8; total_bytes]; let mut scale_bytes = vec![0u8; total_scale_bytes]; let hbm_clone = &hbm_clone; @@ -1387,6 +1390,8 @@ async fn start() { }, intsram: vec![0; 1024], fpsram: vec![f16::ZERO; 1024], + total_hbm_bytes_read: 0, + total_hbm_bytes_written: 0, }; use std::fs; @@ -1474,6 +1479,11 @@ async fn start() { let mut vram_file = std::fs::File::create(vram_dump_path).unwrap(); vram_file.write_all(&vram_bytes).unwrap(); eprintln!("Dumped VRAM content to: {:?}", vram_dump_path); + + eprintln!( + "HBM Utilization - Bytes read: {:?} | Bytes written: {:?}", + accelerator.total_hbm_bytes_read, accelerator.total_hbm_bytes_written + ); } #[tokio::main] From 13bb972ca62e4b2ad0e504c526d7cd067b74fad8 Mon Sep 17 00:00:00 2001 From: Timi Adeniran Date: Tue, 4 Nov 2025 23:18:10 +0000 Subject: [PATCH 03/34] Add memory usage statistics to memory traits --- behavioral_simulator/lib/memory/src/lib.rs | 66 +++++++++++++++---- .../lib/ramulator/src/model.rs | 24 ++++++- behavioral_simulator/src/main.rs | 36 ++++++---- 3 files changed, 97 insertions(+), 29 deletions(-) diff --git a/behavioral_simulator/lib/memory/src/lib.rs b/behavioral_simulator/lib/memory/src/lib.rs index ee0bba36..3b908641 100644 --- a/behavioral_simulator/lib/memory/src/lib.rs +++ b/behavioral_simulator/lib/memory/src/lib.rs @@ -1,51 +1,76 @@ use std::mem::ManuallyDrop; use std::sync::Mutex; +#[derive(Copy, Clone)] +pub struct Statistics { + pub total_bytes_read: u32, + pub total_bytes_written: u32, +} + #[async_trait::async_trait] pub trait MemoryTimingModel: Send + Sync { /// Read 64-bytes of memory. /// /// We fix to 64-bytes to accomodate memory emulators. - async fn read(&self, addr: u64); + async fn read(&mut self, addr: u64); /// Write 64-bytes of memory. - async fn write(&self, addr: u64); + async fn write(&mut self, addr: u64); + + /// Get statistics on the memory usage. + fn statistics(&self) -> Statistics; } #[async_trait::async_trait] impl MemoryTimingModel for ManuallyDrop { - async fn read(&self, addr: u64) { + async fn read(&mut self, addr: u64) { T::read(self, addr).await } - async fn write(&self, addr: u64) { + async fn write(&mut self, addr: u64) { T::write(self, addr).await } + + fn statistics(&self) -> Statistics { + T::statistics(self) + } } #[async_trait::async_trait] pub trait MemoryModel: Send + Sync { /// Read 64-bytes of memory. - async fn read(&self, addr: u64) -> [u8; 64]; + async fn read(&mut self, addr: u64) -> [u8; 64]; /// Write 64-bytes of memory. - async fn write(&self, addr: u64, bytes: [u8; 64]); + async fn write(&mut self, addr: u64, bytes: [u8; 64]); + + /// Get statistics on the memory usage. + fn statistics(&self) -> Statistics; } /// A memory that discards all written data. /// /// This is useful to just test the timing without caring the actual data. -pub struct NoData; +pub struct NoData { + statistics: Statistics, +} #[async_trait::async_trait] impl MemoryModel for NoData { /// Read 64-bytes of memory. - async fn read(&self, _addr: u64) -> [u8; 64] { + async fn read(&mut self, _addr: u64) -> [u8; 64] { + self.statistics.total_bytes_read += 64; [0; 64] } /// Write 64-bytes of memory. - async fn write(&self, _addr: u64, _bytes: [u8; 64]) {} + async fn write(&mut self, _addr: u64, _bytes: [u8; 64]) { + self.statistics.total_bytes_written += 64; + } + + fn statistics(&self) -> Statistics { + self.statistics + } } /// A simulated memory that is backed by memory. @@ -53,6 +78,7 @@ impl MemoryModel for NoData { /// This is useful to just test the timing without caring the actual data. pub struct MemoryBacked { data: Mutex>, + statistics: Statistics, } impl MemoryBacked { @@ -60,6 +86,10 @@ impl MemoryBacked { assert!(size.is_multiple_of(64)); Self { data: Mutex::new(vec![[0; 64]; size / 64]), + statistics: Statistics { + total_bytes_read: 0, + total_bytes_written: 0, + }, } } @@ -74,14 +104,20 @@ impl MemoryBacked { #[async_trait::async_trait] impl MemoryModel for MemoryBacked { /// Read 64-bytes of memory. - async fn read(&self, addr: u64) -> [u8; 64] { + async fn read(&mut self, addr: u64) -> [u8; 64] { + self.statistics.total_bytes_read += 64; self.data.lock().unwrap()[addr as usize / 64] } /// Write 64-bytes of memory. - async fn write(&self, addr: u64, bytes: [u8; 64]) { + async fn write(&mut self, addr: u64, bytes: [u8; 64]) { + self.statistics.total_bytes_written += 64; self.data.lock().unwrap()[addr as usize / 64] = bytes; } + + fn statistics(&self) -> Statistics { + self.statistics + } } /// Combine a data model with an extra timing model. @@ -103,14 +139,18 @@ impl WithTiming { #[async_trait::async_trait] impl MemoryModel for WithTiming { /// Read 64-bytes of memory. - async fn read(&self, addr: u64) -> [u8; 64] { + async fn read(&mut self, addr: u64) -> [u8; 64] { self.timing.read(addr).await; self.data.read(addr).await } /// Write 64-bytes of memory. - async fn write(&self, addr: u64, bytes: [u8; 64]) { + async fn write(&mut self, addr: u64, bytes: [u8; 64]) { self.timing.write(addr).await; self.data.write(addr, bytes).await } + + fn statistics(&self) -> Statistics { + self.timing.statistics() + } } diff --git a/behavioral_simulator/lib/ramulator/src/model.rs b/behavioral_simulator/lib/ramulator/src/model.rs index ecc66894..c241ec9d 100644 --- a/behavioral_simulator/lib/ramulator/src/model.rs +++ b/behavioral_simulator/lib/ramulator/src/model.rs @@ -5,10 +5,13 @@ use anyhow::Result; use runtime::{Duration, Executor, Instant}; use crate::raw::Ramulator as RawRamulator; +use memory::Statistics; struct State { next_instant: Instant, ramulator: RawRamulator, + // Utilization stats + statistics: Statistics, } struct Inner { @@ -40,8 +43,11 @@ impl Ramulator { mutable: Mutex::new(State { next_instant: Instant::INIT, ramulator, + statistics: Statistics { + total_bytes_read: 0, + total_bytes_written: 0, + }, }), - read_lock: tokio::sync::Mutex::new(()), write_lock: tokio::sync::Mutex::new(()), transfer_size, @@ -162,7 +168,11 @@ impl Ramulator { #[async_trait::async_trait] impl memory::MemoryTimingModel for Ramulator { - async fn read(&self, addr: u64) { + async fn read(&mut self, addr: u64) { + { + let mut guard = self.0.mutable.lock().unwrap(); + guard.statistics.total_bytes_read += 64; + } futures::future::join_all( (0..64) .step_by(self.0.transfer_size as _) @@ -171,7 +181,11 @@ impl memory::MemoryTimingModel for Ramulator { .await; } - async fn write(&self, addr: u64) { + async fn write(&mut self, addr: u64) { + { + let mut guard = self.0.mutable.lock().unwrap(); + guard.statistics.total_bytes_written += 64; + } futures::future::join_all( (0..64) .step_by(self.0.transfer_size as _) @@ -179,4 +193,8 @@ impl memory::MemoryTimingModel for Ramulator { ) .await; } + + fn statistics(&self) -> Statistics { + self.0.mutable.lock().unwrap().statistics + } } diff --git a/behavioral_simulator/src/main.rs b/behavioral_simulator/src/main.rs index bce7b7fb..f12bc92b 100644 --- a/behavioral_simulator/src/main.rs +++ b/behavioral_simulator/src/main.rs @@ -708,7 +708,7 @@ impl VectorMachine { struct Accelerator { m_machine: MatrixMachine, v_machine: VectorMachine, - hbm: Arc, + hbm: Arc>, tile_size: u32, reg_file: AcceeleratorRegFile, intsram: Vec, @@ -842,7 +842,8 @@ impl Accelerator { let chunk_size = std::cmp::min(64, total_bytes - chunk_offset); let addr = element_addr + (i * 64) as u64; futures.push(Box::pin(async move { - let data = hbm_clone.read(addr).await; + let mut shared_hbm = hbm_clone.lock().await; + let data = shared_hbm.read(addr).await; ChunkType::Element(chunk_offset, data, chunk_size) })); } @@ -856,7 +857,8 @@ impl Accelerator { let chunk_offset = scale_byte_offset; // where to write in scale_bytes let chunk_size = std::cmp::min(64, total_scale_bytes - chunk_offset); futures.push(Box::pin(async move { - let data = hbm_clone.read(aligned_scale_addr).await; + let mut shared_hbm = hbm_clone.lock().await; + let data = shared_hbm.read(aligned_scale_addr).await; // Copy out only the relevant bytes for this scale_addr // scale_len_in_bytes_per_load says how many bytes to copy from within the chunk let end_offset = std::cmp::min( @@ -1371,10 +1373,10 @@ async fn start() { }; let v_machine = VectorMachine { vram }; // Share same dim with VSRAM - let hbm = Arc::new(memory::WithTiming::new( + let hbm = Arc::new(Mutex::new(memory::WithTiming::new( ManuallyDrop::new(ramulator::Ramulator::hbm2_preset(8).unwrap()), memory::MemoryBacked::with_capacity(*HBM_SIZE), - )); + ))); let mut accelerator = Accelerator { m_machine: machine, @@ -1411,10 +1413,11 @@ async fn start() { // Memory Initialization // - HBM Preload let hbm_data = std::fs::read(opts.hbm).unwrap(); - - hbm.data().with_data(|f| { - f[..hbm_data.len()].copy_from_slice(&hbm_data); - }); + { + hbm.lock().await.data().with_data(|f| { + f[..hbm_data.len()].copy_from_slice(&hbm_data); + }); + } // Load fpsram and intsram as raw bytes and map to the vector files. // - fpsram Preload @@ -1480,10 +1483,17 @@ async fn start() { vram_file.write_all(&vram_bytes).unwrap(); eprintln!("Dumped VRAM content to: {:?}", vram_dump_path); - eprintln!( - "HBM Utilization - Bytes read: {:?} | Bytes written: {:?}", - accelerator.total_hbm_bytes_read, accelerator.total_hbm_bytes_written - ); + { + eprintln!( + "HBM Utilization - Bytes read: {:?} | Bytes written: {:?}", + accelerator.total_hbm_bytes_read, accelerator.total_hbm_bytes_written + ); + let stats = hbm.lock().await.statistics(); + eprintln!( + "[New] HBM Utilization - Bytes read: {:?} | Bytes written: {:?}", + stats.total_bytes_read, stats.total_bytes_written + ); + } } #[tokio::main] From ff43a8ed24f72b08b63788d7aa6893ab11b47051 Mon Sep 17 00:00:00 2001 From: Timi Adeniran Date: Tue, 4 Nov 2025 23:24:04 +0000 Subject: [PATCH 04/34] Revert changes from initial implementation --- behavioral_simulator/src/main.rs | 94 +++++++++++++++----------------- 1 file changed, 43 insertions(+), 51 deletions(-) diff --git a/behavioral_simulator/src/main.rs b/behavioral_simulator/src/main.rs index f12bc92b..3c39fae1 100644 --- a/behavioral_simulator/src/main.rs +++ b/behavioral_simulator/src/main.rs @@ -713,8 +713,6 @@ struct Accelerator { reg_file: AcceeleratorRegFile, intsram: Vec, fpsram: Vec, - total_hbm_bytes_read: usize, - total_hbm_bytes_written: usize, } struct AcceeleratorRegFile { @@ -764,53 +762,53 @@ impl Accelerator { load_dim }; - let element_ty = hbm_type.element_type(); - let element_bits = element_ty.size_in_bits(); + Executor::current().spawn(async move { + let element_ty = hbm_type.element_type(); + let element_bits = element_ty.size_in_bits(); - // Extract scale bits and block size if Mx type, otherwise use element_bits/1 as default - let (scale_bits, blocksize) = match hbm_type { - MxDataType::Mx { + // Extract scale bits and block size if Mx type, otherwise use element_bits/1 as default + let (scale_bits, blocksize) = match hbm_type { + MxDataType::Mx { + elem: _, + scale, + block, + } => (scale.size_in_bits(), block), + _ => (element_bits, 1), // Plain type: each element is "scaled" by 1 + }; + + let element_scale_ratio = (element_bits * blocksize as u8) / scale_bits; + let stride_scale = stride as f32 / element_scale_ratio as f32; + assert!(element_bits.is_power_of_two()); + + let len_in_bits_per_load = element_bits as u32 * load_dim; + // println!("element_bits = {:?}", element_bits); + // println!("load_dim = {:?}", load_dim); + // println!("len_in_bits_per_load = {:?}", len_in_bits_per_load); + assert!(len_in_bits_per_load.is_multiple_of(8 * 64)); + let len_in_bytes_per_load = len_in_bits_per_load / 8; + + // Calculate scale bytes per load iteration (for Mx types) + let (scale_len_in_bytes_per_load, block) = if let MxDataType::Mx { elem: _, scale, block, - } => (scale.size_in_bits(), block), - _ => (element_bits, 1), // Plain type: each element is "scaled" by 1 - }; + } = hbm_type + { + let scale_bits = scale.size_in_bits(); + assert!(scale_bits.is_power_of_two()); + let scale_len_in_bits_per_load = scale_bits as u32 * (load_dim / block); + assert!(scale_len_in_bits_per_load.is_multiple_of(8)); + (scale_len_in_bits_per_load / 8, block as usize) + } else { + (0, usize::MAX) + }; + + // Total elements/bytes for all writes: + let total_elements = (write_dim * num_writes) as usize; + let total_bytes = (len_in_bytes_per_load * write_amount * num_writes) as usize; + let total_scale_bytes = + (scale_len_in_bytes_per_load * write_amount * num_writes) as usize; - let element_scale_ratio = (element_bits * blocksize as u8) / scale_bits; - let stride_scale = stride as f32 / element_scale_ratio as f32; - assert!(element_bits.is_power_of_two()); - - let len_in_bits_per_load = element_bits as u32 * load_dim; - // println!("element_bits = {:?}", element_bits); - // println!("load_dim = {:?}", load_dim); - // println!("len_in_bits_per_load = {:?}", len_in_bits_per_load); - assert!(len_in_bits_per_load.is_multiple_of(8 * 64)); - let len_in_bytes_per_load = len_in_bits_per_load / 8; - - // Calculate scale bytes per load iteration (for Mx types) - let (scale_len_in_bytes_per_load, block) = if let MxDataType::Mx { - elem: _, - scale, - block, - } = hbm_type - { - let scale_bits = scale.size_in_bits(); - assert!(scale_bits.is_power_of_two()); - let scale_len_in_bits_per_load = scale_bits as u32 * (load_dim / block); - assert!(scale_len_in_bits_per_load.is_multiple_of(8)); - (scale_len_in_bits_per_load / 8, block as usize) - } else { - (0, usize::MAX) - }; - - // Total elements/bytes for all writes: - let total_elements = (write_dim * num_writes) as usize; - let total_bytes = (len_in_bytes_per_load * write_amount * num_writes) as usize; - let total_scale_bytes = (scale_len_in_bytes_per_load * write_amount * num_writes) as usize; - self.total_hbm_bytes_read += total_bytes + total_scale_bytes; - - Executor::current().spawn(async move { let mut bytes = vec![0u8; total_bytes]; let mut scale_bytes = vec![0u8; total_scale_bytes]; let hbm_clone = &hbm_clone; @@ -1392,8 +1390,6 @@ async fn start() { }, intsram: vec![0; 1024], fpsram: vec![f16::ZERO; 1024], - total_hbm_bytes_read: 0, - total_hbm_bytes_written: 0, }; use std::fs; @@ -1484,13 +1480,9 @@ async fn start() { eprintln!("Dumped VRAM content to: {:?}", vram_dump_path); { - eprintln!( - "HBM Utilization - Bytes read: {:?} | Bytes written: {:?}", - accelerator.total_hbm_bytes_read, accelerator.total_hbm_bytes_written - ); let stats = hbm.lock().await.statistics(); eprintln!( - "[New] HBM Utilization - Bytes read: {:?} | Bytes written: {:?}", + "HBM Utilization - Bytes read: {:?} | Bytes written: {:?}", stats.total_bytes_read, stats.total_bytes_written ); } From 09baa0554adaeed96297b4a45216913f109fa18d Mon Sep 17 00:00:00 2001 From: Timi Adeniran Date: Wed, 5 Nov 2025 00:03:56 +0000 Subject: [PATCH 05/34] Report bytes per second --- behavioral_simulator/lib/memory/src/lib.rs | 4 ++-- behavioral_simulator/lib/runtime/src/time.rs | 4 ++++ behavioral_simulator/src/main.rs | 8 +++++--- 3 files changed, 11 insertions(+), 5 deletions(-) diff --git a/behavioral_simulator/lib/memory/src/lib.rs b/behavioral_simulator/lib/memory/src/lib.rs index 3b908641..e4c2f5e8 100644 --- a/behavioral_simulator/lib/memory/src/lib.rs +++ b/behavioral_simulator/lib/memory/src/lib.rs @@ -3,8 +3,8 @@ use std::sync::Mutex; #[derive(Copy, Clone)] pub struct Statistics { - pub total_bytes_read: u32, - pub total_bytes_written: u32, + pub total_bytes_read: u64, + pub total_bytes_written: u64, } #[async_trait::async_trait] diff --git a/behavioral_simulator/lib/runtime/src/time.rs b/behavioral_simulator/lib/runtime/src/time.rs index 333bc367..7cdb9e2c 100644 --- a/behavioral_simulator/lib/runtime/src/time.rs +++ b/behavioral_simulator/lib/runtime/src/time.rs @@ -106,6 +106,10 @@ impl Instant { /// /// Effectively means "never". pub const ETERNITY: Self = Self(u64::MAX); + + pub const fn to_secs(&self) -> f64 { + self.0 as f64 / (1_000_000_000_000_u64 as f64) + } } pub trait Deadline { diff --git a/behavioral_simulator/src/main.rs b/behavioral_simulator/src/main.rs index 3c39fae1..f5c7ddbf 100644 --- a/behavioral_simulator/src/main.rs +++ b/behavioral_simulator/src/main.rs @@ -1480,10 +1480,12 @@ async fn start() { eprintln!("Dumped VRAM content to: {:?}", vram_dump_path); { - let stats = hbm.lock().await.statistics(); + let memory_stats = hbm.lock().await.statistics(); + let utilization = (memory_stats.total_bytes_read + memory_stats.total_bytes_written) as f64 + / Executor::current().now().to_secs(); eprintln!( - "HBM Utilization - Bytes read: {:?} | Bytes written: {:?}", - stats.total_bytes_read, stats.total_bytes_written + "HBM Statistics - Bytes read: {:?} | Bytes written: {:?} | Utilization: {:.2e} bytes/sec", + memory_stats.total_bytes_read, memory_stats.total_bytes_written, utilization ); } } From 39a74c6b2bb57eb9a3e73ebfe7e0871a025868d7 Mon Sep 17 00:00:00 2001 From: Timi Adeniran Date: Wed, 12 Nov 2025 01:07:20 +0000 Subject: [PATCH 06/34] Review comments: Remove HBM lock, introduce WithStats struct --- behavioral_simulator/lib/memory/src/lib.rs | 122 +++++++++++------- .../lib/ramulator/src/model.rs | 24 +--- behavioral_simulator/src/main.rs | 42 +++--- 3 files changed, 99 insertions(+), 89 deletions(-) diff --git a/behavioral_simulator/lib/memory/src/lib.rs b/behavioral_simulator/lib/memory/src/lib.rs index e4c2f5e8..f9b4d526 100644 --- a/behavioral_simulator/lib/memory/src/lib.rs +++ b/behavioral_simulator/lib/memory/src/lib.rs @@ -12,65 +12,46 @@ pub trait MemoryTimingModel: Send + Sync { /// Read 64-bytes of memory. /// /// We fix to 64-bytes to accomodate memory emulators. - async fn read(&mut self, addr: u64); + async fn read(&self, addr: u64); /// Write 64-bytes of memory. - async fn write(&mut self, addr: u64); - - /// Get statistics on the memory usage. - fn statistics(&self) -> Statistics; + async fn write(&self, addr: u64); } #[async_trait::async_trait] impl MemoryTimingModel for ManuallyDrop { - async fn read(&mut self, addr: u64) { + async fn read(&self, addr: u64) { T::read(self, addr).await } - async fn write(&mut self, addr: u64) { + async fn write(&self, addr: u64) { T::write(self, addr).await } - - fn statistics(&self) -> Statistics { - T::statistics(self) - } } #[async_trait::async_trait] pub trait MemoryModel: Send + Sync { /// Read 64-bytes of memory. - async fn read(&mut self, addr: u64) -> [u8; 64]; + async fn read(&self, addr: u64) -> [u8; 64]; /// Write 64-bytes of memory. - async fn write(&mut self, addr: u64, bytes: [u8; 64]); - - /// Get statistics on the memory usage. - fn statistics(&self) -> Statistics; + async fn write(&self, addr: u64, bytes: [u8; 64]); } /// A memory that discards all written data. /// /// This is useful to just test the timing without caring the actual data. -pub struct NoData { - statistics: Statistics, -} +pub struct NoData; #[async_trait::async_trait] impl MemoryModel for NoData { /// Read 64-bytes of memory. - async fn read(&mut self, _addr: u64) -> [u8; 64] { - self.statistics.total_bytes_read += 64; + async fn read(&self, _addr: u64) -> [u8; 64] { [0; 64] } /// Write 64-bytes of memory. - async fn write(&mut self, _addr: u64, _bytes: [u8; 64]) { - self.statistics.total_bytes_written += 64; - } - - fn statistics(&self) -> Statistics { - self.statistics - } + async fn write(&self, _addr: u64, _bytes: [u8; 64]) {} } /// A simulated memory that is backed by memory. @@ -78,7 +59,6 @@ impl MemoryModel for NoData { /// This is useful to just test the timing without caring the actual data. pub struct MemoryBacked { data: Mutex>, - statistics: Statistics, } impl MemoryBacked { @@ -86,10 +66,6 @@ impl MemoryBacked { assert!(size.is_multiple_of(64)); Self { data: Mutex::new(vec![[0; 64]; size / 64]), - statistics: Statistics { - total_bytes_read: 0, - total_bytes_written: 0, - }, } } @@ -104,19 +80,73 @@ impl MemoryBacked { #[async_trait::async_trait] impl MemoryModel for MemoryBacked { /// Read 64-bytes of memory. - async fn read(&mut self, addr: u64) -> [u8; 64] { - self.statistics.total_bytes_read += 64; + async fn read(&self, addr: u64) -> [u8; 64] { self.data.lock().unwrap()[addr as usize / 64] } /// Write 64-bytes of memory. - async fn write(&mut self, addr: u64, bytes: [u8; 64]) { - self.statistics.total_bytes_written += 64; + async fn write(&self, addr: u64, bytes: [u8; 64]) { self.data.lock().unwrap()[addr as usize / 64] = bytes; } +} + +pub struct WithStats { + model: T, + statistics: Mutex, +} + +impl WithStats { + pub fn new(model: T) -> Self { + let stats = Statistics { + total_bytes_read: 0, + total_bytes_written: 0, + }; + WithStats { + model, + statistics: Mutex::new(stats), + } + } - fn statistics(&self) -> Statistics { - self.statistics + pub fn model(&self) -> &T { + &self.model + } +} + +#[async_trait::async_trait] +impl MemoryModel for WithStats { + async fn read(&self, addr: u64) -> [u8; 64] { + { + let mut guard = self.statistics.lock().unwrap(); + guard.total_bytes_read += 64; + } + self.model.read(addr).await + } + + async fn write(&self, addr: u64, bytes: [u8; 64]) { + { + let mut guard = self.statistics.lock().unwrap(); + guard.total_bytes_written += 64; + } + self.model.write(addr, bytes).await + } +} + +#[async_trait::async_trait] +impl MemoryTimingModel for WithStats { + async fn read(&self, addr: u64) { + { + let mut guard = self.statistics.lock().unwrap(); + guard.total_bytes_read += 64; + } + self.model.read(addr).await + } + + async fn write(&self, addr: u64) { + { + let mut guard = self.statistics.lock().unwrap(); + guard.total_bytes_written += 64; + } + self.model.write(addr).await } } @@ -136,21 +166,23 @@ impl WithTiming { } } +impl WithTiming, WithStats> { + pub fn statistics(&self) -> Statistics { + self.timing.statistics.lock().unwrap().clone() + } +} + #[async_trait::async_trait] impl MemoryModel for WithTiming { /// Read 64-bytes of memory. - async fn read(&mut self, addr: u64) -> [u8; 64] { + async fn read(&self, addr: u64) -> [u8; 64] { self.timing.read(addr).await; self.data.read(addr).await } /// Write 64-bytes of memory. - async fn write(&mut self, addr: u64, bytes: [u8; 64]) { + async fn write(&self, addr: u64, bytes: [u8; 64]) { self.timing.write(addr).await; self.data.write(addr, bytes).await } - - fn statistics(&self) -> Statistics { - self.timing.statistics() - } } diff --git a/behavioral_simulator/lib/ramulator/src/model.rs b/behavioral_simulator/lib/ramulator/src/model.rs index c241ec9d..ecc66894 100644 --- a/behavioral_simulator/lib/ramulator/src/model.rs +++ b/behavioral_simulator/lib/ramulator/src/model.rs @@ -5,13 +5,10 @@ use anyhow::Result; use runtime::{Duration, Executor, Instant}; use crate::raw::Ramulator as RawRamulator; -use memory::Statistics; struct State { next_instant: Instant, ramulator: RawRamulator, - // Utilization stats - statistics: Statistics, } struct Inner { @@ -43,11 +40,8 @@ impl Ramulator { mutable: Mutex::new(State { next_instant: Instant::INIT, ramulator, - statistics: Statistics { - total_bytes_read: 0, - total_bytes_written: 0, - }, }), + read_lock: tokio::sync::Mutex::new(()), write_lock: tokio::sync::Mutex::new(()), transfer_size, @@ -168,11 +162,7 @@ impl Ramulator { #[async_trait::async_trait] impl memory::MemoryTimingModel for Ramulator { - async fn read(&mut self, addr: u64) { - { - let mut guard = self.0.mutable.lock().unwrap(); - guard.statistics.total_bytes_read += 64; - } + async fn read(&self, addr: u64) { futures::future::join_all( (0..64) .step_by(self.0.transfer_size as _) @@ -181,11 +171,7 @@ impl memory::MemoryTimingModel for Ramulator { .await; } - async fn write(&mut self, addr: u64) { - { - let mut guard = self.0.mutable.lock().unwrap(); - guard.statistics.total_bytes_written += 64; - } + async fn write(&self, addr: u64) { futures::future::join_all( (0..64) .step_by(self.0.transfer_size as _) @@ -193,8 +179,4 @@ impl memory::MemoryTimingModel for Ramulator { ) .await; } - - fn statistics(&self) -> Statistics { - self.0.mutable.lock().unwrap().statistics - } } diff --git a/behavioral_simulator/src/main.rs b/behavioral_simulator/src/main.rs index f5c7ddbf..922d813d 100644 --- a/behavioral_simulator/src/main.rs +++ b/behavioral_simulator/src/main.rs @@ -708,7 +708,7 @@ impl VectorMachine { struct Accelerator { m_machine: MatrixMachine, v_machine: VectorMachine, - hbm: Arc>, + hbm: Arc, tile_size: u32, reg_file: AcceeleratorRegFile, intsram: Vec, @@ -840,8 +840,7 @@ impl Accelerator { let chunk_size = std::cmp::min(64, total_bytes - chunk_offset); let addr = element_addr + (i * 64) as u64; futures.push(Box::pin(async move { - let mut shared_hbm = hbm_clone.lock().await; - let data = shared_hbm.read(addr).await; + let data = hbm_clone.read(addr).await; ChunkType::Element(chunk_offset, data, chunk_size) })); } @@ -855,8 +854,7 @@ impl Accelerator { let chunk_offset = scale_byte_offset; // where to write in scale_bytes let chunk_size = std::cmp::min(64, total_scale_bytes - chunk_offset); futures.push(Box::pin(async move { - let mut shared_hbm = hbm_clone.lock().await; - let data = shared_hbm.read(aligned_scale_addr).await; + let data = hbm_clone.read(aligned_scale_addr).await; // Copy out only the relevant bytes for this scale_addr // scale_len_in_bytes_per_load says how many bytes to copy from within the chunk let end_offset = std::cmp::min( @@ -1371,10 +1369,12 @@ async fn start() { }; let v_machine = VectorMachine { vram }; // Share same dim with VSRAM - let hbm = Arc::new(Mutex::new(memory::WithTiming::new( - ManuallyDrop::new(ramulator::Ramulator::hbm2_preset(8).unwrap()), - memory::MemoryBacked::with_capacity(*HBM_SIZE), - ))); + let hbm = Arc::new(memory::WithTiming::new( + memory::WithStats::new(ManuallyDrop::new( + ramulator::Ramulator::hbm2_preset(8).unwrap(), + )), + memory::WithStats::new(memory::MemoryBacked::with_capacity(*HBM_SIZE)), + )); let mut accelerator = Accelerator { m_machine: machine, @@ -1409,11 +1409,9 @@ async fn start() { // Memory Initialization // - HBM Preload let hbm_data = std::fs::read(opts.hbm).unwrap(); - { - hbm.lock().await.data().with_data(|f| { - f[..hbm_data.len()].copy_from_slice(&hbm_data); - }); - } + hbm.data().model().with_data(|f| { + f[..hbm_data.len()].copy_from_slice(&hbm_data); + }); // Load fpsram and intsram as raw bytes and map to the vector files. // - fpsram Preload @@ -1479,15 +1477,13 @@ async fn start() { vram_file.write_all(&vram_bytes).unwrap(); eprintln!("Dumped VRAM content to: {:?}", vram_dump_path); - { - let memory_stats = hbm.lock().await.statistics(); - let utilization = (memory_stats.total_bytes_read + memory_stats.total_bytes_written) as f64 - / Executor::current().now().to_secs(); - eprintln!( - "HBM Statistics - Bytes read: {:?} | Bytes written: {:?} | Utilization: {:.2e} bytes/sec", - memory_stats.total_bytes_read, memory_stats.total_bytes_written, utilization - ); - } + let memory_stats = hbm.statistics(); + let utilization = (memory_stats.total_bytes_read + memory_stats.total_bytes_written) as f64 + / Executor::current().now().to_secs(); + eprintln!( + "HBM Statistics - Bytes read: {:?} | Bytes written: {:?} | Utilization: {:.2e} bytes/sec", + memory_stats.total_bytes_read, memory_stats.total_bytes_written, utilization + ); } #[tokio::main] From dadb446ff15869925933ec4ac32e62640527ebab Mon Sep 17 00:00:00 2001 From: GeorgeWu1204 Date: Wed, 12 Nov 2025 12:30:35 +0000 Subject: [PATCH 07/34] fix bmm test bugs --- behavioral_simulator/src/main.rs | 38 ++++++++++---------- behavioral_simulator/testbench/bmm_test.py | 6 +++- compiler/asm_templates/batched_matmul_asm.py | 8 +++-- 3 files changed, 29 insertions(+), 23 deletions(-) diff --git a/behavioral_simulator/src/main.rs b/behavioral_simulator/src/main.rs index 80546f5c..cf4e647f 100644 --- a/behavioral_simulator/src/main.rs +++ b/behavioral_simulator/src/main.rs @@ -534,17 +534,17 @@ impl MatrixMachine { let (vec_base, vec_offset) = v_addr.multiple_and_offset(self.mlen); assert!(vec_offset.is_multiple_of(self.blen)); cycle!(1); - println!("======================== MM_WO =========================="); - println!("m accum = {}", self.m_accum); - println!("vec_base = {}, vec_offset = {}, stride_len = {}", vec_base, vec_offset, stride_len); + // println!("======================== MM_WO =========================="); + // println!("m accum = {}", self.m_accum); + // println!("vec_base = {}, vec_offset = {}, stride_len = {}", vec_base, vec_offset, stride_len); for i in 0..self.blen { let tensor = self.m_accum.i((i as i64, ..)); let old = self.vram.read(vec_base + i * self.mlen * stride_len).await; - println!("old = {}", old.as_tensor()); + // println!("old = {}", old.as_tensor()); let new = old.as_tensor().copy(); new.i(vec_offset as i64..(vec_offset + self.blen) as i64) .copy_(&tensor); - println!("new = {}", new); + // println!("new = {}", new); self.vram .write( vec_base + i * self.mlen * stride_len, @@ -561,7 +561,7 @@ impl MatrixMachine { async fn bmm_wo(&mut self, v_addr: u32) { let (vec_base, vec_offset) = v_addr.multiple_and_offset(self.mlen); - println!("======================== BMM_WO =========================="); + // println!("======================== BMM_WO =========================="); assert!(vec_offset.is_multiple_of(self.mlen)); cycle!(1); for j in 0..self.broadcast_amount { @@ -660,10 +660,10 @@ impl VectorMachine { cycle!(*VECTOR_MUL_CYCLES); self.vram.write(vd, c).await; } else { - println!("======================== V_ADD_VF =========================="); - println!("add: mask = {:?}", mask); - println!("a = {}", a.as_tensor()); - println!("f = {}", f); + // println!("======================== V_ADD_VF =========================="); + // println!("add: mask = {:?}", mask); + // println!("a = {}", a.as_tensor()); + // println!("f = {}", f); let mut result = a.as_tensor().shallow_clone(); let total_heads = self.tile_size / self.mask_unit; for head in 0..total_heads { @@ -689,10 +689,10 @@ impl VectorMachine { cycle!(*VECTOR_ADD_CYCLES); self.vram.write(vd, c).await; } else { - println!("======================== V_ADD =========================="); - println!("add: mask = {:?}", mask); - println!("a = {}", a.as_tensor()); - println!("b = {}", b.as_tensor()); + // println!("======================== V_ADD =========================="); + // println!("add: mask = {:?}", mask); + // println!("a = {}", a.as_tensor()); + // println!("b = {}", b.as_tensor()); let mut result = a.as_tensor().shallow_clone(); let total_heads = self.tile_size / self.mask_unit; for head in 0..total_heads { @@ -826,10 +826,10 @@ impl Accelerator { } else { load_dim }; - println!("Call transfer_from_hbm"); - println!("stride = {:?}", stride); - println!("index = {:?}", index); - println!("scale_index = {:?}", scale_index); + // println!("Call transfer_from_hbm"); + // println!("stride = {:?}", stride); + // println!("index = {:?}", index); + // println!("scale_index = {:?}", scale_index); Executor::current().spawn(async move { let element_ty = hbm_type.element_type(); @@ -896,7 +896,7 @@ impl Accelerator { let load_iter = write_idx * write_amount + block_idx; let element_addr = index + (load_iter * stride) as u64; let scale_addr = scale_index + (load_iter as f32 * stride_scale) as u64; - println!("element_addr = {:?}, scale_addr = {:?}", element_addr, scale_addr); + // println!("element_addr = {:?}, scale_addr = {:?}", element_addr, scale_addr); let byte_offset = (write_idx * write_amount * len_in_bytes_per_load) as usize + block_idx as usize * len_in_bytes_per_load as usize; let scale_byte_offset = (write_idx * write_amount * scale_len_in_bytes_per_load) diff --git a/behavioral_simulator/testbench/bmm_test.py b/behavioral_simulator/testbench/bmm_test.py index af19f0e4..7e30ebe4 100644 --- a/behavioral_simulator/testbench/bmm_test.py +++ b/behavioral_simulator/testbench/bmm_test.py @@ -6,7 +6,7 @@ from torch import Tensor, nn # from acc_simulator.quantize.quantized_layers.linear import MXFPLinearPTQ from test_data_gen import get_weights_path, generate_and_save_random_weights -from compiler.asm_templates import batched_matmul_asm, preload_addr_reg_asm +from compiler.asm_templates import batched_matmul_asm, preload_addr_reg_asm, reset_reg_asm from create_sim_env import create_sim_env from sim_env_utils import build_fake_sim_env @@ -44,6 +44,10 @@ addr_reg_val=[int(m * k * batch_size * real_data_ratio)] ) + gen_assembly_code += reset_reg_asm( + alive_registers=[1] + ) + gen_assembly_code += batched_matmul_asm( mlen=mlen, blen=blen, diff --git a/compiler/asm_templates/batched_matmul_asm.py b/compiler/asm_templates/batched_matmul_asm.py index ff41bf17..8083538d 100644 --- a/compiler/asm_templates/batched_matmul_asm.py +++ b/compiler/asm_templates/batched_matmul_asm.py @@ -40,6 +40,8 @@ def batched_matmul_asm( assert k % mlen == 0, "k must be divisible by mlen" assert m % blen == 0, "m must be divisible by blen" assert n % blen == 0, "n must be divisible by blen" + print(f"b = {b}, m = {m}, k = {k}, n = {n}") + print(f"mlen = {mlen}, blen = {blen}") a_actual_register = alive_registers[0] w_actual_register = alive_registers[1] @@ -47,7 +49,7 @@ def batched_matmul_asm( generated_code += f"S_ADDI_INT gp{result_actual_register}, gp0, {result_base_address} \n" - for batch in range(b): + for batch in range(1, b + 1): # preload the activation matrix for i in range(math.ceil(n // blen)): assert w_prefetch_amount >= k, "w_prefetch_amount must be greater than or equal to k" @@ -56,7 +58,7 @@ def batched_matmul_asm( generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, {n} \n" generated_code += f"C_SET_STRIDE_REG gp{w_actual_register} \n" generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, 0 \n" - generated_code += f"H_PREFETCH_M gp{w_actual_register}, gp{w_actual_register}, a{w_base_hbm_offset_reg}, {w_prefetch_amount // k}, 0 \n" + generated_code += f"H_PREFETCH_M gp{w_actual_register}, gp{w_actual_register}, a{w_base_hbm_offset_reg}, 1, 0 \n" for j in range(math.ceil(m // blen)): if j == 0: @@ -68,7 +70,7 @@ def batched_matmul_asm( if j % math.ceil((a_prefetch_amount * mlen) // k) == 0: generated_code += f"H_PREFETCH_V gp{a_actual_register}, gp{a_actual_register}, a{a_base_hbm_offset_reg}, 1, 0 \n" generated_code += f"S_ADDI_INT gp{a_actual_register}, gp0, 0 \n" - for k in range(math.ceil(k // mlen)): + for g in range(math.ceil(k // mlen)): generated_code += f"M_MM 0, gp{w_actual_register}, gp{a_actual_register} \n" generated_code += f"S_ADDI_INT gp{w_actual_register}, gp{w_actual_register}, {blen * mlen} \n" generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{a_actual_register}, {blen * mlen} \n" From a688bf86d3deae39221379819909ee917c5c8bdd Mon Sep 17 00:00:00 2001 From: Timi Adeniran Date: Wed, 12 Nov 2025 13:50:16 +0000 Subject: [PATCH 08/34] Simplify logic: wrap WithTiming in WithStats --- behavioral_simulator/lib/memory/src/lib.rs | 31 ++++------------------ behavioral_simulator/src/main.rs | 12 ++++----- 2 files changed, 10 insertions(+), 33 deletions(-) diff --git a/behavioral_simulator/lib/memory/src/lib.rs b/behavioral_simulator/lib/memory/src/lib.rs index f9b4d526..1d5e4c50 100644 --- a/behavioral_simulator/lib/memory/src/lib.rs +++ b/behavioral_simulator/lib/memory/src/lib.rs @@ -110,10 +110,14 @@ impl WithStats { pub fn model(&self) -> &T { &self.model } + + pub fn statistics(&self) -> Statistics { + self.statistics.lock().unwrap().clone() + } } #[async_trait::async_trait] -impl MemoryModel for WithStats { +impl MemoryModel for WithStats> { async fn read(&self, addr: u64) -> [u8; 64] { { let mut guard = self.statistics.lock().unwrap(); @@ -131,25 +135,6 @@ impl MemoryModel for WithStats { } } -#[async_trait::async_trait] -impl MemoryTimingModel for WithStats { - async fn read(&self, addr: u64) { - { - let mut guard = self.statistics.lock().unwrap(); - guard.total_bytes_read += 64; - } - self.model.read(addr).await - } - - async fn write(&self, addr: u64) { - { - let mut guard = self.statistics.lock().unwrap(); - guard.total_bytes_written += 64; - } - self.model.write(addr).await - } -} - /// Combine a data model with an extra timing model. pub struct WithTiming { timing: T, @@ -166,12 +151,6 @@ impl WithTiming { } } -impl WithTiming, WithStats> { - pub fn statistics(&self) -> Statistics { - self.timing.statistics.lock().unwrap().clone() - } -} - #[async_trait::async_trait] impl MemoryModel for WithTiming { /// Read 64-bytes of memory. diff --git a/behavioral_simulator/src/main.rs b/behavioral_simulator/src/main.rs index 26988821..e595bd71 100644 --- a/behavioral_simulator/src/main.rs +++ b/behavioral_simulator/src/main.rs @@ -1480,12 +1480,10 @@ async fn start() { }; let v_machine = VectorMachine { vram, tile_size: *VLEN, mask_unit: *HLEN }; // Share same dim with VSRAM - let hbm = Arc::new(memory::WithTiming::new( - memory::WithStats::new(ManuallyDrop::new( - ramulator::Ramulator::hbm2_preset(8).unwrap(), - )), - memory::WithStats::new(memory::MemoryBacked::with_capacity(*HBM_SIZE)), - )); + let hbm = Arc::new(memory::WithStats::new(memory::WithTiming::new( + ManuallyDrop::new(ramulator::Ramulator::hbm2_preset(8).unwrap()), + memory::MemoryBacked::with_capacity(*HBM_SIZE), + ))); let mut accelerator = Accelerator { m_machine: machine, @@ -1523,7 +1521,7 @@ async fn start() { // Memory Initialization // - HBM Preload let hbm_data = std::fs::read(opts.hbm).unwrap(); - hbm.data().model().with_data(|f| { + hbm.model().data().with_data(|f| { f[..hbm_data.len()].copy_from_slice(&hbm_data); }); From 3982e9b6f4fd88b6e80b4b9d2c2853dd08c969bc Mon Sep 17 00:00:00 2001 From: Timi Adeniran Date: Wed, 12 Nov 2025 14:13:57 +0000 Subject: [PATCH 09/34] Address review comment: make WithStats impl more generic --- behavioral_simulator/lib/memory/src/lib.rs | 65 +++++++++++----------- 1 file changed, 33 insertions(+), 32 deletions(-) diff --git a/behavioral_simulator/lib/memory/src/lib.rs b/behavioral_simulator/lib/memory/src/lib.rs index 1d5e4c50..a28cc83b 100644 --- a/behavioral_simulator/lib/memory/src/lib.rs +++ b/behavioral_simulator/lib/memory/src/lib.rs @@ -90,6 +90,38 @@ impl MemoryModel for MemoryBacked { } } +/// Combine a data model with an extra timing model. +pub struct WithTiming { + timing: T, + data: M, +} + +impl WithTiming { + pub fn new(timing: T, data: M) -> Self { + WithTiming { timing, data } + } + + pub fn data(&self) -> &M { + &self.data + } +} + +#[async_trait::async_trait] +impl MemoryModel for WithTiming { + /// Read 64-bytes of memory. + async fn read(&self, addr: u64) -> [u8; 64] { + self.timing.read(addr).await; + self.data.read(addr).await + } + + /// Write 64-bytes of memory. + async fn write(&self, addr: u64, bytes: [u8; 64]) { + self.timing.write(addr).await; + self.data.write(addr, bytes).await + } +} + +// Memory model with utilization statistics pub struct WithStats { model: T, statistics: Mutex, @@ -117,7 +149,7 @@ impl WithStats { } #[async_trait::async_trait] -impl MemoryModel for WithStats> { +impl MemoryModel for WithStats { async fn read(&self, addr: u64) -> [u8; 64] { { let mut guard = self.statistics.lock().unwrap(); @@ -134,34 +166,3 @@ impl MemoryModel for WithStats { - timing: T, - data: M, -} - -impl WithTiming { - pub fn new(timing: T, data: M) -> Self { - WithTiming { timing, data } - } - - pub fn data(&self) -> &M { - &self.data - } -} - -#[async_trait::async_trait] -impl MemoryModel for WithTiming { - /// Read 64-bytes of memory. - async fn read(&self, addr: u64) -> [u8; 64] { - self.timing.read(addr).await; - self.data.read(addr).await - } - - /// Write 64-bytes of memory. - async fn write(&self, addr: u64, bytes: [u8; 64]) { - self.timing.write(addr).await; - self.data.write(addr, bytes).await - } -} From baf4745f4a47d61e0fc5683c7d0e144f41e04515 Mon Sep 17 00:00:00 2001 From: GeorgeWu1204 Date: Thu, 13 Nov 2025 18:23:40 +0000 Subject: [PATCH 10/34] latency evaluation --- behavioral_simulator/readme.md | 1 + behavioral_simulator/testbench/ffn_test.py | 36 +++++ compiler/asm_templates/ffn_asm.py | 145 ++++++++++-------- src/definitions/plena_settings.toml | 8 +- tools/cost_model/latency/latency_model.py | 10 +- .../latency/overall_inference_estimation.py | 84 +++++----- 6 files changed, 171 insertions(+), 113 deletions(-) create mode 100644 behavioral_simulator/testbench/ffn_test.py diff --git a/behavioral_simulator/readme.md b/behavioral_simulator/readme.md index bc2dccc8..72d6dc3f 100644 --- a/behavioral_simulator/readme.md +++ b/behavioral_simulator/readme.md @@ -72,3 +72,4 @@ Writes a (BLEN, BLEN) accumulator matrix (`m_accum`) to the Vector SRAM. This op - **Linear Projection Testing** (`linear`) - **RMSNorm Testing** (`rms`) - **Attention Testing** (`attn`) +- **FFN Testing** (`ffn`) diff --git a/behavioral_simulator/testbench/ffn_test.py b/behavioral_simulator/testbench/ffn_test.py new file mode 100644 index 00000000..9ece99cc --- /dev/null +++ b/behavioral_simulator/testbench/ffn_test.py @@ -0,0 +1,36 @@ +import sys +from pathlib import Path +sys.path.insert(0, str(Path(__file__).parent.parent.parent)) + +import torch +from torch import Tensor, nn +# from acc_simulator.quantize.quantized_layers.linear import MXFPLinearPTQ +from test_data_gen import get_weights_path, generate_and_save_random_weights +from compiler.asm_templates import batched_matmul_asm, preload_addr_reg_asm, reset_reg_asm +from create_sim_env import create_sim_env +from sim_env_utils import build_fake_sim_env + + +class LlamaFeedForward(nn.Module): + """ + Standard FeedForward layer used in Llama architectures: + y = W2(activation(W1(x))) + where activation is typically SwiGLU in Llama2. + + Args: + dim (int): input and output dimension (hidden size) + inter_dim (int): intermediate/fc dimension + activation (callable): nonlinearity to use (default: SwiGLU) + """ + def __init__(self, dim: int, inter_dim: int, activation: str = "silu"): + super().__init__() + # Llama uses SwiGLU: x * silu(x) + self.w1 = nn.Linear(dim, inter_dim, bias=False) + self.w3 = nn.Linear(dim, inter_dim, bias=False) # for SwiGLU, "gate" proj + self.w2 = nn.Linear(inter_dim, dim, bias=False) + self.act = torch.nn.SiLU() if activation == "silu" else getattr(torch.nn, activation)() + + def forward(self, x: Tensor) -> Tensor: + # SwiGLU: (x @ w1) * silu(x @ w3) + return self.w2(self.act(self.w1(x)) * self.w3(x)) + diff --git a/compiler/asm_templates/ffn_asm.py b/compiler/asm_templates/ffn_asm.py index b56fcc61..327947cb 100644 --- a/compiler/asm_templates/ffn_asm.py +++ b/compiler/asm_templates/ffn_asm.py @@ -3,84 +3,109 @@ from pathlib import Path - def ffn_asm( mlen: int, vlen: int, blen: int, batch: int, + seq_len: int, hidden_size: int, - alive_registers: List[int], - weight_hbm_offset_reg: int, intermediate_size: int, - const_address: int, + + alive_registers: List[int], + gate_weight_hbm_offset_reg: int, + up_weight_hbm_offset_reg: int, + down_weight_hbm_offset_reg: int, + const_one_fp_address: int, + + w_prefetch_amount: int, activation_base_address: int, result_base_address: int ) -> str: """ - Generates assembly code for a general matrix multiplication operation. + Generates assembly code for a FFN operation. Args: - mlen (int): The number of rows in the first matrix. + mlen (int): The number of rows in the matrix. + vlen (int): The number of columns in the matrix. blen (int): The number of columns in the second matrix. + batch (int): The number of batches. + hidden_size (int): The number of rows in the hidden size. + intermediate_size (int): The number of rows in the intermediate size. alive_registers (List[int]): List of registers that are alive. - rope_base_address (int): index for the address mapper pointing to the base addr of the rope matrix. + gate_weight_hbm_offset_reg (int): index for the address mapper pointing to the base addr of the gate weight matrix. + up_weight_hbm_offset_reg (int): index for the address mapper pointing to the base addr of the up weight matrix. + down_weight_hbm_offset_reg (int): index for the address mapper pointing to the base addr of the down weight matrix. activation_base_address (int): index for the address mapper pointing to the base addr of the activation matrix. + result_base_address (int): index for the address mapper pointing to the base addr of the result matrix. Functionality: Upsize linear (b, s, hidden_size) @ (hidden_size, intermediate_size) - > (b, s, intermediate_size) - Activation (b, s, intermediate_size) -> (b, s, intermediate_size) + Gate Projection (b, s, hidden_size) @ (hidden_size, intermediate_size) -> (b, s, intermediate_size) + SILU Activation (b, s, intermediate_size) -> (b, s, intermediate_size) Downsize linear (b, s, intermediate_size) @ (intermediate_size, hidden_size) -> (b, s, hidden_size) """ - generated_code = " ; FFN Upsize Linear Generation \n" + # Dot product of weight (Hidden Size, Hidden Size) and activation (Batch, 1, Hidden Size) - assert batch < blen, "Batch size must be greater than blen" # get two registers from alive_registers, 1 as w address, 1 as a address - w_base_register = alive_registers[0] - a_base_register = alive_registers[1] + + w_actual_register = alive_registers[0] + a_actual_register = alive_registers[1] result_register = alive_registers[2] - w_actual_register = alive_registers[3] - a_actual_register = alive_registers[4] - intermediate_register = alive_registers[5] + intermediate_register = alive_registers[3] + # reset the registers - set_w_base_register = f"S_ADDI_INT gp{w_base_register}, gp0, 0 \n" - set_a_base_address = f"S_ADDI_INT gp{a_base_register}, gp0, {activation_base_address} \n" - set_result_address = f"S_ADDI_INT gp{result_register}, gp0, {result_base_address} \n" + generated_code = "; FFN Generation \n" - set_w_actual_address = f"S_ADD_INT gp{w_actual_register}, gp0, {w_base_register} \n" - set_a_actual_address = f"S_ADD_INT gp{a_actual_register}, gp0, {a_base_register} \n" + # Settings + generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, {hidden_size * intermediate_size} \n" + generated_code += f"C_SET_SCALE_REG gp{w_actual_register} \n" + generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, {intermediate_size} \n" + generated_code += f"C_SET_STRIDE_REG gp{w_actual_register} \n" + generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, 0 \n" - increment_result_actual_address = f"S_ADDI_INT gp{result_register}, gp{result_register}, {mlen} \n" + generated_code = " ; FFN Upsize Linear Generation \n" + for weight_row in range (intermediate_size // blen): + if weight_row % (mlen // blen) == 0: + for weight_col in range (hidden_size // mlen): + generated_code += f"H_PREFETCH_M gp{w_actual_register}, gp{w_actual_register}, a{up_weight_hbm_offset_reg}, 1, 0 \n" + generated_code += f"S_ADDI_INT gp{w_actual_register}, gp{w_actual_register}, {mlen * mlen} \n" + generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, 0 \n" + for act_col in range (batch * seq_len // blen): + generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{a_actual_register}, {activation_base_address + act_col * hidden_size * blen} \n" + for inner_loop_index in range (hidden_size // mlen): + generated_code += f"M_MM 0, gp{w_actual_register}, gp{a_actual_register} \n" + generated_code += f"S_ADDI_INT gp{w_actual_register}, gp{w_actual_register}, {mlen * mlen} \n" + generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{a_actual_register}, {mlen * blen} \n" + generated_code += f"M_MM_WO gp{result_register}, 0, 0 \n" + generated_code += f"S_ADDI_INT gp{w_actual_register}, gp{w_actual_register}, {hidden_size * blen} \n" - row_loop_over_hid = intermediate_size // blen - vect_loop_over_hid = intermediate_size // vlen - col_loop_over_hid = hidden_size // mlen - generated_code += set_w_base_register - generated_code += set_a_base_address - generated_code += set_result_address - - generated_code += f"PREFETECH_M {w_actual_register}, gp0, a{weight_hbm_offset_reg}, 1, 0 \n" - for i in range(row_loop_over_hid): - generated_code += f"; <---- Generating New Row Tile at index {i} ----> \n" - for j in range(col_loop_over_hid): - generated_code += f"; <---- Generating New Column Tile at row {i} col {j} ----> \n" - generated_code += f"M_MM 0, {w_actual_register}, {a_actual_register} \n" - generated_code += set_w_actual_address - generated_code += set_a_actual_address - generated_code += f"M_MM_WO {result_register}, 0, 0 \n" - if (i % blen) == 0: - generated_code += increment_result_actual_address - + + generated_code = " ; FFN Gate Projection Generation \n" + for weight_row in range (intermediate_size // blen): + if weight_row % (mlen // blen) == 0: + for weight_col in range (hidden_size // mlen): + generated_code += f"H_PREFETCH_M gp{w_actual_register}, gp{w_actual_register}, a{gate_weight_hbm_offset_reg}, 1, 0 \n" + generated_code += f"S_ADDI_INT gp{w_actual_register}, gp{w_actual_register}, {mlen * mlen} \n" + generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, 0 \n" + for act_col in range (batch * seq_len // blen): + generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{a_actual_register}, {activation_base_address + act_col * hidden_size * blen} \n" + for inner_loop_index in range (hidden_size // mlen): + generated_code += f"M_MM 0, {w_actual_register}, {a_actual_register} \n" + generated_code += f"S_ADDI_INT gp{w_actual_register}, gp{w_actual_register}, {mlen * mlen} \n" + generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{a_actual_register}, {mlen * blen} \n" + generated_code += f"M_MM_WO {result_register}, 0, 0 \n" + generated_code += f"S_ADDI_INT gp{w_actual_register}, gp{w_actual_register}, {hidden_size * blen} \n" + + # Hidden Size SILU Activation Generation generated_code += "; SILU Generation \n" - fp_const_reg = "f1" - generated_code += f"S_LD_FP {fp_const_reg}, gp0, {const_address} \n" - for i in range(vect_loop_over_hid): - generated_code += f"; <---- per VLEN block {i} ----> \n" - generated_code += f"V_SUB_VV {intermediate_register}, {intermediate_register}, {result_register} \n" - generated_code += f"V_EXP_V {intermediate_register}, {intermediate_register} \n" - generated_code += f"V_ADD_VF {intermediate_register}, {intermediate_register}, {fp_const_reg} \n" - generated_code += f"V_ADD_VF {intermediate_register}, {intermediate_register}, {fp_const_reg} \n" - generated_code += f"V_REC_V {intermediate_register}, {intermediate_register} \n" - generated_code += f"V_MUL_VV {result_register}, {result_register}, {intermediate_register} \n" + generated_code += f"S_LD_FP f1, gp0, {const_one_fp_address} \n" + for i in range(intermediate_size // vlen): + generated_code += f"V_SUB_VV gp{intermediate_register}, gp{intermediate_register}, gp{result_register} \n" + generated_code += f"V_EXP_V gp{intermediate_register}, gp{intermediate_register} \n" + generated_code += f"V_ADD_VF gp{intermediate_register}, gp{intermediate_register}, f1 \n" + generated_code += f"V_ADD_VF gp{intermediate_register}, gp{intermediate_register}, f1 \n" + generated_code += f"V_REC_V gp{intermediate_register}, gp{intermediate_register} \n" + generated_code += f"V_MUL_VV gp{result_register}, gp{result_register}, gp{intermediate_register} \n" generated_code += "; FFN Downsize Linear Generation \n" @@ -88,15 +113,15 @@ def ffn_asm( row_loop_over_hid = hidden_size // blen col_loop_over_hid = intermediate_size // mlen - for i in range(row_loop_over_hid): - generated_code += f"; <---- Generating New Row Tile at index {i} ----> \n" - for j in range(col_loop_over_hid): - generated_code += f"; <---- Generating New Column Tile at row {i} col {j} ----> \n" - generated_code += f"M_MM 0, {w_actual_register}, {result_register} \n" - generated_code += set_w_actual_address - generated_code += set_result_address - generated_code += f"M_MM_WO {a_actual_register}, 0, 0 \n" - if (i % blen) == 0: - generated_code += set_a_actual_address + # for i in range(row_loop_over_hid): + # generated_code += f"; <---- Generating New Row Tile at index {i} ----> \n" + # for j in range(col_loop_over_hid): + # generated_code += f"; <---- Generating New Column Tile at row {i} col {j} ----> \n" + # generated_code += f"M_MM 0, {w_actual_register}, {result_register} \n" + # generated_code += set_w_actual_address + # generated_code += set_result_address + # generated_code += f"M_MM_WO {a_actual_register}, 0, 0 \n" + # if (i % blen) == 0: + # generated_code += set_a_actual_address return generated_code diff --git a/src/definitions/plena_settings.toml b/src/definitions/plena_settings.toml index 3f27e552..6d347ea1 100644 --- a/src/definitions/plena_settings.toml +++ b/src/definitions/plena_settings.toml @@ -2,16 +2,16 @@ value = 1 [CONFIG.BLEN] -value = 4 +value = 64 [CONFIG.HLEN] -value = 16 +value = 128 [CONFIG.MLEN] -value = 64 +value = 2048 [CONFIG.VLEN] -value = 64 +value = 2048 [CONFIG.BROADCAST_AMOUNT] value = 4 diff --git a/tools/cost_model/latency/latency_model.py b/tools/cost_model/latency/latency_model.py index 520a16e7..42f48100 100644 --- a/tools/cost_model/latency/latency_model.py +++ b/tools/cost_model/latency/latency_model.py @@ -79,9 +79,9 @@ def obtain_per_instr_alone_latency(self, output_file: str = "instr_alone_latency def obtain_overall_latency(self, updated_config): batch_size = 1024 - input_seq_len = 1024 - output_seq_len = 128 - device_num = 4 + input_seq_len = 5600 + output_seq_len = 85000 + device_num = 8 hardware_settings = self.hardware_config for key, value in updated_config.items(): hardware_settings[key] = value @@ -102,8 +102,8 @@ def obtain_overall_latency(self, updated_config): custom_isa_parent_path = os.path.dirname(os.path.abspath(__file__)) custom_isa_path = os.path.join(custom_isa_parent_path, "customISA_lib.json") # model_config_path = os.path.join(config_parent_path, "doc/Model_Lib/qwen2_5_7b.json") - # model_config_path = os.path.join(config_parent_path, "doc/Model_Lib/llama-3.3-70b.json") - model_config_path = os.path.join(config_parent_path, "doc/Model_Lib/llama-3.1-8b.json") + model_config_path = os.path.join(config_parent_path, "doc/Model_Lib/llama-3.3-70b.json") + # model_config_path = os.path.join(config_parent_path, "doc/Model_Lib/llama-3.1-8b.json") model = instr_latency_model(config_path, custom_isa_path, model_config_path) test_from_toml = load_toml_config(toml_path, "active") diff --git a/tools/cost_model/latency/overall_inference_estimation.py b/tools/cost_model/latency/overall_inference_estimation.py index 793ce487..2c767cda 100644 --- a/tools/cost_model/latency/overall_inference_estimation.py +++ b/tools/cost_model/latency/overall_inference_estimation.py @@ -24,6 +24,7 @@ def __init__(self, model_param_path, hardware_config, batch_size = 1, seq_len = self.theoratical_frequency = 10**9 # 1 GHz self.hardware_config = hardware_config self.batch_size = batch_size + self.device_batch_size = batch_size // device_num self.kv_size = seq_len self.device_num = device_num print("=" * 15, "Model Settings","=" * 15) @@ -31,6 +32,7 @@ def __init__(self, model_param_path, hardware_config, batch_size = 1, seq_len = print("batch size: ", self.batch_size) print("input token: ", self.input_token) print("output token: ", self.output_token) + print("hidden size: ", self.hidden_size) print("head dim: ", self.head_dim) print("num key value heads: ", self.num_key_value_heads) print("num attention heads: ", self.num_attention_heads) @@ -54,30 +56,30 @@ def rms_layer(self, mode = "prefill"): instruction_num = 0 instruction_num += setting_inst_num instruction_num += loop_num * loop_inst_num - return instruction_num * self.batch_size + return instruction_num * self.device_batch_size def projection(self, mode = "prefill"): # Compute Q, K, V Projection + RoPE if mode == "prefill": overall_inst_num = 0 # Q, K Projection + RoPE - overall_inst_num += self.batch_size * (math.ceil(self.hidden_size / self.hardware_config["BLEN"]) * (math.ceil(self.hidden_size / self.hardware_config["MLEN"]) * (math.ceil(self.input_token / self.hardware_config["BLEN"]) * 2 + 10))) - overall_inst_num += self.batch_size * (self.num_attention_heads * (self.input_token // self.hardware_config["VLEN"])) * 3 + overall_inst_num += self.device_batch_size * (math.ceil(self.hidden_size / self.hardware_config["BLEN"]) * (math.ceil(self.hidden_size / self.hardware_config["MLEN"]) * (math.ceil(self.input_token / self.hardware_config["BLEN"]) * 2 + 10))) + overall_inst_num += self.device_batch_size * (self.num_attention_heads * math.ceil(self.input_token / self.hardware_config["VLEN"])) * 3 - overall_inst_num += self.batch_size * (math.ceil((self.num_key_value_heads * self.head_dim) / self.hardware_config["BLEN"]) * (math.ceil(self.hidden_size / self.hardware_config["MLEN"]) * (math.ceil(self.input_token / self.hardware_config["BLEN"]) * 2 + 10))) - overall_inst_num += self.batch_size * (self.num_key_value_heads * (self.input_token // self.hardware_config["VLEN"])) * 3 + overall_inst_num += self.device_batch_size * (math.ceil((self.num_key_value_heads * self.head_dim) / self.hardware_config["BLEN"]) * (math.ceil(self.hidden_size / self.hardware_config["MLEN"]) * (math.ceil(self.input_token / self.hardware_config["BLEN"]) * 2 + 10))) + overall_inst_num += self.device_batch_size * (self.num_key_value_heads * math.ceil(self.input_token / self.hardware_config["VLEN"])) * 3 # V - overall_inst_num += self.batch_size * (math.ceil((self.num_key_value_heads * self.head_dim) / self.hardware_config["BLEN"]) * (math.ceil(self.hidden_size / self.hardware_config["MLEN"]) * (math.ceil(self.input_token / self.hardware_config["BLEN"]) * 2 + 10))) + overall_inst_num += self.device_batch_size * (math.ceil((self.num_key_value_heads * self.head_dim) / self.hardware_config["BLEN"]) * (math.ceil(self.hidden_size / self.hardware_config["MLEN"]) * (math.ceil(self.input_token / self.hardware_config["BLEN"]) * 2 + 10))) elif mode == "decode": overall_inst_num = 0 # Q, K Projection + RoPE overall_inst_num += (math.ceil(self.hidden_size / self.hardware_config["BLEN"]) * (math.ceil(self.hidden_size / self.hardware_config["MLEN"]) * 2 + 10)) - overall_inst_num += self.batch_size * (self.num_attention_heads) * 4 + overall_inst_num += self.device_batch_size * (self.num_attention_heads) * 4 overall_inst_num += (math.ceil((self.num_key_value_heads * self.head_dim) / self.hardware_config["BLEN"]) * (math.ceil(self.hidden_size / self.hardware_config["MLEN"]) * 2 + 10)) - overall_inst_num += self.batch_size * (self.num_key_value_heads) * 4 + overall_inst_num += self.device_batch_size * (self.num_key_value_heads) * 4 # V overall_inst_num += (math.ceil((self.num_key_value_heads * self.head_dim) / self.hardware_config["BLEN"]) * (math.ceil(self.hidden_size / self.hardware_config["MLEN"]) * 2 + 10)) @@ -88,30 +90,31 @@ def flash_attention(self, mode = "prefill"): overall_inst_num = 0 mlen = self.hardware_config["MLEN"] blen = self.hardware_config["BLEN"] - tile_in_atten = min(self.head_dim, mlen) + prefill_len = min(self.input_token, mlen) + decode_len = min(self.kv_size, mlen) + per_head_iter = math.ceil(self.num_attention_heads / (self.hardware_config["MLEN"] / self.head_dim)) if mode == "prefill": # Outer loop for kv_head_index in range(self.num_key_value_heads): - for i in range(math.ceil(self.input_token // self.hardware_config["MLEN"])): - overall_inst_num += mlen * 2 # Reset - for j in range(math.ceil(self.input_token // self.hardware_config["MLEN"])): - overall_inst_num += mlen * 2 # QKT - overall_inst_num += 2 + mlen * 14 # Softmax - overall_inst_num += math.ceil(self.head_dim / blen) * (4 + math.ceil(mlen / blen) * 4) #PV - overall_inst_num += mlen * 5 + 4 #Compute O + for i in range(math.ceil(self.input_token / self.hardware_config["MLEN"])): + overall_inst_num += mlen # Reset + for j in range(math.ceil(self.input_token / self.hardware_config["MLEN"])): + overall_inst_num += prefill_len * 2 # QKT + overall_inst_num += 2 + prefill_len * 5 # Softmax + overall_inst_num += math.ceil(self.head_dim / blen) * (4 + math.ceil(prefill_len / blen) * 4) #PV + overall_inst_num += prefill_len * 5 + 4 #Compute O overall_inst_num += 8 self.kv_size = self.input_token + elif mode == "decode": - for kv_head_index in range(self.num_key_value_heads): - for i in range(math.ceil(self.kv_size // self.hardware_config["MLEN"])): - # overall_inst_num += mlen * 2 # Reset - overall_inst_num += mlen * 2 # QKT - overall_inst_num += 2 + mlen * 14 # Softmax - overall_inst_num += math.ceil(self.head_dim / blen) * (4 + math.ceil(mlen / blen) * 4) #PV - overall_inst_num += mlen * 5 + 4 #Compute O + for kv_head_index in range(per_head_iter): + for i in range(math.ceil(self.kv_size / self.hardware_config["MLEN"])): + # overall_inst_num += decode_len * 2 # QKT + overall_inst_num += 2 + decode_len * 3 # Softmax + overall_inst_num += math.ceil(self.head_dim / blen) * (4 + math.ceil(decode_len / blen) * 4) #PV + overall_inst_num += decode_len * 3 + 4 #Compute O overall_inst_num += 8 - self.kv_size = self.kv_size + 1 - return overall_inst_num * self.batch_size + return overall_inst_num * self.device_batch_size def residual (self, mode = "prefill"): @@ -123,7 +126,7 @@ def residual (self, mode = "prefill"): elif mode == "decode": iteration = self.hidden_size // self.hardware_config["VLEN"] overall_inst_num = 5 * iteration + 3 - return overall_inst_num + return overall_inst_num * self.device_batch_size def feed_forward(self, mode = "prefill"): mlen = self.hardware_config["MLEN"] @@ -133,14 +136,14 @@ def feed_forward(self, mode = "prefill"): overall_inst_num = 0 # -- MLP if mode == "prefill": - overall_inst_num += 2 * math.ceil(self.intermediate_size / blen) * math.ceil(self.hidden_size / mlen) * 4 * (self.input_token // blen) + overall_inst_num += 2 * math.ceil(self.intermediate_size / blen) * math.ceil(self.input_token / blen) * math.ceil(self.hidden_size / mlen) * 4 overall_inst_num += math.ceil(self.intermediate_size / vlen) * 5 - overall_inst_num += math.ceil(self.intermediate_size / blen) * math.ceil(self.hidden_size / mlen) * 4 * (self.input_token // blen) - overall_inst_num = (overall_inst_num) * self.batch_size + overall_inst_num += math.ceil(self.intermediate_size / blen) * math.ceil(self.hidden_size / mlen) * math.ceil(self.input_token / blen) * 4 + overall_inst_num = (overall_inst_num) * self.device_batch_size elif mode == "decode": - overall_inst_num += 2 * math.ceil(self.intermediate_size / blen) * math.ceil(self.hidden_size / mlen) * (math.ceil(self.batch_size // blen) ) * 4 - overall_inst_num += math.ceil(self.intermediate_size / vlen) * 5 - overall_inst_num += math.ceil(self.intermediate_size / blen) * math.ceil(self.hidden_size / mlen) * 4 + overall_inst_num += 2 * math.ceil(self.intermediate_size / blen) * math.ceil(self.hidden_size / mlen) * (math.ceil(self.device_batch_size / blen) ) * 4 + overall_inst_num += math.ceil(self.intermediate_size / vlen) * 5 * self.device_batch_size + overall_inst_num += math.ceil(self.intermediate_size / blen) * math.ceil(self.hidden_size / mlen) * (math.ceil(self.device_batch_size / blen) ) * 4 return overall_inst_num def embeddings(self, mode = "prefill"): @@ -173,7 +176,7 @@ def compute_prefill_time(self): overall_inst_num += self.residual(mode) overall_inst_num += self.rms_layer(mode) overall_inst_num += self.feed_forward(mode) - # overall_inst_num += self.residual(mode) + # overall_inst_num += self.residual(mode) # overall_inst_num += self.rms_layer() overall_inst_num += self.lm_head() overall_exe_cycle = overall_inst_num * 2 @@ -204,6 +207,7 @@ def compute_decode_time(self, output_token_size): residual_count += self.residual(mode) rms_count += self.rms_layer(mode) feed_forward_count += self.feed_forward(mode) + self.kv_size = self.kv_size + 1 overall_inst_num = rms_count + projection_count + flash_attention_count + residual_count + feed_forward_count overall_exe_cycle = overall_inst_num * 2 # avg 2 execution cycles theoratical_execution_time = overall_exe_cycle / self.theoratical_frequency @@ -216,16 +220,8 @@ def compute_decode_time(self, output_token_size): print(f"Feed Forward: {feed_forward_count / overall_inst_num * 100}%") return theoratical_execution_time - def compute_overall_perf(self): + # Per batch performance. ttft = (self.compute_prefill_time() + self.compute_decode_time(1)) / self.device_num - tps = (self.device_num * (self.batch_size * self.output_token)) / self.compute_decode_time(self.output_token // self.device_num) - return ttft, tps - - - - - - - - + tps = (self.batch_size * self.output_token) / self.compute_decode_time(self.output_token) + return ttft, tps \ No newline at end of file From d83a7268d40a3711a437561e233ee702edd176d0 Mon Sep 17 00:00:00 2001 From: GeorgeWu1204 Date: Wed, 19 Nov 2025 00:42:16 +0000 Subject: [PATCH 11/34] cost model for paper --- .gitignore | 3 + src/definitions/plena_settings.toml | 4 +- tools/cost_model/latency/ablation_study.py | 334 ++++++++++++++ tools/cost_model/latency/latency_model.py | 12 +- .../latency/overall_inference_estimation.py | 112 +++-- tools/cost_model/utilisation/attainable.py | 408 ++++++++++++------ .../utilisation/attainable_gemm_test.py | 104 +++++ .../combined_ffn_fa_utilization.py | 178 ++++++++ .../utilisation/fa_utilizaiton_test.py | 85 ++++ .../utilisation/individual_units_lib.json | 6 +- .../utilisation/memory_utilization_plot.py | 179 ++++++++ .../utilisation/utilisation_model.py | 37 +- 12 files changed, 1298 insertions(+), 164 deletions(-) create mode 100644 tools/cost_model/latency/ablation_study.py create mode 100644 tools/cost_model/utilisation/attainable_gemm_test.py create mode 100644 tools/cost_model/utilisation/combined_ffn_fa_utilization.py create mode 100644 tools/cost_model/utilisation/fa_utilizaiton_test.py create mode 100644 tools/cost_model/utilisation/memory_utilization_plot.py diff --git a/.gitignore b/.gitignore index b28d259f..bf1e27b7 100644 --- a/.gitignore +++ b/.gitignore @@ -18,6 +18,9 @@ build/ doc/figure_gen/* +tools/cost_model/latency/*.png +tools/cost_model/utilisation/*.png + moduleparamid default.svf diff --git a/src/definitions/plena_settings.toml b/src/definitions/plena_settings.toml index 6d347ea1..ea863825 100644 --- a/src/definitions/plena_settings.toml +++ b/src/definitions/plena_settings.toml @@ -8,10 +8,10 @@ value = 64 value = 128 [CONFIG.MLEN] -value = 2048 +value = 1024 [CONFIG.VLEN] -value = 2048 +value = 1024 [CONFIG.BROADCAST_AMOUNT] value = 4 diff --git a/tools/cost_model/latency/ablation_study.py b/tools/cost_model/latency/ablation_study.py new file mode 100644 index 00000000..e8fb334e --- /dev/null +++ b/tools/cost_model/latency/ablation_study.py @@ -0,0 +1,334 @@ +import json +import os +from pathlib import Path +import matplotlib.pyplot as plt +import numpy as np +from overall_inference_estimation import model_config + +batch_size = 4 + +# Experiment configurations +# sample_1: baseline, MLEN = BLEN = 64, VLEN = 1024, using selfattention +# sample_2: MLEN = BLEN = 64, VLEN = 1024, using flash attention +# sample_3: MLEN = 1024, BLEN = 4, VLEN = 1024, using flash attention, not partitioned optimised +# sample_4: MLEN = 1024, BLEN = 4, VLEN = 1024, using flash attention, partitioned optimised + + +def compute_attention_ffn_distribution(model_param_path, hardware_config, batch_size, input_seq, decode_seq_lengths, + use_flash_attention=True, partitioned_optimized=False, device_num=1): + """ + Compute attention and FFN computation distribution for prefill and decode stages at different sequence lengths. + + Args: + model_param_path: Path to model parameter JSON file + hardware_config: Dictionary with MLEN, BLEN, VLEN + batch_size: Batch size + input_seq: Input sequence length (for prefill) + decode_seq_lengths: List of total sequence lengths for decode (e.g., [10k, 40k, 80k]) + use_flash_attention: Whether to use flash attention (True) or self attention (False) + partitioned_optimized: Whether to use partitioned optimization (for sample 4) + device_num: Number of devices + + Returns: + dict: Dictionary with 'prefill' and 'decode' keys + 'prefill' contains 'attention' and 'ffn' for input_seq + 'decode' contains list of dicts, each with 'attention' and 'ffn' for each decode_seq_length + """ + print("<", "="*10, "starting computation", "="*10, ">") + print("hardware config: ", hardware_config) + print("use flash attention: ", use_flash_attention) + print("partitioned optimized: ", partitioned_optimized) + + print("starting computation") + # Use max decode length for model initialization + max_decode_len = max(decode_seq_lengths) if decode_seq_lengths else 0 + model = model_config( + model_param_path=model_param_path, + hardware_config=hardware_config, + batch_size=batch_size, + seq_len=int(input_seq), + output_token=1000, + device_num=device_num + ) + + results = {} + + # Prefill stage + mode = "prefill" + if use_flash_attention: + attention_inst = model.flash_attention(mode, partitioned_optimized) + else: + attention_inst = model.self_attention(mode) + + ffn_inst = model.feed_forward(mode) + + results['prefill'] = { + 'attention': attention_inst, + 'ffn': ffn_inst, + 'total': attention_inst + ffn_inst + } + print("completed prefill computation") + print("prefill attention: ", attention_inst) + print("prefill ffn: ", ffn_inst) + + # Decode stage - compute for each sequence length + mode = "decode" + results['decode'] = [] + + for total_seq_len in decode_seq_lengths: + model.kv_size = total_seq_len # Reset KV size to initial value + + # Accumulate instruction counts as KV cache grows + attention_inst_decode = 0 + ffn_inst_decode = 0 + + for token_idx in range(total_seq_len): + if use_flash_attention: + attention_inst_per_token = model.flash_attention(mode, partitioned_optimized) + else: + attention_inst_per_token = model.self_attention(mode) + ffn_inst_per_token = model.feed_forward(mode) + + attention_inst_decode += attention_inst_per_token + ffn_inst_decode += ffn_inst_per_token + + # KV cache grows by 1 for each token + model.kv_size += 1 + + results['decode'].append({ + 'attention': attention_inst_decode, + 'ffn': ffn_inst_decode, + 'total': attention_inst_decode + ffn_inst_decode + }) + print(f"completed decode computation for seq_len={total_seq_len}") + print(f" decode attention: {attention_inst_decode:.2e}, ffn: {ffn_inst_decode:.2e}") + + return results + + +def plot_ablation_study(): + """ + Conduct ablation study experiments and plot bar chart showing + FFN + Attention computation distribution across sequence lengths. + Three plots: Left (Prefill 5.6k), Middle (Decode 5k), Right (Decode 8k) + Y-axis: Execution Time (% relative to S1) + """ + colors = [ + "#762a83", # deep purple + "#af8dc3", # light violet + "#e7d4e8", # very light purple-pink + "#d9f0d3", # pale green + "#7fbf7b", # medium green + # Optional sixth if needed: "#1b7837" # dark green + ] + + # Get model parameter path + current_dir = Path(__file__).resolve().parents[3] + model_param_path = os.path.join(current_dir, "doc/Model_Lib/llama-3.1-8b.json") + + # Sequence lengths: prefill at 5.6k, decode at 5k and 8k total + input_seq = 5.6 * 1000 + decode_seq_lengths = [5 * 1000, 8 * 10000] # Total sequence lengths after decode (5k and 8k) + + # Define experiment samples + experiments = [ + { + 'name': 'Sample 1', + 'hardware_config': {'MLEN': 64, 'BLEN': 64, 'VLEN': 1024}, + 'use_flash_attention': False, + 'partitioned_optimized': False, + 'label': 'Sample 1 (Baseline): MLEN=BLEN=64, VLEN=1024' + }, + { + 'name': 'Sample 2', + 'hardware_config': {'MLEN': 64, 'BLEN': 64, 'VLEN': 1024}, + 'use_flash_attention': True, + 'partitioned_optimized': False, + 'label': 'Sample 2 (Flash Attention): MLEN=BLEN=64, VLEN=1024' + }, + { + 'name': 'Sample 3', + 'hardware_config': {'MLEN': 1024, 'BLEN': 4, 'VLEN': 1024}, + 'use_flash_attention': True, + 'partitioned_optimized': True, + 'label': 'Sample 3 (Flash Attention + Flattened Systolic Array): MLEN=1024, BLEN=4, VLEN=1024' + } + ] + + # Collect data for all sequence lengths + # Structure: data[sample_idx][seq_idx] = {'attention': ..., 'ffn': ...} + all_data = [] + experiment_names = [] + + for exp in experiments: + results = compute_attention_ffn_distribution( + model_param_path=model_param_path, + hardware_config=exp['hardware_config'], + batch_size=batch_size, + input_seq=input_seq, + decode_seq_lengths=decode_seq_lengths, + use_flash_attention=exp['use_flash_attention'], + partitioned_optimized=exp['partitioned_optimized'] + ) + + experiment_names.append(exp['name']) + sample_data = [] + + # Add prefill data (at 5.6k) + sample_data.append({ + 'attention': results['prefill']['attention'], + 'ffn': results['prefill']['ffn'] + }) + + # Add decode data (at 10k, 40k, 80k) + for decode_result in results['decode']: + sample_data.append({ + 'attention': decode_result['attention'], + 'ffn': decode_result['ffn'] + }) + + all_data.append(sample_data) + + print("completed computation") + + # Prepare data for plotting + num_samples = len(experiments) + + # Extract data for each plot: prefill (index 0), decode 5k (index 1), decode 8k (index 2) + prefill_attention = np.array([all_data[s][0]['attention'] for s in range(num_samples)]) + prefill_ffn = np.array([all_data[s][0]['ffn'] for s in range(num_samples)]) + + decode_5k_attention = np.array([all_data[s][1]['attention'] for s in range(num_samples)]) + decode_5k_ffn = np.array([all_data[s][1]['ffn'] for s in range(num_samples)]) + + decode_8k_attention = np.array([all_data[s][2]['attention'] for s in range(num_samples)]) + decode_8k_ffn = np.array([all_data[s][2]['ffn'] for s in range(num_samples)]) + + # Normalize all values relative to S1 (S1 = 100%) + s1_prefill_total = prefill_attention[0] + prefill_ffn[0] + s1_decode_5k_total = decode_5k_attention[0] + decode_5k_ffn[0] + s1_decode_8k_total = decode_8k_attention[0] + decode_8k_ffn[0] + + # Normalize prefill + prefill_attention_norm = (prefill_attention / s1_prefill_total) * 100 + prefill_ffn_norm = (prefill_ffn / s1_prefill_total) * 100 + + # Normalize decode 5k + decode_5k_attention_norm = (decode_5k_attention / s1_decode_5k_total) * 100 + decode_5k_ffn_norm = (decode_5k_ffn / s1_decode_5k_total) * 100 + + # Normalize decode 8k + decode_8k_attention_norm = (decode_8k_attention / s1_decode_8k_total) * 100 + decode_8k_ffn_norm = (decode_8k_ffn / s1_decode_8k_total) * 100 + + # Create figure with 3 subplots, shared y-axis - reduced height for compactness + fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=(10, 3), sharey=True) + + # Colors - using the color palette defined above + attention_color = colors[0] # dark blue + ffn_color = colors[4] # teal + + # Bar positioning - make more compact + x = np.arange(num_samples) + bar_width = 0.5 # Slightly wider bars for better visibility + + # Find maximum value across all plots for consistent y-axis + max_total = max( + (prefill_attention_norm + prefill_ffn_norm).max(), + (decode_5k_attention_norm + decode_5k_ffn_norm).max(), + (decode_8k_attention_norm + decode_8k_ffn_norm).max() + ) + y_max = max_total * 1.05 # Reduced padding to 5% for more compact y-axis + + # Data for each plot (normalized) + plot_data = [ + { + 'ax': ax1, + 'xlabel': 'Prefill (5.6k)', + 'attention': prefill_attention_norm, + 'ffn': prefill_ffn_norm + }, + { + 'ax': ax2, + 'xlabel': 'Decode (10k)', + 'attention': decode_5k_attention_norm, + 'ffn': decode_5k_ffn_norm + }, + { + 'ax': ax3, + 'xlabel': 'Decode (80k)', + 'attention': decode_8k_attention_norm, + 'ffn': decode_8k_ffn_norm + } + ] + + # Plot each subplot + for plot_info in plot_data: + ax = plot_info['ax'] + attn_data = plot_info['attention'] + ffn_data = plot_info['ffn'] + + # Plot stacked bars + ax.bar(x, attn_data, bar_width, label='Attention', color=attention_color, edgecolor='black', alpha=0.8) + ax.bar(x, ffn_data, bar_width, bottom=attn_data, label='FFN', color=ffn_color, edgecolor='black', alpha=0.8) + + # X-axis - make more compact + ax.set_xticks(x) + ax.set_xticklabels([f'S{i+1}' for i in range(num_samples)], fontsize=13) + # Set title at the top instead of x-axis label + ax.set_title(plot_info['xlabel'], fontsize=14) + + # Set x-axis limits to reduce space between bars + ax.set_xlim(-0.5, num_samples - 0.5) + + # Set shared y-axis limits with minimal padding + ax.set_ylim(0, y_max) + ax.grid(axis='y', alpha=0.3, linestyle='--') + + # Reduce tick label sizes for more compact appearance + ax.tick_params(axis='both', labelsize=12) + + # Add legend only to the first plot + if ax == ax1: + ax.legend(loc='upper right', fontsize=13, framealpha=0.9) + # Y-axis label only on the leftmost plot - positioned lower + ax.set_ylabel('Exec Time (% relative to S1)', fontsize=12) + # Move y-axis label down by adjusting its position + ax.yaxis.set_label_coords(-0.15, 0.50) + else: + # Remove y-axis labels from middle and right plots + ax.set_ylabel('') + + # Add overall title + # fig.suptitle('Ablation Study for LLaMA 3.1 8B', fontsize=18, y=1.02) + + # Add sample legend at the bottom in three rows (one per sample) + sample_legend = [exp["label"] for exp in experiments] + # Arrange in three rows: one sample per row + legend_text = '\n'.join(sample_legend) + + # More compact vertical spacing + plt.subplots_adjust(wspace=0.15, left=0.12, right=0.95, top=0.88, bottom=0.25) + + # Add sample legend at the bottom in three rows + # fig.text(0.5, 0.12, legend_text, + # fontsize=12, va='top', ha='center', + # bbox=dict(facecolor='white', alpha=0.8, edgecolor='none', pad=4)) + + plt.savefig('ablation_study_attention_ffn.png', dpi=300, bbox_inches='tight') + print("Plot saved as 'ablation_study_attention_ffn.png'") + + # Print summary + print("\n" + "="*70) + print("Ablation Study Results Summary") + print("="*70) + for i, exp in enumerate(experiments): + print(f"\n{exp['label']}:") + print(f" Prefill (5.6k) - Attention: {prefill_attention[i]:.2e}, FFN: {prefill_ffn[i]:.2e}") + print(f" Decode (5k) - Attention: {decode_5k_attention[i]:.2e}, FFN: {decode_5k_ffn[i]:.2e}") + print(f" Decode (8k) - Attention: {decode_8k_attention[i]:.2e}, FFN: {decode_8k_ffn[i]:.2e}") + print("="*70) + + +if __name__ == "__main__": + plot_ablation_study() diff --git a/tools/cost_model/latency/latency_model.py b/tools/cost_model/latency/latency_model.py index 42f48100..769004e3 100644 --- a/tools/cost_model/latency/latency_model.py +++ b/tools/cost_model/latency/latency_model.py @@ -78,10 +78,10 @@ def obtain_per_instr_alone_latency(self, output_file: str = "instr_alone_latency print(f"Alone latency model saved to {output_file}") def obtain_overall_latency(self, updated_config): - batch_size = 1024 - input_seq_len = 5600 - output_seq_len = 85000 - device_num = 8 + batch_size = 4 + input_seq_len = 2048 + output_seq_len = 1024 + device_num = 1 hardware_settings = self.hardware_config for key, value in updated_config.items(): hardware_settings[key] = value @@ -102,8 +102,8 @@ def obtain_overall_latency(self, updated_config): custom_isa_parent_path = os.path.dirname(os.path.abspath(__file__)) custom_isa_path = os.path.join(custom_isa_parent_path, "customISA_lib.json") # model_config_path = os.path.join(config_parent_path, "doc/Model_Lib/qwen2_5_7b.json") - model_config_path = os.path.join(config_parent_path, "doc/Model_Lib/llama-3.3-70b.json") - # model_config_path = os.path.join(config_parent_path, "doc/Model_Lib/llama-3.1-8b.json") + # model_config_path = os.path.join(config_parent_path, "doc/Model_Lib/llama-3.3-70b.json") + model_config_path = os.path.join(config_parent_path, "doc/Model_Lib/llama-3.1-8b.json") model = instr_latency_model(config_path, custom_isa_path, model_config_path) test_from_toml = load_toml_config(toml_path, "active") diff --git a/tools/cost_model/latency/overall_inference_estimation.py b/tools/cost_model/latency/overall_inference_estimation.py index 2c767cda..ff80547f 100644 --- a/tools/cost_model/latency/overall_inference_estimation.py +++ b/tools/cost_model/latency/overall_inference_estimation.py @@ -86,36 +86,98 @@ def projection(self, mode = "prefill"): return overall_inst_num - def flash_attention(self, mode = "prefill"): + def flash_attention(self, mode = "prefill", partitoned_optimised = True): overall_inst_num = 0 mlen = self.hardware_config["MLEN"] blen = self.hardware_config["BLEN"] prefill_len = min(self.input_token, mlen) decode_len = min(self.kv_size, mlen) - per_head_iter = math.ceil(self.num_attention_heads / (self.hardware_config["MLEN"] / self.head_dim)) + max_head_per_mlen = math.ceil(mlen / self.head_dim) + + per_head_iter = math.ceil(self.num_attention_heads / max_head_per_mlen) + if mode == "prefill": # Outer loop - for kv_head_index in range(self.num_key_value_heads): - for i in range(math.ceil(self.input_token / self.hardware_config["MLEN"])): - overall_inst_num += mlen # Reset - for j in range(math.ceil(self.input_token / self.hardware_config["MLEN"])): - overall_inst_num += prefill_len * 2 # QKT + if partitoned_optimised: + for i in range(math.ceil(self.input_token / mlen)): + overall_inst_num += mlen * max_head_per_mlen # Reset + for j in range(math.ceil(self.input_token / mlen)): + # overall_inst_num += math.ceil(prefill_len / blen) * blen * math.ceil(prefill_len / blen) * 5 # QKT + overall_inst_num += (2 + prefill_len * 5) * max_head_per_mlen # Softmax + overall_inst_num += math.ceil(mlen / blen) * blen * math.ceil(self.head_dim / blen) * max_head_per_mlen + overall_inst_num += (prefill_len * 5 + 4) * max_head_per_mlen# Compute O + overall_inst_num += 8 * max_head_per_mlen + self.kv_size = self.input_token + return overall_inst_num * per_head_iter * self.device_batch_size + else: + for i in range(math.ceil(self.input_token / mlen)): + # overall_inst_num += mlen # Reset + for j in range(math.ceil(self.input_token / mlen)): + # (mlen, head_dim) @ (head_dim, mlen) + overall_inst_num += math.ceil(prefill_len / blen) * (math.ceil(self.head_dim / mlen)) * blen * math.ceil(prefill_len / blen) overall_inst_num += 2 + prefill_len * 5 # Softmax - overall_inst_num += math.ceil(self.head_dim / blen) * (4 + math.ceil(prefill_len / blen) * 4) #PV - overall_inst_num += prefill_len * 5 + 4 #Compute O + overall_inst_num += math.ceil(mlen / blen) * blen * math.ceil(self.head_dim / blen) + overall_inst_num += prefill_len * 5 + 4 # Compute O overall_inst_num += 8 - self.kv_size = self.input_token - + return overall_inst_num * self.device_batch_size * self.num_attention_heads elif mode == "decode": - for kv_head_index in range(per_head_iter): - for i in range(math.ceil(self.kv_size / self.hardware_config["MLEN"])): - # overall_inst_num += decode_len * 2 # QKT - overall_inst_num += 2 + decode_len * 3 # Softmax - overall_inst_num += math.ceil(self.head_dim / blen) * (4 + math.ceil(decode_len / blen) * 4) #PV - overall_inst_num += decode_len * 3 + 4 #Compute O - overall_inst_num += 8 - return overall_inst_num * self.device_batch_size + if partitoned_optimised: + overall_inst_num = decode_len + overall_inst_num += (2 + decode_len * 3) * max_head_per_mlen # Softmax + # PV (1, decode_len) @ (decode_len, head_dim) + overall_inst_num += math.ceil(decode_len / mlen) * blen * math.ceil(self.head_dim / blen) * max_head_per_mlen #PV + overall_inst_num += (decode_len * 3 + 4) * max_head_per_mlen #Compute O + overall_inst_num += 8 * max_head_per_mlen + return overall_inst_num * self.device_batch_size * math.ceil(self.kv_size / self.hardware_config["MLEN"]) * (per_head_iter) + else: + overall_inst_num = math.ceil(mlen / blen) * (math.ceil(self.head_dim / mlen) * blen) + overall_inst_num += (2 + decode_len * 3) + overall_inst_num += math.ceil(self.head_dim / blen) * (4 + math.ceil(decode_len / blen) * 4) + overall_inst_num += (decode_len * 3 + 4) + overall_inst_num += 8 + return overall_inst_num * self.device_batch_size * math.ceil(self.kv_size / self.hardware_config["MLEN"]) * self.num_attention_heads + + def self_attention(self, mode = "prefill"): + # Note: this it the latency estimation model for self-attention without using the flash attention algorithm + overall_inst_num = 0 + mlen = self.hardware_config["MLEN"] + blen = self.hardware_config["BLEN"] + vlen = self.hardware_config["VLEN"] + if mode == "prefill": + for i in range(self.num_attention_heads): + # QKT (batch, s, h, d) @ (batch, s, h, d) + overall_inst_num += (math.ceil(self.input_token / blen)) * math.ceil(self.head_dim / mlen) * blen * math.ceil(self.input_token / blen) * 4 + # Store QKT (h, s, s) + overall_inst_num += self.input_token * 4 * (math.ceil(self.input_token / mlen)) * self.num_attention_heads * 2 + # Softmax (batch, s, h, d) + overall_inst_num += self.input_token * (math.ceil(self.input_token / mlen)) * 30 * self.num_attention_heads + # PV (h, s, s) @ (h, s, d) + for i in range(self.num_attention_heads): + overall_inst_num += math.ceil(self.input_token / blen) * math.ceil(self.input_token / mlen) * blen * math.ceil(self.head_dim / blen) + # Store PV (h, s, d) + overall_inst_num += self.head_dim * 4 * (math.ceil(self.input_token / mlen)) * self.num_attention_heads* 2 + # Compute O (batch, s, h, d) + for i in range(self.num_attention_heads): + overall_inst_num += (math.ceil(self.input_token / vlen)) * 5 + 4 + overall_inst_num += 8 + overall_inst_num = overall_inst_num * self.device_batch_size + elif mode == "decode": + for i in range(self.num_attention_heads): + # QKT (batch, 1, h, d) @ (batch, kv, h, d) + overall_inst_num += 4 * math.ceil(self.head_dim / mlen) * math.ceil(self.kv_size / blen) * blen + # Store QKT (h, s, s) + overall_inst_num += 4 * math.ceil(self.kv_size / mlen) * self.num_attention_heads * 36 + # Softmax (batch, s, h, d) + overall_inst_num += (math.ceil(self.kv_size / mlen)) * 30 * self.num_attention_heads + # PV (h, 1, s) @ (h, s, d) + for i in range(self.num_attention_heads): + overall_inst_num += math.ceil(self.kv_size / mlen) * blen * (self.head_dim // blen) * 2 + # Compute O (batch, s, h, d) + for i in range(self.num_attention_heads): + overall_inst_num += (math.ceil(self.kv_size / mlen)) * 5 + 4 + overall_inst_num = overall_inst_num * self.device_batch_size + return overall_inst_num def residual (self, mode = "prefill"): overall_inst_num = 0 @@ -136,14 +198,16 @@ def feed_forward(self, mode = "prefill"): overall_inst_num = 0 # -- MLP if mode == "prefill": - overall_inst_num += 2 * math.ceil(self.intermediate_size / blen) * math.ceil(self.input_token / blen) * math.ceil(self.hidden_size / mlen) * 4 - overall_inst_num += math.ceil(self.intermediate_size / vlen) * 5 - overall_inst_num += math.ceil(self.intermediate_size / blen) * math.ceil(self.hidden_size / mlen) * math.ceil(self.input_token / blen) * 4 + # Upprojection, and Gate (seq, hidden) @ (hidden, intermediate) + overall_inst_num += 2 * math.ceil(self.intermediate_size / blen) * math.ceil(self.input_token / blen) * math.ceil(self.hidden_size / mlen) * blen + overall_inst_num += math.ceil(self.intermediate_size / vlen) * 3 * self.input_token + # Downprojection (seq, intermediate) @ (intermediate, hidden) + overall_inst_num += math.ceil(self.intermediate_size / blen) * math.ceil(self.hidden_size / mlen) * math.ceil(self.input_token / blen) * blen overall_inst_num = (overall_inst_num) * self.device_batch_size elif mode == "decode": - overall_inst_num += 2 * math.ceil(self.intermediate_size / blen) * math.ceil(self.hidden_size / mlen) * (math.ceil(self.device_batch_size / blen) ) * 4 + overall_inst_num += 2 * math.ceil(self.intermediate_size / blen) * math.ceil(self.hidden_size / mlen) * (math.ceil(self.device_batch_size / blen) ) * blen * math.ceil(self.device_batch_size / blen) overall_inst_num += math.ceil(self.intermediate_size / vlen) * 5 * self.device_batch_size - overall_inst_num += math.ceil(self.intermediate_size / blen) * math.ceil(self.hidden_size / mlen) * (math.ceil(self.device_batch_size / blen) ) * 4 + overall_inst_num += math.ceil(self.intermediate_size / blen) * math.ceil(self.hidden_size / mlen) * (math.ceil(self.device_batch_size / blen) ) * blen * math.ceil(self.device_batch_size / blen) return overall_inst_num def embeddings(self, mode = "prefill"): diff --git a/tools/cost_model/utilisation/attainable.py b/tools/cost_model/utilisation/attainable.py index f3392ad9..82993e0f 100644 --- a/tools/cost_model/utilisation/attainable.py +++ b/tools/cost_model/utilisation/attainable.py @@ -1,9 +1,11 @@ import os from typing import Dict, List, Any, Optional import json +import math class attn_model_config: - def __init__(self, model_param_path, hardware_config, batch_size = 1, seq_len = 2048, output_token = 128, device_num = 1): + def __init__(self, MLEN, BLEN, VLEN,partitioned_matrix, model_param_path, batch_size = 1, seq_len = 2048, output_token = 128, device_num = 1): + print(f"Model param path: {model_param_path}") model_param = json.load(open(model_param_path)) self.hidden_size = model_param["hidden_size"] self.num_attention_heads = model_param["num_attention_heads"] @@ -17,15 +19,20 @@ def __init__(self, model_param_path, hardware_config, batch_size = 1, seq_len = self.vocab_size = model_param["vocab_size"] self.DataTypeSize = 2 self.theoratical_frequency = 10**9 # 1 GHz - self.hardware_config = hardware_config self.output_token = output_token self.batch_size = batch_size + print("=" * 25) + print(f"MLEN: {MLEN}, BLEN: {BLEN}, VLEN: {VLEN}") print(f"Batch size: {self.batch_size}") + print(f"num_attention_heads: {self.num_attention_heads}, num_key_value_heads: {self.num_key_value_heads}") + print(f"head_dim: {self.head_dim}, num_head_groups: {self.num_head_groups}") self.kv_size = seq_len self.device_num = device_num - self.M = hardware_config["BLEN"] - self.K = hardware_config["MLEN"] - self.N = self.M + self.MLEN = MLEN + self.BLEN = BLEN + self.VLEN = VLEN + self.partitioned_matrix = partitioned_matrix + self.max_head_per_mlen = math.ceil(self.MLEN / self.head_dim) def _report_flash_attn_utilization(self, mode = "prefill") -> None: """ @@ -40,47 +47,220 @@ def _report_flash_attn_utilization(self, mode = "prefill") -> None: input_token_size = self.input_seq_len theoretical_operation = 0 attainable_operation = 0 + attainable_gemm_operation_amount = 0 overall_operation_amount = 0 + bubble_ratio = 1.2 # Decoding if mode == "prefill": # Projection - operation_amount = ((head_dim * num_attn_heads) // self.M) * ( hidden_size // self.K) * (self.input_seq_len // self.M) + ((head_dim * num_kv_heads) // self.M) * ( hidden_size// self.K) * (self.input_seq_len // self.M) * 2 - overall_operation_amount += operation_amount - attainable_operation += operation_amount * (self.M * self.K * self.N) - theoretical_operation += operation_amount * (self.M * self.K * self.N) - # QKT - operation_amount = batch_size * num_attn_heads * (head_dim // self.K) * (self.input_seq_len // self.N) - overall_operation_amount += operation_amount - attainable_operation += operation_amount * (self.M * min(self.K, head_dim)) - theoretical_operation += operation_amount * (self.M * self.K * self.N) - - # PV - operation_amount = batch_size * num_attn_heads * (input_token_size // self.K) * (head_dim // self.N) - overall_operation_amount += operation_amount - attainable_operation += operation_amount * (self.M * min(self.K, head_dim)) - theoretical_operation += operation_amount * (self.M * self.K * self.N) + gemm_operation_amount = (self.input_seq_len // self.BLEN) * ( hidden_size / self.MLEN) * self.BLEN * ((head_dim * num_attn_heads) // self.BLEN) + ((head_dim * num_kv_heads) // self.BLEN) * ( hidden_size / self.MLEN) * self.BLEN * (self.input_seq_len // self.BLEN) * 2 + attainable_gemm_operation_amount = gemm_operation_amount * batch_size + attention_operation_amount = gemm_operation_amount * bubble_ratio * batch_size + + # Flash Attention + if self.partitioned_matrix: + for b in range(batch_size): + for i in range(math.ceil(num_attn_heads / self.max_head_per_mlen)): + for j in range(math.ceil(self.input_seq_len / self.MLEN)): + for k in range(math.ceil(self.input_seq_len / self.MLEN)): + # QKT (MLEN, Hq) * (Hq, MLEN), (num_attn_heads / (self.MLEN // self.head_dim)) in parallel + gemm_operation_amount = self.head_dim # need to clarify this, coz the actual implementation is not like this + attention_operation_amount += gemm_operation_amount * bubble_ratio # Prefetch Data + if num_attn_heads > self.max_head_per_mlen: + attainable_gemm_operation_amount += gemm_operation_amount * 1 + else: + attainable_gemm_operation_amount += gemm_operation_amount * (head_dim * num_attn_heads / self.MLEN) + # Softmax + attention_operation_amount += self.MLEN * 10 + # PV (MLEN, MLEN) @ (MLEN, head_dim), (num_attn_heads / (self.MLEN // self.head_dim)) in parallel + gemm_operation_amount = math.ceil(self.MLEN / self.BLEN) * self.BLEN * math.ceil(head_dim / self.BLEN) * self.num_head_groups + attention_operation_amount += gemm_operation_amount * bubble_ratio + if head_dim > self.BLEN: + attainable_gemm_operation_amount += gemm_operation_amount * 1 + else: + attainable_gemm_operation_amount += gemm_operation_amount * (head_dim / self.BLEN) + attention_operation_amount += self.MLEN * self.num_head_groups * (head_dim // self.BLEN) * bubble_ratio + # Compute O + attention_operation_amount += self.MLEN * 5 + 4 + else: + for b in range(batch_size): + for i in range(math.ceil(num_attn_heads)): + for j in range(math.ceil(self.input_seq_len / self.MLEN)): + for k in range(math.ceil(self.input_seq_len / self.MLEN)): + # QKT (MLEN, head_dim) @ (head_dim, MLEN) Full Utilization + gemm_operation_amount = (self.MLEN // self.BLEN) * math.ceil(head_dim / self.MLEN) * self.BLEN * (self.MLEN // self.BLEN) + attention_operation_amount += gemm_operation_amount * bubble_ratio # Prefetch Data + if head_dim > self.MLEN: + attainable_gemm_operation_amount += gemm_operation_amount * 1 + else: + attainable_gemm_operation_amount += gemm_operation_amount * (head_dim / self.MLEN) + # Softmax + attention_operation_amount += self.MLEN * 10 + # PV (MLEN, MLEN) @ (MLEN, head_dim) + gemm_operation_amount = (self.MLEN // self.BLEN) * math.ceil(head_dim / self.BLEN) * self.BLEN + attention_operation_amount += gemm_operation_amount * bubble_ratio + if head_dim > self.BLEN: + attainable_gemm_operation_amount += gemm_operation_amount * 1 + else: + attainable_gemm_operation_amount += gemm_operation_amount * (head_dim / self.BLEN) + # Compute O + attention_operation_amount += self.MLEN * 5 + 4 + + attainable_operation = attainable_gemm_operation_amount + theoretical_operation = attention_operation_amount + elif mode == "decode": + # Projection Q, K, V + gemm_operation_amount = ((head_dim * num_attn_heads) // self.BLEN) * ( hidden_size // self.MLEN) * self.BLEN * (math.ceil(self.batch_size / self.BLEN)) + ((head_dim * num_kv_heads) // self.BLEN) * ( hidden_size// self.MLEN) * self.BLEN * (math.ceil(self.batch_size / self.BLEN)) * 2 + attention_operation_amount = gemm_operation_amount * bubble_ratio + if batch_size > self.BLEN: + attainable_gemm_operation_amount = gemm_operation_amount * 1 + else: + attainable_gemm_operation_amount = gemm_operation_amount * (batch_size / self.BLEN) + + # Flash Attention + if self.partitioned_matrix: + for b in range(batch_size): + for i in range(math.ceil(num_kv_heads / self.max_head_per_mlen)): + for j in range(math.ceil(self.kv_size / self.MLEN)): + # QKT (self.num_head_groups, 1, Hq) * (Hq, MLEN), (num_kv_heads / (self.MLEN // self.head_dim)) in parallel + gemm_operation_amount = math.ceil(self.num_head_groups / self.BLEN) * self.head_dim * math.ceil(self.MLEN / self.BLEN) + attention_operation_amount += gemm_operation_amount * bubble_ratio + if num_kv_heads > (self.MLEN // self.head_dim): + if self.num_head_groups > self.BLEN: + attainable_gemm_operation_amount += gemm_operation_amount * 1 + else: + attainable_gemm_operation_amount += gemm_operation_amount * (self.num_head_groups / self.BLEN) + else: + if self.num_head_groups > self.BLEN: + attainable_gemm_operation_amount += gemm_operation_amount * 1 * (num_kv_heads / (self.MLEN // self.head_dim)) + else: + attainable_gemm_operation_amount += gemm_operation_amount * (self.num_head_groups / self.BLEN) * (num_kv_heads / (self.MLEN // self.head_dim)) + # Softmax + attention_operation_amount += self.MLEN * 10 + # PV (self.num_head_groups, 1, MLEN) @ (MLEN, head_dim), (num_kv_heads / (self.MLEN // self.head_dim)) in parallel + gemm_operation_amount = math.ceil(self.num_head_groups / self.BLEN) * math.ceil(head_dim / self.BLEN) * self.num_head_groups + attention_operation_amount += gemm_operation_amount * bubble_ratio + if self.num_head_groups > self.BLEN: + attainable_gemm_operation_amount += gemm_operation_amount * 1 + else: + attainable_gemm_operation_amount += gemm_operation_amount * (self.num_head_groups / self.BLEN) + # Compute O + attention_operation_amount += self.MLEN * 5 + 4 + + else: + for b in range(batch_size): + for i in range(math.ceil(num_attn_heads)): + for j in range(math.ceil(self.kv_size / self.MLEN)): + # QKT (1, head_dim) * (head_dim, MLEN), (num_kv_heads / (self.MLEN // self.head_dim)) + gemm_operation_amount = math.ceil(self.head_dim / self.MLEN) * math.ceil(self.MLEN / self.BLEN) * self.BLEN + attention_operation_amount += gemm_operation_amount * bubble_ratio # Prefetch Data + attainable_gemm_operation_amount += gemm_operation_amount * 1 * (1 / self.BLEN) * (self.head_dim / self.MLEN) + + # Softmax + attention_operation_amount += self.MLEN * 10 + + # PV (1, MLEN) * (MLEN, head_dim), + gemm_operation_amount = (head_dim // self.BLEN) * self.BLEN + attention_operation_amount += gemm_operation_amount * bubble_ratio + attainable_gemm_operation_amount += gemm_operation_amount * 1 * (1 / self.BLEN) + + # Compute O + attention_operation_amount += self.MLEN * 5 + 4 + attainable_operation = attainable_gemm_operation_amount + theoretical_operation = attention_operation_amount + + return [attainable_operation, theoretical_operation] + + def _report_self_attention_utilization(self, mode = "prefill") -> None: + """ + Report the utilization of self attention without using flash attention. + """ + batch_size = self.batch_size + hidden_size = self.hidden_size + num_attn_heads = self.num_attention_heads + num_kv_heads = self.num_key_value_heads + + head_dim = self.head_dim + input_token_size = self.input_seq_len + theoretical_operation = 0 + attainable_operation = 0 + attainable_gemm_operation_amount = 0 + overall_operation_amount = 0 + bubble_ratio = 1.2 + + if mode == "prefill": # Projection - operation_amount = ((head_dim * num_attn_heads) // self.M) * ( hidden_size // self.K) + ((head_dim * num_kv_heads) // self.M) * ( hidden_size// self.K) * 2 - overall_operation_amount += operation_amount - attainable_operation += operation_amount * (self.M * self.K * min(self.batch_size, self.N)) - theoretical_operation += operation_amount * (self.M * self.K * self.N) - - # QKT - operation_amount = batch_size * num_attn_heads * (head_dim // self.K) * (self.kv_size // self.N) - overall_operation_amount += operation_amount - attainable_operation += operation_amount * (self.M * min(self.K, head_dim)) - theoretical_operation += operation_amount * (self.M * self.K * self.N) - - # PV - operation_amount = batch_size * num_attn_heads * (self.kv_size // self.K) * (head_dim // self.N) - overall_operation_amount += operation_amount - attainable_operation += operation_amount * (self.M * min(self.K, head_dim)) - theoretical_operation += operation_amount * (self.M * self.K * self.N) - self.kv_size = self.kv_size + 1 + gemm_operation_amount = ((head_dim * num_attn_heads) // self.BLEN) * ( hidden_size // self.MLEN) * self.BLEN * (self.input_seq_len // self.BLEN) + ((head_dim * num_kv_heads) // self.BLEN) * ( hidden_size// self.MLEN) * self.BLEN * (self.input_seq_len // self.BLEN) * 2 + attainable_gemm_operation_amount = gemm_operation_amount * batch_size + attention_operation_amount = gemm_operation_amount * bubble_ratio * batch_size + + # QKT (batch, input_token_size, num_attn_heads, head_dim) @ (batch, input_token_size, num_attn_heads, head_dim) + for b in range(batch_size): + for i in range(num_attn_heads): + gemm_operation_amount = (self.input_token_size // self.BLEN) * (head_dim // self.MLEN) * (self.input_token_size // self.BLEN) + attention_operation_amount += gemm_operation_amount * bubble_ratio + if head_dim > self.MLEN: + attainable_gemm_operation_amount += gemm_operation_amount * 1 + else: + attainable_gemm_operation_amount += gemm_operation_amount * (head_dim / self.MLEN) + # Storing and fetching QKT (input_token_size, input_token_size) + attention_operation_amount += self.input_token_size * 4 * (self.input_token_size / self.VLEN) + # Softmax + attention_operation_amount += self.input_token_size * (self.input_token_size / self.VLEN) * 30 + # Storing and fetching O (input_token_size, head_dim) + attention_operation_amount += self.input_token_size * (self.input_token_size / self.VLEN) * 5 + + # PV (self.input_token_size, self.input_token_size) @ (self.input_token_size, head_dim) + gemm_operation_amount = math.ceil(self.input_token_size / self.BLEN) * math.ceil(head_dim / self.BLEN) * self.num_head_groups + attention_operation_amount += gemm_operation_amount * bubble_ratio + if head_dim > self.BLEN: + attainable_gemm_operation_amount += gemm_operation_amount * 1 + else: + attainable_gemm_operation_amount += gemm_operation_amount * (head_dim / self.BLEN) + # Compute O + attention_operation_amount += self.input_token_size * (self.input_token_size / self.VLEN) * 5 + 4 + attainable_operation = attainable_gemm_operation_amount + theoretical_operation = attention_operation_amount - return [operation_amount, attainable_operation, theoretical_operation] + elif mode == "decode": + # Projection Q, K, V + gemm_operation_amount = ((head_dim * num_attn_heads) // self.BLEN) * ( hidden_size // self.MLEN) * self.BLEN * (math.ceil(self.batch_size / self.BLEN)) + ((head_dim * num_kv_heads) // self.BLEN) * ( hidden_size// self.MLEN) * self.BLEN * (math.ceil(self.batch_size / self.BLEN)) * 2 + attention_operation_amount = gemm_operation_amount * bubble_ratio + if batch_size > self.BLEN: + attainable_gemm_operation_amount = gemm_operation_amount * 1 + else: + attainable_gemm_operation_amount = gemm_operation_amount * (batch_size / self.BLEN) + # QKT (batch, 1, num_attn_heads, head_dim) @ (batch, kv_size, num_attn_heads, head_dim) + for b in range(batch_size): + for i in range(num_attn_heads): + gemm_operation_amount = (self.kv_size // self.BLEN) * (head_dim // self.MLEN) * (self.kv_size // self.BLEN) + attention_operation_amount += gemm_operation_amount * bubble_ratio + if head_dim > self.MLEN: + attainable_gemm_operation_amount += gemm_operation_amount * 1 + else: + attainable_gemm_operation_amount += gemm_operation_amount * (head_dim / self.MLEN) + # Storing and fetching QKT (1, kv_size) + attention_operation_amount += 4 * (self.kv_size / self.VLEN) + # Softmax + attention_operation_amount += (self.kv_size / self.VLEN) * 30 + # Storing and fetching O (input_token_size, head_dim) + attention_operation_amount += (self.input_token_size / self.VLEN) * 5 + + # PV (1 self.kv_size) @ (self.kv_size, head_dim) + gemm_operation_amount = math.ceil(self.kv_size / self.MLEN) * math.ceil(head_dim / self.BLEN) + attention_operation_amount += gemm_operation_amount * bubble_ratio + if head_dim > self.BLEN: + attainable_gemm_operation_amount += gemm_operation_amount * 1 + else: + attainable_gemm_operation_amount += gemm_operation_amount * (head_dim / self.BLEN) + # Compute O + attention_operation_amount += (self.kv_size / self.VLEN) * 15 + 4 + + attainable_operation = attainable_gemm_operation_amount + theoretical_operation = attention_operation_amount + return [attainable_operation, theoretical_operation] def _report_embedding_utilization(self, mode = "prefill") -> None: """ @@ -95,13 +275,13 @@ def _report_embedding_utilization(self, mode = "prefill") -> None: if mode == "prefill": # Assuming Decoding only - operation_amount = (hidden_size // self.M) * (hidden_size // self.K) * (self.input_seq_len // self.N) - attainable_operation += operation_amount * (self.M * self.K * self.N) - theoretical_operation += operation_amount * (self.M * self.K * self.N) + operation_amount = (hidden_size // self.BLEN) * (hidden_size // self.MLEN) * (self.input_seq_len // self.BLEN) + attainable_operation += operation_amount * (self.BLEN * self.MLEN * self.BLEN) + theoretical_operation += operation_amount * (self.BLEN * self.MLEN * self.BLEN) elif mode == "decode": - operation_amount = (hidden_size // self.M) * (hidden_size // self.K) - attainable_operation += operation_amount * (self.M * self.K * min(self.batch_size, self.N)) - theoretical_operation += operation_amount * (self.M * self.K * self.N) + operation_amount = (hidden_size // self.BLEN) * (hidden_size // self.MLEN) + attainable_operation += operation_amount * (self.BLEN * self.MLEN * min(self.batch_size, self.BLEN)) + theoretical_operation += operation_amount * (self.BLEN * self.MLEN * self.BLEN) return [operation_amount, attainable_operation, theoretical_operation] @@ -109,115 +289,84 @@ def _report_ffn_utilization(self, mode = "prefill") -> None: """ Report the utilization of flash attention for a given node. """ - batch_size = self.batch_size hidden_size = self.hidden_size intermediate_size = self.intermediate_size - overall_operation_amount = 0 theoretical_operation = 0 attainable_operation = 0 + bubble_ratio = 1.2 if mode == "prefill": # Up Projection - operation_amount = (intermediate_size // self.M) * (hidden_size // self.K) * (self.input_seq_len // self.N) - overall_operation_amount += operation_amount - attainable_operation += operation_amount * (self.M * self.K * self.N) - theoretical_operation += operation_amount * (self.M * self.K * self.N) + operation_amount = (intermediate_size // self.BLEN) * (hidden_size // self.MLEN) * (self.input_seq_len // self.BLEN) + attainable_operation += operation_amount + theoretical_operation += operation_amount * bubble_ratio # Gate Projection - operation_amount = (intermediate_size // self.M) * (hidden_size // self.K) * (self.input_seq_len // self.N) - overall_operation_amount += operation_amount - attainable_operation += operation_amount * (self.M * self.K * self.N) - theoretical_operation += operation_amount * (self.M * self.K * self.N) + operation_amount = (intermediate_size // self.BLEN) * (hidden_size // self.MLEN) * (self.input_seq_len // self.BLEN) + attainable_operation += operation_amount + theoretical_operation += operation_amount * bubble_ratio + + # SiLU + theoretical_operation += (intermediate_size // self.VLEN) * 5 # Down Projection - operation_amount = (hidden_size // self.M) * (intermediate_size // self.K) * (self.input_seq_len // self.N) - overall_operation_amount += operation_amount - attainable_operation += operation_amount * (self.M * self.K * self.N) - theoretical_operation += operation_amount * (self.M * self.K * self.N) + operation_amount = (hidden_size // self.BLEN) * (intermediate_size // self.MLEN) * (self.input_seq_len // self.BLEN) + attainable_operation += operation_amount + theoretical_operation += operation_amount * bubble_ratio elif mode == "decode": - operation_amount = (intermediate_size // self.M) * (hidden_size // self.K) - overall_operation_amount += operation_amount - attainable_operation += operation_amount * (self.M * self.K * min(self.batch_size, self.N)) - theoretical_operation += operation_amount * (self.M * self.K * self.N) + + # Up projection + fill_bubble_time = self.BLEN * (intermediate_size // self.BLEN) * math.ceil(self.batch_size / self.BLEN) + operation_amount = (intermediate_size // self.BLEN) * self.BLEN * (hidden_size // self.MLEN) * math.ceil(self.batch_size / self.BLEN) + attainable_operation += operation_amount * (min(self.batch_size, self.BLEN) / self.BLEN) + theoretical_operation += operation_amount + fill_bubble_time # Gate Projection - operation_amount = (intermediate_size // self.M) * (hidden_size // self.K) - overall_operation_amount += operation_amount - attainable_operation += operation_amount * (self.M * self.K * min(self.batch_size, self.N)) - theoretical_operation += operation_amount * (self.M * self.K * self.N) + fill_bubble_time = self.BLEN * (intermediate_size // self.BLEN) * math.ceil(self.batch_size / self.BLEN) + operation_amount = (intermediate_size // self.BLEN) * self.BLEN * (hidden_size // self.MLEN) * math.ceil(self.batch_size / self.BLEN) + attainable_operation += operation_amount * (min(self.batch_size, self.BLEN) / self.BLEN) + theoretical_operation += operation_amount + fill_bubble_time - # Down Projection - operation_amount = (hidden_size // self.M) * (intermediate_size // self.K) - overall_operation_amount += operation_amount - attainable_operation += operation_amount * (self.M * self.K * min(self.batch_size, self.N)) - theoretical_operation += operation_amount * (self.M * self.K * self.N) + # SiLU + theoretical_operation += (intermediate_size // self.VLEN) * 5 + # Down Projection + fill_bubble_time = self.BLEN * (hidden_size // self.BLEN) * math.ceil(self.batch_size / self.BLEN) + operation_amount = (hidden_size // self.BLEN) * self.BLEN * (intermediate_size // self.MLEN) * math.ceil(self.batch_size / self.BLEN) + attainable_operation += operation_amount * (min(self.batch_size, self.BLEN) / self.BLEN) + theoretical_operation += operation_amount+ fill_bubble_time - return [overall_operation_amount, attainable_operation, theoretical_operation] + return [attainable_operation, theoretical_operation] def _report_prefill_utilization(self): - overall_operations = {"embedding": 0, "attention": 0, "ffn": 0, "lm_head": 0} - overall_attainable_FLOPS = {"embedding": 0, "attention": 0, "ffn": 0, "lm_head": 0} - overall_theoretical_FLOPS = {"embedding": 0, "attention": 0, "ffn": 0, "lm_head": 0} - single_op_operation = self._report_embedding_utilization("prefill") - overall_operations["embedding"] += (single_op_operation[0] * 2) / (10 ** 9) - overall_attainable_FLOPS["embedding"] += (single_op_operation[1] * 2) / (10 ** 9) - overall_theoretical_FLOPS["embedding"] += (single_op_operation[2] * 2) / (10 ** 9) - for i in range(self.num_hidden_layers): - single_op_operation = self._report_flash_attn_utilization("prefill") - overall_operations["attention"] += single_op_operation[0] / (10 ** 9) - overall_attainable_FLOPS["attention"] += single_op_operation[1] / (10 ** 9) - overall_theoretical_FLOPS["attention"] += single_op_operation[2] / (10 ** 9) - - single_op_operation = self._report_ffn_utilization("prefill") - overall_operations["ffn"] += single_op_operation[0] / (10 ** 9) - overall_attainable_FLOPS["ffn"] += single_op_operation[1] / (10 ** 9) - overall_theoretical_FLOPS["ffn"] += single_op_operation[2] / (10 ** 9) + overall_attainable_GEMM = {"attention": 0, "ffn": 0} + overall_theoretical_GEMM = {"attention": 0, "ffn": 0} + overall_attainable_GEMM["attention"] = self._report_flash_attn_utilization("prefill")[0] + overall_theoretical_GEMM["attention"] = self._report_flash_attn_utilization("prefill")[1] + overall_attainable_GEMM["ffn"] = self._report_ffn_utilization("prefill")[0] + overall_theoretical_GEMM["ffn"] = self._report_ffn_utilization("prefill")[1] return { - "operations": overall_operations, - "attainable_FLOPS": overall_attainable_FLOPS, - "theoretical_FLOPS": overall_theoretical_FLOPS + "attainable_FLOPS": overall_attainable_GEMM, + "theoretical_FLOPS": overall_theoretical_GEMM } def _report_decode_utilization(self): - overall_operations = {"embedding": 0, "attention": 0, "ffn": 0, "lm_head": 0} - overall_attainable_FLOPS = {"embedding": 0, "attention": 0, "ffn": 0, "lm_head": 0} - overall_theoretical_FLOPS = {"embedding": 0, "attention": 0, "ffn": 0, "lm_head": 0} - for j in range (self.output_token): - per_token_operations = {"embedding": 0, "attention": 0, "ffn": 0, "lm_head": 0} - per_token_attainable_FLOPS = {"embedding": 0, "attention": 0, "ffn": 0, "lm_head": 0} - per_token_theoretical_FLOPS = {"embedding": 0, "attention": 0, "ffn": 0, "lm_head": 0} - - for i in range(self.num_hidden_layers): - single_op_operation = self._report_flash_attn_utilization("decode") - per_token_operations["attention"] += single_op_operation[0] - per_token_attainable_FLOPS["attention"] += single_op_operation[1] - per_token_theoretical_FLOPS["attention"] += single_op_operation[2] - - single_op_operation = self._report_ffn_utilization("decode") - per_token_operations["ffn"] += single_op_operation[0] - per_token_attainable_FLOPS["ffn"] += single_op_operation[1] - per_token_theoretical_FLOPS["ffn"] += single_op_operation[2] - - overall_operations["attention"] += per_token_operations["attention"] / (10**9) - overall_operations["ffn"] += per_token_operations["ffn"] / (10**9) - - overall_attainable_FLOPS["attention"] += per_token_attainable_FLOPS["attention"] / (10**9) - overall_theoretical_FLOPS["attention"] += per_token_theoretical_FLOPS["attention"] / (10**9) - - overall_attainable_FLOPS["ffn"] += per_token_attainable_FLOPS["ffn"] / (10**9) - overall_theoretical_FLOPS["ffn"] += per_token_theoretical_FLOPS["ffn"] / (10**9) + overall_attainable_GEMM = {"attention": 0, "ffn": 0} + overall_theoretical_GEMM = {"attention": 0, "ffn": 0} + overall_attainable_GEMM["attention"] = self._report_flash_attn_utilization("decode")[0] + overall_theoretical_GEMM["attention"] = self._report_flash_attn_utilization("decode")[1] + overall_attainable_GEMM["ffn"] = self._report_ffn_utilization("decode")[0] + overall_theoretical_GEMM["ffn"] = self._report_ffn_utilization("decode")[1] return { - "operations": overall_operations, - "attainable_FLOPS": overall_attainable_FLOPS, - "theoretical_FLOPS": overall_theoretical_FLOPS + "attainable_FLOPS": overall_attainable_GEMM, + "theoretical_FLOPS": overall_theoretical_GEMM } def compute_overall_perf(self): @@ -226,4 +375,17 @@ def compute_overall_perf(self): decode_perf = self._report_decode_utilization() print(f"Decode Performance: {decode_perf}") utilization = (prefill_perf["attainable_FLOPS"]["ffn"] + prefill_perf["attainable_FLOPS"]["attention"] + decode_perf["attainable_FLOPS"]["ffn"] + decode_perf["attainable_FLOPS"]["attention"]) / (prefill_perf["theoretical_FLOPS"]["ffn"] + prefill_perf["theoretical_FLOPS"]["attention"] + decode_perf["theoretical_FLOPS"]["ffn"] + decode_perf["theoretical_FLOPS"]["attention"]) - return utilization \ No newline at end of file + return utilization + + +if __name__ == "__main__": + import os + from pathlib import Path + # Get the absolute path of the current file's directory + current_dir = Path(__file__).resolve().parents[3] + model_param_path = os.path.join(current_dir, "doc/Model_Lib/llama-3.1-8b.json") + model = attn_model_config(MLEN=1024, BLEN=4, VLEN=1024, partitioned_matrix=False, model_param_path=model_param_path, batch_size=1, seq_len=5600, output_token=1000, device_num=1) + model.kv_size = 80000 + actual, theoretical = model._report_flash_attn_utilization("decode") + print(f"Actual: {actual}, Theoretical: {theoretical}") + print(f"Utilization: {actual / theoretical}") \ No newline at end of file diff --git a/tools/cost_model/utilisation/attainable_gemm_test.py b/tools/cost_model/utilisation/attainable_gemm_test.py new file mode 100644 index 00000000..4dec52c2 --- /dev/null +++ b/tools/cost_model/utilisation/attainable_gemm_test.py @@ -0,0 +1,104 @@ +from attainable import attn_model_config + +import os +from pathlib import Path +# Get the absolute path of the current file's directory +current_dir = Path(__file__).resolve().parents[3] +model_param_path = os.path.join(current_dir, "doc/Model_Lib/llama-3.1-8b.json") + +import matplotlib.pyplot as plt +import numpy as np + +# Define MLEN and BLEN configurations and their labels +configs = [ + {"MLEN": 64, "BLEN": 64, "label": "MLEN=64, BLEN=64"}, + {"MLEN": 128, "BLEN": 32, "label": "MLEN=128, BLEN=32"}, + {"MLEN": 256, "BLEN": 16, "label": "MLEN=256, BLEN=16"}, + {"MLEN": 512, "BLEN": 8, "label": "MLEN=512, BLEN=8"}, + {"MLEN": 1024, "BLEN": 4, "label": "MLEN=1024, BLEN=4"} +] + +batch_sizes = [1, 2, 4, 8, 16] + +bar_width = 0.06 # Further reduced bar width +bar_gap = 0.02 # Gap between bars in same group +group_gap = 0 # Gap between groups of bars + +# Create x positions with spacing between groups (reduced distance) +x = np.arange(len(batch_sizes)) * (0.5 + group_gap) + +# Create two subplots side by side: utilization (left) and execution time (right) +fig, (ax1_ffn, ax2_time) = plt.subplots(1, 2, figsize=(16, 4)) + +colors = [ + (240/255, 249/255, 232/255), # light green + (186/255, 228/255, 188/255), # green + (123/255, 204/255, 196/255), # teal + (67/255, 162/255, 202/255), # blue + (8/255, 104/255, 172/255) # dark blue +] + +# Process each configuration +for idx, cfg in enumerate(configs): + # FFN data + ffn_utilizations = [] + execution_times = [] # Execution time in microseconds (cycles / 1GHz) + + for bs in batch_sizes: + # reload model with required settings + model = attn_model_config( + MLEN=cfg["MLEN"], BLEN=cfg["BLEN"], VLEN=1024, partitioned_matrix=True, + model_param_path=model_param_path, batch_size=bs, seq_len=8024, output_token=100, device_num=1 + ) + + # FFN stage + ffn_attainable, ffn_theoretical = model._report_ffn_utilization("decode")[0:2] + ffn_utilization = ffn_attainable / ffn_theoretical if ffn_theoretical != 0 else 0 + ffn_utilizations.append(ffn_utilization) + + # Execution time: theoretical cycles / 1GHz = time in seconds, convert to microseconds + execution_time_us = (ffn_theoretical / 1e9) * 1e6 # Convert to microseconds + execution_times.append(execution_time_us) + + # Bar positions for configs - with spacing between groups + pos = x + (idx - len(configs)/2) * (bar_width + bar_gap) + bar_width/2 + + # Plot FFN utilization (left) + rects_ffn = ax1_ffn.bar(pos, ffn_utilizations, width=bar_width, label=cfg["label"], color=colors[idx], alpha=0.7) + + # Plot execution time (right) + rects_time = ax2_time.bar(pos, execution_times, width=bar_width, label=cfg["label"], color=colors[idx], alpha=0.7) + +# Configure FFN utilization plot (left) +ax1_ffn.set_xlabel("Batch Size", fontsize=18) +ax1_ffn.set_ylabel("Utilization", fontsize=18) +ax1_ffn.tick_params(axis='both', which='major', labelsize=14) +ax1_ffn.set_title("FFN Decode Systolic Array Utilization vs. Batch Size", fontsize=20) +ax1_ffn.set_xticks(x) +ax1_ffn.set_xticklabels([str(bs) for bs in batch_sizes]) +ax1_ffn.set_ylim(0, 1) + +# Set x-axis limits to show all bars with spacing +total_config_width = len(configs) * (bar_width + bar_gap) - bar_gap +ax1_ffn.set_xlim(x[0] - total_config_width/2 - group_gap/2, x[-1] + total_config_width/2 + group_gap/2) + +# Configure execution time plot (right) +ax2_time.set_xlabel("Batch Size", fontsize=18) +ax2_time.set_ylabel("Execution Time (μs)", fontsize=18) +ax2_time.tick_params(axis='both', which='major', labelsize=14) +ax2_time.set_title("FFN Decode Execution Time vs. Batch Size (1GHz)", fontsize=20) +ax2_time.set_xticks(x) +ax2_time.set_xticklabels([str(bs) for bs in batch_sizes]) + +# Set x-axis limits to show all bars with spacing (same as left plot) +ax2_time.set_xlim(x[0] - total_config_width/2 - group_gap/2, x[-1] + total_config_width/2 + group_gap/2) + +# Create shared legend at the bottom in two rows (placed lower) +handles1, labels1 = ax1_ffn.get_legend_handles_labels() +# Create legend with all config labels - transparent with no frame, arranged in 2 rows +# With 5 configs, use ncol=3 to get 2 rows (3 in first row, 2 in second row) +# Push legend lower using bbox_to_anchor +fig.legend(handles1, labels1, loc='upper center', bbox_to_anchor=(0.5, 0.1), ncol=5, fontsize=15, frameon=False) + +plt.tight_layout(rect=[0, 0.08, 1, 0.98]) +plt.savefig("decode_FFN_GEMM_Utilization_and_Latency.png", dpi=300, bbox_inches='tight') diff --git a/tools/cost_model/utilisation/combined_ffn_fa_utilization.py b/tools/cost_model/utilisation/combined_ffn_fa_utilization.py new file mode 100644 index 00000000..59258b22 --- /dev/null +++ b/tools/cost_model/utilisation/combined_ffn_fa_utilization.py @@ -0,0 +1,178 @@ +from attainable import attn_model_config + +import os +from pathlib import Path +# Get the absolute path of the current file's directory +current_dir = Path(__file__).resolve().parents[3] +model_param_path = os.path.join(current_dir, "doc/Model_Lib/llama-3.1-8b.json") + +import matplotlib.pyplot as plt +import matplotlib.ticker as ticker +import numpy as np + +# Define MLEN and BLEN configurations and their labels +configs = [ + {"MLEN": 64, "BLEN": 64, "label": "MLEN=64, BLEN=64"}, + {"MLEN": 128, "BLEN": 32, "label": "MLEN=128, BLEN=32"}, + {"MLEN": 256, "BLEN": 16, "label": "MLEN=256, BLEN=16"}, + {"MLEN": 512, "BLEN": 8, "label": "MLEN=512, BLEN=8"}, + {"MLEN": 1024, "BLEN": 4, "label": "MLEN=1024, BLEN=4"} +] + +batch_sizes = [1, 2, 4, 8, 16] +sequence_lengths = [128, 1024, 5120, 20480, 81920] # 128, 1k, 5k, 20k, 80k +seq_labels = ["128", "1k", "5k", "20k", "80k"] + +bar_width = 0.06 # Bar width +bar_gap = 0.02 # Gap between bars in same group +group_gap = 0 # Gap between groups of bars + +# Create x positions with spacing between groups +x_ffn = np.arange(len(batch_sizes)) * (0.5 + group_gap) +x_fa = np.arange(len(sequence_lengths)) * (0.5 + group_gap) + +# Create 1x4 grid: FFN util, FFN time, FA util, FA time +fig, (ax1_ffn_util, ax2_ffn_time, ax1_fa_util, ax2_fa_time) = plt.subplots(1, 4, figsize=(20, 4)) + +colors = [ + "#762a83", # deep purple + "#af8dc3", # light violet + "#e7d4e8", # very light purple-pink + "#d9f0d3", # pale green + "#7fbf7b", # medium green + # Optional sixth if needed: "#1b7837" # dark green +] + +# Process FFN data +for idx, cfg in enumerate(configs): + # FFN data + ffn_utilizations = [] + ffn_execution_times = [] # Execution time in microseconds (cycles / 1GHz) + + for bs in batch_sizes: + # reload model with required settings + model = attn_model_config( + MLEN=cfg["MLEN"], BLEN=cfg["BLEN"], VLEN=1024, partitioned_matrix=True, + model_param_path=model_param_path, batch_size=bs, seq_len=8024, output_token=100, device_num=1 + ) + + # FFN stage + ffn_attainable, ffn_theoretical = model._report_ffn_utilization("decode")[0:2] + ffn_utilization = (ffn_attainable / ffn_theoretical if ffn_theoretical != 0 else 0) * 100 # Convert to percentage + ffn_utilizations.append(ffn_utilization) + + # Execution time: theoretical cycles / 1GHz = time in seconds, convert to milliseconds + execution_time_ms = (ffn_theoretical / 1e9) * 1e3 # Convert to milliseconds + ffn_execution_times.append(execution_time_ms) + + # Bar positions for configs - with spacing between groups + pos_ffn = x_ffn + (idx - len(configs)/2) * (bar_width + bar_gap) + bar_width/2 + + # Plot FFN utilization + rects_ffn_util = ax1_ffn_util.bar(pos_ffn, ffn_utilizations, width=bar_width, label=cfg["label"], color=colors[idx], alpha=0.7) + + # Plot FFN execution time + rects_ffn_time = ax2_ffn_time.bar(pos_ffn, ffn_execution_times, width=bar_width, label=cfg["label"], color=colors[idx], alpha=0.7) + +# Process Flash Attention data +for idx, cfg in enumerate(configs): + # Flash Attention data + fa_utilizations = [] + fa_execution_times = [] # Execution time in microseconds (cycles / 1GHz) + + for seq_len in sequence_lengths: + # reload model with required settings + model = attn_model_config( + MLEN=cfg["MLEN"], BLEN=cfg["BLEN"], VLEN=1024, partitioned_matrix=True, + model_param_path=model_param_path, batch_size=1, seq_len=seq_len, output_token=100, device_num=1 + ) + + # Flash Attention stage + fa_attainable, fa_theoretical = model._report_flash_attn_utilization("decode")[0:2] + fa_utilization = (fa_attainable / fa_theoretical if fa_theoretical != 0 else 0) * 100 # Convert to percentage + fa_utilizations.append(fa_utilization) + + # Execution time: theoretical cycles / 1GHz = time in seconds, convert to milliseconds + execution_time_ms = (fa_theoretical / 1e9) * 1e3 # Convert to milliseconds + fa_execution_times.append(execution_time_ms) + + # Bar positions for configs - with spacing between groups + pos_fa = x_fa + (idx - len(configs)/2) * (bar_width + bar_gap) + bar_width/2 + + # Plot Flash Attention utilization + rects_fa_util = ax1_fa_util.bar(pos_fa, fa_utilizations, width=bar_width, label=cfg["label"], color=colors[idx], alpha=0.7) + + # Plot Flash Attention execution time + rects_fa_time = ax2_fa_time.bar(pos_fa, fa_execution_times, width=bar_width, label=cfg["label"], color=colors[idx], alpha=0.7) + +# Configure FFN utilization plot +ax1_ffn_util.set_xlabel("Batch Size", fontsize=18) +ax1_ffn_util.set_ylabel("Utilization (%)", fontsize=18) +ax1_ffn_util.tick_params(axis='both', which='major', labelsize=14) +ax1_ffn_util.set_title("FFN Decode SA Utilization", fontsize=20) +ax1_ffn_util.set_xticks(x_ffn) +ax1_ffn_util.set_xticklabels([str(bs) for bs in batch_sizes]) +ax1_ffn_util.set_ylim(0, 100) +# Format y-axis to show only integer ticks +ax1_ffn_util.yaxis.set_major_locator(ticker.MaxNLocator(integer=True)) +ax1_ffn_util.yaxis.set_major_formatter(ticker.FuncFormatter(lambda x, p: f'{int(x)}' if x == int(x) else '')) + +# Set x-axis limits to show all bars with spacing +total_config_width_ffn = len(configs) * (bar_width + bar_gap) - bar_gap +ax1_ffn_util.set_xlim(x_ffn[0] - total_config_width_ffn/2 - group_gap/2, x_ffn[-1] + total_config_width_ffn/2 + group_gap/2) + +# Configure FFN execution time plot +ax2_ffn_time.set_xlabel("Batch Size", fontsize=18) +ax2_ffn_time.set_ylabel("Exec Time (ms)", fontsize=18) +ax2_ffn_time.tick_params(axis='both', which='major', labelsize=14) +ax2_ffn_time.set_title("FFN Decode Exec Time", fontsize=20) +ax2_ffn_time.set_xticks(x_ffn) +ax2_ffn_time.set_xticklabels([str(bs) for bs in batch_sizes]) +# Format y-axis to show only integer ticks +ax2_ffn_time.yaxis.set_major_locator(ticker.MaxNLocator(integer=True)) +ax2_ffn_time.yaxis.set_major_formatter(ticker.FuncFormatter(lambda x, p: f'{int(x)}' if x == int(x) else '')) + +# Set x-axis limits to show all bars with spacing (same as FFN util) +ax2_ffn_time.set_xlim(x_ffn[0] - total_config_width_ffn/2 - group_gap/2, x_ffn[-1] + total_config_width_ffn/2 + group_gap/2) + +# Configure Flash Attention utilization plot +ax1_fa_util.set_xlabel("Number of Tokens", fontsize=18) +ax1_fa_util.set_ylabel("Utilization (%)", fontsize=18) +ax1_fa_util.tick_params(axis='both', which='major', labelsize=14) +ax1_fa_util.set_title("FA Decode SA Utilization", fontsize=20) +ax1_fa_util.set_xticks(x_fa) +ax1_fa_util.set_xticklabels(seq_labels) +ax1_fa_util.set_ylim(0, 100) +# Format y-axis to show only integer ticks +ax1_fa_util.yaxis.set_major_locator(ticker.MaxNLocator(integer=True)) +ax1_fa_util.yaxis.set_major_formatter(ticker.FuncFormatter(lambda x, p: f'{int(x)}' if x == int(x) else '')) + +# Set x-axis limits to show all bars with spacing +total_config_width_fa = len(configs) * (bar_width + bar_gap) - bar_gap +ax1_fa_util.set_xlim(x_fa[0] - total_config_width_fa/2 - group_gap/2, x_fa[-1] + total_config_width_fa/2 + group_gap/2) + +# Configure Flash Attention execution time plot +ax2_fa_time.set_xlabel("Number of Tokens", fontsize=18) +ax2_fa_time.set_ylabel("Exec Time (ms)", fontsize=18) +ax2_fa_time.tick_params(axis='both', which='major', labelsize=14) +ax2_fa_time.set_title("FA Decode Exec Time", fontsize=20) +ax2_fa_time.set_xticks(x_fa) +ax2_fa_time.set_xticklabels(seq_labels) +# Format y-axis to show only integer ticks +ax2_fa_time.yaxis.set_major_locator(ticker.MaxNLocator(integer=True)) +ax2_fa_time.yaxis.set_major_formatter(ticker.FuncFormatter(lambda x, p: f'{int(x)}' if x == int(x) else '')) + +# Set x-axis limits to show all bars with spacing (same as FA util) +ax2_fa_time.set_xlim(x_fa[0] - total_config_width_fa/2 - group_gap/2, x_fa[-1] + total_config_width_fa/2 + group_gap/2) + +# Create shared legend at the bottom in two rows (placed lower) +handles1, labels1 = ax1_ffn_util.get_legend_handles_labels() +# Create legend with all config labels - transparent with no frame, arranged in 2 rows +# With 5 configs, use ncol=3 to get 2 rows (3 in first row, 2 in second row) +# Push legend lower using bbox_to_anchor +fig.legend(handles1, labels1, loc='upper center', bbox_to_anchor=(0.5, 0.1), ncol=5, fontsize=15, frameon=False) + +plt.tight_layout(rect=[0, 0.1, 1, 0.98]) +plt.savefig("combined_FFN_FA_Utilization.png", dpi=300, bbox_inches='tight') +print("Plot saved as 'combined_FFN_FA_Utilization.png'") + diff --git a/tools/cost_model/utilisation/fa_utilizaiton_test.py b/tools/cost_model/utilisation/fa_utilizaiton_test.py new file mode 100644 index 00000000..e6511d42 --- /dev/null +++ b/tools/cost_model/utilisation/fa_utilizaiton_test.py @@ -0,0 +1,85 @@ +from attainable import attn_model_config + +import os +from pathlib import Path +# Get the absolute path of the current file's directory +current_dir = Path(__file__).resolve().parents[3] +model_param_path = os.path.join(current_dir, "doc/Model_Lib/llama-3.1-8b.json") + +import matplotlib.pyplot as plt +import numpy as np + +# Define MLEN and BLEN configurations and their labels +configs = [ + {"MLEN": 64, "BLEN": 64, "label": "MLEN=64, BLEN=64"}, + {"MLEN": 128, "BLEN": 32, "label": "MLEN=128, BLEN=32"}, + {"MLEN": 256, "BLEN": 16, "label": "MLEN=256, BLEN=16"}, + {"MLEN": 512, "BLEN": 8, "label": "MLEN=512, BLEN=8"}, + {"MLEN": 1024, "BLEN": 4, "label": "MLEN=1024, BLEN=4"} +] + +sequence_lengths = [128, 1024, 5120, 20480, 81920] # 128, 1k, 5k, 20k, 80k +seq_labels = ["128", "1k", "5k", "20k", "80k"] + +bar_width = 0.1 # Thinner bars +bar_gap = 0.02 # Gap between bars in same group +group_gap = 0 # Gap between groups of bars + +# Create x positions with spacing between groups +x = np.arange(len(sequence_lengths)) * (0.7 + group_gap) + +# Create single plot for Flash Attention +fig, ax1_fa = plt.subplots(1, 1, figsize=(12, 4)) + +colors = [ + (240/255, 249/255, 232/255), # light green + (186/255, 228/255, 188/255), # green + (123/255, 204/255, 196/255), # teal + (67/255, 162/255, 202/255), # blue + (8/255, 104/255, 172/255) # dark blue +] + +# Process each configuration +for idx, cfg in enumerate(configs): + # Flash Attention data + fa_utilizations = [] + + for seq_len in sequence_lengths: + # reload model with required settings + model = attn_model_config( + MLEN=cfg["MLEN"], BLEN=cfg["BLEN"], VLEN=1024, partitioned_matrix=cfg["partitioned_matrix"], + model_param_path=model_param_path, batch_size=1, seq_len=seq_len, output_token=100, device_num=1 + ) + + # Flash Attention stage + fa_attainable, fa_theoretical = model._report_flash_attn_utilization("decode")[0:2] + fa_utilization = fa_attainable / fa_theoretical if fa_theoretical != 0 else 0 + fa_utilizations.append(fa_utilization) + + # Bar positions for configs - with spacing between groups + pos = x + (idx - len(configs)/2) * (bar_width + bar_gap) + bar_width/2 + + # Plot Flash Attention + rects_fa = ax1_fa.bar(pos, fa_utilizations, width=bar_width, label=cfg["label"], color=colors[idx], alpha=0.7) + +# Configure Flash Attention plot +ax1_fa.set_xlabel("Sequence Length", fontsize=18) +ax1_fa.set_ylabel("Utilization", fontsize=18) +ax1_fa.tick_params(axis='both', which='major', labelsize=14) +ax1_fa.set_title("Flash Attention Decode Systolic Array Utilization vs. Sequence Length", fontsize=20) +ax1_fa.set_xticks(x) +ax1_fa.set_xticklabels(seq_labels) +ax1_fa.set_ylim(0, 1) + +# Set x-axis limits to show all bars with spacing +total_config_width = len(configs) * (bar_width + bar_gap) - bar_gap +ax1_fa.set_xlim(x[0] - total_config_width/2 - group_gap/2, x[-1] + total_config_width/2 + group_gap/2) + +# Create shared legend at the bottom in two rows +handles1, labels1 = ax1_fa.get_legend_handles_labels() +# Create legend with all config labels - transparent with no frame, arranged in 2 rows +# With 5 configs, use ncol=3 to get 2 rows (3 in first row, 2 in second row) +fig.legend(handles1, labels1, loc='lower center', ncol=3, fontsize=15, frameon=False) + +plt.tight_layout(rect=[0, 0.05, 1, 0.98]) +plt.savefig("decode_FA_GEMM_Utilization.png", dpi=300, bbox_inches='tight') diff --git a/tools/cost_model/utilisation/individual_units_lib.json b/tools/cost_model/utilisation/individual_units_lib.json index 6de0e56d..69cdd9b1 100644 --- a/tools/cost_model/utilisation/individual_units_lib.json +++ b/tools/cost_model/utilisation/individual_units_lib.json @@ -2,10 +2,10 @@ "MatrixMachine": { "Coefficients": { "P1": 50, - "P2": 500, + "P2": 100, "P3": 10, "P4": 10, - "P5": 400 + "P5": 200 }, "Relationship": "P1 * MLEN * BLEN * (FP_EXP_WIDTH + FP_MANT_WIDTH + 1) + P2 * MLEN * BLEN * (WT_MX_MANT_WIDTH + WT_MX_EXP_WIDTH + 1 + MX_SCALE_WIDTH) + P3 * (MLEN // BLOCK_DIM) * BLEN * (BLOCK_DIM * ACT_ELEMENT_WIDTH + MX_SCALE_WIDTH) + P4 * MLEN * (FP_EXP_WIDTH + FP_MANT_WIDTH + 1) + P5" @@ -45,7 +45,7 @@ "Coefficients": { "P1": 300, "P2": 300, - "P3": 800, + "P3": 100, "P4": 100, "P5": 100, "P6": 100 diff --git a/tools/cost_model/utilisation/memory_utilization_plot.py b/tools/cost_model/utilisation/memory_utilization_plot.py new file mode 100644 index 00000000..2200e138 --- /dev/null +++ b/tools/cost_model/utilisation/memory_utilization_plot.py @@ -0,0 +1,179 @@ +import json +import re + +import os +from pathlib import Path + +def compute_hbm_storage(batch_size, model_param_path, kv_size, kv_precision=2, act_precision=2, wt_precision=2): + """ + Compute HBM storage requirements for model weights, kv_cache, and other elements. + + Args: + batch_size (int): The batch size. + model_param_path (str or Path): Path to the model parameter JSON file. + kv_size (int): The kv cache sequence length. + kv_precision (int, optional): Byte size of KV cache data type (default: 2 for fp16). + act_precision (int, optional): Byte size of activation data type (unused; default: 2). + wt_precision (int, optional): Byte size of weight data type (default: 2 for fp16). + + Returns: + dict: Dictionary with keys 'weights', 'kv_cache', and 'other' with size in GB. + """ + model_param = json.load(open(model_param_path)) + hidden_size = model_param["hidden_size"] + num_attention_heads = model_param["num_attention_heads"] + num_hidden_layers = model_param["num_hidden_layers"] + intermediate_size = model_param["intermediate_size"] + num_key_value_heads = model_param["num_key_value_heads"] + repeat_layer = model_param["num_hidden_layers"] + vocab_size = model_param["vocab_size"] + head_dim = hidden_size // num_attention_heads + + # Extract model size in Billion parameters from the file name (e.g., ".../llama-3.1-8b.json" -> 8) + model_filename = str(model_param_path).split('/')[-1] + match = re.search(r'(\d+)[bB]', model_filename) + if match: + model_size_b = int(match.group(1)) + else: + model_size_b = None # Optionally raise error + + if model_size_b is not None: + weight_size_gb = model_size_b * 1e9 * wt_precision / (1024 ** 3) + else: + weight_size_gb = 0 # Or raise error + + kv_cache_gb = (2 * head_dim * num_key_value_heads * kv_size * batch_size * kv_precision * repeat_layer) / 1024 / 1024 / 1024 + + other_gb = (hidden_size * 4 + intermediate_size * 4 + vocab_size * 4) / 1024 / 1024 / 1024 + + hbm_storage = { + "weights": weight_size_gb, + "kv_cache": kv_cache_gb, + "other": other_gb + } + return hbm_storage + + + +if __name__ == "__main__": + current_dir = Path(__file__).resolve().parents[3] + model_param_path = os.path.join(current_dir, "doc/Model_Lib/llama-3.1-8b.json") + batch_size = 1 + kv_size = 8024 + kv_precision = 2 + act_precision = 2 + wt_precision = 2 + hbm_storage = compute_hbm_storage(batch_size, model_param_path, kv_size, kv_precision, act_precision, wt_precision) + print(hbm_storage) + + import matplotlib.pyplot as plt + import numpy as np + + # kv_size values + kv_sizes = [1000, 10000, 50000, 100000] # [1k, 10k, 50k, 100k] + kv_size_labels = ['1k', '10k', '50k', '100k'] + + colors = [ + (240/255, 249/255, 232/255), # light green + (186/255, 228/255, 188/255), # green + (123/255, 204/255, 196/255), # teal + (67/255, 162/255, 202/255), # blue + (8/255, 104/255, 172/255) # dark blue + ] + + + # Define the four precision settings (kv, act, wt) + # Format: (kv_precision, act_precision, wt_precision) + precision_settings = [ + (2, 2, 2), + (1.125, 1.125, 1.125), + (1.125, 1.125, 0.625), + (0.625, 1.125, 0.625), + ] + precision_labels = [ + "KV=BF16 ,ACT=BF16, Weight=BF16", + "KV=MXINT E4M3 ,ACT=MXINT E4M3, Weight= MXINT E4M3", + "KV=MXINT E4M3 ,ACT=MXINT E4M3, Weight= MXINT E2M1", + "KV=MXINT E2M1 ,ACT=MXINT E4M3, Weight= MXINT E2M1", + ] + num_settings = len(precision_settings) + + # Colors for each memory type + color_dict = { + "weights": colors[2], + "kv_cache": colors[4], + } + mem_keys = ["weights", "kv_cache"] + + # Collect data: result[setting_idx][kv_idx][mem_type] + results = [] # (settings, kv_size, mem_type) + for setting in precision_settings: + setting_results = [] + for kv in kv_sizes: + kv_precision, act_precision, wt_precision = setting + hbm_storage = compute_hbm_storage(batch_size, model_param_path, kv, kv_precision, act_precision, wt_precision) + vals = [hbm_storage[k] for k in mem_keys] + setting_results.append(vals) + results.append(setting_results) + results = np.array(results) # shape: (num_settings, num_kv, 3) + + # Plotting + fig, ax = plt.subplots(figsize=(12, 6)) + + bar_width = 0.1 # Thinner bars + spacing = 0.02 # Space between each precision group + group_width = num_settings * bar_width + (num_settings - 1) * spacing + x = np.arange(len(kv_sizes)) * 0.6 # Reduce spacing between bar groups + + # For bar placement (center group at each kv_size) + for sidx, setting_label in enumerate(precision_labels): + for kidx in range(len(kv_sizes)): + left = x[kidx] - group_width/2 + sidx*(bar_width + spacing) + bar_width/2 + bottom = 0 + prev_bottom = 0 + # Plot stacked: weights, kv_cache, other + for midx, mem in enumerate(mem_keys): + val = results[sidx, kidx, midx] + color = color_dict[mem] + bar = ax.bar(left, val, bar_width, bottom=prev_bottom, color=color, edgecolor="black" if sidx==0 else None) + prev_bottom += val + # Put legend entries only for the "weights" bar for nice stacking + if sidx == 0: + # Just for the legend + handles = [plt.Rectangle((0,0),1,1, color=color_dict[m]) for m in mem_keys] + ax.legend(handles, mem_keys, title="Memory Type", loc='upper right', fontsize=13, title_fontsize=14) + + # X-axis + ax.set_xticks(x) + xtick_labels = [] + for k in kv_size_labels: + label = f"{k}" + xtick_labels.append(label) + ax.set_xticklabels(xtick_labels, fontsize=12) + + # Put precision labels (P1,P2,...) centered on each bar + for kidx, xpos in enumerate(x): + for sidx in range(num_settings): + left = xpos - group_width/2 + sidx*(bar_width + spacing) + bar_width/2 + height = results[sidx, kidx].sum() + # Center P-label on top part of each stacked bar + ax.text( + left, height + 0.03*max(results.flatten()), + f"P{sidx+1}", + ha='center', + va='bottom', + fontsize=10, + rotation=0 + ) + + # Additional legend for precision index -> actual precision setting (top left) + precision_legend = [f"P{i+1}: {lbl}" for i, lbl in enumerate(precision_labels)] + ax.text(0.02, 0.98, "Precisions:\n" + "\n".join(precision_legend), + transform=ax.transAxes, fontsize=13, va='top', ha='left', bbox=dict(facecolor='white', alpha=0.7, edgecolor='gray')) + + ax.set_xlabel("KV Size", fontsize=15) + ax.set_ylabel("Memory Usage (GB, Stacked)", fontsize=15) + ax.set_title("Single Batch Memory Distribution vs KV Size", fontsize=17) + ax.set_ylim(bottom=0, top=40) + plt.tight_layout() + plt.savefig("memory_utilization.png", dpi=300, bbox_inches='tight') diff --git a/tools/cost_model/utilisation/utilisation_model.py b/tools/cost_model/utilisation/utilisation_model.py index ba31ca6d..5377c2b2 100644 --- a/tools/cost_model/utilisation/utilisation_model.py +++ b/tools/cost_model/utilisation/utilisation_model.py @@ -3,27 +3,52 @@ import os import json from pathlib import Path -from ...utils import load_toml_config, load_json, load_svh_settings +from utils import load_toml_config, load_json, load_svh_settings from attainable import attn_model_config class utilisation_model: def __init__(self, hardware_settings_file: str = "plena_settings.toml", precision_settings_file: str = "precision.toml", unit_info_file: str = "unit_info.json"): self.unit_info = load_json(unit_info_file) + config_settings = load_svh_settings(hardware_settings_file) precision_settings = load_svh_settings(precision_settings_file) self.hardware_settings = {**config_settings, **precision_settings} def obtain_resource_utilisation(self, updated_config): resource_utilisation = 0 hardware_settings = self.hardware_settings - for key, value in updated_config.items(): - hardware_settings[key] = value + #TODOs + # for key, value in updated_config.items(): + # hardware_settings[key] = value + custom_config= { + "MLEN": 1024, + "BLEN": 64, + "VLEN": 1024, + "WT_MX_MANT_WIDTH": 1, + "WT_MX_EXP_WIDTH": 2, + "KV_ELEMENT_WIDTH": 8, + "BLOCK_DIM": 8, + "MX_SCALE_WIDTH" : 8, + # "ACT_MXFP_MANT_WIDTH": 3, + # "ACT_MXFP_EXP_WIDTH": 4, + "ACT_ELEMENT_WIDTH" : 8, + "FP_EXP_WIDTH": 3, + "FP_MANT_WIDTH": 4, + "HBM_ELE_WIDTH": 512, + "HBM_SCALE_WIDTH": 512, + "MATRIX_SRAM_DEPTH": 4096, + "VECTOR_SRAM_DEPTH": 4096, + "INT_DATA_WIDTH": 32, + "INT_SRAM_DEPTH": 256, + "FP_SRAM_DEPTH": 256, + + } for unit, info in self.unit_info.items(): - if "Coefficients" in info and "Relationship" in info: - hardware_settings.update(info["Coefficients"]) + if "Coefficients" in info and "Relationship" in info: + custom_config.update(info["Coefficients"]) relationship = info["Relationship"] - resource_utilisation += eval(relationship, {}, hardware_settings) + resource_utilisation += eval(relationship, {}, custom_config) return resource_utilisation From 7f82284012992a5b4eac16253829cb6c9a85e804 Mon Sep 17 00:00:00 2001 From: GeorgeWu1204 Date: Wed, 19 Nov 2025 03:30:59 +0000 Subject: [PATCH 12/34] fix simulation bugs for linear and flashattn --- behavioral_simulator/src/main.rs | 4 +++- compiler/asm_templates/flash_attn_asm.py | 3 +-- src/definitions/plena_settings.toml | 8 ++++---- 3 files changed, 8 insertions(+), 7 deletions(-) diff --git a/behavioral_simulator/src/main.rs b/behavioral_simulator/src/main.rs index cf4e647f..5558de62 100644 --- a/behavioral_simulator/src/main.rs +++ b/behavioral_simulator/src/main.rs @@ -255,6 +255,8 @@ impl VectorSram { write_amount: u32, tensor: Receiver, ) { + // println!("continous_write_delayed: addr = {:?}, write_amount = {:?}", addr, write_amount); + // println!("tile_size = {:?}", self.tile_size); let addr_in_tiles = addr.assert_multiple_of(self.tile_size); // Await the tensor from the channel (blocks until data arrives) if let Ok(tensor) = tensor.await { @@ -357,7 +359,7 @@ impl MatrixMachine { let (mat_base, mat_offset) = m_addr.multiple_and_offset(self.mlen * self.mlen); // println!("mat_offset = {:?}", mat_offset); // println!("mat_base = {:?}", mat_base); - assert!(mat_offset.is_multiple_of(self.mlen)); + assert!(mat_offset.is_multiple_of(self.blen)); let mat_row_offset = mat_offset as i64 / self.mlen as i64; let full_mat = self.mram.read(mat_base).await; diff --git a/compiler/asm_templates/flash_attn_asm.py b/compiler/asm_templates/flash_attn_asm.py index 17dc6ae9..134ea75e 100644 --- a/compiler/asm_templates/flash_attn_asm.py +++ b/compiler/asm_templates/flash_attn_asm.py @@ -199,7 +199,7 @@ def _computing_pv_code( # Update v_base_register and reset p_base_register and pv_base_register generated_code += f"S_ADDI_INT gp{p_base_register}, gp0, {p_base_address + q_head_index * mlen * mlen} \n" generated_code += f"S_ADDI_INT gp{pv_base_register}, gp0, {pv_base_address + q_head_index * head_dim + (i+1) * blen} \n" - generated_code += f"S_ADDI_INT gp{v_base_register}, gp{v_base_register}, {blen} \n" + generated_code += f"S_ADDI_INT gp{v_base_register}, gp{v_base_register}, {mlen} \n" return generated_code @@ -296,7 +296,6 @@ def _computing_row_wise_scaling_code( generated_code += f"S_RECI_FP f{l_old_fp_register}, f{l_old_fp_register}, 0 \n" # multiply o_old with the inverse of l_old generated_code += f"V_MUL_VF gp{o_old_vector_address_register}, gp{o_old_vector_address_register}, f{l_old_fp_register} \n" - # update o_old base address generated_code += f"S_ADDI_INT gp{o_old_vector_address_register}, gp{o_old_vector_address_register}, {mlen} \n" diff --git a/src/definitions/plena_settings.toml b/src/definitions/plena_settings.toml index ea863825..3f27e552 100644 --- a/src/definitions/plena_settings.toml +++ b/src/definitions/plena_settings.toml @@ -2,16 +2,16 @@ value = 1 [CONFIG.BLEN] -value = 64 +value = 4 [CONFIG.HLEN] -value = 128 +value = 16 [CONFIG.MLEN] -value = 1024 +value = 64 [CONFIG.VLEN] -value = 1024 +value = 64 [CONFIG.BROADCAST_AMOUNT] value = 4 From 7e52514ef2673a90b3507457def20726bef47133 Mon Sep 17 00:00:00 2001 From: GeorgeWu1204 Date: Wed, 19 Nov 2025 17:27:00 +0000 Subject: [PATCH 13/34] S_MAP_V_TEST --- behavioral_simulator/src/main.rs | 37 ++++++---- behavioral_simulator/src/op.rs | 2 +- .../testbench/s_map_v_test.py | 68 +++++++++++++++++++ src/definitions/operation.svh | 2 +- 4 files changed, 93 insertions(+), 16 deletions(-) create mode 100644 behavioral_simulator/testbench/s_map_v_test.py diff --git a/behavioral_simulator/src/main.rs b/behavioral_simulator/src/main.rs index 5558de62..dff73e72 100644 --- a/behavioral_simulator/src/main.rs +++ b/behavioral_simulator/src/main.rs @@ -741,18 +741,17 @@ impl VectorMachine { self.vram.write(vd, c).await; } - // async fn broadcast(&mut self, vd: u32, f: f32) { - // let c = QuantTensor::quantize( - // Tensor::full( - // [self.vram.tile_size as i64], - // f as f64, - // (tch::Kind::Float, tch::Device::Cpu), - // ), - // self.vram.ty, - // ); - // cycle!(*VECTOR_BASIC_CYCLES); - // self.vram.write(vd, c).await; - // } + async fn vector_transfer_fp(&mut self, vd: u32, f: &[f16]) { + assert_eq!(f.len(), self.vram.tile_size as usize, "Input vector length must match tile_size"); + // Convert f16 slice to f32 vector + let f32_vec: Vec = f.iter().map(|x| f32::from(*x)).collect(); + // Create tensor from f32 vector + let tensor = tch::Tensor::from_slice(&f32_vec); + // Quantize the tensor according to vram data type + let c = QuantTensor::quantize(tensor, self.vram.ty); + cycle!(*VLEN); + self.vram.write(vd, c).await; + } async fn reduce_sum(&mut self, vs1: u32, f: f32, rmask: u8, mask: u32) -> f32 { let a = self.vram.read(vs1).await; @@ -1287,8 +1286,18 @@ impl Accelerator { self.reg_file.fp_reg[rd as usize]; cycle!(1); } - op::Opcode::S_MAP_V_FP { rd, rs1, imm } => todo!(), - + op::Opcode::S_MAP_V_FP { rd, rs1, imm } => { + let start_idx = (self.reg_file.gp_reg[rs1 as usize] + imm) as usize; + let end_idx = start_idx + *VLEN as usize; + let f = &self.fpsram[start_idx..end_idx]; + self.v_machine + .vector_transfer_fp( + self.reg_file.gp_reg[rd as usize], + f, + ) + .await; + cycle!(*VLEN); + } op::Opcode::S_ADD_INT { rd, rs1, rs2 } => { self.reg_file.gp_reg[rd as usize] = self.reg_file.gp_reg[rs1 as usize] .wrapping_add(self.reg_file.gp_reg[rs2 as usize]); diff --git a/behavioral_simulator/src/op.rs b/behavioral_simulator/src/op.rs index 81446b54..89b31f4d 100644 --- a/behavioral_simulator/src/op.rs +++ b/behavioral_simulator/src/op.rs @@ -356,4 +356,4 @@ impl Opcode { } } } -} +} \ No newline at end of file diff --git a/behavioral_simulator/testbench/s_map_v_test.py b/behavioral_simulator/testbench/s_map_v_test.py new file mode 100644 index 00000000..f194e331 --- /dev/null +++ b/behavioral_simulator/testbench/s_map_v_test.py @@ -0,0 +1,68 @@ +from re import I +import sys +from pathlib import Path +sys.path.insert(0, str(Path(__file__).parent.parent.parent)) + +import torch +from torch import Tensor, nn +# from acc_simulator.quantize.quantized_layers.linear import MXFPLinearPTQ +from test_data_gen import get_weights_path, generate_and_save_random_weights +from compiler.asm_templates import rms_norm_asm, projection_asm, preload_act_asm, reset_reg_asm, preload_addr_reg_asm +from create_sim_env import create_sim_env +from sim_env_utils import build_fake_sim_env +from tools.memory_mapping.hbm_addr_map import align_addr_to_hbm_bandwidth +import torch.nn.functional as F + +if __name__ == "__main__": + + + vlen = 64 + load_len = vlen * vlen + batch_size = 1 + real_data_ratio = (8*8 + 8) / (8 * 8) + fp_preload = [0.0, 1e-6] + preload_amount = 4 + hbm_data_width = 64 + + torch.manual_seed(42) + fp_preload = [0, 1] + # fp_sram = [fp_preload[0]] * vlen + + gen_assembly_code = "; S MAP V Test Generation \n" + torch.manual_seed(42) + + input_tensor1 = torch.randn(batch_size, load_len) + input_tensor = { + "input_tensor1": input_tensor1, + "input_tensor2": input_tensor1 + } + weights = input_tensor1 + original_output = torch.max(input_tensor1.reshape(batch_size, vlen, vlen), dim=-1).values + + golden_result = { + "input_tensor": input_tensor, + "weights": weights, + "original_output": original_output + } + # Gen Activation Preload for fp0 + gen_assembly_code += preload_act_asm( + vlen=vlen, + preload_len=preload_amount, + batch=batch_size, + hidden_size=load_len, + alive_registers=[1,2,3], # [a_actual_register, set_stride_register, result_register] + act_vram_offset=0, + activation_offset_reg=0 + ) + gen_assembly_code += f"S_ADDI_INT gp1, gp0, 0 \n" + for i in range(vlen): + gen_assembly_code += f"V_RED_MAX f1, gp1, 0 \n" + gen_assembly_code += f"S_ADDI_INT gp1, gp1, {vlen} \n" + gen_assembly_code += f"S_ST_FP f1, gp0, {i} \n" + + # gen_assembly_code += "; V_RED_MAX f1, gp0, 0 \n" + gen_assembly_code += f"S_ADDI_INT gp1, gp0, 0 \n" + gen_assembly_code += f"S_MAP_V_FP gp1, gp0, 0 \n" + + create_sim_env(input_tensor, weights, gen_assembly_code, golden_result, fp_preload) + build_fake_sim_env(data_size=256, mode="behave_sim", asm="dllm", data=None, specified_data_order = ["input_tensor1","input_tensor2"]) \ No newline at end of file diff --git a/src/definitions/operation.svh b/src/definitions/operation.svh index 19dfb6bd..1a974fce 100644 --- a/src/definitions/operation.svh +++ b/src/definitions/operation.svh @@ -218,4 +218,4 @@ typedef struct { logic update_v_waddr; } OP_BUNDLE; -`endif +`endif \ No newline at end of file From e64a5f1989a815d302ab6e9a37dddbb74265c91d Mon Sep 17 00:00:00 2001 From: GeorgeWu1204 Date: Thu, 20 Nov 2025 11:54:10 +0000 Subject: [PATCH 14/34] mask for rest vec operation --- behavioral_simulator/src/main.rs | 92 +++++++++++++++++++++++++++----- 1 file changed, 80 insertions(+), 12 deletions(-) diff --git a/behavioral_simulator/src/main.rs b/behavioral_simulator/src/main.rs index dff73e72..a53f5bcc 100644 --- a/behavioral_simulator/src/main.rs +++ b/behavioral_simulator/src/main.rs @@ -715,30 +715,98 @@ impl VectorMachine { async fn sub(&mut self, vd: u32, vs1: u32, vs2: u32, rmask: u8, mask: u32) { let (a, b) = tokio::join!(self.vram.read(vs1), self.vram.read(vs2)); - let c = QuantTensor::quantize(a.as_tensor() - b.as_tensor(), a.data_type()); - cycle!(*VECTOR_ADD_CYCLES); - self.vram.write(vd, c).await; + if rmask == 0 { + let c = QuantTensor::quantize(a.as_tensor() - b.as_tensor(), a.data_type()); + cycle!(*VECTOR_ADD_CYCLES); + self.vram.write(vd, c).await; + } else { + let mut result = a.as_tensor().shallow_clone(); + let total_heads = self.tile_size / self.mask_unit; + for head in 0..total_heads { + if (mask & (1 << head)) != 0 { + let start = (head * self.mask_unit) as i64; + let end = ((head + 1) * self.mask_unit) as i64; + let sliced = result.narrow(0, start, end - start); + let updated = &sliced - b.as_tensor().narrow(0, start, end - start); + result.narrow(0, start, end - start).copy_(&updated); + } + } + let c = QuantTensor::quantize(result, a.data_type()); + cycle!(*VECTOR_ADD_CYCLES); + self.vram.write(vd, c).await; + } } async fn mul(&mut self, vd: u32, vs1: u32, vs2: u32, rmask: u8, mask: u32) { let (a, b) = tokio::join!(self.vram.read(vs1), self.vram.read(vs2)); - let c = QuantTensor::quantize(a.as_tensor() * b.as_tensor(), a.data_type()); - cycle!(*VECTOR_MUL_CYCLES); - self.vram.write(vd, c).await; + if rmask == 0 { + let c = QuantTensor::quantize(a.as_tensor() * b.as_tensor(), a.data_type()); + cycle!(*VECTOR_MUL_CYCLES); + self.vram.write(vd, c).await; + } else { + let mut result = a.as_tensor().shallow_clone(); + let total_heads = self.tile_size / self.mask_unit; + for head in 0..total_heads { + if (mask & (1 << head)) != 0 { + let start = (head * self.mask_unit) as i64; + let end = ((head + 1) * self.mask_unit) as i64; + let sliced = result.narrow(0, start, end - start); + let updated = &sliced * b.as_tensor().narrow(0, start, end - start); + result.narrow(0, start, end - start).copy_(&updated); + } + } + let c = QuantTensor::quantize(result, a.data_type()); + cycle!(*VECTOR_MUL_CYCLES); + self.vram.write(vd, c).await; + } } async fn exp(&mut self, vd: u32, vs1: u32, rmask: u8, mask: u32) { let a = self.vram.read(vs1).await; - let c = QuantTensor::quantize(a.as_tensor().exp(), a.data_type()); - cycle!(*VECTOR_EXP_CYCLES); - self.vram.write(vd, c).await; + if rmask == 0 { + let c = QuantTensor::quantize(a.as_tensor().exp(), a.data_type()); + cycle!(*VECTOR_EXP_CYCLES); + self.vram.write(vd, c).await; + } else { + let mut result = a.as_tensor().shallow_clone(); + let total_heads = self.tile_size / self.mask_unit; + for head in 0..total_heads { + if (mask & (1 << head)) != 0 { + let start = (head * self.mask_unit) as i64; + let end = ((head + 1) * self.mask_unit) as i64; + let sliced = result.narrow(0, start, end - start); + let updated = &sliced.exp(); + result.narrow(0, start, end - start).copy_(&updated); + } + } + let c = QuantTensor::quantize(result, a.data_type()); + cycle!(*VECTOR_EXP_CYCLES); + self.vram.write(vd, c).await; + } } async fn reciprocal(&mut self, vd: u32, vs1: u32, rmask: u8, mask: u32) { let a = self.vram.read(vs1).await; - let c = QuantTensor::quantize(a.as_tensor().reciprocal(), a.data_type()); - cycle!(*VECTOR_RECI_CYCLES); - self.vram.write(vd, c).await; + if rmask == 0 { + let c = QuantTensor::quantize(a.as_tensor().reciprocal(), a.data_type()); + cycle!(*VECTOR_RECI_CYCLES); + self.vram.write(vd, c).await; + } else { + let mut result = a.as_tensor().shallow_clone(); + let total_heads = self.tile_size / self.mask_unit; + for head in 0..total_heads { + if (mask & (1 << head)) != 0 { + let start = (head * self.mask_unit) as i64; + let end = ((head + 1) * self.mask_unit) as i64; + let sliced = result.narrow(0, start, end - start); + let updated = &sliced.reciprocal(); + result.narrow(0, start, end - start).copy_(&updated); + } + } + let c = QuantTensor::quantize(result, a.data_type()); + cycle!(*VECTOR_RECI_CYCLES); + self.vram.write(vd, c).await; + } } async fn vector_transfer_fp(&mut self, vd: u32, f: &[f16]) { From d46df13c9dc155a94c67550410a64f42829dcfc3 Mon Sep 17 00:00:00 2001 From: GeorgeWu1204 Date: Sat, 22 Nov 2025 09:40:23 +0000 Subject: [PATCH 15/34] fix preload multi preload_len load bugs --- behavioral_simulator/testbench/dllm1_test.py | 83 ++++++++++++++++++++ compiler/asm_templates/preload_act.py | 2 +- 2 files changed, 84 insertions(+), 1 deletion(-) create mode 100644 behavioral_simulator/testbench/dllm1_test.py diff --git a/behavioral_simulator/testbench/dllm1_test.py b/behavioral_simulator/testbench/dllm1_test.py new file mode 100644 index 00000000..0b7ce6e9 --- /dev/null +++ b/behavioral_simulator/testbench/dllm1_test.py @@ -0,0 +1,83 @@ +import sys +from pathlib import Path +sys.path.insert(0, str(Path(__file__).parent.parent.parent)) + +import torch +from torch import Tensor, nn +# from acc_simulator.quantize.quantized_layers.linear import MXFPLinearPTQ +from test_data_gen import get_weights_path, generate_and_save_random_weights +from compiler.asm_templates import preload_act_asm, reset_reg_asm, preload_addr_reg_asm +from create_sim_env import create_sim_env +from sim_env_utils import build_fake_sim_env +import torch.nn.functional as F + +from tools.memory_mapping.hbm_addr_map import align_addr_to_hbm_bandwidth +from transformers import AutoTokenizer + + +if __name__ == "__main__": + + + # Testing the operation (hidden_size, hidden_size) @ (hidden_size, batch_size) + vocal_size = 2 + hidden_size = 64 + vlen = 64 + batch_size = 4 + preload_amount = 4 + real_data_ratio = (8*8 + 8) / (8 * 8) + hbm_data_width = 64 + fp_preload = [0.0, 1e-6] + + torch.manual_seed(42) + logits = torch.randn(batch_size, hidden_size * vocal_size) + weights = logits + original_output = logits + + print('logits.shape= ', logits.shape) + print('original_output.shape= ', original_output.shape) + + + input_tensor = { + "logits": logits, + } + + golden_result = { + "input_tensor": input_tensor, + "weights": weights, + "original_output": original_output + } + print('original_output.shape = ',original_output.shape) + print('original_output = ',original_output) + + gen_assembly_code = "; DLLM Test Generation \n" + + # Set the addr offset for mask + + gen_assembly_code += preload_addr_reg_asm( + addr_reg_to_set=[1,2], + available_registers=[1,2], + addr_reg_val=[int(align_addr_to_hbm_bandwidth(batch_size * hidden_size * vocal_size * real_data_ratio, hbm_data_width)),int(2*align_addr_to_hbm_bandwidth(batch_size * hidden_size * vocal_size * real_data_ratio, hbm_data_width))] + ) + + + # Reset the registers + gen_assembly_code += reset_reg_asm( + alive_registers=[1,2,3,4,5,6] + ) + + # Gen logtis Preload (B,L,V) + gen_assembly_code += preload_act_asm( + vlen=vlen, + preload_len=preload_amount, + batch=batch_size, + hidden_size=hidden_size*vocal_size, + alive_registers=[1,2,3], # [a_actual_register, set_stride_register, result_register] + act_vram_offset=0, + activation_offset_reg=0 + ) + + + + + create_sim_env(input_tensor, weights, gen_assembly_code, golden_result, fp_preload) + build_fake_sim_env(data_size=256, mode="behave_sim", asm="dllm", data=None, specified_data_order = ["logits"]) \ No newline at end of file diff --git a/compiler/asm_templates/preload_act.py b/compiler/asm_templates/preload_act.py index fa2c1bfc..cd48957d 100644 --- a/compiler/asm_templates/preload_act.py +++ b/compiler/asm_templates/preload_act.py @@ -43,7 +43,7 @@ def preload_act_asm( generated_code += f"C_SET_STRIDE_REG gp{set_stride_register} \n" for i in range((batch * hidden_size) // (vlen * preload_len)): generated_code += f"H_PREFETCH_V gp{result_register}, gp{a_actual_register}, a{activation_offset_reg}, 1, 0 \n" - generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{a_actual_register}, {vlen} \n" + generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{a_actual_register}, {vlen * preload_len} \n" generated_code += f"S_ADDI_INT gp{result_register}, gp{result_register}, {vlen * preload_len} \n" return generated_code \ No newline at end of file From ef340a73b773cbadd794276f2643aa15e474a392 Mon Sep 17 00:00:00 2001 From: GeorgeWu1204 Date: Mon, 24 Nov 2025 15:17:21 +0000 Subject: [PATCH 16/34] test for ffn --- behavioral_simulator/src/main.rs | 45 ++++++++- behavioral_simulator/src/op.rs | 43 ++++++++- .../testbench/create_sim_env.py | 13 +-- behavioral_simulator/testbench/ffn_test.py | 94 ++++++++++++++++++- compiler/asm_templates/ffn_asm.py | 81 +++++++++------- tools/assembler/assembly_to_binary.py | 4 +- tools/assembler/parser.py | 4 +- 7 files changed, 225 insertions(+), 59 deletions(-) diff --git a/behavioral_simulator/src/main.rs b/behavioral_simulator/src/main.rs index df0b1db3..660f7796 100644 --- a/behavioral_simulator/src/main.rs +++ b/behavioral_simulator/src/main.rs @@ -655,6 +655,44 @@ impl VectorMachine { } } + async fn sub_scalar(&mut self, vd: u32, vs1: u32, f: f32, rmask: u8, mask: u32, rorder: op::VectorOrder) { + let a = self.vram.read(vs1).await; + if rmask == 0 { + if matches!(rorder, op::VectorOrder::Normal) { + let c = QuantTensor::quantize(a.as_tensor() - (f as f64), a.data_type()); + cycle!(*VECTOR_ADD_CYCLES); + self.vram.write(vd, c).await; + } else { + let c = QuantTensor::quantize((f as f64) - a.as_tensor(), a.data_type()); + cycle!(*VECTOR_ADD_CYCLES); + self.vram.write(vd, c).await; + } + } else { + // mask is a bitmask; each bit controls whether to apply 'f' to corresponding mask_unit-section + let mut result = a.as_tensor().shallow_clone(); + let total_heads = self.tile_size / self.mask_unit; + for head in 0..total_heads { + if (mask & (1 << head)) != 0 { + // Mask is set for this head + let start = (head * self.mask_unit) as i64; + let end = ((head + 1) * self.mask_unit) as i64; + let sliced = result.narrow(0, start, end - start); + let updated = if matches!(rorder, op::VectorOrder::Normal) { + &sliced - (f as f64) + } else { + (f as f64) - &sliced + }; + // Overwrite this section with calculated values + result.narrow(0, start, end - start).copy_(&updated); + } + // else leave unchanged + } + let c = QuantTensor::quantize(result, a.data_type()); + cycle!(*VECTOR_ADD_CYCLES); + self.vram.write(vd, c).await; + } + } + async fn mul_scalar(&mut self, vd: u32, vs1: u32, f: f32, rmask: u8, mask: u32) { let a = self.vram.read(vs1).await; if rmask == 0 { @@ -1209,15 +1247,16 @@ impl Accelerator { ) .await; } - op::Opcode::V_SUB_VF { rd, rs1, rs2, rmask } => { + op::Opcode::V_SUB_VF { rd, rs1, rs2, rmask, rorder} => { let mask = if rmask == 0 { (1 << *HLEN as u32) - 1 } else {self.reg_file.v_mask }; self.v_machine - .add_scalar( + .sub_scalar( self.reg_file.gp_reg[rd as usize], self.reg_file.gp_reg[rs1 as usize], - (-self.reg_file.fp_reg[rs2 as usize]).into(), + self.reg_file.fp_reg[rs2 as usize].into(), rmask, mask, + rorder, ) .await; } diff --git a/behavioral_simulator/src/op.rs b/behavioral_simulator/src/op.rs index 89b31f4d..a02e693e 100644 --- a/behavioral_simulator/src/op.rs +++ b/behavioral_simulator/src/op.rs @@ -10,6 +10,13 @@ pub enum VectorPrecision { KeyValue, } +#[derive(Debug, Clone, Copy)] +pub enum VectorOrder { + Normal, + Reverse, +} + + #[allow(non_camel_case_types)] #[derive(Debug)] pub enum Opcode { @@ -90,6 +97,7 @@ pub enum Opcode { rs1: u8, rs2: u8, rmask: u8, + rorder: VectorOrder, }, V_MUL_VV { rd: u8, @@ -260,6 +268,32 @@ const fn mask(width: u32) -> u32 { } impl Opcode { + #[inline] + fn matrix_precision_from(funct1: u8) -> MatrixPrecision { + if funct1 == 0 { + MatrixPrecision::Weights + } else { + MatrixPrecision::KeyValue + } + } + + #[inline] + fn vector_precision_from(funct1: u8) -> VectorPrecision { + if funct1 == 0 { + VectorPrecision::Activation + } else { + VectorPrecision::KeyValue + } + } + + #[inline] + fn vector_order_from(funct1: u8) -> VectorOrder { + if funct1 == 0 { + VectorOrder::Normal + } else { + VectorOrder::Reverse + } + } pub fn decode(instr: u32) -> Self { // eprintln!( // "decode(): instr = 0x{instr:08X} ({instr:032b})" @@ -269,6 +303,7 @@ impl Opcode { let rs1 = ((instr >> (OPCODE_WIDTH + OPERAND_WIDTH)) & mask(OPERAND_WIDTH)) as u8; let rs2 = ((instr >> (OPCODE_WIDTH + OPERAND_WIDTH * 2)) & mask(OPERAND_WIDTH)) as u8; let rs3 = ((instr >> (OPCODE_WIDTH + OPERAND_WIDTH * 3)) & mask(OPERAND_WIDTH)) as u8; + let funct1 = ((instr >> (OPCODE_WIDTH + OPERAND_WIDTH * 4)) & mask(OPERAND_WIDTH)) as u8; let imm = ((instr >> (OPCODE_WIDTH + OPERAND_WIDTH)) & mask(IMM_WIDTH)) as u32; let imm2 = ((instr >> (OPCODE_WIDTH + OPERAND_WIDTH * 2)) & mask(IMM_2_WIDTH)) as u32; @@ -292,7 +327,7 @@ impl Opcode { 0x0D => Self::V_ADD_VV { rd, rs1, rs2, rmask: rs3 }, 0x0E => Self::V_ADD_VF { rd, rs1, rs2, rmask: rs3 }, 0x0F => Self::V_SUB_VV { rd, rs1, rs2, rmask: rs3 }, - 0x10 => Self::V_SUB_VF { rd, rs1, rs2, rmask: rs3 }, + 0x10 => Self::V_SUB_VF { rd, rs1, rs2, rmask: rs3, rorder: Self::vector_order_from(funct1) }, 0x11 => Self::V_MUL_VV { rd, rs1, rs2, rmask: rs3 }, 0x12 => Self::V_MUL_VF { rd, rs1, rs2, rmask: rs3 }, 0x13 => Self::V_EXP_V { rd, rs1, rmask: rs3 }, @@ -326,7 +361,7 @@ impl Opcode { rs1, rs2, rstride: rs3, - precision: MatrixPrecision::Weights, + precision: Self::matrix_precision_from(funct1), }, // 0x29 => Self::H_PREFETCH_M { rd, rs1, rs2, rstride: rs3, precision: MatrixPrecision::KeyValue }, 0x29 => Self::H_PREFETCH_V { @@ -334,7 +369,7 @@ impl Opcode { rs1, rs2, rstride: rs3, - precision: VectorPrecision::KeyValue, + precision: Self::vector_precision_from(funct1), }, // 0x2A => Self::H_PREFETCH_V { rd, rs1, rs2, rstride: rs3, precision: VectorPrecision::KeyValue }, 0x2A => Self::H_STORE_V { @@ -342,7 +377,7 @@ impl Opcode { rs1, rs2, rstride: rs3, - precision: VectorPrecision::Activation, + precision: Self::vector_precision_from(funct1), }, // 0x2B => Self::H_STORE_V { rd, rs1, rs2, rstride: rs3, precision: VectorPrecision::KeyValue }, 0x2B => Self::C_SET_ADDR_REG { rd, rs1, rs2 }, diff --git a/behavioral_simulator/testbench/create_sim_env.py b/behavioral_simulator/testbench/create_sim_env.py index c13a157d..f87da7d7 100644 --- a/behavioral_simulator/testbench/create_sim_env.py +++ b/behavioral_simulator/testbench/create_sim_env.py @@ -14,7 +14,7 @@ def np_array_to_str_2f(arr): # For higher dimensions, default to numpy's print (rare for this context) return np.array2string(arr, formatter={'float_kind':lambda x: "%.2f" % x}) -def create_sim_env(input_tensor, input_weight, generated_code, golden_result, fp_preload = None, int_preload = None): +def create_sim_env(input_tensor, generated_code, golden_result, fp_preload = None, int_preload = None): build_dir = os.path.join(os.path.dirname(__file__), "build") os.makedirs(build_dir, exist_ok=True) if isinstance(input_tensor, dict): @@ -24,8 +24,6 @@ def create_sim_env(input_tensor, input_weight, generated_code, golden_result, fp else: with open(os.path.join(build_dir, "input_tensor.pt"), "wb") as f: torch.save(input_tensor, f) - with open(os.path.join(build_dir, "model_weights.pt"), "wb") as f: - torch.save(input_weight, f) with open(os.path.join(build_dir, "generated_asm_code.asm"), "w") as f: f.write(generated_code) # Store golden_result in a readable format, including tensor contents. @@ -46,15 +44,6 @@ def create_sim_env(input_tensor, input_weight, generated_code, golden_result, fp else: value_np = input_tensor.detach().cpu().float().numpy() f.write(np_array_to_str_2f(value_np)) - f.write("\n\nWeights (state_dict):\n") - if isinstance(golden_result["weights"], dict): - for key, value in golden_result["weights"].items(): - # Convert BFloat16 to float32 before converting to numpy - value_np = value.detach().cpu().float().numpy() - f.write(f"{key}:\n{np_array_to_str_2f(value_np)}\n") - else: - value_np = golden_result["weights"].detach().cpu().float().numpy() - f.write(np_array_to_str_2f(value_np)) f.write("\n\nOriginal Output:\n") # Convert BFloat16 to float32 before converting to numpy output_np = golden_result["original_output"].detach().cpu().float().numpy() diff --git a/behavioral_simulator/testbench/ffn_test.py b/behavioral_simulator/testbench/ffn_test.py index 9ece99cc..e20af183 100644 --- a/behavioral_simulator/testbench/ffn_test.py +++ b/behavioral_simulator/testbench/ffn_test.py @@ -6,7 +6,7 @@ from torch import Tensor, nn # from acc_simulator.quantize.quantized_layers.linear import MXFPLinearPTQ from test_data_gen import get_weights_path, generate_and_save_random_weights -from compiler.asm_templates import batched_matmul_asm, preload_addr_reg_asm, reset_reg_asm +from compiler.asm_templates import ffn_asm, preload_addr_reg_asm, reset_reg_asm from create_sim_env import create_sim_env from sim_env_utils import build_fake_sim_env @@ -26,11 +26,97 @@ def __init__(self, dim: int, inter_dim: int, activation: str = "silu"): super().__init__() # Llama uses SwiGLU: x * silu(x) self.w1 = nn.Linear(dim, inter_dim, bias=False) - self.w3 = nn.Linear(dim, inter_dim, bias=False) # for SwiGLU, "gate" proj - self.w2 = nn.Linear(inter_dim, dim, bias=False) + print("w1:", self.w1.weight.shape) + self.w2 = nn.Linear(dim, inter_dim, bias=False) + print("w2:", self.w2.weight.shape) + self.w3 = nn.Linear(inter_dim, dim, bias=False) + print("w3:", self.w3.weight.shape) self.act = torch.nn.SiLU() if activation == "silu" else getattr(torch.nn, activation)() def forward(self, x: Tensor) -> Tensor: # SwiGLU: (x @ w1) * silu(x @ w3) - return self.w2(self.act(self.w1(x)) * self.w3(x)) + print("input_tensor:\n", x) + print("up projection:\n", self.w1(x)) + print("gate projection:\n", self.w2(x)) + print("silu activation:\n", self.act(self.w1(x))) + print("product of silu activation and gate projection:\n", self.act(self.w1(x)) * self.w2(x)) + print("final output:\n", self.w3(self.act(self.w1(x)) * self.w2(x))) + return self.w3(self.act(self.w1(x)) * self.w2(x)) +if __name__ == "__main__": + # Testing the operation (hidden_size, hidden_size) @ (hidden_size, batch_size) + hidden_size = 128 + inter_dim = 256 + batch_size = 4 + real_data_ratio = (8*8 + 8) / (8 * 8) + fp_preload = [0.0, 1] + mlen = 64 + blen = 4 + vlen = 64 + seq_len = 128 + + torch.manual_seed(42) + act_tensor = torch.randn(batch_size, hidden_size) + + original_layer = LlamaFeedForward(dim=hidden_size, inter_dim=inter_dim) + weight_up_layer = torch.randn(inter_dim, hidden_size) + weight_gate_layer = torch.randn(inter_dim, hidden_size) + weight_down_layer = torch.randn(hidden_size, inter_dim) + + + # Set weights for w1, w2, w3 to the generated tensors + with torch.no_grad(): + original_layer.w1.weight.copy_(weight_up_layer) + original_layer.w2.weight.copy_(weight_gate_layer) + original_layer.w3.weight.copy_(weight_down_layer) + + original_output = original_layer(act_tensor) + print("original_output:", original_output) + + input_tensor = { + "act_tensor": act_tensor, + "weight_up_layer": weight_up_layer, + "weight_gate_layer": weight_gate_layer, + "weight_down_layer": weight_down_layer, + } + + golden_result = { + "input_tensor": input_tensor, + "original_output": original_output + } + + gen_assembly_code = "; FFN Test Generation \n" + + # Set the addr offset for weight and bias + gen_assembly_code += preload_addr_reg_asm( + addr_reg_to_set=[1, 2, 3], + available_registers=[1, 2, 3], + addr_reg_val=[int(hidden_size * batch_size * real_data_ratio), + int(hidden_size * batch_size * real_data_ratio) + int(hidden_size * inter_dim * real_data_ratio), + int(hidden_size * batch_size * real_data_ratio) + int(hidden_size * inter_dim * real_data_ratio) + int(inter_dim * hidden_size * real_data_ratio)] + ) + + # Reset the registers + gen_assembly_code += reset_reg_asm( + alive_registers=[1,2,3] + ) + + # FFN Generation + gen_assembly_code += ffn_asm( + mlen=mlen, + vlen=vlen, + blen=blen, + batch=batch_size, + seq_len=seq_len, + hidden_size=hidden_size, + intermediate_size=inter_dim, + alive_registers=[1,2,3,4,5,6], + gate_weight_hbm_offset_reg=1, + up_weight_hbm_offset_reg=2, + down_weight_hbm_offset_reg=3, + const_one_fp_address=1, + activation_base_address=0 + ) + + create_sim_env(input_tensor, gen_assembly_code, golden_result, fp_preload) + build_fake_sim_env(data_size=256, mode="behave_sim", asm=None, data=None, specified_data_order = ["act_tensor", "weight_up_layer", "weight_gate_layer", "weight_down_layer"]) \ No newline at end of file diff --git a/compiler/asm_templates/ffn_asm.py b/compiler/asm_templates/ffn_asm.py index 327947cb..5c68a17d 100644 --- a/compiler/asm_templates/ffn_asm.py +++ b/compiler/asm_templates/ffn_asm.py @@ -2,6 +2,7 @@ from typing import Dict, List, Any, Optional from pathlib import Path +IMM2_BOUND = 2**18 def ffn_asm( mlen: int, @@ -18,9 +19,7 @@ def ffn_asm( down_weight_hbm_offset_reg: int, const_one_fp_address: int, - w_prefetch_amount: int, - activation_base_address: int, - result_base_address: int + activation_base_address: int ) -> str: """ Generates assembly code for a FFN operation. @@ -45,23 +44,31 @@ def ffn_asm( Downsize linear (b, s, intermediate_size) @ (intermediate_size, hidden_size) -> (b, s, hidden_size) """ - # Dot product of weight (Hidden Size, Hidden Size) and activation (Batch, 1, Hidden Size) - # get two registers from alive_registers, 1 as w address, 1 as a address + + # memory assignment + # 0 -> activation + # b * s * hidden_size -> upsize intermediate results + # b * s * (hidden_size + intermediate_size) -> gate projection results w_actual_register = alive_registers[0] a_actual_register = alive_registers[1] - result_register = alive_registers[2] + up_result_register = alive_registers[2] intermediate_register = alive_registers[3] + gate_result_register = alive_registers[4] # reset the registers generated_code = "; FFN Generation \n" - # Settings + # Settings for up and gate weight matrices prefetching + assert hidden_size * intermediate_size < IMM2_BOUND, f"hidden_size * hidden_size must be less than {IMM2_BOUND}" generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, {hidden_size * intermediate_size} \n" generated_code += f"C_SET_SCALE_REG gp{w_actual_register} \n" generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, {intermediate_size} \n" generated_code += f"C_SET_STRIDE_REG gp{w_actual_register} \n" generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, 0 \n" + assert hidden_size * batch * seq_len < IMM2_BOUND, f"hidden_size * batch * seq_len must be less than {IMM2_BOUND}" + generated_code += f"S_ADDI_INT gp{up_result_register}, gp0, {batch * seq_len * hidden_size} \n" + generated_code += f"S_ADDI_INT gp{gate_result_register}, gp{up_result_register}, {intermediate_size * batch * seq_len} \n" generated_code = " ; FFN Upsize Linear Generation \n" for weight_row in range (intermediate_size // blen): @@ -71,13 +78,14 @@ def ffn_asm( generated_code += f"S_ADDI_INT gp{w_actual_register}, gp{w_actual_register}, {mlen * mlen} \n" generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, 0 \n" for act_col in range (batch * seq_len // blen): - generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{a_actual_register}, {activation_base_address + act_col * hidden_size * blen} \n" + generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{up_result_register}, {activation_base_address + act_col * hidden_size * blen} \n" for inner_loop_index in range (hidden_size // mlen): generated_code += f"M_MM 0, gp{w_actual_register}, gp{a_actual_register} \n" generated_code += f"S_ADDI_INT gp{w_actual_register}, gp{w_actual_register}, {mlen * mlen} \n" generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{a_actual_register}, {mlen * blen} \n" - generated_code += f"M_MM_WO gp{result_register}, 0, 0 \n" - generated_code += f"S_ADDI_INT gp{w_actual_register}, gp{w_actual_register}, {hidden_size * blen} \n" + generated_code += f"S_ADDI_INT gp{intermediate_register}, gp{up_result_register}, {act_col * batch * seq_len * intermediate_size + weight_row * blen} \n" + generated_code += f"M_MM_WO gp{intermediate_register}, gp0, 0 \n" + generated_code += f"S_ADDI_INT gp{w_actual_register}, gp{w_actual_register}, {weight_row * blen} \n" generated_code = " ; FFN Gate Projection Generation \n" @@ -90,38 +98,47 @@ def ffn_asm( for act_col in range (batch * seq_len // blen): generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{a_actual_register}, {activation_base_address + act_col * hidden_size * blen} \n" for inner_loop_index in range (hidden_size // mlen): - generated_code += f"M_MM 0, {w_actual_register}, {a_actual_register} \n" + generated_code += f"M_MM gp0, gp{w_actual_register}, gp{a_actual_register} \n" generated_code += f"S_ADDI_INT gp{w_actual_register}, gp{w_actual_register}, {mlen * mlen} \n" generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{a_actual_register}, {mlen * blen} \n" - generated_code += f"M_MM_WO {result_register}, 0, 0 \n" + generated_code += f"S_ADDI_INT gp{intermediate_register}, gp{gate_result_register}, {act_col * batch * seq_len * intermediate_size + weight_row * blen} \n" + generated_code += f"M_MM_WO gp{intermediate_register}, gp0, 0 \n" generated_code += f"S_ADDI_INT gp{w_actual_register}, gp{w_actual_register}, {hidden_size * blen} \n" # Hidden Size SILU Activation Generation generated_code += "; SILU Generation \n" generated_code += f"S_LD_FP f1, gp0, {const_one_fp_address} \n" + generated_code += f"S_ADDI_INT gp{intermediate_register}, gp{gate_result_register}, 0 \n" for i in range(intermediate_size // vlen): - generated_code += f"V_SUB_VV gp{intermediate_register}, gp{intermediate_register}, gp{result_register} \n" - generated_code += f"V_EXP_V gp{intermediate_register}, gp{intermediate_register} \n" - generated_code += f"V_ADD_VF gp{intermediate_register}, gp{intermediate_register}, f1 \n" - generated_code += f"V_ADD_VF gp{intermediate_register}, gp{intermediate_register}, f1 \n" - generated_code += f"V_REC_V gp{intermediate_register}, gp{intermediate_register} \n" - generated_code += f"V_MUL_VV gp{result_register}, gp{result_register}, gp{intermediate_register} \n" + # 0 : -x + generated_code += f"V_SUB_VF gp{intermediate_register}, gp{intermediate_register}, f0, 0, 1 \n" + # 1 : exp(-x) + generated_code += f"V_EXP_V gp{intermediate_register}, gp{intermediate_register}, 0 \n" + # 2 : 1 + exp(-x) + generated_code += f"V_ADD_VF gp{intermediate_register}, gp{intermediate_register}, f1, 0 \n" + # 3 : 1 / (1 + exp(-x)) + generated_code += f"V_RECI_V gp{intermediate_register}, gp{intermediate_register}, 0 \n" + # 4 : (1 / (1 + exp(-x))) * gate_result + generated_code += f"V_MUL_VV gp{intermediate_register}, gp{intermediate_register}, gp{gate_result_register}, 0 \n" + generated_code += f"S_ADDI_INT gp{gate_result_register}, gp{gate_result_register}, {vlen} \n" generated_code += "; FFN Downsize Linear Generation \n" - - row_loop_over_hid = hidden_size // blen - col_loop_over_hid = intermediate_size // mlen - - # for i in range(row_loop_over_hid): - # generated_code += f"; <---- Generating New Row Tile at index {i} ----> \n" - # for j in range(col_loop_over_hid): - # generated_code += f"; <---- Generating New Column Tile at row {i} col {j} ----> \n" - # generated_code += f"M_MM 0, {w_actual_register}, {result_register} \n" - # generated_code += set_w_actual_address - # generated_code += set_result_address - # generated_code += f"M_MM_WO {a_actual_register}, 0, 0 \n" - # if (i % blen) == 0: - # generated_code += set_a_actual_address + generated_code += f"S_ADDI_INT gp{intermediate_register}, gp{gate_result_register}, 0 \n" + for weight_row in range (hidden_size // blen): + if weight_row % (mlen // blen) == 0: + for weight_col in range (intermediate_size // mlen): + generated_code += f"H_PREFETCH_M gp{w_actual_register}, gp{w_actual_register}, a{down_weight_hbm_offset_reg}, 1, 0 \n" + generated_code += f"S_ADDI_INT gp{w_actual_register}, gp{w_actual_register}, {mlen * mlen} \n" + generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, 0 \n" + for act_col in range (batch * seq_len // blen): + generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{intermediate_register}, {act_col * hidden_size * blen} \n" + for inner_loop_index in range (intermediate_size // mlen): + generated_code += f"M_MM 0, gp{w_actual_register}, gp{a_actual_register} \n" + generated_code += f"S_ADDI_INT gp{w_actual_register}, gp{w_actual_register}, {mlen * mlen} \n" + generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{a_actual_register}, {mlen * blen} \n" + generated_code += f"S_ADDI_INT gp{intermediate_register}, gp0, {activation_base_address + act_col * batch * seq_len * hidden_size + weight_row * blen} \n" + generated_code += f"M_MM_WO gp{intermediate_register}, gp0, 0 \n" + generated_code += f"S_ADDI_INT gp{w_actual_register}, gp{w_actual_register}, {intermediate_register * blen} \n" return generated_code diff --git a/tools/assembler/assembly_to_binary.py b/tools/assembler/assembly_to_binary.py index 1c38e77e..fd4c1b03 100644 --- a/tools/assembler/assembly_to_binary.py +++ b/tools/assembler/assembly_to_binary.py @@ -69,7 +69,7 @@ def _convert_to_binary(self, instruction): (rd << opw) + opcode ) - elif instruction.opcode in [ "H_PREFETCH_M", "H_PREFETCH_V", "H_STORE_V"]: + elif instruction.opcode in [ "H_PREFETCH_M", "H_PREFETCH_V", "H_STORE_V", "V_SUB_VF"]: binary_instruction = ( (funct1 << (opw + 4 * ow)) + (rstride << (opw + 3 * ow)) + @@ -78,7 +78,7 @@ def _convert_to_binary(self, instruction): (rd << opw) + opcode ) - elif instruction.opcode in ["V_ADD_VV", "V_ADD_VF", "V_SUB_VV", "V_SUB_VF", "V_MUL_VV", "V_MUL_VF", "V_EXP_V", "V_RECI_V", "V_RED_SUM", "V_RED_MAX"]: + elif instruction.opcode in ["V_ADD_VV", "V_ADD_VF", "V_MUL_VV", "V_SUB_VV", "V_MUL_VF", "V_EXP_V", "V_RECI_V", "V_RED_SUM", "V_RED_MAX"]: binary_instruction = ( (rmask << (opw + 3 * ow)) + (rs2 << (opw + 2 * ow)) + diff --git a/tools/assembler/parser.py b/tools/assembler/parser.py index 5c31533e..2cd772b0 100644 --- a/tools/assembler/parser.py +++ b/tools/assembler/parser.py @@ -62,7 +62,7 @@ def load_isa_settings(file_path: str) -> dict: class Instruction: - def __init__(self, opcode: str, rd: str, rs1: Optional[str], rs2: Optional[str], rstride: Optional[str], funct1: Optional[int], imm: Optional[int] = None): + def __init__(self, opcode: str, rd: str, rs1: Optional[str], rs2: Optional[str], rstride: Optional[str], funct1: Optional[int], imm: Optional[int] = None, rflag: Optional[int] = None): self.opcode = opcode self.rd = rd @@ -74,7 +74,7 @@ def __init__(self, opcode: str, rd: str, rs1: Optional[str], rs2: Optional[str], self.rmask = rstride def __repr__(self): - return f"Instruction(opcode='{self.opcode}', rd='{self.rd}', rs1='{self.rs1}', rs2='{self.rs2}', rstride = '{self.rstride}', funct1={self.funct1}, imm={self.imm})" + return f"Instruction(opcode='{self.opcode}', rd='{self.rd}', rs1='{self.rs1}', rs2='{self.rs2}', rstride = '{self.rstride}', funct1={self.funct1}, imm={self.imm}, rflag={self.rflag})" def parse_asm_file(file_path: str) -> List[Instruction]: From c79d74e3a9e268fd296baa839373a682734234b0 Mon Sep 17 00:00:00 2001 From: GeorgeWu1204 Date: Mon, 24 Nov 2025 21:39:57 +0000 Subject: [PATCH 17/34] fix weight memory mismatch bugs in matmul, pass check for linear --- behavioral_simulator/src/main.rs | 22 +++---- behavioral_simulator/testbench/attn_test.py | 4 +- behavioral_simulator/testbench/bmm_test.py | 4 +- behavioral_simulator/testbench/dllm1_test.py | 4 +- behavioral_simulator/testbench/ffn_test.py | 4 +- behavioral_simulator/testbench/linear_test.py | 61 ++++++++--------- behavioral_simulator/testbench/rms_test.py | 4 +- .../testbench/s_map_v_test.py | 4 +- .../testbench/two_input_test.py | 4 +- behavioral_simulator/testbench/view_mem.py | 8 +-- compiler/asm_templates/preload_act.py | 3 +- compiler/asm_templates/projection_asm.py | 66 ++++--------------- justfile | 2 +- tools/sim_env_utils/__init__.py | 2 +- tools/sim_env_utils/build_env.py | 4 +- 15 files changed, 79 insertions(+), 117 deletions(-) diff --git a/behavioral_simulator/src/main.rs b/behavioral_simulator/src/main.rs index 660f7796..2b3b866f 100644 --- a/behavioral_simulator/src/main.rs +++ b/behavioral_simulator/src/main.rs @@ -353,15 +353,15 @@ struct MatrixMachine { impl MatrixMachine { async fn mm(&mut self, m_addr: u32, v_addr: u32) { - // println!("======================== M_MM =========================="); - // println!("m_addr = {:?}", m_addr); - // println!("v_addr = {:?}", v_addr); + println!("======================== M_MM =========================="); + println!("m_addr = {:?}", m_addr); + println!("v_addr = {:?}", v_addr); let (mat_base, mat_offset) = m_addr.multiple_and_offset(self.mlen * self.mlen); - // println!("mat_offset = {:?}", mat_offset); - // println!("mat_base = {:?}", mat_base); + println!("mat_offset = {:?}", mat_offset); + println!("mat_base = {:?}", mat_base); assert!(mat_offset.is_multiple_of(self.blen)); - let mat_row_offset = mat_offset as i64 / self.mlen as i64; - + assert!(mat_offset <= self.mlen); + let mat_row_offset = mat_offset as i64; let full_mat = self.mram.read(mat_base).await; // Slice columns instead of rows: [mlen, blen] let mat = full_mat @@ -381,6 +381,8 @@ impl MatrixMachine { } // Stack along dimension 0 to get [blen, mlen] let vec = tch::Tensor::stack(&tensors, 0); + // println!("vec = {}", vec); + println!("mat = {}", mat); // Now vec @ mat: [blen, mlen] @ [mlen, blen] = [blen, blen] self.m_accum += vec.matmul(&mat); } @@ -933,10 +935,6 @@ impl Accelerator { } else { load_dim }; - // println!("Call transfer_from_hbm"); - // println!("stride = {:?}", stride); - // println!("index = {:?}", index); - // println!("scale_index = {:?}", scale_index); Executor::current().spawn(async move { let element_ty = hbm_type.element_type(); @@ -1624,7 +1622,7 @@ async fn start() { use std::fs; let op_file = fs::read_to_string(opts.opcode).unwrap(); - eprintln!("Loaded opcode file: {:?}", op_file); + // eprintln!("Loaded opcode file: {:?}", op_file); let op: Vec = op_file .split_whitespace() // split by spaces/newlines diff --git a/behavioral_simulator/testbench/attn_test.py b/behavioral_simulator/testbench/attn_test.py index 3fc54f38..74ea3c48 100644 --- a/behavioral_simulator/testbench/attn_test.py +++ b/behavioral_simulator/testbench/attn_test.py @@ -14,7 +14,7 @@ from create_sim_env import create_sim_env from transformers.modeling_flash_attention_utils import _flash_attention_forward as _flash_attention_forward_ref from aria_lm_ops.models.llama import flash_attn2_gemv -from sim_env_utils import build_fake_sim_env +from sim_env_utils import build_sim_env # Reshape q, k, v so that each batch's data stays together and rows are of length mem_row_size @@ -145,4 +145,4 @@ ) create_sim_env(input_tensor, weights, gen_assembly_code, golden_result, fp_preload) - build_fake_sim_env(data_size=256, mode="behave_sim", asm="attn", data=None, specified_data_order = ["q", "k", "v"]) \ No newline at end of file + build_sim_env(data_size=256, mode="behave_sim", asm="attn", data=None, specified_data_order = ["q", "k", "v"]) \ No newline at end of file diff --git a/behavioral_simulator/testbench/bmm_test.py b/behavioral_simulator/testbench/bmm_test.py index 7e30ebe4..b48a26b2 100644 --- a/behavioral_simulator/testbench/bmm_test.py +++ b/behavioral_simulator/testbench/bmm_test.py @@ -8,7 +8,7 @@ from test_data_gen import get_weights_path, generate_and_save_random_weights from compiler.asm_templates import batched_matmul_asm, preload_addr_reg_asm, reset_reg_asm from create_sim_env import create_sim_env -from sim_env_utils import build_fake_sim_env +from sim_env_utils import build_sim_env if __name__ == "__main__": # Testing the operation (hidden_size, hidden_size) @ (hidden_size, batch_size) @@ -65,7 +65,7 @@ create_sim_env(input_tensor, weight_2_tensor, gen_assembly_code, golden_result, fp_preload) - build_fake_sim_env(data_size=256, mode="behave_sim", asm="linear", data=None, specified_data_order = ["input_tensor", "model_weights"]) + build_sim_env(data_size=256, mode="behave_sim", asm="linear", data=None, specified_data_order = ["input_tensor", "model_weights"]) print("================================================") print("Finished generating assembly code") diff --git a/behavioral_simulator/testbench/dllm1_test.py b/behavioral_simulator/testbench/dllm1_test.py index 0b7ce6e9..814bd454 100644 --- a/behavioral_simulator/testbench/dllm1_test.py +++ b/behavioral_simulator/testbench/dllm1_test.py @@ -8,7 +8,7 @@ from test_data_gen import get_weights_path, generate_and_save_random_weights from compiler.asm_templates import preload_act_asm, reset_reg_asm, preload_addr_reg_asm from create_sim_env import create_sim_env -from sim_env_utils import build_fake_sim_env +from sim_env_utils import build_sim_env import torch.nn.functional as F from tools.memory_mapping.hbm_addr_map import align_addr_to_hbm_bandwidth @@ -80,4 +80,4 @@ create_sim_env(input_tensor, weights, gen_assembly_code, golden_result, fp_preload) - build_fake_sim_env(data_size=256, mode="behave_sim", asm="dllm", data=None, specified_data_order = ["logits"]) \ No newline at end of file + build_sim_env(data_size=256, mode="behave_sim", asm="dllm", data=None, specified_data_order = ["logits"]) \ No newline at end of file diff --git a/behavioral_simulator/testbench/ffn_test.py b/behavioral_simulator/testbench/ffn_test.py index e20af183..a5c3423e 100644 --- a/behavioral_simulator/testbench/ffn_test.py +++ b/behavioral_simulator/testbench/ffn_test.py @@ -8,7 +8,7 @@ from test_data_gen import get_weights_path, generate_and_save_random_weights from compiler.asm_templates import ffn_asm, preload_addr_reg_asm, reset_reg_asm from create_sim_env import create_sim_env -from sim_env_utils import build_fake_sim_env +from sim_env_utils import build_sim_env class LlamaFeedForward(nn.Module): @@ -119,4 +119,4 @@ def forward(self, x: Tensor) -> Tensor: ) create_sim_env(input_tensor, gen_assembly_code, golden_result, fp_preload) - build_fake_sim_env(data_size=256, mode="behave_sim", asm=None, data=None, specified_data_order = ["act_tensor", "weight_up_layer", "weight_gate_layer", "weight_down_layer"]) \ No newline at end of file + build_sim_env(data_size=256, mode="behave_sim", asm=None, data=None, specified_data_order = ["act_tensor", "weight_up_layer", "weight_gate_layer", "weight_down_layer"]) \ No newline at end of file diff --git a/behavioral_simulator/testbench/linear_test.py b/behavioral_simulator/testbench/linear_test.py index 990de79f..fe757e24 100644 --- a/behavioral_simulator/testbench/linear_test.py +++ b/behavioral_simulator/testbench/linear_test.py @@ -8,7 +8,7 @@ from test_data_gen import get_weights_path, generate_and_save_random_weights from compiler.asm_templates import projection_asm, preload_act_asm, reset_reg_asm, preload_addr_reg_asm from create_sim_env import create_sim_env -from sim_env_utils import build_fake_sim_env +from sim_env_utils import build_sim_env # TODOs: Need to integrate the MX quantizer here. @@ -34,13 +34,11 @@ # generate_and_save_random_weights(hidden_size, hidden_size, get_weights_path('model_weights.pt')) torch.manual_seed(42) - input_tensor = torch.randn(batch_size, hidden_size) - # Print input_tensor split in half along columns, as two (4, 64) tensors - print("input_tensor lhs (4, 64):\n", input_tensor[:, :64]) - print("input_tensor rhs (4, 64):\n", input_tensor[:, 64:]) - + act_tensor = torch.randn(batch_size, hidden_size) original_layer = nn.Linear(in_features=hidden_size, out_features=hidden_size, bias=False) weights = original_layer.state_dict() + + original_output = original_layer(act_tensor) # Print weight k, (128, 128) -> print 4 quadrants of (64, 64) each # Quadrant indices: @@ -48,36 +46,41 @@ # 1: [0:64, 64:128] # 2: [64:128, 0:64] # 3: [64:128, 64:128] - w_k = weights['weight'] if isinstance(weights, dict) else weights + w_k = weights['weight'].t() if isinstance(weights, dict) else weights print("Weight k shape:", w_k.shape) - for idx, (r_slice, c_slice) in enumerate([ - (slice(0, 64), slice(0, 64)), - (slice(0, 64), slice(64, 128)), - (slice(64, 128), slice(0, 64)), - (slice(64, 128), slice(64, 128)), - ]): - print(f"---------- Quadrant {idx}: Rows {r_slice}, Cols {c_slice} ----------") - print(w_k[r_slice, c_slice]) + # for idx, (r_slice, c_slice) in enumerate([ + # (slice(0, 64), slice(0, 64)), + # (slice(0, 64), slice(64, 128)), + # (slice(64, 128), slice(0, 64)), + # (slice(64, 128), slice(64, 128)), + # ]): + # print(f"---------- Quadrant {idx}: Rows {r_slice}, Cols {c_slice} ----------") + # print(w_k[r_slice, c_slice]) # Print the matmul result of input_tensor[:, :64] and weight[0:64, 0:4] - matmul_result_11 = input_tensor[:, :64] @ w_k[:64, :4] - print("Matmul result of input_tensor[:, :64] @ weight[0:64, 0:4]:") - print(matmul_result_11) + matmul_result_11 = act_tensor[:, :64] @ w_k[:64, :4] + # print("act_tensor[:, :64]: \n", act_tensor[:, :64]) + # print("w_k[:64, :4]: \n", w_k[:64, :4]) + # print("Matmul result of input_tensor[:, :64] @ weight[0:64, 0:4]:") + # print(matmul_result_11) - matmul_result_22 = input_tensor[:, 64:] @ w_k[64:, :4] + matmul_result_22 = act_tensor[:, 64:] @ w_k[64:, :4] print("Matmul result of input_tensor[:, 64:] @ weight[64:, :4]:") + print("act_tensor[:, 64:]: \n", act_tensor[:, 64:]) + print("w_k[64:, :4]: \n", w_k[64:, :4]) print(matmul_result_22) print ("sum of two matmul results:", matmul_result_11 + matmul_result_22) - original_output = original_layer(input_tensor) - print ("original_output:", original_output) + input_tensor = { + "act_tensor": act_tensor, + "weights": weights['weight'].t(), + } golden_result = { "input_tensor": input_tensor, - "weights": weights, "original_output": original_output } @@ -90,8 +93,8 @@ addr_reg_val=[int(hidden_size * batch_size * real_data_ratio), int((hidden_size * (batch_size + 1) + hidden_size * hidden_size) * real_data_ratio)] ) - print("hidden_size * batch_size * real_data_ratio", hidden_size * batch_size * real_data_ratio) - print("(hidden_size * (batch_size + 1) + hidden_size * hidden_size) * real_data_ratio", (hidden_size * (batch_size + 1) + hidden_size * hidden_size) * real_data_ratio) + # print("hidden_size * batch_size * real_data_ratio", hidden_size * batch_size * real_data_ratio) + # print("(hidden_size * (batch_size + 1) + hidden_size * hidden_size) * real_data_ratio", (hidden_size * (batch_size + 1) + hidden_size * hidden_size) * real_data_ratio) # Reset the registers gen_assembly_code += reset_reg_asm( @@ -106,7 +109,8 @@ hidden_size=128, alive_registers=[1,2,3], act_vram_offset=0, - activation_offset_reg=0 + activation_offset_reg=0, + stride_size=hidden_size ) # Reset the registers @@ -120,17 +124,14 @@ batch=4, hidden_size=128, alive_registers=[1,2,3,4], - head_dim=128, w_base_hbm_offset_reg=1, - rope_hbm_offset_reg=0, - rope_on_chip_address=0, activation_base_address=0, result_base_address=hidden_size * batch_size, rope_enabled=False ) - create_sim_env(input_tensor, weights['weight'].t(), gen_assembly_code, golden_result, fp_preload) - build_fake_sim_env(data_size=256, mode="behave_sim", asm="linear", data=None, specified_data_order = ["input_tensor", "model_weights"]) + create_sim_env(input_tensor, gen_assembly_code, golden_result, fp_preload) + build_sim_env(data_size=256, mode="behave_sim", asm="linear", data=None, specified_data_order = ["act_tensor", "weights"]) print("================================================") print("Finished generating assembly code") diff --git a/behavioral_simulator/testbench/rms_test.py b/behavioral_simulator/testbench/rms_test.py index f1e77e6e..3357dd5a 100644 --- a/behavioral_simulator/testbench/rms_test.py +++ b/behavioral_simulator/testbench/rms_test.py @@ -8,7 +8,7 @@ from test_data_gen import get_weights_path, generate_and_save_random_weights from compiler.asm_templates import rms_norm_asm, preload_act_asm, reset_reg_asm, preload_addr_reg_asm from create_sim_env import create_sim_env -from sim_env_utils import build_fake_sim_env +from sim_env_utils import build_sim_env # Taken from LLAMA RMSNorm implementation @@ -135,4 +135,4 @@ def forward(self, x): ) create_sim_env(input_tensor, weights['weight'].t(), gen_assembly_code, golden_result, fp_preload) - build_fake_sim_env(data_size=256, mode="behave_sim", asm="rms", data=None, specified_data_order = ["input_tensor", "model_weights"]) \ No newline at end of file + build_sim_env(data_size=256, mode="behave_sim", asm="rms", data=None, specified_data_order = ["input_tensor", "model_weights"]) \ No newline at end of file diff --git a/behavioral_simulator/testbench/s_map_v_test.py b/behavioral_simulator/testbench/s_map_v_test.py index f194e331..0cfc4148 100644 --- a/behavioral_simulator/testbench/s_map_v_test.py +++ b/behavioral_simulator/testbench/s_map_v_test.py @@ -9,7 +9,7 @@ from test_data_gen import get_weights_path, generate_and_save_random_weights from compiler.asm_templates import rms_norm_asm, projection_asm, preload_act_asm, reset_reg_asm, preload_addr_reg_asm from create_sim_env import create_sim_env -from sim_env_utils import build_fake_sim_env +from sim_env_utils import build_sim_env from tools.memory_mapping.hbm_addr_map import align_addr_to_hbm_bandwidth import torch.nn.functional as F @@ -65,4 +65,4 @@ gen_assembly_code += f"S_MAP_V_FP gp1, gp0, 0 \n" create_sim_env(input_tensor, weights, gen_assembly_code, golden_result, fp_preload) - build_fake_sim_env(data_size=256, mode="behave_sim", asm="dllm", data=None, specified_data_order = ["input_tensor1","input_tensor2"]) \ No newline at end of file + build_sim_env(data_size=256, mode="behave_sim", asm="dllm", data=None, specified_data_order = ["input_tensor1","input_tensor2"]) \ No newline at end of file diff --git a/behavioral_simulator/testbench/two_input_test.py b/behavioral_simulator/testbench/two_input_test.py index f4d9e978..811f7942 100644 --- a/behavioral_simulator/testbench/two_input_test.py +++ b/behavioral_simulator/testbench/two_input_test.py @@ -9,7 +9,7 @@ from test_data_gen import get_weights_path, generate_and_save_random_weights from compiler.asm_templates import rms_norm_asm, projection_asm, preload_act_asm, reset_reg_asm, preload_addr_reg_asm from create_sim_env import create_sim_env -from sim_env_utils import build_fake_sim_env +from sim_env_utils import build_sim_env from tools.memory_mapping.hbm_addr_map import align_addr_to_hbm_bandwidth import torch.nn.functional as F @@ -82,4 +82,4 @@ ) create_sim_env(input_tensor, weights, gen_assembly_code, golden_result, fp_preload) - build_fake_sim_env(data_size=256, mode="behave_sim", asm="dllm", data=None, specified_data_order = ["input_tensor1","input_tensor2"]) \ No newline at end of file + build_sim_env(data_size=256, mode="behave_sim", asm="dllm", data=None, specified_data_order = ["input_tensor1","input_tensor2"]) \ No newline at end of file diff --git a/behavioral_simulator/testbench/view_mem.py b/behavioral_simulator/testbench/view_mem.py index 5674434a..7f76b4af 100644 --- a/behavioral_simulator/testbench/view_mem.py +++ b/behavioral_simulator/testbench/view_mem.py @@ -117,8 +117,8 @@ def raw_to_fp(bits_val): # print("\nViewing VRAM dump from O_Old Base Address") # view_bin_file_by_row(vram_file, exp_width=8, man_width=7, row_dim=64, num_bytes_per_val=2, start_row_idx=576, load_row_size=16) - # print("Viewing MRAM dump 0 to 7 rows (BF16 format)") - # view_bin_file_by_row(mram_file, exp_width=8, man_width=7, row_dim=64, num_bytes_per_val=2, start_row_idx=0, load_row_size=8) + print("Viewing MRAM dump 0 to 7 rows (BF16 format)") + view_bin_file_by_row(mram_file, exp_width=8, man_width=7, row_dim=64, num_bytes_per_val=2, start_row_idx=0, load_row_size=8) - # print("Viewing MRAM dump 64 to 71 rows (BF16 format)") - # view_bin_file_by_row(mram_file, exp_width=8, man_width=7, row_dim=64, num_bytes_per_val=2, start_row_idx=64, load_row_size=8) \ No newline at end of file + print("Viewing MRAM dump 64 to 71 rows (BF16 format)") + view_bin_file_by_row(mram_file, exp_width=8, man_width=7, row_dim=64, num_bytes_per_val=2, start_row_idx=64, load_row_size=8) \ No newline at end of file diff --git a/compiler/asm_templates/preload_act.py b/compiler/asm_templates/preload_act.py index cd48957d..3b3df2e0 100644 --- a/compiler/asm_templates/preload_act.py +++ b/compiler/asm_templates/preload_act.py @@ -14,6 +14,7 @@ def preload_act_asm( ) -> str: """ Generates assembly code for preloading activation. + Memory Layout: Here we assume the activation is stored in (Hidden // MLEN , Batch, MLEN) """ generated_code = "; Preload Activation Generation \n" # get two registers from alive_registers, 1 as a address @@ -43,7 +44,7 @@ def preload_act_asm( generated_code += f"C_SET_STRIDE_REG gp{set_stride_register} \n" for i in range((batch * hidden_size) // (vlen * preload_len)): generated_code += f"H_PREFETCH_V gp{result_register}, gp{a_actual_register}, a{activation_offset_reg}, 1, 0 \n" - generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{a_actual_register}, {vlen * preload_len} \n" + generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{a_actual_register}, {vlen} \n" generated_code += f"S_ADDI_INT gp{result_register}, gp{result_register}, {vlen * preload_len} \n" return generated_code \ No newline at end of file diff --git a/compiler/asm_templates/projection_asm.py b/compiler/asm_templates/projection_asm.py index 3483cc5c..a7c89a1f 100644 --- a/compiler/asm_templates/projection_asm.py +++ b/compiler/asm_templates/projection_asm.py @@ -1,27 +1,27 @@ import os from typing import Dict, List, Any, Optional from pathlib import Path - +import math IMM2_BOUND = 2**18 + def projection_asm( mlen: int, blen: int, batch: int, hidden_size: int, alive_registers: List[int], - head_dim: int, w_base_hbm_offset_reg: int, - rope_hbm_offset_reg: int, - rope_on_chip_address: int, activation_base_address: int, result_base_address: int, - rope_enabled: bool = True + rope_enabled: bool = True, + rope_hbm_offset_reg: int = 0, + rope_on_chip_address: int = 0 ) -> str: """ Generates assembly code for a general matrix multiplication operation. (Batch, Hidden Size) @ (Hidden Size, Hidden Size) -> (Batch, Hidden Size) - + assume Batch = BLEN for this version. Args: mlen (int): The number of rows in the first matrix. blen (int): The number of columns in the second matrix. @@ -32,16 +32,14 @@ def projection_asm( Returns: str: Generated assembly code for projection, including dot product and RoPE(cond) """ - generated_code = "" - assert batch <= blen, "Batch size must be less than blen" - # get two registers from alive_registers, 1 as w address, 1 as a address - result_register = alive_registers[0] - w_actual_register = alive_registers[1] - w_hbm_offset_register = alive_registers[2] - a_actual_register = alive_registers[3] + generated_code = "; Projection Generation \n" + + result_register = alive_registers[0] + w_actual_register = alive_registers[1] + w_hbm_offset_register = alive_registers[2] + a_actual_register = alive_registers[3] # Set scale offset - #TODO: when hidden is large, cannot use addi command. assert hidden_size * hidden_size < IMM2_BOUND, f"hidden_size * hidden_size must be less than {IMM2_BOUND}" generated_code += f"S_ADDI_INT gp{a_actual_register}, gp0, {hidden_size * hidden_size} \n" generated_code += f"C_SET_SCALE_REG gp{a_actual_register} \n" @@ -49,7 +47,6 @@ def projection_asm( generated_code += f"C_SET_STRIDE_REG gp{a_actual_register} \n" # reset the registers - row_loop_over_hid = hidden_size // blen col_loop_over_hid = hidden_size // mlen generated_code += f"S_ADDI_INT gp{a_actual_register}, gp0, {activation_base_address} \n" @@ -60,56 +57,21 @@ def projection_asm( # Load a complete col of hidden size into on-chip memory (hidden_size, mlen) generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, 0 \n" generated_code += f"S_ADDI_INT gp{w_hbm_offset_register}, gp0, {i * blen} \n" + generated_code += f"S_ADDI_INT gp{result_register}, gp0, {result_base_address + (math.floor(i / (mlen // blen))) * mlen * blen} \n" for k in range (hidden_size // mlen): - generated_code += f"; <---- Generating New Row Tile at index {i} col {k} ----> \n" generated_code += f"H_PREFETCH_M gp{w_actual_register}, gp{w_hbm_offset_register}, a{w_base_hbm_offset_reg}, 1, 0 \n" generated_code += f"S_ADDI_INT gp{w_hbm_offset_register}, gp{w_hbm_offset_register}, {mlen * hidden_size} \n" generated_code += f"S_ADDI_INT gp{w_actual_register}, gp{w_actual_register}, {mlen * mlen} \n" generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, 0 \n" else: generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, {(i % (mlen // blen)) * blen} \n" - + generated_code += f"S_ADDI_INT gp{result_register}, gp{result_register}, {blen} \n" for j in range(col_loop_over_hid): # Loop over the hidden size dimension - generated_code += f"; <---- Generating New Column Tile at row {i} col {j} \n" generated_code += f"M_MM 0, gp{w_actual_register}, gp{a_actual_register} \n" generated_code += f"S_ADDI_INT gp{w_actual_register}, gp{w_actual_register}, {mlen * mlen} \n" generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{a_actual_register}, {mlen * blen} \n" generated_code += f"M_MM_WO {result_register}, gp0, 0 \n" generated_code += f"S_ADDI_INT gp{a_actual_register}, gp0, {activation_base_address} \n" - generated_code += f"S_ADDI_INT gp{result_register}, gp{result_register}, {blen} \n" - - - # RoPE - if rope_enabled: - generated_code += "; Generating RoPE code here \n" - generated_code += f"S_ADDI_INT gp{result_register}, gp0, {result_base_address} \n" - upper_base_register = alive_registers[0] - lower_base_register = alive_registers[1] - roped_upper_base_register = alive_registers[2] - roped_lower_base_register = alive_registers[3] - cos_base_register = alive_registers[4] - sin_base_register = alive_registers[5] - intermediate_1_register = alive_registers[6] - intermediate_2_register = alive_registers[7] - - generated_code += f"S_ADDI_INT gp{cos_base_register}, gp0, {rope_on_chip_address} \n" - generated_code += f"H_PREFETCH_V gp{cos_base_register}, gp{rope_on_chip_address}, a{rope_hbm_offset_reg}, 0, 0 \n" - generated_code += f"S_ADDI_INT gp{sin_base_register}, gp{cos_base_register}, {head_dim} \n" - generated_code += f"H_PREFETCH_V gp{sin_base_register}, gp{rope_on_chip_address}, a{rope_hbm_offset_reg}, 0, 0 \n" - - for i in range(batch * head_dim): - generated_code += f"; <---- Generating RoPE code for batch {i // head_dim} head {i % head_dim} ----> \n" - generated_code += f"V_MUL_VV gp{intermediate_1_register}, gp{upper_base_register}, gp{cos_base_register} \n" - generated_code += f"V_MUL_VV gp{intermediate_2_register}, gp{lower_base_register}, gp{sin_base_register} \n" - generated_code += f"V_SUB_VV gp{roped_upper_base_register}, gp{intermediate_1_register}, gp{intermediate_2_register} \n" - generated_code += f"V_MUL_VV gp{intermediate_1_register}, gp{upper_base_register}, gp{sin_base_register} \n" - generated_code += f"V_MUL_VV gp{intermediate_2_register}, gp{lower_base_register}, gp{cos_base_register} \n" - generated_code += f"V_ADD_VV gp{roped_lower_base_register}, gp{intermediate_1_register}, gp{intermediate_2_register} \n" - generated_code += f"S_ADDI_INT gp{upper_base_register}, gp{upper_base_register}, {mlen} \n" - generated_code += f"S_ADDI_INT gp{lower_base_register}, gp{lower_base_register}, {mlen} \n" - generated_code += f"S_ADDI_INT gp{roped_upper_base_register}, gp{roped_upper_base_register}, {mlen} \n" - generated_code += f"S_ADDI_INT gp{roped_lower_base_register}, gp{roped_lower_base_register}, {mlen} \n" - return generated_code diff --git a/justfile b/justfile index 2231845a..2e4aa42e 100644 --- a/justfile +++ b/justfile @@ -44,7 +44,7 @@ build-behave-sim-debug arg: cd behavioral_simulator && \ RUST_BACKTRACE=1 cargo run --release -- --opcode "$asm_path" --hbm "$data_path" --fpsram "$fp_sram_path" python3 behavioral_simulator/testbench/view_mem.py - python3 tools/utils/compare_match.py + # python3 tools/utils/compare_match.py build-rtl-sim arg: rm -rf test/Instr_Level_Benchmark/build/{{arg}} diff --git a/tools/sim_env_utils/__init__.py b/tools/sim_env_utils/__init__.py index 54f36c84..fe36eb7d 100644 --- a/tools/sim_env_utils/__init__.py +++ b/tools/sim_env_utils/__init__.py @@ -1 +1 @@ -from .build_env import build_fake_sim_env \ No newline at end of file +from .build_env import build_sim_env \ No newline at end of file diff --git a/tools/sim_env_utils/build_env.py b/tools/sim_env_utils/build_env.py index 63d473f3..c5e086ec 100644 --- a/tools/sim_env_utils/build_env.py +++ b/tools/sim_env_utils/build_env.py @@ -9,7 +9,7 @@ logger = get_logger("testbench") logger.setLevel(logging.DEBUG) -def build_fake_sim_env(data_size=256, mode="behave_sim", asm="attn", data=None, specified_data_order = None): +def build_sim_env(data_size=256, mode="behave_sim", asm="attn", data=None, specified_data_order = None): config_settings = load_svh_settings(str(SRC_PATH / "definitions" / "configuration.svh")) precision_settings = load_svh_settings(str(SRC_PATH / "definitions" / "precision.svh")) @@ -77,5 +77,5 @@ def build_fake_sim_env(data_size=256, mode="behave_sim", asm="attn", data=None, env_setup(grp_blocks, grp_bias, asm_file.parent, data_config, quant_config, hbm_row_width=config_settings["HBM_WIDTH"]) if __name__ == "__main__": - build_fake_sim_env() + build_sim_env() pass \ No newline at end of file From fd54752c5d8851e1db2010d51d2c0488748c2894 Mon Sep 17 00:00:00 2001 From: GeorgeWu1204 Date: Tue, 25 Nov 2025 01:09:07 +0000 Subject: [PATCH 18/34] fix rms norm bugs caused by new preload function --- behavioral_simulator/src/main.rs | 18 +++++----- behavioral_simulator/testbench/rms_test.py | 38 +++++++++++---------- compiler/asm_templates/normalization_asm.py | 11 +++--- tools/assembler/assembly_to_binary.py | 2 +- 4 files changed, 35 insertions(+), 34 deletions(-) diff --git a/behavioral_simulator/src/main.rs b/behavioral_simulator/src/main.rs index 2b3b866f..24fc1ea9 100644 --- a/behavioral_simulator/src/main.rs +++ b/behavioral_simulator/src/main.rs @@ -353,12 +353,12 @@ struct MatrixMachine { impl MatrixMachine { async fn mm(&mut self, m_addr: u32, v_addr: u32) { - println!("======================== M_MM =========================="); - println!("m_addr = {:?}", m_addr); - println!("v_addr = {:?}", v_addr); + // println!("======================== M_MM =========================="); + // println!("m_addr = {:?}", m_addr); + // println!("v_addr = {:?}", v_addr); let (mat_base, mat_offset) = m_addr.multiple_and_offset(self.mlen * self.mlen); - println!("mat_offset = {:?}", mat_offset); - println!("mat_base = {:?}", mat_base); + // println!("mat_offset = {:?}", mat_offset); + // println!("mat_base = {:?}", mat_base); assert!(mat_offset.is_multiple_of(self.blen)); assert!(mat_offset <= self.mlen); let mat_row_offset = mat_offset as i64; @@ -382,20 +382,20 @@ impl MatrixMachine { // Stack along dimension 0 to get [blen, mlen] let vec = tch::Tensor::stack(&tensors, 0); // println!("vec = {}", vec); - println!("mat = {}", mat); + // println!("mat = {}", mat); // Now vec @ mat: [blen, mlen] @ [mlen, blen] = [blen, blen] self.m_accum += vec.matmul(&mat); } async fn bmm(&mut self, m_addr: u32, v_addr: u32, stride_len: u32, bmm_scale: f32) { - println!("m_addr = {:?}", m_addr); - println!("v_addr = {:?}", v_addr); + // println!("m_addr = {:?}", m_addr); + // println!("v_addr = {:?}", v_addr); assert!(self.broadcast_amount * self.hlen == self.mlen); // Load matrix from matrix SRAM. let (mat_base, mat_offset) = m_addr.multiple_and_offset(self.mlen * self.blen); let (mat_offset, head_offset) = mat_offset.multiple_and_offset(self.mlen); - println!("mat_offset = {:?}", mat_offset); + // println!("mat_offset = {:?}", mat_offset); assert!(mat_offset.is_multiple_of(self.blen)); assert!(head_offset.is_multiple_of(self.hlen)); let full_mat = self.mram.read(mat_base).await; diff --git a/behavioral_simulator/testbench/rms_test.py b/behavioral_simulator/testbench/rms_test.py index 3357dd5a..467cd8c8 100644 --- a/behavioral_simulator/testbench/rms_test.py +++ b/behavioral_simulator/testbench/rms_test.py @@ -73,33 +73,30 @@ def forward(self, x): # generate_and_save_random_weights(hidden_size, hidden_size, get_weights_path('model_weights.pt')) torch.manual_seed(42) - input_tensor = torch.randn(batch_size, hidden_size) + act_tensor = torch.randn(batch_size, hidden_size) # Print input_tensor split in half along columns, as two (4, 64) tensors - print("input_tensor lhs (4, 64):\n", input_tensor[:, :64]) - print("input_tensor rhs (4, 64):\n", input_tensor[:, 64:]) + print("act_tensor lhs (4, 64):\n", act_tensor[:, :64]) + print("act_tensor rhs (4, 64):\n", act_tensor[:, 64:]) original_layer = RMSNorm(dim=hidden_size) weights = original_layer.state_dict() - original_output = original_layer(input_tensor) + input_tensor = { + "act_tensor": act_tensor, + "weights": weights['weight'].t(), + } + + original_output = original_layer(act_tensor) golden_result = { "input_tensor": input_tensor, - "weights": weights, "original_output": original_output } gen_assembly_code = "; RMSNorm Test Generation \n" - - # Set the addr offset for weight and bias - gen_assembly_code += preload_addr_reg_asm( - addr_reg_to_set=[1, 2], - available_registers=[1, 2], - addr_reg_val=[int(hidden_size * batch_size * real_data_ratio), int((hidden_size * (batch_size + 1) + hidden_size * hidden_size) * real_data_ratio)] - ) - print("hidden_size * batch_size * real_data_ratio", hidden_size * batch_size * real_data_ratio) - print("(hidden_size * (batch_size + 1) + hidden_size * hidden_size) * real_data_ratio", (hidden_size * (batch_size + 1) + hidden_size * hidden_size) * real_data_ratio) + # print("hidden_size * batch_size * real_data_ratio", hidden_size * batch_size * real_data_ratio) + # print("(hidden_size * (batch_size + 1) + hidden_size * hidden_size) * real_data_ratio", (hidden_size * (batch_size + 1) + hidden_size * hidden_size) * real_data_ratio) # Reset the registers @@ -115,12 +112,13 @@ def forward(self, x): hidden_size=128, alive_registers=[1,2,3], act_vram_offset=0, - activation_offset_reg=0 + activation_offset_reg=0, + stride_size=hidden_size ) # Reset the registers gen_assembly_code += reset_reg_asm( - alive_registers=[1,2,3,4] + alive_registers=[1,2,3] ) gen_assembly_code += rms_norm_asm( @@ -134,5 +132,9 @@ def forward(self, x): hidden_dim=128 ) - create_sim_env(input_tensor, weights['weight'].t(), gen_assembly_code, golden_result, fp_preload) - build_sim_env(data_size=256, mode="behave_sim", asm="rms", data=None, specified_data_order = ["input_tensor", "model_weights"]) \ No newline at end of file + create_sim_env(input_tensor, gen_assembly_code, golden_result, fp_preload) + build_sim_env(data_size=256, mode="behave_sim", asm="rms", data=None, specified_data_order = ["act_tensor", "weights"]) + + print("================================================") + print("Finished generating assembly code") + print("================================================") diff --git a/compiler/asm_templates/normalization_asm.py b/compiler/asm_templates/normalization_asm.py index da489889..dec889bd 100644 --- a/compiler/asm_templates/normalization_asm.py +++ b/compiler/asm_templates/normalization_asm.py @@ -32,7 +32,7 @@ def rms_norm_asm( for batch in range(batch_size): for i in range(hidden_dim // vlen): # Compute square of the activation vector and summation - generated_code += f"V_MUL_VV gp{scratchpad_addr}, gp{act_addr}, gp{act_addr} \n" + generated_code += f"V_MUL_VV gp{scratchpad_addr}, gp{act_addr}, gp{act_addr}, 0 \n" generated_code += f"V_RED_SUM f2, gp{scratchpad_addr} \n" # Move to next vector @@ -41,23 +41,22 @@ def rms_norm_asm( # Taking the avg generated_code += f"S_MUL_FP f2, f2, f3 \n" - # # Plus epsilon + # Plus epsilon generated_code += f"S_ADD_FP f2, f2, f1 \n" - # # Compute square root + # Compute square root generated_code += "S_SQRT_FP f2, f2 \n" - # # Compute reciprocal + # Compute reciprocal generated_code += "S_RECI_FP f2, f2 \n" for i in range(hidden_dim // vlen): # Normalize the activation vector - generated_code += f"V_MUL_VF gp{act_addr}, gp{act_addr}, f2 \n" + generated_code += f"V_MUL_VF gp{act_addr}, gp{act_addr}, f2, 0 \n" # Move to next vector generated_code += f"S_ADDI_INT gp{act_addr}, gp{act_addr}, {vlen * batch_size} \n" generated_code += "S_ADD_FP f2, f0, f0 \n" generated_code += f"S_ADDI_INT gp{act_addr}, gp0, {activation_base_address + vlen * batch} \n" - return generated_code \ No newline at end of file diff --git a/tools/assembler/assembly_to_binary.py b/tools/assembler/assembly_to_binary.py index fd4c1b03..27a642d4 100644 --- a/tools/assembler/assembly_to_binary.py +++ b/tools/assembler/assembly_to_binary.py @@ -41,7 +41,7 @@ def _convert_to_binary(self, instruction): imm = instruction.imm rmask = instruction.rmask binary_instruction = 0 - # print(f"Converting instruction: {instruction.opcode} with opcode={hex(opcode)}, rd={rd}, rs1={rs1}, rs2={rs2}, rstride={rstride}, funct1={funct1}, imm={imm}") + print(f"Converting instruction: {instruction.opcode} with opcode={hex(opcode)}, rd={rd}, rs1={rs1}, rs2={rs2}, rstride={rstride}, funct1={funct1}, imm={imm}") ow = self.operands_width opw = self.opcode_width From de7e2212788e7c3e7fd63480dd0a75ce3e69fe05 Mon Sep 17 00:00:00 2001 From: GeorgeWu1204 Date: Wed, 26 Nov 2025 03:35:28 +0000 Subject: [PATCH 19/34] fix up and gate proj asm bug in ffn --- behavioral_simulator/src/main.rs | 12 +- behavioral_simulator/testbench/ffn_test.py | 65 +++++++--- behavioral_simulator/testbench/view_mem.py | 3 + compiler/asm_templates/ffn_asm.py | 136 ++++++++++++--------- compiler/asm_templates/preload_act.py | 24 ++-- tools/assembler/assembly_to_binary.py | 2 +- 6 files changed, 151 insertions(+), 91 deletions(-) diff --git a/behavioral_simulator/src/main.rs b/behavioral_simulator/src/main.rs index 24fc1ea9..da2172ef 100644 --- a/behavioral_simulator/src/main.rs +++ b/behavioral_simulator/src/main.rs @@ -1673,11 +1673,13 @@ async fn start() { } // - Execute Instructions - accelerator - .do_ops(&dbg!( - op.into_iter().map(op::Opcode::decode).collect::>() - )) - .await; + // accelerator + // .do_ops(&dbg!( + // op.into_iter().map(op::Opcode::decode).collect::>() + // )) + // .await; + let decoded_ops = op.into_iter().map(op::Opcode::decode).collect::>(); + accelerator.do_ops(&decoded_ops).await; println!("gp1 = {:x}", accelerator.reg_file.gp_reg[1]); println!("scale = {}", accelerator.reg_file.scale); diff --git a/behavioral_simulator/testbench/ffn_test.py b/behavioral_simulator/testbench/ffn_test.py index a5c3423e..03376f6b 100644 --- a/behavioral_simulator/testbench/ffn_test.py +++ b/behavioral_simulator/testbench/ffn_test.py @@ -6,7 +6,7 @@ from torch import Tensor, nn # from acc_simulator.quantize.quantized_layers.linear import MXFPLinearPTQ from test_data_gen import get_weights_path, generate_and_save_random_weights -from compiler.asm_templates import ffn_asm, preload_addr_reg_asm, reset_reg_asm +from compiler.asm_templates import ffn_asm, preload_addr_reg_asm, reset_reg_asm, preload_act_asm from create_sim_env import create_sim_env from sim_env_utils import build_sim_env @@ -36,11 +36,23 @@ def __init__(self, dim: int, inter_dim: int, activation: str = "silu"): def forward(self, x: Tensor) -> Tensor: # SwiGLU: (x @ w1) * silu(x @ w3) print("input_tensor:\n", x) - print("up projection:\n", self.w1(x)) - print("gate projection:\n", self.w2(x)) - print("silu activation:\n", self.act(self.w1(x))) - print("product of silu activation and gate projection:\n", self.act(self.w1(x)) * self.w2(x)) - print("final output:\n", self.w3(self.act(self.w1(x)) * self.w2(x))) + print("up weight:\n", self.w1.weight.t()) + up_proj = self.w1(x) + print("up projection (all elements):") + # INSERT_YOUR_CODE + w1_out = self.w1(x) + print("self.w1(x) (all elements):") + print(w1_out) + # INSERT_YOUR_CODE + # Print all elements in w1_out without truncation + with torch.no_grad(): + torch.set_printoptions(profile="full") + print("w1_out (all elements):\n", w1_out) + torch.set_printoptions(profile="default") + # print("gate projection:\n", self.w2(x)) + # print("silu activation:\n", self.act(self.w1(x))) + # print("product of silu activation and gate projection:\n", self.act(self.w1(x)) * self.w2(x)) + # print("final output:\n", self.w3(self.act(self.w1(x)) * self.w2(x))) return self.w3(self.act(self.w1(x)) * self.w2(x)) if __name__ == "__main__": @@ -53,17 +65,16 @@ def forward(self, x: Tensor) -> Tensor: mlen = 64 blen = 4 vlen = 64 - seq_len = 128 + seq_len = 2 torch.manual_seed(42) - act_tensor = torch.randn(batch_size, hidden_size) + act_tensor = torch.randn(batch_size, seq_len, hidden_size) original_layer = LlamaFeedForward(dim=hidden_size, inter_dim=inter_dim) weight_up_layer = torch.randn(inter_dim, hidden_size) weight_gate_layer = torch.randn(inter_dim, hidden_size) weight_down_layer = torch.randn(hidden_size, inter_dim) - # Set weights for w1, w2, w3 to the generated tensors with torch.no_grad(): original_layer.w1.weight.copy_(weight_up_layer) @@ -74,10 +85,10 @@ def forward(self, x: Tensor) -> Tensor: print("original_output:", original_output) input_tensor = { - "act_tensor": act_tensor, - "weight_up_layer": weight_up_layer, - "weight_gate_layer": weight_gate_layer, - "weight_down_layer": weight_down_layer, + "act_tensor": act_tensor.reshape(batch_size * seq_len, hidden_size), + "weight_up_layer": weight_up_layer.t(), + "weight_gate_layer": weight_gate_layer.t(), + "weight_down_layer": weight_down_layer.t(), } golden_result = { @@ -91,16 +102,32 @@ def forward(self, x: Tensor) -> Tensor: gen_assembly_code += preload_addr_reg_asm( addr_reg_to_set=[1, 2, 3], available_registers=[1, 2, 3], - addr_reg_val=[int(hidden_size * batch_size * real_data_ratio), - int(hidden_size * batch_size * real_data_ratio) + int(hidden_size * inter_dim * real_data_ratio), - int(hidden_size * batch_size * real_data_ratio) + int(hidden_size * inter_dim * real_data_ratio) + int(inter_dim * hidden_size * real_data_ratio)] + addr_reg_val=[int(hidden_size * batch_size * seq_len * real_data_ratio), + int(hidden_size * batch_size * seq_len * real_data_ratio) + int(hidden_size * inter_dim * real_data_ratio), + int(hidden_size * batch_size * seq_len * real_data_ratio) + int(hidden_size * inter_dim * real_data_ratio) + int(inter_dim * hidden_size * real_data_ratio)] ) + print("up_addr_hbm_val:", int(hidden_size * batch_size * seq_len * real_data_ratio)) + print("gate_addr_hbm_val:", int(hidden_size * batch_size * seq_len * real_data_ratio) + int(hidden_size * inter_dim * real_data_ratio)) + print("down_addr_hbm_val:", int(hidden_size * batch_size * seq_len * real_data_ratio) + int(hidden_size * inter_dim * real_data_ratio) + int(inter_dim * hidden_size * real_data_ratio)) + # Reset the registers gen_assembly_code += reset_reg_asm( alive_registers=[1,2,3] ) + # Preload Activation + gen_assembly_code += preload_act_asm( + vlen=vlen, + preload_len=4, + batch=batch_size * seq_len, + hidden_size=hidden_size, + alive_registers=[1,2,3], + act_vram_offset=0, + activation_offset_reg=0, + stride_size=hidden_size + ) + # FFN Generation gen_assembly_code += ffn_asm( mlen=mlen, @@ -110,9 +137,9 @@ def forward(self, x: Tensor) -> Tensor: seq_len=seq_len, hidden_size=hidden_size, intermediate_size=inter_dim, - alive_registers=[1,2,3,4,5,6], - gate_weight_hbm_offset_reg=1, - up_weight_hbm_offset_reg=2, + alive_registers=[1,2,3,4,5,6,7], + up_weight_hbm_offset_reg=1, + gate_weight_hbm_offset_reg=2, down_weight_hbm_offset_reg=3, const_one_fp_address=1, activation_base_address=0 diff --git a/behavioral_simulator/testbench/view_mem.py b/behavioral_simulator/testbench/view_mem.py index 7f76b4af..7b952461 100644 --- a/behavioral_simulator/testbench/view_mem.py +++ b/behavioral_simulator/testbench/view_mem.py @@ -105,6 +105,9 @@ def raw_to_fp(bits_val): print("Viewing VRAM dump from 0 Base Address") view_bin_file_by_row(vram_file, exp_width=8, man_width=7, row_dim=64, num_bytes_per_val=2, start_row_idx=0, load_row_size=16) + print("Viewing VRAM dump from 16 Base Address") + view_bin_file_by_row(vram_file, exp_width=8, man_width=7, row_dim=64, num_bytes_per_val=2, start_row_idx=16, load_row_size=32) + # print("Viewing VRAM dump from Q Base Address") # view_bin_file_by_row(vram_file, exp_width=8, man_width=7, row_dim=64, num_bytes_per_val=2, start_row_idx=0, load_row_size=16) diff --git a/compiler/asm_templates/ffn_asm.py b/compiler/asm_templates/ffn_asm.py index 5c68a17d..d888918c 100644 --- a/compiler/asm_templates/ffn_asm.py +++ b/compiler/asm_templates/ffn_asm.py @@ -1,7 +1,7 @@ import os from typing import Dict, List, Any, Optional from pathlib import Path - +import math IMM2_BOUND = 2**18 def ffn_asm( @@ -51,10 +51,12 @@ def ffn_asm( # b * s * (hidden_size + intermediate_size) -> gate projection results w_actual_register = alive_registers[0] - a_actual_register = alive_registers[1] - up_result_register = alive_registers[2] - intermediate_register = alive_registers[3] - gate_result_register = alive_registers[4] + w_temp_register = alive_registers[1] + a_actual_register = alive_registers[2] + up_result_register = alive_registers[3] + intermediate_register = alive_registers[4] + gate_result_register = alive_registers[5] + w_hbm_offset_register = alive_registers[6] # reset the registers generated_code = "; FFN Generation \n" @@ -62,83 +64,105 @@ def ffn_asm( # Settings for up and gate weight matrices prefetching assert hidden_size * intermediate_size < IMM2_BOUND, f"hidden_size * hidden_size must be less than {IMM2_BOUND}" generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, {hidden_size * intermediate_size} \n" + print("Scale Amount", hidden_size * intermediate_size) generated_code += f"C_SET_SCALE_REG gp{w_actual_register} \n" generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, {intermediate_size} \n" generated_code += f"C_SET_STRIDE_REG gp{w_actual_register} \n" generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, 0 \n" assert hidden_size * batch * seq_len < IMM2_BOUND, f"hidden_size * batch * seq_len must be less than {IMM2_BOUND}" + # Set the address for on-chip sram generated_code += f"S_ADDI_INT gp{up_result_register}, gp0, {batch * seq_len * hidden_size} \n" generated_code += f"S_ADDI_INT gp{gate_result_register}, gp{up_result_register}, {intermediate_size * batch * seq_len} \n" - generated_code = " ; FFN Upsize Linear Generation \n" + generated_code += " ; FFN Upsize Linear Generation \n" for weight_row in range (intermediate_size // blen): if weight_row % (mlen // blen) == 0: + generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, 0 \n" + generated_code += f"S_ADDI_INT gp{w_hbm_offset_register}, gp0, {weight_row * blen} \n" + generated_code += f"S_ADDI_INT gp{intermediate_register}, gp{up_result_register}, 0 \n" + for weight_col in range (hidden_size // mlen): - generated_code += f"H_PREFETCH_M gp{w_actual_register}, gp{w_actual_register}, a{up_weight_hbm_offset_reg}, 1, 0 \n" + generated_code += f"H_PREFETCH_M gp{w_actual_register}, gp{w_hbm_offset_register}, a{up_weight_hbm_offset_reg}, 1, 0 \n" generated_code += f"S_ADDI_INT gp{w_actual_register}, gp{w_actual_register}, {mlen * mlen} \n" + generated_code += f"S_ADDI_INT gp{w_hbm_offset_register}, gp{w_hbm_offset_register}, {mlen * intermediate_size} \n" generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, 0 \n" - for act_col in range (batch * seq_len // blen): - generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{up_result_register}, {activation_base_address + act_col * hidden_size * blen} \n" + else: + generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, {(weight_row % (mlen // blen)) * blen} \n" + generated_code += f"S_ADDI_INT gp{intermediate_register}, gp{up_result_register}, {(weight_row % (mlen // blen)) * blen} \n" + for act_col in range ((batch * seq_len) // blen): + generated_code += f"S_ADDI_INT gp{a_actual_register}, gp0, {activation_base_address + act_col * mlen * blen} \n" + generated_code += f"S_ADDI_INT gp{w_temp_register}, gp{w_actual_register}, 0 \n" for inner_loop_index in range (hidden_size // mlen): - generated_code += f"M_MM 0, gp{w_actual_register}, gp{a_actual_register} \n" - generated_code += f"S_ADDI_INT gp{w_actual_register}, gp{w_actual_register}, {mlen * mlen} \n" - generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{a_actual_register}, {mlen * blen} \n" - generated_code += f"S_ADDI_INT gp{intermediate_register}, gp{up_result_register}, {act_col * batch * seq_len * intermediate_size + weight_row * blen} \n" + generated_code += f"M_MM 0, gp{w_temp_register}, gp{a_actual_register} \n" + generated_code += f"S_ADDI_INT gp{w_temp_register}, gp{w_temp_register}, {mlen * mlen} \n" + generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{a_actual_register}, {mlen * batch * seq_len} \n" generated_code += f"M_MM_WO gp{intermediate_register}, gp0, 0 \n" - generated_code += f"S_ADDI_INT gp{w_actual_register}, gp{w_actual_register}, {weight_row * blen} \n" + generated_code += f"S_ADDI_INT gp{intermediate_register}, gp{intermediate_register}, {blen * mlen} \n" # generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{a_actual_register}, {activation_base_address} \n" + if (weight_row + 1) % (mlen // blen) == 0 and weight_row != intermediate_size // blen - 1: + generated_code += f"S_ADDI_INT gp{up_result_register}, gp{up_result_register}, {mlen * batch * seq_len} \n" - - generated_code = " ; FFN Gate Projection Generation \n" + generated_code += " ; FFN Gate Projection Generation \n" for weight_row in range (intermediate_size // blen): if weight_row % (mlen // blen) == 0: + generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, 0 \n" + generated_code += f"S_ADDI_INT gp{w_hbm_offset_register}, gp0, {weight_row * blen} \n" + generated_code += f"S_ADDI_INT gp{intermediate_register}, gp{gate_result_register}, 0 \n" + for weight_col in range (hidden_size // mlen): - generated_code += f"H_PREFETCH_M gp{w_actual_register}, gp{w_actual_register}, a{gate_weight_hbm_offset_reg}, 1, 0 \n" + generated_code += f"H_PREFETCH_M gp{w_actual_register}, gp{w_hbm_offset_register}, a{gate_weight_hbm_offset_reg}, 1, 0 \n" generated_code += f"S_ADDI_INT gp{w_actual_register}, gp{w_actual_register}, {mlen * mlen} \n" + generated_code += f"S_ADDI_INT gp{w_hbm_offset_register}, gp{w_hbm_offset_register}, {mlen * intermediate_size} \n" generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, 0 \n" - for act_col in range (batch * seq_len // blen): - generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{a_actual_register}, {activation_base_address + act_col * hidden_size * blen} \n" + else: + generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, {(weight_row % (mlen // blen)) * blen} \n" + generated_code += f"S_ADDI_INT gp{intermediate_register}, gp{gate_result_register}, {(weight_row % (mlen // blen)) * blen} \n" + for act_col in range ((batch * seq_len) // blen): + generated_code += f"S_ADDI_INT gp{a_actual_register}, gp0, {activation_base_address + act_col * mlen * blen} \n" + generated_code += f"S_ADDI_INT gp{w_temp_register}, gp{w_actual_register}, 0 \n" for inner_loop_index in range (hidden_size // mlen): - generated_code += f"M_MM gp0, gp{w_actual_register}, gp{a_actual_register} \n" - generated_code += f"S_ADDI_INT gp{w_actual_register}, gp{w_actual_register}, {mlen * mlen} \n" - generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{a_actual_register}, {mlen * blen} \n" - generated_code += f"S_ADDI_INT gp{intermediate_register}, gp{gate_result_register}, {act_col * batch * seq_len * intermediate_size + weight_row * blen} \n" + generated_code += f"M_MM 0, gp{w_temp_register}, gp{a_actual_register} \n" + generated_code += f"S_ADDI_INT gp{w_temp_register}, gp{w_temp_register}, {mlen * mlen} \n" + generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{a_actual_register}, {mlen * batch * seq_len} \n" generated_code += f"M_MM_WO gp{intermediate_register}, gp0, 0 \n" - generated_code += f"S_ADDI_INT gp{w_actual_register}, gp{w_actual_register}, {hidden_size * blen} \n" + generated_code += f"S_ADDI_INT gp{intermediate_register}, gp{intermediate_register}, {blen * mlen} \n" # generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{a_actual_register}, {activation_base_address} \n" + if (weight_row + 1) % (mlen // blen) == 0 and weight_row != intermediate_size // blen - 1: + generated_code += f"S_ADDI_INT gp{gate_result_register}, gp{gate_result_register}, {mlen * batch * seq_len} \n" + # Hidden Size SILU Activation Generation - generated_code += "; SILU Generation \n" - generated_code += f"S_LD_FP f1, gp0, {const_one_fp_address} \n" - generated_code += f"S_ADDI_INT gp{intermediate_register}, gp{gate_result_register}, 0 \n" - for i in range(intermediate_size // vlen): - # 0 : -x - generated_code += f"V_SUB_VF gp{intermediate_register}, gp{intermediate_register}, f0, 0, 1 \n" - # 1 : exp(-x) - generated_code += f"V_EXP_V gp{intermediate_register}, gp{intermediate_register}, 0 \n" - # 2 : 1 + exp(-x) - generated_code += f"V_ADD_VF gp{intermediate_register}, gp{intermediate_register}, f1, 0 \n" - # 3 : 1 / (1 + exp(-x)) - generated_code += f"V_RECI_V gp{intermediate_register}, gp{intermediate_register}, 0 \n" - # 4 : (1 / (1 + exp(-x))) * gate_result - generated_code += f"V_MUL_VV gp{intermediate_register}, gp{intermediate_register}, gp{gate_result_register}, 0 \n" - generated_code += f"S_ADDI_INT gp{gate_result_register}, gp{gate_result_register}, {vlen} \n" + # generated_code += "; SILU Generation \n" + # generated_code += f"S_LD_FP f1, gp0, {const_one_fp_address} \n" + # generated_code += f"S_ADDI_INT gp{intermediate_register}, gp{gate_result_register}, 0 \n" + # for i in range(intermediate_size // vlen): + # # 0 : -x + # generated_code += f"V_SUB_VF gp{intermediate_register}, gp{intermediate_register}, f0, 0, 1 \n" + # # 1 : exp(-x) + # generated_code += f"V_EXP_V gp{intermediate_register}, gp{intermediate_register}, 0 \n" + # # 2 : 1 + exp(-x) + # generated_code += f"V_ADD_VF gp{intermediate_register}, gp{intermediate_register}, f1, 0 \n" + # # 3 : 1 / (1 + exp(-x)) + # generated_code += f"V_RECI_V gp{intermediate_register}, gp{intermediate_register}, 0 \n" + # # 4 : (1 / (1 + exp(-x))) * gate_result + # generated_code += f"V_MUL_VV gp{intermediate_register}, gp{intermediate_register}, gp{gate_result_register}, 0 \n" + # generated_code += f"S_ADDI_INT gp{gate_result_register}, gp{gate_result_register}, {vlen} \n" - generated_code += "; FFN Downsize Linear Generation \n" - generated_code += f"S_ADDI_INT gp{intermediate_register}, gp{gate_result_register}, 0 \n" - for weight_row in range (hidden_size // blen): - if weight_row % (mlen // blen) == 0: - for weight_col in range (intermediate_size // mlen): - generated_code += f"H_PREFETCH_M gp{w_actual_register}, gp{w_actual_register}, a{down_weight_hbm_offset_reg}, 1, 0 \n" - generated_code += f"S_ADDI_INT gp{w_actual_register}, gp{w_actual_register}, {mlen * mlen} \n" - generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, 0 \n" - for act_col in range (batch * seq_len // blen): - generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{intermediate_register}, {act_col * hidden_size * blen} \n" - for inner_loop_index in range (intermediate_size // mlen): - generated_code += f"M_MM 0, gp{w_actual_register}, gp{a_actual_register} \n" - generated_code += f"S_ADDI_INT gp{w_actual_register}, gp{w_actual_register}, {mlen * mlen} \n" - generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{a_actual_register}, {mlen * blen} \n" - generated_code += f"S_ADDI_INT gp{intermediate_register}, gp0, {activation_base_address + act_col * batch * seq_len * hidden_size + weight_row * blen} \n" - generated_code += f"M_MM_WO gp{intermediate_register}, gp0, 0 \n" - generated_code += f"S_ADDI_INT gp{w_actual_register}, gp{w_actual_register}, {intermediate_register * blen} \n" + # generated_code += "; FFN Downsize Linear Generation \n" + # generated_code += f"S_ADDI_INT gp{intermediate_register}, gp{gate_result_register}, 0 \n" + # for weight_row in range (hidden_size // blen): + # if weight_row % (mlen // blen) == 0: + # for weight_col in range (intermediate_size // mlen): + # generated_code += f"H_PREFETCH_M gp{w_actual_register}, gp{w_actual_register}, a{down_weight_hbm_offset_reg}, 1, 0 \n" + # generated_code += f"S_ADDI_INT gp{w_actual_register}, gp{w_actual_register}, {mlen * mlen} \n" + # generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, 0 \n" + # for act_col in range (batch * seq_len // blen): + # generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{intermediate_register}, {act_col * hidden_size * blen} \n" + # for inner_loop_index in range (intermediate_size // mlen): + # generated_code += f"M_MM 0, gp{w_actual_register}, gp{a_actual_register} \n" + # generated_code += f"S_ADDI_INT gp{w_actual_register}, gp{w_actual_register}, {mlen * mlen} \n" + # generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{a_actual_register}, {mlen * blen} \n" + # generated_code += f"S_ADDI_INT gp{intermediate_register}, gp0, {activation_base_address + act_col * batch * seq_len * hidden_size + weight_row * blen} \n" + # generated_code += f"M_MM_WO gp{intermediate_register}, gp0, 0 \n" + # generated_code += f"S_ADDI_INT gp{w_actual_register}, gp{w_actual_register}, {intermediate_register * blen} \n" return generated_code diff --git a/compiler/asm_templates/preload_act.py b/compiler/asm_templates/preload_act.py index 3b3df2e0..b67cf3cb 100644 --- a/compiler/asm_templates/preload_act.py +++ b/compiler/asm_templates/preload_act.py @@ -1,6 +1,9 @@ import os from typing import Dict, List, Any, Optional from pathlib import Path +import math + +IMM2_BOUND = 2**18 def preload_act_asm( vlen: int, @@ -14,7 +17,7 @@ def preload_act_asm( ) -> str: """ Generates assembly code for preloading activation. - Memory Layout: Here we assume the activation is stored in (Hidden // MLEN , Batch, MLEN) + Memory Layout: Here we assume the activation is stored in (Hidden // MLEN , Batch (Integrate with Seq Len), MLEN) """ generated_code = "; Preload Activation Generation \n" # get two registers from alive_registers, 1 as a address @@ -28,10 +31,9 @@ def preload_act_asm( generated_code += f"C_SET_SCALE_REG gp{a_actual_register} \n" # reset the registers - set_a_base_address = f"S_ADDI_INT gp{a_actual_register}, gp0, 0 \n" - set_act_vram_base_address = f"S_ADDI_INT gp{result_register}, gp0, {act_vram_offset} \n" - generated_code += set_a_base_address - generated_code += set_act_vram_base_address + generated_code += f"S_ADDI_INT gp{a_actual_register}, gp0, 0 \n" + generated_code += f"S_ADDI_INT gp{result_register}, gp0, {act_vram_offset} \n" + load_amount_per_hidden = math.ceil(hidden_size / vlen) if batch == 1: generated_code += f"S_ADDI_INT gp{set_stride_register}, gp0, {stride_len} \n" @@ -42,9 +44,11 @@ def preload_act_asm( else: generated_code += f"S_ADDI_INT gp{set_stride_register}, gp0, {stride_len} \n" generated_code += f"C_SET_STRIDE_REG gp{set_stride_register} \n" - for i in range((batch * hidden_size) // (vlen * preload_len)): - generated_code += f"H_PREFETCH_V gp{result_register}, gp{a_actual_register}, a{activation_offset_reg}, 1, 0 \n" - generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{a_actual_register}, {vlen} \n" - generated_code += f"S_ADDI_INT gp{result_register}, gp{result_register}, {vlen * preload_len} \n" - + assert batch * hidden_size <= IMM2_BOUND, "batch * hidden_size must be less than {IMM2_BOUND}" + for i in range(load_amount_per_hidden): + for j in range(math.ceil(batch / preload_len)): + generated_code += f"H_PREFETCH_V gp{result_register}, gp{a_actual_register}, a{activation_offset_reg}, 1, 0 \n" + generated_code += f"S_ADDI_INT gp{result_register}, gp{result_register}, {vlen * preload_len} \n" + generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{a_actual_register}, {hidden_size * preload_len} \n" + generated_code += f"S_ADDI_INT gp{a_actual_register}, gp0, {(i + 1) * vlen} \n" return generated_code \ No newline at end of file diff --git a/tools/assembler/assembly_to_binary.py b/tools/assembler/assembly_to_binary.py index 27a642d4..fd4c1b03 100644 --- a/tools/assembler/assembly_to_binary.py +++ b/tools/assembler/assembly_to_binary.py @@ -41,7 +41,7 @@ def _convert_to_binary(self, instruction): imm = instruction.imm rmask = instruction.rmask binary_instruction = 0 - print(f"Converting instruction: {instruction.opcode} with opcode={hex(opcode)}, rd={rd}, rs1={rs1}, rs2={rs2}, rstride={rstride}, funct1={funct1}, imm={imm}") + # print(f"Converting instruction: {instruction.opcode} with opcode={hex(opcode)}, rd={rd}, rs1={rs1}, rs2={rs2}, rstride={rstride}, funct1={funct1}, imm={imm}") ow = self.operands_width opw = self.opcode_width From 2958d301628ae1ebfde5bc41729dd4f4bf5b84ff Mon Sep 17 00:00:00 2001 From: GeorgeWu1204 Date: Wed, 26 Nov 2025 18:30:23 +0000 Subject: [PATCH 20/34] silu pass acc check --- behavioral_simulator/testbench/ffn_test.py | 59 +++++++++---- behavioral_simulator/testbench/view_mem.py | 4 + compiler/asm_templates/ffn_asm.py | 96 ++++++++++++++-------- 3 files changed, 109 insertions(+), 50 deletions(-) diff --git a/behavioral_simulator/testbench/ffn_test.py b/behavioral_simulator/testbench/ffn_test.py index 03376f6b..88190be2 100644 --- a/behavioral_simulator/testbench/ffn_test.py +++ b/behavioral_simulator/testbench/ffn_test.py @@ -35,23 +35,50 @@ def __init__(self, dim: int, inter_dim: int, activation: str = "silu"): def forward(self, x: Tensor) -> Tensor: # SwiGLU: (x @ w1) * silu(x @ w3) - print("input_tensor:\n", x) - print("up weight:\n", self.w1.weight.t()) - up_proj = self.w1(x) + # print("input_tensor:\n", x) + # print("up weight:\n", self.w1.weight.t()) + # up_proj = self.w1(x) print("up projection (all elements):") - # INSERT_YOUR_CODE w1_out = self.w1(x) - print("self.w1(x) (all elements):") - print(w1_out) - # INSERT_YOUR_CODE - # Print all elements in w1_out without truncation - with torch.no_grad(): - torch.set_printoptions(profile="full") - print("w1_out (all elements):\n", w1_out) - torch.set_printoptions(profile="default") + w1_out = w1_out.reshape(w1_out.shape[0] * w1_out.shape[1], w1_out.shape[-1]) + print("self.w1(x) (upper all elements):") + print(w1_out[:, :8]) + print("self.w1(x) (mid upper all elements):") + print(w1_out[:, 64:72]) + print("self.w1(x) (mid lower all elements):") + print(w1_out[:, 128:135]) + print("self.w1(x) (lower all elements):") + print(w1_out[:, 192:199]) + + # w2_out = self.w2(x) + # w2_out = w2_out.reshape(w2_out.shape[0] * w2_out.shape[1], w2_out.shape[-1]) + # print("self.w2(x) (upper all elements):") + # print(w2_out[:, :8]) + # print("self.w2(x) (mid upper all elements):") + # print(w2_out[:, 64:72]) + # print("self.w2(x) (mid lower all elements):") + # print(w2_out[:, 128:135]) + # print("self.w2(x) (lower all elements):") # print("gate projection:\n", self.w2(x)) + # print("silu activation:\n", self.act(self.w1(x))) # print("product of silu activation and gate projection:\n", self.act(self.w1(x)) * self.w2(x)) + silu_mixed_out = self.act(self.w1(x)) * self.w2(x) + # silu_mixed_out = silu_mixed_out.reshape(silu_mixed_out.shape[0] * silu_mixed_out.shape[1], silu_mixed_out.shape[-1]) + # print(f"silu mixed out of shape {silu_mixed_out.shape}: \n") + # print("silu mixed out (upper all elements):") + # print(silu_mixed_out[:, :8]) + # print("silu mixed out (mid upper all elements):") + # print(silu_mixed_out[:, 64:72]) + # print("silu mixed out (mid lower all elements):") + # print(silu_mixed_out[:, 128:135]) + # print("silu mixed out (lower all elements):") + # print(silu_mixed_out[:, 192:199]) + outcome = self.w3(silu_mixed_out) + print("final output (upper all elements):") + print(outcome[:, :8]) + print("final output (mid upper all elements):") + print(outcome[:, 64:72]) # print("final output:\n", self.w3(self.act(self.w1(x)) * self.w2(x))) return self.w3(self.act(self.w1(x)) * self.w2(x)) @@ -73,7 +100,12 @@ def forward(self, x: Tensor) -> Tensor: original_layer = LlamaFeedForward(dim=hidden_size, inter_dim=inter_dim) weight_up_layer = torch.randn(inter_dim, hidden_size) weight_gate_layer = torch.randn(inter_dim, hidden_size) - weight_down_layer = torch.randn(hidden_size, inter_dim) + weight_down_layer = torch.ones(hidden_size, inter_dim) + print(f"weight_down_layer of shape {weight_down_layer.t().shape}: \n") + print("upper all elements:") + print(weight_down_layer.t()[:, :8]) + print("lower all elements:") + print(weight_down_layer.t()[:, 64:72]) # Set weights for w1, w2, w3 to the generated tensors with torch.no_grad(): @@ -82,7 +114,6 @@ def forward(self, x: Tensor) -> Tensor: original_layer.w3.weight.copy_(weight_down_layer) original_output = original_layer(act_tensor) - print("original_output:", original_output) input_tensor = { "act_tensor": act_tensor.reshape(batch_size * seq_len, hidden_size), diff --git a/behavioral_simulator/testbench/view_mem.py b/behavioral_simulator/testbench/view_mem.py index 7b952461..4846ad36 100644 --- a/behavioral_simulator/testbench/view_mem.py +++ b/behavioral_simulator/testbench/view_mem.py @@ -108,6 +108,10 @@ def raw_to_fp(bits_val): print("Viewing VRAM dump from 16 Base Address") view_bin_file_by_row(vram_file, exp_width=8, man_width=7, row_dim=64, num_bytes_per_val=2, start_row_idx=16, load_row_size=32) + + print("Viewing VRAM dump from 48 Base Address") + view_bin_file_by_row(vram_file, exp_width=8, man_width=7, row_dim=64, num_bytes_per_val=2, start_row_idx=48, load_row_size=32) + # print("Viewing VRAM dump from Q Base Address") # view_bin_file_by_row(vram_file, exp_width=8, man_width=7, row_dim=64, num_bytes_per_val=2, start_row_idx=0, load_row_size=16) diff --git a/compiler/asm_templates/ffn_asm.py b/compiler/asm_templates/ffn_asm.py index d888918c..626f41da 100644 --- a/compiler/asm_templates/ffn_asm.py +++ b/compiler/asm_templates/ffn_asm.py @@ -64,7 +64,6 @@ def ffn_asm( # Settings for up and gate weight matrices prefetching assert hidden_size * intermediate_size < IMM2_BOUND, f"hidden_size * hidden_size must be less than {IMM2_BOUND}" generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, {hidden_size * intermediate_size} \n" - print("Scale Amount", hidden_size * intermediate_size) generated_code += f"C_SET_SCALE_REG gp{w_actual_register} \n" generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, {intermediate_size} \n" generated_code += f"C_SET_STRIDE_REG gp{w_actual_register} \n" @@ -129,40 +128,65 @@ def ffn_asm( generated_code += f"S_ADDI_INT gp{gate_result_register}, gp{gate_result_register}, {mlen * batch * seq_len} \n" - # Hidden Size SILU Activation Generation - # generated_code += "; SILU Generation \n" - # generated_code += f"S_LD_FP f1, gp0, {const_one_fp_address} \n" - # generated_code += f"S_ADDI_INT gp{intermediate_register}, gp{gate_result_register}, 0 \n" - # for i in range(intermediate_size // vlen): - # # 0 : -x - # generated_code += f"V_SUB_VF gp{intermediate_register}, gp{intermediate_register}, f0, 0, 1 \n" - # # 1 : exp(-x) - # generated_code += f"V_EXP_V gp{intermediate_register}, gp{intermediate_register}, 0 \n" - # # 2 : 1 + exp(-x) - # generated_code += f"V_ADD_VF gp{intermediate_register}, gp{intermediate_register}, f1, 0 \n" - # # 3 : 1 / (1 + exp(-x)) - # generated_code += f"V_RECI_V gp{intermediate_register}, gp{intermediate_register}, 0 \n" - # # 4 : (1 / (1 + exp(-x))) * gate_result - # generated_code += f"V_MUL_VV gp{intermediate_register}, gp{intermediate_register}, gp{gate_result_register}, 0 \n" - # generated_code += f"S_ADDI_INT gp{gate_result_register}, gp{gate_result_register}, {vlen} \n" - - - # generated_code += "; FFN Downsize Linear Generation \n" - # generated_code += f"S_ADDI_INT gp{intermediate_register}, gp{gate_result_register}, 0 \n" - # for weight_row in range (hidden_size // blen): - # if weight_row % (mlen // blen) == 0: - # for weight_col in range (intermediate_size // mlen): - # generated_code += f"H_PREFETCH_M gp{w_actual_register}, gp{w_actual_register}, a{down_weight_hbm_offset_reg}, 1, 0 \n" - # generated_code += f"S_ADDI_INT gp{w_actual_register}, gp{w_actual_register}, {mlen * mlen} \n" - # generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, 0 \n" - # for act_col in range (batch * seq_len // blen): - # generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{intermediate_register}, {act_col * hidden_size * blen} \n" - # for inner_loop_index in range (intermediate_size // mlen): - # generated_code += f"M_MM 0, gp{w_actual_register}, gp{a_actual_register} \n" - # generated_code += f"S_ADDI_INT gp{w_actual_register}, gp{w_actual_register}, {mlen * mlen} \n" - # generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{a_actual_register}, {mlen * blen} \n" - # generated_code += f"S_ADDI_INT gp{intermediate_register}, gp0, {activation_base_address + act_col * batch * seq_len * hidden_size + weight_row * blen} \n" - # generated_code += f"M_MM_WO gp{intermediate_register}, gp0, 0 \n" - # generated_code += f"S_ADDI_INT gp{w_actual_register}, gp{w_actual_register}, {intermediate_register * blen} \n" + # Intermediate Dim SILU Activation Generation, now x in shape of (b, s, intermediate_size) + generated_code += "; SILU Generation \n" + generated_code += f"S_LD_FP f1, gp0, {const_one_fp_address} \n" + # Reset the addr for up and gate result + generated_code += f"S_ADDI_INT gp{up_result_register}, gp0, {batch * seq_len * hidden_size} \n" + generated_code += f"S_ADDI_INT gp{gate_result_register}, gp{up_result_register}, {intermediate_size * batch * seq_len} \n" + generated_code += f"S_ADDI_INT gp{intermediate_register}, gp0, {activation_base_address} \n" + + # Treat the original activation region as the place for scratchpad. + for b in range(batch * seq_len): + for i in range(intermediate_size // vlen): + # 0 : -x + generated_code += f"V_SUB_VF gp{intermediate_register}, gp{up_result_register}, f0, 0, 1 \n" + # 1 : exp(-x) + generated_code += f"V_EXP_V gp{intermediate_register}, gp{intermediate_register}, 0 \n" + # 2 : 1 + exp(-x) + generated_code += f"V_ADD_VF gp{intermediate_register}, gp{intermediate_register}, f1, 0 \n" + # 3 : 1 / (1 + exp(-x)) + generated_code += f"V_RECI_V gp{intermediate_register}, gp{intermediate_register}, 0 \n" + # 4 : (1 / (1 + exp(-x))) * gate_result + generated_code += f"V_MUL_VV gp{intermediate_register}, gp{intermediate_register}, gp{up_result_register}, 0 \n" + # 5: multiply by gate result and store to the up result region + generated_code += f"V_MUL_VV gp{up_result_register}, gp{intermediate_register}, gp{gate_result_register}, 0 \n" + generated_code += f"S_ADDI_INT gp{gate_result_register}, gp{gate_result_register}, {vlen} \n" + generated_code += f"S_ADDI_INT gp{up_result_register}, gp{up_result_register}, {vlen} \n" + generated_code += "; FFN Downsize Linear Generation \n" + # Reset the addr for up and gate result + generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, {hidden_size * intermediate_size} \n" + generated_code += f"C_SET_SCALE_REG gp{w_actual_register} \n" + generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, {hidden_size} \n" + generated_code += f"C_SET_STRIDE_REG gp{w_actual_register} \n" + generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, 0 \n" + # Storing the results to the activation base region + act_result_register = gate_result_register + generated_code += f"S_ADDI_INT gp{act_result_register}, gp0, {activation_base_address} \n" + generated_code += f"S_ADDI_INT gp{up_result_register}, gp0, {batch * seq_len * hidden_size} \n" + for weight_row in range (hidden_size // blen): + if weight_row % (mlen // blen) == 0: + generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, 0 \n" + generated_code += f"S_ADDI_INT gp{w_hbm_offset_register}, gp0, {weight_row * blen} \n" + generated_code += f"S_ADDI_INT gp{intermediate_register}, gp{act_result_register}, 0 \n" + for weight_col in range (intermediate_size // mlen): + generated_code += f"H_PREFETCH_M gp{w_actual_register}, gp{w_hbm_offset_register}, a{down_weight_hbm_offset_reg}, 1, 0 \n" + generated_code += f"S_ADDI_INT gp{w_actual_register}, gp{w_actual_register}, {mlen * mlen} \n" + generated_code += f"S_ADDI_INT gp{w_hbm_offset_register}, gp{w_hbm_offset_register}, {mlen * hidden_size} \n" + generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, 0 \n" + else: + generated_code += f"S_ADDI_INT gp{w_actual_register}, gp0, {(weight_row % (mlen // blen)) * blen} \n" + generated_code += f"S_ADDI_INT gp{intermediate_register}, gp{act_result_register}, {(weight_row % (mlen // blen)) * blen} \n" + for act_col in range (batch * seq_len // blen): + generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{up_result_register}, {act_col * mlen * blen} \n" + generated_code += f"S_ADDI_INT gp{w_temp_register}, gp{w_actual_register}, 0 \n" + for inner_loop_index in range (intermediate_size // mlen): + generated_code += f"M_MM 0, gp{w_actual_register}, gp{a_actual_register} \n" + generated_code += f"S_ADDI_INT gp{w_temp_register}, gp{w_actual_register}, {mlen * mlen} \n" + generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{a_actual_register}, {mlen * batch * seq_len} \n" + generated_code += f"M_MM_WO gp{intermediate_register}, gp0, 0 \n" + generated_code += f"S_ADDI_INT gp{intermediate_register}, gp{intermediate_register}, {blen * mlen} \n" # generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{a_actual_register}, {activation_base_address} \n" + if (weight_row + 1) % (mlen // blen) == 0 and weight_row != intermediate_size // blen - 1: + generated_code += f"S_ADDI_INT gp{act_result_register}, gp{act_result_register}, {mlen * batch * seq_len} \n" return generated_code From 508d58a394b7d88d738f2fc20e6078e5ba6feebf Mon Sep 17 00:00:00 2001 From: GeorgeWu1204 Date: Wed, 26 Nov 2025 19:47:41 +0000 Subject: [PATCH 21/34] preload for batch = 1 --- behavioral_simulator/testbench/dllm1_test.py | 26 ++++++++++---------- compiler/asm_templates/preload_act.py | 9 +++---- 2 files changed, 17 insertions(+), 18 deletions(-) diff --git a/behavioral_simulator/testbench/dllm1_test.py b/behavioral_simulator/testbench/dllm1_test.py index 814bd454..9d10680c 100644 --- a/behavioral_simulator/testbench/dllm1_test.py +++ b/behavioral_simulator/testbench/dllm1_test.py @@ -22,14 +22,14 @@ vocal_size = 2 hidden_size = 64 vlen = 64 - batch_size = 4 - preload_amount = 4 + batch_size = 1 + preload_amount = 1 real_data_ratio = (8*8 + 8) / (8 * 8) hbm_data_width = 64 fp_preload = [0.0, 1e-6] torch.manual_seed(42) - logits = torch.randn(batch_size, hidden_size * vocal_size) + logits = torch.randn(batch_size, hidden_size) weights = logits original_output = logits @@ -53,24 +53,24 @@ # Set the addr offset for mask - gen_assembly_code += preload_addr_reg_asm( - addr_reg_to_set=[1,2], - available_registers=[1,2], - addr_reg_val=[int(align_addr_to_hbm_bandwidth(batch_size * hidden_size * vocal_size * real_data_ratio, hbm_data_width)),int(2*align_addr_to_hbm_bandwidth(batch_size * hidden_size * vocal_size * real_data_ratio, hbm_data_width))] - ) + # gen_assembly_code += preload_addr_reg_asm( + # addr_reg_to_set=[1,2], + # available_registers=[1,2], + # addr_reg_val=[int(align_addr_to_hbm_bandwidth(batch_size * hidden_size * vocal_size * real_data_ratio, hbm_data_width)),int(2*align_addr_to_hbm_bandwidth(batch_size * hidden_size * vocal_size * real_data_ratio, hbm_data_width))] + # ) # Reset the registers - gen_assembly_code += reset_reg_asm( - alive_registers=[1,2,3,4,5,6] - ) + # gen_assembly_code += reset_reg_asm( + # alive_registers=[1,2,3,4,5,6] + # ) # Gen logtis Preload (B,L,V) gen_assembly_code += preload_act_asm( vlen=vlen, preload_len=preload_amount, batch=batch_size, - hidden_size=hidden_size*vocal_size, + hidden_size=hidden_size, alive_registers=[1,2,3], # [a_actual_register, set_stride_register, result_register] act_vram_offset=0, activation_offset_reg=0 @@ -79,5 +79,5 @@ - create_sim_env(input_tensor, weights, gen_assembly_code, golden_result, fp_preload) + create_sim_env(input_tensor, gen_assembly_code, golden_result, fp_preload) build_sim_env(data_size=256, mode="behave_sim", asm="dllm", data=None, specified_data_order = ["logits"]) \ No newline at end of file diff --git a/compiler/asm_templates/preload_act.py b/compiler/asm_templates/preload_act.py index b67cf3cb..b5e5a0d2 100644 --- a/compiler/asm_templates/preload_act.py +++ b/compiler/asm_templates/preload_act.py @@ -36,11 +36,10 @@ def preload_act_asm( load_amount_per_hidden = math.ceil(hidden_size / vlen) if batch == 1: - generated_code += f"S_ADDI_INT gp{set_stride_register}, gp0, {stride_len} \n" - generated_code += f"C_SET_STRIDE_REG gp{set_stride_register} \n" - for i in range((hidden_size + (vlen * preload_len) - 1) // (vlen * preload_len)): - generated_code += f"H_PREFETCH_V gp{a_actual_register}, gp{a_actual_register}, a{activation_offset_reg}, 0, 0 \n" - generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{a_actual_register}, {vlen * preload_len} \n" + for i in range(math.ceil(hidden_size / (vlen * preload_len))): + generated_code += f"H_PREFETCH_V gp{result_register}, gp{a_actual_register}, a{activation_offset_reg}, 0, 0 \n" + generated_code += f"S_ADDI_INT gp{result_register}, gp{result_register}, {vlen * preload_len} \n" + generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{a_actual_register}, {load_amount_per_hidden} \n" else: generated_code += f"S_ADDI_INT gp{set_stride_register}, gp0, {stride_len} \n" generated_code += f"C_SET_STRIDE_REG gp{set_stride_register} \n" From cd52ed540ac5b49a461872c08a4578499a84ebba Mon Sep 17 00:00:00 2001 From: GeorgeWu1204 Date: Mon, 1 Dec 2025 03:18:30 +0000 Subject: [PATCH 22/34] update plena isa spec --- compiler/asm_templates/projection_asm.py | 2 +- compiler/doc/plena_isa_spec.md | 542 +++++++++++++++++++++++ 2 files changed, 543 insertions(+), 1 deletion(-) create mode 100644 compiler/doc/plena_isa_spec.md diff --git a/compiler/asm_templates/projection_asm.py b/compiler/asm_templates/projection_asm.py index a7c89a1f..e534547e 100644 --- a/compiler/asm_templates/projection_asm.py +++ b/compiler/asm_templates/projection_asm.py @@ -14,7 +14,7 @@ def projection_asm( w_base_hbm_offset_reg: int, activation_base_address: int, result_base_address: int, - rope_enabled: bool = True, + rope_enabled: bool = False, rope_hbm_offset_reg: int = 0, rope_on_chip_address: int = 0 ) -> str: diff --git a/compiler/doc/plena_isa_spec.md b/compiler/doc/plena_isa_spec.md new file mode 100644 index 00000000..92e24c2d --- /dev/null +++ b/compiler/doc/plena_isa_spec.md @@ -0,0 +1,542 @@ +# PLENA Instruction Set Architecture (ISA) Specification + +## Register Types + +The PLENA architecture supports four types of registers: + +- **gp_reg** (`gp0` to `gp15`): General-purpose integer registers +- **fp_reg** (`f0` to `f7`): Floating-point registers +- **hbm_addr_reg** (`a0` to `a7`): HBM address registers + +## Instruction Format + +Instructions follow one of the following formats: + +- `opcode, rd, rs1, rs2, rstride, precision` +- `opcode, rd, rs1, rs2, rmask` (V two sources instructions) +- `opcode, rd, rs1, rmask` (V one source instructions) +- `opcode, 0, rs1, rs2` (M_MM, M_TMM, M_MV, M_TMV) +- `opcode, rd, rs1, rs2` +- `opcode, rd, imm` (M_BMM_WO, M_MV_WO) +- `opcode, rd, rs1, imm` + +## Parameters +Refer to `plena_settings.toml` for the detailed parameters. +- **MLEN**: Tile size used in matrix machine +- **BLEN**: Tile size used in systolic array +- **HLEN**: Tile size used in partitioned systolic array +- **VLEN**: Tile size used in vector machine +- **HBM_M_Prefetch_Amount**: Number of MLEN rows fetched from HBM +- **HBM_V_Prefetch_Amount**: Number of VLEN rows fetched from HBM + + +## Matrix (M-Type) Instructions + +### Notation + +| Notation | Description | +|----------|-------------| +| **Matrix[i]** | i-th entry of the Matrix SRAM | +| **Vector[i]** | i-th entry of the Vector SRAM | + +### M_MM + +**Format:** `opcode, rd, rs1, rs2` + +**Operation:** `Systolic Array = Vector_SRAM[gp_reg] @ Matrix_SRAM[gp_reg]` + +**Description:** + +Fetch an (BLEN,MLEN) vector from the Vector SRAM using the address provided by `rs1` and an (MLEN, BLEN) matrix from the Matrix SRAM using the address provided by `rs2`. Then, perform an array of dot products. The result matrix (MLEN, BLEN) is internally accumulated in every PE of the systolic array. + +### M_TMM + +**Format:** `opcode, x, rs1, rs2` + +**Operation:** `Systolic Array = Vector[gp_reg] @ Matrix[gp_reg]^T` + +**Description:** + +Similar to `M_MM`, but transposes the matrix. + +### M_BMM + +**Format:** `opcode, rd, rs1, rs2` + +**Operation:** `Systolic Array = Per Head (Vector_SRAM[gp_reg] @ Matrix_SRAM[gp_reg + gp_reg])` + +`[MLEN // HLEN, MLEN, HLEN] @ [HLEN, MLEN] = [MLEN // HLEN, MLEN, MLEN]` + +**Description:** + +Only take the sliced (HLEN, MLEN) matrix from the Matrix SRAM using the address provided by `gp_reg + gp_reg`, and the vector of shape (MLEN // HLEN, MLEN, HLEN) from the Vector SRAM using the address provided by `gp_reg`. Then, perform an array of dot products. The result matrix [MLEN // HLEN, MLEN, MLEN] is internally accumulated in every PE of the systolic array. + +### M_BTMM + +**Format:** `opcode, rd, rs1, rs2` + +**Operation:** `Systolic Array = Per Head (Vector_SRAM[gp_reg] @ Matrix_SRAM[gp_reg + gp_reg])^T` + +**Description:** + +Similar to `M_BMM`, but the matrix from Matrix SRAM is transposed before the operation. + +### M_BMM_WO + +**Format:** `opcode, rd, imm` + +**Description:** + +Store the accumulated result [MLEN // HLEN, MLEN, MLEN] to the Vector SRAM at the address specified by `gp_reg + imm` with stride `MLEN // HLEN` and precision `Weights` or `KeyValue` depending on the precision of the MXFP data. + +### M_MM_WO + +**Format:** `opcode, rd, rstride, imm` + +**Description:** + +Output the accumulated result (BLEN, BLEN) stored in the first row of the systolic array to the Vector SRAM at the address specified by `gp_reg` with stride `rstride`. + +### M_MV + +**Format:** `opcode, rd, rs1, x` + +**Operation:** `First Row of Sys Array = Vector[gp_reg] @ Matrix[gp_reg]` +**Description:** + +Fetch an (MLEN, MLEN) matrix from the Matrix SRAM using the address provided by `gp_reg`, and an (MLEN, 1) vector from the Vector SRAM using the address provided by `gp_reg`. Then, perform a dot product and store the resulting (MLEN, 1) vector in the **First Row of Sys Array**. + +### M_TMV + +**Format:** `opcode, rd, rs1, x` + +**Operation:** `First Row of Sys Array = Vector[gp_reg] @ Matrix[gp_reg]^T` + +**Description:** + +This instruction is similar to `M_MV`, but transposes the Matrix when fetching from the Matrix SRAM at the address set by `rs2`. + +### M_BMV (TODO: Implement) + +### M_BTMV (TODO: Implement) + +### M_MV_WO + +**Format:** `opcode, rd, imm` + +**Description:** + +Store the accumulated result (MLEN, 1) stored in the first row of the systolic array to the Vector SRAM at the address specified by `gp_reg + imm` + +### M_BMV_WO (TODO: Implement) + +--- + +## Vector (V-Type) Instructions + +### Notation + +| Notation | Description | +|----------|-------------| +| **Vector[i]** | i-th entry of the Vector SRAM | + +`rmask` is a binary flag indicating whether to apply the mask to the result. The mask is set by the `C_SET_V_MASK_REG` instruction. + +### V_ADD_VV + +**Format:** `opcode, rd, rs1, rs2, rmask` + +**Operation:** `Vector[gp_reg] & gp_rmask = (Vector[gp_reg] & gp_reg) + (Vector[gp_reg]) & gp_rmask` + +**Description:** + +Fetch two (MLEN, 1) vectors from the Vector SRAM using the addresses provided by `rs2` and `rs1`, and then perform element-wise addition. Store the resulting vector back to the Vector SRAM at the address provided by `rd`. + +### V_ADD_VF + +**Format:** `opcode, rd, rs1, rs2, rmask` + +**Operation:** `Vector[gp_reg] & gp_rmask = (Vector[gp_reg] & gp_reg) + Broadcast(fp_reg) & gp_reg` + +**Description:** + +Fetch an (MLEN, 1) vector from the Vector SRAM using the address provided by `rs1`, then fetch a single floating-point value from the FP register file using the index provided by `rs2`. Broadcast this value by duplicating it to form an (MLEN, 1) vector, and then perform element-wise addition. Store the resulting vector back to Vector SRAM at the address provided by `rd`. + +### V_SUB_VV + +**Format:** `opcode, rd, rs1, rs2, rmask` + +**Operation:** `Vector[gp_reg] & gp_rmask = (Vector[gp_reg] & gp_reg) - (Vector[gp_reg] & gp_reg)` + +**Description:** + +Similar to `V_ADD_VV`, but performs element-wise subtraction. + +### V_SUB_VF + +**Format:** `opcode, rd, rs1, fp2, rmask, rorder` + +**Operation:** +- If `rorder = Normal`: `Vector[gp_reg] & gp_rmask = (Vector[gp_reg] & gp_reg) - Broadcast(fp_reg) & gp_reg` +- If `rorder = Reverse`: `Vector[gp_reg] & gp_rmask = Broadcast(fp_reg) & gp_reg - (Vector[gp_reg] & gp_reg)` + +**Description:** + +Similar to `V_ADD_VF`, but performs element-wise subtraction. The `rorder` parameter (decoded from `funct1`) controls the order of subtraction: +- `0` (Normal): vector - scalar +- `>0` (Reverse): scalar - vector + +### V_MUL_VV + +**Format:** `opcode, rd, rs1, rs2, rmask` + +**Operation:** `Vector[gp_reg] & gp_rmask = (Vector[gp_reg] & gp_reg) * (Vector[gp_reg] & gp_reg)` + +**Description:** + +Similar to `V_ADD_VV`, but performs element-wise multiplication. + +### V_MUL_VF + +**Format:** `opcode, rd, rs1, fp2, rmask` + +**Operation:** `Vector[gp_reg] & gp_rmask = (Vector[gp_reg] & gp_reg) * Broadcast(fp_reg) & gp_reg` + +**Description:** + +Similar to `V_ADD_VF`, but performs element-wise multiplication. + +### V_EXP_V + +**Format:** `opcode, rd, rs1, x, rmask` + +**Operation:** `Vector[gp_reg] = exp(Vector[gp_reg])` + +**Description:** + +Fetch an (MLEN, 1) vector from the Vector SRAM using the address provided by `rs1`, perform element-wise exponentiation, and store the resulting vector back into the Vector SRAM at the address specified by `rd`. + +### V_RECI_V + +**Format:** `opcode, rd, rs1, x` + +**Operation:** `Vector[gp_reg] = reciprocal(Vector[gp_reg])` + +**Description:** + +Fetch an (MLEN, 1) vector from the Vector SRAM using the address provided by `rs1`, perform element-wise reciprocal, and store the resulting vector back into the Vector SRAM at the address specified by `rd`. + +### V_RED_SUM + +**Format:** `opcode, rd, rs1, 0` + +**Operation:** `fp_reg = sum(Vector[gp_reg], fp_reg)` + +**Description:** + +Fetch an (MLEN, 1) vector from the Vector SRAM using the address provided by `rs1`, and a single floating-point value from the FP register file using the index specified by `rd`. Perform element-wise addition on the combined vector, and store the resulting sum back into the FP register file at the index specified by `rd`. This instruction is designed to facilitate continuous summation along a high-dimension vector. + +### V_RED_MAX + +**Format:** `opcode, rd, rs1, 0` + +**Operation:** `fp_reg = max(Vector[gp_reg], fp_reg)` + +**Description:** + +Similar to `V_RED_SUM` but performs the max value selection operation. + +--- + +## Scalar (S-Type) Instructions + +### Integer Operations + +#### Notation + +| Notation | Description | +|----------|-------------| +| **INT_MEM[i]** | i-th entry of the SRAM within the scalar machine specifically designed for integer operations | + +#### S_ADD_INT + +**Format:** `opcode, rd, rs1, rs2` + +**Operation:** `gp_reg = gp_reg + gp_reg` + +#### S_ADDI_INT + +**Format:** `opcode, rd, rs1, imm2` + +**Operation:** `gp_reg = gp_reg + imm2` + +#### S_SUB_INT + +**Format:** `opcode, rd, rs1, rs2` + +**Operation:** `gp_reg = gp_reg - gp_reg` + +#### S_MUL_INT + +**Format:** `opcode, rd, rs1, rs2` + +**Operation:** `gp_reg = gp_reg * gp_reg` + +#### S_LUI_INT + +**Format:** `opcode, rd, imm` + +**Operation:** `gp_reg = imm << 12` + +**Description:** + +Load upper immediate value into the integer register. + +#### S_LD_INT + +**Format:** `opcode, rd, rs1, imm2` + +**Operation:** `gp_reg = INT_MEM[gp_reg + imm2]` + +#### S_ST_INT + +**Format:** `opcode, rd, rs1, imm2` + +**Operation:** `INT_MEM[gp_reg + imm2] = gp_reg` + +### Floating-Point Operations + +#### Notation + +| Notation | Description | +|----------|-------------| +| **FP_MEM[i]** | i-th entry of the SRAM within the scalar machine specifically designed for floating-point operations | + +#### S_ADD_FP + +**Format:** `opcode, rd, rs1, rs2` + +**Operation:** `fp_reg = fp_reg + fp_reg` + +#### S_SUB_FP + +**Format:** `opcode, rd, rs1, rs2` + +**Operation:** `fp_reg = fp_reg - fp_reg` + +#### S_MAX_FP + +**Format:** `opcode, rd, rs1, rs2` + +**Operation:** `fp_reg = max(fp_reg, fp_reg)` + +#### S_MUL_FP + +**Format:** `opcode, rd, rs1, rs2` + +**Operation:** `fp_reg = fp_reg * fp_reg` + +#### S_EXP_FP + +**Format:** `opcode, rd, rs1, x` + +**Operation:** `fp_reg = exp(fp_reg)` + +#### S_RECI_FP + +**Format:** `opcode, rd, rs1, x` + +**Operation:** `fp_reg = reciprocal(fp_reg)` + +#### S_SQRT_FP + +**Format:** `opcode, rd, rs1, x` + +**Operation:** `fp_reg = sqrt(fp_reg)` + +#### S_LD_FP + +**Format:** `opcode, rd, rs1, imm2` + +**Operation:** `fp_reg = FP_MEM[gp_reg + imm2]` + +#### S_ST_FP + +**Format:** `opcode, rd, rs1, imm2` + +**Operation:** `FP_MEM[gp_reg + imm2] = fp_reg` + +#### S_MAP_V_FP + +**Format:** `opcode, rd, rs1, imm2` + +**Operation:** `Vector[gp_reg :+ VLEN] = FP_MEM[gp_reg + imm2 :+ VLEN]` + +**Description:** + +Copy a vector of length VLEN from FP_MEM to Vector SRAM. + +--- + +## Memory (H-Type) Instructions + +### Notation + +| Notation | Description | +|----------|-------------| +| **Matrix[i]** | The i-th entry of the Matrix SRAM | +| **Vector[i]** | The i-th entry of the Vector SRAM | +| **HBM[i]** | The i-th entry of the HBM | + +### H_PREFETCH_M + +**Format:** `opcode, rd, rs1, rs2, rstride, precision` + +**Operation:** `Matrix[gp_reg] = HBM[gp_reg + hbm_addr_reg_]` + +**Description:** + +Prefetch a matrix of size **HBM_M_Prefetch_Amount × MLEN** from the HBM to the Matrix SRAM, with a stride width specified by **hbm_stride_reg[rstride]**. + +The `precision` field (decoded from `funct1`) determines the data precision: +- `0` (Weights): High precision weights +- `>0` (KeyValue): Lower precision key-value data + +### H_PREFETCH_V + +**Format:** `opcode, rd, rs1, rs2, rstride, precision` + +**Operation:** `Vector[gp_reg] = HBM[gp_reg + hbm_addr_reg_]` + +**Description:** + +Prefetch a matrix of size **HBM_V_Prefetch_Amount × VLEN** from the HBM to the Vector SRAM, with a stride width specified by **hbm_stride_reg[rstride]**. + +The `precision` field (decoded from `funct1`) determines the data precision: +- `0` (Activation): High precision activation data +- `>0` (KeyValue): Lower precision key-value data + +### H_STORE_V + +**Format:** `opcode, rd, rs1, rs2, rstride, precision` + +**Operation:** `HBM[gp_reg + hbm_addr_reg_] = Vector[gp_reg]` + +**Description:** + +Store a matrix of size **HBM_V_Writeback_Amount × VLEN** from Vector SRAM to HBM, with a stride width specified by **hbm_stride_reg[rstride]**. + +The `precision` field (decoded from `funct1`) determines the data precision: +- `0` (Activation): High precision activation data +- `>0` (KeyValue): Lower precision key-value data + +--- + +## Control and Status Register (C-Type) Instructions + +### C_SET_ADDR_REG + +**Format:** `opcode, rd, rs1, rs2` + +**Operation:** `hbm_addr_reg_ = {gp_reg, gp_reg}` + +**Description:** + +This instruction is used to set the value of `hbm_addr_reg[rd]`, assuming it has double the bit width of `fix_reg`, by concatenating two `fix_reg` entries and storing the result in `hbm_addr_reg[rd]`. The concatenation order is `{rs2, rs1}`. + +### C_SET_SCALE_REG + +**Format:** `opcode, rd` + +**Operation:** `SCALE_OFFSET = rd` + +**Description:** + +This instruction is used to set the scale offset. The blocks and scales of the MXFP data are stored separately in HBM for memory alignment purposes. Their distance is set by the scale offset. This value differs depending on the precision of the MXFP and the data size. For example, for Q(B, S, H, D), the scales are stored after the blocks, and the offset is `B × S × H × D × (EXP_WIDTH + MANT_WIDTH + 1) / 8`. + +### C_SET_STRIDE_REG + +**Format:** `opcode, rd` + +**Operation:** `STRIDE_SIZE = rd` + +**Description:** + +This instruction is used to set the stride size for the prefetch instructions. + +### C_SET_V_MASK_REG + +**Format:** `opcode, rd` + +**Operation:** `V_MASK = rd` + +**Description:** + +This instruction is used to set the vector mask register for masked vector operations. + +### C_BREAK + +**Format:** `opcode, 0, 0, 0` + +**Operation:** Breakpoint exception + +**Description:** + +Triggers a breakpoint exception, typically used for debugging purposes. + +--- + +## Instruction Encoding Summary + +### Opcode Map + +| Opcode | Instruction | Type | +|--------|-------------|------| +| 0x00 | Invalid | - | +| 0x01 | M_MM | M-Type | +| 0x02 | M_TMM | M-Type | +| 0x03 | M_BMM | M-Type | +| 0x04 | M_BTMM | M-Type | +| 0x05 | M_BMM_WO | M-Type | +| 0x06 | M_MM_WO | M-Type | +| 0x07 | M_MV | M-Type | +| 0x08 | M_TMV | M-Type | +| 0x09 | M_BMV | M-Type | +| 0x0A | M_BTMV | M-Type | +| 0x0B | M_MV_WO | M-Type | +| 0x0C | M_BMV_WO | M-Type | +| 0x0D | V_ADD_VV | V-Type | +| 0x0E | V_ADD_VF | V-Type | +| 0x0F | V_SUB_VV | V-Type | +| 0x10 | V_SUB_VF | V-Type | +| 0x11 | V_MUL_VV | V-Type | +| 0x12 | V_MUL_VF | V-Type | +| 0x13 | V_EXP_V | V-Type | +| 0x14 | V_RECI_V | V-Type | +| 0x15 | V_RED_SUM | V-Type | +| 0x16 | V_RED_MAX | V-Type | +| 0x17 | S_ADD_FP | S-Type | +| 0x18 | S_SUB_FP | S-Type | +| 0x19 | S_MAX_FP | S-Type | +| 0x1A | S_MUL_FP | S-Type | +| 0x1B | S_EXP_FP | S-Type | +| 0x1C | S_RECI_FP | S-Type | +| 0x1D | S_SQRT_FP | S-Type | +| 0x1E | S_LD_FP | S-Type | +| 0x1F | S_ST_FP | S-Type | +| 0x20 | S_MAP_V_FP | S-Type | +| 0x21 | S_ADD_INT | S-Type | +| 0x22 | S_ADDI_INT | S-Type | +| 0x23 | S_SUB_INT | S-Type | +| 0x24 | S_MUL_INT | S-Type | +| 0x25 | S_LUI_INT | S-Type | +| 0x26 | S_LD_INT | S-Type | +| 0x27 | S_ST_INT | S-Type | +| 0x28 | H_PREFETCH_M | H-Type | +| 0x29 | H_PREFETCH_V | H-Type | +| 0x2A | H_STORE_V | H-Type | +| 0x2B | C_SET_ADDR_REG | C-Type | +| 0x2C | C_SET_SCALE_REG | C-Type | +| 0x2D | C_SET_STRIDE_REG | C-Type | +| 0x2E | C_SET_V_MASK_REG | C-Type | +| 0x2F | C_BREAK | C-Type | \ No newline at end of file From ca3c7cec142cf9eca53ce8800508ff4ffcb221f4 Mon Sep 17 00:00:00 2001 From: GeorgeWu1204 Date: Mon, 1 Dec 2025 20:38:52 +0000 Subject: [PATCH 23/34] hbm int fetch and extend operand function field for prefetch_v --- behavioral_simulator/src/main.rs | 148 ++++++++++++++++++++------ behavioral_simulator/src/op.rs | 9 ++ tools/assembler/assembly_to_binary.py | 11 +- tools/assembler/parser.py | 43 +++++++- 4 files changed, 176 insertions(+), 35 deletions(-) diff --git a/behavioral_simulator/src/main.rs b/behavioral_simulator/src/main.rs index da2172ef..1db81205 100644 --- a/behavioral_simulator/src/main.rs +++ b/behavioral_simulator/src/main.rs @@ -14,7 +14,7 @@ use futures::StreamExt; use futures::stream::FuturesUnordered; use half::f16; use memory::MemoryModel; -use quantize::{MxDataType, QuantTensor}; +use quantize::{DataType, MxDataType, QuantTensor}; use runtime::{Duration, Executor, Instant}; use tch::{IndexOp, Tensor}; @@ -908,7 +908,7 @@ impl Accelerator { /// - stride: Byte offset between consecutive loads /// - load_dim: Number of elements to load per iteration /// - load_amount: Number of strided loads to perform - fn transfer_from_hbm( + fn transfer_mx_from_hbm( &mut self, index: u64, scale_index: u64, @@ -1124,6 +1124,86 @@ impl Accelerator { receiver } + fn transfer_normal_from_hbm( + &mut self, + index: u64, + hbm_type: DataType, + sram_type: MxDataType, + rstride: u8, + load_dim: u32, + load_amount: u32, + ) -> Receiver { + // Simplified version: only loads elements from HBM based on index and stride + // No scale operations, no MX data type complexity + + let (sender, receiver) = oneshot::channel(); + + let hbm_clone = self.hbm.clone(); + let stride = if rstride == 1 { + self.reg_file.stride + } else { + load_dim + }; + + Executor::current().spawn(async move { + let element_bits = hbm_type.size_in_bits(); + 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)); + let len_in_bytes_per_load = len_in_bits_per_load / 8; + + // Total elements/bytes for all loads: + let total_elements = (load_dim * load_amount) as usize; + let total_bytes = (len_in_bytes_per_load * load_amount) as usize; + + let mut bytes = vec![0u8; total_bytes]; + let hbm_clone = &hbm_clone; + + enum ChunkType { + Element(usize, [u8; 64], usize), // (offset, data, size) + } + let mut futures = + FuturesUnordered:: + Send>>>::new(); + + // Gather all HBM reads + for load_idx in 0..load_amount { + let element_addr = index + (load_idx * stride) as u64; + let byte_offset = (load_idx * 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) + })); + } + } + + // Collect all HBM reads + while let Some(chunk_result) = futures.next().await { + match chunk_result { + ChunkType::Element(offset, data, size) => { + bytes[offset..offset + size].copy_from_slice(&data[..size]); + } + } + } + + // Convert bytes to f32 and quantize + let mut vec = vec![0f32; total_elements]; + hbm_type.convert_bytes_to_f32_vec(&bytes, &mut vec); + + let tensor = tch::Tensor::from_slice(&vec); + let _ = sender.send(QuantTensor::quantize(tensor, sram_type)); + }); + + receiver + } + async fn do_ops(&mut self, ops: &[op::Opcode]) { for op in ops { println!("execute op = {:?}", op); @@ -1459,7 +1539,7 @@ impl Accelerator { / (elem.size_in_bits() as u32 * block / scale.size_in_bits() as u32) } // Element addr shifted by (element to scale ratio) }; - let xfer = self.transfer_from_hbm( + let xfer = self.transfer_mx_from_hbm( addr + offset as u64, addr + self.reg_file.scale as u64 + scale as u64, dtype, @@ -1485,39 +1565,45 @@ impl Accelerator { rs2, rstride, precision, + datatype, } => { // TODO: rstride support to be added let offset = self.reg_file.gp_reg[rs1 as usize]; let addr = self.reg_file.hbm_addr_reg[rs2 as usize]; + if matches! (datatype, op::HBM_LOAD_TYPE::MX) { + let dtype = match precision { + op::VectorPrecision::Activation => *VECTOR_ACTIVATION_TYPE, + op::VectorPrecision::KeyValue => *VECTOR_KV_TYPE, + }; + + let scale = match dtype { + MxDataType::Plain(_) => 0, + MxDataType::Mx { elem, scale, block } => { + offset + / (elem.size_in_bits() as u32 * block / scale.size_in_bits() as u32) + } + }; + let xfer = self.transfer_mx_from_hbm( + addr + offset as u64, + addr + self.reg_file.scale as u64 + scale as u64, + dtype, + self.v_machine.vram.ty, + rstride, + *VLEN, + *PREFETCH_V_AMOUNT, + 1, + ); + + let dest = self.reg_file.gp_reg[rd as usize]; + self.v_machine + .vram + .continous_write_delayed(dest, *PREFETCH_V_AMOUNT, xfer) + .await; + } else { + + } - let dtype = match precision { - op::VectorPrecision::Activation => *VECTOR_ACTIVATION_TYPE, - op::VectorPrecision::KeyValue => *VECTOR_KV_TYPE, - }; - - let scale = match dtype { - MxDataType::Plain(_) => 0, - MxDataType::Mx { elem, scale, block } => { - offset - / (elem.size_in_bits() as u32 * block / scale.size_in_bits() as u32) - } - }; - let xfer = self.transfer_from_hbm( - addr + offset as u64, - addr + self.reg_file.scale as u64 + scale as u64, - dtype, - self.v_machine.vram.ty, - rstride, - *VLEN, - *PREFETCH_V_AMOUNT, - 1, - ); - - let dest = self.reg_file.gp_reg[rd as usize]; - self.v_machine - .vram - .continous_write_delayed(dest, *PREFETCH_V_AMOUNT, xfer) - .await; + } op::Opcode::H_STORE_V { rd, diff --git a/behavioral_simulator/src/op.rs b/behavioral_simulator/src/op.rs index a02e693e..2cc86a1a 100644 --- a/behavioral_simulator/src/op.rs +++ b/behavioral_simulator/src/op.rs @@ -16,6 +16,12 @@ pub enum VectorOrder { Reverse, } +#[derive(Debug, Clone, Copy)] +pub enum HBM_LOAD_TYPE { + MX, + Normal, +} + #[allow(non_camel_case_types)] #[derive(Debug)] @@ -232,6 +238,7 @@ pub enum Opcode { rs2: u8, rstride: u8, precision: VectorPrecision, + datatype: HBM_LOAD_TYPE, }, H_STORE_V { rd: u8, @@ -304,6 +311,7 @@ impl Opcode { let rs2 = ((instr >> (OPCODE_WIDTH + OPERAND_WIDTH * 2)) & mask(OPERAND_WIDTH)) as u8; let rs3 = ((instr >> (OPCODE_WIDTH + OPERAND_WIDTH * 3)) & mask(OPERAND_WIDTH)) as u8; let funct1 = ((instr >> (OPCODE_WIDTH + OPERAND_WIDTH * 4)) & mask(OPERAND_WIDTH)) as u8; + let funct2 = ((instr >> (OPCODE_WIDTH + OPERAND_WIDTH * 5)) & mask(OPERAND_WIDTH)) as u8; let imm = ((instr >> (OPCODE_WIDTH + OPERAND_WIDTH)) & mask(IMM_WIDTH)) as u32; let imm2 = ((instr >> (OPCODE_WIDTH + OPERAND_WIDTH * 2)) & mask(IMM_2_WIDTH)) as u32; @@ -370,6 +378,7 @@ impl Opcode { rs2, rstride: rs3, precision: Self::vector_precision_from(funct1), + datatype: Self::hbm_load_type_from(funct2), }, // 0x2A => Self::H_PREFETCH_V { rd, rs1, rs2, rstride: rs3, precision: VectorPrecision::KeyValue }, 0x2A => Self::H_STORE_V { diff --git a/tools/assembler/assembly_to_binary.py b/tools/assembler/assembly_to_binary.py index fd4c1b03..7eeec2fb 100644 --- a/tools/assembler/assembly_to_binary.py +++ b/tools/assembler/assembly_to_binary.py @@ -69,7 +69,16 @@ def _convert_to_binary(self, instruction): (rd << opw) + opcode ) - elif instruction.opcode in [ "H_PREFETCH_M", "H_PREFETCH_V", "H_STORE_V", "V_SUB_VF"]: + elif instruction.opcode in ["H_PREFETCH_M", "V_SUB_VF"]: + binary_instruction = ( + (funct1 << (opw + 4 * ow)) + + (rstride << (opw + 3 * ow)) + + (rs2 << (opw + 2 * ow)) + + (rs1 << (opw + ow)) + + (rd << opw) + + opcode + ) + elif instruction.opcode in ["H_PREFETCH_V", "H_STORE_V"]: binary_instruction = ( (funct1 << (opw + 4 * ow)) + (rstride << (opw + 3 * ow)) + diff --git a/tools/assembler/parser.py b/tools/assembler/parser.py index 2cd772b0..e81d6d55 100644 --- a/tools/assembler/parser.py +++ b/tools/assembler/parser.py @@ -62,7 +62,7 @@ def load_isa_settings(file_path: str) -> dict: class Instruction: - def __init__(self, opcode: str, rd: str, rs1: Optional[str], rs2: Optional[str], rstride: Optional[str], funct1: Optional[int], imm: Optional[int] = None, rflag: Optional[int] = None): + def __init__(self, opcode: str, rd: str, rs1: Optional[str], rs2: Optional[str], rstride: Optional[str], funct1: Optional[int], funct2: Optional[int], imm: Optional[int] = None, rflag: Optional[int] = None): self.opcode = opcode self.rd = rd @@ -70,11 +70,12 @@ def __init__(self, opcode: str, rd: str, rs1: Optional[str], rs2: Optional[str], self.rs2 = rs2 self.rstride = rstride self.funct1 = funct1 + self.funct2 = funct2 self.imm = imm self.rmask = rstride def __repr__(self): - return f"Instruction(opcode='{self.opcode}', rd='{self.rd}', rs1='{self.rs1}', rs2='{self.rs2}', rstride = '{self.rstride}', funct1={self.funct1}, imm={self.imm}, rflag={self.rflag})" + return f"Instruction(opcode='{self.opcode}', rd='{self.rd}', rs1='{self.rs1}', rs2='{self.rs2}', rstride = '{self.rstride}', funct1={self.funct1}, funct2={self.funct2}, imm={self.imm}, rflag={self.rflag})" def parse_asm_file(file_path: str) -> List[Instruction]: @@ -117,6 +118,7 @@ def parse_asm_file(file_path: str) -> List[Instruction]: rs2 = None rstride = None funct1 = None + funct2 = None imm = None # Helper to parse a register or int operand @@ -221,9 +223,44 @@ def parse_reg_or_int(operand): funct1 = int(funct1_raw) except ValueError: funct1 = funct1_raw # fallback, if not int, keep as string + elif len(operands) == 6: + operand_0, operand_1, operand_2, operand_3, operand_4, operand_5 = operands + rd = parse_reg_or_int(operand_0) + if operand_1.strip().startswith(('gp','f','a')): + rs1 = parse_reg_or_int(operand_1) + else: + try: + imm = int(operand_1) + except ValueError: + imm = None + if operand_2.strip().startswith(('gp','f','a')): + rs2 = parse_reg_or_int(operand_2) + else: + try: + imm = int(operand_2) + except ValueError: + pass + try: + rstride = int(operand_3) + except ValueError: + rstride = None + funct1_raw = operand_4.strip() + if funct1_raw.endswith(';'): + funct1_raw = funct1_raw[:-1] + try: + funct1 = int(funct1_raw) + except ValueError: + funct1 = funct1_raw # fallback, if not int, keep as string + funct2_raw = operand_5.strip() + if funct2_raw.endswith(';'): + funct2_raw = funct2_raw[:-1] + try: + funct2 = int(funct2_raw) + except ValueError: + funct2 = funct2_raw # fallback, if not int, keep as string - instructions.append(Instruction(opcode, rd, rs1, rs2, rstride, funct1, imm)) + instructions.append(Instruction(opcode, rd, rs1, rs2, rstride, funct1, funct2, imm)) return instructions From 30ecd71eedb812ff8f9f879b29a1585bbcd73094 Mon Sep 17 00:00:00 2001 From: GeorgeWu1204 Date: Tue, 2 Dec 2025 01:27:59 +0000 Subject: [PATCH 24/34] replace original fp vect based vsram with binary based one --- behavioral_simulator/Cargo.lock | 10 + behavioral_simulator/Cargo.toml | 2 + .../lib/vector_sram/Cargo.toml | 10 + .../lib/vector_sram/src/lib.rs | 434 ++++++++++++++++++ behavioral_simulator/src/main.rs | 155 +------ behavioral_simulator/src/op.rs | 10 + compiler/asm_templates/load_int.py | 0 compiler/asm_templates/preload_act.py | 4 +- 8 files changed, 477 insertions(+), 148 deletions(-) create mode 100644 behavioral_simulator/lib/vector_sram/Cargo.toml create mode 100644 behavioral_simulator/lib/vector_sram/src/lib.rs create mode 100644 compiler/asm_templates/load_int.py diff --git a/behavioral_simulator/Cargo.lock b/behavioral_simulator/Cargo.lock index 5053feac..55b0f1c3 100644 --- a/behavioral_simulator/Cargo.lock +++ b/behavioral_simulator/Cargo.lock @@ -146,6 +146,7 @@ dependencies = [ "tch", "tokio", "toml", + "vector_sram", ] [[package]] @@ -1437,6 +1438,15 @@ version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "06abde3611657adf66d383f00b093d7faecc7fa57071cce2578660c9f1010821" +[[package]] +name = "vector_sram" +version = "0.1.0" +dependencies = [ + "quantize", + "tch", + "tokio", +] + [[package]] name = "version_check" version = "0.9.5" diff --git a/behavioral_simulator/Cargo.toml b/behavioral_simulator/Cargo.toml index 792ad426..d9d0d3d5 100644 --- a/behavioral_simulator/Cargo.toml +++ b/behavioral_simulator/Cargo.toml @@ -9,6 +9,7 @@ memory = { path = "lib/memory" } ramulator = { path = "lib/ramulator" } buddy = { path = "lib/buddy" } quantize = { path = "lib/quantize" } +vector_sram = { path = "lib/vector_sram" } tokio = { version = "1.45.1", features = ["macros", "rt", "rt-multi-thread", "sync"] } anyhow = "1" async-trait = "0.1.88" @@ -26,6 +27,7 @@ members = [ "lib/ramulator", "lib/buddy", "lib/quantize", + "lib/vector_sram", ] [profile.release] diff --git a/behavioral_simulator/lib/vector_sram/Cargo.toml b/behavioral_simulator/lib/vector_sram/Cargo.toml new file mode 100644 index 00000000..ae74dc10 --- /dev/null +++ b/behavioral_simulator/lib/vector_sram/Cargo.toml @@ -0,0 +1,10 @@ +[package] +name = "vector_sram" +version = "0.1.0" +edition = "2021" + +[dependencies] +tokio = { version = "1.45.1", features = ["macros", "rt", "rt-multi-thread", "sync"] } +quantize = { path = "../quantize" } +tch = { version = "0.20.0", features = ["download-libtorch"] } + diff --git a/behavioral_simulator/lib/vector_sram/src/lib.rs b/behavioral_simulator/lib/vector_sram/src/lib.rs new file mode 100644 index 00000000..e7f93d9c --- /dev/null +++ b/behavioral_simulator/lib/vector_sram/src/lib.rs @@ -0,0 +1,434 @@ +use quantize::{MxDataType, QuantTensor}; +use tokio::sync::Mutex; +use tokio::sync::oneshot::Receiver; +use tch::{IndexOp, Tensor}; + +/// Vector SRAM that stores data in binary format with row-based storage. +/// +/// The SRAM is organized as rows, where each row has a width of `vlen * element_size_in_bytes`. +/// During read/write operations, data is clipped to VLEN-sized vectors and converted +/// between storage precision and access precision. +pub struct VectorSram { + /// Vector length (VLEN) - determines the size of each vector operation + vlen: u32, + /// Number of rows in the SRAM + depth: usize, + /// Storage precision - how data is stored internally + storage_precision: MxDataType, + /// Access precision - how data is read/written (defaults to storage_precision if None) + access_precision: Option, + /// Raw binary storage: each row is stored as bytes + /// Row width = vlen * element_size_in_bytes(storage_precision) + rows: Vec>, +} + +/// Represents a row of data, either ready or pending from a delayed write +enum RowData { + Ready(Vec), + Pending(Receiver), +} + +impl VectorSram { + /// Create a new Vector SRAM with given vector length, depth, and storage precision. + /// + /// # Arguments + /// * `vlen` - Vector length (VLEN) + /// * `depth` - Number of rows in the SRAM + /// * `storage_precision` - Precision used for internal storage + /// * `access_precision` - Optional precision for read/write operations (defaults to storage_precision) + pub fn new( + vlen: u32, + depth: usize, + storage_precision: MxDataType, + access_precision: Option, + ) -> Self { + let element_ty = storage_precision.element_type(); + let element_size = element_ty.size_in_bits() as usize / 8; + let row_width = vlen as usize * element_size; + + // Handle Mx format: need to account for scale bytes + let row_width = match storage_precision { + MxDataType::Plain(_) => row_width, + MxDataType::Mx { block, scale, .. } => { + let num_blocks = (vlen as usize + block as usize - 1) / block as usize; + let scale_bytes = scale.size_in_bits() as usize / 8; + row_width + num_blocks * scale_bytes + } + }; + + let rows = (0..depth) + .map(|_| { + Mutex::new(RowData::Ready(vec![0u8; row_width])) + }) + .collect(); + + Self { + vlen, + depth, + storage_precision, + access_precision, + rows, + } + } + + /// Get the vector length (VLEN) + pub fn tile_size(&self) -> u32 { + self.vlen + } + + /// Get the storage precision type + pub fn ty(&self) -> MxDataType { + self.access_precision.unwrap_or(self.storage_precision) + } + + /// Get the size of the SRAM in bytes + pub fn size_in_bytes(&self) -> usize { + let element_ty = self.storage_precision.element_type(); + let element_size = element_ty.size_in_bits() as usize / 8; + let base_row_width = self.vlen as usize * element_size; + + let row_width = match self.storage_precision { + MxDataType::Plain(_) => base_row_width, + MxDataType::Mx { block, scale, .. } => { + let num_blocks = (self.vlen as usize + block as usize - 1) / block as usize; + let scale_bytes = scale.size_in_bits() as usize / 8; + base_row_width + num_blocks * scale_bytes + } + }; + + row_width * self.depth + } + + /// Read a vector from the SRAM at the given address. + /// + /// The address must be a multiple of vlen (in element units). + /// Data is read from storage precision and converted to access precision. + pub async fn read(&self, addr: u32) -> QuantTensor { + let row_idx = self.addr_to_row_idx(addr); + assert!(row_idx < self.depth, "Address out of bounds"); + + let mut guard = self.rows[row_idx].lock().await; + + // Handle pending writes + if let RowData::Pending(ref mut receiver) = *guard { + let tensor = receiver.await.unwrap(); + let row_bytes = self.quant_tensor_to_bytes(&tensor, self.storage_precision); + *guard = RowData::Ready(row_bytes); + } + + // Read the row data + let row_bytes = match &*guard { + RowData::Ready(bytes) => bytes.clone(), + RowData::Pending(_) => unreachable!(), + }; + + // Convert from storage precision to access precision + let access_precision = self.access_precision.unwrap_or(self.storage_precision); + let tensor = self.bytes_to_quant_tensor(&row_bytes, self.storage_precision, self.vlen); + + // Convert to access precision if different + if access_precision != self.storage_precision { + let tensor_f32 = tensor.as_tensor(); + QuantTensor::quantize(tensor_f32.shallow_clone(), access_precision) + } else { + tensor + } + } + + /// Write a vector to the SRAM at the given address. + /// + /// The address must be a multiple of vlen (in element units). + /// Data is converted from access precision to storage precision before writing. + pub async fn write(&self, addr: u32, tensor: QuantTensor) { + let row_idx = self.addr_to_row_idx(addr); + assert!(row_idx < self.depth, "Address out of bounds"); + + // Convert to storage precision if needed + let storage_tensor = if tensor.data_type() != self.storage_precision { + let tensor_f32 = tensor.as_tensor(); + QuantTensor::quantize(tensor_f32.shallow_clone(), self.storage_precision) + } else { + tensor + }; + + // Clip to VLEN + let clipped = self.clip_to_vlen(&storage_tensor); + + // Convert to bytes + let row_bytes = self.quant_tensor_to_bytes(&clipped, self.storage_precision); + + *self.rows[row_idx].lock().await = RowData::Ready(row_bytes); + } + + /// Write a vector with delayed delivery (from a channel). + pub async fn write_delayed(&self, addr: u32, tensor: Receiver) { + let row_idx = self.addr_to_row_idx(addr); + assert!(row_idx < self.depth, "Address out of bounds"); + + *self.rows[row_idx].lock().await = RowData::Pending(tensor); + } + + /// Continuous write delayed - writes multiple rows from a single tensor. + pub async fn continous_write_delayed( + &self, + addr: u32, + write_amount: u32, + tensor: Receiver, + ) { + let start_row_idx = self.addr_to_row_idx(addr); + + // Await the tensor from the channel + let tensor = tensor.await.unwrap(); + let tensor_data = tensor.as_tensor(); + let total_elements = tensor_data.size1().unwrap() as i64; + + // Convert to storage precision if needed + let storage_tensor = if tensor.data_type() != self.storage_precision { + QuantTensor::quantize(tensor_data.shallow_clone(), self.storage_precision) + } else { + tensor + }; + + let chunk_size = self.vlen as i64; + let num_chunks = write_amount.min(((total_elements + chunk_size - 1) / chunk_size) as u32); + + for i in 0..num_chunks { + let row_idx = start_row_idx + i as usize; + if row_idx >= self.depth { + break; + } + + let start = i as i64 * chunk_size; + let end = (start + chunk_size).min(total_elements); + let chunk = storage_tensor + .as_tensor() + .narrow(0, start, end - start) + .shallow_clone(); + + // Pad to VLEN if needed + let padded = if (end - start) < chunk_size { + let mut padded_tensor = Tensor::zeros( + [chunk_size], + (tch::Kind::Float, tch::Device::Cpu), + ); + padded_tensor + .i(0..(end - start)) + .copy_(&chunk); + padded_tensor + } else { + chunk + }; + + let chunk_qt = QuantTensor::quantize(padded, self.storage_precision); + let row_bytes = self.quant_tensor_to_bytes(&chunk_qt, self.storage_precision); + *self.rows[row_idx].lock().await = RowData::Ready(row_bytes); + } + } + + /// Load data from bytes into the SRAM. + /// + /// This is used for preloading the SRAM with test data. + pub async fn load_from_bytes(&self, bytes: &[u8]) { + let element_ty = self.storage_precision.element_type(); + let element_bits = element_ty.size_in_bits(); + let bytes_per_element = (element_bits / 8) as usize; + + // Calculate row width + let base_row_width = self.vlen as usize * bytes_per_element; + let _row_width = match self.storage_precision { + MxDataType::Plain(_) => base_row_width, + MxDataType::Mx { block, scale, .. } => { + let num_blocks = (self.vlen as usize + block as usize - 1) / block as usize; + let scale_bytes = scale.size_in_bits() as usize / 8; + base_row_width + num_blocks * scale_bytes + } + }; + + let total_elements = bytes.len() / bytes_per_element; + let num_rows = (total_elements + self.vlen as usize - 1) / self.vlen as usize; + + for row_idx in 0..num_rows.min(self.depth) { + let start_element = row_idx * self.vlen as usize; + let end_element = (start_element + self.vlen as usize).min(total_elements); + let elements_in_row = end_element - start_element; + + let start_byte = start_element * bytes_per_element; + let end_byte = end_element * bytes_per_element; + + // Convert bytes to f32 values + let mut vec = vec![0f32; elements_in_row]; + element_ty.convert_bytes_to_f32_vec(&bytes[start_byte..end_byte], &mut vec); + + // Pad with zeros if needed + if elements_in_row < self.vlen as usize { + vec.resize(self.vlen as usize, 0.0f32); + } + + // Create QuantTensor and convert to bytes + let tensor = Tensor::from_slice(&vec); + let quant_tensor = QuantTensor::quantize(tensor, self.storage_precision); + let row_bytes = self.quant_tensor_to_bytes(&quant_tensor, self.storage_precision); + *self.rows[row_idx].lock().await = RowData::Ready(row_bytes); + } + } + + /// Dump the entire SRAM content as bytes. + /// + /// This returns the raw binary representation of all stored data. + pub async fn as_bytes(&self) -> Vec { + let mut result = Vec::new(); + + for row_mutex in &self.rows { + let mut guard = row_mutex.lock().await; + + // Handle pending writes + if let RowData::Pending(ref mut receiver) = *guard { + let tensor = receiver.await.unwrap(); + let row_bytes = self.quant_tensor_to_bytes(&tensor, self.storage_precision); + *guard = RowData::Ready(row_bytes); + } + + // Read the row data + let row_bytes = match &*guard { + RowData::Ready(bytes) => bytes.clone(), + RowData::Pending(_) => unreachable!(), + }; + + result.extend_from_slice(&row_bytes); + } + + result + } + + // Helper methods + + /// Convert address (in element units) to row index + fn addr_to_row_idx(&self, addr: u32) -> usize { + assert!(addr % self.vlen == 0, "Address must be multiple of vlen"); + (addr / self.vlen) as usize + } + + /// Clip a tensor to VLEN size + fn clip_to_vlen(&self, tensor: &QuantTensor) -> QuantTensor { + let tensor_data = tensor.as_tensor(); + let len = tensor_data.size1().unwrap() as i64; + + if len <= self.vlen as i64 { + tensor.clone() + } else { + let clipped = tensor_data.narrow(0, 0, self.vlen as i64); + QuantTensor::quantize(clipped, tensor.data_type()) + } + } + + /// Convert QuantTensor to bytes according to the given precision + fn quant_tensor_to_bytes(&self, tensor: &QuantTensor, precision: MxDataType) -> Vec { + let tensor_data = tensor.as_tensor(); + let len = tensor_data.size1().unwrap() as usize; + let f32_slice = unsafe { + core::slice::from_raw_parts(tensor_data.data_ptr() as *const f32, len) + }; + + match precision { + MxDataType::Plain(elem_ty) => { + let total_bits = len * elem_ty.size_in_bits() as usize; + let bytes_needed = (total_bits + 7) / 8; + let mut bytes = vec![0u8; bytes_needed]; + elem_ty.bytes_from_f32(f32_slice, &mut bytes); + bytes + } + MxDataType::Mx { elem, scale, block } => { + // Serialize elements + let total_bits = len * elem.size_in_bits() as usize; + let elem_bytes_needed = (total_bits + 7) / 8; + let mut elem_bytes = vec![0u8; elem_bytes_needed]; + elem.bytes_from_f32(f32_slice, &mut elem_bytes); + + // Serialize scales (one per block) + let num_blocks = (len + block as usize - 1) / block as usize; + // For now, use scale of 1.0 for each block + // TODO: implement proper scale calculation + let scale_vec = vec![1.0f32; num_blocks]; + let scale_bytes_needed = num_blocks * (scale.size_in_bits() as usize / 8); + let mut scale_bytes = vec![0u8; scale_bytes_needed]; + scale.bytes_from_f32(&scale_vec, &mut scale_bytes); + + // Combine: elements first, then scales + let mut result = elem_bytes; + result.extend_from_slice(&scale_bytes); + result + } + } + } + + /// Convert bytes to QuantTensor according to the given precision + fn bytes_to_quant_tensor( + &self, + bytes: &[u8], + precision: MxDataType, + expected_len: u32, + ) -> QuantTensor { + match precision { + MxDataType::Plain(elem_ty) => { + let bytes_per_element = elem_ty.size_in_bits() as usize / 8; + let num_elements = bytes.len() / bytes_per_element; + let actual_len = num_elements.min(expected_len as usize); + + let mut vec = vec![0f32; actual_len]; + elem_ty.convert_bytes_to_f32_vec(&bytes[..actual_len * bytes_per_element], &mut vec); + + // Pad to expected_len if needed + if actual_len < expected_len as usize { + vec.resize(expected_len as usize, 0.0f32); + } + + let tensor = Tensor::from_slice(&vec); + QuantTensor::quantize(tensor, precision) + } + MxDataType::Mx { elem, scale, block } => { + // Parse elements and scales + let elem_bytes_per_element = elem.size_in_bits() as usize / 8; + let num_blocks = (expected_len as usize + block as usize - 1) / block as usize; + let scale_bytes_per_scale = scale.size_in_bits() as usize / 8; + + let elem_bytes_len = expected_len as usize * elem_bytes_per_element; + let scale_bytes_len = num_blocks * scale_bytes_per_scale; + let total_expected = elem_bytes_len + scale_bytes_len; + + let (elem_bytes, scale_bytes) = if bytes.len() >= total_expected { + bytes.split_at(elem_bytes_len) + } else { + // Handle incomplete data + let elem_len = bytes.len().min(elem_bytes_len); + (&bytes[..elem_len], &bytes[elem_len..]) + }; + + // Convert elements + let mut vec = vec![0f32; expected_len as usize]; + if !elem_bytes.is_empty() { + let actual_elem_len = (elem_bytes.len() / elem_bytes_per_element).min(expected_len as usize); + elem.convert_bytes_to_f32_vec( + &elem_bytes[..actual_elem_len * elem_bytes_per_element], + &mut vec[..actual_elem_len], + ); + } + + // Apply scales if available + if !scale_bytes.is_empty() && scale_bytes.len() >= scale_bytes_per_scale { + let mut scale_vec = vec![0f32; num_blocks]; + scale.convert_bytes_to_f32_vec(scale_bytes, &mut scale_vec); + + for (chunk, scale_val) in vec.chunks_mut(block as usize).zip(scale_vec.iter().copied()) { + for elem in chunk.iter_mut() { + *elem *= scale_val; + } + } + } + + let tensor = Tensor::from_slice(&vec); + QuantTensor::quantize(tensor, precision) + } + } + } +} + diff --git a/behavioral_simulator/src/main.rs b/behavioral_simulator/src/main.rs index 1db81205..b0db5e12 100644 --- a/behavioral_simulator/src/main.rs +++ b/behavioral_simulator/src/main.rs @@ -17,6 +17,7 @@ use memory::MemoryModel; use quantize::{DataType, MxDataType, QuantTensor}; use runtime::{Duration, Executor, Instant}; use tch::{IndexOp, Tensor}; +use vector_sram::VectorSram; use tokio::sync::Mutex; use tokio::sync::oneshot::{self, Receiver}; @@ -199,145 +200,7 @@ impl MatrixSram { } } -/// Behaviour modelling of vector SRAM. -/// -/// The timing aspect is to be considered by the matrix and vector machines themselves. -struct VectorSram { - tile_size: u32, - tiles: Vec>>>, - ty: MxDataType, -} - -impl VectorSram { - /// Creata a matrix SRAM with given tile size and depth. - fn new(tile_size: u32, depth: usize, ty: MxDataType) -> Self { - let tiles = (0..depth) - .map(|_| Mutex::new(Ok(QuantTensor::zeros(tile_size as usize, ty)))) - .collect(); - Self { - tile_size, - tiles, - ty, - } - } - - fn size_in_bytes(&self) -> usize { - self.tile_size as usize * self.tiles.len() - } - - async fn read(&self, addr: u32) -> QuantTensor { - let addr_in_tiles = addr.assert_multiple_of(self.tile_size); - - let mut guard = self.tiles[addr_in_tiles as usize].lock().await; - if let Err(ref mut fut) = *guard { - *guard = Ok(fut.await.unwrap()); - } - - guard.as_ref().map_err(|_| ()).unwrap().clone() - } - - async fn write(&self, addr: u32, tensor: QuantTensor) { - let addr_in_tiles = addr.assert_multiple_of(self.tile_size); - - assert_eq!(tensor.data_type(), self.ty); - *self.tiles[addr_in_tiles as usize].lock().await = Ok(tensor); - } - - async fn write_delayed(&self, addr: u32, tensor: Receiver) { - let addr_in_tiles = addr.assert_multiple_of(self.tile_size); - - *self.tiles[addr_in_tiles as usize].lock().await = Err(tensor); - } - - async fn continous_write_delayed( - &self, - addr: u32, - write_amount: u32, - tensor: Receiver, - ) { - // println!("continous_write_delayed: addr = {:?}, write_amount = {:?}", addr, write_amount); - // println!("tile_size = {:?}", self.tile_size); - let addr_in_tiles = addr.assert_multiple_of(self.tile_size); - // Await the tensor from the channel (blocks until data arrives) - if let Ok(tensor) = tensor.await { - let dims = tensor.as_tensor().size(); - let chunk_size = self.tile_size as i64; - let total = dims[0]; - - // Split the tensor into chunks of self.tile_size and store each in self.tiles. - for i in 0..write_amount.min((total as u32 + self.tile_size - 1) / self.tile_size) { - let start = (i as i64) * chunk_size; - let end = ((i as i64 + 1) * chunk_size).min(total); - let chunk = tensor - .as_tensor() - .narrow(0, start, end - start) - .shallow_clone(); - let chunk_qt = QuantTensor::quantize(chunk, self.ty); - *self.tiles[(addr_in_tiles + i) as usize].lock().await = Ok(chunk_qt); - } - } - } - - /// TODO: used to preload Vector SRAM to facilitate testing. - async fn load_from_bytes(&self, bytes: &[u8]) { - let element_ty = self.ty.element_type(); - let element_bits = element_ty.size_in_bits(); - let bytes_per_element = (element_bits / 8) as usize; - - // Total number of elements that can be loaded - let total_elements = bytes.len() / bytes_per_element; - let tile_size = self.tile_size as usize; - let num_tiles = (total_elements + tile_size - 1) / tile_size; // Round up - - for tile_idx in 0..num_tiles.min(self.tiles.len()) { - let start_element = tile_idx * tile_size; - let end_element = (start_element + tile_size).min(total_elements); - let elements_in_tile = end_element - start_element; - - let start_byte = start_element * bytes_per_element; - let end_byte = end_element * bytes_per_element; - - // Convert bytes to f32 values - let mut vec = vec![0f32; elements_in_tile]; - element_ty.convert_bytes_to_f32_vec(&bytes[start_byte..end_byte], &mut vec); - - // Pad with zeros if needed - if elements_in_tile < tile_size { - vec.resize(tile_size, 0.0f32); - } - - // Create QuantTensor and store it - let tensor = tch::Tensor::from_slice(&vec); - let quant_tensor = QuantTensor::quantize(tensor, self.ty); - *self.tiles[tile_idx].lock().await = Ok(quant_tensor); - } - } - - async fn as_bytes(&self) -> Vec { - let element_ty = self.ty.element_type(); - let mut result = Vec::new(); - - for tile_mutex in &self.tiles { - let mut guard = tile_mutex.lock().await; - if let Err(ref mut fut) = *guard { - *guard = Ok(fut.await.unwrap()); - } - let tensor = guard.as_ref().map_err(|_| ()).unwrap(); - let tensor_data = tensor.as_tensor(); - let len = tensor_data.size1().unwrap() as usize; - let f32_slice = - unsafe { core::slice::from_raw_parts(tensor_data.data_ptr() as *const f32, len) }; - // Calculate bytes needed for THIS tile's actual size - let total_bits = len * element_ty.size_in_bits() as usize; - let bytes_needed = (total_bits + 7) / 8; - let mut tile_bytes = vec![0u8; bytes_needed]; - element_ty.bytes_from_f32(f32_slice, &mut tile_bytes); - result.extend_from_slice(&tile_bytes); - } - - result - } -} +// VectorSram is now imported from the vector_sram library struct MatrixMachine { mram: Arc, @@ -571,7 +434,7 @@ impl MatrixMachine { for j in 0..self.broadcast_amount { for i in 0..self.mlen { let tensor = self.h_accum.i((j as i64, i as i64, ..)); - self.vram.write(vec_base + (j * self.mlen + i) * self.mlen, QuantTensor::quantize(tensor, self.vram.ty)).await; + self.vram.write(vec_base + (j * self.mlen + i) * self.mlen, QuantTensor::quantize(tensor, self.vram.ty())).await; } } self.h_accum = Tensor::zeros( @@ -615,7 +478,7 @@ impl MatrixMachine { } async fn mv_wo(&mut self, v_addr: u32) { - let quant = QuantTensor::quantize(self.v_accum.shallow_clone(), self.vram.ty); + let quant = QuantTensor::quantize(self.v_accum.shallow_clone(), self.vram.ty()); self.vram.write(v_addr, quant).await; cycle!(1); self.v_accum = Tensor::zeros([self.mlen as i64], (tch::Kind::Float, tch::Device::Cpu)); @@ -850,13 +713,13 @@ impl VectorMachine { } async fn vector_transfer_fp(&mut self, vd: u32, f: &[f16]) { - assert_eq!(f.len(), self.vram.tile_size as usize, "Input vector length must match tile_size"); + assert_eq!(f.len(), self.vram.tile_size() as usize, "Input vector length must match tile_size"); // Convert f16 slice to f32 vector let f32_vec: Vec = f.iter().map(|x| f32::from(*x)).collect(); // Create tensor from f32 vector let tensor = tch::Tensor::from_slice(&f32_vec); // Quantize the tensor according to vram data type - let c = QuantTensor::quantize(tensor, self.vram.ty); + let c = QuantTensor::quantize(tensor, self.vram.ty()); cycle!(*VLEN); self.vram.write(vd, c).await; } @@ -1587,7 +1450,7 @@ impl Accelerator { addr + offset as u64, addr + self.reg_file.scale as u64 + scale as u64, dtype, - self.v_machine.vram.ty, + self.v_machine.vram.ty(), rstride, *VLEN, *PREFETCH_V_AMOUNT, @@ -1662,7 +1525,7 @@ struct Opts { async fn start() { let opts = Opts::parse(); let mram = Arc::new(MatrixSram::new(*MLEN, *MATRIX_SRAM_SIZE, *MATRIX_SRAM_TYPE)); // Matrix SRAM - let vram = Arc::new(VectorSram::new(*VLEN, *VECTOR_SRAM_SIZE, *VECTOR_SRAM_TYPE)); // Vector SRAM + let vram = Arc::new(VectorSram::new(*VLEN, *VECTOR_SRAM_SIZE, *VECTOR_SRAM_TYPE, None)); // Vector SRAM let machine = MatrixMachine { mram, vram: vram.clone(), @@ -1777,7 +1640,7 @@ async fn start() { "Matrix SRAM Contents: \n {}", accelerator.m_machine.mram.read(0x0000).await.as_tensor() ); - println!("FP SRAM Contents: \n {:?}", accelerator.fpsram); + // println!("FP SRAM Contents: \n {:?}", accelerator.fpsram); // Dump MRAM let mram_dump_path = "mram_dump.bin"; diff --git a/behavioral_simulator/src/op.rs b/behavioral_simulator/src/op.rs index 2cc86a1a..bd3c95a4 100644 --- a/behavioral_simulator/src/op.rs +++ b/behavioral_simulator/src/op.rs @@ -301,6 +301,16 @@ impl Opcode { VectorOrder::Reverse } } + + #[inline] + fn hbm_load_type_from(funct2: u8) -> HBM_LOAD_TYPE { + if funct2 == 0 { + HBM_LOAD_TYPE::MX + } else { + HBM_LOAD_TYPE::Normal + } + } + pub fn decode(instr: u32) -> Self { // eprintln!( // "decode(): instr = 0x{instr:08X} ({instr:032b})" diff --git a/compiler/asm_templates/load_int.py b/compiler/asm_templates/load_int.py new file mode 100644 index 00000000..e69de29b diff --git a/compiler/asm_templates/preload_act.py b/compiler/asm_templates/preload_act.py index b5e5a0d2..cf5851ae 100644 --- a/compiler/asm_templates/preload_act.py +++ b/compiler/asm_templates/preload_act.py @@ -37,7 +37,7 @@ def preload_act_asm( if batch == 1: for i in range(math.ceil(hidden_size / (vlen * preload_len))): - generated_code += f"H_PREFETCH_V gp{result_register}, gp{a_actual_register}, a{activation_offset_reg}, 0, 0 \n" + generated_code += f"H_PREFETCH_V gp{result_register}, gp{a_actual_register}, a{activation_offset_reg}, 0, 0, 0 \n" generated_code += f"S_ADDI_INT gp{result_register}, gp{result_register}, {vlen * preload_len} \n" generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{a_actual_register}, {load_amount_per_hidden} \n" else: @@ -46,7 +46,7 @@ def preload_act_asm( assert batch * hidden_size <= IMM2_BOUND, "batch * hidden_size must be less than {IMM2_BOUND}" for i in range(load_amount_per_hidden): for j in range(math.ceil(batch / preload_len)): - generated_code += f"H_PREFETCH_V gp{result_register}, gp{a_actual_register}, a{activation_offset_reg}, 1, 0 \n" + generated_code += f"H_PREFETCH_V gp{result_register}, gp{a_actual_register}, a{activation_offset_reg}, 1, 0, 0 \n" generated_code += f"S_ADDI_INT gp{result_register}, gp{result_register}, {vlen * preload_len} \n" generated_code += f"S_ADDI_INT gp{a_actual_register}, gp{a_actual_register}, {hidden_size * preload_len} \n" generated_code += f"S_ADDI_INT gp{a_actual_register}, gp0, {(i + 1) * vlen} \n" From a0f1517120e87585c8e2e9a08ca06805f29a53e1 Mon Sep 17 00:00:00 2001 From: GeorgeWu1204 Date: Tue, 2 Dec 2025 02:05:59 +0000 Subject: [PATCH 25/34] vsram support write int --- .../lib/vector_sram/src/lib.rs | 417 +++++++++--------- behavioral_simulator/src/main.rs | 2 +- 2 files changed, 205 insertions(+), 214 deletions(-) diff --git a/behavioral_simulator/lib/vector_sram/src/lib.rs b/behavioral_simulator/lib/vector_sram/src/lib.rs index e7f93d9c..f4cc7855 100644 --- a/behavioral_simulator/lib/vector_sram/src/lib.rs +++ b/behavioral_simulator/lib/vector_sram/src/lib.rs @@ -1,24 +1,27 @@ -use quantize::{MxDataType, QuantTensor}; +use quantize::{DataType, MxDataType, QuantTensor}; use tokio::sync::Mutex; use tokio::sync::oneshot::Receiver; -use tch::{IndexOp, Tensor}; +use tch::Tensor; -/// Vector SRAM that stores data in binary format with row-based storage. +/// Vector SRAM that stores data in pure binary format with row-based storage. +/// +/// The SRAM supports two data types: +/// - FP (Floating Point): Stored as binary representation of the FP type +/// - INT (Integer): Stored as binary representation of integers /// /// The SRAM is organized as rows, where each row has a width of `vlen * element_size_in_bytes`. -/// During read/write operations, data is clipped to VLEN-sized vectors and converted -/// between storage precision and access precision. +/// During read/write operations, data is clipped to VLEN-sized vectors. pub struct VectorSram { /// Vector length (VLEN) - determines the size of each vector operation vlen: u32, /// Number of rows in the SRAM depth: usize, - /// Storage precision - how data is stored internally - storage_precision: MxDataType, - /// Access precision - how data is read/written (defaults to storage_precision if None) - access_precision: Option, + /// Data type for FP storage (used when writing QuantTensor) + fp_type: DataType, + /// Size of integer in bytes (used when writing integer vectors) + int_size_bytes: usize, /// Raw binary storage: each row is stored as bytes - /// Row width = vlen * element_size_in_bytes(storage_precision) + /// Row width = vlen * element_size_in_bytes rows: Vec>, } @@ -29,33 +32,23 @@ enum RowData { } impl VectorSram { - /// Create a new Vector SRAM with given vector length, depth, and storage precision. + /// Create a new Vector SRAM with given vector length, depth, and data types. /// /// # Arguments /// * `vlen` - Vector length (VLEN) /// * `depth` - Number of rows in the SRAM - /// * `storage_precision` - Precision used for internal storage - /// * `access_precision` - Optional precision for read/write operations (defaults to storage_precision) + /// * `fp_type` - Floating point data type for FP operations + /// * `int_size_bytes` - Size of integer in bytes (typically 4 for i32) pub fn new( vlen: u32, depth: usize, - storage_precision: MxDataType, - access_precision: Option, + fp_type: DataType, + int_size_bytes: usize, ) -> Self { - let element_ty = storage_precision.element_type(); - let element_size = element_ty.size_in_bits() as usize / 8; + // Use FP type size for row width (can be changed if needed) + let element_size = fp_type.size_in_bits() as usize / 8; let row_width = vlen as usize * element_size; - // Handle Mx format: need to account for scale bytes - let row_width = match storage_precision { - MxDataType::Plain(_) => row_width, - MxDataType::Mx { block, scale, .. } => { - let num_blocks = (vlen as usize + block as usize - 1) / block as usize; - let scale_bytes = scale.size_in_bits() as usize / 8; - row_width + num_blocks * scale_bytes - } - }; - let rows = (0..depth) .map(|_| { Mutex::new(RowData::Ready(vec![0u8; row_width])) @@ -65,44 +58,48 @@ impl VectorSram { Self { vlen, depth, - storage_precision, - access_precision, + fp_type, + int_size_bytes, rows, } } + /// Create a new Vector SRAM from MxDataType (for backward compatibility). + /// + /// Extracts the Plain DataType from MxDataType and uses it for FP storage. + pub fn from_mx_type( + vlen: u32, + depth: usize, + mx_type: MxDataType, + ) -> Self { + let fp_type = match mx_type { + MxDataType::Plain(dt) => dt, + MxDataType::Mx { elem, .. } => elem, + }; + Self::new(vlen, depth, fp_type, 4) // Default to 4 bytes for int (i32) + } + /// Get the vector length (VLEN) pub fn tile_size(&self) -> u32 { self.vlen } - /// Get the storage precision type + /// Get the data type (for backward compatibility) pub fn ty(&self) -> MxDataType { - self.access_precision.unwrap_or(self.storage_precision) + MxDataType::Plain(self.fp_type) } /// Get the size of the SRAM in bytes pub fn size_in_bytes(&self) -> usize { - let element_ty = self.storage_precision.element_type(); - let element_size = element_ty.size_in_bits() as usize / 8; - let base_row_width = self.vlen as usize * element_size; - - let row_width = match self.storage_precision { - MxDataType::Plain(_) => base_row_width, - MxDataType::Mx { block, scale, .. } => { - let num_blocks = (self.vlen as usize + block as usize - 1) / block as usize; - let scale_bytes = scale.size_in_bits() as usize / 8; - base_row_width + num_blocks * scale_bytes - } - }; - + let element_size = self.fp_type.size_in_bits() as usize / 8; + let row_width = self.vlen as usize * element_size; row_width * self.depth } - /// Read a vector from the SRAM at the given address. + /// Read a vector from the SRAM at the given address as FP (QuantTensor). /// /// The address must be a multiple of vlen (in element units). - /// Data is read from storage precision and converted to access precision. + /// Data is read from binary storage and converted to QuantTensor. pub async fn read(&self, addr: u32) -> QuantTensor { let row_idx = self.addr_to_row_idx(addr); assert!(row_idx < self.depth, "Address out of bounds"); @@ -112,7 +109,7 @@ impl VectorSram { // Handle pending writes if let RowData::Pending(ref mut receiver) = *guard { let tensor = receiver.await.unwrap(); - let row_bytes = self.quant_tensor_to_bytes(&tensor, self.storage_precision); + let row_bytes = self.quant_tensor_to_bytes(&tensor); *guard = RowData::Ready(row_bytes); } @@ -122,40 +119,65 @@ impl VectorSram { RowData::Pending(_) => unreachable!(), }; - // Convert from storage precision to access precision - let access_precision = self.access_precision.unwrap_or(self.storage_precision); - let tensor = self.bytes_to_quant_tensor(&row_bytes, self.storage_precision, self.vlen); + // Convert from binary to QuantTensor + self.bytes_to_quant_tensor(&row_bytes, self.vlen) + } + + /// Read a vector from the SRAM at the given address as integers. + /// + /// The address must be a multiple of vlen (in element units). + /// Returns a vector of i32 values. + pub async fn read_int(&self, addr: u32) -> Vec { + let row_idx = self.addr_to_row_idx(addr); + assert!(row_idx < self.depth, "Address out of bounds"); + + let mut guard = self.rows[row_idx].lock().await; - // Convert to access precision if different - if access_precision != self.storage_precision { - let tensor_f32 = tensor.as_tensor(); - QuantTensor::quantize(tensor_f32.shallow_clone(), access_precision) - } else { - tensor + // Handle pending writes (convert to bytes first) + if let RowData::Pending(ref mut receiver) = *guard { + let tensor = receiver.await.unwrap(); + let row_bytes = self.quant_tensor_to_bytes(&tensor); + *guard = RowData::Ready(row_bytes); } + + // Read the row data + let row_bytes = match &*guard { + RowData::Ready(bytes) => bytes.clone(), + RowData::Pending(_) => unreachable!(), + }; + + // Convert from binary to integers + self.bytes_to_int_vec(&row_bytes, self.vlen) } - /// Write a vector to the SRAM at the given address. + /// Write a vector to the SRAM at the given address as FP (QuantTensor). /// /// The address must be a multiple of vlen (in element units). - /// Data is converted from access precision to storage precision before writing. + /// Data is converted from QuantTensor to binary storage. pub async fn write(&self, addr: u32, tensor: QuantTensor) { let row_idx = self.addr_to_row_idx(addr); assert!(row_idx < self.depth, "Address out of bounds"); - // Convert to storage precision if needed - let storage_tensor = if tensor.data_type() != self.storage_precision { - let tensor_f32 = tensor.as_tensor(); - QuantTensor::quantize(tensor_f32.shallow_clone(), self.storage_precision) - } else { - tensor - }; - // Clip to VLEN - let clipped = self.clip_to_vlen(&storage_tensor); + let clipped = self.clip_to_vlen(&tensor); // Convert to bytes - let row_bytes = self.quant_tensor_to_bytes(&clipped, self.storage_precision); + let row_bytes = self.quant_tensor_to_bytes(&clipped); + + *self.rows[row_idx].lock().await = RowData::Ready(row_bytes); + } + + /// Write a vector to the SRAM at the given address as integers. + /// + /// The address must be a multiple of vlen (in element units). + /// Data is converted from integers to binary storage. + pub async fn write_int(&self, addr: u32, int_vec: &[i32]) { + let row_idx = self.addr_to_row_idx(addr); + assert!(row_idx < self.depth, "Address out of bounds"); + assert!(int_vec.len() <= self.vlen as usize, "Vector too long"); + + // Convert integers to bytes + let row_bytes = self.int_vec_to_bytes(int_vec, self.vlen); *self.rows[row_idx].lock().await = RowData::Ready(row_bytes); } @@ -177,19 +199,19 @@ impl VectorSram { ) { let start_row_idx = self.addr_to_row_idx(addr); - // Await the tensor from the channel + // Await the tensor from the channel and extract data immediately to make it Send let tensor = tensor.await.unwrap(); let tensor_data = tensor.as_tensor(); - let total_elements = tensor_data.size1().unwrap() as i64; - - // Convert to storage precision if needed - let storage_tensor = if tensor.data_type() != self.storage_precision { - QuantTensor::quantize(tensor_data.shallow_clone(), self.storage_precision) - } else { - tensor + let total_elements = tensor_data.size1().unwrap() as usize; + + // Extract f32 data from tensor to make it Send-safe + let len = total_elements; + let f32_slice = unsafe { + core::slice::from_raw_parts(tensor_data.data_ptr() as *const f32, len) }; + let data_vec: Vec = f32_slice.to_vec(); - let chunk_size = self.vlen as i64; + let chunk_size = self.vlen as usize; let num_chunks = write_amount.min(((total_elements + chunk_size - 1) / chunk_size) as u32); for i in 0..num_chunks { @@ -198,29 +220,50 @@ impl VectorSram { break; } - let start = i as i64 * chunk_size; + let start = (i as usize) * chunk_size; let end = (start + chunk_size).min(total_elements); - let chunk = storage_tensor - .as_tensor() - .narrow(0, start, end - start) - .shallow_clone(); + let chunk_data = &data_vec[start..end]; // Pad to VLEN if needed - let padded = if (end - start) < chunk_size { - let mut padded_tensor = Tensor::zeros( - [chunk_size], - (tch::Kind::Float, tch::Device::Cpu), - ); - padded_tensor - .i(0..(end - start)) - .copy_(&chunk); - padded_tensor - } else { - chunk - }; + let mut padded_data = vec![0.0f32; chunk_size]; + let chunk_len = end - start; + padded_data[..chunk_len].copy_from_slice(chunk_data); + + // Create tensor from padded data and convert to bytes + let padded_tensor = Tensor::from_slice(&padded_data); + let chunk_qt = QuantTensor::quantize(padded_tensor, MxDataType::Plain(self.fp_type)); + let row_bytes = self.quant_tensor_to_bytes(&chunk_qt); + *self.rows[row_idx].lock().await = RowData::Ready(row_bytes); + } + } + + /// Continuous write delayed for integers - writes multiple rows from a single integer vector. + pub async fn continous_write_delayed_int( + &self, + addr: u32, + write_amount: u32, + int_vec: Receiver>, + ) { + let start_row_idx = self.addr_to_row_idx(addr); + + // Await the integer vector from the channel + let int_vec = int_vec.await.unwrap(); + let total_elements = int_vec.len(); + let chunk_size = self.vlen as usize; + let num_chunks = write_amount.min(((total_elements + chunk_size - 1) / chunk_size) as u32); + + for i in 0..num_chunks { + let row_idx = start_row_idx + i as usize; + if row_idx >= self.depth { + break; + } - let chunk_qt = QuantTensor::quantize(padded, self.storage_precision); - let row_bytes = self.quant_tensor_to_bytes(&chunk_qt, self.storage_precision); + let start = (i as usize) * chunk_size; + let end = (start + chunk_size).min(total_elements); + let chunk = &int_vec[start..end]; + + // Convert to bytes (will pad to VLEN if needed) + let row_bytes = self.int_vec_to_bytes(chunk, self.vlen); *self.rows[row_idx].lock().await = RowData::Ready(row_bytes); } } @@ -229,21 +272,8 @@ impl VectorSram { /// /// This is used for preloading the SRAM with test data. pub async fn load_from_bytes(&self, bytes: &[u8]) { - let element_ty = self.storage_precision.element_type(); - let element_bits = element_ty.size_in_bits(); - let bytes_per_element = (element_bits / 8) as usize; - - // Calculate row width - let base_row_width = self.vlen as usize * bytes_per_element; - let _row_width = match self.storage_precision { - MxDataType::Plain(_) => base_row_width, - MxDataType::Mx { block, scale, .. } => { - let num_blocks = (self.vlen as usize + block as usize - 1) / block as usize; - let scale_bytes = scale.size_in_bits() as usize / 8; - base_row_width + num_blocks * scale_bytes - } - }; - + let element_size = self.fp_type.size_in_bits() as usize / 8; + let bytes_per_element = element_size; let total_elements = bytes.len() / bytes_per_element; let num_rows = (total_elements + self.vlen as usize - 1) / self.vlen as usize; @@ -257,7 +287,7 @@ impl VectorSram { // Convert bytes to f32 values let mut vec = vec![0f32; elements_in_row]; - element_ty.convert_bytes_to_f32_vec(&bytes[start_byte..end_byte], &mut vec); + self.fp_type.convert_bytes_to_f32_vec(&bytes[start_byte..end_byte], &mut vec); // Pad with zeros if needed if elements_in_row < self.vlen as usize { @@ -266,8 +296,8 @@ impl VectorSram { // Create QuantTensor and convert to bytes let tensor = Tensor::from_slice(&vec); - let quant_tensor = QuantTensor::quantize(tensor, self.storage_precision); - let row_bytes = self.quant_tensor_to_bytes(&quant_tensor, self.storage_precision); + let quant_tensor = QuantTensor::quantize(tensor, MxDataType::Plain(self.fp_type)); + let row_bytes = self.quant_tensor_to_bytes(&quant_tensor); *self.rows[row_idx].lock().await = RowData::Ready(row_bytes); } } @@ -284,7 +314,7 @@ impl VectorSram { // Handle pending writes if let RowData::Pending(ref mut receiver) = *guard { let tensor = receiver.await.unwrap(); - let row_bytes = self.quant_tensor_to_bytes(&tensor, self.storage_precision); + let row_bytes = self.quant_tensor_to_bytes(&tensor); *guard = RowData::Ready(row_bytes); } @@ -321,114 +351,75 @@ impl VectorSram { } } - /// Convert QuantTensor to bytes according to the given precision - fn quant_tensor_to_bytes(&self, tensor: &QuantTensor, precision: MxDataType) -> Vec { + /// Convert QuantTensor to bytes (FP format) + fn quant_tensor_to_bytes(&self, tensor: &QuantTensor) -> Vec { let tensor_data = tensor.as_tensor(); let len = tensor_data.size1().unwrap() as usize; let f32_slice = unsafe { core::slice::from_raw_parts(tensor_data.data_ptr() as *const f32, len) }; - match precision { - MxDataType::Plain(elem_ty) => { - let total_bits = len * elem_ty.size_in_bits() as usize; - let bytes_needed = (total_bits + 7) / 8; - let mut bytes = vec![0u8; bytes_needed]; - elem_ty.bytes_from_f32(f32_slice, &mut bytes); - bytes - } - MxDataType::Mx { elem, scale, block } => { - // Serialize elements - let total_bits = len * elem.size_in_bits() as usize; - let elem_bytes_needed = (total_bits + 7) / 8; - let mut elem_bytes = vec![0u8; elem_bytes_needed]; - elem.bytes_from_f32(f32_slice, &mut elem_bytes); - - // Serialize scales (one per block) - let num_blocks = (len + block as usize - 1) / block as usize; - // For now, use scale of 1.0 for each block - // TODO: implement proper scale calculation - let scale_vec = vec![1.0f32; num_blocks]; - let scale_bytes_needed = num_blocks * (scale.size_in_bits() as usize / 8); - let mut scale_bytes = vec![0u8; scale_bytes_needed]; - scale.bytes_from_f32(&scale_vec, &mut scale_bytes); - - // Combine: elements first, then scales - let mut result = elem_bytes; - result.extend_from_slice(&scale_bytes); - result - } + let total_bits = len * self.fp_type.size_in_bits() as usize; + let bytes_needed = (total_bits + 7) / 8; + let mut bytes = vec![0u8; bytes_needed]; + self.fp_type.bytes_from_f32(f32_slice, &mut bytes); + bytes + } + + /// Convert bytes to QuantTensor (FP format) + fn bytes_to_quant_tensor(&self, bytes: &[u8], expected_len: u32) -> QuantTensor { + let bytes_per_element = self.fp_type.size_in_bits() as usize / 8; + let num_elements = bytes.len() / bytes_per_element; + let actual_len = num_elements.min(expected_len as usize); + + let mut vec = vec![0f32; actual_len]; + self.fp_type.convert_bytes_to_f32_vec( + &bytes[..actual_len * bytes_per_element], + &mut vec, + ); + + // Pad to expected_len if needed + if actual_len < expected_len as usize { + vec.resize(expected_len as usize, 0.0f32); } + + let tensor = Tensor::from_slice(&vec); + QuantTensor::quantize(tensor, MxDataType::Plain(self.fp_type)) } - /// Convert bytes to QuantTensor according to the given precision - fn bytes_to_quant_tensor( - &self, - bytes: &[u8], - precision: MxDataType, - expected_len: u32, - ) -> QuantTensor { - match precision { - MxDataType::Plain(elem_ty) => { - let bytes_per_element = elem_ty.size_in_bits() as usize / 8; - let num_elements = bytes.len() / bytes_per_element; - let actual_len = num_elements.min(expected_len as usize); - - let mut vec = vec![0f32; actual_len]; - elem_ty.convert_bytes_to_f32_vec(&bytes[..actual_len * bytes_per_element], &mut vec); - - // Pad to expected_len if needed - if actual_len < expected_len as usize { - vec.resize(expected_len as usize, 0.0f32); - } - - let tensor = Tensor::from_slice(&vec); - QuantTensor::quantize(tensor, precision) - } - MxDataType::Mx { elem, scale, block } => { - // Parse elements and scales - let elem_bytes_per_element = elem.size_in_bits() as usize / 8; - let num_blocks = (expected_len as usize + block as usize - 1) / block as usize; - let scale_bytes_per_scale = scale.size_in_bits() as usize / 8; - - let elem_bytes_len = expected_len as usize * elem_bytes_per_element; - let scale_bytes_len = num_blocks * scale_bytes_per_scale; - let total_expected = elem_bytes_len + scale_bytes_len; - - let (elem_bytes, scale_bytes) = if bytes.len() >= total_expected { - bytes.split_at(elem_bytes_len) - } else { - // Handle incomplete data - let elem_len = bytes.len().min(elem_bytes_len); - (&bytes[..elem_len], &bytes[elem_len..]) - }; - - // Convert elements - let mut vec = vec![0f32; expected_len as usize]; - if !elem_bytes.is_empty() { - let actual_elem_len = (elem_bytes.len() / elem_bytes_per_element).min(expected_len as usize); - elem.convert_bytes_to_f32_vec( - &elem_bytes[..actual_elem_len * elem_bytes_per_element], - &mut vec[..actual_elem_len], - ); - } - - // Apply scales if available - if !scale_bytes.is_empty() && scale_bytes.len() >= scale_bytes_per_scale { - let mut scale_vec = vec![0f32; num_blocks]; - scale.convert_bytes_to_f32_vec(scale_bytes, &mut scale_vec); - - for (chunk, scale_val) in vec.chunks_mut(block as usize).zip(scale_vec.iter().copied()) { - for elem in chunk.iter_mut() { - *elem *= scale_val; - } - } - } - - let tensor = Tensor::from_slice(&vec); - QuantTensor::quantize(tensor, precision) + /// Convert integer vector to bytes + fn int_vec_to_bytes(&self, int_vec: &[i32], expected_len: u32) -> Vec { + let mut bytes = Vec::with_capacity(expected_len as usize * self.int_size_bytes); + + // Write the actual integers + for &val in int_vec.iter() { + bytes.extend_from_slice(&val.to_le_bytes()); + } + + // Pad with zeros to expected_len + bytes.resize(expected_len as usize * self.int_size_bytes, 0); + + bytes + } + + /// Convert bytes to integer vector + fn bytes_to_int_vec(&self, bytes: &[u8], expected_len: u32) -> Vec { + let mut result = Vec::with_capacity(expected_len as usize); + let mut offset = 0; + + for _ in 0..expected_len as usize { + if offset + self.int_size_bytes <= bytes.len() { + let mut int_bytes = [0u8; 4]; + let copy_len = self.int_size_bytes.min(4); + int_bytes[..copy_len].copy_from_slice(&bytes[offset..offset + copy_len]); + let val = i32::from_le_bytes(int_bytes); + result.push(val); + offset += self.int_size_bytes; + } else { + result.push(0); } } + + result } } - diff --git a/behavioral_simulator/src/main.rs b/behavioral_simulator/src/main.rs index b0db5e12..0589a36d 100644 --- a/behavioral_simulator/src/main.rs +++ b/behavioral_simulator/src/main.rs @@ -1525,7 +1525,7 @@ struct Opts { async fn start() { let opts = Opts::parse(); let mram = Arc::new(MatrixSram::new(*MLEN, *MATRIX_SRAM_SIZE, *MATRIX_SRAM_TYPE)); // Matrix SRAM - let vram = Arc::new(VectorSram::new(*VLEN, *VECTOR_SRAM_SIZE, *VECTOR_SRAM_TYPE, None)); // Vector SRAM + let vram = Arc::new(VectorSram::from_mx_type(*VLEN, *VECTOR_SRAM_SIZE, *VECTOR_SRAM_TYPE)); // Vector SRAM let machine = MatrixMachine { mram, vram: vram.clone(), From ab6cba80ca04aafaf3194e1ff0ab0f0315d5c5f6 Mon Sep 17 00:00:00 2001 From: GeorgeWu1204 Date: Tue, 2 Dec 2025 17:30:34 +0000 Subject: [PATCH 26/34] include int type in behave simulator --- .../lib/quantize/src/dtype.rs | 49 +++++++++++++ behavioral_simulator/lib/quantize/src/lib.rs | 2 +- behavioral_simulator/src/load_config.rs | 31 +++++++- behavioral_simulator/src/main.rs | 72 +++++++++++++++---- behavioral_simulator/src/op.rs | 5 +- src/definitions/plena_settings.toml | 6 ++ 6 files changed, 148 insertions(+), 17 deletions(-) diff --git a/behavioral_simulator/lib/quantize/src/dtype.rs b/behavioral_simulator/lib/quantize/src/dtype.rs index eb535767..00084f51 100644 --- a/behavioral_simulator/lib/quantize/src/dtype.rs +++ b/behavioral_simulator/lib/quantize/src/dtype.rs @@ -162,9 +162,44 @@ fn test_f16() { ); } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub struct IntType { + pub width: u32, +} + +impl IntType { + pub const fn size_in_bits(self) -> u8 { + self.width as u8 + } + + /// Convert f32 to integer bits. Truncates the float to an integer. + pub const fn bits_from_f32(self, float: f32) -> u32 { + let int_val = float as i32; + let mask = if self.width >= 32 { + 0xFFFFFFFFu32 + } else { + ((1u64 << self.width) - 1) as u32 + }; + (int_val as u32) & mask + } + + /// Convert integer bits to f32. Interprets bits as unsigned integer. + pub const fn convert_bits_to_f32(self, bits: u32) -> f32 { + let mask = if self.width >= 32 { + 0xFFFFFFFFu32 + } else { + ((1u64 << self.width) - 1) as u32 + }; + let masked_bits = bits & mask; + masked_bits as f32 + } +} + + #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub enum DataType { Fp(FpType), + Int(IntType), } impl From for DataType { @@ -173,22 +208,27 @@ impl From for DataType { } } + + impl DataType { pub fn size_in_bits(self) -> u8 { match self { DataType::Fp(fp_type) => fp_type.size_in_bits(), + DataType::Int(int_type) => int_type.size_in_bits(), } } pub const fn bits_from_f32(self, float: f32) -> u32 { match self { DataType::Fp(fp_type) => fp_type.bits_from_f32(float), + DataType::Int(int_type) => int_type.bits_from_f32(float), } } pub const fn convert_bits_to_f32(self, bits: u32) -> f32 { match self { DataType::Fp(fp_type) => fp_type.convert_bits_to_f32(bits), + DataType::Int(int_type) => int_type.convert_bits_to_f32(bits), } } @@ -259,6 +299,15 @@ impl MxDataType { MxDataType::Mx { elem, .. } => elem, } } + + /// Returns the size in bits of the element type + /// Works for both Plain (FP and Int) and Mx variants + pub fn size_in_bits(self) -> u8 { + match self { + MxDataType::Plain(data_type) => data_type.size_in_bits(), + MxDataType::Mx { elem, .. } => elem.size_in_bits(), + } + } } impl From for MxDataType { diff --git a/behavioral_simulator/lib/quantize/src/lib.rs b/behavioral_simulator/lib/quantize/src/lib.rs index 4bc0342a..45967a9d 100644 --- a/behavioral_simulator/lib/quantize/src/lib.rs +++ b/behavioral_simulator/lib/quantize/src/lib.rs @@ -1,5 +1,5 @@ mod dtype; mod tensor; -pub use dtype::{DataType, FpType, MxDataType}; +pub use dtype::{DataType, FpType, IntType, MxDataType}; pub use tensor::QuantTensor; diff --git a/behavioral_simulator/src/load_config.rs b/behavioral_simulator/src/load_config.rs index 4791835d..dcaf09bc 100644 --- a/behavioral_simulator/src/load_config.rs +++ b/behavioral_simulator/src/load_config.rs @@ -4,7 +4,7 @@ use std::time::Duration; use std::{fs, sync::LazyLock}; // Import the types from your main module -use quantize::{DataType, FpType, MxDataType}; +use quantize::{DataType, FpType, IntType, MxDataType}; #[derive(Debug, Serialize, Deserialize, Clone)] pub struct ConfigValue { @@ -29,10 +29,16 @@ pub struct FpTypeConfig { pub mantissa: u8, } +#[derive(Debug, Serialize, Deserialize, Clone)] +pub struct IntTypeConfig { + pub width: u32, +} + #[derive(Debug, Serialize, Deserialize, Clone)] #[serde(tag = "type")] pub enum DataTypeConfig { Fp(FpTypeConfig), + Int(IntTypeConfig), } #[derive(Debug, Serialize, Deserialize, Clone)] @@ -110,6 +116,8 @@ pub struct PrecisionSection { pub hbm_v_act_type: MxDataTypeConfig, #[serde(rename = "HBM_V_KV_TYPE")] pub hbm_v_kv_type: MxDataTypeConfig, + #[serde(rename = "HBM_V_INT_TYPE")] + pub hbm_v_int_type: MxDataTypeConfig, #[serde(rename = "SCALAR_FP")] pub scalar_fp: DataTypeConfig, } @@ -250,6 +258,14 @@ impl Default for AcceleratorConfig { }), }, }, + hbm_v_int_type: MxDataTypeConfig { + format: "Plain".to_string(), + data: MxDataTypeData::Plain { + data_type: DataTypeConfig::Int(IntTypeConfig { + width: 32, + }), + }, + }, scalar_fp: DataTypeConfig::Fp(FpTypeConfig { sign: true, exponent: 8, @@ -333,10 +349,19 @@ impl From for FpType { } } +impl From for IntType { + fn from(config: IntTypeConfig) -> Self { + IntType { + width: config.width, + } + } +} + impl From for DataType { fn from(config: DataTypeConfig) -> Self { match config { DataTypeConfig::Fp(fp_config) => DataType::Fp(fp_config.into()), + DataTypeConfig::Int(int_config) => DataType::Int(int_config.into()), } } } @@ -446,6 +471,10 @@ pub fn vector_kv_type() -> MxDataType { CONFIG.precision.hbm_v_kv_type.clone().into() } +pub fn vector_int_type() -> MxDataType { + CONFIG.precision.hbm_v_int_type.clone().into() +} + // Additional accessor functions for new parameters pub fn mlen() -> u32 { CONFIG.config.mlen.value diff --git a/behavioral_simulator/src/main.rs b/behavioral_simulator/src/main.rs index 0589a36d..f0f17504 100644 --- a/behavioral_simulator/src/main.rs +++ b/behavioral_simulator/src/main.rs @@ -57,6 +57,7 @@ static MATRIX_WEIGHT_TYPE: LazyLock = LazyLock::new(|| matrix_weight static MATRIX_KV_TYPE: LazyLock = LazyLock::new(|| matrix_kv_type()); static VECTOR_ACTIVATION_TYPE: LazyLock = LazyLock::new(|| vector_activation_type()); static VECTOR_KV_TYPE: LazyLock = LazyLock::new(|| vector_kv_type()); +static VECTOR_INT_TYPE: LazyLock = LazyLock::new(|| vector_int_type()); static PREFETCH_M_AMOUNT: LazyLock = LazyLock::new(|| hbm_m_prefetch_amount()); static PREFETCH_V_AMOUNT: LazyLock = LazyLock::new(|| hbm_v_prefetch_amount()); static WRITEBACK_V_AMOUNT: LazyLock = LazyLock::new(|| hbm_v_writeback_amount()); @@ -987,17 +988,17 @@ impl Accelerator { receiver } - fn transfer_normal_from_hbm( + fn transfer_int_from_hbm( &mut self, index: u64, - hbm_type: DataType, + hbm_type: MxDataType, sram_type: MxDataType, rstride: u8, load_dim: u32, load_amount: u32, - ) -> Receiver { - // Simplified version: only loads elements from HBM based on index and stride - // No scale operations, no MX data type complexity + ) -> Receiver> { + // Transfer integer data from HBM: loads elements based on index and stride + // Converts bytes directly to Vec for writing to vector SRAM let (sender, receiver) = oneshot::channel(); @@ -1056,17 +1057,43 @@ impl Accelerator { } } - // Convert bytes to f32 and quantize - let mut vec = vec![0f32; total_elements]; - hbm_type.convert_bytes_to_f32_vec(&bytes, &mut vec); + // Convert bytes to Vec + // For integer types, interpret bytes as little-endian integer values + let element_size_bytes = element_bits as usize / 8; + let mut int_vec = Vec::with_capacity(total_elements); + + for i in 0..total_elements { + let byte_offset = i * element_size_bytes; + if byte_offset + element_size_bytes <= bytes.len() { + // Read bytes as little-endian and convert to i32 + let int_value = match element_size_bytes { + 1 => bytes[byte_offset] as i8 as i32, + 2 => { + let bytes_slice = &bytes[byte_offset..byte_offset + 2]; + i16::from_le_bytes([bytes_slice[0], bytes_slice[1]]) as i32 + } + 4 => { + let bytes_slice = &bytes[byte_offset..byte_offset + 4]; + i32::from_le_bytes([ + bytes_slice[0], + bytes_slice[1], + bytes_slice[2], + bytes_slice[3], + ]) + } + _ => panic!("Unsupported integer size: {} bytes (must be 1, 2, or 4)", element_size_bytes), + }; + int_vec.push(int_value); + } + } - let tensor = tch::Tensor::from_slice(&vec); - let _ = sender.send(QuantTensor::quantize(tensor, sram_type)); + let _ = sender.send(int_vec); }); receiver } + async fn do_ops(&mut self, ops: &[op::Opcode]) { for op in ops { println!("execute op = {:?}", op); @@ -1428,15 +1455,16 @@ impl Accelerator { rs2, rstride, precision, - datatype, + loadtype, } => { // TODO: rstride support to be added let offset = self.reg_file.gp_reg[rs1 as usize]; let addr = self.reg_file.hbm_addr_reg[rs2 as usize]; - if matches! (datatype, op::HBM_LOAD_TYPE::MX) { + if matches! (loadtype, op::HBM_LOAD_TYPE::MX) { let dtype = match precision { op::VectorPrecision::Activation => *VECTOR_ACTIVATION_TYPE, op::VectorPrecision::KeyValue => *VECTOR_KV_TYPE, + op::VectorPrecision::INT => *VECTOR_INT_TYPE, }; let scale = match dtype { @@ -1463,7 +1491,25 @@ impl Accelerator { .continous_write_delayed(dest, *PREFETCH_V_AMOUNT, xfer) .await; } else { - + let dtype = match precision { + // TODO: Left for future support FP load. + op::VectorPrecision::Activation => *VECTOR_ACTIVATION_TYPE, + op::VectorPrecision::KeyValue => *VECTOR_KV_TYPE, + op::VectorPrecision::INT => *VECTOR_INT_TYPE, + }; + let xfer = self.transfer_int_from_hbm( + addr + offset as u64, + dtype, + self.v_machine.vram.ty(), + rstride, + *VLEN, + *PREFETCH_V_AMOUNT, + ); + let dest = self.reg_file.gp_reg[rd as usize]; + self.v_machine + .vram + .continous_write_delayed_int(dest, *PREFETCH_V_AMOUNT, xfer) + .await; } diff --git a/behavioral_simulator/src/op.rs b/behavioral_simulator/src/op.rs index bd3c95a4..4acf1903 100644 --- a/behavioral_simulator/src/op.rs +++ b/behavioral_simulator/src/op.rs @@ -8,6 +8,7 @@ pub enum MatrixPrecision { pub enum VectorPrecision { Activation, KeyValue, + INT, } #[derive(Debug, Clone, Copy)] @@ -238,7 +239,7 @@ pub enum Opcode { rs2: u8, rstride: u8, precision: VectorPrecision, - datatype: HBM_LOAD_TYPE, + loadtype: HBM_LOAD_TYPE, }, H_STORE_V { rd: u8, @@ -388,7 +389,7 @@ impl Opcode { rs2, rstride: rs3, precision: Self::vector_precision_from(funct1), - datatype: Self::hbm_load_type_from(funct2), + loadtype: Self::hbm_load_type_from(funct2), }, // 0x2A => Self::H_PREFETCH_V { rd, rs1, rs2, rstride: rs3, precision: VectorPrecision::KeyValue }, 0x2A => Self::H_STORE_V { diff --git a/src/definitions/plena_settings.toml b/src/definitions/plena_settings.toml index 3f27e552..18845241 100644 --- a/src/definitions/plena_settings.toml +++ b/src/definitions/plena_settings.toml @@ -106,6 +106,12 @@ sign = false exponent = 8 mantissa = 0 +[PRECISION.HBM_V_INT_TYPE] +format = "Plain" +[PRECISION.HBM_V_INT_TYPE.DATA_TYPE] +type = "Int" +width = 32 + [PRECISION.SCALAR_FP] type = "Fp" sign = true From 799c2ad75c09e091589264dbeec8c09b2abec17e Mon Sep 17 00:00:00 2001 From: GeorgeWu1204 Date: Tue, 2 Dec 2025 20:02:41 +0000 Subject: [PATCH 27/34] scripts to create hbm sim mem for int --- behavioral_simulator/testbench/attn_test.py | 4 +- behavioral_simulator/testbench/bmm_test.py | 4 +- behavioral_simulator/testbench/dllm1_test.py | 21 ++-- behavioral_simulator/testbench/ffn_test.py | 4 +- behavioral_simulator/testbench/linear_test.py | 4 +- behavioral_simulator/testbench/rms_test.py | 4 +- .../testbench/s_map_v_test.py | 4 +- .../testbench/two_input_test.py | 4 +- behavioral_simulator/testbench/view_mem.py | 118 +++++++++++++++--- src/definitions/plena_settings.toml | 3 + tools/assembler/assembly_to_binary.py | 2 + tools/memory_mapping/memory_map.py | 27 +++- tools/sim_env_utils/__init__.py | 2 +- tools/sim_env_utils/build_env.py | 58 +++++---- tools/sim_env_utils/build_sys_tools.py | 61 +++++---- tools/utils/load_config.py | 26 +--- 16 files changed, 234 insertions(+), 112 deletions(-) diff --git a/behavioral_simulator/testbench/attn_test.py b/behavioral_simulator/testbench/attn_test.py index 74ea3c48..d05b5630 100644 --- a/behavioral_simulator/testbench/attn_test.py +++ b/behavioral_simulator/testbench/attn_test.py @@ -14,7 +14,7 @@ from create_sim_env import create_sim_env from transformers.modeling_flash_attention_utils import _flash_attention_forward as _flash_attention_forward_ref from aria_lm_ops.models.llama import flash_attn2_gemv -from sim_env_utils import build_sim_env +from sim_env_utils import create_mem_for_sim # Reshape q, k, v so that each batch's data stays together and rows are of length mem_row_size @@ -145,4 +145,4 @@ ) create_sim_env(input_tensor, weights, gen_assembly_code, golden_result, fp_preload) - build_sim_env(data_size=256, mode="behave_sim", asm="attn", data=None, specified_data_order = ["q", "k", "v"]) \ No newline at end of file + create_mem_for_sim(data_size=256, mode="behave_sim", asm="attn", data=None, specified_data_order = ["q", "k", "v"]) \ No newline at end of file diff --git a/behavioral_simulator/testbench/bmm_test.py b/behavioral_simulator/testbench/bmm_test.py index b48a26b2..0a123b06 100644 --- a/behavioral_simulator/testbench/bmm_test.py +++ b/behavioral_simulator/testbench/bmm_test.py @@ -8,7 +8,7 @@ from test_data_gen import get_weights_path, generate_and_save_random_weights from compiler.asm_templates import batched_matmul_asm, preload_addr_reg_asm, reset_reg_asm from create_sim_env import create_sim_env -from sim_env_utils import build_sim_env +from sim_env_utils import create_mem_for_sim if __name__ == "__main__": # Testing the operation (hidden_size, hidden_size) @ (hidden_size, batch_size) @@ -65,7 +65,7 @@ create_sim_env(input_tensor, weight_2_tensor, gen_assembly_code, golden_result, fp_preload) - build_sim_env(data_size=256, mode="behave_sim", asm="linear", data=None, specified_data_order = ["input_tensor", "model_weights"]) + create_mem_for_sim(data_size=256, mode="behave_sim", asm="linear", data=None, specified_data_order = ["input_tensor", "model_weights"]) print("================================================") print("Finished generating assembly code") diff --git a/behavioral_simulator/testbench/dllm1_test.py b/behavioral_simulator/testbench/dllm1_test.py index 9d10680c..7003494a 100644 --- a/behavioral_simulator/testbench/dllm1_test.py +++ b/behavioral_simulator/testbench/dllm1_test.py @@ -8,7 +8,7 @@ from test_data_gen import get_weights_path, generate_and_save_random_weights from compiler.asm_templates import preload_act_asm, reset_reg_asm, preload_addr_reg_asm from create_sim_env import create_sim_env -from sim_env_utils import build_sim_env +from sim_env_utils import create_mem_for_sim import torch.nn.functional as F from tools.memory_mapping.hbm_addr_map import align_addr_to_hbm_bandwidth @@ -20,10 +20,10 @@ # Testing the operation (hidden_size, hidden_size) @ (hidden_size, batch_size) vocal_size = 2 - hidden_size = 64 + hidden_size = 128 vlen = 64 - batch_size = 1 - preload_amount = 1 + batch_size = 4 + preload_amount = 4 real_data_ratio = (8*8 + 8) / (8 * 8) hbm_data_width = 64 fp_preload = [0.0, 1e-6] @@ -33,12 +33,17 @@ weights = logits original_output = logits + # Generate vlen random int32 data + int_preload = torch.randint(low=0, high=2**31, size=(vlen,), dtype=torch.int32) + + print('logits.shape= ', logits.shape) print('original_output.shape= ', original_output.shape) input_tensor = { "logits": logits, + "int": int_preload, } golden_result = { @@ -75,9 +80,11 @@ act_vram_offset=0, activation_offset_reg=0 ) - - + # Preload Integer Activation to the later section. + gen_assembly_code += f"S_ADDI_INT gp1, gp0, {hidden_size * batch_size} \n" + gen_assembly_code += f"S_ADDI_INT gp2, gp0, {int(hidden_size * batch_size * real_data_ratio)} \n" + gen_assembly_code += "H_PREFETCH_V gp1, gp2, a0, 0, 0, 1 \n" create_sim_env(input_tensor, gen_assembly_code, golden_result, fp_preload) - build_sim_env(data_size=256, mode="behave_sim", asm="dllm", data=None, specified_data_order = ["logits"]) \ No newline at end of file + create_mem_for_sim(data_size=256, mode="behave_sim", asm="dllm", data=None, specified_data_order = ["logits", "int"]) \ No newline at end of file diff --git a/behavioral_simulator/testbench/ffn_test.py b/behavioral_simulator/testbench/ffn_test.py index 88190be2..791cb583 100644 --- a/behavioral_simulator/testbench/ffn_test.py +++ b/behavioral_simulator/testbench/ffn_test.py @@ -8,7 +8,7 @@ from test_data_gen import get_weights_path, generate_and_save_random_weights from compiler.asm_templates import ffn_asm, preload_addr_reg_asm, reset_reg_asm, preload_act_asm from create_sim_env import create_sim_env -from sim_env_utils import build_sim_env +from sim_env_utils import create_mem_for_sim class LlamaFeedForward(nn.Module): @@ -177,4 +177,4 @@ def forward(self, x: Tensor) -> Tensor: ) create_sim_env(input_tensor, gen_assembly_code, golden_result, fp_preload) - build_sim_env(data_size=256, mode="behave_sim", asm=None, data=None, specified_data_order = ["act_tensor", "weight_up_layer", "weight_gate_layer", "weight_down_layer"]) \ No newline at end of file + create_mem_for_sim(data_size=256, mode="behave_sim", asm=None, data=None, specified_data_order = ["act_tensor", "weight_up_layer", "weight_gate_layer", "weight_down_layer"]) \ No newline at end of file diff --git a/behavioral_simulator/testbench/linear_test.py b/behavioral_simulator/testbench/linear_test.py index fe757e24..a9bc1aff 100644 --- a/behavioral_simulator/testbench/linear_test.py +++ b/behavioral_simulator/testbench/linear_test.py @@ -8,7 +8,7 @@ from test_data_gen import get_weights_path, generate_and_save_random_weights from compiler.asm_templates import projection_asm, preload_act_asm, reset_reg_asm, preload_addr_reg_asm from create_sim_env import create_sim_env -from sim_env_utils import build_sim_env +from sim_env_utils import create_mem_for_sim # TODOs: Need to integrate the MX quantizer here. @@ -131,7 +131,7 @@ ) create_sim_env(input_tensor, gen_assembly_code, golden_result, fp_preload) - build_sim_env(data_size=256, mode="behave_sim", asm="linear", data=None, specified_data_order = ["act_tensor", "weights"]) + create_mem_for_sim(data_size=256, mode="behave_sim", asm="linear", data=None, specified_data_order = ["act_tensor", "weights"]) print("================================================") print("Finished generating assembly code") diff --git a/behavioral_simulator/testbench/rms_test.py b/behavioral_simulator/testbench/rms_test.py index 467cd8c8..8a9a8644 100644 --- a/behavioral_simulator/testbench/rms_test.py +++ b/behavioral_simulator/testbench/rms_test.py @@ -8,7 +8,7 @@ from test_data_gen import get_weights_path, generate_and_save_random_weights from compiler.asm_templates import rms_norm_asm, preload_act_asm, reset_reg_asm, preload_addr_reg_asm from create_sim_env import create_sim_env -from sim_env_utils import build_sim_env +from sim_env_utils import create_mem_for_sim # Taken from LLAMA RMSNorm implementation @@ -133,7 +133,7 @@ def forward(self, x): ) create_sim_env(input_tensor, gen_assembly_code, golden_result, fp_preload) - build_sim_env(data_size=256, mode="behave_sim", asm="rms", data=None, specified_data_order = ["act_tensor", "weights"]) + create_mem_for_sim(data_size=256, mode="behave_sim", asm="rms", data=None, specified_data_order = ["act_tensor", "weights"]) print("================================================") print("Finished generating assembly code") diff --git a/behavioral_simulator/testbench/s_map_v_test.py b/behavioral_simulator/testbench/s_map_v_test.py index 0cfc4148..51bdec7f 100644 --- a/behavioral_simulator/testbench/s_map_v_test.py +++ b/behavioral_simulator/testbench/s_map_v_test.py @@ -9,7 +9,7 @@ from test_data_gen import get_weights_path, generate_and_save_random_weights from compiler.asm_templates import rms_norm_asm, projection_asm, preload_act_asm, reset_reg_asm, preload_addr_reg_asm from create_sim_env import create_sim_env -from sim_env_utils import build_sim_env +from sim_env_utils import create_mem_for_sim from tools.memory_mapping.hbm_addr_map import align_addr_to_hbm_bandwidth import torch.nn.functional as F @@ -65,4 +65,4 @@ gen_assembly_code += f"S_MAP_V_FP gp1, gp0, 0 \n" create_sim_env(input_tensor, weights, gen_assembly_code, golden_result, fp_preload) - build_sim_env(data_size=256, mode="behave_sim", asm="dllm", data=None, specified_data_order = ["input_tensor1","input_tensor2"]) \ No newline at end of file + create_mem_for_sim(data_size=256, mode="behave_sim", asm="dllm", data=None, specified_data_order = ["input_tensor1","input_tensor2"]) \ No newline at end of file diff --git a/behavioral_simulator/testbench/two_input_test.py b/behavioral_simulator/testbench/two_input_test.py index 811f7942..5eb3e14a 100644 --- a/behavioral_simulator/testbench/two_input_test.py +++ b/behavioral_simulator/testbench/two_input_test.py @@ -9,7 +9,7 @@ from test_data_gen import get_weights_path, generate_and_save_random_weights from compiler.asm_templates import rms_norm_asm, projection_asm, preload_act_asm, reset_reg_asm, preload_addr_reg_asm from create_sim_env import create_sim_env -from sim_env_utils import build_sim_env +from sim_env_utils import create_mem_for_sim from tools.memory_mapping.hbm_addr_map import align_addr_to_hbm_bandwidth import torch.nn.functional as F @@ -82,4 +82,4 @@ ) create_sim_env(input_tensor, weights, gen_assembly_code, golden_result, fp_preload) - build_sim_env(data_size=256, mode="behave_sim", asm="dllm", data=None, specified_data_order = ["input_tensor1","input_tensor2"]) \ No newline at end of file + create_mem_for_sim(data_size=256, mode="behave_sim", asm="dllm", data=None, specified_data_order = ["input_tensor1","input_tensor2"]) \ No newline at end of file diff --git a/behavioral_simulator/testbench/view_mem.py b/behavioral_simulator/testbench/view_mem.py index 4846ad36..d4516d15 100644 --- a/behavioral_simulator/testbench/view_mem.py +++ b/behavioral_simulator/testbench/view_mem.py @@ -3,7 +3,97 @@ import os import struct -def view_bin_file_by_row(bin_file, +def view_bin_file_by_row_int(bin_file, + row_dim, + int_width=32, + num_bytes_per_val=None, + start_row_idx=0, + load_row_size=None, + signed=True): + """ + Reads a binary file and parses each value as an integer. + + - bin_file: Path to the binary file + - row_dim: Number of values per row to group/print + - int_width: Bit width of the integer (default 32 for i32) + Common values: 8 (i8), 16 (i16), 32 (i32) + - num_bytes_per_val:Number of bytes per value (auto-calculated from int_width if None) + - start_row_idx: Starting row index + - load_row_size: If set, limit number of rows to print + - signed: If True, interpret as signed integer (default True) + + Note: The Rust code uses little-endian byte order. + If int_width is less than num_bytes_per_val * 8, the value will be masked. + """ + # Auto-calculate bytes per value if not specified + if num_bytes_per_val is None: + num_bytes_per_val = (int_width + 7) // 8 # Round up to nearest byte + + # Validate that num_bytes_per_val is sufficient + if int_width > num_bytes_per_val * 8: + raise ValueError(f"int_width ({int_width}) exceeds num_bytes_per_val * 8 ({num_bytes_per_val * 8})") + + # Create mask for the specified bit width + if int_width >= 64: + mask = (1 << 64) - 1 + else: + mask = (1 << int_width) - 1 + + with open(bin_file, "rb") as f: + data = f.read() + + num_vals = len(data) // num_bytes_per_val + total_rows = (num_vals + row_dim - 1) // row_dim + + for row_idx in range(total_rows): + if row_idx < start_row_idx: + continue + if load_row_size is not None and row_idx >= load_row_size + start_row_idx: + print("... (truncated)") + break + vals = [] + for col_idx in range(row_dim): + val_idx = row_idx * row_dim + col_idx + if val_idx >= num_vals: + break + chunk = data[val_idx * num_bytes_per_val : (val_idx + 1) * num_bytes_per_val] + if not chunk or len(chunk) < num_bytes_per_val: + vals.append(None) + continue + # Use little-endian byte order to match Rust's byte packing + if signed: + int_val = int.from_bytes(chunk, byteorder='little', signed=True) + else: + int_val = int.from_bytes(chunk, byteorder='little', signed=False) + + # Apply bit width mask + int_val = int_val & mask + + # Sign extend if signed and the sign bit is set + if signed and int_width < num_bytes_per_val * 8: + sign_bit = (int_val >> (int_width - 1)) & 1 + if sign_bit == 1: + # Sign extend: set all upper bits to 1 + sign_extend_mask = ((1 << (num_bytes_per_val * 8 - int_width)) - 1) << int_width + int_val = int_val | sign_extend_mask + # Convert to signed integer + if num_bytes_per_val == 1: + int_val = int_val if int_val < 128 else int_val - 256 + elif num_bytes_per_val == 2: + int_val = int_val if int_val < 32768 else int_val - 65536 + elif num_bytes_per_val == 4: + int_val = int_val if int_val < 2147483648 else int_val - 4294967296 + + vals.append(int_val) + print(f"Row {row_idx:3d}: ", end="") + for v in vals: + if v is not None: + print(f"{v:8d}", end=" ") + else: + print(" ", end=" ") + print() + +def view_bin_file_by_row_fp(bin_file, exp_width, man_width, row_dim, @@ -103,29 +193,29 @@ def raw_to_fp(bits_val): # VRAM uses BF16 format by default: sign=1, exponent=8, mantissa=7 (16 bits total = 2 bytes) print("Viewing VRAM dump from 0 Base Address") - view_bin_file_by_row(vram_file, exp_width=8, man_width=7, row_dim=64, num_bytes_per_val=2, start_row_idx=0, load_row_size=16) + view_bin_file_by_row_fp(vram_file, exp_width=8, man_width=7, row_dim=64, num_bytes_per_val=2, start_row_idx=0, load_row_size=8) - print("Viewing VRAM dump from 16 Base Address") - view_bin_file_by_row(vram_file, exp_width=8, man_width=7, row_dim=64, num_bytes_per_val=2, start_row_idx=16, load_row_size=32) + print("Viewing VRAM dump from 8 Base Address") + view_bin_file_by_row_int(vram_file, row_dim=64, int_width=32, start_row_idx=8, load_row_size=1) - print("Viewing VRAM dump from 48 Base Address") - view_bin_file_by_row(vram_file, exp_width=8, man_width=7, row_dim=64, num_bytes_per_val=2, start_row_idx=48, load_row_size=32) + # print("Viewing VRAM dump from 48 Base Address") + # view_bin_file_by_row_fp(vram_file, exp_width=8, man_width=7, row_dim=64, num_bytes_per_val=2, start_row_idx=48, load_row_size=32) # print("Viewing VRAM dump from Q Base Address") - # view_bin_file_by_row(vram_file, exp_width=8, man_width=7, row_dim=64, num_bytes_per_val=2, start_row_idx=0, load_row_size=16) + # view_bin_file_by_row_fp(vram_file, exp_width=8, man_width=7, row_dim=64, num_bytes_per_val=2, start_row_idx=0, load_row_size=16) # print("\nViewing VRAM dump from S Base Address") - # view_bin_file_by_row(vram_file, exp_width=8, man_width=7, row_dim=64, num_bytes_per_val=2, start_row_idx=64, load_row_size=16) + # view_bin_file_by_row_fp(vram_file, exp_width=8, man_width=7, row_dim=64, num_bytes_per_val=2, start_row_idx=64, load_row_size=16) # print("\nViewing VRAM dump from PV Base Address") - # view_bin_file_by_row(vram_file, exp_width=8, man_width=7, row_dim=64, num_bytes_per_val=2, start_row_idx=320, load_row_size=16) + # view_bin_file_by_row_fp(vram_file, exp_width=8, man_width=7, row_dim=64, num_bytes_per_val=2, start_row_idx=320, load_row_size=16) # print("\nViewing VRAM dump from O_Old Base Address") - # view_bin_file_by_row(vram_file, exp_width=8, man_width=7, row_dim=64, num_bytes_per_val=2, start_row_idx=576, load_row_size=16) + # view_bin_file_by_row_fp(vram_file, exp_width=8, man_width=7, row_dim=64, num_bytes_per_val=2, start_row_idx=576, load_row_size=16) - print("Viewing MRAM dump 0 to 7 rows (BF16 format)") - view_bin_file_by_row(mram_file, exp_width=8, man_width=7, row_dim=64, num_bytes_per_val=2, start_row_idx=0, load_row_size=8) + # print("Viewing MRAM dump 0 to 7 rows (BF16 format)") + # view_bin_file_by_row_fp(mram_file, exp_width=8, man_width=7, row_dim=64, num_bytes_per_val=2, start_row_idx=0, load_row_size=8) - print("Viewing MRAM dump 64 to 71 rows (BF16 format)") - view_bin_file_by_row(mram_file, exp_width=8, man_width=7, row_dim=64, num_bytes_per_val=2, start_row_idx=64, load_row_size=8) \ No newline at end of file + # print("Viewing MRAM dump 64 to 71 rows (BF16 format)") + # view_bin_file_by_row_fp(mram_file, exp_width=8, man_width=7, row_dim=64, num_bytes_per_val=2, start_row_idx=64, load_row_size=8) \ No newline at end of file diff --git a/src/definitions/plena_settings.toml b/src/definitions/plena_settings.toml index 18845241..f355eecf 100644 --- a/src/definitions/plena_settings.toml +++ b/src/definitions/plena_settings.toml @@ -19,6 +19,9 @@ value = 4 [CONFIG.HBM_SIZE] value = 1073741824 +[CONFIG.HBM_WIDTH] +value = 512 + [CONFIG.MATRIX_SRAM_SIZE] value = 1024 diff --git a/tools/assembler/assembly_to_binary.py b/tools/assembler/assembly_to_binary.py index 7eeec2fb..36b12aab 100644 --- a/tools/assembler/assembly_to_binary.py +++ b/tools/assembler/assembly_to_binary.py @@ -38,6 +38,7 @@ def _convert_to_binary(self, instruction): rs2 = instruction.rs2 rstride = instruction.rstride funct1 = instruction.funct1 + funct2 = instruction.funct2 imm = instruction.imm rmask = instruction.rmask binary_instruction = 0 @@ -80,6 +81,7 @@ def _convert_to_binary(self, instruction): ) elif instruction.opcode in ["H_PREFETCH_V", "H_STORE_V"]: binary_instruction = ( + (funct2 << (opw + 5 * ow)) + (funct1 << (opw + 4 * ow)) + (rstride << (opw + 3 * ow)) + (rs2 << (opw + 2 * ow)) + diff --git a/tools/memory_mapping/memory_map.py b/tools/memory_mapping/memory_map.py index 3b0e9282..e77e8764 100644 --- a/tools/memory_mapping/memory_map.py +++ b/tools/memory_mapping/memory_map.py @@ -111,7 +111,7 @@ def map_data_to_fake_hbm_for_rtl_sim(blocks, element_width, block_width, bias, b f.write("0x" + insert_bias_row + "\n") -def map_data_to_fake_hbm_for_behave_sim(blocks, element_width, block_width, bias, bias_width, directory, append=True, hbm_row_width=64): +def map_mx_data_to_hbm_for_behave_sim(blocks, element_width, block_width, bias, bias_width, directory, append=True, hbm_row_width=64): """ Maps the quantized blocks and bias to binary memory file for fake HBM memory. Writes raw bytes instead of ASCII hex text. @@ -124,6 +124,7 @@ def map_data_to_fake_hbm_for_behave_sim(blocks, element_width, block_width, bias output_file = os.path.join(directory, "hbm_for_behave_sim.bin") mode = 'ab' if append else 'wb' + for row_idx, row in enumerate(blocks): hex_row = " ".join(f"0x{val:02X}" for val in row) @@ -176,9 +177,31 @@ def map_data_to_fake_hbm_for_behave_sim(blocks, element_width, block_width, bias padding_needed = hbm_row_bias_num - len(row_buffer) row_buffer.extend(b'\x00' * padding_needed) f.write(row_buffer) - print_outputfile_contents(output_file) + # print_outputfile_contents(output_file) +def map_normal_data_to_hbm_for_behave_sim(data, data_width, directory, append=True, hbm_row_width=64): + """ + Maps the normal data to binary memory file for fake HBM memory. + """ + if not os.path.exists(directory): + os.makedirs(directory) + output_file = os.path.join(directory, "hbm_for_behave_sim.bin") + mode = 'ab' if append else 'wb' + with open(output_file, mode) as f: + row_buffer = bytearray() + for i, element in enumerate(data): + hex_str = map_scale_to_value(element, data_width) + data_bytes = hex_to_bytes(hex_str) + row_buffer.extend(data_bytes) + if len(row_buffer) >= hbm_row_width: + f.write(row_buffer[:hbm_row_width]) + row_buffer = bytearray() + if len(row_buffer) > 0: + f.write(row_buffer) + print_outputfile_contents(output_file) + + if __name__ == "__main__": directory = "../../test/weight" fake_hbm_dir = "../../test/load_mem" diff --git a/tools/sim_env_utils/__init__.py b/tools/sim_env_utils/__init__.py index fe36eb7d..3b52183e 100644 --- a/tools/sim_env_utils/__init__.py +++ b/tools/sim_env_utils/__init__.py @@ -1 +1 @@ -from .build_env import build_sim_env \ No newline at end of file +from .build_env import create_mem_for_sim \ No newline at end of file diff --git a/tools/sim_env_utils/build_env.py b/tools/sim_env_utils/build_env.py index c5e086ec..9881075b 100644 --- a/tools/sim_env_utils/build_env.py +++ b/tools/sim_env_utils/build_env.py @@ -3,16 +3,17 @@ from cfl_cocotb import SRC_PATH from cfl_tools.logger import get_logger from memory_mapping.rand_gen import Random_MXFP_Tensor_Generator -from utils.load_config import load_svh_settings +from utils.load_config import load_toml_config from pathlib import Path logger = get_logger("testbench") logger.setLevel(logging.DEBUG) -def build_sim_env(data_size=256, mode="behave_sim", asm="attn", data=None, specified_data_order = None): +def create_mem_for_sim(data_size=256, mode="behave_sim", asm="attn", data=None, specified_data_order = None): - config_settings = load_svh_settings(str(SRC_PATH / "definitions" / "configuration.svh")) - precision_settings = load_svh_settings(str(SRC_PATH / "definitions" / "precision.svh")) + plena_toml_path = str(SRC_PATH / "definitions" / "plena_settings.toml") + config_settings = load_toml_config(plena_toml_path, "CONFIG") + precision_settings = load_toml_config(plena_toml_path, "PRECISION") if mode == "behave_sim": asm_file = Path(PROJECT_PATH / "behavioral_simulator" / "testbench" / "build" / "generated_asm_code.asm") else: @@ -22,14 +23,15 @@ def build_sim_env(data_size=256, mode="behave_sim", asm="attn", data=None, speci data_config = { "tensor_size": [1, data_size], - "block_size" : [1, precision_settings["BLOCK_DIM"]], + "block_size" : [1, precision_settings["HBM_M_WEIGHT_TYPE"]["block"]], } quant_config = { - "exp_width": precision_settings["ACT_MXFP_EXP_WIDTH"], - "man_width": precision_settings["ACT_MXFP_MANT_WIDTH"], - "exp_bias_width": precision_settings["MX_SCALE_WIDTH"], + "exp_width": precision_settings["HBM_V_ACT_TYPE"]["ELEM"]["exponent"], + "man_width": precision_settings["HBM_V_ACT_TYPE"]["ELEM"]["mantissa"], + "exp_bias_width": precision_settings["HBM_V_ACT_TYPE"]["SCALE"]["exponent"], "block_size": data_config["block_size"], + "int_width": precision_settings["HBM_V_INT_TYPE"]["DATA_TYPE"]["width"], "skip_first_dim": False, } @@ -57,25 +59,31 @@ def build_sim_env(data_size=256, mode="behave_sim", asm="attn", data=None, speci else: pt_files = list(target_dir.glob("*.pt")) + list(target_dir.glob("*.pth")) - grp_blocks = [] - grp_bias = [] + memory_data_dict = {} for pt_file in pt_files: - print("loading file", pt_file) - file_raw_data = Random_MXFP_Tensor_Generator( - shape = tuple(data_config["tensor_size"]), - quant_config = quant_config, - config_settings = config_settings, - directory = Path(asm_file).parent, - filename = pt_file - ) - file_tensor = file_raw_data.tensor_load() - blocks, bias = file_raw_data.quantize_tensor(file_tensor) - for block, b in zip(blocks, bias): - grp_blocks.append(block) - grp_bias.append(b) + if pt_file.stem != "int": + print("loading file", pt_file) + file_raw_data = Random_MXFP_Tensor_Generator( + shape = tuple(data_config["tensor_size"]), + quant_config = quant_config, + config_settings = config_settings, + directory = Path(asm_file).parent, + filename = pt_file + ) + file_tensor = file_raw_data.tensor_load().squeeze(0) + blocks, bias = file_raw_data.quantize_tensor(file_tensor) + memory_data_dict["mx"] = { + "blocks": blocks, + "bias": bias + } + else: + print("loading file", pt_file) + memory_data_dict["normal"] = { + "data": torch.load(pt_file) + } # generate_golden_result(data, logger, precision_settings, data_config) - env_setup(grp_blocks, grp_bias, asm_file.parent, data_config, quant_config, hbm_row_width=config_settings["HBM_WIDTH"]) + env_setup(memory_data_dict, asm_file.parent, data_config, quant_config, hbm_row_width=config_settings["HBM_WIDTH"]["value"]) if __name__ == "__main__": - build_sim_env() + create_mem_for_sim() pass \ No newline at end of file diff --git a/tools/sim_env_utils/build_sys_tools.py b/tools/sim_env_utils/build_sys_tools.py index 81b6db14..c65afc04 100644 --- a/tools/sim_env_utils/build_sys_tools.py +++ b/tools/sim_env_utils/build_sys_tools.py @@ -8,7 +8,7 @@ from cfl_cocotb.torch_fp_conversion import pack_fp_to_bin, fp_2_bin from cfl_tools import PROJECT_PATH -from memory_mapping.memory_map import map_fp_data_to_fake_hbm, map_data_to_fake_hbm_for_rtl_sim, map_data_to_fake_hbm_for_behave_sim +from memory_mapping.memory_map import map_fp_data_to_fake_hbm, map_data_to_fake_hbm_for_rtl_sim, map_mx_data_to_hbm_for_behave_sim, map_normal_data_to_hbm_for_behave_sim from assembler.assembly_to_binary import AssemblyToBinary def generate_golden_result(data, logger, precision_settings, data_config): @@ -51,7 +51,7 @@ def generate_golden_result(data, logger, precision_settings, data_config): return qdata -def env_setup(grp_blocks, grp_bias, build_path: str, data_config, quant_config, hbm_row_width=256, test_file_name=None): +def env_setup(memory_data_dict, build_path: str, data_config, quant_config, hbm_row_width=256, test_file_name=None): isa_file_path = PROJECT_PATH / 'src' / 'definitions' / 'operation.svh' config_file_path = PROJECT_PATH / 'src' / 'definitions' / 'configuration.svh' @@ -62,29 +62,40 @@ def env_setup(grp_blocks, grp_bias, build_path: str, data_config, quant_config, assembler = AssemblyToBinary(str(isa_file_path), str(config_file_path)) assembler.generate_binary(build_path / f'{test_file_name}.asm', build_path / f'{test_file_name}.mem') - for blocks, bias in zip(grp_blocks, grp_bias): - print("blocks", blocks) - print("bias", bias) - map_data_to_fake_hbm_for_rtl_sim( - blocks =blocks, - element_width =quant_config["exp_width"] + quant_config["man_width"] + 1, - block_width =data_config["block_size"][1], - bias =bias, - bias_width =quant_config["exp_bias_width"], - combined_blk_dim=hbm_row_width // data_config["block_size"][1], - directory =build_path, - append =True, - hbm_row_width =hbm_row_width) - - map_data_to_fake_hbm_for_behave_sim( - blocks =blocks, - element_width =quant_config["exp_width"] + quant_config["man_width"] + 1, - block_width =data_config["block_size"][1], - bias =bias, - bias_width =quant_config["exp_bias_width"], - directory =build_path, - append =True, - hbm_row_width =hbm_row_width) + for type, value in memory_data_dict.items(): + if type == "mx": + grp_blocks = value["blocks"] + grp_bias = value["bias"] + for blocks, bias in zip(grp_blocks, grp_bias): + # map_data_to_fake_hbm_for_rtl_sim( + # blocks =blocks, + # element_width =quant_config["exp_width"] + quant_config["man_width"] + 1, + # block_width =data_config["block_size"][1], + # bias =bias, + # bias_width =quant_config["exp_bias_width"], + # combined_blk_dim=hbm_row_width // data_config["block_size"][1], + # directory =build_path, + # append =True, + # hbm_row_width =hbm_row_width) + + map_mx_data_to_hbm_for_behave_sim( + blocks =blocks, + element_width =quant_config["exp_width"] + quant_config["man_width"] + 1, + block_width =data_config["block_size"][1], + bias =bias, + bias_width =quant_config["exp_bias_width"], + directory =build_path, + append =True, + hbm_row_width =hbm_row_width) + elif type == "normal": + data = value["data"] + print("data", data) + map_normal_data_to_hbm_for_behave_sim( + data =data, + data_width =quant_config["int_width"], + directory =build_path, + append =True, + hbm_row_width =hbm_row_width) def parse_args(): diff --git a/tools/utils/load_config.py b/tools/utils/load_config.py index f97490de..e812d234 100644 --- a/tools/utils/load_config.py +++ b/tools/utils/load_config.py @@ -34,29 +34,7 @@ def load_json(file_path): return ml_config -def load_toml_config(file_path, mode=None): - section_to_load = ["CONFIG", "PRECISION", "INSTR"] - config = {} - +def load_toml_config(file_path, section_to_load=None): with open(file_path, "r") as f: full_toml = toml.load(f) - for section in section_to_load: - toml_config = full_toml.get(section, {}) - if not isinstance(toml_config, dict): - continue - - for param, values in toml_config.items(): - if mode is None: - config[param] = values - continue - - if isinstance(values, dict): - if mode in values: - config[param] = values[mode] - elif "value" in values: - config[param] = values["value"] - else: - config[param] = values - else: - config[param] = values - return config \ No newline at end of file + return full_toml.get(section_to_load, {}) \ No newline at end of file From 541539665b45852fe10af763821b38d1c75bf61e Mon Sep 17 00:00:00 2001 From: GeorgeWu1204 Date: Tue, 2 Dec 2025 20:48:23 +0000 Subject: [PATCH 28/34] pass check int load --- behavioral_simulator/lib/vector_sram/src/lib.rs | 5 ++++- behavioral_simulator/src/main.rs | 10 +++++----- behavioral_simulator/src/op.rs | 4 +++- behavioral_simulator/testbench/dllm1_test.py | 4 ++-- behavioral_simulator/testbench/view_mem.py | 9 ++++++--- 5 files changed, 20 insertions(+), 12 deletions(-) diff --git a/behavioral_simulator/lib/vector_sram/src/lib.rs b/behavioral_simulator/lib/vector_sram/src/lib.rs index f4cc7855..b37570a1 100644 --- a/behavioral_simulator/lib/vector_sram/src/lib.rs +++ b/behavioral_simulator/lib/vector_sram/src/lib.rs @@ -248,6 +248,8 @@ impl VectorSram { // Await the integer vector from the channel let int_vec = int_vec.await.unwrap(); + // println!("addr = {:?}", addr); + // println!("in write int_vec = {:?}", int_vec); let total_elements = int_vec.len(); let chunk_size = self.vlen as usize; let num_chunks = write_amount.min(((total_elements + chunk_size - 1) / chunk_size) as u32); @@ -307,6 +309,7 @@ impl VectorSram { /// This returns the raw binary representation of all stored data. pub async fn as_bytes(&self) -> Vec { let mut result = Vec::new(); + let mut row_idx = 0; for row_mutex in &self.rows { let mut guard = row_mutex.lock().await; @@ -323,7 +326,7 @@ impl VectorSram { RowData::Ready(bytes) => bytes.clone(), RowData::Pending(_) => unreachable!(), }; - + row_idx += 1; result.extend_from_slice(&row_bytes); } diff --git a/behavioral_simulator/src/main.rs b/behavioral_simulator/src/main.rs index f0f17504..ecebeb7f 100644 --- a/behavioral_simulator/src/main.rs +++ b/behavioral_simulator/src/main.rs @@ -1070,15 +1070,16 @@ impl Accelerator { 1 => bytes[byte_offset] as i8 as i32, 2 => { let bytes_slice = &bytes[byte_offset..byte_offset + 2]; - i16::from_le_bytes([bytes_slice[0], bytes_slice[1]]) as i32 + i16::from_le_bytes([bytes_slice[1], bytes_slice[0]]) as i32 } 4 => { let bytes_slice = &bytes[byte_offset..byte_offset + 4]; + // println!("bytes_slice = {:?}", bytes_slice.iter().map(|b| format!("{:02x}", b)).collect::>()); i32::from_le_bytes([ - bytes_slice[0], - bytes_slice[1], - bytes_slice[2], bytes_slice[3], + bytes_slice[2], + bytes_slice[1], + bytes_slice[0], ]) } _ => panic!("Unsupported integer size: {} bytes (must be 1, 2, or 4)", element_size_bytes), @@ -1086,7 +1087,6 @@ impl Accelerator { int_vec.push(int_value); } } - let _ = sender.send(int_vec); }); diff --git a/behavioral_simulator/src/op.rs b/behavioral_simulator/src/op.rs index 4acf1903..cc49ccd8 100644 --- a/behavioral_simulator/src/op.rs +++ b/behavioral_simulator/src/op.rs @@ -289,8 +289,10 @@ impl Opcode { fn vector_precision_from(funct1: u8) -> VectorPrecision { if funct1 == 0 { VectorPrecision::Activation - } else { + } else if funct1 == 1 { VectorPrecision::KeyValue + } else { + VectorPrecision::INT } } diff --git a/behavioral_simulator/testbench/dllm1_test.py b/behavioral_simulator/testbench/dllm1_test.py index 7003494a..18c23099 100644 --- a/behavioral_simulator/testbench/dllm1_test.py +++ b/behavioral_simulator/testbench/dllm1_test.py @@ -34,7 +34,7 @@ original_output = logits # Generate vlen random int32 data - int_preload = torch.randint(low=0, high=2**31, size=(vlen,), dtype=torch.int32) + int_preload = torch.randint(low=0, high=10, size=(vlen,), dtype=torch.int32) print('logits.shape= ', logits.shape) @@ -84,7 +84,7 @@ # Preload Integer Activation to the later section. gen_assembly_code += f"S_ADDI_INT gp1, gp0, {hidden_size * batch_size} \n" gen_assembly_code += f"S_ADDI_INT gp2, gp0, {int(hidden_size * batch_size * real_data_ratio)} \n" - gen_assembly_code += "H_PREFETCH_V gp1, gp2, a0, 0, 0, 1 \n" + gen_assembly_code += "H_PREFETCH_V gp1, gp2, a0, 0, 2, 1 \n" create_sim_env(input_tensor, gen_assembly_code, golden_result, fp_preload) create_mem_for_sim(data_size=256, mode="behave_sim", asm="dllm", data=None, specified_data_order = ["logits", "int"]) \ No newline at end of file diff --git a/behavioral_simulator/testbench/view_mem.py b/behavioral_simulator/testbench/view_mem.py index d4516d15..611c68bc 100644 --- a/behavioral_simulator/testbench/view_mem.py +++ b/behavioral_simulator/testbench/view_mem.py @@ -9,6 +9,7 @@ def view_bin_file_by_row_int(bin_file, num_bytes_per_val=None, start_row_idx=0, load_row_size=None, + offset=0, signed=True): """ Reads a binary file and parses each value as an integer. @@ -41,7 +42,6 @@ def view_bin_file_by_row_int(bin_file, with open(bin_file, "rb") as f: data = f.read() - num_vals = len(data) // num_bytes_per_val total_rows = (num_vals + row_dim - 1) // row_dim @@ -53,10 +53,11 @@ def view_bin_file_by_row_int(bin_file, break vals = [] for col_idx in range(row_dim): - val_idx = row_idx * row_dim + col_idx + val_idx = row_idx * row_dim + col_idx + offset if val_idx >= num_vals: break chunk = data[val_idx * num_bytes_per_val : (val_idx + 1) * num_bytes_per_val] + print("chunk = {:?}", chunk) if not chunk or len(chunk) < num_bytes_per_val: vals.append(None) continue @@ -99,6 +100,7 @@ def view_bin_file_by_row_fp(bin_file, row_dim, num_bytes_per_val=2, start_row_idx=0, + offset=0, load_row_size=None): """ Reads a binary file generated by Rust behavioral simulator, parses each value as custom FP. @@ -196,7 +198,8 @@ def raw_to_fp(bits_val): view_bin_file_by_row_fp(vram_file, exp_width=8, man_width=7, row_dim=64, num_bytes_per_val=2, start_row_idx=0, load_row_size=8) print("Viewing VRAM dump from 8 Base Address") - view_bin_file_by_row_int(vram_file, row_dim=64, int_width=32, start_row_idx=8, load_row_size=1) + # Note: TODO: now 32 bits address require offset, that's why start_row_idx is 4 instead of 8. + view_bin_file_by_row_int(vram_file, row_dim=64, int_width=32, start_row_idx=4) # print("Viewing VRAM dump from 48 Base Address") From d2f2fbd999d8b0f5c50864c2693d86fb74687844 Mon Sep 17 00:00:00 2001 From: GeorgeWu1204 Date: Fri, 5 Dec 2025 09:51:05 +0000 Subject: [PATCH 29/34] dllm1_test case update --- behavioral_simulator/testbench/view_mem.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/behavioral_simulator/testbench/view_mem.py b/behavioral_simulator/testbench/view_mem.py index 611c68bc..3445e733 100644 --- a/behavioral_simulator/testbench/view_mem.py +++ b/behavioral_simulator/testbench/view_mem.py @@ -57,7 +57,7 @@ def view_bin_file_by_row_int(bin_file, if val_idx >= num_vals: break chunk = data[val_idx * num_bytes_per_val : (val_idx + 1) * num_bytes_per_val] - print("chunk = {:?}", chunk) + # print("chunk = {:?}", chunk) if not chunk or len(chunk) < num_bytes_per_val: vals.append(None) continue @@ -199,7 +199,7 @@ def raw_to_fp(bits_val): print("Viewing VRAM dump from 8 Base Address") # Note: TODO: now 32 bits address require offset, that's why start_row_idx is 4 instead of 8. - view_bin_file_by_row_int(vram_file, row_dim=64, int_width=32, start_row_idx=4) + view_bin_file_by_row_int(vram_file, row_dim=64, int_width=32, start_row_idx=4, load_row_size=4) # print("Viewing VRAM dump from 48 Base Address") From eb11dbd2a6607f0c4c442c65ac172fb05cd1c8c3 Mon Sep 17 00:00:00 2001 From: GeorgeWu1204 Date: Sat, 6 Dec 2025 17:11:34 +0000 Subject: [PATCH 30/34] fix bugs related to mx data gen --- behavioral_simulator/testbench/linear_test.py | 56 ++++++------- behavioral_simulator/testbench/view_mem.py | 6 +- tools/memory_mapping/memory_map.py | 2 +- tools/memory_mapping/rand_gen.py | 83 ++++++++----------- tools/sim_env_utils/build_env.py | 72 +++++++++++++--- tools/sim_env_utils/build_sys_tools.py | 75 ++++++++++------- 6 files changed, 172 insertions(+), 122 deletions(-) diff --git a/behavioral_simulator/testbench/linear_test.py b/behavioral_simulator/testbench/linear_test.py index a9bc1aff..a0d4c770 100644 --- a/behavioral_simulator/testbench/linear_test.py +++ b/behavioral_simulator/testbench/linear_test.py @@ -87,19 +87,19 @@ gen_assembly_code = "; Linear Test Generation \n" # Set the addr offset for weight and bias - gen_assembly_code += preload_addr_reg_asm( - addr_reg_to_set=[1, 2], - available_registers=[1, 2], - addr_reg_val=[int(hidden_size * batch_size * real_data_ratio), int((hidden_size * (batch_size + 1) + hidden_size * hidden_size) * real_data_ratio)] - ) - - # print("hidden_size * batch_size * real_data_ratio", hidden_size * batch_size * real_data_ratio) - # print("(hidden_size * (batch_size + 1) + hidden_size * hidden_size) * real_data_ratio", (hidden_size * (batch_size + 1) + hidden_size * hidden_size) * real_data_ratio) + # gen_assembly_code += preload_addr_reg_asm( + # addr_reg_to_set=[1, 2], + # available_registers=[1, 2], + # addr_reg_val=[int(hidden_size * batch_size * real_data_ratio), int((hidden_size * (batch_size + 1) + hidden_size * hidden_size) * real_data_ratio)] + # ) + + # # print("hidden_size * batch_size * real_data_ratio", hidden_size * batch_size * real_data_ratio) + # # print("(hidden_size * (batch_size + 1) + hidden_size * hidden_size) * real_data_ratio", (hidden_size * (batch_size + 1) + hidden_size * hidden_size) * real_data_ratio) - # Reset the registers - gen_assembly_code += reset_reg_asm( - alive_registers=[1,2,3] - ) + # # Reset the registers + # gen_assembly_code += reset_reg_asm( + # alive_registers=[1,2,3] + # ) # Gen Activation Preload gen_assembly_code += preload_act_asm( @@ -113,22 +113,22 @@ stride_size=hidden_size ) - # Reset the registers - gen_assembly_code += reset_reg_asm( - alive_registers=[1,2,3,4] - ) - - gen_assembly_code += projection_asm( - mlen=64, - blen=4, - batch=4, - hidden_size=128, - alive_registers=[1,2,3,4], - w_base_hbm_offset_reg=1, - activation_base_address=0, - result_base_address=hidden_size * batch_size, - rope_enabled=False - ) + # # Reset the registers + # gen_assembly_code += reset_reg_asm( + # alive_registers=[1,2,3,4] + # ) + + # gen_assembly_code += projection_asm( + # mlen=64, + # blen=4, + # batch=4, + # hidden_size=128, + # alive_registers=[1,2,3,4], + # w_base_hbm_offset_reg=1, + # activation_base_address=0, + # result_base_address=hidden_size * batch_size, + # rope_enabled=False + # ) create_sim_env(input_tensor, gen_assembly_code, golden_result, fp_preload) create_mem_for_sim(data_size=256, mode="behave_sim", asm="linear", data=None, specified_data_order = ["act_tensor", "weights"]) diff --git a/behavioral_simulator/testbench/view_mem.py b/behavioral_simulator/testbench/view_mem.py index 3445e733..e67016fd 100644 --- a/behavioral_simulator/testbench/view_mem.py +++ b/behavioral_simulator/testbench/view_mem.py @@ -197,9 +197,9 @@ def raw_to_fp(bits_val): print("Viewing VRAM dump from 0 Base Address") view_bin_file_by_row_fp(vram_file, exp_width=8, man_width=7, row_dim=64, num_bytes_per_val=2, start_row_idx=0, load_row_size=8) - print("Viewing VRAM dump from 8 Base Address") - # Note: TODO: now 32 bits address require offset, that's why start_row_idx is 4 instead of 8. - view_bin_file_by_row_int(vram_file, row_dim=64, int_width=32, start_row_idx=4, load_row_size=4) + # print("Viewing VRAM dump from 8 Base Address") + # # Note: TODO: now 32 bits address require offset, that's why start_row_idx is 4 instead of 8. + # view_bin_file_by_row_int(vram_file, row_dim=64, int_width=32, start_row_idx=4, load_row_size=4) # print("Viewing VRAM dump from 48 Base Address") diff --git a/tools/memory_mapping/memory_map.py b/tools/memory_mapping/memory_map.py index e77e8764..22ca2e7f 100644 --- a/tools/memory_mapping/memory_map.py +++ b/tools/memory_mapping/memory_map.py @@ -177,7 +177,7 @@ def map_mx_data_to_hbm_for_behave_sim(blocks, element_width, block_width, bias, padding_needed = hbm_row_bias_num - len(row_buffer) row_buffer.extend(b'\x00' * padding_needed) f.write(row_buffer) - # print_outputfile_contents(output_file) + print_outputfile_contents(output_file) def map_normal_data_to_hbm_for_behave_sim(data, data_width, directory, append=True, hbm_row_width=64): diff --git a/tools/memory_mapping/rand_gen.py b/tools/memory_mapping/rand_gen.py index bc91174b..f392ea0d 100644 --- a/tools/memory_mapping/rand_gen.py +++ b/tools/memory_mapping/rand_gen.py @@ -139,58 +139,41 @@ def quantize_tensor(self, tensor): # INSERT_YOUR_CODE # Accept tensor as either a torch.Tensor or an OrderedDict of tensors # If it's an OrderedDict, quantize each value (assuming each is a tensor), otherwise just quantize the single tensor - - block_list = [] - scaling_list = [] - tensors = [] - - if isinstance(tensor, collections.OrderedDict): - tensors = list(tensor.values()) - elif isinstance(tensor, dict): # in case sometimes dict is used, not strictly OrderedDict - tensors = list(tensor.values()) - elif tensor == None: - return block_list, scaling_list - else: - tensors = [tensor] - - for t in tensors: - # print("quantizing tensor", t.shape) - # If the input is 1D, add a dimension to make it 2D (row vector) - if t.ndim == 1: - t = t.unsqueeze(0) - print("reshaped to", t.shape) - - bm_x, per_block_exponent, per_block_mantissa, per_block_scaling = _mx_fp_quantize_hardware( - t, - width = self.quant_config["exp_width"] + self.quant_config["man_width"] + 1, - exponent_width = self.quant_config["exp_width"], - exponent_bias_width = self.quant_config["exp_bias_width"], - block_size = self.quant_config["block_size"], - skip_first_dim = self.quant_config["skip_first_dim"], + if tensor.ndim == 1: + tensor = tensor.unsqueeze(0) + print("reshaped to", tensor.shape) + + bm_x, per_block_exponent, per_block_mantissa, per_block_scaling = _mx_fp_quantize_hardware( + tensor, + width = self.quant_config["exp_width"] + self.quant_config["man_width"] + 1, + exponent_width = self.quant_config["exp_width"], + exponent_bias_width = self.quant_config["exp_bias_width"], + block_size = self.quant_config["block_size"], + skip_first_dim = self.quant_config["skip_first_dim"], + ) + + logger.debug(f"per_block_mantissa: {per_block_mantissa.shape}") + logger.debug(f"per_block_exponent: {per_block_exponent.shape}") + logger.debug(f"per_block_quant_bias: {per_block_scaling.shape}") + + inner_block_list = [] + inner_scaling_list = [] + + for i in range(per_block_mantissa.shape[0]): + bin_block = pack_fp_to_bin( + per_block_exponent[i], + per_block_mantissa[i], + self.quant_config["exp_width"], + self.quant_config["man_width"], ) + inner_block_list.append(bin_block.tolist()) + inner_scaling_list.append(int(per_block_scaling[i])) + # note here the block_mantissa was represented as unsigned integer + # the exponent was represented as signed integer - logger.debug(f"per_block_mantissa: {per_block_mantissa.shape}") - logger.debug(f"per_block_exponent: {per_block_exponent.shape}") - logger.debug(f"per_block_quant_bias: {per_block_scaling.shape}") - - inner_block_list = [] - inner_scaling_list = [] - - for i in range(per_block_mantissa.shape[0]): - bin_block = pack_fp_to_bin( - per_block_exponent[i], - per_block_mantissa[i], - self.quant_config["exp_width"], - self.quant_config["man_width"], - ) - inner_block_list.append(bin_block.tolist()) - inner_scaling_list.append(int(per_block_scaling[i])) - # note here the block_mantissa was represented as unsigned integer - # the exponent was represented as signed integer - - block_list.append(inner_block_list) - scaling_list.append(inner_scaling_list) - return block_list, scaling_list + # block_list.append(inner_block_list) + # scaling_list.append(inner_scaling_list) + return inner_block_list, inner_scaling_list if __name__ == "__main__": diff --git a/tools/sim_env_utils/build_env.py b/tools/sim_env_utils/build_env.py index 9881075b..09e6fab8 100644 --- a/tools/sim_env_utils/build_env.py +++ b/tools/sim_env_utils/build_env.py @@ -5,10 +5,63 @@ from memory_mapping.rand_gen import Random_MXFP_Tensor_Generator from utils.load_config import load_toml_config from pathlib import Path +import torch logger = get_logger("testbench") logger.setLevel(logging.DEBUG) +class MemoryDataManager: + """Manages memory data from pt files, supporting multiple mx and int entries.""" + def __init__(self): + self.mx_entries = [] # Can have multiple mx entries + self.int_entries = [] # Can have multiple int entries + + def add_mx_file(self, filename, blocks, bias): + """Add an mx type data entry.""" + self.mx_entries.append({ + "filename": filename, + "type": "mx", + "blocks": blocks, + "bias": bias + }) + + def add_int_file(self, filename, data): + """Add an int type data entry.""" + self.int_entries.append({ + "filename": filename, + "type": "int", + "data": data + }) + + def get_all_entries(self): + """Get all entries as a list for iteration.""" + entries = [] + entries.extend(self.mx_entries) + entries.extend(self.int_entries) + return entries + + def to_dict(self): + """Convert to dictionary format for backward compatibility if needed.""" + result = {} + if self.mx_entries: + result["mx"] = { + "blocks": [entry["blocks"] for entry in self.mx_entries], + "bias": [entry["bias"] for entry in self.mx_entries] + } + if self.int_entries: + # For backward compatibility, use "normal" key + # If multiple int entries, combine them or use the last one + if len(self.int_entries) == 1: + result["normal"] = { + "data": self.int_entries[0]["data"] + } + else: + # If multiple int entries, use the last one (or could combine) + result["normal"] = { + "data": self.int_entries[-1]["data"] + } + return result + def create_mem_for_sim(data_size=256, mode="behave_sim", asm="attn", data=None, specified_data_order = None): plena_toml_path = str(SRC_PATH / "definitions" / "plena_settings.toml") @@ -52,14 +105,14 @@ def create_mem_for_sim(data_size=256, mode="behave_sim", asm="attn", data=None, grp_bias.append(bias) else: # The provided path (args.data) is a directory. Enumerate all .pt and .pth files within, - # then load and quantize all of them. Collect the results in dictionaries keyed by filename. + # then load and quantize all of them. Collect the results in a MemoryDataManager. target_dir = PROJECT_PATH / "behavioral_simulator" / "testbench" / "build" if specified_data_order is not None: pt_files = [target_dir / f"{data}.pt" for data in specified_data_order] else: pt_files = list(target_dir.glob("*.pt")) + list(target_dir.glob("*.pth")) - memory_data_dict = {} + memory_data_manager = MemoryDataManager() for pt_file in pt_files: if pt_file.stem != "int": print("loading file", pt_file) @@ -70,19 +123,16 @@ def create_mem_for_sim(data_size=256, mode="behave_sim", asm="attn", data=None, directory = Path(asm_file).parent, filename = pt_file ) - file_tensor = file_raw_data.tensor_load().squeeze(0) + file_tensor = file_raw_data.tensor_load() blocks, bias = file_raw_data.quantize_tensor(file_tensor) - memory_data_dict["mx"] = { - "blocks": blocks, - "bias": bias - } + # Multiple mx files are all kept + memory_data_manager.add_mx_file(pt_file.name, blocks, bias) else: print("loading file", pt_file) - memory_data_dict["normal"] = { - "data": torch.load(pt_file) - } + int_data = torch.load(pt_file) + memory_data_manager.add_int_file(pt_file.name, int_data) # generate_golden_result(data, logger, precision_settings, data_config) - env_setup(memory_data_dict, asm_file.parent, data_config, quant_config, hbm_row_width=config_settings["HBM_WIDTH"]["value"]) + env_setup(memory_data_manager, asm_file.parent, data_config, quant_config, hbm_row_width=config_settings["HBM_WIDTH"]["value"]) if __name__ == "__main__": create_mem_for_sim() diff --git a/tools/sim_env_utils/build_sys_tools.py b/tools/sim_env_utils/build_sys_tools.py index c65afc04..cec85d2e 100644 --- a/tools/sim_env_utils/build_sys_tools.py +++ b/tools/sim_env_utils/build_sys_tools.py @@ -51,7 +51,19 @@ def generate_golden_result(data, logger, precision_settings, data_config): return qdata -def env_setup(memory_data_dict, build_path: str, data_config, quant_config, hbm_row_width=256, test_file_name=None): +def env_setup(memory_data_manager, build_path: str, data_config, quant_config, hbm_row_width=256, test_file_name=None): + """ + Setup environment for simulation using MemoryDataManager. + Each pt file entry is processed based on its type (mx or int). + + Args: + memory_data_manager: MemoryDataManager instance or dict (for backward compatibility) + build_path: Path to build directory + data_config: Data configuration dictionary + quant_config: Quantization configuration dictionary + hbm_row_width: HBM row width + test_file_name: Optional test file name + """ isa_file_path = PROJECT_PATH / 'src' / 'definitions' / 'operation.svh' config_file_path = PROJECT_PATH / 'src' / 'definitions' / 'configuration.svh' @@ -62,34 +74,39 @@ def env_setup(memory_data_dict, build_path: str, data_config, quant_config, hbm_ assembler = AssemblyToBinary(str(isa_file_path), str(config_file_path)) assembler.generate_binary(build_path / f'{test_file_name}.asm', build_path / f'{test_file_name}.mem') - for type, value in memory_data_dict.items(): - if type == "mx": - grp_blocks = value["blocks"] - grp_bias = value["bias"] - for blocks, bias in zip(grp_blocks, grp_bias): - # map_data_to_fake_hbm_for_rtl_sim( - # blocks =blocks, - # element_width =quant_config["exp_width"] + quant_config["man_width"] + 1, - # block_width =data_config["block_size"][1], - # bias =bias, - # bias_width =quant_config["exp_bias_width"], - # combined_blk_dim=hbm_row_width // data_config["block_size"][1], - # directory =build_path, - # append =True, - # hbm_row_width =hbm_row_width) - - map_mx_data_to_hbm_for_behave_sim( - blocks =blocks, - element_width =quant_config["exp_width"] + quant_config["man_width"] + 1, - block_width =data_config["block_size"][1], - bias =bias, - bias_width =quant_config["exp_bias_width"], - directory =build_path, - append =True, - hbm_row_width =hbm_row_width) - elif type == "normal": - data = value["data"] - print("data", data) + entries = memory_data_manager.get_all_entries() + + # Process each entry based on its type + for entry in entries: + if entry["type"] == "mx": + blocks = entry["blocks"] + bias = entry["bias"] + print("blocks", blocks) + print("bias", bias) + print(f"Processing mx file: {entry.get('filename', 'unknown')}") + # map_data_to_fake_hbm_for_rtl_sim( + # blocks =blocks, + # element_width =quant_config["exp_width"] + quant_config["man_width"] + 1, + # block_width =data_config["block_size"][1], + # bias =bias, + # bias_width =quant_config["exp_bias_width"], + # combined_blk_dim=hbm_row_width // data_config["block_size"][1], + # directory =build_path, + # append =True, + # hbm_row_width =hbm_row_width) + + map_mx_data_to_hbm_for_behave_sim( + blocks =blocks, + element_width =quant_config["exp_width"] + quant_config["man_width"] + 1, + block_width =data_config["block_size"][1], + bias =bias, + bias_width =quant_config["exp_bias_width"], + directory =build_path, + append =True, + hbm_row_width =hbm_row_width) + elif entry["type"] == "int": + data = entry["data"] + print(f"Processing int file: {entry.get('filename', 'unknown')}") map_normal_data_to_hbm_for_behave_sim( data =data, data_width =quant_config["int_width"], From 99b00761371f0d66239d8bc34e960403ceeb5b50 Mon Sep 17 00:00:00 2001 From: GeorgeWu1204 Date: Sat, 6 Dec 2025 20:36:35 +0000 Subject: [PATCH 31/34] linear test pass --- behavioral_simulator/testbench/linear_test.py | 56 +++++++++---------- 1 file changed, 28 insertions(+), 28 deletions(-) diff --git a/behavioral_simulator/testbench/linear_test.py b/behavioral_simulator/testbench/linear_test.py index a0d4c770..a9bc1aff 100644 --- a/behavioral_simulator/testbench/linear_test.py +++ b/behavioral_simulator/testbench/linear_test.py @@ -87,19 +87,19 @@ gen_assembly_code = "; Linear Test Generation \n" # Set the addr offset for weight and bias - # gen_assembly_code += preload_addr_reg_asm( - # addr_reg_to_set=[1, 2], - # available_registers=[1, 2], - # addr_reg_val=[int(hidden_size * batch_size * real_data_ratio), int((hidden_size * (batch_size + 1) + hidden_size * hidden_size) * real_data_ratio)] - # ) - - # # print("hidden_size * batch_size * real_data_ratio", hidden_size * batch_size * real_data_ratio) - # # print("(hidden_size * (batch_size + 1) + hidden_size * hidden_size) * real_data_ratio", (hidden_size * (batch_size + 1) + hidden_size * hidden_size) * real_data_ratio) + gen_assembly_code += preload_addr_reg_asm( + addr_reg_to_set=[1, 2], + available_registers=[1, 2], + addr_reg_val=[int(hidden_size * batch_size * real_data_ratio), int((hidden_size * (batch_size + 1) + hidden_size * hidden_size) * real_data_ratio)] + ) + + # print("hidden_size * batch_size * real_data_ratio", hidden_size * batch_size * real_data_ratio) + # print("(hidden_size * (batch_size + 1) + hidden_size * hidden_size) * real_data_ratio", (hidden_size * (batch_size + 1) + hidden_size * hidden_size) * real_data_ratio) - # # Reset the registers - # gen_assembly_code += reset_reg_asm( - # alive_registers=[1,2,3] - # ) + # Reset the registers + gen_assembly_code += reset_reg_asm( + alive_registers=[1,2,3] + ) # Gen Activation Preload gen_assembly_code += preload_act_asm( @@ -113,22 +113,22 @@ stride_size=hidden_size ) - # # Reset the registers - # gen_assembly_code += reset_reg_asm( - # alive_registers=[1,2,3,4] - # ) - - # gen_assembly_code += projection_asm( - # mlen=64, - # blen=4, - # batch=4, - # hidden_size=128, - # alive_registers=[1,2,3,4], - # w_base_hbm_offset_reg=1, - # activation_base_address=0, - # result_base_address=hidden_size * batch_size, - # rope_enabled=False - # ) + # Reset the registers + gen_assembly_code += reset_reg_asm( + alive_registers=[1,2,3,4] + ) + + gen_assembly_code += projection_asm( + mlen=64, + blen=4, + batch=4, + hidden_size=128, + alive_registers=[1,2,3,4], + w_base_hbm_offset_reg=1, + activation_base_address=0, + result_base_address=hidden_size * batch_size, + rope_enabled=False + ) create_sim_env(input_tensor, gen_assembly_code, golden_result, fp_preload) create_mem_for_sim(data_size=256, mode="behave_sim", asm="linear", data=None, specified_data_order = ["act_tensor", "weights"]) From ec16ccce594939b1693b70edfb5455e44fa34418 Mon Sep 17 00:00:00 2001 From: GeorgeWu1204 Date: Mon, 8 Dec 2025 14:10:10 +0000 Subject: [PATCH 32/34] solve bugs for memory alignment --- behavioral_simulator/testbench/dllm1_test.py | 5 ++- .../testbench/two_input_test.py | 6 +-- tools/memory_mapping/addr_align.py | 43 +++++++++++++++++++ 3 files changed, 47 insertions(+), 7 deletions(-) create mode 100644 tools/memory_mapping/addr_align.py diff --git a/behavioral_simulator/testbench/dllm1_test.py b/behavioral_simulator/testbench/dllm1_test.py index 18c23099..736f550b 100644 --- a/behavioral_simulator/testbench/dllm1_test.py +++ b/behavioral_simulator/testbench/dllm1_test.py @@ -12,6 +12,7 @@ import torch.nn.functional as F from tools.memory_mapping.hbm_addr_map import align_addr_to_hbm_bandwidth +from tools.memory_mapping.addr_align import align_addr_up from transformers import AutoTokenizer @@ -20,7 +21,7 @@ # Testing the operation (hidden_size, hidden_size) @ (hidden_size, batch_size) vocal_size = 2 - hidden_size = 128 + hidden_size = 64 vlen = 64 batch_size = 4 preload_amount = 4 @@ -83,7 +84,7 @@ # Preload Integer Activation to the later section. gen_assembly_code += f"S_ADDI_INT gp1, gp0, {hidden_size * batch_size} \n" - gen_assembly_code += f"S_ADDI_INT gp2, gp0, {int(hidden_size * batch_size * real_data_ratio)} \n" + gen_assembly_code += f"S_ADDI_INT gp2, gp0, {align_addr_up(hidden_size * batch_size * real_data_ratio, hbm_data_width)} \n" gen_assembly_code += "H_PREFETCH_V gp1, gp2, a0, 0, 2, 1 \n" create_sim_env(input_tensor, gen_assembly_code, golden_result, fp_preload) diff --git a/behavioral_simulator/testbench/two_input_test.py b/behavioral_simulator/testbench/two_input_test.py index 5eb3e14a..704f488f 100644 --- a/behavioral_simulator/testbench/two_input_test.py +++ b/behavioral_simulator/testbench/two_input_test.py @@ -76,10 +76,6 @@ activation_offset_reg=1 ) - # Reset the registers - gen_assembly_code += reset_reg_asm( - alive_registers=[1,2,3,4] - ) - create_sim_env(input_tensor, weights, gen_assembly_code, golden_result, fp_preload) + create_sim_env(input_tensor, gen_assembly_code, golden_result, fp_preload) create_mem_for_sim(data_size=256, mode="behave_sim", asm="dllm", data=None, specified_data_order = ["input_tensor1","input_tensor2"]) \ No newline at end of file diff --git a/tools/memory_mapping/addr_align.py b/tools/memory_mapping/addr_align.py new file mode 100644 index 00000000..db6f61b9 --- /dev/null +++ b/tools/memory_mapping/addr_align.py @@ -0,0 +1,43 @@ +""" +Address alignment utilities. +""" + + +def align_addr_up(addr: int, multiple: int) -> int: + """ + Round an address up to the next multiple of the given value. + + Args: + addr: The input address to align + multiple: The alignment multiple (e.g., 64 for 64-byte alignment) + + Returns: + The address rounded up to the next multiple + """ + if multiple <= 0: + raise ValueError("multiple must be positive") + + # Round up: (addr + multiple - 1) // multiple * multiple + return int(((addr + multiple - 1) // multiple) * multiple) + + +if __name__ == "__main__": + # Test cases + test_cases = [ + (32, 64, 64), + (16, 64, 64), + (64, 64, 64), + (65, 64, 128), + (100, 64, 128), + (0, 64, 0), + (1, 64, 64), + (63, 64, 64), + (128, 64, 128), + ] + + print("Testing align_addr_up:") + for addr, multiple, expected in test_cases: + result = align_addr_up(addr, multiple) + status = "✓" if result == expected else "✗" + print(f"{status} align_addr_up({addr}, {multiple}) = {result} (expected {expected})") + From f566d8c7db1bcd4932beb675292311deec8e830f Mon Sep 17 00:00:00 2001 From: GeorgeWu1204 Date: Mon, 8 Dec 2025 16:30:15 +0000 Subject: [PATCH 33/34] fix view mem bugs for int --- behavioral_simulator/testbench/dllm1_test.py | 1 + behavioral_simulator/testbench/view_mem.py | 15 ++++++++------- 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/behavioral_simulator/testbench/dllm1_test.py b/behavioral_simulator/testbench/dllm1_test.py index 736f550b..6fe64f1a 100644 --- a/behavioral_simulator/testbench/dllm1_test.py +++ b/behavioral_simulator/testbench/dllm1_test.py @@ -36,6 +36,7 @@ # Generate vlen random int32 data int_preload = torch.randint(low=0, high=10, size=(vlen,), dtype=torch.int32) + print("int_preload", int_preload) print('logits.shape= ', logits.shape) diff --git a/behavioral_simulator/testbench/view_mem.py b/behavioral_simulator/testbench/view_mem.py index e67016fd..1113bf76 100644 --- a/behavioral_simulator/testbench/view_mem.py +++ b/behavioral_simulator/testbench/view_mem.py @@ -4,12 +4,11 @@ import struct def view_bin_file_by_row_int(bin_file, - row_dim, int_width=32, num_bytes_per_val=None, + row_size = 64 * 2, start_row_idx=0, load_row_size=None, - offset=0, signed=True): """ Reads a binary file and parses each value as an integer. @@ -42,6 +41,8 @@ def view_bin_file_by_row_int(bin_file, with open(bin_file, "rb") as f: data = f.read() + + row_dim = row_size // num_bytes_per_val num_vals = len(data) // num_bytes_per_val total_rows = (num_vals + row_dim - 1) // row_dim @@ -53,10 +54,10 @@ def view_bin_file_by_row_int(bin_file, break vals = [] for col_idx in range(row_dim): - val_idx = row_idx * row_dim + col_idx + offset + val_idx = row_idx * row_size + col_idx * num_bytes_per_val if val_idx >= num_vals: break - chunk = data[val_idx * num_bytes_per_val : (val_idx + 1) * num_bytes_per_val] + chunk = data[val_idx : val_idx + num_bytes_per_val] # print("chunk = {:?}", chunk) if not chunk or len(chunk) < num_bytes_per_val: vals.append(None) @@ -195,11 +196,11 @@ def raw_to_fp(bits_val): # VRAM uses BF16 format by default: sign=1, exponent=8, mantissa=7 (16 bits total = 2 bytes) print("Viewing VRAM dump from 0 Base Address") - view_bin_file_by_row_fp(vram_file, exp_width=8, man_width=7, row_dim=64, num_bytes_per_val=2, start_row_idx=0, load_row_size=8) + view_bin_file_by_row_fp(vram_file, exp_width=8, man_width=7, row_dim=64, num_bytes_per_val=2, start_row_idx=0, load_row_size=4) # print("Viewing VRAM dump from 8 Base Address") - # # Note: TODO: now 32 bits address require offset, that's why start_row_idx is 4 instead of 8. - # view_bin_file_by_row_int(vram_file, row_dim=64, int_width=32, start_row_idx=4, load_row_size=4) + # Note: TODO: now 32 bits address require offset, that's why start_row_idx is 4 instead of 8. + view_bin_file_by_row_int(vram_file, row_size=64 * 2, int_width=32, start_row_idx=4, num_bytes_per_val=4, load_row_size=4) # print("Viewing VRAM dump from 48 Base Address") From f90b6b539d2cf72b78e70d45ca44d09fd7e0524b Mon Sep 17 00:00:00 2001 From: GeorgeWu1204 Date: Tue, 9 Dec 2025 12:29:04 +0000 Subject: [PATCH 34/34] fix error for int load, add strided load for dllm_test --- behavioral_simulator/testbench/dllm1_test.py | 15 ++++++++------- behavioral_simulator/testbench/view_mem.py | 4 ++-- tools/memory_mapping/memory_map.py | 1 + 3 files changed, 11 insertions(+), 9 deletions(-) diff --git a/behavioral_simulator/testbench/dllm1_test.py b/behavioral_simulator/testbench/dllm1_test.py index 6fe64f1a..0c466729 100644 --- a/behavioral_simulator/testbench/dllm1_test.py +++ b/behavioral_simulator/testbench/dllm1_test.py @@ -21,7 +21,7 @@ # Testing the operation (hidden_size, hidden_size) @ (hidden_size, batch_size) vocal_size = 2 - hidden_size = 64 + hidden_size = 128 vlen = 64 batch_size = 4 preload_amount = 4 @@ -35,7 +35,7 @@ original_output = logits # Generate vlen random int32 data - int_preload = torch.randint(low=0, high=10, size=(vlen,), dtype=torch.int32) + int_preload = torch.randint(low=0, high=10, size=(batch_size, hidden_size), dtype=torch.int32) print("int_preload", int_preload) @@ -80,13 +80,14 @@ hidden_size=hidden_size, alive_registers=[1,2,3], # [a_actual_register, set_stride_register, result_register] act_vram_offset=0, - activation_offset_reg=0 + activation_offset_reg=0, + stride_size=hidden_size ) - # Preload Integer Activation to the later section. - gen_assembly_code += f"S_ADDI_INT gp1, gp0, {hidden_size * batch_size} \n" - gen_assembly_code += f"S_ADDI_INT gp2, gp0, {align_addr_up(hidden_size * batch_size * real_data_ratio, hbm_data_width)} \n" - gen_assembly_code += "H_PREFETCH_V gp1, gp2, a0, 0, 2, 1 \n" + # # Preload Integer Activation to the later section. + # gen_assembly_code += f"S_ADDI_INT gp1, gp0, {hidden_size * batch_size} \n" + # gen_assembly_code += f"S_ADDI_INT gp2, gp0, {align_addr_up(hidden_size * batch_size * real_data_ratio, hbm_data_width)} \n" + # gen_assembly_code += "H_PREFETCH_V gp1, gp2, a0, 0, 2, 1 \n" create_sim_env(input_tensor, gen_assembly_code, golden_result, fp_preload) create_mem_for_sim(data_size=256, mode="behave_sim", asm="dllm", data=None, specified_data_order = ["logits", "int"]) \ No newline at end of file diff --git a/behavioral_simulator/testbench/view_mem.py b/behavioral_simulator/testbench/view_mem.py index 1113bf76..c1dbedad 100644 --- a/behavioral_simulator/testbench/view_mem.py +++ b/behavioral_simulator/testbench/view_mem.py @@ -196,11 +196,11 @@ def raw_to_fp(bits_val): # VRAM uses BF16 format by default: sign=1, exponent=8, mantissa=7 (16 bits total = 2 bytes) print("Viewing VRAM dump from 0 Base Address") - view_bin_file_by_row_fp(vram_file, exp_width=8, man_width=7, row_dim=64, num_bytes_per_val=2, start_row_idx=0, load_row_size=4) + view_bin_file_by_row_fp(vram_file, exp_width=8, man_width=7, row_dim=64, num_bytes_per_val=2, start_row_idx=0, load_row_size=8) # print("Viewing VRAM dump from 8 Base Address") # Note: TODO: now 32 bits address require offset, that's why start_row_idx is 4 instead of 8. - view_bin_file_by_row_int(vram_file, row_size=64 * 2, int_width=32, start_row_idx=4, num_bytes_per_val=4, load_row_size=4) + view_bin_file_by_row_int(vram_file, row_size=64 * 2, int_width=32, start_row_idx=8, num_bytes_per_val=4, load_row_size=4) # print("Viewing VRAM dump from 48 Base Address") diff --git a/tools/memory_mapping/memory_map.py b/tools/memory_mapping/memory_map.py index 22ca2e7f..b6831801 100644 --- a/tools/memory_mapping/memory_map.py +++ b/tools/memory_mapping/memory_map.py @@ -188,6 +188,7 @@ def map_normal_data_to_hbm_for_behave_sim(data, data_width, directory, append=Tr os.makedirs(directory) output_file = os.path.join(directory, "hbm_for_behave_sim.bin") mode = 'ab' if append else 'wb' + data = data.flatten() with open(output_file, mode) as f: row_buffer = bytearray() for i, element in enumerate(data):