fix: preserve dynamic LDS in group_segment_size on instrumented dispatch

src/interceptor.cc

@@ -812,7 +812,18 @@ hsa_kernel_dispatch_packet_t * hsaInterceptor::fixupPacket(const hsa_kernel_disp
                             fixupKernArgs(new_kernargs, packet->kernarg_address, comms->get_dev_rsrc_ptr(), args);
                             dispatch->kernarg_address = new_kernargs;
                             dispatch->private_segment_size = args.private_segment_size;
-                            dispatch->group_segment_size = args.group_segment_size;
+                            // The original packet's group_segment_size includes both the
+                            // kernel's fixed LDS and any dynamic shared memory requested
+                            // at launch (e.g. extern __shared__). The instrumented clone's
+                            // descriptor only has its fixed portion. Preserve the dynamic
+                            // part by computing: instrumented_fixed + dynamic_from_launch.
+                            {
+                                arg_descriptor_t orig_args = {};
+                                uint32_t dynamic_lds = packet->group_segment_size;
+                                if (kernel_cache_.getArgDescriptor(queues_[queue], it->second.name_, orig_args, false))
+                                    dynamic_lds = packet->group_segment_size - orig_args.group_segment_size;
+                                dispatch->group_segment_size = args.group_segment_size + dynamic_lds;
+                            }
                             // Store the new kernarg address so we can free it up at kernel completion
                             pending_kernargs_[sig] = new_kernargs;
                         }

tests/run_dynamic_lds_tests.sh

@@ -0,0 +1,96 @@
+#!/bin/bash
+################################################################################
+# Dynamic LDS preservation tests for omniprobe instrumentation.
+#
+# Verifies that instrumented dispatches preserve the dynamic shared memory
+# portion of group_segment_size requested at launch (extern __shared__ with
+# dynamicSharedMemoryBytes). Without the corresponding fix in
+# interceptor.cc::fixupPacket the instrumented dispatch only allocates the
+# instrumented clone's fixed LDS, the kernel reads past its allocated LDS,
+# and the output mismatches.
+################################################################################
+
+set -e
+
+# Source common utilities
+SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+source "${SCRIPT_DIR}/test_common.sh"
+
+check_omniprobe
+
+DYNAMIC_LDS_TEST="${BUILD_DIR}/tests/test_kernels/dynamic_lds_test"
+
+if [ ! -x "$DYNAMIC_LDS_TEST" ]; then
+    echo -e "\n${YELLOW}SKIP: Dynamic LDS tests (dynamic_lds_test not built)${NC}"
+    export TESTS_RUN TESTS_PASSED TESTS_FAILED
+    return 0 2>/dev/null || exit 0
+fi
+
+echo ""
+echo "================================================================================"
+echo "Dynamic LDS Preservation Tests"
+echo "================================================================================"
+echo "  Test binary: $DYNAMIC_LDS_TEST"
+echo "================================================================================"
+
+################################################################################
+# Test: Native run (no instrumentation) produces correct output.
+#
+# Sanity check that the test kernel itself is correct and that the dynamic
+# LDS request is honoured by the runtime under normal dispatch.
+################################################################################
+
+TESTS_RUN=$((TESTS_RUN + 1))
+TEST_NAME="dynamic_lds_native"
+echo -e "\n${YELLOW}[TEST $TESTS_RUN]${NC} $TEST_NAME"
+echo "  Run native (no instrumentation): kernel uses extern __shared__ with"
+echo "  dynamicSharedMemoryBytes set at launch"
+
+OUTPUT_FILE="$OUTPUT_DIR/${TEST_NAME}.out"
+
+ROCR_VISIBLE_DEVICES="$ROCR_VISIBLE_DEVICES" \
+    "$DYNAMIC_LDS_TEST" > "$OUTPUT_FILE" 2>&1 \
+    && run_ok=true || run_ok=true
+
+if grep -q "dynamic_lds_test: PASS" "$OUTPUT_FILE"; then
+    echo -e "  ${GREEN}✓ PASS${NC} - Native execution correct"
+    TESTS_PASSED=$((TESTS_PASSED + 1))
+else
+    echo -e "  ${RED}✗ FAIL${NC} - Native execution incorrect (test kernel itself is broken)"
+    tail -20 "$OUTPUT_FILE"
+    echo "  Output saved to: $OUTPUT_FILE"
+    TESTS_FAILED=$((TESTS_FAILED + 1))
+fi
+
+################################################################################
+# Test: Instrumented run preserves dynamic LDS.
+#
+# Without the fixupPacket fix, the instrumented dispatch overwrites
+# group_segment_size with the instrumented kernel's fixed-LDS value, dropping
+# the dynamic portion requested at launch. The kernel then reads/writes past
+# its allocated LDS region, producing incorrect output.
+################################################################################
+
+TESTS_RUN=$((TESTS_RUN + 1))
+TEST_NAME="dynamic_lds_instrumented"
+echo -e "\n${YELLOW}[TEST $TESTS_RUN]${NC} $TEST_NAME"
+echo "  Run under omniprobe -i: instrumented dispatch must preserve dynamic LDS"
+
+OUTPUT_FILE="$OUTPUT_DIR/${TEST_NAME}.out"
+
+ROCR_VISIBLE_DEVICES="$ROCR_VISIBLE_DEVICES" \
+    "$OMNIPROBE" -i -a Heatmap -- "$DYNAMIC_LDS_TEST" > "$OUTPUT_FILE" 2>&1 \
+    && run_ok=true || run_ok=true
+
+if grep -q "dynamic_lds_test: PASS" "$OUTPUT_FILE"; then
+    echo -e "  ${GREEN}✓ PASS${NC} - Instrumented execution correct (dynamic LDS preserved)"
+    TESTS_PASSED=$((TESTS_PASSED + 1))
+else
+    echo -e "  ${RED}✗ FAIL${NC} - Instrumented execution incorrect (dynamic LDS dropped?)"
+    tail -20 "$OUTPUT_FILE"
+    echo "  Output saved to: $OUTPUT_FILE"
+    TESTS_FAILED=$((TESTS_FAILED + 1))
+fi
+
+# Export updated counters for parent script
+export TESTS_RUN TESTS_PASSED TESTS_FAILED

