""" c_lab.py —— 在 Python 里编译、加载和调用 C 代码（前置讲 P1 的工具箱） ===================================================================== 【这个文件要回答的问题】 写 CUDA 内核之前，先要能读懂、写出、调用最普通的 C 代码。这个模块把"编译 C → 得到共享库 / 可执行文件 → 从 Python 里调用"这一串步骤包装起来，让 P1 的脚本只关心概念本身。 【本文件提供】 find_compiler() —— 找本机的 C 编译器（cc / clang / gcc），找不到返回 None build(name, kind) —— 把 csrc/.c 编译成共享库（kind="shared"）或可执行文件（kind="exe"）， 结果放在 build/，源代码没变就直接复用（按源代码 + 编译选项的指纹缓存） compile_stages(name) —— 把编译拆成预处理 / 编译成汇编 / 汇编成目标文件 / 链接四步，统计每一步的产物 run_exe(name) —— 编译并运行一个 C 程序，返回它打印的内容 load_p1() —— 加载 p1lib，并给每个函数声明参数和返回值的类型（ctypes 需要知道） np_ptr(arr) —— 把 NumPy 数组的内存地址变成 C 的 float* / long* 等指针 Tensor2D —— 和 p1lib.c 里 struct Tensor2D 内存布局完全相同的 ctypes 结构体 hexbytes(b) —— 把一串字节写成 "2a 00 00 00" 的形式 【没有 C 编译器怎么办】 macOS 上运行一次 `xcode-select --install`；Linux 上安装 gcc（例如 `sudo apt install build-essential`）。 没有编译器时 build() 会抛出 CompilerMissing，脚本会跳过需要 C 的部分并提示。 【系统联系】 PyTorch、NumPy 的核心都是 C/C++（和 CUDA）写的，Python 只是"指挥官"。原课的 miniTorch 也是这样： Python 里定义张量和自动微分，CUDA 内核编译成共享库，再通过 ctypes 之类的机制调用（HW1）。 """ import ctypes import hashlib import os import shutil import subprocess import sys HERE = os.path.dirname(os.path.abspath(__file__)) CSRC = os.path.join(HERE, "csrc") BUILD = os.path.join(HERE, "build") SHARED_EXT = ".dylib" if sys.platform == "darwin" else ".so" CFLAGS = ["-O2", "-std=c11", "-Wall", "-Wextra"] class CompilerMissing(RuntimeError): """本机没有可用的 C 编译器。""" def find_compiler(): """ 找一个 C 编译器。 【返回】编译器的完整路径（字符串），找不到时返回 None 【Python 小知识】shutil.which("cc") 在 PATH 里查找名为 cc 的可执行文件，相当于终端里的 `which cc`。 """ for name in ("cc", "clang", "gcc"): path = shutil.which(name) if path: return path return None def _fingerprint(src_path, flags): """源代码内容 + 编译选项的 SHA-256 前 10 位：任何一样变了，就要重新编译。""" h = hashlib.sha256() with open(src_path, "rb") as f: h.update(f.read()) h.update(" ".join(flags).encode()) return h.hexdigest()[:10] def build(name, kind="shared", extra_flags=None, csrc=None, build_dir=None): """ 编译 csrc/.c。 【参数】 name : 源文件名（不含 .c） kind : "shared"（共享库，供 ctypes 加载）或 "exe"（可执行文件） extra_flags : 额外的编译选项列表，例如 ["-O0"]；会放在默认选项之后（后出现的同类选项生效） csrc, build_dir : 源文件目录和产物目录；None 表示本文件旁边的 csrc/ 和 build/（P4 用它编译自己的 C 代码） 【返回】产物的路径 【边界】 - 没有编译器时抛出 CompilerMissing；编译出错时抛出 RuntimeError（附编译器的错误信息）。 - 同样的源代码和选项只编译一次（产物名里带指纹）。 【Python 小知识】subprocess.run([...], capture_output=True, text=True) 运行一个外部命令， 并把它的标准输出 / 标准错误收集成字符串。 """ cc = find_compiler() if cc is None: raise CompilerMissing("找不到 C 编译器（cc / clang / gcc）") src = os.path.join(csrc or CSRC, name + ".c") flags = CFLAGS + list(extra_flags or []) if kind == "shared": flags = flags + ["-shared", "-fPIC"] ext = SHARED_EXT elif kind == "exe": ext = "" else: raise ValueError("kind 必须是 'shared' 或 'exe'") out_dir = build_dir or BUILD os.makedirs(out_dir, exist_ok=True) out = os.path.join(out_dir, f"{name}-{_fingerprint(src, flags)}{ext}") if not os.path.exists(out): r = subprocess.run([cc] + flags + ["-o", out, src], capture_output=True, text=True) if r.returncode != 0: raise RuntimeError(f"编译 {name}.c 失败：\n{r.stderr}") return out def compile_stages(name): """ 把"源代码 → 可执行文件"拆成四步，分别运行，统计每一步产物的大小。 【四步】 1. 预处理（cc -E）：展开 #include、#define，得到一个很长的纯 C 文件 2. 编译（cc -S）：把 C 翻译成这台机器的汇编语言（人能勉强读懂的机器指令） 3. 汇编（cc -c）：把汇编变成二进制的目标文件（.o） 4. 链接（cc 目标文件）：把目标文件和 C 标准库（printf、malloc 所在的库）接起来，得到可执行文件 【返回】列表 [(步骤名, 命令行, 产物路径, 产物描述)] 【边界】需要 C 编译器；产物放在 build/stages/ 下，每次调用都会重新生成。 """ cc = find_compiler() if cc is None: raise CompilerMissing("找不到 C 编译器") src = os.path.join(CSRC, name + ".c") d = os.path.join(BUILD, "stages") os.makedirs(d, exist_ok=True) pre, asm, obj, exe = (os.path.join(d, name + ext) for ext in (".i", ".s", ".o", "")) steps = [("预处理", [cc, "-E", src, "-o", pre], pre), ("编译成汇编", [cc, "-O2", "-S", src, "-o", asm], asm), ("汇编成目标文件", [cc, "-O2", "-c", src, "-o", obj], obj), ("链接", [cc, obj, "-o", exe], exe)] out = [] for title, cmd, path in steps: r = subprocess.run(cmd, capture_output=True, text=True) if r.returncode != 0: raise RuntimeError(f"{title}失败：\n{r.stderr}") size = os.path.getsize(path) if path.endswith((".i", ".s")): with open(path, errors="replace") as f: n_lines = sum(1 for _ in f) desc = f"{n_lines:,} 行文本" else: desc = f"{size:,} 字节二进制" short = " ".join(os.path.basename(c) if os.path.isabs(c) else c for c in cmd) out.append((title, short, path, desc)) return out def run_exe(name): """ 编译并运行 csrc/.c，返回它打印到标准输出的文字。 【边界】程序返回非零时抛出 RuntimeError。 """ exe = build(name, kind="exe") r = subprocess.run([exe], capture_output=True, text=True) if r.returncode != 0: raise RuntimeError(f"{name} 运行失败：{r.stderr}") return r.stdout # --------------------------------------------------------------------- # ctypes：告诉 Python 每个 C 函数的参数和返回值是什么类型 # --------------------------------------------------------------------- c_float_p = ctypes.POINTER(ctypes.c_float) c_double_p = ctypes.POINTER(ctypes.c_double) c_int_p = ctypes.POINTER(ctypes.c_int) c_size_t_p = ctypes.POINTER(ctypes.c_size_t) class Tensor2D(ctypes.Structure): """ 和 p1lib.c 里的 struct Tensor2D 一一对应：字段的顺序和类型必须完全相同， 这样 Python 构造出来的这块内存，C 函数才能按同样的布局去读。 【Python 小知识】ctypes.Structure 的子类用 _fields_ 列出 (字段名, C 类型)。 """ _fields_ = [("data", c_float_p), ("rows", ctypes.c_long), ("cols", ctypes.c_long), ("stride_row", ctypes.c_long), ("stride_col", ctypes.c_long)] _SIGNATURES = { # 函数名: (参数类型列表, 返回类型) "type_sizes": ([c_size_t_p], None), "ptr_diff_bytes_float": ([c_float_p, ctypes.c_long], ctypes.c_long), "ptr_diff_bytes_double": ([c_double_p, ctypes.c_long], ctypes.c_long), "read_by_index": ([c_float_p, ctypes.c_long], ctypes.c_float), "read_by_pointer": ([c_float_p, ctypes.c_long], ctypes.c_float), "sum_floats": ([c_float_p, ctypes.c_long], ctypes.c_float), "sum_floats_double_acc": ([c_float_p, ctypes.c_long], ctypes.c_double), "do_nothing": ([ctypes.c_int], ctypes.c_int), "swap_by_value": ([ctypes.c_int, ctypes.c_int], None), "swap_by_pointer": ([c_int_p, c_int_p], None), "my_malloc": ([ctypes.POINTER(ctypes.c_void_p), ctypes.c_size_t], ctypes.c_int), "my_free": ([ctypes.c_void_p], None), "fill_squares": ([c_float_p, ctypes.c_long], None), "touch_alloc": ([ctypes.c_size_t], ctypes.c_void_p), "get2d": ([c_float_p, ctypes.c_long, ctypes.c_long, ctypes.c_long, ctypes.c_long], ctypes.c_float), "tensor_get": ([ctypes.POINTER(Tensor2D), ctypes.c_long, ctypes.c_long], ctypes.c_float), "tensor_row_sums": ([ctypes.POINTER(Tensor2D), c_float_p], None), "matmul_naive": ([c_float_p, c_float_p, c_float_p, ctypes.c_long, ctypes.c_long, ctypes.c_long], None), "sum_by_rows": ([c_float_p, ctypes.c_long, ctypes.c_long], ctypes.c_double), "sum_by_cols": ([c_float_p, ctypes.c_long, ctypes.c_long], ctypes.c_double), } def load_p1(extra_flags=None): """ 编译（如有必要）并加载 p1lib，声明好每个函数的签名。 【为什么要声明签名】ctypes 默认把所有参数当 int、返回值当 int。如果不声明， 传一个 float 进去、或者让它返回 double，结果就是一串莫名其妙的数字（甚至崩溃）。 【参数】extra_flags：额外编译选项（例如 ["-O0"] 编译一个不优化的版本作对比） 【返回】ctypes.CDLL 对象，用法：lib.sum_floats(ptr, n) """ lib = ctypes.CDLL(build("p1lib", "shared", extra_flags)) for fname, (args, ret) in _SIGNATURES.items(): f = getattr(lib, fname) f.argtypes = args f.restype = ret return lib def np_ptr(arr, ctype=ctypes.c_float): """ 取 NumPy 数组数据区的地址，包装成 C 指针（不复制数据）。 【参数】arr：NumPy 数组；ctype：元素的 C 类型（float32 用 c_float，int64 用 c_long……） 【返回】ctypes 指针对象 【边界】 - 调用者要保证 arr 的 dtype 和 ctype 一致，而且在 C 用完之前 arr 不被释放。 - 只传首地址：C 函数默认按"连续存放"去读，非连续的数组（例如转置）要么先变成连续的， 要么把步长也一起传进去（03 号脚本的 Tensor2D）。 【Python 小知识】arr.ctypes.data_as(...) 是 NumPy 提供的接口，arr.ctypes.data 是地址本身（一个整数）。 """ return arr.ctypes.data_as(ctypes.POINTER(ctype)) def hexbytes(b): """把 bytes 写成 "2a 00 00 00"。""" return " ".join(f"{x:02x}" for x in b)