Skip to content

test(tensilelite): fix stale Stream-K parallel workspace-starved expectations #11635

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
jaopaulolc merged 1 commit into from
Sep 3, 2026
Merged

test(tensilelite): fix stale Stream-K parallel workspace-starved expectations #11635

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 14 additions & 16 deletions projects/hipblaslt/tensilelite/include/Tensile/ContractionSolution.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -475,25 +475,23 @@ namespace TensileLite
// reduction it sizes with requiredWorkspaceSizeGsu(problem, hardware,
// grid / tiles) instead of partialTileSize(grid).
//
// The two can therefore disagree about WHETHER a workspace is needed, not
// just about how many bytes. The case where they do is parallel reduction
// whose selected grid comes back equal to tiles: the k-split factor
// grid / tiles is then 1, and requiredWorkspaceSizeGsu() short-circuits to 0
// for a split of 1, while this field still reserves partialTileSize(tiles)
// because parallel reduction always needs a partials region. Concretely, a
// 256x4096x4096 SK3 launch with a 128x128x64 macro tile on a 256-CU device,
// handed a workspace between 4 MiB and 16 MiB, reports 4 MiB here and 0 from
// requiredWorkspaceSize().
// The two could in principle disagree about WHETHER a workspace is needed,
// not just about how many bytes: at a k-split factor grid / tiles of 1,
// requiredWorkspaceSizeGsu() short-circuits to 0 while partialTileSize(grid)
// does not, so a parallel reduction whose grid came back equal to tiles
// would reserve here and not there. That case is unreachable -- both call
// sites run streamKReconcileReduction() on the same (reduction, grid, tiles)
// triple immediately after getSKGridImpl(), and it demotes parallel to tree
// whenever the split factor is below 2, so neither sizing ever sees parallel
// at a split of 1. The formulas differ; the reserve-or-not answer does not.
//
// That divergence is an over-reservation of a buffer the caller already
// supplied, never an under-allocation, so it is benign in the real
// allocate-then-launch flow: the allocator sizes from
// That agreement is load-bearing rather than incidental: it is what lets the
// allocate-then-launch flow close. The allocator sizes from
// requiredWorkspaceSize(), that size is what problem.workspaceSize() reports
// on the subsequent launch, and re-deriving this snapshot against it reaches
// a self-consistent fixed point (0 bytes -> workspace-DP fallback -> 0 bytes
// reserved). It only becomes visible when a caller hands in a workspace it
// sized some other way. Both implementations still encode the same intended
// rule -- change one, check the other.
// a self-consistent fixed point. Both implementations encode the same
// intended rule, by way of the same reconcile helper -- change one, check
// the other two.
size_t requiredWorkspaceBytes = 0;
// recomputed: partials(+work-queue) bytes wanted, before the fit check against
// givenWorkspaceBytes. Non-zero even when the fallback fires, which is what
Expand Down
143 changes: 116 additions & 27 deletions projects/hipblaslt/tensilelite/tests/StreamKLaunchSummary_test.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -934,8 +934,10 @@ TEST(StreamKLaunchSummaryTest, ClusterDpClampWinsAttributionOverFixedGrid)
// grid/tiles) -- two different formulas for the same region. They agree here
// because the parallel grid is an exact multiple of tiles (grid == tiles*skSplit
// == 64*4), which makes requiredWorkspaceSizeGsu's tiles*gsu equal
// partialTileSize's grid. They do NOT agree when the parallel grid comes back
// equal to tiles; see the comment on StreamKDecisions::requiredWorkspaceBytes.
// partialTileSize's grid. The one triple where the formulas would disagree --
// parallel at a split factor of 1 -- never reaches either of them, because
// streamKReconcileReduction() demotes it to tree first; see the comment on
// StreamKDecisions::requiredWorkspaceBytes.
// ---------------------------------------------------------------------------
TEST(StreamKLaunchSummaryTest, Sk3ParallelReductionReservesPartialsWorkspace)
{
Expand Down Expand Up @@ -1015,18 +1017,29 @@ TEST(StreamKLaunchSummaryTest, Sk3ParallelReductionReservesPartialsWorkspace)
}

