embedl-deploy-tensorrt 0.4.1__tar.gz → 0.6.0__tar.gz

{embedl_deploy_tensorrt-0.4.1 → embedl_deploy_tensorrt-0.6.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: embedl-deploy-tensorrt
-Version: 0.4.1
+Version: 0.6.0
 Summary: TensorRT backend for embedl-deploy.
 Author-email: Embedl AB <support@embedl.com>
 Project-URL: Homepage, https://www.embedl.com/
@@ -54,9 +54,10 @@ hardware target ensuring correct quantization and compilation.
 
 ## Supported Backends
 
-| Backend                 | Status      |
-|-------------------------|-------------|
-| NVIDIA TensorRT (v10.3) | Supported   |
+| Backend                   | Status          |
+|---------------------------|-----------------|
+| NVIDIA TensorRT  (v10.3)  | Supported       |
+| Lattice SensAI (v8.0)     | In Development  |
 
 Contact Embedl for other backends.
 
@@ -71,7 +72,7 @@ intermediate.
 
 ---
 
-## Quick Start
+## Quick Start for TensorRT Backend
 
 ```python
 import torch
@@ -85,7 +86,7 @@ model = Model().eval()
 example_input = torch.randn(1, 3, 224, 224)
 
 # 2. Transform — fuse and optimize for TensorRT in one call
-# For more compatibilty you can trace your model with torch.export.export
+# For more compatibility you can trace your model with torch.export.export
 # as follows:
 # model = torch.export.export(model, (example_input)).module()
 res = transform(model, patterns=TENSORRT_PATTERNS)

{embedl_deploy_tensorrt-0.4.1 → embedl_deploy_tensorrt-0.6.0}/README.md

@@ -35,9 +35,10 @@ hardware target ensuring correct quantization and compilation.
 
 ## Supported Backends
 
-| Backend                 | Status      |
-|-------------------------|-------------|
-| NVIDIA TensorRT (v10.3) | Supported   |
+| Backend                   | Status          |
+|---------------------------|-----------------|
+| NVIDIA TensorRT  (v10.3)  | Supported       |
+| Lattice SensAI (v8.0)     | In Development  |
 
 Contact Embedl for other backends.
 
@@ -52,7 +53,7 @@ intermediate.
 
 ---
 
-## Quick Start
+## Quick Start for TensorRT Backend
 
 ```python
 import torch
@@ -66,7 +67,7 @@ model = Model().eval()
 example_input = torch.randn(1, 3, 224, 224)
 
 # 2. Transform — fuse and optimize for TensorRT in one call
-# For more compatibilty you can trace your model with torch.export.export
+# For more compatibility you can trace your model with torch.export.export
 # as follows:
 # model = torch.export.export(model, (example_input)).module()
 res = transform(model, patterns=TENSORRT_PATTERNS)

{embedl_deploy_tensorrt-0.4.1 → embedl_deploy_tensorrt-0.6.0}/src/embedl_deploy/__init__.py

@@ -3,6 +3,7 @@
 """Python package to make AI models deployment-ready for any hardware."""
 
 from embedl_deploy._internal.core.plan import (
+    Trace,
     TransformationPlan,
     TransformationResult,
     apply_transformation_plan,
@@ -14,6 +15,7 @@ from embedl_deploy.version.public import PUBLIC_VERSION
 __version__ = PUBLIC_VERSION
 
 __all__ = [
+    "Trace",
     "TransformationPlan",
     "TransformationResult",
     "__version__",

{embedl_deploy_tensorrt-0.4.1 → embedl_deploy_tensorrt-0.6.0}/src/embedl_deploy/_internal/tensorrt/backend.py

@@ -11,6 +11,7 @@ from embedl_deploy._internal.tensorrt.plan import (
 )
 
 BACKEND = Backend(
+    name="tensorrt",
     conversion_patterns=TENSORRT_CONVERSION_PATTERNS,
     fusion_patterns=TENSORRT_FUSION_PATTERNS,
     smooth_patterns=TENSORRT_SMOOTH_PATTERNS,

embedl_deploy_tensorrt-0.6.0/src/embedl_deploy/_internal/tensorrt/modules/AGENTS.md

@@ -0,0 +1,408 @@
+# TensorRT Fused Modules — Subsystem Guide
+
+## Role & Boundary
+
+**This subsystem owns:** concrete `FusedModule` (and `ConvertedModule`) implementations
+for TensorRT. Every class in `modules/` corresponds to a hardware-fusible or
+structurally-decomposed operation the TensorRT backend can emit as a single engine
+layer or kernel.
+
+**This subsystem does NOT own:**
+
+- When or how to match these modules in a graph — that lives in `patterns/`.
+- Q/DQ stub insertion or calibration logic — that lives in `core/quantize/`.
+- The pattern-matching engine itself — that lives in `core/`.
+
+---
+
+## Position in the System
+
+```
+patterns/fusions.py                # creates FusedModule instances
+patterns/conversions/general.py    # creates ConvertedModule instances (erasure, flatten-linear)
+patterns/conversions/attention.py  # creates ConvertedModule instances (MHA/Swin/SDPA)
+        │
+        ▼
+modules/ (this subsystem)
+        │
+        ▼
+core/quantize/prepare.py       # walks the graph, uses isinstance(mod, FusedModule)
+                               # and mod.inputs_to_quantize to insert Q/DQ stubs
+```
+
+**Instantiated by:** pattern grafts and `replace()` methods in `patterns/fusions.py`
+and `patterns/conversions/`. Most fusion patterns declare a `graft` attribute
+pointing to the fused module class; the graft system calls `_collect_modules()`
+to gather matched modules in tree order and passes them as positional arguments
+to the constructor. For example, `ConvBNActPattern` grafts `FusedConvBNAct`;
+`DecomposeMultiheadAttentionPattern.replace()` constructs `MHAInProjection` and
+`ScaledDotProductAttention`.
+
+**Used by:** The Q/DQ insertion pass in `core/quantize/prepare.py` via
+`isinstance(mod, FusedModule)` and `mod.inputs_to_quantize`, to determine which
+inputs of each `call_module` node should receive a `QuantStub`.
+
+**Also used as intermediate graph nodes by:** conversion patterns, which produce
+`ConvertedModule` subclasses (`MHAInProjection`, `ScaledDotProductAttention`,
+`SwinWindowPartition`, `SwinAttention`, `SwinWindowReverse`). These stay opaque
+during re-tracing so that downstream fusion patterns can see them as atomic nodes.
+
+---
+
+## FusedModule Contract
+
+Every `FusedModule` subclass must satisfy three requirements:
+
+1. **Set `inputs_to_quantize` as a class attribute** (not an instance attribute).
+   The Q/DQ insertion pass reads this before the module is instantiated to decide
+   which positional arguments of the `call_module` FX node will have a `QuantStub`
+   inserted in front of them.
+
+2. **Call `super().__init__()`**, which creates `self.input_quant_stubs` — a dict
+   mapping each index in `inputs_to_quantize` to a fresh `QuantStub`. The Q/DQ
+   pass later enables and configures these stubs during `prepare_qdq()`. It also
+   initialises `self.surrounded = False`, which is later set to `True` by
+   `SurroundWithQuantStubsPattern` to mark modules that have been surrounded
+   with input `QuantStub` entries.
+
+3. **Implement `forward()`** with the fused computation the module represents.
+
+### Graft compatibility
+
+When a pattern uses `graft = FusedFoo` (bare class), the graft system calls
+`_collect_modules()` to walk the matched tree and collect the `nn.Module`
+instances corresponding to trunk and fork nodes (nested branches first, then
+trunk nodes). These are passed as positional arguments to the constructor.
+Therefore the constructor signature must accept modules in the same order
+they appear in the pattern tree.
+
+### What `inputs_to_quantize` means
+
+`inputs_to_quantize` is a set of *positional argument indices* corresponding to
+the `call_module` FX node's arguments — that is, the arguments visible in the
+traced graph, not necessarily the Python keyword positions in `forward()`.
+
+For example, `FusedConvBN.inputs_to_quantize = {0}` means the first tensor
+argument to the fused node (the image tensor) gets a `QuantStub`. The convolution
+weight is quantized separately via `WeightFakeQuantize`, which is attached to the
+module directly rather than via `inputs_to_quantize`.
+
+`FusedConvBNAddAct.inputs_to_quantize = {0, 1}` means both the main feature map
+and the residual skip tensor each get their own `QuantStub`.
+
+### Why `inputs_to_quantize` is a class attribute
+
+Because it describes the module *type's* quantization contract, not any particular
+instance's configuration. The Q/DQ insertion pass queries it once per class during
+graph preparation, before modules are constructed. Making it a class attribute
+(rather than an instance attribute set in `__init__`) ensures the value is
+available from the class itself without constructing an instance.
+
+---
+
+## ConvertedModule vs FusedModule
+
+Both base classes are defined in `core/modules.py`. They serve different roles:
+
+**`ConvertedModule` subclasses** are *intermediate decomposition products*. They
+replace high-level opaque ops (e.g. `nn.MultiheadAttention`,
+`shifted_window_attention`) with sub-modules that downstream *fusion* patterns can
+then match individually. The custom tracer (`_LeafTracer`) treats them as leaf
+nodes — they are never unwrapped during re-tracing. Examples: `MHAInProjection`,
+`ScaledDotProductAttention`, `SwinWindowPartition`, `SwinAttention`,
+`SwinWindowReverse`.
+
+**`FusedModule` subclasses** are the *final quantization target*. The Q/DQ
+insertion pass identifies them via `isinstance(mod, FusedModule)` and wraps their
+inputs/weights with fake-quantization stubs. Examples: `FusedConvBN`,
+`FusedConvBNAct`, `FusedLinear`, `FusedSwinAttention`, `FusedMHAInProjection`.
+
+A `FusedModule` may wrap a `ConvertedModule` (e.g. `FusedMHAInProjection` wraps
+`MHAInProjection`, `FusedSwinAttention` wraps `SwinAttention`). In that case the
+conversion pass runs first, then the fusion pass replaces the `ConvertedModule`
+with its `FusedModule` counterpart.
+
+---
+
+## Module Catalog
+
+### Conv family — `conv.py`
+
+Pattern that creates them: `patterns/fusions.py` (`ConvBNActPattern`,
+`ConvBNPattern`, `StemConvBNActMaxPoolPattern`, `ConvBNAddActPattern`).
+
+All conv fused modules store a reference to the original `nn.Conv2d` (and
+optionally `nn.BatchNorm2d`, `ActivationLike`, `nn.MaxPool2d`) — they do not
+pre-compute fused weights. This lets the same module object work correctly in both
+training (with live BN statistics) and eval (where BN is folded by TensorRT at
+export time).
+
+**BN folding indicator:** The presence of `self.bn` (not `None`) implicitly
+indicates that BN can be folded into the convolution at export time. The
+`__repr__` methods print `"(foldable)"` when `self.bn is not None`.
+
+**`FusedConvBN`** — `Conv2d → [BatchNorm2d]`, `inputs_to_quantize = {0}`.
+
+**`FusedConvBNAct`** — `Conv2d → [BatchNorm2d] → Activation`, `inputs_to_quantize = {0}`.
+
+**`FusedConvBNActMaxPool`** — `Conv2d → [BatchNorm2d] → Activation → MaxPool2d`,
+`inputs_to_quantize = {0}`. MaxPool is included here (rather than fused
+separately) because TensorRT supports fusing the full stem sequence
+(`Conv → BN → Act → Pool`) as a single engine node in the common 7x7 stem pattern.
+Fusing them together avoids an extra quantized activation between Conv and Pool.
+
+**`FusedConvBNAddAct`** — `Conv2d → BatchNorm2d → add(·, residual) → Activation`,
+`inputs_to_quantize = {0, 1}`. The residual tensor is the second input (index 1),
+hence both inputs are quantized. This is the ResNet skip-connection block tail.
+
+**INT8 compatibility guard:** grouped convolutions where `in_channels / groups` or
+`out_channels / groups` is not a multiple of 4 cannot be quantized to INT8 in
+TensorRT. For those cases `_is_int8_compatible_conv()` returns `False`, and the
+module sets `self.input_quant_stubs = {}` (overriding the `super().__init__()`
+default), effectively opting out of quantization.
+
+### Linear family — `linear.py`
+
+Pattern that creates them: `patterns/fusions.py`.
+
+**`FusedLinear`** — wraps a single `nn.Linear`, `inputs_to_quantize = {0}`.
+
+**`FusedLinearAct`** — wraps `nn.Linear → Activation`, `inputs_to_quantize = {0}`.
+
+**`FusedLayerNorm`** — wraps `nn.LayerNorm`, `inputs_to_quantize = set()`.
+LayerNorm is placed in the linear family because it appears in transformer
+architectures immediately before or after linear layers, and the SmoothQuant
+pass needs to reason about LayerNorm and the following linear together.
+Input quantization is disabled (`inputs_to_quantize = set()`, `prefers_fp_input =
+True`) because LayerNorm normalises its input internally — quantizing the input
+before LayerNorm would destroy the statistical properties that normalization relies
+on.
+
+**`FusedLayerNorm.smooth_quant_observer`** — holds a `SmoothQuantObserver`
+instance. It is created in `__init__` but is populated (i.e., scale factors are
+computed and migrated) by the SmoothQuant calibration pass in
+`core/quantize/calibrate.py`. The module itself does not use the observer
+during `forward()` — it only provides a hook for the calibration pass to attach to.
+
+### Attention family — `attention.py`
+
+Pattern that creates the `ConvertedModule` subclasses:
+`patterns/conversions/attention.py` (`DecomposeMultiheadAttentionPattern`).
+Pattern that creates the `FusedModule` wrappers: `patterns/fusions.py`.
+
+**Decomposition model:** `nn.MultiheadAttention` is opaque to FX and to TensorRT.
+The conversion pass decomposes it into three explicit sub-modules:
+
+1. `MHAInProjection` (in-projection) — packed `Linear(E, 3E)` followed by
+   chunking and head splitting.
+2. `ScaledDotProductAttention` (SDPA) — `softmax(Q·Kᵀ / √H) · V`.
+3. `nn.Linear` (out-projection) — the original `mha.out_proj`.
+
+**Why in-proj returns a tuple:** The in-projection produces three independent
+tensors (Q, K, V). TensorRT can fuse the single packed `Linear(E, 3E)` and the
+split into an efficient multi-head projection kernel. Returning a tuple of
+`(Q, K, V)` lets the graph represent the split explicitly — the conversion pass
+inserts `operator.getitem` nodes to fan out the tuple into three separate feeds
+for SDPA.
+
+**Head splitting:** Inside `MHAInProjection.forward()`, after the packed linear
+op, each of Q, K, V is reshaped from `[B, S, E]` to `[B, num_heads, S, head_dim]`
+via `view` + `transpose`. This puts the head dimension before the sequence
+dimension, which is the layout expected by `F.scaled_dot_product_attention` and
+by matrix-multiply kernels.
+
+**`FusedMHAInProjection`** — wraps `MHAInProjection`, `inputs_to_quantize = {0}`.
+Adds a `WeightFakeQuantize` for the packed linear weight. Only `query` (index 0)
+is quantized; `_key` and `_value` are accepted to match the self-attention
+call-site but ignored.
+
+**`FusedScaledDotProductAttention`** — wraps `ScaledDotProductAttention`,
+`inputs_to_quantize = set()`. Adds an internal `softmax_quant` stub with a fixed
+calibration of `(1/127, 0)` — i.e., 8-bit symmetric with a fixed scale matched to
+the softmax output range `[0, 1]`. When the stub is disabled the module delegates
+to the plain SDPA; when enabled it performs manual attention with the quantization
+step between softmax and the second batched matrix multiply (BMM2).
+
+### Swin Attention family — `swin_attention.py`
+
+Pattern that creates the `ConvertedModule` subclasses:
+`patterns/conversions/attention.py` (`DecomposeSwinAttentionPattern`). Pattern
+that creates `FusedSwinAttention`: `patterns/fusions.py`.
+
+**Decomposition model:** `torchvision`'s `shifted_window_attention` free function
+is an opaque `fx.wrap`-ped call that includes spatial padding, cyclic shifting,
+window partitioning, QKV projection, attention, output projection, and window
+reversal — all in one node. The conversion pass splits it into five sub-modules
+that downstream fusion patterns can match individually:
+
+1. `SwinWindowPartition` — pad, cyclic-shift, partition to windows.
+2. `MHAInProjection` — QKV projection (shared with the MHA attention family).
+3. `SwinAttention` — windowed attention with relative position bias and shifted-window mask.
+4. `nn.Linear` — output projection.
+5. `SwinWindowReverse` — reverse partition, shift, and unpad.
+
+**`SwinSpatialState`** — a mutable dataclass shared by all three spatial modules
+(`SwinWindowPartition`, `SwinAttention`, `SwinWindowReverse`). During each
+`forward()` call:
+
+- `SwinWindowPartition.forward()` **writes** `batch_size`, `height`, `width`,
+  `pad_height`, `pad_width`, and `effective_shift_size` into the state.
+- `SwinAttention.forward()` **reads** `pad_height`, `pad_width`,
+  `effective_shift_size`, and `batch_size` to compute the attention mask.
+- `SwinWindowReverse.forward()` **reads** all fields to undo the partition and
+  remove padding.
+
+The state must be shared (not copied) because `SwinWindowPartition` computes the
+actual padded dimensions and effective shift at runtime — these depend on the input
+spatial size, which is not known until `forward()` is called.
+
+**`deepcopy` note:** `copy.deepcopy` preserves the shared reference among the
+three modules. This is intentional: the three modules reference the *same*
+`SwinSpatialState` object; `deepcopy` copies the object once and updates all
+referencing attributes to point to the copy. However, the sharing is not
+thread-safe (see Gotchas).
+
+**`FusedSwinAttention`** — wraps `SwinAttention`, `inputs_to_quantize = set()`.
+Mirrors `FusedScaledDotProductAttention`: adds an internal `softmax_quant` stub
+with fixed calibration. When enabled it manually expands the attention computation
+to insert the quantization step between softmax and BMM2.
+
+### Pointwise family — `pointwise.py`
+
+Pattern that creates it: `patterns/fusions.py` (`ActAddPattern`).
+
+**`FusedActAdd`** — `Activation → add(·, residual)`, `inputs_to_quantize = {0, 1}`.
+Constructor takes `(act: ActivationLike)`. `forward(x, residual)` applies
+`act(x) + residual`. This prevents TensorRT from merging the upstream convolution
+into an activation-fused kernel when the activation output is consumed by a
+subsequent add. Both the activated feature map and the skip-connection tensor are
+quantized at their respective scales.
+
+### Pool family — `pool.py`
+
+Pattern that creates it: `patterns/fusions.py`.
+
+**`FusedAdaptiveAvgPool2d`** — wraps `nn.AdaptiveAvgPool2d`,
+`inputs_to_quantize = set()`. No quantization is applied (pooling is a
+linear operation that does not benefit from separate quantization). Exists as a
+`FusedModule` so the Q/DQ pass treats it uniformly without special-casing it.
+
+---
+
+## BN Folding
+
+Batch normalisation folding is the process of absorbing the BN scale and bias into
+the preceding convolution weight and bias at inference time:
+
+```
+w_folded = w * (γ / σ)
+b_folded = (b - μ) * (γ / σ) + β
+```
+
+where `γ`, `β` are BN's learned weight/bias and `μ`, `σ` are running statistics.
+
+The fused modules store the original `nn.Conv2d` and `nn.BatchNorm2d` as separate
+sub-modules. During training, the BN is applied as a normal operation (live
+statistics). At export, TensorRT performs the folding in the engine. Whether a
+fused module carries a BN is determined implicitly: `self.bn is not None` means
+the module was created from a pattern that included a `BatchNorm2d`, and folding
+is safe. When `self.bn is None`, the convolution weight is used as-is.
+
+---
+
+## Design Decisions
+
+**Why modules store original sub-modules rather than pre-computing fused weights:**
+Because the modules must be correct in both training and eval modes. Training uses
+live BN statistics; the fused weights can only be computed accurately in eval mode
+with frozen running mean/variance. Storing the originals defers the decision to
+the export step while keeping `forward()` correct throughout.
+
+**Why `inputs_to_quantize` is a class attribute:** The Q/DQ insertion pass queries
+the set before constructing any instances — it reads it from the class via
+`type(mod).inputs_to_quantize`. Making it a class attribute also prevents
+accidental per-instance divergence.
+
+**Why intermediate decomposition modules are `ConvertedModule`:** The custom tracer
+(`_LeafTracer` in `core/modules.py`) checks `isinstance(m, (ConvertedModule,
+FusedModule))` to decide whether to treat a module as a leaf. `ConvertedModule` is
+the marker that tells the tracer "do not recurse into this module's `forward()`".
+Without this, re-tracing the graph after conversion would unwrap the decomposed
+sub-modules, defeating the purpose of decomposition.
+
+---
+
+## Gotchas & Pitfalls
+
+**`SwinSpatialState` thread-safety:** The state is a mutable shared object written
+during every `SwinWindowPartition.forward()` call and read by the other two
+modules in the same forward pass. If the model is run from multiple threads
+concurrently (e.g. in a data-parallel setup without model replication), the writes
+from one thread will corrupt the reads of another. Use `torch.nn.DataParallel` (not
+`DistributedDataParallel`) only if each replica gets its own model copy via
+`deepcopy`, which does preserve sharing within a single copy.
+
+**`inputs_to_quantize` indices vs. Python `forward()` args:** The indices refer to
+the positional arguments of the FX `call_module` node in the traced graph, not to
+the `forward()` Python signature. In the normal case they are identical. They
+diverge if the FX graph is produced from a non-trivial call-site (e.g. keyword
+arguments re-ordered). Always verify against the actual graph node when debugging
+Q/DQ placement.
+
+**Weight sharing:** The conv/linear sub-modules stored inside fused modules hold
+references to the tensors from the *original* model. If the original model's
+weights are modified after fusion, the fused module sees the change. This is
+usually desirable (e.g. for QAT gradient updates), but can be surprising if the
+original model is used independently.
+
+**Grouped conv INT8 opt-out:** When `_is_int8_compatible_conv()` returns `False`
+the `__init__` of the conv fused modules sets `self.input_quant_stubs = {}`,
+overriding the dict populated by `FusedModule.__init__()`. This means the module
+is effectively excluded from quantization despite being a `FusedModule`. The
+`weight_fake_quant` attribute is also not created in this path.
+
+---
+
+## Adding a New Fused Module
+
+1. **Subclass `FusedModule`** (from `core/modules.py`).
+2. **Set `inputs_to_quantize`** as a class attribute — a `set[int]` of positional
+   argument indices that should receive activation `QuantStub`s.
+3. **Optionally add `self.weight_fake_quant = WeightFakeQuantize({self})`** in
+   `__init__` if the module has a learnable weight that should be fake-quantized.
+4. **Implement `forward()`** with the fused computation.
+5. **Write a `Pattern` subclass** in `patterns/fusions.py`. Prefer declaring
+   `graft = FusedFoo` (bare class) so the graft system handles replacement
+   automatically. The constructor must accept modules in tree order (nested
+   branches first, then trunk nodes). If the replacement logic cannot be
+   expressed as a bare-class graft, provide a `ReplacementMaker` or a custom
+   `replace()` method instead.
+6. **Add the pattern to `TENSORRT_PATTERNS`** (or the appropriate pattern list) in
+   `tensorrt/plan.py`.
+7. **Write tests** in `tests/tensorrt/patterns/fusions/` covering: pattern match,
+   pattern replace, correct quantisation stub placement.
+
+---
+
+## Testing
+
+Test models for these modules live in two locations:
+
+- `tests/models/conv.py` — `ConvBnRelu`, `ConvBnSiLU`, `ConvBn`, `ConvOnly`,
+  `ConvBnAddRelu`, `StemConvBnReluMaxPool`, etc.
+- `tests/models/attention.py` — `SimpleSelfAttention`, `SimpleSwinAttention`, etc.
+- `tests/models/linear.py`, `tests/models/pool.py` — linear and pool variants.
+- `tests/tensorrt/models/attention.py` — `LayerNormMHAInProjection`,
+  `SDPALinearOutProjection`, `SwinAttentionLinearOutProjection`.
+
+Pattern tests live in `tests/tensorrt/patterns/fusions/`.
+
+**What to verify for a new fused module:**
+
+- The pattern matches exactly the expected nodes (check `tree_match.serialize()`).
+- The fused node's `target` name is the expected auto-generated name.
+- `resolve_module(fused_node, FusedXxx)` succeeds.
+- `bool(fused_module.input_quant_stubs)` matches `inputs_to_quantize` expectations.
+- `hasattr(fused_module, 'weight_fake_quant')` matches expectations.
+- Graph equivalence: fused model output matches original model output on the same
+  input (see `tests/tensorrt/test_equivalence.py`).