torchpipe 0.1.23__py3-none-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl

torchpipe/__init__.py

@@ -0,0 +1,99 @@
+# isort: skip_file
+
+from packaging import version
+import logging
+logger = logging.getLogger(__name__)  # type: ignore
+
+import ctypes, os
+
+ORI_TVM_FFI_DISABLE_TORCH_C_DLPACK = os.environ.get(
+    "TVM_FFI_DISABLE_TORCH_C_DLPACK", "0")
+if ORI_TVM_FFI_DISABLE_TORCH_C_DLPACK == "0":
+    os.environ["TVM_FFI_DISABLE_TORCH_C_DLPACK"] = "1"
+
+import omniback
+
+import torch
+
+
+try:
+    from importlib.metadata import version as _get_version
+    __version__ = _get_version("torchpipe")
+except Exception:
+    __version__ = "0.0.0-dev"  
+
+# -----------------------
+# assert omniback.compiled_with_cxx11_abi() == torch.compiled_with_cxx11_abi()
+
+logger.info(f'torch.cuda.is_available() = {torch.cuda.is_available()}')
+
+torch.set_num_threads(torch.get_num_threads())
+
+# -----------------------
+from .load_libs import _load_or_build_lib, _load_or_build_lib_skip_if_error  # nosort
+from .load_libs import _setting_group_handle  # nosort
+
+SKIP_ALL=os.environ.get("TORCHPIPE_SKIP_ALL", "0")
+
+if SKIP_ALL != "1":
+    try:
+        _load_or_build_lib("torchpipe_core")
+        if torch.cuda.is_available():
+            _load_or_build_lib("torchpipe_core_cuda")
+    except Exception as e:
+        logger.warning(f'Failed to load or JIT compile builtin extensions: \n{e}')
+        SKIP_ALL = "1"
+    else:
+        SKIP_TENSORRT=os.environ.get("TORCHPIPE_SKIP_TENSORRT", "0")
+                
+        if torch.cuda.is_available():
+            # _load_or_build_lib_skip_if_error("torchpipe_core_cuda")
+            if SKIP_TENSORRT != "1":
+                _load_or_build_lib_skip_if_error("torchpipe_tensorrt")
+
+            _load_or_build_lib_skip_if_error("torchpipe_nvjpeg")
+        else:
+            logger.warning("[JIT] CUDA is not available, skip loading CUDA extensions.")
+
+        SKIP_OPENCV=os.environ.get("TORCHPIPE_SKIP_OPENCV", "0")
+        if SKIP_OPENCV != "1":
+            _load_or_build_lib_skip_if_error("torchpipe_opencv")
+        
+grp_config = os.path.join(os.path.dirname(__file__), "group-torchpipe.toml")
+assert os.path.exists(grp_config), grp_config
+_setting_group_handle(grp_config)
+logger.info(f"Loaded group config from {grp_config}")
+
+
+# -----------------------
+pipe = omniback.pipe
+Dict = omniback.Dict
+register = omniback.register
+
+
+# -----------------------
+def set_fast_dlpack():
+    import tvm_ffi
+    tvm_ffi._optional_torch_c_dlpack.load_torch_c_dlpack_extension()
+    tvm_ffi._optional_torch_c_dlpack.patch_torch_cuda_stream_protocol()
+    if hasattr(torch.Tensor, "__dlpack_c_exchange_api__"):
+        # type: ignore[attr-defined]
+        api_attr = torch.Tensor.__dlpack_c_exchange_api__
+        if api_attr:
+            # PyCapsule - extract the pointer as integer
+            pythonapi = ctypes.pythonapi
+            # Set restype to c_size_t to get integer directly (avoids c_void_p quirks)
+            pythonapi.PyCapsule_GetPointer.restype = ctypes.c_size_t
+            pythonapi.PyCapsule_GetPointer.argtypes = [
+                ctypes.py_object, ctypes.c_char_p]
+            capsule_name = b"dlpack_exchange_api"
+            api_ptr = pythonapi.PyCapsule_GetPointer(api_attr, capsule_name)
+            assert api_ptr != 0, "API pointer from PyCapsule should not be NULL"
+            omniback.ffi.set_dlpack_exchange_api(api_ptr)
+
+
+if torch.__version__ >= torch.torch_version.TorchVersion("2.3.0"):
+    if ORI_TVM_FFI_DISABLE_TORCH_C_DLPACK == "0":
+        os.environ["TVM_FFI_DISABLE_TORCH_C_DLPACK"] = "0"
+        
+    set_fast_dlpack()

