google / heir

A compiler for homomorphic encryption
https://heir.dev/
Apache License 2.0
330 stars 48 forks source link

reduce one multiplication in tosa to arith sigmoid #982

Open ai-mannamalai opened 1 month ago

ai-mannamalai commented 1 month ago

In recent implementation of lib/Conversion/TosaToSecretArith/TosaToSecretArith.cpp we lower Tosa::sigmoid() to degree-3 polynomial approximation as $-0.004x^3 + 0.197 * x + 0.5$. It could have used Horner's rule instead of direct computation.

Further, while current implementation computes the polynomial directly a better way would be to compute it as, $x *(0.197 - 0.004 x^2) + 0.5 $ where the number of arith::mulop is fewer by 1.

If we had easy ability to compute complex numbers we can further break down the compute as 2 products from applying the fundamental theorem of algebra.

j2kun commented 1 month ago

I think we will ultimately do a lot better by just handling polynomial approximation and evaluation more generally. My rough plan is to add a polynomial.evaluate op, which can then be lowered via horner (or better!), and would replace this hard coded lowering. Mostly we have this implemented so that we could start collecting some basic benchmarks of neural network lowering.

And then, even later, the hard coded approximation would be replaced by a remez solver so the accuracy and domain may vary as inputs.

j2kun commented 1 month ago

In other words, while I wouldn't object to a patch for this, we may as well focus our efforts on the longer term plan.

ai-mannamalai commented 1 month ago

Sounds good. Just it was a bit obvious so I had to go and say it. Your plan sounds great

Get Outlook for iOShttps://aka.ms/o0ukef


From: Jeremy Kun @.> Sent: Saturday, September 21, 2024 11:58:11 AM To: google/heir @.> Cc: Muthu Annamalai @.>; Author @.> Subject: [EXTERNAL] Re: [google/heir] reduce one multiplication in tosa to arith sigmoid (Issue #982)

In other words, while I wouldn't object to a patch for this, we may as well focus our efforts on the longer term plan.

— Reply to this email directly, view it on GitHubhttps://github.com/google/heir/issues/982#issuecomment-2365288274, or unsubscribehttps://github.com/notifications/unsubscribe-auth/APGLTDEON5PTRDUJFCLZESLZXW6UHAVCNFSM6AAAAABOTZ7FLCVHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMZDGNRVGI4DQMRXGQ. You are receiving this because you authored the thread.Message ID: @.***>

asraa commented 1 month ago

Further, while current implementation computes the polynomial directly a better way would be to compute it as, $x *(0.197 - 0.004 x^2) + 0.5 $ where the number of arith::mulop is fewer by 1.

The operation balancer pass rebalances this multiplication btw - https://github.com/google/heir/commit/704df21aa9afcdc325afd94496926b1c0b197138

In generality, Horner's method is really cool though! I think integrating that into a generic solver would be great.