Track the conv2d operations in stable-diffusion. (#431)

* Track the conv2d operations in stable-diffusion.

* Add more tracing to stable-diffusion.

* Also trace the resnet bits.

* Trace the attention blocks.

* Also trace the attention inner part.

* Small tweak.

candle-examples/examples/stable-diffusion/utils.rs

@@ -29,3 +29,29 @@ pub fn save_image<P: AsRef<std::path::Path>>(img: &Tensor, p: P) -> Result<()> {
     image.save(p).map_err(candle::Error::wrap)?;
     Ok(())
 }
+
+// Wrap the conv2d op to provide some tracing.
+#[derive(Debug)]
+pub struct Conv2d {
+    inner: candle_nn::Conv2d,
+    span: tracing::Span,
+}
+
+impl Conv2d {
+    pub fn forward(&self, x: &Tensor) -> Result<Tensor> {
+        let _enter = self.span.enter();
+        self.inner.forward(x)
+    }
+}
+
+pub fn conv2d(
+    in_channels: usize,
+    out_channels: usize,
+    kernel_size: usize,
+    cfg: candle_nn::Conv2dConfig,
+    vs: candle_nn::VarBuilder,
+) -> Result<Conv2d> {
+    let span = tracing::span!(tracing::Level::TRACE, "conv2d");
+    let inner = candle_nn::conv2d(in_channels, out_channels, kernel_size, cfg, vs)?;
+    Ok(Conv2d { inner, span })
+}