Add with_signpost_interval layer (#711)

* Add with_signpost_interval layer

This layer wraps a layer, adding macOS interval signposts for the
forward and backward pass. These intervals can then be visualized
in the macOS Instruments.app timeline.

* Fix reference in api-layers.md

Co-authored-by: Madeesh Kannan <[email protected]>

* End message is optional since signpost 0.0.3

* with_signpost_interval: also wrap init callback

* docs: we wrap init as well

* Add documentation fixes

Suggested by @svlandeg.

Co-authored-by: Madeesh Kannan <[email protected]>

Author: danieldk

thinc/api.py

@@ -40,6 +40,7 @@
 from .layers import with_reshape, with_getitem, strings2arrays, list2array
 from .layers import list2ragged, ragged2list, list2padded, padded2list, remap_ids
 from .layers import array_getitem, with_cpu, with_debug, with_nvtx_range
+from .layers import with_signpost_interval
 from .layers import tuplify
 
 from .layers import reduce_first, reduce_last, reduce_max, reduce_mean, reduce_sum

thinc/compat.py

@@ -78,4 +78,13 @@
     h5py = None
 
 
+try:  # pragma: no cover
+    import os_signpost
+
+    has_os_signpost = True
+except ImportError:
+    os_signpost = None
+    has_os_signpost = False
+
+
 has_gpu = has_cupy_gpu or has_torch_mps_gpu

thinc/layers/__init__.py

@@ -71,6 +71,7 @@
 from .with_getitem import with_getitem
 from .with_debug import with_debug
 from .with_nvtx_range import with_nvtx_range
+from .with_signpost_interval import with_signpost_interval
 
 
 __all__ = [
@@ -128,5 +129,6 @@
     "with_flatten",
     "with_debug",
     "with_nvtx_range",
+    "with_signpost_interval",
     "remap_ids",
 ]

thinc/layers/with_signpost_interval.py

@@ -0,0 +1,50 @@
+from typing import Optional, Callable, Any, Tuple, TypeVar
+
+from ..compat import has_os_signpost, os_signpost
+from ..model import Model
+
+
+_ModelT = TypeVar("_ModelT", bound=Model)
+
+
+def with_signpost_interval(
+    layer: _ModelT,
+    signposter: "os_signpost.Signposter",
+    name: Optional[str] = None,
+) -> _ModelT:
+    """Wraps any layer and marks the init, forward and backprop phases using
+    signpost intervals for macOS Instruments profiling
+
+    By default, the name of the layer is used as the name of the range,
+    followed by the name of the pass.
+    """
+    if not has_os_signpost:
+        raise ValueError(
+            "with_signpost_interval layer requires the 'os_signpost' package"
+        )
+
+    name = layer.name if name is None else name
+
+    orig_forward = layer._func
+    orig_init = layer.init
+
+    def forward(model: Model, X: Any, is_train: bool) -> Tuple[Any, Callable]:
+        with signposter.use_interval(f"{name} forward"):
+            layer_Y, layer_callback = orig_forward(model, X, is_train=is_train)
+
+        def backprop(dY: Any) -> Any:
+            with signposter.use_interval(f"{name} backprop"):
+                return layer_callback(dY)
+
+        return layer_Y, backprop
+
+    def init(_model: Model, X: Any, Y: Any) -> Model:
+        if orig_init is not None:
+            with signposter.use_interval(f"{name} init"):
+                return orig_init(layer, X, Y)
+        else:
+            return layer
+
+    layer.replace_callbacks(forward, init=init)
+
+    return layer

website/docs/api-layers.md

@@ -835,8 +835,8 @@ https://github.com/explosion/thinc/blob/master/thinc/layers/reduce_last.py
 </inline-list>
 
 Pooling layer that reduces the dimensions of the data by selecting the maximum
-value for each feature. A `ValueError` is raised if any element in `lengths`
-is zero.
+value for each feature. A `ValueError` is raised if any element in `lengths` is
+zero.
 
 | Argument    | Type                             | Description                |
 | ----------- | -------------------------------- | -------------------------- |
@@ -1531,6 +1531,45 @@ model.initialize()
 https://github.com/explosion/thinc/blob/master/thinc/layers/with_nvtx_range.py
 ```
 
+### with_signpost_interval {#with_signpost_interval tag="function" new="8.1.1"}
+
+<inline-list>
+
+- **Input:** <tt>Any</tt>
+- **Output:** <tt>Any</tt>
+
+</inline-list>
+
+Layer that wraps any layer and marks the init, forward and backprop passes as a
+(macOS) signpost interval. This can be helpful when profiling the performance of
+a layer using macOS
+[Instruments.app](https://help.apple.com/instruments/mac/current/). Use of this
+layer requires that the
+[`os-signpost`](https://github.com/explosion/os-signpost) package is installed.
+
+```python
+### Example
+from os_signpost import Signposter
+from thinc.api import Linear, with_signpost_interval
+
+signposter = Signposter("com.example.my_subsystem",
+    Signposter.Category.DynamicTracing)
+
+model = with_signpost_interval(Linear(2, 5), signposter)
+model.initialize()
+```
+
+| Argument     | Type                              | Description                                                                     |
+| ------------ | --------------------------------- | ------------------------------------------------------------------------------- |
+| `layer`      | <tt>Model</tt>                    | The layer to wrap.                                                              |
+| `signposter` | <tt>os_signposter.Signposter</tt> | `Signposter` object to log the interval with.                                   |
+| `name`       | <tt>Optional[str]</tt>            | Optional name for the wrapped layer. Defaults to the name of the wrapped layer. |
+| **RETURNS**  | <tt>Model</tt>                    | The wrapped layer.                                                              |
+
+```python
+https://github.com/explosion/thinc/blob/master/thinc/layers/with_signpost_interval.py
+```
+
 ---
 
 ## Wrappers {#wrappers}