add rust upow solver

Author: vans163

ex/config/runtime.exs

@@ -87,6 +87,8 @@ config :ama, :pruner_enabled, not archival_node and history_keep_epochs > 0
 config :ama, :statepeerdownload, System.get_env("STATEPEERDOWNLOAD") in ["true", "y", "yes"]
 config :ama, :autoupdate, System.get_env("AUTOUPDATE") in ["true", "y", "yes"]
 config :ama, :computor_type, (case System.get_env("COMPUTOR") do nil -> nil; "trainer" -> :trainer; _ -> :default end)
+computor_upow_threads_default = max(1, div(:erlang.system_info(:schedulers_online), 2) - 2)
+config :ama, :computor_upow_threads, (case System.get_env("COMPUTOR_UPOW_THREADS") do nil -> computor_upow_threads_default; v -> :erlang.binary_to_integer(v) end)
 
 config :ama, :max_peers, (System.get_env("MAX_PEERS") || "300") |> :erlang.binary_to_integer()
 config :ama, :buy_peer_sol, System.get_env("BUY_PEER_SOL") in ["true", "y", "yes"]

ex/lib/native/rdb.ex

@@ -53,6 +53,8 @@ defmodule RDB do
 
   def freivalds(_tensor, _vr), do: :erlang.nif_error(:nif_not_loaded)
 
+  def compute_upow(_epoch, _segment_vr_hash, _trainer, _pop, _computor, _diff_bits, _iterations, _threads), do: :erlang.nif_error(:nif_not_loaded)
+
   def bintree_root(_propslist), do: :erlang.nif_error(:nif_not_loaded)
   def bintree_root_prove(_propslist, _ns, _key), do: :erlang.nif_error(:nif_not_loaded)
   def bintree_root_verify(_expected_root, _proof, _ns, _key, _value), do: :erlang.nif_error(:nif_not_loaded)

ex/lib/node/computor_gen.ex

@@ -1,6 +1,8 @@
 defmodule ComputorGen do
   use GenServer
 
+  @batch_iterations 2000
+
   def start(type \\ nil) do
     send(__MODULE__, {:start, type})
   end
@@ -24,22 +26,22 @@ defmodule ComputorGen do
   end
 
   def handle_info(:tick, state) do
-    state = cond do
-      !state[:enabled] -> state
-      !Application.fetch_env!(:ama, :testnet) ->
-        IO.puts "Computor currently cannot find sols on mainnet due to difficulty. Do not waste CPU running it."
-        state
+    next_ms = cond do
+      !state[:enabled] -> 1000
       !FabricSyncAttestGen.isQuorumIsInEpoch() ->
         IO.puts "🔴 cannot compute: out_of_sync"
-        state
+        1000
       true ->
         tick(state)
+        0
     end
-    :erlang.send_after(1000, self(), :tick)
+    :erlang.send_after(next_ms, self(), :tick)
     {:noreply, state}
   end
 
   def handle_info({:start, type}, state) do
+    threads = Application.get_env(:ama, :computor_upow_threads, 0)
+    IO.puts "🔢 computor enabled (type=#{inspect type}, upow_threads=#{if threads == 0, do: "auto", else: threads})"
     state = Map.put(state, :enabled, true)
     state = Map.put(state, :type, type)
     {:noreply, state}
@@ -53,23 +55,28 @@ defmodule ComputorGen do
   def handle_info(_msg, state), do: {:noreply, state}
 
   def tick(state) do
-    IO.puts "computor running #{DateTime.utc_now()}"
     pk = Application.fetch_env!(:ama, :trainer_pk)
     pop = Application.fetch_env!(:ama, :trainer_pop)
+    threads = Application.get_env(:ama, :computor_upow_threads, 0)
 
     coins = DB.Chain.balance(pk)
     epoch = DB.Chain.epoch()
+    segment_vr_hash = DB.Chain.segment_vr_hash()
+    diff_bits = DB.Chain.diff_bits()
     hasExecCoins = coins >= BIC.Coin.to_cents(100)
     cond do
+        is_nil(segment_vr_hash) -> :ok # epoch segment_vr not set yet; nothing to compute against
+
         (state.type == :trainer and !hasExecCoins) or state.type == nil ->
-          sol = UPOW.compute_for(epoch, EntryGenesis.signer(), EntryGenesis.pop(), pk, :crypto.strong_rand_bytes(96), 100)
+          # Compute to ourselves as a node: own pk is the computor (reward recipient).
+          sol = UPOW.compute(epoch, EntryGenesis.signer(), EntryGenesis.pop(), pk, segment_vr_hash, diff_bits, @batch_iterations, threads)
           if sol do
             IO.puts "🔢 tensor matmul complete! broadcasting sol.."
             NodeGen.broadcast(%{op: :sol, sol: sol})
           end
 
         true ->
