poly: add initial lowering for NTT op and add root parameter to rings by inbelic · Pull Request #642 · google/heir

inbelic · 2024-04-24T04:43:23Z

- add a standard iterative implementation of the NTT Op
- add testing for lowering and a simple case runner
- add performance test that uses random coefficients to prevent
  optimizations on dense values that artifically improves
  compile/runtime

include an optional parameter that computes the 2n-th primitive
root of a rings ideal degree and its coefficient modulus
allow for this to be retrieved from a pre-computed or to be
user-specified
Add a script that will compute the 2n-th roots of unity for the
specified polynomial degree values and n. The values are only
outputted if the 2n-th root of unity exists.

Resolve #543

inbelic · 2024-04-24T04:59:39Z

@j2kun Ready for review.

Sorry for letting it grow so much. I should have created a first pull request with the changes adding the root parameter. Will try to reduce the size of commits in the future.

Hmm, I am not able to reproduce the build failure locally. Can take another look when I am more awake.

j2kun · 2024-04-24T18:14:29Z

Hmm, I am not able to reproduce the build failure locally.

I also cannot reproduce the failure locally. I'll take a look at the CI.

j2kun

LGTM! Just a few optional improvements. When you're ready please manually squash your commits and I'll get it merged.

j2kun · 2024-04-24T21:55:40Z

My guess is that this will fix the build failure:

diff --git a/lib/Dialect/Polynomial/IR/BUILD b/lib/Dialect/Polynomial/IR/BUILD
index 32666254..ac791511 100644
--- a/lib/Dialect/Polynomial/IR/BUILD
+++ b/lib/Dialect/Polynomial/IR/BUILD
@@ -38,12 +38,12 @@ cc_library(
         "PolynomialAttributes.cpp",
     ],
     hdrs = [
-        "@heir//include/Dialect/Polynomial/IR:NumberTheory.h",
         "@heir//include/Dialect/Polynomial/IR:PolynomialAttributes.h",
         "@heir//include/Dialect/Polynomial/IR:PolynomialDialect.h",
         "@heir//include/Dialect/Polynomial/IR:StaticRoots.h",
     ],
     deps = [
+        ":NumberTheory",
         ":Polynomial",
         "@heir//include/Dialect/Polynomial/IR:attributes_inc_gen",
         "@heir//include/Dialect/Polynomial/IR:dialect_inc_gen",

It's weird that I can't reproduce it locally, but I think it may be because I'm using clang + lld and the CI is using gcc + ld. I don't have time to do a detailed test/fix, but just based on my intuition this is the problem: you have the header there but no dependency on the implementation directly on the PolynomialAttributes target.

inbelic · 2024-04-24T22:27:22Z

That fixed the issue. I had accidentally added the dep to the PolynomialOp build rule instead of PolynomialAttributes.

Just need to rebase on main again to correct the failing canonicalize_ntt test-case.

- add a standard iterative implementation of the NTT Op - add performance test that uses random coefficients to prevent optimizations on dense values that artifically improves compile/runtime - include an optional parameter that computes the 2n-th primitive root of a rings ideal degree and its coefficient modulus - add a script that will compute the 2n-th roots of unity for the specified polynomial degree values and n. The values are only outputted if the 2n-th root of unity exists.

inbelic marked this pull request as ready for review April 24, 2024 04:59

inbelic mentioned this pull request Apr 24, 2024

poly: add initial lowering of NTT op #590

Closed