Hugging Face's logo

Kernels:
phanerozoic
/
physarum

Kernel card Files
xet
 Community
physarum / physarum_csrc /flow_cpu.cpp
phanerozoic's picture
phanerozoic
add Physarum flow solver (shortest path / network); card leads with it
0a16898 verified
Raw
History Blame
7.25 kB
// physarum-flow: the Tero adaptive-conductivity Physarum model, whose compute
// core is a weighted-grid-Laplacian solve. Given per-edge conductivities on a
// masked grid and a balanced source/sink injection, flow_cg solves L(cE,cS) p = b
// for the node pressures p by Jacobi-preconditioned conjugate gradients. One node
// is grounded with a diagonal penalty so the Laplacian is positive-definite (no
// singular nullspace to project out); the gauge choice does not affect edge flux.
// The caller iterates: solve -> flux Q = c*grad(p) -> reinforce conductive edges,
// which converges onto the shortest path (single source/sink) or an efficient
// network (multiple terminals). The 5-point weighted mat-vec is in gather form
// (each node reads its neighbours, no scatter) and vectorises with NEON; CG on a
// modest grid is latency-bound and sequential, so it is not threaded.
#include <torch/torch.h>
#include <cstdint>
#include <cmath>
#include <vector>
namespace {
const float GROUND = 1e3f; // diagonal penalty pinning the ground node's gauge
 // (kept modest: too large floors the f32 CG residual)
#if defined(__ARM_NEON)
#include <arm_neon.h>
#endif
// Ap = (L + GROUND e_g e_g^T) p. cE[y*W+x] is the edge (y,x)-(y,x+1); cS the edge
// (y,x)-(y+1,x); both zero on non-open edges, so wall nodes yield 0.
void matvec(float* Ap, const float* p, const float* cE, const float* cS,
 int H, int W, int g){
 for (int y=0; y<H; y++){
 const float* pr = p + (size_t)y*W;
 const float* pu = (y>0) ? p + (size_t)(y-1)*W : nullptr;
 const float* pd = (y<H-1) ? p + (size_t)(y+1)*W : nullptr;
 const float* cEr = cE + (size_t)y*W;
 const float* cSr = cS + (size_t)y*W;
 const float* cSu = (y>0) ? cS + (size_t)(y-1)*W : nullptr;
 float* o = Ap + (size_t)y*W;
 int x = 0;
#if defined(__ARM_NEON)
 if (y>0 && y<H-1){ // interior rows: all 4 edges exist
 { float pc=pr[0], s=cEr[0]*(pc-pr[1])+cSr[0]*(pc-pd[0])+cSu[0]*(pc-pu[0]); o[0]=s; }
 for (x=1; x+4<=W-1; x+=4){ // vectorise interior cols [1, W-2]
 float32x4_t pc=vld1q_f32(pr+x);
 float32x4_t ce=vld1q_f32(cEr+x), cw=vld1q_f32(cEr+x-1);
 float32x4_t cs=vld1q_f32(cSr+x), cn=vld1q_f32(cSu+x);
 float32x4_t s=vmulq_f32(ce, vsubq_f32(pc, vld1q_f32(pr+x+1)));
 s=vmlaq_f32(s, cw, vsubq_f32(pc, vld1q_f32(pr+x-1)));
 s=vmlaq_f32(s, cs, vsubq_f32(pc, vld1q_f32(pd+x)));
 s=vmlaq_f32(s, cn, vsubq_f32(pc, vld1q_f32(pu+x)));
 vst1q_f32(o+x, s);
 }
 for (; x<W; x++){ // scalar tail (incl. col W-1)
 float pc=pr[x], s=cEr[x-1]*(pc-pr[x-1])+cSr[x]*(pc-pd[x])+cSu[x]*(pc-pu[x]);
 if (x<W-1) s+=cEr[x]*(pc-pr[x+1]);
 o[x]=s;
 }
 continue;
 }
#endif
 for (x=0; x<W; x++){ // boundary rows / scalar fallback
 float pc=pr[x], s=0.0f;
 if (x<W-1) s+=cEr[x]*(pc-pr[x+1]);
 if (x>0) s+=cEr[x-1]*(pc-pr[x-1]);
 if (y<H-1) s+=cSr[x]*(pc-pd[x]);
 if (y>0) s+=cSu[x]*(pc-pu[x]);
 o[x]=s;
 }
 }
 Ap[g] += GROUND * p[g];
}
#if defined(__ARM_NEON)
inline double hadd(float32x4_t v){ float32x2_t s=vadd_f32(vget_low_f32(v),vget_high_f32(v)); return (double)vget_lane_f32(vpadd_f32(s,s),0); }
#endif
double ddot(const float* a, const float* b, int N){
#if defined(__ARM_NEON)
 float32x4_t acc=vdupq_n_f32(0); int i=0;
 for (; i+4<=N; i+=4) acc=vmlaq_f32(acc, vld1q_f32(a+i), vld1q_f32(b+i));
 double s=hadd(acc); for (; i<N; i++) s+=(double)a[i]*b[i]; return s;
#else
 double s=0; for (int i=0;i<N;i++) s += (double)a[i]*(double)b[i]; return s;
#endif
}
// p += a*d ; r -= a*Ad ; return sum r^2
double axpy_rr(float* p, float* r, const float* d, const float* Ad, float a, int N){
#if defined(__ARM_NEON)
 float32x4_t va=vdupq_n_f32(a), acc=vdupq_n_f32(0); int i=0;
 for (; i+4<=N; i+=4){
 float32x4_t di=vld1q_f32(d+i), adi=vld1q_f32(Ad+i);
 float32x4_t pi=vmlaq_f32(vld1q_f32(p+i), va, di); vst1q_f32(p+i, pi);
 float32x4_t ri=vmlsq_f32(vld1q_f32(r+i), va, adi); vst1q_f32(r+i, ri);
 acc=vmlaq_f32(acc, ri, ri);
 }
 double s=hadd(acc);
 for (; i<N; i++){ p[i]+=a*d[i]; r[i]-=a*Ad[i]; s+=(double)r[i]*r[i]; }
 return s;
#else
 double s=0; for (int i=0;i<N;i++){ p[i]+=a*d[i]; r[i]-=a*Ad[i]; s+=(double)r[i]*r[i]; } return s;
#endif
}
void precond(float* z, const float* r, const float* diag, int N){ // z = r/diag (0 where diag~0)
#if defined(__ARM_NEON)
 float32x4_t eps=vdupq_n_f32(1e-20f); int i=0;
 for (; i+4<=N; i+=4){
 float32x4_t dg=vld1q_f32(diag+i);
 float32x4_t zz=vdivq_f32(vld1q_f32(r+i), dg);
 vst1q_f32(z+i, vbslq_f32(vcgtq_f32(dg,eps), zz, vdupq_n_f32(0)));
 }
 for (; i<N; i++) z[i]= diag[i]>1e-20f ? r[i]/diag[i] : 0.f;
#else
 for (int i=0;i<N;i++) z[i]= diag[i]>1e-20f ? r[i]/diag[i] : 0.f;
#endif
}
void xpby(float* d, const float* z, float beta, int N){ // d = z + beta*d
#if defined(__ARM_NEON)
 float32x4_t vb=vdupq_n_f32(beta); int i=0;
 for (; i+4<=N; i+=4) vst1q_f32(d+i, vmlaq_f32(vld1q_f32(z+i), vb, vld1q_f32(d+i)));
 for (; i<N; i++) d[i]=z[i]+beta*d[i];
#else
 for (int i=0;i<N;i++) d[i]=z[i]+beta*d[i];
#endif
}
} // namespace
// Solve L p = b with node `ground` pinned, by Jacobi-preconditioned CG; p is
// warm-started and overwritten with the solution.
void flow_cg(torch::Tensor cE, torch::Tensor cS, torch::Tensor b,
 torch::Tensor p, int64_t ground, int64_t max_iters, double tol){
 TORCH_CHECK(p.dtype()==torch::kFloat32 && p.is_contiguous(), "p must be f32 contiguous");
 int H=(int)p.size(0), W=(int)p.size(1), N=H*W, g=(int)ground;
 const float* cEp=cE.data_ptr<float>(); const float* cSp=cS.data_ptr<float>();
 const float* bp=b.data_ptr<float>();
 float* pp=p.data_ptr<float>();
 std::vector<float> diag(N,0.f), r(N), z(N), d(N), Ad(N);
 for (int y=0;y<H;y++) for (int x=0;x<W;x++){ // diag = sum of incident conductivities
 int i=y*W+x; float s=0.f;
 if (x<W-1) s+=cEp[i]; if (x>0) s+=cEp[i-1];
 if (y<H-1) s+=cSp[i]; if (y>0) s+=cSp[i-W];
 diag[i]=s;
 }
 diag[g]+=GROUND;
matvec(Ad.data(), pp, cEp, cSp, H, W, g);
 for (int i=0;i<N;i++) r[i]=bp[i]-Ad[i];
 precond(z.data(), r.data(), diag.data(), N);
 for (int i=0;i<N;i++) d[i]=z[i];
 double rz = ddot(r.data(), z.data(), N), tol2 = tol*tol;
 for (int it=0; it<max_iters; it++){
 matvec(Ad.data(), d.data(), cEp, cSp, H, W, g);
 double dAd = ddot(d.data(), Ad.data(), N);
 if (dAd<=1e-30) break;
 float alpha=(float)(rz/dAd);
 double rr = axpy_rr(pp, r.data(), d.data(), Ad.data(), alpha, N);
 if (rr < tol2) break;
 precond(z.data(), r.data(), diag.data(), N);
 double rzn = ddot(r.data(), z.data(), N);
 float beta=(float)(rzn/rz);
 xpby(d.data(), z.data(), beta, N);
 rz=rzn;
 }
}