-          sol = UPOW.compute_for(epoch, pk, pop, pk, :crypto.strong_rand_bytes(96), 100)
+          sol = UPOW.compute(epoch, pk, pop, pk, segment_vr_hash, diff_bits, @batch_iterations, threads)
           if sol do
             sk = Application.fetch_env!(:ama, :trainer_sk)
             packed_tx = TX.build(sk, "Epoch", "submit_sol", [sol])

ex/lib/node/upow.ex

@@ -1,6 +1,9 @@
 defmodule UPOW do
     def compute_for(epoch, trainer, pop, computor, segment_vr, itrs \\ 30)
-    def compute_for(epoch, trainer, pop, computor, segment_vr, 0), do: nil
+    def compute_for(epoch, trainer, pop, computor, segment_vr, itrs) when epoch >= 156 do
+        compute(epoch, trainer, pop, computor, Blake3.hash(segment_vr), DB.Chain.diff_bits(), itrs, 0)
+    end
+    def compute_for(_epoch, _trainer, _pop, _computor, _segment_vr, 0), do: nil
     def compute_for(epoch, trainer, pop, computor, segment_vr, itrs) do
         {hash, sol} = branch_sol(epoch, trainer, pop, computor, segment_vr)
         valid = BIC.Sol.verify_hash(epoch, hash)
@@ -11,6 +14,13 @@ defmodule UPOW do
         end
     end
 
+    def compute(epoch, trainer, pop, computor, segment_vr_hash, diff_bits, iterations \\ 100, threads \\ 0) do
+        case RDB.compute_upow(epoch, segment_vr_hash, trainer, pop, computor, diff_bits, iterations, threads) do
+            {:ok, sol} -> sol
+            nil -> nil
+        end
+    end
+
     def branch_sol(epoch, trainer, pop, computor, segment_vr) do
         cond do
             epoch >= 156 -> UPOW2.tensormath(epoch, Blake3.hash(segment_vr), trainer, pop, computor)

ex/native/rdb/src/lib.rs

@@ -2,6 +2,7 @@ pub mod atoms;
 pub mod consensus;
 pub mod model;
 pub mod tx_filter;
+pub mod upow;
 
 use rustler::types::{Binary, OwnedBinary};
 use rustler::{Atom, Encoder, Env, Error, NifResult, NifTaggedEnum, ResourceArc, Term};
@@ -863,6 +864,50 @@ fn freivalds(tensor: Binary, vr_b3: Binary) -> bool {
     crate::consensus::bic::sol_freivalds::freivalds(tensor.as_slice(), vr_b3.as_slice())
 }
 
+/// UPOW2 computer. Computes up to `iterations` nonce attempts (each a real 16x50240
+/// * 50240x16 matmul), split across `threads` workers (min 1), and returns the first
+/// 1264-byte sol whose Blake3 hash has at least `diff_bits` leading zero bits (the
+/// network difficulty), or nil. DirtyCpu.
+#[rustler::nif(schedule = "DirtyCpu")]
+fn compute_upow<'a>(
+    env: Env<'a>,
+    epoch: u64,
+    segment_vr_hash: Binary,
+    trainer: Binary,
+    pop: Binary,
+    computor: Binary,
+    diff_bits: u64,
+    iterations: u64,
+    threads: u64,
+) -> NifResult<Term<'a>> {
+    use crate::upow;
+    if segment_vr_hash.len() != 32 || trainer.len() != 48 || pop.len() != 96 || computor.len() != 48 {
+        return Err(Error::BadArg);
+    }
+
+    // Build the 240-byte seed template; nonce bytes [228..240] are filled per attempt.
+    let mut seed = [0u8; upow::PREAMBLE];
+    seed[0..4].copy_from_slice(&(epoch as u32).to_le_bytes());
+    seed[4..36].copy_from_slice(segment_vr_hash.as_slice());
+    seed[36..84].copy_from_slice(trainer.as_slice());
+    seed[84..180].copy_from_slice(pop.as_slice());
+    seed[180..228].copy_from_slice(computor.as_slice());
+
+    let nthreads = (threads as usize).max(1);
+    let rng_seeds: Vec<u64> = (0..nthreads).map(|_| rand::random::<u64>()).collect();
+
+    let found = std::panic::catch_unwind(|| upow::compute(&seed, diff_bits as u32, iterations, nthreads, &rng_seeds)).unwrap_or(None);
+
+    match found {
+        Some(sol) => {
+            let mut ob = OwnedBinary::new(sol.len()).ok_or_else(|| Error::Term(Box::new("alloc failed")))?;
+            ob.as_mut_slice().copy_from_slice(&sol);
+            Ok((atoms::ok(), Binary::from_owned(ob, env)).encode(env))
+        }
+        None => Ok(atoms::nil().encode(env)),
+    }
+}
+
 #[rustler::nif]
 fn bintree_root<'a>(env: Env<'a>, proplist: Vec<(Option<Binary<'a>>, Binary<'a>, Binary<'a>)>) -> Term<'a> {
     let mut ops = Vec::with_capacity(100);