Merge branch 'Iainmon:main' into main

Author: BrandynTucknott

CMakeLists.txt

@@ -131,8 +131,22 @@ target_include_directories(
 )
 
 
+add_library(bridge_objs STATIC $<TARGET_OBJECTS:bridge>)
+set_target_properties(bridge_objs
+    PROPERTIES ARCHIVE_OUTPUT_DIRECTORY "${CMAKE_BINARY_DIR}"
+)
+
+
 set(BRIDGE_OBJECT_FILES $<TARGET_OBJECTS:bridge>)
 
+# add_custom_command(
+#     TARGET bridge
+#     POST_BUILD
+#     COMMAND ${CMAKE_COMMAND} -E copy_directory
+#             "${CMAKE_CURRENT_SOURCE_DIR}/style-transfer/models"
+#             "$<TARGET_FILE_DIR:StyleTransfer>/style-transfer/models"
+#     COMMENT "NOT! Copying ${PROJECT_ROOT_DIR}/examples/vgg/images to $<TARGET_FILE_DIR:vgg>/images"
+# )
 
 
 
@@ -214,12 +228,8 @@ add_dependencies(TinyLayerTest ChAI)
 target_link_options(TinyLayerTest
     PRIVATE
         --main-module layer_test.chpl
-        -M ${PROJECT_ROOT_DIR}/lib
-        ${BRIDGE_DIR}/include/bridge.h
-        ${BRIDGE_OBJECT_FILES}
-        -L ${LIBTORCH_DIR}/lib
-        ${LIBTORCH_LIBS_LINKER_ARGS}
-        --ldflags "-Wl,-rpath,${LIBTORCH_DIR}/lib"
+        # -M ${PROJECT_ROOT_DIR}/lib
+        ${CHAI_LINKER_ARGS}
 )
 # chpl test/tiny/layer_test.chpl -M lib bridge/include/bridge.h build/CMakeFiles/bridge.dir/bridge/lib/bridge.cpp.o -L libtorch/lib -ltorch -ltorch_cpu -lc10 -ltorch_global_deps --ldflags "-Wl,-rpath,libtorch/lib"
 
@@ -237,7 +247,8 @@ set(CHAI_LINKER_ARGS
     ${BRIDGE_OBJECT_FILES}
     -L ${LIBTORCH_DIR}/lib
     ${LIBTORCH_LIBS_LINKER_ARGS}
-    --ldflags "-Wl,-rpath,${LIBTORCH_DIR}/lib"
+    --ccflags "-I${BRIDGE_DIR}/include -L${PROJECT_ROOT_DIR}/build"
+    --ldflags "-L${PROJECT_ROOT_DIR}/build -Wl,-rpath,${LIBTORCH_DIR}/lib"
 )
 
 

bridge/.DS_Store

@@ -15,6 +15,9 @@ typedef double float64_t;
 typedef char bool_t;
 typedef unsigned char uint8_t;
 typedef unsigned int uint32_t;
