make lattigo pipeline aware of lintrans rotations

Author: AlexanderViand

lib/Dialect/Lattigo/Transforms/ConfigureCryptoContext.cpp

@@ -223,6 +223,47 @@ struct LattigoCKKSScheme {
   }
 };
 
+// Returns the best baby-step size N1 for BSGS given diagonal indices, slot
+// count, and log2 of the target baby/giant ratio. Mirrors Lattigo's
+// lintrans.FindBestBSGSRatio.
+static int64_t findBestBSGSRatio(ArrayRef<int32_t> diags, int64_t slots,
+                                 int64_t logMaxRatio) {
+  int64_t maxRatio = 1LL << logMaxRatio;
+  for (int64_t N1 = 1; N1 < slots; N1 <<= 1) {
+    DenseSet<int64_t> rotN1Set, rotN2Set;
+    for (auto rot : diags) {
+      int64_t r = (int64_t)rot & (slots - 1);
+      rotN1Set.insert(((r / N1) * N1) & (slots - 1));
+      rotN2Set.insert(r & (N1 - 1));
+    }
+    int64_t nbN1 = (int64_t)rotN1Set.size() - 1;
+    int64_t nbN2 = (int64_t)rotN2Set.size() - 1;
+    if (nbN1 > 0) {
+      if (nbN2 == maxRatio * nbN1) return N1;
+      if (nbN2 > maxRatio * nbN1) return N1 / 2;
+    }
+  }
+  return 1;
+}
+
+// Returns all non-zero rotation indices needed by lintrans.EvaluateNew for a
+// CKKSLinearTransformOp with the given diagonal indices, slot count, and BSGS
+// ratio. Mirrors Lattigo's lintrans.GaloisElements().
+static DenseSet<int64_t> lintransRotationIndices(ArrayRef<int32_t> diags,
+                                                 int64_t slots,
+                                                 int64_t logBSGS) {
+  DenseSet<int64_t> result;
+  int64_t N1 = (logBSGS < 0) ? slots : findBestBSGSRatio(diags, slots, logBSGS);
+  for (auto rot : diags) {
+    int64_t r = (int64_t)rot & (slots - 1);
+    int64_t giant = ((r / N1) * N1) & (slots - 1);
+    int64_t baby = r & (N1 - 1);
+    if (giant != 0) result.insert(giant);
+    if (baby != 0) result.insert(baby);
+  }
+  return result;
+}
+
 template <typename LattigoScheme>
 LogicalResult convertFuncForScheme(func::FuncOp op) {
   using EvaluatorType = typename LattigoScheme::EvaluatorType;
@@ -314,6 +355,26 @@ LogicalResult convertFuncForScheme(func::FuncOp op) {
 
   auto setIndices = analysis.getRotationIndices();
   SmallVector<int64_t> rotIndices(setIndices.begin(), setIndices.end());
+
+  // Supplement RotationAnalysis with galois indices needed by
+  // CKKSLinearTransformOp (lintrans.EvaluateNew BSGS rotations).
+  // RotationAnalysis only sees explicit ckks.rotate ops.
+  {
+    int64_t slots = 1LL << (logN - 1);  // CKKS: N/2 slots
+    DenseSet<int64_t> seen(rotIndices.begin(), rotIndices.end());
+    walkFuncAndCallees(op, [&](Operation* innerOp) {
+      auto ltOp = dyn_cast<CKKSLinearTransformOp>(innerOp);
+      if (!ltOp) return WalkResult::advance();
+      int64_t logBSGS = ltOp.getLogBabyStepGiantStepRatio().getInt();
+      auto newRots = lintransRotationIndices(
+          ltOp.getDiagonalIndicesAttr().asArrayRef(), slots, logBSGS);
+      for (auto r : newRots) {
+        if (seen.insert(r).second) rotIndices.push_back(r);
+      }
+      return WalkResult::advance();
+    });
+  }
+
   LLVM_DEBUG({
     llvm::dbgs() << "Finished rotation analysis; found " << rotIndices.size()
                  << " rotations which were=\n";

tests/Examples/orion/linear_transform/BUILD

@@ -17,9 +17,6 @@ go_test(
     srcs = ["linear_transform_test.go"],
     embed = [":lineartransform"],
     deps = [
-        "@com_github_tuneinsight_lattigo_v6//circuits/ckks/lintrans",
-        "@com_github_tuneinsight_lattigo_v6//core/rlwe",
-        "@com_github_tuneinsight_lattigo_v6//ring",
         "@com_github_tuneinsight_lattigo_v6//schemes/ckks",
     ],
 )

tests/Examples/orion/linear_transform/linear_transform_test.go

@@ -1,140 +1,67 @@
 package lineartransform
 
 import (
-	"fmt"
-	"github.com/tuneinsight/lattigo/v6/circuits/ckks/lintrans"
-	"github.com/tuneinsight/lattigo/v6/core/rlwe"
-	"github.com/tuneinsight/lattigo/v6/ring"
-	"github.com/tuneinsight/lattigo/v6/schemes/ckks"
 	"math"
 	"testing"
+
+	"github.com/tuneinsight/lattigo/v6/schemes/ckks"
 )
 
 func TestLinearTransform(t *testing.T) {
-	// This test is a bit weird because we're skipping most of the pipeline
-	// to just test the emitter for linear_transform which is needed in the
-	// context of comparisons with the orion compiler. This requires us to
-	// manually set up the crypto parameters and encode/encrypt stuff.
+	evaluator, params, encoder, encryptor, decryptor := linear_transform__configure()
+	numSlots := params.MaxSlots()
 
-	// Input vector of all 1s
-	numSlots := 4096
 	inputClear := make([]float64, numSlots)
 	for i := range inputClear {
 		inputClear[i] = 1.0
 	}
 
-	// Matrix of weights. This represents two nonzero diagonals on an
-	// otherwise all-zero matrix:
-	//
-	// 0    4096
-	//         1  4097
-	//               2  4098
-	//                     ...
-	//                        4094  8190
-	// 8191                         4095
-	//
-	// In this way, the nonzero diagonals become
-	// diagonal 0: range(0, 4096)
-	// diagonal 1: range(4096, 8192)
+	// Matrix of weights: two nonzero diagonals on an otherwise all-zero matrix.
 	//
-	// and the expected values are (noting lattigo does
-	// left-multiplication of the cleartext matrix and the vector is all
-	// 1s):
+	// diagonal 0: [0, 1, 2, ..., 4095]
+	// diagonal 1: [4096, 4097, ..., 8191]
 	//
-	// [4096 + 0, 4097 + 1, ..., 8191 + 4095]
+	// With an all-1s input vector, expected output = diagonal0 + diagonal1:
+	// [4096+0, 4097+1, ..., 8191+4095]
 	diagonals := 2
-	cols := numSlots
-	// Matrix is flattened 2 x numSlots
-	matrix := make([]float64, diagonals*cols)
+	matrix := make([]float64, diagonals*numSlots)
 	value := 0.0
 	for r := 0; r < diagonals; r++ {
-		for c := 0; c < cols; c++ {
-			matrix[r*cols+c] = value
+		for c := 0; c < numSlots; c++ {
+			matrix[r*numSlots+c] = value
 			value++
 		}
 	}
 
 	expectedClear := make([]float64, numSlots)
 	for i := range expectedClear {
-		expectedClear[i] = float64(4096 + 2*i)
+		expectedClear[i] = float64(numSlots + 2*i)
 	}
 
-	// These parameters should match linear_transform.mlir, though due to
-	// the weird nature of this test, this is the source of truth for what
-	// is used, not the mlir file.
-	param, err := ckks.NewParametersFromLiteral(ckks.ParametersLiteral{
-		LogN:            13,
-		Q:               []uint64{536903681, 67043329, 66994177, 67239937, 66961409, 66813953},
-		P:               []uint64{536952833, 536690689},
-		LogDefaultScale: 26,
-	})
-	if err != nil {
-		panic(err)
-	}
-
-	encoder := ckks.NewEncoder(param)
-	kgen := rlwe.NewKeyGenerator(param)
-	sk, pk := kgen.GenKeyPairNew()
-	encryptor := rlwe.NewEncryptor(param, pk)
-	decryptor := rlwe.NewDecryptor(param, sk)
-
-	// This is copied from the generated code so we can get access to the
-	// Lattigo-produced Galois key set to generate... ideally this is moved
-	// to a shared client helper generated by HEIR, but the client
-	// interface generation is before lowering to scheme, and Orion enters
-	// post-lowering-to-scheme. For this test it's OK and, if nothing else,
-	// a good reference.
-	ct1diags := make(lintrans.Diagonals[float64])
-	for i := 0; i < 2; i++ {
-		ct1diags[i] = matrix[i*numSlots : (i+1)*numSlots]
+	pt := ckks.NewPlaintext(params, params.MaxLevel())
+	pt.Scale = params.DefaultScale()
+	if err := encoder.Encode(inputClear, pt); err != nil {
+		t.Fatal(err)
 	}
-	ct1params := lintrans.Parameters{
-		DiagonalsIndexList:        ct1diags.DiagonalsIndexList(),
-		LevelQ:                    5,
-		LevelP:                    param.MaxLevelP(),
-		Scale:                     rlwe.NewScale(param.Q()[5]),
-		LogDimensions:             ring.Dimensions{Rows: 0, Cols: 12}, // 1x4096
-		LogBabyStepGiantStepRatio: 2,
-	}
-	ct1lt := lintrans.NewTransformation(param, ct1params)
-	galEls := ct1lt.GaloisElements(param)
-
-	// Manually add Galois key for rotation index 2048
-	rotIndex := 2048
-	logN := 13
-	galoisElement := uint64(1)
-	for i := 0; i < rotIndex; i++ {
-		galoisElement = (galoisElement * 5) % (1 << (logN + 1))
-	}
-	galEls = append(galEls, galoisElement)
-	fmt.Printf("Final galEls: %v\n", galEls)
-
-	evk := rlwe.NewMemEvaluationKeySet(nil, kgen.GenGaloisKeysNew(galEls, sk)...)
-	evaluator := ckks.NewEvaluator(param, evk)
-
-	pt := ckks.NewPlaintext(param, param.MaxLevel())
-	pt.LogDimensions = ring.Dimensions{Rows: 0, Cols: 12} // 2^(0+12) = 4096 slots
-	encoder.Encode(inputClear, pt)
-	ctInput, err25 := encryptor.EncryptNew(pt)
-	if err25 != nil {
-		panic(err25)
+	ct, err := encryptor.EncryptNew(pt)
+	if err != nil {
+		t.Fatal(err)
 	}
 
-	resultCt := linear_transform(evaluator, param, encoder, ctInput, matrix)
+	resultCt := linear_transform(evaluator, params, encoder, ct, matrix)
 	resultPt := decryptor.DecryptNew(resultCt)
-	resultFloat64 := make([]float64, 4096)
-	encoder.Decode(resultPt, resultFloat64)
+	resultFloat64 := make([]float64, numSlots)
+	if err := encoder.Decode(resultPt, resultFloat64); err != nil {
+		t.Fatal(err)
+	}
 
-	// We need such a large epsilon because scale 26 is not very precise,
-	// increasing scale to 40 produces errors of about 1e-04.
+	// Scale 26 is not very precise; epsilon of 1.5 is sufficient here.
 	epsilon := 1.5
 	for i := 0; i < numSlots; i++ {
 		diff := math.Abs(resultFloat64[i] - expectedClear[i])
 		if diff > epsilon {
 			t.Errorf("Mismatch at index %d: got %f, expected %f (diff: %e)",
 				i, resultFloat64[i], expectedClear[i], diff)
-
-			// Fail fast to avoid spamming 4096 errors
 			if i > 10 {
 				t.Fatal("Too many errors, stopping verification.")
 			}

tests/Examples/orion/mlp/BUILD

@@ -17,9 +17,6 @@ go_test(
     srcs = ["mlp_test.go"],
     embed = [":mlp"],
     deps = [
-        "@com_github_tuneinsight_lattigo_v6//circuits/ckks/lintrans",
-        "@com_github_tuneinsight_lattigo_v6//core/rlwe",
-        "@com_github_tuneinsight_lattigo_v6//ring",
         "@com_github_tuneinsight_lattigo_v6//schemes/ckks",
     ],
 )