tests/run_handler_tests.sh

@@ -43,5 +43,8 @@ source "${SCRIPT_DIR}/run_scope_filter_tests.sh"
 # Module-load kernel discovery tests (hipModuleLoad .hsaco)
 source "${SCRIPT_DIR}/run_module_load_tests.sh"
 
+# Dynamic LDS preservation tests (extern __shared__ at launch)
+source "${SCRIPT_DIR}/run_dynamic_lds_tests.sh"
+
 # Print summary
 print_summary

tests/test_kernels/CMakeLists.txt

@@ -82,6 +82,18 @@ target_compile_options(block_filter_test PRIVATE
 target_link_options(block_filter_test PRIVATE -fgpu-rdc)
 add_dependencies(block_filter_test copy_bitcode AMDGCNSubmitAddressMessages-rocm)
 
+# Dynamic LDS test kernel - uses extern __shared__ with dynamicSharedMemoryBytes
+# at launch to verify that instrumented dispatches preserve the dynamic portion
+# of group_segment_size.
+add_executable(dynamic_lds_test dynamic_lds_test.cpp)
+set_source_files_properties(dynamic_lds_test.cpp PROPERTIES LANGUAGE HIP)
+target_compile_options(dynamic_lds_test PRIVATE
+    ${TEST_KERNEL_COMPILE_FLAGS}
+    -fpass-plugin=${INST_PLUGIN}
+)
+target_link_options(dynamic_lds_test PRIVATE -fgpu-rdc)
+add_dependencies(dynamic_lds_test copy_bitcode AMDGCNSubmitAddressMessages-rocm)
+
 # Module-load test: kernel compiled to standalone .hsaco, loaded via hipModuleLoad at runtime.
 # This exercises the code path where instrumented kernels live in a dynamically-loaded code
 # object rather than being embedded in the executable's .hip_fatbin section.

tests/test_kernels/dynamic_lds_test.cpp

@@ -0,0 +1,83 @@
+// Test kernel exercising dynamic shared memory (extern __shared__).
+//
+// Verifies that omniprobe's instrumented dispatch preserves the dynamic
+// portion of group_segment_size requested at launch via the third launch
+// argument (dynamicSharedMemoryBytes). Without that preservation the
+// instrumented clone allocates only its own fixed LDS, the kernel reads
+// past the end of LDS, and the output mismatches.
+
+#include <hip/hip_runtime.h>
+#include <iostream>
+#include <vector>
+#include "hip_test_utils.h"
+
+// Each thread writes its global index into shared memory at threadIdx.x,
+// then reads from a permuted slot and writes that to global memory.
+// All accesses are within blockDim.x ints of LDS, requested at launch.
+__global__ void dynamic_lds_kernel(int* output, size_t n) {
+    extern __shared__ int sdata[];
+
+    int tid = threadIdx.x;
+    size_t idx = static_cast<size_t>(blockIdx.x) * blockDim.x + tid;
+
+    sdata[tid] = static_cast<int>(idx);
+    __syncthreads();
+
+    int permuted_tid = blockDim.x - 1 - tid;
+    if (idx < n) {
+        output[idx] = sdata[permuted_tid];
+    }
+}
+
+int main() {
+    std::cerr << "Starting dynamic_lds_test" << std::endl;
+
+    constexpr int blocksize = 256;
+    constexpr int no_blocks = 4;
+    constexpr size_t size = static_cast<size_t>(blocksize) * no_blocks;
+
+    int* d_output = nullptr;
+    CHECK_HIP(hipMalloc(&d_output, size * sizeof(int)));
+    CHECK_HIP(hipMemset(d_output, -1, size * sizeof(int)));
+
+    // Request enough dynamic LDS for one int per thread in the block.
+    // Without the fixupPacket fix, the instrumented dispatch will only
+    // allocate the instrumented kernel's fixed LDS (typically much smaller
+    // or zero), and the kernel will read past its allocated LDS region.
+    const size_t dyn_lds_bytes = blocksize * sizeof(int);
+
+    dynamic_lds_kernel<<<no_blocks, blocksize, dyn_lds_bytes>>>(d_output, size);
+    CHECK_HIP(hipGetLastError());
+    CHECK_HIP(hipDeviceSynchronize());
+
+    std::vector<int> h_output(size);
+    CHECK_HIP(hipMemcpy(h_output.data(), d_output,
+                        size * sizeof(int), hipMemcpyDeviceToHost));
+    CHECK_HIP(hipFree(d_output));
+
+    int errors = 0;
+    constexpr int max_errors_to_report = 5;
+    for (int b = 0; b < no_blocks; ++b) {
+        for (int t = 0; t < blocksize; ++t) {
+            size_t idx = static_cast<size_t>(b) * blocksize + t;
+            int permuted_tid = blocksize - 1 - t;
+            int expected = b * blocksize + permuted_tid;
+            if (h_output[idx] != expected) {
+                if (errors < max_errors_to_report) {
+                    std::cerr << "  MISMATCH at idx=" << idx
+                              << ": expected " << expected
+                              << ", got " << h_output[idx] << std::endl;
+                }
+                ++errors;
+            }
+        }
+    }
+
+    if (errors == 0) {
+        std::cout << "dynamic_lds_test: PASS" << std::endl;
+        return 0;
+    }
+    std::cout << "dynamic_lds_test: FAIL (" << errors
+              << " mismatches out of " << size << ")" << std::endl;
+    return 1;
+}