From ae8294da3082d1ed830e9ffb5637b17eecbbf10e Mon Sep 17 00:00:00 2001 From: Arthur Meyre Date: Wed, 8 Dec 2021 15:44:59 +0100 Subject: [PATCH] docs: update compilation with torch compilation flow - update frontend flow - added torch to numpy flow refs #1080 --- .../compilation-pipeline/frontend_flow.svg | 3 ++- .../torch_to_numpy_flow.svg | 4 ++++ docs/dev/explanation/COMPILATION.md | 22 ++++++++++++++----- 3 files changed, 23 insertions(+), 6 deletions(-) create mode 100644 docs/_static/compilation-pipeline/torch_to_numpy_flow.svg diff --git a/docs/_static/compilation-pipeline/frontend_flow.svg b/docs/_static/compilation-pipeline/frontend_flow.svg index 28235ea17..6b03dec9e 100644 --- a/docs/_static/compilation-pipeline/frontend_flow.svg +++ b/docs/_static/compilation-pipeline/frontend_flow.svg @@ -1,3 +1,4 @@ + -
Input Program
v0: numpy
Input Program...
Tracing
Tracing
Data
Data
Algorithm/Function/Transform
Algorithm/Function/Transform
Operator DAG:
"Base Graph"
Operator DAG:...
Topological transform
Topological transform
Operator DAG:
"Candidate Graph"
Operator DAG:...
Constraints check
Constraints check
Input/Intermediate/Output values: 7b unsigned int
Constants: 7+1 = 8b signless int

AND

Opset v0: ADD, MUL, DOT, TLU
Input/Intermediate/Output values: 7b unsigned int...
UI/UX
UI/UX
Error Message + Debug Infos
Error Message + Debug Infos
Bounds Measurement
Bounds Measurement
Dataset + Evaluation
Dataset + Evaluation
Input Bounds + Propagation
Input Bounds + Propagation
OR
OR
Operator DAG + Width:
"Compilable Graph"
Operator DAG + Width:...
MLIR Lowering
MLIR Lowering
MLIR
MLIR
Compiler "Backend"
Compiler "Backend"
Input/Intermediate/Output values: unsigned int
Constants: signless int
Input/Intermediate/Output values: unsigned int...
Error Message + Debug Infos
Error Message + Debug InfosViewer does not support full SVG 1.1 \ No newline at end of file +
Input Program
NumPy function
Input Program...
Tracing
Tracing
Data
Data
Algorithm/Function/Transform
Algorithm/Function/Transform
Operator DAG:
"Base Graph"
Operator DAG:...
Topological transform
Topological transform
Operator DAG:
"Candidate Graph"
Operator DAG:...
Constraints check
Constraints check
Input/Intermediate/Output values: 7b unsigned int
Constants: 7+1 = 8b signless int
Input/Intermediate/Output values: 7b unsigned int...
UI/UX
UI/UX
Error Message + Debug Infos
Error Message + Debug Infos
Bounds Measurement
Bounds Measurement
Dataset + Evaluation
Dataset + Evaluation
Operator DAG + Width:
"Compilable Graph"
Operator DAG + Width:...
MLIR Lowering
MLIR Lowering
MLIR
MLIR
Compiler "Backend"
Compiler "Backend"
Input/Intermediate/Output values: unsigned int
Constants: signless int
Input/Intermediate/Output values: unsigned int...
Error Message + Debug Infos
Error Message + Debug Infos
Executable FHE program
Executable FHE programViewer does not support full SVG 1.1 \ No newline at end of file diff --git a/docs/_static/compilation-pipeline/torch_to_numpy_flow.svg b/docs/_static/compilation-pipeline/torch_to_numpy_flow.svg new file mode 100644 index 000000000..ea6710b1f --- /dev/null +++ b/docs/_static/compilation-pipeline/torch_to_numpy_flow.svg @@ -0,0 +1,4 @@ + + + +
Data
Data
Algorithm/Function/Transform
Algorithm/Function/Transform
Constraints check
Constraints check
Input
torch Module
Input...
Post training-quantization
Post training-quantization
QuantizedModule
QuantizedModule
Supported layers
Well formed model definition
Supported layers...
Representative dataset
Representative dataset
Convert to NumPy equivalent
Convert to NumPy equivalent
NumpyModule
NumpyModule
Quantize with QuantizedModule parameters
Quantize with QuantizedMod...
Quantized dataset
Quantized dataset
NumPy compilation flow
NumPy compilation flow
Executable FHE program
Executable FHE programViewer does not support full SVG 1.1 \ No newline at end of file diff --git a/docs/dev/explanation/COMPILATION.md b/docs/dev/explanation/COMPILATION.md index be0eee789..266f82f5c 100644 --- a/docs/dev/explanation/COMPILATION.md +++ b/docs/dev/explanation/COMPILATION.md @@ -28,7 +28,7 @@ circuit = compiler.get_compiled_fhe_circuit() circuit.run(1, 0) ``` -## Overview +## Overview of the numpy compilation process The compilation journey begins with tracing to get an easy to understand and manipulate representation of the function. We call this representation `Operation Graph` which is basically a Directed Acyclic Graph (DAG) containing nodes representing the computations done in the function. @@ -51,12 +51,24 @@ Once the MLIR is prepared, the rest of the stack, which you can learn more about Here is the visual representation of the pipeline: -```{warning} -FIXME(arthur): check the graph, update with what is missing, notably: torch to numpy, quantization. Maybe an independant graph for quantization may be clearer. -``` - ![Frontend Flow](../../_static/compilation-pipeline/frontend_flow.svg) +## Overview of the torch compilation process + +Compiling a torch Module is pretty straightforward. + +The torch Module is first converted to a NumPy equivalent we call `NumpyModule` if all the layers in the torch Module are supported. + +Then the module is quantized post training to be compatible with our compiler which only works on integers. The post training quantization uses the provided dataset for calibration. + +The dataset is then quantized to be usable for compilation with the QuantizedModule. + +The QuantizedModule is compiled yielding an executable FHECircuit. + +Here is the visual representation of the different steps: + +![Torch compilation flow](../../_static/compilation-pipeline/torch_to_numpy_flow.svg) + ## Tracing Given a Python function `f` such as this one,