iassiour · Oct 28, 2026 · Oct 26, 2026 · Oct 27, 2026
@@ -285,14 +285,21 @@ hipError_t ihipGraphAddMemsetNode(hip::GraphNode** pGraphNode, hip::Graph* graph
   return status;
 }
 
-hipError_t capturehipLaunchKernel(hipStream_t& stream, const void*& hostFunction, dim3& gridDim,
-                                  dim3& blockDim, void**& args, size_t& sharedMemBytes) {
+hipError_t capturehipLaunchKernelCommon(hipStream_t& stream, const void*& hostFunction,
+                                        dim3& gridDim, dim3& blockDim, void**& args,
+                                        size_t& sharedMemBytes, dim3& clusterDim,
+                                        bool& clusterDimsSpecified) {
   ClPrint(amd::LOG_DETAIL_DEBUG, amd::LOG_API,
-          "[hipGraph] Current capture node LaunchKernel on stream : %p", stream);
+          "[hipGraph] Current capture node LaunchKernel on stream : %p (cluster %u,%u,%u%s)",
+          stream, clusterDim.x, clusterDim.y, clusterDim.z,
+          clusterDimsSpecified ? "" : ", not specified");
 
   if (!hip::isValid(stream)) {
     return hipErrorContextIsDestroyed;
   }
+  // {0,0,0} tells the kernel node that no cluster dimensions were given, so it resolves the
+  // kernel's own requirement (__cluster_dims__ or a runtime cluster attribute) instead.
+  const dim3 nodeClusterDim = clusterDimsSpecified ? clusterDim : dim3(0, 0, 0);
   hip::Stream* s = reinterpret_cast<hip::Stream*>(stream);
   hipKernelNodeParams nodeParams;
   nodeParams.func = const_cast<void*>(hostFunction);
@@ -303,9 +310,10 @@ hipError_t capturehipLaunchKernel(hipStream_t& stream, const void*& hostFunction
   nodeParams.sharedMemBytes = sharedMemBytes;
 
   hip::GraphNode* pGraphNode;
-  hipError_t status = ihipGraphAddKernelNode(
-      &pGraphNode, s->GetCaptureGraph(), s->GetLastCapturedNodes().data(),
-      s->GetLastCapturedNodes().size(), &nodeParams, nullptr, true, 0, s->DeviceId());
+  hipError_t status =
+      ihipGraphAddKernelNode(&pGraphNode, s->GetCaptureGraph(), s->GetLastCapturedNodes().data(),
+                             s->GetLastCapturedNodes().size(), &nodeParams, nullptr, true, 0,
+                             s->DeviceId(), 0, 0, 0, nodeClusterDim);
   if (status != hipSuccess) {
     return status;
   }

@@ -7,8 +7,13 @@
 #pragma once
 namespace hip {
 // forward declaration of capture methods
-hipError_t capturehipLaunchKernel(hipStream_t& stream, const void*& hostFunction, dim3& gridDim,
-                                  dim3& blockDim, void**& args, size_t& sharedMemBytes);
+
+// Capture path of hipLaunchKernel_common, shared by hipLaunchKernel and hipLaunchKernelExC.
+// Keeps the cluster dimensions given with hipLaunchAttributeClusterDimension on the kernel node.
+hipError_t capturehipLaunchKernelCommon(hipStream_t& stream, const void*& hostFunction,
+                                        dim3& gridDim, dim3& blockDim, void**& args,
+                                        size_t& sharedMemBytes, dim3& clusterDim,
+                                        bool& clusterDimsSpecified);
 
 hipError_t capturehipExtModuleLaunchKernel(hipStream_t& stream, hipFunction_t& f,
                                            uint32_t& globalWorkSizeX, uint32_t& globalWorkSizeY,

@@ -1112,7 +1112,8 @@ hipError_t hipLaunchKernel_common(
     return hipErrorInvalidValue;
   }
   CHECK_STREAM_DETACHED(stream);
-  STREAM_CAPTURE(hipLaunchKernel, stream, hostFunction, gridDim, blockDim, args, sharedMemBytes);
+  STREAM_CAPTURE(hipLaunchKernelCommon, stream, hostFunction, gridDim, blockDim, args,
+                 sharedMemBytes, clusterDim, clusterDimsSpecified);
   return ihipLaunchKernel(hostFunction, gridDim, blockDim, args, sharedMemBytes, stream, nullptr,
                           nullptr, 0, clusterDim, clusterDimsSpecified, dynDataPrefetchConfig);
 }

@@ -22,3 +22,7 @@ cluster:
     <<: *level_default
   Unit_hipClusterLaunch_LaunchApi_Basic:
     <<: *level_default
+  Unit_hipClusterLaunch_StreamCapture_LaunchAttribute:
+    <<: *level_default
+  Unit_hipClusterLaunch_StreamCapture_CompilerDirective:
+    <<: *level_default
@@ -27,6 +27,7 @@ set(TEST_SRC
     hipClusterLaunch.cc
     cluster_class.cc
     cluster_split_barrier.cc
+    hipClusterStreamCapture.cc
 )
 
 # Append it to the list

@@ -0,0 +1,229 @@
+/*
+ * Copyright (c) Advanced Micro Devices, Inc., or its affiliates.
+ *
+ * SPDX-License-Identifier: MIT
+ */
+
+// Cluster dimensions of kernel launches captured into a graph by stream capture.
+
+#include <hip_test_common.hh>
+#include <resource_guards.hh>
+#include <hip/hip_cooperative_groups.h>
+
+#include <memory>
+#include <vector>
+
+#include "ClusterHelper.hpp"
+
+namespace cg = cooperative_groups;
+
+namespace {
+
+constexpr int kBlocks = 4;
+constexpr int kThreads = 64;
+
+// Every block records the cluster shape it runs in.
+__global__ void RecordClusterDims(dim3* out) {
+  if (threadIdx.x == 0) out[blockIdx.x] = cg::this_cluster().dim_blocks();
+}
+
+__global__ void CLUSTER_DIMS(2, 1, 1) RecordClusterDimsCD(dim3* out) {
+  if (threadIdx.x == 0) out[blockIdx.x] = cg::this_cluster().dim_blocks();
+}
+
+// Global-mode stream capture that is ended, and its graph destroyed, on every exit path.
+class CaptureGuard {
+ public:
+  explicit CaptureGuard(hipStream_t stream) : stream_(stream) {
+    HIP_CHECK(hipStreamBeginCapture(stream_, hipStreamCaptureModeGlobal));
+    capturing_ = true;
+  }
+  CaptureGuard(const CaptureGuard&) = delete;
+  CaptureGuard& operator=(const CaptureGuard&) = delete;
+  ~CaptureGuard() {
+    if (capturing_) {
+      hipGraph_t partial = nullptr;
+      if (hipStreamEndCapture(stream_, &partial) == hipSuccess && partial != nullptr) {
+        static_cast<void>(hipGraphDestroy(partial));
+      }
+    }
+    if (graph_ != nullptr) static_cast<void>(hipGraphDestroy(graph_));
+  }
+
+  hipGraph_t End() {
+    capturing_ = false;
+    HIP_CHECK(hipStreamEndCapture(stream_, &graph_));
+    return graph_;
+  }
+
+ private:
+  hipStream_t stream_;
+  bool capturing_ = false;
+  hipGraph_t graph_ = nullptr;
+};
+
+// The cluster dimensions stored on the graph's only kernel node.
+dim3 KernelNodeClusterDims(hipGraph_t graph) {
+  size_t num_nodes = 0;
+  HIP_CHECK(hipGraphGetNodes(graph, nullptr, &num_nodes));
+  std::vector<hipGraphNode_t> nodes(num_nodes);
+  HIP_CHECK(hipGraphGetNodes(graph, nodes.data(), &num_nodes));
+  hipGraphNode_t kernel_node = nullptr;
+  for (auto node : nodes) {
+    hipGraphNodeType type;
+    HIP_CHECK(hipGraphNodeGetType(node, &type));
+    if (type == hipGraphNodeTypeKernel) {
+      REQUIRE(kernel_node == nullptr);
+      kernel_node = node;
+    }
+  }
+  REQUIRE(kernel_node != nullptr);
+  hipKernelNodeAttrValue value{};
+  HIP_CHECK(
+      hipGraphKernelNodeGetAttribute(kernel_node, hipLaunchAttributeClusterDimension, &value));
+  return dim3(value.clusterDim.x, value.clusterDim.y, value.clusterDim.z);
+}
+
+// Instantiates and launches the graph, then checks the cluster shape every block ran in.
+void ReplayAndCheck(hipGraph_t graph, hipStream_t stream, dim3* d_out, dim3 expected) {
+  HIP_CHECK(hipMemset(d_out, 0, kBlocks * sizeof(dim3)));
+  HIP_CHECK(hipDeviceSynchronize());
+  hipGraphExec_t exec = nullptr;
+  HIP_CHECK(hipGraphInstantiate(&exec, graph, nullptr, nullptr, 0));
+  std::unique_ptr<hipGraphExec, decltype(&hipGraphExecDestroy)> exec_guard(exec,
+                                                                           hipGraphExecDestroy);
+  HIP_CHECK(hipGraphLaunch(exec, stream));
+  HIP_CHECK(hipStreamSynchronize(stream));
+  std::vector<dim3> out(kBlocks);
+  HIP_CHECK(hipMemcpy(out.data(), d_out, kBlocks * sizeof(dim3), hipMemcpyDeviceToHost));
+  for (int b = 0; b < kBlocks; ++b) {
+    INFO("block " << b);
+    REQUIRE(out[b].x == expected.x);
+    REQUIRE(out[b].y == expected.y);
+    REQUIRE(out[b].z == expected.z);
+  }
+}
+
+}  // namespace
+
+/**
+ * @addtogroup cluster
+ * @{
+ * @ingroup ClusterTest
+ */
+
+/**
+ * Test Description
+ * ------------------------
+ *  - A hipLaunchKernelEx launch with hipLaunchAttributeClusterDimension = {2, 1, 1} is captured
+ *    into a graph. The captured kernel node must carry the cluster dimensions, and every block of
+ *    the replayed graph must run in a 2x1x1 cluster.
+ *  - Without the attribute, the captured node must have no cluster ({1, 1, 1}).
+ *
+ * Test source
+ * ------------------------
+ *  - catch/unit/cluster/hipClusterStreamCapture.cc
+ *
+ * Test requirements
+ * ------------------------
+ *  - HIP_VERSION >= 7.0
+ */
+HIP_TEST_CASE(Unit_hipClusterLaunch_StreamCapture_LaunchAttribute) {
+  if (!CheckTargetSupport()) {
+    HIP_SKIP_TEST("Cluster launch is not supported on this device");
+    return;
+  }
+  LinearAllocGuard<dim3> out_guard(LinearAllocs::hipMalloc, kBlocks * sizeof(dim3));
+  dim3* d_out = out_guard.ptr();
+  StreamGuard stream_guard(Streams::withFlags, hipStreamNonBlocking);
+  hipStream_t stream = stream_guard.stream();
+
+  hipLaunchAttribute attr{};
+  attr.id = hipLaunchAttributeClusterDimension;
+  attr.val.clusterDim.x = 2;
+  attr.val.clusterDim.y = 1;
+  attr.val.clusterDim.z = 1;
+  hipLaunchConfig_t config{};
+  config.gridDim = dim3(kBlocks);
+  config.blockDim = dim3(kThreads);
+  config.stream = stream;
+  const bool with_attribute = GENERATE(true, false);
+  INFO("cluster attribute " << (with_attribute ? "{2, 1, 1}" : "absent"));
+  config.attrs = with_attribute ? &attr : nullptr;
+  config.numAttrs = with_attribute ? 1 : 0;
+  const dim3 expected = with_attribute ? dim3(2, 1, 1) : dim3(1, 1, 1);
+
+  CaptureGuard capture(stream);
+  HIP_CHECK(hipLaunchKernelEx(&config, RecordClusterDims, d_out));
+  hipGraph_t graph = capture.End();
+
+  const dim3 node_dims = KernelNodeClusterDims(graph);
+  REQUIRE(node_dims.x == expected.x);
+  REQUIRE(node_dims.y == expected.y);
+  REQUIRE(node_dims.z == expected.z);
+  ReplayAndCheck(graph, stream, d_out, expected);
+}
+
+/**
+ * Test Description
+ * ------------------------
+ *  - A kernel declared with __cluster_dims__(2, 1, 1) is captured from launches that give no
+ *    cluster dimensions (hipLaunchKernelGGL, and hipLaunchKernelEx without attributes). The
+ *    captured kernel node must take the kernel's own cluster dimensions, and the replayed blocks
+ *    must run in 2x1x1 clusters. A launch attribute that matches the kernel's dimensions must
+ *    also be accepted.
+ *
+ * Test source
+ * ------------------------
+ *  - catch/unit/cluster/hipClusterStreamCapture.cc
+ *
+ * Test requirements
+ * ------------------------
+ *  - HIP_VERSION >= 7.0
+ */
+HIP_TEST_CASE(Unit_hipClusterLaunch_StreamCapture_CompilerDirective) {
+  if (!CheckTargetSupport()) {
+    HIP_SKIP_TEST("Cluster launch is not supported on this device");
+    return;
+  }
+  LinearAllocGuard<dim3> out_guard(LinearAllocs::hipMalloc, kBlocks * sizeof(dim3));
+  dim3* d_out = out_guard.ptr();
+  StreamGuard stream_guard(Streams::withFlags, hipStreamNonBlocking);
+  hipStream_t stream = stream_guard.stream();
+
+  hipLaunchAttribute attr{};
+  attr.id = hipLaunchAttributeClusterDimension;
+  attr.val.clusterDim.x = 2;
+  attr.val.clusterDim.y = 1;
+  attr.val.clusterDim.z = 1;
+  hipLaunchConfig_t config{};
+  config.gridDim = dim3(kBlocks);
+  config.blockDim = dim3(kThreads);
+  config.stream = stream;
+
+  CaptureGuard capture(stream);
+  SECTION("hipLaunchKernelGGL") {
+    hipLaunchKernelGGL(RecordClusterDimsCD, dim3(kBlocks), dim3(kThreads), 0, stream, d_out);
+    HIP_CHECK(hipGetLastError());
+  }
+  SECTION("hipLaunchKernelEx without attributes") {
+    HIP_CHECK(hipLaunchKernelEx(&config, RecordClusterDimsCD, d_out));
+  }
+  SECTION("hipLaunchKernelEx with a matching cluster attribute") {
+    config.attrs = &attr;
+    config.numAttrs = 1;
+    HIP_CHECK(hipLaunchKernelEx(&config, RecordClusterDimsCD, d_out));
+  }
+  hipGraph_t graph = capture.End();
+
+  const dim3 node_dims = KernelNodeClusterDims(graph);
+  REQUIRE(node_dims.x == 2);
+  REQUIRE(node_dims.y == 1);
+  REQUIRE(node_dims.z == 1);
+  ReplayAndCheck(graph, stream, d_out, dim3(2, 1, 1));
+}
+
+/**
+ * End doxygen group ClusterTest.
+ * @}
+ */