publish buildable source (threaded + tile-sorted kernel)

Browse files

Files changed (9) hide show

build.toml +22 -0
flake.nix +17 -0
load_local.py +35 -0
local_bind.cpp +11 -0
physarum_csrc/physarum_cpu.cpp +225 -0
tests/test_physarum.py +28 -0
torch-ext/physarum/__init__.py +85 -0
torch-ext/torch_binding.cpp +15 -0
torch-ext/torch_binding.h +7 -0

build.toml ADDED Viewed

	@@ -0,0 +1,22 @@

+[general]
+name = "physarum"
+version = 1
+edition = 5
+license = "Apache-2.0"
+backends = [
+  "cpu",
+]
+[general.hub]
+repo-id = "phanerozoic/physarum"
+[torch]
+src = [
+  "torch-ext/torch_binding.cpp",
+  "torch-ext/torch_binding.h",
+]
+[kernel.physarum]
+backend = "cpu"
+depends = ["torch"]
+src = ["physarum_csrc/physarum_cpu.cpp"]

flake.nix ADDED Viewed

	@@ -0,0 +1,17 @@

+{
+  description = "Flake for physarum kernel";
+  inputs = {
+    builder.url = "github:huggingface/kernels";
+  };
+  outputs =
+    {
+      self,
+      builder,
+    }:
+    builder.lib.genKernelFlakeOutputs {
+      inherit self;
+      path = ./.;
+    };
+}

load_local.py ADDED Viewed

	@@ -0,0 +1,35 @@

+"""Load physarum natively on hosts without a published variant (e.g. a Pi).
+JIT-builds the local source with `torch.utils.cpp_extension` and exposes the
+identical `physarum` module API.
+    import load_local
+    physarum = load_local.load()
+    sim = physarum.Physarum(width=1024, height=1024, agents=200000)
+"""
+import os
+import sys
+import types
+_cached = None
+def load(verbose=False):
+    global _cached
+    if _cached is not None:
+        return _cached
+    from torch.utils.cpp_extension import load as _jit
+    root = os.path.dirname(os.path.abspath(__file__))
+    ext = _jit(name="physarum_ext",
+               sources=[os.path.join(root, "local_bind.cpp"),
+                        os.path.join(root, "physarum_csrc", "physarum_cpu.cpp")],
+               extra_cflags=["-O3"],
+               verbose=verbose)
+    ops_mod = types.ModuleType("physarum._ops")
+    ops_mod.ops = type("_Ops", (), {"physarum_step": staticmethod(ext.physarum_step)})()
+    sys.modules["physarum._ops"] = ops_mod
+    sys.path.insert(0, os.path.join(root, "torch-ext"))
+    import physarum
+    _cached = physarum
+    return physarum

local_bind.cpp ADDED Viewed

	@@ -0,0 +1,11 @@

+// JIT binding for load_local.py; the published build registers via torch_binding.cpp.
+#include <torch/extension.h>
+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);
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("physarum_step", &physarum_step);
+}

physarum_csrc/physarum_cpu.cpp ADDED Viewed

	@@ -0,0 +1,225 @@

+// 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);
+}

tests/test_physarum.py ADDED Viewed

	@@ -0,0 +1,28 @@

+import torch
+import kernels
+physarum = kernels.get_kernel("phanerozoic/physarum", version=1, trust_remote_code=True)
+def test_network_develops():
+    sim = physarum.Physarum(width=256, height=256, agents=20000, seed=1)
+    sim.step(120)
+    t = sim.trail
+    assert torch.isfinite(t).all()
+    assert float(t.max()) > 0.0
+    # a network is high-contrast, not a uniform wash
+    assert float(t.std()) > 0.4 * float(t.mean())
+def test_image_shape():
+    sim = physarum.Physarum(width=128, height=128, agents=5000, seed=2)
+    sim.step(60)
+    img = sim.image()
+    assert tuple(img.shape) == (128, 128, 3)
+    assert img.dtype == torch.uint8
+def test_deterministic():
+    a = physarum.Physarum(width=128, height=128, agents=5000, seed=7); a.step(40)
+    b = physarum.Physarum(width=128, height=128, agents=5000, seed=7); b.step(40)
+    assert torch.equal(a.trail, b.trail)

torch-ext/physarum/__init__.py ADDED Viewed

	@@ -0,0 +1,85 @@

