From 9149eea87cf2f450f5bb71615689e375df0f1009 Mon Sep 17 00:00:00 2001 From: Brandon Ros Date: Mon, 14 Sep 2026 17:52:48 -0400 Subject: [PATCH] Add verified LLVM cleanup passes with retention and replay regressions --- .github/workflows/ptx_export.yml | 8 + Cargo.lock | 7 + Cargo.toml | 1 + crates/cuda_builder/src/lib.rs | 57 ++++++ crates/rustc_codegen_nvvm/build.rs | 6 +- .../rustc_llvm_wrapper/PassWrapper.cpp | 73 ++++++++ crates/rustc_codegen_nvvm/src/back.rs | 41 +++++ crates/rustc_codegen_nvvm/src/context.rs | 30 ++++ crates/rustc_codegen_nvvm/src/llvm.rs | 11 ++ crates/rustc_codegen_nvvm/src/nvvm.rs | 46 +++-- examples/ptx_export/README.md | 12 ++ examples/ptx_export/check_cleanup_ir.py | 98 ++++++++++ examples/ptx_export/check_default_dce.py | 84 +++++++++ .../ptx_export/check_integrated_cleanup.py | 67 +++++++ examples/ptx_export/cleanup_ir_oracle.c | 30 ++++ examples/ptx_export/inspect_codegen.py | 167 ++++++++++++++++++ .../ptx_export/kernels/src/guarded_select.rs | 112 ++++++++++++ examples/ptx_export/kernels/src/lib.rs | 56 ++++++ examples/ptx_export/optimization_pipelines.py | 48 +++++ examples/ptx_export/replay_cleanup.py | 139 +++++++++++++++ .../ptx_export/retention-kernels/Cargo.toml | 11 ++ .../ptx_export/retention-kernels/src/lib.rs | 70 ++++++++ examples/ptx_export/src/main.rs | 48 ++++- examples/ptx_export/test_inspect_codegen.py | 67 +++++++ examples/ptx_export/test_inspection_reuse.py | 34 ++++ 25 files changed, 1308 insertions(+), 15 deletions(-) create mode 100644 examples/ptx_export/check_cleanup_ir.py create mode 100644 examples/ptx_export/check_default_dce.py create mode 100644 examples/ptx_export/check_integrated_cleanup.py create mode 100644 examples/ptx_export/cleanup_ir_oracle.c create mode 100644 examples/ptx_export/inspect_codegen.py create mode 100644 examples/ptx_export/kernels/src/guarded_select.rs create mode 100644 examples/ptx_export/optimization_pipelines.py create mode 100644 examples/ptx_export/replay_cleanup.py create mode 100644 examples/ptx_export/retention-kernels/Cargo.toml create mode 100644 examples/ptx_export/retention-kernels/src/lib.rs create mode 100644 examples/ptx_export/test_inspect_codegen.py create mode 100644 examples/ptx_export/test_inspection_reuse.py diff --git a/.github/workflows/ptx_export.yml b/.github/workflows/ptx_export.yml index d2363406..4b57305b 100644 --- a/.github/workflows/ptx_export.yml +++ b/.github/workflows/ptx_export.yml @@ -24,6 +24,14 @@ jobs: git rev-parse HEAD > artifacts/ptx/source-commit.txt nix develop .#v21 --command rustc -Vv > artifacts/ptx/rustc-version.txt sha256sum artifacts/ptx/rust_kernels.ptx > artifacts/ptx/SHA256SUMS + - name: Check Python regressions + run: python3 -B -m unittest discover -s examples/ptx_export -p 'test_*.py' + - name: Check guarded-select semantics + run: | + nix develop .#v21 --command rustc --edition=2024 --test examples/ptx_export/kernels/src/guarded_select.rs -o /tmp/guarded-select + /tmp/guarded-select + - name: Check default DCE retention + run: nix develop .#v21 --command python3 examples/ptx_export/check_default_dce.py artifacts/ptx - uses: actions/upload-artifact@v4 with: name: rust-ptx diff --git a/Cargo.lock b/Cargo.lock index 5e53c7b5..a4fa73c7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2592,6 +2592,13 @@ dependencies = [ "sha2", ] +[[package]] +name = "ptx-retention-kernels" +version = "0.1.0" +dependencies = [ + "cuda_std", +] + [[package]] name = "ptx_compiler" version = "0.1.1" diff --git a/Cargo.toml b/Cargo.toml index 77ceaa42..d872f361 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -53,6 +53,7 @@ members = [ "examples/vecadd/kernels", "examples/ptx_export", "examples/ptx_export/kernels", + "examples/ptx_export/retention-kernels", "samples/introduction/async_api", "samples/introduction/async_api/kernels", diff --git a/crates/cuda_builder/src/lib.rs b/crates/cuda_builder/src/lib.rs index 17a4b893..6d5962a5 100644 --- a/crates/cuda_builder/src/lib.rs +++ b/crates/cuda_builder/src/lib.rs @@ -51,6 +51,18 @@ impl DebugInfo { } } +/// Experimental pre-NVVM optimization. Requires the LLVM 21 backend. +/// The historical LLVM 19 API names are retained for caller compatibility. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Llvm19Cleanup { + /// Remove unreachable internal definitions without rewriting live function bodies. + GlobalDce, + /// Inline internal calls, then run scalar cleanup without correlated propagation. + InlineScalar, + Scalar, + Inline, +} + pub enum EmitOption { LlvmIr, Bitcode, @@ -194,6 +206,12 @@ pub struct CudaBuilder { /// An optional path where to dump LLVM IR of the final output the codegen will feed to libnvvm. Usually /// used for debugging. pub final_module_path: Option, + /// Whether the modern backend removes unreachable definitions at the merged handoff. + pub llvm19_global_dce: bool, + /// Additional opt-in modern LLVM cleanup; disabled by default. + pub llvm19_cleanup: Option, + /// Experimental scalar cleanup of each codegen unit before serialization. + pub llvm19_module_cleanup: bool, } impl CudaBuilder { @@ -216,9 +234,33 @@ impl CudaBuilder { debug: DebugInfo::None, build_args: vec![], final_module_path: None, + llvm19_global_dce: true, + llvm19_cleanup: None, + llvm19_module_cleanup: false, } } + /// Enable or disable the default modern LLVM merged-module GlobalDCE pass. + /// Disabling is intended for compiler-output comparisons; LLVM 7 is unchanged. + pub fn llvm19_global_dce(mut self, enabled: bool) -> Self { + self.llvm19_global_dce = enabled; + self + } + + /// Enable verified scalar cleanup before each codegen unit is serialized. + /// Disabled by default; independent of merged-module cleanup. + pub fn llvm19_module_cleanup(mut self, enabled: bool) -> Self { + self.llvm19_module_cleanup = enabled; + self + } + + /// Enable a bounded modern LLVM cleanup pipeline before NVVM compilation. + /// This is experimental; compare numerical results and generated code. + pub fn llvm19_cleanup(mut self, cleanup: Llvm19Cleanup) -> Self { + self.llvm19_cleanup = Some(cleanup); + self + } + /// Additional arguments passed to cargo during `cargo build`. pub fn build_args(mut self, args: &[impl AsRef]) -> Self { self.build_args @@ -723,6 +765,21 @@ fn invoke_rustc(builder: &CudaBuilder) -> Result { } let mut llvm_args = vec![NvvmOption::Arch(builder.arch).to_string()]; + if !builder.llvm19_global_dce { + llvm_args.push("--disable-llvm19-global-dce".to_string()); + } + if builder.llvm19_module_cleanup { + llvm_args.push("--llvm19-module-cleanup".to_string()); + } + if let Some(mode) = builder.llvm19_cleanup { + let mode = match mode { + Llvm19Cleanup::GlobalDce => "dce", + Llvm19Cleanup::InlineScalar => "inline-scalar", + Llvm19Cleanup::Scalar => "scalar", + Llvm19Cleanup::Inline => "inline", + }; + llvm_args.push(format!("--llvm19-cleanup={mode}")); + } if !builder.nvvm_opts { llvm_args.push("-opt=0".to_string()); diff --git a/crates/rustc_codegen_nvvm/build.rs b/crates/rustc_codegen_nvvm/build.rs index aed3a7c8..e7654546 100644 --- a/crates/rustc_codegen_nvvm/build.rs +++ b/crates/rustc_codegen_nvvm/build.rs @@ -306,7 +306,11 @@ fn rustc_llvm_build(flavor: &LlvmFlavor) { configure_libintrinsics(&llvm_config, flavor); - let required_components = &["ipo", "bitreader", "bitwriter", "lto", "nvptx"]; + let required_components: &[&str] = if flavor.major >= 19 { + &["ipo", "bitreader", "bitwriter", "lto", "nvptx", "passes"] + } else { + &["ipo", "bitreader", "bitwriter", "lto", "nvptx"] + }; let components = output(Command::new(&llvm_config).arg("--components")); let mut components = components.split_whitespace().collect::>(); diff --git a/crates/rustc_codegen_nvvm/rustc_llvm_wrapper/PassWrapper.cpp b/crates/rustc_codegen_nvvm/rustc_llvm_wrapper/PassWrapper.cpp index d6c67a38..fa9aa23a 100644 --- a/crates/rustc_codegen_nvvm/rustc_llvm_wrapper/PassWrapper.cpp +++ b/crates/rustc_codegen_nvvm/rustc_llvm_wrapper/PassWrapper.cpp @@ -13,6 +13,7 @@ #include #include #include +#include #include #include "rustllvm.h" @@ -25,6 +26,9 @@ #include "llvm/Support/FileSystem.h" #if LLVM_VERSION_MAJOR >= 19 #include "llvm/Transforms/IPO/Internalize.h" +#include "llvm/Passes/PassBuilder.h" +#include "llvm/IR/Verifier.h" +#include "llvm/Support/Error.h" #endif #if LLVM_VERSION_MAJOR >= 19 #include "llvm/TargetParser/Host.h" @@ -166,6 +170,75 @@ extern "C" void LLVMPassManagerBuilderPopulateLTOPassManager( } #endif +// Explicit, bounded modern-PM cleanup at the final NVVM handoff. This is +// separate from the legacy compatibility builder and remains opt-in. +// Keep discriminants in sync with llvm::NvvmCleanup on the Rust side. +enum class LLVMRustNvvmCleanup : uint32_t { Scalar = 0, Inline = 1, GlobalDce = 2, InlineScalar = 3 }; + +extern "C" LLVMRustResult LLVMRustRunNvvmCleanup(LLVMModuleRef M, LLVMRustNvvmCleanup Mode) +{ +#if LLVM_VERSION_MAJOR >= 19 + Module &Mod = *unwrap(M); + std::string Diagnostics; + raw_string_ostream OS(Diagnostics); + if (verifyModule(Mod, &OS)) { + LLVMRustSetLastError(OS.str().c_str()); + return LLVMRustResult::Failure; + } + LoopAnalysisManager LAM; + FunctionAnalysisManager FAM; + CGSCCAnalysisManager CGAM; + ModuleAnalysisManager MAM; + // Match opt's target-aware analyses. Without a TargetMachine the inliner + // uses generic costs and can disagree with replay even on identical IR. + // Use the module's NVPTX triple and generic CPU, as the existing backend + // does; NVVM remains responsible for the selected compute architecture. + std::string TargetError; + const Target *T = TargetRegistry::lookupTarget(Mod.getTargetTriple(), TargetError); + if (!T) { + LLVMRustSetLastError(TargetError.c_str()); + return LLVMRustResult::Failure; + } + std::unique_ptr TM(T->createTargetMachine( + Mod.getTargetTriple(), "", "", TargetOptions(), std::nullopt)); + if (!TM) { + LLVMRustSetLastError("Could not create cleanup TargetMachine"); + return LLVMRustResult::Failure; + } + PassBuilder PB(TM.get()); + PB.registerModuleAnalyses(MAM); + PB.registerCGSCCAnalyses(CGAM); + PB.registerFunctionAnalyses(FAM); + PB.registerLoopAnalyses(LAM); + PB.crossRegisterProxies(LAM, FAM, CGAM, MAM); + ModulePassManager PM; + // Prune unreachable functions before inlining, then expose and simplify + // branch-correlated iterator values. Keep this identical to offline replay. + const char *Pipeline = Mode == LLVMRustNvvmCleanup::GlobalDce + ? "globaldce,verify" + : Mode == LLVMRustNvvmCleanup::InlineScalar + ? "globaldce,cgscc(inline),function(sroa,instcombine,simplifycfg,adce),globaldce,verify" + : Mode == LLVMRustNvvmCleanup::Inline + ? "globaldce,cgscc(inline),function(sroa,instcombine,simplifycfg,adce)," + "globaldce,function(correlated-propagation,instcombine,simplifycfg,adce),verify" + : "function(sroa,instcombine,simplifycfg,adce),verify"; + if (auto Error = PB.parsePassPipeline(PM, Pipeline)) { + LLVMRustSetLastError(toString(std::move(Error)).c_str()); + return LLVMRustResult::Failure; + } + PM.run(Mod, MAM); + Diagnostics.clear(); + if (verifyModule(Mod, &OS)) { + LLVMRustSetLastError(OS.str().c_str()); + return LLVMRustResult::Failure; + } + return LLVMRustResult::Success; +#else + LLVMRustSetLastError("NVVM cleanup requires LLVM 19"); + return LLVMRustResult::Failure; +#endif +} + extern "C" void LLVMInitializePasses() { #if LLVM_VERSION_MAJOR >= 19 diff --git a/crates/rustc_codegen_nvvm/src/back.rs b/crates/rustc_codegen_nvvm/src/back.rs index a3ce18f1..a7997328 100644 --- a/crates/rustc_codegen_nvvm/src/back.rs +++ b/crates/rustc_codegen_nvvm/src/back.rs @@ -328,6 +328,47 @@ pub fn compile_codegen_unit(tcx: TyCtxt<'_>, cgu_name: Symbol) -> (ModuleCodegen } } + // Run only after definitions, used globals and debug metadata are + // finalized. Scalar cleanup preserves cross-module linkage; no DCE or + // inlining is performed here. Merged-module cleanup remains independent. + let args = crate::context::CodegenArgs::from_session(tcx.sess); + if args.llvm19_module_cleanup { + let llmod = unsafe { &*llvm_module.llmod }; + let dump = |stage: &str| { + if let Some(final_path) = &args.final_module_path { + let directory = final_path.parent().unwrap().join("per-module"); + std::fs::create_dir_all(&directory).unwrap_or_else(|error| { + tcx.sess + .dcx() + .fatal(format!("cannot create module IR directory: {error}")) + }); + let path = directory.join(format!("{cgu_name}.{stage}.ll")); + let path = path.to_str().unwrap(); + unsafe { + llvm::LLVMRustPrintModule( + llmod, + path.as_c_char_ptr(), + path.len(), + demangle_callback, + ) + .into_result() + .unwrap_or_else(|_| { + llvm_err(tcx.sess.dcx(), "cannot save per-module IR").raise(); + }); + } + } + }; + dump("before"); + unsafe { + llvm::LLVMRustRunNvvmCleanup(llmod, llvm::NvvmCleanup::Scalar) + .into_result() + .unwrap_or_else(|_| { + llvm_err(tcx.sess.dcx(), "LLVM 19 per-module cleanup failed").raise(); + }); + } + dump("after"); + } + ModuleCodegen::new_regular(cgu_name.to_string(), llvm_module) } diff --git a/crates/rustc_codegen_nvvm/src/context.rs b/crates/rustc_codegen_nvvm/src/context.rs index 4087f7a7..9d5d9628 100644 --- a/crates/rustc_codegen_nvvm/src/context.rs +++ b/crates/rustc_codegen_nvvm/src/context.rs @@ -651,6 +651,10 @@ pub struct CodegenArgs { pub override_libm: bool, pub use_constant_memory_space: bool, pub final_module_path: Option, + // None leaves the existing NVVM handoff unchanged. + pub llvm19_cleanup: Option, + pub llvm19_module_cleanup: bool, + pub disable_llvm19_global_dce: bool, pub disassemble: Option, } @@ -675,6 +679,32 @@ impl CodegenArgs { cg_args.override_libm = true; } else if arg == "--use-constant-memory-space" { cg_args.use_constant_memory_space = true; + } else if arg == "--disable-llvm19-global-dce" { + if !cfg!(feature = "llvm21") { + sess.dcx() + .fatal("--disable-llvm19-global-dce requires the llvm21 backend feature"); + } + cg_args.disable_llvm19_global_dce = true; + } else if arg == "--llvm19-module-cleanup" { + if !cfg!(feature = "llvm21") { + sess.dcx() + .fatal("--llvm19-module-cleanup requires the llvm21 backend feature"); + } + cg_args.llvm19_module_cleanup = true; + } else if let Some(mode) = arg.strip_prefix("--llvm19-cleanup=") { + if !cfg!(feature = "llvm21") { + sess.dcx() + .fatal("--llvm19-cleanup requires the llvm21 backend feature"); + } + cg_args.llvm19_cleanup = Some(match mode { + "scalar" => crate::llvm::NvvmCleanup::Scalar, + "inline" => crate::llvm::NvvmCleanup::Inline, + "dce" => crate::llvm::NvvmCleanup::GlobalDce, + "inline-scalar" => crate::llvm::NvvmCleanup::InlineScalar, + _ => sess + .dcx() + .fatal("--llvm19-cleanup expects scalar, inline, inline-scalar, or dce"), + }); } else if arg == "--final-module-path" { let path = match args.get(idx + 1) { Some(p) => p, diff --git a/crates/rustc_codegen_nvvm/src/llvm.rs b/crates/rustc_codegen_nvvm/src/llvm.rs index a51641eb..5160079f 100644 --- a/crates/rustc_codegen_nvvm/src/llvm.rs +++ b/crates/rustc_codegen_nvvm/src/llvm.rs @@ -28,6 +28,16 @@ use std::ptr::{self}; use crate::{builder::unnamed, common::AsCCharPtr}; pub use debuginfo::*; +// Keep discriminants in sync with LLVMRustNvvmCleanup in PassWrapper.cpp. +#[repr(u32)] +#[derive(Clone, Copy)] +pub enum NvvmCleanup { + Scalar = 0, + Inline = 1, + GlobalDce = 2, + InlineScalar = 3, +} + impl PartialEq for Value { fn eq(&self, other: &Self) -> bool { ptr::eq(self, other) @@ -1331,6 +1341,7 @@ unsafe extern "C" { ) -> Option<&'a DILocation>; pub(crate) fn LLVMRustRunFunctionPassManager(PM: &PassManager, M: &Module); + pub(crate) fn LLVMRustRunNvvmCleanup(M: &Module, mode: NvvmCleanup) -> LLVMRustResult; pub(crate) fn LLVMRustAddAlwaysInlinePass(P: &PassManagerBuilder, AddLifetimes: bool); pub(crate) fn LLVMRustAddBuilderLibraryInfo( diff --git a/crates/rustc_codegen_nvvm/src/nvvm.rs b/crates/rustc_codegen_nvvm/src/nvvm.rs index 83f66fe5..1237deb8 100644 --- a/crates/rustc_codegen_nvvm/src/nvvm.rs +++ b/crates/rustc_codegen_nvvm/src/nvvm.rs @@ -95,6 +95,7 @@ pub fn codegen_bitcode_modules( unsafe { LLVMRustRestoreNvvmKernelAnnotations(module); internalize_pass(module, llcx); + #[cfg(not(feature = "llvm21"))] dce_pass(module); if sess.opts.debuginfo != DebugInfo::None { @@ -115,6 +116,39 @@ pub fn codegen_bitcode_modules( LLVMAddNamedMetadataOperand(module, c"nvvmir.version".as_ptr().cast(), node); + // Inline pipelines already contain GlobalDCE; avoid running it twice. + // Keep the default pass at the verified handoff, after debug/IR metadata + // is finalized, exactly where the opt-in implementation was validated. + let run_default_dce = cfg!(feature = "llvm21") + && !args.disable_llvm19_global_dce + && matches!(args.llvm19_cleanup, None | Some(NvvmCleanup::Scalar)); + if run_default_dce || args.llvm19_cleanup.is_some() { + if let Some(path) = &args.final_module_path { + let before = path.with_extension("before-cleanup.ll"); + let before = before.to_str().unwrap(); + LLVMRustPrintModule( + module, + before.as_c_char_ptr(), + before.len(), + demangle_callback, + ) + .into_result() + .expect("failed to write pre-cleanup LLVM IR"); + } + let modes = run_default_dce + .then_some(NvvmCleanup::GlobalDce) + .into_iter() + .chain(args.llvm19_cleanup); + for mode in modes { + if LLVMRustRunNvvmCleanup(module, mode).into_result().is_err() { + sess.dcx().fatal(format!( + "LLVM 19 cleanup failed: {}", + crate::llvm::last_error().unwrap_or_else(|| "unknown LLVM error".into()) + )); + } + } + } + if let Some(path) = &args.final_module_path { let out = path.to_str().unwrap(); let result = @@ -355,18 +389,8 @@ unsafe fn internalize_pass(module: &Module, cx: &Context) { } } +#[cfg(not(feature = "llvm21"))] unsafe fn dce_pass(module: &Module) { - #[cfg(feature = "llvm21")] - { - // The legacy C API entrypoint used below (`LLVMAddGlobalDCEPass`) is not - // available on our current LLVM 21 runtime path. Keep the backend loadable - // by skipping this cleanup for now; revisit if LLVM 21 smoke tests show we - // need an explicit replacement pass. - let _ = module; - return; - } - - #[cfg(not(feature = "llvm21"))] unsafe { let pass_manager = LLVMCreatePassManager(); diff --git a/examples/ptx_export/README.md b/examples/ptx_export/README.md index e227dde7..99fee586 100644 --- a/examples/ptx_export/README.md +++ b/examples/ptx_export/README.md @@ -14,3 +14,15 @@ The kernels accept explicit pointer/count arguments. SHA-256 reads and writes 32 bytes per work item. Input and output buffers must not overlap. Export success does not establish another PTX consumer's numerical correctness. + +The exporter also includes guarded-select regression kernels. Modern LLVM +GlobalDCE runs by default; use the second argument `none` to disable it. +Experimental modes are `dce`, `scalar`, `inline-scalar`, `inline`, +`module-scalar`, `module-inline`, `size-s`, and `size-z`. Optional third and +fourth arguments select a kernel crate and its features. + +Historical `llvm19_*` builder settings control the LLVM 21 backend on this +stack. These names and pipelines are preserved from the original branch. +The DCE retention check covers used globals, external functions, initialized +data, and unreachable negative controls. Replay checks compare integrated +cleanup with standalone LLVM processing. GPU runtime validation is separate. diff --git a/examples/ptx_export/check_cleanup_ir.py b/examples/ptx_export/check_cleanup_ir.py new file mode 100644 index 00000000..8888b065 --- /dev/null +++ b/examples/ptx_export/check_cleanup_ir.py @@ -0,0 +1,98 @@ +#!/usr/bin/env python3 +"""Check extracted integer-only LLVM helpers on the host, not on a GPU. + +llvm-extract preserves their bodies and makes their linkage usable by the C +oracle. Only the target triple/data layout are changed in the host copy. +This supplements, and never substitutes for, validation of NVIDIA PTX/SASS. +""" +import argparse +import hashlib +import json +from pathlib import Path +import re +import shutil +import subprocess + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('ir', type=Path) + parser.add_argument('--out', type=Path, required=True) + parser.add_argument('--llvm-bin', type=Path) + parser.add_argument('--cc', default='cc') + args = parser.parse_args() + out = args.out.resolve(); out.mkdir(parents=True, exist_ok=True) + llvm = args.llvm_bin or Path(shutil.which('llvm-extract') or '').parent + commands = [] + + def run(command, filename): + commands.append([str(x) for x in command]) + (out/'commands.json').write_text(json.dumps(commands, indent=2)+'\n') + with (out/filename).open('w') as log: + subprocess.run(command, stdout=log, stderr=subprocess.STDOUT, check=True) + return (out/filename).read_text() + + version = run([llvm/'llc', '--version'], 'llvm-version.txt') + if not re.search(r'LLVM version 21\.', version): raise RuntimeError('requires LLVM 21') + triple = re.search(r'Default target:\s*(\S+)', version)[1] + run([args.cc, '--version'], 'cc-version.txt') + source = args.ir.read_text() + symbols = {} + for name in ('filtered', 'stepped', 'preserved', 'direct', 'observed'): + found = re.findall(r'^define [^\n]*@([\w]+guarded_select\d+'+name+r')\(', source, re.M) + if len(found) != 1: raise RuntimeError(f'expected one {name} definition') + symbols[name] = found[0] + # --recursive follows function calls, but does not retain referenced global + # initializers. Before scalar cleanup, panic paths still reference source + # locations. Include those globals and their transitive data dependencies + # unchanged instead of substituting dummy definitions or dropping paths. + globals_to_keep = set() + while True: + run([llvm/'llvm-extract', *['--func='+s for s in symbols.values()], + *['--glob='+s for s in sorted(globals_to_keep)], '--recursive', + '-S', args.ir.resolve(), '-o', out/'extracted-gpu.ll'], 'extract.log') + extracted = (out/'extracted-gpu.ll').read_text() + external = re.findall(r'^@("[^"\\]*"|[-\w.$]+) = external ', extracted, re.M) + if not external: + break + names = {s.strip('"') for s in external} + if names <= globals_to_keep: + raise RuntimeError(f'host extraction contains unresolved globals: {sorted(names)}') + globals_to_keep.update(names) + # These five helpers use only integer arithmetic and ordinary table loads. + # Reject a changed reproducer that acquires GPU-specific or external calls. + declarations = re.findall(r'^declare [^\n]*@([^ (]+)\(', extracted, re.M) + # DCE-only retains ordinary lifetime markers and assumptions that scalar + # cleanup removes. Preserve their semantics in the host copy; these are + # target-independent LLVM intrinsics, not GPU operations or external calls. + allowed = {'llvm.trap', 'llvm.umin.i32', 'llvm.umin.i64', 'llvm.assume', 'llvm.expect.i1', + 'llvm.lifetime.start.p0', 'llvm.lifetime.end.p0', + 'llvm.experimental.noalias.scope.decl'} + if any(s not in allowed for s in declarations): + raise RuntimeError(f'host extraction contains unexpected declarations: {declarations}') + host = re.sub(r'^target datalayout = .*$', 'target datalayout = ""', extracted, flags=re.M) + host = re.sub(r'^target triple = .*$', f'target triple = "{triple}"', host, flags=re.M) + (out/'host.ll').write_text(host) + run([llvm/'opt', '-passes=verify', '-disable-output', out/'host.ll'], 'verify.log') + # Linux's host compiler links PIE by default. Referenced source-location + # data requires PIC relocations; this affects only the host oracle object. + run([llvm/'llc', '-relocation-model=pic', '-filetype=obj', out/'host.ll', '-o', out/'helpers.o'], 'llc.log') + header = [] + for name, symbol in symbols.items(): + parameters = 'const uint64_t *, uint32_t' + if name in ('preserved', 'observed'): parameters += ', uint64_t' + header += [f'#define {name} {symbol}', f'extern uint64_t {name}({parameters});'] + (out/'helpers.h').write_text('\n'.join(header)+'\n') + oracle = Path(__file__).with_name('cleanup_ir_oracle.c').resolve() + run([args.cc, '-I', out, oracle, out/'helpers.o', '-o', out/'oracle'], 'link.log') + output = run([out/'oracle'], 'numerical.log') + if 'HOST_IR_NUMERICAL_PASS: 1608 cases' not in output: raise RuntimeError('missing numerical result') + result = {'input_sha256': hashlib.sha256(args.ir.read_bytes()).hexdigest(), + 'oracle_sha256': hashlib.sha256(oracle.read_bytes()).hexdigest(), + 'host_triple': triple, 'gpu_execution': False, 'result': output.strip()} + (out/'result.json').write_text(json.dumps(result, indent=2)+'\n') + print(output, end='') + + +if __name__ == '__main__': + main() diff --git a/examples/ptx_export/check_default_dce.py b/examples/ptx_export/check_default_dce.py new file mode 100644 index 00000000..dd7a1909 --- /dev/null +++ b/examples/ptx_export/check_default_dce.py @@ -0,0 +1,84 @@ +#!/usr/bin/env python3 +"""Validate default GlobalDCE through Rust codegen, internalization and NVVM.""" +import argparse +import hashlib +import json +from pathlib import Path +import re +import subprocess +import sys +from replay_cleanup import normalized_functions + + +def definitions(ir): + return set(re.findall(r'^define [^\n]*@([\w.$]+)\(', ir, re.M)) + + +def globals_defined(ir): + return {name: body for name, body in re.findall(r'^@([\w.$]+) = ([^\n]+)', ir, re.M) + if not body.startswith('external ')} + + +def check_retention(before, after): + live_functions = {'retention_probe', 'retention_external', 'retention_table_target', 'retention_used_target', 'retention_linker_target'} + live_globals = {'RETENTION_TABLE', 'RETENTION_USED', 'RETENTION_DATA', 'RETENTION_DATA_REF', 'RETENTION_LINKER_USED'} + for phase, ir in [('before', before), ('after', after)]: + if not live_functions <= definitions(ir): + raise RuntimeError(f'{phase}: missing retained functions: {live_functions - definitions(ir)}') + if not live_globals <= globals_defined(ir).keys(): + raise RuntimeError(f'{phase}: missing retained initialized data') + for symbol in live_globals: + if globals_defined(before)[symbol] != globals_defined(after)[symbol]: + raise RuntimeError(f'GlobalDCE changed initializer/linkage: {symbol}') + if 'retention_unreachable' not in definitions(before) or 'RETENTION_UNREACHABLE_DATA' not in globals_defined(before): + raise RuntimeError('negative controls were not emitted before GlobalDCE') + if 'retention_unreachable' in definitions(after) or 'RETENTION_UNREACHABLE_DATA' in globals_defined(after): + raise RuntimeError('default GlobalDCE did not remove unreachable controls') + for used in ('llvm.used', 'llvm.compiler.used'): + if not re.search(r'^@'+re.escape(used)+r' = appending ', after, re.M): + raise RuntimeError(f'missing appending-linkage {used}') + if not re.search(r'^define (?!internal\b)[^\n]*@retention_external\(', after, re.M): + raise RuntimeError('explicit external function was internalized') + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('artifacts', type=Path) + args = parser.parse_args() + root = args.artifacts.resolve()/'default-dce' + scripts = Path(__file__).resolve().parent + root.mkdir(parents=True, exist_ok=True) + results = [] + for fixture, extra in [('small', []), ('retention', [str(scripts/'retention-kernels')])]: + modules = {} + for mode in ('none', 'default'): + out = root/fixture/mode; out.mkdir(parents=True, exist_ok=True) + command = ['cargo', 'run', '-vv', '-p', 'ptx_export', '--features', 'llvm21', '--', str(out), mode, *extra] + (out/'compiler-command.json').write_text(json.dumps(command, indent=2)+'\n') + with (out/'build.log').open('w') as log: + subprocess.run(command, stdout=log, stderr=subprocess.STDOUT, check=True) + subprocess.run(['opt-21', '-passes=verify', '-disable-output', str(out/'final-module.ll')], check=True) + subprocess.run([sys.executable, str(scripts/'inspect_codegen.py'), str(out)], check=True) + modules[mode] = out + raw = (modules['none']/'final-module.ll').read_text() + default = (modules['default']/'final-module.ll').read_text() + before = (modules['default']/'final-module.before-cleanup.ll').read_text() + if not len(definitions(default)) < len(definitions(raw)): + raise RuntimeError('default did not prune definitions relative to explicit disable') + for left, right in [(modules['none'], modules['default'])]: + if normalized_functions((left/'rust_kernels.ptx').read_text()) != normalized_functions((right/'rust_kernels.ptx').read_text()): + raise RuntimeError(f'{fixture}: default changes baseline PTX function bodies') + if fixture == 'retention': + check_retention(before, default) + else: + subprocess.run([sys.executable, str(scripts/'check_cleanup_ir.py'), str(modules['default']/'final-module.ll'), + '--out', str(modules['default']/'host-ir-check')], check=True) + results.append({'fixture':fixture, 'definitions_before_after':[len(definitions(raw)), len(definitions(default))], + 'matches_disabled_ptx_function_bodies':True, + 'ptx_sha256':{m:hashlib.sha256((p/'rust_kernels.ptx').read_bytes()).hexdigest() for m,p in modules.items()}}) + (root/'comparison.json').write_text(json.dumps(results, indent=2)+'\n') + print(f'{fixture}: default GlobalDCE verified against explicit disable', flush=True) + + +if __name__ == '__main__': + main() diff --git a/examples/ptx_export/check_integrated_cleanup.py b/examples/ptx_export/check_integrated_cleanup.py new file mode 100644 index 00000000..f8c6d9f6 --- /dev/null +++ b/examples/ptx_export/check_integrated_cleanup.py @@ -0,0 +1,67 @@ +#!/usr/bin/env python3 +"""Compile through the real backend and compare against the offline pass replay.""" +import argparse +import json +import re +from pathlib import Path +import subprocess +import sys +from replay_cleanup import normalized_functions + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('artifacts', type=Path) + args = parser.parse_args() + root = args.artifacts.resolve() + results = [] + for mode, replay in [('inline-scalar', 'inline-only'), ('dce', 'dce-only'), ('scalar', 'local-cleanup'), ('inline', 'inline-cleanup')]: + dest = root/'integrated-cleanup'/mode + dest.mkdir(parents=True, exist_ok=True) + command = ['cargo', 'run', '-vv', '-p', 'ptx_export', '--features', 'llvm21', '--', str(dest), mode] + (dest/'compiler-command.json').write_text(json.dumps(command, indent=2)+'\n') + with (dest/'build.log').open('w') as log: + subprocess.run(command, stdout=log, stderr=subprocess.STDOUT, check=True) + subprocess.run(['opt-21', '-passes=verify', '-disable-output', str(dest/'final-module.ll')], check=True) + before = dest/'final-module.before-cleanup.ll' + if not before.is_file(): raise RuntimeError('backend did not record the pre-cleanup IR') + actual = normalized_functions((dest/'rust_kernels.ptx').read_text()) + expected = normalized_functions((root/'cleanup-experiment'/replay/'rust_kernels.ptx').read_text()) + matches = actual == expected + required = {'rust_vecadd', 'rust_sha256_32', 'rust_guarded_select', 'rust_filtered_select'} + if not required <= actual.keys(): raise RuntimeError('cleanup removed an exported kernel') + result = {'mode': mode, 'matches_replay_function_bodies': matches} + if mode == 'dce': + baseline = normalized_functions((root/'rust_kernels.ptx').read_text()) + if actual != baseline: + raise RuntimeError('DCE-only changed baseline PTX function bodies') + before_count = len(re.findall(r'^define ', before.read_text(), re.M)) + after_count = len(re.findall(r'^define ', (dest/'final-module.ll').read_text(), re.M)) + if not 0 < after_count < before_count: + raise RuntimeError('DCE-only did not prune unused definitions') + result['definitions_before_after'] = [before_count, after_count] + if mode == 'inline': + # Check the real compiler output, not a hand-written substitute. + # The negative control must still carry its observable dependency. + ir = (dest/'final-module.ll').read_text() + counts = {} + for helper in ('filtered', 'observed'): + bodies = re.findall(r'^define [^\n]*guarded_select\d+' + helper + + r'\([^\n]*\{\n(.*?)^}', ir, re.M | re.S) + if len(bodies) != 1: raise RuntimeError(f'expected one {helper} IR definition') + counts[helper] = len(re.findall(r'= select ', bodies[0])) + result['ir_select_counts'] = counts + if counts != {'filtered': 0, 'observed': 1}: + raise RuntimeError(f'guarded-select cleanup regression: {counts}') + results.append(result) + (root/'integrated-cleanup'/'comparison.json').write_text(json.dumps(results, indent=2)+'\n') + if not matches: raise RuntimeError(f'{mode}: integrated cleanup differs from replay') + subprocess.run([sys.executable, str(Path(__file__).with_name('inspect_codegen.py')), str(dest)], check=True) + if mode in ('dce', 'inline', 'inline-scalar'): + subprocess.run([sys.executable, str(Path(__file__).with_name('check_cleanup_ir.py')), + str(dest/'final-module.ll'), '--out', str(dest/'host-ir-check')], check=True) + print(f'{mode}: real backend matches replay; LLVM verified and PTX assembled', flush=True) + + +if __name__ == '__main__': + main() diff --git a/examples/ptx_export/cleanup_ir_oracle.c b/examples/ptx_export/cleanup_ir_oracle.c new file mode 100644 index 00000000..cf255da9 --- /dev/null +++ b/examples/ptx_export/cleanup_ir_oracle.c @@ -0,0 +1,30 @@ +// Independent index oracle for the five actual compiler-emitted helpers. +// The generated header binds their Rust symbol names; no Rust source is rebuilt. +#include +#include +#include "helpers.h" + +int main(void) { + uint64_t seeds[] = {0, 1, UINT64_C(1)<<63, UINT64_MAX}; + uint64_t initials[] = {0,1,17,63,64,UINT64_MAX}; + unsigned cases=0; + for (unsigned t=0;t<4;t++) { + uint64_t table[64]; + for (unsigned i=0;i<64;i++) table[i]=UINT64_C(0x9e3779b97f4a7c15)*i+seeds[t]; + for (unsigned l=0;l<67;l++) for (unsigned a=0;a<6;a++) { + uint32_t limit=l==66?UINT32_MAX:l; + unsigned n=limit<64?limit:64; + uint64_t odd=0,live=0; + for (unsigned i=1;i u64 { + let mut remembered = initial; + let mut sum = 0u64; + for i in 0..limit.min(64) { + let odd = i & 1 != 0; + if odd { + remembered = u64::from(i); + } + if odd { + sum = sum.wrapping_add(table[(remembered & 63) as usize]); + } + } + sum +} + +#[inline(never)] +pub fn direct(table: &[u64; 64], limit: u32) -> u64 { + let mut sum = 0u64; + for i in 0..limit.min(64) { + if i & 1 != 0 { + sum = sum.wrapping_add(table[i as usize]); + } + } + sum +} + +/// Negative control: the false-path value is observed, including at i=0. +/// Replacing the conditional assignment with `remembered = i` is incorrect. +#[inline(never)] +pub fn observed(table: &[u64; 64], limit: u32, initial: u64) -> u64 { + let mut remembered = initial; + let mut sum = 0u64; + for i in 0..limit.min(64) { + if i & 1 != 0 { + remembered = u64::from(i); + } + sum = sum.wrapping_add(table[(remembered & 63) as usize]); + } + sum +} + +/// Preserve Dalek's filtered-range idiom; strip curve arithmetic first. +#[inline(never)] +pub fn filtered(table: &[u64; 64], limit: u32) -> u64 { + let mut sum = 0u64; + for i in (0..limit.min(64) as usize).filter(|x| x % 2 == 1) { + sum = sum.wrapping_add(table[i]); + } + sum +} + +/// Same odd-index accesses, without Filter::next's retained result. +#[inline(never)] +pub fn stepped(table: &[u64; 64], limit: u32) -> u64 { + let mut sum = 0u64; + for i in (1..limit.min(64) as usize).step_by(2) { + sum = sum.wrapping_add(table[i]); + } + sum +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn compare_against_independent_index_oracle() { + for seed in [0u64, 1, 0x8000_0000_0000_0000, u64::MAX] { + let table = core::array::from_fn(|i| { + (i as u64) + .wrapping_mul(0x9e37_79b9_7f4a_7c15) + .wrapping_add(seed) + }); + for limit in (0..=65).chain([u32::MAX]) { + for initial in [0u64, 1, 17, 63, 64, u64::MAX] { + let n = limit.min(64) as usize; + let expected = (1..n) + .step_by(2) + .fold(0u64, |sum, i| sum.wrapping_add(table[i])); + // Last odd index <= i; on the first iteration use initial. + let observed_expected = (0..n).fold(0u64, |sum, i| { + let index = if i == 0 { + (initial & 63) as usize + } else { + (i - 1) | 1 + }; + sum.wrapping_add(table[index]) + }); + assert_eq!(preserved(&table, limit, initial), expected); + assert_eq!(direct(&table, limit), expected); + assert_eq!(filtered(&table, limit), expected); + assert_eq!(stepped(&table, limit), expected); + assert_eq!(observed(&table, limit, initial), observed_expected); + } + } + } + } + + #[test] + fn control_detects_discarding_a_live_false_arm() { + let table = core::array::from_fn(|i| i as u64); + assert_eq!(preserved(&table, 1, 17), 0); + assert_eq!(direct(&table, 1), 0); + assert_eq!(observed(&table, 1, 17), 17); + assert_ne!(observed(&table, 1, 17), table[0]); + } +} diff --git a/examples/ptx_export/kernels/src/lib.rs b/examples/ptx_export/kernels/src/lib.rs index 9f02d2b3..29f199be 100644 --- a/examples/ptx_export/kernels/src/lib.rs +++ b/examples/ptx_export/kernels/src/lib.rs @@ -27,3 +27,59 @@ pub unsafe fn rust_sha256_32(input: *const u8, out: *mut u8, count: u32) { unsafe { core::ptr::copy_nonoverlapping(digest.as_ptr(), out.add(offset), 32) }; } } + +mod guarded_select; + +/// Compare guarded loop forms with runtime inputs; write three u64s per case: +/// preserved, direct, and the observable-false-arm negative control. +/// +/// # Safety +/// `table` addresses 64 readable u64s, `limits` and `initials` address `count` +/// readable elements, and `out` addresses `count * 3` writable u64s. +/// Output must not overlap any input. The loop limit is capped at 64. +#[kernel] +pub unsafe fn rust_guarded_select( + table: *const u64, + limits: *const u32, + initials: *const u64, + out: *mut u64, + count: u32, +) { + let i = thread::index_1d(); + if i < count { + let i = i as usize; + let table = unsafe { &*table.cast::<[u64; 64]>() }; + let limit = unsafe { *limits.add(i) }; + let initial = unsafe { *initials.add(i) }; + unsafe { + *out.add(i * 3) = guarded_select::preserved(table, limit, initial); + *out.add(i * 3 + 1) = guarded_select::direct(table, limit); + *out.add(i * 3 + 2) = guarded_select::observed(table, limit, initial); + } + } +} + +/// Compare the filtered iterator in Dalek's basepoint multiplication with a +/// stepped iterator, using runtime table contents and loop limits. +/// +/// # Safety +/// `table` addresses 64 readable u64s, `limits` addresses `count` readable u32s, +/// and `out` addresses `count * 2` writable u64s. Output must not overlap inputs. +#[kernel] +pub unsafe fn rust_filtered_select( + table: *const u64, + limits: *const u32, + out: *mut u64, + count: u32, +) { + let i = thread::index_1d(); + if i < count { + let i = i as usize; + let table = unsafe { &*table.cast::<[u64; 64]>() }; + let limit = unsafe { *limits.add(i) }; + unsafe { + *out.add(i * 2) = guarded_select::filtered(table, limit); + *out.add(i * 2 + 1) = guarded_select::stepped(table, limit); + } + } +} diff --git a/examples/ptx_export/optimization_pipelines.py b/examples/ptx_export/optimization_pipelines.py new file mode 100644 index 00000000..d6e56cb9 --- /dev/null +++ b/examples/ptx_export/optimization_pipelines.py @@ -0,0 +1,48 @@ +"""Bounded LLVM 21 experiments. Explicit options are saved with each result.""" +IC = 'instcombine' +SCALAR = f'function(sroa,{IC},simplifycfg,adce)' +INLINE = f'globaldce,cgscc(inline),{SCALAR},globaldce' +CORRELATED = f'function(correlated-propagation,{IC},simplifycfg,adce)' + + +def experiments(extended=False): + pipelines = { + 'baseline': ('verify', []), + 'dce-only': ('globaldce,verify', []), + 'local-cleanup': (f'{SCALAR},verify', []), + 'inline-only': (f'{INLINE},verify', []), + 'inline-cleanup': (f'{INLINE},{CORRELATED},verify', []), + 'constrained-cleanup': (f'{INLINE},{CORRELATED},function(constraint-elimination,{IC},simplifycfg,adce),verify', []), + } + if extended: + pipelines.update({ + 'cfg-no-final': (f'{INLINE},function(correlated-propagation,{IC},adce),verify', []), + 'cfg-before-combine': (f'{INLINE},function(correlated-propagation,simplifycfg,{IC},adce),verify', []), + 'cfg-no-both': (f'globaldce,cgscc(inline),function(sroa,{IC},adce),globaldce,function(correlated-propagation,{IC},adce),verify', []), + 'dce-scalar': (f'globaldce,{SCALAR},globaldce,verify', []), + 'memory-early-cse': (f'{INLINE},function(early-cse,{IC},adce),verify', []), + 'memory-gvn': (f'{INLINE},function(gvn,{IC},adce),verify', []), + 'memory-stores': (f'{INLINE},function(memcpyopt,dse,{IC},adce),verify', []), + 'memory-combined': (f'{INLINE},function(early-cse,gvn,memcpyopt,dse,{IC},adce),verify', []), + 'inline-threshold-0': (f'{INLINE},{CORRELATED},verify', ['-inline-threshold=0', '-inlinehint-threshold=0']), + 'inline-threshold-50': (f'{INLINE},{CORRELATED},verify', ['-inline-threshold=50', '-inlinehint-threshold=50']), + 'inline-threshold-450': (f'{INLINE},{CORRELATED},verify', ['-inline-threshold=450', '-inlinehint-threshold=450']), + }) + return pipelines + + +# Small constant-trip loops are a concrete source of residual RNG stack buffers. +# Disable runtime/partial/peeling expansion and cap full unrolling at four trips. +TINY_UNROLL = 'loop-unroll' +TINY_LOOPS = f'loop-simplify,lcssa,loop(indvars),{TINY_UNROLL},sroa,{IC},simplifycfg,adce' + + +def wave2_experiments(): + return { + 'baseline': ('verify', []), + 'dce-only': ('globaldce,verify', []), + 'inline-only': (f'{INLINE},verify', []), + 'tiny-loops': (f'{INLINE},function({TINY_LOOPS}),verify', []), + 'correlated-tiny-loops': (f'{INLINE},{CORRELATED},function({TINY_LOOPS}),verify', []), + 'loop-idiom': (f'{INLINE},function(loop-simplify,lcssa,loop(loop-idiom),memcpyopt,sroa,{IC},simplifycfg,adce),verify', []), + } diff --git a/examples/ptx_export/replay_cleanup.py b/examples/ptx_export/replay_cleanup.py new file mode 100644 index 00000000..ac386940 --- /dev/null +++ b/examples/ptx_export/replay_cleanup.py @@ -0,0 +1,139 @@ +#!/usr/bin/env python3 +"""Opt-in pre-NVVM cleanup experiment; never changes Rust-CUDA defaults.""" +import argparse +import ctypes as c +import hashlib +import json +import os +from pathlib import Path +import re +import subprocess +import sys +import time +from optimization_pipelines import experiments, wave2_experiments +from inspect_codegen import ptx_functions, ptx_summary + + +def compile_nvvm(bitcode, libraries, output): + lib = c.CDLL('libnvvm.so') + ptr = c.c_void_p + signatures = { + 'nvvmCreateProgram': [c.POINTER(ptr)], 'nvvmDestroyProgram': [c.POINTER(ptr)], + 'nvvmAddModuleToProgram': [ptr, ptr, c.c_size_t, c.c_char_p], + 'nvvmLazyAddModuleToProgram': [ptr, ptr, c.c_size_t, c.c_char_p], + 'nvvmCompileProgram': [ptr, c.c_int, c.POINTER(c.c_char_p)], + 'nvvmVerifyProgram': [ptr, c.c_int, c.POINTER(c.c_char_p)], + 'nvvmGetProgramLogSize': [ptr, c.POINTER(c.c_size_t)], + 'nvvmGetProgramLog': [ptr, ptr], + 'nvvmGetCompiledResultSize': [ptr, c.POINTER(c.c_size_t)], + 'nvvmGetCompiledResult': [ptr, ptr], + } + for name, types in signatures.items(): + fn = getattr(lib, name); fn.argtypes = types; fn.restype = c.c_int + program = ptr() + def call(name, *args): + status = getattr(lib, name)(*args) + if status: raise RuntimeError(f'{name}: NVVM status {status}') + def log(): + size = c.c_size_t() + if lib.nvvmGetProgramLogSize(program, c.byref(size)) or not size.value: return '' + data = c.create_string_buffer(size.value) + if lib.nvvmGetProgramLog(program, data): return 'cannot read NVVM log' + return data.value.decode(errors='replace') + call('nvvmCreateProgram', c.byref(program)) + # Keep all backing buffers alive until destruction of the NVVM program. + buffers = [] + try: + for index, path in enumerate([bitcode, *libraries]): + data = path.read_bytes(); buf = c.create_string_buffer(data); buffers.append(buf) + call('nvvmAddModuleToProgram' if index == 0 else 'nvvmLazyAddModuleToProgram', + program, buf, len(data), [b'merged', b'libdevice', b'libintrinsics'][index]) + options = (c.c_char_p * 1)(b'-arch=compute_100') + status = lib.nvvmVerifyProgram(program, 1, options) + verification_log = log() + (output / 'nvvm-verify.log').write_text(verification_log) + if status: + # Match the existing LLVM 21 backend's narrowly recognized verifier + # false negative; retain the log and let compilation decide. + known = all(x in verification_log for x in + ("Producer: 'LLVM21", "Reader: 'LLVM 7.0.1'", 'parse Invalid value')) + if not known: raise RuntimeError(f'NVVM verification failed: {status}') + status = lib.nvvmCompileProgram(program, 1, options) + (output / 'nvvm-compile.log').write_text(log()) + if status: raise RuntimeError(f'NVVM compilation failed: {status}') + size = c.c_size_t(); call('nvvmGetCompiledResultSize', program, c.byref(size)) + data = c.create_string_buffer(size.value) + call('nvvmGetCompiledResult', program, data) + (output / 'rust_kernels.ptx').write_bytes(data.raw.rstrip(b'\0')) + finally: + call('nvvmDestroyProgram', c.byref(program)) + + +def normalized_functions(ptx): + return {name: re.sub(r'\s+', ' ', re.sub(r'//[^\n]*', '', body)).strip() + for name, body in ptx_functions(ptx)} + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument('artifacts', type=Path) + parser.add_argument('--extended', action='store_true') + parser.add_argument('--wave2', action='store_true', help='Record inlining remarks and bounded loop experiments') + parser.add_argument('--only', help='Comma-separated experiment names; baseline is always included') + args = parser.parse_args() + root = args.artifacts.resolve() + out = root / 'cleanup-experiment'; out.mkdir(exist_ok=True) + # Require unique content when build caches contain multiple backend hashes. + candidates = list(Path('target/cuda-builder-codegen').rglob('libintrinsics_v21.bc')) + unique = {hashlib.sha256(p.read_bytes()).hexdigest(): p for p in candidates} + if len(unique) != 1: raise RuntimeError(f'expected one distinct LLVM 21 intrinsic library, got {len(unique)}') + intrinsics = next(iter(unique.values())).resolve() + libdevice = Path(os.environ['CUDA_HOME']) / 'nvvm/libdevice/libdevice.10.bc' + libraries = [libdevice, intrinsics] + metadata = {'nvvm_options': ['-arch=compute_100'], + 'libraries': [{'path': str(p), 'sha256': hashlib.sha256(p.read_bytes()).hexdigest()} for p in libraries], + 'input_sha256': hashlib.sha256((root/'final-module.ll').read_bytes()).hexdigest(), + 'results': []} + pipelines = wave2_experiments() if args.wave2 else experiments(args.extended) + if args.only: + requested = set(args.only.split(',')) | {'baseline'} + unknown = requested - pipelines.keys() + if unknown: raise ValueError(f'unknown experiments: {sorted(unknown)}') + pipelines = {name: spec for name, spec in pipelines.items() if name in requested} + baseline_matches = False + for name, (passes, options) in pipelines.items(): + dest = out/name; dest.mkdir(exist_ok=True) + if args.wave2: + options = [*options, '-pass-remarks=inline|loop-unroll|sroa', '-pass-remarks-missed=inline|loop-unroll|sroa', + '-pass-remarks-analysis=inline|loop-unroll|sroa', '-pass-remarks-output='+str(dest/'remarks.yaml')] + command = ['opt-21', '-passes='+passes, '-verify-each', *options, str(root/'final-module.ll'), '-o', str(dest/'module.bc')] + item = {'name': name, 'passes': passes, 'opt_options': options, 'opt_command': command} + started = time.perf_counter() + try: + with (dest/'opt.log').open('w') as log: + subprocess.run(command, stdout=log, stderr=subprocess.STDOUT, check=True) + item['opt_seconds'] = time.perf_counter() - started + subprocess.run(['llvm-dis-21', str(dest/'module.bc'), '-o', str(dest/'final-module.ll')], check=True) + compile_nvvm(dest/'module.bc', libraries, dest) + source = (dest/'rust_kernels.ptx').read_text() + if name == 'baseline': + baseline_matches = normalized_functions(source) == normalized_functions((root/'rust_kernels.ptx').read_text()) + item['matches_original_function_bodies'] = baseline_matches + if not baseline_matches: raise RuntimeError('baseline replay differs; do not attribute differences to cleanup') + subprocess.run([sys.executable, str(Path(__file__).with_name('inspect_codegen.py')), str(dest), + '--reuse-from',str(root),'--reuse-from',str(root.parent/'inline-scalar')], check=True) + item['ptx_helpers'] = {n:v for n,v in ptx_summary(source).items() if 'guarded_select' in n} + item['status'] = 'compiled_and_assembled' + except (RuntimeError, subprocess.CalledProcessError) as error: + item['status'] = 'failed'; item['error'] = str(error) + item['total_seconds'] = time.perf_counter() - started + metadata['results'].append(item) + (out/'experiment.json').write_text(json.dumps(metadata, indent=2)+'\n') + print(name, item['status'], item.get('error',''), flush=True) + if name == 'baseline' and not baseline_matches: return 1 + # Unsupported optimized IR is an experiment result, not a silent success. + return int(any(item['status'] == 'failed' for item in metadata['results'])) + + +if __name__ == '__main__': + raise SystemExit(main()) diff --git a/examples/ptx_export/retention-kernels/Cargo.toml b/examples/ptx_export/retention-kernels/Cargo.toml new file mode 100644 index 00000000..207b50b7 --- /dev/null +++ b/examples/ptx_export/retention-kernels/Cargo.toml @@ -0,0 +1,11 @@ +[package] +name = "ptx-retention-kernels" +version = "0.1.0" +edition = "2021" +publish = false + +[dependencies] +cuda_std = { path = "../../../crates/cuda_std" } + +[lib] +crate-type = ["cdylib", "rlib"] diff --git a/examples/ptx_export/retention-kernels/src/lib.rs b/examples/ptx_export/retention-kernels/src/lib.rs new file mode 100644 index 00000000..288bf899 --- /dev/null +++ b/examples/ptx_export/retention-kernels/src/lib.rs @@ -0,0 +1,70 @@ +//! Exercise actual Rust-CUDA internalization and DCE root discovery. +#![feature(used_with_arg)] + +use cuda_std::{externally_visible, kernel}; + +type Callback = extern "C" fn(u32) -> u32; + +#[no_mangle] +#[inline(never)] +pub extern "C" fn retention_table_target(x: u32) -> u32 { + x.wrapping_mul(7) +} + +#[no_mangle] +#[inline(never)] +pub extern "C" fn retention_used_target(x: u32) -> u32 { + x.wrapping_add(19) +} + +#[no_mangle] +#[inline(never)] +pub extern "C" fn retention_linker_target(x: u32) -> u32 { + x.wrapping_sub(13) +} + +#[used(linker)] +#[no_mangle] +pub static RETENTION_LINKER_USED: [Callback; 1] = [retention_linker_target]; + +#[externally_visible] +#[no_mangle] +pub extern "C" fn retention_external(x: u32) -> u32 { + x ^ 0x5a5a_5a5a +} + +#[no_mangle] +pub static RETENTION_TABLE: [Callback; 1] = [retention_table_target]; + +// Never read by a kernel. Its initializer must nevertheless retain its target. +#[used(compiler)] +#[no_mangle] +pub static RETENTION_USED: [Callback; 1] = [retention_used_target]; + +#[no_mangle] +pub static RETENTION_DATA: [u32; 3] = [11, 29, 47]; + +#[no_mangle] +pub static RETENTION_DATA_REF: &[u32; 3] = &RETENTION_DATA; + +// Emitted as externally named items, then internalized and eligible for DCE. +#[no_mangle] +pub extern "C" fn retention_unreachable(x: u32) -> u32 { + x.wrapping_add(23) +} + +#[no_mangle] +pub static RETENTION_UNREACHABLE_DATA: [u32; 3] = [101, 103, 107]; + +/// # Safety +/// `out` addresses three writable u64s. No inputs overlap the output. +#[kernel] +pub unsafe fn retention_probe(out: *mut u64, index: u32) { + // Volatile reads keep the initializer dependencies visible in the emitted IR. + let callback = + core::ptr::read_volatile(core::ptr::addr_of!(RETENTION_TABLE).cast::()); + let data = core::ptr::read_volatile(core::ptr::addr_of!(RETENTION_DATA_REF)); + *out = callback as usize as u64; + *out.add(1) = core::ptr::read_volatile(data.as_ptr().add((index % 3) as usize)) as u64; + *out.add(2) = retention_external(index) as u64; +} diff --git a/examples/ptx_export/src/main.rs b/examples/ptx_export/src/main.rs index b8b82079..706a7681 100644 --- a/examples/ptx_export/src/main.rs +++ b/examples/ptx_export/src/main.rs @@ -1,4 +1,4 @@ -use cuda_builder::CudaBuilder; +use cuda_builder::{CudaBuilder, Llvm19Cleanup}; use std::{env, fs, path::PathBuf}; fn main() -> Result<(), Box> { @@ -9,8 +9,50 @@ fn main() -> Result<(), Box> { ); fs::create_dir_all(&output)?; let output = output.canonicalize()?; - let kernels = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("kernels"); - let ptx = CudaBuilder::new(kernels) + let kernels = env::args_os() + .nth(3) + .map(PathBuf::from) + .unwrap_or_else(|| PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("kernels")); + let mut builder = CudaBuilder::new(kernels); + if let Some(features) = env::args().nth(4) { + builder = builder.build_args(&["--no-default-features", "--features", &features]); + } + let mode = env::args().nth(2).unwrap_or_else(|| "default".into()); + // Historical wave-1 experiments explicitly isolate their selected pipeline. + // The default mode exercises the production default without overrides. + if cfg!(feature = "llvm21") && mode != "default" { + builder = builder.llvm19_global_dce(false); + } + match mode.as_str() { + "default" => {} + "none" => {} + "size-s" => { + builder = builder + .llvm19_cleanup(Llvm19Cleanup::Inline) + .build_args(&["--config", "profile.release.opt-level=\"s\""]) + } + "size-z" => { + builder = builder + .llvm19_cleanup(Llvm19Cleanup::Inline) + .build_args(&["--config", "profile.release.opt-level=\"z\""]) + } + "module-scalar" => builder = builder.llvm19_module_cleanup(true), + "module-inline" => { + builder = builder + .llvm19_module_cleanup(true) + .llvm19_cleanup(Llvm19Cleanup::Inline) + } + "inline-scalar" => builder = builder.llvm19_cleanup(Llvm19Cleanup::InlineScalar), + "dce" => builder = builder.llvm19_cleanup(Llvm19Cleanup::GlobalDce), + "scalar" => builder = builder.llvm19_cleanup(Llvm19Cleanup::Scalar), + "inline" => builder = builder.llvm19_cleanup(Llvm19Cleanup::Inline), + _ => { + return Err( + "cleanup mode must be default, none, dce, scalar, inline, inline-scalar, module-scalar, module-inline, size-s, or size-z".into(), + ); + } + } + let ptx = builder .copy_to(output.join("rust_kernels.ptx")) .final_module_path(output.join("final-module.ll")) .emit_llvm_ir(true) diff --git a/examples/ptx_export/test_inspect_codegen.py b/examples/ptx_export/test_inspect_codegen.py new file mode 100644 index 00000000..394a4bbf --- /dev/null +++ b/examples/ptx_export/test_inspect_codegen.py @@ -0,0 +1,67 @@ +import unittest +from inspect_codegen import ptx_summary, sass_summary, sass_symbols + + +class InventoryTests(unittest.TestCase): + def test_ptx_prototypes_nested_calls_and_self_select(self): + source = ''' +.extern .func (.param .b64 retval) prototype(.param .b64 arg); +.visible .func (.param .b64 retval) preserved(.param .b64 arg) { + .reg .b64 %rd<3>; + { .param .b64 slot; } + selp.b64 %rd1, %rd2, %rd1, %p1; + @%p1 bra DONE; +DONE: + ret; +} +.visible .entry direct() { + selp.b32 %r1, %r2, %r3, %p1; + ret; +} +''' + result = ptx_summary(source) + self.assertEqual(set(result), {'preserved', 'direct'}) + self.assertEqual(len(result['preserved']['self_false_selects']), 1) + self.assertEqual(result['direct']['self_false_selects'], []) + self.assertEqual(result['preserved']['opcode_histogram']['bra'], 1) + + def test_sass_counts_instructions_not_encoding_continuations(self): + result = sass_summary('''Function : preserved + /*0000*/ @!P0 LDG.E R2, [R4]; /* 0x123 */ + /*0010*/ SEL R3, R2, R1, P0; + /* 0x000000 */ +Function : direct + /*0000*/ EXIT; +''') + self.assertEqual(result['preserved'], {'LDG.E': 1, 'SEL': 1}) + self.assertEqual(result['direct'], {'EXIT': 1}) + + def test_helper_extents_exclude_adjacent_functions(self): + source = """ + .size kernel,(END - kernel) +kernel: + /*0000*/ CALL.REL helper; + .size $kernel$first,($kernel$second - $kernel$first) +$kernel$first: + /*0010*/ @P0 SEL R1, R2, R3, P0; +LOCAL: + /*0020*/ RET.REL.NODEC; + .size $kernel$second,(END - $kernel$second) +$kernel$second: + /*0030*/ LDG.E R2, [R4]; + /*0040*/ NOP; +END: +""" + result = sass_symbols(source) + self.assertEqual(result['$kernel$first']['opcode_histogram'], {'SEL': 1, 'RET.REL.NODEC': 1}) + self.assertEqual(result['$kernel$second']['non_nop_instructions'], 1) + self.assertEqual(result['kernel']['non_nop_instructions'], 4) + self.assertEqual(result['kernel']['scope'], 'entry_including_helpers') + + def test_missing_symbol_end_is_an_error(self): + with self.assertRaises(ValueError): + sass_symbols('.size helper,(MISSING - helper)\nhelper:\n /*0000*/ RET;\n') + + +if __name__ == '__main__': + unittest.main() diff --git a/examples/ptx_export/test_inspection_reuse.py b/examples/ptx_export/test_inspection_reuse.py new file mode 100644 index 00000000..f5615e7d --- /dev/null +++ b/examples/ptx_export/test_inspection_reuse.py @@ -0,0 +1,34 @@ +import json +from pathlib import Path +import tempfile +import unittest +from inspect_codegen import INSPECTION_OUTPUTS, reuse_inspection + + +class InspectionReuseTests(unittest.TestCase): + def test_identical_inputs_reuse_with_explicit_provenance(self): + with tempfile.TemporaryDirectory() as directory: + source, dest = self.fixtures(Path(directory)) + self.assertTrue(reuse_inspection(source,dest)) + for name in INSPECTION_OUTPUTS: + self.assertEqual((source/name).read_bytes(),(dest/name).read_bytes()) + self.assertIn('assembly/disassembly not rerun',json.loads((dest/'reused-inspection.json').read_text())['note']) + self.assertFalse(reuse_inspection(source,source)) + + def test_changed_ptx_tool_or_inspector_never_reuses(self): + for name in ['rust_kernels.ptx','inspection-producer-sha256.txt','inspection-tools-sha256.json','ptxas-version.txt']: + with self.subTest(name=name), tempfile.TemporaryDirectory() as directory: + source,dest=self.fixtures(Path(directory)) + (dest/name).write_text('different') + self.assertFalse(reuse_inspection(source,dest)) + self.assertFalse((dest/'rust_kernels.cubin').exists()) + + @staticmethod + def fixtures(root): + source,dest=root/'source',root/'dest';source.mkdir();dest.mkdir() + for name in ['rust_kernels.ptx','inspection-producer-sha256.txt','inspection-tools-sha256.json', + 'ptxas-version.txt','nvdisasm-version.txt','cuobjdump-version.txt']: + for path in [source,dest]:(path/name).write_text(name) + for name in INSPECTION_OUTPUTS:(source/name).write_text('offline output '+name) + (source/'inspection-commands.json').write_text('[]') + return source,dest