Improve python frontend (incl. basic type inference)

WORKSPACE

@@ -43,13 +43,13 @@ new_git_repository(
 load("@rules_python//python:repositories.bzl", "python_register_toolchains")
 
 python_register_toolchains(
-    name = "python3_10",
+    name = "python3_11",
     # Available versions are listed at
     # https://github.com/bazelbuild/rules_python/blob/main/python/versions.bzl
-    python_version = "3.10",
+    python_version = "3.11",
 )
 
-load("@python3_10//:defs.bzl", "interpreter")
+load("@python3_11//:defs.bzl", "interpreter")
 load("@rules_python//python:pip.bzl", "pip_parse")
 
 # Download the Go rules.
@@ -231,14 +231,14 @@ load("@heir_pip_deps//:requirements.bzl", "install_deps")
 
 install_deps()
 
-# separate pip deps for heir_py
+# separate pip deps for frontend
 pip_parse(
-    name = "heir_py_pip_deps",
+    name = "frontend_pip_deps",
     python_interpreter_target = interpreter,
-    requirements_lock = "//heir_py:requirements.txt",
+    requirements_lock = "//frontend:requirements.txt",
 )
 
-load("@heir_py_pip_deps//:requirements.bzl", "install_deps")  # buildifier: disable=load
+load("@frontend_pip_deps//:requirements.bzl", "install_deps")  # buildifier: disable=load
 
 install_deps()
 

bazel/lit.bzl

@@ -56,6 +56,7 @@ def lit_test(name = None, src = None, size = "small", tags = None, data = None):
         data = data,
         srcs = ["@llvm-project//llvm:lit"],
         main = "lit.py",
+        deps = [Label("@llvm-project//llvm:lit")],
         tags = tags,
     )
 

heir_py/BUILD → frontend/BUILD

