Hugging Face's logo

Kernels:
phanerozoic
/
physarum

Kernel card Files
xet
 Community
physarum / physarum_csrc /physarum_cpu.cpp
phanerozoic's picture
phanerozoic
publish buildable source (threaded + tile-sorted kernel)
bfdbd33 verified
Raw
History Blame
12 kB
// physarum-cpu: one step of a Physarum (slime-mould) transport-network sim.
// Synchronous update: all agents sense the start-of-step field along three
// sensors, steer up-gradient and move (parallel over agents, each agent's jitter
// from an index-seeded RNG); deposits are then applied in a serial index-order
// pass, and the field is diffused (separable 3-tap) and decayed. On a cadence
// the agents are counting-sorted into tile order so that consecutive (and
// same-thread) agents sense a small hot window instead of scattering three
// gathers across the whole field -- the gathers dominate and otherwise miss.
// Read-only sensing + ordered deposit + serial sort make the step deterministic
// for any thread count. Agent trig/arithmetic and the diffuse stencil vectorise
// on NEON (aarch64); a portable scalar path covers everything else.
#include <torch/torch.h>
#include <ATen/Parallel.h>
#include <cstdint>
#include <cstring>
#include <cmath>
#include <thread>
#include <vector>
static const int TILE = 32; // sort granularity (px); ~8-step drift stays inside
static const int RESORT = 8; // re-sort every RESORT steps (amortises the sort)
namespace {
// Fan a range [0,total) across at::get_num_threads() worker threads, each given
// a contiguous chunk; f(begin,end) must touch only its own slice. std::thread
// (not OpenMP) so it parallelises under the stock kernel-builder compile, which
// exposes no way to pass -fopenmp to a CPU kernel. Falls back to serial for a
// single thread or a small range.
template<class F>
inline void parallel_chunks(int total, F f){
 int T = at::get_num_threads(); if (T < 1) T = 1;
 if (T == 1 || total < 4096){ f(0, total); return; }
 if (T > total) T = total;
 int chunk = (total + T - 1) / T;
 std::vector<std::thread> pool; pool.reserve(T - 1);
 for (int t=1; t<T; t++){
 int b=t*chunk, e=b+chunk>total?total:b+chunk;
 if (b>=e) break;
 pool.emplace_back([&,b,e]{ f(b,e); });
 }
 f(0, chunk>total?total:chunk); // calling thread runs chunk 0
 for (auto& th : pool) th.join();
}
inline uint32_t xr(uint32_t &s){ s ^= s<<13; s ^= s>>17; s ^= s<<5; return s; }
inline float frand(uint32_t &s){ return (xr(s) >> 8) * (1.0f/16777216.0f); }
// per-agent, per-step independent seed: the step's `seed` argument already
// encodes base_seed+t, so hashing it with the agent index gives a stream that
// varies per agent and per step yet does not depend on iteration order.
inline uint32_t agent_seed(uint32_t base, int i){
 uint32_t s = (uint32_t)i*2654435761u ^ base*2246822519u;
 s ^= s>>15; s *= 2246822519u; s ^= s>>13; s *= 3266489917u; s ^= s>>16;
 return s ? s : 1u;
}
inline float fsin(float x){ // parabola approx, |err|<1e-3
 float k = x * 0.15915494f;
 k = x - 6.2831853f * (float)((int)(k + (k < 0 ? -0.5f : 0.5f)));
 float y = 1.27323954f*k - 0.405284735f*k*(k<0?-k:k);
 return 0.225f*(y*(y<0?-y:y) - y) + y;
}
inline float fcos(float x){ return fsin(x + 1.5707963f); }
inline float samp(const float* t, float x, float y, int W, int H){
 int xi = ((int)x % W + W) % W, yi = ((int)y % H + H) % H;
 return t[(size_t)yi*W + xi];
}
void agents(const float* t, float* ax, float* ay, float* ah, int W, int H, int N,
 float sd, float sa, float turn, float speed, uint32_t base);
void deposit(float* t, const float* ax, const float* ay, int W, int H, int N, float dep);
void diffuse(float* t, float* tmp, int W, int H, float decay);
#if defined(__ARM_NEON)
#include <arm_neon.h>
inline float32x4_t vsin(float32x4_t x){
 float32x4_t k = vrndnq_f32(vmulq_n_f32(x, 0.15915494f));
 x = vmlsq_n_f32(x, k, 6.2831853f);
 float32x4_t ax = vabsq_f32(x);
 float32x4_t y = vsubq_f32(vmulq_n_f32(x,1.27323954f),
 vmulq_n_f32(vmulq_f32(x,ax),0.405284735f));
 float32x4_t ay = vabsq_f32(y);
 return vmlaq_n_f32(y, vsubq_f32(vmulq_f32(y,ay), y), 0.225f);
}
inline float32x4_t vcos(float32x4_t x){ return vsin(vaddq_f32(x, vdupq_n_f32(1.5707963f))); }
// sense + steer + move for the agent quads [q0,q1); reads the field only, writes
// ax/ay/ah. No field writes, so any set of quads runs on any thread.
inline void move_quads(const float* t, float* ax, float* ay, float* ah,
 int q0, int q1, int W, int H, float sd, float sa,
 float turn, float speed, uint32_t base){
 float32x4_t vsd=vdupq_n_f32(sd), vsp=vdupq_n_f32(speed), vsa=vdupq_n_f32(sa);
 float32x4_t vW=vdupq_n_f32((float)W), vH=vdupq_n_f32((float)H);
 float32x4_t iW=vdupq_n_f32(1.0f/W), iH=vdupq_n_f32(1.0f/H);
 for (int q=q0; q<q1; q++){
 int i=q*4;
 float32x4_t x=vld1q_f32(ax+i), y=vld1q_f32(ay+i), h=vld1q_f32(ah+i);
 float32x4_t hl=vsubq_f32(h,vsa), hr=vaddq_f32(h,vsa);
 float cx[4],cy[4],lx[4],ly[4],rx[4],ry[4];
 vst1q_f32(cx, vmlaq_f32(x, vcos(h), vsd)); vst1q_f32(cy, vmlaq_f32(y, vsin(h), vsd));
 vst1q_f32(lx, vmlaq_f32(x, vcos(hl),vsd)); vst1q_f32(ly, vmlaq_f32(y, vsin(hl),vsd));
 vst1q_f32(rx, vmlaq_f32(x, vcos(hr),vsd)); vst1q_f32(ry, vmlaq_f32(y, vsin(hr),vsd));
 float FC[4],FL[4],FR[4];
 for (int l=0;l<4;l++){ FC[l]=samp(t,cx[l],cy[l],W,H);
 FL[l]=samp(t,lx[l],ly[l],W,H); FR[l]=samp(t,rx[l],ry[l],W,H); }
 float32x4_t fc=vld1q_f32(FC), fl=vld1q_f32(FL), fr=vld1q_f32(FR);
 float jit[4]; for(int l=0;l<4;l++){ uint32_t s=agent_seed(base,i+l); jit[l]=0.5f+0.5f*frand(s); }
 float32x4_t vj=vld1q_f32(jit), vturn=vmulq_n_f32(vj,turn);
 uint32x4_t straight=vandq_u32(vcgtq_f32(fc,fl), vcgtq_f32(fc,fr));
 float32x4_t dh=vbslq_f32(vcgtq_f32(fl,fr), vnegq_f32(vturn),
 vbslq_f32(vcgtq_f32(fr,fl), vturn,
 vmulq_n_f32(vsubq_f32(vj,vdupq_n_f32(0.5f)), 2.0f*turn)));
 dh=vbslq_f32(straight, vdupq_n_f32(0.0f), dh);
 h=vaddq_f32(h,dh);
 x=vmlaq_f32(x, vcos(h), vsp); y=vmlaq_f32(y, vsin(h), vsp);
 x=vsubq_f32(x, vmulq_f32(vW, vrndmq_f32(vmulq_f32(x,iW))));
 y=vsubq_f32(y, vmulq_f32(vH, vrndmq_f32(vmulq_f32(y,iH))));
 vst1q_f32(ax+i,x); vst1q_f32(ay+i,y); vst1q_f32(ah+i,h);
 }
}
void agents(const float* t, float* ax, float* ay, float* ah, int W, int H, int N,
 float sd, float sa, float turn, float speed, uint32_t base){
 int nq = N/4;
 parallel_chunks(nq, [&](int q0, int q1){
 move_quads(t, ax, ay, ah, q0, q1, W, H, sd, sa, turn, speed, base);
 });
 for (int i=nq*4; i<N; i++){ // scalar tail, same sync rule
 float x=ax[i], y=ay[i], h=ah[i];
 float fc=samp(t,x+fcos(h)*sd,y+fsin(h)*sd,W,H);
 float fl=samp(t,x+fcos(h-sa)*sd,y+fsin(h-sa)*sd,W,H);
 float fr=samp(t,x+fcos(h+sa)*sd,y+fsin(h+sa)*sd,W,H);
 uint32_t s=agent_seed(base,i); float j=0.5f+0.5f*frand(s);
 if(fc>fl&&fc>fr){} else if(fl>fr) h-=turn*j; else if(fr>fl) h+=turn*j; else h+=(j-0.5f)*2.0f*turn;
 x=fmodf(x+fcos(h)*speed+W,W); y=fmodf(y+fsin(h)*speed+H,H);
 ax[i]=x; ay[i]=y; ah[i]=h;
 }
}
void diffuse(float* t, float* tmp, int W, int H, float dec){
 for (int y=0;y<H;y++){ const float* r=t+(size_t)y*W; float* o=tmp+(size_t)y*W;
 o[0]=0.25f*r[W-1]+0.5f*r[0]+0.25f*r[1]; int x=1;
 for (; x+4<=W-1; x+=4){ float32x4_t L=vld1q_f32(r+x-1),C=vld1q_f32(r+x),R=vld1q_f32(r+x+1);
 vst1q_f32(o+x, vmlaq_n_f32(vmlaq_n_f32(vmulq_n_f32(C,0.5f),L,0.25f),R,0.25f)); }
 for (; x<W-1; x++) o[x]=0.25f*r[x-1]+0.5f*r[x]+0.25f*r[x+1];
 o[W-1]=0.25f*r[W-2]+0.5f*r[W-1]+0.25f*r[0]; }
 float32x4_t vd=vdupq_n_f32(dec);
 for (int y=0;y<H;y++){ int yu=y?y-1:H-1, yd=y+1<H?y+1:0;
 const float* U=tmp+(size_t)yu*W,*C=tmp+(size_t)y*W,*D=tmp+(size_t)yd*W;
 float* o=t+(size_t)y*W; int x=0;
 for (; x+4<=W; x+=4){ float32x4_t s=vmlaq_n_f32(vmlaq_n_f32(vmulq_n_f32(vld1q_f32(C+x),0.5f),
 vld1q_f32(U+x),0.25f), vld1q_f32(D+x),0.25f); vst1q_f32(o+x, vmulq_f32(s,vd)); }
 for (; x<W; x++) o[x]=(0.25f*U[x]+0.5f*C[x]+0.25f*D[x])*dec; }
}
#else // ---------- portable scalar ----------
void agents(const float* t, float* ax, float* ay, float* ah, int W, int H, int N,
 float sd, float sa, float turn, float speed, uint32_t base){
 parallel_chunks(N, [&](int i0, int i1){
 for (int i=i0;i<i1;i++){
 float x=ax[i], y=ay[i], h=ah[i];
 float fc=samp(t,x+fcos(h)*sd,y+fsin(h)*sd,W,H);
 float fl=samp(t,x+fcos(h-sa)*sd,y+fsin(h-sa)*sd,W,H);
 float fr=samp(t,x+fcos(h+sa)*sd,y+fsin(h+sa)*sd,W,H);
 uint32_t s=agent_seed(base,i); float j=0.5f+0.5f*frand(s);
 if(fc>fl&&fc>fr){} else if(fl>fr) h-=turn*j; else if(fr>fl) h+=turn*j; else h+=(j-0.5f)*2.0f*turn;
 x=fmodf(x+fcos(h)*speed+W,W); y=fmodf(y+fsin(h)*speed+H,H);
 ax[i]=x; ay[i]=y; ah[i]=h;
 }
 });
}
void diffuse(float* t, float* tmp, int W, int H, float dec){
 for (int y=0;y<H;y++){ const float* r=t+(size_t)y*W; float* o=tmp+(size_t)y*W;
 for (int x=0;x<W;x++){ int xl=x?x-1:W-1, xr=x+1<W?x+1:0;
 o[x]=0.25f*r[xl]+0.5f*r[x]+0.25f*r[xr]; } }
 for (int y=0;y<H;y++){ int yu=y?y-1:H-1, yd=y+1<H?y+1:0;
 const float* U=tmp+(size_t)yu*W,*C=tmp+(size_t)y*W,*D=tmp+(size_t)yd*W;
 float* o=t+(size_t)y*W;
 for (int x=0;x<W;x++) o[x]=(0.25f*U[x]+0.5f*C[x]+0.25f*D[x])*dec; }
}
#endif
// deposit is a serial index-order scatter so the field sum is identical for any
// thread count; positions are already wrapped into range by the move.
void deposit(float* t, const float* ax, const float* ay, int W, int H, int N, float dep){
 for (int i=0;i<N;i++){ int xi=(int)ax[i], yi=(int)ay[i];
 if(xi>=0&&xi<W&&yi>=0&&yi<H) t[(size_t)yi*W+xi]+=dep; }
}
// counting-sort agents into row-major tile order (serial, deterministic). Tile
// coords are clamped because the position wrap can round to exactly x==W.
void bin_sort(float* ax, float* ay, float* ah, int W, int H, int N, int tile){
 int TW=(W+tile-1)/tile, TH=(H+tile-1)/tile, nt=TW*TH;
 std::vector<int> off(nt+1, 0), cur(nt);
 auto tid=[&](int i){ int xi=(int)ax[i]; xi=xi<0?0:(xi>=W?W-1:xi);
 int yi=(int)ay[i]; yi=yi<0?0:(yi>=H?H-1:yi);
 return (yi/tile)*TW + xi/tile; };
 for (int i=0;i<N;i++) off[tid(i)+1]++;
 for (int t=1;t<=nt;t++) off[t]+=off[t-1];
 for (int t=0;t<nt;t++) cur[t]=off[t];
 std::vector<float> sx(N), sy(N), sh(N);
 for (int i=0;i<N;i++){ int p=cur[tid(i)]++; sx[p]=ax[i]; sy[p]=ay[i]; sh[p]=ah[i]; }
 memcpy(ax, sx.data(), (size_t)N*4);
 memcpy(ay, sy.data(), (size_t)N*4);
 memcpy(ah, sh.data(), (size_t)N*4);
}
} // namespace
void physarum_step(torch::Tensor trail, torch::Tensor tmp,
 torch::Tensor ax, torch::Tensor ay, torch::Tensor ah,
 int64_t seed, double sense_dist, double sense_ang,
 double turn, double speed, double deposit_, double decay){
 TORCH_CHECK(trail.dtype()==torch::kFloat32 && trail.is_contiguous(), "trail must be f32 contiguous");
 TORCH_CHECK(trail.dim()==2 && tmp.sizes()==trail.sizes(), "trail/tmp must be [H,W]");
 int H=(int)trail.size(0), W=(int)trail.size(1), N=(int)ax.size(0);
 uint32_t base = (uint32_t)(seed ^ 0x9e3779b9u); if(!base) base=1;
 float* t = trail.data_ptr<float>();
 // seed increments per step (base_seed + t), so this re-sorts every RESORT steps
 if ((uint64_t)seed % (uint64_t)RESORT == 0)
 bin_sort(ax.data_ptr<float>(), ay.data_ptr<float>(), ah.data_ptr<float>(), W,H,N, TILE);
 agents(t, ax.data_ptr<float>(), ay.data_ptr<float>(), ah.data_ptr<float>(),
 W,H,N, (float)sense_dist,(float)sense_ang,(float)turn,(float)speed, base);
 deposit(t, ax.data_ptr<float>(), ay.data_ptr<float>(), W,H,N, (float)deposit_);
 diffuse(t, tmp.data_ptr<float>(), W,H,(float)decay);
}