Using word_by_word_montgomery with one-word prime modulus

[divergentdave]

I have a use case where I'd like to try using fiat-crypto with smaller prime moduli, such that the modulus can fit in one machine word. (see divviup/libprio-rs#757) Currently, this doesn't work. Running ./src/ExtractionOCaml/word_by_word_montgomery --lang Rust --inline field64 64 18446744069414584321 opp returns an error and dumps some IR. The issue is that Z.zselect appears deep inside an expression, when it can only be synthesized if it's the top-level node of a let binding.

Comparing the IR to successful output from a similar invocation with a two-word prime, ./src/ExtractionOCaml/word_by_word_montgomery --lang Rust --inline --debug on-success field128 64 340282366920938462946865773367900766209 opp, I noticed that the bad substitution happened following the line After rewriting subst01 for RewriteArith_0. It appears this is a particular case of #1477. I made the following patch, and it fixed the issue with synthesizing opp() (though this sort of workaround would be a game of whack-a-mole).

diff --git a/src/PushButtonSynthesis/WordByWordMontgomery.v b/src/PushButtonSynthesis/WordByWordMontgomery.v
index 74888dc11..01149e3c1 100644
--- a/src/PushButtonSynthesis/WordByWordMontgomery.v
+++ b/src/PushButtonSynthesis/WordByWordMontgomery.v
@@ -356,7 +356,7 @@ Section __.
 
   Definition opp
     := Pipeline.BoundsPipeline
-         true (* subst01 *)
+         false (* subst01 *)
          possible_values
          (reified_opp_gen
             @ GallinaReify.Reify (machine_wordsize:Z) @ GallinaReify.Reify n @ GallinaReify.Reify m)

[JasonGross]

It appears this is a particular case of #1477.
Indeed

I made the following patch, and it fixed the issue with synthesizing opp() (though this sort of workaround would be a game of whack-a-mole).

Let me know if you're interested in trying the more ambitious fix, of writing the pass that fixes #1477. I'm happy to advise on how to implement this if you want to give it a shot. (Note that SubstVarLike already exists, so the only pass to write is the one that let-binds applications of selected idents, and then piping the machinery that invokes the pass and records the idents to be bound for each backend)

[Rexicon226]

Hello, I usually don't like to comment non-productively, but I'm curious about the progress of this issue. I'm developing a library for writing zkSnarks in Zig https://github.com/Rexicon226/boson, and for some testing, I'm using a small field with an order of 641. I worked around this locally with a similar fix to the one in the OP, but having this fixed upstream would be nice.

Is there anything I could do to help? I'm always happy to learn more about this project and understand its internals better.

[JasonGross]

I don't currently have the time to work on this at all, but if you (or anybody else) wants to take on #1477 I'm happy to advise (I expect it to be about 50--150 lines of code total, you might even be able to massage DeepSeek r1 or o3-mini-high into writing a decent chunk of the code)

[Rexicon226]

Thank you. I am not interested in using AI for this, but I will take a look at solving that issue tonight. Could you elaborate a bit more on the strategy you're looking for in a fix?

From my limited understanding of the subject, I think what you're saying is we stay in "return the value" format until a final pass which converts it into "pass in a pointer to the return value" format for specific backends?

Also, could you explain what this subst01 thing is actually doing? I'm still not really sure I understand :)

[JasonGross]

From my limited understanding of the subject, I think what you're saying is we stay in "return the value" format until a final pass which converts it into "pass in a pointer to the return value" format for specific backends?

That's roughly right, yes. If you point at something that seems uncertain in this picture, I can elaborate more.

Also, could you explain what this subst01 thing is actually doing? I'm still not really sure I understand :)

subst01 says "any variable that is used 0 or 1 times can have its definition inlined into the use site". This does not work if the definition is a "return via pointer".

Instead of the kludge of turning off this inlining, we should instead re-factor-out the "return-via-pointer" calls (cf LetBindReturn, which does this for the return values to ensure that inputs are fully read by the time outputs are written). (Note that which functions return things via pointers differs depending on the output language, so the actual code transformation should be parametrized over a "does this function return via pointers" check that is hardcoded per language backend.) This pass may create a couple of duplicative variable bindings (things like int x2 = x1;), so we may need to substitute these at the end with SubstVarLike.

[Rexicon226]

Just wanted to note that I haven't forgotten about this, just quite busy with work and school the last couple of weeks. Since I don't have much familiarity with the codebase, it's a challenge for sure. I plan to get back to this either this weekend or the next.