google / google/heir

Convert to ciphertext semantics incorrectly converts rotation shift to a tensor

Open
#3,189 0 comments 0 reactions 0 assignees View on GitHub
Dominant language
MLIR
Stars
906
Forks
171
Avg merge
4d 12h
Merged PRs (30d)
32

Description

```
module {
func.func @rotate_i32_shift(
%arg0: !secret.secret>)
-> !secret.secret> {
%shift = arith.constant 1 : i32

%0 = secret.generic(%arg0: !secret.secret>) {
^body(%input: tensor<1x64xi32>):
%r = tensor_ext.rotate %input, %shift : tensor<1x64xi32>, i32
secret.yield %r : tensor<1x64xi32>
} -> !secret.secret>

return %0 : !secret.secret>
}
}
```
With the above code as an input snippet -> `bazel run //tools:heir-opt -- /heir/biscotti-benchmarks/matrix_multiplication/test.mlir --layout-propagation="ciphertext-size=64" --convert-to-ciphertext-semantics="ciphertext-size=64"`

Output
```
module {
func.func @rotate_i32_shift(%arg0: !secret.secret> {tensor_ext.original_type = #tensor_ext.original_type, layout = #tensor_ext.layout<"{ [i0, i1] -> [ct, slot] : i0 = 0 and ct = 0 and (-i1 + slot) mod 64 = 0 and 0 <= i1 <= 63 and 0 <= slot <= 63 }">>}) -> !secret.secret> {
%cst = arith.constant dense<1> : tensor<1x64xi32>
%0 = secret.generic(%arg0: !secret.secret>) {
^body(%input0: tensor<1x64xi32>):
%1 = tensor_ext.rotate %input0, %cst : tensor<1x64xi32>, tensor<1x64xi32>
secret.yield %1 : tensor<1x64xi32>
} -> !secret.secret>
return %0 : !secret.secret>
}
}
```

The convert-to-ciphertext-semantics pass does a wrong transformation on `tensor_ext.rotate`, earlier the shift constant argument was a `i32` but after conversion to ciphertext semantics it became `tensor<1x64xi32>`. It should remain `i32` as a tensor argument on a shift constant argument is semantically wrong.

Contributor guide

Open the contributing guide

Research direction

Start with the tools:heir-opt invocation and the convert-to-ciphertext-semantics pass using the supplied MLIR snippet and ciphertext-size=64 options. Trace how tensor_ext.rotate operands are converted; done means the shift remains i32 while tensor operands are converted appropriately, with the shown output no longer containing a tensor shift constant.

Written by the indexing model from the issue text.

Assessment

Domain
compilers
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Quiet
Clarity
Clearly specified
Newbie friendliness
68/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.