Skip to content

Commit 59efbfe

Browse files
skyler-ruiterskyler-ruiter
committed
updates graph mode tests to catch more explicitly that graph mode doesnt work on vgpu since the fallback doesnt get to use the stream ordered memory allocator
1 parent a0d69b2 commit 59efbfe

4 files changed

Lines changed: 212 additions & 0 deletions

File tree

‎src/pipeline/compressor.cpp‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -345,6 +345,12 @@ void Pipeline::finalize() {
345345
"Graph mode requires PREALLOCATE memory strategy. "
346346
"Call setMemoryStrategy(MemoryStrategy::PREALLOCATE) before finalize().");
347347
}
348+
if (mem_pool_ && mem_pool_->isFallbackMode()) {
349+
throw std::runtime_error(
350+
"Graph mode is not supported when the CUDA memory pool is unavailable "
351+
"(cudaMalloc fallback mode; common on vGPU or when FZ_FORCE_MEMPOOL_FALLBACK is set). "
352+
"Disable graph mode or use a GPU that supports cudaMemPool.");
353+
}
348354
if (input_size_hint_ == 0) {
349355
throw std::runtime_error(
350356
"Graph mode requires a non-zero input size hint. "

‎tests/pipeline/test_buffer_coloring.cpp‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -265,6 +265,12 @@ TEST(BufferColoring, RepeatedCompressStaysCorrect) {
265265
// (PREALLOCATE + coloring enabled + graph mode; 5 compress replays)
266266
// ─────────────────────────────────────────────────────────────────────────────
267267
TEST(BufferColoring, ColoredBuffersStableAcrossGraphReplays) {
268+
{
269+
Pipeline probe(kNBytes, MemoryStrategy::PREALLOCATE, 1.0f);
270+
if (probe.isMemPoolFallbackMode()) {
271+
GTEST_SKIP() << "Graph mode unsupported in cudaMalloc fallback mode";
272+
}
273+
}
268274
auto h_input = make_test_data();
269275

270276
CudaStream stream;

‎tests/pipeline/test_graph_capture.cpp‎

Lines changed: 193 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -15,13 +15,20 @@
1515
* GC9 - isGraphCompatible() false for RZE inverse, true for all forward stages
1616
* GC10 - graph mode with coloring disabled still produces correct output
1717
* GC11 - warmup() before captureGraph() does not break subsequent graph compress
18+
* GC12 - graph mode with mempool fallback throws at finalize()
19+
* GC13 - captureGraph() without enableGraphMode() throws
20+
* GC14 - captureGraph() after compress() throws
21+
* GC15 - enableGraphMode() after finalize() throws
22+
* GC16 - graph capture works with a Lorenzo -> RLE pipeline
1823
*/
1924

2025
#include <gtest/gtest.h>
2126
#include "helpers/fz_test_utils.h"
2227
#include "fzgpumodules.h"
2328

2429
#include <cmath>
30+
#include <cstdlib>
31+
#include <string>
2532
#include <vector>
2633

2734
using namespace fz;
@@ -39,6 +46,29 @@ static std::vector<float> make_smooth(size_t n) {
3946
return v;
4047
}
4148

49+
struct ScopedEnvVar {
50+
std::string name;
51+
bool had_old = false;
52+
std::string old_value;
53+
54+
ScopedEnvVar(const char* var_name, const char* value) : name(var_name) {
55+
const char* old = std::getenv(var_name);
56+
if (old) {
57+
had_old = true;
58+
old_value = old;
59+
}
60+
setenv(var_name, value, /*overwrite=*/1);
61+
}
62+
63+
~ScopedEnvVar() {
64+
if (had_old) {
65+
setenv(name.c_str(), old_value.c_str(), /*overwrite=*/1);
66+
} else {
67+
unsetenv(name.c_str());
68+
}
69+
}
70+
};
71+
4272
// Build a finalized graph-mode pipeline with a single Lorenzo stage.
4373
// Caller owns the pipeline; the graph stream must be non-default.
4474
static std::unique_ptr<Pipeline> make_graph_pipeline(size_t in_bytes, bool disable_coloring = false) {
@@ -54,10 +84,34 @@ static std::unique_ptr<Pipeline> make_graph_pipeline(size_t in_bytes, bool disab
5484
return p;
5585
}
5686

87+
static bool is_graph_supported() {
88+
Pipeline p(1024, MemoryStrategy::PREALLOCATE, 1.0f);
89+
return !p.isMemPoolFallbackMode();
90+
}
91+
92+
static std::unique_ptr<Pipeline> make_graph_pipeline_with_rle(size_t in_bytes) {
93+
auto p = std::make_unique<Pipeline>(in_bytes, MemoryStrategy::PREALLOCATE, 4.0f);
94+
auto* lrz = p->addStage<LorenzoQuantStage<float, uint16_t>>();
95+
lrz->setErrorBound(1e-2f);
96+
lrz->setQuantRadius(512);
97+
lrz->setOutlierCapacity(0.2f);
98+
99+
auto* rle = p->addStage<RLEStage<uint16_t>>();
100+
p->connect(rle, lrz, "codes");
101+
102+
p->enableGraphMode(true);
103+
p->setPoolManagedDecompOutput(false);
104+
p->finalize();
105+
return p;
106+
}
107+
57108
// ─────────────────────────────────────────────────────────────────────────────
58109
// GC1: captureGraph() produces correct compressed + decompressed output
59110
// ─────────────────────────────────────────────────────────────────────────────
60111
TEST(GraphCapture, CorrectOutput) {
112+
if (!is_graph_supported()) {
113+
GTEST_SKIP() << "Graph mode unsupported in cudaMalloc fallback mode";
114+
}
61115
constexpr size_t N = 1 << 14;
62116
constexpr float EB = 1e-2f;
63117
const size_t in_bytes = N * sizeof(float);
@@ -105,6 +159,9 @@ TEST(GraphCapture, CorrectOutput) {
105159
// semantically meaningful invariant is data correctness, not byte identity.
106160
// ─────────────────────────────────────────────────────────────────────────────
107161
TEST(GraphCapture, ReplayCorrect) {
162+
if (!is_graph_supported()) {
163+
GTEST_SKIP() << "Graph mode unsupported in cudaMalloc fallback mode";
164+
}
108165
constexpr size_t N = 1 << 14;
109166
constexpr float EB = 1e-2f;
110167
constexpr int ITERS = 12;
@@ -151,6 +208,9 @@ TEST(GraphCapture, ReplayCorrect) {
151208
// filler may differ. The meaningful invariant is data correctness.
152209
// ─────────────────────────────────────────────────────────────────────────────
153210
TEST(GraphCapture, MatchesPreallocate) {
211+
if (!is_graph_supported()) {
212+
GTEST_SKIP() << "Graph mode unsupported in cudaMalloc fallback mode";
213+
}
154214
constexpr size_t N = 1 << 14;
155215
constexpr float EB = 1e-2f;
156216
const size_t in_bytes = N * sizeof(float);
@@ -268,6 +328,9 @@ TEST(GraphCapture, CaptureBeforeFinalizeThrows) {
268328
// kernels. The pipeline must handle this gracefully and produce valid output.
269329
// ─────────────────────────────────────────────────────────────────────────────
270330
TEST(GraphCapture, CompressBeforeCaptureSucceeds) {
331+
if (!is_graph_supported()) {
332+
GTEST_SKIP() << "Graph mode unsupported in cudaMalloc fallback mode";
333+
}
271334
constexpr size_t N = 1 << 13;
272335
constexpr float EB = 1e-2f;
273336
const size_t in_bytes = N * sizeof(float);
@@ -380,6 +443,9 @@ TEST(GraphCapture, IsGraphCompatiblePerStage) {
380443
// GC10: Graph mode with buffer coloring disabled still produces correct output
381444
// ─────────────────────────────────────────────────────────────────────────────
382445
TEST(GraphCapture, ColoringDisabledCorrectOutput) {
446+
if (!is_graph_supported()) {
447+
GTEST_SKIP() << "Graph mode unsupported in cudaMalloc fallback mode";
448+
}
383449
constexpr size_t N = 1 << 13;
384450
constexpr float EB = 1e-2f;
385451
const size_t in_bytes = N * sizeof(float);
@@ -426,6 +492,9 @@ TEST(GraphCapture, ColoringDisabledCorrectOutput) {
426492
// still succeed and compress must produce correct output.
427493
// ─────────────────────────────────────────────────────────────────────────────
428494
TEST(GraphCapture, WarmupBeforeCaptureWorks) {
495+
if (!is_graph_supported()) {
496+
GTEST_SKIP() << "Graph mode unsupported in cudaMalloc fallback mode";
497+
}
429498
constexpr size_t N = 1 << 14;
430499
constexpr float EB = 1e-2f;
431500
const size_t in_bytes = N * sizeof(float);
@@ -466,3 +535,127 @@ TEST(GraphCapture, WarmupBeforeCaptureWorks) {
466535
EXPECT_LE(max_abs_error(h_input, h_recon), EB * 1.01f)
467536
<< "GC11: warmup→capture round-trip error exceeds bound";
468537
}
538+
539+
// ─────────────────────────────────────────────────────────────────────────────
540+
// GC12: Graph mode must throw when mempool fallback is active (vGPU / forced)
541+
// ─────────────────────────────────────────────────────────────────────────────
542+
TEST(GraphCapture, MempoolFallbackThrowsAtFinalize) {
543+
ScopedEnvVar force_fallback("FZ_FORCE_MEMPOOL_FALLBACK", "1");
544+
constexpr size_t N = 1024;
545+
const size_t in_bytes = N * sizeof(float);
546+
547+
Pipeline p(in_bytes, MemoryStrategy::PREALLOCATE, 4.0f);
548+
auto* lrz = p.addStage<LorenzoQuantStage<float, uint16_t>>();
549+
lrz->setErrorBound(1e-2f);
550+
p.enableGraphMode(true);
551+
552+
EXPECT_THROW(p.finalize(), std::runtime_error)
553+
<< "GC12: graph mode must be rejected when mempool fallback is active";
554+
}
555+
556+
// ─────────────────────────────────────────────────────────────────────────────
557+
// GC13: captureGraph() without enableGraphMode() throws
558+
// ─────────────────────────────────────────────────────────────────────────────
559+
TEST(GraphCapture, CaptureWithoutEnableThrows) {
560+
constexpr size_t N = 1024;
561+
const size_t in_bytes = N * sizeof(float);
562+
563+
Pipeline p(in_bytes, MemoryStrategy::PREALLOCATE, 4.0f);
564+
auto* lrz = p.addStage<LorenzoQuantStage<float, uint16_t>>();
565+
lrz->setErrorBound(1e-2f);
566+
p.finalize();
567+
568+
CudaStream stream;
569+
EXPECT_THROW(p.captureGraph(stream), std::runtime_error)
570+
<< "GC13: captureGraph() without enableGraphMode() must throw";
571+
}
572+
573+
// ─────────────────────────────────────────────────────────────────────────────
574+
// GC14: captureGraph() after compress() throws (must capture before first run)
575+
// ─────────────────────────────────────────────────────────────────────────────
576+
TEST(GraphCapture, CaptureAfterCompressThrows) {
577+
if (!is_graph_supported()) {
578+
GTEST_SKIP() << "Graph mode unsupported in cudaMalloc fallback mode";
579+
}
580+
constexpr size_t N = 1 << 12;
581+
const size_t in_bytes = N * sizeof(float);
582+
583+
auto h_input = make_smooth(N);
584+
CudaStream stream;
585+
CudaBuffer<float> d_in(N);
586+
d_in.upload(h_input, stream);
587+
stream.sync();
588+
589+
auto p = make_graph_pipeline(in_bytes);
590+
591+
void* d_comp = nullptr;
592+
size_t comp_sz = 0;
593+
p->compress(d_in.void_ptr(), in_bytes, &d_comp, &comp_sz, stream);
594+
stream.sync();
595+
ASSERT_GT(comp_sz, 0u);
596+
597+
EXPECT_THROW(p->captureGraph(stream), std::runtime_error)
598+
<< "GC14: captureGraph() after compress() must throw";
599+
}
600+
601+
// ─────────────────────────────────────────────────────────────────────────────
602+
// GC15: enableGraphMode() after finalize() throws
603+
// ─────────────────────────────────────────────────────────────────────────────
604+
TEST(GraphCapture, EnableAfterFinalizeThrows) {
605+
constexpr size_t N = 1024;
606+
Pipeline p(N * sizeof(float), MemoryStrategy::PREALLOCATE, 4.0f);
607+
auto* lrz = p.addStage<LorenzoQuantStage<float, uint16_t>>();
608+
lrz->setErrorBound(1e-2f);
609+
p.finalize();
610+
611+
EXPECT_THROW(p.enableGraphMode(true), std::runtime_error)
612+
<< "GC15: enableGraphMode() after finalize() must throw";
613+
}
614+
615+
// ─────────────────────────────────────────────────────────────────────────────
616+
// GC16: Graph capture works for a Lorenzo → RLE pipeline
617+
// ─────────────────────────────────────────────────────────────────────────────
618+
TEST(GraphCapture, RlePipelineCorrectOutput) {
619+
if (!is_graph_supported()) {
620+
GTEST_SKIP() << "Graph mode unsupported in cudaMalloc fallback mode";
621+
}
622+
constexpr size_t N = 1 << 14;
623+
constexpr float EB = 1e-2f;
624+
const size_t in_bytes = N * sizeof(float);
625+
626+
auto h_input = make_smooth(N);
627+
CudaStream stream;
628+
CudaBuffer<float> d_in(N);
629+
d_in.upload(h_input, stream);
630+
stream.sync();
631+
632+
auto p = make_graph_pipeline_with_rle(in_bytes);
633+
634+
// Warmup to pre-allocate stage scratch and host buffers outside the capture window.
635+
p->warmup(stream);
636+
stream.sync();
637+
638+
p->captureGraph(stream);
639+
ASSERT_TRUE(p->isGraphCaptured());
640+
641+
void* d_comp = nullptr;
642+
size_t comp_sz = 0;
643+
p->compress(d_in.void_ptr(), in_bytes, &d_comp, &comp_sz, stream);
644+
stream.sync();
645+
ASSERT_GT(comp_sz, 0u);
646+
647+
void* d_dec = nullptr;
648+
size_t dec_sz = 0;
649+
p->decompress(nullptr, 0, &d_dec, &dec_sz, stream);
650+
stream.sync();
651+
ASSERT_NE(d_dec, nullptr);
652+
ASSERT_EQ(dec_sz, in_bytes);
653+
654+
std::vector<float> h_recon(N);
655+
cudaMemcpy(h_recon.data(), d_dec, in_bytes, cudaMemcpyDeviceToHost);
656+
cudaFree(d_dec);
657+
658+
float max_err = max_abs_error(h_input, h_recon);
659+
EXPECT_LE(max_err, EB * 1.01f)
660+
<< "GC16: Lorenzo->RLE graph round-trip max_err=" << max_err;
661+
}

‎tests/pipeline/test_liveness.cpp‎

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -259,6 +259,13 @@ TEST(Liveness, LivenessApiConsistentUnderGraphMode) {
259259
constexpr size_t N = 1 << 12;
260260
const size_t in_bytes = N * sizeof(float);
261261

262+
{
263+
Pipeline probe(in_bytes, MemoryStrategy::PREALLOCATE, 1.0f);
264+
if (probe.isMemPoolFallbackMode()) {
265+
GTEST_SKIP() << "Graph mode unsupported in cudaMalloc fallback mode";
266+
}
267+
}
268+
262269
// Build a 2-stage PREALLOCATE + graph-mode pipeline.
263270
auto p = std::make_unique<Pipeline>(in_bytes, MemoryStrategy::PREALLOCATE, 5.0f);
264271

0 commit comments

Comments
 (0)