Lowering a matmul loop - plaintext/ciphertext type agreement
- Dominant language
- MLIR
- Stars
- 906
- Forks
- 171
- Avg merge
- 4d 12h
- Merged PRs (30d)
- 32
Description
This is a small question about how to approach type agreement when lowering a loop that has a secret loop-carried variable with an initial plaintext value. This particular loop input is an output of the Halevi-Shoup diagonalization, so I expect it to be a somewhat common pattern for lowering.
Here's the secret / arith model performing the diagonalized matmul:
```
func.func @mlp(%arg0: !secret.secret>) -> !secret.secret> {
%cst_0 = arith.constant dense_resource<__elided__> : tensor<1x1024xf32>
%c1 = arith.constant 1 : index
%cst_2 = arith.constant dense_resource<__elided__> : tensor<1024x1024xf32>
%cst_3 = arith.constant dense<0.000000e+00> : tensor<1x1024xf32>
%0 = secret.generic ins(%arg0 : !secret.secret>) {
^body(%input0: tensor<1x1024xf32>):
%1:2 = affine.for %arg1 = 0 to 1023 iter_args(%arg2 = %cst_3, %arg3 = %input0) -> (tensor<1x1024xf32>, tensor<1x1024xf32>) {
%extracted_slice = tensor.extract_slice %cst_2[%arg1, 0] [1, 1024] [1, 1] : tensor<1024x1024xf32> to tensor<1x1024xf32>
%8 = arith.mulf %arg3, %extracted_slice : tensor<1x1024xf32>
%9 = arith.addf %arg2, %8 : tensor<1x1024xf32>
%10 = tensor_ext.rotate %arg3, %c1 : tensor<1x1024xf32>, index
affine.yield %9, %10 : tensor<1x1024xf32>, tensor<1x1024xf32>
}
%2 = arith.mulf %1#1, %cst_0 : tensor<1x1024xf32>
%3 = arith.addf %1#0, %2 : tensor<1x1024xf32>
secret.yield %7 : tensor<1x1024xf32>
} -> !secret.secret>
return %0 : !secret.secret>
}
```
The main type issue is that we are starting with a cleartext bias vector, and then multiplying it with a ciphertext and accumulating it as a ciphertext result. The question is:
*What is the type of the initial loop carried variable? An LWE plaintext, or an LWE ciphertext?*
Conumdrum:
* If it is an LWE plaintext, then the affine for loop doesn't like that the init_arg is a plaintext, but the loop yielded value is a ciphertext
* If it is an LWE ciphertext, then how do we perform type conversion from a cleartext to a ciphertext, with no info on the context level through the mgmt attr?
Note that `--secret-distribute-generic` handles this type issue by `secret.conceal`ing the initial iter_arg, and that makes everything happy as secret types, but at the scheme level, we have a plaintext/ciphertext distrinction. Obviously secret.conceal should lower to an encode operation at the scheme level, but then we have conumdrum point 1. Here's the result after distribute generic and the insert mgmt pass:
```
func.func @mlp(%arg0: !secret.secret> {mgmt.mgmt = #mgmt.mgmt}) -> !secret.secret> attributes {llvm.emit_c_interface} {
%cst = arith.constant dense_resource<__elided__> : tensor<1x1024xf32>
%c1 = arith.constant 1 : index
%cst_0 = arith.constant dense_resource<__elided__> : tensor<1024x1024xf32>
%cst_1 = arith.constant dense<0.000000e+00> : tensor<1x1024xf32>
%0 = secret.conceal %cst_1 : tensor<1x1024xf32> -> >
%1:2 = affine.for %arg1 = 0 to 1023 iter_args(%arg2 = %0, %arg3 = %arg0) -> (!secret.secret>, !secret.secret>) {
%extracted_slice = tensor.extract_slice %cst_0[%arg1, 0] [1, 1024] [1, 1] : tensor<1024x1024xf32> to tensor<1x1024xf32>
%5 = secret.generic ins(%arg3 : !secret.secret>) attrs = {mgmt.mgmt = #mgmt.mgmt} {
^body(%input0: tensor<1x1024xf32>):
%8 = arith.mulf %input0, %extracted_slice : tensor<1x1024xf32>
secret.yield %8 : tensor<1x1024xf32>
} -> !secret.secret>
%6 = secret.generic ins(%arg2, %5 : !secret.secret>, !secret.secret>) attrs = {mgmt.mgmt = #mgmt.mgmt} {
^body(%input0: tensor<1x1024xf32>, %input1: tensor<1x1024xf32>):
%8 = arith.addf %input0, %input1 : tensor<1x1024xf32>
secret.yield %8 : tensor<1x1024xf32>
} -> !secret.secret>
%7 = secret.generic ins(%arg3 : !secret.secret>) attrs = {mgmt.mgmt = #mgmt.mgmt} {
^body(%input0: tensor<1x1024xf32>):
%8 = tensor_ext.rotate %input0, %c1 : tensor<1x1024xf32>, index
secret.yield %8 : tensor<1x1024xf32>
} -> !secret.secret>
affine.yield %6, %7 : !secret.secret>, !secret.secret>
} {mgmt.mgmt = #mgmt.mgmt}
%2 = secret.generic ins(%1#1 : !secret.secret>) attrs = {mgmt.mgmt = #mgmt.mgmt} {
^body(%input0: tensor<1x1024xf32>):
%5 = arith.mulf %input0, %cst : tensor<1x1024xf32>
secret.yield %5 : tensor<1x1024xf32>
} -> !secret.secret>
%3 = secret.generic ins(%1#0, %2 : !secret.secret>, !secret.secret>) attrs = {mgmt.mgmt = #mgmt.mgmt} {
^body(%input0: tensor<1x1024xf32>, %input1: tensor<1x1024xf32>):
%5 = arith.addf %input0, %input1 : tensor<1x1024xf32>
secret.yield %5 : tensor<1x1024xf32>
} -> !secret.secret>
%4 = secret.generic ins(%3 : !secret.secret>) attrs = {mgmt.mgmt = #mgmt.mgmt} {
^body(%input0: tensor<1x1024xf32>):
%5 = mgmt.modreduce %input0 : tensor<1x1024xf32>
secret.yield %5 : tensor<1x1024xf32>
} -> !secret.secret>
return %4 : !secret.secret>
}
```
This is the IR input for secret-to-ckks which will then fail because
```
error: failed to legalize operation 'affine.for' that was explicitly marked illegal
%1:2 = affine.for %arg1 = 0 to 1023 iter_args(%arg2 = %cst_3, %arg3 = %input0) -> (tensor<1x1024xf32>, tensor<1x1024xf32>) {
^
test.mlir:10:14: note: see current operation:
%6:2 = "affine.for"(%5, %arg0) <{lowerBoundMap = affine_map<() -> (0)>, operandSegmentSizes = array, step = 1 : index, upperBoundMap = affine_map<() -> (1023)>}> ({
^bb0(%arg5: index, %arg6: !secret.secret>, %arg7: !secret.secret>):
%13 = "tensor.extract_slice"(%2, %arg5) <{operandSegmentSizes = array, static_offsets = array, static_sizes = array, static_strides = array}> : (tensor<1024x1024xf32>, index) -> tensor<1x1024xf32>
%14 = "secret.generic"(%arg7) ({
^bb0(%arg11: tensor<1x1024xf32>):
%19 = "arith.mulf"(%arg11, %13) <{fastmath = #arith.fastmath}> : (tensor<1x1024xf32>, tensor<1x1024xf32>) -> tensor<1x1024xf32>
"secret.yield"(%19) : (tensor<1x1024xf32>) -> ()
}) {mgmt.mgmt = #mgmt.mgmt} : (!secret.secret>) -> !secret.secret>
%15 = "secret.generic"(%arg6, %14) ({
^bb0(%arg9: tensor<1x1024xf32>, %arg10: tensor<1x1024xf32>):
%18 = "arith.addf"(%arg9, %arg10) <{fastmath = #arith.fastmath}> : (tensor<1x1024xf32>, tensor<1x1024xf32>) -> tensor<1x1024xf32>
"secret.yield"(%18) : (tensor<1x1024xf32>) -> ()
}) {mgmt.mgmt = #mgmt.mgmt} : (!secret.secret>, !secret.secret>) -> !secret.secret>
%16 = "secret.generic"(%arg7) ({
^bb0(%arg8: tensor<1x1024xf32>):
%17 = "tensor_ext.rotate"(%arg8, %1) : (tensor<1x1024xf32>, index) -> tensor<1x1024xf32>
"secret.yield"(%17) : (tensor<1x1024xf32>) -> ()
}) {mgmt.mgmt = #mgmt.mgmt} : (!secret.secret>) -> !secret.secret>
"affine.yield"(%15, %16) : (!secret.secret>, !secret.secret>) -> ()
}) {mgmt.mgmt = #mgmt.mgmt} : (!secret.secret>, !lwe.new_lwe_ciphertext<...>) -> (!secret.secret>, !secret.secret>)
```
Contributor guide
Assessment
This issue has not been assessed yet.