fix(tensilelite): refuse kernel launches whose grid does not fit in 32 bits

projects/hipblaslt/tensilelite/Tensile/Contractions.py

@@ -534,14 +534,25 @@ def CompoundPredicates(cls, state, problemType):
 
                 rv += [cls('SynchronizerSizeCheck', index=0, value=valuepredicates)]
 
-        if state["InternalSupportParams"]["KernArgsVersion"] >= 1 and \
-                 not (('StreamK' in state) and (state['StreamK'] > 0)):
+        isStreamK = ('StreamK' in state) and (state['StreamK'] > 0)
+        if state["InternalSupportParams"]["KernArgsVersion"] >= 1 and not isStreamK:
             valuepredicates = []
             valuepredicates.append(state["MacroTile0"])
             valuepredicates.append(state["MacroTile1"])
             valuepredicates.append(state["GlobalSplitU"])
             rv += [cls('WorkgroupNumberCheck', index=0, value=valuepredicates)]
 
+        if isStreamK:
+            # Stream-K's normal grid is CU-scaled, not tile-scaled, so it is
+            # exempt from WorkgroupNumberCheck above. But several launch-time
+            # fallbacks (workspace too small for the ideal partial-tile
+            # reduction, the tree-fixup 24-bit bounds guard, a fixed-grid
+            # debug override, ...) hand Stream-K a one-workgroup-per-tile grid
+            # instead, same shape as the tile-scaled grid WorkgroupNumberCheck
+            # already bounds. Bound that worst case too.
+            rv += [cls('StreamKWorkgroupNumberCheck',
+                      value=[state["MacroTile0"], state["MacroTile1"]])]
+
         if not problemType.aType.isInt8x4():
             # calculate the minimum supported free dimension size
             TLUA = state['ProblemType']['TLUA']

projects/hipblaslt/tensilelite/Tensile/Tests/unit/characterization/Contractions/test_contractions_char.py

@@ -51,6 +51,48 @@ def solution_state():
     return dict(sol._state)
 
 
+# Real gfx950 Stream-K (static, StreamK: 3) solution, already vendored for the
+# BenchmarkProblems solution-pool tests (test_benchmarkProblems_solution_pool.py,
+# test_LibraryIO.py). Reused here, rather than patching a synthetic 'StreamK'
+# key onto the gfx942 fixture above, so the isStreamK branch in
+# CompoundPredicates() (Contractions.py) is exercised against a solution
+# state a real Stream-K kernel actually has, not a hand-edited one. Kept as
+# its own separate fixture chain (rather than sharing helpers with raw/
+# problem_type/solution_state above) so this addition cannot change the
+# behavior of the existing fixtures/tests.
+_STREAMK_FIXTURE = (
+    Path(__file__).parent.parent.parent / "test_data" / "solution_pool_gfx950.yaml"
+)
+
+
+@pytest.fixture(scope="module")
+def raw_streamk():
+    doc = _STREAMK_FIXTURE.read_text()
+    return yaml.load(doc, Loader=L.StrictTypeLoader)  # nosec B506
+
+
+@pytest.fixture
+def problem_type_streamk(raw_streamk):
+    return C.ProblemType.FromOriginalState(raw_streamk[4])
+
+
+@pytest.fixture(scope="module")
+def solution_state_streamk():
+    from Tensile.Common.Architectures import SUPPORTED_ISA
+    from Tensile.Common.Capabilities import makeIsaInfoMap
+    from Tensile.Toolchain.Assembly import makeAssemblyToolchain
+    from Tensile.Toolchain.Validators import validateToolchain, ToolchainDefaults
+
+    cxx = validateToolchain("amdclang++")
+    iim = makeIsaInfoMap(SUPPORTED_ISA, cxx)
+    bundler = validateToolchain(ToolchainDefaults.OFFLOAD_BUNDLER)
+    asm = makeAssemblyToolchain(cxx, bundler, "default").assembler
+    sol = L.parseLibraryLogicFile(
+        str(_STREAMK_FIXTURE), asm, False, False, False, iim, False
+    ).solutions[0]
+    return dict(sol._state)
+
+
 # --- index value classes ----------------------------------------------------
 
 def test_index_classes(snapshot):
@@ -99,3 +141,18 @@ def test_internal_args_support(solution_state):
 def test_problem_predicate_compound(problem_type, solution_state, snapshot):
     preds = C.ProblemPredicate.CompoundPredicates(solution_state, problem_type)
     assert {"count": len(preds), "tags": sorted({p.tag for p in preds})} == snapshot