+typedef unsigned long long uint64_t;
+
+void debug_cpu_only_mode(bool_t mode);
 
 typedef struct bridge_tensor_t {
     float* data;
@@ -24,6 +27,17 @@ typedef struct bridge_tensor_t {
 } bridge_tensor_t;
 
 
+typedef struct bridge_pt_model_t {
+    void* pt_module;
+} bridge_pt_model_t;
+
+typedef struct test_struct_t {
+    int* field;
+} test_struct_t;
+
+
+void hello_world(void);
+
 typedef struct nil_scalar_tensor_t {
     float scalar;
     bridge_tensor_t tensor;
@@ -36,6 +50,14 @@ float* unsafe(const float* arr);
 bridge_tensor_t load_tensor_from_file(const uint8_t* file_path);
 bridge_tensor_t load_tensor_dict_from_file(const uint8_t* file_path,const uint8_t* tensor_key);
 bridge_tensor_t load_run_model(const uint8_t* model_path, bridge_tensor_t input);
+
+bridge_pt_model_t load_model(const uint8_t* model_path);
+
+bool_t accelerator_available(void);
+
+bridge_tensor_t model_forward(bridge_pt_model_t model, bridge_tensor_t input);
+bridge_tensor_t model_forward_style_transfer(bridge_pt_model_t model, bridge_tensor_t input);
+
 bridge_tensor_t resize(bridge_tensor_t input,int height,int width);
 bridge_tensor_t imagenet_normalize(bridge_tensor_t input);
 
@@ -107,7 +129,6 @@ proto_bridge_simple(tanhshrink);
 void split_loop(int64_t idx, int64_t n);
 void split_loop_filler(int64_t n,int64_t* ret);
 
-void show_webcam(void);
 
 
 // bridge_tensor_t conv2d(

bridge/include/bridge.h

@@ -1,6 +1,8 @@
 #include <bridge.h>
 
 #include <torch/torch.h>
+#include <Aten/ATen.h>
+
 #include <torch/script.h>
 
 // #include <torch/script.h>
@@ -16,7 +18,6 @@
 #include <chrono>
 #include <thread>
 
-#include <opencv2/opencv.hpp>
 
 
 #define def_bridge_simple(Name) \
@@ -28,6 +29,58 @@
 
 
 
+// Globals
+
+
+torch::Device get_best_device();
+torch::ScalarType get_best_dtype();
+
+auto best_device = get_best_device();
+auto best_dtype = get_best_dtype();
+
+torch::NoGradGuard no_grad;
+torch::AutoGradMode enable_grad(false);
+
+bool debug_cpu_only = false;
+
+
+
+torch::Device get_best_device() {
+    if (debug_cpu_only) 
+        return torch::Device(torch::kCPU);
+    
+    if (torch::hasMPS()) {
+        return torch::Device(torch::kMPS);
+    } else if (torch::hasCUDA()) {
+        return torch::Device(torch::kCUDA);
+    } else {
+        return torch::Device(torch::kCPU);
+    }
+}
+
+extern "C" void debug_cpu_only_mode(bool_t mode) {
+    debug_cpu_only = mode;
+    if (debug_cpu_only) {
+        best_device = torch::Device(torch::kCPU);
+    } else {
+        best_device = get_best_device();
+    }
+}
+
+extern "C" bool_t accelerator_available() {
+    return (best_device == torch::Device(torch::kCUDA) || best_device == torch::Device(torch::kMPS));
+}
+
+torch::ScalarType get_best_dtype() {
+    if (torch::hasMPS()) {
+        return torch::kFloat16;
+    } else if (torch::hasCUDA()) {
+        return torch::kFloat16;
+    } else {
+        return torch::kFloat32;
+    }
+}
+
 int bridge_tensor_elements(bridge_tensor_t &bt) {
     int size = 1;
     for (int i = 0; i < bt.dim; ++i) {
@@ -40,14 +93,14 @@ size_t bridge_tensor_size(bridge_tensor_t &bt) {
     return sizeof(float32_t) * bridge_tensor_elements(bt);
 }
 
-void store_tensor(torch::Tensor &input, float32_t* dest) {
+void store_tensor(at::Tensor &input, float32_t* dest) {
     float32_t * data = input.data_ptr<float32_t>();
     size_t bytes_size = sizeof(float32_t) * input.numel();
     // std::memmove(dest,data,bytes_size);
     std::memcpy(dest,data,bytes_size);
 }
 
-bridge_tensor_t torch_to_bridge(torch::Tensor &tensor) {
+bridge_tensor_t torch_to_bridge(at::Tensor &tensor) {
     bridge_tensor_t result;
     result.created_by_c = true;
     result.dim = tensor.dim();
@@ -60,12 +113,25 @@ bridge_tensor_t torch_to_bridge(torch::Tensor &tensor) {
     return result;
 }
 
-torch::Tensor bridge_to_torch(bridge_tensor_t &bt) {
+at::Tensor bridge_to_torch(bridge_tensor_t &bt) {
     std::vector<int64_t> sizes_vec(bt.sizes, bt.sizes + bt.dim);
     auto shape = torch::IntArrayRef(sizes_vec);
     return torch::from_blob(bt.data, shape, torch::kFloat);
 }
 
+at::Tensor bridge_to_torch(bridge_tensor_t &bt,torch::Device device, bool copy,torch::ScalarType dtype = torch::kFloat32) {
+    std::vector<int64_t> sizes_vec(bt.sizes, bt.sizes + bt.dim);
+    auto shape = torch::IntArrayRef(sizes_vec);
+    auto t = torch::from_blob(bt.data, shape, torch::kFloat);
+    if (device != torch::kCPU)
+        copy = true;
+    if (copy)
+        return t.to(device, dtype, /*non_blocking=*/false, /*copy=*/true);
+    else
+        return t.to(device, dtype, /*non_blocking=*/false, /*copy=*/false);
+    
+}
+
 extern "C" float32_t* unsafe(const float32_t* arr) {
     return const_cast<float32_t*>(arr);
 }
@@ -131,6 +197,92 @@ extern "C" bridge_tensor_t load_run_model(const uint8_t* model_path, bridge_tens
     return torch_to_bridge(output);
 }
 
+
+
+
+extern "C" bridge_pt_model_t load_model(const uint8_t* model_path) {
+
+    std::cout << "Begin loading model from path: " << model_path << std::endl;
+    std::cout.flush();
+    std::string path(reinterpret_cast<const char*>(model_path));
+    std::cout << "Loading model from path: " << path << std::endl;
+    std::cout.flush();
+
+    try {
+        auto* module = new torch::jit::Module(torch::jit::load(path));
+        module->to(best_device,best_dtype,false);
+        module->eval();
+        std::cout << "Model loaded successfully!" << std::endl;
+        std::cout.flush();
+        return { static_cast<void*>(module) };
+    } catch (const c10::Error& e) {
+        std::cerr << "error loading the model\n" << e.msg();
+        std::cout << "error loading the model\n" << e.msg();
+        std::cout.flush();
+        std::cerr.flush();
+    }
+    std::cout << "Model loading failed!" << std::endl;
+    std::cout.flush();
+
+    return { nullptr };
+}
+
+
+
+bridge_tensor_t model_forward(bridge_pt_model_t model, bridge_tensor_t input, bool is_vgg_based_model) {
+    auto tn_mps = bridge_to_torch(input,best_device,true,best_dtype);
+    // tn_mps = tn_mps.permute({2, 0, 1}).contiguous();
+    // tn_mps.unsqueeze_(0);//.contiguous();
+    auto tn = tn_mps.permute({2, 0, 1}).unsqueeze(0).contiguous();
+
+    std::vector<torch::jit::IValue> ins;
+    ins.push_back(tn);
+
+    auto* module = static_cast<torch::jit::Module*>(model.pt_module);
+    auto o = module->forward(ins).toTensor();
+    // auto tn_out = o.squeeze(0).permute({1, 2, 0}).contiguous();
+    auto tn_out = o.squeeze(0).contiguous().permute({1, 2, 0}).contiguous();
+
+    if (is_vgg_based_model) {
+        tn_out.div_(255.0);
+    }
+
+    auto tn_out_cpu = tn_out.to(torch::kCPU,torch::kFloat32,false,true);
+
+    return torch_to_bridge(tn_out_cpu);
+
+}
+
+extern "C" bridge_tensor_t model_forward(bridge_pt_model_t model, bridge_tensor_t input) {
+    return model_forward(model, input, false);
+}
+
+extern "C" bridge_tensor_t model_forward_style_transfer(bridge_pt_model_t model, bridge_tensor_t input) {
+    return model_forward(model, input, true);
+}
+
+// std::tuple<uint64_t, uint64_t> get_cpu_frame_size(uint64_t width, uint64_t height, float32_t scale_factor) {
+//     // if (best_device == torch::kMPS || best_device == torch::kCUDA)
+//     if (accelerator_available())
+//         return std::make_tuple(width, height);
+//     uint64_t new_width = static_cast<uint64_t>(width * scale_factor);
+//     uint64_t new_height = static_cast<uint64_t>(height * scale_factor);
+//     return std::make_tuple(new_width, new_height);
+// }
+
+// extern "C" uint64_t get_cpu_frame_width(uint64_t width,float32_t scale_factor) {
+//     return std::get<0>(get_cpu_frame_size(width, 0, scale_factor));
+// }
+// extern "C" uint64_t get_cpu_frame_height(uint64_t height,float32_t scale_factor) {
+//     return std::get<1>(get_cpu_frame_size(0, height, scale_factor));
+// }
+
+
+extern "C" void hello_world(void) {
+    std::cout << "Hello from C++!" << std::endl;
+    std::cout.flush();
+}
+
 extern "C" bridge_tensor_t increment3(bridge_tensor_t arr) {
     auto t = bridge_to_torch(arr);
     // Increment the tensor
@@ -404,37 +556,37 @@ extern "C" void split_loop_filler(int64_t n,int64_t* ret) {
 
 
 
-cv::VideoCapture open_camera(int cam_index) {
-    cv::VideoCapture cap(cam_index, cv::CAP_AVFOUNDATION);
-    if (!cap.isOpened()) {
-        std::cerr << "Could not open camera index " << cam_index << std::endl;
-        return cv::VideoCapture();
-    }
-    cap.set(cv::CAP_PROP_BUFFERSIZE, 1); // minimal internal buffering
-    cap.set(cv::CAP_PROP_FPS, 60);       // request higher FPS if possible
-    return cap;
-}
+// cv::VideoCapture open_camera(int cam_index) {
+//     cv::VideoCapture cap(cam_index, cv::CAP_AVFOUNDATION);
+//     if (!cap.isOpened()) {
+//         std::cerr << "Could not open camera index " << cam_index << std::endl;
+//         return cv::VideoCapture();
+//     }
+//     cap.set(cv::CAP_PROP_BUFFERSIZE, 1); // minimal internal buffering
+//     cap.set(cv::CAP_PROP_FPS, 60);       // request higher FPS if possible
+//     return cap;
+// }
 
 
-extern "C" void show_webcam(void) {
-    cv::VideoCapture cap;
-    cap = open_camera(0);
+// extern "C" void show_webcam(void) {
+//     cv::VideoCapture cap;
+//     cap = open_camera(0);
 
-    cv::Mat frame_bgr;
+//     cv::Mat frame_bgr;
 
-    while (true) {
-        if (!cap.read(frame_bgr) || frame_bgr.empty()) {
-            std::cerr << "[WARN] Empty frame, exiting" << std::endl;
-            break;
-        }
+//     while (true) {
+//         if (!cap.read(frame_bgr) || frame_bgr.empty()) {
+//             std::cerr << "[WARN] Empty frame, exiting" << std::endl;
+//             break;
+//         }
 
-        cv::imshow("webcam", frame_bgr);
+//         cv::imshow("webcam", frame_bgr);
 
-        if (cv::waitKey(1) == 27) { // ESC key
-            break;
-        }
-    }
+//         if (cv::waitKey(1) == 27) { // ESC key
+//             break;
+//         }
+//     }
 
-    cap.release();
-    cv::destroyAllWindows();
-}
+//     cap.release();
+//     cv::destroyAllWindows();
+// }

bridge/lib/bridge.cpp

@@ -0,0 +1 @@
+This folder contains the model architectures used in the demos.