Skip to content

Commit 2be3fe2

Browse files
sjain-stanfordcodex
sjain-stanford
and
codex
authored
[1/N Dynamic Shapes]: Defer workspace size queries until execute for access to runtime input shapes (#393)
For the static case, workspace size was a constant and independent of the runtime input sizes so it was OK to query and cache it during VM context creation. However for the dynamic case it is dependent on the runtime input shapes. This PR moves workspace-size querying out of VM context creation and into the explicit `getWorkspaceSize()` call. This keeps compile/load from querying transient size and makes `execute()` require callers to query workspace size before allocating and passing a workspace buffer. For this PR, `getWorkspaceSize()` remains independent of runtime inputs: it queries the currently supported constant transient size, stores the maximum queried size on the graph, and returns `std::nullopt` when no runtime artifact is loaded. Variant-pack/handle-dependent dynamic workspace sizing is intentionally deferred until the dynamic size function path is added. Updates samples, benchmarks, and tests to call `getWorkspaceSize()` before `allocateWorkspace(...)` and then pass the resulting workspace to `execute()`. Also covers the AOT compile/load path. --------- Signed-off-by: Sambhav Jain <sambhav@alumni.stanford.edu> Co-authored-by: GPT 5.5 <codex@openai.com>
1 parent dc198af commit 2be3fe2

51 files changed

Lines changed: 319 additions & 209 deletions

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

‎benchmarks/driver.cpp‎

Lines changed: 12 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -235,8 +235,9 @@ static ErrorObject benchmarkConvFprop(const ConvOptions &opts,
235235
}
236236

237237
// Allocate workspace buffer if needed.
238+
FUSILLI_ASSIGN_OR_RETURN(auto workspaceSize, graph.getWorkspaceSize());
238239
FUSILLI_ASSIGN_OR_RETURN(auto workspace,
239-
allocateWorkspace(handle, graph.getWorkspaceSize()));
240+
allocateWorkspace(handle, workspaceSize));
240241

241242
// Execute graph a few times.
242243
for (size_t i = 0; i < iter; ++i)
@@ -366,8 +367,9 @@ static ErrorObject benchmarkConvWGrad(const ConvOptions &opts,
366367
}
367368

368369
// Allocate workspace buffer if needed.
370+
FUSILLI_ASSIGN_OR_RETURN(auto workspaceSize, graph.getWorkspaceSize());
369371
FUSILLI_ASSIGN_OR_RETURN(auto workspace,
370-
allocateWorkspace(handle, graph.getWorkspaceSize()));
372+
allocateWorkspace(handle, workspaceSize));
371373

372374
// Execute graph a few times.
373375
for (size_t i = 0; i < iter; ++i)
@@ -497,8 +499,9 @@ static ErrorObject benchmarkConvDGrad(const ConvOptions &opts,
497499
}
498500

499501
// Allocate workspace buffer if needed.
502+
FUSILLI_ASSIGN_OR_RETURN(auto workspaceSize, graph.getWorkspaceSize());
500503
FUSILLI_ASSIGN_OR_RETURN(auto workspace,
501-
allocateWorkspace(handle, graph.getWorkspaceSize()));
504+
allocateWorkspace(handle, workspaceSize));
502505

503506
// Execute graph a few times.
504507
for (size_t i = 0; i < iter; ++i)
@@ -593,8 +596,9 @@ static ErrorObject benchmarkLayerNormFwd(const LayerNormOptions &opts,
593596
}
594597

595598
// Allocate workspace buffer if needed.
599+
FUSILLI_ASSIGN_OR_RETURN(auto workspaceSize, graph.getWorkspaceSize());
596600
FUSILLI_ASSIGN_OR_RETURN(auto workspace,
597-
allocateWorkspace(handle, graph.getWorkspaceSize()));
601+
allocateWorkspace(handle, workspaceSize));
598602

599603
// Execute graph `iter` times.
600604
for (size_t i = 0; i < iter; ++i)
@@ -722,8 +726,9 @@ static ErrorObject benchmarkSdpaFwd(const SdpaOptions &opts,
722726
}
723727

