From 86916a8fa05ad77576b5488a01660d701a9a0efb Mon Sep 17 00:00:00 2001 From: NiccoloN Date: Thu, 9 Apr 2026 14:19:16 +0200 Subject: [PATCH] add softmax support to pim-simulator --- .../src/lib/instruction_set/isa.rs | 42 +++++++++++++++++++ .../src/lib/json_to_instruction/json_isa.rs | 22 ++++++++++ .../src/lib/memory_manager/type_traits.rs | 17 ++++++++ .../pim-simulator/src/lib/tracing/disable.rs | 16 +++++++ .../src/lib/tracing/tracing_isa.rs | 29 +++++++++++++ 5 files changed, 126 insertions(+) diff --git a/backend-simulators/pim/pim-simulator/src/lib/instruction_set/isa.rs b/backend-simulators/pim/pim-simulator/src/lib/instruction_set/isa.rs index fc9d5be..2eff04b 100644 --- a/backend-simulators/pim/pim-simulator/src/lib/instruction_set/isa.rs +++ b/backend-simulators/pim/pim-simulator/src/lib/instruction_set/isa.rs @@ -52,6 +52,7 @@ static NAMES: LazyLock> = LazyLock::new(|| { add_name_simd!(hash, vrelu); add_name_simd!(hash, vtanh); add_name_simd!(hash, vsigm); + add_name_simd!(hash, vsoftmax); add_name!(hash, vmv); add_name!(hash, vrsu); add_name!(hash, vrsl); @@ -177,6 +178,7 @@ static SIMD: LazyLock>> add_simd_to_map!(storage, vrelu); add_simd_to_map!(storage, vtanh); add_simd_to_map!(storage, vsigm); + add_simd_to_map!(storage, vsoftmax); add_simd_to_map!(storage, mvmul); storage }); @@ -626,6 +628,46 @@ where Ok(InstructionStatus::Completed) } +pub fn vsoftmax(cores: &mut CPU, data: InstructionData) -> Result { + panic!("You are calling a placeholder, the real call is the generic version"); +} + +pub(super) fn vsoftmax_impl(cores: &mut CPU, data: InstructionData) -> Result +where + [F]: UpcastSlice, + T: UpcastDestTraits + MemoryStorable, + F: UpcastDestTraits + MemoryStorable + From, +{ + TRACER.lock().unwrap().pre_vsoftmax::(cores, data); + let (core_indx, rd, r1, r2, imm_len, offset_select, offset_value) = + data.get_core_rd_r1_r2_immlen_offset(); + let core = cores.core(core_indx); + let r1_val = core.register(r1); + let rd_val = core.register(rd); + let r1_val = add_offset_r1(r1_val, offset_select, offset_value); + let rd_val = add_offset_rd(rd_val, offset_select, offset_value); + let loads = core.reserve_load(r1_val, imm_len)?.execute_load::()?; + let load1 = loads[0]; + ensure!(!load1.is_empty(), "vsoftmax does not support empty vectors"); + let max_val = load1 + .iter() + .copied() + .reduce(|a, b| if a > b { a } else { b }) + .unwrap(); + let exp_values: Vec = load1.iter().map(|&a| (a - max_val).exp()).collect(); + let sum = exp_values + .iter() + .copied() + .reduce(|a, b| a + b) + .unwrap(); + ensure!(sum > 0.0.into(), "vsoftmax normalization sum must be positive"); + let res: Vec = exp_values.iter().map(|&a| a / sum).collect(); + let res_up: Cow<[T]> = res.as_slice().up(); + core.execute_store(rd_val, res_up.as_ref()); + TRACER.lock().unwrap().post_vsoftmax::(cores, data); + Ok(InstructionStatus::Completed) +} + pub fn vmv(cores: &mut CPU, data: InstructionData) -> Result { todo!() } diff --git a/backend-simulators/pim/pim-simulator/src/lib/json_to_instruction/json_isa.rs b/backend-simulators/pim/pim-simulator/src/lib/json_to_instruction/json_isa.rs index f2885b8..e20d3c9 100644 --- a/backend-simulators/pim/pim-simulator/src/lib/json_to_instruction/json_isa.rs +++ b/backend-simulators/pim/pim-simulator/src/lib/json_to_instruction/json_isa.rs @@ -40,6 +40,7 @@ static SIMD: LazyLock> = LazyLock::new(|| { add_to_json_map!(storage, vrelu); add_to_json_map!(storage, vtanh); add_to_json_map!(storage, vsigm); + add_to_json_map!(storage, vsoftmax); add_to_json_map!(storage, vmv); add_to_json_map!(storage, vrsu); add_to_json_map!(storage, vrsl); @@ -417,6 +418,27 @@ fn json_to_vsigm( Ok(()) } +fn json_to_vsoftmax( + inst_builder: &mut InstructionsBuilder, + inst_data_builder: &mut InstructionDataBuilder, + json: &Value, +) -> Result<()> { + let json = json.as_object().expect("Not an object"); + assert_eq!("vsoftmax", json_str!(json, "op")); + let rd = json_i64!(json, "rd") as i32; + let rs1 = json_i64!(json, "rs1") as i32; + let len = json_i64!(json, "len") as i32; + let (offset_select, offset_value) = json_to_offset(json.get("offset").unwrap()); + inst_data_builder + .set_rd(rd) + .set_r1(rs1) + .set_imm_len(len) + .set_offset_select(offset_select) + .set_offset_value(offset_value); + inst_builder.make_inst(vsoftmax, inst_data_builder.build()); + Ok(()) +} + fn json_to_vmv( inst_builder: &mut InstructionsBuilder, inst_data_builder: &mut InstructionDataBuilder, diff --git a/backend-simulators/pim/pim-simulator/src/lib/memory_manager/type_traits.rs b/backend-simulators/pim/pim-simulator/src/lib/memory_manager/type_traits.rs index 151209e..6778dfa 100644 --- a/backend-simulators/pim/pim-simulator/src/lib/memory_manager/type_traits.rs +++ b/backend-simulators/pim/pim-simulator/src/lib/memory_manager/type_traits.rs @@ -67,6 +67,22 @@ impl HasSigm for f64 { } } +pub trait HasExp { + fn exp(self) -> Self; +} + +impl HasExp for f32 { + fn exp(self) -> Self { + self.exp() + } +} + +impl HasExp for f64 { + fn exp(self) -> Self { + self.exp() + } +} + pub trait TryToUsize: TryInto @@ -112,6 +128,7 @@ pub trait UpcastDestTraits: + PartialOrd + HasTanh + HasSigm + + HasExp + FromUsize { } diff --git a/backend-simulators/pim/pim-simulator/src/lib/tracing/disable.rs b/backend-simulators/pim/pim-simulator/src/lib/tracing/disable.rs index 47729ad..f31bf38 100644 --- a/backend-simulators/pim/pim-simulator/src/lib/tracing/disable.rs +++ b/backend-simulators/pim/pim-simulator/src/lib/tracing/disable.rs @@ -248,6 +248,22 @@ impl Trace { { } + pub fn pre_vsoftmax(&mut self, cores: &mut CPU, data: InstructionData) + where + [F]: UpcastSlice, + T: UpcastDestTraits + MemoryStorable, + F: UpcastDestTraits + MemoryStorable + From, + { + } + + pub fn post_vsoftmax(&mut self, cores: &mut CPU, data: InstructionData) + where + [F]: UpcastSlice, + T: UpcastDestTraits + MemoryStorable, + F: UpcastDestTraits + MemoryStorable + From, + { + } + ///////////////////////////////////////////////////////////////// /////Communication/synchronization Instructions///////////////// ///////////////////////////////////////////////////////////////// diff --git a/backend-simulators/pim/pim-simulator/src/lib/tracing/tracing_isa.rs b/backend-simulators/pim/pim-simulator/src/lib/tracing/tracing_isa.rs index 55026ca..6d0177b 100644 --- a/backend-simulators/pim/pim-simulator/src/lib/tracing/tracing_isa.rs +++ b/backend-simulators/pim/pim-simulator/src/lib/tracing/tracing_isa.rs @@ -956,6 +956,35 @@ impl Trace { // Ok(InstructionStatus::Completed) } + pub fn pre_vsoftmax(&mut self, cores: &mut CPU, data: InstructionData) + where + [F]: UpcastSlice, + T: UpcastDestTraits + MemoryStorable, + F: UpcastDestTraits + MemoryStorable + From, + { + let (core_indx, rd, r1, r2, imm_len, offset_select, offset_value) = + data.get_core_rd_r1_r2_immlen_offset(); + let file: &mut File = self + .out_files + .get_mut(core_indx as usize) + .expect("File at index not found"); + writeln!(file, "\t\tVSOFTMAX\t\t"); + } + + pub fn post_vsoftmax(&mut self, cores: &mut CPU, data: InstructionData) + where + [F]: UpcastSlice, + T: UpcastDestTraits + MemoryStorable, + F: UpcastDestTraits + MemoryStorable + From, + { + let (core_indx, rd, r1, r2, imm_len, offset_select, offset_value) = + data.get_core_rd_r1_r2_immlen_offset(); + let file: &mut File = self + .out_files + .get_mut(core_indx as usize) + .expect("File at index not found"); + } + ///////////////////////////////////////////////////////////////// /////Communication/synchronization Instructions///////////////// /////////////////////////////////////////////////////////////////