add lola e2e test

[AlexanderViand]

This is the model that kept blowing up for us :)

The PR CI passes trivially because the test is "large" and not run, but it currently fails in --convert-to-ciphertext-semantics, choking on this:

https://github.com/google/heir/blob/main/lib/Transforms/ConvertToCiphertextSemantics/TypeConversion.cpp#L28-L29

I added some debug output, so here is what it's actually trying to do:

materializeLayout: computing upper bound for range dimension 0 (ct), layout=#tensor_ext.layout<"{ [i0, i1, i2, i3] -> [ct, slot] : exists (e1, e2, e3, e4, e5, e6: i1 = 0 and 1024e5 = -i0 - 28i2 - i3 + ct + 784*floor((i0)/4) + 2e1 - 56e2 - 2e3 + 28e4 and 2048e6 = i0 + slot - 784*floor((i0)/4) - 2e1 - 28e4 and 0 <= i0 <= 7 and 0 <= i2 <= 1 and 0 <= i3 <= 1 and 0 <= ct <= 1023 and 0 <= slot <= 2047 and i0 <= 2e1 <= 27 + i0 and 0 <= e2 <= 13 and 0 <= e3 <= 13 and -1 - i0 + 2e1 <= 2e3 <= -i0 + 2e1 and 0 <= e4 <= 27 and -1 + i0 - 4*floor((i0)/4) + 4e2 <= 2e4 <= i0 - 4*floor((i0)/4) + 4e2) }">

[AlexanderViand]

According to our 🤖 overlords, IntegerRelation::getConstantBound64 is a bad idea here (projects out every variable, and somehow blows up exponentially) and we should use presburger::Simplex's computeIntegerBounds instead.

With that, it succeeds at finding the upper bound:

materializeLayout: computing upper bound for range dimension 0 (ct), layout=#tensor_ext.layout<"{ [i0, i1, i2, i3] -> [ct, slot] : exists (e1, e2, e3, e4, e5, e6: i1 = 0 and 1024e5 = -i0 - 28i2 - i3 + ct + 784floor((i0)/4) + 2e1 - 56e2 - 2e3 + 28e4 and 2048e6 = i0 + slot - 784floor((i0)/4) - 2e1 - 28e4 and 0 <= i0 <= 7 and 0 <= i2 <= 1 and 0 <= i3 <= 1 and 0 <= ct <= 1023 and 0 <= slot <= 2047 and i0 <= 2e1 <= 27 + i0 and 0 <= e2 <= 13 and 0 <= e3 <= 13 and -1 - i0 + 2e1 <= 2e3 <= -i0 + 2e1 and 0 <= e4 <= 27 and

but then it just slows down to a crawl a bit later:

