TPU-MLIR 正式开源!跟算能一起打造下一代 AI 编译器

近几年,以深度学习为代表的AI技术发展迅猛,深度神经网络广泛应用于各行各业。

算能TPU编译器开源,架起算力硬件和AI框架的桥梁

开发者使用TensorFlow、PyTorch、PaddlePaddle等深度学习框架可以快速训练出自己的AI模型,但这些模型想要部署到TPU、GPU等终端算力设备时,需要将模型中的计算图/算子等输出为低级的机器语言更快地执行,这时AI编译器就有用武之地了。作为框架和硬件之间的桥梁,AI编译器可以实现一次代码开发,各种算力芯片复用的目标。最近,算能也将自己研发的TPU编译工具TPU-MLIR开源了。

TPU-MLIR是一个专注于 AI 芯片的TPU编译器开源工程,提供了一套完整的工具链,可以将不同框架下预训练过的深度学习模型,转化为可以在TPU上运行的二进制文件bmodel,从而实现更高效的推理。

MLIR(Multi-Level Intermediate Representation)是一种用来构建可重用和可扩展编译基础设施的新方法,由LLVM原作者Chris Lattner在Google工作时期开发,MLIR目的是做一个通用、可复用的编译器框架,解决软件碎片化,改进异构硬件的编译,显著减少构建特定领域编译器的成本。

TPU-MLIR站在MLIR巨人的肩膀上打造,现在整个工程的所有代码都已开源,向所有用户免费开放

精度与效率并存,TPU-MLIR实现主流框架全支持

TPU-MLIR相对其他编译工具,有以下几个优势:

简单:通过阅读开发手册与工程中已包含的样例,用户可以了解模型转化流程与原理,快速上手。并且,TPU-MLIR基于当前主流的编译器工具库MLIR进行设计,用户也可以通过它来学习MLIR的应用。

便捷:该工程已经提供了一套完备的工具链,用户可以直接通过现有接口快速地完成模型的转化工作,不用自己适配不同的网络。

通用:目前TPU-MLIR已经支持TFLite以及onnx格式,这两种格式的模型可以直接转化为TPU可用的bmodel。如果不是这两种格式呢?实际上onnx提供了一套转换工具,可以将现在市面上主流深度学习框架编写的模型转为onnx格式,然后就能继续转为bmodel了。

 

精度与效率并存:模型转换的过程中有时会产生精度损失,TPU-MLIR支持INT8对称和非对称量化,在大幅提高性能的同时又结合原始开发企业的Calibration与Tune等技术保证了模型的高精度。不仅如此,TPU- MLIR中还运用了大量图优化和算子切分优化技术,以保证模型的高效运行。

目前TPU-MLIR工程已被应用在算能研发的最新一代人工智能芯片BM1684X 上,搭配上芯片本身的高性能ARM内核以及相应的SDK,能够实现深度学习算法的快速部署。

实现极致性价比,打造下一代AI编译器

神经网络模型想要支持GPU计算,就要把神经网络模型里面的算子开发一个GPU版本,如果要适配TPU,又要把每个算子开发一个TPU的版本,还有的场景需要适配同种算力芯片不同型号的产品,每次都要手工编译的话,会非常耗时耗力。

 AI编译器正是为了解决上面的问题而生,TPU-MLIR的一系列自动优化工具可以节省大量的人工优化时间,让CPU上开发的模型可以平滑、零成本地移植到TPU上,实现极致性价比。

 随着Transformer等神经网络结构的出现,新的算子不断增加,不仅需要实现出来,还需要结合后端硬件的特点进行优化和测试,尽量将硬件的性能压榨到极致,这也导致算子的复杂度越来越高,调优难度变大,而且并不是所有的算子都能通过一个工具有效生成,整个AI编译器领域仍然处于一个不断完善的状态。

TPU- MLIR也需要持续的研发投入,更多的AI芯片支持,代码生成性能优化,运行时调度优化等环节还有很大提升空间,用户在使用本工程的同时也可以参与到该项目的完善与改进工作中,跟算能一起打造一款站在时代前沿的AI编译器!

猜你喜欢

转载自www.oschina.net/news/206300/tpu-mlir-released