Repository: coreylowman/dfdx (language: rust)

Issue:
Reuse allocations during inference forward
Recent changes #664 have added in place kernel methods to a bunch of kernel operations, that attempt to reuse the input's allocation if possible. The changes were mainly aimed at training, but there is an additional optimization we can add for inference. We can check if the tape is tracing `T::OWNS_TAPE`, and if so, try to reuse the memory.


Discussion / hints:

Fix patch:
diff --git a/benches/batchnorm2d.rs b/benches/batchnorm2d.rs
index 412171154..292e03a06 100644
--- a/benches/batchnorm2d.rs
+++ b/benches/batchnorm2d.rs
@@ -26,6 +26,11 @@ fn main() {
     loop {
         let img: Tensor<InputShape, Dtype, _> = dev.sample_normal();
 
+        let start = Instant::now();
+        let _ = m.forward(img.clone());
+        dev.synchronize();
+        let infer_dur = start.elapsed();
+
         let start = Instant::now();
         let out = m.forward_mut(img.traced(grads));
         let loss = out.square().mean();
@@ -36,6 +41,6 @@ fn main() {
         grads = loss.backward();
         dev.synchronize();
         let bwd_dur = start.elapsed();
-        println!("fwd={:?} bwd={:?}", fwd_dur, bwd_dur);
+        println!("infer={infer_dur:?}, fwd={fwd_dur:?} bwd={bwd_dur:?}");
     }
 }
diff --git a/benches/conv2d.rs b/benches/conv2d.rs
index 58eaf686a..bf9f44389 100644
--- a/benches/conv2d.rs
+++ b/benches/conv2d.rs
@@ -28,6 +28,11 @@ fn main() {
     loop {
         let img: Tensor<InputShape, Dtype, _> = dev.sample_normal();
 
+        let start = Instant::now();
+        let _ = m.forward(img.clone());
+        dev.synchronize();
+        let infer_dur = start.elapsed();
+
         let start = Instant::now();
         let out = m.forward_mut(img.leaky_traced());
         let loss = out.square().mean();
@@ -38,7 +43,8 @@ fn main() {
         let _ = loss.backward();
         dev.synchronize();
         let bwd_dur = start.elapsed();
-        println!("fwd={:?} bwd={:?}", fwd_dur, bwd_dur);
+
+        println!("infer={infer_dur:?}, fwd={fwd_dur:?} bwd={bwd_dur:?}");
     }
 }
 
diff --git a/benches/softmax.rs b/benches/softmax.rs
index 517615aa1..0bd83f6da 100644
--- a/benches/softmax.rs
+++ b/benches/softmax.rs
@@ -23,10 +23,16 @@ fn main() {
 
     loop {
         let img: Tensor<InputShape, Dtype, _> = dev.sample_normal();
-        let grads = Gradients::leaky();
 
         let start = Instant::now();
-        let y = img.traced(grads).softmax::<Ax>();
+        let _ = img.softmax::<Ax>();
+        dev.synchronize();
+        let infer_dur = start.elapsed();
+
+        let img: Tensor<InputShape, Dtype, _> = dev.sample_normal();
+
+        let start = Instant::now();
+        let y = img.leaky_traced().softmax::<Ax>();
         dev.synchronize();
         let fwd_dur = start.elapsed();
 
@@ -34,6 +40,7 @@ fn main() {
         let _ = y.sum().backward();
         dev.synchronize();
         let bwd_dur = start.elapsed();
-        println!("fwd={:?} bwd={:?}", fwd_dur, bwd_dur);
+
+        println!("infer={infer_dur:?}, fwd={fwd_dur:?} bwd={bwd_dur:?}");
     }
 }
diff --git a/benches/sum.rs b/benches/sum.rs
index 2e490d5d1..c91e7a14b 100644
--- a/benches/sum.rs
+++ b/benches/sum.rs
@@ -24,6 +24,11 @@ fn main() {
         let img: Tensor<InputShape, Dtype, _> = dev.sample_normal();
         let grads = Gradients::leaky();
 
+        let start = Instant::now();
+        let _ = img.clone().sum::<(), _>();
+        dev.synchronize();
+        let infer_dur = start.elapsed();
+
         let start = Instant::now();
         let y = img.traced(grads).sum();
         dev.synchronize();
@@ -33,6 +38,7 @@ fn main() {
         let _ = y.backward();
         dev.synchronize();
         let bwd_dur = start.elapsed();
-        println!("fwd={:?} bwd={:?}", fwd_dur, bwd_dur);
+
+        println!("infer={infer_dur:?}, fwd={fwd_dur:?} bwd={bwd_dur:?}");
     }
 }
diff --git a/src/nn/batchnorm2d.rs b/src/nn/batchnorm2d.rs
index 32d0cee22..5640ac623 100644
--- a/src/nn/batchnorm2d.rs
+++ b/src/nn/batchnorm2d.rs
@@ -79,12 +79,13 @@ where
     let shape = *x.shape();
 
     // statistics for normalizing
-    let std = (var.clone().try_add(epsilon)?).try_sqrt()?;
+    let std = var.clone().try_add(epsilon)?.try_sqrt()?;
+
+    let scale = scale.clone().try_div(std)?.try_broadcast_like(&shape)?;
 
     // normalize & affine
     let x = x.try_sub(mean.clone().try_broadcast_like(&shape)?)?;
-    let x = x.try_div(std.try_broadcast_like(&shape)?)?;
-    let x = x.try_mul(scale.clone().try_broadcast_like(&shape)?)?;
+    let x = x.try_mul(scale)?;
     x.try_add(bias.clone().try_broadcast_like(&shape)?)
 }
 
diff --git a/src/tensor_ops/utilities/ops.rs b/src/tensor_ops/utilities/ops.rs
index 9f9e4d056..374ef8cfd 100644
--- a/src/tensor_ops/utilities/ops.rs
+++ b/src/tensor_ops/utilities/ops.rs
@@ -53,7 +53,7 @@ pub(crate) fn try_unary_op<
     let (inp, mut tape) = inp.split_tape();
     let inp_ghost = inp.ghost();
     let dev = inp.device.clone();
-    if D::BACKWARD_WITHOUT_DATA {
+    if !T::OWNS_TAPE || D::BACKWARD_WITHOUT_DATA {
         let out = inp_ghost.dev.forward(op.clone(), Err(inp))?;
         let out_ghost = out.ghost();
         tape.add_backward_op(move |grads| {
@@ -105,7 +105,7 @@ pub(crate) fn try_binary_op<
     let lhs_ghost = lhs.ghost();
     let rhs_ghost = rhs.ghost();
     let mut tape = ltape.merge(rtape);
-    if D::BACKWARD_WITHOUT_DATA {
+    if !LhsTape::OWNS_TAPE || D::BACKWARD_WITHOUT_DATA {
         let out = lhs_ghost.dev.forward(op, Err(lhs), Err(rhs))?;
         let out_ghost = out.ghost();
         tape.add_backward_op(move |grads| {


Test patch: