[Optimizer] support external partitions #915

youben11 · 2024-06-25T10:56:56Z

Add a change_partition operator to manually go from one partition to the other. We currently allow to specify both partitions, but our main usecase would only require setting one of them.
Take into account external partitions during partitioning and optim

github-actions

Benchmark

Benchmark suite	Current: `b81e13d`	Previous: `43e66cf`	Ratio
`v0 PBS table generation`	`64900958` ns/iter (`± 1665664`)	`64765442` ns/iter (`± 2184591`)	`1.00`
`v0 PBS simulate dag table generation`	`39494744` ns/iter (`± 434852`)	`39744131` ns/iter (`± 385074`)	`0.99`
`v0 WoP-PBS table generation`	`68580766` ns/iter (`± 2367553`)	`68542012` ns/iter (`± 2031379`)	`1.00`

This comment was automatically generated by workflow using github-action-benchmark.

...concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/partitionning.rs

compilers/concrete-optimizer/concrete-optimizer/src/dag/operator/operator.rs

...ilers/concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/analyze.rs

...concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/partition_cut.rs

...concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/partitionning.rs

aPere3

Here are a few comments / question, since I noticed you made some changes 😄 . Keep up the good work !

aPere3 · 2024-07-08T09:18:43Z

compilers/concrete-optimizer/concrete-optimizer/src/dag/operator/operator.rs

+        src_partition: Option<ExternalPartition>,
+        dst_partition: Option<ExternalPartition>,


Is it really possible to have None, None ? Or Some(a), Some(a) ?

If so, maybe it would be worth having an assert when adding the op !

I agree that (None, None), doesn't make sense, but having them both is something we decided not to restrict for now, but we will see. I will add an assert for the first case.

Hmm actually it is not clear for me how this is supposed to work. Only one of them is expected to be something (though you support having both set) ? And the other (none) is supposed to signal the fact that it is free to have any partition ?

My original point was checking that both were something different. Maybe we could have a quick chat, so that I am sure I got everything right ?

I think the "change partition" is just an informative flag of the origin/destination partition for circuit inputs and results. We choosed to keep it more generic and compatible with eventual future use in the middle of a circuit.
Current used form are
ChangePartition(Some, None) on inputs
ChangePartition(None, Some) on results
E.g. future use
ChangePartition(None, None) could hint that if a conflict exists the partition change should happen here and not somewhere else.
ChangePartition(None, Some), could be use to express manually where a transition occurs, which is forced automatically today by the optimizer policy but better manual decision could be taken.
Also it's translated from MLIR, so it could be used there to have multi partition annoted code instead of relying on optimizer choosing partition and parameters.

compilers/concrete-optimizer/concrete-optimizer/src/dag/unparametrized.rs

aPere3 · 2024-07-08T09:21:07Z

...concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/partition_cut.rs

+#[derive(Hash, Eq, PartialEq, Clone, Debug)]
+pub struct ExternalPartition {
+    pub name: String,
+    // TODO add params (maybe just macros)
+}


Do we need a string here ? If used for debugging, it indeed make sense, but maybe we could also have something equivalent for classic partitions ? Is that what you mean in your TODO below ?

I don't think we have classic partitions here. Partitions will be coming from the compiler with a name as an identifier. And yes, it's helpful for debugging and visualizing the graph. The TODO is for parameters to be specified for the partition. Those should be coming from the compiler.

Yes, I understand that this is not supposed to represent a classic partition. I was referring to the fact that for classic partition we refer to them with an index that refers to something like LOW_PPRECISION_PARTITION and maybe it would be worth having named partitions for the classic one as well. But yes, it was more thinking out loud than anything 🤣

And I was referring to todo line 32, sorry for not being clear.

aPere3 · 2024-07-08T09:22:29Z

...concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/partition_cut.rs

+    // Whether it has internal partitions or not
+    pub has_internal_partitions: bool,
+


Reading below, I have a hard time understanding why this is needed. Why isn't it enough to check the length of p_cut ?

the test_tfhers_in_out_dot_compute test showcase the issue. p_cut could actually be empty if there is no lut, and this was translated to having an internal partition. However, we wanted to ignore this default partition if change partition is used.

