numpy_bridge.hpp

Go to the documentation of this file.

// SPDX-License-Identifier: AGPL-3.0-or-later
// Copyright 2026 Johnson Ogundeji
#pragma once
 
 
#include "signet/ai/tensor_bridge.hpp"
#include "signet/error.hpp"
 
#include <cstdint>
#include <cstdlib>
#include <cstring>
#include <string>
#include <utility>
#include <vector>
 
namespace signet::forge {
 
 
enum class DLDeviceType : int32_t {
    kDLCPU      = 1,   
    kDLCUDA     = 2,   
    kDLCUDAHost = 3,   
    kDLROCM     = 10,  
    kDLMetal    = 8,   
    kDLVulkan   = 7    
};
 
enum class DLDataTypeCode : uint8_t {
    kDLInt    = 0,   
    kDLUInt   = 1,   
    kDLFloat  = 2,   
    kDLBfloat = 4    
};
 
struct DLDataType {
    DLDataTypeCode code;       
    uint8_t        bits;       
    uint16_t       lanes = 1;  
};
 
struct DLDevice {
    DLDeviceType device_type = DLDeviceType::kDLCPU; 
    int32_t      device_id   = 0;                     
};
 
struct DLTensor {
    void*       data;        
    DLDevice    device;      
    int32_t     ndim;        
    DLDataType  dtype;       
    int64_t*    shape;       
    int64_t*    strides;     
    uint64_t    byte_offset; 
};
 
struct DLManagedTensor {
    DLTensor dl_tensor;                   
    void*    manager_ctx;                 
    void     (*deleter)(DLManagedTensor*); 
};
 
 
 
inline DLDataType to_dlpack_dtype(TensorDataType dtype) {
    DLDataType dt;
    dt.lanes = 1;
 
    switch (dtype) {
        case TensorDataType::FLOAT32:
            dt.code = DLDataTypeCode::kDLFloat;
            dt.bits = 32;
            break;
        case TensorDataType::FLOAT64:
            dt.code = DLDataTypeCode::kDLFloat;
            dt.bits = 64;
            break;
        case TensorDataType::FLOAT16:
            dt.code = DLDataTypeCode::kDLFloat;
            dt.bits = 16;
            break;
        case TensorDataType::INT32:
            dt.code = DLDataTypeCode::kDLInt;
            dt.bits = 32;
            break;
        case TensorDataType::INT64:
            dt.code = DLDataTypeCode::kDLInt;
            dt.bits = 64;
            break;
        case TensorDataType::INT16:
            dt.code = DLDataTypeCode::kDLInt;
            dt.bits = 16;
            break;
        case TensorDataType::INT8:
            dt.code = DLDataTypeCode::kDLInt;
            dt.bits = 8;
            break;
        case TensorDataType::UINT8:
            dt.code = DLDataTypeCode::kDLUInt;
            dt.bits = 8;
            break;
        case TensorDataType::BOOL:
            // DLPack convention: booleans are represented as uint8
            dt.code = DLDataTypeCode::kDLUInt;
            dt.bits = 8;
            break;
    }
 
    return dt;
}
 
inline expected<TensorDataType> from_dlpack_dtype(DLDataType dl_dtype) {
    if (dl_dtype.lanes != 1) {
        return Error{ErrorCode::UNSUPPORTED_TYPE,
                     "multi-lane DLPack dtypes are not supported"};
    }
 
    switch (dl_dtype.code) {
        case DLDataTypeCode::kDLFloat:
            switch (dl_dtype.bits) {
                case 16: return TensorDataType::FLOAT16;
                case 32: return TensorDataType::FLOAT32;
                case 64: return TensorDataType::FLOAT64;
                default:
                    return Error{ErrorCode::UNSUPPORTED_TYPE,
                                 "unsupported DLPack float bit width: "
                                 + std::to_string(dl_dtype.bits)};
            }
 
        case DLDataTypeCode::kDLInt:
            switch (dl_dtype.bits) {
                case 8:  return TensorDataType::INT8;
                case 16: return TensorDataType::INT16;
                case 32: return TensorDataType::INT32;
                case 64: return TensorDataType::INT64;
                default:
                    return Error{ErrorCode::UNSUPPORTED_TYPE,
                                 "unsupported DLPack int bit width: "
                                 + std::to_string(dl_dtype.bits)};
            }
 
        case DLDataTypeCode::kDLUInt:
            switch (dl_dtype.bits) {
                case 8:  return TensorDataType::UINT8;
                default:
                    return Error{ErrorCode::UNSUPPORTED_TYPE,
                                 "unsupported DLPack uint bit width: "
                                 + std::to_string(dl_dtype.bits)
                                 + " (only uint8 is supported)"};
            }
 
        case DLDataTypeCode::kDLBfloat:
            return Error{ErrorCode::UNSUPPORTED_TYPE,
                         "bfloat16 is not supported by Signet TensorDataType"};
 
        default:
            return Error{ErrorCode::UNSUPPORTED_TYPE,
                         "unknown DLPack data type code"};
    }
}
 
 
namespace detail {
 
struct DLPackViewCtx {
    std::vector<int64_t> shape;   
    // strides left empty for contiguous tensors (dl_tensor.strides = nullptr)
};
 
struct DLPackOwnedCtx {
    OwnedTensor          owned_tensor;   
    std::vector<int64_t> shape;          
};
 
inline void dlpack_view_deleter(DLManagedTensor* self) {
    if (self == nullptr) return;
    if (self->manager_ctx != nullptr) {
        delete static_cast<DLPackViewCtx*>(self->manager_ctx);
        self->manager_ctx = nullptr;
    }
    delete self;
}
 
inline void dlpack_owned_deleter(DLManagedTensor* self) {
    if (self == nullptr) return;
    if (self->manager_ctx != nullptr) {
        delete static_cast<DLPackOwnedCtx*>(self->manager_ctx);
        self->manager_ctx = nullptr;
    }
    delete self;
}
 
} // namespace detail
 
class NumpyBridge {
public:
    static inline expected<DLManagedTensor*> export_tensor(const TensorView& tensor) {
        if (!tensor.is_valid()) {
            return Error{ErrorCode::INTERNAL_ERROR,
                         "cannot export invalid tensor to DLPack"};
        }
 
        if (!tensor.is_contiguous()) {
            return Error{ErrorCode::UNSUPPORTED_TYPE,
                         "non-contiguous tensor cannot be exported to DLPack; "
                         "call clone() first"};
        }
 
        // Build the manager context (owns shape array, NOT the data)
        auto* ctx = new detail::DLPackViewCtx();
        ctx->shape = tensor.shape().dims;
 
        // Build the DLManagedTensor
        auto* managed = new DLManagedTensor();
        managed->manager_ctx = ctx;
        managed->deleter     = detail::dlpack_view_deleter;
 
        DLTensor& dl = managed->dl_tensor;
        dl.data        = const_cast<void*>(tensor.data());
        dl.device      = DLDevice{DLDeviceType::kDLCPU, 0};
        dl.ndim        = static_cast<int32_t>(tensor.shape().ndim());
        dl.dtype       = to_dlpack_dtype(tensor.dtype());
        dl.shape       = ctx->shape.data();
        dl.strides     = nullptr; // contiguous C-order
        dl.byte_offset = 0;
 
        return managed;
    }
 
