From 63934a9a552cbb255845190a079ebd90f48892a8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Jean-Fran=C3=A7ois=20Reboud?=
 <jean-francois.reboud@owkin.com>
Date: Fri, 8 Dec 2023 10:00:55 +0100
Subject: [PATCH] =?UTF-8?q?=F0=9F=90=9B=20fix:=20run=20on=20Apple=20Silico?=
 =?UTF-8?q?n=20(#110)?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 CHANGELOG.md                              |  1 +
 Sources/GrAIdient/Metal/MetalKernel.swift | 11 +++-
 Sources/GrAIdient/Utils/Image.swift       | 72 +++++++++--------------
 Tests/GrAIExamples/Base/Utils.swift       |  2 +-
 Tests/GrAITorchTests/Base/Utils.swift     |  2 +-
 5 files changed, 41 insertions(+), 47 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 8aed98a3..ca6b982a 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -4,6 +4,7 @@ All notable changes to this project will be documented in this file.
 
 ## [unreleased]
 
+🐛 **fix:** run on Apple Silicon ([110](https://github.com/owkin/GrAIdient/pull/110))\
 ⚙️ **core:** initForward,Backward model API ([109](https://github.com/owkin/GrAIdient/pull/109))\
 🪜 **layer_1d:** Dropout1D ([#108](https://github.com/owkin/GrAIdient/pull/108))\
 🪜 **feat:** VQGrad, VQGradSeq ([#107](https://github.com/owkin/GrAIdient/pull/107))
diff --git a/Sources/GrAIdient/Metal/MetalKernel.swift b/Sources/GrAIdient/Metal/MetalKernel.swift
index 7228653c..5425b42c 100644
--- a/Sources/GrAIdient/Metal/MetalKernel.swift
+++ b/Sources/GrAIdient/Metal/MetalKernel.swift
@@ -969,7 +969,16 @@ public class MetalCommand
     public func setBytes<T>(_ data: [T], atIndex index: Int)
     {
         let byteLength = data.count * MemoryLayout<T>.size
-        _encoder.setBytes(data, length: byteLength, index: index)
+        data.withUnsafeBufferPointer
+        {
+            dataPtr in
+            
+            _encoder.setBytes(
+                UnsafeRawPointer(dataPtr.baseAddress)!,
+                length: byteLength,
+                index: index
+            )
+        }
     }
     
     ///
diff --git a/Sources/GrAIdient/Utils/Image.swift b/Sources/GrAIdient/Utils/Image.swift
index 2450a321..9c24c81d 100644
--- a/Sources/GrAIdient/Utils/Image.swift
+++ b/Sources/GrAIdient/Utils/Image.swift
@@ -6,7 +6,7 @@
 //
 
 import Foundation
-import Cocoa
+import AppKit
 
 /// Error occuring when processing images.
 public enum ImageError: Error
@@ -107,42 +107,14 @@ public class Image
         let bufferPtr = metalBuffer.download()
         let nbImages = metalBuffer.nbElems / (width * height * 3)
         
-        var output = [[UInt8]]()
-        for elem in 0..<nbImages
+        var images = [[Float]]()
+        for i in 0..<nbImages
         {
-            var grid: [UInt8] = [UInt8](repeating: 0, count: width * height * 3)
-            grid.withUnsafeMutableBufferPointer { gridPtr in
-            Concurrency.slice(gridPtr.count)
-            {
-                (index: Int) in
-                
-                let depth = index / (width * height)
-                let i = (index - depth * width * height) / width
-                let j = (index - depth * width * height) % width
-                
-                let offsetGet = elem * 3 * height * width
-                let offsetSet = j + i * width
-                
-                let valTmp = bufferPtr[index + offsetGet] * 255.0
-                let val: UInt8
-                if valTmp < 0
-                {
-                    val = 0
-                }
-                else if valTmp > 255.0
-                {
-                    val = 255
-                }
-                else
-                {
-                    val = UInt8(valTmp)
-                }
-                    
-                gridPtr[3 * offsetSet + depth] = val
-            }}
-            output.append(grid)
+            images.append([Float](
+                bufferPtr[i * 3 * height * width..<(i+1) * 3 * height * width]
+            ))
         }
-        return output
+        return toRGB(toPixel(images), width: width, height: height)
     }
 
     ///
@@ -157,7 +129,8 @@ public class Image
         var output = [[UInt8]]()
         for elem in 0..<images.count
         {
-            output.append(images[elem].map {
+            output.append(images[elem].map 
+            {
                 let valTmp = $0 * T(255.0)
                 let val: UInt8
                 if valTmp < 0
@@ -385,15 +358,26 @@ public extension NSImage
     ///
     func extractPixels() throws -> [UInt8]
     {
-        if let imageData = tiffRepresentation,
-           let imageRep = NSBitmapImageRep(data: imageData),
-           let dataPtr = imageRep.bitmapData
+        if let pixelData = (cgImage(
+            forProposedRect: nil, context: nil, hints: nil)!).dataProvider?.data
         {
-            let bufferPtr = UnsafeBufferPointer(
-                start: dataPtr,
-                count: Int(3 * size.height * size.width)
-            )
-            return [UInt8](bufferPtr)
+            let data: UnsafePointer<UInt8> = CFDataGetBytePtr(pixelData)
+            
+            var pixels = [UInt8]()
+            for i in 0..<Int(size.height) {
+            for j in 0..<Int(size.width)
+            {
+                let pos = CGPoint(x: j, y: i)
+                
+                let pixelInfo: Int = (Int(size.width) * Int(pos.y) * 4) +
+                    Int(pos.x) * 4
+                
+                let r = data[pixelInfo]
+                let g = data[pixelInfo + 1]
+                let b = data[pixelInfo + 2]
+                pixels += [r, g, b]
+            }}
+            return pixels
         }
         else
         {
diff --git a/Tests/GrAIExamples/Base/Utils.swift b/Tests/GrAIExamples/Base/Utils.swift
index bbc1e9f7..5f46f133 100644
--- a/Tests/GrAIExamples/Base/Utils.swift
+++ b/Tests/GrAIExamples/Base/Utils.swift
@@ -12,7 +12,7 @@ import GrAIdient
 /// Python library default path.
 let PYTHON_LIB =
     FileManager.default.homeDirectoryForCurrentUser.path +
-    "/opt/anaconda3/envs/graiexamples/lib/libpython3.9.dylib"
+    "/miniconda3/envs/graiexamples/lib/libpython3.9.dylib"
 
 /// Set the Python library path.
 func setPythonLib()
diff --git a/Tests/GrAITorchTests/Base/Utils.swift b/Tests/GrAITorchTests/Base/Utils.swift
index 439b0e13..9c80f4ec 100644
--- a/Tests/GrAITorchTests/Base/Utils.swift
+++ b/Tests/GrAITorchTests/Base/Utils.swift
@@ -12,7 +12,7 @@ import GrAIdient
 /// Python library default path.
 let PYTHON_LIB =
     FileManager.default.homeDirectoryForCurrentUser.path +
-    "/opt/anaconda3/envs/graitorch/lib/libpython3.9.dylib"
+    "/miniconda3/envs/graitorch/lib/libpython3.9.dylib"
 
 /// Set the Python library path.
 func setPythonLib()