Python Frontend: Use Numba Type Inference

[AlexanderViand-Intel]

Very much WIP,

This does a few things:

• Enable Numba type inference and uses the inferred types to decide the MLIR type that gets emitted
• Add an @mlir decorator that is a slightly modified version of numba's @intrinsic, with the main difference being that one can provide a pure python implementation and run it under python. This allows us to define, e.g., linalg.matmul and assign it np.matmul as its python implementation.
• Changes a few defaults/adds a few options to the pipeline here and there. Not really related to type inference, might pull out into another PR
• Adds dummy classes for MLR types (e.g., Secret[...]), in preparation for grabbing signature types from type annotations (rathter than from a string passed into the decorator, as is Numba's approach for ahead-of-time compilation)
• TODO: lots of cleanup, but most importantly we need to define MLIR-style integers and Tensor types (see also Decide on (Cleartext) Integer Semantics for HEIR's Python Frontend #1252)

PS: fun fact - in numba/numpy world, i4 is not a 4-bit integer (as in MLIR), but shorthand for the 4-byte type int32.

[AlexanderViand-Intel]

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}")

[AlexanderViand-Intel]

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).

[j2kun]

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.

[j2kun]

Part of #1105

[AlexanderViand-Intel]

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 😅

[AlexanderViand-Intel]

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

[j2kun]

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.

[j2kun]

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

[AlexanderViand-Intel]

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.

[AlexanderViand-Intel]

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!

[j2kun]

Amazing work!

[j2kun]

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.

[j2kun]

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.

[AlexanderViand-Intel]

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

[AlexanderViand-Intel]

@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)