diff --git a/src/interceptor.cc b/src/interceptor.cc index 4161444..b29d86e 100644 --- a/src/interceptor.cc +++ b/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; } diff --git a/tests/run_dynamic_lds_tests.sh b/tests/run_dynamic_lds_tests.sh new file mode 100755 index 0000000..5305597 --- /dev/null +++ b/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 diff --git a/tests/run_handler_tests.sh b/tests/run_handler_tests.sh index 3e8e0bd..ffb5c85 100755 --- a/tests/run_handler_tests.sh +++ b/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 diff --git a/tests/test_kernels/CMakeLists.txt b/tests/test_kernels/CMakeLists.txt index 89d93fb..3db462e 100644 --- a/tests/test_kernels/CMakeLists.txt +++ b/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. diff --git a/tests/test_kernels/dynamic_lds_test.cpp b/tests/test_kernels/dynamic_lds_test.cpp new file mode 100644 index 0000000..05828d7 --- /dev/null +++ b/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 +#include +#include +#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(blockIdx.x) * blockDim.x + tid; + + sdata[tid] = static_cast(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(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<<>>(d_output, size); + CHECK_HIP(hipGetLastError()); + CHECK_HIP(hipDeviceSynchronize()); + + std::vector 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(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; +}