google / heir Goto Github PK

View Code? Open in Web Editor NEW

196.0 14.0 24.0 7.24 MB

A compiler for homomorphic encryption

Home Page: https://heir.dev/

License: Apache License 2.0

Starlark 9.46% C++ 43.80% C 0.98% MLIR 34.72% Python 2.26% Verilog 6.32% Rust 0.67% Jinja 1.79%

homomorphic-encryption mlir

heir's Introduction

HEIR: Homomorphic Encryption Intermediate Representation

An MLIR-based toolchain for homomorphic encryption compilers. Read the docs at the HEIR website.

For more information on MLIR, see the MLIR homepage.

Contributing

There are many ways to contribute to HEIR:

Come to our twice-monthly meetings to discuss active work on HEIR and future project directions. The meetings are recorded and posted to our blog
Work on an issue marked "good first issue" or browse issues labeled by topic.
Tell us about what you'd like to use HEIR for in an offline discussion thread.
Help us understand new FHE research: either
- Read a paper tagged under research synthesis and summarize the novel techniques that could be ported to HEIR.
- File new issues under research synthesis to alert us of papers that should be investigated and incorporated into HEIR.

Support disclaimer

This is not an officially supported Google product.

heir's People

Contributors

Stargazers

Watchers

heir's Issues

poly: add `ElementwiseMappable` trait to ops.

please assign to me

poly: add correctness tests for lowerings to LLVM

Here's a skeleton of a test.

It seems tricky to get arith ops on tensors lowered. (It converts to memref, but now I need to expand the memrefs).

// Tests correctness of polynomial addition lowering.
//
// RUN: heir-opt --lower-poly %s > %t
// RUN: FileCheck %s < %t

#cycl_2048 = #poly.polynomial<1 + x**1024>
#ring_prime = #poly.ring<cmod=4294967291, ideal=#cycl_2048>

// RUN: heir-opt --lower-poly %s | \
// RUN:   mlir-opt -pass-pipeline="builtin.module( \
// RUN:     affine-expand-index-ops, \
// RUN:     lower-affine, \
// RUN:     finalize-memref-to-llvm, \
// RUN:     func.func( \
// RUN:       convert-scf-to-cf, \
// RUN:       affine-expand-index-ops, \
// RUN:       convert-arith-to-llvm, \
// RUN:       lower-affine), \
// RUN:     convert-func-to-llvm, \
// RUN:     convert-cf-to-llvm, \
// RUN:     reconcile-unrealized-casts)" | \
// RUN:   mlir-cpu-runner -e test_lower_poly_add -entry-point-result=void \
// RUN:      --shared-libs="%mlir_lib_dir/libmlir_c_runner_utils%shlibext,%mlir_runner_utils" | \
// RUN:   FileCheck %s --allow-empty

// CHECK-NOT: MISMATCH

