Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

🐛 fix: use buffers for neuron selection in SelectNeurons1D #50

Merged
merged 1 commit into from
Feb 6, 2023
Merged

🐛 fix: use buffers for neuron selection in SelectNeurons1D #50

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions CHANGELOG.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ All notable changes to this project will be documented in this file.

## [unreleased]

🐛 **fix:** use buffers for neuron selection in SelectNeurons1D ([#50](https://github.com/owkin/GrAIdient/pull/50))\
🪜 **feat:** Softmax1D, DotProduct1D & Constant1D ([#49](https://github.com/owkin/GrAIdient/pull/49))\
🪜 **feat:** remove activation from layer ([#47](https://github.com/owkin/GrAIdient/pull/47))\
🎉 **refactor:** logo ([#46](https://github.com/owkin/GrAIdient/pull/46))\
Expand Down
84 changes: 59 additions & 25 deletions Sources/GrAIdient/Layer1D/SelectNeurons1D.swift
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,17 @@ public class SelectNeurons1D: Layer1D
/// List of coefficients to scale each selected neuron.
let _coeffs: [Double]

///
/// Indices of selected neurons.
/// Shape ~ (nbNeurons,).
///
var _neuronsBuffer: MetalPrivateBuffer<UInt32>! = nil
///
/// Coefficients of selected neurons.
/// Shape ~ (nbNeurons,).
///
var _coeffsBuffer: MetalPrivateBuffer<Float>! = nil

private enum Keys: String, CodingKey
{
case neurons
Expand Down Expand Up @@ -110,6 +121,50 @@ public class SelectNeurons1D: Layer1D
return layer
}

///
/// Clean state resources in the GPU execution context.
///
/// We clean the neurons' state (forward and backward).
///
public override func resetKernelGPU()
{
super.resetKernelGPU()
_neuronsBuffer = nil
_coeffsBuffer = nil
}

///
/// Initialize state resources in the GPU execution context.
///
/// We initialize the neurons' forward state.
///
public override func checkStateForwardGPU(batchSize: Int) throws
{
try super.checkStateForwardGPU(batchSize: batchSize)

if _neuronsBuffer == nil
{
_neuronsBuffer = MetalPrivateBuffer<UInt32>(
nbNeurons, deviceID: deviceID
)
_coeffsBuffer = MetalPrivateBuffer<Float>(
nbNeurons, deviceID: deviceID
)

let neuronsPtr = _neuronsBuffer.shared.buffer
let coeffsPtr = _coeffsBuffer.shared.buffer

for (num, neuron) in _neurons.enumerated()
{
neuronsPtr[num] = UInt32(neuron)
coeffsPtr[num] = Float(_coeffs[num])
}

MetalKernel.get.upload([_neuronsBuffer])
MetalKernel.get.upload([_coeffsBuffer])
}
}

///
/// Apply the forward pass of the Gradient Checking in CPU execution context.
///
Expand Down Expand Up @@ -189,25 +244,15 @@ public class SelectNeurons1D: Layer1D
let pNbNeurons: [UInt32] = [UInt32(nbNeurons)]
let pNbNeuronsPrev: [UInt32] = [UInt32(layerPrev.nbNeurons)]
let pNbBatch: [UInt32] = [UInt32(batchSize)]
var pNeurons = [UInt32]()
for neuron in _neurons
{
pNeurons.append(UInt32(neuron))
}
var pCoeffs = [Float]()
for coeff in _coeffs
{
pCoeffs.append(Float(coeff))
}

let command = MetalKernel.get.createCommand(
"selectNeurons1DForward", deviceID: deviceID
)
command.setBuffer(layerPrev.outs.metal, atIndex: 0)
command.setBytes(pNbNeurons, atIndex: 1)
command.setBytes(pNbNeuronsPrev, atIndex: 2)
command.setBytes(pNeurons, atIndex: 3)
command.setBytes(pCoeffs, atIndex: 4)
command.setBuffer(_neuronsBuffer.metal, atIndex: 3)
command.setBuffer(_coeffsBuffer.metal, atIndex: 4)
command.setBytes(pNbBatch, atIndex: 5)
command.setBuffer(outs.metal, atIndex: 6)

Expand Down Expand Up @@ -260,19 +305,8 @@ public class SelectNeurons1D: Layer1D
let pNbNeurons: [UInt32] = [UInt32(nbNeurons)]
let pNbNeuronsPrev: [UInt32] = [UInt32(layerPrev.nbNeurons)]
let pNbBatch: [UInt32] = [UInt32(batchSize)]
var pNeurons = [UInt32]()
for neuron in _neurons
{
pNeurons.append(UInt32(neuron))
}
var pCoeffs = [Float]()
for coeff in _coeffs
{
pCoeffs.append(Float(coeff))
}

var command: MetalCommand

if layerPrev.dirty
{
let nbElems = layerPrev.delta.nbElems
Expand All @@ -294,8 +328,8 @@ public class SelectNeurons1D: Layer1D
command.setBuffer(delta.metal, atIndex: 0)
command.setBytes(pNbNeurons, atIndex: 1)
command.setBytes(pNbNeuronsPrev, atIndex: 2)
command.setBytes(pNeurons, atIndex: 3)
command.setBytes(pCoeffs, atIndex: 4)
command.setBuffer(_neuronsBuffer.metal, atIndex: 3)
command.setBuffer(_coeffsBuffer.metal, atIndex: 4)
command.setBytes(pNbBatch, atIndex: 5)
command.setBuffer(layerPrev.delta.metal, atIndex: 6)

Expand Down