如何绑定local dailect的python接口

Viewed 419

问题描述

mlir提供了绑定工具,可以根据本地定义的dailect文件生成对应的python接口。方便直接调用。

可以查看官方文档实现该问题

1 Answers

以top dailect的生成python接口为说明:

一、新建ods.py文件

首先新建文件top.py,文件内容为
python from ._top_ops_gen import *
这个目的是将top作为输入接口,在python代码中可以直接调用top,方法如下:

import mlir.dialects.top as top
new_op = top.TransposeOp(output, op, dim0, dim1, loc=Loc, ip=insert_point)

二、新建TopOps.td

新建python文件的TopOps.td。用来绑定对应的Top dailect的tablegen文件。

include "mlir/Bindings/Python/Attributes.td"
include "tpu_mlir/Dialect/Top/IR/TopOps.td"

三、修改CMakeList.txt

这一步非常重要,但是官方文档并无详细描述。通过查看torch-mlir的cmake文件得到了绑定dailect的方法。

1. include文件

首先就要include(AddMLIRPython),这里的AddMLIRPython是mlir的makefile,绑定python,必须include该makefile。

2. Structural groupings

这里需要声明source路径:

declare_mlir_python_sources(TopMLIRPythonSources)
declare_mlir_python_sources(TopMLIRPythonSources.Dialects
  ROOT_DIR "${TOP_MLIR_PYTHON_ROOT_DIR}"
  ADD_TO_PARENT TopMLIRPythonSources
)

此处使用了declare_mlir_python_sources来声明代码的目录,此处分别声明了TopMLIRPythonSources和TopMLIRPythonSources.Dialects。

3. Dialect bindings

这里就可以生成top dailect的python代码。

declare_mlir_dialect_python_bindings(
  ADD_TO_PARENT TopMLIRPythonSources.Dialects
  ROOT_DIR "${CMAKE_CURRENT_SOURCE_DIR}/tpu_mlir"
  TD_FILE dialects/TopOps.td
  SOURCES dialects/top.py
  DIALECT_NAME top)

这里声明了top dailect绑定需要的TopOps.td和top.py代码。他们都位于./tpu_mlir/dialects目录下。生成的将生成的python代码和top.py代码都放入TopMLIRPythonSources.Dialects路径下。
此处会生成_DIALECT_NAME_ops_gen.py文件,是根据TD_FILE dialects/TopOps.td来生成的。由于这里DIALECT_NAME top,因此文件的名字为_top_ops_gen.py

4. The fully assembled package of modules

将python代码打包,放入build的目的地址

set(_source_components
  TopMLIRPythonSources
)
add_mlir_python_modules(TopMLIRPythonModules
  ROOT_PREFIX "${TOP_MLIR_PYTHON_PACKAGES_DIR}/mlir"
  INSTALL_PREFIX "python/mlir/"
  DECLARED_SOURCES ${_source_components}
  )

5. 结果

通过如上步骤,build代码,可以看到install/python/mlir目录下已经有了top.py和_top_ops_gen.py文件。
WeChatd3b1dd773db7ae58a71d3876778ae6cc.png
同时,在build目录下会看到对应的文件,如下所示:
WeChat4dc183a3185a48d4540867fe5af46d01.png

该问题已经解决,代码为 top dialects bind python