============
Architecture
============

Brevitas is organized around a few different concepts listed below.

Functions
---------

Functions include both Autograd functions and TorchScript functions and can be found under `brevitas.function`.

Autograd functions implement straight-through estimators.
All Autograd functions are implemented both in Python, found under `brevitas.ops.autograd_ste_ops`,
and C++, found under *brevitas/csrc*.
This is because to date Python Autograd functions cannot be compiled
by Pytorch's JIT compiler, but C++ Autograd functions can. Compiling the C++ extension is performed at runtime
in order to simplify packaging and distribution (only a single .cpp source file is distributed).
This requires the user to have an appropriate C++ compiler, so a Python implemention is provided as fallback.

Because of this, as long as `BREVITAS_JIT=0` (which it is by default), the C++ backend is not loaded.
Wrapping and switching between the two implementations happens in `brevitas.function.ops_ste`.

TorchScript functions implements various bits and pieces used commonly across Brevitas that are
required to be JIT-compilable. They can be found under `brevitas.function.ops` and `brevitas.function.shape`.

Core ScriptModules
------------------

The algorithmic core of Brevitas can be found under `brevitas.core`.
The core implements all the various building blocks required to assemble an affine quantizer.
All building blocks are implemented as a ScriptModule in the *old-style* TorchScript scriping API
(i.e. inheriting from `torch.jit.ScriptModule`). Many of the functions described in the section above
are called from within one of the core ScriptModules.
Any module within `brevitas.core` is required to be a ScriptModule compatible with Pytorch 1.3.0 JIT compiler,
with most building blocks being compatible with Pytorch 1.1.0 JIT compiler.
Implementing everything in TorchScript enables compute and memory optimizations at training time,
which for more complicated quantization pipelines can be quite significant, thus reducing the intrinsic
training-time cost of quantization-aware training.

Adhering to the restriction imposed by TorchScript pose a challange in terms of how to achieve flexibility
while minimizing redundancy. Because the flavour of TorchScript adopted by Brevitas does not allow
for inheritance, Brevitas' core is highly biased towards leveraging composition.
In particular, the implementation favours inversion of control through dependency injection (DI),
whether manually or through the DI library *dependencies* (as explained in the Injectors section below.)

Injectors and Quantizers
------------------------

Auto-wiring dependency injection (DI) with *dependencies* is the machinery at the heart of Brevitas.
If you have ever used fixtures in pytest, you already know (high-level) how auto-wiring DI works.
In the case of *dependencies*, the idea is that the objects to be instantiated and wired together are
declared as attributes of an Injector class. The driving mechanism behind the auto-wiring process is
to match the name of an attribute of an Injector with them name of arguments required to build other
attributes of the same Injectors. \
When applied to Brevitas then, the idea is to throw together a bunch of components
appropriately chosen from `brevitas.core` as attributes of an Injector, and have the DI library automatically
assemble them together. A quantizer then is an Injector with a `tensor_quant` attribute.
A `tensor_quant` object is expected to be a torch *Module* (either a `torch.nn.Module` or a `torch.jit.ScriptModule`)
that takes in as input a torch.Tensor to quantize, and return as output a tuple containing four
`torch.Tensor`, representing respectively the output *quantized tensor* in *dequantized format*,
its *scale factor*, its *zero-point*, and it's *bit-width*.

Injectors are a powerful way to express quantizers: being standard Python classes, they lend themselves
to both inheritance and composition (through multiple inheritance in a mixin style).
That means for example that a new quantizer can be defined simply by inheriting from an existing one
and overriding some of its attributes, or by inheriting from multiple smaller Injector that
declare only some of the components required to assemble a quantizer. Pre-built quantizers can be found under
`brevitas.quant`. Specifically, `brevitas.quant.shifted_scaled_int` holds quantizers with *zero-point != 0.*
and *floating-point* scale factors, `brevitas.quant.scaled_int` holds quantizers with *zero-point == 0.* and
*floating-point* scale factors, and `brevitas.quant.fixed_point` holds fixed-point, i.e. quantizers
with *zero-point == 0* and *power-of-two* scale factors. Pre-built quantizers are assembled together
from smaller Injectors that can be found in `brevitas.quant.base`.

Brevitas depends on a specific older version of *dependencies* (*v2.0.1*) that plays well with ScriptModules,
and it extends it with an `ExtendedInjector` class that provides support for some additional features.
Specifically, `ExtendedInjector` allows to dynamically declare an object to auto-wire by returning a class
from a `@value` function. This style of syntax allows to create intermediate abstractions between the
`brevitas.core` components and the `brevitas.quant` quantizers, as explained in the section below.

Enums, Shapes and Solvers
-------------------------

