Skip to content

Commit 1cc2618

Browse files
gongchensugongchensu
committed
Add python interface documentation.
1 parent 61e1fe2 commit 1cc2618

File tree

16 files changed

+671
-2
lines changed

16 files changed

+671
-2
lines changed

‎README.md‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
## 文档目录
1010

1111
### InfiniCore
12-
- [`Python APIs`]
12+
- [`Python APIs`](python/README.md)
1313

1414
- [`C++ APIs`]
1515

@@ -22,7 +22,7 @@
2222
- [`InfiniCCL`]：统一集合通信库，提供常用的集合通信功能，包括点对点、广播、聚合等。
2323

2424

25-
[`Python APIs`]:README.md
25+
[`Python APIs`]:python/README.md
2626
[`C++ APIs`]:README.md
2727
[`InfiniRT`]:/infinirt/README.md
2828
[`InfiniOP`]:/infiniop/README.md

‎python/README.md‎

Lines changed: 115 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,115 @@
1+
# `infinicore` Python 前端
2+
3+
*InfiniCore* 提供了与 C++ 前端一致的 Python 封装，位于 `python/infinicore/`。该模块通过 `pybind11` 将核心张量、算子与设备上下文暴露给 Python，便于在推理框架或调试脚本中快速集成。
4+
5+
## 模块结构
6+
7+
| 符号 | 说明 |
8+
| --- | --- |
9+
| `device` | 设备句柄类 (`python/infinicore/device.py`)，支持 `"cuda:0"`、`device(\"cpu\", 0)` 等写法或复用已有实例。 |
10+
| `dtype` / `float16` 等 | 数据类型枚举 (`python/infinicore/dtype.py`)。 |
11+
| `Tensor` | 张量包装类 (`python/infinicore/tensor.py`)，内部封装底层 `_infinicore` 对象。 |
12+
| `empty` / `zeros` / `ones` / `empty_like` 等 | 张量构造函数（`python/infinicore/tensor.py`），默认要求显式传入 `dtype` 与 `device`。 |
13+
| 顶层算子（`add`、`matmul`、`rearrange`、`attention`） | 暴露在 `infinicore` 命名空间下，对应实现位于 `python/infinicore/ops/`。 |
14+
| `infinicore.nn` | 神经网络相关模块集合，未来可扩展更多组件。 |
15+
| `infinicore.nn.functional` | 函数式算子集合 (`python/infinicore/nn/functional.py`)。 |
16+
| `use_ntops` / `infinicore.ntops` | 若系统安装 `ntops` 包，将自动置位 `use_ntops=True` 并暴露原始模块。 |
17+
18+
所有符号在包的 `__init__.py` 中进行了显式导出，可直接通过 `import infinicore as ic` 后使用。
19+
20+
相关导出定义见 `InfiniCore/python/infinicore/__init__.py`。
21+
22+
## API 索引
23+
24+
- **核心对象**
25+
- [`device`](device/README.md)
26+
- [`dtype`](dtype/README.md)
27+
- [`Tensor`](tensor/README.md)
28+
- 构造函数：[`empty`](tensor/README.md#构造函数)、[`zeros`](tensor/README.md#构造函数)、[`ones`](tensor/README.md#构造函数)、[`empty_like`](tensor/README.md#构造函数)、[`from_blob`](tensor/README.md#构造函数)
29+
- 常用方法：[`copy_`](tensor/README.md#tensor-类)、[`to`](tensor/README.md#tensor-类)、[`permute`](tensor/README.md#tensor-类)、[`view`](tensor/README.md#tensor-类)、[`debug`](tensor/README.md#tensor-类)
30+
- [`use_ntops` 协作说明](#与-ntops-的协作)
31+
- **顶层算子**
32+
- [`add`](ops/add/README.md)
33+
- [`matmul`](ops/matmul/README.md)
34+
- [`rearrange`](ops/rearrange/README.md)
35+
- [`attention`](ops/attention/README.md)
36+
- **函数式算子**
37+
- [`causal_softmax`](nn/functional/causal_softmax/README.md)
38+
- [`rms_norm`](nn/functional/rms_norm/README.md)
39+
- [`silu`](nn/functional/silu/README.md)
40+
- [`swiglu`](nn/functional/swiglu/README.md)
41+
- **更多参考**
42+
- [`infinicore.ops` 索引](ops/README.md)
43+
- [`nn` 模块概览](nn/README.md)
44+
45+
## 张量与构造函数
46+
47+
`Tensor` 是对底层 `_infinicore.Tensor` 的 Python 包装，常用接口包括：
48+
49+
- `shape` / `ndim` / `size(dim)` / `stride(dim)`：获取张量维度与步长信息。
50+
- `dtype` / `device`：返回 `dtype` 与 `device` 包装类。
51+
- `numel()` / `is_contiguous()`：查看张量元素数量与存储布局。
52+
- `copy_(src)` / `to(...)`：执行数据拷贝与跨设备搬运。
53+
- `contiguous()` / `permute(dims)` / `view(shape)` / `as_strided(size, stride)`：布局调整与视图操作。
54+
- `debug(filename=None)`：将张量内容打印或输出到二进制文件。
55+
56+
常用构造函数包括 `empty`、`strided_empty`、`zeros`、`ones`、`from_blob`、`strided_from_blob`、`empty_like` 等：
57+
58+
```python
59+
import infinicore as ic
60+
61+
cpu = ic.device("cpu")
62+
a = ic.empty((4, 8), dtype=ic.float16, device=cpu)
63+
b = ic.ones((4, 8), dtype=ic.float16, device=cpu)
64+
a.copy_(b)
65+
```
66+
67+
> 注意：这些函数要求显式传入 `dtype` 与 `device`，避免隐式从 PyTorch/TensorFlow 对象推断。
68+
69+
## 顶层算子 (`infinicore.*`)
70+
71+
详见 [`ops` 文档索引`](ops/README.md) 及各算子文档。
72+
73+
## 函数式算子 (`infinicore.nn.functional`)
74+
75+
详见 [`nn.functional` 文档](nn/functional/README.md) 及子目录。
76+
77+
## 运行时上下文
78+
79+
- `_infinicore` 在进程内维护运行时状态；创建张量时请显式传入 `device`，并保持算子的所有输入位于同一设备。
80+
- 如需强制同步，可调用 `infinicore.lib._infinicore.sync_stream()`、`sync_device()` 等底层绑定。
81+
- 在同一执行流内串行调用算子通常无需额外同步。
82+
83+
## 与 `ntops` 的协作
84+
85+
- 导入 `ntops` 成功后，`infinicore.use_ntops` 会被设置为 `True`，并可通过 `infinicore.ntops` 访问原始模块。
86+
- `nn.functional.silu` 在 `use_ntops=True` 且设备类型为 `"cuda"`/`"musa"` 且未传 `out` 时，会委托 `ntops.torch.silu`。
87+
- 若想强制禁用，可直接设置 `infinicore.use_ntops = False`。
88+
89+
## 端到端示例
90+
91+
```python
92+
import infinicore as ic
93+
94+
device = ic.device("cuda:0")
95+
96+
q = ic.empty((8, 1, 128), dtype=ic.float16, device=device)
97+
k = ic.empty((2, 1, 128), dtype=ic.float16, device=device)
98+
v = ic.empty((2, 1, 128), dtype=ic.float16, device=device)
99+
k_cache = ic.empty((2, 128, 128), dtype=ic.float16, device=device)
100+
v_cache = ic.empty((2, 128, 128), dtype=ic.float16, device=device)
101+
102+
out = ic.attention(q, k, v, k_cache, v_cache, pos=0)
103+
104+
if ic.use_ntops:
105+
# 在部分设备上，SiLU 会委托给 ntops 的高性能实现
106+
out = ic.nn.functional.silu(out)
107+
108+
out.debug()
109+
```
110+
111+
## 相关链接
112+
113+
- [`infinicore.ops 顶层算子`](ops/README.md)
114+
- [`nn.functional 函数式文档`](nn/functional/README.md)
115+
- [`InfiniOP` 统一算子库](/infiniop/README.md)

‎python/device/README.md‎

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,35 @@
1+
# `infinicore.device`
2+
3+
设备句柄类，定义于 `InfiniCore/python/infinicore/device.py`，用于在 Python 端选择和描述运行时设备。
4+
5+
## 构造方式
6+
7+
```python
8+
from infinicore import device
9+
10+
cpu = device() # 默认 "cpu"
11+
cuda0 = device("cuda:0") # 字符串形式指定
12+
cuda1 = device("cuda", 1) # 类型 + index
13+
clone = device(cuda0) # 从已有实例拷贝
14+
```
15+
16+
- `type`：支持 `"cpu"`、`"cuda"`、`"mlu"`、`"npu"`、`"musa"` 等，具体取决于 `_infinicore` 编译时支持。
17+
- `index`：可选整型索引，字符串中已包含 `":"` 时禁止再传入。
18+
- 传入已有 `device` 实例时会拷贝其 `type`/`index`。
19+
20+
## 属性与方法
21+
22+
- `type` / `index`：公开的设备类型与序号。
23+
- `__repr__()` / `__str__()`：打印友好格式，如 `device(type='cuda', index=0)` 或 `"cuda:0"`。
24+
- `_underlying`：内部 `_infinicore.Device` 对象，供底层 API 使用。
25+
26+
## 与运行时的关系
27+
28+
- `device` 实例用于所有张量构造函数及顶层算子，确保 `_infinicore` 在正确的设备上执行。
29+
- 若需要从底层 `_infinicore.Device` 转换为 Python 对象，可使用 `device._from_infinicore_device`（内部方法）。
30+
31+
## 相关链接
32+
33+
- [`Tensor` 构造函数](../tensor/README.md#构造函数)
34+
- [`运行时上下文`](../README.md#运行时上下文)
35+

‎python/dtype/README.md‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
# `infinicore.dtype` 与标量类型
2+
3+
`InfiniCore/python/infinicore/dtype.py` 导出与 C++ 端一致的标量类型枚举。通过 `from infinicore import dtype, float16, int32, ...` 可直接访问。
4+
5+
## 常用类型列表
6+
7+
- `dtype`：枚举工厂，可用于创建自定义 dtype 或从 `_underlying` 还原。
8+
- 浮点类型：`float`, `float16`, `float32`, `float64`, `half`, `bfloat16`, `double`.
9+
- 复数类型：`cfloat`, `cdouble`, `complex32`, `complex64`, `complex128`.
10+
- 整型：`int`, `int8`, `int16`, `int32`, `int64`, `short`, `long`, `uint8`.
11+
- 布尔：`bool`.
12+
13+
所有类型对象都暴露 `_underlying` 属性，用于在调用底层 `_infinicore` 接口时传递。
14+
15+
## 示例
16+
17+
```python
18+
import infinicore as ic
19+
20+
cpu = ic.device("cpu")
21+
a = ic.empty((4, 8), dtype=ic.float16, device=cpu)
22+
23+
if a.dtype is ic.float16:
24+
print("half precision tensor")
25+
```
26+
27+
## 相关链接
28+
29+
- [`Tensor` 构造函数](../tensor/README.md#构造函数)
30+
- [`infinicore` 模块概览](../README.md)

‎python/nn/README.md‎

Lines changed: 25 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,25 @@
1+
# `infinicore.nn` 模块
2+
3+
`infinicore.nn` 聚合了面向神经网络的辅助模块，目前主要暴露函数式算子集合 `infinicore.nn.functional`，并预留位置扩展其他组件（如模块化层、优化器等）。
4+
5+
## 模块结构
6+
7+
| 子模块 | 说明 |
8+
| --- | --- |
9+
| `functional` | 函数式算子集合，接口文档见 [`functional/README.md`](functional/README.md)。 |
10+
11+
## 使用示例
12+
13+
```python
14+
import infinicore as ic
15+
from infinicore.nn import functional as F
16+
17+
x = ic.ones((4, 1024), dtype=ic.float16, device=ic.device("cuda:0"))
18+
normed = F.rms_norm(x, normalized_shape=[1024], weight=ic.ones((1024,), dtype=x.dtype, device=x.device))
19+
activated = F.silu(normed)
20+
```
21+
22+
## 相关链接
23+
24+
- [`nn.functional 函数式文档`](functional/README.md)
25+
- [`Python API 总览`](../README.md)

‎python/nn/functional/README.md‎

Lines changed: 37 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,37 @@
1+
# `infinicore.nn.functional` 函数式接口
2+
3+
`infinicore.nn.functional` 集中收录 PyTorch 风格的函数式算子封装。实现位于 `InfiniCore/python/infinicore/nn/functional.py`，依赖 `_infinicore` C++ 绑定并复用运行时上下文。
4+
5+
## 公共约定
6+
7+
- 所有函数都返回 `infinicore.Tensor`；当提供 `out`/`inplace` 等参数时会复用已有缓冲区。
8+
- 输入张量需由 `infinicore` 创建（或至少携带 `_underlying` 指针），否则无法与底层运行时交互。
9+
- 若函数内部调用 `_infinicore.*_` 原位接口，需确保输出张量与输入形状、dtype 一致。
10+
11+
## API 详情
12+
13+
- [`causal_softmax`](causal_softmax/README.md)：因果掩码 Softmax。
14+
- [`rms_norm`](rms_norm/README.md)：Root Mean Square LayerNorm。
15+
- [`silu`](silu/README.md)：SiLU（Sigmoid Linear Unit）激活。
16+
- [`swiglu`](swiglu/README.md)：SwiGLU 前向门控。
17+
18+
## 示例
19+
20+
```python
21+
import infinicore as ic
22+
from infinicore.nn import functional as F
23+
24+
x = ic.empty((4, 1024), dtype=ic.float16, device=ic.device("cuda:0"))
25+
w = ic.empty((1024,), dtype=ic.float16, device=x.device)
26+
27+
normed = F.rms_norm(x, normalized_shape=list(w.shape), weight=w)
28+
activated = F.silu(normed)
29+
gated = F.swiglu(activated, ic.empty_like(activated))
30+
31+
probs = F.causal_softmax(gated, out=ic.empty_like(gated))
32+
```
33+
34+
## 相关链接
35+
36+
- [`Python API 总览`](../../README.md)
37+
- [`ntops` 协作接口说明](../../README.md#与-ntops-的协作)

‎python/nn/functional/causal_softmax/README.md‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
# `nn.functional.causal_softmax`
2+
3+
对最后一维应用因果掩码 Softmax，用于自回归注意力场景。函数定义位于 `InfiniCore/python/infinicore/nn/functional.py`。
4+
5+
## 函数签名
6+
7+
```python
8+
def causal_softmax(input: Tensor, out: Optional[Tensor] = None) -> Tensor
9+
```
10+
11+
- `input`：任意维度张量，末维视为序列维，将应用因果掩码。
12+
- `out`：可选输出张量，若提供需与 `input` 形状、`dtype`、`device` 完全一致。
13+
14+
默认返回新张量；提供 `out` 时调用 `_infinicore.causal_softmax_` 原位写入。
15+
16+
## 示例
17+
18+
```python
19+
import infinicore as ic
20+
from infinicore.nn import functional as F
21+
22+
device = ic.device("cuda:0")
23+
logits = ic.empty((4, 128), dtype=ic.float16, device=device)
24+
25+
probs = F.causal_softmax(logits)
26+
F.causal_softmax(logits, out=logits) # 原位写回
27+
```
28+
29+
## 相关链接
30+
31+
- [`nn.functional` 文档](../README.md)
32+
- [`infinicore.attention` 算子](../../../ops/attention/README.md)

‎python/nn/functional/rms_norm/README.md‎

Lines changed: 43 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,43 @@
1+
# `nn.functional.rms_norm`
2+
3+
实现 Root Mean Square LayerNorm。函数定义位于 `InfiniCore/python/infinicore/nn/functional.py`。
4+
5+
## 函数签名
6+
7+
```python
8+
def rms_norm(
9+
input: Tensor,
10+
normalized_shape: list[int],
11+
weight: Tensor,
12+
eps: float = 1e-5,
13+
*,
14+
out: Optional[Tensor] = None,
15+
) -> Tensor
16+
```
17+
18+
- `input`：待归一化张量，末维通常为隐藏维度。
19+
- `normalized_shape`：期望归一化的维度大小列表，会与 `weight.shape` 进行严格比较。
20+
- `weight`：缩放系数张量，维度需与 `normalized_shape` 匹配。
21+
- `eps`：数值稳定项，默认 `1e-5`。
22+
- `out`：可选输出张量；若提供需与 `input` 形状、`dtype`、`device` 一致。
23+
24+
函数首先断言 `normalized_shape == weight.shape`，然后调用 `_infinicore.rms_norm` / `_infinicore.rms_norm_` 完成计算。
25+
26+
## 示例
27+
28+
```python
29+
import infinicore as ic
30+
from infinicore.nn import functional as F
31+
32+
device = ic.device("cuda:0")
33+
x = ic.ones((4, 1024), dtype=ic.float16, device=device)
34+
gamma = ic.ones((1024,), dtype=ic.float16, device=device)
35+
36+
y = F.rms_norm(x, normalized_shape=list(gamma.shape), weight=gamma, eps=1e-5)
37+
F.rms_norm(x, normalized_shape=[1024], weight=gamma, out=x) # 原位写回
38+
```
39+
40+
## 相关链接
41+
42+
- [`nn.functional` 文档](../README.md)
43+
- [`Tensor` 构造函数](../../../README.md#张量与构造函数)

‎python/nn/functional/silu/README.md‎

Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
# `nn.functional.silu`
2+
3+
Sigmoid Linear Unit (SiLU) 激活函数。实现位于 `InfiniCore/python/infinicore/nn/functional.py`。
4+
5+
## 函数签名
6+
7+
```python
8+
def silu(
9+
input: Tensor,
10+
inplace: bool = False,
11+
*,
12+
out: Optional[Tensor] = None,
13+
) -> Tensor
14+
```
15+
16+
- `input`：待激活张量。
17+
- `inplace`：是否原地写回到 `input`。
18+
- `out`：可选输出张量，若提供需与 `input` 形状、`dtype`、`device` 一致。
19+
20+
### 行为说明
21+
22+
- 当 `inplace=True` 时，直接调用 `_infinicore.silu_` 写回 `input` 并返回。
23+
- 当未提供 `out` 且 `infinicore.use_ntops=True` 且设备类型为 `"cuda"` 或 `"musa"` 时，会委托 `ntops.torch.silu` 以复用优化实现。
24+
- 其他情况下调用 `_infinicore.silu` 或 `_infinicore.silu_` 完成计算。
25+
26+
## 示例
27+
28+
```python
29+
import infinicore as ic
30+
from infinicore.nn import functional as F
31+
32+
device = ic.device("cuda:0")
33+
x = ic.ones((4, 8), dtype=ic.float16, device=device)
34+
35+
y = F.silu(x) # 返回新张量
36+
F.silu(x, inplace=True) # 原地更新
37+
```
38+
39+
## 相关链接
40+
41+
- [`nn.functional` 文档](../README.md)
42+
- [`use_ntops` 协作说明](../../../README.md#与-ntops-的协作)

0 commit comments

Comments
 (0)