torchpipe/csrc/core_cuda/SyncTensor.cpp

@@ -0,0 +1,148 @@
+#include "c10/cuda/CUDAFunctions.h"
+#include "c10/cuda/CUDAStream.h"
+
+#include "ATen/cuda/CUDAEvent.h"
+#include "core_cuda/SyncTensor.hpp"
+#include "helper/torch.hpp"
+#include "helper_cuda/torch.hpp"
+
+#include "omniback/core/helper.hpp"
+#include "omniback/helper/resource_pool.hpp"
+
+#include <tvm/ffi/extra/c_env_api.h>
+#include "helper/dlpack_helper.hpp"
+
+namespace torchpipe {
+void SyncTensor::impl_init(
+    const std::unordered_map<std::string, std::string>& config,
+    const dict& kwargs) {
+  auto dep = om::parser_v2::get_opt_dependency_name(this, config);
+
+  const auto cls_name = om::get_cls_name(this, "SyncTensor");
+
+  auto iter = config.find(TASK_INDEX_KEY);
+  OMNI_ASSERT(
+      iter != config.end(),
+      "You are not in an independent thread mode(TASK_INDEX_KEY was not detected. Maybe use `With[StreamPool, *]` instead");
+
+  independent_thread_index_ = std::stoi(iter->second);
+  OMNI_ASSERT(independent_thread_index_ >= 0);
+
+  if (config.find("device_id") != config.end()) {
+    throw std::runtime_error(
+        "SyncTensor: device_id is not supported by SyncTensor yet.");
+  }
+
+  const auto device_id_int = c10::cuda::current_device(); //- 1;
+
+  c10::cuda::getCurrentCUDAStream().synchronize();
+
+  bNeedSync_ = torch_not_use_default_stream(device_id_int, true);
+  // Schedule保证了init和forward在同一个线程
+  OMNI_ASSERT(
+      bNeedSync_,
+      "This backend can only be used in default current stream. may be use `With[StreamPool,*]` instead.");
+
+  TVMFFIStreamHandle out_original_stream{nullptr};
+  TVMFFIStreamHandle in_stream = c10::cuda::getCurrentCUDAStream().stream();
+  TVM_FFI_ICHECK(
+      0 ==
+      TVMFFIEnvSetStream(
+          kDLCUDA, device_id_int, in_stream, &out_original_stream));
+  DLPackManagedTensorAllocator opt_out_original_allocator{nullptr};
+  TVM_FFI_ICHECK(nullptr == TVMFFIEnvGetDLPackManagedTensorAllocator());
+  // https: //
+  // github.com/apache/tvm-ffi/blob/6e7cafab78cb007d066bc860c600e2ba80b4d1a7/python/tvm_ffi/utils/_build_optional_torch_c_dlpack.py#L535
+  TVM_FFI_ICHECK(
+      0 ==
+      TVMFFIEnvSetDLPackManagedTensorAllocator(
+          torch_allocator(), 0, &opt_out_original_allocator));
+  TVM_FFI_ICHECK(nullptr == out_original_stream);
+  // TVM_FFI_ICHECK(nullptr == opt_out_original_allocator);
+
+  if (dep && !owned_backend_) {
+    owned_backend_ = om::init_backend(*dep, config, kwargs);
+  }
+  // ManagedTensorAllocator
+
+  c10::cuda::getCurrentCUDAStream().synchronize();
+
+  OMNI_ASSERT(c10::cuda::device_count() >= 1);
+
+  return;
+}
+
+void SyncTensor::impl_forward(const std::vector<dict>& ios) {
+  // std::string sync_stream = dict_get<std::string>(ios[0],
+  // "sync_stream",true);
+
+  if (owned_backend_) {
+    static auto curr_stream = c10::cuda::getCurrentCUDAStream(-1);
+    static auto default_stream = c10::cuda::getDefaultCUDAStream(-1);
+
+    event_.record(default_stream);
+    event_.block(curr_stream);
+  }
+  impl_dep_forward(ios);
+
+  c10::cuda::getCurrentCUDAStream().synchronize();
+}
+
+OMNI_REGISTER(om::Backend, SyncTensor, "SyncTensor,StreamGuard");
+
+class TorchStreamPool : public om::Backend {
+  void impl_init(
+      const std::unordered_map<std::string, std::string>& params,
+      const dict& options) override {
+    auto [args, kwargs] =
+        om::parser_v2::get_args_kwargs(this, "TorchStreamPool", params);
+    om::str::try_update<size_t>(kwargs, "max_stream", max_stream_count_);
+    OMNI_ASSERT(max_stream_count_ > 0 && max_stream_count_ < 32);
+    stream_pool_ =
+        std::make_unique<om::pool::ResourcePool<size_t>>(max_stream_count_);
+    for (size_t i = 0; i < max_stream_count_; ++i) {
+      auto stream = c10::cuda::getStreamFromPool(true, -1);
+      stream_event_.emplace_back(
+          StreamWithEvent{std::move(stream), at::cuda::CUDAEvent()});
+    }
+  }
+
+  void impl_forward_with_dep(const std::vector<om::dict>& ios, om::Backend& dep)
+      override {
+    auto index = stream_pool_->acquire();
+    om::pool::ResourcePool<size_t>::lease_guard guard(
+        stream_pool_.get(), index);
+    auto& se = stream_event_.at(index);
+    auto original_stream = c10::cuda::getCurrentCUDAStream(-1);
+    if (se.stream != original_stream) {
+      c10::cuda::CUDAStreamGuard s(se.stream);
+      // https://stackoverflow.com/questions/15501699/cudastreamwaitevent-does-not-seem-to-wait
+      se.event.record(original_stream);
+      se.event.block(se.stream);
+
+      dep.safe_forward(ios);
+
+      se.event.record(se.stream);
+      se.event.block(original_stream);
+    } else {
+      dep.safe_forward(ios);
+    }
+  }
+
+  [[nodiscard]] uint32_t impl_max() const override {
+    return std::numeric_limits<uint32_t>::max();
+  }
+
+ private:
+  struct StreamWithEvent {
+    c10::cuda::CUDAStream stream;
+    at::cuda::CUDAEvent event;
+    // https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__EVENT.html#group__CUDART__EVENT_1gf4fcb74343aa689f4159791967868446
+  };
+  size_t max_stream_count_{1};
+  std::vector<StreamWithEvent> stream_event_;
+  std::unique_ptr<om::pool::ResourcePool<size_t>> stream_pool_;
+};
+
+OMNI_REGISTER(om::Backend, TorchStreamPool, "TorchStreamPool, StreamPool");
+} // namespace torchpipe