    static inline expected<DLManagedTensor*> export_owned_tensor(OwnedTensor&& tensor) {
        TensorView view = tensor.view();
 
        if (!view.is_valid()) {
            return Error{ErrorCode::INTERNAL_ERROR,
                         "cannot export invalid OwnedTensor to DLPack"};
        }
 
        if (!view.is_contiguous()) {
            return Error{ErrorCode::UNSUPPORTED_TYPE,
                         "non-contiguous OwnedTensor cannot be exported to DLPack"};
        }
 
        // Build the manager context (owns the tensor data + shape array)
        auto* ctx = new detail::DLPackOwnedCtx();
        ctx->owned_tensor = std::move(tensor);
        ctx->shape = view.shape().dims;
 
        // Build the DLManagedTensor
        auto* managed = new DLManagedTensor();
        managed->manager_ctx = ctx;
        managed->deleter     = detail::dlpack_owned_deleter;
 
        DLTensor& dl = managed->dl_tensor;
        dl.data        = const_cast<void*>(view.data());
        dl.device      = DLDevice{DLDeviceType::kDLCPU, 0};
        dl.ndim        = static_cast<int32_t>(view.shape().ndim());
        dl.dtype       = to_dlpack_dtype(view.dtype());
        dl.shape       = ctx->shape.data();
        dl.strides     = nullptr; // contiguous C-order
        dl.byte_offset = 0;
 
        return managed;
    }
 
