Skip to content

Commit 3d8e93b

Browse files
yoshifuminakamurayoshifuminakamura
committed
Adjust SBD NCU profile defaults
Signed-off-by: yoshifuminakamura <nakamura@riken.jp>
1 parent d1b97c2 commit 3d8e93b

3 files changed

Lines changed: 156 additions & 31 deletions

File tree

‎programs/sbd/profile.sh‎

Lines changed: 114 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,43 @@ sbd_ncu_profile_mode() {
1313
printf '%s\n' "${BK_SBD_NCU_PROFILE_MODE:-discovery}"
1414
}
1515

16+
sbd_ncu_profile_timeout_seconds() {
17+
local timeout_seconds="${BK_SBD_NCU_PROFILE_TIMEOUT_SECONDS:-0}"
18+
19+
case "$timeout_seconds" in
20+
''|*[!0-9]*)
21+
echo "SBD NCU profile timeout must be a non-negative integer: ${timeout_seconds}" >&2
22+
return 1
23+
;;
24+
esac
25+
26+
printf '%s\n' "$timeout_seconds"
27+
}
28+
29+
sbd_strip_ncu_launch_count_args() {
30+
local filtered_args=()
31+
32+
while [ "$#" -gt 0 ]; do
33+
case "$1" in
34+
--launch-count)
35+
shift
36+
if [ "$#" -gt 0 ]; then
37+
shift
38+
fi
39+
;;
40+
--launch-count=*)
41+
shift
42+
;;
43+
*)
44+
filtered_args+=("$1")
45+
shift
46+
;;
47+
esac
48+
done
49+
50+
printf '%s\n' "${filtered_args[*]}"
51+
}
52+
1653
sbd_supports_ncu_profile() {
1754
case "$1" in
1855
RIKYU|RC_DGXSP|RC_GH200) return 0 ;;
@@ -24,6 +61,7 @@ sbd_configure_ncu_profile_from_run_env() {
2461
local system_name="$1"
2562
local profiler_tool
2663
local profiler_level
64+
local default_profiler_level="detailed"
2765

2866
if ! sbd_supports_ncu_profile "$system_name"; then
2967
return 0
@@ -44,8 +82,20 @@ sbd_configure_ncu_profile_from_run_env() {
4482
if [ -z "${BK_SBD_NCU_PROFILE:-}" ]; then
4583
export BK_SBD_NCU_PROFILE=true
4684
fi
85+
if [ "$system_name" = "RIKYU" ]; then
86+
default_profiler_level="single"
87+
if [ -z "${BK_SBD_NCU_PLAN_TOP_K:-}" ]; then
88+
export BK_SBD_NCU_PLAN_TOP_K=1
89+
fi
90+
if [ -z "${BK_SBD_NCU_PLAN_LAUNCH_COUNT:-}" ]; then
91+
export BK_SBD_NCU_PLAN_LAUNCH_COUNT=1
92+
fi
93+
if [ -z "${BK_SBD_NCU_PROFILE_TIMEOUT_SECONDS:-}" ]; then
94+
export BK_SBD_NCU_PROFILE_TIMEOUT_SECONDS=1800
95+
fi
96+
fi
4797
if [ -z "${BK_SBD_NCU_PROFILER_LEVEL:-}" ]; then
48-
profiler_level=$(bk_resolve_profiler_level detailed SBD_PROFILER_LEVEL)
98+
profiler_level=$(bk_resolve_profiler_level "$default_profiler_level" SBD_PROFILER_LEVEL)
4999
export BK_SBD_NCU_PROFILER_LEVEL="$profiler_level"
50100
fi
51101
}
@@ -280,6 +330,8 @@ sbd_run_rank0_ncu_profile() {
280330
local report_file
281331
local profiler_status
282332
local archive_status
333+
local profile_timeout_seconds
334+
local profile_cmd=()
283335

284336
results_dir=$(sbd_profile_results_dir)
285337
archive_path="${results_dir}/padata_${profile_slug}.tgz"
@@ -298,39 +350,56 @@ sbd_run_rank0_ncu_profile() {
298350

299351
ncu_level_arg_text=$(bk_profiler_ncu_level_args "$profiler_level") || return 1
300352
read -r -a ncu_level_args <<< "$ncu_level_arg_text"
353+
ncu_level_arg_text=$(sbd_strip_ncu_launch_count_args "${ncu_level_args[@]}") || return 1
354+
read -r -a ncu_level_args <<< "$ncu_level_arg_text"
355+
profile_timeout_seconds=$(sbd_ncu_profile_timeout_seconds) || return 1
301356

302357
rm -rf "$raw_dir" "$stage_dir"
303358
mkdir -p "$rep_dir" "$stage_dir/raw" "$stage_dir/reports"
304359

305360
echo "SBD NCU profile: profile='${profile_name}' kernel='${kernel_regex}' skip=${launch_skip} count=${launch_count}" >&2
306361
echo "bk_profiler[ncu]: starting ${rep_name} level=${profiler_level} on rank 0" >&2
307-
set +e
308-
mpirun -np "$n_ranks" bash -lc '
309-
rank=${OMPI_COMM_WORLD_RANK:-${PMIX_RANK:-${SLURM_PROCID:-0}}}
310-
local_rank=${OMPI_COMM_WORLD_LOCAL_RANK:-${SLURM_LOCALID:-0}}
311-
export CUDA_VISIBLE_DEVICES="${local_rank}"
312-
profile_base="$1"
313-
level_argc="$2"
314-
shift 2
315-
level_args=()
316-
i=0
317-
while [ "$i" -lt "$level_argc" ]; do
318-
level_args+=("$1")
319-
shift
320-
i=$((i + 1))
321-
done
322-
kernel_regex="$1"
323-
launch_skip="$2"
324-
launch_count="$3"
325-
shift 3
326-
if [ "$rank" = 0 ]; then
327-
exec ncu -o "$profile_base" --target-processes all "${level_args[@]}" \
328-
--kernel-name-base demangled --kernel-name "$kernel_regex" \
329-
--launch-skip "$launch_skip" --launch-count "$launch_count" ./diag "$@"
362+
if [ "$profile_timeout_seconds" -gt 0 ]; then
363+
echo "SBD NCU profile timeout: ${profile_timeout_seconds}s" >&2
364+
if ! command -v timeout >/dev/null 2>&1; then
365+
echo "SBD NCU profile timeout requested but timeout command is not available." >&2
366+
return 1
330367
fi
331-
exec ./diag "$@"
332-
' bash "$profile_base" "${#ncu_level_args[@]}" "${ncu_level_args[@]}" \
333-
"$kernel_regex" "$launch_skip" "$launch_count" "$@" > "$profile_log" 2>&1
368+
fi
369+
profile_cmd=(
370+
mpirun -np "$n_ranks" bash -lc '
371+
rank=${OMPI_COMM_WORLD_RANK:-${PMIX_RANK:-${SLURM_PROCID:-0}}}
372+
local_rank=${OMPI_COMM_WORLD_LOCAL_RANK:-${SLURM_LOCALID:-0}}
373+
export CUDA_VISIBLE_DEVICES="${local_rank}"
374+
profile_base="$1"
375+
level_argc="$2"
376+
shift 2
377+
level_args=()
378+
i=0
379+
while [ "$i" -lt "$level_argc" ]; do
380+
level_args+=("$1")
381+
shift
382+
i=$((i + 1))
383+
done
384+
kernel_regex="$1"
385+
launch_skip="$2"
386+
launch_count="$3"
387+
shift 3
388+
if [ "$rank" = 0 ]; then
389+
exec ncu -o "$profile_base" --target-processes all "${level_args[@]}" \
390+
--kernel-name-base demangled --kernel-name "$kernel_regex" \
391+
--launch-skip "$launch_skip" --launch-count "$launch_count" ./diag "$@"
392+
fi
393+
exec ./diag "$@"
394+
' bash "$profile_base" "${#ncu_level_args[@]}" "${ncu_level_args[@]}" \
395+
"$kernel_regex" "$launch_skip" "$launch_count" "$@"
396+
)
397+
set +e
398+
if [ "$profile_timeout_seconds" -gt 0 ]; then
399+
timeout --kill-after=60s "$profile_timeout_seconds" "${profile_cmd[@]}" > "$profile_log" 2>&1
400+
else
401+
"${profile_cmd[@]}" > "$profile_log" 2>&1
402+
fi
334403
profiler_status=$?
335404
set -e
336405

@@ -455,6 +524,24 @@ PY
455524
fi
456525
}
457526

527+
sbd_run_optional_ncu_profiles() {
528+
local system_name="$1"
529+
local n_ranks="$2"
530+
local profile_status
531+
shift 2
532+
533+
if ! sbd_ncu_profile_enabled; then
534+
return 0
535+
fi
536+
537+
profile_status=0
538+
sbd_run_configured_ncu_profiles "$system_name" "$n_ranks" "$@" || profile_status=$?
539+
if [ "$profile_status" -ne 0 ]; then
540+
echo "SBD NCU profile acquisition failed with status ${profile_status}; continuing with benchmark result." >&2
541+
fi
542+
return 0
543+
}
544+
458545
sbd_run_configured_ncu_profiles() {
459546
local system_name="$1"
460547
local n_ranks="$2"

‎programs/sbd/run.sh‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -136,11 +136,11 @@ if ! awk -v actual="${energy}" -v reference="${reference_energy}" \
136136
fi
137137

138138
cp diag.log "${RESULTS_DIR}/"
139-
sbd_run_configured_ncu_profiles "${system}" "${n_ranks}" "${diag_args[@]}"
140139
bk_emit_result --fom "${davidson_time}" --fom-unit s \
141140
--fom-version "davidson_internal_s" --exp "${experiment}" \
142141
--nodes "${nodes}" --numproc-node "${numproc_node}" \
143142
--nthreads "${nthreads}" >> "${RESULTS_DIR}/result"
143+
sbd_run_optional_ncu_profiles "${system}" "${n_ranks}" "${diag_args[@]}"
144144
if [[ -n "${mult_time}" ]]; then
145145
bk_emit_section mult "${mult_time}" "" "${SBD_MULT_SECTION_ARTIFACTS:-}" >> "${RESULTS_DIR}/result"
146146
fi

‎scripts/tests/test_sbd_ncu_profile.sh‎

Lines changed: 41 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,8 @@ fi
103103
104104
output=""
105105
kernel=""
106+
launch_count=""
107+
launch_count_seen=0
106108
while [ $# -gt 0 ]; do
107109
case "$1" in
108110
-o|--output)
@@ -113,7 +115,16 @@ while [ $# -gt 0 ]; do
113115
shift
114116
kernel="$1"
115117
;;
116-
--kernel-name-base|--launch-skip|--launch-count|--target-processes|--set)
118+
--launch-count)
119+
launch_count_seen=$((launch_count_seen + 1))
120+
shift
121+
launch_count="$1"
122+
;;
123+
--launch-count=*)
124+
launch_count_seen=$((launch_count_seen + 1))
125+
launch_count="${1#--launch-count=}"
126+
;;
127+
--kernel-name-base|--launch-skip|--target-processes|--set)
117128
shift
118129
;;
119130
--nvtx)
@@ -127,8 +138,11 @@ while [ $# -gt 0 ]; do
127138
done
128139
129140
test -n "$output"
141+
if [ "${FAKE_NCU_FAIL:-}" = "1" ]; then
142+
exit 7
143+
fi
130144
mkdir -p "$(dirname "$output")"
131-
printf '%s\n' "$kernel" >> "$FAKE_NCU_LOG"
145+
printf '%s\tlaunch_count=%s\tlaunch_count_seen=%s\n' "$kernel" "$launch_count" "$launch_count_seen" >> "$FAKE_NCU_LOG"
132146
printf 'fake ncu report\n' > "${output}.ncu-rep"
133147
EOF
134148

@@ -154,8 +168,30 @@ export PATH="${FAKE_BIN}:${PATH}"
154168
test "${BK_SBD_NCU_PROFILE:-}" = "false"
155169
)
156170

