keplemon 3.5.0

Expanded functionality for the Standardized Astrodynamics Algorithms Library (SAAL)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140

use std::env;
use std::ffi::OsStr;
use std::fs;
use std::path::{Path, PathBuf};
#[cfg(feature = "cuda")]
use std::process::Command;

fn target_dir(out_dir: &Path) -> PathBuf {
    out_dir
        .ancestors()
        .nth(3)
        .expect("Couldn't determine target directory")
        .to_path_buf()
}

fn main() {
    println!("cargo:rerun-if-changed=build.rs");

    // CUDA kernel compilation
    #[cfg(feature = "cuda")]
    compile_cuda_kernels();

    // Python wheel build
    if env::var("CARGO_FEATURE_PYTHON").is_err() {
        return;
    }

    let out_dir = PathBuf::from(env::var("OUT_DIR").expect("OUT_DIR not set"));
    let target_dir = target_dir(&out_dir);

    let python_pkg_dir = Path::new("python").join("keplemon");
    fs::create_dir_all(&python_pkg_dir).expect("Failed to create python/keplemon directory");

    for entry in fs::read_dir(&target_dir).expect("Failed to read target directory") {
        let entry = entry.expect("Failed to access entry in target directory");
        let path = entry.path();
        if !path.is_file() {
            continue;
        }
        let filename = path.file_name().expect("Invalid target file name");
        if filename == "Cargo.lock" || filename == ".cargo-lock" || filename == "libkeplemon.d" {
            continue;
        }

        let dest_path = python_pkg_dir.join(filename);
        fs::copy(&path, &dest_path)
            .unwrap_or_else(|_| panic!("Failed to copy {} to {}", path.display(), dest_path.display()));
    }

    let stubs_dir = Path::new("stubs").join("keplemon");
    if stubs_dir.is_dir() {
        for entry in fs::read_dir(&stubs_dir).expect("Failed to read stubs/keplemon directory") {
            let entry = entry.expect("Failed to access entry in stubs/keplemon");
            let path = entry.path();
            if path.extension() != Some(OsStr::new("pyi")) {
                continue;
            }
            println!("cargo:rerun-if-changed={}", path.display());
            let filename = path.file_name().expect("Invalid stub file name");
            let dest_path = python_pkg_dir.join(filename);
            fs::copy(&path, &dest_path)
                .unwrap_or_else(|_| panic!("Failed to copy stub {} to {}", path.display(), dest_path.display()));
        }
    }
}

#[cfg(feature = "cuda")]
fn compile_cuda_kernels() {
    println!("cargo:rerun-if-changed=kernels/sgp4_init.cu");
    println!("cargo:rerun-if-changed=kernels/sgp4_batch.cu");
    println!("cargo:rerun-if-changed=kernels/sgp4_types.cuh");
    println!("cargo:rerun-if-changed=kernels/sgp4_constants.cuh");

    let out_dir = env::var("OUT_DIR").expect("OUT_DIR not set");

    // Find nvcc - check CUDA_PATH or common locations
    let cuda_path = env::var("CUDA_PATH")
        .or_else(|_| env::var("CUDA_HOME"))
        .unwrap_or_else(|_| "/usr/local/cuda".to_string());

    let nvcc = PathBuf::from(&cuda_path).join("bin").join("nvcc");

    // Check if nvcc exists
    if !nvcc.exists() && !Command::new("nvcc").arg("--version").output().is_ok() {
        println!(
            "cargo:warning=nvcc not found. CUDA kernels will not be compiled. \
             CUDA features will be unavailable at runtime. \
             To enable CUDA: install CUDA Toolkit or set CUDA_PATH environment variable. \
             Looked in: {}",
            nvcc.display()
        );
        println!("cargo:warning=Skipping CUDA kernel compilation");

        // Create empty stub PTX files so include_str! doesn't fail
        let stub_ptx = "// CUDA kernels not compiled - nvcc not available\n";
        fs::write(format!("{}/sgp4_init.ptx", out_dir), stub_ptx).expect("Failed to write stub sgp4_init.ptx");
        fs::write(format!("{}/sgp4_batch.ptx", out_dir), stub_ptx).expect("Failed to write stub sgp4_batch.ptx");

        return;
    }

    let nvcc_cmd = if nvcc.exists() { nvcc.to_str().unwrap() } else { "nvcc" };

    // Compile initialization kernel
    compile_kernel(nvcc_cmd, "kernels/sgp4_init.cu", &format!("{}/sgp4_init.ptx", out_dir));

    // Compile batch propagation kernel
    compile_kernel(
        nvcc_cmd,
        "kernels/sgp4_batch.cu",
        &format!("{}/sgp4_batch.ptx", out_dir),
    );

    println!("cargo:info=CUDA kernels compiled successfully");
}

#[cfg(feature = "cuda")]
fn compile_kernel(nvcc: &str, input: &str, output: &str) {
    let status = Command::new(nvcc)
        .args(&[
            "-ptx",            // Compile to PTX
            "-O3",             // Optimization level 3
            "--use_fast_math", // Use fast math operations
            "-arch=sm_50",     // Target compute capability 5.0+ (Maxwell and newer)
            "--std=c++14",     // C++14 standard
            "-I",
            "kernels", // Include directory for headers
            "-o",
            output, // Output PTX file
            input,  // Input CUDA source
        ])
        .status()
        .unwrap_or_else(|e| panic!("Failed to execute nvcc: {}", e));

    if !status.success() {
        panic!("nvcc compilation failed for {}", input);
    }

    println!("cargo:info=Compiled {} to {}", input, output);
}