    static inline expected<TensorView> import_tensor(const DLManagedTensor* managed) {
        if (managed == nullptr) {
            return Error{ErrorCode::INTERNAL_ERROR,
                         "null DLManagedTensor pointer"};
        }
 
        const DLTensor& dl = managed->dl_tensor;
 
        // Only CPU tensors are supported
        if (dl.device.device_type != DLDeviceType::kDLCPU &&
            dl.device.device_type != DLDeviceType::kDLCUDAHost) {
            return Error{ErrorCode::UNSUPPORTED_TYPE,
                         "only CPU/CUDAHost DLPack tensors can be imported"};
        }
 
        if (dl.data == nullptr) {
            return Error{ErrorCode::INTERNAL_ERROR,
                         "DLTensor data pointer is null"};
        }
 
        if (dl.ndim <= 0) {
            return Error{ErrorCode::INTERNAL_ERROR,
                         "DLTensor has zero or negative ndim"};
        }
 
        if (dl.shape == nullptr) {
            return Error{ErrorCode::INTERNAL_ERROR,
                         "DLTensor shape pointer is null"};
        }
 
        // Reject strided (non-contiguous) tensors for TensorView import.
        // TensorView assumes C-contiguous layout.
        if (dl.strides != nullptr) {
            // Verify strides match C-contiguous layout
            int64_t expected_stride = 1;
            for (int32_t d = dl.ndim - 1; d >= 0; --d) {
                if (dl.strides[d] != expected_stride) {
                    return Error{ErrorCode::UNSUPPORTED_TYPE,
                                 "strided (non-contiguous) DLPack tensor cannot "
                                 "be imported as TensorView; use import_tensor_copy()"};
                }
                expected_stride *= dl.shape[d];
            }
        }
 
        auto dtype_result = from_dlpack_dtype(dl.dtype);
        if (!dtype_result) {
            return dtype_result.error();
        }
 
        // CWE-20: Improper Input Validation (DLPack §3.2 max_ndim)
        // L14: Reject unreasonable ndim values (DLPack spec uses int32_t)
        if (dl.ndim > 32) {
            return Error{ErrorCode::INVALID_ARGUMENT,
                         "DLTensor ndim exceeds reasonable limit (32)"};
        }
 
        TensorShape shape;
        shape.dims.reserve(static_cast<size_t>(dl.ndim));
 
        // CWE-190: Integer Overflow — validate byte_offset is within tensor data bounds
        const size_t elem_size = tensor_element_size(*dtype_result);
        size_t total_elements = 1;
        for (int32_t d = 0; d < dl.ndim; ++d) {
            if (dl.shape[d] <= 0) {
                return Error{ErrorCode::INVALID_ARGUMENT,
                             "DLTensor shape dimension must be positive"};
            }
            const size_t dim_val = static_cast<size_t>(dl.shape[d]);
            if (total_elements > SIZE_MAX / dim_val) {
                return Error{ErrorCode::INVALID_ARGUMENT,
                             "DLTensor shape product overflows size_t"};
            }
            total_elements *= dim_val;
            shape.dims.push_back(dl.shape[d]);
        }
        // CWE-190: Integer Overflow — num_elements*elem_size overflow check
        if (elem_size > 0 && total_elements > SIZE_MAX / elem_size) {
            return Error{ErrorCode::INVALID_ARGUMENT,
                         "DLTensor total byte size overflows size_t"};
        }
        const size_t total_size = total_elements * elem_size;
        if (dl.byte_offset > total_size) {
            return Error{ErrorCode::INVALID_ARGUMENT,
                         "DLPack byte_offset out of range"};
        }
 
        // Apply byte_offset
        void* data_ptr = static_cast<uint8_t*>(dl.data) + dl.byte_offset;
 
        return TensorView(data_ptr, shape, *dtype_result);
    }
 