+
+
+def test_problem_predicate_compound_streamk(problem_type_streamk, solution_state_streamk):
+    # test_problem_predicate_compound above never exercises
+    # CompoundPredicates()'s isStreamK branch (Contractions.py): the vendored
+    # gfx942 fixture's solution has no StreamK field. This drives the same
+    # call from a real Stream-K solution's derived state instead, which must
+    # skip the ordinary WorkgroupNumberCheck and attach
+    # StreamKWorkgroupNumberCheck in its place.
+    tags = {
+        p.tag
+        for p in C.ProblemPredicate.CompoundPredicates(solution_state_streamk, problem_type_streamk)
+    }
+    assert "StreamKWorkgroupNumberCheck" in tags
+    assert "WorkgroupNumberCheck" not in tags

projects/hipblaslt/tensilelite/Tensile/Tests/unit/characterization/DECISIONS.md

@@ -315,3 +315,43 @@ available where this change was authored.
 **Decision:** Sort every code object's input paths immediately before linking,
 so default and explicitly grouped code objects share one deterministic physical
 kernel order even though their inputs originate from different collection types.
+
+## D24 — `Contractions.py`: first arm of the deferred predicate/state matrix (isStreamK)
+
+**Context:** ADR 0007 (D10) deferred the `ProblemPredicate.CompoundPredicates`
+arms that only fire for other-problem-configurations — including the
+`isStreamK` branch, which swaps the ordinary `WorkgroupNumberCheck` for
+`StreamKWorkgroupNumberCheck` — for a future "Contractions matrix" slice, and
+explicitly rejected hand-synthesizing a derived-solution state for that slice
+("fragile, since they must match the full post-derivation key set exactly").
+While adding `StreamKWorkgroupNumberCheck` for ROCM-31290, this branch
+initially got exactly that rejected treatment: a synthetic
+`dict(solution_state, StreamK=1)` patched onto the vendored gfx942-HSS
+fixture's derived state.
+
+**Decision:** Replaced the synthetic patch with a second fixture chain
+(`raw_streamk` / `problem_type_streamk` / `solution_state_streamk`) parsed via
+the same `parseLibraryLogicFile` path, pointed at
+`Tensile/Tests/unit/test_data/solution_pool_gfx950.yaml` — a real gfx950
+static Stream-K (`StreamK: 3`) solution already vendored for the
+BenchmarkProblems solution-pool tests (`test_LibraryIO.py`,
+`test_benchmarkProblems_solution_pool.py`). `test_problem_predicate_compound_streamk`
+now asserts the `isStreamK` swap against a real derived state instead of a
+hand-edited one. Kept as a fully separate fixture chain from
+`raw`/`problem_type`/`solution_state` (no shared helpers) so the addition
+cannot change the existing fixtures' behavior.
+
+**Why:** This reuses an already-vendored fixture rather than vendoring a new
+logic YAML or hand-synthesizing state, so it avoids both alternatives ADR
+0007 rejected, while covering the first of the deferred arms. It does not
+close ADR 0007 — sparse, activation, bias, batched, and dtype arms remain
+deferred — so ADR 0007 stays `Accepted` as-is.
+
+**Alternatives rejected:**
+- Keep the synthetic `dict(solution_state, StreamK=1)` patch — rejected per
+  ADR 0007's own reasoning: a hand-edited state isn't guaranteed to match
+  what `Solution` construction actually derives for a real Stream-K solution,
+  so it could pass while a real one fails (or vice versa).
+- Vendor a new, purpose-built gfx950 StreamK logic fixture for this test —
+  rejected as unnecessary; a suitable one is already vendored and exercised
+  elsewhere in the suite.

projects/hipblaslt/tensilelite/include/Tensile/ContractionProblemPredicates.hpp

@@ -39,6 +39,7 @@
 #include <array>
 #include <cmath>
 #include <cstddef>
+#include <cstdint>
 #include <limits>
 #include <sstream>
 #include <vector>
@@ -1636,6 +1637,118 @@ namespace TensileLite
                 }
             };
 
