From cc8259469aab6c2cf55edf05d02ff38fa6762b6b Mon Sep 17 00:00:00 2001
From: Andrzej Janik <vosen@vosen.pl>
Date: Wed, 24 Sep 2025 17:08:11 +0000
Subject: [PATCH] Add output

---
 Cargo.lock               |  1 +
 zluda_replay/Cargo.toml  |  1 +
 zluda_replay/src/main.rs | 11 +++++++++++
 3 files changed, 13 insertions(+)

diff --git a/Cargo.lock b/Cargo.lock
index 31ec2d8..bc96359 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -3833,6 +3833,7 @@ dependencies = [
  "cuda_macros",
  "cuda_types",
  "libloading",
+ "rustc-hash 2.0.0",
  "zluda_trace_common",
 ]
 
diff --git a/zluda_replay/Cargo.toml b/zluda_replay/Cargo.toml
index 73295d4..7ccdeb1 100644
--- a/zluda_replay/Cargo.toml
+++ b/zluda_replay/Cargo.toml
@@ -12,6 +12,7 @@ zluda_trace_common = { path = "../zluda_trace_common" }
 cuda_macros = { path = "../cuda_macros" }
 cuda_types = { path = "../cuda_types" }
 libloading = "0.8"
+rustc-hash = "2.0.0"
 
 [package.metadata.zluda]
 debug_only = true
diff --git a/zluda_replay/src/main.rs b/zluda_replay/src/main.rs
index 47b3ccd..928ec73 100644
--- a/zluda_replay/src/main.rs
+++ b/zluda_replay/src/main.rs
@@ -1,6 +1,7 @@
 use std::mem;
 
 use cuda_types::cuda::{CUdeviceptr_v2, CUstream};
+use rustc_hash::FxHashMap;
 
 struct CudaDynamicFns {
     handle: libloading::Library,
@@ -36,6 +37,7 @@ fn main() {
     unsafe { libcuda.cuCtxCreate_v2(&mut mem::zeroed(), 0, 0) }.unwrap();
     let reader = std::fs::File::open(&args[2]).unwrap();
     let (mut manifest, mut source, mut buffers) = zluda_trace_common::replay::load(reader);
+    let mut outputs = FxHashMap::default();
     let mut args = manifest
         .parameters
         .iter()
@@ -53,6 +55,10 @@ fn main() {
                     ))
                     .unwrap();
                 unsafe { libcuda.cuMemAlloc_v2(&mut dev_ptr, host_buffer.len()) }.unwrap();
+                outputs.insert(
+                    format!("param_{i}_ptr_{}_post.bin", param_ptr.offset_in_param),
+                    (dev_ptr, host_buffer.len()),
+                );
                 unsafe {
                     libcuda.cuMemcpyHtoD_v2(dev_ptr, host_buffer.as_ptr().cast(), host_buffer.len())
                 }
@@ -100,4 +106,9 @@ fn main() {
     }
     .unwrap();
     unsafe { libcuda.cuCtxSynchronize() }.unwrap();
+    for (path, (dev_ptr, len)) in outputs {
+        let mut host_buffer = vec![0u8; len];
+        unsafe { libcuda.cuMemcpyDtoH_v2(host_buffer.as_mut_ptr().cast(), dev_ptr, len) }.unwrap();
+        std::fs::write(path, host_buffer).unwrap();
+    }
 }