Trying to navigate all the various components in `brevitas.core` that can be assembled together in a
quantizer can be confusing. The lack a of clear object hierarchy between various ScriptModules means
that it's not obvious which components adhere to a certain interface. Additionally, while some components
might fit together from a purely software perspective, they might not work together
from a machine learning perspective. In order to provide a simplified interface that abstracts away some
of these details, Brevitas provides various `Enum` on top of the components found in `brevitas.core`.

Some examples of enums are: `QuantType` specifies the type of quantization, e.g. `QuantType.BINARY`
for binary quantization. `ScalingImplType`, for specifying the type (algorithmic-wise) of scaling,
e.g. `ScalingImplType.PARAMETER_FROM_STATS` to specify that the scale factor should be a learned
`torch.nn.Parameter` initialized from statistics of the tensor to quantize.
Enums can currently be found under `brevitas.inject.enum`.

Depending on the kind of tensor to quantize, say *weights vs activations*, the same enum value is gonna
translate to different `brevitas.core` components. So for example `QuantType.BINARY` translates to
`brevitas.core.quant.BinaryQuant` for weights and to `brevitas.core.quant.ClampedBinaryQuant` for activations.
This way enums allows to declare quantizers while abstract away from the specifics of a target tensor.
In general there can be a *1-to-1*, *many-to-1* or *many-to-many* relationship between enums and
`brevitas.core` components and their hyperparameters.

The translation between enums and `brevitas.core` is performed by a *solver*, which can be found under
`brevitas.quant.solver`.\
Solvers are really just an `ExtendedInjector` that take advantage of the extended syntax of Injectors implemented in Brevitas
to translate e.g. `quant_type = QuantType.BINARY` to `tensor_quant = BinaryQuant` within the scope of
a quantizer at *dependency-injection time*. That means that solvers are as composable as quantizers are,
so for example to solve enums against weight quantization, it's enough to create a quantizer that inherits from
`brevitas.quant.solver.weight.WeightQuantSolver`, which is itself just the collection of various solvers for individual tasks.
Looking at the quantizers found under brevitas.quant, it can be seen that for the most part they actually specify
enums rather than directly `brevitas.core` components. Then enums are then solved to different core components
depending on which solver is applied. This is meant to provide a blueprint for users to understand which enums
are supported and how they go together.

A similar mechanism applies when for example the directive `scaling_per_output_channel=True` is specified.
In the case of `ScalingImplType.PARAMETER_FROM_STATS`, that means that a `torch.nn.Parameter` with
size the number of channels of the tensor to quantize has to be allocated. Because TorchScript does not allow for much dynamic behaviour,
the shape of the parameter has to be known at dependency-injection time. Rather than forcing the user
to specify the appropriate shape in a quantizer, Brevitas is capable of inferring it from the `nn.Module`
whose weights the quantizer is applied to. This again happens by means of a solver included as part of
`WeightQuantSolver`.

Thanks to how dependencies works, solvers are invoked only whenever their output is actually required to build
an attribute of an `ExtendedInjector`, which in the case of Brevitas is `tensor_quant`. Additionally,
by specifying a solver as last in the list of classes from which a quantizer inherits from,
it's always possible to override its behaviour and directly declare its output. So for example it's
possible to directly declare `tensor_quant = BinaryQuant` instead of `quant_type = QuantType.BINARY` even
when `WeightQuantSolver` is applied to a quantizer. This allows more advanced users to mix-and-match
enums with custom components. Finally, it's always possibly to just not apply a solver to a quantizer
and simply declare everything manually.

QuantTensor
-----------

A `QuantTensor` is a custom data structure for representing a uniform, affine quantized tensor.
It can be found under `brevitas.quant_tensor`.
It can be *valid* or *non-valid*. A non-valid `QuantTensor` is simply a wrapper around a torch.Tensor that
had been previously quantized and is now in dequantized format. The `QuantTensor` is marked as non-valid because it doesn't carry enough information
to derive its quantized representation back. A valid QuantTensor carries *scale*, *zero-point*, *bit-width*, *sign*,
and whether it was generated in training or inference mode.

The arithmetic of QuantTensors implments a generalized version of fixed-point arithmetic, with the main assumption
being that only two QuantTensor with the same scale factor can be summed together. This constrain is enforced
when the QuantTensors involved in a sum have been generated in inference mode, but it's not enforced in training mode.
This is because when dealing with e.g. `ScalingImplType.PARAMETER_FROM_STATS`, the activation tensors in a residual topology
can have different scale factors along the skip connection at training time, but not at inference time.

Proxies
-------

A proxy is a `nn.Module` that wraps a quantizer. Proxies can be found under `brevitas.proxy`.
Proxies are specialized w.r.t. the kind of tensor they quantize, such as *weights*, *biases*, or *activations*.

