尧图网站设计 尧图网站设计YAOTU DESIGN
ARTICLE DETAIL

资讯详情

深耕网站设计与一线实操的经验洞察。

10分钟写出第一个StableHLO程序:Python API快速上手教程(含完整代码)

10分钟写出第一个StableHLO程序:Python API快速上手教程(含完整代码) 10分钟写出第一个StableHLO程序Python API快速上手教程含完整代码【免费下载链接】stablehloBackward compatible ML compute opset inspired by HLO/MHLO项目地址: https://gitcode.com/gh_mirrors/st/stablehloStableHLO是一套向后兼容的机器学习计算算子集Backward compatible ML compute opset定位为连接机器学习框架TensorFlow、JAX、PyTorch与机器学习编译器XLA、IREE 等的可移植性中间层。本文将带你用官方StableHLO Python API基于 MLIR Python 绑定在 10 分钟内写出并运行第一个 StableHLO 程序附完整可运行代码。 先搞懂StableHLO 是什么StableHLO 本质上是 ML 领域的通用指令集框架把模型导出为 StableHLO 程序编译器再消费它从而实现一次导出、多端运行。它的核心亮点是版本化与向后兼容——每个算子都带版本号如vhlo.add_v2新版本编译器可以读懂旧版本导出的程序旧版本也能被升级到新格式后消费正向场景同样成立StableHLO 程序经 legalize 转为带版本标记的 VHLO 格式再降级到目标版本保证不同版本工具链之间的无缝互通 机制细节可阅读仓库文档 docs/compatibility.md 与 docs/vhlo.md。 第 1 步环境准备与安装 Python 绑定StableHLO 的 Python 绑定需要先构建 MLIR再启用 Python 绑定构建 StableHLO步骤如下完整说明见根目录README.md的Python章节1️⃣ 安装构建工具# Linux sudo apt install cmake ninja-build lld ccache2️⃣ 克隆仓库并准备 LLVM 源码git clone --depth1 --single-branch https://gitcode.com/gh_mirrors/st/stablehlo cd stablehlo git clone --depth1 --single-branch https://github.com/llvm/llvm-project.git # 检出 build_tools/llvm_version.txt 中指定的 LLVM 提交 hash$(cat ./build_tools/llvm_version.txt) cd llvm-project git fetch origin $hash git checkout $hash cd ..3️⃣ 构建 MLIR耗时较长M1 Pro 上约 10 分钟MLIR_ENABLE_BINDINGS_PYTHONON build_tools/build_mlir.sh ${PWD}/llvm-project/ ${PWD}/llvm-build4️⃣ 安装 Python 依赖并构建 StableHLO开启 Python 绑定pip install -r ./llvm-project/mlir/python/requirements.txt STABLEHLO_ENABLE_BINDINGS_PYTHONON ./build_tools/github_actions/ci_build_cmake.sh ${PWD}/llvm-build ${PWD}/build5️⃣ 验证绑定是否可用PYTHONPATH./build/python_packages/stablehlo python3 -c import mlir.dialects.stablehlo; print(OK)✅ 也可以跳过编译直接使用官方通过setup.py位于stablehlo/integrations/python/setup.py发布的预构建 wheel 包。 第 2 步写出第一个 StableHLO 程序完整代码下面是最经典的入门示例——用 Python 构造两个 3×4 的常量张量并求和完整源码见examples/python/stablehlo_add.pyfrom mlir import ir import mlir.dialects.stablehlo as stablehlo import mlir.dialects.func as func from mlir.ir import Context, Location, InsertionPoint, Module import numpy as np with Context() as ctx, Location.unknown(): stablehlo.register_dialect(ctx) # 1. 注册 StableHLO 方言 module Module.create() # 2. 创建 MLIR Module with InsertionPoint(module.body): func.func() def main(): a_value ir.DenseElementsAttr.get(np.zeros(shape[3,4], dtypenp.int64)) b_value ir.DenseElementsAttr.get(np.zeros(shape[3,4], dtypenp.int64)) a stablehlo.constant(a_value) # 3. 构造常量张量 b stablehlo.constant(b_value) add stablehlo.add(a, b) # 4. 插入 stablehlo.add 算子 return add assert main.func_op.verify() # 5. 验证 IR 合法性 print(str(module)) # 6. 打印生成的 StableHLO 程序 代码逐段解读步骤说明Contextregister_dialectMLIR 的编程模型要求先创建上下文并注册方言之后才能使用stablehlo.*算子DenseElementsAttr用 NumPy 数组生成常量张量的 MLIR 属性stablehlo.constant/stablehlo.add直接以 Python 函数调用的方式插入 StableHLO 算子func.func装饰器自动完成func方言函数的创建与收尾verify()在不运行编译器的情况下快速校验 IR 是否合法 第 3 步运行并检查生成的 StableHLO 程序执行后终端会打印出生成的 MLIR 文本表示module { func.func main() - tensor3x4xi64 { %0 stablehlo.constant dense0 : tensor3x4xi64 %1 stablehlo.constant dense0 : tensor3x4xi64 %2 stablehlo.add %0, %1 : tensor3x4xi64 return %2 : tensor3x4xi64 } } 恭喜你已经写出了一个合法的 StableHLO 程序。这段 IR 可以直接被 XLA、IREE 等编译器消费也可以经 MLIR 字节码格式序列化后跨版本传递序列化机制见docs/bytecode.md格式示例见docs/images/vhlo/中的兼容流程图。 学完之后进阶学习路径算子规范所有 StableHLO 算子的语义定义在 docs/spec.md机器可读的算子定义位于stablehlo/dialect/StablehloOps.tdChlo 扩展方言超越稳定算子集的扩展运算erf、atanh 等见stablehlo/dialect/ChloOps.td转 TensorFlow SavedModel仓库提供了将 StableHLO 程序含权重元数据转换为 TF SavedModel 的工具见stablehlo/integrations/python/stablehlo/savedmodel/stablehlo_to_tf_saved_model.pyC Builder API偏好 C 的话stablehlo/integrations/cpp/builder/下提供了与 Python 绑对等的流式构建接口官方文档站openxla.org/stablehlo按照以上路径从会写第一个程序到能把真实模型接入 StableHLO 生态你已经在正确的路线上了。【免费下载链接】stablehloBackward compatible ML compute opset inspired by HLO/MHLO项目地址: https://gitcode.com/gh_mirrors/st/stablehlo创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考
返回列表