// CHECK-LABEL: module
module {
  // Test helpers
  llvm.mlir.global internal constant @str_fail("MISMATCH\0A")
  func.func private @printCString(!llvm.ptr<i8>) -> ()
  // Prints 'MISMATCH' to stdout.
  func.func @printMismatch() -> () {
    %0 = llvm.mlir.addressof @str_fail : !llvm.ptr<array<9 x i8>>
    %1 = llvm.mlir.constant(0 : index) : i64
    %2 = llvm.getelementptr %0[%1, %1]
      : (!llvm.ptr<array<9 x i8>>, i64, i64) -> !llvm.ptr<i8>
    func.call @printCString(%2) : (!llvm.ptr<i8>) -> ()
    return
  }

  func.func @test_lower_poly_add() -> () {
    // CHECK-NOT: poly.add
    %coeffs2 = arith.constant dense<2> : tensor<1024xi64>
    %coeffs3 = arith.constant dense<3> : tensor<1024xi64>

    // CHECK: arith.constant dense<5>
    %poly0 = poly.from_tensor %coeffs2 : tensor<1024xi64> -> !poly.poly<#ring_prime>
    %poly1 = poly.from_tensor %coeffs3 : tensor<1024xi64> -> !poly.poly<#ring_prime>
    %poly2 = poly.add(%poly0, %poly1) {ring = #ring_prime} : !poly.poly<#ring_prime>
    %tensor2 = poly.to_tensor %poly2 : !poly.poly<#ring_prime> -> tensor<1024xi64>
    %coeffs5 = arith.constant dense<5> : tensor<1024xi64>

    // check tensor
    %mismatches = arith.cmpi ne, %tensor2, %coeffs5 : tensor<1024xi64>
    %c0 = arith.constant 0 : index
    %size = tensor.dim %mismatches, %c0 : tensor<1024xi1>
    affine.for %i = 0 to %size {
        %mismatch = tensor.extract %mismatches[%i] : tensor<1024xi1>
        scf.if %mismatch -> () {
            func.call @printMismatch() : () -> ()
        }
    }

    // CHECK: return
    return
  }
}

Generealize SameOperandsAndResultRings to work for multiple dialects

SameOperandsAndResultRings is useful for BFV/CKKS as well, and so we should generalize the trait. My idea is as follows:

Add a type (or op?) interface that says "this type/op has rings associated with it"
Make it somehow easy for a dialect to specialize this to its particular types/ops
Rewrite SameOperandsAndResultRings in terms of the interface and put it in include/Dialect/CommonTraits.td.

feat: dialect to represent combinational circuits

Some programs may benefit from translation into a combinational circuit with logic gates that can then be optimized by tools (like Yosys).

These programs can eventually target FHE schemes that implement boolean logic / LUTs (such as TFHE).

In order to integrate combinational lgoic optimizations, we need a dialect or IR that can represent these programs, and then create passes that may invoke programs like ABC through Yosys to optimize gates.

The ops involved would be LUTs (perhaps with rewrite patterns that would be able to optimize these LUTs (related: google/fully-homomorphic-encryption#62 @johannesmono) or fix the LUTs to certain sizes) and boolean gates.

@j2kun

[tfhe-rs] Add tfhe-rs dialect

This is an exit dialect for the tfhe-rs rust API. Only planning to support the shortint API

bgv: Add verifiers for BGV ops

It would be nice to add verifiers to BGV ops that can perform functionality like:

Verify that a multiplication operation's inputs have the same number of polys, and its output's number of polys is increased.
Verify that addition operations' inputs have the same number of polys
Verify that a relinearize inputs and outputs are valid based on the from_basis and to_basis.
Verify that a modulus switch's from_level and to_level are decrementing and also valid indices in the #rings of the input ciphertext.

I'm not sure whether these exist better and verifiers on the types, or belong to a single high level BGV validation pass.

feat: create a pass to wrap a program with the Secret dialect

Inputs: high-level programs, for example, in built-in dialects or other supported dialects, with some marking or indication of encrypted or secret values.

Outputs: A program that converts encrypted data to secret and propogates the operations using secret data to secret.generic ops.

feat: Add a CGGI dialect

poly: add function materialization patterns to lower-poly

I am planning to do this, just wanted to split it off from the initial #94 in light of https://reviews.llvm.org/D158107

bgv: Add a plaintext type and plaintext-ciphertext ops

We need to add a bgv plaintext type:

def Plaintext : BGV_Type<"Plaintext", "plaintext"> {
  let summary = "A type for BGV plaintext";

  let description = [{
    A type for BGV plaintext.
  }];

  let parameters = (ins
    BGVRingAttrArray:$rings,
  );

  let assemblyFormat = "`<` `rings` = $rings `>`";
}

That also likely needs to hold on to plaintext bits. (see #99)

Which would enable representation of muls, or other ops that can occur between ciphertext and plaintext:

def BGV_MulPlaintextOp : BGV_Op<"mul_plaintext"> {
  let summary = "Multiplication by a plaintext element.";

  let arguments = (ins
    Ciphertext:$x,
    Plaintext:$y
  );

  let results = (outs
    Ciphertext:$output
  );
}

This can even be an input to a trivial encryption.

bgv: Add a attribute to ciphertext to track plaintext bits

In the last HEIR meeting, we noted that rather than holding onto the underlying encrypted type, that it would be more useful to keep track of the number of plaintext bits that are encoded (potentially in each slot) of a BGV ciphertext. This would make the error analysis better. Likewise, at the BGV layer, it really is no longer relevant whether something was SIMD or not - this can occur at a higher level pass.

It may also be useful to have some notions of which slots are being used? Although not sure. That can be a separate issue.

poly: sometimes PolyToStandard pass crashes due to APInts of different widths being compared

Sometimes (nondeterministically) I'll get a crash when running --poly-to-standard on lower_mul.mlir. Stack trace below. It appears to be happening when two RingAttrs are constructed with cmod's with different bit widths, and the comparison test in the AttributeUniquer complains because two APInt's with different bit widths cannot be compared. Maybe changing the ringattr to take an element type container for the coefficients would fix this, since the APInt can always have that same container type, and two RingAttrs' cmods wouldn't be compared unless the first test for equality of the container type passes.

# | Stack dump:
# | 0.  Program arguments: heir-opt --poly-to-standard /home/j2kun/.cache/bazel/_bazel_j2kun/b07a48cf9da8c73a462d3608e5c508dd/sandbox/linux-sandbox/9014/execroot/heir/bazel-out/k8-dbg/bin/tests/poly/lower_mul.mlir.test.runfiles/heir/tests/poly/lower_mul.mlir
# |  #0 0x000055c0c1fc4dd8 llvm::sys::PrintStackTrace(llvm::raw_ostream&, int) /proc/self/cwd/external/llvm-project/llvm/lib/Support/Unix/Signals.inc:723:22
# |  #1 0x000055c0c1fc51b5 PrintStackTraceSignalHandler(void*) /proc/self/cwd/external/llvm-project/llvm/lib/Support/Unix/Signals.inc:798:1
# |  #2 0x000055c0c1fc26b7 llvm::sys::RunSignalHandlers() /proc/self/cwd/external/llvm-project/llvm/lib/Support/Signals.cpp:105:20
# |  #3 0x000055c0c1fc46b4 SignalHandler(int) /proc/self/cwd/external/llvm-project/llvm/lib/Support/Unix/Signals.inc:413:1
# |  #4 0x00007f2123591520 (/lib/x86_64-linux-gnu/libc.so.6+0x42520)
# |  #5 0x00007f21235e5a7c __pthread_kill_implementation ./nptl/pthread_kill.c:44:76
# |  #6 0x00007f21235e5a7c __pthread_kill_internal ./nptl/pthread_kill.c:78:10
# |  #7 0x00007f21235e5a7c pthread_kill ./nptl/pthread_kill.c:89:10
# |  #8 0x00007f2123591476 gsignal ./signal/../sysdeps/posix/raise.c:27:6
# |  #9 0x00007f21235777f3 abort ./stdlib/abort.c:81:7
# | #10 0x00007f212357771b _nl_load_domain ./intl/loadmsgcat.c:1177:9
# | #11 0x00007f2123588e96 (/lib/x86_64-linux-gnu/libc.so.6+0x39e96)
# | #12 0x000055c0bb2e7d04 llvm::APInt::operator==(llvm::APInt const&) const /proc/self/cwd/external/llvm-project/llvm/include/llvm/ADT/APInt.h:1030:21
# | #13 0x000055c0bb38a9e8 mlir::heir::poly::detail::RingAttrStorage::operator==(std::tuple<llvm::APInt, mlir::heir::poly::Polynomial> const&) const /proc/self/cwd/bazel-out/k8-dbg/bin/include/Dialect/Poly/IR/PolyAttributes.cpp.inc:103:45
# | #14 0x000055c0bb39250c mlir::heir::poly::detail::RingAttrStorage* mlir::StorageUniquer::get<mlir::heir::poly::detail::RingAttrStorage, llvm::APInt&, mlir::heir::poly::Polynomial&>(llvm::function_ref<void (mlir::heir::poly::detail::RingAttrStorage*)>, mlir::TypeID, llvm::APInt&, mlir::heir::poly::Polynomial&)::'lambda'(mlir::StorageUniquer::BaseStorage const*)::operator()(mlir::StorageUniquer::BaseStorage const*) const /proc/self/cwd/external/llvm-project/mlir/include/mlir/Support/StorageUniquer.h:205:20
# | #15 0x000055c0bb39702d bool llvm::function_ref<bool (mlir::StorageUniquer::BaseStorage const*)>::callback_fn<mlir::heir::poly::detail::RingAttrStorage* mlir::StorageUniquer::get<mlir::heir::poly::detail::RingAttrStorage, llvm::APInt&, mlir::heir::poly::Polynomial&>(llvm::function_ref<void (mlir::heir::poly::detail::RingAttrStorage*)>, mlir::TypeID, llvm::APInt&, mlir::heir::poly::Polynomial&)::'lambda'(mlir::StorageUniquer::BaseStorage const*)>(long, mlir::StorageUniquer::BaseStorage const*) /proc/self/cwd/external/llvm-project/llvm/include/llvm/ADT/STLFunctionalExtras.h:47:3
# | #16 0x000055c0c1ebf16f llvm::function_ref<bool (mlir::StorageUniquer::BaseStorage const*)>::operator()(mlir::StorageUniquer::BaseStorage const*) const /proc/self/cwd/external/llvm-project/llvm/include/llvm/ADT/STLFunctionalExtras.h:69:3
# | #17 0x000055c0c1ebe6c3 (anonymous namespace)::ParametricStorageUniquer::StorageKeyInfo::isEqual((anonymous namespace)::ParametricStorageUniquer::LookupKey const&, (anonymous namespace)::ParametricStorageUniquer::HashedStorage const&) /proc/self/cwd/external/llvm-project/mlir/lib/Support/StorageUniquer.cpp:73:37
# | #18 0x000055c0c1ec25dd bool llvm::DenseMapBase<llvm::DenseMap<(anonymous namespace)::ParametricStorageUniquer::HashedStorage, llvm::detail::DenseSetEmpty, (anonymous namespace)::ParametricStorageUniquer::StorageKeyInfo, llvm::detail::DenseSetPair<(anonymous namespace)::ParametricStorageUniquer::HashedStorage>>, (anonymous namespace)::ParametricStorageUniquer::HashedStorage, llvm::detail::DenseSetEmpty, (anonymous namespace)::ParametricStorageUniquer::StorageKeyInfo, llvm::detail::DenseSetPair<(anonymous namespace)::ParametricStorageUniquer::HashedStorage>>::LookupBucketFor<(anonymous namespace)::ParametricStorageUniquer::LookupKey>((anonymous namespace)::ParametricStorageUniquer::LookupKey const&, llvm::detail::DenseSetPair<(anonymous namespace)::ParametricStorageUniquer::HashedStorage> const*&) const /proc/self/cwd/external/llvm-project/llvm/include/llvm/ADT/DenseMap.h:658:11
# | #19 0x000055c0c1ec0ce9 bool llvm::DenseMapBase<llvm::DenseMap<(anonymous namespace)::ParametricStorageUniquer::HashedStorage, llvm::detail::DenseSetEmpty, (anonymous namespace)::ParametricStorageUniquer::StorageKeyInfo, llvm::detail::DenseSetPair<(anonymous namespace)::ParametricStorageUniquer::HashedStorage>>, (anonymous namespace)::ParametricStorageUniquer::HashedStorage, llvm::detail::DenseSetEmpty, (anonymous namespace)::ParametricStorageUniquer::StorageKeyInfo, llvm::detail::DenseSetPair<(anonymous namespace)::ParametricStorageUniquer::HashedStorage>>::LookupBucketFor<(anonymous namespace)::ParametricStorageUniquer::LookupKey>((anonymous namespace)::ParametricStorageUniquer::LookupKey const&, llvm::detail::DenseSetPair<(anonymous namespace)::ParametricStorageUniquer::HashedStorage>*&) /proc/self/cwd/external/llvm-project/llvm/include/llvm/ADT/DenseMap.h:689:24
# | #20 0x000055c0c1ebff31 llvm::DenseMapIterator<(anonymous namespace)::ParametricStorageUniquer::HashedStorage, llvm::detail::DenseSetEmpty, (anonymous namespace)::ParametricStorageUniquer::StorageKeyInfo, llvm::detail::DenseSetPair<(anonymous namespace)::ParametricStorageUniquer::HashedStorage>, false> llvm::DenseMapBase<llvm::DenseMap<(anonymous namespace)::ParametricStorageUniquer::HashedStorage, llvm::detail::DenseSetEmpty, (anonymous namespace)::ParametricStorageUniquer::StorageKeyInfo, llvm::detail::DenseSetPair<(anonymous namespace)::ParametricStorageUniquer::HashedStorage>>, (anonymous namespace)::ParametricStorageUniquer::HashedStorage, llvm::detail::DenseSetEmpty, (anonymous namespace)::ParametricStorageUniquer::StorageKeyInfo, llvm::detail::DenseSetPair<(anonymous namespace)::ParametricStorageUniquer::HashedStorage>>::find_as<(anonymous namespace)::ParametricStorageUniquer::LookupKey>((anonymous namespace)::ParametricStorageUniquer::LookupKey const&) /proc/self/cwd/external/llvm-project/llvm/include/llvm/ADT/DenseMap.h:182:5
# | #21 0x000055c0c1ebf2cd llvm::detail::DenseSetImpl<(anonymous namespace)::ParametricStorageUniquer::HashedStorage, llvm::DenseMap<(anonymous namespace)::ParametricStorageUniquer::HashedStorage, llvm::detail::DenseSetEmpty, (anonymous namespace)::ParametricStorageUniquer::StorageKeyInfo, llvm::detail::DenseSetPair<(anonymous namespace)::ParametricStorageUniquer::HashedStorage>>, (anonymous namespace)::ParametricStorageUniquer::StorageKeyInfo>::Iterator llvm::detail::DenseSetImpl<(anonymous namespace)::ParametricStorageUniquer::HashedStorage, llvm::DenseMap<(anonymous namespace)::ParametricStorageUniquer::HashedStorage, llvm::detail::DenseSetEmpty, (anonymous namespace)::ParametricStorageUniquer::StorageKeyInfo, llvm::detail::DenseSetPair<(anonymous namespace)::ParametricStorageUniquer::HashedStorage>>, (anonymous namespace)::ParametricStorageUniquer::StorageKeyInfo>::find_as<(anonymous namespace)::ParametricStorageUniquer::LookupKey>((anonymous namespace)::ParametricStorageUniquer::LookupKey const&) /proc/self/cwd/external/llvm-project/llvm/include/llvm/ADT/DenseSet.h:196:35
# | #22 0x000055c0c1ebe8b5 (anonymous namespace)::ParametricStorageUniquer::getOrCreate(bool, unsigned int, llvm::function_ref<bool (mlir::StorageUniquer::BaseStorage const*)>, llvm::function_ref<mlir::StorageUniquer::BaseStorage* ()>) /proc/self/cwd/external/llvm-project/mlir/lib/Support/StorageUniquer.cpp:151:40
# | #23 0x000055c0c1ebec04 mlir::detail::StorageUniquerImpl::getOrCreate(mlir::TypeID, unsigned int, llvm::function_ref<bool (mlir::StorageUniquer::BaseStorage const*)>, llvm::function_ref<mlir::StorageUniquer::BaseStorage* (mlir::StorageUniquer::StorageAllocator&)>) /proc/self/cwd/external/llvm-project/mlir/lib/Support/StorageUniquer.cpp:270:38
# | #24 0x000055c0c1eba39f mlir::StorageUniquer::getParametricStorageTypeImpl(mlir::TypeID, unsigned int, llvm::function_ref<bool (mlir::StorageUniquer::BaseStorage const*)>, llvm::function_ref<mlir::StorageUniquer::BaseStorage* (mlir::StorageUniquer::StorageAllocator&)>) /proc/self/cwd/external/llvm-project/mlir/lib/Support/StorageUniquer.cpp:375:27
# | #25 0x000055c0bb39267d mlir::heir::poly::detail::RingAttrStorage* mlir::StorageUniquer::get<mlir::heir::poly::detail::RingAttrStorage, llvm::APInt&, mlir::heir::poly::Polynomial&>(llvm::function_ref<void (mlir::heir::poly::detail::RingAttrStorage*)>, mlir::TypeID, llvm::APInt&, mlir::heir::poly::Polynomial&) /proc/self/cwd/external/llvm-project/mlir/include/mlir/Support/StorageUniquer.h:219:37
# | #26 0x000055c0bb38f5e2 _ZN4mlir6detail16AttributeUniquer13getWithTypeIDINS_4heir4poly8RingAttrEJRN4llvm5APIntERNS4_10PolynomialEEEENSt9enable_ifIXntsrSt7is_sameINT_8ImplTypeENS_16AttributeStorageEE5valueESD_E4typeEPNS_11MLIRContextENS_6TypeIDEDpOT0_ /proc/self/cwd/external/llvm-project/mlir/include/mlir/IR/AttributeSupport.h:227:64
# | #27 0x000055c0bb38d3bc mlir::heir::poly::RingAttr mlir::detail::AttributeUniquer::get<mlir::heir::poly::RingAttr, llvm::APInt&, mlir::heir::poly::Polynomial&>(mlir::MLIRContext*, llvm::APInt&, mlir::heir::poly::Polynomial&) /proc/self/cwd/external/llvm-project/mlir/include/mlir/IR/AttributeSupport.h:210:3
# | #28 0x000055c0bb38b6fc mlir::heir::poly::RingAttr mlir::detail::StorageUserBase<mlir::heir::poly::RingAttr, mlir::Attribute, mlir::heir::poly::detail::RingAttrStorage, mlir::detail::AttributeUniquer>::get<llvm::APInt&, mlir::heir::poly::Polynomial&>(mlir::MLIRContext*, llvm::APInt&, mlir::heir::poly::Polynomial&) /proc/self/cwd/external/llvm-project/mlir/include/mlir/IR/StorageUniquerSupport.h:183:3
# | #29 0x000055c0bb378114 mlir::heir::poly::RingAttr::get(llvm::APInt, mlir::heir::poly::Polynomial) /proc/self/cwd/bazel-out/k8-dbg/bin/include/Dialect/Poly/IR/PolyAttributes.cpp.inc:122:1
# | #30 0x000055c0bb2abaa4 mlir::heir::poly::PolyToStandard::buildPolyModFunc(mlir::FunctionType, mlir::heir::poly::RingAttr) /proc/self/cwd/lib/Conversion/PolyToStandard/PolyToStandard.cpp:607:31
# | #31 0x000055c0bb2ab24a mlir::heir::poly::PolyToStandard::generateOpImplementations()::'lambda'(mlir::Operation*)::operator()(mlir::Operation*) const::'lambda'(mlir::heir::poly::MulOp)::operator()(mlir::heir::poly::MulOp) const /proc/self/cwd/lib/Conversion/PolyToStandard/PolyToStandard.cpp:521:46

poly: investigate ConvertToLLVMPatternInterface

From review on

Reading the MLIR newsletter made me aware of ConvertToLLVMPatternInterface and I'd recommend we add it to the poly dialect :)

As I understand it, this interface is already implemented for math, arith, func, complex, cf, index, and memref, so I'm not entirely sure what benefit we'd get from implementing it vs lowering to those dialects. But it's worth learning more about this interface to see how appropriate it is for poly.

CI: try configuring include-what-you-use

https://github.com/pocc/pre-commit-hooks has a pre-commit hook we could try.

Also worth checking out this project which integrates directly with bazel https://github.com/storypku/bazel_iwyu

poly: add poly.constant folding and materialization

From #94 (comment)

we should add a poly.constant op to directly convert a #poly.polynomial to a !poly.poly without the need to go via tensors

The idea is to have an op that is analogous to arith.constant, which, by the use of an array attribute, can populate a tensor of constants, not just a single value, as in

%coeffs1 = "arith.constant"() <{value = dense<[2, 2, 5]> : tensor<3xi32>}> : () -> tensor<3xi32>

We could have something similar like

#cycl_2048 = #poly.polynomial<1 + x**1024>
#ring = #poly.ring<cmod=4294967296, ideal=#cycl_2048>
%poly1 = "poly.constant"() <{coefficients = dense<[2, 2, 5]> : tensor<3xi32>}> : !poly.poly<#ring> : () -> !poly.poly<#ring>

The attributes are already implemented, and so we would just need the op and lowerings.

website: reorganize dialect docs

The dialect docs are organized in kind of a haphazard way, we should reorganize it by dialect

feat: Integrate tools to optimize circuits

Related: #110

We can take a high-level arithmetic program and create optimization passes that can invoke tools like Yosys and ABC to optimize and output a combinational circuit.

The combinational circuit can then be lowered to boolean FHE schemes.

Configure a custom LSP

Analogous to

https://github.com/MarbleHE/HECO/blob/7cc9d14ca93e94303357dd8306dbd1ace79b3bab/src/tools/mlir-lsp-server/mlir-lsp-server.cpp#L3

website: Pick a custom domain name for the website

heir.dev is taken, so I browsed a list https://domains.google.com/registrar/search?searchTerm=heir.dev&hl=en&tab=1

Some fun ones:

heir.day
heir.band
heir.software
heir.tools
heir.style

Issues with Bazel & ccache

If, like me, you're using ccache to significantly speed up clean rebuilds and other edge cases, you'll probably run into this problem: https://stackoverflow.com/questions/52370202/how-to-get-bazel-ccache-and-sandboxing-to-work-together-ccache-read-only-file

For me, in fact this, required also adding another directory, which you can get from either the error messages you'll, or by running ccache -p and looking what is set as temporary_dir.

For me, I then also had issues with bazel crashing because $HOME failed it's internal "must be an absolute path" which made bazel throw an exception and dump the stack 🤪

This is what my .bazelrc ended up looking like

common --sandbox_writable_path=/run/user/1000/ccache-tmp --sandbox_writable_path=/home/vianda/.cache

chore: Rename poly dialect to polynomial dialect

per feedback!

secret: add attribute designating which entity a secret datum is known to

I think it would make sense to have an opaque integer attribute on secret.secret to identify an "entity" (whether that's a person or a secret key) for which the secret is known.

The idea is that we'd have an optimization pass that determines what part of a program is server vs client, and decide where certain ops should live to minimize total latency. I.e., do some more preprocessing on the client side, or maybe even have the client encrypt additional data to send up to the server, compared to the cost of bandwidth, server side processing, etc.

Compile information about the different kinds of encodings used across FHE schemes

types: simplify assembly formats by using type constraints

Most of our ops do not utilize Constrains which can be used to verify operations.

For example, it may enforce that for BGV mul ops, that their ring parameter must be the same across inputs and outputs, or that the inputs size must be the same, while their output size is predicatable one greater.

These also allow simplification of the assembly format.

poly: add a verifier for `poly.constant`

See https://github.com/google/heir/pull/188/files/359edd3c30ab619d4bb5f74e001cc66e8ebaf892#r1348820744

The coefficients should be small enough for the output type's cmod, and the degree should be small enough for the ideal

Add pre-commit hook configuration to enforce style guide

https://pre-commit.com/hooks.html

e.g., under hooks it supports stuff like bazel-buildifier, clang-format, etc.

[lwe]: does inverse canonical embedding need non-power-of-two scaling factors?

See TODO in LWEAttributes.td

poly: write verifiers to enforce coefficient sizes below coefficient modulus

cc @j2kun

The coefficient modulus may be much smaller than the APInt size that contains it. Because of this, an integer may validly be represented as a value much larger than cmod.

However, if we can verify that all poly's that are constructed are constructed with coeffs smaller than cmod, and all ops between poly's keep this invariant, than ops can use certain optimizations (e.g. barret reduction, or select and subtract) rather handling more remainder operations.

website: add mathjax support

In some docs we have a need for LaTeX, see https://google.github.io/heir/docs/dialects/polyattributes/

bgv: add lowering to poly

process: Improve process around PR merges

Right now, it is challenging for an external contributor to note when a PR should get merged. I think we should do the following changes:

To enforce external review:
(1) Require an external LGTM for each PR
(2) Enable a process that indicates whether the mirrored internal PR has that external LGTM, and only merge if that condition is satisfied

JAX does something similar, where once a PR gets an LGTM from GitHub, they update the reviewers to something like jax-mergeready.

To help external contributors to understand the process, we should document the merge / contribution process.

This can also be used to solve #107; we can gate the -mergeready review on having the commits squashed.

"Optional" bison/flex dependency potentially not optional?

I recently had to set up my environment again, and I realized that building HEIR (current main) with bazel build @heir//... as described in the getting started document fails with bison/flex not found if either of them isn't installed.

However, the documentation specifies that they are optional: https://github.com/google/heir/blob/a59ec085b3aea604348a33f1b8e2f9343ba7141e/docs/content/en/docs/getting_started.md?plain=1#L16C1-L18

chore: Make the Yosys dependency optional

We had some initial prototypes involving yosys for optimization, and this adds flex and bison as user-facing dependencies since there's a yosys CLI invocation in some tests.

If we migrate to using Yosys as a library we can avoid those deps, and in the mean time we can make the yosys dep optional in the workspace file, so lower the barrier to entry.

[poly]: Use a default builder to make the encoding parameter optional

Right now poly.poly has a custom builder to get the defaults working properly, but tablegen is supposed to have native support for default valued attributes. I couldn't get it to work. See PolyTypes.td

[poly]: add poly.ntt/intt ops

This would require converting a polynomial.polynomial to an appropriate tensor, and using the encoding attribute of tensor to store the information about which polynomial ring it came from.

secret: Add basic folders and verifiers for secret ops

See TODOs in code linking to this issue created in #78

[poly] Add more robust end-to-end tests for the lowering of poly_mul to standard

Tests to add:

https://github.com/google/heir/pull/203/files/2b456dc526c2c5a3b95fdfc4777c61b6cc582d24#r1378833393

Flip usePropertiesForAttributes to 1 after migrating

See #124 (comment)

poly: extract_slice `InferTypeOpAdaptor` is wrong

error: 'poly.extract_slice' op inferred type(s) 'tensor<1024xui64, #poly.ring<cmod=2837465, ideal=#poly.polynomial<1 + x**1024>>>' are incompatible with return type(s) of operation 'tensor<2xi32>'
    %slice = poly.extract_slice(%p0, %c0, %c2) : (!poly.poly<#ring1>, index, index) -> tensor<2xi32>

It uses poly.ring.ideal.degree instead of difference of indexes to deduce shape.

Code: https://github.com/google/heir/blob/main/lib/Dialect/Poly/IR/PolyOps.cpp#L7

poly: add a test for lower_mul with a divisor that has a non-1 leading coefficient

The initial lowering for poly.mul involves computing a remainder of the product modulo the ring's polynomial modulus. This in turn computes an inverse of the leading coefficient of the divisor. E.g., for a poly mod of 3x^1024 + 1, this would be the inverse of 3 modulo the coefficient modulus of the ring.

There are two more cases to test: that the division uses the correct inverse when an inverse exists (e.g., 3 mod 1024), and that the lowering fails with a sensible error message when the leading coefficient has no inverse (e.g., 2 mod 1024).

poly: make the dense elements attr more robust

In ConvertMul::matchAndRewrite there are a few places in which we create a arith.constant op using a dense elements attr. The number type for this op has caused me some pain, in that if you get it wrong by having a tensor element type that disagrees with the value type passed in, you get a nondescript assertion failure that is hard to debug.

The current code kind of punts because the tests are all 64-bit, and so it passes in something like 0L and it fits. But in general we'd want to inspect the types and set up the dense elements attr to be large enough to fit. For example, if the two input polynomials have coefficients that fit in i64, we'd need i128, if they're i16 we could use i32. Setting up tests that simply lower poly.mul with arguments using such coefficients should demonstrate the problem.

poly: lower poly.add and poly.mul

Initial poly lowering pass shell added in #94

poly: add formal proof documentation to poly.add lowering

PolyToStandard: handling tensors of poly?

@mr0re1 brought up a great review comment about lowering tensors of poly's to standard. https://github.com/google/heir/pull/134/files#r1312057016

Tensors of polys arise from BGV lowerings (and likely all other scheme lowerings to poly). When lowering to standard, we convert poly's to a tensor of ints.

Tensors of tensors are disallowed in MLIR - they are invalid tensor element types.

We have a few options here:

(1) Flatten the tensor: e.g. if we have a tensor polynomials, each of which are lowered to tensor<1024xi64>, then a size 2 tensor would become tensor<2048xi64>. However, this would appear to severely complicate the lowering math.

(2) Detensorize/lower to linalg and then lower to standard? E.g. inputs to PolyToStandard would not include tensors.

(3) It feels like I'm missing something - surely type conversion with tensors has been done before?

Re-enable micro_speech test

Re-enable the tests of the ML workflow once HEIR is in a place to execute it without it taking ~hours to run.

[poly] FromTensorOp's argument values may be larger than coefficient modulus

In the case that we can't statically determine poly.from_tensors argument value, we would need to include an arith::remui (remainder modulo the coefficient modulus) to ensure that the poly values are smaller than the coefficient modulus of the ring. Poly lowerings rely on that.

FromTensorOp::verify (

heir/lib/Dialect/Poly/IR/PolyOps.cpp

Line 38 in ce10f51

if (inputBitWidth > cmodBitWidth) {

) ensures that the element type of the argument is less than or equal to the coefficient modulus, but doesn't address the case where both are the same bitwidth and the element type is larger.

Currently, ConvertFromTensor (

heir/lib/Conversion/PolyToStandard/PolyToStandard.cpp

Line 48 in ce10f51

struct ConvertFromTensor : public OpConversionPattern<FromTensorOp> {

) does not insert a remainder operation.

I think it makes sense to add, right?

poly: add an extra modulus for lowering mul with non-machine-word sized coefficients

See TODO linking to this issue in PolyToStandard.cpp

secret: distribute generic op through the IR

The idea here is that a front-end will, via #112, take an input program with secrecy annotations, and use that to wrap large regions of code inside secret.generic. Lowering secret.generic would be easier, however, if there were only a single op in each region.

This ticket is to write a pass that splits a secret.generic with a large body into secret.generic with atomic bodies. Usually this would be a single instruction followed by a yield. One major use case here is to split a secret.generic(loop) into a loop over secret ops. Then we would have a canonicalization pass that unwraps any secret.generic blocks that involve no secret input data. This would effectively provide a dataflow of the secret data through a program.

We could also have an attribute on a secret.generic that instructs the pass to treat some ops (by name) as atomic even if they have nested regions. This might be useful to avoid mangling something like a linalg.generic op.

Some special care will likely need to be taken for memory-impacting operations, e.g., memrefs of cleartext values would need to be upgraded to memrefs of secret values. This might reveal some obstacles.

chore: use PR labels to indicate final merge status

Because of internal requirements, we require that PRs be squashed into a single commit before being merged. Currently, this is implemented with a pr_style workflow that checks the number of commits. However, enforcing this during the review process makes it difficult to see changes after comments are addressed.

Instead, we may want to allow the author to keep multiple commits (and potentially disallow merge commits to make this easier), and simply require a squash at the end of the review process.

poly: determine if the poly_mod function name needs to include the input tensor type

See the TODO referring to this issue in PolyToStandard.cpp