The main responsability of a proxy is to make sure a QuantTensor is returned as output of quantization, which wouldn't be
possible in TorchScript. Additionally, it has to make sure a quantizer is re-initialized any time it is necessary.

For example, when performing `ScalingImplType.PARAMETER_FROM_STATS` scaling on a weight tensor,
statistics of the weight tensor are computed at dependency-injection time and used to initialize a
learned parameter. However, it's not possible to know a-priori whether a pretrained floating-point state-dict
will be later on loaded on top of a quantized model definition. In that case, any initialization logic that
depends on the state-dict of the model that is being quantized has to be recomputed.
Injectors invoked by proxies on state_dict changes allow to do so automatically,
providing a mechanism to reconcile the inherent rigidity of TorchScript with the typical define-by-run execution of Pytorch models.

Proxies also allow to support more complex quantization scenarios, such as when the same quantizer has to be shared between
different layers. A typical situation where that happens is when the output of multiple branches of a residual topology are summed together
at high-precision, without requantizing first. In that case, the weight tensors of the layers that feed into the accumulator
need to have the same scale factor. That can be accomplished by declaring a single `WeightQuantProxy` that is shared among multiple layers.
What happens is that - for example for `ScalingImplType.STATS`, `ScalingImplType.AFFINE_STATS` or `ScalingImplType.PARAMETER_FROM_STATS` -
the scale factor is computed as a statistics of the concatenation of the weight tensors to be quantized.
Because it can't be known a-priori to between how many layers the same WeightQuantProxy is shared, again every
time a WeightQuantProxy starts tracking a new weight tensor, the underlying quantizer has to be re-initialized.

Quant Layers and Mixins
-----------------------

A `QuantLayer` is the quantized variant of a `torch.nn` layer, and can be found under `brevitas.nn`.
Typically a QuantLayer inherits from both its floating-point variant (e.g. `QuantConv2d` inherits from `Conv2d`),
plus a serie of mixins, each responsibile for instantiating a proxy within the QuantLayer.
A mixin is more specialized than a proxy, so for example both `QuantInputMixin` and `QuantOutputMixin`
instantiate an activation quantization proxy. So-called _QuantWBIOL_ layers (such as `QuantConv2d`)
inherit from `QuantWeightMixin`, `QuantBiasMixin`, `QuantInputMixin` and `QuantOutputMixin`.
That means that they can quantize respectively *weight*, *bias*, *input* and *output*.
Quantizers can be passed to a *QuantWBIOL* layer by setting respectively `quant_weight=`, `quant_bias=`, `quant_input=`
and `quant_output=`. If an `ExtendedInjector` is passed it, a proxy will be allocated by the mixin to deal with its
initialization. Otherwise, if a proxy is passed in, it will be set as-is. Setting e.g. `quant_weight=None`
will disable quantization for weights. A layer where quantization is disabled is supposed to act exactly like
its floating-point counterpart, so `QuantConv2d(..., weight_quant=None, bias_quant=None, input_quant=None, output_quant=None)`
behaves like a `Conv2d` layer.

Typically `torch.nn` layers expose a flat interface. To support a similar *UX* in Brevitas, QuantLayers
support setting attributes in a quantizer by passing keyword arguments with an appropriate prefix.
For *QuantWBIOL* layers, keyword arguments with prefix `weight_` are passed to the `weight_quant` quantizer, `bias_` to `bias_quant`,
`input_` to `input_quant`, and `output_` to `output_quant`.
For `quantized activation` layers, like `QuantReLU`, a prefix is not required, and keyword arguments are directly passed to output quantization.
In case an `ExtendedInjector` is not set, e.g. if `weight_quant=None`, but additional keyword arguments are passed in,
an empty `ExtendedInjector` is automatically allocated and keyword arguments are set as its attribute according to the their prefix.
Keyword arguments have priority over pre-existing attribute of a quantizer, so passing a keyword argument is a way
to override the attribute of a quantizer on an individual-layer level.

In many real-life scenarios, a user might want to first quantize only certain layers or certain parts
of certain layers to perform some exploratory analysis in terms of accuracy.
Correctness w.r.t. the specifics of a target hardware, if any, might not be a concern. To minimize friction with adoption then,
Brevitas is designed to remain functional as much as possible under partially specified information.
With the exception of `TruncAvgPool2d`, a `QuantLayer` is not expected to receive a `QuantTensor` as input
(altough doing so enables more scenarios), nor it returns one by default (i.e. `return_quant_tensor=False` by default).
Specifically, The output of a QuantLayer is always in de-quantized format
(whether wrapped in a valid or non-valid `QuantTensor` or not).
That means that QuantLayers can be easily mixed with standard `torch.nn` layers.


Export
------

TODO

FX graph tracing and transformations
------------------------------------

TODO

Losses
------

TODO
