🔨 refactor: handle float16 along float on GPU (#120)

owkin · May 12, 2024 · 52ab4df · 52ab4df
1 parent a9d176c
commit 52ab4df
Show file tree

Hide file tree

Showing 126 changed files with 16,078 additions and 1,557 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -4,6 +4,7 @@ All notable changes to this project will be documented in this file.
 
 ## [unreleased]
 
+🔨 **refactor:** handle float16 along float on GPU ([#120](https://github.com/owkin/GrAIdient/pull/120))\
 🚀 **perf:** copy & generate weights faster ([119](https://github.com/owkin/GrAIdient/pull/119))\
 🚀 **perf:** Convolution2D ([118](https://github.com/owkin/GrAIdient/pull/118))\
 🪜 **feat:** LayerCAM2D -> VQGrad2D, LayerCAMSeq -> VQGradSeq ([#117](https://github.com/owkin/GrAIdient/pull/117))\

diff --git a/Package.swift b/Package.swift
@@ -7,7 +7,7 @@ import PackageDescription
 let package = Package(
     name: "GrAIdient",
     platforms: [
-        .macOS(.v10_15)
+        .macOS(.v13)
     ],
     products: [
         .library(

diff --git a/Sources/GrAIdient/Core/Function/Activation.swift b/Sources/GrAIdient/Core/Function/Activation.swift
@@ -307,8 +307,8 @@ open class ActivationFunction: Codable
     ///     - deviceID: GPU device where to execute the operation.
     ///
     private func _forwardGPU(
-        tmp: MetalBuffer<Float>,
-        outs: MetalBuffer<Float>,
+        tmp: FloatBuffer,
+        outs: FloatBuffer,
         deviceID: Int)
     {
         let nbElems = outs.nbElems
@@ -335,8 +335,9 @@ open class ActivationFunction: Codable
         let nbElems = layer.outs.nbElems
         if layer._tmp == nil
         {
-            layer._tmp = MetalPrivateBuffer<Float>(
-                nbElems, deviceID: layer.deviceID)
+            layer._tmp = FloatBuffer(
+                nbElems: nbElems, deviceID: layer.deviceID
+            )
         }
         _forwardGPU(
             tmp: layer._tmp,
@@ -355,7 +356,7 @@ open class ActivationFunction: Codable
         let nbElems = layer.outs.nbElems
         if layer._tmp == nil
         {
-            layer._tmp = MetalPrivateBuffer<Float>(
+            layer._tmp = FloatBuffer(nbElems: 
                 nbElems, deviceID: layer.deviceID)
         }
         _forwardGPU(
@@ -375,8 +376,9 @@ open class ActivationFunction: Codable
         let nbElems = layer.outs.nbElems
         if layer._tmp == nil
         {
-            layer._tmp = MetalPrivateBuffer<Float>(
-                nbElems, deviceID: layer.deviceID)
+            layer._tmp = FloatBuffer(
+                nbElems: nbElems, deviceID: layer.deviceID
+            )
         }
         _forwardGPU(
             tmp: layer._tmp,
@@ -394,8 +396,8 @@ open class ActivationFunction: Codable
     ///     - deviceID: GPU device where to execute the operation.
     ///
     private func _backwardGPU(
-        tmp: MetalBuffer<Float>,
-        delta: MetalBuffer<Float>,
+        tmp: FloatBuffer,
+        delta: FloatBuffer,
         deviceID: Int)
     {
         let nbElems = delta.nbElems

diff --git a/Sources/GrAIdient/Core/Layer/LayerInput.swift b/Sources/GrAIdient/Core/Layer/LayerInput.swift
@@ -105,14 +105,13 @@ class InputBuffers<T: Layer>
 {
     /// The link to the layer.
     unowned let _layer: T
-    /// Number of elements in the different buffers.
-    let nbElems: Int
-    /// GPU device where the buffers are sent.
-    let deviceID: Int
 
-    var _m: MetalBuffer<Float>! = nil
-    var _v: MetalBuffer<Float>! = nil
-    var _vHat: MetalBuffer<Float>! = nil
+    /// Momentum buffer.
+    public let m: FloatBuffer
+    /// Velocity buffer.
+    public let v: FloatBuffer
+    /// Velocity normalized buffer.
+    public let vHat: FloatBuffer
 
     ///
     /// Create a container of buffers.
@@ -127,51 +126,16 @@ class InputBuffers<T: Layer>
          deviceID: Int)
     {
         _layer = layer
-        self.nbElems = nbElems
-        self.deviceID = deviceID
-    }
-
-    /// Momentum buffer.
-    var m: MetalBuffer<Float>
-    {
-        get {
-            if _m == nil
-            {
-                _m = MetalPrivateBuffer<Float>(nbElems, deviceID: deviceID)
-            }
-            return _m
-        }
-    }
-
-    /// Velocity buffer.
-    var v: MetalBuffer<Float>
-    {
-        get {
-            if _v == nil
-            {
-                _v = MetalPrivateBuffer<Float>(nbElems, deviceID: deviceID)
-            }
-            return _v
-        }
-    }
-
-    /// Velocity normalized buffer.
-    var vHat: MetalBuffer<Float>
-    {
-        get {
-            if _vHat == nil
-            {
-                _vHat = MetalPrivateBuffer<Float>(nbElems, deviceID: deviceID)
-            }
-            return _vHat
-        }
+        m = FloatBuffer(nbElems: nbElems, deviceID: deviceID)
+        v = FloatBuffer(nbElems: nbElems, deviceID: deviceID)
+        vHat = FloatBuffer(nbElems: nbElems, deviceID: deviceID)
     }
 
     /// Clean the momentum..., preserving the weights.
     func reset()
     {
-        _m = nil
-        _v = nil
-        _vHat = nil
+        m.reset()
+        v.reset()
+        vHat.reset()
     }
 }