/
ops.mlir
48 lines (40 loc) · 1.88 KB
/
ops.mlir
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
// RUN: heir-opt %s | FileCheck %s
// RUN: heir-translate --emit-tfhe-rust %s | FileCheck --check-prefix=RS %s
// This simply tests for syntax.
!sks = !tfhe_rust.server_key
module {
// CHECK-LABEL: func @test_create_trivial
// RS-LABEL: pub fn test_create_trivial
func.func @test_create_trivial(%sks : !sks) {
%0 = arith.constant 1 : i8
%1 = arith.constant 1 : i3
%2 = arith.constant 1 : i128
%e1 = tfhe_rust.create_trivial %sks, %0 : (!sks, i8) -> !tfhe_rust.ei8
%eu1 = tfhe_rust.create_trivial %sks, %1 : (!sks, i3) -> !tfhe_rust.eui8
%e2 = tfhe_rust.create_trivial %sks, %2 : (!sks, i128) -> !tfhe_rust.ei128
return
}
// CHECK-LABEL: func @test_bitand
// RS-LABEL: pub fn test_bitand
func.func @test_bitand(%sks : !sks) {
%0 = arith.constant 1 : i1
%1 = arith.constant 1 : i1
%e1 = tfhe_rust.create_trivial %sks, %0 : (!sks, i1) -> !tfhe_rust.eui8
%e2 = tfhe_rust.create_trivial %sks, %1 : (!sks, i1) -> !tfhe_rust.eui8
%out = tfhe_rust.bitand %sks, %e1, %e2: (!sks, !tfhe_rust.eui8, !tfhe_rust.eui8) -> !tfhe_rust.eui8
return
}
// CHECK-LABEL: func @test_apply_lookup_table
// RS-LABEL: pub fn test_apply_lookup_table
func.func @test_apply_lookup_table(%sks : !sks, %lut: !tfhe_rust.lookup_table) {
%0 = arith.constant 1 : i3
%1 = arith.constant 2 : i3
%e1 = tfhe_rust.create_trivial %sks, %0 : (!sks, i3) -> !tfhe_rust.eui3
%e2 = tfhe_rust.create_trivial %sks, %0 : (!sks, i3) -> !tfhe_rust.eui3
%shiftAmount = arith.constant 1 : i8
%e2Shifted = tfhe_rust.scalar_left_shift %sks, %e2, %shiftAmount : (!sks, !tfhe_rust.eui3, i8) -> !tfhe_rust.eui3
%eCombined = tfhe_rust.add %sks, %e1, %e2Shifted : (!sks, !tfhe_rust.eui3, !tfhe_rust.eui3) -> !tfhe_rust.eui3
%out = tfhe_rust.apply_lookup_table %sks, %eCombined, %lut : (!sks, !tfhe_rust.eui3, !tfhe_rust.lookup_table) -> !tfhe_rust.eui3
return
}
}