test_inference.cpp

#include <iostream>
#include <vector>
#include <future>
#include <chrono>
#include "o2_hal_filter.h"

// Define the namespace alias for convenience
using namespace ALICE::O2::Tracking;

int main() {
    std::cout << "=========================================================\n";
    std::cout << " ALICE ITS - C++ Hardware Abstraction Layer Test Harness \n";
    std::cout << "=========================================================\n\n";

    try {
        // 1. Initialize the Filter with the exported ONNX model
        std::string model_path = "helical_flow_tracker_fp8.onnx";
        NeuralODEFilter tracker(model_path);

        // 2. Create a mock batch of 10,000 tracks (4 floats per track: pT, pz, R, q)
        size_t num_tracks = 10000;
        size_t num_elements = num_tracks * 4;
        
        std::cout << "Allocating pinned host memory for " << num_tracks << " tracks...\n";
        // We use the custom HAL allocator to ensure zero-copy PCIe transfers
        float* pinned_input_data = tracker.allocatePinned<float>(num_elements);

        // Populate with dummy kinematic data roughly matching ALICE distributions
        for (size_t i = 0; i < num_elements; i += 4) {
            pinned_input_data[i]   = 1.5f;   // pT (GeV/c)
            pinned_input_data[i+1] = 0.5f;   // pz (GeV/c)
            pinned_input_data[i+2] = 39.5f;  // hit_r (cm)
            pinned_input_data[i+3] = 1.0f;   // charge (q)
        }

        // Convert raw pointer to std::vector for the async interface
        std::vector<float> seed_kinematics(pinned_input_data, pinned_input_data + num_elements);

        std::cout << "Dispatching asynchronous FP8 inference to GPU stream...\n";
        
        auto start_time = std::chrono::high_resolution_clock::now();

        // 3. Fire the Asynchronous Inference
        // This returns immediately, allowing the main CPU thread to continue
        std::future<std::vector<float>> future_residuals = tracker.evaluateBatchAsync(seed_kinematics);

        // ... O2 framework does other tasks here ...

        // 4. Await the GPU results
        std::vector<float> results = future_residuals.get();
        
        auto end_time = std::chrono::high_resolution_clock::now();
        auto duration = std::chrono::duration_cast<std::chrono::microseconds>(end_time - start_time);

        std::cout << "\nSUCCESS: Processed " << num_tracks << " tracks.\n";
        std::cout << "Total Async Turnaround Time: " << duration.count() / 1000.0 << " ms\n";
        std::cout << "Sample Residual [Delta Phi, Delta Z]: [" << results[0] << ", " << results[1] << "]\n";

        // Free the pinned memory
        tracker.freePinned(pinned_input_data);

    } catch (const std::exception& e) {
        std::cerr << "CRITICAL ERROR: " << e.what() << "\n";
        return 1;
    }

    return 0;
}