tinygrad 多后端运行时深度指南DEV 变量、设备接口与跨框架零拷贝互操作【免费下载链接】tinygradYou like pytorch? You like micrograd? You love tinygrad! ❤️项目地址: https://gitcode.com/GitHub_Trending/tiny/tinygradtinygrad 内置了面向 NVIDIA/AMD/高通/Apple/NVIDIA-CUDA/OpenCL/CPU/WebGPU 等多类硬件的运行时Runtime后端本文以 docs/runtime.md 为主线系统讲解各后端的编译选项与硬件要求、DEV环境变量的完整语法、AMD/NV 设备接口KFD/PCI/USB/NVK的选型以及通过Tensor.from_blob与 PyTorch、OpenCL 进行零拷贝互操作的方法。读完本文你将能够根据硬件环境正确选择并强制指定 tinygrad 运行时与渲染器为 CPU 后端配置精确的架构特性并在不复制数据的前提下让 tinygrad 与 PyTorch 共享 GPU 张量。运行时总览tinygrad 支持哪些后端tinygrad 通过tinygrad/runtime/目录下的ops_*.py模块为每种硬件提供独立运行时。默认情况下tinygrad 会根据可用硬件自动选择运行时你也可以用环境变量强制指定默认运行时例如DEVCPU。运行时说明编译选项依赖要求NV为 NVIDIA GPU 提供加速nvrtc默认PTXDEVNV:PTXAmpere/Ada/Blackwell 系列 GPU。可通过DEV变量选择接口详见下文 NV 接口。AMD为 AMD GPU 提供加速LLVMDEVAMD:LLVMHIP/COMGRDEVAMD:HIPCDNA3、CDNA4、RDNA3 或 RDNA4 GPU。可通过DEV变量选择接口详见下文 AMD 接口。QCOM为高通QualcommGPU 提供加速—6xx 系列 GPUMETAL利用 Metal 在 Apple 设备上加速—M1 Macbfloat支持需要 Metal 3.0CUDA利用 CUDA 在 NVIDIA GPU 上加速nvrtc默认PTXDEVCUDA:PTX支持 CUDA 的 NVIDIA GPUCL使用 OpenCL 在 GPU 上加速—兼容 OpenCL 2.0 的设备CPU使用 clang 或 llvm 编译器在 CPU 上运行Clang JIT默认LLVM IRDEVCPU:LLVM系统PATH中存在clang编译器。可通过DEV变量附加架构参数详见下文 CPU 架构。WEBGPU使用 Dawn WebGPU 引擎在 GPU 上运行用于 Google Chrome—安装并可在系统中发现的 Dawn 库对应二进制pydawn v0.3.0除上述后端外tinygrad/runtime/还包含面向其他场景的运行时例如ops_disk.py磁盘张量、ops_npy.pynpy 张量、ops_python.py纯 Python 解释执行、ops_null.py空设备用于调度/编译调试与ops_dsp.py等它们共同组成了完整的运行时生态。从源码结构看设备注册表Device见 tinygrad/device.py维护了一个ALL_DEVICES列表[METAL, AMD, NV, CUDA, QCOM, CL, CPU, DSP, WEBGPU]并据此实现默认设备的自动探测_select_device会按顺序尝试打开每个可用设备并返回第一个成功者当DEV变量被显式设置时Device.DEFAULT会直接返回DEV.device从而覆盖自动选择。开发者还可通过get_available_devices()枚举当前机器上真正可用的设备列表。DEV 变量后端、渲染器、架构与接口的统一选择语法DEV是 tinygrad 中最核心的运行时选择入口其语法在 docs/env_vars.md 中有完整说明并在 tinygrad/helpers.py 的Target.parse中实际解析。DEV用于指定目标设备、目标渲染器以及该设备的架构三者以冒号:分隔渲染器和架构可以省略省略时 tinygrad 会自动确定合适的取值。此外DEV还可以通过加号前缀指定访问设备的接口interface如PCI、USB接口可位于目标三元组之前。接口、设备索引与设备三元组整体构成一条 target 字符串多条 target 之间以分号;分隔。DEV取值示例与含义DEV内容含义AMD使用 AMD 设备AMD:LLVM使用 AMD 设备渲染器为 LLVMNV:CUDA:sm_70使用 NV 设备渲染器为 CUDA目标架构 sm_70AMD::gfx950使用 AMD 设备目标架构 gfx950USBAMD通过 USB 接口使用 AMD 设备CPU:LLVM使用 CPU 设备渲染器为 LLVMCPU:LLVM:x86_64,znver2,avx2,-avx512f使用 CPU 设备、LLVM 渲染器并附加架构参数详见 CPU 架构对应的解析逻辑在Target.parsetinygrad/helpers.py先按切分接口与设备三元组再按:切分device[:renderer[:arch]]最终生成包含device、renderer、arch、interface、indices五个字段的Target。设备名与渲染器名会被自动转为大写架构名保持原样。DEV同时是一个ContextVartinygrad/helpers.py因此除了在 shell 中设置环境变量如DEVCL DEBUG4 python3 -m pytest还可以在 Python 代码里用Context临时切换而无需修改全局环境from tinygrad.helpers import Context # 仅在该 with 块内使用 CPU 后端 with Context(DEVCPU): a Tensor.ones(10, 10) a * 2 # 或者用装饰器限定某个函数 Context(DEVCUDA) def run_on_cuda(): ...与之配套的调试开关是DEBUG同为ContextVar默认 0级别越高输出越详细DEBUG1列出正在使用的设备DEBUG2输出每个 kernel 的执行时间、内存占用与带宽等性能指标DEBUG3输出 kernel 级优化结果DEBUG4输出生成的 kernel 代码DEBUG5/6输出计算中间表示 UOps普通/线性化两种视图DEBUG7输出针对目标硬件生成的汇编代码。渲染器与接口的实际选择发生在 tinygrad/device.py 的Compiled._select_renderer与_select_iface中它们根据DEV.target()解析出的renderer/interface字段从设备注册的候选列表中按名称筛选再通过select_first_inited依次尝试初始化取第一个成功者作为默认这也解释了为什么省略渲染器/接口时 tinygrad 能自动挑选可用的实现。跨框架互操作Tensor.from_blob 与零拷贝数据共享tinygrad 提供了与 OpenCL 和 PyTorch 的互操作能力通过Tensor.from_blobAPI 直接基于外部内存指针创建张量实现框架间的零拷贝数据共享。其实现位于 tinygrad/tensor.py先按 shape 分配一个空张量再通过Buffer.allocate(external_ptrptr)将该指针直接挂接到张量底层缓冲上tinygrad 不会取得数据所有权。重要约束使用外部内存指针创建 tinygrad 张量时必须保证这些指针在整个 tinygrad 张量生命周期内保持有效否则会造成内存损坏即悬垂指针访问。CUDA/METAL 与 PyTorch 互操作可以在 PyTorch 与 tinygrad 之间直接共享 CUDA/MPS 张量无需数据拷贝from tinygrad.dtype import _from_torch_dtype tensor1 torch.tensor([1.0, 2.0, 3.0], devicetorch.device(cuda)) tiny_tensor1 Tensor.from_blob(tensor1.data_ptr(), tensor1.shape, dtype_from_torch_dtype(tensor1.dtype), deviceCUDA) # tinygrad 计算前需要同步 mps确保数据有效。 if data.device.type mps: torch.mps.synchronize() else: torch.cuda.synchronize() x (tiny_tensor1 1).realize()要点说明_from_torch_dtype定义于 tinygrad/dtype.py负责把 PyTorch 的torch.dtype映射为 tinygrad 的DType避免手工转换出错deviceCUDA也可替换为METAL以共享 MPS 张量此时同步用torch.mps.synchronize()由于 GPU 上 PyTorch 的操作是异步提交的在 tinygrad 读取数据前必须先调用对应设备的同步函数否则读到的可能是未完成的旧数据调用.realize()强制 tinygrad 实际执行计算并落盘结果。QCOM OpenCL 互操作QCOM 后端支持 OpenCL 互操作可以直接访问 OpenCL 内存缓冲# 创建原始 opencl buffer。 cl_buf cl.clCreateBuffer(cl_context, cl.CL_MEM_READ_WRITE, 0x100, None, status : ctypes.c_int32()) # 提取指针 cl_buf_desc_ptr to_mv(ctypes.addressof(cl_buf), 8).cast(Q)[0] rawbuf_ptr to_mv(cl_buf_desc_ptr, 0x100).cast(Q)[20] # offset 0xA0 处是原始 gpu 指针。 # 创建 tiny tensor tiny Tensor.from_blob(rawbuf_ptr, (8, 8), dtypedtypes.int, deviceQCOM)图像Image2D同样支持互操作# 创建 cl image。 cl_img cl.clCreateImage2D(cl_context, cl.CL_MEM_READ_WRITE, cl.cl_image_format(cl.CL_RGBA, cl.CL_FLOAT), w, h, 0, None, status : ctypes.c_int32()) # 提取指针 cl_buf_desc_ptr to_mv(ctypes.addressof(cl_img), 8).cast(Q)[0] rawbuf_ptr to_mv(cl_buf_desc_ptr, 0x100).cast(Q)[20] # offset 0xA0 处是原始 gpu 指针。 # 创建 tiny tensor tiny Tensor.from_blob(rawbuf_ptr, (h*w*4,), dtypedtypes.imagef((h,w)), deviceQCOM)这里通过 ctypes 读取 OpenCL 缓冲描述符的内存布局在偏移 0xA0 处取出真正的 GPU 原始指针再交给Tensor.from_blobdtypes.imagef((h,w))表示按 (h,w) 的 image 格式解释数据。注意上述指针偏移量是针对特定 OpenCL 实现与平台的布局假设实际使用时需以目标平台的描述符布局为准。AMD 接口KFD / PCI / USBAMD 后端支持多种与设备通信的接口定义于 tinygrad/runtime/ops_amd.pyKFD使用 amdgpu 内核驱动PCI使用 AM 驱动详见 docs/developer/am.md该文档介绍底层驱动实现USB面向 asm24xx 芯片的 USB3 接口。通过把DEV环境变量的接口部分设置为上述值之一即可强制指定接口例如DEVPCIAMD表示通过 PCI 接口访问 AMD 设备。注意当设置为PCI时可能会将 GPU 从 amdgpu 驱动上解绑unbind。这意味着该 GPU 将不再被系统的 amdgpu 驱动管理操作前需确认不影响同机其他用途。NV 接口NVK / PCINV 后端支持多种与设备通信的接口定义于 tinygrad/runtime/ops_nv.pyNVK使用 nvidia 驱动PCI使用 NV 驱动底层实现见 tinygrad/runtime/support/nv/nvdev.py。与 AMD 类似接口通过DEV变量指定例如DEVPCINV。CPU 架构配置CPU 渲染器Clang JIT 与 LLVM IR可以使用DEV环境变量的架构arch部分进行附加配置。CPU 架构必须指定为逗号分隔的参数列表且至少包含两个值架构族architecture family如x86_64、arm64或riscv64CPU 类型cpu type即clang的-march可接受的取值。其余逗号分隔的值被解释为CPU 特性标志feature flags当某个值以-字符开头时对应的特性标志被禁用否则该标志被启用。注意启用状态的标志不应以开头。如果 CPU 类型指定为nativetinygrad或委托的编译器会查询宿主机 CPU 类型来实际确定目标。这与 CPU 后端的默认行为一致从 tinygrad/runtime/ops_cpu.py 的源码可以看到CPUDevice初始化时默认 arch 为{amd64:x86_64, aarch64:arm64}.get(platform.machine().lower(), platform.machine()) ,native——即把宿主机器架构映射为 clang 约定名称如 amd64→x86_64、aarch64→arm64并追加native让编译器探测宿主机 CPU 特性。组合示例来自 docs/env_vars.md# x86_64 架构族、znver2 CPU 类型启用 avx2、禁用 avx512f DEVCPU:LLVM:x86_64,znver2,avx2,-avx512f# 直接用 native 让编译器探测宿主机再显式启用 fma、禁用 avx512f DEVCPU:x86_64,native,fma,-avx512f python3 -m pytest这种配置让开发者可以精确控制 CPU kernel 的指令集生成例如在支持 AVX-512 的机器上为兼容性显式禁用该特性或在交叉编译/容器场景下指定确定的微架构目标。实践小结选择 tinygrad 运行时的通用流程可归纳为四步探测可用设备不设置DEV直接运行tinygrad 会按ALL_DEVICES顺序自动选择第一个可用的后端Device.DEFAULT见 tinygrad/device.py显式指定后端DEVAMD、DEVNV、DEVCUDA、DEVCPU等避免自动探测的不确定性按需指定渲染器/架构如DEVAMD:LLVM、DEVNV:CUDA:sm_70、DEVCPU:LLVM:x86_64,znver2,avx2,-avx512f必要时切换接口与做互操作DEVUSBAMD、DEVPCINV选择设备访问通道用Tensor.from_blob与 PyTorch/OpenCL 共享内存注意指针生命周期与设备同步。完整的DEV语法与DEBUG级别对照表可参考 docs/env_vars.md运行时模块的完整实现位于 tinygrad/runtime 目录test/device下的测试用例如 test/test_device、test/test_device/test_ocl.py可作为各后端实际行为的验证参考。【免费下载链接】tinygradYou like pytorch? You like micrograd? You love tinygrad! ❤️项目地址: https://gitcode.com/GitHub_Trending/tiny/tinygrad创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考