@@ -1,4 +1,4 @@
-load("@heir//heir_py:testing.bzl", "heir_py_test")
+load("@heir//frontend:testing.bzl", "frontend_test")
 load("@rules_python//python:py_library.bzl", "py_library")
 
 package(
@@ -25,31 +25,22 @@ DATA_DEPS = [
 # note we chose the style of putting all test rules below, because glob does
 # not recurse into subdirectories with BUILD files in them.
 py_library(
-    name = "heir_py",
+    name = "frontend",
     srcs = glob(
-        ["**/*.py"],
+        ["heir/**/*.py"],
         exclude = [
             "**/*_test.py",
         ],
     ),
     data = DATA_DEPS,
     deps = [
-        "@heir_py_pip_deps_numba//:pkg",
-        "@heir_py_pip_deps_pybind11//:pkg",
-        "@heir_py_pip_deps_pybind11_global//:pkg",
+        "@frontend_pip_deps_numba//:pkg",
+        "@frontend_pip_deps_pybind11//:pkg",
+        "@frontend_pip_deps_pybind11_global//:pkg",
     ],
 )
 
-heir_py_test(
-    name = "pipeline_test",
-    srcs = ["pipeline_test.py"],
-    tags = [
-        # copybara: manual
-        "notap",
-    ],
-)
-
-heir_py_test(
+frontend_test(
     name = "e2e_test",
     srcs = ["e2e_test.py"],
     tags = [

heir_py/README.md → frontend/README.md

@@ -1,10 +1,4 @@
-# heir_py Developer's Notes
-
-This Python frontend is currently _experimental_. Current **missing** features
-include:
-
-- Being able to `pip install` the local development tree
-- Being able to `pip install` from PyPI.
+# HEIR Frontend Developer's Notes
 
 ## Overview
 
@@ -23,32 +17,40 @@ C++ standard libraries available on you system.
 You should be able to run the frontend directly from python (without using
 bazel) by (from the HEIR project root):
 
-- `pip install -r heir_py/requirements.txt`
-- `python -m heir_py.example` or `python -i -m heir_py.pipeline`
+- `python frontend/example.py` This will use `heir_opt` and `heir_translate`
+  from your bazel-bin build folder, and assumes OpenFHE is installed.
+
+You can also install the frontend:
 
-This will attempt to discover everything needed on your system from the
-operating system.
+- `pip install -e frontend/example` In this case, `heir_opt` and
+  `heir_translate` should be on your PATH, or the appropriate environment
+  variables (see below) should be set.
 
 You will need the following system-level dependencies to do this:
 
-- `clang` (not sure what versions are supported, has been tested on `clang-17`)
-- For C++ backends like OpenFHE, `python3.11-dev` (or similar for your python
-  version) in order to build generated `pybind11` bindings.
-- C/c++ standard libraries that come with your system's compiler toolchain, but
-  may be in a nonstandard location. Ensure they can be discovered by a call to
-  `clang` without any special flags. (`clang -v` will show you which paths it
-  considers).
+- Python 3.11 or newer
+- For C++ backends like OpenFHE:
+  - `clang` (not sure what versions are supported, has been tested on
+    `clang-17`)
+  - `python3.11-dev` (or similar for your python version) in order to build
+    generated `pybind11` bindings.
+  - C/c++ standard libraries that come with your system's compiler toolchain,
+    but may be in a nonstandard location. Ensure they can be discovered by a
+    call to `clang` without any special flags. (`clang -v` will show you which
+    paths it considers).
 
-If this doesn't work, you may need to set up the following environment
-variables:
+## Environment Variables:
 
 - `HEIR_OPT_PATH` and `HEIR_TRANSLATE_PATH`: to the location of the `heir-opt`
   and `heir-translate` binaries on your system.
+
   - Uses `shutil.which` to find the binary on your path if not set.
   - Defaults to `bazel-bin/tools/heir-{opt,translate}`.
-  - Cf. `heir_py/heir_config.py` for more details.
+  - Cf. `heir/backends/openfhe/config.py` for more details.
+
 - OpenFHE installation locations (default to where `cmake` installs them in the
   OpenFHE development repo).
+
   - `OPENFHE_LIB_DIR`: a string containing the directory containing the OpenFHE
     .so files. Usually `/usr/local/lib`
   - `OPENFHE_INCLUDE_DIR`: a colon-separated string of directories containing
@@ -68,5 +70,5 @@ variables:
 ## Running from bazel
 
 `bazel test` should work out of the box. If it does not, file a bug.
-`heir_py/testing.bzl` contains the environment variable setup required to tell
-the frontend where to find OpenFHE and related backend shared libraries.
+`testing.bzl` contains the environment variable setup required to tell the
+frontend where to find OpenFHE and related backend shared libraries.

heir_py/e2e_test.py → frontend/e2e_test.py

@@ -1,20 +1,18 @@
-from heir_py import compile
+from heir import compile
+from heir.mlir import *
 
 from absl.testing import absltest  # fmt: skip
-
-
 class EndToEndTest(absltest.TestCase):
 
   def test_simple_arithmetic(self):
-    @compile(backend="openfhe")
-    def foo(a, b):
+    @compile() # defaults to BGV and OpenFHE
+    def foo(a : Secret[I16], b : Secret[I16]):
       return a * a - b * b
 
-    # Test the e2e path
+    # Test plaintext functionality
     self.assertEqual(-15, foo(7, 8))
 
-    # Also test the manual way with crypto_context and keys threaded in
-    # automatically
+    # Test FHE functionality
     foo.setup()
     enc_a = foo.encrypt_a(7)
     enc_b = foo.encrypt_b(8)

frontend/example.py

@@ -0,0 +1,101 @@
+"""Example of HEIR Python usage."""
+
+from heir import compile
+from heir.mlir import *
+
+# TODO (#1162): Also add the tensorflow-to-tosa-to-HEIR example in example.py, even it doesn't use the main Python frontend?
+
+### Simple Example
+@compile() # defaults to scheme="bgv", OpenFHE backend, and debug=False
+def foo(x : Secret[I16], y : Secret[I16]):
+  sum = x + y
+  diff = x - y
+  mul = x * y
+  expression = sum * diff + mul
+  deadcode = expression * mul
+  return expression
+
+foo.setup() # runs keygen/etc
+enc_x = foo.encrypt_x(7)
+enc_y = foo.encrypt_y(8)
+result_enc = foo.eval(enc_x, enc_y)
+result = foo.decrypt_result(result_enc)
+print(f"Expected result for `foo`: {foo(7,8)}, decrypted result: {result}")
+
+
+# TODO (#1162) : Fix "ImportError: generic_type: type "PublicKey" is already registered!" when doing setup twice. (Required to allow multiple compilations in same python file)
+### CKKS Example
+# @compile(scheme="ckks")
+# def bar(x : Secret[F32], y : Secret[F32]):
+#   sum = x + y
+#   diff = x - y
+#   mul = x * y
+#   expression = sum * diff + mul
+#   deadcode = expression * mul
+#   return expression
+
+# bar.setup() # runs keygen/etc
+# enc_x = bar.encrypt_x(7)
+# enc_y = bar.encrypt_y(8)
+# result_enc = bar.eval(enc_x, enc_y)
+# result = bar.decrypt_result(result_enc)
+# print(f"Expected result for `bar`: {bar(7,8)}, decrypted result: {result}")
+
+
+# ### Ciphertext-Plaintext Example
+# @compile()
+# def baz2(x: Secret[I16], y : Secret[I16], z : I16):
+#   ptxt_mul = x * z
+#   ctxt_mul = x * x
+#   ctxt_mul2 = y * y
+#   add = ctxt_mul + ctxt_mul2
+#   return  ptxt_mul + add
+
+# baz2.setup() # runs keygen/etc
+# enc_x = baz2.encrypt_x(7)
+# enc_y = baz2.encrypt_y(8)
+# result_enc = baz2.eval(enc_x, enc_y, 9)
+# result = baz2.decrypt_result(result_enc)
+# print(f"Expected result for `baz2`: {baz2(7,8,9)}, decrypted result: {result}")
+
+
+
+# ### Custom Pipeline Example
+# @compile(heir_opt_options=["--mlir-to-secret-arithmetic", "--canonicalize", "--cse"], backend=None, debug=True)
+# def foo(x : Secret[I16], y : Secret[I16]):
+#   sum = x + y
+#   diff = x - y
+#   mul = x * y
+#   expression = sum * diff + mul
+#   deadcode = expression * mul
+#   return expression
+
+# # The below are basically no-ops/plain python with the Dummy Backend
+# foo.setup()
+# enc_x = foo.encrypt_x(7)
+# enc_y = foo.encrypt_y(8)
+# result_enc = foo.eval(enc_x, enc_y)
+# result = foo.decrypt_result(result_enc)
+# print(f"Expected result for `foo`: {foo(7,8)}, decrypted result: {result}")
+
+
+# ### Matmul Example (WIP)
+# # #TODO (#1330): Implement Shape Inference and support, e.g., matmul in Frontend)
+# from heir.mlir.linalg import matmul
+# import numpy as np
+# @compile(scheme='ckks', debug=True)
+# def qux(a : Secret[Tensor[4,4,F32]], b : Secret[Tensor[4,4,F32]]):
+#  AB = matmul(a,b)
+#  AABB = matmul(a+a, b+b)
+#  return AB + AABB
+
+# a = np.array([[1,2],[3,4]])
+# b = np.array([[5,6],[7,8]])
+# print(qux(a,b))
+
+# qux.setup()
+# enc_a = qux.encrypt_a(a)
+# enc_b = qux.encrypt_b(b)
+# result_enc = qux.eval(enc_a, enc_b)
+# result = qux.decrypt_result(result_enc)
+# print(f"Expected result: {np.matmul(a,b)}, decrypted result: {result}")

frontend/heir/__init__.py

@@ -0,0 +1,3 @@
+from .pipeline import compile
+
+__all__ = ["compile"]

frontend/heir/backends/dummy.py

@@ -0,0 +1,52 @@
+"""Dummy Backend."""
+
+from colorama import Fore, Style, init
+
+from heir.core import BackendInterface, CompilationResult,  ClientInterface
+
+class DummyClientInterface(ClientInterface):
+
+  def __init__(self, compilation_result: CompilationResult):
+    self.compilation_result = compilation_result
+
+  def setup(self):
+    print("HEIR Warning (Dummy Backend): " + Fore.YELLOW + Style.BRIGHT + f"{self.compilation_result.func_name}.setup() is a no-op in the Dummy Backend")
+
+  def decrypt_result(self, result):
+    print("HEIR Warning (Dummy Backend): " + Fore.YELLOW + Style.BRIGHT + f"{self.compilation_result.func_name}.decrypt() is a no-op in the Dummy Backend")
+    return result
+
+  def __getattr__(self, key):
+
+    if key.startswith("encrypt_"):
+      arg_name = key[len("encrypt_") :]
+
+      def wrapper(arg):
+        print("HEIR Warning (Dummy Backend): " + Fore.YELLOW + Style.BRIGHT + f"{self.compilation_result.func_name}.{key}() is a no-op in the Dummy Backend")
+        return arg
+
+      return wrapper
+
+    if key == self.compilation_result.func_name or key == "eval":
+      print("HEIR Warning (Dummy Backend): " + Fore.YELLOW + Style.BRIGHT + f"{self.compilation_result.func_name}.eval() is the same as {self.compilation_result.func_name}() in the Dummy Backend.")
+      return self.func
+
+    raise AttributeError(f"Attribute {key} not found")
+
+
+class DummyBackend(BackendInterface):
+
+    def run_backend(self, workspace_dir, heir_opt, heir_translate, func_name, arg_names, secret_args, heir_opt_output, debug):
+
+        result = CompilationResult(
+              module=None,
+              func_name=func_name,
+              secret_args=secret_args,
+              arg_names=arg_names,
+              arg_enc_funcs=None,
+              result_dec_func=None,
+              main_func=None,
+              setup_funcs=None
+          )
+
+        return DummyClientInterface(result)

frontend/heir/backends/openfhe/__init__.py

@@ -0,0 +1,4 @@
+from .backend import OpenFHEBackend
+from .config import OpenFHEConfig, DEFAULT_INSTALLED_OPENFHE_CONFIG, from_os_env
+
+__all__ = ["OpenFHEBackend", "OpenFHEConfig", "DEFAULT_INSTALLED_OPENFHE_CONFIG", "from_os_env"]