ton/IronpenGPU
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

version 0.0.5-alpha

Browse Source
This commit is contained in:
ton
2023-08-25 09:40:43 +07:00
parent 0af39ee09e
commit c74eea9cdf
14 changed files with 4155 additions and 150 deletions
Download Patch File Download Diff File Expand all files Collapse all files

96
src/forward.jl
View File

@@ -55,6 +55,7 @@ function (kfn::kfn_1)(input::AbstractArray)
# project 3D kfn zit into 4D lif zit
i1, i2, i3, i4 = size(kfn.lif_zit)
kfn.lif_zit .= reshape(kfn.zit, (i1, i2, 1, i4)) .* kfn.lif_arrayProjection4d
kfn.lif_exInType .= kfn.exInType .* kfn.lif_arrayProjection4d
lifForward( kfn.lif_zit,
kfn.lif_wRec,
@@ -70,13 +71,16 @@ function (kfn::kfn_1)(input::AbstractArray)
kfn.lif_gammaPd,
kfn.lif_firingCounter,
kfn.lif_recSignal,
kfn.lif_subscription,
kfn.lif_exInType,
kfn.lif_neuronInactivityCounter,
kfn.lif_synapticInactivityCounter,
)
end
@async begin
# project 3D kfn zit into 4D alif zit
i1, i2, i3, i4 = size(kfn.alif_zit)
kfn.alif_zit .= reshape(kfn.zit, (i1, i2, 1, i4)) .* kfn.alif_arrayProjection4d
kfn.alif_exInType .= kfn.exInType .* kfn.alif_arrayProjection4d
alifForward(kfn.alif_zit,
kfn.alif_wRec,
@@ -92,12 +96,15 @@ function (kfn::kfn_1)(input::AbstractArray)
kfn.alif_gammaPd,
kfn.alif_firingCounter,
kfn.alif_recSignal,
kfn.alif_subscription,
kfn.alif_exInType,
kfn.alif_neuronInactivityCounter,
kfn.alif_synapticInactivityCounter,
kfn.alif_epsilonRecA,
kfn.alif_a,
kfn.alif_avth,
kfn.alif_beta,
kfn.alif_rho,)
kfn.alif_rho,
)
end
end
@@ -132,13 +139,11 @@ function (kfn::kfn_1)(input::AbstractArray)
kfn.on_gammaPd,
kfn.on_firingCounter,
kfn.on_recSignal,
kfn.on_subscription,
)
# get on_zt4d to on_zt
kfn.on_zt .= reduce(max, kfn.on_zt4d, dims=(1,2))
logit = reshape(kfn.on_zt, (size(input, 1), :))
# error("DEBUG -> kfn forward")
return logit,
kfn.zit
end
@@ -158,7 +163,9 @@ function lifForward( zit::CuArray,
gammaPd::CuArray,
firingCounter::CuArray,
recSignal::CuArray,
subscription::CuArray,
exInType::CuArray,
neuronInactivityCounter::CuArray,
synapticInactivityCounter::CuArray,
)
kernel = @cuda launch=false lifForward( zit,
@@ -175,7 +182,9 @@ function lifForward( zit::CuArray,
gammaPd,
firingCounter,
recSignal,
subscription,
exInType,
neuronInactivityCounter,
synapticInactivityCounter,
GeneralUtils.linear_to_cartesian,
)
config = launch_configuration(kernel.fun)
@@ -206,7 +215,9 @@ function lifForward( zit::CuArray,
gammaPd,
firingCounter,
recSignal,
subscription,
exInType,
neuronInactivityCounter,
synapticInactivityCounter,
GeneralUtils.linear_to_cartesian; threads, blocks)
end
end
@@ -226,7 +237,9 @@ function lifForward( zit,
gammaPd,
firingCounter,
recSignal,
subscription,
exInType,
neuronInactivityCounter,
synapticInactivityCounter,
linear_to_cartesian,
)
i = (blockIdx().x - 1) * blockDim().x + threadIdx().x # gpu threads index
@@ -247,7 +260,8 @@ function lifForward( zit,
epsilonRec[i1,i2,i3,i4] = (alpha[i1,i2,i3,i4] * epsilonRec[i1,i2,i3,i4])
else # refractory period is inactive
recSignal[i1,i2,i3,i4] = wRec[i1,i2,i3,i4] * zit[i1,i2,i3,i4]
recSignal[i1,i2,i3,i4] = wRec[i1,i2,i3,i4] * zit[i1,i2,i3,i4] *
exInType[i1,i2,i3,i4]
vt[i1,i2,i3,i4] = (alpha[i1,i2,i3,i4] * vt[i1,i2,i3,i4]) +
sum(@view(recSignal[:,:,i3,i4]))
@@ -257,8 +271,12 @@ function lifForward( zit,
refractoryCounter[i1,i2,i3,i4] = refractoryDuration[i1,i2,i3,i4]
firingCounter[i1,i2,i3,i4] += 1
vt[i1,i2,i3,i4] = vRest[i1,i2,i3,i4]
# reset counter if neuron fires
neuronInactivityCounter[i1,i2,i3,i4] = 10000
else
zt[i1,i2,i3,i4] = 0
neuronInactivityCounter[i1,i2,i3,i4] -= 1
end
# compute phi, there is a difference from lif formula
@@ -267,7 +285,18 @@ function lifForward( zit,
# compute epsilonRec
epsilonRec[i1,i2,i3,i4] = (alpha[i1,i2,i3,i4] * epsilonRec[i1,i2,i3,i4]) +
(zit[i1,i2,i3,i4] * subscription[i1,i2,i3,i4])
(zit[i1,i2,i3,i4] * !iszero(wRec[i1,i2,i3,i4]))
# !iszero indicates synaptic subscription
# count synaptic inactivity
if !iszero(wRec[i1,i2,i3,i4]) # check if this is wRec subscription
if !iszero(zit[i1,i2,i3,i4]) # synapse is active, reset counter
synapticInactivityCounter[i1,i2,i3,i4] = 10000
else # synapse is inactive, counting
synapticInactivityCounter[i1,i2,i3,i4] -= 1
end
end
end
end
return nothing
@@ -288,7 +317,9 @@ function alifForward( zit::CuArray,
gammaPd::CuArray,
firingCounter::CuArray,
recSignal::CuArray,
subscription::CuArray,
exInType::CuArray,
neuronInactivityCounter::CuArray,
synapticInactivityCounter::CuArray,
epsilonRecA::CuArray,
a::CuArray,
avth::CuArray,
@@ -310,7 +341,9 @@ function alifForward( zit::CuArray,
gammaPd,
firingCounter,
recSignal,
subscription,
exInType,
neuronInactivityCounter,
synapticInactivityCounter,
epsilonRecA,
a,
avth,
@@ -345,7 +378,9 @@ function alifForward( zit::CuArray,
gammaPd,
firingCounter,
recSignal,
subscription,
exInType,
neuronInactivityCounter,
synapticInactivityCounter,
epsilonRecA,
a,
avth,
@@ -370,7 +405,9 @@ function alifForward( zit,
gammaPd,
firingCounter,
recSignal,
subscription,
exInType,
neuronInactivityCounter,
synapticInactivityCounter,
epsilonRecA,
a,
avth,
@@ -404,7 +441,8 @@ function alifForward( zit,
avth[i1,i2,i3,i4] = vth[i1,i2,i3,i4] + (beta[i1,i2,i3,i4] * a[i1,i2,i3,i4])
else # refractory period is inactive
recSignal[i1,i2,i3,i4] = zit[i1,i2,i3,i4] * wRec[i1,i2,i3,i4]
recSignal[i1,i2,i3,i4] = wRec[i1,i2,i3,i4] * zit[i1,i2,i3,i4] *
exInType[i1,i2,i3,i4]
vt[i1,i2,i3,i4] = (alpha[i1,i2,i3,i4] * vt[i1,i2,i3,i4]) +
sum(@view(recSignal[:,:,i3,i4]))
@@ -418,9 +456,11 @@ function alifForward( zit,
firingCounter[i1,i2,i3,i4] += 1
vt[i1,i2,i3,i4] = vRest[i1,i2,i3,i4]
a[i1,i2,i3,i4] = (rho[i1,i2,i3,i4] * a[i1,i2,i3,i4]) + 1
neuronInactivityCounter[i1,i2,i3,i4] = 10000
else
zt[i1,i2,i3,i4] = 0
a[i1,i2,i3,i4] = (rho[i1,i2,i3,i4] * a[i1,i2,i3,i4])
neuronInactivityCounter[i1,i2,i3,i4] -= 1
end
# compute phi, there is a difference from alif formula
@@ -429,11 +469,20 @@ function alifForward( zit,
# compute epsilonRec
epsilonRec[i1,i2,i3,i4] = (alpha[i1,i2,i3,i4] * epsilonRec[i1,i2,i3,i4]) +
(zit[i1,i2,i3,i4] * subscription[i1,i2,i3,i4])
(zit[i1,i2,i3,i4] * !iszero(wRec[i1,i2,i3,i4]))
# compute epsilonRecA use eq.26
epsilonRecA[i1,i2,i3,i4] = (rho[i1,i2,i3,i4] *
(phi[i1,i2,i3,i4] * epsilonRec[i1,i2,i3,i4])) +
(zit[i1,i2,i3,i4] * subscription[i1,i2,i3,i4])
(phi[i1,i2,i3,i4] * epsilonRec[i1,i2,i3,i4])) +
(zit[i1,i2,i3,i4] * !iszero(wRec[i1,i2,i3,i4]))
# count synaptic inactivity
if !iszero(wRec[i1,i2,i3,i4]) # check if this is wRec subscription
if !iszero(zit[i1,i2,i3,i4]) # synapse is active, reset counter
synapticInactivityCounter[i1,i2,i3,i4] = 10000
else # synapse is inactive, counting
synapticInactivityCounter[i1,i2,i3,i4] -= 1
end
end
end
end
return nothing
@@ -454,7 +503,6 @@ function onForward( zit::CuArray,
gammaPd::CuArray,
firingCounter::CuArray,
recSignal::CuArray,
subscription::CuArray,
)
kernel = @cuda launch=false onForward( zit,
@@ -471,7 +519,6 @@ function onForward( zit::CuArray,
gammaPd,
firingCounter,
recSignal,
subscription,
GeneralUtils.linear_to_cartesian,
)
config = launch_configuration(kernel.fun)
@@ -501,7 +548,6 @@ function onForward( zit::CuArray,
gammaPd,
firingCounter,
recSignal,
subscription,
GeneralUtils.linear_to_cartesian; threads, blocks)
end
end
@@ -521,7 +567,6 @@ function onForward( zit,
gammaPd,
firingCounter,
recSignal,
subscription,
linear_to_cartesian,
)
i = (blockIdx().x - 1) * blockDim().x + threadIdx().x # gpu threads index
@@ -556,11 +601,12 @@ function onForward( zit,
end
# compute phi, there is a difference from on formula
phi[i1,i2,i3,i4] = (gammaPd[i1,i2,i3,i4] / vth[i1,i2,i3,i4]) * max(0, 1 - ((vt[i1,i2,i3,i4] - vth[i1,i2,i3,i4]) / vth[i1,i2,i3,i4]))
phi[i1,i2,i3,i4] = (gammaPd[i1,i2,i3,i4] / vth[i1,i2,i3,i4]) *
max(0, 1 - ((vt[i1,i2,i3,i4] - vth[i1,i2,i3,i4]) / vth[i1,i2,i3,i4]))
# compute epsilonRec
epsilonRec[i1,i2,i3,i4] = (alpha[i1,i2,i3,i4] * epsilonRec[i1,i2,i3,i4]) +
(zit[i1,i2,i3,i4] * subscription[i1,i2,i3,i4])
(zit[i1,i2,i3,i4] * !iszero(wOut[i1,i2,i3,i4]))
end
end
return nothing