Improve chapel/pytorch interop with unsafe pointers.

Author: Iainmon

bridge/.DS_Store

@@ -1,4 +1,5 @@
 
+all: bridge.o Bridge
 
 bridge.o: lib/bridge.cpp include/bridge.h 
 	g++ -c lib/bridge.cpp -I include -o bridge.o -I /Users/iainmoncrief/Documents/Github/ChAI/bridge/libtorch/include/torch/csrc/api/include -I /Users/iainmoncrief/Documents/Github/ChAI/bridge/libtorch/include --std=c++17

bridge/Bridge

@@ -13,6 +13,9 @@ typedef struct bridge_tensor_t {
     int dim;
 } bridge_tensor_t;
 
+typedef float float32_t;
+typedef double float64_t;
+
 int baz(void);
 
 void wrHello(void);
@@ -25,17 +28,16 @@ void increment(float* arr, int* sizes, int dim, float* output);
 bridge_tensor_t increment2(float* arr, int* sizes, int dim);
 bridge_tensor_t increment3(bridge_tensor_t arr);
 
-// void convolve(
-//     float* input,
-//     int* input_sizes,
-//     int input_dim,
-//     float* kernel,
-//     int* kernel_sizes,
-//     int kernel_dim,
-//     float* output,
-//     int* output_sizes,
-//     int output_dim
-// );
+bridge_tensor_t convolve2d(
+    bridge_tensor_t input,
+    bridge_tensor_t kernel,
+    bridge_tensor_t bias,
+    int stride,
+    int padding
+);
+
+float* unsafe(const float* arr);
+
 
 #ifdef __cplusplus
 }

bridge/Makefile

@@ -23,6 +23,14 @@ extern record bridge_tensor_t {
 extern proc increment2(arr: [] real(32), sizes: [] int(32), dim: int(32)): bridge_tensor_t;
 extern proc increment3(in arr: bridge_tensor_t): bridge_tensor_t;
 
+extern proc convolve2d(
+    in input: bridge_tensor_t, 
+    in kernel: bridge_tensor_t, 
+    in bias: bridge_tensor_t, 
+    in stride: int(32), 
+    in padding: int(32)): bridge_tensor_t;
+
+extern proc unsafe(const ref arr: [] real(32)): c_ptr(real(32));
 
 // baz();
 
@@ -127,9 +135,9 @@ proc bridgeTensorToArray(param rank: int, package: bridge_tensor_t): [] real(32)
 }
 
 
-proc createBridgeTensor(ref data: [] real(32)): bridge_tensor_t {
+proc createBridgeTensor(const ref data: [] real(32)): bridge_tensor_t {
     var result: bridge_tensor_t;
-    result.data = c_ptrTo(data);
+    result.data = c_ptrToConst(data) : c_ptr(real(32));
     result.sizes = allocate(int(32),data.rank);
     const sizeArr = getSizeArray(data);
     for i in 0..<data.rank do
@@ -169,4 +177,19 @@ writeln(a);
 writeln("----------");
 writeln(chplIncrement(a));
 writeln("----------");
-writeln(a);
+writeln(a);
+
+
+var input: [domainFromShape(2,64,28,28)] real(32) = 1.0;
+var kernel: [domainFromShape(128,64,3,3)] real(32) = 2.0;
+var bias: [domainFromShape(128)] real(32) = 3.0;
+var stride: int(32) = 1;
+var padding: int(32) = 1;
+writeln("Begin.");
+var resultBT = convolve2d(createBridgeTensor(input), createBridgeTensor(kernel), createBridgeTensor(bias), stride, padding);
+var result = bridgeTensorToArray(4, resultBT);
+// writeln("Input: ", input);
+// writeln("Kernel: ", kernel);
+// writeln("Bias: ", bias);
+// writeln("Result: ", result);
+writeln("Result: ", result.size);

bridge/bridge.o

@@ -8,9 +8,12 @@
 
 #include <cstdint>
 
-using float32_t = float;
 
 
+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;
     for (int i = 0; i < bt.dim; ++i) {
@@ -56,6 +59,20 @@ extern "C" bridge_tensor_t increment3(bridge_tensor_t arr) {
     return tensor_result_convert(incremented_tensor);
 }
 
+extern "C" bridge_tensor_t convolve2d(
+    bridge_tensor_t input,
+    bridge_tensor_t kernel,
+    bridge_tensor_t bias,
+    int stride,
+    int padding
+) {
+    auto t_input = bridge_to_torch(input);
+    auto t_kernel = bridge_to_torch(kernel);
+    auto t_bias = bridge_to_torch(bias);
+    auto output = torch::conv2d(t_input, t_kernel, t_bias, stride, padding);
+    return tensor_result_convert(output);
+}
+
 
 extern "C" int baz(void) {
     printf("Hello from baz!\n");