supports expanded conv2d as matvec with halevi-shoup: isPowerOfTwoDims=1 isConv2dAsMatvec=1
Converting linalg.conv_2d_nchw_fchw op with halevi shoup kernel: %18 = "linalg.conv_2d_nchw_fchw"(<>, %16, %15) <{dilations = dense<1> : vector<2xi64>, operandSegmentSizes = array<i32: 2, 1>, strides = dense<2> : vector<2xi64>}> ({
^bb0(%arg8: f32, %arg9: f32, %arg10: f32):
%36 = "arith.mulf"(%arg8, %arg9) <{fastmath = #arith.fastmath}> : (f32, f32) -> f32
%37 = "arith.addf"(%arg10, %36) <{fastmath = #arith.fastmath}> : (f32, f32) -> f32
"linalg.yield"(%37) : (f32) -> ()
}) {heir.kernel_info = {gap_factor = 2 : i64, input_shape = array<i64: 1, 1, 28, 28>, result_shape = array<i64: 1, 2, 28, 28>}, linalg.memoized_indexing_maps = [affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d4, d2 * 2 + d5, d3 * 2 + d6)>, affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d1, d4, d5, d6)>, affine_map<(d0, d1, d2, d3, d4, d5, d6) -> (d0, d1, d2, d3)>], secret.kernel = #secret.kernel<name = "MatvecDiagonal", force = false>, tensor_ext.layout = [#tensor_ext.layout<"{ [i0, i1, i2, i3] -> [ct, slot] : i0 = 0 and ct = 0 and (-196i1 - 14i2 - i3 + slot) mod 2048 = 0 and 0 <= i1 <= 7 and 0 <= i2 <= 13 and 0 <= i3 <= 13 and 0 <= slot <= 2047 }">, #tensor_ext.layout<"{ [i0, i1] -> [ct, slot] : exists (e0, e1, e2, e3, e4, e5: i0 = 0 and ct = 0 and 0 <= i1 <= 2047 and 0 <= slot <= 2047 and e0 <= 0 and 2048e0 >= -1567 + i1 and 2048e0 >= -1567 + slot and 0 <= e1 <= 13 and 2e2 >= slot and -13 + i1 + 196slot - 2048e0 - 14e1 <= 392e2 <= i1 + 196slot - 2048e0 - 14e1 and -3 - slot + 2e2 <= 4e3 <= -slot + 2e2 and -1 - slot + 4e1 + 2e2 - 4e3 <= 2e4 <= -slot + 4e1 + 2e2 - 4e3 and -3 - slot + 2e2 <= 4e5 <= -slot + 2e2 and -1 - 2i1 - 391slot + 2048e0 + 28e1 + 784e2 - 28e4 <= 784e5 <= -2i1 - 391slot + 2048e0 + 28e1 + 784e2 - 28e4) }">]} : (tensor<1x1x28x28xf32>, tensor<8x1x2x2xf32>, tensor<1x8x14x14xf32>) -> tensor<1x8x14x14xf32>

Specifically, it's our old friend getCtComplementPoints again, and --convert-to-ciphertext-semantics takes about 5 minutes, and the full compilation pipeline about 20min :(

[AlexanderViand]

Interestingly, even if you wait the full 20+ min for the compile, it produces an incorrect/invalid IR:

FAIL: //tests/Examples/lattigo/ckks/lola:lola_test (Exit 2) (see /private/var/tmp/_bazel_a.viand/00b31163a90b5ef5a4c32391f713b490/execroot/_main/bazel-out/darwin_arm64-fastbuild-ST-2c5af747e1f8/testlogs/tests/Examples/lattigo/ckks/lola/lola_test/test.log)
INFO: From Testing //tests/Examples/lattigo/ckks/lola:lola_test:
==================== Test output for //tests/Examples/lattigo/ckks/lola:lola_test:
--- FAIL: TestLoLA (20.83s)
    lola_test.go:65: Lola__preprocessing took 522.491416ms
panic: cannot Rescale: input Ciphertext level is too low [recovered, repanicked]

goroutine 37 [running]:
testing.tRunner.func1.2({0x10091f660, 0x71fa6d064020})
        GOROOT/src/testing/testing.go:1974 +0x1a0
testing.tRunner.func1()
        GOROOT/src/testing/testing.go:1977 +0x318
panic({0x10091f660?, 0x71fa6d064020?})
        GOROOT/src/runtime/panic.go:860 +0x12c
tests/Examples/lattigo/ckks/lola/lola_test.Lola__preprocessed(0x71f8f94be008, 0x71f8fce5d8c0, {{0xf, {0x71f8f943a060, 0xb, 0xb}, {0x71f8f9364180, 0x3, 0x3}, {{0x10099e1d8, ...}, ...}, ...}}, ...)
        bazel-out/darwin_arm64-fastbuild/bin/tests/Examples/lattigo/ckks/lola/lola_lib.go:879 +0x9d5c
tests/Examples/lattigo/ckks/lola/lola_test.TestLoLA(0x71f8f935eb48)
        tests/Examples/lattigo/ckks/lola/lola_test.go:68 +0x278
testing.tRunner(0x71f8f935eb48, 0x100998360)
        GOROOT/src/testing/testing.go:2036 +0xc4
created by testing.(*T).Run in goroutine 1
        GOROOT/src/testing/testing.go:2101 +0x3a8
================================================================================

[AlexanderViand]

With #3226, it does now compile successfully (in 20 minutes) and run, but the results are wrong!
EDIT: the reason seems to be that we're asking Lattigo to bootstrap fairly large intermediate values that are way outside the range of that it's designed to accurately bootstrap. Unfortunately, HEIR doesn't really have an answer to this problem yet.

FAIL: //tests/Examples/lattigo/ckks/lola:lola_test (Exit 1) (see /private/var/tmp/_bazel_a.viand/00b31163a90b5ef5a4c32391f713b490/execroot/_main/bazel-out/darwin_arm64-fastbuild-ST-2c5af747e1f8/testlogs/tests/Examples/lattigo/ckks/lola/lola_test/test.log)
INFO: From Testing //tests/Examples/lattigo/ckks/lola:lola_test:
==================== Test output for //tests/Examples/lattigo/ckks/lola:lola_test:
--- FAIL: TestLoLA (31.47s)
    lola_test.go:65: Lola__preprocessing took 543.372333ms
    lola_test.go:70: Lola__preprocessed took 15.194155416s
    lola_test.go:99: logit 0: got -0.549884, want -8.300922 (+/- 0.10), delta 7.751038
    lola_test.go:104: logit 0: got -0.549884, plaintext -8.300922
    lola_test.go:99: logit 1: got 0.537425, want -5.588467 (+/- 0.10), delta 6.125892
    lola_test.go:104: logit 1: got 0.537425, plaintext -5.588467
    lola_test.go:99: logit 2: got -1.570775, want -3.635293 (+/- 0.10), delta 2.064519
    lola_test.go:104: logit 2: got -1.570775, plaintext -3.635293
    lola_test.go:99: logit 3: got -2.705697, want -7.802109 (+/- 0.10), delta 5.096412
    lola_test.go:104: logit 3: got -2.705697, plaintext -7.802109
    lola_test.go:99: logit 4: got -2.340250, want -3.840308 (+/- 0.10), delta 1.500059
    lola_test.go:104: logit 4: got -2.340250, plaintext -3.840308
    lola_test.go:99: logit 5: got -0.855760, want -5.856731 (+/- 0.10), delta 5.000972
    lola_test.go:104: logit 5: got -0.855760, plaintext -5.856731
    lola_test.go:99: logit 6: got 1.037338, want -5.982526 (+/- 0.10), delta 7.019863
    lola_test.go:104: logit 6: got 1.037338, plaintext -5.982526
    lola_test.go:99: logit 7: got 2.594647, want 11.231959 (+/- 0.10), delta 8.637312
    lola_test.go:104: logit 7: got 2.594647, plaintext 11.231959
    lola_test.go:99: logit 8: got -1.040165, want -3.031523 (+/- 0.10), delta 1.991359
    lola_test.go:104: logit 8: got -1.040165, plaintext -3.031523
    lola_test.go:99: logit 9: got -1.765203, want -3.831371 (+/- 0.10), delta 2.066167
    lola_test.go:104: logit 9: got -1.765203, plaintext -3.831371
FAIL
================================================================================

EDIT2: The more interesting issue is: why does this tiny quad-activation model bootstrap at all? And the answer is because there's apparently a really deep layout conversion network being materialised?

[asraa]

you could increase the level-budget to avoid bootstrap

[AlexanderViand]

you could increase the level-budget to avoid bootstrap

Looking at it, there should be no need to bootstrap - the conversion could be pushed into the plaintext encode for free, but it's not picked up by our system. Working on a fix for that now :)

[AlexanderViand]

I rebased and cleaned up the churn a bit (the bootstrap fix, for example, while discovered because of Lola, isn't even necessary anymore after we stopped the deep layout conversion from being materialised)

[mdgrs]

I changed the way the bounding box is computed. It now scans the rows to check if there is a relation that allows to deduce a bounding box immediately for each variable. I think this is mostly the case. It falls back to the previous (expensive) way to compute this if a row giving an immediate bound is not available.

Caveat: there could be cases where many inequalities, when taken together, give a tighter bounding box for a variable, in which case this new way would give the looser bound.

[j2kun]

So, I love the change in c8efe4b, the idea of eagerly repacking a layout conversion from the vector input to a matrix input is awesome.

However, we already implemented the underlying plaintext-repacking logic as part of the "hoisting" idea: hoisting the output layout conversion through to the input and (essentially) permuting the columns of the matrix to match.

heir/lib/Utils/Layout/Hoisting.h

Lines 13 to 24 in 13f16c8

	// Hoist the conversion of a vector layout through a matrix-vector multiply
	// operation.
	//
	// Returns a new layout for the matrix argument of the matvec op.
	//
	// Note: this function requires the assumption that the chosen packing for the
	// vector (and the corresponding matvec kernel) packs the vector into a single
	// ciphertext.
	presburger::IntegerRelation hoistConversionThroughMatvec(
	const presburger::IntegerRelation& matrixLayout,
	const presburger::IntegerRelation& fromVecLayout,
	const presburger::IntegerRelation& toVecLayout);

Is it too much to ask to find the core subset of logic that applies to both situations, and refactor appropriately? The two implementations look eerily similar to me, though I am reviewing this code while drinking sake and my reasoning may be clouded as such.

[j2kun]

Looks great, lots of nice improvements here. Just a few comments.

[j2kun]

Also, if the lola test is still failing, I think it would be good to pull some of these incremental improvements out into separate PRs and we can get them merged.

[mdgrs]

c8efe4b

See #3241

[mdgrs]

Also, if the lola test is still failing, I think it would be good to pull some of these incremental improvements out into separate PRs and we can get them merged.

I pulled two commits into their own PR and expanded them. Once we agree on these two pieces of work, I can rebase and drop the corresponding commits from this branch.

[mdgrs]

With the other two PRs green lit, I don't think much else is going on here other than adding the test. Let me know if I should squash, or if you would like to see some more changes

[j2kun]

Rebase and squash and let's see what's left.

[AlexanderViand]

Rebase and squash and let's see what's left.

Rebased, and it still needs both the fold-dense-resources and simplex-in-materialize-layout changes to compile (I tried w/o either, or just one, and all got stuck in heir-opt for 30+ min before I killed them).

There's an (I assume) unrelated CI failure, but if the test were to run on CI, it'd have also failed after rebasing because the accuracy was still off (convolution init bug, fix cherry-picked over from our fork)

EDIT: to be discussed at office hours, but I guess I'll open three new PRs for those three changes/fixes

[j2kun]

LGTM! Can you squash and we'll get it in?

[AlexanderViand]

LGTM! Can you squash and we'll get it in?

@j2kun Done!