updated to new gosl generation that should work on nvidia

emer · Jan 9, 2025 · 9d9bdc3 · 9d9bdc3
1 parent 20c1d3d
commit 9d9bdc3
Show file tree

Hide file tree

Showing 37 changed files with 1,149 additions and 2,641 deletions.
diff --git a/axon/act-net.go b/axon/act-net.go
diff --git a/axon/act-net.goal b/axon/act-net.goal
@@ -60,8 +60,12 @@ func (nt *Network) NewState(mode enums.Enum, testing bool) {
 	ctx := nt.Context()
 	nd := int(nix.NNeurons * ctx.NData)
 	ctx.NewState(mode, testing)
-	ToGPUCtxGlobal()
-	RunNewStateLayer(int(nix.NLayers))
+	for li := range nix.NLayers {
+		NewStateLayer(li)
+	}
+	ToGPULayers()
+	// ToGPUCtxGlobal()
+	// RunNewStateLayer(int(nix.NLayers))
 	RunNewStateNeuron(nd)
 	RunInitGBuffsPath(int(nix.NPaths))
 	// note: not completed until run cycles
@@ -272,7 +276,7 @@ func CyclePost(i uint32) { //gosl:kernel
 }
 
 // CycleInc is the kernel over 1 call to increment the cycle counter.
-func CycleInc(i uint32) { //gosl:kernel
+func CycleInc(i uint32) { //gosl:kernel read-write:Ctx
 	if i != 0 {
 		return
 	}
@@ -295,7 +299,7 @@ func ApplyExtsNeuron(i uint32) { //gosl:kernel
 
 // NewStateLayer is the kernel over Layers (not Data)
 // which does new state on pools as well.
-func NewStateLayer(li uint32) { //gosl:kernel
+func NewStateLayer(li uint32) { // note: not a kernel at this point
 	ctx := GetCtx(0)
 	Layers[li].NewStateLayer(ctx)
 }
@@ -376,7 +380,7 @@ func MinusPhasePost(li uint32) { //gosl:kernel
 }
 
 // PlusPhaseStartContext is the kernel over 1 call to call PlusPhaseStart on context.
-func PlusPhaseStartContext(i uint32) { //gosl:kernel
+func PlusPhaseStartContext(i uint32) { //gosl:kernel read-write:Ctx
 	if i != 0 {
 		return
 	}

diff --git a/axon/context.go b/axon/context.go
@@ -117,12 +117,14 @@ func (ctx *Context) DataIndex(idx uint32) uint32 {
 }
 
 // CycleInc increments at the cycle level
+//
+//gosl:pointer-receiver
 func (ctx *Context) CycleInc() {
 	ctx.PhaseCycle++
 	ctx.Cycle++
 	ctx.CyclesTotal++
 	ctx.Time += ctx.TimePerCycle
-	ctx.RandCounter.Add(uint32(RandFunIndexN))
+	// ctx.RandCounter.Add(uint32(RandFunIndexN))
 }
 
 // SlowInc increments the Slow counter and returns true if time

diff --git a/axon/gosl.go b/axon/gosl.go
diff --git a/axon/shaders/ApplyExtsNeuron.wgsl b/axon/shaders/ApplyExtsNeuron.wgsl
@@ -25,7 +25,7 @@ var<storage, read> RecvPathIxs: array<u32>;
 var<storage, read> PathRecvCon: array<u32>;
 @group(1) @binding(4)
 var<storage, read> RecvSynIxs: array<u32>;
-// // Ctx is the current context state (one only). 
+// // Ctx is the current context state (one only). This is read-only except in // specific kernels. 
 @group(2) @binding(0)
 var<storage, read_write> Ctx: array<Context>;
 @group(2) @binding(1)
@@ -78,10 +78,10 @@ fn Index3D(s0: u32, s1: u32, s2: u32, i0: u32, i1: u32, i2: u32) -> u32 {
 //////// import: "vars.go"
 
 //////// import: "act-layer.go"
-fn LayerParams_ApplyExtFlags(ly: ptr<function,LayerParams>, clearMask: ptr<function,NeuronFlags>,setMask: ptr<function,NeuronFlags>, toTarg: ptr<function,bool>) {
+fn LayerParams_ApplyExtFlags(ly: LayerParams, clearMask: ptr<function,NeuronFlags>,setMask: ptr<function,NeuronFlags>, toTarg: ptr<function,bool>) {
 	*clearMask = NeuronHasExt | NeuronHasTarg | NeuronHasCmpr;
 	*toTarg = false;
-	switch ((*ly).Type) {
+	switch (ly.Type) {
 	case TargetLayer: {
 		*setMask = NeuronHasTarg;
 		*toTarg = true;
@@ -95,12 +95,12 @@ fn LayerParams_ApplyExtFlags(ly: ptr<function,LayerParams>, clearMask: ptr<funct
 	}
 	}return;
 }
-fn LayerParams_InitExt(ly: ptr<function,LayerParams>, ni: u32,di: u32) {
+fn LayerParams_InitExt(ly: LayerParams, ni: u32,di: u32) {
 	Neurons[Index3D(TensorStrides[70], TensorStrides[71], TensorStrides[72], u32(ni), u32(di), u32(Ext))] = 0.0;
 	Neurons[Index3D(TensorStrides[70], TensorStrides[71], TensorStrides[72], u32(ni), u32(di), u32(Target))] = 0.0;
 	NeuronClearFlag(NeuronHasExt|NeuronHasTarg|NeuronHasCmpr, ni, di);
 }
-fn LayerParams_ApplyExtValue(ly: ptr<function,LayerParams>, ni: u32,di: u32, val: f32) {
+fn LayerParams_ApplyExtValue(ly: LayerParams, ni: u32,di: u32, val: f32) {
 	if (val < 0) {
 		return;
 	}
@@ -116,27 +116,26 @@ fn LayerParams_ApplyExtValue(ly: ptr<function,LayerParams>, ni: u32,di: u32, val
 	NeuronClearFlag(clearMask, ni, di);
 	NeuronSetFlag(setMask, ni, di);
 }
-fn LayerParams_ApplyExtsNeuron(ly: ptr<function,LayerParams>, ni: u32,di: u32) {
-	var lni = ni - (*ly).Indexes.NeurSt; // layer-based
+fn LayerParams_ApplyExtsNeuron(ly: LayerParams, ni: u32,di: u32) {
+	var lni = ni - ly.Indexes.NeurSt; // layer-based
 	LayerParams_InitExt(ly, ni, di);
-	if (IsExtLayerType((*ly).Type)) {
-		var ei = (*ly).Indexes.ExtsSt + lni;
+	if (IsExtLayerType(ly.Type)) {
+		var ei = ly.Indexes.ExtsSt + lni;
 		LayerParams_ApplyExtValue(ly, ni, di, Exts[Index2D(TensorStrides[120], TensorStrides[121],
 		u32(ei), u32(di))]);
 	}
 }
 
 //////// import: "act-net.go"
 fn ApplyExtsNeuron(i: u32) { //gosl:kernel
-	var ctx = Ctx[0];
-	var ni = Context_ItemIndex(&ctx, i);
+	let ctx = Ctx[0];
+	var ni = Context_ItemIndex(ctx, i);
 	if (ni >= NetworkIxs[0].NNeurons) {
 		return;
 	}
-	var di = Context_DataIndex(&ctx, i);
+	var di = Context_DataIndex(ctx, i);
 	var li = NeuronIxs[Index2D(TensorStrides[10], TensorStrides[11], u32(ni), u32(NrnLayIndex))];
-	var layers=Layers[li]; LayerParams_ApplyExtsNeuron(&layers, ni, di);
-	Ctx[0] = ctx;
+	let layers=Layers[li]; LayerParams_ApplyExtsNeuron(layers, ni, di);
 }
 
 //////// import: "act-path.go"
@@ -438,11 +437,11 @@ struct Context { //types:add -setters
 	SlowCounter: i32,
 	RandCounter: RandCounter,
 }
-fn Context_ItemIndex(ctx: ptr<function,Context>, idx: u32) -> u32 {
-	return idx / (*ctx).NData;
+fn Context_ItemIndex(ctx: Context, idx: u32) -> u32 {
+	return idx / ctx.NData;
 }
-fn Context_DataIndex(ctx: ptr<function,Context>, idx: u32) -> u32 {
-	return idx % (*ctx).NData;
+fn Context_DataIndex(ctx: Context, idx: u32) -> u32 {
+	return idx % ctx.NData;
 }
 
 //////// import: "deep-layer.go"

diff --git a/axon/shaders/Beta1Neuron.wgsl b/axon/shaders/Beta1Neuron.wgsl
@@ -25,7 +25,7 @@ var<storage, read> RecvPathIxs: array<u32>;
 var<storage, read> PathRecvCon: array<u32>;
 @group(1) @binding(4)
 var<storage, read> RecvSynIxs: array<u32>;
-// // Ctx is the current context state (one only). 
+// // Ctx is the current context state (one only). This is read-only except in // specific kernels. 
 @group(2) @binding(0)
 var<storage, read_write> Ctx: array<Context>;
 @group(2) @binding(1)
@@ -78,21 +78,20 @@ fn Index3D(s0: u32, s1: u32, s2: u32, i0: u32, i1: u32, i2: u32) -> u32 {
 //////// import: "vars.go"
 
 //////// import: "act-layer.go"
-fn LayerParams_Beta1Neuron(ly: ptr<function,LayerParams>, ctx: ptr<function,Context>, ni: u32,di: u32) {
+fn LayerParams_Beta1Neuron(ly: LayerParams, ctx: Context, ni: u32,di: u32) {
 	Neurons[Index3D(TensorStrides[70], TensorStrides[71], TensorStrides[72], u32(ni), u32(di), u32(Beta1))] = Neurons[Index3D(TensorStrides[70], TensorStrides[71], TensorStrides[72], u32(ni), u32(di), u32(CaP))];
 }
 
 //////// import: "act-net.go"
 fn Beta1Neuron(i: u32) { //gosl:kernel
-	var ctx = Ctx[0];
-	var ni = Context_ItemIndex(&ctx, i);
+	let ctx = Ctx[0];
+	var ni = Context_ItemIndex(ctx, i);
 	if (ni >= NetworkIxs[0].NNeurons) {
 		return;
 	}
-	var di = Context_DataIndex(&ctx, i);
+	var di = Context_DataIndex(ctx, i);
 	var li = NeuronIxs[Index2D(TensorStrides[10], TensorStrides[11], u32(ni), u32(NrnLayIndex))];
-	var layers=Layers[li]; LayerParams_Beta1Neuron(&layers, &ctx, ni, di);
-	Ctx[0] = ctx;
+	let layers=Layers[li]; LayerParams_Beta1Neuron(layers, ctx, ni, di);
 }
 
 //////// import: "act-path.go"
@@ -388,11 +387,11 @@ struct Context { //types:add -setters
 	SlowCounter: i32,
 	RandCounter: RandCounter,
 }
-fn Context_ItemIndex(ctx: ptr<function,Context>, idx: u32) -> u32 {
-	return idx / (*ctx).NData;
+fn Context_ItemIndex(ctx: Context, idx: u32) -> u32 {
+	return idx / ctx.NData;
 }
-fn Context_DataIndex(ctx: ptr<function,Context>, idx: u32) -> u32 {
-	return idx % (*ctx).NData;
+fn Context_DataIndex(ctx: Context, idx: u32) -> u32 {
+	return idx % ctx.NData;
 }
 
 //////// import: "deep-layer.go"

diff --git a/axon/shaders/Beta2Neuron.wgsl b/axon/shaders/Beta2Neuron.wgsl
@@ -25,7 +25,7 @@ var<storage, read> RecvPathIxs: array<u32>;
 var<storage, read> PathRecvCon: array<u32>;
 @group(1) @binding(4)
 var<storage, read> RecvSynIxs: array<u32>;
-// // Ctx is the current context state (one only). 
+// // Ctx is the current context state (one only). This is read-only except in // specific kernels. 
 @group(2) @binding(0)
 var<storage, read_write> Ctx: array<Context>;
 @group(2) @binding(1)
@@ -78,21 +78,20 @@ fn Index3D(s0: u32, s1: u32, s2: u32, i0: u32, i1: u32, i2: u32) -> u32 {
 //////// import: "vars.go"
 
 //////// import: "act-layer.go"
-fn LayerParams_Beta2Neuron(ly: ptr<function,LayerParams>, ctx: ptr<function,Context>, ni: u32,di: u32) {
+fn LayerParams_Beta2Neuron(ly: LayerParams, ctx: Context, ni: u32,di: u32) {
 	Neurons[Index3D(TensorStrides[70], TensorStrides[71], TensorStrides[72], u32(ni), u32(di), u32(Beta2))] = Neurons[Index3D(TensorStrides[70], TensorStrides[71], TensorStrides[72], u32(ni), u32(di), u32(CaP))];
 }
 
 //////// import: "act-net.go"
 fn Beta2Neuron(i: u32) { //gosl:kernel
-	var ctx = Ctx[0];
-	var ni = Context_ItemIndex(&ctx, i);
+	let ctx = Ctx[0];
+	var ni = Context_ItemIndex(ctx, i);
 	if (ni >= NetworkIxs[0].NNeurons) {
 		return;
 	}
-	var di = Context_DataIndex(&ctx, i);
+	var di = Context_DataIndex(ctx, i);
 	var li = NeuronIxs[Index2D(TensorStrides[10], TensorStrides[11], u32(ni), u32(NrnLayIndex))];
-	var layers=Layers[li]; LayerParams_Beta2Neuron(&layers, &ctx, ni, di);
-	Ctx[0] = ctx;
+	let layers=Layers[li]; LayerParams_Beta2Neuron(layers, ctx, ni, di);
 }
 
 //////// import: "act-path.go"
@@ -388,11 +387,11 @@ struct Context { //types:add -setters
 	SlowCounter: i32,
 	RandCounter: RandCounter,
 }
-fn Context_ItemIndex(ctx: ptr<function,Context>, idx: u32) -> u32 {
-	return idx / (*ctx).NData;
+fn Context_ItemIndex(ctx: Context, idx: u32) -> u32 {
+	return idx / ctx.NData;
 }
-fn Context_DataIndex(ctx: ptr<function,Context>, idx: u32) -> u32 {
-	return idx % (*ctx).NData;
+fn Context_DataIndex(ctx: Context, idx: u32) -> u32 {
+	return idx % ctx.NData;
 }
 
 //////// import: "deep-layer.go"