// ---------------------------------------------------------------------------
// The parallel branch's workspace-DP fallback.
// A workspace-starved parallel scenario that never reaches the workspace-DP
// fallback at all, because grid selection already collapsed the grid.
//
// Same shape as the test above, but with no workspace at all. getSKReduction()
// still says parallel -- it never looks at the workspace -- while the snapshot
// demotes the launch to tree with grid = tiles because the partials region does
// not fit. This is the parallel-side complement of WorkspaceDpFallbackFires
// (which covers the tree / indivisible-tiles side), and it is also a second
// demonstration that changedBy is about ATTRIBUTION rather than about the numbers
// differing: with no workspace, origami's grid selection already returns tiles, so
// selectedGrid == finalGrid while the fallback is still correctly credited.
// still says parallel -- it never looks at the workspace -- but origami's
// k_split_aware grid selection does: with zero workspace no split factor F >= 2
// is admissible, so it returns grid == tiles (64) instead of 256. The snapshot
// then runs streamKReconcileReduction() on that (parallel, grid, tiles) triple
// BEFORE the workspace-fit guard, and F = grid / tiles = 1 is not expressible in
// parallel reduction, so the reduction is demoted to tree right there. The guard
// below it therefore sees tree with tiles % grid == 0, decides no partials are
// needed, and never runs its body: idealWorkspaceBytes stays 0 and
// workspaceDPFallbackFired stays false.
//
// So this is NOT a fallback -- it is the selected launch, and the summary
// correctly attributes the grid to nothing ("changedBy = none"). The launch
// itself (grid = tiles = 64, tree, zero workspace) is the same one the fallback
// would have produced; only the attribution differs, and it differs because
// nothing was clamped. The genuine parallel-side workspace-DP fallback is
// covered by the test below, which uses a fixed grid to get past origami's own
// workspace clamp.
// ---------------------------------------------------------------------------
TEST(StreamKLaunchSummaryTest, Sk3ParallelReductionWorkspaceStarvedFallsBackToTree)
TEST(StreamKLaunchSummaryTest, Sk3ParallelWorkspaceStarvedGridCollapsesToTilesAndReconcilesToTree)
{
AnalyticalEnv env;
ContractionSolution solution;
Expand All @@ -1039,43 +1052,119 @@ TEST(StreamKLaunchSummaryTest, Sk3ParallelReductionWorkspaceStarvedFallsBackToTr
ASSERT_FALSE(Debug::Instance().useStreamKDataParrallel())
<< "unset TENSILE_STREAMK_DATA_PARALLEL before running this suite";

// Anti-vacuity: the pre-fallback reduction really is parallel. (getSKReduction
// Anti-vacuity: the pre-reconcile reduction really is parallel. (getSKReduction
// is workspace-independent, so it reports what the snapshot started from.)
ASSERT_EQ(solution.getSKReduction(problem, env.device), origami::reduction_t::parallel)
<< "scenario must start from parallel reduction for the fallback to be the "
"parallel-path fallback";
<< "scenario must start from parallel reduction, otherwise the reconcile "
"step below is not what demoted it";

auto d = solution.computeStreamKDecisions(problem, env.device);

EXPECT_EQ(d.reduction, origami::reduction_t::tree) << "the fallback demotes to tree";
EXPECT_TRUE(d.workspaceDPFallbackFired);
EXPECT_TRUE(d.dpOnly);
EXPECT_FALSE(d.forceDPOnly) << "runtime fallback, not the compile-time param";
EXPECT_EQ(d.reduction, origami::reduction_t::tree)
<< "streamKReconcileReduction demotes parallel at a split factor of 1";
EXPECT_FALSE(d.workspaceDPFallbackFired)
<< "the reconcile ran first, so the workspace guard saw tree and divisible "
"tiles and never fired";
EXPECT_FALSE(d.dpOnly) << "no DP trigger is set: this is a plain tree launch";
EXPECT_FALSE(d.forceDPOnly);

EXPECT_EQ(d.tiles, 64u);
EXPECT_EQ(d.finalGrid, d.tiles) << "the DP fallback sets grid = tiles";
// With no workspace the grid selection already returns tiles, so the fallback is
// credited even though it did not move the number.
EXPECT_EQ(d.finalGrid, d.tiles) << "grid selection, not a clamp, produced grid = tiles";
// Nothing moved the grid: selection already landed on tiles.
EXPECT_EQ(d.selectedGrid, d.tiles);
EXPECT_EQ(d.selectedGrid, d.finalGrid);

EXPECT_EQ(d.idealWorkspaceBytes, solution.partialTileSize(d.tiles))
<< "the launch wanted a partials region sized by the parallel grid";
EXPECT_EQ(d.idealWorkspaceBytes, 4194304u);
EXPECT_EQ(d.idealWorkspaceBytes, 0u)
<< "the partials guard never ran, so no ideal size was ever computed";
EXPECT_EQ(d.requiredWorkspaceBytes, 0u);
EXPECT_FALSE(d.workspaceAllocated);
// Both sizings agree that nothing is reserved.
EXPECT_EQ(solution.requiredWorkspaceSize(problem, env.device), 0u);

std::ostringstream os;
solution.printStreamKLaunchSummary(os, problem, d);
const std::string line = collapseSpaces(os.str());
EXPECT_NE(line.find("changedBy = none"), std::string::npos)
<< "no clamp produced the launch grid; grid selection did";
EXPECT_NE(line.find("source = none"), std::string::npos);
EXPECT_NE(line.find("reduction = tree"), std::string::npos);
EXPECT_EQ(line.find("preFallback"), std::string::npos);
}

// ---------------------------------------------------------------------------
// The parallel branch's workspace-DP fallback, for real.
//
// The test above cannot reach it: origami applies the same workspace predicate
// during grid selection, so a zero-workspace parallel scenario never survives to
// the guard as parallel. skFixedGrid bypasses origami's selection entirely
// (getSKGridImpl takes the user-override branch before it consults
// skDynamicGrid), so a fixed grid of 2*tiles keeps the split factor at 2 --
// which streamKReconcileReduction accepts -- while the workspace stays at zero.
// That is the one shape where the guard sees parallel, computes a non-zero
// partials size, finds it does not fit, and demotes the launch itself.
//
// This is the parallel-side complement of WorkspaceDpFallbackFires (which covers
// the tree / indivisible-tiles side), and unlike that test the fallback here
// genuinely moves the grid: selected 128 -> final 64.
// ---------------------------------------------------------------------------
TEST(StreamKLaunchSummaryTest, Sk3ParallelFixedGridWorkspaceDpFallbackFires)
{
AnalyticalEnv env;
ContractionSolution solution;
solution.kernelName = "test_streamk_parallel_fixed_starved";
initStreamKSolution(solution, 3);

auto problem = makeGemmProblem(256, 4096, 4096);
problem.setWorkspaceSize(0); // no workspace at all

// 2 * tiles, so the split factor is exactly 2 -- the smallest value parallel
// reduction can express, and enough for streamKReconcileReduction to leave it
// alone.
env.device.skFixedGrid = 128;

ASSERT_FALSE(Debug::Instance().useStreamKDataParrallel())
<< "unset TENSILE_STREAMK_DATA_PARALLEL before running this suite";

ASSERT_EQ(solution.getSKReduction(problem, env.device), origami::reduction_t::parallel)
<< "scenario must start from parallel reduction for the fallback to be the "
"parallel-path fallback";

auto d = solution.computeStreamKDecisions(problem, env.device);

EXPECT_EQ(d.tiles, 64u);
ASSERT_TRUE(d.fixedGridUsed) << "the fixed-grid override is what bypasses origami's clamp";
ASSERT_EQ(d.selectedGrid, 128u) << "2 * tiles, so the reconcile keeps parallel";

EXPECT_TRUE(d.workspaceDPFallbackFired);
EXPECT_EQ(d.reduction, origami::reduction_t::tree) << "the fallback demotes to tree";
EXPECT_TRUE(d.dpOnly);
EXPECT_FALSE(d.forceDPOnly) << "runtime fallback, not the compile-time param";
EXPECT_FALSE(d.treeBoundsFallbackFired) << "tree-bounds fixup is tree-reduction only";
EXPECT_FALSE(d.clusterDPGridClamped);

EXPECT_EQ(d.finalGrid, d.tiles) << "the DP fallback sets grid = tiles";
EXPECT_EQ(d.skGridPreFallback, d.selectedGrid) << "nothing inside getSKGridImpl moved it";

// Sized by the PARALLEL grid the launch wanted, not by the grid it settled for.
EXPECT_EQ(d.idealWorkspaceBytes, solution.partialTileSize(d.selectedGrid));
EXPECT_EQ(d.idealWorkspaceBytes, 8388608u) << "128*128*4 bytes * 128 work-groups";
EXPECT_GT(d.idealWorkspaceBytes, d.givenWorkspaceBytes) << "why the fallback fired";
EXPECT_EQ(d.requiredWorkspaceBytes, 0u) << "the fallback reserves nothing";
EXPECT_FALSE(d.workspaceAllocated);

std::ostringstream os;
solution.printStreamKLaunchSummary(os, problem, d);
const std::string line = collapseSpaces(os.str());
EXPECT_NE(line.find("changedBy = workspaceDP"), std::string::npos)
<< "attribution is by which clamp produced the launch grid, not by whether "
"the number changed";
<< "the workspace fallback, not the fixed-grid override, produced the launch grid";
EXPECT_EQ(line.find("changedBy = fixedGrid"), std::string::npos);
EXPECT_NE(line.find("source = workspaceDP(runtime)"), std::string::npos);
EXPECT_NE(line.find("reduction = tree"), std::string::npos);
EXPECT_EQ(line.find("preFallback"), std::string::npos);
EXPECT_EQ(line.find("preFallback"), std::string::npos)
<< "preFallback only prints when a clamp INSIDE getSKGridImpl moved the grid";
// Both are still reported as having fired in the fallbacks section.
EXPECT_NE(line.find("fixedGrid = yes"), std::string::npos);
EXPECT_NE(line.find("workspaceDPFallback = yes"), std::string::npos);
}

// ---------------------------------------------------------------------------
Expand Down
Loading