torchpipe/csrc/core_cuda/SyncTensor.hpp

@@ -0,0 +1,134 @@
+#pragma once
+
+#include <omniback/extension.hpp>
+#include <string>
+#include <unordered_set>
+
+#include <ATen/cuda/CUDAEvent.h>
+#include "c10/cuda/CUDAGuard.h"
+
+using om::dict;
+
+namespace torchpipe {
+class SyncTensor : public om::Backend {
+ public:
+  /**
+   * @brief
+   * Initialization, determines whether the default stream is being used; if
+   * so, and in independent thread mode, it will bind a new CUDA stream to the
+   * current thread for GPU asynchronous execution. (Renamed from Torch)
+   *
+   * @param TASK_INDEX_KEY When the parameter is not null, it represents
+   * independent thread mode, at this time it can be assumed that init and
+   * forward are running in the same independent thread. Check if the CUDA
+   * stream is the default stream, if not, we bind the thread to a new stream
+   * and set bNeedSync_ to true, otherwise do nothing.
+   * @param SyncTensor::backend Default is Identity. The backend it forwards
+   * execution to.
+   * @note Usage: SyncTensor[A], SequentialV0[A,B,C,SyncTensor] or
+   * SequentialV0[A,B,SyncTensor[C]]. For serial units, such as
+   * SequentialV0[SyncTensor[A],SyncTensor[B]], it will initialize in reverse
+   * order and forward in order: SyncTensor[B].init -> SyncTensor[A].init ->
+   * SyncTensor[A].forward
+   * -> SyncTensor[B].forward;  SyncTensor[A] is not the default stream at
+   * initialization, so it does not need to set a new stream, and it does not
+   * need to be responsible for stream synchronization during forward, at this
+   * time if SyncTensor[B] has set a new stream, then SyncTensor[B] is
+   * responsible for stream synchronization;
+   *
+   *  @ref SequentialV0 and other containers can ensure that the
+   * initialization order of their child backends is opposite to the forward
+   * order, therefore, even in complex situations, the correct stream
+   * synchronization timing can be obtained.
+   */
+  virtual void impl_init(
+      const std::unordered_map<std::string, std::string>&,
+      const dict&) override;
+
+  //   virtual void post_init(
+  //       const std::unordered_map<std::string, std::string>&,
+  //       const dict&) override;
+  /**
+   * @brief
+   * 如果init时绑定了新的流（也就是bNeedSync_==true），则forward时当子后端执行完毕后执行当前流上的同步。
+   *
+   */
+  virtual void impl_forward(const std::vector<dict>& ios) override;
+
+  [[nodiscard]] uint32_t impl_max() const override {
+    if (owned_backend_) {
+      return owned_backend_->max();
+    } else {
+      return std::numeric_limits<uint32_t>::max(); // default
+    }
+  }
+
+  [[nodiscard]] uint32_t impl_min() const override {
+    if (owned_backend_) {
+      return owned_backend_->min();
+    } else {
+      return 1; // default
+    }
+  }
+
+ private:
+  void impl_dep_forward(const std::vector<dict>& ios) {
+    if (owned_backend_)
+      owned_backend_->forward(ios);
+    else {
+      for (const auto& io : ios) {
+        (*io)[TASK_RESULT_KEY] = io->at(TASK_DATA_KEY);
+      }
+    }
+  }
+
+  bool bNeedSync_ = false;
+  std::unique_ptr<Backend> owned_backend_;
+  int independent_thread_index_{-1};
+  std::optional<c10::cuda::CUDAStream> stream_;
+  at::cuda::CUDAEvent event_;
+};
+
+// class StreamGuard : public om::DependencyV0 {
+//  public:
+//   /**
+//    * @brief
+//    * Initialization, determines whether the default stream is being used; if
+//    * so, and in independent thread mode, it will bind a new CUDA stream to
+//    the
+//    * current thread for GPU asynchronous execution. (Renamed from Torch)
+//    *
+//    * @param TASK_INDEX_KEY When the parameter is not null, it represents
+//    * independent thread mode, at this time it can be assumed that init and
+//    * forward are running in the same independent thread. Check if the CUDA
+//    * stream is the default stream, if not, we bind the thread to a new stream
+//    * and set bNeedSync_ to true, otherwise do nothing.
+//    * @param StreamGuard::backend Default is Identity. The backend it forwards
+//    * execution to.
+//    *
+//    *  @ref SequentialV0 and other containers can ensure that the
+//    * initialization order of their child backends is opposite to the forward
+//    * order, therefore, even in complex situations, the correct stream
+//    * synchronization timing can be obtained.
+//    */
+//   virtual void pre_init(
+//       const std::unordered_map<std::string, std::string>&,
+//       const dict&) override;
+
+//   virtual void post_init(
+//       const std::unordered_map<std::string, std::string>&,
+//       const dict&) override;
+//   /**
+//    * @brief
+//    * 如果init时绑定了新的流（也就是bNeedSync_==true），则forward时当子后端执行完毕后执行当前流上的同步。
+//    *
+//    */
+//   virtual void custom_forward_with_dep(const std::vector<dict>&, Backend*)
+//       override;
+
+//  private:
+//   //   bool bNeedSync_ = false;
+//   std::optional<c10::cuda::CUDAStream> stream_;
+//   std::unique_ptr<c10::cuda::CUDAStreamGuard> stream_guard_;
+// };
+} // namespace torchpipe