    static inline expected<OwnedTensor> import_tensor_copy(const DLManagedTensor* managed) {
        if (managed == nullptr) {
            return Error{ErrorCode::INTERNAL_ERROR,
                         "null DLManagedTensor pointer"};
        }
 
        const DLTensor& dl = managed->dl_tensor;
 
        if (dl.device.device_type != DLDeviceType::kDLCPU &&
            dl.device.device_type != DLDeviceType::kDLCUDAHost) {
            return Error{ErrorCode::UNSUPPORTED_TYPE,
                         "only CPU/CUDAHost DLPack tensors can be imported"};
        }
 
        if (dl.data == nullptr || dl.ndim <= 0 || dl.shape == nullptr) {
            return Error{ErrorCode::INTERNAL_ERROR,
                         "invalid DLTensor (null data/shape or non-positive ndim)"};
        }
 
        auto dtype_result = from_dlpack_dtype(dl.dtype);
        if (!dtype_result) {
            return dtype_result.error();
        }
 
        TensorDataType dtype = *dtype_result;
        const size_t elem_size = tensor_element_size(dtype);
 
        TensorShape shape;
        shape.dims.assign(dl.shape, dl.shape + dl.ndim);
        const size_t num_elements = shape.num_elements();
 
        // CWE-190: Integer Overflow — check for multiplication overflow before allocating
        if (elem_size != 0 && num_elements > SIZE_MAX / elem_size) {
            return Error{ErrorCode::INVALID_ARGUMENT,
                         "DLPack tensor size overflow (num_elements * elem_size)"};
        }
 
        // Source data pointer with byte offset applied
        const uint8_t* src_base = static_cast<const uint8_t*>(dl.data)
                                + dl.byte_offset;
 
        // Check if contiguous -- if so, fast memcpy path
        bool is_contiguous = (dl.strides == nullptr);
        if (!is_contiguous && dl.strides != nullptr) {
            // Check if strides match C-contiguous
            int64_t expected_stride = 1;
            is_contiguous = true;
            for (int32_t d = dl.ndim - 1; d >= 0; --d) {
                if (dl.strides[d] != expected_stride) {
                    is_contiguous = false;
                    break;
                }
                expected_stride *= dl.shape[d];
            }
        }
 
        if (is_contiguous) {
            // Fast path: contiguous data, direct memcpy
            OwnedTensor result(shape, dtype);
            std::memcpy(result.data(), src_base, num_elements * elem_size);
            return result;
        }
 
        // Slow path: strided data, element-by-element copy.
        // Walk the multi-dimensional index space and compute source offsets
        // from strides.
        OwnedTensor result(shape, dtype);
        uint8_t* dst = static_cast<uint8_t*>(result.data());
 
        // Multi-index iteration
        std::vector<int64_t> idx(static_cast<size_t>(dl.ndim), 0);
        for (size_t flat = 0; flat < num_elements; ++flat) {
            // Compute source byte offset from strides
            int64_t src_elem_offset = 0;
            for (int32_t d = 0; d < dl.ndim; ++d) {
                src_elem_offset += idx[static_cast<size_t>(d)]
                                 * dl.strides[static_cast<size_t>(d)];
            }
 
            const uint8_t* src_elem = src_base
                + static_cast<size_t>(src_elem_offset) * elem_size;
            std::memcpy(dst + flat * elem_size, src_elem, elem_size);
 
            // Increment multi-index (row-major / C-order)
            for (int32_t d = dl.ndim - 1; d >= 0; --d) {
                auto ud = static_cast<size_t>(d);
                idx[ud]++;
                if (idx[ud] < dl.shape[d]) break;
                idx[ud] = 0;
            }
        }
 
        return result;
    }
};
 
struct BufferInfo {
    void*                data;      
    size_t               itemsize;  
    std::string          format;    
    int64_t              ndim;      
    std::vector<int64_t> shape;     
    std::vector<int64_t> strides;   
};
 
namespace detail {
 
inline const char* tensor_dtype_to_pybuf_format(TensorDataType dtype) {
    switch (dtype) {
        case TensorDataType::FLOAT32: return "f";
        case TensorDataType::FLOAT64: return "d";
        case TensorDataType::INT32:   return "i";
        case TensorDataType::INT64:   return "l";
        case TensorDataType::INT8:    return "b";
        case TensorDataType::UINT8:   return "B";
        case TensorDataType::INT16:   return "h";
        case TensorDataType::FLOAT16: return "e";
        case TensorDataType::BOOL:    return "?";
    }
    return nullptr; // unreachable
}
 
} // namespace detail
 
inline expected<BufferInfo> to_buffer_info(const TensorView& tensor) {
    if (!tensor.is_valid()) {
        return Error{ErrorCode::INTERNAL_ERROR,
                     "cannot create BufferInfo from invalid tensor"};
    }
 
    if (!tensor.is_contiguous()) {
        return Error{ErrorCode::UNSUPPORTED_TYPE,
                     "non-contiguous tensor cannot be described by BufferInfo; "
                     "call clone() first"};
    }
 
    const char* fmt = detail::tensor_dtype_to_pybuf_format(tensor.dtype());
    if (fmt == nullptr) {
        return Error{ErrorCode::UNSUPPORTED_TYPE,
                     "tensor dtype has no Python buffer format mapping"};
    }
 
    const size_t elem_size = tensor.element_size();
    const auto& dims = tensor.shape().dims;
    const int64_t ndim = static_cast<int64_t>(dims.size());
 
    // Compute C-contiguous strides (in bytes) from innermost to outermost
    std::vector<int64_t> strides(static_cast<size_t>(ndim));
    if (ndim > 0) {
        strides[static_cast<size_t>(ndim - 1)] = static_cast<int64_t>(elem_size);
        for (int64_t d = ndim - 2; d >= 0; --d) {
            auto ud  = static_cast<size_t>(d);
            auto ud1 = static_cast<size_t>(d + 1);
            strides[ud] = strides[ud1] * dims[ud1];
        }
    }
 
    BufferInfo info;
    info.data     = const_cast<void*>(tensor.data());
    info.itemsize = elem_size;
    info.format   = fmt;
    info.ndim     = ndim;
    info.shape    = dims;
    info.strides  = std::move(strides);
 
    return info;
}
 
} // namespace signet::forge