171+
(
172+
unset BK_PROFILER BK_PROFILER_LEVEL BK_SBD_NCU_PROFILE BK_SBD_NCU_PROFILER_LEVEL
173+
unset BK_SBD_NCU_PLAN_TOP_K BK_SBD_NCU_PLAN_LAUNCH_COUNT BK_SBD_NCU_PROFILE_TIMEOUT_SECONDS
174+
unset SBD_PROFILER_TOOL
175+
sbd_configure_ncu_profile_from_run_env RIKYU
176+
test "${BK_SBD_NCU_PROFILE}" = "true"
177+
test "${BK_SBD_NCU_PROFILER_LEVEL}" = "single"
178+
test "${BK_SBD_NCU_PLAN_TOP_K}" = "1"
179+
test "${BK_SBD_NCU_PLAN_LAUNCH_COUNT}" = "1"
180+
case "${BK_SBD_NCU_PROFILE_TIMEOUT_SECONDS}" in
181+
''|0|*[!0-9]*) exit 1 ;;
182+
esac
183+
)
184+
185+
test "$(sbd_strip_ncu_launch_count_args --set basic --launch-count 1 --nvtx --launch-count=3)" = "--set basic --nvtx"
186+
187+
(
188+
export BK_SBD_NCU_PROFILE=true
189+
sbd_run_configured_ncu_profiles() { return 7; }
190+
sbd_run_optional_ncu_profiles RIKYU 4
191+
)
192+
157193
pushd "${RUN_DIR}" >/dev/null
158-
unset BK_PROFILER SBD_PROFILER_TOOL BK_SBD_NCU_PROFILE BK_SBD_NCU_PROFILER_LEVEL
194+
unset BK_PROFILER SBD_PROFILER_TOOL BK_SBD_NCU_PROFILE BK_SBD_NCU_PROFILER_LEVEL BK_SBD_NCU_PLAN_LAUNCH_COUNT
159195
export BK_SBD_NCU_PROFILE_MODE=discovery
160196
export BK_SBD_NCU_PLAN_TOP_K=5
161197
export BK_PROFILER_LEVEL=single
@@ -174,6 +210,7 @@ test -f "${RESULTS_DIR}/sbd_ncu_plan.json"
174210

175211
jq -e '
176212
(.profiles | length) == 5 and
213+
([.profiles[].launch_count] | unique) == [1] and
177214
([.profiles[].kernel_match.pattern] | unique | length) == 5 and
178215
.profiles[1].kernel_match.pattern == "regex:.*sbd::MultUnified<double,[[:space:]]*\\(int\\)0,[[:space:]]*\\(int\\)1>.*" and
179216
.profiles[4].kernel_match.pattern == "regex:.*sbd::MultUnified<double,[[:space:]]*\\(int\\)1,[[:space:]]*\\(int\\)0>.*"
@@ -208,5 +245,6 @@ for archive in "${profile_archives[@]}"; do
208245
done
209246

210247
grep -Fq 'regex:.*sbd::MultUnified<double,[[:space:]]*\(int\)0,[[:space:]]*\(int\)1>.*' "$FAKE_NCU_LOG"
248+
awk -F '\t' '$2 != "launch_count=1" || $3 != "launch_count_seen=1" { bad = 1 } END { exit bad }' "$FAKE_NCU_LOG"
211249

212250
echo "SBD NCU profile test passed"

0 commit comments

Comments
 (0)