1. Builder Python API¶

Warning

Builder Python API 当前为 Beta 版本，仅用于适配 OpenSora 模型（提供 codegen pass 功能）。

Builder Python API 提供了以下三个模块，通过 Python 接口，实现与 Builder C++ API 相同的功能:

Builder: 全局构图管理对象，用于管理构造 HLIR Module 的整个过程；
Op: 构造 HLIR Module 中算子的接口；
Type: 封装算子输入和输出的 shape 和 data type。

1.1. Builder¶

GetBuilder()¶

def GetBuilder()

简介:

全局 Builder 创建函数

返回值:

Builder 类对象

CreateInput¶

def CreateInput(type: builder::Type, func_name: str["main"])

简介:

给函数 func_name 设置输入 op

参数:

type: func_name 对应输入的 type
func_name: 创建 input op 对应的 func，默认为 main

返回值:

创建的输入 op, shape 和 dtype 对应的类型为 type

SetOutput¶

def SetOutput(outputs: list[builder::Op], func_name: str["main"])

简介:

给函数 func_name 设置输出 op

参数:

outputs: 要设置为输出的 ops
func_name: func 函数，默认为 main

Dump¶

def Dump()

简介:

在终端打印当前 HLIR Module, flags 选项固定为 PrintingFlags::ElideLargeElementsAttrs, 即 largeElementLimit = 16

Save¶

def Save()

简介:

序列化当前构造的 HLIR Module，用于 TopsGraphCompiler 编译

返回值:

字符串类型的 HLIR Module

以下是用 Python API 完成构图的简单示例：

import numpy as np
import topsgraph as tg
import topsgraph.builder as tgb

builder = tgb.GetBuilder()
type = tgb.Type([3,2], tgb.PrimitiveType.F32())
arg0 = builder.CreateInput(type)
arg1 = builder.CreateInput(type)
out  = tgb.add(arg0, arg1)
builder.SetOutput([out])
builder.Dump()

生成的 IR 如下所示：

module @hlir_module attributes {dtu_hlir.module_id = 1 : i64} {
  func.func @main(%arg0: tensor<3x2xf32>, %arg1: tensor<3x2xf32>) -> tensor<3x2xf32> {
    %0 = "dtu_hlir.add"(%arg0, %arg1) : (tensor<3x2xf32>, tensor<3x2xf32>) -> tensor<3x2xf32>
    "dtu_hlir.return"(%0) : (tensor<3x2xf32>) -> ()
  }
}

1.2. PrimitiveType¶

static member function¶

PrimitiveType	Basic data type in python
`NONE()`	`none(for unknown data type)`
`PRED()`	`bool`
`S8()`	`int8_t`
`U8()`	`uint8_t`
`S16()`	`int16_t`
`U16()`	`uint16_t`
`F16()`	`half(not a standard python type)`
`BF16()`	`not a standard python type`
`S32()`	`int32_t`
`U32()`	`uint32_t`
`F32()`	`float`
`S64()`	`int64_t`
`U64()`	`uint64_t`
`F64()`	`double`
`TUPLE()`	`tuple(a list of above types)`

1.3. Type¶

Type 封装了 Op 的 shape 和 PrimitiveType.

Type¶

def Type(primitive_type: PrimitiveType[builder::PrimitiveType::NONE()])
def Type(shape: list[int], primitive_type: PrimitiveType)

简介:

根据不同参数类型，构造对应的 Type 对象。
CASE #1: 标量类型
CASE #2: 构造 RankedTensorType

返回值:

Type 对象

1.4. Op¶

算子 API 返回类型，封装了构图过程中具体的算子信息，当前未对外开放 API 接口。