torchpipe/csrc/core_cuda/kvcacheTensor.cpp

@@ -0,0 +1,160 @@
+#include <numeric>
+#include <unordered_map>
+
+#include <cuda_runtime_api.h>
+#include "c10/cuda/CUDAStream.h"
+#include "omniback/builtin/page_table.hpp"
+#include "omniback/core/backend.hpp"
+
+#include <torch/torch.h>
+#include "omniback/helper/base_logging.hpp"
+
+namespace torchpipe {
+class LocationManager {
+ public:
+};
+
+using namespace om;
+class FIAppendTensor : public om::BackendOne {
+ private:
+  size_t max_num_req_{16};
+  size_t max_num_page_{0};
+  size_t max_context_len_{4096};
+  size_t num_layer_{32};
+  size_t head_num_{32};
+  size_t head_dim_{128};
+  size_t page_size_{16};
+
+ private:
+  size_t max_num_page_per_seq_{0};
+  bool inited_{false};
+  // std::vector<torch::Tensor> k_;
+  // std::vector<torch::Tensor> v_;
+  std::unique_ptr<PageTable> pool_;
+
+ private:
+  void impl_init(
+      const std::unordered_map<string, string>& params,
+      const dict& options) override {
+    str::try_update(params, "max_num_req", max_num_req_);
+    max_num_page_ = str::get<size_t>(params, "max_num_page");
+    // str::try_update(params, "max_context_len", max_context_len_);
+    // str::try_update(params, "num_layer", num_layer_);
+    // str::try_update(params, "head_num", head_num_);
+    // str::try_update(params, "head_dim", head_dim_);
+    OMNI_ASSERT(max_context_len_ % page_size_ == 0);
+    max_num_page_per_seq_ = max_context_len_ / page_size_;
+  }
+
+  // void get(torch::Tensor kv_append_length) {
+  //   for (size_t i = 0; i < kv_append_length.size(0); ++i) {
+  //   }
+
+  //   torch::Tensor batch_indices;
+  //   torch::Tensor positions;
+
+  //   torch::Tensor kv_page_indptr;
+  //   torch::Tensor kv_last_page_len;
+  //   torch::Tensor kv_append_indptr;
+  // }
+  torch::Tensor vec2tensor(const std::vector<int>& data) {
+    thread_local auto options = torch::TensorOptions()
+                                    .device(torch::kCUDA, -1)
+                                    .dtype(torch::kInt) // torch::kByte
+                                    .layout(torch::kStrided)
+                                    .requires_grad(false);
+    torch::Tensor re =
+        torch::empty({static_cast<int64_t>(data.size())}, options);
+    cudaError_t cuda_status = cudaMemcpyAsync(
+        re.data_ptr(), // 目标设备指针
+        data.data(), // 主机源指针
+        data.size() * sizeof(int), // 字节大小
+        cudaMemcpyHostToDevice, // 传输方向
+        c10::cuda::getCurrentCUDAStream());
+    if (cuda_status != cudaSuccess) {
+      throw std::runtime_error(
+          "CUDA 拷贝失败: " + std::string(cudaGetErrorString(cuda_status)));
+    }
+    return re;
+  }
+
+  void forward(const dict& io) override {
+    if (!inited_)
+      lazy_init();
+    // in
+    // id, type(prefill, decode), seq_len,
+    // out
+    bool success = true;
+    std::vector<id_type> id = dict_gets<id_type>(io, "request_ids");
+    // https://docs.flashinfer.ai/generated/flashinfer.page.append_paged_kv_cache.html#flashinfer.page.append_paged_kv_cache
+    torch::Tensor seq_lens = dict_get<torch::Tensor>(io, "kv_append_length");
+    OMNI_ASSERT(id.size() == seq_lens.size(0) && seq_lens.is_cpu());
+    size_t total{0};
+
+    for (size_t i = 0; i < id.size(); ++i) {
+      SPDLOG_INFO("id = {}", id[i]);
+      auto seq_len = seq_lens[i].item<int>();
+      success = success && pool_->alloc(id[i], seq_len);
+      OMNI_ASSERT(success);
+      const auto& infor = pool_->page_info(id[i]);
+      total += infor.kv_page_indices.size();
+    }
+    std::vector<int> kv_page_indices;
+    kv_page_indices.reserve(total);
+
+    std::vector<int> kv_page_indptr(1 + id.size(), 0);
+    std::vector<int> kv_last_page_len(id.size());
+    for (size_t i = 0; i < id.size(); ++i) {
+      const auto& infor = pool_->page_info(id[i]);
+      kv_page_indices.insert(
+          kv_page_indices.end(),
+          infor.kv_page_indices.begin(),
+          infor.kv_page_indices.end());
+      kv_page_indptr[i + 1] = kv_page_indptr[i] + infor.kv_page_indices.size();
+      kv_last_page_len[i] = infor.kv_last_page_len;
+    }
+
+    auto kv_page_indices_CUDA = vec2tensor(kv_page_indices);
+    auto kv_page_indptr_CUDA = vec2tensor(kv_page_indptr);
+    auto kv_last_page_len_CUDA = vec2tensor(kv_last_page_len);
+
+    (*io)["kv_page_indices"] = kv_page_indices_CUDA;
+    (*io)["kv_page_indptr"] = kv_page_indptr_CUDA;
+    (*io)["kv_last_page_len"] = kv_last_page_len_CUDA;
+  }
+
+  void lazy_init() {
+    // if (max_num_page_ == 0) {
+    //   auto stats = torch::cuda::memory_stats(-1);
+    //   int64_t free_memory = stats.free_bytes; // 剩余显存字节数
+
+    //   max_num_page_ = static_cast<size_t>(
+    //       (free_memory * 0.9) /
+    //       (page_size_ * head_num_ * head_dim_ * 2 /*kv*/ * num_layer_ *
+    //        2 /*fp16 */));
+    // }
+
+    // k_.resize(num_layer_);
+    // v_.resize(num_layer_);
+    // auto options = torch::TensorOptions()
+    //                    .device(torch::kCUDA, -1)
+    //                    .dtype(torch::kFloat16)
+    //                    .layout(torch::kStrided)
+    //                    .requires_grad(false);
+    // for (size_t layer_index = 0; layer_index < num_layer_; ++layer_index) {
+    //   k_[layer_index] = torch::empty(
+    //       {max_num_page_, page_size_, head_num_, head_dim_},
+    //       options,
+    //       torch::MemoryFormat::Contiguous);
+    //   v_[layer_index] = torch::empty(
+    //       {max_num_page_, page_size_, head_num_, head_dim_},
+    //       options,
+    //       torch::MemoryFormat::Contiguous);
+    // }
+    pool_ =
+        std::make_unique<PageTable>(max_num_req_, max_num_page_, page_size_);
+    inited_ = true;
+  }
+};
+OMNI_REGISTER_BACKEND(FIAppendTensor);
+} // namespace torchpipe

