diff --git a/include/spot-observer.h b/include/spot-observer.h index 634ff06..87bde61 100644 --- a/include/spot-observer.h +++ b/include/spot-observer.h @@ -94,11 +94,34 @@ UNITY_INTERFACE_EXPORT bool UNITY_INTERFACE_API SOb_PushNextVisionPipelineImageSetToUnityBuffers(int32_t robot_id, int32_t cam_stream_id); // Model stuff + +// Model families. The family cannot be inferred from the file: both the +// single-shot and the streaming (KV-cache) depth models are .onnx. +// SOb_MODEL_SINGLE_SHOT - stateless model; one instance may be shared by +// several pipelines (path-cached internally). +// SOb_MODEL_STREAMING - autoregressive model with per-sequence state; every +// load returns a fresh instance and each instance can +// drive exactly one pipeline at a time. +#define SOb_MODEL_SINGLE_SHOT 0 +#define SOb_MODEL_STREAMING 1 + UNITY_INTERFACE_EXPORT SObModel UNITY_INTERFACE_API SOb_LoadModel(const char* modelPath, const char* backend); +// Like SOb_LoadModel, with the model family stated explicitly (SOb_MODEL_*). +// SOb_LoadModel is equivalent to kind = SOb_MODEL_SINGLE_SHOT. +UNITY_INTERFACE_EXPORT +SObModel UNITY_INTERFACE_API SOb_LoadModelEx(const char* modelPath, const char* backend, int32_t kind); UNITY_INTERFACE_EXPORT void UNITY_INTERFACE_API SOb_UnloadModel(SObModel model); +// Live model switch: stops the vision pipeline on the given camera stream (if +// one is running) and relaunches it with `model`. The camera stream keeps +// running throughout; only the inference side is swapped. Intended pattern is to +// load every selectable model once at startup and flip between the handles -- +// no load stall and no unload hazard at switch time. +UNITY_INTERFACE_EXPORT +bool UNITY_INTERFACE_API SOb_SwitchVisionPipelineModel(int32_t robot_id, int32_t cam_stream_id, SObModel model); + // Config calls UNITY_INTERFACE_EXPORT bool UNITY_INTERFACE_API SOb_ToggleDepthCompletion(bool enable); diff --git a/onnx/set_window.py b/onnx/set_window.py new file mode 100644 index 0000000..1a1e4fc --- /dev/null +++ b/onnx/set_window.py @@ -0,0 +1,129 @@ +"""Retarget the streaming model's KV-cache retention window. + +The window is baked into the graph as 48 Constant nodes feeding the `starts` +input of the Slice that trims each new_* cache output. Rewriting them is a +graph-only edit: the 7.5 GB .onnx.data blob is referenced by filename and is +neither read nor rewritten, so each variant costs ~8 MB and shares the weights. + +VRAM saved is linear in the window, and doubled at peak because the graph reads +`past_*` while writing `new_*`: + + per retained frame = 48 * 16 * 1041 * 64 * 4 B = 195.2 MiB (fp32) + peak cache = 2 * window * per_frame + +Usage: + python set_window.py 8 + python set_window.py 8 --in mae_model_step_consolidated.onnx +""" +import argparse, pathlib, sys +import onnx +from onnx import numpy_helper +import numpy as np + +HEADS, TOKENS, HEAD_DIM, N_CACHE = 16, 1041, 64, 48 +BYTES_PER_FRAME = N_CACHE * HEADS * TOKENS * HEAD_DIM * 4 + + +def main() -> int: + ap = argparse.ArgumentParser() + ap.add_argument("window", type=int, help="frames of KV cache to retain") + ap.add_argument("--in", dest="src", default="mae_model_step_consolidated.onnx") + ap.add_argument("--out", dest="dst", default=None) + args = ap.parse_args() + + if args.window < 1: + print("window must be >= 1", file=sys.stderr) + return 2 + + here = pathlib.Path(__file__).parent + src = here / args.src + dst = here / (args.dst or f"{src.stem}_w{args.window}.onnx") + + # load_external_data=False keeps the 7.5 GB blob out of memory entirely. + model = onnx.load(src, load_external_data=False) + graph = model.graph + + # Preflight: every external-data file the graph references must sit next to + # the graph, under the exact name recorded inside it. Fail with the list + # rather than letting the checker produce a cryptic CWD-relative error. + referenced = {e.value for i in graph.initializer + for e in i.external_data if e.key == "location"} + missing = sorted(loc for loc in referenced if not (src.parent / loc).exists()) + if missing: + print("graph references external data files that are not next to it:", file=sys.stderr) + for loc in missing: + print(f" expected: {src.parent / loc}", file=sys.stderr) + print("copy the .data file(s) into that directory with exactly those names", file=sys.stderr) + return 1 + + cache_outputs = {v.name for v in graph.output if v.name.startswith("new_")} + producer = {o: n for n in graph.node for o in n.output} + + # The Slice may not feed the graph output directly: fp16 conversion and + # re-exports append pass-through ops (Cast, Identity) after it. Walk back + # from each new_* output through those to the producing Slice. + PASSTHROUGH = {"Identity", "Cast", "Squeeze", "Unsqueeze"} + def find_slice(output_name): + node, hops = producer.get(output_name), 0 + while node is not None and node.op_type in PASSTHROUGH and hops < 8: + node = producer.get(node.input[0]) + hops += 1 + return node if node is not None and node.op_type == "Slice" else None + + slice_nodes = [s for s in (find_slice(name) for name in sorted(cache_outputs)) if s is not None] + starts_inputs = {n.input[1] for n in slice_nodes} + + if len(starts_inputs) != N_CACHE: + # Fallback: find Slices by pattern -- 1-element negative `starts` -- and + # report what actually feeds the outputs so a failure is diagnosable. + feeders = sorted({(producer[n].op_type if n in producer else "") + for n in cache_outputs}) + print(f"expected {N_CACHE} cache Slice nodes, found {len(starts_inputs)}", file=sys.stderr) + print(f"ops feeding new_* outputs: {feeders}", file=sys.stderr) + print("(if these are not Slice/Cast/Identity, the export's window structure " + "changed -- send this output back)", file=sys.stderr) + return 1 + + patched, old_values = 0, set() + new_val = np.array([-args.window], dtype=np.int64) + + for node in graph.node: + if node.op_type == "Constant" and node.output and node.output[0] in starts_inputs: + for attr in node.attribute: + if attr.name == "value": + old_values.add(int(numpy_helper.to_array(attr.t)[0])) + attr.t.CopyFrom(numpy_helper.from_array(new_val, attr.t.name)) + patched += 1 + for init in graph.initializer: + if init.name in starts_inputs: + old_values.add(int(numpy_helper.to_array(init)[0])) + init.CopyFrom(numpy_helper.from_array(new_val, init.name)) + patched += 1 + + if patched != N_CACHE: + print(f"patched {patched} of {N_CACHE} window constants -- aborting", file=sys.stderr) + return 1 + + # The checker resolves external-data locations relative to the process CWD, + # not the graph -- pin it to the graph's directory so running the script + # from anywhere behaves the same. + import os + prev_cwd = os.getcwd() + os.chdir(src.parent) + try: + onnx.checker.check_model(model) + finally: + os.chdir(prev_cwd) + onnx.save(model, dst) + + old = ", ".join(str(-v) for v in sorted(old_values)) + peak = 2 * args.window * BYTES_PER_FRAME / (1024 ** 3) + print(f"patched {patched} window constants: {old} -> {args.window}") + print(f"wrote {dst.name} ({dst.stat().st_size / 1024**2:.1f} MB; shares {src.stem}.onnx.data)") + print(f"fp32 peak KV cache at window {args.window}: {peak:.2f} GiB " + f"(steady {peak/2:.2f} GiB)") + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/src/cuda_kernels.cu b/src/cuda_kernels.cu index d18e7f0..0d5a4db 100644 --- a/src/cuda_kernels.cu +++ b/src/cuda_kernels.cu @@ -1421,4 +1421,116 @@ cudaError_t postprocess_depth_image( return cudaGetLastError(); } +/////////////////////////////////////////////////////////////////////////////////////////////////// +// Resampling between camera resolution and a model's native input size. + +__global__ void resize_bilinear_chw_kernel( + const float* __restrict__ src, + float* __restrict__ dst, + int in_h, int in_w, + int out_h, int out_w, + int channels +) { + const int x = blockIdx.x * blockDim.x + threadIdx.x; + const int y = blockIdx.y * blockDim.y + threadIdx.y; + const int c = blockIdx.z; + if (x >= out_w || y >= out_h || c >= channels) return; + + // Half-pixel centres (equivalent to align_corners=False), so the resample + // stays centred and does not shift the image by half a source pixel. + const float fy = (y + 0.5f) * in_h / out_h - 0.5f; + const float fx = (x + 0.5f) * in_w / out_w - 0.5f; + + int y0 = static_cast(floorf(fy)); + int x0 = static_cast(floorf(fx)); + const float wy = fy - y0; + const float wx = fx - x0; + const int y1 = min(max(y0 + 1, 0), in_h - 1); + const int x1 = min(max(x0 + 1, 0), in_w - 1); + y0 = min(max(y0, 0), in_h - 1); + x0 = min(max(x0, 0), in_w - 1); + + const float* plane = src + static_cast(c) * in_h * in_w; + const float v00 = plane[static_cast(y0) * in_w + x0]; + const float v01 = plane[static_cast(y0) * in_w + x1]; + const float v10 = plane[static_cast(y1) * in_w + x0]; + const float v11 = plane[static_cast(y1) * in_w + x1]; + + dst[static_cast(c) * out_h * out_w + static_cast(y) * out_w + x] = + (1.f - wy) * ((1.f - wx) * v00 + wx * v01) + + wy * ((1.f - wx) * v10 + wx * v11); +} + +__global__ void resize_sparse_depth_kernel( + const float* __restrict__ src, + float* __restrict__ dst, + int in_h, int in_w, + int out_h, int out_w +) { + const int x = blockIdx.x * blockDim.x + threadIdx.x; + const int y = blockIdx.y * blockDim.y + threadIdx.y; + if (x >= out_w || y >= out_h) return; + + // Source footprint of this output pixel, plus its centre for tie-breaking. + const float cy = (y + 0.5f) * in_h / out_h; + const float cx = (x + 0.5f) * in_w / out_w; + int y0 = max(static_cast(floorf(static_cast(y) * in_h / out_h)), 0); + int y1 = min(static_cast(ceilf (static_cast(y + 1) * in_h / out_h)), in_h); + int x0 = max(static_cast(floorf(static_cast(x) * in_w / out_w)), 0); + int x1 = min(static_cast(ceilf (static_cast(x + 1) * in_w / out_w)), in_w); + // Upscaling can leave an empty footprint; always sample at least one pixel. + if (y1 <= y0) y1 = min(y0 + 1, in_h); + if (x1 <= x0) x1 = min(x0 + 1, in_w); + + float best = 0.f; + float best_dist2 = CUDART_INF_F; + for (int yy = y0; yy < y1; ++yy) { + for (int xx = x0; xx < x1; ++xx) { + const float v = src[static_cast(yy) * in_w + xx]; + // 0 means "no sample". Negated compare so NaN is rejected too. + // Range validation is the model's, not the resample's. + if (!(v > 0.f)) continue; + const float dy = (yy + 0.5f) - cy; + const float dx = (xx + 0.5f) - cx; + const float d2 = dy * dy + dx * dx; + if (d2 < best_dist2) { best_dist2 = d2; best = v; } + } + } + dst[static_cast(y) * out_w + x] = best; +} + +cudaError_t resize_bilinear_chw( + const float* d_in, + float* d_out, + int in_h, int in_w, + int out_h, int out_w, + int channels, + cudaStream_t stream +) { + dim3 block(32, 8); + dim3 grid((out_w + block.x - 1) / block.x, + (out_h + block.y - 1) / block.y, + channels); + resize_bilinear_chw_kernel<<>>( + d_in, d_out, in_h, in_w, out_h, out_w, channels + ); + return cudaGetLastError(); +} + +cudaError_t resize_sparse_depth( + const float* d_in, + float* d_out, + int in_h, int in_w, + int out_h, int out_w, + cudaStream_t stream +) { + dim3 block(32, 8); + dim3 grid((out_w + block.x - 1) / block.x, + (out_h + block.y - 1) / block.y); + resize_sparse_depth_kernel<<>>( + d_in, d_out, in_h, in_w, out_h, out_w + ); + return cudaGetLastError(); +} + } // namespace SOb \ No newline at end of file diff --git a/src/include/cuda_kernels.cuh b/src/include/cuda_kernels.cuh index f49aa27..e9a237a 100644 --- a/src/include/cuda_kernels.cuh +++ b/src/include/cuda_kernels.cuh @@ -67,6 +67,34 @@ void convert_uint8_img_to_float_img( void loadImageToCudaFloatRGB(const std::string& path, int& outW, int& outH, float* d_image); +// Resampling between the camera resolution and a model's native input size. +// Neither axis is an integer ratio (e.g. 480->392, 640->518), so the integer +// downscale in preprocess_depth_image2 does not cover this. + +// Planar CHW bilinear resample. +cudaError_t resize_bilinear_chw( + const float* d_in, + float* d_out, + int in_h, int in_w, + int out_h, int out_w, + int channels, + cudaStream_t stream = 0 +); + +// Resample sparse metric depth. Bilinear is wrong here: 0 means "no sample", and +// blending it with valid neighbours fabricates depth. Each output pixel takes the +// nearest sampled pixel in its source footprint, which preserves sparse points +// that plain nearest-neighbour would drop when downscaling. +// Only 0 (and NaN) are treated as "no sample"; range validation belongs to the +// consuming model, not to the resample. +cudaError_t resize_sparse_depth( + const float* d_in, + float* d_out, + int in_h, int in_w, + int out_h, int out_w, + cudaStream_t stream = 0 +); + // Running average depth maintenance cudaError_t prefill_invalid_depth( float* d_depth_data, diff --git a/src/include/model.h b/src/include/model.h index 282335c..1d31a1b 100644 --- a/src/include/model.h +++ b/src/include/model.h @@ -6,6 +6,7 @@ #include "utils.h" +#include #include #include #include @@ -32,6 +33,27 @@ class MLModel { TensorShape depth_shape, TensorShape output_shape ) = 0; + + // Streaming models carry per-frame state (e.g. a KV cache) that is only valid + // for a contiguous frame sequence. The pipeline calls this whenever the + // sequence restarts. No-op for stateless models. + virtual void resetState() {} + + // True when the model consumes full-resolution sparse metric depth directly + // instead of the pipeline's downscaled depth. + virtual bool wantsFullResDepth() const { return false; } + + // Models holding per-sequence state can serve exactly one pipeline at a time; + // a second pipeline driving the same instance would interleave two camera + // streams into one cache. Stateless models are freely shareable, so the + // default always succeeds. Returns false if another owner holds the instance. + virtual bool acquire(const void* owner) { (void)owner; return true; } + virtual void release(const void* owner) { (void)owner; } + + // Whether the model can run a batch of n images per step. Streaming models + // constrain this to their graph's batch dim (or accept any n when the export + // declares it symbolic); stateless models take whatever they're given. + virtual bool supportsBatch(int32_t n) const { return n >= 1; } }; /////////////////////////////////////////////////////////////////////////////////////////////////// @@ -86,10 +108,14 @@ class ONNXModel : public MLModel { Ort::Allocator* alloc_; }; - std::unique_ptr m_session; + // Declaration order is load-bearing: members are destroyed in reverse, and + // Ort::Env must outlive every Session created from it (and the Session must + // outlive the Allocator/IoBinding that reference it). Keep m_env first and + // m_session ahead of m_allocator/m_binding. Ort::Env m_env; Ort::SessionOptions m_sess_options; Ort::MemoryInfo m_memory_info; + std::unique_ptr m_session; std::unique_ptr m_allocator; std::vector m_input_names; @@ -134,4 +160,122 @@ class ONNXModel : public MLModel { std::string getDevice() const; }; -} // namespace UB \ No newline at end of file +/////////////////////////////////////////////////////////////////////////////////////////////////// + +// Autoregressive depth model with a transformer KV cache. +// +// Graph I/O (50 in, 50 out): +// in : rgb [1,3,H,W], sparse_depth [1,1,H,W], then past_k_00, past_v_00, +// past_k_01, ... interleaved K/V per layer, each [1,16,n_frames,P,64] +// out : depth, depth_conf, then new_k_00, new_v_00, ... in the same interleaved +// order, already sliced to the model's retention window +// +// The caches never leave the GPU. Each frame binds the previous step's outputs +// straight back as inputs via IoBinding, and binds the new outputs by MemoryInfo +// so ORT sizes them to whatever n_frames_out the graph produced -- they cannot be +// pre-allocated because the sequence axis grows before the window clamps it. +// Frame 0 feeds zero-length caches. +// +// The cache is per-frame-sequence state, so an instance must not be shared +// between pipelines; interleaving two camera streams into one cache silently +// corrupts it. Load a separate instance per pipeline. +class StreamingONNXModel : public MLModel { + static constexpr int32_t kNumCacheTensors = 48; // 24 layers x {K, V} + static constexpr int32_t kNumFixedInputs = 2; // rgb, sparse_depth + static constexpr int32_t kNumFixedOutputs = 2; // depth, depth_conf + + // Declaration order is load-bearing; see the note in ONNXModel. + Ort::Env m_env; + Ort::SessionOptions m_sess_options; + Ort::MemoryInfo m_memory_info; + std::unique_ptr m_session; + std::unique_ptr m_binding; + + // Index-aligned by construction: m_past_names[i] pairs with m_new_names[i], + // which is session output kNumFixedOutputs + i, which is m_cache[i]. + std::vector m_past_names; + std::vector m_new_names; + std::vector m_past_cstr; + std::vector m_new_cstr; + + // Previous step's caches, device-resident and owned by ORT's CUDA allocator. + std::vector m_cache; + // Backing pointer for the zero-length frame-0 caches (no elements are read). + void* m_d_empty{nullptr}; + + // Model-native input geometry, read from the graph rather than assumed. + int64_t m_model_h{0}; + int64_t m_model_w{0}; + int64_t m_num_heads{0}; + int64_t m_num_tokens{0}; + int64_t m_head_dim{0}; + + // Batch handling. m_graph_batch is what the graph declares for the batch dim: + // a positive value pins it; 0 means symbolic ("B"), so the model takes + // whatever batch each pipeline stream delivers (one cache sequence per batch + // slot -- slots are independent views, e.g. the two front cameras). + // m_cur_batch is the batch of the active sequence; changing it invalidates + // the cache, so a mid-stream change forces a sequence restart. + int64_t m_graph_batch{1}; + int64_t m_cur_batch{1}; + int64_t m_alloc_batch{0}; // scratch capacity, grown on demand + + // Element type the graph declares for the caches. Their contents are never + // read or written here -- they leave ORT and come straight back in -- so this + // only decides how the zero-length frame-0 tensors are created. An fp16 cache + // halves the largest allocation this model makes. + ONNXTensorElementDataType m_cache_type{ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT}; + + // Scratch at model resolution. + float* m_d_rgb{nullptr}; + float* m_d_depth{nullptr}; + + int64_t m_frames_seen{0}; + bool m_use_cuda{false}; + // The pipeline currently driving this instance, or null when free. + std::atomic m_owner{nullptr}; + + void _setDevice(const std::string& device_type); + void _buildCacheNames(); + void _readGeometry(); + void _ensureScratch(int64_t batch); + void _freeScratch(); + std::vector _makeEmptyCaches() const; + +public: + explicit StreamingONNXModel(const std::string& model_path, const std::string& device_type = "cuda"); + ~StreamingONNXModel() override; + + bool runInference( + const float* input_data, + const float* depth_data, + float* output_data, + TensorShape input_shape, + TensorShape depth_shape, + TensorShape output_shape + ) override; + + bool runInference( + const uint8_t* input_data, + const float* depth_data, + float* output_data, + TensorShape input_shape, + TensorShape depth_shape, + TensorShape output_shape + ) override; + + void resetState() override; + bool wantsFullResDepth() const override { return true; } + bool acquire(const void* owner) override; + void release(const void* owner) override; + bool supportsBatch(int32_t n) const override { + if (n < 1) return false; + return m_graph_batch == 0 || n == m_graph_batch; + } + + std::string getDevice() const { return m_use_cuda ? "cuda" : "cpu"; } + int64_t getModelHeight() const { return m_model_h; } + int64_t getModelWidth() const { return m_model_w; } +}; + +} // namespace SOb \ No newline at end of file diff --git a/src/model.cpp b/src/model.cpp index dc587a8..af8e188 100644 --- a/src/model.cpp +++ b/src/model.cpp @@ -877,4 +877,516 @@ std::string ONNXModel::getDevice() const { return m_use_cuda ? "cuda" : "cpu"; } +/////////////////////////////////////////////////////////////////////////////////////////////////// + +// Bytes per element for the types a KV cache is plausibly carried in. +static size_t onnxElementSize(ONNXTensorElementDataType type) { + switch (type) { + case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT: return 4; + case ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16: + case ONNX_TENSOR_ELEMENT_DATA_TYPE_BFLOAT16: return 2; + case ONNX_TENSOR_ELEMENT_DATA_TYPE_INT8: + case ONNX_TENSOR_ELEMENT_DATA_TYPE_UINT8: return 1; + default: return 0; // unknown; reported as 0 MB + } +} + +StreamingONNXModel::StreamingONNXModel(const std::string& model_path, const std::string& device_type) + : m_env(ORT_LOGGING_LEVEL_WARNING, "StreamingONNXModel") + , m_sess_options() + , m_memory_info(Ort::MemoryInfo::CreateCpu(OrtAllocatorType::OrtArenaAllocator, OrtMemTypeDefault)) +{ + if (!fs::exists(model_path)) { + throw std::runtime_error("Model " + model_path + " does not exist."); + } + + try { + _setDevice(device_type); + if (!m_use_cuda) { + throw std::runtime_error("StreamingONNXModel requires CUDA; the KV cache must stay device-resident."); + } + + LogMessage("ONNXModel loading from: {}", model_path); + + // Weights live in a sibling .onnx.data file, so this reads several GB. + auto load_start = std::chrono::high_resolution_clock::now(); + std::wstring wide_model_path = std::wstring(model_path.begin(), model_path.end()); + m_session = std::make_unique(m_env, wide_model_path.c_str(), m_sess_options); + + LogMessage("Ort session created successfully."); + auto load_end = std::chrono::high_resolution_clock::now(); + + LogMessage("Streaming ONNX model loaded from {} in {} ms", model_path, + std::chrono::duration_cast(load_end - load_start).count()); + + _buildCacheNames(); + _readGeometry(); + _ensureScratch(m_cur_batch); + + m_binding = std::make_unique(*m_session); + resetState(); + + const double cache_mb_per_frame = + static_cast(kNumCacheTensors) * m_num_heads * m_num_tokens * m_head_dim + * onnxElementSize(m_cache_type) / (1024.0 * 1024.0); + // The .onnx file is just the graph; the weights live in the external-data + // sibling, so include it or the log understates the footprint ~400x. + uintmax_t weight_bytes = fs::file_size(model_path); + const fs::path external_data = fs::path(model_path).concat(".data"); + if (fs::exists(external_data)) { + weight_bytes += fs::file_size(external_data); + } + LogPerf("[mem] StreamingONNXModel: graph + weights {:.2f} MB on disk, KV cache {:.2f} MB per retained frame " + "(x2 live during Run)", + weight_bytes / (1024.0 * 1024.0), cache_mb_per_frame); + + } catch (const Ort::Exception& e) { + // The destructor does not run for a partially constructed object, so any + // scratch already allocated has to be released here. + _freeScratch(); + throw std::runtime_error("Error loading streaming ONNX model: " + std::string(e.what())); + } catch (...) { + _freeScratch(); + throw; + } +} + +StreamingONNXModel::~StreamingONNXModel() { + // m_cache is declared after m_session, so the cache tensors are released + // before the session and env they were allocated from. + _freeScratch(); + LogMessage("StreamingONNXModel destroyed."); +} + +void StreamingONNXModel::_freeScratch() { + if (m_d_rgb) { cudaFree(m_d_rgb); m_d_rgb = nullptr; } + if (m_d_depth) { cudaFree(m_d_depth); m_d_depth = nullptr; } + if (m_d_empty) { cudaFree(m_d_empty); m_d_empty = nullptr; } +} + +void StreamingONNXModel::_setDevice(const std::string& device_type) { + std::string device_type_lower = device_type; + std::transform(device_type_lower.begin(), device_type_lower.end(), device_type_lower.begin(), ::tolower); + + if (device_type_lower != "cuda" && device_type_lower != "gpu") { + LogMessage("StreamingONNXModel: unsupported device '{}'", device_type); + m_use_cuda = false; + return; + } + + int device_count = 0; + cudaError_t err = cudaGetDeviceCount(&device_count); + if (err != cudaSuccess || device_count == 0) { + LogMessage("StreamingONNXModel: CUDA not available"); + m_use_cuda = false; + return; + } + + const auto& ort_api = Ort::GetApi(); + + // Every OrtStatus* must be checked and released: silently ignoring them would + // leave m_use_cuda true while the session actually ran on CPU, and device + // pointers would then be bound with CUDA MemoryInfo. + OrtCUDAProviderOptionsV2* cuda_options = nullptr; + if (OrtStatus* status = ort_api.CreateCUDAProviderOptions(&cuda_options)) { + LogMessage("StreamingONNXModel: CreateCUDAProviderOptions failed: {}", + ort_api.GetErrorMessage(status)); + ort_api.ReleaseStatus(status); + m_use_cuda = false; + return; + } + std::unique_ptr + rel_cuda_options(cuda_options, ort_api.ReleaseCUDAProviderOptions); + + // Arena strategy stays at the default (kNextPowerOfTwo) deliberately. The + // cache tensors grow to a NEW size every frame until the window clamps, and + // exact-size allocation (kSameAsRequested) can never reuse a freed chunk for + // the next, larger request -- the arena extends every frame and VRAM climbs + // ~quadratically until WDDM starts evicting (measured: 375 ms -> 14.8 s per + // frame at the cliff). Power-of-two rounding puts consecutive sizes in the + // same bucket, so freed generations actually get reused during growth. + + // Deliberately no enable_cuda_graph here: graph capture requires static + // shapes, and the cache sequence axis grows for the first frames before the + // retention window clamps it. + if (OrtStatus* status = ort_api.SessionOptionsAppendExecutionProvider_CUDA_V2( + static_cast(m_sess_options), + rel_cuda_options.get())) { + LogMessage("StreamingONNXModel: registering the CUDA EP failed: {}", + ort_api.GetErrorMessage(status)); + ort_api.ReleaseStatus(status); + m_use_cuda = false; + return; + } + + m_memory_info = Ort::MemoryInfo("Cuda", OrtAllocatorType::OrtArenaAllocator, 0, OrtMemTypeDefault); + m_use_cuda = true; + LogMessage("StreamingONNXModel device set to: CUDA"); +} + +void StreamingONNXModel::_buildCacheNames() { + Ort::AllocatorWithDefaultOptions alloc; + + const size_t num_inputs = m_session->GetInputCount(); + const size_t num_outputs = m_session->GetOutputCount(); + const size_t expected = kNumFixedInputs + kNumCacheTensors; + + if (num_inputs != expected || num_outputs != expected) { + throw std::runtime_error("Expected " + std::to_string(expected) + " inputs and outputs, got " + + std::to_string(num_inputs) + " and " + std::to_string(num_outputs)); + } + + auto input_name = [&](size_t i) { return std::string(m_session->GetInputNameAllocated(i, alloc).get()); }; + auto output_name = [&](size_t i) { return std::string(m_session->GetOutputNameAllocated(i, alloc).get()); }; + + if (input_name(0) != "rgb" || input_name(1) != "sparse_depth") { + throw std::runtime_error("Expected inputs [rgb, sparse_depth], got [" + + input_name(0) + ", " + input_name(1) + "]"); + } + if (output_name(0) != "depth" || output_name(1) != "depth_conf") { + throw std::runtime_error("Expected outputs [depth, depth_conf], got [" + + output_name(0) + ", " + output_name(1) + "]"); + } + + // Take the pairing from the graph's own ordering rather than reconstructing + // names, then assert the past_/new_ suffixes line up. A silent misalignment + // here would route one layer's K into another's V: no error, just wrong depth. + m_past_names.reserve(kNumCacheTensors); + m_new_names.reserve(kNumCacheTensors); + for (int32_t i = 0; i < kNumCacheTensors; ++i) { + std::string in_name = input_name(kNumFixedInputs + i); + std::string out_name = output_name(kNumFixedOutputs + i); + + if (in_name.rfind("past_", 0) != 0 || out_name.rfind("new_", 0) != 0) { + throw std::runtime_error("Cache tensor " + std::to_string(i) + " naming unexpected: " + + in_name + " / " + out_name); + } + if (in_name.substr(5) != out_name.substr(4)) { + throw std::runtime_error("Cache tensor " + std::to_string(i) + " misaligned: " + + in_name + " does not pair with " + out_name); + } + + m_past_names.push_back(std::move(in_name)); + m_new_names.push_back(std::move(out_name)); + } + + // Built only after the name vectors are final, so the pointers stay valid. + m_past_cstr.reserve(kNumCacheTensors); + m_new_cstr.reserve(kNumCacheTensors); + for (int32_t i = 0; i < kNumCacheTensors; ++i) { + m_past_cstr.push_back(m_past_names[i].c_str()); + m_new_cstr.push_back(m_new_names[i].c_str()); + } + + LogMessage("StreamingONNXModel: {} cache tensors, {} <-> {} ... {} <-> {}", + kNumCacheTensors, m_past_names.front(), m_new_names.front(), + m_past_names.back(), m_new_names.back()); +} + +void StreamingONNXModel::_readGeometry() { + // TypeInfo owns the shape/type view, so each must outlive the view taken from it. + Ort::TypeInfo rgb_type_info = m_session->GetInputTypeInfo(0); + Ort::TypeInfo depth_in_type_info = m_session->GetInputTypeInfo(1); + Ort::TypeInfo depth_out_type_info = m_session->GetOutputTypeInfo(0); + Ort::TypeInfo cache_type_info = m_session->GetInputTypeInfo(kNumFixedInputs); + + // rgb: [B, 3, H, W]. H and W must be static; the batch dim may be a fixed + // value or symbolic (reported as -1), in which case any batch is accepted + // and each batch slot carries its own independent cache sequence. + auto rgb_info = rgb_type_info.GetTensorTypeAndShapeInfo(); + auto rgb_shape = rgb_info.GetShape(); + if (rgb_shape.size() != 4 || rgb_shape[2] <= 0 || rgb_shape[3] <= 0) { + throw std::runtime_error("rgb input must have [B,3,H,W] shape with static H and W"); + } + m_model_h = rgb_shape[2]; + m_model_w = rgb_shape[3]; + m_graph_batch = rgb_shape[0] > 0 ? rgb_shape[0] : 0; // 0 = symbolic = any + + // The image tensors are read and written directly by the fp32 resize kernels, + // so those must be fp32. Fail loudly here rather than reinterpret half floats + // as single and hand back plausible-looking garbage. + auto require_float = [](const auto& info, const char* what) { + if (info.GetElementType() != ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT) { + throw std::runtime_error(std::string(what) + " must be float32, got ONNX element type " + + std::to_string(static_cast(info.GetElementType())) + + " (convert with keep_io_types=True to leave the image tensors fp32)"); + } + }; + require_float(rgb_info, "rgb input"); + require_float(depth_in_type_info.GetTensorTypeAndShapeInfo(), "sparse_depth input"); + require_float(depth_out_type_info.GetTensorTypeAndShapeInfo(), "depth output"); + + // past_k_00: [1, heads, n_frames, tokens, head_dim]; n_frames is symbolic (-1). + auto cache_info = cache_type_info.GetTensorTypeAndShapeInfo(); + auto cache_shape = cache_info.GetShape(); + if (cache_shape.size() != 5 || cache_shape[1] <= 0 || cache_shape[3] <= 0 || cache_shape[4] <= 0) { + throw std::runtime_error("cache input must have shape [1,heads,n_frames,tokens,head_dim]"); + } + m_num_heads = cache_shape[1]; + m_num_tokens = cache_shape[3]; + m_head_dim = cache_shape[4]; + // Any element type is fine here: the caches are opaque to this class. + m_cache_type = cache_info.GetElementType(); + + // The cache batch dim must agree with rgb's -- a graph that batches images + // but not cache sequences cannot stream per-view. + const int64_t cache_batch = cache_shape[0] > 0 ? cache_shape[0] : 0; + if (cache_batch != m_graph_batch) { + throw std::runtime_error("rgb and cache batch dims disagree (" + + std::to_string(m_graph_batch) + " vs " + std::to_string(cache_batch) + ")"); + } + + m_cur_batch = m_graph_batch > 0 ? m_graph_batch : 1; + + LogMessage("StreamingONNXModel geometry: input {}x{}, batch {}, cache [B,{},n_frames,{},{}] element type {} ({} bytes)", + m_model_h, m_model_w, m_graph_batch == 0 ? "dynamic" : std::to_string(m_graph_batch), + m_num_heads, m_num_tokens, m_head_dim, + static_cast(m_cache_type), onnxElementSize(m_cache_type)); +} + +void StreamingONNXModel::_ensureScratch(int64_t batch) { + if (batch <= m_alloc_batch) { + return; + } + // Scratch is only live within a single runInference call (which syncs before + // returning), so growing it between calls is safe. + if (m_d_rgb) { cudaFree(m_d_rgb); m_d_rgb = nullptr; } + if (m_d_depth) { cudaFree(m_d_depth); m_d_depth = nullptr; } + checkCudaError(cudaMalloc(&m_d_rgb, batch * 3 * m_model_h * m_model_w * sizeof(float)), + "cudaMalloc streaming model rgb scratch"); + checkCudaError(cudaMalloc(&m_d_depth, batch * m_model_h * m_model_w * sizeof(float)), + "cudaMalloc streaming model depth scratch"); + if (!m_d_empty) { + // Backing for the zero-length frame-0 caches. No elements are read; ORT + // just wants a valid device pointer. + checkCudaError(cudaMalloc(&m_d_empty, 256), "cudaMalloc streaming model empty cache backing"); + } + m_alloc_batch = batch; +} + +std::vector StreamingONNXModel::_makeEmptyCaches() const { + std::vector shape{m_cur_batch, m_num_heads, 0, m_num_tokens, m_head_dim}; + std::vector caches; + caches.reserve(kNumCacheTensors); + for (int32_t i = 0; i < kNumCacheTensors; ++i) { + // Untyped overload so the caches are created in whatever type the graph + // declares. This is the only place cache dtype matters -- every later + // generation is an ORT-owned value we pass straight back in -- which is + // what makes an fp16 cache a ~10 line change rather than a rewrite. + // Byte count is 0: no element is ever read from m_d_empty. + caches.push_back(Ort::Value::CreateTensor( + m_memory_info, + m_d_empty, + 0, + shape.data(), + shape.size(), + m_cache_type + )); + } + return caches; +} + +bool StreamingONNXModel::acquire(const void* owner) { + const void* expected = nullptr; + if (m_owner.compare_exchange_strong(expected, owner)) { + return true; + } + // Re-acquiring from the same pipeline (e.g. stop/start) is fine. + return expected == owner; +} + +void StreamingONNXModel::release(const void* owner) { + const void* expected = owner; + m_owner.compare_exchange_strong(expected, nullptr); +} + +void StreamingONNXModel::resetState() { + if (m_binding) { + m_binding->ClearBoundInputs(); + m_binding->ClearBoundOutputs(); + } + m_cache = _makeEmptyCaches(); + m_frames_seen = 0; + LogMessage("StreamingONNXModel: KV cache reset (next frame feeds zero-length caches)"); +} + +bool StreamingONNXModel::runInference( + const float* input_data, + const float* depth_data, + float* output_data, + TensorShape input_shape, + TensorShape depth_shape, + TensorShape output_shape +) { + if (!m_session || !m_binding) { + LogMessage("StreamingONNXModel: session not initialized"); + return false; + } + if (!input_data || !depth_data || !output_data) { + LogMessage("StreamingONNXModel: null input, depth, or output pointer"); + return false; + } + + auto time_start = std::chrono::high_resolution_clock::now(); + + // Batch = images per step, one independent cache sequence per slot. + const int64_t batch = static_cast(input_shape.N); + if (!supportsBatch(static_cast(batch))) { + LogMessage("StreamingONNXModel: batch {} not supported (graph batch: {})", + batch, m_graph_batch == 0 ? std::string("dynamic") : std::to_string(m_graph_batch)); + return false; + } + if (batch != m_cur_batch) { + // The cache's batch dim is part of the sequence state; a different batch + // means a different set of view sequences, so restart rather than feed a + // mismatched cache into the graph. + LogMessage("StreamingONNXModel: batch changed {} -> {}, restarting sequence", + m_cur_batch, batch); + m_cur_batch = batch; + resetState(); + } + + try { + _ensureScratch(batch); + + // Camera resolution -> model resolution, all batch slots in one pass: + // [B,3,H,W] is B*3 contiguous planes to the resampler. Depth is resampled + // sparsely so invalid (zero) pixels are never blended into valid ones. + // Range validation and orientation are handled inside the graph. + checkCudaError(resize_bilinear_chw( + input_data, m_d_rgb, + static_cast(input_shape.H), static_cast(input_shape.W), + static_cast(m_model_h), static_cast(m_model_w), + static_cast(batch * 3), 0 + ), "resize rgb to model resolution"); + + const size_t depth_in_elems = depth_shape.H * depth_shape.W; + const size_t depth_mdl_elems = static_cast(m_model_h * m_model_w); + for (int64_t b = 0; b < batch; ++b) { + checkCudaError(resize_sparse_depth( + depth_data + b * depth_in_elems, m_d_depth + b * depth_mdl_elems, + static_cast(depth_shape.H), static_cast(depth_shape.W), + static_cast(m_model_h), static_cast(m_model_w), + 0 + ), "resize sparse depth to model resolution"); + } + + // ORT runs on its own stream; make sure the resamples are visible first. + checkCudaError(cudaStreamSynchronize(0), "sync before streaming Run"); + + std::vector rgb_tensor_shape{batch, 3, m_model_h, m_model_w}; + std::vector depth_tensor_shape{batch, 1, m_model_h, m_model_w}; + + Ort::Value rgb_tensor = Ort::Value::CreateTensor( + m_memory_info, m_d_rgb, static_cast(batch * 3 * m_model_h * m_model_w), + rgb_tensor_shape.data(), rgb_tensor_shape.size() + ); + Ort::Value depth_tensor = Ort::Value::CreateTensor( + m_memory_info, m_d_depth, static_cast(batch * m_model_h * m_model_w), + depth_tensor_shape.data(), depth_tensor_shape.size() + ); + + m_binding->ClearBoundInputs(); + m_binding->ClearBoundOutputs(); + + m_binding->BindInput("rgb", rgb_tensor); + m_binding->BindInput("sparse_depth", depth_tensor); + for (int32_t i = 0; i < kNumCacheTensors; ++i) { + m_binding->BindInput(m_past_cstr[i], m_cache[i]); + } + + // Bound by MemoryInfo rather than to our own buffers: the graph decides + // n_frames_out, so ORT must allocate these device-side at the size it + // produced. The returned values are fed straight back in next frame -- + // the cache never touches host memory. + m_binding->BindOutput("depth", m_memory_info); + m_binding->BindOutput("depth_conf", m_memory_info); + for (int32_t i = 0; i < kNumCacheTensors; ++i) { + m_binding->BindOutput(m_new_cstr[i], m_memory_info); + } + + m_session->Run(Ort::RunOptions{nullptr}, *m_binding); + + // Bind order, so index kNumFixedOutputs + i is m_new_names[i]. + std::vector outputs = m_binding->GetOutputValues(); + if (outputs.size() != static_cast(kNumFixedOutputs + kNumCacheTensors)) { + LogMessage("StreamingONNXModel: expected {} outputs, got {}", + kNumFixedOutputs + kNumCacheTensors, outputs.size()); + return false; + } + + // Drop the previous generation's bindings before those tensors are freed. + m_binding->ClearBoundInputs(); + m_binding->ClearBoundOutputs(); + + // depth is [B, H, W] in whatever orientation the graph emits; take the + // dims it reports rather than assuming them. + auto depth_out_shape = outputs[0].GetTensorTypeAndShapeInfo().GetShape(); + if (depth_out_shape.size() != 3 || depth_out_shape[0] != batch) { + LogMessage("StreamingONNXModel: unexpected depth output shape (rank {}, batch {})", + depth_out_shape.size(), depth_out_shape.empty() ? -1 : depth_out_shape[0]); + return false; + } + const int64_t src_h = depth_out_shape[1]; + const int64_t src_w = depth_out_shape[2]; + + // All batch slots in one pass: [B,H,W] in and [B,H',W'] out are both + // contiguous single-channel planes. + checkCudaError(resize_bilinear_chw( + outputs[0].GetTensorData(), output_data, + static_cast(src_h), static_cast(src_w), + static_cast(output_shape.H), static_cast(output_shape.W), + static_cast(batch), 0 + ), "resize model depth to output resolution"); + checkCudaError(cudaStreamSynchronize(0), "sync after streaming depth resize"); + + auto cache_shape = outputs[kNumFixedOutputs].GetTensorTypeAndShapeInfo().GetShape(); + if (cache_shape.size() != 5 || cache_shape[0] != batch) { + LogMessage("StreamingONNXModel: cache output shape unexpected (rank {}, batch {} vs {})", + cache_shape.size(), cache_shape.empty() ? -1 : cache_shape[0], batch); + resetState(); + return false; + } + // The graph slices to its own retention window, so this is just reported. + const int64_t retained = cache_shape[2]; + + // This step's caches become the next step's inputs; assigning here frees + // the previous generation back to ORT's arena for reuse. + m_cache.clear(); + m_cache.reserve(kNumCacheTensors); + for (int32_t i = 0; i < kNumCacheTensors; ++i) { + m_cache.push_back(std::move(outputs[kNumFixedOutputs + i])); + } + + m_frames_seen++; + + auto time_end = std::chrono::high_resolution_clock::now(); + LogMessage("StreamingONNXModel frame {}: batch {}, {} frames retained, depth out [{}, {}], {} ms", + m_frames_seen, batch, retained, src_h, src_w, + std::chrono::duration_cast(time_end - time_start).count()); + + } catch (const std::exception& e) { + LogMessage("Error during streaming ONNX inference: {}", e.what()); + // The cache generation is now indeterminate; restart the sequence rather + // than feeding a half-updated cache into the next frame. + resetState(); + return false; + } + + return true; +} + +bool StreamingONNXModel::runInference( + const uint8_t* /*input_data*/, + const float* /*depth_data*/, + float* /*output_data*/, + TensorShape /*input_shape*/, + TensorShape /*depth_shape*/, + TensorShape /*output_shape*/ +) { + LogMessage("StreamingONNXModel: uint8 input overload is not implemented; " + "the pipeline converts to float before inference."); + return false; +} + }; \ No newline at end of file diff --git a/src/spot-observer.cpp b/src/spot-observer.cpp index 3a34a41..d56aff7 100644 --- a/src/spot-observer.cpp +++ b/src/spot-observer.cpp @@ -284,6 +284,29 @@ static SObModel loadONNXModel(const std::string& modelPath, const std::string& b return ret; } +// Autoregressive KV-cache model. Deliberately not registered in +// s_path_to_model_map: that map hands the same instance to every pipeline asking +// for a path, and a streaming model's cache belongs to one camera's frame +// sequence. Each load returns a fresh instance; s_model_to_path_map is still +// populated so unloadModel can find it. +static SObModel loadStreamingONNXModel(const std::string& modelPath, const std::string& device) { + LogMessage("Loading streaming ONNX model: {}", modelPath); + LogMessage("Using Provider: {}", device); + + SObModel ret = nullptr; + try { + auto* model = new StreamingONNXModel(modelPath, device); + ret = reinterpret_cast(model); + } catch (const std::exception& e) { + LogMessage("Exception while loading streaming ONNX model: {}", e.what()); + return nullptr; + } + + s_model_to_path_map[ret] = modelPath; + LogMessage("Successfully loaded streaming ONNX model: {}", modelPath); + return ret; +} + static void unloadModel(SObModel model) { if (!model) { LogMessage("SOb::unloadModel: Model is null, nothing to unload"); @@ -297,7 +320,14 @@ static void unloadModel(SObModel model) { } std::string modelPath = it->second; - s_path_to_model_map.erase(modelPath); + // Streaming models are deliberately absent from the path map (each pipeline + // gets its own instance), so only erase an entry that maps back to this exact + // model. Erasing by path alone would deregister a different, still-live model + // loaded from the same file and orphan its session. + auto path_it = s_path_to_model_map.find(modelPath); + if (path_it != s_path_to_model_map.end() && path_it->second == model) { + s_path_to_model_map.erase(path_it); + } s_model_to_path_map.erase(it); delete reinterpret_cast(model); @@ -412,29 +442,50 @@ bool UNITY_INTERFACE_API SOb_DestroyCameraStream(int32_t robot_id, int32_t cam_s } +// Model-family values, mirrored from include/spot-observer.h (this file follows +// the existing convention of not including the public header). +#define SOb_MODEL_SINGLE_SHOT 0 +#define SOb_MODEL_STREAMING 1 + UNITY_INTERFACE_EXPORT -SObModel UNITY_INTERFACE_API SOb_LoadModel(const char* modelPath, const char* backend) { +SObModel UNITY_INTERFACE_API SOb_LoadModelEx(const char* modelPath, const char* backend, int32_t kind) { if (!modelPath || !backend) { - SOb::LogMessage("SOb_LoadModel: Invalid null pointer parameters"); + SOb::LogMessage("SOb_LoadModelEx: Invalid null pointer parameters"); return nullptr; } - - // If model filename ends with .onnx, use ONNX model loader + std::string model_path_str(modelPath); + // The family is an explicit parameter rather than inferred: both families + // are .onnx files, and "backend" keeps its execution-provider meaning. SObModel ret = nullptr; - if (model_path_str.ends_with(".onnx")) { - ret = SOb::loadONNXModel(modelPath, backend); - } else { - ret = SOb::loadTorchModel(modelPath, backend); + switch (kind) { + case SOb_MODEL_SINGLE_SHOT: + if (model_path_str.ends_with(".onnx")) { + ret = SOb::loadONNXModel(modelPath, backend); + } else { + ret = SOb::loadTorchModel(modelPath, backend); + } + break; + case SOb_MODEL_STREAMING: + ret = SOb::loadStreamingONNXModel(model_path_str, backend); + break; + default: + SOb::LogMessage("SOb_LoadModelEx: Unknown model kind {}", kind); + return nullptr; } + if (!ret) { - SOb::LogMessage("Failed to load model: {} with backend: {}", modelPath, backend); + SOb::LogMessage("Failed to load model: {} (backend: {}, kind: {})", modelPath, backend, kind); } - return ret; } +UNITY_INTERFACE_EXPORT +SObModel UNITY_INTERFACE_API SOb_LoadModel(const char* modelPath, const char* backend) { + return SOb_LoadModelEx(modelPath, backend, SOb_MODEL_SINGLE_SHOT); +} + UNITY_INTERFACE_EXPORT void UNITY_INTERFACE_API SOb_UnloadModel(SObModel model) { try { @@ -512,6 +563,53 @@ bool UNITY_INTERFACE_API SOb_StopVisionPipeline(int32_t robot_id, int32_t cam_st } } +// Live model switch. Tears down the running pipeline (worker joined and model +// ownership released before this returns) and relaunches against `model`. The +// camera stream keeps running throughout. With all selectable models preloaded, +// this is the whole switch -- no load stall, nothing unloaded. +UNITY_INTERFACE_EXPORT +bool UNITY_INTERFACE_API SOb_SwitchVisionPipelineModel(int32_t robot_id, int32_t cam_stream_id, SObModel model) { + using namespace SOb; + try { + auto robot_it = __robot_connections.find(robot_id); + if (robot_it == __robot_connections.end()) { + LogMessage("SOb_SwitchVisionPipelineModel: Robot ID {} not found", robot_id); + return false; + } + if (!model) { + LogMessage("SOb_SwitchVisionPipelineModel: Invalid model provided"); + return false; + } + + SpotConnection& spot_connection = *robot_it->second; + + // Absent pipeline is fine -- then this degenerates into a launch. + // ~VisionPipeline stops the worker and releases the old model. + if (spot_connection.getVisionPipeline(cam_stream_id) != nullptr && + !spot_connection.removeVisionPipeline(cam_stream_id)) { + LogMessage("SOb_SwitchVisionPipelineModel: Failed to stop existing pipeline for robot ID {} @ stream-ID {}", + robot_id, cam_stream_id); + return false; + } + + auto* ml_model = reinterpret_cast(model); + if (!spot_connection.createVisionPipeline(*ml_model, cam_stream_id)) { + LogMessage("SOb_SwitchVisionPipelineModel: Failed to relaunch pipeline for robot ID {} @ stream-ID {}", + robot_id, cam_stream_id); + return false; + } + + LogMessage("SOb_SwitchVisionPipelineModel: Switched model for robot ID {} @ stream-ID {}", + robot_id, cam_stream_id); + return true; + + } catch (const std::exception& e) { + LogMessage("SOb_SwitchVisionPipelineModel: Exception for robot ID {} @ stream-ID {}: {}", + robot_id, cam_stream_id, e.what()); + return false; + } +} + UNITY_INTERFACE_EXPORT bool UNITY_INTERFACE_API SOb_RegisterUnityReadbackBuffers( int32_t robot_id, diff --git a/src/vision-pipeline.cpp b/src/vision-pipeline.cpp index f4c1afc..f1d6be2 100644 --- a/src/vision-pipeline.cpp +++ b/src/vision-pipeline.cpp @@ -65,11 +65,34 @@ bool VisionPipeline::start() { return false; } + // A streaming model's cache has one sequence per batch slot, so the stream's + // camera count must be a batch size the graph accepts (fixed-batch exports + // pin it; dynamic-batch exports take any). Refusing here beats the silent + // cache corruption that a mismatched batch would cause. + if (!model_.supportsBatch(static_cast(output_shape_.N))) { + LogMessage("Model does not support {} images per step for this stream", output_shape_.N); + return false; + } + if (!allocateCudaBuffers()) { LogMessage("Failed to allocate CUDA buffers"); return false; } + // A streaming model's KV cache is one camera sequence's state, so refuse to + // let a second pipeline drive the same instance -- the two streams would + // interleave into one cache and corrupt it silently. + if (!model_.acquire(this)) { + LogMessage("Model instance is already driven by another pipeline"); + deallocateCudaBuffers(); + return false; + } + + // Restarting the pipeline restarts the frame sequence, so any cached model + // state from a previous run is stale. + model_.resetState(); + first_run_ = true; + read_idx_.store(0); write_idx_ = 0; new_data_.store(false); @@ -95,6 +118,10 @@ void VisionPipeline::stop() { pipeline_thread_->join(); LogMessage("Vision pipeline thread joined"); } + + // Released only after the worker is joined, so the model is free for another + // pipeline exactly when nothing can still be calling into it. + model_.release(this); } bool VisionPipeline::allocateCudaBuffers() { @@ -270,13 +297,23 @@ void VisionPipeline::pipelineWorker(std::stop_token stop_token) { input_shape_float.W = input_shape_.H; } + // Streaming models take the raw sparse metric depth at sensor + // resolution and do their own resampling, range validation and + // orientation handling; the EMA prefill and the /4 downscale below + // are for the single-shot models only. The dims handed over describe + // the buffer exactly as it sits in memory -- what the model does with + // it from there is the model's business. + const bool full_res_depth = model_.wantsFullResDepth(); + TensorShape depth_shape = depth_shape_; - if (do_rotate_90_cw) { - depth_shape.H = depth_shape_.W / depth_scale_factor; - depth_shape.W = depth_shape_.H / depth_scale_factor; - } else { - depth_shape.H = depth_shape_.H / depth_scale_factor; - depth_shape.W = depth_shape_.W / depth_scale_factor; + if (!full_res_depth) { + if (do_rotate_90_cw) { + depth_shape.H = depth_shape_.W / depth_scale_factor; + depth_shape.W = depth_shape_.H / depth_scale_factor; + } else { + depth_shape.H = depth_shape_.H / depth_scale_factor; + depth_shape.W = depth_shape_.W / depth_scale_factor; + } } TensorShape output_shape = output_shape_; @@ -302,7 +339,9 @@ void VisionPipeline::pipelineWorker(std::stop_token stop_token) { LogMessage("Starting pipeline for image {}. cur_rgb_ptr = {:#x}, cur_depth_ptr = {:#x}, cur_depth_output_ptr = {:#x}", i, size_t(cur_rgb_input_ptr), size_t(cur_depth_input_ptr), size_t(cur_depth_output_ptr)); - if (ema_enabled && !first_run_) { + if (full_res_depth) { + // Nothing to do: the model reads d_depth_data_ directly. + } else if (ema_enabled && !first_run_) { checkCudaError(prefill_invalid_depth( cur_depth_input_ptr, cur_preprocessed_depth_ptr, @@ -366,7 +405,7 @@ void VisionPipeline::pipelineWorker(std::stop_token stop_token) { // TODO: Support running models on cuda_stream_ bool inference_success = model_.runInference( cuda_ws_.d_rgb_float_data_, - cuda_ws_.d_preprocessed_depth_data_, + full_res_depth ? cuda_ws_.d_depth_data_ : cuda_ws_.d_preprocessed_depth_data_, d_depth_output_ptr, input_shape_float, depth_shape, diff --git a/tests/integ-test/integ-test.cpp b/tests/integ-test/integ-test.cpp index fbef4e0..1ad2f38 100644 --- a/tests/integ-test/integ-test.cpp +++ b/tests/integ-test/integ-test.cpp @@ -4,6 +4,7 @@ #include "spot-observer.h" +#include #include #include #include @@ -56,8 +57,13 @@ static int32_t disconnect_from_spots(const int32_t spot_ids[], size_t num_spots) int main(int argc, char* argv[]) { using namespace std::chrono; - if (argc < 5 || argc > 6) { - std::cerr << "Usage: " << argv[0] << " [model_path]" << std::endl; + if (argc < 5 || argc > 9) { + std::cerr << "Usage: " << argv[0] + << " " + << " [model_path] [model_kind] [model2_path] [model2_kind]\n" + << " model_kind: 0 = single-shot (default), 1 = streaming (KV cache)\n" + << " With model2 given, the pipeline live-switches between the two models\n" + << " every 15 s via SOb_SwitchVisionPipelineModel (both loaded up front)." << std::endl; return 1; } @@ -66,6 +72,10 @@ int main(int argc, char* argv[]) { std::string password = argv[4]; //SOb_ToggleDebugDumps("./spot_dump"); + // PERF (1): timing + memory lines only. Level 2 (ALL) additionally prints + // several LogMessage lines per camera frame -- useful when diagnosing a + // launch failure, but console I/O at that volume measurably drags the + // streaming threads. Bump to 2 temporarily when something fails silently. SOb_SetLogLevel(1); int32_t spot_ids[2] = {-1, -1}; @@ -102,34 +112,58 @@ int main(int argc, char* argv[]) { // TODO: Setup a listener for ctrl-c to gracefully stop the connection // std::cout << "Press Ctrl-C to stop reading camera feeds..." << std::endl; - bool using_vision_pipeline = (argc == 6); - SObModel model = nullptr; + bool using_vision_pipeline = (argc >= 6); + const int32_t model_kind = (argc >= 7) ? std::atoi(argv[6]) : SOb_MODEL_SINGLE_SHOT; + const bool switching = (argc >= 8); + const int32_t model2_kind = (argc == 9) ? std::atoi(argv[8]) : SOb_MODEL_SINGLE_SHOT; + // Preferred pipeline stream: the two-camera front stream (index 0). + // Batch-capable streaming exports (dynamic or batch-2) run there directly; a + // fixed batch-1 export is refused by the native supportsBatch guard, and the + // launch loop falls back to the single-camera HAND stream (index 1). + int32_t vp_stream_idx = 0; + + SObModel models[2] = {nullptr, nullptr}; + int32_t active_model = 0; if (using_vision_pipeline) { - const char* model_path = argv[5]; - std::cout << "Loading model from: " << model_path << std::endl; - model = SOb_LoadModel(model_path, "cuda"); - if (!model) { - std::cerr << "Failed to load model from: " << argv[5] << std::endl; - disconnect_from_spots(spot_ids, 2); - cv::destroyAllWindows(); - - return -1; + // Preload every selectable model up front; switching later is + // handle-to-handle with no load in the hot path. + const int32_t n_models = switching ? 2 : 1; + for (int32_t m = 0; m < n_models; m++) { + const char* path = argv[5 + 2 * m]; + const int32_t kind = (m == 0) ? model_kind : model2_kind; + std::cout << "Loading model " << m << " from: " << path << " (kind " << kind << ")" << std::endl; + models[m] = SOb_LoadModelEx(path, "cuda", kind); + if (!models[m]) { + std::cerr << "Failed to load model from: " << path << std::endl; + disconnect_from_spots(spot_ids, 2); + if (models[0]) SOb_UnloadModel(models[0]); + cv::destroyAllWindows(); + return -1; + } } - std::cout << "Model loaded successfully!" << std::endl; + std::cout << "Model(s) loaded successfully!" << std::endl; - // Launch vision pipeline on both robots + // Launch vision pipeline on both robots. Try the front stream first; + // fall back to the HAND stream if the model can't take its batch (e.g. + // a fixed batch-1 streaming export against the two-camera front pair). for (size_t i = 0; i < 2; i++) { if (spot_ids[i] < 0) continue; - // Launch vision pipeline only on the first camera stream - bool ret = SOb_LaunchVisionPipeline(spot_ids[i], cam_stream_ids[spot_ids[i]][0], model); + bool ret = SOb_LaunchVisionPipeline(spot_ids[i], cam_stream_ids[spot_ids[i]][vp_stream_idx], models[0]); + if (!ret && vp_stream_idx == 0 && cam_stream_ids[spot_ids[i]].size() > 1) { + std::cout << "Front-stream launch refused (see native log); trying HAND stream" << std::endl; + vp_stream_idx = 1; + ret = SOb_LaunchVisionPipeline(spot_ids[i], cam_stream_ids[spot_ids[i]][vp_stream_idx], models[0]); + } if (!ret) { std::cerr << "Failed to launch vision pipeline on robot " << i << std::endl; disconnect_from_spots(spot_ids, 2); - SOb_UnloadModel(model); + for (auto m : models) if (m) SOb_UnloadModel(m); cv::destroyAllWindows(); return -1; } - std::cout << "Vision pipeline launched on robot " << i << std::endl; + SOb_SetDepthAveraging(spot_ids[i], cam_stream_ids[spot_ids[i]][vp_stream_idx], true); + std::cout << "Vision pipeline launched on robot " << i + << " (stream idx " << vp_stream_idx << ")" << std::endl; } } @@ -150,8 +184,24 @@ int main(int argc, char* argv[]) { bool new_images = false; time_point start_time = high_resolution_clock::now(); + time_point last_switch_time = high_resolution_clock::now(); + constexpr auto switch_interval = seconds(15); bool exit_requested = false; while (!exit_requested) { + // Live model switch: flip between the preloaded handles on a timer. + // The camera stream keeps running; only the inference side swaps. + if (using_vision_pipeline && switching && + high_resolution_clock::now() - last_switch_time >= switch_interval) { + active_model ^= 1; + for (int32_t spot = 0; spot < 2; spot++) { + if (spot_ids[spot] < 0) continue; + bool ok = SOb_SwitchVisionPipelineModel( + spot_ids[spot], cam_stream_ids[spot_ids[spot]][vp_stream_idx], models[active_model]); + std::cout << "Switched robot " << spot << " to model " << active_model + << (ok ? "" : " -- FAILED") << std::endl; + } + last_switch_time = high_resolution_clock::now(); + } if (new_images) { time_point end_time = high_resolution_clock::now(); auto duration = duration_cast(end_time - start_time); @@ -187,7 +237,7 @@ int main(int argc, char* argv[]) { uint8_t** images_set = images[stream]; float** depths_set = depths[stream]; - if (using_vision_pipeline && stream == 0) { + if (using_vision_pipeline && stream == vp_stream_idx) { if (!SOb_GetNextVisionPipelineImageSet(spot_id, cam_stream_id, int32_t(num_images_requested), images_set, depths_set)) { std::this_thread::sleep_for(std::chrono::milliseconds(10)); continue; @@ -220,8 +270,13 @@ int main(int argc, char* argv[]) { cv::cvtColor(image, image, cv::COLOR_RGBA2BGR); cv::normalize(depth, depth, 0, 1, cv::NORM_MINMAX); + // Make it unmistakable which depth is model output and which is + // raw sensor depth -- the pipeline stream moves depending on + // model kind, and raw registered depth masquerades convincingly. + const bool is_model_stream = using_vision_pipeline && stream == vp_stream_idx; + const std::string depth_tag = is_model_stream ? " Depth[MODEL]" : " Depth[RAW]"; cv::imshow("SPOT " + std::to_string(spot) + " Stream " + std::to_string(stream) + " RGB" + std::to_string(i), image); - cv::imshow("SPOT " + std::to_string(spot) + " Stream " + std::to_string(stream) + " Depth" + std::to_string(i), depth); + cv::imshow("SPOT " + std::to_string(spot) + " Stream " + std::to_string(stream) + depth_tag + std::to_string(i), depth); } if (cv::waitKey(1) == 'q') { exit_requested = true; @@ -234,7 +289,7 @@ int main(int argc, char* argv[]) { disconnect_from_spots(spot_ids, 2); cv::destroyAllWindows(); - if (model) SOb_UnloadModel(model); + for (auto m : models) if (m) SOb_UnloadModel(m); for (auto image_set : images) delete[] image_set; for (auto depth_set : depths) delete[] depth_set; diff --git a/tests/streaming-harness/gen_integ_test.py b/tests/streaming-harness/gen_integ_test.py new file mode 100644 index 0000000..cfc7324 --- /dev/null +++ b/tests/streaming-harness/gen_integ_test.py @@ -0,0 +1,346 @@ +"""Integration harness: runs the REAL StreamingONNXModel (verbatim from model.h / +model.cpp) plus the REAL resize kernels against a tiny stub model with the same IO +contract, on the ONNX Runtime CPU EP. + +Only two substitutions are made, both documented in the output: + * _setDevice -> CPU EP (the CUDA EP is the one function not covered here) + * cudaMalloc/Free -> host malloc/free, so the kernels operate on host buffers +Everything else -- name pairing, geometry, zero-length frame 0, the bind/Run/ +ping-pong loop, resetState -- is the shipping code. +""" +import pathlib as _pl +_HERE = _pl.Path(__file__).resolve().parent +_REPO = _HERE.parents[1] +_BUILD = _HERE / "build" +_BUILD.mkdir(exist_ok=True) +import re, pathlib + +SRC = pathlib.Path(str(_REPO / "src")) +OUT = pathlib.Path(str(_BUILD)) +OUT.mkdir(parents=True, exist_ok=True) + +header = (SRC / "include/model.h").read_text(encoding="utf-8", errors="replace") +impl = (SRC / "model.cpp").read_text(encoding="utf-8", errors="replace") +cu = (SRC / "cuda_kernels.cu").read_text(encoding="utf-8", errors="replace") + +h_start = header.index("class StreamingONNXModel") +h_rest = header[h_start:] +class_decl = h_rest[: re.search(r"^\};$", h_rest, re.M).end()] + +i_start = impl.index("// Bytes per element for the types a KV cache") +i_rest = impl[i_start:] +class_impl = i_rest[: re.search(r"^\};\s*$", i_rest, re.M).start()] + +kernels = cu[cu.index("__global__ void resize_bilinear_chw_kernel"): + cu.index("cudaError_t resize_bilinear_chw(")] + +# --- substitution 1: CPU execution provider ----------------------------------- +sd_start = class_impl.index("void StreamingONNXModel::_setDevice") +sd_end = class_impl.index("void StreamingONNXModel::_buildCacheNames") +class_impl = class_impl[:sd_start] + """void StreamingONNXModel::_setDevice(const std::string&) { + // TEST SUBSTITUTION: CPU EP instead of the CUDA EP. + m_use_cuda = true; // gate in the ctor; memory info stays CPU +} + +""" + class_impl[sd_end:] + +# --- test accessors (injected into the copy only) ------------------------------ +class_decl = class_decl.replace( + "public:\n explicit StreamingONNXModel(", + "public:\n" + " const std::vector& testCache() const { return m_cache; }\n" + " int64_t testFrames() const { return m_frames_seen; }\n" + " explicit StreamingONNXModel(", 1) + +tu = r"""// GENERATED integration harness -- not part of the build. +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +// ---- CUDA stubs: host memory, so the real kernels run on host buffers -------- +using cudaError_t = int; +using cudaStream_t = void*; +constexpr cudaError_t cudaSuccess = 0; +static size_t g_alloc_bytes = 0, g_alloc_count = 0, g_free_count = 0; +static cudaError_t cudaMalloc(void** p, size_t n) { *p = malloc(n); g_alloc_bytes += n; g_alloc_count++; return *p ? 0 : 1; } +template static cudaError_t cudaMalloc(T** p, size_t n) { return cudaMalloc(reinterpret_cast(p), n); } +static cudaError_t cudaFree(void* p) { if (p) { free(p); g_free_count++; } return 0; } +static cudaError_t cudaGetDeviceCount(int* c) { *c = 1; return 0; } +static cudaError_t cudaStreamSynchronize(cudaStream_t) { return 0; } +static const char* cudaGetErrorString(cudaError_t) { return "stub"; } + +#define __global__ +#define __device__ +#define __forceinline__ inline +#define __restrict__ +struct dim3 { unsigned x, y, z; dim3(unsigned a=1,unsigned b=1,unsigned c=1):x(a),y(b),z(c){} }; +static dim3 blockIdx, threadIdx, blockDim; +static const float CUDART_INF_F = HUGE_VALF; +using std::min; using std::max; + +namespace SOb { +namespace fs = std::filesystem; + +struct TensorShape { size_t N, C, H, W; }; + +static bool g_quiet = true; +template +void LogMessage(const std::format_string fmt, Args&&... args) { + if (!g_quiet) { printf(" [log] %s\n", std::format(fmt, std::forward(args)...).c_str()); } +} +inline void LogMessage(const std::string& s) { if (!g_quiet) printf(" [log] %s\n", s.c_str()); } +template +void LogPerf(const std::format_string fmt, Args&&... args) { + if (!g_quiet) { printf(" [perf] %s\n", std::format(fmt, std::forward(args)...).c_str()); } +} +inline void checkCudaError(cudaError_t error, const std::string& operation) { + if (error != cudaSuccess) throw std::runtime_error(operation); +} + +// ---- verbatim kernels from src/cuda_kernels.cu ------------------------------ +KERNELS_HERE + +// Host launchers matching the real .cu launcher signatures. +cudaError_t resize_bilinear_chw(const float* d_in, float* d_out, int in_h, int in_w, + int out_h, int out_w, int channels, cudaStream_t) { + blockDim = dim3(1,1,1); + for (int c = 0; c < channels; ++c) + for (int y = 0; y < out_h; ++y) + for (int x = 0; x < out_w; ++x) { + blockIdx = dim3(x,y,c); threadIdx = dim3(0,0,0); + resize_bilinear_chw_kernel(d_in, d_out, in_h, in_w, out_h, out_w, channels); + } + return 0; +} +cudaError_t resize_sparse_depth(const float* d_in, float* d_out, int in_h, int in_w, + int out_h, int out_w, cudaStream_t) { + blockDim = dim3(1,1,1); + for (int y = 0; y < out_h; ++y) + for (int x = 0; x < out_w; ++x) { + blockIdx = dim3(x,y,0); threadIdx = dim3(0,0,0); + resize_sparse_depth_kernel(d_in, d_out, in_h, in_w, out_h, out_w); + } + return 0; +} + +class MLModel { +public: + virtual ~MLModel() = default; + virtual bool runInference(const float*, const float*, float*, TensorShape, TensorShape, TensorShape) = 0; + virtual bool runInference(const uint8_t*, const float*, float*, TensorShape, TensorShape, TensorShape) = 0; + virtual void resetState() {} + virtual bool wantsFullResDepth() const { return false; } + virtual bool acquire(const void* owner) { (void)owner; return true; } + virtual void release(const void* owner) { (void)owner; } + virtual bool supportsBatch(int32_t n) const { return n >= 1; } +}; + +// ---- verbatim class declaration from src/include/model.h -------------------- +CLASS_DECL_HERE + +// ---- verbatim implementation from src/model.cpp ----------------------------- +CLASS_IMPL_HERE + +} // namespace SOb + +static int g_fail = 0; +static void check(bool ok, const char* name, const std::string& detail = "") { + printf("%-56s %s %s\n", name, ok ? "PASS" : "FAIL", detail.c_str()); + if (!ok) g_fail++; +} + +static float half_to_float(uint16_t h) { + uint32_t sign = (h >> 15) & 1u, exp = (h >> 10) & 0x1Fu, man = h & 0x3FFu, f; + if (exp == 0) { + if (man == 0) { f = sign << 31; } + else { uint32_t e = 127 - 15 + 1; while (!(man & 0x400u)) { man <<= 1; e--; } + man &= 0x3FFu; f = (sign << 31) | (e << 23) | (man << 13); } + } else if (exp == 31) { f = (sign << 31) | (0xFFu << 23) | (man << 13); } + else { f = (sign << 31) | ((exp - 15 + 127) << 23) | (man << 13); } + float out; std::memcpy(&out, &f, 4); return out; +} + +// Reads cache element 0 whatever the graph carries it as. +static float cache_elem0(const Ort::Value& v) { + if (v.GetTensorTypeAndShapeInfo().GetElementType() == ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16) + return half_to_float(*v.GetTensorData()); + return *v.GetTensorData(); +} + +static void run_suite(const char* model, bool fp16_cache) { + using namespace SOb; + const int WINDOW = 4; // matches make_stub_model.py + const size_t CH = 12, CW = 16; // "camera" resolution; model is 8x10 + const char* tag = fp16_cache ? "[fp16 cache] " : "[fp32 cache] "; + printf("\n--- %s ---\n", fp16_cache ? "fp16 caches, fp32 image IO (hybrid recipe)" + : "fp32 throughout (current export)"); + + std::unique_ptr m; + try { + m = std::make_unique(model, "cuda"); + } catch (const std::exception& e) { + printf("%sconstruction threw: %s\n", tag, e.what()); + g_fail++; return; + } + check(true, (std::string(tag) + "construct: graph accepted").c_str()); + + const auto want = fp16_cache ? ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT16 + : ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT; + check(m->testCache().size() == 48 && + m->testCache()[0].GetTensorTypeAndShapeInfo().GetElementType() == want, + (std::string(tag) + "reset: empty caches match the graph's element type").c_str()); + check(m->testCache()[0].GetTensorTypeAndShapeInfo().GetShape()[2] == 0, + (std::string(tag) + "reset: frame-0 cache has n_frames == 0").c_str()); + + TensorShape in_s{1, 3, CH, CW}, dp_s{1, 1, CH, CW}, out_s{1, 1, CH, CW}; + std::vector rgb(3 * CH * CW, 1.0f), depth(CH * CW, 2.0f), out(CH * CW, -1.f); + + bool growth_ok = true, content_ok = true, out_ok = true, type_ok = true; + std::string seq; + for (int f = 0; f < 7; ++f) { + std::fill(out.begin(), out.end(), -1.f); + if (!m->runInference(rgb.data(), depth.data(), out.data(), in_s, dp_s, out_s)) { + check(false, (std::string(tag) + "runInference returned false").c_str(), + std::format("frame {}", f)); + return; + } + int64_t n = m->testCache()[0].GetTensorTypeAndShapeInfo().GetShape()[2]; + seq += std::to_string(n) + (f < 6 ? "," : ""); + if (n != std::min(f + 1, WINDOW)) growth_ok = false; + + for (int i = 0; i < 48; ++i) { + if (cache_elem0(m->testCache()[i]) != float(i)) content_ok = false; + if (m->testCache()[i].GetTensorTypeAndShapeInfo().GetElementType() != want) type_ok = false; + } + for (size_t i = 0; i < out.size(); ++i) + if (std::fabs(out[i] - 1.0f) > 1e-4f) out_ok = false; + } + check(growth_ok, (std::string(tag) + "stream: n_frames grows then clamps").c_str(), "seq=" + seq); + check(content_ok, (std::string(tag) + "stream: cache[i] carries layer i").c_str()); + check(type_ok, (std::string(tag) + "stream: cache element type stable across frames").c_str()); + check(out_ok, (std::string(tag) + "stream: depth round-trips camera->model->camera").c_str()); + + m->resetState(); + check(m->testCache()[0].GetTensorTypeAndShapeInfo().GetShape()[2] == 0 && m->testFrames() == 0, + (std::string(tag) + "reset: mid-stream reset returns to zero-length").c_str()); + m->runInference(rgb.data(), depth.data(), out.data(), in_s, dp_s, out_s); + check(m->testCache()[0].GetTensorTypeAndShapeInfo().GetShape()[2] == 1, + (std::string(tag) + "reset: sequence restarts cleanly").c_str()); + + int pipeline_a = 0, pipeline_b = 0; + check(m->acquire(&pipeline_a), (std::string(tag) + "ownership: first pipeline acquires").c_str()); + check(!m->acquire(&pipeline_b), (std::string(tag) + "ownership: second pipeline refused").c_str()); + m->release(&pipeline_a); + check(m->acquire(&pipeline_b), (std::string(tag) + "ownership: released is reusable").c_str()); + m->release(&pipeline_b); + + size_t before = g_free_count; + m.reset(); + check(g_free_count > before, (std::string(tag) + "teardown: scratch freed").c_str()); +} + +static void run_batch_suite(const char* dyn_model, const char* static_model) { + using namespace SOb; + const size_t CH = 12, CW = 16; // "camera" resolution; model is 8x10 + printf("\n--- dynamic batch (symbolic B, per-slot sequences) ---\n"); + + std::unique_ptr m; + try { + m = std::make_unique(dyn_model, "cuda"); + } catch (const std::exception& e) { + printf("[dyn] construction threw: %s\n", e.what()); + g_fail++; return; + } + check(true, "[dyn] construct: symbolic-batch graph accepted"); + check(m->supportsBatch(1) && m->supportsBatch(2) && !m->supportsBatch(0), + "[dyn] supportsBatch: any positive batch accepted"); + + // Phase 1: batch 1, two frames. + TensorShape in1{1, 3, CH, CW}, dp1{1, 1, CH, CW}, out1{1, 1, CH, CW}; + std::vector rgb1(3 * CH * CW, 1.0f), depth1(CH * CW, 2.0f), out_b1(CH * CW, -1.f); + bool ok = m->runInference(rgb1.data(), depth1.data(), out_b1.data(), in1, dp1, out1) + && m->runInference(rgb1.data(), depth1.data(), out_b1.data(), in1, dp1, out1); + auto cshape = m->testCache()[0].GetTensorTypeAndShapeInfo().GetShape(); + check(ok && cshape[0] == 1 && cshape[2] == 2, + "[dyn] batch 1: runs, cache [1,..,2,..]", + std::format("batch={} frames={}", cshape[0], cshape[2])); + + // Phase 2: switch to batch 2 mid-stream -> sequence must restart cleanly. + TensorShape in2{2, 3, CH, CW}, dp2{2, 1, CH, CW}, out2{2, 1, CH, CW}; + std::vector rgb2(2 * 3 * CH * CW), depth2(2 * CH * CW, 2.0f), out_b2(2 * CH * CW, -1.f); + std::fill(rgb2.begin(), rgb2.begin() + 3 * CH * CW, 1.0f); // slot 0 + std::fill(rgb2.begin() + 3 * CH * CW, rgb2.end(), 3.0f); // slot 1 + ok = m->runInference(rgb2.data(), depth2.data(), out_b2.data(), in2, dp2, out2); + cshape = m->testCache()[0].GetTensorTypeAndShapeInfo().GetShape(); + check(ok && cshape[0] == 2 && cshape[2] == 1 && m->testFrames() == 1, + "[dyn] batch change 1->2: restarts sequence, cache [2,..,1,..]", + std::format("batch={} frames={} counter={}", cshape[0], cshape[2], m->testFrames())); + + // Phase 3: batch-2 streaming -- growth, per-slot output, cache identity. + bool growth_ok = true, slot_ok = true, content_ok = true; + for (int f = 0; f < 5; ++f) { + std::fill(out_b2.begin(), out_b2.end(), -1.f); + if (!m->runInference(rgb2.data(), depth2.data(), out_b2.data(), in2, dp2, out2)) { + check(false, "[dyn] batch-2 runInference returned false"); return; + } + int64_t n = m->testCache()[0].GetTensorTypeAndShapeInfo().GetShape()[2]; + if (n != std::min(f + 2, 4)) growth_ok = false; // continues from frame 1 + for (size_t i = 0; i < CH * CW; ++i) { + if (std::fabs(out_b2[i] - 1.0f) > 1e-4f) slot_ok = false; // slot 0 + if (std::fabs(out_b2[CH * CW + i] - 3.0f) > 1e-4f) slot_ok = false; // slot 1 + } + for (int i = 0; i < 48; ++i) + if (cache_elem0(m->testCache()[i]) != float(i)) content_ok = false; + } + check(growth_ok, "[dyn] batch 2: n_frames grows then clamps"); + check(slot_ok, "[dyn] batch 2: slot outputs independent (1.0 / 3.0)"); + check(content_ok, "[dyn] batch 2: cache[i] carries layer i"); + + // Phase 4: back to batch 1 -> restart again. + ok = m->runInference(rgb1.data(), depth1.data(), out_b1.data(), in1, dp1, out1); + cshape = m->testCache()[0].GetTensorTypeAndShapeInfo().GetShape(); + check(ok && cshape[0] == 1 && cshape[2] == 1, + "[dyn] batch change 2->1: restarts sequence, cache [1,..,1,..]"); + + // Fixed-batch graph must refuse a batch-2 call cleanly, not crash. + std::unique_ptr ms; + try { + ms = std::make_unique(static_model, "cuda"); + } catch (const std::exception& e) { + printf("[static] construction threw: %s\n", e.what()); g_fail++; return; + } + check(!ms->supportsBatch(2), "[static] fixed batch-1 graph reports supportsBatch(2)==false"); + ok = ms->runInference(rgb2.data(), depth2.data(), out_b2.data(), in2, dp2, out2); + check(!ok, "[static] batch-2 call on fixed-batch-1 graph fails cleanly"); +} + +int main() { + run_suite("STUB_PATH", false); + run_suite("STUB_PATH_FP16", true); + run_batch_suite("STUB_PATH_DYN", "STUB_PATH"); + printf("\n%s (%d failure(s))\n", + g_fail ? "FAILURES PRESENT" : "ALL INTEGRATION TESTS PASSED", g_fail); + return g_fail ? 1 : 0; +} +""" + +tu = tu.replace("KERNELS_HERE", kernels) +tu = tu.replace("CLASS_DECL_HERE", class_decl) +tu = tu.replace("CLASS_IMPL_HERE", class_impl) +tu = tu.replace("STUB_PATH_FP16", str(OUT / "stub_stream_fp16cache.onnx").replace("\\", "/")) +tu = tu.replace("STUB_PATH_DYN", str(OUT / "stub_stream_dyn.onnx").replace("\\", "/")) +tu = tu.replace("STUB_PATH", str(OUT / "stub_stream.onnx").replace("\\", "/")) +(OUT / "integ_test.cpp").write_text(tu, encoding="utf-8") +print("wrote", OUT / "integ_test.cpp") diff --git a/tests/streaming-harness/gen_kernel_test.py b/tests/streaming-harness/gen_kernel_test.py new file mode 100644 index 0000000..560bd97 --- /dev/null +++ b/tests/streaming-harness/gen_kernel_test.py @@ -0,0 +1,168 @@ +"""Extract the two resize kernels verbatim from cuda_kernels.cu and wrap them in a +host-side harness so their math can be executed and checked without nvcc/a GPU.""" +import pathlib as _pl +_HERE = _pl.Path(__file__).resolve().parent +_REPO = _HERE.parents[1] +_BUILD = _HERE / "build" +_BUILD.mkdir(exist_ok=True) +import pathlib + +SRC = pathlib.Path(str(_REPO / "src/cuda_kernels.cu")) +OUT = pathlib.Path(str(_BUILD)) +OUT.mkdir(parents=True, exist_ok=True) + +txt = SRC.read_text(encoding="utf-8", errors="replace") +start = txt.index("__global__ void resize_bilinear_chw_kernel") +end = txt.index("cudaError_t resize_bilinear_chw(") +kernels = txt[start:end] +print(f"extracted {kernels.count(chr(10))} lines of kernel source") + +harness = r"""// GENERATED kernel test harness -- not part of the build. +// Executes the real kernel bodies on the host with CUDA builtins stubbed, so the +// resize math can be validated on a machine with no CUDA toolkit and no GPU. +#include +#include +#include +#include +#include +#include + +#define __global__ +#define __device__ +#define __forceinline__ inline +#define __restrict__ + +struct dim3 { unsigned x, y, z; dim3(unsigned a=1,unsigned b=1,unsigned c=1):x(a),y(b),z(c){} }; +static dim3 blockIdx, threadIdx, blockDim; +static const float CUDART_INF_F = HUGE_VALF; +using std::min; using std::max; + +// ---- verbatim kernels from src/cuda_kernels.cu ------------------------------ +KERNELS_HERE +// ----------------------------------------------------------------------------- + +static int g_failures = 0; +static void check(bool ok, const char* name, const char* detail = "") { + printf("%-58s %s %s\n", name, ok ? "PASS" : "FAIL", detail); + if (!ok) g_failures++; +} + +// Serial stand-ins for the <<>> launches. +static void launch_bilinear(const float* src, float* dst, int in_h, int in_w, + int out_h, int out_w, int ch) { + blockDim = dim3(1,1,1); + for (int c = 0; c < ch; ++c) + for (int y = 0; y < out_h; ++y) + for (int x = 0; x < out_w; ++x) { + blockIdx = dim3(x, y, c); threadIdx = dim3(0,0,0); + resize_bilinear_chw_kernel(src, dst, in_h, in_w, out_h, out_w, ch); + } +} +static void launch_sparse(const float* src, float* dst, int in_h, int in_w, + int out_h, int out_w) { + blockDim = dim3(1,1,1); + for (int y = 0; y < out_h; ++y) + for (int x = 0; x < out_w; ++x) { + blockIdx = dim3(x, y, 0); threadIdx = dim3(0,0,0); + resize_sparse_depth_kernel(src, dst, in_h, in_w, out_h, out_w); + } +} + +int main() { + // 1. Identity resize must be exact (half-pixel centres collapse to integers). + { + const int H = 37, W = 53; + std::vector src(3*H*W), dst(3*H*W, -1.f); + for (size_t i = 0; i < src.size(); ++i) src[i] = float(i % 97) * 0.37f; + launch_bilinear(src.data(), dst.data(), H, W, H, W, 3); + float worst = 0.f; + for (size_t i = 0; i < src.size(); ++i) worst = max(worst, fabsf(src[i]-dst[i])); + char d[64]; snprintf(d, sizeof d, "max|diff|=%.3e", worst); + check(worst < 1e-4f, "bilinear: identity resize is exact", d); + } + + // 2. Bilinear reproduces a linear ramp exactly (non-circular analytic check). + // Interior only: edge clamping intentionally breaks linearity at borders. + { + const int IH = 480, IW = 640, OH = 392, OW = 518; + const float a = 0.013f, b = -0.007f, c0 = 1.5f; + std::vector src(IH*IW), dst(OH*OW, -1.f); + for (int y = 0; y < IH; ++y) for (int x = 0; x < IW; ++x) src[y*IW+x] = a*x + b*y + c0; + launch_bilinear(src.data(), dst.data(), IH, IW, OH, OW, 1); + float worst = 0.f; + for (int y = 1; y < OH-1; ++y) for (int x = 1; x < OW-1; ++x) { + float sy = (y+0.5f)*IH/OH - 0.5f, sx = (x+0.5f)*IW/OW - 0.5f; + worst = max(worst, fabsf(dst[y*OW+x] - (a*sx + b*sy + c0))); + } + char d[64]; snprintf(d, sizeof d, "max|diff|=%.3e", worst); + check(worst < 1e-3f, "bilinear: reproduces linear ramp (480x640->392x518)", d); + } + + // 4. Sparse depth must never invent a value: every output is 0 or a real input. + { + const int IH = 480, IW = 640, OH = 392, OW = 518; + std::vector src(IH*IW, 0.f), dst(OH*OW, -1.f); + unsigned s = 12345; + int n_valid_in = 0; + for (int i = 0; i < IH*IW; ++i) { + s = s*1664525u + 1013904223u; + if ((s >> 16) % 100 < 30) { src[i] = 0.5f + float((s>>8)%1000)*0.0075f; n_valid_in++; } + } + std::set allowed(src.begin(), src.end()); + allowed.insert(0.f); + launch_sparse(src.data(), dst.data(), IH, IW, OH, OW); + bool invented = false; int n_valid_out = 0; + for (int i = 0; i < OH*OW; ++i) { + if (dst[i] != 0.f) n_valid_out++; + if (!allowed.count(dst[i])) invented = true; + } + char d[96]; + snprintf(d, sizeof d, "valid in=%d out=%d (%.0f%% kept)", n_valid_in, n_valid_out, + 100.0*n_valid_out/ (double)(OH*OW) / 0.30); + check(!invented, "sparse depth: never interpolates an invented value", d); + check(n_valid_out > 0.6*OH*OW*0.30, "sparse depth: retains most valid samples", d); + } + + // 5. Degenerate inputs: 0 and NaN are the only "no sample" markers. + { + const int IH = 64, IW = 64, OH = 50, OW = 50; + std::vector src(IH*IW, 0.f), dst(OH*OW, 9.f); + launch_sparse(src.data(), dst.data(), IH, IW, OH, OW); + bool all_zero = std::all_of(dst.begin(), dst.end(), [](float v){ return v == 0.f; }); + check(all_zero, "sparse depth: all-invalid input yields all-zero output"); + + std::fill(src.begin(), src.end(), NAN); + std::fill(dst.begin(), dst.end(), 9.f); + launch_sparse(src.data(), dst.data(), IH, IW, OH, OW); + all_zero = std::all_of(dst.begin(), dst.end(), [](float v){ return v == 0.f; }); + check(all_zero, "sparse depth: NaN rejected (negated compare)"); + } + + // 6. Upscaling must still write every output pixel (empty-footprint guard). + { + const int IH = 392, IW = 518, OH = 480, OW = 640; + std::vector src(IH*IW, 2.5f), dst(OH*OW, -1.f); + launch_sparse(src.data(), dst.data(), IH, IW, OH, OW); + bool all_set = std::all_of(dst.begin(), dst.end(), [](float v){ return v == 2.5f; }); + check(all_set, "sparse depth: upscale 392x518->480x640 fills every pixel"); + } + + // 7. Canary check for out-of-bounds writes on the model-resolution path. + { + const int IH = 480, IW = 640, OH = 392, OW = 518; + std::vector src(3*IH*IW, 1.f); + std::vector buf(3*OH*OW + 64, -12345.f); + launch_bilinear(src.data(), buf.data(), IH, IW, OH, OW, 3); + bool canary_ok = true; + for (int i = 0; i < 64; ++i) if (buf[3*OH*OW + i] != -12345.f) canary_ok = false; + check(canary_ok, "bilinear: no writes past the output buffer"); + } + + printf("\n%s (%d failure(s))\n", g_failures ? "FAILURES PRESENT" : "ALL KERNEL TESTS PASSED", g_failures); + return g_failures ? 1 : 0; +} +""" + +harness = harness.replace("KERNELS_HERE", kernels) +(OUT / "kern_test.cpp").write_text(harness, encoding="utf-8") +print("wrote", OUT / "kern_test.cpp") diff --git a/tests/streaming-harness/gen_ort_check.py b/tests/streaming-harness/gen_ort_check.py new file mode 100644 index 0000000..983e4f9 --- /dev/null +++ b/tests/streaming-harness/gen_ort_check.py @@ -0,0 +1,102 @@ +"""Extract StreamingONNXModel verbatim from the real sources into a single TU +that can be compiled against the real ONNX Runtime headers, with CUDA/torch/log +dependencies stubbed. Purpose: type-check the ~350 lines of Ort C++ API usage on +a machine that has no CUDA toolkit and no libtorch.""" +import pathlib as _pl +_HERE = _pl.Path(__file__).resolve().parent +_REPO = _HERE.parents[1] +_BUILD = _HERE / "build" +_BUILD.mkdir(exist_ok=True) +import re, pathlib + +SRC = pathlib.Path(str(_REPO / "src")) +OUT = pathlib.Path(str(_BUILD)) +OUT.mkdir(parents=True, exist_ok=True) + +header = (SRC / "include/model.h").read_text(encoding="utf-8", errors="replace") +impl = (SRC / "model.cpp").read_text(encoding="utf-8", errors="replace") + +# class decl: from "class StreamingONNXModel" up to the first line that is exactly "};" +h_start = header.index("class StreamingONNXModel") +h_rest = header[h_start:] +m = re.search(r"^\};$", h_rest, re.M) +class_decl = h_rest[: m.end()] + +# impl: from the ctor to just before the final namespace close +i_start = impl.index("// Bytes per element for the types a KV cache") +i_rest = impl[i_start:] +m2 = re.search(r"^\};\s*$", i_rest, re.M) +class_impl = i_rest[: m2.start()] + +print(f"extracted class decl: {class_decl.count(chr(10))} lines") +print(f"extracted impl : {class_impl.count(chr(10))} lines") + +tu = f"""// GENERATED compile-check harness -- not part of the build. +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +// ---- stubs standing in for the real project headers ------------------------- +using cudaError_t = int; +using cudaStream_t = void*; +constexpr cudaError_t cudaSuccess = 0; +cudaError_t cudaMalloc(void** p, size_t n); +template cudaError_t cudaMalloc(T** p, size_t n) {{ return cudaMalloc(reinterpret_cast(p), n); }} +cudaError_t cudaFree(void* p); +cudaError_t cudaGetDeviceCount(int* c); +cudaError_t cudaStreamSynchronize(cudaStream_t s); +const char* cudaGetErrorString(cudaError_t e); + +namespace SOb {{ + +namespace fs = std::filesystem; + +struct TensorShape {{ + size_t N, C, H, W; +}}; + +template +void LogMessage(const std::format_string fmt, Args&&... args) {{ (void)fmt; }} +inline void LogMessage(const std::string& s) {{ (void)s; }} +template +void LogPerf(const std::format_string fmt, Args&&... args) {{ (void)fmt; }} + +inline void checkCudaError(cudaError_t error, const std::string& operation) {{ + if (error != cudaSuccess) throw std::runtime_error(operation); +}} + +cudaError_t resize_bilinear_chw(const float* d_in, float* d_out, int in_h, int in_w, + int out_h, int out_w, int channels, cudaStream_t stream); +cudaError_t resize_sparse_depth(const float* d_in, float* d_out, int in_h, int in_w, + int out_h, int out_w, cudaStream_t stream); + +class MLModel {{ +public: + virtual ~MLModel() = default; + virtual bool runInference(const float*, const float*, float*, TensorShape, TensorShape, TensorShape) = 0; + virtual bool runInference(const uint8_t*, const float*, float*, TensorShape, TensorShape, TensorShape) = 0; + virtual void resetState() {{}} + virtual bool wantsFullResDepth() const {{ return false; }} + virtual bool acquire(const void* owner) {{ (void)owner; return true; }} + virtual void release(const void* owner) {{ (void)owner; }} + virtual bool supportsBatch(int32_t n) const {{ return n >= 1; }} +}}; + +// ---- verbatim class declaration from src/include/model.h -------------------- +{class_decl} + +// ---- verbatim implementation from src/model.cpp ----------------------------- +{class_impl} + +}} // namespace SOb +""" + +(OUT / "ort_check.cpp").write_text(tu, encoding="utf-8") +print("wrote", OUT / "ort_check.cpp") diff --git a/tests/streaming-harness/make_stub_model.py b/tests/streaming-harness/make_stub_model.py new file mode 100644 index 0000000..4b4ca24 --- /dev/null +++ b/tests/streaming-harness/make_stub_model.py @@ -0,0 +1,102 @@ +"""Build tiny ONNX models with the exact IO contract of the real streaming model +(2 + 48 in, 2 + 48 out, growing n_frames axis with a sliding window) so the real +C++ StreamingONNXModel can be exercised without loading the multi-GB export. + +Variants: + (default) fp32 caches, batch fixed at 1 + --fp16 fp16 caches, fp32 image IO (the hybrid conversion recipe) + --dyn symbolic batch dim "B" everywhere, like the dynamic re-export; + each batch slot is an independent sequence + +Each cache i appends a slab filled with the constant float(i), so the C++ harness +can verify that cache[i] really carries layer i -- i.e. that GetOutputValues() +comes back in bind order and nothing is cross-wired. In the --dyn variant the +slab is Expand-ed to the runtime batch, so the check holds per slot. + +depth = per-image mean over rgb channels, so output depends on the real input and +batch slots stay distinguishable. +""" +import pathlib as _pl +import sys + +import numpy as np +import onnx +from onnx import helper, TensorProto, numpy_helper + +_HERE = _pl.Path(__file__).resolve().parent +_BUILD = _HERE / "build" +_BUILD.mkdir(exist_ok=True) + +H, W, HEADS, P, HDIM, WINDOW, LAYERS = 8, 10, 2, 3, 4, 4, 24 + +CACHE_FP16 = "--fp16" in sys.argv +DYN_BATCH = "--dyn" in sys.argv +CACHE_T = TensorProto.FLOAT16 if CACHE_FP16 else TensorProto.FLOAT +CACHE_NP = np.float16 if CACHE_FP16 else np.float32 +BATCH = "B" if DYN_BATCH else 1 + +name = "stub_stream" +if CACHE_FP16: + name += "_fp16cache" +if DYN_BATCH: + name += "_dyn" +OUT = _BUILD / f"{name}.onnx" + +inputs, outputs, nodes, inits = [], [], [], [] + +inputs.append(helper.make_tensor_value_info("rgb", TensorProto.FLOAT, [BATCH, 3, H, W])) +inputs.append(helper.make_tensor_value_info("sparse_depth", TensorProto.FLOAT, [BATCH, 1, H, W])) +outputs.append(helper.make_tensor_value_info("depth", TensorProto.FLOAT, [BATCH, H, W])) +outputs.append(helper.make_tensor_value_info("depth_conf", TensorProto.FLOAT, [BATCH, H, W])) + +nodes.append(helper.make_node("ReduceMean", ["rgb"], ["rgb_mean"], axes=[1], keepdims=0)) +nodes.append(helper.make_node("Identity", ["rgb_mean"], ["depth"])) +nodes.append(helper.make_node("ReduceMean", ["sparse_depth"], ["sd_mean"], axes=[1], keepdims=0)) +nodes.append(helper.make_node("Identity", ["sd_mean"], ["depth_conf"])) + +# Slice(concat, starts=[-WINDOW], ends=[huge], axes=[2]) grows then clamps, +# exactly like the real graph's retention window. +inits.append(numpy_helper.from_array(np.array([-WINDOW], dtype=np.int64), "slice_start")) +inits.append(numpy_helper.from_array(np.array([2**31], dtype=np.int64), "slice_end")) +inits.append(numpy_helper.from_array(np.array([2], dtype=np.int64), "slice_axis")) + +if DYN_BATCH: + # Runtime slab shape [B, HEADS, 1, P, HDIM], derived from rgb's batch dim. + inits.append(numpy_helper.from_array(np.array([0], dtype=np.int64), "b_idx")) + inits.append(numpy_helper.from_array( + np.array([HEADS, 1, P, HDIM], dtype=np.int64), "slab_tail")) + nodes.append(helper.make_node("Shape", ["rgb"], ["rgb_shape"])) + nodes.append(helper.make_node("Gather", ["rgb_shape", "b_idx"], ["b_1d"], axis=0)) + nodes.append(helper.make_node("Concat", ["b_1d", "slab_tail"], ["slab_shape"], axis=0)) + +idx = 0 +for layer in range(LAYERS): + for kind in ("k", "v"): + pin = f"past_{kind}_{layer:02d}" + pout = f"new_{kind}_{layer:02d}" + inputs.append(helper.make_tensor_value_info( + pin, CACHE_T, [BATCH, HEADS, "n_frames", P, HDIM])) + outputs.append(helper.make_tensor_value_info( + pout, CACHE_T, [BATCH, HEADS, "n_frames_out", P, HDIM])) + + slab = np.full((1, HEADS, 1, P, HDIM), float(idx), dtype=CACHE_NP) + inits.append(numpy_helper.from_array(slab, f"slab_{idx}")) + slab_src = f"slab_{idx}" + if DYN_BATCH: + nodes.append(helper.make_node( + "Expand", [f"slab_{idx}", "slab_shape"], [f"slab_b_{idx}"])) + slab_src = f"slab_b_{idx}" + nodes.append(helper.make_node("Concat", [pin, slab_src], [f"cat_{idx}"], axis=2)) + nodes.append(helper.make_node( + "Slice", [f"cat_{idx}", "slice_start", "slice_end", "slice_axis"], [pout])) + idx += 1 + +graph = helper.make_graph(nodes, name, inputs, outputs, initializer=inits) +model = helper.make_model(graph, opset_imports=[helper.make_opsetid("", 17)]) +model.ir_version = 8 +onnx.checker.check_model(model) +onnx.save(model, OUT) +print(f"wrote {OUT}") +print(f"inputs={len(inputs)} outputs={len(outputs)} window={WINDOW} " + f"geometry H={H} W={W} heads={HEADS} P={P} head_dim={HDIM} " + f"cache={'FLOAT16' if CACHE_FP16 else 'FLOAT'} batch={BATCH}") diff --git a/tests/streaming-harness/run_all.bat b/tests/streaming-harness/run_all.bat new file mode 100644 index 0000000..b5d965f --- /dev/null +++ b/tests/streaming-harness/run_all.bat @@ -0,0 +1,57 @@ +@echo off +REM Host-side test harness for StreamingONNXModel. Runs without a GPU or the CUDA +REM toolkit: it extracts the real class and the real kernels from src/ and compiles +REM them against the real ONNX Runtime headers, stubbing only CUDA allocation and +REM the CUDA execution provider. +REM +REM Requires: MSVC, python with onnx + numpy, extern/onnxruntime-win-x64-gpu-1.22.0. +REM Usage: run_all.bat [path-to-python] + +setlocal +set HERE=%~dp0 +set REPO=%HERE%..\.. +set ORT=%REPO%\extern\onnxruntime-win-x64-gpu-1.22.0 +set PY=%~1 +if "%PY%"=="" set PY=python + +if not exist "%ORT%\include\onnxruntime_cxx_api.h" ( + echo ERROR: ONNX Runtime not found at %ORT% + echo Download onnxruntime-win-x64-gpu-1.22.0 and extract it under extern\. + exit /b 1 +) + +call "C:\Program Files\Microsoft Visual Studio\2022\Community\VC\Auxiliary\Build\vcvars64.bat" >nul +if errorlevel 1 ( echo ERROR: vcvars64.bat failed & exit /b 1 ) + +echo === generating harnesses from src/ === +"%PY%" "%HERE%make_stub_model.py" || exit /b 1 +"%PY%" "%HERE%make_stub_model.py" --fp16 || exit /b 1 +"%PY%" "%HERE%make_stub_model.py" --dyn || exit /b 1 +"%PY%" "%HERE%gen_ort_check.py" || exit /b 1 +"%PY%" "%HERE%gen_kernel_test.py" || exit /b 1 +"%PY%" "%HERE%gen_integ_test.py" || exit /b 1 + +set B=%HERE%build + +echo. +echo === 1/3 ORT API compile check === +cl /c /nologo /std:c++20 /EHsc /permissive- /I "%ORT%\include" ^ + /Fo:"%B%\ort_check.obj" "%B%\ort_check.cpp" || exit /b 1 + +echo. +echo === 2/3 kernel math tests === +cl /nologo /std:c++20 /EHsc /O2 /fp:precise ^ + /Fe:"%B%\kern_test.exe" /Fo:"%B%\kern_test.obj" "%B%\kern_test.cpp" || exit /b 1 +"%B%\kern_test.exe" || exit /b 1 + +echo. +echo === 3/3 integration tests (fp32 + fp16 caches) === +if not exist "%B%\onnxruntime.dll" copy /y "%ORT%\lib\onnxruntime.dll" "%B%\" >nul +cl /nologo /std:c++20 /EHsc /O2 /fp:precise /I "%ORT%\include" ^ + /Fe:"%B%\integ_test.exe" /Fo:"%B%\integ_test.obj" "%B%\integ_test.cpp" ^ + /link "%ORT%\lib\onnxruntime.lib" || exit /b 1 +"%B%\integ_test.exe" || exit /b 1 + +echo. +echo ALL HARNESS SUITES PASSED +endlocal