github/concrete
1
1
Fork 0
mirror of https://github.com/zama-ai/concrete.git synced 2026-04-17 03:00:54 -04:00
Code Issues Packages Projects Releases Wiki Activity

test(compiler): Fix the size of lookup tables used in tests with bootstrap operations

Browse Source 
Some of the tests use lookup tables whose numbers of elements do not
match the sizes of the polynoms of the bootstrap operations they are
passed to. This commit replaces these lookup tables with tables of the
right size.
This commit is contained in:
Andi Drebes
2024-01-29 15:04:05 +01:00
parent f113e8bcfe
commit a133407035
4 changed files with 72 additions and 25 deletions
Download Patch File Download Diff File Expand all files Collapse all files

8
compilers/concrete-compiler/compiler/tests/check_tests/Conversion/TFHEGlobalParametrization/pbs_ks_bs.mlir
View File

@@ -1,14 +1,14 @@
// RUN: concretecompiler --action=dump-parametrized-tfhe --optimizer-strategy=V0 --v0-parameter=2,10,750,1,23,3,4 --v0-constraint=4,0 %s 2>&1| FileCheck %s
//CHECK: func.func @main(%[[A0:.*]]: !TFHE.glwe<sk<0,1,2048>>) -> !TFHE.glwe<sk<0,1,2048>> {
//CHECK-NEXT: %cst = arith.constant dense<[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]> : tensor<16xi64>
//CHECK-NEXT: %cst = arith.constant dense<0> : tensor<1024xi64>
//CHECK-NEXT: %[[V1:.*]] = "TFHE.keyswitch_glwe"(%[[A0]]) {key = #TFHE.ksk<sk<0,1,2048>, sk<1,1,750>, 3, 4>} : (!TFHE.glwe<sk<0,1,2048>>) -> !TFHE.glwe<sk<1,1,750>>
//CHECK-NEXT: %[[V2:.*]] = "TFHE.bootstrap_glwe"(%[[V1]], %cst) {key = #TFHE.bsk<sk<1,1,750>, sk<0,1,2048>, 1024, 2, 1, 23>} : (!TFHE.glwe<sk<1,1,750>>, tensor<16xi64>) -> !TFHE.glwe<sk<0,1,2048>>
//CHECK-NEXT: %[[V2:.*]] = "TFHE.bootstrap_glwe"(%[[V1]], %cst) {key = #TFHE.bsk<sk<1,1,750>, sk<0,1,2048>, 1024, 2, 1, 23>} : (!TFHE.glwe<sk<1,1,750>>, tensor<1024xi64>) -> !TFHE.glwe<sk<0,1,2048>>
//CHECK-NEXT: return %[[V2]] : !TFHE.glwe<sk<0,1,2048>>
//CHECK-NEXT: }
func.func @main(%arg0: !TFHE.glwe<sk?>) -> !TFHE.glwe<sk?> {
%cst = arith.constant dense<[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]> : tensor<16xi64>
%cst = arith.constant dense<0> : tensor<1024xi64>
%1 = "TFHE.keyswitch_glwe"(%arg0) {key = #TFHE.ksk<sk?, sk?, -1, -1>} : (!TFHE.glwe<sk?>) -> !TFHE.glwe<sk?>
%2 = "TFHE.bootstrap_glwe"(%1, %cst) {key = #TFHE.bsk<sk?, sk?, -1, -1, -1, -1>} : (!TFHE.glwe<sk?>, tensor<16xi64>) -> !TFHE.glwe<sk?>
%2 = "TFHE.bootstrap_glwe"(%1, %cst) {key = #TFHE.bsk<sk?, sk?, -1, -1, -1, -1>} : (!TFHE.glwe<sk?>, tensor<1024xi64>) -> !TFHE.glwe<sk?>
return %2 : !TFHE.glwe<sk?>
}

8
compilers/concrete-compiler/compiler/tests/check_tests/Conversion/TFHEToConcrete/bootstrap.mlir
View File

File diff suppressed because one or more lines are too long

8
compilers/concrete-compiler/compiler/tests/check_tests/Dialect/TFHE/ops.mlir
View File

@@ -8,11 +8,11 @@ func.func @keyswitch_glwe(%arg0: !TFHE.glwe<sk[1]<1024,1>>) -> !TFHE.glwe<sk[1]<
return %0: !TFHE.glwe<sk[1]<527,1>>
}
// CHECK: func.func @bootstrap_glwe(%[[GLWE:.*]]: !TFHE.glwe<sk[1]<527,1>>, %[[LUT:.*]]: tensor<128xi64>) -> !TFHE.glwe<sk[1]<1024,1>> {
func.func @bootstrap_glwe(%glwe: !TFHE.glwe<sk[1]<527,1>>, %lut: tensor<128xi64>) -> !TFHE.glwe<sk[1]<1024,1>> {
// CHECK-NEXT: %[[V0:.*]] = "TFHE.bootstrap_glwe"(%[[GLWE]], %[[LUT]]) {key = #TFHE.bsk<sk[1]<527,1>, sk[1]<1024,1>, 512, 2, 4, 4>} : (!TFHE.glwe<sk[1]<527,1>>, tensor<128xi64>) -> !TFHE.glwe<sk[1]<1024,1>>
// CHECK: func.func @bootstrap_glwe(%[[GLWE:.*]]: !TFHE.glwe<sk[1]<527,1>>, %[[LUT:.*]]: tensor<512xi64>) -> !TFHE.glwe<sk[1]<1024,1>> {
func.func @bootstrap_glwe(%glwe: !TFHE.glwe<sk[1]<527,1>>, %lut: tensor<512xi64>) -> !TFHE.glwe<sk[1]<1024,1>> {
// CHECK-NEXT: %[[V0:.*]] = "TFHE.bootstrap_glwe"(%[[GLWE]], %[[LUT]]) {key = #TFHE.bsk<sk[1]<527,1>, sk[1]<1024,1>, 512, 2, 4, 4>} : (!TFHE.glwe<sk[1]<527,1>>, tensor<512xi64>) -> !TFHE.glwe<sk[1]<1024,1>>
// CHECK-NEXT: return %[[V0]] : !TFHE.glwe<sk[1]<1024,1>>
%0 = "TFHE.bootstrap_glwe"(%glwe, %lut) {key=#TFHE.bsk<sk[1]<527,1>,sk[1]<1024,1>,512,2,4,4>} : (!TFHE.glwe<sk[1]<527,1>>, tensor<128xi64>) -> !TFHE.glwe<sk[1]<1024,1>>
%0 = "TFHE.bootstrap_glwe"(%glwe, %lut) {key=#TFHE.bsk<sk[1]<527,1>,sk[1]<1024,1>,512,2,4,4>} : (!TFHE.glwe<sk[1]<527,1>>, tensor<512xi64>) -> !TFHE.glwe<sk[1]<1024,1>>
return %0 : !TFHE.glwe<sk[1]<1024,1>>
}

73
compilers/concrete-compiler/compiler/tests/check_tests/Transforms/batching.mlir
View File

@@ -197,8 +197,8 @@ func.func @batch_offset_shifted_bounds_nonunitstep_extract_keyswitch(%arg0: tens
// CHECK: (%arg0: tensor<2x3x4x!TFHE.glwe<sk{{\[}}[[SK_IN:.*]]{{\]}}<1,2048>>>, %arg1: tensor<4xi64>) -> tensor<2x3x4x!TFHE.glwe<sk{{\[}}[[SK_IN]]{{\]}}<1,2048>>> {
// CHECK: %[[Vcollapsed:.*]] = tensor.collapse_shape %[[Varg0:.*]] {{\[\[0, 1, 2\]\]}} : tensor<2x3x4x!TFHE.glwe<sk{{\[}}[[SK_IN]]{{\]}}<1,2048>>> into tensor<24x!TFHE.glwe<sk{{\[}}[[SK_IN]]{{\]}}<1,2048>>>
// CHECK-NEXT: %[[V1:.*]] = "TFHE.batched_keyswitch_glwe"(%[[Vcollapsed]]) {key = #TFHE<ksk{{\[}}[[KSK:.*]]{{\]}}<sk{{\[}}[[SK_IN]]{{\]}}<1,2048>, sk{{\[}}[[SK_OUT:.*]]{{\]}}<1,750>, 3, 4>>} : (tensor<24x!TFHE.glwe<sk{{\[}}[[SK_IN]]{{\]}}<1,2048>>>) -> tensor<24x!TFHE.glwe<sk{{\[}}[[SK_OUT]]{{\]}}<1,750>>>
// CHECK-NEXT: %[[V2:.*]] = "TFHE.encode_expand_lut_for_bootstrap"(%[[Varg1:.*]]) {isSigned = false, outputBits = 2 : i32, polySize = 256 : i32} : (tensor<4xi64>) -> tensor<256xi64>
// CHECK-NEXT: %[[V3:.*]] = "TFHE.batched_bootstrap_glwe"(%[[V1]], %[[V2]]) {key = #TFHE<bsk{{\[}}[[BSK:.*]]{{\]}}<sk{{\[}}[[SK_IN]]{{\]}}<1,2048>, sk{{\[}}[[SK_IN]]{{\]}}<1,2048>, 1024, 2, 1, 23>>} : (tensor<24x!TFHE.glwe<sk{{\[}}[[SK_OUT]]{{\]}}<1,750>>>, tensor<256xi64>) -> tensor<24x!TFHE.glwe<sk{{\[}}[[SK_IN]]{{\]}}<1,2048>>>
// CHECK-NEXT: %[[V2:.*]] = "TFHE.encode_expand_lut_for_bootstrap"(%[[Varg1:.*]]) {isSigned = false, outputBits = 2 : i32, polySize = 1024 : i32} : (tensor<4xi64>) -> tensor<1024xi64>
// CHECK-NEXT: %[[V3:.*]] = "TFHE.batched_bootstrap_glwe"(%[[V1]], %[[V2]]) {key = #TFHE<bsk{{\[}}[[BSK:.*]]{{\]}}<sk{{\[}}[[SK_IN]]{{\]}}<1,2048>, sk{{\[}}[[SK_IN]]{{\]}}<1,2048>, 1024, 2, 1, 23>>} : (tensor<24x!TFHE.glwe<sk{{\[}}[[SK_OUT]]{{\]}}<1,750>>>, tensor<1024xi64>) -> tensor<24x!TFHE.glwe<sk{{\[}}[[SK_IN]]{{\]}}<1,2048>>>
// CHECK-NEXT: %[[Vexpanded:.*]] = tensor.expand_shape %[[V3]] {{\[\[0, 1, 2\]\]}} : tensor<24x!TFHE.glwe<sk{{\[}}[[SK_IN]]{{\]}}<1,2048>>> into tensor<2x3x4x!TFHE.glwe<sk{{\[}}[[SK_IN]]{{\]}}<1,2048>
func.func @apply_lookup_table_contiguous(%arg0: tensor<2x3x4x!TFHE.glwe<sk<0,1,2048>>>, %arg1: tensor<4xi64>) -> tensor<2x3x4x!TFHE.glwe<sk<0,1,2048>>> {
%c0 = arith.constant 0 : index
@@ -211,9 +211,9 @@ func.func @apply_lookup_table_contiguous(%arg0: tensor<2x3x4x!TFHE.glwe<sk<0,1,2
%2 = scf.for %arg4 = %c0 to %c3 step %c1 iter_args(%arg5 = %arg3) -> (tensor<2x3x4x!TFHE.glwe<sk<0,1,2048>>>) {
%3 = scf.for %arg6 = %c0 to %c4 step %c1 iter_args(%arg7 = %arg5) -> (tensor<2x3x4x!TFHE.glwe<sk<0,1,2048>>>) {
%extracted = tensor.extract %arg0[%arg2, %arg4, %arg6] : tensor<2x3x4x!TFHE.glwe<sk<0,1,2048>>>
%4 = "TFHE.encode_expand_lut_for_bootstrap"(%arg1) {isSigned = false, outputBits = 2 : i32, polySize = 256 : i32} : (tensor<4xi64>) -> tensor<256xi64>
%4 = "TFHE.encode_expand_lut_for_bootstrap"(%arg1) {isSigned = false, outputBits = 2 : i32, polySize = 1024 : i32} : (tensor<4xi64>) -> tensor<1024xi64>
%5 = "TFHE.keyswitch_glwe"(%extracted) {key = #TFHE.ksk<sk<0,1,2048>, sk<1,1,750>, 3, 4>} : (!TFHE.glwe<sk<0,1,2048>>) -> !TFHE.glwe<sk<1,1,750>>
%6 = "TFHE.bootstrap_glwe"(%5, %4) {key = #TFHE.bsk<sk<0,1,2048>, sk<0,1,2048>, 1024, 2, 1, 23>} : (!TFHE.glwe<sk<1,1,750>>, tensor<256xi64>) -> !TFHE.glwe<sk<0,1,2048>>
%6 = "TFHE.bootstrap_glwe"(%5, %4) {key = #TFHE.bsk<sk<0,1,2048>, sk<0,1,2048>, 1024, 2, 1, 23>} : (!TFHE.glwe<sk<1,1,750>>, tensor<1024xi64>) -> !TFHE.glwe<sk<0,1,2048>>
%inserted = tensor.insert %6 into %arg7[%arg2, %arg4, %arg6] : tensor<2x3x4x!TFHE.glwe<sk<0,1,2048>>>
scf.yield %inserted : tensor<2x3x4x!TFHE.glwe<sk<0,1,2048>>>
}
@@ -229,14 +229,13 @@ func.func @apply_lookup_table_contiguous(%arg0: tensor<2x3x4x!TFHE.glwe<sk<0,1,2
// CHECK-LABEL: func.func @apply_lookup_table_contiguous_hoistable_nonbatchable
// CHECK: (%arg0: tensor<2x3x4x!TFHE.glwe<sk{{\[}}[[SK_IN:.*]]{{\]}}<1,2048>>>, %arg1: tensor<4xi64>) -> tensor<2x3x4x!TFHE.glwe<sk{{\[}}[[SK_IN]]{{\]}}<1,2048>>> {
// CHECK: %[[V1:.*]] = "TFHE.batched_keyswitch_glwe"(%[[Vcollapsed]]) {key = #TFHE<ksk{{\[}}[[KSK:.*]]{{\]}}<sk{{\[}}[[SK_IN]]{{\]}}<1,2048>, sk{{\[}}[[SK_OUT:.*]]{{\]}}<1,750>, 3, 4>>} : (tensor<24x!TFHE.glwe<sk{{\[}}[[SK_IN]]{{\]}}<1,2048>>>) -> tensor<24x!TFHE.glwe<sk{{\[}}[[SK_OUT]]{{\]}}<1,750>>>
// CHECK: %[[V3:.*]] = "TFHE.batched_bootstrap_glwe"(%[[V1]], %[[V2:.*]]) {key = #TFHE<bsk{{\[}}[[BSK:.*]]{{\]}}<sk{{\[}}[[SK_IN]]{{\]}}<1,2048>, sk{{\[}}[[SK_IN]]{{\]}}<1,2048>, 1024, 2, 1, 23>>} : (tensor<24x!TFHE.glwe<sk{{\[}}[[SK_OUT]]{{\]}}<1,750>>>, tensor<256xi64>) -> tensor<24x!TFHE.glwe<sk{{\[}}[[SK_IN]]{{\]}}<1,2048>>>
// CHECK: %[[V3:.*]] = "TFHE.batched_bootstrap_glwe"(%[[V1]], %[[V2:.*]]) {key = #TFHE<bsk{{\[}}[[BSK:.*]]{{\]}}<sk{{\[}}[[SK_IN]]{{\]}}<1,2048>, sk{{\[}}[[SK_IN]]{{\]}}<1,2048>, 1024, 2, 1, 23>>} : (tensor<24x!TFHE.glwe<sk{{\[}}[[SK_OUT]]{{\]}}<1,750>>>, tensor<1024xi64>) -> tensor<24x!TFHE.glwe<sk{{\[}}[[SK_IN]]{{\]}}<1,2048>>>
func.func @apply_lookup_table_contiguous_hoistable_nonbatchable(%arg0: tensor<2x3x4x!TFHE.glwe<sk<0,1,2048>>>, %arg1: tensor<4xi64>) -> tensor<2x3x4x!TFHE.glwe<sk<0,1,2048>>> {
%c0 = arith.constant 0 : index
%c2 = arith.constant 2 : index
%c1 = arith.constant 1 : index
%c3 = arith.constant 3 : index
%c4 = arith.constant 4 : index
%c256 = arith.constant 256 : index
%c1_i64 = arith.constant 1 : i64
%c0_i64 = arith.constant 0 : i64
%0 = bufferization.alloc_tensor() : tensor<2x3x4x!TFHE.glwe<sk<0,1,2048>>>
@@ -244,8 +243,56 @@ func.func @apply_lookup_table_contiguous_hoistable_nonbatchable(%arg0: tensor<2x
%2 = scf.for %arg4 = %c0 to %c3 step %c1 iter_args(%arg5 = %arg3) -> (tensor<2x3x4x!TFHE.glwe<sk<0,1,2048>>>) {
%3 = scf.for %arg6 = %c0 to %c4 step %c1 iter_args(%arg7 = %arg5) -> (tensor<2x3x4x!TFHE.glwe<sk<0,1,2048>>>) {
%extracted = tensor.extract %arg0[%arg2, %arg4, %arg6] : tensor<2x3x4x!TFHE.glwe<sk<0,1,2048>>>
%4 = "TFHE.encode_expand_lut_for_bootstrap"(%arg1) {isSigned = false, outputBits = 2 : i32, polySize = 256 : i32} : (tensor<4xi64>) -> tensor<256xi64>
%4 = "TFHE.encode_expand_lut_for_bootstrap"(%arg1) {isSigned = false, outputBits = 2 : i32, polySize = 1024 : i32} : (tensor<4xi64>) -> tensor<1024xi64>
%aa = tensor.from_elements
%c1_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c1_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c1_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c1_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
@@ -262,13 +309,13 @@ func.func @apply_lookup_table_contiguous_hoistable_nonbatchable(%arg0: tensor<2x
%c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64,
%c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64, %c0_i64
: tensor<256xi64>
%b = tensor.extract_slice %4[0][255][1] : tensor<256xi64> to tensor<255xi64>
%x = tensor.extract %4[%c0] : tensor<256xi64>
%c = tensor.insert_slice %b into %aa[1][255][1] : tensor<255xi64> into tensor<256xi64>
%d = tensor.insert %x into %c[%c0] : tensor<256xi64>
: tensor<1024xi64>
%b = tensor.extract_slice %4[0][255][1] : tensor<1024xi64> to tensor<255xi64>
%x = tensor.extract %4[%c0] : tensor<1024xi64>
%c = tensor.insert_slice %b into %aa[1][255][1] : tensor<255xi64> into tensor<1024xi64>
%d = tensor.insert %x into %c[%c0] : tensor<1024xi64>
%5 = "TFHE.keyswitch_glwe"(%extracted) {key = #TFHE.ksk<sk<0,1,2048>, sk<1,1,750>, 3, 4>} : (!TFHE.glwe<sk<0,1,2048>>) -> !TFHE.glwe<sk<1,1,750>>
%6 = "TFHE.bootstrap_glwe"(%5, %d) {key = #TFHE.bsk<sk<0,1,2048>, sk<0,1,2048>, 1024, 2, 1, 23>} : (!TFHE.glwe<sk<1,1,750>>, tensor<256xi64>) -> !TFHE.glwe<sk<0,1,2048>>
%6 = "TFHE.bootstrap_glwe"(%5, %d) {key = #TFHE.bsk<sk<0,1,2048>, sk<0,1,2048>, 1024, 2, 1, 23>} : (!TFHE.glwe<sk<1,1,750>>, tensor<1024xi64>) -> !TFHE.glwe<sk<0,1,2048>>
%inserted = tensor.insert %6 into %arg7[%arg2, %arg4, %arg6] : tensor<2x3x4x!TFHE.glwe<sk<0,1,2048>>>
scf.yield %inserted : tensor<2x3x4x!TFHE.glwe<sk<0,1,2048>>>
}