torchpipe/csrc/helper/dlpack_helper.hpp

@@ -0,0 +1,26 @@
+#pragma once
+#include <string>
+#include <unordered_set>
+
+#include <tvm/ffi/extra/c_env_api.h>
+#include <tvm/ffi/c_api.h>
+#include <tvm/ffi/error.h>
+#include "omniback/ffi/types.hpp"
+
+
+namespace torchpipe {
+
+inline DLPackManagedTensorAllocator torch_allocator() {
+  DLPackExchangeAPI* api =
+      reinterpret_cast<DLPackExchangeAPI*>(om::ffi::dlpack_exchange_api());
+  if (api)
+   { 
+    static DLPackManagedTensorAllocator& alloc = api->managed_tensor_allocator;
+    TVM_FFI_ICHECK(alloc);
+    return alloc;
+  }else{
+    return nullptr;
+  }
+}
+
+}

torchpipe/csrc/helper/mat.cpp

@@ -0,0 +1,18 @@
+#include "helper/mat.hpp"
+// #include <torch/torch.h>
+#include "helper/torch.hpp"
+
+namespace torchpipe::convert{
+ImageData TorchAny2ImageData(om::any tensor) {
+  torch::Tensor data = tensor.cast<torch::Tensor>();
+  return torchpipe::torch2ImageData(data);
+}
+
+om::any imageDataToAnyTorchCPU(const convert::ImageData& img) {
+  return torchpipe::imageDataToTorchCPU(img);
+}
+
+om::any imageDataToAnyTorchGPU(const convert::ImageData& img) {
+  return torchpipe::imageDataToTorchCPU(img).cuda();
+}
+}

