第一轮分析工作暂存

2026-06-09 18:46:59 +08:00
parent 4612928398
commit d1b9922c6b
111 changed files with 24287 additions and 1581 deletions
@@ -0,0 +1,87 @@
+# CelerisLab/kernels/compiler.py
+
+import subprocess
+import re
+import os
+
+from .utils import FlowFieldConfig, CudaConfig
+
+
+def kernel_path(file_name: str) -> str:
+    current_dir = os.path.dirname(os.path.abspath(__file__))
+    return os.path.join(current_dir, "kernels", file_name)
+
+
+def read_lines(file_path):
+    with open(file_path, "r") as file:
+        lines = file.readlines()
+    return lines
+
+
+def write_lines(file_path, lines):
+    with open(file_path, "w") as file:
+        file.writelines(lines)
+
+
+def modify_macro(lines, macro_name, new_value):
+    pattern = re.compile(rf"(#define\s+{macro_name}\s+)(\S+)")
+    for i, line in enumerate(lines):
+        if pattern.match(line):
+            lines[i] = pattern.sub(rf"\g<1>{new_value}", line)
+            break
+    return lines
+
+def modify_const(lines, const_name, new_type, new_value):
+    pattern = re.compile(rf"(__constant__\s+)(\S+\s+{const_name}\s*=\s*)([^;]+)(;)")
+    for i, line in enumerate(lines):
+        if pattern.match(line):
+            lines[i] = pattern.sub(rf"\g<1>{new_type} {const_name} = {new_value}\4", line)
+            break
+    return lines
+
+def compile_kernel():
+    subprocess.run(
+        [
+            "nvcc",
+            "-ptx",
+            kernel_path("kernel.cu"),
+            "-o",
+            kernel_path("kernel.ptx"),
+        ]
+    )
+
+def config_kernal(config_cuda: CudaConfig, config_field: FlowFieldConfig):
+    lines = read_lines(kernel_path("macros.h"))
+    lines = modify_macro(lines, "MULT_GPU", config_cuda.multi_gpu)
+    lines = modify_macro(lines, "NT", config_cuda.threads_per_block)
+    lines = modify_macro(lines, "X_1U", config_cuda.unit_dimensions[0])
+    lines = modify_macro(lines, "Y_1U", config_cuda.unit_dimensions[1])
+    lines = modify_macro(lines, "Z_1U", config_cuda.unit_dimensions[2])
+
+    if config_field.data_type == "FP32":
+        lb_type = "float"
+    else:
+        raise ValueError(f"Unsupported data type {config_field.data_type}.")
+    lines = modify_macro(lines, "LBtype", lb_type)
+    lines = modify_macro(lines, "UX", config_field.field_dim_in_U[0])
+    lines = modify_macro(lines, "UY", config_field.field_dim_in_U[1])
+    lines = modify_macro(lines, "UZ", config_field.field_dim_in_U[2])
+    lines = modify_macro(lines, "NX", config_field.field_dim_in_U[0] * config_cuda.unit_dimensions[0])
+    lines = modify_macro(lines, "NY", config_field.field_dim_in_U[1] * config_cuda.unit_dimensions[1])
+    lines = modify_macro(lines, "NZ", config_field.field_dim_in_U[2] * config_cuda.unit_dimensions[2])
+    lines = modify_macro(lines, "DIM", config_field.dimensionality)
+    lines = modify_macro(lines, "NQ", config_field.lattice)
+    lines = modify_macro(lines, "VIS", config_field.viscosity)
+    lines = modify_macro(lines, "U0", config_field.velocity)
+
+    write_lines(kernel_path("macros.h"), lines)
+
+def config_object(n_obj: int):
+    lines = read_lines(kernel_path("macros.h"))
+    lines = modify_macro(lines, "N_OBJS", n_obj)
+    write_lines(kernel_path("macros.h"), lines)
+
+def config_sensor(n_sen: int):
+    lines = read_lines(kernel_path("macros.h"))
+    lines = modify_macro(lines, "N_SENS", n_sen)
+    write_lines(kernel_path("macros.h"), lines)