par.bend

bend-prot@0.1.0.0 · 36 lines · 3 definitions

1import Base
2import bend-prot@0.1.0.0/geom.bend as G
3import bend-prot@0.1.0.0/protein.bend as P
4import bend-prot@0.1.0.0/topology.bend as T
5import bend-prot@0.1.0.0/force.bend as F
6
7# --- Parallel kernels v0: 4-wide unrolled fork-join. ---
8# Bend has no forward references, so divide-and-conquer halving (which must
9# match a computed split) is inexpressible: it needs a helper cycle. Instead
10# each step runs 4 independent head computations as parallel calls and folds
11# one tail, which is a single self-recursion. Call an entry with `!` to run
12# every nested parallel call on the GPU (falls back to CPU without one;
13# verify codegen with `bend <file> -o <file>.c`). Remainders under 4 fall
14# back to the sequential kernels.
15
16def count_par4(xs: List<&2, P.Atom>, +qp: G.Vec3, +c2: F32) -> Nat:
17  match xs:
18    case h1 <> h2 <> h3 <> h4 <> t:
19      a b = T.head_hit_nat(h1, qp, c2) T.head_hit_nat(h2, qp, c2)
20      c d = T.head_hit_nat(h3, qp, c2) T.head_hit_nat(h4, qp, c2)
21      Nat.add(Nat.add(a, b), Nat.add(c, Nat.add(d, count_par4(t, qp, c2))))
22    case rest:
23      T.count_below(rest, qp, c2)
24
25# Plain-cutoff wrapper, mirroring T.count_within.
26def count_within_par(xs: List<&2, P.Atom>, +qp: G.Vec3, +cutoff: F32) -> Nat:
27  count_par4(xs, qp, (cutoff * cutoff : F32))
28
29def lj_row_par4(+pi: G.Vec3, +eps: F32, +sig2: F32, +c2: F32, ys: List<&2, P.Atom>) -> F32:
30  match ys:
31    case h1 <> h2 <> h3 <> h4 <> t:
32      a b = F.lj_head(pi, eps, sig2, c2, h1) F.lj_head(pi, eps, sig2, c2, h2)
33      c d = F.lj_head(pi, eps, sig2, c2, h3) F.lj_head(pi, eps, sig2, c2, h4)
34      ((a + b : F32) + ((c + d : F32) + lj_row_par4(pi, eps, sig2, c2, t) : F32) : F32)
35    case rest:
36      F.lj_row(pi, eps, sig2, c2, rest)