Skip to content

Commit 2ad57be

Browse files
lowdy1lowdy1
committed
add conv1d and conv2d docs
1 parent 38b42fe commit 2ad57be

6 files changed

Lines changed: 507 additions & 0 deletions

File tree

‎docs/zh/python-api/_ascend_constraints.py‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -807,6 +807,34 @@
807807
"example":
808808
"triton.language.extra.cann.extension.dot",
809809
},
810+
"triton.language.extra.cann.extension.conv1d": {
811+
"constraints": [
812+
"DataType: Ascend supports fp16, bf16, fp32.",
813+
"``input``: 2D ``(iC, iW)`` or 3D ``(N, iC, iW)``.",
814+
"``weight``: 3D ``(oC, iC // groups, wW)``; ``iC % groups == 0`` and ``oC % groups == 0``.",
815+
"``bias``: optional, 1D ``(oC)``.",
816+
"``groups``: int; must divide both ``iC`` and ``oC`` (``iC % groups == 0`` and ``oC % groups == 0``).",
817+
"``padding``: int (symmetric on both sides) or 2-element tuple ``(padding_left, padding_right)`` (asymmetric).",
818+
"``stride``: int, the stride of the convolution kernel.",
819+
"``dilation``: only ``dilation=1`` is currently supported.",
820+
],
821+
"example":
822+
"triton.language.extra.cann.extension.conv1d",
823+
},
824+
"triton.language.extra.cann.extension.conv2d": {
825+
"constraints": [
826+
"DataType: Ascend supports fp16, bf16, fp32.",
827+
"``input``: 3D ``(iC, iH, iW)`` or 4D ``(N, iC, iH, iW)``.",
828+
"``weight``: 4D ``(oC, iC // groups, wH, wW)``; ``iC % groups == 0`` and ``oC % groups == 0``.",
829+
"``bias``: optional, 1D ``(oC)``.",
830+
"``groups``: int; must divide both ``iC`` and ``oC`` (``iC % groups == 0`` and ``oC % groups == 0``).",
831+
"``padding``: int (symmetric on all sides), 2-element tuple ``(padding_h, padding_w)`` (symmetric per dimension), or 4-element tuple ``(padding_top, padding_bottom, padding_left, padding_right)`` (asymmetric).",
832+
"``stride``: int or 2-element tuple ``(stride_h, stride_w)``.",
833+
"``dilation``: only ``dilation=1`` is currently supported.",
834+
],
835+
"example":
836+
"triton.language.extra.cann.extension.conv2d",
837+
},
810838
"triton.language.split": {
811839
"constraints": [
812840
"DataType: Ascend A2/A3 does not support fp64, fp8e4, fp8e5, uint16, uint32, uint64 (hardware limitation).",

‎docs/zh/python-api/_examples/triton.language.extra.cann.extension.conv1d.py‎

Lines changed: 78 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,78 @@
1+
import torch
2+
import triton
3+
import triton.language as tl
4+
import triton.language.extra.cann.extension as al
5+
6+
7+
@triton.jit
8+
def conv1d_kernel(
9+
input_ptr,
10+
weight_ptr,
11+
bias_ptr,
12+
output_ptr,
13+
N: tl.constexpr,
14+
C_in: tl.constexpr,
15+
L_in: tl.constexpr,
16+
C_out: tl.constexpr,
17+
L_out: tl.constexpr,
18+
K: tl.constexpr,
19+
stride: tl.constexpr,
20+
padding: tl.constexpr,
21+
groups: tl.constexpr,
22+
):
23+
# Load input: (N, C_in, L_in)
24+
n_offs = tl.arange(0, N)[:, None, None]
25+
c_offs = tl.arange(0, C_in)[None, :, None]
26+
l_offs = tl.arange(0, L_in)[None, None, :]
27+
input_tile = tl.load(input_ptr + n_offs * (C_in * L_in) + c_offs * L_in + l_offs)
28+
29+
# Load weight: (C_out, C_in // groups, K)
30+
co_offs = tl.arange(0, C_out)[:, None, None]
31+
ci_offs = tl.arange(0, C_in // groups)[None, :, None]
32+
k_offs = tl.arange(0, K)[None, None, :]
33+
weight_tile = tl.load(weight_ptr + co_offs * ((C_in // groups) * K) + ci_offs * K + k_offs)
34+
35+
# Load bias: (C_out,)
36+
bias_tile = tl.load(bias_ptr + tl.arange(0, C_out))
37+
38+
output = al.conv1d(
39+
input_tile,
40+
weight_tile,
41+
bias_tile,
42+
groups=groups,
43+
padding=padding,
44+
stride=stride,
45+
dilation=1,
46+
)
47+
48+
# Store output: (N, C_out, L_out)
49+
no_offs = tl.arange(0, N)[:, None, None]
50+
co_offs = tl.arange(0, C_out)[None, :, None]
51+
lo_offs = tl.arange(0, L_out)[None, None, :]
52+
tl.store(output_ptr + no_offs * (C_out * L_out) + co_offs * L_out + lo_offs, output)
53+
54+
55+
def test_conv1d():
56+
N, C_in, L_in = 2, 16, 32
57+
C_out, K = 32, 3
58+
stride, padding, groups = 1, 1, 1
59+
L_out = (L_in + 2 * padding - (K - 1) - 1) // stride + 1
60+
61+
x = torch.randn(N, C_in, L_in, dtype=torch.float16)
62+
w = torch.randn(C_out, C_in // groups, K, dtype=torch.float16)
63+
b = torch.randn(C_out, dtype=torch.float16)
64+
65+
x_npu = x.npu()
66+
w_npu = w.npu()
67+
b_npu = b.npu()
68+
out_npu = torch.empty(N, C_out, L_out, dtype=torch.float16).npu()
69+
70+
conv1d_kernel[(1, )](x_npu, w_npu, b_npu, out_npu, N=N, C_in=C_in, L_in=L_in, C_out=C_out, L_out=L_out, K=K,
71+
stride=stride, padding=padding, groups=groups)
72+
73+
gold = torch.nn.functional.conv1d(x, w, b, stride=stride, padding=padding, groups=groups)
74+
torch.testing.assert_close(out_npu.cpu(), gold, rtol=1e-2, atol=1e-2)
75+
76+
77+
if __name__ == "__main__":
78+
test_conv1d()

‎docs/zh/python-api/_examples/triton.language.extra.cann.extension.conv2d.py‎

Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,92 @@
1+
import torch
2+
import triton
3+
import triton.language as tl
4+
import triton.language.extra.cann.extension as al
5+
6+
7+
@triton.jit
8+
def conv2d_kernel(
9+
input_ptr,
10+
weight_ptr,
11+
bias_ptr,
12+
output_ptr,
13+
N: tl.constexpr,
14+
C_in: tl.constexpr,
15+
H_in: tl.constexpr,
16+
W_in: tl.constexpr,
17+
C_out: tl.constexpr,
18+
H_out: tl.constexpr,
19+
W_out: tl.constexpr,
20+
K_h: tl.constexpr,
21+
K_w: tl.constexpr,
22+
stride_h: tl.constexpr,
23+
stride_w: tl.constexpr,
24+
padding_h: tl.constexpr,
25+
padding_w: tl.constexpr,
26+
groups: tl.constexpr,
27+
):
28+
# Load input: (N, C_in, H_in, W_in)
29+
n_offs = tl.arange(0, N)[:, None, None, None]
30+
c_offs = tl.arange(0, C_in)[None, :, None, None]
31+
h_offs = tl.arange(0, H_in)[None, None, :, None]
32+
w_offs = tl.arange(0, W_in)[None, None, None, :]
33+
input_tile = tl.load(input_ptr + n_offs * (C_in * H_in * W_in) + c_offs * (H_in * W_in) + h_offs * W_in + w_offs)
34+
35+
# Load weight: (C_out, C_in // groups, K_h, K_w)
36+
co_offs = tl.arange(0, C_out)[:, None, None, None]
37+
ci_offs = tl.arange(0, C_in // groups)[None, :, None, None]
38+
kh_offs = tl.arange(0, K_h)[None, None, :, None]
39+
kw_offs = tl.arange(0, K_w)[None, None, None, :]
40+
weight_tile = tl.load(weight_ptr + co_offs * ((C_in // groups) * K_h * K_w) + ci_offs * (K_h * K_w) +
41+
kh_offs * K_w + kw_offs)
42+
43+
# Load bias: (C_out,)
44+
bias_tile = tl.load(bias_ptr + tl.arange(0, C_out))
45+
46+
output = al.conv2d(
47+
input_tile,
48+
weight_tile,
49+
bias_tile,
50+
groups=groups,
51+
padding=(padding_h, padding_w),
52+
stride=(stride_h, stride_w),
53+
dilation=1,
54+
)
55+
56+
# Store output: (N, C_out, H_out, W_out)
57+
no_offs = tl.arange(0, N)[:, None, None, None]
58+
co_offs = tl.arange(0, C_out)[None, :, None, None]
59+
ho_offs = tl.arange(0, H_out)[None, None, :, None]
60+
wo_offs = tl.arange(0, W_out)[None, None, None, :]
61+
tl.store(output_ptr + no_offs * (C_out * H_out * W_out) + co_offs * (H_out * W_out) + ho_offs * W_out + wo_offs,
62+
output)
63+
64+
65+
def test_conv2d():
66+
N, C_in, H_in, W_in = 2, 16, 32, 32
67+
C_out, K_h, K_w = 32, 3, 3
68+
stride = (1, 1)
69+
padding = (1, 1)
70+
groups = 1
71+
H_out = (H_in + 2 * padding[0] - (K_h - 1) - 1) // stride[0] + 1
72+
W_out = (W_in + 2 * padding[1] - (K_w - 1) - 1) // stride[1] + 1
73+
74+
x = torch.randn(N, C_in, H_in, W_in, dtype=torch.float16)
75+
w = torch.randn(C_out, C_in // groups, K_h, K_w, dtype=torch.float16)
76+
b = torch.randn(C_out, dtype=torch.float16)
77+
78+
x_npu = x.npu()
79+
w_npu = w.npu()
80+
b_npu = b.npu()
81+
out_npu = torch.empty(N, C_out, H_out, W_out, dtype=torch.float16).npu()
82+
83+
conv2d_kernel[(1, )](x_npu, w_npu, b_npu, out_npu, N=N, C_in=C_in, H_in=H_in, W_in=W_in, C_out=C_out, H_out=H_out,
84+
W_out=W_out, K_h=K_h, K_w=K_w, stride_h=stride[0], stride_w=stride[1], padding_h=padding[0],
85+
padding_w=padding[1], groups=groups)
86+
87+
gold = torch.nn.functional.conv2d(x, w, b, stride=stride, padding=padding, groups=groups)
88+
torch.testing.assert_close(out_npu.cpu(), gold, rtol=1e-2, atol=1e-2)
89+
90+
91+
if __name__ == "__main__":
92+
test_conv2d()

‎docs/zh/python-api/triton.language.extra.cann.extension.rst‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -51,6 +51,7 @@ Vector Operations
5151
sub_vec_id
5252
sub_vec_num
5353
conv1d
54+
conv2d
5455
dot
5556

5657
Enums

‎docs/zh/triton_api_extension/al/conv1d.md‎

Lines changed: 149 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,149 @@
1+
# al.conv1d 接口文档
2+
3+
## 1. 背景
4+
5+
al.conv1d 在输入信号上执行一维卷积，支持可选偏置（bias）与分组卷积（groups），stride 与 padding 支持标量或元组形式，接口语义对齐 torch.nn.functional.conv1d。
6+
7+
## 2. 接口说明
8+
9+
<table>
10+
<tr>
11+
<td>Python<br>output = al.conv1d(<br> input,<br> weight,<br> bias,<br> groups=1,<br> padding=0,<br> stride=1,<br> dilation=1,<br>) -&gt; tensor :</td>
12+
</tr>
13+
</table>
14+
15+
### 参数
16+
17+
<table>
18+
<tr>
19+
<td>参数名</td>
20+
<td>类型</td>
21+
<td>必需</td>
22+
<td>说明</td>
23+
</tr>
24+
<tr>
25+
<td>input</td>
26+
<td>tensor</td>
27+
<td>是</td>
28+
<td>输入张量，形状 [N, iC, iW] 或 [iC, iW]，N 为 batch size，iC 为输入通道数，iW 为输入宽度</td>
29+
</tr>
30+
<tr>
31+
<td>weight</td>
32+
<td>tensor</td>
33+
<td>是</td>
34+
<td>权重张量，形状 [oC, iC / groups, wW]，oC 为输出通道数，wW 为卷积核宽度，要求 oC % groups == 0</td>
35+
</tr>
36+
<tr>
37+
<td>bias</td>
38+
<td>tensor</td>
39+
<td>否</td>
40+
<td>偏置张量，形状 [oC]，默认 None</td>
41+
</tr>
42+
<tr>
43+
<td>groups</td>
44+
<td>int</td>
45+
<td>否</td>
46+
<td>输入到输出通道的分组数，默认 1</td>
47+
</tr>
48+
<tr>
49+
<td>padding</td>
50+
<td>int / tuple</td>
51+
<td>否</td>
52+
<td>输入两侧的填充，支持 int（两侧对称）或 2 元组 (paddingLeft, paddingRight)（非对称），默认 0</td>
53+
</tr>
54+
<tr>
55+
<td>stride</td>
56+
<td>int</td>
57+
<td>否</td>
58+
<td>卷积核的步长，默认 1</td>
59+
</tr>
60+
<tr>
61+
<td>dilation</td>
62+
<td>int</td>
63+
<td>否</td>
64+
<td>卷积核元素之间的间距，暂未支持非 1，默认 1</td>
65+
</tr>
66+
</table>
67+
68+
### 返回值
69+
70+
输出张量，形状 [N, oC, oW] 或 [oC, oW]。
71+
72+
### 2.3 支持规格
73+
74+
#### 2.3.1 DataType 支持
75+
76+
| 输入类型 | int8 | int16 | int32 | uint8 | uint16 | uint32 | uint64 | int64 | fp16 | fp32 | fp64 | bf16 | bool |
77+
| ------ | ---- | ----- | ----- | ----- | ------ | ------ | ------ | ----- | ---- | ---- | ---- | ---- | ---- |
78+
| Ascend A2/A3 | × | × | × | × | × | × | × | × | √ | √ | × | √ | × |
79+
| Ascend A5 | × | × | × | × | × | × | × | × | √ | √ | × | √ | × |
80+
81+
结论：al.conv1d 支持 fp16、bf16、fp32 三种浮点数据类型。
82+
83+
### 2.4 约束说明
84+
85+
- groups 必须同时整除 iC 与 oC（oC % groups == 0）。
86+
87+
- bias 为可选参数，形状必须为 [oC]。
88+
89+
- dilation 暂未支持非 1 取值。
90+
91+
- padding 支持 int 或 2 元组 (paddingLeft, paddingRight)。
92+
93+
- 默认值：groups=1、padding=0、stride=1、dilation=1。
94+
95+
## 3. 用例示例
96+
97+
```python
98+
import triton
99+
import triton.language as tl
100+
import triton.language.extra.cann.extension as al
101+
102+
103+
@triton.jit
104+
def conv1d_kernel(
105+
input_ptr,
106+
weight_ptr,
107+
bias_ptr,
108+
output_ptr,
109+
N: tl.constexpr,
110+
C_in: tl.constexpr,
111+
L_in: tl.constexpr,
112+
C_out: tl.constexpr,
113+
L_out: tl.constexpr,
114+
K: tl.constexpr,
115+
stride: tl.constexpr,
116+
padding: tl.constexpr,
117+
groups: tl.constexpr,
118+
):
119+
# Load input: (N, C_in, L_in)
120+
n_offs = tl.arange(0, N)[:, None, None]
121+
c_offs = tl.arange(0, C_in)[None, :, None]
122+
l_offs = tl.arange(0, L_in)[None, None, :]
123+
input_tile = tl.load(input_ptr + n_offs * (C_in * L_in) + c_offs * L_in + l_offs)
124+
125+
# Load weight: (C_out, C_in // groups, K)
126+
co_offs = tl.arange(0, C_out)[:, None, None]
127+
ci_offs = tl.arange(0, C_in // groups)[None, :, None]
128+
k_offs = tl.arange(0, K)[None, None, :]
129+
weight_tile = tl.load(weight_ptr + co_offs * ((C_in // groups) * K) + ci_offs * K + k_offs)
130+
131+
# Load bias: (C_out,)
132+
bias_tile = tl.load(bias_ptr + tl.arange(0, C_out))
133+
134+
output = al.conv1d(
135+
input_tile,
136+
weight_tile,
137+
bias_tile,
138+
groups=groups,
139+
padding=padding,
140+
stride=stride,
141+
dilation=1,
142+
)
143+
144+
# Store output: (N, C_out, L_out)
145+
no_offs = tl.arange(0, N)[:, None, None]
146+
co_offs = tl.arange(0, C_out)[None, :, None]
147+
lo_offs = tl.arange(0, L_out)[None, None, :]
148+
tl.store(output_ptr + no_offs * (C_out * L_out) + co_offs * L_out + lo_offs, output)
149+
```

0 commit comments

Comments
 (0)