torchpipe/csrc/helper/mat.hpp

@@ -0,0 +1,22 @@
+#pragma once
+#include <cstddef>
+#include <cstdint>
+#include <functional>
+#include "omniback/core/any.hpp"
+
+namespace torchpipe::convert{
+
+struct ImageData {
+  void* data = nullptr;
+  size_t rows = 0;
+  size_t cols = 0;
+  size_t channels = 0;
+  bool is_float = false; // true: float32, false: uint8
+  std::function<void(void*)> deleter;
+};
+ImageData TorchAny2ImageData(om::any tensor);
+
+om::any imageDataToAnyTorchCPU(const convert::ImageData& img);
+
+om::any imageDataToAnyTorchGPU(const convert::ImageData& img);
+}

torchpipe/csrc/helper/net_info.cpp

@@ -0,0 +1,57 @@
+#include <optional>
+#include "helper/net_info.hpp"
+
+namespace torchpipe {
+inline bool is_all_positive(const NetIOInfo::Dims64& dims) {
+    if (dims.nbDims <= 0) return false;
+    for (size_t index = 0; index < dims.nbDims; ++index) {
+        if (dims.d[index] <= 0) return false;
+    }
+    return true;
+}
+
+bool is_all_positive(NetIOInfos& info) {
+    for (const auto& item : info.first) {
+        if (!is_all_positive(item.min)) return false;
+    }
+    for (const auto& item : info.second) {
+        if (!is_all_positive(item.min)) return false;
+    }
+    for (const auto& item : info.first) {
+        if (!is_all_positive(item.max)) return false;
+    }
+    for (const auto& item : info.second) {
+        if (!is_all_positive(item.max)) return false;
+    }
+
+    return true;
+}
+
+size_t elementSize(NetIOInfo::DataType info) {
+    switch (info) {
+        case NetIOInfo::DataType::INT4:
+        case NetIOInfo::DataType::FP4:
+            return 1;  // 4 bits = 0.5 bytes
+        case NetIOInfo::DataType::INT8:
+        case NetIOInfo::DataType::UINT8:
+        case NetIOInfo::DataType::BOOL:
+        case NetIOInfo::DataType::FP8:
+            return 1;  // 8 bits = 1 byte
+        case NetIOInfo::DataType::INT32:
+        case NetIOInfo::DataType::FP32:
+        case NetIOInfo::DataType::BF32:
+            return 4;  // 32 bits = 4 bytes
+        case NetIOInfo::DataType::INT64:
+            return 8;  // 64 bits = 8 bytes
+        case NetIOInfo::DataType::FP16:
+        case NetIOInfo::DataType::BF16:
+            return 2;  // 16 bits = 2 bytes
+        case NetIOInfo::DataType::RESERVED_INT:
+        case NetIOInfo::DataType::RESERVED_FP:
+        case NetIOInfo::DataType::RESERVED_BF:
+        case NetIOInfo::DataType::UNKNOWN:
+        default:
+            return 0;  // Unknown or reserved types
+    }
+}
+}  // namespace torchpipe