Merge pull request #103 from outbrain/caching-forward-pass-with-readi…

…ng-no-simd

Introduces the concept of a forward pass cache

benchmark/run_with_plots_intel.sh

@@ -1,4 +1,6 @@
 #!/bin/bash
+
 export RUSTFLAGS="-C opt-level=3 -C target-cpu=skylake"
 cargo build --release
+
 python3 benchmark.py fw all True

benchmark/run_without_plots.sh

@@ -1,3 +1,5 @@
 #!/bin/bash
+
 cargo build --release
+
 python3 benchmark.py fw all False

examples/ffm/run_fw_with_prediction_tests.sh

@@ -72,8 +72,8 @@ mkdir -p predictions
 
 echo "Building FW"
 (
-	cd $PROJECT_ROOT
-	cargo build --release
+  cd $PROJECT_ROOT
+  cargo build --release
 )
 
 # Change this to your preference if required - this is tailored for the toy example

src/block_helpers.rs

@@ -7,7 +7,7 @@ use crate::feature_buffer;
 use crate::graph;
 use crate::optimizer::OptimizerSGD;
 use crate::port_buffer;
-use crate::regressor::BlockTrait;
+use crate::regressor::{BlockCache, BlockTrait};
 use std::cmp::min;
 use std::mem;
 use std::slice;
@@ -168,6 +168,33 @@ pub fn slearn2<'a>(
     pb.observations[0]
 }
 
+pub fn ssetup_cache2<'a>(
+    bg: &mut graph::BlockGraph,
+    cache_fb: &feature_buffer::FeatureBuffer,
+    caches: &mut Vec<BlockCache>,
+) {
+    let (create_block_run, create_further_blocks) = bg.blocks_final.split_at_mut(1);
+    create_block_run[0].create_forward_cache(create_further_blocks, caches);
+
+    let (prepare_block_run, prepare_further_blocks) = bg.blocks_final.split_at_mut(1);
+    prepare_block_run[0].prepare_forward_cache(prepare_further_blocks, cache_fb, caches.as_mut_slice());
+}
+
+pub fn spredict2_with_cache<'a>(
+    bg: &mut graph::BlockGraph,
+    cache_fb: &feature_buffer::FeatureBuffer,
+    fb: &feature_buffer::FeatureBuffer,
+    pb: &mut port_buffer::PortBuffer,
+    caches: &[BlockCache],
+) -> f32 {
+    pb.reset();
+    let (block_run, further_blocks) = bg.blocks_final.split_at(1);
+    block_run[0].forward_with_cache(further_blocks, fb, pb, caches);
+
+    let prediction_probability = pb.observations[0];
+    return prediction_probability;
+}
+
 pub fn spredict2<'a>(
     bg: &mut graph::BlockGraph,
     fb: &feature_buffer::FeatureBuffer,
@@ -206,3 +233,45 @@ pub fn forward(
         None => {}
     }
 }
+
+#[inline(always)]
+pub fn forward_with_cache(
+    further_blocks: &[Box<dyn BlockTrait>],
+    fb: &feature_buffer::FeatureBuffer,
+    pb: &mut port_buffer::PortBuffer,
+    caches: &[BlockCache],
+) {
+    match further_blocks.split_first() {
+        Some((next_regressor, further_blocks)) => next_regressor
+            .forward_with_cache(further_blocks, fb, pb, caches),
+        None => {}
+    }
+}
+
+
+#[inline(always)]
+pub fn prepare_forward_cache(
+    further_blocks: &mut [Box<dyn BlockTrait>],
+    fb: &feature_buffer::FeatureBuffer,
+    caches: &mut [BlockCache],
+) {
+    match further_blocks.split_first_mut() {
+        Some((next_regressor, further_blocks)) => next_regressor
+            .prepare_forward_cache(further_blocks, fb, caches),
+        None => {}
+    }
+}
+
+#[inline(always)]
+pub fn create_forward_cache(
+    further_blocks: &mut [Box<dyn BlockTrait>],
+    caches: &mut Vec<BlockCache>
+) {
+    match further_blocks.split_first_mut() {
+        Some((next_regressor, further_blocks)) => next_regressor
+            .create_forward_cache(further_blocks, caches),
+        None => {}
+    }
+}
+
+