+"""physarum: a Physarum (slime-mould) transport-network simulation on the CPU.
+Agents diffuse a chemoattractant trail and steer up its gradient; the emergent
+network is the same structure a real slime mould forms solving a maze. The
+agent step (three sensor look-ups, steering, motion, deposition) and the trail
+diffusion vectorise with NEON on aarch64; the agent step is threaded across
+cores (set torch.set_num_threads) and works on tile-sorted agents for cache
+locality, with a portable scalar fallback. The trail is an ordinary CPU tensor
+you read and render each step.
+    from kernels import get_kernel
+    physarum = get_kernel("phanerozoic/physarum", version=1, trust_remote_code=True)
+    sim = physarum.Physarum(width=1024, height=1024, agents=200_000)
+    for _ in range(400):
+        sim.step()
+    img = sim.image()                 # [H, W, 3] uint8, inferno colormap
+"""
+import torch
+from ._ops import ops
+# inferno-ish colormap stops (black -> violet -> magenta -> orange -> pale yellow)
+_CMAP = torch.tensor(
+    [[0.0, 0.0, 0.0], [0.12, 0.02, 0.28], [0.52, 0.08, 0.42],
+     [0.88, 0.28, 0.16], [0.99, 0.68, 0.12], [1.0, 1.0, 0.86]],
+    dtype=torch.float32)
+class Physarum:
+    """A Physarum transport-network simulation.
+    Attributes:
+        trail: [H, W] float32 chemoattractant field (read/modify/render freely).
+        ax, ay, ah: [N] float32 agent x, y, heading.
+    Parameters follow Jones (2010): sensor offset `sense_dist` (px), sensor
+    angle `sense_ang` (rad, 22.5 deg default), rotation `turn` (rad, 45 deg),
+    step `speed` (px), `deposit` per agent per step, and per-step `decay` of the
+    field (after a 3-tap diffusion).
+    """
+    def __init__(self, width=1024, height=1024, agents=200_000, seed=0,
+                 sense_dist=9.0, sense_ang=0.3927, turn=0.7854, speed=1.0,
+                 deposit=5.0, decay=0.88):
+        g = torch.Generator().manual_seed(seed)
+        self.W, self.H, self.N = int(width), int(height), int(agents)
+        self.trail = torch.zeros(self.H, self.W, dtype=torch.float32)
+        self._tmp = torch.zeros(self.H, self.W, dtype=torch.float32)
+        self.ax = (torch.rand(self.N, generator=g) * self.W).contiguous()
+        self.ay = (torch.rand(self.N, generator=g) * self.H).contiguous()
+        self.ah = (torch.rand(self.N, generator=g) * 6.2831853).contiguous()
+        self.params = dict(sense_dist=float(sense_dist), sense_ang=float(sense_ang),
+                           turn=float(turn), speed=float(speed),
+                           deposit=float(deposit), decay=float(decay))
+        self.base_seed = int(seed)
+        self.t = 0
+    def step(self, n: int = 1) -> torch.Tensor:
+        """Advance the simulation `n` steps in place; returns the trail field."""
+        p = self.params
+        for _ in range(n):
+            ops.physarum_step(self.trail, self._tmp, self.ax, self.ay, self.ah,
+                              self.base_seed + self.t, p["sense_dist"], p["sense_ang"],
+                              p["turn"], p["speed"], p["deposit"], p["decay"])
+            self.t += 1
+        return self.trail
+    def image(self, gamma: float = 0.62, percentile: float = 0.995) -> torch.Tensor:
+        """Colormap the trail field to an [H, W, 3] uint8 image.
+        Normalised by a high percentile (not the max) so the network shows
+        rather than a few hot cells crushing it to black.
+        """
+        t = self.trail.flatten()
+        stride = max(1, t.numel() // 1_000_000)
+        mx = torch.quantile(t[::stride], percentile).clamp(min=1e-6)
+        v = (self.trail / mx).clamp(0, 1).pow(gamma) * 5.0
+        s = v.floor().clamp(max=4).long()
+        f = (v - s).unsqueeze(-1)
+        rgb = (_CMAP[s] + (_CMAP[s + 1] - _CMAP[s]) * f).clamp(0, 1)
+        return (rgb * 255 + 0.5).to(torch.uint8)
+__all__ = ["Physarum"]

torch-ext/torch_binding.cpp ADDED Viewed

	@@ -0,0 +1,15 @@

+#include <torch/library.h>
+#include "registration.h"
+#include "torch_binding.h"
+TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
+  ops.def("physarum_step(Tensor! trail, Tensor! tmp, Tensor! ax, Tensor! ay, "
+          "Tensor! ah, int seed, float sense_dist, float sense_ang, float turn, "
+          "float speed, float deposit, float decay) -> ()");
+#if defined(CPU_KERNEL)
+  ops.impl("physarum_step", torch::kCPU, &physarum_step);
+#endif
+}
+REGISTER_EXTENSION(TORCH_EXTENSION_NAME)

torch-ext/torch_binding.h ADDED Viewed

	@@ -0,0 +1,7 @@

+#pragma once
+#include <torch/torch.h>
+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);