|
| 1 | +# DeepX IR(deepxir)规范 |
| 2 | + |
| 3 | +## 1. 类型系统 |
| 4 | + |
| 5 | +### 基础数据类型 |
| 6 | +``` |
| 7 | +type f16, f32, f64, bf16, bf8 // 浮点类型 |
| 8 | +type i8, i16, i32, i64, u8 // 整数类型 |
| 9 | +type bool // 布尔类型 |
| 10 | +``` |
| 11 | + |
| 12 | +### 动态长度类型 |
| 13 | +``` |
| 14 | +list<type> // list 可以和以上基础类型组合 |
| 15 | +``` |
| 16 | + |
| 17 | +### 类型约束 |
| 18 | +``` |
| 19 | +f32|f64 // 支持两种/多种 类型之一 |
| 20 | +``` |
| 21 | + |
| 22 | +### Tensor 类型模板 |
| 23 | +``` |
| 24 | +type tensor<shape, elem_type> |
| 25 | +``` |
| 26 | +- shape 格式:dim1xdim2x...xdimN,或使用 `?` 表示动态维度。 最后一个x后的是精度。 |
| 27 | +- 示例:`tensor<10x20xf32>`, `tensor<?x?xi32>` |
| 28 | + |
| 29 | +tensor 也可以没有 shape 和 dtype 的约束,例如: |
| 30 | +``` |
| 31 | +deepxir addscalar(A:tensor, b:i8|i16|i32|i64) -> (c:tensor) { ... } |
| 32 | +``` |
| 33 | +表示任意 shape、任意 dtype 的 tensor 都可作为参数。 |
| 34 | + |
| 35 | +### 动态维度变量 |
| 36 | +- `?` 任意数字 |
| 37 | +- `?1` 动态维度变量 1 |
| 38 | +- `?2` 动态维度变量 2(用于表示同名变量处维度需一致) |
| 39 | +- 示例:`tensor<?1x?2xf32>` |
| 40 | + |
| 41 | +## 2. IR 定义格式 |
| 42 | + |
| 43 | +语法示例: |
| 44 | +``` |
| 45 | +deepxir ir_name(ro_p1:type1, ro_param2:type2, ...) -> (w_p1:type3, w_p2:type4, ...) |
| 46 | +{ |
| 47 | + // 函数体:IR 操作序列 |
| 48 | + operation_name(ro_p1, ro_p2) -> w_p1 |
| 49 | + operation_name(ro_p2, ro_p2) -> w_p2 |
| 50 | +} |
| 51 | +``` |
| 52 | +- `deepxir` 为关键字,也可使用 `function`、`func` 等。 |
| 53 | +- 参数遵循“左读右写”规则(无返回值;通过写入参数实现输出)。 |
| 54 | +- 参数类型支持:`tensor`、`list<tensor>`、基础类型,以及基础类型的 list。 |
| 55 | + |
| 56 | +## 3. 设计思考 |
| 57 | +DeepX IR 采用简洁的文本格式表示张量类型约束、运算定义与运算体,便于阅读与解析。 |
| 58 | +deepx不是ssa,调用时,依然遵循左读右写的参数列表原则,右写的参数列表支持多个。 |
| 59 | + |
| 60 | +## 4. 具体示例 |
| 61 | + |
| 62 | +### 示例 1:融合 Linear + 归一化 |
| 63 | +``` |
| 64 | +deepxir fused_linear_norm( |
| 65 | + A: tensor<?1x?2xf32>, |
| 66 | + W: tensor<?2x?3xf32>, |
| 67 | + b: tensor<?3xf32>, |
| 68 | + axis: i32, |
| 69 | + keepdims: bool |
| 70 | +) -> (out: tensor<?1x?3xf32>) { |
| 71 | + newtensor(?1x?3, f32)->(mm) |
| 72 | + matmul(A, W)-> (mm) |
| 73 | + newtensor(?1x?3, f32)-> bias |
| 74 | + add(mm, b)-> bias |
| 75 | + deltensor(mm)-> mm |
| 76 | + newtensor(?1, f32)-> mean |
| 77 | + sum(bias, axis, keepdims)-> mean |
| 78 | + newtensor(?1x?3, f32)-> centered |
| 79 | + sub(bias, mean)-> centered |
| 80 | + deltensor(bias)-> bias |
| 81 | + deltensor(mean)-> mean |
| 82 | + newtensor(?1x?3, f32)-> sq |
| 83 | + mul(centered, centered)-> sq |
| 84 | + deltensor(centered)-> centered |
| 85 | + newtensor(?1, f32)-> var |
| 86 | + sum(sq, axis, keepdims)-> var |
| 87 | + deltensor(sq)-> sq |
| 88 | + constant(1e-5)-> eps |
| 89 | + newtensor(?1, f32)-> var_eps |
| 90 | + add(var, eps)-> var_eps |
| 91 | + deltensor(var)-> var |
| 92 | + deltensor(eps)-> eps |
| 93 | + newtensor(?1, f32)-> std |
| 94 | + sqrt(var_eps)-> std |
| 95 | + deltensor(var_eps)-> var_eps |
| 96 | + div(std, std)-> std |
| 97 | + deltensor(std)-> std |
| 98 | + div(centered, std)-> out |
| 99 | +} |
| 100 | +``` |
| 101 | + |
| 102 | +下面给出一个完整的 `deepxir` 调用示例:在一个 IR 中先构造输入张量和辅助参数,然后调用 `fused_linear_norm`,输出 `out`。 |
| 103 | + |
| 104 | +``` |
| 105 | +deepxir example_use_fused_linear_norm() -> (out: tensor<2x3xf32>) { |
| 106 | + newtensor([2,4], f32)-> A |
| 107 | + newtensor([4,3], f32)-> W |
| 108 | + newtensor([3], f32)-> b |
| 109 | + fused_linear_norm(A, W, b, 1, false) -> out |
| 110 | +} |
| 111 | +``` |
| 112 | + |
| 113 | +该示例展示了如何在 IR 中构造必要的张量/参数并调用 `fused_linear_norm`,其中 `out` 的类型为 `tensor<2x3xf32>`,与 `W` 的列数和 `A` 的行数对应。 |
| 114 | + |
| 115 | +### 示例 2:融合 Attention score + Softmax |
| 116 | +``` |
| 117 | +deepxir fused_attention_scores( |
| 118 | + Q: tensor<?x?xf32>, |
| 119 | + K: tensor<?x?xf32>, |
| 120 | + axis: list<i32>, |
| 121 | + keepdims: bool, |
| 122 | + shape_scores: list<i32>, |
| 123 | + shape_sum: list<i32> |
| 124 | +) -> (out: tensor<?x?xf32>) { |
| 125 | + newtensor(shape_scores, f32)-> scores_tmp |
| 126 | + matmul(Q, K)-> scores_tmp |
| 127 | + newtensor(shape_scores, f32)-> exp_tmp |
| 128 | + exp(scores_tmp)-> exp_tmp |
| 129 | + deltensor(scores_tmp)-> scores_tmp |
| 130 | + newtensor(shape_sum, f32)-> sum_tmp |
| 131 | + sum(exp_tmp, axis, keepdims)-> sum_tmp |
| 132 | + div(exp_tmp, sum_tmp)-> out |
| 133 | + deltensor(exp_tmp)-> exp_tmp |
| 134 | + deltensor(sum_tmp)-> sum_tmp |
| 135 | +} |
| 136 | +``` |
0 commit comments