724728
// Allocate workspace buffer if needed.
729+
FUSILLI_ASSIGN_OR_RETURN(auto workspaceSize, graph.getWorkspaceSize());
725730
FUSILLI_ASSIGN_OR_RETURN(auto workspace,
726-
allocateWorkspace(handle, graph.getWorkspaceSize()));
731+
allocateWorkspace(handle, workspaceSize));
727732

728733
// Execute graph `iter` times.
729734
for (size_t i = 0; i < iter; ++i)
@@ -851,8 +856,9 @@ static ErrorObject benchmarkMatmul(const MatmulOptions &opts, DataType aType,
851856
}
852857

853858
// Allocate workspace buffer if needed.
859+
FUSILLI_ASSIGN_OR_RETURN(auto workspaceSize, graph.getWorkspaceSize());
854860
FUSILLI_ASSIGN_OR_RETURN(auto workspace,
855-
allocateWorkspace(handle, graph.getWorkspaceSize()));
861+
allocateWorkspace(handle, workspaceSize));
856862

857863
// Execute graph `iter` times.
858864
for (size_t i = 0; i < iter; ++i)

‎include/fusilli/backend/runtime.h‎

Lines changed: 17 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -311,11 +311,19 @@ inline ErrorObject Graph::createVmContext(const Handle &handle) {
311311
if (executeAsync)
312312
vmInputListCapacity_ += 2;
313313

314-
// Query the required workspace size from the compiled module.
314+
return ok();
315+
}
316+
317+
inline ErrorOr<std::optional<size_t>> Graph::getWorkspaceSize() {
318+
if (vmContext_ == nullptr)
319+
return ok(std::optional<size_t>());
320+
315321
FUSILLI_LOG_LABEL_ENDL("INFO: Querying workspace size from compiled module");
316-
FUSILLI_ASSIGN_OR_RETURN(workspaceSize_, queryTransientSize());
322+
FUSILLI_ASSIGN_OR_RETURN(size_t workspaceSize, queryTransientSize());
323+
if (!workspaceSize_.has_value() || workspaceSize > *workspaceSize_)
324+
workspaceSize_ = workspaceSize;
317325

318-
return ok();
326+
return ok(workspaceSize_);
319327
}
320328

