WIP: polygeist integration

tests/polygeist/BUILD

@@ -0,0 +1,13 @@
+load("//bazel:lit.bzl", "glob_lit_tests")
+
+package(
+    default_applicable_licenses = ["@heir//:license"],
+    default_visibility = ["//visibility:public"],
+)
+
+glob_lit_tests(
+    name = "all_tests",
+    data = ["@heir//tests:test_utilities"],
+    driver = "@heir//tests:run_lit.sh",
+    test_file_exts = ["mlir"],
+)

tests/polygeist/README.md

@@ -0,0 +1,13 @@
+# Polygeist
+
+In lieu of actually integrating Polygeist into HEIR (they have incompatibly
+pinned upstream LLVM hashes), run polygeist as a separate tool and check in both
+the C++ and output MLIR sources into this directory.
+
+```bash
+Polygeist/build/bin/cgeist \
+  '-function=*' \
+  -raise-scf-to-affine <... more passes ...> \
+  -S -O3 \
+  input_cc_file.cpp > test_file.mlir
+```

tests/polygeist/add_one.cpp

@@ -0,0 +1,6 @@
+#include <cstdint>
+
+uint8_t add_one(uint8_t x) {
+  uint8_t one = 1;
+  return x + one;
+}

tests/polygeist/add_one.mlir

@@ -0,0 +1,22 @@
+// RUN: heir-opt --secretize=entry-function=_Z7add_onec --wrap-generic --secret-distribute-generic %s
+
+module {
+  func.func @_Z7add_onec(%arg0: i8) -> i8 {
+    // Note: polygeist does this unnecessary extsi/trunci because it defines
+    // the constant as an i32. We may want to add a pass to narrow this.
+    // Note that `arith-int-narrowing` helps:
+    //
+    //    '--arith-int-narrowing=int-bitwidths-supported=1,2,4,8,16,32' --canonicalize
+    //
+    // results in the constant 1 being defined as an i16, presumably because
+    // adding 1 to an 8-bit integer can cause it to overflow, so it has to
+    // guarantee the same results of having an i32 addition truncated, and this
+    // probably must be fixed at the polygeist level by having the constant be
+    // defined as an i8. See https://github.com/llvm/Polygeist/issues/388
+    %c1_i32 = arith.constant 1 : i32
+    %0 = arith.extsi %arg0 : i8 to i32
+    %1 = arith.addi %0, %c1_i32 : i32
+    %2 = arith.trunci %1 : i32 to i8
+    return %2 : i8
+  }
+}

tests/polygeist/sqrt.cpp

@@ -0,0 +1,29 @@
+/*
+* Compile this file to mlir using
+*
+
+   Polygeist/build/bin/cgeist \
+    '-function=*' \
+    -raise-scf-to-affine \
+    --memref-fullrank \
+    -S \
+    -O3 \
+    sqrt.cc
+
+*/
+
+int isqrt(int num) {
+  int res = 0;
+  int bit = 1 << 14;  // ((unsigned) INT16_MAX + 1) / 2.
+
+  for (int i = 0; i < 8; ++i) {
+    if (num >= res + bit) {
+      num -= res + bit;
+      res = (res >> 1) + bit;
+    } else {
+      res >>= 1;
+    }
+    bit >>= 2;
+  }
+  return res;
+}

tests/polygeist/sqrt.mlir

@@ -0,0 +1,36 @@
+// RUN: heir-opt --secretize=entry-function=_Z5isqrti --wrap-generic --secret-distribute-generic %s
+
+
+module {
+  // Note: this should be i16, but polygeist failed to emit valid MLIR when the
+  // input is a short instead of an int see
+  // https://github.com/llvm/Polygeist/issues/387
+  //
+  // Also note that this could not be raised to affine because it is naturally
+  // iteration-dependent. This makes it a bad candidate for SIMD-style FHE,
+  // so we're mainly using it to exercise the compilation chain.
+  //
+  // The secret annotation is added manually, not by polygeist.
+  func.func @_Z5isqrti(%arg0: i32) -> i32 {
+    %c2_i32 = arith.constant 2 : i32
+    %c1_i32 = arith.constant 1 : i32
+    %c16384_i32 = arith.constant 16384 : i32
+    %c0_i32 = arith.constant 0 : i32
+    %0:3 = affine.for %arg1 = 0 to 8 iter_args(%arg2 = %c16384_i32, %arg3 = %c0_i32, %arg4 = %arg0) -> (i32, i32, i32) {
+      %1 = arith.addi %arg3, %arg2 : i32
+      %2 = arith.cmpi sge, %arg4, %1 : i32
+      %3:2 = scf.if %2 -> (i32, i32) {
+        %5 = arith.subi %arg4, %1 : i32
+        %6 = arith.shrsi %arg3, %c1_i32 : i32
+        %7 = arith.addi %6, %arg2 : i32
+        scf.yield %7, %5 : i32, i32
+      } else {
+        %5 = arith.shrsi %arg3, %c1_i32 : i32
+        scf.yield %5, %arg4 : i32, i32
+      }
+      %4 = arith.shrsi %arg2, %c2_i32 : i32
+      affine.yield %4, %3#0, %3#1 : i32, i32, i32
+    }
+    return %0#1 : i32
+  }
+}