Adding write control of s_mem for first iteration and adding a new stage in priority mux.

1 files changed, 4 insertions, 2 deletions

diff --git a/src/rtl/montprod.v b/src/rtl/montprod.v
index 4f4a9aa..00227d0 100644
--- a/src/rtl/montprod.v
+++ b/src/rtl/montprod.v
@@ -354,7 +354,7 @@ module montprod(
         CTRL_L_CALC_SA:
           begin
             //s = (s + q*M + b*A) >>> 1;, if(b==1) S+= A. Takes (1..length) cycles.
-            s_mem_we_new = b_reg | q_reg;
+            s_mem_we_new = b_reg | q_reg | first_iteration_reg;
           end
 
         CTRL_L_CALC_SDIV2:
@@ -392,8 +392,10 @@ module montprod(
 
             if (b_reg)
               s_mem_new = add_result_sa;
-            else
+            else if (q_reg)
               s_mem_new = add_result_sm;
+            else
+              s_mem_new = muxed_s_mem_read_data;
 
             add_carry_in_sa_new = add_carry_out_sa;
             add_carry_in_sm_new = add_carry_out_sm;