+            struct StreamKWorkgroupNumberCheck
+                : public Predicate_CRTP<StreamKWorkgroupNumberCheck, ContractionProblemGemm>
+            {
+                enum
+                {
+                    HasIndex = false,
+                    HasValue = true
+                };
+                std::array<int, 2> value; // [MacroTile0, MacroTile1]
+
+                StreamKWorkgroupNumberCheck() = default;
+                StreamKWorkgroupNumberCheck(std::array<int, 2> value)
+                    : value(value)
+                {
+                }
+
+                static std::string Type()
+                {
+                    return "StreamKWorkgroupNumberCheck";
+                }
+
+                // WorkgroupNumberCheck above skips Stream-K solutions because
+                // Stream-K's normal grid is CU-scaled, not tile-scaled. But
+                // several launch-time fallbacks (workspace too small for the
+                // ideal partial-tile reduction, the tree-fixup 24-bit bounds
+                // guard, a fixed-grid debug override, ...) hand Stream-K a
+                // one-workgroup-per-tile grid instead, same as the plain
+                // tile-scaled grid WorkgroupNumberCheck already bounds. Any of
+                // those fallbacks can fire at launch time depending on
+                // workspace and hardware, so this is checked unconditionally
+                // for every Stream-K solution rather than only when a
+                // fallback is predicted: if the tile grid itself is unsafe,
+                // dispatch cannot rely on always landing on the CU-scaled
+                // path.
+
+                // Integer ceiling division, exact for every input; a float
+                // ceil() here would lose precision once a dimension exceeds
+                // 2^24 and could undercount tiles right at the boundary this
+                // predicate exists to guard.
+                static size_t ceilDiv(size_t numerator, size_t denominator)
+                {
+                    return (numerator + denominator - 1) / denominator;
+                }
+
+                // The count solve() builds the fallback grid from is
+                // ContractionProblemGemm::getNumTiles(sizeMapping, 1). This is
+                // computed separately because a predicate is handed
+                // MacroTile0/MacroTile1 as its value rather than the solution's
+                // sizeMapping, and it agrees with getNumTiles only for the
+                // single-free-index, unpacked-batch case. getNumTiles multiplies
+                // over every free and batch index and, when packBatchDims is set,
+                // folds batch into the M or N extent before dividing by the macro
+                // tile; this uses index 0 of each and multiplies batch in
+                // afterwards. WorkgroupNumberCheck makes the same simplification,
+                // so the two predicates agree with each other on every shape;
+                // revisit both together if a multi-free-index or packed-batch
+                // contraction ever has to be bounded here.
+                static size_t tiles(ContractionProblemGemm const& problem,
+                                    std::array<int, 2> const&     value)
+                {
+                    return ceilDiv(problem.freeSizeA(0), value[0])
+                           * ceilDiv(problem.freeSizeB(0), value[1])
+                           * problem.batchSize(0);
+                }
+
+                // The launch narrows the 64-bit work-item count
+                // (workGroupSize * numWorkGroups, both size_t) to the
+                // `unsigned int` globalWorkSize parameters of
+                // hipExtModuleLaunchKernel, so a tile-scaled Stream-K grid is
+                // only safe while that product still fits in 32 bits.
+                // Measured on gfx950 at 256 threads: 2^24 + 16 workgroups
+                // wraps to a 16-workgroup grid and leaves the other
+                // 16,777,216 workgroups' worth of D unwritten with no error
+                // reported, and 2^24 workgroups lands on exactly 2^32, which
+                // narrows to 0 and HIP rejects outright.
+                //
+                // NumThreads is not carried in this predicate's value, so the
+                // bound assumes the 256-thread configuration that dominates
+                // the shipped Stream-K solutions. That is exact at 256
+                // threads, conservative at 64 and 128 (whose real limits are
+                // 2^26 and 2^25), and still too loose for the one 512-thread
+                // family, MacroTile 448x128, which wraps at 2^23 but needs
+                // roughly 896 GiB of D to get there. Deliberately *not*
+                // MAX_WORKGROUP_NUMBER: that constant happens to have the
+                // same magnitude but guards an unrelated fp32 conversion.
+                static constexpr size_t MAX_STREAMK_TILE_COUNT
+                    = std::numeric_limits<uint32_t>::max() / 256;
+
+                virtual bool operator()(ContractionProblemGemm const& problem) const override
+                {
+                    return tiles(problem, value) <= MAX_STREAMK_TILE_COUNT;
+                }
+
+                virtual std::string toString() const override
+                {
+                    return concatenate(this->type(), "(MT0:", value[0], ",MT1:", value[1], ")");
+                }
+
+                virtual bool debugEval(ContractionProblemGemm const& problem,
+                                       std::ostream&                 stream) const override
+                {
+                    return debugEvalCmp(problem,
+                                        stream,
+                                        "sk_tiles",
+                                        tiles(problem, value),
+                                        "<=",
+                                        "max",
+                                        MAX_STREAMK_TILE_COUNT);
+                }
+            };
+
+
             struct PersistentKernelCheck
                 : public Predicate_CRTP<PersistentKernelCheck, ContractionProblemGemm>
             {

projects/hipblaslt/tensilelite/include/Tensile/Serialization/ContractionPredicates.hpp

@@ -96,6 +96,7 @@ namespace TensileLite
                          Predicates::Contraction::BufferLoadOffsetLimitCheck_Beta>(),
                      Base::template Pair<Predicates::Contraction::BufferStoreOffsetLimitCheck>(),
                      Base::template Pair<Predicates::Contraction::WorkgroupNumberCheck>(),
+                     Base::template Pair<Predicates::Contraction::StreamKWorkgroupNumberCheck>(),
                      Base::template Pair<Predicates::Contraction::PersistentKernelCheck>(),
                      Base::template Pair<Predicates::Contraction::GlobalSplitUCheckMinK>(),
                      Base::template Pair<Predicates::Contraction::CDStridesEqual>(),
@@ -364,6 +365,12 @@ namespace TensileLite
         {
         };
 
+        template <typename IO>
+        struct MappingTraits<Predicates::Contraction::StreamKWorkgroupNumberCheck, IO>
+            : public AutoMappingTraits<Predicates::Contraction::StreamKWorkgroupNumberCheck, IO>
+        {
+        };
+
         template <typename IO>
         struct MappingTraits<Predicates::Contraction::PersistentKernelCheck, IO>
             : public AutoMappingTraits<Predicates::Contraction::PersistentKernelCheck, IO>

projects/hipblaslt/tensilelite/src/hip/HipSolutionAdapter.cpp

@@ -28,6 +28,7 @@
 #include <hip/hip_runtime.h>
 
 #include <cstddef>
+#include <limits>
 
 #include <Tensile/Debug.hpp>
 #include <Tensile/EmbeddedData.hpp>
@@ -514,6 +515,28 @@ namespace TensileLite
             return hipSuccess;
         }
 
