132 lines
4.8 KiB
Python
132 lines
4.8 KiB
Python
# CelerisLab/cuda/compiler.py
|
|
|
|
import subprocess
|
|
import re
|
|
import os
|
|
|
|
from ..common.utils import FlowFieldConfig, CudaConfig
|
|
|
|
|
|
def kernel_path(file_name: str) -> str:
|
|
# kernels are in lbm/kernels/
|
|
current_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
|
|
return os.path.join(current_dir, "lbm", "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 compile_kernel_v2():
|
|
"""Compile the new modular kernel (kernel_v2.cu → kernel_v2.ptx)."""
|
|
subprocess.run(
|
|
[
|
|
"nvcc",
|
|
"-ptx",
|
|
kernel_path("kernel_v2.cu"),
|
|
"-o",
|
|
kernel_path("kernel_v2.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_kernal_v2(config_cuda: CudaConfig, config_field: FlowFieldConfig,
|
|
collision_model: int = 2,
|
|
streaming_model: int = 0,
|
|
store_precision: int = 0,
|
|
use_ddf_shifting: int = 0,
|
|
use_les: int = 0,
|
|
les_cs: float = 0.16):
|
|
"""Configure macros.h for the new modular kernel architecture.
|
|
|
|
Args:
|
|
collision_model: 0=SRT, 1=TRT, 2=MRT (default)
|
|
streaming_model: 0=double-buffer (default), 1=Esoteric-Pull
|
|
store_precision: 0=FP32 (default), 1=FP16S, 2=FP16C
|
|
use_ddf_shifting: 0=off (default), 1=on
|
|
use_les: 0=off (default), 1=Smagorinsky LES
|
|
les_cs: Smagorinsky constant C_s
|
|
"""
|
|
# First apply legacy config
|
|
config_kernal(config_cuda, config_field)
|
|
|
|
# Then apply new architecture macros
|
|
lines = read_lines(kernel_path("macros.h"))
|
|
lines = modify_macro(lines, "COLLISION_MODEL", collision_model)
|
|
lines = modify_macro(lines, "STREAMING_MODEL", streaming_model)
|
|
lines = modify_macro(lines, "STORE_PRECISION", store_precision)
|
|
lines = modify_macro(lines, "USE_DDF_SHIFTING", use_ddf_shifting)
|
|
lines = modify_macro(lines, "USE_LES", use_les)
|
|
lines = modify_macro(lines, "LES_CS", f"{les_cs:.6f}f")
|
|
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) |