Ok, but is this something we should have a public field for then ? I mean, my understanding is that the PartitionCut object is kind of an ensemble of parameters used to perform the cut. But this is more of a derived value ?

Maybe I miss something; my point is essentially:

avoid having the information twice if it is already available inside in p_cut and external_parttions

avoid making PartitionCut too intricate.

p_cut empty used to mean 1 internal partition only. Each member of the vec introduce the limit of the partitiion, the last partition always accept all the remaining and is not part of the vec.
since in theory there could be circuit with just 1 external partition, we would be unable to say the p_cut object has only 1 partition.
if it's too confusive we can always add a last member in the p_cut fields explicitely with very high limits.

...concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/partition_cut.rs

compilers/concrete-optimizer/concrete-optimizer/src/dag/operator/operator.rs

compilers/concrete-optimizer/concrete-optimizer/src/dag/unparametrized.rs

...concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/partitionning.rs

compilers/concrete-optimizer/concrete-optimizer/src/dag/unparametrized.rs

...ilers/concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/analyze.rs

...rete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/symbolic_variance.rs

...concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/partition_cut.rs

BourgerieQuentin

looks good to me

BourgerieQuentin · 2024-08-08T14:02:43Z

...concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/partitionning.rs

+        PartitionIndex(default)
+    }
+
+    pub const GLWE_PARAMS: GlweParameters = GlweParameters {


maybe just prefix with TEST_ just to be sure is well used.

Could be great also to name it like in tfhe-rs and give a pointer to that code. And maybe add more than one thfers partition to test?

It's under the tests module already, so not sure it needs a prefix.
Partitioning isn't affected by those parameters, so we shouldn't need to try different ones. Optimization might need that in the future however.

...concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/partitionning.rs

rudy-6-4 · 2024-08-09T11:38:34Z

...oncrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/optimize/tests.rs

@@ -843,10 +843,12 @@ fn test_bug_with_zero_noise() {
 fn test_optimize_tfhers_in_out_dot_compute() {
    let mut dag = unparametrized::Dag::new();
    let input1 = dag.add_input(16, Shape::number());
+    let variance = get_tfhers_max_noise();


max_variance

rudy-6-4 · 2024-08-09T11:38:51Z

...oncrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/optimize/tests.rs

    let tfhers_partition = ExternalPartition {
        name: String::from("tfhers"),
        macro_params: TFHERS_MACRO_PARAMS,
-        max_variance: get_tfhers_max_noise(),
+        max_variance: variance * 2.0,


variance : max_variance / 2.0

rudy-6-4 · 2024-08-09T11:46:11Z

...concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/partitionning.rs

+        internal_dim: 841,
+    };
+
+    pub fn get_tfhers_max_noise() -> f64 {


it's not the max noise. It's the br noise.
Max noise is somefactor * br noise.
So either there is a missing factor here (e.g. 4) or it should be renamed get_tfhers_pbs_noise

rudy-6-4 · 2024-08-09T11:49:36Z

...concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/partitionning.rs

+        assert!(partitions.nb_partitions == 1);
+        let instrs_partition = partitions.instrs_partition;
+        show_partitionning(&dag, &instrs_partition);
+        assert!(instrs_partition[0].instruction_partition == LOW_PRECISION_PARTITION);


Maybe add an assert to show that the exterenal partition is used and no internal partition exists.

rudy-6-4 · 2024-08-09T11:51:00Z

...concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/partitionning.rs

+
+        let partitions = partitionning(&dag);
+        assert!(partitions.nb_partitions == 2);
+        let tfhers_partition_index = PartitionIndex(1);


add an assert veryfing it's considered as external partition

rudy-6-4 · 2024-08-09T11:53:25Z

...concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/partitionning.rs

+
+    #[test]
+    fn test_tfhers_different_in_out_lut_compute() {
+        let mut dag = unparametrized::Dag::new();


put dag construction after partition definitions

rudy-6-4 · 2024-08-09T11:53:52Z

...concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/partitionning.rs

+        let mut dag = unparametrized::Dag::new();
+        let input = dag.add_input(16, Shape::number());
+        let variance = get_tfhers_max_noise();
+        let tfhers_partition = ExternalPartition {


put partition definition at start before dag construction

rudy-6-4 · 2024-08-09T11:54:06Z

...concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/partitionning.rs

+        let mut dag = unparametrized::Dag::new();
+        let input1 = dag.add_input(16, Shape::number());
+        let variance = get_tfhers_max_noise();
+        let tfhers_partition = ExternalPartition {


put partition definition at start before dag construction

rudy-6-4 · 2024-08-09T11:56:55Z

...concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/partitionning.rs

+    }
+
+    #[test]
+    fn test_tfhers_different_in_out_2lut_compute() {


I think you can keep this test and the one just before. not sure what is the added value.

rudy-6-4 · 2024-08-09T12:01:03Z

...concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/partitionning.rs

+                })]
+        );
+    }
+


Maybe adding more cases just to verify no patitionning bug occurs.
Like what happens when add something from external partition with something not.
Eg external_input + lut(other_input) to verify the conversion is added on the lut output.

rudy-6-4 · 2024-08-09T12:03:19Z

...concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/partition_cut.rs

+    pub name: String,
+    pub macro_params: MacroParameters,
+    pub max_variance: f64,
+    pub variance: f64,


should be named pbs_variance

rudy-6-4 · 2024-08-09T12:04:26Z

...concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/partition_cut.rs

+pub struct ExternalPartition {
+    pub name: String,
+    pub macro_params: MacroParameters,
+    pub max_variance: f64,


maybe a comment saying it's the maximimum allowed variance for any ciphertext in that partition in absolute scale

we block them by reducing the search space to a single set of parameters which is provided by the external partition

youben11 requested review from rudy-6-4 and aPere3 June 25, 2024 10:56

cla-bot bot added the cla-signed label Jun 25, 2024

github-actions bot reviewed Jun 25, 2024

View reviewed changes

aPere3 reviewed Jun 25, 2024

View reviewed changes

...concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/partitionning.rs Outdated Show resolved Hide resolved

youben11 force-pushed the optimizer-change-partition branch from 681ad02 to 9b0df72 Compare June 25, 2024 14:48

youben11 requested a review from aPere3 June 28, 2024 09:15

youben11 changed the title ~~feat(optimizer): add change_partition operator~~ [Optimizer] support external partitions Jun 28, 2024

rudy-6-4 reviewed Jun 28, 2024

View reviewed changes

compilers/concrete-optimizer/concrete-optimizer/src/dag/operator/operator.rs Outdated Show resolved Hide resolved

rudy-6-4 reviewed Jun 28, 2024

View reviewed changes

...ilers/concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/analyze.rs Outdated Show resolved Hide resolved

rudy-6-4 reviewed Jun 28, 2024

View reviewed changes

...concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/partition_cut.rs Outdated Show resolved Hide resolved

rudy-6-4 reviewed Jun 28, 2024

View reviewed changes

...concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/partition_cut.rs Outdated Show resolved Hide resolved

rudy-6-4 reviewed Jun 28, 2024

View reviewed changes

...concrete-optimizer/concrete-optimizer/src/optimization/dag/multi_parameters/partitionning.rs Outdated Show resolved Hide resolved

youben11 force-pushed the optimizer-change-partition branch 2 times, most recently from 573183c to 56c7bc1 Compare June 28, 2024 13:56

youben11 force-pushed the optimizer-change-partition branch from 56c7bc1 to f12f7e5 Compare July 8, 2024 09:15

aPere3 reviewed Jul 8, 2024

View reviewed changes

youben11 force-pushed the optimizer-change-partition branch from f12f7e5 to 9f8d1b3 Compare July 8, 2024 10:29

youben11 force-pushed the optimizer-change-partition branch from 1f765f3 to 3fa927e Compare July 24, 2024 11:01

youben11 force-pushed the optimizer-change-partition branch 2 times, most recently from 8cfbc41 to 0e442af Compare August 7, 2024 07:48