321329
// Queries the required transient/workspace buffer size from the compiled
@@ -324,7 +332,7 @@ inline ErrorObject Graph::createVmContext(const Handle &handle) {
324332
// for the constant workspace size case, or an "iree.abi.transients.size"
325333
// function for the data-dependent workspace size case. Only the former is
326334
// supported by Fusilli at the moment.
327-
inline ErrorOr<size_t> Graph::queryTransientSize() {
335+
inline ErrorOr<size_t> Graph::queryTransientSize() const {
328336
// Always resolve the async function for attribute queries. The
329337
// iree.abi.transients.size.constant attribute is stored in the
330338
// iree.reflection dict on the @main$async entry point. The sync wrapper
@@ -394,6 +402,10 @@ Graph::execute(const Handle &handle,
394402
", but the loaded artifact uses backend " +
395403
kBackendToStr.at(*loadedBackend_));
396404
bool executeAsync = kBackendExecuteAsync.at(handle.getBackend());
405+
FUSILLI_RETURN_ERROR_IF(
406+
!workspaceSize_.has_value(), ErrorCode::InvalidArgument,
407+
"Graph::execute requires getWorkspaceSize() to be called before "
408+
"workspace allocation and execution");
397409

398410
iree_allocator_t allocator = iree_allocator_system();
399411

@@ -448,7 +460,7 @@ Graph::execute(const Handle &handle,
448460
// adds a !hal.buffer argument to the generated function signature, even when
449461
// no transient storage is needed (size = 0). We must always push a buffer
450462
// (or null ref when size = 0) to satisfy the function signature.
451-
if (workspaceSize_.value_or(0) > 0) {
463+
if (*workspaceSize_ > 0) {
452464
FUSILLI_RETURN_ERROR_IF(
453465
workspace == nullptr, ErrorCode::InvalidArgument,
454466
"Workspace buffer required but not provided (size=" +

‎include/fusilli/graph/graph.h‎

Lines changed: 24 additions & 21 deletions
Original file line numberDiff line numberDiff line change
@@ -166,19 +166,19 @@ class Graph : public INode {
166166
}
167167

168168
// Loads a compiled VMFB artifact onto the device owned by `handle`, creating
169-
// per-graph runtime state and querying workspace size. This does not require
170-
// the artifact to have been produced by this `Graph` instance.
169+
// per-graph runtime state. This does not require the artifact to have been
170+
// produced by this `Graph` instance.
171171
//
172172
// The artifact must be compatible with `handle.getBackend()`. A VMFB compiled
173173
// for one backend is not a portable artifact for another backend.
174174
//
175175
// A `Graph` owns one loaded runtime state at a time. Loading a new artifact
176-
// replaces any previously loaded VM context, function, workspace size, and VM
177-
// input list capacity. If loading fails, the `Graph` is left without loaded
178-
// runtime state. To keep multiple backends loaded concurrently, use one
179-
// `Graph` instance per backend artifact. To switch backends on one `Graph`,
180-
// call `loadFromArtifact()` with the selected backend artifact before
181-
// `execute()`.
176+
// replaces any previously loaded VM context, function, workspace query cache,
177+
// and VM input list capacity. If loading fails, the `Graph` is left without
178+
// loaded runtime state. To keep multiple backends loaded concurrently, use
179+
// one `Graph` instance per backend artifact. To switch backends on one
180+
// `Graph`, call `loadFromArtifact()` with the selected backend artifact
181+
// before `execute()`.
182182
ErrorObject loadFromArtifact(const Handle &handle,
183183
std::span<const uint8_t> vmfbBytes) {
184184
FUSILLI_LOG_LABEL_ENDL("INFO: Loading compiled artifact into VM context");
@@ -279,14 +279,16 @@ class Graph : public INode {
279279
// doSomethingWith(hostData);
280280
//
281281
// Workspace Buffer Usage:
282-
// After calling compile(), query getWorkspaceSize() to determine if a
283-
// workspace buffer is needed. If size > 0, allocate using
284-
// Buffer::allocateRaw() and pass it to execute(). The same workspace
285-
// buffer can be reused across multiple execute() calls.
282+
// After calling compile(), query getWorkspaceSize() to determine
283+
// if a workspace buffer is needed. If size > 0, allocate using
284+
// Buffer::allocateRaw() and pass it to execute(). Calling
285+
// getWorkspaceSize() before execute() is required, and the same
286+
// workspace buffer can be reused across multiple execute() calls.
286287
//
287288
// Example:
288289
// graph.compile(handle);
289-
// auto wsSize = graph.getWorkspaceSize();
290+
// FUSILLI_ASSIGN_OR_RETURN(auto wsSize,
291+
// graph.getWorkspaceSize());
290292
// std::shared_ptr<Buffer> workspace = nullptr;
291293
// if (wsSize.value_or(0) > 0) {
292294
// FUSILLI_ASSIGN_OR_RETURN(auto wsBuf,
@@ -388,10 +390,11 @@ class Graph : public INode {
388390
customOp(std::vector<std::shared_ptr<TensorAttr>> inputs,
389391
CustomOpAttr &customOpAttr);
390392

391-
// Query required workspace buffer size.
392-
// Returns std::nullopt if not compiled, 0 if no workspace needed,
393-
// or the required size in bytes.
394-
std::optional<size_t> getWorkspaceSize() const { return workspaceSize_; }
393+
// Query required workspace buffer size. Returns std::nullopt if no runtime
394+
// artifact is loaded, 0 if no workspace is needed, or the maximum required
395+
// size in bytes seen for the currently loaded runtime state. Dynamic
396+
// workspace sizes are not supported yet and return an error.
397+
ErrorOr<std::optional<size_t>> getWorkspaceSize();
395398

396399
// ASM emitter driver method.
397400
//
@@ -477,7 +480,7 @@ class Graph : public INode {
477480
// module. Returns the size in bytes, or 0 if no transients are needed.
478481
// Returns an error if the module requires dynamic transient sizes.
479482
// Definition in `fusilli/backend/runtime.h`.
480-
ErrorOr<size_t> queryTransientSize();
483+
ErrorOr<size_t> queryTransientSize() const;
481484

482485
// Create compiled artifacts from graph writing results to the cache. Set
483486
// `remove = true` to remove cache files when returned `CachedAssets` lifetime
@@ -708,9 +711,9 @@ class Graph : public INode {
708711
// Avoids repeated function lookup on every execute() call.
709712
std::optional<iree_vm_function_t> vmFunction_;
710713

711-
// Required workspace buffer size in bytes. Set during createVmContext()
712-
// by querying the iree.abi.transients.size.constant attribute.
713-
// std::nullopt indicates the graph has not been compiled yet.
714+
// Maximum required workspace buffer size in bytes seen by getWorkspaceSize().
715+
// std::nullopt indicates the workspace size has not been queried for the
716+
// currently loaded runtime state.
714717
std::optional<size_t> workspaceSize_;
715718

716719
// Pre-computed VM input list capacity for iree_vm_list_create().

‎samples/aot/multi_backend.cpp‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -54,7 +54,6 @@ std::vector<uint8_t> compileArtifact(Backend backend) {
5454
compileGraph.graph->compileToArtifact(backend, /*remove=*/true));
5555

5656
REQUIRE(!compiledArtifactBytes.empty());
57-
REQUIRE(!compileGraph.graph->getWorkspaceSize().has_value());
5857
return compiledArtifactBytes;
5958
}
6059

@@ -66,8 +65,6 @@ void loadAndExecuteArtifact(Backend backend,
6665
auto runtimeGraph = buildPointwiseAddGraph(graphName);
6766
FUSILLI_REQUIRE_OK(runtimeGraph.graph->loadFromArtifact(handle, vmfbBytes));
6867

69-
REQUIRE(runtimeGraph.graph->getWorkspaceSize().has_value());
70-
7168
FUSILLI_REQUIRE_ASSIGN(
7269
auto x0Buf,
7370
allocateBufferOfType(handle, runtimeGraph.x0, DataType::Int32, 2));
@@ -84,9 +81,12 @@ void loadAndExecuteArtifact(Backend backend,
8481
{runtimeGraph.y, yBuf},
8582
};
8683

87-
FUSILLI_REQUIRE_ASSIGN(
88-
auto workspace,
89-
allocateWorkspace(handle, runtimeGraph.graph->getWorkspaceSize()));
84+
FUSILLI_REQUIRE_ASSIGN(auto workspaceSize,
85+
runtimeGraph.graph->getWorkspaceSize());
86+
REQUIRE(workspaceSize.has_value());
87+
88+
FUSILLI_REQUIRE_ASSIGN(auto workspace,
89+
allocateWorkspace(handle, workspaceSize));
9090

9191
FUSILLI_REQUIRE_OK(
9292
runtimeGraph.graph->execute(handle, variantPack, workspace));

‎samples/aot/single_backend.cpp‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -64,7 +64,6 @@ TEST_CASE("AOT single-backend artifact compile/load/execute round trip",
6464
compileGraph.graph->compileToArtifact(kDefaultBackend,
6565
/*remove=*/true));
6666
REQUIRE(!callerOwnedArtifactBytes.empty());
67-
REQUIRE(!compileGraph.graph->getWorkspaceSize().has_value());
6867
}
6968

7069
// Caller code may serialize and de-serialize the compiled artifacts here.
@@ -78,8 +77,6 @@ TEST_CASE("AOT single-backend artifact compile/load/execute round trip",
7877
FUSILLI_REQUIRE_OK(
7978
runtimeGraph.graph->loadFromArtifact(handle, callerOwnedArtifactBytes));
8079

81-
REQUIRE(runtimeGraph.graph->getWorkspaceSize().has_value());
82-
8380
FUSILLI_REQUIRE_ASSIGN(
8481
auto x0Buf,
8582
allocateBufferOfType(handle, runtimeGraph.x0, DataType::Int32, 2));
@@ -96,9 +93,12 @@ TEST_CASE("AOT single-backend artifact compile/load/execute round trip",
9693
{runtimeGraph.y, yBuf},
9794
};
9895

99-
FUSILLI_REQUIRE_ASSIGN(
100-
auto workspace,
101-
allocateWorkspace(handle, runtimeGraph.graph->getWorkspaceSize()));
96+
FUSILLI_REQUIRE_ASSIGN(auto workspaceSize,
97+
runtimeGraph.graph->getWorkspaceSize());
98+
REQUIRE(workspaceSize.has_value());
99+
100+
FUSILLI_REQUIRE_ASSIGN(auto workspace,
101+
allocateWorkspace(handle, workspaceSize));
102102

103103
FUSILLI_REQUIRE_OK(
104104
runtimeGraph.graph->execute(handle, variantPack, workspace));

‎samples/batchnorm/batchnorm_infer_nchw.cpp‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -86,8 +86,9 @@ TEST_CASE("Batch normalization; inference mode; NCHW layout; no scale/bias",
8686
{yT, yBuf},
8787
};
8888

89+
FUSILLI_REQUIRE_ASSIGN(auto workspaceSize, graph->getWorkspaceSize());
8990
FUSILLI_REQUIRE_ASSIGN(auto workspace,
90-
allocateWorkspace(handle, graph->getWorkspaceSize()));
91+
allocateWorkspace(handle, workspaceSize));
9192

9293
FUSILLI_REQUIRE_OK(graph->execute(handle, variantPack, workspace));
9394

‎samples/batchnorm/batchnorm_infer_nchw_scale_bias.cpp‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -91,8 +91,9 @@ TEST_CASE("Batch normalization; inference mode; NCHW layout; scale, bias",
9191
{meanT, meanBuf}, {varT, varBuf}, {yT, yBuf},
9292
};
9393

94+
FUSILLI_REQUIRE_ASSIGN(auto workspaceSize, graph->getWorkspaceSize());
9495
FUSILLI_REQUIRE_ASSIGN(auto workspace,
95-
allocateWorkspace(handle, graph->getWorkspaceSize()));
96+
allocateWorkspace(handle, workspaceSize));
9697

9798
FUSILLI_REQUIRE_OK(graph->execute(handle, variantPack, workspace));
9899

‎samples/batchnorm/batchnorm_infer_nhwc_scale_bias.cpp‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -92,8 +92,9 @@ TEST_CASE("Batch normalization; inference mode; NHWC layout; scale, bias",
9292
{meanT, meanBuf}, {varT, varBuf}, {yT, yBuf},
9393
};
9494

95+
FUSILLI_REQUIRE_ASSIGN(auto workspaceSize, graph->getWorkspaceSize());
9596
FUSILLI_REQUIRE_ASSIGN(auto workspace,
96-
allocateWorkspace(handle, graph->getWorkspaceSize()));
97+
allocateWorkspace(handle, workspaceSize));
9798

9899
FUSILLI_REQUIRE_OK(graph->execute(handle, variantPack, workspace));
99100

‎samples/batchnorm/batchnorm_train_nchw.cpp‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -85,8 +85,9 @@ TEST_CASE("Batch normalization; training mode; NCHW layout; no scale/bias",
8585
{sivT, sivBuf},
8686
};
8787

88+
FUSILLI_REQUIRE_ASSIGN(auto workspaceSize, graph->getWorkspaceSize());
8889
FUSILLI_REQUIRE_ASSIGN(auto workspace,
89-
allocateWorkspace(handle, graph->getWorkspaceSize()));
90+
allocateWorkspace(handle, workspaceSize));
9091

9192
FUSILLI_REQUIRE_OK(graph->execute(handle, variantPack, workspace));
9293

‎samples/batchnorm/batchnorm_train_nchw_scale_bias.cpp‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -92,8 +92,9 @@ TEST_CASE("Batch normalization; training mode; NCHW layout; scale, bias",
9292
{yT, yBuf}, {smT, smBuf}, {sivT, sivBuf},
9393
};
9494

95+
FUSILLI_REQUIRE_ASSIGN(auto workspaceSize, graph->getWorkspaceSize());
9596
FUSILLI_REQUIRE_ASSIGN(auto workspace,
96-
allocateWorkspace(handle, graph->getWorkspaceSize()));
97+
allocateWorkspace(handle, workspaceSize));
9798

9899
FUSILLI_REQUIRE_OK(graph->execute(handle, variantPack, workspace));
99100

0 commit comments

Comments
 (0)