Add pytorch file format loading functionality for ndarrays.

Author: Iainmon

CMakeLists.txt

@@ -29,7 +29,7 @@ if(APPLE)
     set(CMAKE_C_COMPILER "/usr/bin/clang")
     set(CMAKE_CXX_COMPILER "/usr/bin/clang++")
 endif()
-set(CMAKE_CXX_STANDARD 17)
+set(CMAKE_CXX_STANDARD 23)
 
 
 include(LibTorchDL)

bridge/include/bridge.h

@@ -10,6 +10,8 @@ extern "C" {
 typedef float float32_t;
 typedef double float64_t;
 typedef char bool_t;
+typedef unsigned char uint8_t;
+typedef unsigned int uint32_t;
 
 typedef struct bridge_tensor_t {
     float* data;
@@ -27,6 +29,10 @@ typedef struct nil_scalar_tensor_t {
     bool_t is_tensor;
 } nil_scalar_tensor_t;
 
+float* unsafe(const float* arr);
+bridge_tensor_t load_tensor_from_file(const uint8_t* file_path);
+
+
 int baz(void);
 
 void wrHello(void);
@@ -65,6 +71,7 @@ bridge_tensor_t max_pool2d(
     int dilation
 );
 
+
 // bridge_tensor_t conv2d(
 //     bridge_tensor_t input,
 //     bridge_tensor_t kernel,
@@ -73,7 +80,6 @@ bridge_tensor_t max_pool2d(
 //     nil_scalar_tensor_t padding
 // );
 
-float* unsafe(const float* arr);
 
 
 #ifdef __cplusplus

bridge/lib/bridge.cpp

@@ -4,15 +4,16 @@
 // #include <torch/script.h>
 // #include <Aten/ATen.h>
 #include <iostream>
+#include <fstream>
+#include <string>
+#include <cstring>
+#include <sstream>
+#include <cstdlib>
 #include <vector>
-
 #include <cstdint>
 
 
 
-extern "C" float32_t* unsafe(const float32_t* arr) {
-    return const_cast<float32_t*>(arr);
-}
 
 int bridge_tensor_elements(bridge_tensor_t &bt) {
     int size = 1;
@@ -51,6 +52,62 @@ torch::Tensor bridge_to_torch(bridge_tensor_t &bt) {
     return torch::from_blob(bt.data, shape, torch::kFloat);
 }
 
+
+
+
+
+
+
+
+
+
+
+
+
+extern "C" float32_t* unsafe(const float32_t* arr) {
+    return const_cast<float32_t*>(arr);
+}
+
+std::vector<char> get_the_bytes(std::string filename) {
+    std::ifstream input(filename, std::ios::binary);
+    std::vector<char> bytes((std::istreambuf_iterator<char>(input)),(std::istreambuf_iterator<char>()));
+    input.close();
+    return bytes;
+}
+
+extern "C" bridge_tensor_t load_tensor_from_file(const uint8_t* file_path) {
+    // // Load the tensor from a file
+    // torch::Tensor tensor;
+    // // torch::load(tensor,file_path);
+
+    // std::cout << "Tensor loaded from file: " << tensor.sizes() << std::endl;
+
+    // // Convert the tensor to a bridge_tensor_t
+
+    std::string fp(reinterpret_cast<const char*>(file_path));
+    std::cout << "File path: " << fp << std::endl;
+
+    std::vector<char> f = get_the_bytes(fp);
+    std::cout << "File size: " << f.size() << std::endl;
+
+    torch::IValue x = torch::pickle_load(f);
+    // std::cout << "IValue loaded from file: " << x << std::endl;
+
+    torch::Tensor t = x.toTensor();
+    std::cout << "Tensor loaded from IValue: " << t.sizes() << std::endl;
+    std::cout << "Tensor sum: " << t.sum() << std::endl;
+
+    return torch_to_bridge(t);
+}
+
+
+
+
+
+
+
+
+
 extern "C" bridge_tensor_t increment3(bridge_tensor_t arr) {
     auto t = bridge_to_torch(arr);
     // Increment the tensor

examples/vgg/images/my_tensor.pt

@@ -0,0 +1,109 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "873dd3b8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import torch"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "a07c23ff",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "def find_factors(num):\n",
+    "    \"\"\"\n",
+    "    Finds all factors of a given number.\n",
+    "\n",
+    "    Args:\n",
+    "        num: An integer.\n",
+    "\n",
+    "    Returns:\n",
+    "        A list of integers representing the factors of num.\n",
+    "    \"\"\"\n",
+    "    factors = []\n",
+    "    for i in range(1, num + 1):\n",
+    "        if num % i == 0:\n",
+    "            factors.append(i)\n",
+    "    return factors"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 19,
+   "id": "131adc46",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "f1: 10000, f2: 5\n"
+     ]
+    }
+   ],
+   "source": [
+    "num_elt = 50000\n",
+    "f1 = find_factors(num_elt)[-4]\n",
+    "f2 = num_elt // f1\n",
+    "print(f\"f1: {f1}, f2: {f2}\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 20,
+   "id": "d4aed442",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "x = torch.arange(0,num_elt)\n",
+    "x = x.reshape(f1,f2).to(torch.float32)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 21,
+   "id": "481d4709",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "torch.save(x, 'my_tensor.pt')"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "1a17ec15",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": ".venv",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

examples/vgg/mktensor.ipynb

@@ -53,29 +53,32 @@ proc run(model: shared VGG16(real(32)), file: string) {
     return (topPredictions.data, percentTopk);
 }
 
+import Path;
+
+
 proc main(args: [] string) {
-  writeln("Loading labels from ", labelFile);
-  const labels = getLabels();
-  writeln("Loaded ", labels.size, " labels.");
+    writeln("Loading labels from ", labelFile);
+    const labels = getLabels();
+    writeln("Loaded ", labels.size, " labels.");
 
-  writeln("Constructing VGG16 model.");
-  const vgg = new shared VGG16(real(32));
-  writeln("Constructed VGG16 model.");
+    writeln("Constructing VGG16 model.");
+    const vgg = new shared VGG16(real(32));
+    writeln("Constructed VGG16 model.");
 
-  writeln("Loading VGG16 model weights.");
-  vgg.loadPyTorchDump(modelDir, false);
-  writeln("Loaded VGG16 model.");
+    writeln("Loading VGG16 model weights.");
+    vgg.loadPyTorchDump(modelDir, false);
+    writeln("Loaded VGG16 model.");
 
 
-  var files = args[1..];
+    var files = args[1..];
 
-  for f in files {
-    var (topArr, percent) = run(vgg, f);
-    writeln("For '", f, "' the top ", k, " predictions are: ");
-    for i in 0..<k {
-      writef("  %?: label=%?; confidence=%2.2r%%\n", i, labels[topArr[i]], percent[i]);
+    for f in files {
+        var (topArr, percent) = run(vgg, f);
+        writeln("For '", f, "' the top ", k, " predictions are: ");
+        for i in 0..<k {
+        writef("  %?: label=%?; confidence=%2.2r%%\n", i, labels[topArr[i]], percent[i]);
+        }
+        writeln();
     }
-    writeln();
-  }
 
 }

examples/vgg/my_tensor.pt

@@ -6,6 +6,8 @@ module Bridge {
     use Utilities.Standard;
     use Allocators;
 
+    use CTypes;
+
 
     extern record bridge_tensor_t {
         var data: c_ptr(real(32));
@@ -31,6 +33,17 @@ module Bridge {
         }
     }
 
+    extern proc unsafe(const ref arr: [] real(32)): c_ptr(real(32));
+
+    // extern proc load_tensor_from_file(file_path: c_ptrConst(u_char)): bridge_tensor_t; // Working
+
+    // extern proc load_tensor_from_file(const ref file_path: uint(8)): bridge_tensor_t;
+    // extern proc load_tensor_from_file(file_path: c_ptrConst(c_uchar)): bridge_tensor_t;
+    
+    // extern proc load_tensor_from_file(const file_path: c_ptr(uint(8))): bridge_tensor_t; // also working
+
+    extern proc load_tensor_from_file(const file_path: c_ptr(uint(8))): bridge_tensor_t;
+
 
     extern proc convolve2d(
         in input: bridge_tensor_t, 
@@ -56,8 +69,6 @@ module Bridge {
         in dilation: int(32)): bridge_tensor_t;
 
 
-    extern proc unsafe(const ref arr: [] real(32)): c_ptr(real(32));
-
 
     proc getSizeArray(const ref arr: [] ?eltType): [] int(32) {
         var sizes: [0..<arr.rank] int(32);

examples/vgg/test.chpl

@@ -4,6 +4,7 @@ import ChapelArray;
 import Math;
 import Random;
 import IO;
+import Path;
 
 use Env;
 
@@ -2284,6 +2285,27 @@ proc ref ndarray.saveImage(imagePath: string) where rank == 3 {
     Image.writeImage(imagePath,format=imgType,pixels=pixelData);
 }
 
+proc ref ndarray.loadChData(fr: IO.fileReader(?)) throws {
+    var r = fr.read(int);
+    if r != rank then
+        err("Error reading tensor: rank mismatch.", r , " != this." , rank);
+    var s = this.shape;
+    for i in 0..#rank do
+        s[i] = fr.read(int);
+    var d = util.domainFromShape((...s));
+    this._domain = d;
+    // for i in d do
+    //     this.data[i] = fr.read(eltType);
+    fr.read(this.data);
+}
+
+proc type ndarray.loadPyTorchTensor(param rank: int,in filePath: string,type eltType = defaultEltType): ndarray(rank,eltType) {
+    use CTypes;
+    const fpPtr: c_ptr(uint(8)) = c_ptrTo(filePath);
+    var th = Bridge.load_tensor_from_file(fpPtr);
+    return ndarray.fromBridgeTensor(rank,th) : ndarray(rank,eltType);
+}
+
 // For printing. 
 proc ndarray.serialize(writer: IO.fileWriter(locking=false, IO.defaultSerializer),ref serializer: IO.defaultSerializer) throws {
 
@@ -2300,17 +2322,26 @@ proc ndarray.serialize(writer: IO.fileWriter(locking=false, IO.defaultSerializer
 }
 
 proc ref ndarray.read(fr: IO.fileReader(?)) throws {
-    var r = fr.read(int);
-    if r != rank then
-        err("Error reading tensor: rank mismatch.", r , " != this." , rank);
-    var s = this.shape;
-    for i in 0..#rank do
-        s[i] = fr.read(int);
-    var d = util.domainFromShape((...s));
-    this._domain = d;
-    // for i in d do
-    //     this.data[i] = fr.read(eltType);
-    fr.read(this.data);
+
+    const file = fr.getFile();
+    const filePath: string = file.path;
+    const (_,fileName,fileExt) = util.splitPathParts(filePath);
+
+    select fileExt {
+        when "chdata" do
+            this.loadChData(fr);
+        when "png" do
+            this = ndarray.loadImage(filePath,eltType);
+        when "jpg" do
+            this = ndarray.loadImage(filePath,eltType);
+        when "jpeg" do
+            this = ndarray.loadImage(filePath,eltType);
+        when "bmp" do
+            this = ndarray.loadImage(filePath,eltType);
+    }
+
+
+
 }