Sorry for the lack of clean-up/updates on this one, I was feeling a bit under the weather.

The big ToDo, as discussed last week, to move from Numba-style inference to something that supports statically shaped tensors is still open, but I did do some other smaller cleanup

• You can now select between schemes (only, bgv and ckks so far, but it'd be interesting to see how much effort it would be to support the boolean pipeline (at least to the heir-translate output)).
• The emitter now supports floats (i.e., emits arith.addf instead of arith.addi)
• You can now use Python type annotations to specify types (including ranked tensors, though the sizes will be converted to "?")
• The frontend now respects Secret[...] annotations and no longer relies on --secretize. This also means it will only generate enc helpers for the secret arguments.

Here's a current example:

from heir_py import compile
from heir_py.mlir import *
from heir_py.types import *

@compile() # defaults to scheme="bgv"", backend="openfhe", debug=False
def foo(x : Secret[I16], y : I16):
  sum = x + y
  mul = x * y
  expression = sum * mul
  return expression

foo.setup() # runs keygen/etc
enc_x = foo.encrypt_x(7)
result_enc = foo.eval(enc_x, 8)
result = foo.decrypt_result(result_enc)
print(f"Expected result for `foo`: {foo(7,8)}, decrypted result: {result}")

As mentioned in yesterday's meeting, there's been a bit more progress, in addition to what I mentioned above:

• Lots of debugging improvements, including

  • Support for custom pipelines , making the frontend a convenient tool for debugging/testing HEIR itself

  @compile(heir_opt_options=["--mlir-to-bgv", "--canonicalize", "--cse"], backend=None, debug=True)

  • Debug mode now also creates --view-op-graph Graphviz/*.dot files so you can inspect pre/post HEIR IR visually
  • Python source locations are now correctly forwarded to MLIR, e.g., you'll now get
    /home/aviand/code/heir/heir_py/example.py:19:0: error: failed to legalize operation 'secret.generic' that was explicitly marked illegal , and most decent IDEs should let you click on the path:line:col to jump to it directly.

• Initial work on supporting different backends, but there's really only openfhe or none at the moment, with the latter simply meaning that no heir-translate happens and that no actual fhe eval and helper functions get generated.

Since this PR is getting a bit large already, I'm leaning towards cleaning this up (once #1302 is merged, allowing for full scheme/backend separation) and then making a new PR for the tensor shape inference stuff (which will probably be an MLIR pass after all).

Python source locations are now correctly forwarded to MLIR

❤️ !!

Feel free to clean/split it up as you see fit. I have no problems reviewing a large PR for this one.

Part of #1105

This is essentially ready to review, I just need some help getting the test setup working properly, as the e2e test currenlty fails if run via bazel.

Since the last update, major changes include

• Some heavy re-factoring/cleanup, including lots and lots of new checks and error messages (now in color 🌈!)
• You can now pip install the frontend (well, locally, as it's not yet on pypi)
• OpenFHE is now no longer special, and instead there are pluggable backends, which can do whatever heir-opt, heir-translate, clang/etc calls they want to do, as long as they wire up everything to the expected python client functions. The DummyBackend for use with custom heir_opt_options simply makes all of these no-ops/plain python.
  EDIT: I expect the design of the interfaces around this (espcially run_backend(...)) to improve once we actually start implementing "real" alternative backends, so I've left this pretty rough for now
• Probably some other stuff I forgot 😅

I have no problems reviewing a large PR for this one.

@j2kun : let me now if this is too large, and I'll figure out how to split it into more sensible chunks

The testing issue is because you put the python package one step further nested under frontend/heir, and bazel's working directory is the base directory of the project (rather, a hermetic clone of that directory inside bazel-bin/frontend/e2e_test.runfiles/heir/. So all the imports would need to be changed from import heir to import frontend.heir. Ideally a published python package would not have that more verbose import path. But we'd need to figure out how to make it nice for development and publishing.

Edit: see next comment What do you think about leaving the paths as more verbose now, and having a followup PR focus on the PyPI publishing part? I think what a lot of people do in this case (e.g., jax, tensorflow) is call a bunch of mv commands from bazel to make the python import paths work out. It would be nice though, if instead we could find a way to tell bazel to set its working directory to frontend/ for these tests.

Ohh! Researching this I found on https://bazel.build/reference/be/python

imports	
List of strings; default is []

List of import directories to be added to the PYTHONPATH. Subject to "Make variable" substitution. These import directories will be added for this rule and all rules that depend on it (note: not the rules this rule depends on. Each directory will be added to `PYTHONPATH` by `py_binary` rules that depend on this rule. The strings are repo-runfiles-root relative, Absolute paths (paths that start with `/`) and paths that references a path above the execution root are not allowed and will result in an error.

So with this patch I can get it to run in bazel, although the test fails

diff --git a/frontend/BUILD b/frontend/BUILD
index 190aee59..342bf248 100644
--- a/frontend/BUILD
+++ b/frontend/BUILD
@@ -34,6 +34,7 @@ py_library(
     ),
     data = DATA_DEPS,
     deps = [
+        "@frontend_pip_deps_colorama//:pkg",
         "@frontend_pip_deps_numba//:pkg",
         "@frontend_pip_deps_pybind11//:pkg",
         "@frontend_pip_deps_pybind11_global//:pkg",
diff --git a/frontend/testing.bzl b/frontend/testing.bzl
index 3d24a318..902f7dd6 100644
--- a/frontend/testing.bzl
+++ b/frontend/testing.bzl
@@ -29,6 +29,7 @@ def frontend_test(name, srcs, deps = [], data = [], tags = []):
             ":frontend",
             "@com_google_absl_py//absl/testing:absltest",
         ],
+        imports = ["."],
         data = data,
         tags = tags,
         env = {

Test failure: https://gist.github.com/j2kun/40fe73429938e814cf3b08beb514dd44

Awesome, thanks! I just realized the test failure is likely because my current approach still relies on a nasty local patch to numba to support sizes smaller than i64. I'll need to see if there's a way to replace that with something nice tomorrow, or if I do need to bite the bullet and re-implement the inference after all.

Now the bazel tests work locally but fail on the CI, probably because I messed something up with openfhe_config...

ValueError: Error running /usr/bin/clang. stderr was:

/tmp/tmpm1pni9oh/foo.cpp:2:10: fatal error: 'openfhe/pke/openfhe.h' file not found
    2 | #include "openfhe/pke/openfhe.h"  // from @openfhe
      |          ^~~~~~~~~~~~~~~~~~~~~~~
1 error generated.


options were:

/usr/bin/clang /tmp/tmpm1pni9oh/foo.cpp -o /tmp/tmpm1pni9oh/foo.so -O3 -fPIC -shared -std=c++17 -I/usr/local/include/openfhe -I/usr/local/include/openfhe/binfhe -I/usr/local/include/openfhe/core -I/usr/local/include/openfhe/pke -L/usr/local/lib -lOPENFHEbinfhe -lOPENFHEcore -lOPENFHEpke 

I'll look into it tomorrow!

Now the bazel tests work locally but fail on the CI, probably because I messed something up with openfhe_config...

ValueError: Error running /usr/bin/clang. stderr was:

/tmp/tmpm1pni9oh/foo.cpp:2:10: fatal error: 'openfhe/pke/openfhe.h' file not found
    2 | #include "openfhe/pke/openfhe.h"  // from @openfhe
      |          ^~~~~~~~~~~~~~~~~~~~~~~
1 error generated.


options were:

/usr/bin/clang /tmp/tmpm1pni9oh/foo.cpp -o /tmp/tmpm1pni9oh/foo.so -O3 -fPIC -shared -std=c++17 -I/usr/local/include/openfhe -I/usr/local/include/openfhe/binfhe -I/usr/local/include/openfhe/core -I/usr/local/include/openfhe/pke -L/usr/local/lib -lOPENFHEbinfhe -lOPENFHEcore -lOPENFHEpke 

I'll look into it tomorrow!

The command is using -I/usr/local/include/openfhe when it should be using a bazel-ified path. I think this is because frontend/heir/pipeline.py:compile uses DEFAULT_INSTALLED_OPENFHE_CONFIG as its kwarg default, rather than calling from_os_env(). I think we should always use from_os_env since it internally defaults to the DEFAULT_INSTALLED_OPENFHE_CONFIG when env vars are missing.

That said, I am trying to wrap your example.py in a py_binary and for whatever reason that doesn't set a RUNFILES_DIR env var... so I will have to improve that in future work.

That said, I am trying to wrap your example.py in a py_binary and for whatever reason that doesn't set a RUNFILES_DIR env var... so I will have to improve that in future work.

I was able to figure out why this is: bazel doesn't set the same environment variables for binary targets vs test targets, so I need to add another env var check for BINDIR. This should be easy but I will wait for this PR to be submitted.

Thanks for the detailed review and all the bazel-debugging! I'll try and get this pull-ready later today :)

FYI @AlexanderViand-Intel Asra's loop support is in review, so you may have a big rebase ahead of you.

FYI @AlexanderViand-Intel Asra's loop support is in review, so you may have a big rebase ahead of you.

Sorry for letting this sit for so long, as you've probably noticed I've not had a lot of bandwidth the last few weeks. I guess the best move is to wait for #1307 to be merged, then rebase and split this PR into two, one with the various QoL features and another one for the actual type inference?

FYI @AlexanderViand-Intel Asra's loop support is in review, so you may have a big rebase ahead of you.

Sorry for letting this sit for so long, as you've probably noticed I've not had a lot of bandwidth the last few weeks. I guess the best move is to wait for #1307 to be merged, then rebase and split this PR into two, one with the various QoL features and another one for the actual type inference?

SGTM. I'd give it maybe two weeks before we have to take it over and get it merged so we can have the issues ironed out before the conferences :)

I just checked, and there's actually not much to do for the rebase, just needed to replace a hardcoded (i64) with a type lookup. However, there is a more subtle issue, in that the iterator variable will have HW-bitwidth, which we probably don't want, since that might force a large ptxt modulus to be selected in the future. I guess we want to round up to the next power of two bitwidth that can fit range_end?

More generally, HEIR doesn't currently support doing things such as i8 + i16 (more specifically, I don't think our lowerings handle the arith.truncate/arith.extui this would require. I feel like this will be more and more relevant as we get code not from hand-written MLIR but autogenerated from frontend(s).

I've integrated the review feedback and removed all the matmul/intrinsic related things for now, as I need to dig deeper into numba for that.

I've also found a workaround for the issue mentioned in #1252 (numba will upcast, e.g., int8 + int8 to int32/int64). It's not very pretty, so there's a good chance some internal tool will not accept it (I recall some C++ linking hack also getting stopped by internal checks), but I guess loading a modified version of a module via importlib is technically "normal" python?

EDIT: Any idea why the test fails due to an openfhe header not being found? (The test works fine locally)
It's because I'm an idiot and either forgot to switch the default OpenFHE config to from_os_env or lost the change during a rebase.... Should be fixed now.

@j2kun: I think everything I was planning to do (and a lot of semi-related QoL things...feature creep I guess 😓) are in now. :)

EDIT: I just found one more annoying thing - numba assumes all literal ints are intp sized (int32 or int64) which means foo(x: I16): return x + 5 generates an invalid arith.addi with mismatched types. Looking into how to fix that now.
EDIT2: So I can easily use the same overriding trick to make numba just pick the smallest type in int8/16/32/64 that will fit the constant (specifically, adjusting @typeof_impl.register(int) def _typeof_int(val, c) in numba.core.typing.typeof).
However, then we might still get a mismatch in the example above, just between i8 and i16 instead of between i64 and i8. I think solving this might require either adding support for extui/truncate to HEIR or a very different type inference approach.
EDIT3: There was also an unrelated issue where the emitter didn't even use the inferred type for constant and used hardcoded i64, pushing the fix for this one already since it's just a straight-up bugfix. I'll revisit the second numba override and how to get valid mlir from i16 + i8 tomorrow.

There is one more "QoL" thing I'd love to tackle, and that's the inability to have multiple instances of the frontend/HEIR (the "PublicKey registerd" error), but I think I might need your assistance to tackle that one as it's less about the frontend and more about the bindings.

Next steps would be

• Enabling "intrinsics" (e.g., linalg.matmul) while preserving pure-python functionality,
• Tensor type/shape inference (either in Python or MLIR)