Use ISL equality check

[mdgrs]

When running the tc-resnet model, we could never finish running heir-opt because the following check was blowing up

heir/lib/Transforms/ConvertToCiphertextSemantics/ConvertToCiphertextSemantics.cpp

Lines 2377 to 2378 in de797a2

	if (!expandedRelation.isEqual(resultLayout.getIntegerRelation())) {
	return rewriter.notifyMatchFailure(

Here, we replace this check with an enumeration of the points to make it faster.

Alternatively, we could simply omit this check and assume that the expandedRelation is indeed equal to the input relation.

@j2kun Do I remember correctly that you said you could compile this model in 30 seconds? It could be that we are hitting this check because of some other changes we introduced and you avoided it? These checks are certainly not free (we compile in 320 seconds without the checks, and 420 seconds with the checks)

Our TC-Resnet has dimension 40x101, and I suspect that the 101 dimension is particularly nasty for ISL, introducing a lot of floor/divs

I am looking for feedback if you think this approach is sensible.

[j2kun]

Do I remember correctly that you said you could compile this model in 30 seconds?

That was the LRM model. For tc_resnet I am compiling with these flags

bazel run -c opt //tools:heir-opt -- "--annotate-module=backend=lattigo scheme=ckks" \
    "--torch-linalg-to-ckks=ciphertext-degree=4096 modulus-switch-after-mul=true experimental-disable-loop-unroll=true level-budget=40 ckks-bootstrap-waterline=40" \
    "--scheme-to-lattigo" \
    --mlir-timing \
    --mlir-print-ir-before-all \
    --mlir-print-ir-tree-dir=/tmp/mlir \
    $PWD/path/to/hotword.mlir

and the version of tc_resnet I'm using has this signature

func.func @run_inference(%arg0: tensor<1x40x98xf32> {secret.secret}) -> tensor<1x12xf32>

It compiles in 75 seconds with 19 seconds in convert-to-ciphertext-semantics and 33 seconds in configure-crypto-context.

That said, I'm currently working on a bug where if I omit the mlir-print-ir-before-all part it hits an error like error: 'mgmt.adjust_scale' op Dataflow analysis failed to populate scale lattice for result >_<

[j2kun]

Generally speaking, I think enumeration of points will not scale well and I've observed performance issues with it.

Can you do the following:

1. Add the slow relations to the benchmarking target in benchmark/isl (maybe add an equality check benchmark).
2. Add a step to this that estimates the size of the two relations using a fast approximate method. You could try the recently added relationSize (which is exact) or isl's box_hull or isl_map_polyhedral_hull. Use this with a decent threshold to determine if enumeration will be too slow.
3. Use the tryProveUnequal test before enumeration, since that one should be much faster and can short-circuit the rest of the function. If you use relationSize and it's fast enough, you could also use this to short-circuit the full enumeration (if it's false, the two relations are unequal).

[mdgrs]

Generally speaking, I think enumeration of points will not scale well and I've observed performance issues with it.

Can you do the following:

1. Add the slow relations to the benchmarking target in benchmark/isl (maybe add an equality check benchmark).
2. Add a step to this that estimates the size of the two relations using a fast approximate method. You could try the recently added relationSize (which is exact) or isl's box_hull or isl_map_polyhedral_hull. Use this with a decent threshold to determine if enumeration will be too slow.
3. Use the tryProveUnequal test before enumeration, since that one should be much faster and can short-circuit the rest of the function. If you use relationSize and it's fast enough, you could also use this to short-circuit the full enumeration (if it's false, the two relations are unequal).

I changed the name of the PR because I tried the ISL version and it is much faster on my example and does not hang. As such, this solution might be a bit over engineered: I could also change it to just using the ISL version and be done with it, without doing the quick checks first.

[j2kun]

yes, I am fine to just using the ISL version, provided all the benchmarks align. Can you also add a comment line there to remark which specific Presburger API call was the one that was slow? Eventually I would like to improve the performance of the upstream Presburger library, and it would help to have some bread-crumbs (this PR and code comments)

[mdgrs]

yes, I am fine to just using the ISL version, provided all the benchmarks align. Can you also add a comment line there to remark which specific Presburger API call was the one that was slow? Eventually I would like to improve the performance of the upstream Presburger library, and it would help to have some bread-crumbs (this PR and code comments)

I ended up keeping the volume check since it is very fast (and left a comment as you requested). If you could have a look, I think it is good to go (modulo the missing hash. Should I reopen a new PR from the google repo instead of our Belfortlabs fork?)