+        namespace
+        {
+            // Both launch APIs below take 32-bit grid parameters, but the values
+            // handed to them come from dim3, which is vector3<size_t>. A grid that
+            // does not fit is narrowed silently rather than rejected: measured on
+            // gfx950, a 256-thread kernel at 2^24 + 16 workgroups wraps to a
+            // 16-workgroup launch and leaves the other 16,777,216 workgroups' worth
+            // of D holding whatever was there before, with no error raised.
+            //
+            // StreamKWorkgroupNumberCheck keeps dispatch away from the tuned
+            // Stream-K solutions that reach this, but it cannot see every path:
+            // skFixedGrid and skGridMultiplier rewrite the grid after selection, the
+            // predicate's bound assumes 256 threads, and the hand-written custom
+            // kernels carry no such predicate. Refusing the launch is what makes
+            // those remaining paths loud instead of silent.
+            bool fitsLaunchDim(TensileLite::dim3 const& dim)
+            {
+                constexpr size_t limit = std::numeric_limits<unsigned int>::max();
+                return dim.x <= limit && dim.y <= limit && dim.z <= limit;
+            }
+        }
+
         hipError_t SolutionAdapter::launchKernel(KernelInvocation const& kernel)
         {
             return launchKernel(kernel, nullptr, nullptr, nullptr);
@@ -547,6 +570,43 @@ namespace TensileLite
                 return hipSuccess;
             }
 
+#ifdef HIP_HAS_CLUSTER_LAUNCH
+            const bool enableCluster = (kernel.clusterDim.x > 1 || kernel.clusterDim.y > 1);
+#else
+            const bool enableCluster = false;
+#endif
+
+            // Checked before the code object is looked up: a grid that cannot be
+            // expressed is impossible whether or not the kernel loads, and this is
+            // the last point where the 64-bit values are still intact. The two
+            // launch APIs below narrow different quantities, so each is bounded
+            // against the one it actually passes.
+            if(enableCluster)
+            {
+                // hipDrvLaunchKernelEx enumerates the grid in workgroups.
+                if(!fitsLaunchDim(kernel.numWorkGroups))
+                {
+                    std::cerr << "hipDrvLaunchKernelEx: workgroup count exceeds the 32-bit grid "
+                              << "dimensions (numWorkGroups " << kernel.numWorkGroups
+                              << ") for kernel: " << kernel.kernelName << std::endl;
+                    return hipErrorInvalidValue;
+                }
+            }
+            else
+            {
+                // hipExtModuleLaunchKernel's globalWorkSize is in work items, so
+                // the workGroupSize * numWorkGroups product is what has to fit.
+                if(!fitsLaunchDim(kernel.numWorkItems))
+                {
+                    std::cerr << "hipExtModuleLaunchKernel: work-item count exceeds the 32-bit "
+                              << "globalWorkSize parameters (numWorkItems " << kernel.numWorkItems
+                              << ", from numWorkGroups " << kernel.numWorkGroups
+                              << " of workgroup size " << kernel.workGroupSize
+                              << ") for kernel: " << kernel.kernelName << std::endl;
+                    return hipErrorInvalidValue;
+                }
+            }
+
             hipFunction_t function;
             HIP_CHECK_RETURN_WITH_LOG(getKernel(function, kernel.kernelName),
                 [&](hipError_t error) {
@@ -571,7 +631,6 @@ namespace TensileLite
                 HIP_CHECK_RETURN(hipEventRecord(startEvent, stream));
 
 #ifdef HIP_HAS_CLUSTER_LAUNCH
-            bool enableCluster = (kernel.clusterDim.x > 1 || kernel.clusterDim.y > 1);
             if(enableCluster)
             {
                 if(kernel.clusterDim.x == 0 || kernel.clusterDim.y == 0)