[Docs](feat) Add conv1d and conv2d API doc

docs/zh/python-api/_ascend_constraints.py

@@ -807,6 +807,34 @@
         "example":
         "triton.language.extra.cann.extension.dot",
     },
+    "triton.language.extra.cann.extension.conv1d": {
+        "constraints": [
+            "DataType: Ascend supports fp16, bf16, fp32.",
+            "``input``: 2D ``(iC, iW)`` or 3D ``(N, iC, iW)``.",
+            "``weight``: 3D ``(oC, iC // groups, wW)``; ``iC % groups == 0`` and ``oC % groups == 0``.",
+            "``bias``: optional, 1D ``(oC)``.",
+            "``groups``: int; must divide both ``iC`` and ``oC`` (``iC % groups == 0`` and ``oC % groups == 0``).",
+            "``padding``: int (symmetric on both sides) or 2-element tuple ``(padding_left, padding_right)`` (asymmetric).",
+            "``stride``: int, the stride of the convolution kernel.",
+            "``dilation``: only ``dilation=1`` is currently supported.",
+        ],
+        "example":
+        "triton.language.extra.cann.extension.conv1d",
+    },
+    "triton.language.extra.cann.extension.conv2d": {
+        "constraints": [
+            "DataType: Ascend supports fp16, bf16, fp32.",
+            "``input``: 3D ``(iC, iH, iW)`` or 4D ``(N, iC, iH, iW)``.",
+            "``weight``: 4D ``(oC, iC // groups, wH, wW)``; ``iC % groups == 0`` and ``oC % groups == 0``.",
+            "``bias``: optional, 1D ``(oC)``.",
+            "``groups``: int; must divide both ``iC`` and ``oC`` (``iC % groups == 0`` and ``oC % groups == 0``).",
+            "``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).",
+            "``stride``: int or 2-element tuple ``(stride_h, stride_w)``.",
+            "``dilation``: only ``dilation=1`` is currently supported.",
+        ],
+        "example":
+        "triton.language.extra.cann.extension.conv2d",
+    },
     "triton.language.split": {
         "constraints": [
             "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

@@ -0,0 +1,78 @@
+import torch
+import triton
+import triton.language as tl
+import triton.language.extra.cann.extension as al
+
+
+@triton.jit
+def conv1d_kernel(
+    input_ptr,
+    weight_ptr,
+    bias_ptr,
+    output_ptr,
+    N: tl.constexpr,
+    C_in: tl.constexpr,
+    L_in: tl.constexpr,
+    C_out: tl.constexpr,
+    L_out: tl.constexpr,
+    K: tl.constexpr,
+    stride: tl.constexpr,
+    padding: tl.constexpr,
+    groups: tl.constexpr,
+):
+    # Load input: (N, C_in, L_in)
+    n_offs = tl.arange(0, N)[:, None, None]
+    c_offs = tl.arange(0, C_in)[None, :, None]
+    l_offs = tl.arange(0, L_in)[None, None, :]
+    input_tile = tl.load(input_ptr + n_offs * (C_in * L_in) + c_offs * L_in + l_offs)
+
+    # Load weight: (C_out, C_in // groups, K)
+    co_offs = tl.arange(0, C_out)[:, None, None]
+    ci_offs = tl.arange(0, C_in // groups)[None, :, None]
+    k_offs = tl.arange(0, K)[None, None, :]
+    weight_tile = tl.load(weight_ptr + co_offs * ((C_in // groups) * K) + ci_offs * K + k_offs)
+
+    # Load bias: (C_out,)
+    bias_tile = tl.load(bias_ptr + tl.arange(0, C_out))
+
+    output = al.conv1d(
+        input_tile,
+        weight_tile,
+        bias_tile,
+        groups=groups,
+        padding=padding,
+        stride=stride,
+        dilation=1,
+    )
+
+    # Store output: (N, C_out, L_out)
+    no_offs = tl.arange(0, N)[:, None, None]
+    co_offs = tl.arange(0, C_out)[None, :, None]
+    lo_offs = tl.arange(0, L_out)[None, None, :]
+    tl.store(output_ptr + no_offs * (C_out * L_out) + co_offs * L_out + lo_offs, output)
+
+
+def test_conv1d():
+    N, C_in, L_in = 2, 16, 32
+    C_out, K = 32, 3
+    stride, padding, groups = 1, 1, 1
+    L_out = (L_in + 2 * padding - (K - 1) - 1) // stride + 1
+
+    x = torch.randn(N, C_in, L_in, dtype=torch.float16)
+    w = torch.randn(C_out, C_in // groups, K, dtype=torch.float16)
+    b = torch.randn(C_out, dtype=torch.float16)
+
+    x_npu = x.npu()
+    w_npu = w.npu()
+    b_npu = b.npu()
+    out_npu = torch.empty(N, C_out, L_out, dtype=torch.float16).npu()
+
+    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,
+                         stride=stride, padding=padding, groups=groups)
+
+    gold = torch.nn.functional.conv1d(x, w, b, stride=stride, padding=padding, groups=groups)
+    torch.testing.assert_close(out_npu.cpu(), gold, rtol=1e-2, atol=1e-2)
+
+
+if __name__ == "__main__":
+    test_conv1d()

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

@@ -0,0 +1,92 @@
+import torch
+import triton
+import triton.language as tl
+import triton.language.extra.cann.extension as al
+
+
+@triton.jit
+def conv2d_kernel(
+    input_ptr,
+    weight_ptr,
+    bias_ptr,
+    output_ptr,
+    N: tl.constexpr,
+    C_in: tl.constexpr,
+    H_in: tl.constexpr,
+    W_in: tl.constexpr,
+    C_out: tl.constexpr,
+    H_out: tl.constexpr,
+    W_out: tl.constexpr,
+    K_h: tl.constexpr,
+    K_w: tl.constexpr,
+    stride_h: tl.constexpr,
+    stride_w: tl.constexpr,
+    padding_h: tl.constexpr,
+    padding_w: tl.constexpr,
+    groups: tl.constexpr,
+):
+    # Load input: (N, C_in, H_in, W_in)
+    n_offs = tl.arange(0, N)[:, None, None, None]
+    c_offs = tl.arange(0, C_in)[None, :, None, None]
+    h_offs = tl.arange(0, H_in)[None, None, :, None]
+    w_offs = tl.arange(0, W_in)[None, None, None, :]
+    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)
+
+    # Load weight: (C_out, C_in // groups, K_h, K_w)
+    co_offs = tl.arange(0, C_out)[:, None, None, None]
+    ci_offs = tl.arange(0, C_in // groups)[None, :, None, None]
+    kh_offs = tl.arange(0, K_h)[None, None, :, None]
+    kw_offs = tl.arange(0, K_w)[None, None, None, :]
+    weight_tile = tl.load(weight_ptr + co_offs * ((C_in // groups) * K_h * K_w) + ci_offs * (K_h * K_w) +
+                          kh_offs * K_w + kw_offs)
+
+    # Load bias: (C_out,)
+    bias_tile = tl.load(bias_ptr + tl.arange(0, C_out))
+
+    output = al.conv2d(
+        input_tile,
+        weight_tile,
+        bias_tile,
+        groups=groups,
+        padding=(padding_h, padding_w),
+        stride=(stride_h, stride_w),
+        dilation=1,
+    )
+
+    # Store output: (N, C_out, H_out, W_out)
+    no_offs = tl.arange(0, N)[:, None, None, None]
+    co_offs = tl.arange(0, C_out)[None, :, None, None]
+    ho_offs = tl.arange(0, H_out)[None, None, :, None]
+    wo_offs = tl.arange(0, W_out)[None, None, None, :]
+    tl.store(output_ptr + no_offs * (C_out * H_out * W_out) + co_offs * (H_out * W_out) + ho_offs * W_out + wo_offs,
+             output)
+
+
+def test_conv2d():
+    N, C_in, H_in, W_in = 2, 16, 32, 32
+    C_out, K_h, K_w = 32, 3, 3
+    stride = (1, 1)
+    padding = (1, 1)
+    groups = 1
+    H_out = (H_in + 2 * padding[0] - (K_h - 1) - 1) // stride[0] + 1
+    W_out = (W_in + 2 * padding[1] - (K_w - 1) - 1) // stride[1] + 1
+
+    x = torch.randn(N, C_in, H_in, W_in, dtype=torch.float16)
+    w = torch.randn(C_out, C_in // groups, K_h, K_w, dtype=torch.float16)
+    b = torch.randn(C_out, dtype=torch.float16)
+
+    x_npu = x.npu()
+    w_npu = w.npu()
+    b_npu = b.npu()
+    out_npu = torch.empty(N, C_out, H_out, W_out, dtype=torch.float16).npu()
+
+    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,
+                         W_out=W_out, K_h=K_h, K_w=K_w, stride_h=stride[0], stride_w=stride[1], padding_h=padding[0],
+                         padding_w=padding[1], groups=groups)
+
+    gold = torch.nn.functional.conv2d(x, w, b, stride=stride, padding=padding, groups=groups)
+    torch.testing.assert_close(out_npu.cpu(), gold, rtol=1e-2, atol=1e-2)
+
+
+if __name__ == "__main__":
+    test_conv2d()

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

@@ -51,6 +51,7 @@ Vector Operations
     sub_vec_id
     sub_vec_num
     conv1d
+    conv2d
     dot
 
 Enums

docs/zh/triton_api_extension/al/conv1d.md

@@ -0,0 +1,149 @@
+# al.conv1d 接口文档
+
+## 1. 背景
+
+al.conv1d 在输入信号上执行一维卷积，支持可选偏置（bias）与分组卷积（groups），padding 支持标量或元组形式，接口语义对齐 torch.nn.functional.conv1d。
+
+## 2. 接口说明
+
+<table>
+  <tr>
+    <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>
+  </tr>
+</table>
+
+### 参数
+
+<table>
+  <tr>
+    <td>参数名</td>
+    <td>类型</td>
+    <td>必需</td>
+    <td>说明</td>
+  </tr>
+  <tr>
+    <td>input</td>
+    <td>tensor</td>
+    <td>是</td>
+    <td>输入张量，形状 [N, iC, iW] 或 [iC, iW]，N 为 batch size，iC 为输入通道数，iW 为输入宽度</td>
+  </tr>
+  <tr>
+    <td>weight</td>
+    <td>tensor</td>
+    <td>是</td>
+    <td>权重张量，形状 [oC, iC / groups, wW]，oC 为输出通道数，wW 为卷积核宽度，要求 oC % groups == 0</td>
+  </tr>
+  <tr>
+    <td>bias</td>
+    <td>tensor</td>
+    <td>否</td>
+    <td>偏置张量，形状 [oC]，默认 None</td>
+  </tr>
+  <tr>
+    <td>groups</td>
+    <td>int</td>
+    <td>否</td>
+    <td>输入到输出通道的分组数，默认 1</td>
+  </tr>
+  <tr>
+    <td>padding</td>
+    <td>int / tuple</td>
+    <td>否</td>
+    <td>输入两侧的填充，支持 int（两侧对称）或 2 元组 (paddingLeft, paddingRight)（非对称），默认 0</td>
+  </tr>
+  <tr>
+    <td>stride</td>
+    <td>int</td>
+    <td>否</td>
+    <td>卷积核的步长，默认 1</td>
+  </tr>
+  <tr>
+    <td>dilation</td>
+    <td>int</td>
+    <td>否</td>
+    <td>卷积核元素之间的间距，暂未支持非 1，默认 1</td>
+  </tr>
+</table>
+
+### 返回值
+
+输出张量，形状 [N, oC, oW] 或 [oC, oW]。
+
+### 2.3 支持规格
+
+#### 2.3.1 DataType 支持
+
+| 输入类型 | int8 | int16 | int32 | uint8 | uint16 | uint32 | uint64 | int64 | fp16 | fp32 | fp64 | bf16 | bool |
+| ------ | ---- | ----- | ----- | ----- | ------ | ------ | ------ | ----- | ---- | ---- | ---- | ---- | ---- |
+| Ascend A2/A3 | ×    | ×     | ×     | ×     | ×      | ×      | ×      | ×     | √    | √    | ×    | √    | ×    |
+| Ascend A5 | ×    | ×     | ×     | ×     | ×      | ×      | ×      | ×     | √    | √    | ×    | √    | ×    |
+
+结论：al.conv1d 支持 fp16、bf16、fp32 三种浮点数据类型。
+
+### 2.4 约束说明
+
+- groups 必须同时整除 iC 与 oC（oC % groups == 0）。
+
+- bias 为可选参数，形状必须为 [oC]。
+
+- dilation 暂未支持非 1 取值。
+
+- padding 支持 int 或 2 元组 (paddingLeft, paddingRight)。
+
+- 默认值：groups=1、padding=0、stride=1、dilation=1。
+
+## 3. 用例示例
+
+```python
+import triton
+import triton.language as tl
+import triton.language.extra.cann.extension as al
+
+
+@triton.jit
+def conv1d_kernel(
+    input_ptr,
+    weight_ptr,
+    bias_ptr,
+    output_ptr,
+    N: tl.constexpr,
+    C_in: tl.constexpr,
+    L_in: tl.constexpr,
+    C_out: tl.constexpr,
+    L_out: tl.constexpr,
+    K: tl.constexpr,
+    stride: tl.constexpr,
+    padding: tl.constexpr,
+    groups: tl.constexpr,
+):
+    # Load input: (N, C_in, L_in)
+    n_offs = tl.arange(0, N)[:, None, None]
+    c_offs = tl.arange(0, C_in)[None, :, None]
+    l_offs = tl.arange(0, L_in)[None, None, :]
+    input_tile = tl.load(input_ptr + n_offs * (C_in * L_in) + c_offs * L_in + l_offs)
+
+    # Load weight: (C_out, C_in // groups, K)
+    co_offs = tl.arange(0, C_out)[:, None, None]
+    ci_offs = tl.arange(0, C_in // groups)[None, :, None]
+    k_offs = tl.arange(0, K)[None, None, :]
+    weight_tile = tl.load(weight_ptr + co_offs * ((C_in // groups) * K) + ci_offs * K + k_offs)
+
+    # Load bias: (C_out,)
+    bias_tile = tl.load(bias_ptr + tl.arange(0, C_out))
+
+    output = al.conv1d(
+        input_tile,
+        weight_tile,
+        bias_tile,
+        groups=groups,
+        padding=padding,
+        stride=stride,
+        dilation=1,
+    )
+
+    # Store output: (N, C_out, L_out)
+    no_offs = tl.arange(0, N)[:, None, None]
+    co_offs = tl.arange(0, C_out)[None, :, None]
+    lo_offs = tl.arange(0, L_out)[None, None, :]
+    tl.store(output_ptr + no_offs * (C_out * L_out) + co_offs * L_out + lo_offs, output)
+```