src/block_loss_functions.rs

@@ -7,6 +7,9 @@ use crate::graph;
 use crate::port_buffer;
 use crate::regressor;
 use regressor::BlockTrait;
+use crate::feature_buffer::FeatureBuffer;
+use crate::port_buffer::PortBuffer;
+use crate::regressor::BlockCache;
 
 #[inline(always)]
 pub fn logistic(t: f32) -> f32 {
@@ -37,6 +40,47 @@ pub fn new_logloss_block(
     Ok(block_outputs.pop().unwrap())
 }
 
+impl BlockSigmoid {
+
+    #[inline(always)]
+    fn internal_forward(
+        &self,
+        fb: &feature_buffer::FeatureBuffer,
+        pb: &mut port_buffer::PortBuffer,
+    ) {
+        unsafe {
+            debug_assert!(self.input_offset != usize::MAX);
+            debug_assert!(self.output_offset != usize::MAX);
+            let wsum: f32 = {
+                let myslice = pb
+                    .tape
+                    .get_unchecked(self.input_offset..(self.input_offset + self.num_inputs));
+                myslice.iter().sum()
+            };
+
+            let prediction_probability: f32;
+            if wsum.is_nan() {
+                log::warn!(
+                        "NAN prediction in example {}, forcing 0.0",
+                        fb.example_number
+                    );
+                prediction_probability = logistic(0.0);
+            } else if wsum < -50.0 {
+                prediction_probability = logistic(-50.0);
+            } else if wsum > 50.0 {
+                prediction_probability = logistic(50.0);
+            } else {
+                prediction_probability = logistic(wsum);
+            }
+
+            pb.tape[self.output_offset] = prediction_probability;
+            if self.copy_to_result {
+                pb.observations.push(prediction_probability);
+            }
+        }
+    }
+}
+
 impl BlockTrait for BlockSigmoid {
     fn as_any(&mut self) -> &mut dyn Any {
         self
@@ -122,39 +166,19 @@ impl BlockTrait for BlockSigmoid {
         fb: &feature_buffer::FeatureBuffer,
         pb: &mut port_buffer::PortBuffer,
     ) {
-        unsafe {
-            /*        if further_blocks.len() != 0 {
-                panic!("RegSigmoid can only be at the end of the chain!");
-            }*/
-            debug_assert!(self.input_offset != usize::MAX);
-            debug_assert!(self.output_offset != usize::MAX);
-            let wsum: f32 = {
-                let myslice = &pb
-                    .tape
-                    .get_unchecked(self.input_offset..(self.input_offset + self.num_inputs));
-                myslice.iter().sum()
-            };
-
-            let prediction_probability: f32;
-            if wsum.is_nan() {
-                log::warn!(
-                    "NAN prediction in example {}, forcing 0.0",
-                    fb.example_number
-                );
-                prediction_probability = logistic(0.0);
-            } else if wsum < -50.0 {
-                prediction_probability = logistic(-50.0);
-            } else if wsum > 50.0 {
-                prediction_probability = logistic(50.0);
-            } else {
-                prediction_probability = logistic(wsum);
-            }
+        self.internal_forward(fb, pb);
+        block_helpers::forward(further_blocks, fb, pb);
+    }
 
-            pb.tape[self.output_offset] = prediction_probability;
-            if self.copy_to_result {
-                pb.observations.push(prediction_probability);
-            }
-            block_helpers::forward(further_blocks, fb, pb);
-        }
+    fn forward_with_cache(
+        &self,
+        further_blocks: &[Box<dyn BlockTrait>],
+        fb: &FeatureBuffer,
+        pb: &mut PortBuffer,
+        caches: &[BlockCache],
+    ) {
+        self.internal_forward(fb, pb);
+        block_helpers::forward_with_cache(further_blocks, fb, pb, caches);
     }
+
 }