derek.holloway
7502bf64a4
CodeQL / Analyze (actions) (push) Canceled after 0s
CodeQL / Analyze (python) (push) Canceled after 0s
pre-commit / pre-commit (push) Canceled after 0s
Multi-Arch CI / setup (push) Canceled after 0s
Publish build_manylinux_x86_64 images / publish_build_manylinux_x86_64 (push) Canceled after 0s
Setup (Multi-Arch) / setup (push) Canceled after 0s
Multi-Arch CI / Linux::${{ fromJSON(needs.setup.outputs.linux_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubi10 image / publish_no_rocm_image_ubi10 (push) Canceled after 0s
Multi-Arch CI / Windows::${{ fromJSON(needs.setup.outputs.windows_build_config || '{}').build_variant_label || 'skip' }} (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Multi-Arch CI / Manifest Diff (push) Canceled after 0s
TheRock Manifest Diff Report / Generate Manifest Diff Report (push) Canceled after 0s
Multi-Arch CI / CI Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_ocl_rt images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04_rocgdb images / publish_no_rocm_image_ubuntu24_04_rocgdb (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_ocl_rt (push) Canceled after 0s
Publish TheRock Docker image / build-and-push-image (push) Canceled after 0s
Unit Tests / Unit Tests :: ubuntu-24.04 (push) Canceled after 0s
Unit Tests / Unit Tests :: windows-2022 (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_media (push) Canceled after 0s
Unit Tests / Unit Tests Summary (push) Canceled after 0s
Publish no_rocm_image_ubuntu24_04 images / publish_no_rocm_image_ubuntu24_04_openmpi (push) Canceled after 0s
790 lines
28 KiB
Python
Executable File
790 lines
28 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
"""RCCL device compilation driver.
|
|
|
|
A compiler/linker driver for the RCCL assembly-extract device pipeline.
|
|
Presents a standard compiler CLI to CMake's custom language machinery.
|
|
|
|
Compile mode (--compile):
|
|
rccl-device-compile --compile --arch=gfx942 -DFOO -Ipath -o out.o in.cpp
|
|
1. Compiles .cpp to assembly via amdclang++
|
|
2. Extracts device function + resource JSON (inline)
|
|
3. Assembles extracted .s to .o via amdclang++
|
|
Side effect: writes out.resources.json alongside out.o
|
|
|
|
Link mode (--link):
|
|
rccl-device-compile --link --arch=gfx942 --dispatcher=common.cu.cpp -o device.elf obj1.o obj2.o
|
|
1. Aggregates resource JSON files from all input objects
|
|
2. Compiles dispatcher source to assembly
|
|
3. Patches dispatcher with aggregated resources
|
|
4. Assembles patched dispatcher
|
|
5. Optionally compiles --rocshmem-bitcode=<path>.bc to a device object
|
|
6. Links all objects + dispatcher (+ optional rocSHMEM object) into device.elf via ld.lld
|
|
|
|
Tool discovery:
|
|
amdclang++, ld.lld are found relative to --clang=<path> or by searching
|
|
PATH and the directory containing the system C++ compiler.
|
|
Uses --offload-device-only for device compilation. HIP path flags
|
|
(--hip-path, --hip-device-lib-path) are forwarded from CMake when present.
|
|
"""
|
|
|
|
import argparse
|
|
import json
|
|
import os
|
|
import re
|
|
import subprocess
|
|
import sys
|
|
import tempfile
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Assembly extraction (from extract_device_function.py)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _parse_metadata_resources(lines):
|
|
"""Parse .amdgpu_metadata YAML to extract resource usage."""
|
|
metadata_start = None
|
|
metadata_end = None
|
|
for i, line in enumerate(lines):
|
|
if line.strip() == '.amdgpu_metadata':
|
|
metadata_start = i + 1
|
|
elif line.strip() == '.end_amdgpu_metadata':
|
|
metadata_end = i
|
|
|
|
if metadata_start is None or metadata_end is None:
|
|
return {}
|
|
|
|
yaml_text = ''.join(lines[metadata_start:metadata_end])
|
|
resources = {}
|
|
for key, pattern in [
|
|
('vgpr_count', r'\.vgpr_count:\s+(\d+)'),
|
|
('agpr_count', r'\.agpr_count:\s+(\d+)'),
|
|
('sgpr_count', r'\.sgpr_count:\s+(\d+)'),
|
|
('group_segment_fixed_size', r'\.group_segment_fixed_size:\s+(\d+)'),
|
|
('private_segment_fixed_size', r'\.private_segment_fixed_size:\s+(\d+)'),
|
|
('sgpr_spill_count', r'\.sgpr_spill_count:\s+(\d+)'),
|
|
('vgpr_spill_count', r'\.vgpr_spill_count:\s+(\d+)'),
|
|
]:
|
|
m = re.search(pattern, yaml_text)
|
|
if m:
|
|
resources[key] = int(m.group(1))
|
|
return resources
|
|
|
|
|
|
def _find_devfunc_symbol(lines):
|
|
for line in lines:
|
|
m = re.match(r'\s+\.type\s+(.*ncclDevFunc[^,]+),\s*@function', line)
|
|
if m:
|
|
mangled = m.group(1)
|
|
dm = re.search(r'ncclDevFunc_\w+', mangled)
|
|
unmangled = dm.group(0) if dm else mangled
|
|
return mangled, unmangled
|
|
return None, None
|
|
|
|
|
|
def _find_kernel_symbol(lines):
|
|
for line in lines:
|
|
m = re.match(r'\s+\.globl\s+(.*ncclDevKernel\S+)', line)
|
|
if m:
|
|
return m.group(1).strip()
|
|
return None
|
|
|
|
|
|
def _find_amdhsa_kernel_range(lines):
|
|
amdhsa_start = None
|
|
for i, line in enumerate(lines):
|
|
if line.strip().startswith('.amdhsa_kernel'):
|
|
start = i
|
|
for j in range(i - 1, max(i - 5, -1), -1):
|
|
s = lines[j].strip()
|
|
if s.startswith('.section') and '.rodata' in s:
|
|
start = j
|
|
break
|
|
if s.startswith('.p2align') or s == '':
|
|
start = j
|
|
else:
|
|
break
|
|
amdhsa_start = start
|
|
if line.strip() == '.end_amdhsa_kernel':
|
|
end = i + 1
|
|
if end < len(lines) and lines[end].strip() == '.text':
|
|
end += 1
|
|
return (amdhsa_start, end)
|
|
return None
|
|
|
|
|
|
def _find_metadata_range(lines):
|
|
start = None
|
|
for i, line in enumerate(lines):
|
|
if line.strip() == '.amdgpu_metadata':
|
|
start = i
|
|
elif line.strip() == '.end_amdgpu_metadata':
|
|
if start is not None:
|
|
return (start, i + 1)
|
|
return None
|
|
|
|
|
|
def extract_device_function(asm_lines):
|
|
"""Extract device function from assembly lines.
|
|
Returns (extracted_asm_lines, resources_dict) or raises on error."""
|
|
mangled, unmangled = _find_devfunc_symbol(asm_lines)
|
|
if not mangled:
|
|
raise RuntimeError("No ncclDevFunc found in assembly")
|
|
|
|
kernel_sym = _find_kernel_symbol(asm_lines)
|
|
resources = _parse_metadata_resources(asm_lines)
|
|
resources['function_name'] = unmangled
|
|
resources['mangled_name'] = mangled
|
|
|
|
skip_ranges = []
|
|
amdhsa_range = _find_amdhsa_kernel_range(asm_lines)
|
|
if amdhsa_range:
|
|
skip_ranges.append(amdhsa_range)
|
|
metadata_range = _find_metadata_range(asm_lines)
|
|
if metadata_range:
|
|
skip_ranges.append(metadata_range)
|
|
|
|
def in_skip_range(idx):
|
|
return any(s <= idx < e for s, e in skip_ranges)
|
|
|
|
out_lines = []
|
|
globl_inserted = False
|
|
for i, line in enumerate(asm_lines):
|
|
if in_skip_range(i):
|
|
continue
|
|
stripped = line.strip()
|
|
if '__hip_cuid' in stripped:
|
|
continue
|
|
if kernel_sym and re.match(r'\s*\.globl\s+' + re.escape(kernel_sym), stripped):
|
|
continue
|
|
if not globl_inserted and mangled in line and '.type' in line and '@function' in line:
|
|
out_lines.append(f'\t.globl\t{mangled}\n')
|
|
globl_inserted = True
|
|
out_lines.append(line)
|
|
|
|
return out_lines, resources
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Resource aggregation (from aggregate_resources.py)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
RESOURCE_KEYS = [
|
|
'vgpr_count', 'agpr_count', 'sgpr_count',
|
|
'private_segment_fixed_size', 'group_segment_fixed_size',
|
|
]
|
|
|
|
|
|
def _align_up(val, alignment):
|
|
return ((val + alignment - 1) // alignment) * alignment
|
|
|
|
|
|
def _has_unified_vgpr_agpr(gpu_target):
|
|
return gpu_target in ('gfx90a', 'gfx942', 'gfx950')
|
|
|
|
|
|
def aggregate_resources(resource_list, gpu_target):
|
|
"""Aggregate a list of resource dicts into max values.
|
|
Returns the aggregated dict with derived fields."""
|
|
max_vals = {k: 0 for k in RESOURCE_KEYS}
|
|
max_regular_vgpr = 0
|
|
|
|
for res in resource_list:
|
|
for key in RESOURCE_KEYS:
|
|
max_vals[key] = max(max_vals[key], res.get(key, 0))
|
|
regular = res.get('vgpr_count', 0) - res.get('agpr_count', 0)
|
|
max_regular_vgpr = max(max_regular_vgpr, regular)
|
|
|
|
accum_offset = _align_up(max_regular_vgpr, 4)
|
|
|
|
if _has_unified_vgpr_agpr(gpu_target):
|
|
next_free_vgpr = max(max_vals['vgpr_count'],
|
|
accum_offset + max_vals['agpr_count'])
|
|
else:
|
|
next_free_vgpr = max_vals['vgpr_count']
|
|
|
|
max_vals['accum_offset'] = accum_offset
|
|
max_vals['next_free_vgpr'] = next_free_vgpr
|
|
|
|
_MAX_NUMBERED_SGPR = {9: 102, 10: 106, 11: 106, 12: 106}
|
|
gen = _gfx_generation(gpu_target)
|
|
if gen not in _MAX_NUMBERED_SGPR:
|
|
raise RuntimeError(f"Unknown SGPR limit for generation gfx{gen}xx (target '{gpu_target}')")
|
|
max_vals['next_free_sgpr'] = min(max_vals['sgpr_count'], _MAX_NUMBERED_SGPR[gen])
|
|
|
|
return max_vals
|
|
|
|
|
|
def _gfx_generation(gpu_target):
|
|
m = re.match(r'gfx(\d+)', gpu_target)
|
|
if not m:
|
|
raise RuntimeError(f"Cannot parse GPU target '{gpu_target}'")
|
|
gen = int(m.group(1))
|
|
if gen >= 100:
|
|
gen //= 100
|
|
elif gen >= 10:
|
|
gen //= 10
|
|
return gen
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Dispatcher patching (from patch_dispatcher.py)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def _patch_abs_symbols(lines, max_res):
|
|
result = []
|
|
for line in lines:
|
|
m = re.match(r'(\s*\.set\s+\S+)\.(num_vgpr),\s*(.*)', line)
|
|
if m:
|
|
result.append(f'{m.group(1)}.num_vgpr, {max_res["next_free_vgpr"]}\n')
|
|
continue
|
|
m = re.match(r'(\s*\.set\s+\S+)\.(num_agpr),\s*(.*)', line)
|
|
if m:
|
|
result.append(f'{m.group(1)}.num_agpr, {max_res["agpr_count"]}\n')
|
|
continue
|
|
m = re.match(r'(\s*\.set\s+\S+)\.(numbered_sgpr),\s*(.*)', line)
|
|
if m:
|
|
result.append(f'{m.group(1)}.numbered_sgpr, {max_res["next_free_sgpr"]}\n')
|
|
continue
|
|
m = re.match(r'(\s*\.set\s+\S+)\.(num_named_barrier),\s*(.*)', line)
|
|
if m:
|
|
result.append(f'{m.group(1)}.num_named_barrier, 0\n')
|
|
continue
|
|
m = re.match(r'(\s*\.set\s+\S+)\.(private_seg_size),\s*(.*)', line)
|
|
if m:
|
|
result.append(f'{m.group(1)}.private_seg_size, {max_res["private_segment_fixed_size"]}\n')
|
|
continue
|
|
m = re.match(r'(\s*\.set\s+\S+)\.(uses_flat_scratch),\s*(.*)', line)
|
|
if m:
|
|
result.append(f'{m.group(1)}.uses_flat_scratch, 1\n')
|
|
continue
|
|
m = re.match(r'(\s*\.set\s+\S+)\.(has_dyn_sized_stack),\s*(.*)', line)
|
|
if m:
|
|
result.append(f'{m.group(1)}.has_dyn_sized_stack, 1\n')
|
|
continue
|
|
m = re.match(r'(\s*\.set\s+\S+)\.(has_recursion),\s*(.*)', line)
|
|
if m:
|
|
result.append(f'{m.group(1)}.has_recursion, 1\n')
|
|
continue
|
|
m = re.match(r'(\s*\.set\s+\S+)\.(has_indirect_call),\s*(.*)', line)
|
|
if m:
|
|
result.append(f'{m.group(1)}.has_indirect_call, 1\n')
|
|
continue
|
|
result.append(line)
|
|
return result
|
|
|
|
|
|
def _patch_amdhsa_kernel_directives(lines, max_res):
|
|
result = []
|
|
in_kernel = False
|
|
for line in lines:
|
|
if re.match(r'\s*\.amdhsa_kernel\s+', line):
|
|
in_kernel = True
|
|
result.append(line)
|
|
continue
|
|
if in_kernel and re.match(r'\s*\.end_amdhsa_kernel', line):
|
|
in_kernel = False
|
|
result.append(line)
|
|
continue
|
|
if not in_kernel:
|
|
result.append(line)
|
|
continue
|
|
|
|
m = re.match(r'(\s*)(\.amdhsa_accum_offset)\s+(.*)', line)
|
|
if m:
|
|
result.append(f'{m.group(1)}.amdhsa_accum_offset {max_res["accum_offset"]}\n')
|
|
continue
|
|
m = re.match(r'(\s*)(\.amdhsa_next_free_vgpr)\s+(.*)', line)
|
|
if m:
|
|
result.append(f'{m.group(1)}.amdhsa_next_free_vgpr {max_res["next_free_vgpr"]}\n')
|
|
continue
|
|
m = re.match(r'(\s*)(\.amdhsa_next_free_sgpr)\s+(.*)', line)
|
|
if m:
|
|
result.append(f'{m.group(1)}.amdhsa_next_free_sgpr {max_res["next_free_sgpr"]}\n')
|
|
continue
|
|
m = re.match(r'(\s*)(\.amdhsa_private_segment_fixed_size)\s+(.*)', line)
|
|
if m:
|
|
result.append(f'{m.group(1)}.amdhsa_private_segment_fixed_size {max_res["private_segment_fixed_size"]}\n')
|
|
continue
|
|
m = re.match(r'(\s*)(\.amdhsa_uses_dynamic_stack)\s+(.*)', line)
|
|
if m:
|
|
result.append(f'{m.group(1)}.amdhsa_uses_dynamic_stack 1\n')
|
|
continue
|
|
m = re.match(r'(\s*)(\.amdhsa_enable_private_segment)\s+(.*)', line)
|
|
if m:
|
|
result.append(f'{m.group(1)}.amdhsa_enable_private_segment 1\n')
|
|
continue
|
|
result.append(line)
|
|
return result
|
|
|
|
|
|
def _patch_gpr_maximums(lines, max_res):
|
|
result = []
|
|
for line in lines:
|
|
m = re.match(r'(\s*\.set\s+amdgpu\.max_num_vgpr,\s*)\d+', line)
|
|
if m:
|
|
result.append(f'{m.group(1)}{max_res["next_free_vgpr"]}\n')
|
|
continue
|
|
m = re.match(r'(\s*\.set\s+amdgpu\.max_num_agpr,\s*)\d+', line)
|
|
if m:
|
|
result.append(f'{m.group(1)}{max_res["agpr_count"]}\n')
|
|
continue
|
|
m = re.match(r'(\s*\.set\s+amdgpu\.max_num_sgpr,\s*)\d+', line)
|
|
if m:
|
|
result.append(f'{m.group(1)}{max_res["next_free_sgpr"]}\n')
|
|
continue
|
|
result.append(line)
|
|
return result
|
|
|
|
|
|
def _patch_amdgpu_metadata(lines, max_res):
|
|
result = []
|
|
in_metadata = False
|
|
for line in lines:
|
|
stripped = line.strip()
|
|
if stripped == '.amdgpu_metadata':
|
|
in_metadata = True
|
|
result.append(line)
|
|
continue
|
|
if stripped == '.end_amdgpu_metadata':
|
|
in_metadata = False
|
|
result.append(line)
|
|
continue
|
|
if not in_metadata:
|
|
result.append(line)
|
|
continue
|
|
|
|
m = re.match(r'(\s+\.vgpr_count:\s+)\d+', line)
|
|
if m:
|
|
result.append(f'{m.group(1)}{max_res["next_free_vgpr"]}\n')
|
|
continue
|
|
m = re.match(r'(\s+(?:-\s+)?\.agpr_count:\s+)\d+', line)
|
|
if m:
|
|
result.append(f'{m.group(1)}{max_res["agpr_count"]}\n')
|
|
continue
|
|
m = re.match(r'(\s+\.sgpr_count:\s+)\d+', line)
|
|
if m:
|
|
result.append(f'{m.group(1)}{max_res["sgpr_count"]}\n')
|
|
continue
|
|
m = re.match(r'(\s+\.private_segment_fixed_size:\s+)\d+', line)
|
|
if m:
|
|
result.append(f'{m.group(1)}{max_res["private_segment_fixed_size"]}\n')
|
|
continue
|
|
m = re.match(r'(\s+\.uses_dynamic_stack:\s+)(true|false)', line)
|
|
if m:
|
|
result.append(f'{m.group(1)}true\n')
|
|
continue
|
|
result.append(line)
|
|
return result
|
|
|
|
|
|
def _normalize_file_directives(lines):
|
|
return [re.sub(r'(\s+\.file\s+\d+\s+"[^"]*"\s+"[^"]*")\s+md5\s+0x[0-9a-fA-F]+', r'\1', line)
|
|
for line in lines]
|
|
|
|
|
|
def patch_dispatcher(asm_lines, max_res):
|
|
"""Patch dispatcher assembly with aggregated resource values.
|
|
Returns patched lines."""
|
|
lines = _patch_abs_symbols(asm_lines, max_res)
|
|
lines = _patch_amdhsa_kernel_directives(lines, max_res)
|
|
lines = _patch_gpr_maximums(lines, max_res)
|
|
lines = _patch_amdgpu_metadata(lines, max_res)
|
|
lines = _normalize_file_directives(lines)
|
|
return lines
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Tool discovery
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def find_tool(name, hints, required=True):
|
|
"""Find an executable by name, searching hints first, then PATH."""
|
|
for d in hints:
|
|
if not d:
|
|
continue
|
|
path = os.path.join(d, name)
|
|
if os.path.isfile(path) and os.access(path, os.X_OK):
|
|
return path
|
|
# Fall back to PATH
|
|
import shutil
|
|
path = shutil.which(name)
|
|
if path:
|
|
return path
|
|
if required:
|
|
raise FileNotFoundError(
|
|
f"Required tool '{name}' not found. Searched: {hints} and PATH")
|
|
return None
|
|
|
|
|
|
|
|
def discover_tools(clang_path):
|
|
"""Discover amdclang++ and ld.lld from the given compiler path."""
|
|
if clang_path:
|
|
clang = clang_path
|
|
else:
|
|
import shutil
|
|
clang = shutil.which('amdclang++') or shutil.which('clang++')
|
|
if not clang:
|
|
raise FileNotFoundError("Cannot find amdclang++ or clang++ on PATH")
|
|
|
|
compiler_dir = os.path.dirname(os.path.realpath(clang))
|
|
llvm_bin = os.path.join(compiler_dir, '..', 'lib', 'llvm', 'bin')
|
|
hints = [compiler_dir, llvm_bin]
|
|
|
|
lld = find_tool('ld.lld', hints)
|
|
return clang, lld
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Run subprocess with error handling
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def run(cmd, description=""):
|
|
"""Run a command, raising on failure with the full command for diagnosis."""
|
|
try:
|
|
subprocess.check_call(cmd)
|
|
except subprocess.CalledProcessError as e:
|
|
label = f" ({description})" if description else ""
|
|
print(f"ERROR{label}: command failed with exit code {e.returncode}:",
|
|
file=sys.stderr)
|
|
print(f" {' '.join(cmd)}", file=sys.stderr)
|
|
sys.exit(e.returncode)
|
|
except FileNotFoundError:
|
|
print(f"ERROR: executable not found: {cmd[0]}", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Compile mode: .cpp -> .o + .resources.json
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def do_compile(args, forwarded_flags):
|
|
clang, _ = discover_tools(args.clang)
|
|
arch = args.arch
|
|
source = args.source
|
|
output = args.output
|
|
keep_temps = args.keep_temps
|
|
|
|
base = os.path.splitext(output)[0]
|
|
json_out = base + '.resources.json'
|
|
|
|
# Determine temp directory: alongside output if --keep-temps, else truly temp
|
|
if keep_temps:
|
|
asm_file = base + '.full.s'
|
|
ext_file = base + '.extracted.s'
|
|
cleanup = []
|
|
else:
|
|
tmpdir = tempfile.mkdtemp(prefix='rccl-dc-')
|
|
asm_file = os.path.join(tmpdir, 'full.s')
|
|
ext_file = os.path.join(tmpdir, 'extracted.s')
|
|
cleanup = [asm_file, ext_file, tmpdir]
|
|
|
|
try:
|
|
# Step 1: Compile to assembly
|
|
compile_cmd = [
|
|
clang,
|
|
'-x', 'hip',
|
|
'--offload-device-only',
|
|
f'--offload-arch={arch}',
|
|
'-gline-tables-only',
|
|
'-w',
|
|
] + forwarded_flags + [
|
|
'-S', '-o', asm_file, source,
|
|
]
|
|
run(compile_cmd, f"compile {os.path.basename(source)}")
|
|
|
|
# Step 2: Extract device function (inline)
|
|
with open(asm_file) as f:
|
|
asm_lines = f.readlines()
|
|
|
|
try:
|
|
extracted_lines, resources = extract_device_function(asm_lines)
|
|
except RuntimeError as e:
|
|
print(f"ERROR: extraction failed for {source}: {e}", file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
extracted_lines = _normalize_file_directives(extracted_lines)
|
|
|
|
with open(ext_file, 'w') as f:
|
|
f.writelines(extracted_lines)
|
|
|
|
with open(json_out, 'w') as f:
|
|
json.dump(resources, f, indent=2)
|
|
|
|
# Step 3: Assemble to object
|
|
assemble_cmd = [
|
|
clang,
|
|
'-x', 'assembler',
|
|
'-target', 'amdgcn-amd-amdhsa',
|
|
f'-mcpu={arch}',
|
|
'-c', '-o', output, ext_file,
|
|
]
|
|
run(assemble_cmd, f"assemble {os.path.basename(source)}")
|
|
|
|
finally:
|
|
if not keep_temps:
|
|
for f in cleanup:
|
|
try:
|
|
if os.path.isfile(f):
|
|
os.unlink(f)
|
|
elif os.path.isdir(f):
|
|
os.rmdir(f)
|
|
except OSError:
|
|
pass
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Link mode: objects -> device.elf
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def do_link(args, forwarded_flags):
|
|
clang, lld = discover_tools(args.clang)
|
|
arch = args.arch
|
|
objects = args.objects
|
|
output = args.output
|
|
dispatcher = args.dispatcher
|
|
keep_temps = args.keep_temps
|
|
rocshmem_bitcode = getattr(args, 'rocshmem_bitcode', None)
|
|
|
|
if not dispatcher:
|
|
raise SystemExit("ERROR: --link requires --dispatcher=<source>")
|
|
|
|
# Collect resource JSONs from alongside each .o
|
|
resource_list = []
|
|
for obj in objects:
|
|
json_path = os.path.splitext(obj)[0] + '.resources.json'
|
|
if os.path.exists(json_path):
|
|
with open(json_path) as f:
|
|
resource_list.append(json.load(f))
|
|
|
|
if not resource_list:
|
|
raise SystemExit(f"ERROR: no .resources.json files found alongside input objects")
|
|
|
|
# Step 1: Aggregate resources
|
|
max_res = aggregate_resources(resource_list, arch)
|
|
print(f"[{arch}] Aggregated {len(resource_list)} resources: "
|
|
f"VGPR={max_res['vgpr_count']}, AGPR={max_res['agpr_count']}, "
|
|
f"SGPR={max_res['sgpr_count']}, "
|
|
f"scratch={max_res['private_segment_fixed_size']}")
|
|
|
|
output_dir = os.path.dirname(output) or '.'
|
|
if keep_temps:
|
|
disp_asm = os.path.join(output_dir, f'dispatcher_{arch}.s')
|
|
disp_patched = os.path.join(output_dir, f'dispatcher_{arch}_patched.s')
|
|
disp_obj = os.path.join(output_dir, f'dispatcher_{arch}.o')
|
|
rocshmem_obj = os.path.join(output_dir, f'rocshmem_device_{arch}.o') if rocshmem_bitcode else None
|
|
cleanup = []
|
|
else:
|
|
tmpdir = tempfile.mkdtemp(prefix='rccl-dl-')
|
|
disp_asm = os.path.join(tmpdir, 'dispatcher.s')
|
|
disp_patched = os.path.join(tmpdir, 'dispatcher_patched.s')
|
|
disp_obj = os.path.join(tmpdir, 'dispatcher.o')
|
|
rocshmem_obj = os.path.join(tmpdir, 'rocshmem_device.o') if rocshmem_bitcode else None
|
|
cleanup = [disp_asm, disp_patched, disp_obj, tmpdir]
|
|
if rocshmem_obj:
|
|
cleanup.insert(-1, rocshmem_obj) # remove before rmdir
|
|
|
|
try:
|
|
# Step 2: Compile dispatcher to assembly
|
|
disp_compile_cmd = [
|
|
clang,
|
|
'-x', 'hip',
|
|
'--offload-device-only',
|
|
f'--offload-arch={arch}',
|
|
'-g', '-w',
|
|
] + forwarded_flags + [
|
|
'-S', '-o', disp_asm, dispatcher,
|
|
]
|
|
run(disp_compile_cmd, f"compile dispatcher for {arch}")
|
|
|
|
# Step 3: Patch dispatcher with aggregated resources
|
|
with open(disp_asm) as f:
|
|
disp_lines = f.readlines()
|
|
patched_lines = patch_dispatcher(disp_lines, max_res)
|
|
with open(disp_patched, 'w') as f:
|
|
f.writelines(patched_lines)
|
|
|
|
# Step 4: Assemble patched dispatcher
|
|
disp_assemble_cmd = [
|
|
clang,
|
|
'-x', 'assembler',
|
|
'-target', 'amdgcn-amd-amdhsa',
|
|
f'-mcpu={arch}',
|
|
'-c', '-o', disp_obj, disp_patched,
|
|
]
|
|
run(disp_assemble_cmd, f"assemble dispatcher for {arch}")
|
|
|
|
# Step 5: Compile rocSHMEM device bitcode to an amdgcn object (if provided).
|
|
# rocSHMEM device API symbols (rocshmem::*_wg, rocshmem_n_pes, …) have
|
|
# hidden visibility and cannot be imported from a shared lib — they must
|
|
# be statically present in the device ELF. The bitcode path is supplied
|
|
# by CMake via --rocshmem-bitcode when ENABLE_ROCSHMEM is on.
|
|
if rocshmem_bitcode:
|
|
if not os.path.exists(rocshmem_bitcode):
|
|
print(f" rocSHMEM: no device bitcode for {arch} ({rocshmem_bitcode}), skipping",
|
|
file=sys.stderr)
|
|
rocshmem_obj = None
|
|
else:
|
|
bc_compile_cmd = [
|
|
clang,
|
|
'-target', 'amdgcn-amd-amdhsa',
|
|
f'-mcpu={arch}',
|
|
'-c', '-x', 'ir',
|
|
'-o', rocshmem_obj, rocshmem_bitcode,
|
|
]
|
|
run(bc_compile_cmd, f"compile rocSHMEM device bitcode for {arch}")
|
|
|
|
# Step 6: Link all objects + dispatcher into device.elf
|
|
# Use a response file for potentially long object lists
|
|
rsp_path = os.path.join(output_dir, f'link_{arch}.rsp')
|
|
with open(rsp_path, 'w') as f:
|
|
f.write(disp_obj + '\n')
|
|
for obj in objects:
|
|
f.write(obj + '\n')
|
|
if rocshmem_obj:
|
|
f.write(rocshmem_obj + '\n')
|
|
|
|
link_cmd = [
|
|
lld, '-shared',
|
|
'-o', output,
|
|
f'@{rsp_path}',
|
|
]
|
|
run(link_cmd, f"link device.elf for {arch}")
|
|
|
|
if not keep_temps:
|
|
try:
|
|
os.unlink(rsp_path)
|
|
except OSError:
|
|
pass
|
|
|
|
finally:
|
|
if not keep_temps:
|
|
for f in cleanup:
|
|
try:
|
|
if os.path.isfile(f):
|
|
os.unlink(f)
|
|
elif os.path.isdir(f):
|
|
os.rmdir(f)
|
|
except OSError:
|
|
pass
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Argument parsing
|
|
# ---------------------------------------------------------------------------
|
|
|
|
def parse_compiler_flags(argv):
|
|
"""Separate our flags from forwarded compiler flags (-D, -I, -isystem, -std, -O, etc.)"""
|
|
our_args = []
|
|
forwarded = []
|
|
sources = []
|
|
i = 0
|
|
while i < len(argv):
|
|
arg = argv[i]
|
|
# Our flags
|
|
if arg in ('--compile', '--link', '--keep-temps', '--version'):
|
|
our_args.append(arg)
|
|
elif arg.startswith('--arch='):
|
|
our_args.append(arg)
|
|
elif arg == '--arch' and i + 1 < len(argv):
|
|
our_args.extend([arg, argv[i + 1]])
|
|
i += 1
|
|
elif arg.startswith('--clang='):
|
|
our_args.append(arg)
|
|
elif arg == '--clang' and i + 1 < len(argv):
|
|
our_args.extend([arg, argv[i + 1]])
|
|
i += 1
|
|
elif arg.startswith('--dispatcher='):
|
|
our_args.append(arg)
|
|
elif arg == '--dispatcher' and i + 1 < len(argv):
|
|
our_args.extend([arg, argv[i + 1]])
|
|
i += 1
|
|
elif arg.startswith('--rocshmem-bitcode='):
|
|
our_args.append(arg)
|
|
elif arg == '--rocshmem-bitcode' and i + 1 < len(argv):
|
|
our_args.extend([arg, argv[i + 1]])
|
|
i += 1
|
|
elif arg == '-o' and i + 1 < len(argv):
|
|
our_args.extend(['-o', argv[i + 1]])
|
|
i += 1
|
|
# Forwarded compiler flags
|
|
elif arg.startswith('-D') or arg.startswith('-I') or arg.startswith('-isystem'):
|
|
forwarded.append(arg)
|
|
elif arg.startswith('-std=') or arg.startswith('-O') or arg == '-g':
|
|
forwarded.append(arg)
|
|
elif arg == '-D' and i + 1 < len(argv):
|
|
forwarded.extend(['-D', argv[i + 1]])
|
|
i += 1
|
|
elif arg == '-I' and i + 1 < len(argv):
|
|
forwarded.extend(['-I', argv[i + 1]])
|
|
i += 1
|
|
elif arg.startswith('-'):
|
|
forwarded.append(arg)
|
|
else:
|
|
# Positional: source file or object file
|
|
sources.append(arg)
|
|
i += 1
|
|
return our_args, forwarded, sources
|
|
|
|
|
|
def expand_response_files(argv):
|
|
"""Expand @file arguments: replace with one-per-line entries from file."""
|
|
expanded = []
|
|
for arg in argv:
|
|
if arg.startswith('@') and os.path.isfile(arg[1:]):
|
|
with open(arg[1:]) as f:
|
|
expanded.extend(line.strip() for line in f if line.strip())
|
|
else:
|
|
expanded.append(arg)
|
|
return expanded
|
|
|
|
|
|
def main():
|
|
raw_args, forwarded_flags, positional = parse_compiler_flags(
|
|
expand_response_files(sys.argv[1:]))
|
|
|
|
parser = argparse.ArgumentParser(prog='rccl-device-compile', add_help=False)
|
|
parser.add_argument('--compile', action='store_true')
|
|
parser.add_argument('--link', action='store_true')
|
|
parser.add_argument('--arch', required=False)
|
|
parser.add_argument('--clang', default=None)
|
|
parser.add_argument('--dispatcher', default=None)
|
|
parser.add_argument('--rocshmem-bitcode', default=None, dest='rocshmem_bitcode')
|
|
parser.add_argument('-o', '--output', required=False)
|
|
parser.add_argument('--keep-temps', action='store_true')
|
|
parser.add_argument('--version', action='store_true')
|
|
|
|
args = parser.parse_args(raw_args)
|
|
|
|
if args.version:
|
|
print("rccl-device-compile 1.0")
|
|
sys.exit(0)
|
|
|
|
if args.compile:
|
|
if not positional:
|
|
raise SystemExit("ERROR: --compile requires a source file")
|
|
if not args.arch:
|
|
raise SystemExit("ERROR: --compile requires --arch=<target>")
|
|
if not args.output:
|
|
raise SystemExit("ERROR: --compile requires -o <output>")
|
|
args.source = positional[0]
|
|
do_compile(args, forwarded_flags)
|
|
|
|
elif args.link:
|
|
if not positional:
|
|
raise SystemExit("ERROR: --link requires object files")
|
|
if not args.arch:
|
|
raise SystemExit("ERROR: --link requires --arch=<target>")
|
|
if not args.output:
|
|
raise SystemExit("ERROR: --link requires -o <output>")
|
|
args.objects = positional
|
|
do_link(args, forwarded_flags)
|
|
|
|
else:
|
|
raise SystemExit("ERROR: specify --compile or --link mode")
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|